From 5624dbb4796124e2b23d6a295783a3158944758d Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Thu, 16 Oct 2025 16:36:24 -0700 Subject: [PATCH 001/521] Changed VERSION to 2.10.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 8bfb1cae85..c7f2fd9b8e 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.9.0.dev0 +2.10.0.dev0 From 9dd619222283da058f28c6c81451f2c361434c98 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 16 Oct 2025 20:45:47 -0700 Subject: [PATCH 002/521] [JAX] Fix imports in test for deprecated jax.experimental.pjit (#2274) * Fix imports in test for deprecated jax.experimental.pjit Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix: Pass NamedSharding instead of PartitionSpec to compare_ops() so that when the in and out sharding is used to create a jitted function, it has the mesh info Signed-off-by: Kshitij Janardan Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- 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 --- tests/jax/distributed_test_base.py | 14 +++++++------ tests/jax/test_distributed_layernorm.py | 26 ++++++++++++++++--------- tests/jax/test_distributed_softmax.py | 10 ++++++---- 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/tests/jax/distributed_test_base.py b/tests/jax/distributed_test_base.py index 4693086b83..137fa480dd 100644 --- a/tests/jax/distributed_test_base.py +++ b/tests/jax/distributed_test_base.py @@ -8,7 +8,7 @@ import pytest import jax -from jax.experimental.pjit import pjit, _UNSPECIFIED +from jax._src.sharding_impls import UNSPECIFIED as _UNSPECIFIED from transformer_engine.jax.sharding import MeshResource @@ -154,13 +154,15 @@ def compare_ops( grad_args = tuple(range(len(inputs))) target_grad_func = jax.value_and_grad(target_func, argnums=grad_args) - target_pjitter = pjit(target_grad_func, in_shardings=in_shardings, out_shardings=out_shardings) - target_fwd, target_grads = target_pjitter(*inputs, **kwargs) - target_hlo = target_pjitter.lower(*inputs, **kwargs).compile().as_text() + target_jitter = jax.jit( + target_grad_func, in_shardings=in_shardings, out_shardings=out_shardings + ) + target_fwd, target_grads = target_jitter(*inputs, **kwargs) + target_hlo = target_jitter.lower(*inputs, **kwargs).compile().as_text() ref_grad_func = jax.value_and_grad(ref_func, argnums=grad_args) - ref_pjitter = pjit(ref_grad_func, in_shardings=in_shardings, out_shardings=out_shardings) - ref_fwd, ref_grads = ref_pjitter(*inputs, **kwargs) + ref_jitter = jax.jit(ref_grad_func, in_shardings=in_shardings, out_shardings=out_shardings) + ref_fwd, ref_grads = ref_jitter(*inputs, **kwargs) assert_allclose(target_fwd, ref_fwd, dtype=metric_fwd_dtype) diff --git a/tests/jax/test_distributed_layernorm.py b/tests/jax/test_distributed_layernorm.py index 977d010afd..d551b73905 100644 --- a/tests/jax/test_distributed_layernorm.py +++ b/tests/jax/test_distributed_layernorm.py @@ -134,9 +134,12 @@ def ref_func(x, gamma, beta): devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) mesh = Mesh(devices, mesh_axes) with mesh, autocast(enabled=True, recipe=fp8_recipe, mesh_resource=mesh_resource): - x_ = jax.device_put(x, NamedSharding(mesh, x_pspec)) - gamma_ = jax.device_put(gamma, NamedSharding(mesh, g_pspec)) - beta_ = jax.device_put(beta, NamedSharding(mesh, b_pspec)) + x_named_sharding = NamedSharding(mesh, x_pspec) + g_named_sharding = NamedSharding(mesh, g_pspec) + b_named_sharding = NamedSharding(mesh, b_pspec) + x_ = jax.device_put(x, x_named_sharding) + gamma_ = jax.device_put(gamma, g_named_sharding) + beta_ = jax.device_put(beta, b_named_sharding) with warnings.catch_warnings(record=True) as warns: try: @@ -148,8 +151,11 @@ def ref_func(x, gamma, beta): grad_args=(0, 1, 2), metric_fwd_dtype=q_dtype, metric_bwd_dtype=q_dtype, - in_shardings=(x_pspec, g_pspec, b_pspec), - out_shardings=(None, (x_pspec, g_pspec, b_pspec)), + in_shardings=(x_named_sharding, g_named_sharding, b_named_sharding), + out_shardings=( + None, + (x_named_sharding, g_named_sharding, b_named_sharding), + ), ) except AssertionError as err: # Layernorm should still produce the correct numerical result with @@ -210,8 +216,10 @@ def ref_func(x, gamma): devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) mesh = Mesh(devices, mesh_axes) with mesh, autocast(enabled=True, recipe=fp8_recipe, mesh_resource=mesh_resource): - x_ = jax.device_put(x, NamedSharding(mesh, x_pspec)) - gamma_ = jax.device_put(gamma, NamedSharding(mesh, g_pspec)) + x_named_sharding = NamedSharding(mesh, x_pspec) + g_named_sharding = NamedSharding(mesh, g_pspec) + x_ = jax.device_put(x, x_named_sharding) + gamma_ = jax.device_put(gamma, g_named_sharding) with warnings.catch_warnings(record=True) as warns: try: @@ -223,8 +231,8 @@ def ref_func(x, gamma): grad_args=(0, 1), metric_fwd_dtype=q_dtype, metric_bwd_dtype=q_dtype, - in_shardings=(x_pspec, g_pspec), - out_shardings=(None, (x_pspec, g_pspec)), + in_shardings=(x_named_sharding, g_named_sharding), + out_shardings=(None, (x_named_sharding, g_named_sharding)), ) except AssertionError as err: # RmsNorm should still produce the correct numerical result with diff --git a/tests/jax/test_distributed_softmax.py b/tests/jax/test_distributed_softmax.py index 2bd4d862a6..f1ae6c9e49 100644 --- a/tests/jax/test_distributed_softmax.py +++ b/tests/jax/test_distributed_softmax.py @@ -103,8 +103,10 @@ def impl_test_softmax( devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) mesh = Mesh(devices, mesh_axes) with mesh, autocast(mesh_resource=mesh_resource): - x_ = jax.device_put(x, NamedSharding(mesh, x_pspec)) - mask_ = jax.device_put(mask, NamedSharding(mesh, mask_pspec)) + x_named_sharding = NamedSharding(mesh, x_pspec) + mask_named_sharding = NamedSharding(mesh, mask_pspec) + x_ = jax.device_put(x, x_named_sharding) + mask_ = jax.device_put(mask, mask_named_sharding) with warnings.catch_warnings(record=True) as warns: try: @@ -116,8 +118,8 @@ def impl_test_softmax( grad_args=(0,), metric_fwd_dtype=dtype, metric_bwd_dtype=dtype, - in_shardings=(x_pspec, mask_pspec), - out_shardings=(None, (x_pspec,)), + in_shardings=(x_named_sharding, mask_named_sharding), + out_shardings=(None, x_named_sharding), ) except AssertionError as err: # Softmax should still produce the correct numerical result with From 05dc1e624386b38bafbd2c8bba55e3ea8eb16a49 Mon Sep 17 00:00:00 2001 From: Kevin Tong Date: Fri, 17 Oct 2025 06:48:31 -0700 Subject: [PATCH 003/521] NVFP4 Move RHT BLAS to GPU (#2275) * CUDA RHT Signed-off-by: Kevin Tong * Fix cuda graphs Signed-off-by: Kirthi Shankar Sivamani * Fix bug where RHT mask is tensor instead of int Signed-off-by: Tim Moon --------- Signed-off-by: Kevin Tong Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Tim Moon Co-authored-by: Kirthi Shankar Sivamani Co-authored-by: Tim Moon --- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index ca2154f554..5e2eeed726 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -29,7 +29,7 @@ def get_no_random_sign_vector() -> torch.Tensor: """Non-random sign vector for Hadamard transform.""" - return torch.tensor([1], dtype=torch.float32) + return torch.tensor([1], dtype=torch.float32, device="cuda") def get_sign_from_vector(vector: torch.Tensor) -> int: @@ -41,7 +41,7 @@ def get_sign_from_vector(vector: torch.Tensor) -> int: mask = 0 for i, v in enumerate(vector): mask |= (v == -1) << i - return mask + return mask.item() def get_wgrad_sign_vector() -> torch.Tensor: @@ -53,6 +53,7 @@ def get_wgrad_sign_vector() -> torch.Tensor: return torch.tensor( [1, 1, 1, -1, 1, -1, -1, -1, -1, -1, -1, 1, -1, 1, -1, -1], dtype=torch.float32, + device="cuda", ) @@ -81,6 +82,7 @@ def get_hadamard_matrix(hadamard_dimension: int) -> torch.Tensor: [1, -1, -1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, -1, 1], ], dtype=torch.float32, + device="cuda", ) * hadamard_scale ) @@ -94,9 +96,9 @@ def get_rht_matrix(with_random_sign_mask: bool) -> torch.Tensor: signs = get_wgrad_sign_vector() else: signs = get_no_random_sign_vector() - sign_matrix = signs * torch.eye(hadamard_dimension, dtype=torch.float32) + sign_matrix = signs * torch.eye(hadamard_dimension, dtype=torch.float32, device="cuda") rht_matrix = sign_matrix @ get_hadamard_matrix(hadamard_dimension) - return rht_matrix.to(dtype=torch.bfloat16).cuda() + return rht_matrix.to(dtype=torch.bfloat16) @functools.lru_cache(maxsize=None) From bd38004800052f88e9773d1db5c84b596f0f5861 Mon Sep 17 00:00:00 2001 From: Tim Geypens Date: Fri, 17 Oct 2025 19:21:59 +0200 Subject: [PATCH 004/521] fall back after failing ldconfig-based lib loading for cuDNN (#2277) Signed-off-by: Tim Geypens --- transformer_engine/common/__init__.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index dd1ec480b2..134705f600 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -252,9 +252,7 @@ def _load_cudnn(): return handle # Attempt to locate libcudnn via ldconfig - libs = subprocess.check_output( - f"ldconfig -p | grep 'libcudnn{_get_sys_extension()}'", shell=True - ) + libs = subprocess.check_output(["ldconfig", "-p"]) libs = libs.decode("utf-8").split("\n") sos = [] for lib in libs: @@ -284,9 +282,7 @@ def _load_nvrtc(): return handle # Attempt to locate NVRTC via ldconfig - libs = subprocess.check_output( - f"ldconfig -p | grep 'libnvrtc{_get_sys_extension()}'", shell=True - ) + libs = subprocess.check_output(["ldconfig", "-p"]) libs = libs.decode("utf-8").split("\n") sos = [] for lib in libs: @@ -316,9 +312,7 @@ def _load_curand(): return handle # Attempt to locate cuRAND via ldconfig - libs = subprocess.check_output( - f"ldconfig -p | grep 'libcurand{_get_sys_extension()}'", shell=True - ) + libs = subprocess.check_output(["ldconfig", "-p"]) libs = libs.decode("utf-8").split("\n") sos = [] for lib in libs: From a7a69ca61c050df7bb78dc5a7f0a0077f6f57946 Mon Sep 17 00:00:00 2001 From: Haowen Zheng <157908761+Owen1B@users.noreply.github.com> Date: Sat, 18 Oct 2025 01:26:41 +0800 Subject: [PATCH 005/521] Bump up FA to 2.8.3 (#2282) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 将来 Co-authored-by: 将来 Co-authored-by: Kirthi Shankar Sivamani --- qa/L3_pytorch_FA_versions_test/test.sh | 4 ++-- .../pytorch/attention/dot_product_attention/utils.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 7e9616cd03..418e824c10 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -18,10 +18,10 @@ sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); pri export FLASH_ATTN_CUDA_ARCHS=$sm_arch if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.1) + FA_versions=(2.8.3) elif [ $sm_arch -eq 90 ] then - FA_versions=(2.7.3 2.8.1 3.0.0b1) + FA_versions=(2.7.3 2.8.3 3.0.0b1) fi for fa_version in "${FA_versions[@]}" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index b45edc716d..174d7ee9e4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -115,7 +115,7 @@ class FlashAttentionUtils: version = PkgVersion("0") version_required = PkgVersion("2.1.1") version_required_blackwell = PkgVersion("2.7.3") - max_version = PkgVersion("2.8.1") + max_version = PkgVersion("2.8.3") v2_plus = False v2_1_plus = False v2_3_plus = False From c593bcefc1379a5b3d795bf913315cd209c59835 Mon Sep 17 00:00:00 2001 From: Neil Tenenholtz Date: Fri, 17 Oct 2025 13:31:21 -0400 Subject: [PATCH 006/521] Fix test of FSDP2 by correcting init logic and applying autocast (#2105) * Fix test of FSDP2 by correcting init logic and applying autocast This fixes multiple issues in the FSDP2 test, namely 1. Previously fp8 init was performed when `args.fp8_init == False`. I have updated the logic to match what I presume was intended by leveraging the nullcontext context manager. 2. `te.fp8_autocast` was previously not called; the recipe was created but was unused. The autocast context manager now wraps the model's computation. Signed-off-by: Neil Tenenholtz * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix typo Signed-off-by: Neil Tenenholtz * Update tests/pytorch/distributed/run_fsdp2_model.py 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 * Fix bug when constructing context for model init Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Neil Tenenholtz 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> --- tests/pytorch/distributed/run_fsdp2_model.py | 23 +++++++++----------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/run_fsdp2_model.py index 8026fc0a34..d3f8c82baa 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/run_fsdp2_model.py @@ -105,23 +105,19 @@ def _train(args): fp8_format = Format.HYBRID fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max") - if not args.fp8_init: - # Build model context (FP8 init) - build_model_context = nullcontext - build_model_context_args = {} - + # Create build context manager + if args.fp8_init: from transformer_engine.pytorch import quantized_model_init - build_model_context = quantized_model_init - build_model_context_args["enabled"] = True - - # Build the model with the specified context - with build_model_context(**build_model_context_args): - model = SimpleNet(args.input_size, args.hidden_size, args.output_size) + build_model_context = quantized_model_init() else: + build_model_context = nullcontext() + + # Build the model with the specified context + with build_model_context: model = SimpleNet(args.input_size, args.hidden_size, args.output_size) - # Move the model to the correct device + # Move the model to the correct device model.to(device) if LOCAL_RANK == 0: @@ -163,7 +159,8 @@ def _train(args): # Zero the parameter gradients optimizer.zero_grad() input_data = torch.randn(args.batch_size, args.input_size).to(device) - output = model(input_data) + with te.autocast(enabled=True, recipe=fp8_recipe): + output = model(input_data) target = torch.randn(args.batch_size, args.output_size).to(device) loss = F.mse_loss(output, target) loss.backward() From ee384ab566709144e91c2949625bb1f1357bee50 Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Fri, 17 Oct 2025 13:02:08 -0500 Subject: [PATCH 007/521] Make `CanonicalizeGemmInput()` support non-TN layout FP8 GEMM on Blackwell with column-wise/transposed data (#2233) Modified CanonicalizeGemmInput() logic to pull from column-wise data for FP8 GEMM on Blackwell when row-wise is not available. Signed-off-by: Alp Dener --- .../common/gemm/cublaslt_gemm.cu | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 84a1b735a4..97e8ec9a3e 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -140,6 +140,16 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla } else { NVTE_CHECK(!is_fp8_dtype(ret.Atype), "Input A is missing column-wise usage"); } + } else if (nvte_is_non_tn_fp8_gemm_supported() && !A.has_data()) { + // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed + // data with the mirrored transpose-flag if we don't have row-wise data. + NVTE_CHECK(A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype), + "Input A is missing column-wise usage"); + ret.A = A.columnwise_data.dptr; + ret.transA = is_A_transposed ? CUBLAS_OP_N : CUBLAS_OP_T; + ret.Atype = A.columnwise_data.dtype; + ret.A_scale_inv = A.columnwise_scale_inv.dptr; + ret.lda = is_A_transposed ? m : k; } if (is_fp8_dtype(ret.Atype)) { @@ -221,6 +231,16 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla } else { NVTE_CHECK(!is_fp8_dtype(ret.Btype), "Input B is missing column-wise usage"); } + } else if (nvte_is_non_tn_fp8_gemm_supported() && !B.has_data()) { + // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed + // data with the mirrored transpose-flag if we don't have row-wise data. + NVTE_CHECK(B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype), + "Input B is missing column-wise usage"); + ret.B = B.columnwise_data.dptr; + ret.transB = is_B_transposed ? CUBLAS_OP_N : CUBLAS_OP_T; + ret.Btype = B.columnwise_data.dtype; + ret.B_scale_inv = B.columnwise_scale_inv.dptr; + ret.ldb = is_B_transposed ? k : n; } if (is_fp8_dtype(ret.Atype)) { From fd234d8006c9b9fc293f569be36d9278563e99db Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Sat, 18 Oct 2025 00:00:01 -0400 Subject: [PATCH 008/521] Wheels for cuda 13 (#2278) * Support wheel build for cuda 13 Signed-off-by: Kirthi Shankar Sivamani * Fixes Signed-off-by: Kirthi Shankar Sivamani * Fixes for cu13 runtime, format Signed-off-by: Kirthi Shankar Sivamani * Add documentation Signed-off-by: Kirthi Shankar Sivamani * Better error handling Signed-off-by: Kirthi Shankar Sivamani * fix Signed-off-by: Kirthi Shankar Sivamani * fix jax sdist Signed-off-by: Kirthi Shankar Sivamani * Modify function names Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- README.rst | 2 +- build_tools/wheel_utils/Dockerfile.aarch | 29 ++++-- build_tools/wheel_utils/Dockerfile.x86 | 29 ++++-- build_tools/wheel_utils/build_wheels.sh | 18 ++-- build_tools/wheel_utils/launch_aarch.sh | 28 ++++- build_tools/wheel_utils/launch_x86.sh | 28 ++++- docs/installation.rst | 8 ++ setup.py | 5 +- transformer_engine/common/__init__.py | 124 ++++++++++++++++------- transformer_engine/jax/setup.py | 32 +++++- transformer_engine/pytorch/setup.py | 14 ++- 11 files changed, 243 insertions(+), 74 deletions(-) diff --git a/README.rst b/README.rst index 9b65c60ae8..50c1dcd807 100644 --- a/README.rst +++ b/README.rst @@ -205,7 +205,7 @@ pip Installation **Prerequisites for pip installation:** * A compatible C++ compiler -* CUDA Toolkit with cuDNN and NVCC (NVIDIA CUDA Compiler) installed +* CUDA Toolkit with cuDNN and NVCC (NVIDIA CUDA Compiler) if installing from source. To install the latest stable version with pip: diff --git a/build_tools/wheel_utils/Dockerfile.aarch b/build_tools/wheel_utils/Dockerfile.aarch index 223c4a7f1c..404cb941cb 100644 --- a/build_tools/wheel_utils/Dockerfile.aarch +++ b/build_tools/wheel_utils/Dockerfile.aarch @@ -7,23 +7,34 @@ FROM quay.io/pypa/manylinux_2_28_aarch64 WORKDIR /TransformerEngine/ COPY ../.. /TransformerEngine/ -ARG VER="12-3" -ARG ARCH="aarch64" -RUN dnf -y install vim +ARG CUDA_MAJOR="12" +ARG CUDA_MINOR="3" + +# Args for build_wheels.sh +ARG BUILD_METAPACKAGE=true +ARG BUILD_COMMON=true +ARG BUILD_PYTORCH=true +ARG BUILD_JAX=true +ENV BUILD_METAPACKAGE=${BUILD_METAPACKAGE} +ENV BUILD_COMMON=${BUILD_COMMON} +ENV BUILD_PYTORCH=${BUILD_PYTORCH} +ENV BUILD_JAX=${BUILD_JAX} +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-${VER}.${ARCH} \ - cuda-libraries-${VER}.${ARCH} \ - cuda-libraries-devel-${VER}.${ARCH} -RUN dnf -y install --allowerasing cudnn9-cuda-12 +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 +RUN dnf -y install --allowerasing cudnn9-cuda-${CUDA_MAJOR} RUN dnf clean all RUN rm -rf /var/cache/dnf/* RUN echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf -RUN dnf -y install cuda-toolkit +RUN dnf -y install cuda-toolkit-${CUDA_MAJOR} RUN dnf clean all RUN dnf -y install glog.aarch64 glog-devel.aarch64 +RUN dnf -y install libnccl libnccl-devel libnccl-static ENV PATH="/usr/local/cuda/bin:${PATH}" ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" @@ -33,4 +44,4 @@ ENV CUDA_PATH=/usr/local/cuda ENV CUDADIR=/usr/local/cuda ENV NVTE_RELEASE_BUILD=1 -CMD ["/bin/bash", "/TransformerEngine/build_tools/wheel_utils/build_wheels.sh", "manylinux_2_28_aarch64", "true", "true", "false", "false", "false"] +CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_aarch64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"] diff --git a/build_tools/wheel_utils/Dockerfile.x86 b/build_tools/wheel_utils/Dockerfile.x86 index 26122eed9b..daa7f961cd 100644 --- a/build_tools/wheel_utils/Dockerfile.x86 +++ b/build_tools/wheel_utils/Dockerfile.x86 @@ -7,23 +7,34 @@ FROM quay.io/pypa/manylinux_2_28_x86_64 WORKDIR /TransformerEngine/ COPY ../.. /TransformerEngine/ -ARG VER="12-3" -ARG ARCH="x86_64" -RUN dnf -y install vim +ARG CUDA_MAJOR="12" +ARG CUDA_MINOR="3" + +# Args for build_wheels.sh +ARG BUILD_METAPACKAGE=true +ARG BUILD_COMMON=true +ARG BUILD_PYTORCH=true +ARG BUILD_JAX=true +ENV BUILD_METAPACKAGE=${BUILD_METAPACKAGE} +ENV BUILD_COMMON=${BUILD_COMMON} +ENV BUILD_PYTORCH=${BUILD_PYTORCH} +ENV BUILD_JAX=${BUILD_JAX} +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-${VER}.${ARCH} \ - cuda-libraries-${VER}.${ARCH} \ - cuda-libraries-devel-${VER}.${ARCH} -RUN dnf -y install --allowerasing cudnn9-cuda-12 +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 +RUN dnf -y install --allowerasing cudnn9-cuda-${CUDA_MAJOR} RUN dnf clean all RUN rm -rf /var/cache/dnf/* RUN echo "/usr/local/cuda/lib64" >> /etc/ld.so.conf.d/999_nvidia_cuda.conf -RUN dnf -y install cuda-toolkit +RUN dnf -y install cuda-toolkit-${CUDA_MAJOR} RUN dnf clean all RUN dnf -y install glog.x86_64 glog-devel.x86_64 +RUN dnf -y install libnccl libnccl-devel libnccl-static ENV PATH="/usr/local/cuda/bin:${PATH}" ENV LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" @@ -33,4 +44,4 @@ ENV CUDA_PATH=/usr/local/cuda ENV CUDADIR=/usr/local/cuda ENV NVTE_RELEASE_BUILD=1 -CMD ["/bin/bash", "/TransformerEngine/build_tools/wheel_utils/build_wheels.sh", "manylinux_2_28_x86_64", "true", "true", "true", "true", "true"] +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 diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index bf4f9d2bc2..954a8f1c67 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -9,8 +9,10 @@ BUILD_METAPACKAGE=${2:-true} BUILD_COMMON=${3:-true} BUILD_PYTORCH=${4:-true} BUILD_JAX=${5:-true} +CUDA_MAJOR=${6:-12} export NVTE_RELEASE_BUILD=1 +export PIP_CONSTRAINT="" export TARGET_BRANCH=${TARGET_BRANCH:-} mkdir -p /wheelhouse/logs @@ -21,7 +23,7 @@ git checkout $TARGET_BRANCH git submodule update --init --recursive # Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja +/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel nvidia-mathdx==25.1.1 if $BUILD_METAPACKAGE ; then cd /TransformerEngine @@ -36,32 +38,32 @@ if $BUILD_COMMON ; then # Create the wheel. /opt/python/cp310-cp310/bin/python setup.py bdist_wheel --verbose --python-tag=py3 --plat-name=$PLATFORM 2>&1 | tee /wheelhouse/logs/common.txt - # Repack the wheel for cuda specific package, i.e. cu12. + # Repack the wheel for specific cuda version. /opt/python/cp310-cp310/bin/wheel unpack dist/* # From python 3.10 to 3.11, the package name delimiter in metadata got changed from - (hyphen) to _ (underscore). - sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" - sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" - mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" + sed -i "s/Name: transformer-engine/Name: transformer-engine-cu${CUDA_MAJOR}/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" + sed -i "s/Name: transformer_engine/Name: transformer_engine_cu${CUDA_MAJOR}/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" + mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu${CUDA_MAJOR}-${VERSION}.dist-info" /opt/python/cp310-cp310/bin/wheel pack ${WHL_BASE} # Rename the wheel to make it python version agnostic. whl_name=$(basename dist/*) IFS='-' read -ra whl_parts <<< "$whl_name" - whl_name_target="${whl_parts[0]}_cu12-${whl_parts[1]}-py3-none-${whl_parts[4]}" + whl_name_target="${whl_parts[0]}_cu${CUDA_MAJOR}-${whl_parts[1]}-py3-none-${whl_parts[4]}" rm -rf $WHL_BASE dist mv *.whl /wheelhouse/"$whl_name_target" fi if $BUILD_PYTORCH ; then cd /TransformerEngine/transformer_engine/pytorch - /opt/python/cp310-cp310/bin/pip install torch + /opt/python/cp310-cp310/bin/pip install torch /opt/python/cp310-cp310/bin/python setup.py sdist 2>&1 | tee /wheelhouse/logs/torch.txt cp dist/* /wheelhouse/ fi if $BUILD_JAX ; then cd /TransformerEngine/transformer_engine/jax - /opt/python/cp310-cp310/bin/pip install "jax[cuda12_local]" jaxlib + /opt/python/cp310-cp310/bin/pip install "jax[cuda${CUDA_MAJOR}_local]" jaxlib /opt/python/cp310-cp310/bin/python setup.py sdist 2>&1 | tee /wheelhouse/logs/jax.txt cp dist/* /wheelhouse/ fi diff --git a/build_tools/wheel_utils/launch_aarch.sh b/build_tools/wheel_utils/launch_aarch.sh index 04e3cd6916..85f754ca19 100644 --- a/build_tools/wheel_utils/launch_aarch.sh +++ b/build_tools/wheel_utils/launch_aarch.sh @@ -2,7 +2,29 @@ # # See LICENSE for license information. -docker build --no-cache -t "aarch_wheel" -f build_tools/wheel_utils/Dockerfile.aarch . +# Remove leftovers. +rm -rf aarch_wheelhouse_cu12 aarch_wheelhouse_cu13 + +# CUDA 12. +docker build --no-cache \ + --build-arg CUDA_MAJOR=12 \ + --build-arg CUDA_MINOR=3 \ + --build-arg BUILD_METAPACKAGE=false \ + --build-arg BUILD_COMMON=true \ + --build-arg BUILD_PYTORCH=false \ + --build-arg BUILD_JAX=false \ + -t "aarch_wheel" -f build_tools/wheel_utils/Dockerfile.aarch . +docker run --runtime=nvidia --gpus=all --ipc=host "aarch_wheel" +docker cp $(docker ps -aq | head -1):/wheelhouse aarch_wheelhouse_cu12 + +# CUDA 13. +docker build --no-cache \ + --build-arg CUDA_MAJOR=13 \ + --build-arg CUDA_MINOR=0 \ + --build-arg BUILD_METAPACKAGE=false \ + --build-arg BUILD_COMMON=true \ + --build-arg BUILD_PYTORCH=false \ + --build-arg BUILD_JAX=false \ + -t "aarch_wheel" -f build_tools/wheel_utils/Dockerfile.aarch . docker run --runtime=nvidia --gpus=all --ipc=host "aarch_wheel" -rm -rf aarch_wheelhouse -docker cp $(docker ps -aq | head -1):/wheelhouse/ aarch_wheelhouse +docker cp $(docker ps -aq | head -1):/wheelhouse aarch_wheelhouse_cu13 diff --git a/build_tools/wheel_utils/launch_x86.sh b/build_tools/wheel_utils/launch_x86.sh index b0d20be3f4..11fc522947 100644 --- a/build_tools/wheel_utils/launch_x86.sh +++ b/build_tools/wheel_utils/launch_x86.sh @@ -2,7 +2,29 @@ # # See LICENSE for license information. -docker build --no-cache -t "x86_wheel" -f build_tools/wheel_utils/Dockerfile.x86 . +# Remove leftovers. +rm -rf x86_wheelhouse_cu12 x86_wheelhouse_cu13 + +# CUDA 12. +docker build --no-cache \ + --build-arg CUDA_MAJOR=12 \ + --build-arg CUDA_MINOR=3 \ + --build-arg BUILD_METAPACKAGE=true \ + --build-arg BUILD_COMMON=true \ + --build-arg BUILD_PYTORCH=true \ + --build-arg BUILD_JAX=true \ + -t "x86_wheel" -f build_tools/wheel_utils/Dockerfile.x86 . +docker run --runtime=nvidia --gpus=all --ipc=host "x86_wheel" +docker cp $(docker ps -aq | head -1):/wheelhouse x86_wheelhouse_cu12 + +# CUDA 13. +docker build --no-cache \ + --build-arg CUDA_MAJOR=13 \ + --build-arg CUDA_MINOR=0 \ + --build-arg BUILD_METAPACKAGE=false \ + --build-arg BUILD_COMMON=true \ + --build-arg BUILD_PYTORCH=false \ + --build-arg BUILD_JAX=false \ + -t "x86_wheel" -f build_tools/wheel_utils/Dockerfile.x86 . docker run --runtime=nvidia --gpus=all --ipc=host "x86_wheel" -rm -rf x86_wheelhouse -docker cp $(docker ps -aq | head -1):/wheelhouse x86_wheelhouse +docker cp $(docker ps -aq | head -1):/wheelhouse x86_wheelhouse_cu13 diff --git a/docs/installation.rst b/docs/installation.rst index ecb1e9a0dd..a8bb74fd1a 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -38,6 +38,14 @@ Transformer Engine can be directly installed from `our PyPI Tuple[List[str], List[str]]: ext_modules = [] package_data = {} include_package_data = False - install_requires = ([f"transformer_engine_cu12=={__version__}"],) + install_requires = [] extras_require = { + "core": [f"transformer_engine_cu12=={__version__}"], + "core_cu12": [f"transformer_engine_cu12=={__version__}"], + "core_cu13": [f"transformer_engine_cu13=={__version__}"], "pytorch": [f"transformer_engine_torch=={__version__}"], "jax": [f"transformer_engine_jax=={__version__}"], } diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 134705f600..3ffe1c7b1d 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -8,22 +8,18 @@ import functools import glob import importlib -from importlib.metadata import version, metadata, PackageNotFoundError -import logging +from importlib.metadata import version, distribution, PackageNotFoundError import os from pathlib import Path import platform import subprocess import sys import sysconfig -from typing import Optional - - -_logger = logging.getLogger(__name__) +from typing import Optional, Tuple @functools.lru_cache(maxsize=None) -def _is_pip_package_installed(package) -> bool: +def _is_package_installed(package) -> bool: """Check if the given package is installed via pip.""" # This is needed because we only want to return true @@ -31,12 +27,34 @@ def _is_pip_package_installed(package) -> bool: # if it's importable in the current directory due to # the presence of the shared library module. try: - metadata(package) + distribution(package) except PackageNotFoundError: return False return True +@functools.lru_cache(maxsize=None) +def _is_package_installed_from_wheel(package) -> bool: + """Check if the given package is installed via PyPI.""" + + if not _is_package_installed(package): + return False + + te_dist = distribution(package) + te_wheel_file = "" + for file_path in te_dist.files: + if file_path.name == "WHEEL": + te_wheel_file = te_dist.locate_file("") / file_path + if not te_wheel_file: + return False + + with te_wheel_file.open("r") as f: + for line in f: + if line.startswith("Root-Is-Purelib:"): + return line.strip().split(":")[1].strip().lower() == "true" + return False + + @functools.lru_cache(maxsize=None) def _find_shared_object_in_te_dir(te_path: Path, prefix: str) -> Optional[Path]: """ @@ -112,6 +130,19 @@ def _get_shared_object_file(library: str) -> Path: ) +def get_te_core_package_info() -> Tuple[bool, str, str]: + """ + Check if Tranformer Engine core package is installed. + Returns the module name and version if found. + """ + + te_core_packages = ("transformer-engine-cu12", "transformer-engine-cu13") + for package in te_core_packages: + if _is_package_installed(package): + return True, package, version(package) + return False, "", "" + + @functools.lru_cache(maxsize=None) def load_framework_extension(framework: str) -> None: """ @@ -130,39 +161,30 @@ def load_framework_extension(framework: str) -> None: if framework == "torch": extra_dep_name = "pytorch" + # Find the TE packages. The core and framework packages can only be installed via PyPI. + # For the `transformer-engine` package, we need to check explicity. + te_core_installed, te_core_package_name, te_core_version = get_te_core_package_info() + te_framework_installed = _is_package_installed(module_name) + te_installed = _is_package_installed("transformer_engine") + te_installed_via_pypi = _is_package_installed_from_wheel("transformer_engine") + + assert te_installed, "Could not find `transformer_engine`." + # If the framework extension pip package is installed, it means that TE is installed via # PyPI. For this case we need to make sure that the metapackage, the core lib, and framework - # extension are all installed via PyPI and have matching version. - if _is_pip_package_installed(module_name): - assert _is_pip_package_installed( - "transformer_engine" - ), "Could not find `transformer-engine`." - assert _is_pip_package_installed( - "transformer_engine_cu12" - ), "Could not find `transformer-engine-cu12`." - assert ( - version(module_name) - == version("transformer-engine") - == version("transformer-engine-cu12") - ), ( - "TransformerEngine package version mismatch. Found" + # extension are all installed via PyPI and have matching versions. + if te_framework_installed: + assert te_installed_via_pypi, "Could not find `transformer-engine` PyPI package." + assert te_core_installed, "Could not find TE core package `transformer-engine-cu*`." + + assert version(module_name) == version("transformer-engine") == te_core_version, ( + "Transformer Engine package version mismatch. Found" f" {module_name} v{version(module_name)}, transformer-engine" - f" v{version('transformer-engine')}, and transformer-engine-cu12" - f" v{version('transformer-engine-cu12')}. Install transformer-engine using " - f"'pip3 install transformer-engine[{extra_dep_name}]==VERSION'" + f" v{version('transformer-engine')}, and {te_core_package_name}" + f" v{te_core_version}. Install transformer-engine using " + f"'pip3 install --no-build-isolation transformer-engine[{extra_dep_name}]==VERSION'" ) - # If the core package is installed via PyPI, log if - # the framework extension is not found from PyPI. - # Note: Should we error? This is a rare use case. - if _is_pip_package_installed("transformer-engine-cu12"): - if not _is_pip_package_installed(module_name): - _logger.info( - "Could not find package %s. Install transformer-engine using " - f"'pip3 install transformer-engine[{extra_dep_name}]==VERSION'", - module_name, - ) - # After all checks are completed, load the shared object file. spec = importlib.util.spec_from_file_location(module_name, _get_shared_object_file(framework)) solib = importlib.util.module_from_spec(spec) @@ -170,6 +192,35 @@ def load_framework_extension(framework: str) -> None: spec.loader.exec_module(solib) +def sanity_checks_for_pypi_installation() -> None: + """Ensure that package is installed correctly if using PyPI.""" + + te_core_installed, te_core_package_name, te_core_version = get_te_core_package_info() + te_installed = _is_package_installed("transformer_engine") + te_installed_via_pypi = _is_package_installed_from_wheel("transformer_engine") + + assert te_installed, "Could not find `transformer-engine`." + + # If the core package is installed via PyPI. + if te_core_installed: + assert te_installed_via_pypi, "Could not find `transformer-engine` PyPI package." + assert version("transformer-engine") == te_core_version, ( + "Transformer Engine package version mismatch. Found " + f"transformer-engine v{version('transformer-engine')} " + f"and {te_core_package_name} v{te_core_version}." + ) + + # Only the metapackage is found, invalid usecase. + elif te_installed_via_pypi: + raise RuntimeError( + "Found empty `transformer-engine` meta package installed. " + "Install `transformer-engine` with framework extensions via" + "'pip3 install --no-build-isolation transformer-engine[pytorch,jax]==VERSION'" + " or 'pip3 install transformer-engine[core]` for the TE core lib only. The `core_cu12`" + " or `core_cu13` extra deps can be used to specify CUDA version for the TE core lib." + ) + + @functools.lru_cache(maxsize=None) def _get_sys_extension() -> str: """File extension for shared objects.""" @@ -332,6 +383,7 @@ def _load_core_library(): if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): + sanity_checks_for_pypi_installation() _CUDNN_LIB_CTYPES = _load_cudnn() _NVRTC_LIB_CTYPES = _load_nvrtc() _CURAND_LIB_CTYPES = _load_curand() diff --git a/transformer_engine/jax/setup.py b/transformer_engine/jax/setup.py index f83375d821..ccdbcdb529 100644 --- a/transformer_engine/jax/setup.py +++ b/transformer_engine/jax/setup.py @@ -54,6 +54,26 @@ CMakeBuildExtension = get_build_ext(BuildExtension, True) +def get_cuda_major_version() -> int: + """Get CUDA major version using Jax backend.""" + + assert ( + jax._src.lib.cuda_versions is not None + ), "GPU backend is required to build TE jax extensions." + + # Jax currently does not have any stable/public method to get cuda version. + # Try using internal function and default to cuda12 if not found. + try: + cuda_version = jax._src.lib.cuda_versions.cuda_runtime_get_version() + cuda_major_version = cuda_version // 1000 + except AttributeError: + cuda_version = os.getenv("CUDA_VERSION", "12") + cuda_major_version = int(cuda_version.split(".")[0]) + + assert cuda_major_version in (12, 13), f"Unsupported cuda version {cuda_version}." + return cuda_major_version + + if __name__ == "__main__": """Main entry point for JAX extension installation. @@ -93,15 +113,23 @@ ) ] + # Setup version and requirements. + # Having the framework extension depend on the core lib allows + # us to detect CUDA version dynamically during compilation and + # choose the correct wheel for te core lib. + __version__ = te_version() + te_core = f"transformer_engine_cu{get_cuda_major_version()}=={__version__}" + install_requires = install_requirements() + [te_core] + # Configure package setuptools.setup( name="transformer_engine_jax", - version=te_version(), + version=__version__, description="Transformer acceleration library - Jax Lib", ext_modules=ext_modules, cmdclass={"build_ext": CMakeBuildExtension}, python_requires=f">={min_python_version_str()}", - install_requires=install_requirements(), + install_requires=install_requires, tests_require=test_requirements(), ) if any(x in sys.argv for x in (".", "sdist", "bdist_wheel")): diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py index 08870040f3..7a81550047 100644 --- a/transformer_engine/pytorch/setup.py +++ b/transformer_engine/pytorch/setup.py @@ -145,15 +145,25 @@ def run(self): ) ] + # Setup version and requirements. + # Having the framework extension depend on the core lib allows + # us to detect CUDA version dynamically during compilation and + # choose the correct wheel for te core lib. + __version__ = te_version() + cuda_major_version = parse(torch.version.cuda).major + assert cuda_major_version in (12, 13), f"Unsupported cuda version {torch.version.cuda}." + te_core = f"transformer_engine_cu{cuda_major_version}=={__version__}" + install_requires = install_requirements() + [te_core] + # Configure package setuptools.setup( name=PACKAGE_NAME, - version=te_version(), + version=__version__, description="Transformer acceleration library - Torch Lib", ext_modules=ext_modules, cmdclass={"build_ext": CMakeBuildExtension, "bdist_wheel": CachedWheelsCommand}, python_requires=f">={min_python_version_str()}", - install_requires=install_requirements(), + install_requires=install_requires, tests_require=test_requirements(), ) if any(x in sys.argv for x in (".", "sdist", "bdist_wheel")): From dd7ab715a55d18740de5f10546ac71842f832e07 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 20 Oct 2025 23:02:05 +0800 Subject: [PATCH 009/521] Fix error with triton 3.5 (#2286) * Update permutation.py Signed-off-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> * Update permutation.py Signed-off-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> * Update transformer_engine/pytorch/triton/permutation.py Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/pytorch/triton/permutation.py Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/triton/permutation.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 6292acb69b..1474a664cf 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -12,11 +12,16 @@ from triton.language import core from triton.language.standard import _log2 +from packaging import version # The following three argsort related kernels are adapted from # the issue https://github.com/triton-lang/triton/issues/3698 +get_int_dtype = core.get_int_dtype +if version.parse(triton.__version__) >= version.parse("3.5.0"): + get_int_dtype = triton.constexpr_function(get_int_dtype) + @triton.jit def _compare_and_swap(x, indices, flip, i: tl.constexpr, n_dims: tl.constexpr): @@ -37,7 +42,7 @@ def _compare_and_swap(x, indices, flip, i: tl.constexpr, n_dims: tl.constexpr): l_indice = tl.reshape(tl.broadcast_to(tl.sum(z * (1 - mask), 1)[:, None, :], shape), x.shape) r_indice = tl.reshape(tl.broadcast_to(tl.sum(z * mask, 1)[:, None, :], shape), x.shape) - idtype = core.get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) + idtype = get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) il_value = l_value.to(idtype, bitcast=True) ir_value = r_value.to(idtype, bitcast=True) From bd55e7ba5f0235a80eaa63d49adaa8fb7c6ced50 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 20 Oct 2025 16:28:23 -0400 Subject: [PATCH 010/521] [PyTorch] Fix CI failures due to deterministic attention backend (#2288) * Fix CI failures due to deterministic attention Signed-off-by: Kirthi Shankar Sivamani * some more cleanup Signed-off-by: Kirthi Shankar Sivamani * Fix debug test Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- qa/L0_pytorch_debug_unittest/test.sh | 2 +- qa/L0_pytorch_unittest/test.sh | 4 +-- tests/pytorch/test_numerics.py | 30 +------------------ .../attention/dot_product_attention/utils.py | 2 +- 4 files changed, 5 insertions(+), 33 deletions(-) diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index 7f19dda670..9980ccfb05 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -32,6 +32,6 @@ pytest -v -s --junitxml=$XML_LOG_DIR/test_perf.xml $TE_PATH/tests/pytorch/debug/ # 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 || FAIL=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_numerics_2.xml $TE_PATH/tests/pytorch/test_numerics.py || FAIL=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 NVTE_FUSED_ATTN=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics_2.xml $TE_PATH/tests/pytorch/test_numerics.py || FAIL=1 exit $FAIL diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index cdf0df8887..b23ce3b6cf 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -27,8 +27,8 @@ 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" 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 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 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" +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" 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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_nvfp4.xml $TE_PATH/tests/pytorch/nvfp4 || test_fail "test_nvfp4" diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index bef076a385..35698b819c 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -43,11 +43,10 @@ ) from transformer_engine.pytorch import checkpoint as te_checkpoint from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm -from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend from transformer_engine.pytorch.module.base import get_multi_stream_cublas_workspace, get_workspace from transformer_engine.common import recipe import transformer_engine_torch as tex -from utils import ModelConfig, reset_rng_states, get_available_attention_backends +from utils import ModelConfig, reset_rng_states # Only run FP8 tests on supported devices. @@ -130,23 +129,6 @@ use_cutlass_grouped_gemm.append(True) -def is_fused_attn_available( - config: ModelConfig, - dtype: torch.dtype, - qkv_layout="bshd_bshd_bshd", - is_training=True, - deterministic=False, -): - _, _, fused_attn_backends = get_available_attention_backends( - config, - qkv_dtype=dtype, - qkv_layout=qkv_layout, - is_training=is_training, - deterministic=deterministic, - ) - return FusedAttnBackend["F16_arbitrary_seqlen"] in fused_attn_backends - - def get_causal_attn_mask(sq: int) -> torch.Tensor: return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool() @@ -853,8 +835,6 @@ def _test_e2e_checkpointing(bs, dtype, config, checkpoint=False, steps=10, path= @pytest.mark.parametrize("model", ["126m"]) def test_gpt_checkpointing(dtype, bs, model): config = model_configs[model] - if not is_fused_attn_available(config, dtype, deterministic=True): - pytest.skip("No attention backend available.") outputs = _test_e2e_checkpointing(bs, dtype, config, checkpoint=False) outputs_checkpoint = _test_e2e_checkpointing(bs, dtype, config, checkpoint=True) @@ -901,10 +881,6 @@ def _test_e2e_gpt_accuracy(block, bs, dtype, config): @pytest.mark.parametrize("parallel_attention_mlp", all_boolean) def test_gpt_accuracy(dtype, bs, model, parallel_attention_mlp): config = model_configs[model] - if not is_fused_attn_available( - config, dtype, qkv_layout="sb3hd", is_training=True, deterministic=True - ): - pytest.skip("No attention backend available.") te_gpt = TransformerLayer( hidden_size=config.hidden_size, @@ -1016,10 +992,6 @@ def _test_mha_accuracy(block, bs, dtype, config, mask_type, te=True): @pytest.mark.parametrize("mask_type", mask_types) def test_mha_accuracy(dtype, bs, model, mask_type): config = model_configs[model] - if not is_fused_attn_available( - config, dtype, qkv_layout="sb3hd", is_training=True, deterministic=True - ): - pytest.skip("No attention backend available.") te_mha = MultiheadAttention( config.hidden_size, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 174d7ee9e4..4cb39cda09 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -983,7 +983,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False fused_attention_backend = None - if is_training and device_compute_capability >= (10, 0) and cudnn_version <= (9, 14, 0): + if is_training and device_compute_capability >= (10, 0): logger.debug("Disabling FusedAttention for determinism reasons on Blackwell") use_fused_attention = False fused_attention_backend = None From b4a1d4d6f4f00a3b30d305c72cd040ae95ea41e4 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Mon, 20 Oct 2025 21:15:16 -0700 Subject: [PATCH 011/521] [PyTorch][MOE] Support NVFP4 Grouped Linear (#2215) * pipeclean, fix nvfp4 padding of 32 alignment Signed-off-by: Zhongbo Zhu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * numerical test passed Signed-off-by: Zhongbo Zhu * fix CI failure with test_cast_master_weights_to_fp8 (in a hacky way) Signed-off-by: Zhongbo Zhu * found CUDA mis-aligned address error in training in multi-swizzle, hack the vec_load_size to 1 to unblock Signed-off-by: Zhongbo Zhu * leave comments about alignment issue Signed-off-by: Zhongbo Zhu * fused bulk alloc nvfp4 Signed-off-by: Zhongbo Zhu * fix RHT sign mask CPU overhead Signed-off-by: Zhongbo Zhu * fix Signed-off-by: Zhongbo Zhu * resolve comments Signed-off-by: Zhongbo Zhu * Remove incorrect logic that treats 0-D tensor as uninitialized Tensor shape logic still requires treating 0-D tensor as uninitialized. Signed-off-by: Tim Moon * Fix invalid conversion from tensor to int Signed-off-by: Tim Moon --------- Signed-off-by: Zhongbo Zhu Signed-off-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- benchmarks/linear/benchmark_grouped_linear.py | 72 ++++-- tests/pytorch/test_numerics.py | 93 +++++++- transformer_engine/common/common.h | 27 ++- transformer_engine/common/swizzle/swizzle.cu | 83 +++++-- transformer_engine/pytorch/csrc/common.cpp | 5 +- .../pytorch/csrc/extensions/cast.cpp | 213 +++++++++++++++++- .../pytorch/csrc/extensions/recipe.cpp | 3 +- transformer_engine/pytorch/csrc/quantizer.cpp | 10 +- transformer_engine/pytorch/csrc/util.cpp | 59 ++--- .../pytorch/module/fp8_padding.py | 11 +- .../pytorch/module/fp8_unpadding.py | 15 +- 11 files changed, 504 insertions(+), 87 deletions(-) diff --git a/benchmarks/linear/benchmark_grouped_linear.py b/benchmarks/linear/benchmark_grouped_linear.py index 48adb2a10b..d4bbad75cd 100644 --- a/benchmarks/linear/benchmark_grouped_linear.py +++ b/benchmarks/linear/benchmark_grouped_linear.py @@ -8,53 +8,67 @@ import pandas as pd from transformer_engine.pytorch.module import GroupedLinear -from transformer_engine.common.recipe import Float8BlockScaling, MXFP8BlockScaling +from transformer_engine.common.recipe import ( + Float8BlockScaling, + MXFP8BlockScaling, + NVFP4BlockScaling, +) from transformer_engine.pytorch.quantization import autocast, FP8GlobalStateManager from contextlib import nullcontext """ # Profile BF16 recipe with Nsight Systems nsys profile \ - --output=./benchmarks/linear/b200_mkn_4096_4096_4096_numgemm_8_bf16 \ + --output=./benchmarks/linear/b200_numgemm_8_bf16 \ --force-overwrite true \ --trace=cuda,nvtx,cudnn,cublas \ python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe bf16 # Profile FP8 sub-channel recipe with Nsight Systems nsys profile \ - --output=./benchmarks/linear/h100hbm_mkn_4096_4096_4096_numgemm_8_fp8_sub_channel \ + --output=./benchmarks/linear/h100hbm_numgemm_8_fp8_sub_channel \ --force-overwrite true \ --trace=cuda,nvtx,cudnn,cublas \ python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe fp8_sub_channel # Profile MXFP8 recipe with Nsight Systems nsys profile \ - --output=./benchmarks/linear/b200_mkn_4096_4096_4096_numgemm_8_mxfp8 \ + --output=./benchmarks/linear/b200_numgemm_8_mxfp8 \ --force-overwrite true \ --trace=cuda,nvtx,cudnn,cublas \ python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe mxfp8 +# Profile NVFP4 recipe with Nsight Systems +nsys profile \ + --output=./benchmarks/linear/b200_numgemm_8_nvfp4 \ + --force-overwrite true \ + --trace=cuda,nvtx,cudnn,cublas \ + python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe nvfp4 + """ RECIPES = { "bf16": None, "fp8_sub_channel": Float8BlockScaling(), "mxfp8": MXFP8BlockScaling(), + "nvfp4": NVFP4BlockScaling(), } mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available() fp8_block_scaling_available, reason_for_no_fp8_block_scaling = ( FP8GlobalStateManager.is_fp8_block_scaling_available() ) +nvfp4_available, reason_for_no_nvfp4 = FP8GlobalStateManager.is_nvfp4_available() def run_linear_multiple_steps(layer, x, m_splits, mode, gradient, run_num_steps=1, recipe=None): assert mode in ["fwd_only", "fwd_bwd"] - fp8_context = autocast(enabled=True, fp8_recipe=recipe) if recipe is not None else nullcontext() - # print(f"fp8_context: {fp8_context} and is it nullcontext? {isinstance(fp8_context, nullcontext)}") + quantization_context = ( + autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() + ) if mode == "fwd_only": - with torch.no_grad(), fp8_context: + with torch.no_grad(), quantization_context: for i in range(run_num_steps): y_q = layer.forward( x, @@ -67,7 +81,7 @@ def run_linear_multiple_steps(layer, x, m_splits, mode, gradient, run_num_steps= layer.zero_grad() x.grad = None - with fp8_context: + with quantization_context: for i in range(run_num_steps): label = f"step_{i}" torch.cuda.nvtx.range_push(label) @@ -142,7 +156,7 @@ def benchmark_linear( "recipe": recipe, }, num_threads=1, - ).blocked_autorange(min_run_time=5) + ).blocked_autorange(min_run_time=10) print(f"{recipe_name}: {timing} \n") timing_ms = timing.median * 1000 / num_microbatches @@ -225,30 +239,44 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): use_bias = False # Set the MKN values to benchmark + # Deepseek V3 EP64, SEQ_LEN=8192, topK8 + # 256 expert => 4 local experts + # Avg M per expert: AvgM = SEQ_LEN * topK / localExperts = 16384 + # M = AvgM * localExperts = 65536 + # K = 7168 + # N = 2048 + + # Deepseek V3 EP32, SEQ_LEN=8192, topK8 + # 256 expert => 8 local experts + # Avg M per expert: AvgM = SEQ_LEN * topK / localExperts = 8192 + # M = AvgM * localExperts = 65536 + # K = 7168 + # N = 2048 + + # 4 or 8local experts per rank + num_gemms_list = [4, 8] + + # MKN for group linear mkns = [] - for m in [8192]: - # for m in [4096, 8192, 16384]: - # for n in [1024, 2048, 4096, 8192, 16384]: - for n in [8192]: - for k in [4096]: + for m in [65536]: + for k in [7168]: + for n in [2048]: mkns.append((m, k, n)) # default recipes to run if not specified recipe_list = ["bf16"] if args.recipe == "all": - recipe_list = ["bf16", "fp8_sub_channel", "mxfp8"] + recipe_list = ["bf16", "fp8_sub_channel", "mxfp8", "nvfp4"] else: recipe_list = [args.recipe] - num_gemms_list = [8] - if args.profile: - mkns = [(4096 * 8, 4096, 4096)] + mkns = [(8192 * 8, 7168, 2048)] # in profile mode, only run one recipe specified in args.recipe assert args.recipe != "all", ( "In profile mode, only one recipe can be specified, please specify the recipe as" - " fp8_sub_channel, mxfp8, or bf16" + " fp8_sub_channel, mxfp8, nvfp4, or bf16" ) recipe_list = [args.recipe] num_gemms_list = [8] @@ -265,13 +293,17 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): "bf16", "fp8_sub_channel", "mxfp8", - ], "Recipe must be one of bf16, fp8_sub_channel, or mxfp8" + "nvfp4", + ], "Recipe must be one of bf16, fp8_sub_channel, mxfp8, or nvfp4" if recipe_name == "mxfp8" and not mxfp8_available: print(f"MXFP8 is not available, skipping {recipe_name}") continue if recipe_name == "fp8_sub_channel" and not fp8_block_scaling_available: print(f"FP8 block scaling is not available, skipping {recipe_name}") continue + if recipe_name == "nvfp4" and not nvfp4_available: + print(f"NVFP4 is not available, skipping {recipe_name}") + continue df = run_benchmark_linear( mkns, diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 35698b819c..01f1deb983 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -40,6 +40,7 @@ is_mxfp8_available, is_fp8_block_scaling_available, is_bf16_available, + is_nvfp4_available, ) from transformer_engine.pytorch import checkpoint as te_checkpoint from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm @@ -53,6 +54,7 @@ fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) fp8_block_scaling_available = is_fp8_block_scaling_available() +nvfp4_available = is_nvfp4_available() sm_80plus = get_device_compute_capability() >= (8, 0) @@ -114,6 +116,43 @@ ) +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 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 + 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 using RHT, we can add fp32 as well + if not check_rht_usage(recipe): + supported_input_dtypes.append(torch.float32) + return supported_input_dtypes + + fp8_recipes = [] if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) @@ -122,6 +161,8 @@ if fp8_available: fp8_recipes.append(recipe.Float8CurrentScaling()) fp8_recipes.append(recipe.DelayedScaling()) +if nvfp4_available: + fp8_recipes.append(nvfp4_rht_and_2d_quantization()) use_cutlass_grouped_gemm = [False] # Only enable cutlass grouped gemm on Hopper @@ -582,6 +623,11 @@ def _test_e2e_selective_recompute( 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 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__}" + ) config = model_configs[model] @@ -692,6 +738,11 @@ 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 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__}" + ) config = model_configs[model] @@ -1275,6 +1326,12 @@ def test_linear_accuracy_save_original_input(dtype, model, recipe): 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): te_linear_ref = Linear( config.hidden_size, @@ -1718,8 +1775,8 @@ def _test_grouped_linear_accuracy( split_size = 1 if fp8: split_size = 16 - if recipe.mxfp8(): - split_size = 128 + if recipe.mxfp8() or recipe.nvfp4(): + split_size = 32 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 @@ -1791,6 +1848,12 @@ def test_grouped_linear_accuracy( 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, @@ -1927,6 +1990,12 @@ def test_grouped_linear_accuracy_save_original_input( 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, @@ -2014,7 +2083,7 @@ def _test_padding_grouped_linear_accuracy(block, num_gemms, bs, dtype, config, r def _pad_tensor_for_fp8(hidden_states, tokens_per_expert): align_size = 16 - if recipe.mxfp8(): + if recipe.mxfp8() or recipe.nvfp4(): align_size = 32 padded_tokens_per_expert = [ (num_tokens + align_size - 1) // align_size * align_size @@ -2129,6 +2198,12 @@ def test_padding_grouped_linear_accuracy( 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, @@ -2200,6 +2275,12 @@ def test_padding_grouped_linear_accuracy_save_original_input( 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, @@ -2409,6 +2490,12 @@ def test_gpt_fp8_parameters(dtype, bs, model, recipe): if NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") + if 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__}" + ) + config = model_configs[model] outputs = _test_gpt_fp8_parameters(bs, dtype, config, False, recipe) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index bddd9bf194..97b130952d 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -183,21 +183,38 @@ struct Tensor { * https://gcc.gnu.org/bugzilla/show_bug.cgi?id=109569). */ switch (scaling_mode) { - case NVTE_NVFP4_1D_SCALING: case NVTE_DELAYED_TENSOR_SCALING: - if (!has_data() && has_columnwise_data()) { + case NVTE_NVFP4_1D_SCALING: { + // Choose data buffer based on whether it is initialized + // Note: Uninitialized buffers currently have shape=[]. + // However, this is logically incorrect. 0-D tensors have 1 + // entry, and uninitialized tensors should have shape=[0]. + bool use_columnwise_shape = false; + if (data.dptr != nullptr) { + use_columnwise_shape = false; + } else if (columnwise_data.dptr != nullptr) { + use_columnwise_shape = true; + } else if (data.shape.size() != 0) { + use_columnwise_shape = false; + } else if (columnwise_data.shape.size() != 0) { + use_columnwise_shape = true; + } + + // Infer shape based on data + if (use_columnwise_shape) { + // Column-wise data is transposed std::vector 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]); } ret.push_back(columnwise_data.shape.front()); } return ret; - } else { - return data.shape; } - break; + return data.shape; + } case NVTE_MXFP8_1D_SCALING: if (!has_data() && has_columnwise_data()) { return columnwise_data.shape; diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 36e06173d0..06735e3104 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -332,11 +332,9 @@ __global__ void multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_ } // namespace void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t stream) { - NVTE_CHECK(input->scaling_mode == NVTE_MXFP8_1D_SCALING || - input->scaling_mode == NVTE_BLOCK_SCALING_1D || - input->scaling_mode == NVTE_BLOCK_SCALING_2D || - input->scaling_mode == NVTE_NVFP4_1D_SCALING, - "Input tensor has invalid scaling mode (", to_string(input->scaling_mode), ")."); + NVTE_CHECK( + input->scaling_mode == NVTE_MXFP8_1D_SCALING || input->scaling_mode == NVTE_NVFP4_1D_SCALING, + "Input tensor has invalid scaling mode (", to_string(input->scaling_mode), ")."); NVTE_CHECK(is_fp8_dtype(input->dtype()) || is_fp4_dtype(input->dtype()), "Input tensor has invalid dtype (", to_string(input->dtype()), ")."); @@ -583,16 +581,19 @@ void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, NVTE_CHECK_CUDA(cudaGetLastError()); } -// TODO(nvfp4): Add NVFP4 support. void multi_tensor_swizzle_scaling_factors(const std::vector& input, std::vector& output, cudaStream_t stream) { auto num_tensors = input.size(); bool all_has_data = true; bool all_has_columnwise_data = true; + bool all_nvfp4 = true; for (size_t i = 0; i < num_tensors; i++) { - if (!is_fp8_dtype(input[i]->dtype()) || !is_mxfp_scaling(input[i]->scaling_mode)) { - NVTE_ERROR("Not implemented caling mode " + to_string(input[i]->scaling_mode) + "."); - } + auto scaling_mode = input[i]->scaling_mode; + auto is_fp8 = is_fp8_dtype(input[i]->dtype()); + auto is_fp4 = is_fp4_dtype(input[i]->dtype()); + NVTE_CHECK( + (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)), + "Not implemented scaling mode " + to_string(scaling_mode) + "."); // We don't allow empty tensors. They should be filtered out before calling this function. if (input[i]->data.numel() == 0) { NVTE_ERROR("Tensor input[" + std::to_string(i) + "] is empty."); @@ -601,13 +602,17 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, CheckInputTensor(*output[i], "scaling_factor_output[" + std::to_string(i) + "]"); all_has_data &= input[i]->has_data(); all_has_columnwise_data &= input[i]->has_columnwise_data(); + all_nvfp4 &= is_nvfp4_scaling(scaling_mode); } NVTE_CHECK(all_has_data || all_has_columnwise_data, "All tensors should have data or columnwise data."); + const bool rowwise_swizzle = all_has_data || all_nvfp4; + const bool columnwise_swizzle = all_has_columnwise_data && !all_nvfp4; + constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; - if (all_has_data) { + if (rowwise_swizzle) { MultiSwizzleArgs kernel_args; kernel_args.num_tensors = 0; kernel_args.block_range[0] = 0; @@ -623,29 +628,60 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args.num_tensors = 0; vec_load_size = 4; } - const int m = input[i]->scale_inv.shape[0]; - const int k = input[i]->scale_inv.shape[1]; + + int m, k; + + if (all_has_data) { + m = input[i]->scale_inv.shape[0]; + k = input[i]->scale_inv.shape[1]; + } else { + NVTE_CHECK(all_nvfp4, "When doing rowwise swizzle with rowwise data, it has to be NVFP4"); + m = input[i]->columnwise_scale_inv.shape[0]; + k = input[i]->columnwise_scale_inv.shape[1]; + } NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Input should be padded in M/N dimension!"); NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Input should be padded in K dimension!"); NVTE_CHECK(k > 0, "Input scale inverse should be 2D!"); - NVTE_CHECK( - m * k == std::accumulate(output[i]->scale_inv.shape.begin(), - output[i]->scale_inv.shape.end(), 1, std::multiplies()), - "Input.scale_inv size is not equal to Output.scale_inv size!"); + + if (output[i]->has_data()) { + NVTE_CHECK( + m * k == std::accumulate(output[i]->scale_inv.shape.begin(), + output[i]->scale_inv.shape.end(), 1, std::multiplies()), + "Input.scale_inv size is not equal to Output.scale_inv size!"); + } + if (output[i]->has_columnwise_data()) { + NVTE_CHECK(m * k == std::accumulate(output[i]->columnwise_scale_inv.shape.begin(), + output[i]->columnwise_scale_inv.shape.end(), 1, + std::multiplies()), + "Input.columnwise_scale_inv size is not equal to " + "Output.columnwise_scale_inv size!"); + } int num_tiles_k = k / SF_TILE_DIM_K; 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); + // TODO(zhongbo): fix vec_load_size for NVFP4 + // Current unit test won't capture this issue, but in E2E + // using vec_load_size = 1 other than 1 will lead to mis-aligned + // address error in MOE training + vec_load_size = all_nvfp4 ? 1 : std::min(vec_load_size, vec_load_size_i); const int pos = kernel_args.num_tensors; - kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr); - kernel_args.output_list[pos] = output[i]->scale_inv.dptr; kernel_args.m_list[pos] = m; kernel_args.k_list[pos] = k; - kernel_args.original_m_list[pos] = input[i]->flat_first_dim(); - kernel_args.original_k_list[pos] = input[i]->flat_last_dim() / MXFP8_BLOCK_SIZE; + 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; + } 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.num_tensors++; } // Launch the remaining tensors @@ -655,7 +691,10 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args, vec_load_size, true, stream); } - if (all_has_columnwise_data) { + if (columnwise_swizzle) { + // NVFP4 shouldn't end up here because it only needs rowwise swizzle + NVTE_CHECK(!all_nvfp4, "NVFP4 shouldn't end up here because it only needs rowwise swizzle"); + MultiSwizzleArgs kernel_args; kernel_args.num_tensors = 0; kernel_args.block_range[0] = 0; diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index 49ae963d74..e054424dd4 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -190,8 +190,9 @@ transformer_engine::TensorWrapper makeTransformerEngineTensor( const std::vector meta_shape{1}; ret.set_amax(amax_ptr, DType::kFloat32, meta_shape); ret.set_scale(scale_ptr, DType::kFloat32, meta_shape); - auto scale_inv_dtype = - (scaling_mode == NVTE_MXFP8_1D_SCALING) ? DType::kFloat8E8M0 : DType::kFloat32; + auto scale_inv_dtype = (scaling_mode == NVTE_MXFP8_1D_SCALING) ? DType::kFloat8E8M0 + : (scaling_mode == NVTE_NVFP4_1D_SCALING) ? DType::kFloat8E4M3 + : DType::kFloat32; ret.set_rowwise_scale_inv(scale_inv_ptr, scale_inv_dtype, scale_inv_shape); ret.set_columnwise_scale_inv(columnwise_scale_inv_ptr, scale_inv_dtype, columnwise_scale_inv_shape); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index b6e9ef828c..7d15e436ea 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -491,6 +491,207 @@ std::tuple, std::vector> bulk_allocate_mx return retval; } +// allocate fp4 data, fp8 scalings, and amax values +// layout: [fp4_data0, ..., fp4_dataN, fp8_scaling0, ..., fp8_scalingN, amax0, ..., amaxN] +// amax buffer will be zeroed out by later amax kernels, so we can use empty to allocate +std::tuple, std::vector> bulk_allocate_nvfp4_tensors( + std::vector> &shape_list, std::vector &quantizer_py_list, + std::vector &quantizer_cpp_list) { + init_extension(); + std::tuple, std::vector> retval; + auto &tensor_py_list = std::get<0>(retval); + auto &tensor_cpp_list = std::get<1>(retval); + + // Number of tensors + const size_t num_tensors = shape_list.size(); + if (num_tensors == 0) { + return retval; + } + + // Quantization parameters + const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage; + const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; + const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); + const auto fp4_dtype = quantizer_cpp_list[0]->dtype; + 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; + }; + + // 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; + for (size_t i = 0; i < num_tensors; ++i) { + buffer_size = roundup(buffer_size, 256); // align to 256B + data_offsets.push_back(buffer_size); + // Store ceil(product/2) bytes for fp4 (since each element is 4 bits = 0.5 bytes). + // Integer arithmetic: ceil(product / 2) == (product + 1) / 2. + buffer_size += (product(rowwise_data_shapes[i]) + 1) / 2; + } + 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; + } + for (size_t i = 0; i < num_tensors; ++i) { + buffer_size = roundup(buffer_size, 16); // align to 16B + amax_offsets.push_back(buffer_size); + // amax is scalar in fp32, 4 bytes each + buffer_size += 4; + } + + // Allocate full buffer + auto buffer = std::make_shared( + at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); + + // Construct tensor views + 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)); + amax_rowwise_list.emplace_back( + make_torch_view(buffer, std::vector{1}, amax_offsets[i], torch::kUInt8)); + } + } + + // Allocate column-wise data + 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 + columnwise_data_shapes.emplace_back(); + auto &shape = columnwise_data_shapes.back(); + shape.push_back(shape_list[i].back()); + for (size_t j = 0; j < shape_list[i].size() - 1; ++j) { + shape.push_back(shape_list[i][j]); + } + columnwise_scale_shapes.emplace_back( + 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; + for (size_t i = 0; i < num_tensors; ++i) { + buffer_size = roundup(buffer_size, 256); // align to 256B + data_offsets.push_back(buffer_size); + // Store ceil(product/2) bytes for fp4 (since each element is 4 bits = 0.5 bytes). + // Integer arithmetic: ceil(product / 2) == (product + 1) / 2. + buffer_size += (product(columnwise_data_shapes[i]) + 1) / 2; + } + 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; + } + for (size_t i = 0; i < num_tensors; ++i) { + buffer_size = roundup(buffer_size, 16); // align to 16B + amax_offsets.push_back(buffer_size); + // amax is scalar in fp32, 4 bytes each + buffer_size += 4; + } + + // Allocate full buffer + auto buffer = std::make_shared( + at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); + + // Construct tensor views + 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::kUInt8)); + } + } + + // Construct nvfp4 tensors + py::handle NVFP4TensorClass(reinterpret_cast(NVFP4TensorStoragePythonClass)); + for (size_t i = 0; i < num_tensors; ++i) { + // Create tensor objects with proper reference counting + py::object rowwise_data = rowwise_usage ? py::cast(rowwise_data_list[i]) : py::none(); + py::object rowwise_scale = rowwise_usage ? py::cast(rowwise_scale_list[i]) : py::none(); + py::object columnwise_data = + (columnwise_usage ? py::cast(columnwise_data_list[i]) : py::none()); + py::object columnwise_scale = + (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); + py::object amax_rowwise = rowwise_usage ? py::cast(amax_rowwise_list[i]) : py::none(); + 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])); + + // Construct C++ tensor + // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, + // then set the amax and amax_columnwise values. + { + auto tensor_wrapper = makeTransformerEngineTensor( + rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr, + columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr, + rowwise_usage ? rowwise_data_shapes[i] : std::vector{}, + columnwise_usage ? columnwise_data_shapes[i] : std::vector{}, fp4_dtype, + /*amax_ptr=*/nullptr, + /*scale_ptr=*/nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr, + columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, + rowwise_usage ? rowwise_scale_shapes[i] : std::vector{}, + columnwise_usage ? columnwise_scale_shapes[i] : std::vector{}, scaling_mode); + + // 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}); + } + if (columnwise_usage) { + tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, + std::vector{1}); + } + tensor_cpp_list.emplace_back(std::move(tensor_wrapper)); + } + } + + return retval; +} + } // namespace std::vector split_quantize(const at::Tensor &tensor, @@ -549,7 +750,8 @@ std::vector split_quantize(const at::Tensor &tensor, bool use_fused_bulk_alloc = true; for (size_t i = 0; i < quantizer_list.size(); i++) { if (!detail::IsFloat8BlockwiseQuantizers(quantizer_list[i].ptr()) && - !detail::IsMXFP8Quantizers(quantizer_list[i].ptr())) { + !detail::IsMXFP8Quantizers(quantizer_list[i].ptr()) && + !detail::IsNVFP4Quantizers(quantizer_list[i].ptr())) { use_fused_bulk_alloc = false; break; } @@ -570,6 +772,7 @@ std::vector split_quantize(const at::Tensor &tensor, // TODO(zhongbo): make a better api to make this part less hacky bool is_fp8_blockwise = detail::IsFloat8BlockwiseQuantizers(quantizer_list[0].ptr()); bool is_mxfp8 = detail::IsMXFP8Quantizers(quantizer_list[0].ptr()); + bool is_nvfp4 = detail::IsNVFP4Quantizers(quantizer_list[0].ptr()); if (is_fp8_blockwise) { // FP8 block-scaling: construct output tensors with bulk allocations std::vector blockwise_quantizers; @@ -586,6 +789,14 @@ std::vector split_quantize(const at::Tensor &tensor, } std::tie(output_py_list, output_cpp_list) = bulk_allocate_mxfp8_tensors(split_shapes, quantizer_list, mxfp8_quantizers); + } else if (is_nvfp4) { + // NVFP4: construct output tensors with bulk allocations + std::vector nvfp4_quantizers; + for (auto &quantizer : quantizer_cpp_list) { + nvfp4_quantizers.push_back(static_cast(quantizer.get())); + } + std::tie(output_py_list, output_cpp_list) = + bulk_allocate_nvfp4_tensors(split_shapes, quantizer_list, nvfp4_quantizers); } else { NVTE_CHECK(false, "Expected either FP8 block-scaling or MXFP8 quantizer"); } diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp index 3635d4a9c0..8d1d865604 100644 --- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp +++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp @@ -20,10 +20,11 @@ void compute_amax(const at::Tensor& tensor, at::Tensor& amax) { TORCH_CHECK(amax.scalar_type() == at::kFloat, "amax must be a float tensor"); TORCH_CHECK(amax.numel() == 1, "amax must have exactly one element"); + auto* amax_ptr = amax.data_ptr(); TensorWrapper fake_te_output( nullptr, te_input.shape(), DType::kFloat8E4M3, // It doesn't matter because we only compute amax. - amax.data_ptr()); + amax_ptr); nvte_compute_amax(te_input.data(), fake_te_output.data(), at::cuda::getCurrentCUDAStream()); } diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 42ae658f2a..d7e8912ac7 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1200,6 +1200,8 @@ 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); + // 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); } if (columnwise_usage) { @@ -1213,6 +1215,8 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve columnwise_data_tensor = at::empty(convert_shape_for_fp4(transpose_shape_int64), bit8_tensor_opts); columnwise_scale_inv_tensor = at::empty(scale_inv_shape_int64, bit8_tensor_opts); + // hadamard amax kernel will zero out pointer with ZeroAmaxKernel + // nvte_compute_amax_with_config will zero out the pointer if needed amax_columnwise = at::empty({1}, bit32_tensor_opts); } @@ -1352,6 +1356,8 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } if (!amax_rowwise) { 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); tensor.attr("_amax_rowwise") = *amax_rowwise; } @@ -1392,7 +1398,9 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } if (!amax_columnwise) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - amax_columnwise = at::zeros({1}, opts); + // hadamard amax kernel will zero out pointer with ZeroAmaxKernel + // nvte_compute_amax_with_config will zero out the pointer if needed + amax_columnwise = at::empty({1}, opts); tensor.attr("_amax_columnwise") = *amax_columnwise; } } else { // columnwise_usage == false diff --git a/transformer_engine/pytorch/csrc/util.cpp b/transformer_engine/pytorch/csrc/util.cpp index ffba5b2763..134185ac82 100644 --- a/transformer_engine/pytorch/csrc/util.cpp +++ b/transformer_engine/pytorch/csrc/util.cpp @@ -50,8 +50,6 @@ std::optional swizzle_scaling_factors(transformer_engine::TensorWrap void* scale_inv_dptr = scale_inv.data_ptr; void* swizzled_scale_inv_dptr = getDataPtr(swizzled_scale_inv, 0); - // Reconstruct input only to avoid swizzling both directions if not needed. - // The specific dtype used is irrelevant, just needs to be correct bits. transformer_engine::TensorWrapper input_cu(input.scaling_mode()); transformer_engine::TensorWrapper output_cu(input.scaling_mode()); @@ -100,10 +98,14 @@ std::optional multi_tensor_swizzle_scaling_factors( if (tensors.front().scaling_mode() == NVTE_INVALID_SCALING) { NVTE_ERROR("Invalid scaling mode for swizzle."); - } else if (tensors.front().scaling_mode() != NVTE_MXFP8_1D_SCALING) { + } else if (tensors.front().scaling_mode() != NVTE_MXFP8_1D_SCALING && + tensors.front().scaling_mode() != NVTE_NVFP4_1D_SCALING) { return std::nullopt; } + const auto scaling_mode = tensors.front().scaling_mode(); + const auto nvfp4 = scaling_mode == NVTE_NVFP4_1D_SCALING; + std::vector wrappers; std::vector input_tensors, output_tensors; @@ -131,39 +133,44 @@ std::optional multi_tensor_swizzle_scaling_factors( // Allocate full buffer auto buffer = at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8)); + const auto input_dtype = + (nvfp4) ? transformer_engine::DType::kFloat4E2M1 : transformer_engine::DType::kFloat8E4M3; + const auto scale_inv_dtype = + (nvfp4) ? transformer_engine::DType::kFloat8E4M3 : transformer_engine::DType::kFloat8E8M0; + for (size_t i = 0; i < tensors.size(); ++i) { auto& tensor = tensors[i]; void* scale_inv_dptr = scale_inv_dptrs[i]; void* swizzled_scale_inv_dptr = getDataPtr(buffer, scale_inv_offsets[i]); - auto input_shape = nvte_shape_to_vector(tensor.shape()); - + // auto input_shape = nvte_shape_to_vector(tensor.shape()); + NVTEShape nvte_input_shape; + if (rowwise) { + nvte_input_shape = tensor.shape(); + } else { + nvte_input_shape = tensor.get_columnwise_data().shape; + } + auto input_shape = nvte_shape_to_vector(nvte_input_shape); // Reconstruct input only to avoid swizzling both directions if not needed. // Use any 8 bit type, it's irrelevant. - transformer_engine::TensorWrapper input_cu(NVTE_MXFP8_1D_SCALING); - transformer_engine::TensorWrapper output_cu(NVTE_MXFP8_1D_SCALING); + transformer_engine::TensorWrapper input_cu(scaling_mode); + transformer_engine::TensorWrapper output_cu(scaling_mode); if (rowwise) { - input_cu.set_rowwise_data(tensor.dptr(), transformer_engine::DType::kFloat8E4M3, input_shape); - input_cu.set_rowwise_scale_inv(scale_inv_dptr, transformer_engine::DType::kFloat8E8M0, - scale_inv_shapes[i]); - output_cu.set_rowwise_data(tensor.dptr(), transformer_engine::DType::kFloat8E4M3, - input_shape); - output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, - transformer_engine::DType::kFloat8E8M0, scale_inv_shapes[i]); + input_cu.set_rowwise_data(tensor.dptr(), input_dtype, input_shape); + input_cu.set_rowwise_scale_inv(scale_inv_dptr, scale_inv_dtype, scale_inv_shapes[i]); + output_cu.set_rowwise_data(tensor.dptr(), input_dtype, input_shape); + output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, + scale_inv_shapes[i]); // Set the swizzled scaling factor to the original tensor. - tensor.set_rowwise_scale_inv(swizzled_scale_inv_dptr, transformer_engine::DType::kFloat8E8M0, - scale_inv_shapes[i]); + tensor.set_rowwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, scale_inv_shapes[i]); } else { - input_cu.set_columnwise_data(tensor.columnwise_dptr(), transformer_engine::DType::kFloat8E4M3, - input_shape); - input_cu.set_columnwise_scale_inv(scale_inv_dptr, transformer_engine::DType::kFloat8E8M0, - scale_inv_shapes[i]); - output_cu.set_columnwise_data(tensor.columnwise_dptr(), - transformer_engine::DType::kFloat8E4M3, input_shape); - output_cu.set_columnwise_scale_inv( - swizzled_scale_inv_dptr, transformer_engine::DType::kFloat8E8M0, scale_inv_shapes[i]); + input_cu.set_columnwise_data(tensor.columnwise_dptr(), input_dtype, input_shape); + input_cu.set_columnwise_scale_inv(scale_inv_dptr, scale_inv_dtype, scale_inv_shapes[i]); + output_cu.set_columnwise_data(tensor.columnwise_dptr(), input_dtype, input_shape); + output_cu.set_columnwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, + scale_inv_shapes[i]); // Set the swizzled scaling factor to the original tensor. - tensor.set_columnwise_scale_inv(swizzled_scale_inv_dptr, - transformer_engine::DType::kFloat8E8M0, scale_inv_shapes[i]); + tensor.set_columnwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, + scale_inv_shapes[i]); } input_tensors.emplace_back(input_cu.data()); diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index 5d569d59d4..fca89fbaa9 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -78,7 +78,7 @@ class Fp8Padding(torch.nn.Module): number of GEMMs to be performed simultaneously. align_size : int, optional the alignment size for the input tensor. If not provided, the alignment size will - be determined by the FP8 recipe (32 for MXFP8 and 16 for others) in the first + be determined by the FP8/FP4 recipe (32 for MXFP8/NVFP4 and 16 for others) in the first forward pass. """ @@ -111,7 +111,14 @@ def forward( assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." if self.align_size is None: - self.align_size = 32 if FP8GlobalStateManager.get_fp8_recipe().mxfp8() else 16 + self.align_size = ( + 32 + if ( + FP8GlobalStateManager.get_fp8_recipe().mxfp8() + or FP8GlobalStateManager.get_fp8_recipe().nvfp4() + ) + else 16 + ) # FP8 padding calculate padded_m_splits = [ diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index b74395dd8c..7a01f15729 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -75,9 +75,9 @@ class Fp8Unpadding(torch.nn.Module): num_gemms : int number of GEMMs to be performed simultaneously. align_size : int, optional - the alignment size for the input tensor. If not provided, the alignment size will - be determined by the FP8 recipe (32 for MXFP8 and 16 for others) in the first - forward pass. + The alignment size for the input tensor. If not provided, the alignment size will + be automatically determined based on the FP8/FP4 recipe in the first forward pass: + 32 for MXFP8 or NVFP4, otherwise 16. """ def __init__( @@ -109,7 +109,14 @@ def forward( assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." if self.align_size is None: - self.align_size = 32 if FP8GlobalStateManager.get_fp8_recipe().mxfp8() else 16 + self.align_size = ( + 32 + if ( + FP8GlobalStateManager.get_fp8_recipe().mxfp8() + or FP8GlobalStateManager.get_fp8_recipe().nvfp4() + ) + else 16 + ) # FP8 padding calculate padded_m_splits = [ From e90582f2010deae477a71bad0aacf278dd5abfa4 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Tue, 21 Oct 2025 19:52:37 +0200 Subject: [PATCH 012/521] [Common] Removed activations from NVFP4 quantize C++ unit tests (#2289) * Removed activations from NVFP4 CPP tests. Removed CMake debugging flags Signed-off-by: Oleg Goncharov * Better wording Signed-off-by: Oleg Goncharov --------- Signed-off-by: Oleg Goncharov --- tests/cpp/operator/CMakeLists.txt | 6 ------ tests/cpp/operator/test_cast_nvfp4_transpose.cu | 9 ++------- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 479d378ba6..b2f14b1892 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -32,12 +32,6 @@ add_executable(test_operator test_swap_first_dims.cu ../test_common.cu) -# Add profiling and debug flags for CUDA compilation -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -lineinfo") # Generate line info for device code -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -g") # Add debug symbols for host code -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --ptxas-options=-v") # Add info about registers usage -# Note: Using -lineinfo instead of -G to avoid conflicts and get line mapping - # Find required packages find_package(OpenMP REQUIRED) list(APPEND test_operator_LINKER_LIBS CUDA::cudart GTest::gtest_main ${TE_LIB} CUDA::nvrtc CUDNN::cudnn) diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index e905a00640..afd7927da2 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -661,14 +661,9 @@ std::vector> tensor_dims = { {4096, 13312}, }; -// Only GeLU activation tests are supported +// Only the Identity activation is currently supported. std::vector Activation_types = { - ActivationType::Identity, - ActivationType::GeLU, - ActivationType::SiLU, - ActivationType::ReLU, - ActivationType::QGeLU, - ActivationType::SReLU, + ActivationType::Identity }; } // namespace From ce2f9fa4632a688d45efc55586be18cd0931ea50 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 21 Oct 2025 12:29:09 -0700 Subject: [PATCH 013/521] [JAX] HuggingFace login in JAX examples if token is available (#2290) HF login in JAX examples Signed-off-by: Jeremy Berchtold --- examples/jax/encoder/common.py | 11 +++++++++++ examples/jax/encoder/test_model_parallel_encoder.py | 2 ++ examples/jax/encoder/test_multigpu_encoder.py | 7 ++++++- examples/jax/encoder/test_multiprocessing_encoder.py | 2 ++ examples/jax/encoder/test_single_gpu_encoder.py | 7 ++++++- examples/jax/mnist/test_single_gpu_mnist.py | 8 +++++++- 6 files changed, 34 insertions(+), 3 deletions(-) diff --git a/examples/jax/encoder/common.py b/examples/jax/encoder/common.py index 772d5f4c14..9ffcfe57da 100644 --- a/examples/jax/encoder/common.py +++ b/examples/jax/encoder/common.py @@ -118,3 +118,14 @@ def get_quantization_recipe_from_name_string(name: str): return recipe.NVFP4BlockScaling() case _: raise ValueError(f"Invalid quantization_recipe, got {name}") + + +def hf_login_if_available(): + """Login to HF hub if available""" + try: + from huggingface_hub import login + + login() + except Exception as e: + print(e) + pass diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index 7807d1fd96..c6d867ef98 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -23,12 +23,14 @@ is_bf16_supported, get_quantization_recipe_from_name_string, assert_params_sufficiently_sharded, + hf_login_if_available, ) import transformer_engine.jax as te import transformer_engine.jax.cpp_extensions as tex import transformer_engine.jax.flax as te_flax from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode +hf_login_if_available() DEVICE_DP_AXIS = "data" DEVICE_TP_AXIS = "model" diff --git a/examples/jax/encoder/test_multigpu_encoder.py b/examples/jax/encoder/test_multigpu_encoder.py index 8ea1dcde37..1004dd2dd2 100644 --- a/examples/jax/encoder/test_multigpu_encoder.py +++ b/examples/jax/encoder/test_multigpu_encoder.py @@ -19,12 +19,17 @@ from jax.experimental import mesh_utils from jax.sharding import PartitionSpec, NamedSharding -from common import is_bf16_supported, get_quantization_recipe_from_name_string +from common import ( + is_bf16_supported, + get_quantization_recipe_from_name_string, + hf_login_if_available, +) import transformer_engine.jax as te import transformer_engine.jax.cpp_extensions as tex import transformer_engine.jax.flax as te_flax from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode +hf_login_if_available() DEVICE_DP_AXIS = "data" PARAMS_KEY = "params" diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index 7e708466c2..c2e97029b0 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -27,11 +27,13 @@ is_mxfp8_supported, is_nvfp4_supported, get_quantization_recipe_from_name_string, + hf_login_if_available, ) import transformer_engine.jax as te import transformer_engine.jax.cpp_extensions as tex import transformer_engine.jax.flax as te_flax +hf_login_if_available() os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" DEVICE_DP_AXIS = "data" diff --git a/examples/jax/encoder/test_single_gpu_encoder.py b/examples/jax/encoder/test_single_gpu_encoder.py index 79178485c2..1c62de7fa4 100644 --- a/examples/jax/encoder/test_single_gpu_encoder.py +++ b/examples/jax/encoder/test_single_gpu_encoder.py @@ -16,11 +16,16 @@ from flax import linen as nn from flax.training import train_state -from common import is_bf16_supported, get_quantization_recipe_from_name_string +from common import ( + is_bf16_supported, + get_quantization_recipe_from_name_string, + hf_login_if_available, +) import transformer_engine.jax as te import transformer_engine.jax.flax as te_flax from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode +hf_login_if_available() PARAMS_KEY = "params" DROPOUT_KEY = "dropout" diff --git a/examples/jax/mnist/test_single_gpu_mnist.py b/examples/jax/mnist/test_single_gpu_mnist.py index d0aebeb53d..2e9d56e93f 100644 --- a/examples/jax/mnist/test_single_gpu_mnist.py +++ b/examples/jax/mnist/test_single_gpu_mnist.py @@ -22,7 +22,13 @@ DIR = str(Path(__file__).resolve().parents[1]) sys.path.append(str(DIR)) -from encoder.common import is_bf16_supported, get_quantization_recipe_from_name_string +from encoder.common import ( + is_bf16_supported, + get_quantization_recipe_from_name_string, + hf_login_if_available, +) + +hf_login_if_available() IMAGE_H = 28 IMAGE_W = 28 From 2712bb95cb4a4f7d1f2b8b473a2240ac3d6e7e58 Mon Sep 17 00:00:00 2001 From: Kunlun Li <94586211+kunlunl@users.noreply.github.com> Date: Wed, 22 Oct 2025 07:04:04 +0800 Subject: [PATCH 014/521] Add post-processing API for FP8 primary weights to support CUDA Graph (#2266) * Add post-processing API for FP8 primary weights to support CUDA Graph Signed-off-by: kunlunl * Add post-processing support for plain pytorch tensors Signed-off-by: kunlunl * Update type hint Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: kunlunl Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- .../run_cast_master_weights_to_fp8.py | 46 +++++++++++-------- transformer_engine/pytorch/tensor/utils.py | 43 +++++++++-------- 2 files changed, 50 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py index 9769916335..2f11a24ee8 100644 --- a/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py @@ -27,7 +27,7 @@ Float8BlockwiseQTensor, ) from transformer_engine.pytorch.tensor import cast_master_weights_to_fp8 -from transformer_engine.pytorch.tensor.utils import replace_raw_data +from transformer_engine.pytorch.tensor.utils import post_all_gather_processing, replace_raw_data def _get_raw_data(quantized_tensor): @@ -203,12 +203,15 @@ def step(self): # ----------------------------------------------------------------------------------------- # Step 7: Copy the gathered weights from weight buffer to the actual weights # ----------------------------------------------------------------------------------------- + quantized_weights = [] for weight, offset in zip(self.weights, self.offsets[:-1]): start = offset end = offset + weight.numel() if isinstance(weight, QuantizedTensor): + quantized_weights.append(weight) weight = _get_raw_data(weight) weight.view(-1).data.copy_(self.weight_buffer[start:end]) + post_all_gather_processing(quantized_weights) class MiniOptimizer: @@ -252,10 +255,6 @@ def __init__(self, weights, lr, dp_group): self.dp_group = dp_group # Flatten the weights and pad to align with world size - raw_data_list = [ - _get_raw_data(w).view(-1) if isinstance(w, QuantizedTensor) else w.view(-1) - for w in weights - ] if isinstance(weights[0], QuantizedTensor): raw_data_list = [_get_raw_data(w).view(-1) for w in weights] else: @@ -264,7 +263,9 @@ def __init__(self, weights, lr, dp_group): # Split flattened weights into shards self.local_weight_shard = torch.chunk(self.flatten_weight, world_size)[rank] - self.local_main_grad_shard = torch.zeros_like(self.local_weight_shard) + self.local_main_grad_shard = torch.zeros_like( + self.local_weight_shard, dtype=torch.float32, device="cuda" + ) shard_size = self.flatten_weight.size(0) // world_size # Map original tensors to flattened indices @@ -341,9 +342,8 @@ def _flatten_tensors_with_pad(self, tensors): padding_needed = (world_size - original_length % world_size) % world_size if padding_needed > 0: - flatten_tensor = torch.cat( - [flatten_tensor, torch.zeros(padding_needed, dtype=flatten_tensor.dtype)] - ) + zeros = torch.zeros(padding_needed, dtype=flatten_tensor.dtype, device="cuda") + flatten_tensor = torch.cat([flatten_tensor, zeros]) return flatten_tensor, original_length @@ -369,10 +369,10 @@ def step(self): main_grad_buffer, _ = self._flatten_tensors_with_pad( [weight.main_grad.view(-1) for weight in self.weights] ) - main_grad_buffer = main_grad_buffer.to(self.local_main_grad_shard.dtype) dist.reduce_scatter_tensor( self.local_main_grad_shard, main_grad_buffer, group=self.dp_group ) + self.local_main_grad_shard /= dist.get_world_size(self.dp_group) # Step 2: Update the master weights for weight, master_weight, (shard_start, shard_end) in zip( @@ -416,6 +416,11 @@ def step(self): dist.all_gather_into_tensor( self.flatten_weight, self.local_weight_shard, group=self.dp_group ) + quantized_weights = [] + for weight in self.weights: + if isinstance(weight, QuantizedTensor): + quantized_weights.append(weight) + post_all_gather_processing(quantized_weights) def _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group): @@ -435,7 +440,7 @@ def _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group): linear_kwargs = { "params_dtype": torch.bfloat16, "bias": False, - "fuse_wgrad_accumulation": False, + "fuse_wgrad_accumulation": True, } # Create model with FP8 weights @@ -503,14 +508,9 @@ def _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group): torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) - print( - f"✅ Successfully validated FSDP {NUM_STEPS} training steps with" - f" {quantization} quantization" - ) - -def _test_zero_1(dp_group): - """Make sure the implementation of zero-1 optimizer is correct""" +def _test_mini_optimizer(dp_group): + """Make sure the implementation of MiniZero_1 and MiniFSDP is correct""" rank = dist.get_rank(dp_group) world_size = dist.get_world_size(dp_group) @@ -525,13 +525,15 @@ def _test_zero_1(dp_group): weights_1 = weights weights_2 = [weight.clone() for weight in weights] + weights_3 = [weight.clone() for weight in weights] lr = 1.0 optimizer_1 = MiniZero_1(weights_1, lr, dp_group) optimizer_2 = MiniOptimizer(weights_2, lr, dp_group) + optimizer_3 = MiniFSDP(weights_3, lr, dp_group) for _ in range(100): - for w1, w2 in zip(weights_1, weights_2): + for w1, w2, w3 in zip(weights_1, weights_2, weights_3): main_grads = [ torch.randn_like(w1, dtype=torch.float32, device="cuda") for _ in range(world_size) ] @@ -539,12 +541,16 @@ def _test_zero_1(dp_group): main_grad = main_grads[rank] w1.main_grad = main_grad w2.main_grad = main_grad + w3.main_grad = main_grad optimizer_1.step() optimizer_2.step() + optimizer_3.step() for w1, w2 in zip(weights_1, weights_2): torch.testing.assert_close(w1, w2, atol=0, rtol=0) + for w1, w3 in zip(weights_1, weights_3): + torch.testing.assert_close(w1, w3, atol=0, rtol=0) def quantization_recipe(quantization) -> Recipe: @@ -671,7 +677,7 @@ def main(argv=None, namespace=None): args = parser.parse_args(argv, namespace) dp_group = dist.new_group(backend="nccl") - _test_zero_1(dp_group) + _test_mini_optimizer(dp_group) _test_cast_master_weights_to_fp8(args.quantization, dp_group) _test_fsdp_cast_master_weights_to_fp8(args.quantization, dp_group) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index cc02494013..72c465edb2 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -5,7 +5,7 @@ """Helper functions for using fp8 tensors as weights""" import os -from typing import Optional, Union +from typing import Optional, List, Union import torch import transformer_engine_torch as tex from transformer_engine_torch import multi_tensor_scale, multi_tensor_compute_scale_and_scale_inv @@ -15,6 +15,7 @@ from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from ..optimizers.multi_tensor_apply import multi_tensor_applier +from ..utils import is_non_tn_fp8_gemm_supported def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): @@ -159,12 +160,6 @@ def _cast_master_weights_to_fp8_delayed_scaling(params, group, use_fsdp_shard_mo amaxes, scales, scale_invs = [], [], [] for model_weight, master_weight, start_offset, shard_model_weight_raw in params: - # Reset transpose cache for all model weights. - # We cannot create transpose cache here because users (like megatron) may want to overlap - # the all-gather of model weights and forward process, so the model weight is not updated - # currently. - model_weight._reset_caches() - quantizer = model_weight._get_quantizer() amaxes.append(quantizer.amax.view(1)) @@ -302,12 +297,6 @@ def _cast_master_weights_to_fp8_current_scaling(params, group, use_fsdp_shard_mo for (model_weight, master_weight, start_offset, model_weight_fragment), scale in zip( params, scales ): - # Reset transpose cache for all model weights. - # We cannot create transpose cache here because users (like megatron) may want to overlap - # the all-gather of model weights and forward process, so the model weight is not updated - # currently. - model_weight._reset_caches() - # If master weight is None, it means that the master weight of the current model weight # is in other DP ranks. if master_weight is None: @@ -432,12 +421,6 @@ def _cast_master_weights_to_fp8_blockwise_scaling( for (model_weight, master_weight, start_offset, model_weight_fragment), scale in zip( params, scales ): - # Clear columnwise data for all model weights. - # We cannot create columnwise data here because users (like megatron) may want to overlap - # the all-gather of model weights and forward process, so the model weight is not updated - # at this moment. - model_weight.update_usage(rowwise_usage=True, columnwise_usage=False) - # If master weight is None, it means that the master weight of the current model weight # is in other DP ranks. if master_weight is None: @@ -454,6 +437,28 @@ def _cast_master_weights_to_fp8_blockwise_scaling( ) +def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Tensor]]): + """ + Post-processing after all-gather for weights in distributed optimizer. + - Float8Tensor: may need to create a transposed view to match backend GEMM. + - Float8BlockwiseQTensor: create column-wise storage. + - Plain pytorch tensor: noop. + """ + if not isinstance(model_weights, list): + model_weights = [model_weights] + for model_weight in model_weights: + if isinstance(model_weight, Float8Tensor): + # Delayed scaling and per-tensor current scaling: if backend does not support + # non-transposed FP8 GEMM, pre-create the transpose. + if not is_non_tn_fp8_gemm_supported(): + model_weight._create_transpose() + elif isinstance(model_weight, Float8BlockwiseQTensor): + # Blockwise scaling: create column-wise storage. + model_weight._create_columnwise() + elif isinstance(model_weight, QuantizedTensor): + raise ValueError(f"post_processing for {type(model_weight)} is not supported") + + def is_experimental(x: Optional[Union[Quantizer, QuantizedTensorStorage]] = None) -> bool: """Check if an environment or object is using experimental Kitchen middleware. From ce2e8bd12edfe10647bec8f54fedc394d6287b58 Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Wed, 22 Oct 2025 14:58:27 +0200 Subject: [PATCH 015/521] [PyTorch] Decouple python quantization classes and refactor custom quantization (#2276) * rename experimental -> custom_recipes Signed-off-by: Evgeny * Decouple python base classes (api) Signed-off-by: Evgeny * update test_custom_recipe Signed-off-by: Evgeny * Rename experimental -> custom Signed-off-by: Evgeny * Minor Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix import Signed-off-by: Evgeny * Update tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py Co-authored-by: Kirthi Shankar Sivamani Signed-off-by: Evgeny Tsykunov * Update tests/pytorch/test_custom_recipe.py Co-authored-by: Kirthi Shankar Sivamani Signed-off-by: Evgeny Tsykunov * quantization_base -> quantized_tensor rename Signed-off-by: Evgeny --------- Signed-off-by: Evgeny Signed-off-by: Evgeny Tsykunov Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/attention/test_attention.py | 3 +- .../pytorch/distributed/run_numerics_exact.py | 6 +- .../test_fusible_ops_with_userbuffers.py | 1 + .../distributed/test_numerics_exact.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 4 +- .../pytorch/nvfp4/test_nvfp4_module_exact.py | 4 +- .../nvfp4/test_nvfp4_quantize_exact.py | 4 +- .../nvfp4/test_nvfp4_rht_quantize_exact.py | 6 +- tests/pytorch/test_custom_recipe.py | 42 ++++++++++ .../debug/pytorch/debug_quantization.py | 2 +- transformer_engine/pytorch/__init__.py | 10 +-- .../dot_product_attention/backends.py | 2 +- .../dot_product_attention/context_parallel.py | 4 +- .../pytorch/cpp_extensions/fused_attn.py | 2 +- .../pytorch/cpp_extensions/gemm.py | 12 +-- transformer_engine/pytorch/cpu_offload.py | 2 +- .../__init__.py | 0 .../{experimental => custom_recipes}/gemm.py | 12 +-- .../quantization.py | 0 .../quantization_nvfp4.py | 14 ++-- .../{experimental => custom_recipes}/utils.py | 0 transformer_engine/pytorch/distributed.py | 2 +- transformer_engine/pytorch/module/base.py | 2 +- .../pytorch/module/grouped_linear.py | 2 +- .../pytorch/module/layernorm_linear.py | 12 +-- .../pytorch/module/layernorm_mlp.py | 12 +-- transformer_engine/pytorch/module/linear.py | 14 ++-- transformer_engine/pytorch/ops/_common.py | 2 +- .../ops/fused/userbuffers_backward_linear.py | 2 +- .../ops/fused/userbuffers_forward_linear.py | 2 +- transformer_engine/pytorch/ops/fuser.py | 2 +- transformer_engine/pytorch/permutation.py | 2 +- .../pytorch/{tensor => }/quantized_tensor.py | 76 ++--------------- transformer_engine/pytorch/tensor/__init__.py | 2 +- .../pytorch/tensor/_quantization_helpers.py | 84 +++++++++++++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 7 +- .../pytorch/tensor/float8_tensor.py | 7 +- .../pytorch/tensor/mxfp8_tensor.py | 7 +- .../pytorch/tensor/nvfp4_tensor.py | 3 +- .../float8_blockwise_tensor_storage.py | 4 +- .../tensor/storage/float8_tensor_storage.py | 4 +- .../tensor/storage/mxfp8_tensor_storage.py | 4 +- .../tensor/storage/nvfp4_tensor_storage.py | 3 +- transformer_engine/pytorch/tensor/utils.py | 19 ++--- transformer_engine/pytorch/utils.py | 2 +- 45 files changed, 227 insertions(+), 181 deletions(-) rename transformer_engine/pytorch/{experimental => custom_recipes}/__init__.py (100%) rename transformer_engine/pytorch/{experimental => custom_recipes}/gemm.py (90%) rename transformer_engine/pytorch/{experimental => custom_recipes}/quantization.py (100%) rename transformer_engine/pytorch/{experimental => custom_recipes}/quantization_nvfp4.py (98%) rename transformer_engine/pytorch/{experimental => custom_recipes}/utils.py (100%) rename transformer_engine/pytorch/{tensor => }/quantized_tensor.py (89%) create mode 100644 transformer_engine/pytorch/tensor/_quantization_helpers.py diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 7dc6caeb81..3150c06abb 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -45,7 +45,8 @@ ) from transformer_engine.pytorch.utils import get_cudnn_version import transformer_engine_torch as tex -from transformer_engine.pytorch.tensor.quantized_tensor import ( +from transformer_engine.pytorch.quantized_tensor import ( + Quantizer, prepare_for_saving, restore_from_saved, ) diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index ccbc3259bb..3605b3c708 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -22,8 +22,8 @@ ) from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE -from transformer_engine.pytorch.experimental import quantization_nvfp4 -from transformer_engine.pytorch.experimental import utils +from transformer_engine.pytorch.custom_recipes import quantization_nvfp4 +from transformer_engine.pytorch.custom_recipes import utils from run_layer_with_overlap import _compare_tensors @@ -486,7 +486,7 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs): sequence_parallel (bool): Enable sequence parallelism if True. kwargs (dict): Additional arguments for the linear layer. - QUANTIZATION options: nvfp4 <=> experimental nvfp4 as a reference + QUANTIZATION options: nvfp4 <=> custom nvfp4 as a reference """ params_dtype = torch.bfloat16 use_bias = kwargs.get("bias", True) diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index 24112cc9ff..61c813b8f2 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -34,6 +34,7 @@ Float8Tensor, ) + # Import utility functions _current_file = pathlib.Path(__file__).resolve() sys.path.append(str(_current_file.parent.parent)) diff --git a/tests/pytorch/distributed/test_numerics_exact.py b/tests/pytorch/distributed/test_numerics_exact.py index fd6ef65e09..72aa786646 100644 --- a/tests/pytorch/distributed/test_numerics_exact.py +++ b/tests/pytorch/distributed/test_numerics_exact.py @@ -14,7 +14,7 @@ Distributed numerics tests This numerical test aims for zero tolerance test for absolute confidence in numerics. - In the case of NVFP4, with the experimental NVFP4 quantization, we matched bitwise + In the case of NVFP4, with the custom NVFP4 quantization, we matched bitwise result with the native silicon. For distrbuted test cases, we can do the same by thing by comparing BF16 AG results with the low precision AG results at layer level. """ diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 77cfaaffe8..6009643ffa 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -8,8 +8,8 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.experimental.quantization_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.experimental import utils +from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import utils recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index 44f222b9d1..0292063ab9 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -6,8 +6,8 @@ import torch import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.experimental import quantization_nvfp4 -from transformer_engine.pytorch.experimental import utils +from transformer_engine.pytorch.custom_recipes import quantization_nvfp4 +from transformer_engine.pytorch.custom_recipes import utils recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 8c24445573..2467c7e2e1 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -7,10 +7,10 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.experimental.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_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 -from transformer_engine.pytorch.experimental import utils recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index 6f2f846a36..904dfc2eab 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -12,10 +12,10 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.common.recipe import NVFP4BlockScaling +from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.experimental.quantization_nvfp4 import NVFP4QuantizerRef -from transformer_engine.pytorch.experimental import utils +from transformer_engine.common.recipe import NVFP4BlockScaling import pytest import torch diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 516354a34b..64f1c3d159 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -17,6 +17,48 @@ Float8CurrentScalingQuantizer, ) import transformer_engine.pytorch.ops as te_ops +from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import ( + nvfp4_ref_rht_2d_quantizer_factory, +) + + +@pytest.mark.parametrize("module_type", ["Linear", "LayerNormLinear", "OpsLinear"]) +def test_custom_recipe_sanity_modules_nvfp4(module_type): + """Test modules with NVFP4 custom recipe support""" + 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) + + # Simple linear layer with dims divisible by 16 + in_features = 64 + out_features = 64 + batch = 32 + + if module_type == "Linear": + model = Linear(in_features, out_features, params_dtype=torch.bfloat16, bias=False).cuda() + elif module_type == "LayerNormLinear": + model = LayerNormLinear( + in_features, out_features, params_dtype=torch.bfloat16, bias=False + ).cuda() + else: # OpsLinear + model = te_ops.Linear( + in_features, out_features, device="cuda", dtype=torch.bfloat16, bias=False + ) + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + # Use NVFP4 quantizer factory + custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) + + # Execute with custom recipe + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp) + loss = out.float().sum() + loss.backward() + + # Basic sanity: gradients exist + assert inp.grad is not None @pytest.mark.parametrize("module_type", ["Linear", "LayerNormLinear", "OpsLinear", "LayerNormMLP"]) diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 185bf15d05..7f45a24e20 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -15,7 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch.tensor.quantized_tensor import ( +from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensor, Quantizer, QuantizedTensorStorage, diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 77c71b8119..9d894a389b 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -66,24 +66,24 @@ def torch_version() -> tuple[int, ...]: from transformer_engine.pytorch import optimizers from transformer_engine.pytorch.export import onnx_export from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy -from transformer_engine.pytorch.tensor import Quantizer +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +from transformer_engine.pytorch.quantized_tensor import Quantizer +from transformer_engine.pytorch.quantized_tensor import prepare_for_saving +from transformer_engine.pytorch.quantized_tensor import restore_from_saved from transformer_engine.pytorch.tensor import Float8Quantizer from transformer_engine.pytorch.tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor import Float8BlockQuantizer from transformer_engine.pytorch.tensor import NVFP4Quantizer -from transformer_engine.pytorch.tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor import Float8TensorStorage from transformer_engine.pytorch.tensor import MXFP8TensorStorage from transformer_engine.pytorch.tensor import Float8BlockwiseQTensorStorage from transformer_engine.pytorch.tensor import NVFP4TensorStorage -from transformer_engine.pytorch.tensor import QuantizedTensor from transformer_engine.pytorch.tensor import Float8Tensor 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 import prepare_for_saving -from transformer_engine.pytorch.tensor import restore_from_saved try: torch._dynamo.config.error_on_nested_jit_trace = False diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6dfe0d31b3..6c19d868a1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -24,7 +24,7 @@ Float8Quantizer, Float8CurrentScalingQuantizer, ) -from transformer_engine.pytorch.tensor.quantized_tensor import ( +from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensorStorage, prepare_for_saving, restore_from_saved, 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 a474cb809a..e5ee8cc7db 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -21,7 +21,7 @@ ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor -from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensorStorage +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.jit import jit_fuser from transformer_engine.pytorch.constants import ( dist_group_type, @@ -33,7 +33,7 @@ gather_along_first_dim, reduce_scatter_along_first_dim, ) -from transformer_engine.pytorch.tensor.quantized_tensor import ( +from transformer_engine.pytorch.quantized_tensor import ( prepare_for_saving, restore_from_saved, ) diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 94a12c4a09..f80c001a1d 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -15,7 +15,7 @@ NVTE_Softmax_Type, NVTE_Fused_Attn_Backend, ) -from ..tensor.quantized_tensor import Quantizer +from ..quantized_tensor import Quantizer __all__ = [ diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a45fafb68a..dd04112982 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -11,10 +11,10 @@ from ..constants import TE_DType from ..utils import get_sm_count, _empty_tensor -from ..tensor.quantized_tensor import Quantizer +from ..quantized_tensor import Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage -from ..tensor.utils import is_experimental -from ..experimental.gemm import experimental_gemm +from ..tensor.utils import is_custom +from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer __all__ = [ @@ -79,9 +79,9 @@ def general_gemm( if not out.is_contiguous(): raise ValueError("Output tensor is not contiguous.") - # If A or B are experimental tensors -> dispatch to quantizers's qgemm implementation - if is_experimental(A) or is_experimental(B): - return experimental_gemm( + # If A or B are custom tensors -> dispatch to quantizers's qgemm implementation + if is_custom(A) or is_custom(B): + return custom_gemm( A, B, workspace, diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 648b21eb4d..6edc126200 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -10,7 +10,7 @@ import torch from transformer_engine.debug.pytorch.debug_state import TEDebugState -from .tensor.quantized_tensor import QuantizedTensorStorage +from .quantized_tensor import QuantizedTensorStorage from .tensor.float8_tensor import Float8Tensor __all__ = ["get_cpu_offload_context"] diff --git a/transformer_engine/pytorch/experimental/__init__.py b/transformer_engine/pytorch/custom_recipes/__init__.py similarity index 100% rename from transformer_engine/pytorch/experimental/__init__.py rename to transformer_engine/pytorch/custom_recipes/__init__.py diff --git a/transformer_engine/pytorch/experimental/gemm.py b/transformer_engine/pytorch/custom_recipes/gemm.py similarity index 90% rename from transformer_engine/pytorch/experimental/gemm.py rename to transformer_engine/pytorch/custom_recipes/gemm.py index 0bd740d85d..cc98a8a57a 100644 --- a/transformer_engine/pytorch/experimental/gemm.py +++ b/transformer_engine/pytorch/custom_recipes/gemm.py @@ -2,21 +2,21 @@ # # See LICENSE for license information. -"""GEMM API for experimental middleware between Transformer Engine and Kitchen.""" +"""GEMM API that enables custom GEMM logic for custom quantization recipes.""" from typing import Iterable, Optional import torch -from transformer_engine.pytorch.experimental.quantization import ( +from transformer_engine.pytorch.custom_recipes.quantization import ( MMParams, GEMMType, ) -from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensorStorage, Quantizer -from transformer_engine.pytorch.tensor.utils import is_experimental +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer +from transformer_engine.pytorch.tensor.utils import is_custom -def experimental_gemm( +def custom_gemm( A: QuantizedTensorStorage, B: QuantizedTensorStorage, workspace: torch.Tensor, # pylint: disable=unused-argument @@ -32,7 +32,7 @@ def experimental_gemm( grad: bool = False, ) -> Iterable[Optional[torch.Tensor]]: """Dispatch GEMM to quantizer's qgemm method.""" - assert is_experimental(A) and is_experimental(B), "A and B must be experimental tensors" + assert is_custom(A) and is_custom(B), "A and B must be custom tensors" A, B = B, A diff --git a/transformer_engine/pytorch/experimental/quantization.py b/transformer_engine/pytorch/custom_recipes/quantization.py similarity index 100% rename from transformer_engine/pytorch/experimental/quantization.py rename to transformer_engine/pytorch/custom_recipes/quantization.py diff --git a/transformer_engine/pytorch/experimental/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py similarity index 98% rename from transformer_engine/pytorch/experimental/quantization_nvfp4.py rename to transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py index fc50d07424..1ce9079eb1 100644 --- a/transformer_engine/pytorch/experimental/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py @@ -9,9 +9,9 @@ import torch -from transformer_engine.pytorch.experimental import quantization -from transformer_engine.pytorch.experimental import utils -from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensorStorage, Quantizer +from transformer_engine.pytorch.custom_recipes import quantization +from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer def nvfp4_ref_rht_2d_quantizer_factory(role): @@ -229,8 +229,8 @@ class NVFP4TensorRef(QuantizedTensorStorage): _quantizer: Optional[Quantizer] = None @property - def experimental(self) -> bool: - """Flag to indicate this quantizer is using experimental Kitchen middleware.""" + def custom(self) -> bool: + """Flag to indicate this quantized tensor is custom.""" return True def prepare_for_saving( @@ -362,8 +362,8 @@ def __init__( self.with_random_sign_mask = with_random_sign_mask @property - def experimental(self) -> bool: - """Flag to indicate this quantizer is using experimental Kitchen middleware""" + def custom(self) -> bool: + """Flag to indicate this quantizer is custom.""" return True @staticmethod diff --git a/transformer_engine/pytorch/experimental/utils.py b/transformer_engine/pytorch/custom_recipes/utils.py similarity index 100% rename from transformer_engine/pytorch/experimental/utils.py rename to transformer_engine/pytorch/custom_recipes/utils.py diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 5ed73f6783..8c14d5ab7f 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -41,7 +41,7 @@ from .tensor.mxfp8_tensor import MXFP8Quantizer from .tensor.nvfp4_tensor import NVFP4Quantizer from .tensor.float8_blockwise_tensor import Float8BlockQuantizer -from .tensor.quantized_tensor import QuantizedTensorStorage, QuantizedTensor, Quantizer +from .quantized_tensor import QuantizedTensorStorage, QuantizedTensor, Quantizer from .tensor.storage.float8_tensor_storage import Float8TensorStorage from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index d16455b5b4..7f571ce011 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -38,7 +38,7 @@ _fsdp_gather_tensors, ) from ..constants import dist_group_type -from ..tensor.quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer +from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index a5bf21ee17..aae85e2cab 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -43,7 +43,7 @@ from ..cpu_offload import is_cpu_offload_enabled from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer -from ..tensor.quantized_tensor import ( +from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, prepare_for_saving, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 05f2e9cde4..933c7cde53 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -16,7 +16,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch import torch_version -from transformer_engine.pytorch.tensor.utils import is_experimental +from transformer_engine.pytorch.tensor.utils import is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, get_workspace, @@ -56,7 +56,7 @@ from ..jit import no_torch_dynamo from ..graph import is_graph_capturing from ._common import apply_normalization, noop_cat, WeightGradStore -from ..tensor.quantized_tensor import ( +from ..quantized_tensor import ( QuantizedTensor, QuantizedTensorStorage, Quantizer, @@ -194,13 +194,13 @@ def forward( # Avoid quantized norm kernel if norm output will be returned # or if a gather of ln_out must be in high precision. - experimental = is_experimental(input_quantizer) + custom = is_custom(input_quantizer) with_quantized_norm = ( fp8 and not debug and not return_layernorm_output and not return_layernorm_output_gathered - and not experimental # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() + and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() ) # Apply normalization @@ -246,8 +246,8 @@ def forward( quantizer = None if fp8 or debug: quantizer = input_quantizer - # experimental recipe doesn't need to support quantized AG - if not with_quantized_norm and not experimental: + # custom recipe doesn't need to support quantized AG + if not with_quantized_norm and not custom: ln_out = quantizer(ln_out) quantizer.set_usage(rowwise=True, columnwise=False) if ub_overlap_ag_fprop: # Initialize Userbuffers all-gather diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index a2ddb970af..bae0f28251 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -17,7 +17,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch import torch_version -from transformer_engine.pytorch.tensor.utils import is_experimental +from transformer_engine.pytorch.tensor.utils import is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, get_workspace, @@ -70,7 +70,7 @@ from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ._common import apply_normalization, WeightGradStore from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload -from ..tensor.quantized_tensor import ( +from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, prepare_for_saving, @@ -268,13 +268,13 @@ def forward( # high precision layernorm output and output of the linear are returned # for debug: : layernorm output = High precision to enable processing of this norm - experimental = is_experimental(fc1_input_quantizer) + custom = is_custom(fc1_input_quantizer) with_quantized_norm = ( fp8 and not debug and not return_layernorm_output and not return_layernorm_output_gathered - and not experimental + and not custom ) # Apply normalization @@ -314,8 +314,8 @@ def forward( quantizer = None if fp8 or debug: quantizer = fc1_input_quantizer - # experimental recipe doesn't need to support quantized AG - if not with_quantized_norm and not experimental: + # custom recipe doesn't need to support quantized AG + if not with_quantized_norm and not custom: ln_out = fc1_input_quantizer(ln_out) fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) if ub_overlap_ag: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3069c21d9f..ccb84e6642 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -57,7 +57,7 @@ from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing -from ..tensor.quantized_tensor import ( +from ..quantized_tensor import ( QuantizedTensor, QuantizedTensorStorage, Quantizer, @@ -66,7 +66,7 @@ ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer -from ..tensor.utils import is_experimental +from ..tensor.utils import is_custom from ..export import is_in_onnx_export_mode, assert_warmed_up from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...debug.pytorch.debug_state import TEDebugState @@ -153,8 +153,8 @@ def forward( ub_obj = get_ub(ub_name + "_fprop", fp8) ub_type = tex.CommOverlapType.AG - # experimental recipe check - experimental = is_experimental(input_quantizer) or is_experimental(weight_quantizer) + # custom recipe check + custom = is_custom(input_quantizer) or is_custom(weight_quantizer) # ------------------------------------------------------ # Prepare input tensor @@ -178,7 +178,7 @@ def forward( if fp8 or debug: if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") - if not isinstance(inputmat, QuantizedTensorStorage) and not experimental: + if not isinstance(inputmat, QuantizedTensorStorage) and not custom: own_quantized_input = True input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) if isinstance( @@ -448,7 +448,7 @@ def forward( ctx.main_grad_func = lambda: weight.main_grad ctx.debug = debug - ctx.experimental = experimental + ctx.custom = custom ctx.cpu_offloading = cpu_offloading ctx.is_first_microbatch = is_first_microbatch ctx.use_bias = bias is not None @@ -616,7 +616,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if isinstance(inputmat, QuantizedTensorStorage): # Input tensor is already quantized pass - elif ctx.debug or ctx.experimental: + elif ctx.debug or ctx.custom: # Debug quantizer will be applied immediately before wgrad GEMM pass else: diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 52ca84b5df..a07ffea43f 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -13,7 +13,7 @@ from .. import torch_version from ..quantization import FP8GlobalStateManager from ..tensor.float8_tensor import Float8Tensor -from ..tensor.quantized_tensor import QuantizedTensorStorage +from ..quantized_tensor import QuantizedTensorStorage from ..utils import canonicalize_dtype diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index d95b2298fe..fd1820d15d 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -21,7 +21,7 @@ get_ub, get_workspace, ) -from ...tensor.quantized_tensor import Quantizer +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 diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index e20de53da3..057eb576d7 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -21,7 +21,7 @@ get_workspace, _2X_ACC_FPROP, ) -from ...tensor.quantized_tensor import Quantizer +from ...quantized_tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ...tensor.storage.float8_tensor_storage import Float8TensorStorage from .._common import maybe_dequantize, is_quantized_tensor diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 8ae112022c..6026a40b65 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -28,7 +28,7 @@ fuse_userbuffers_backward_linear, fuse_userbuffers_forward_linear, ) -from transformer_engine.pytorch.tensor.quantized_tensor import ( +from transformer_engine.pytorch.quantized_tensor import ( prepare_for_saving, restore_from_saved, ) diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index ea3e67a57c..f73bc9a966 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex import transformer_engine.pytorch.triton.permutation as triton_permutation from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.tensor.quantized_tensor import QuantizedTensor +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor diff --git a/transformer_engine/pytorch/tensor/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py similarity index 89% rename from transformer_engine/pytorch/tensor/quantized_tensor.py rename to transformer_engine/pytorch/quantized_tensor.py index a524d5c8de..15f5b6bd5e 100644 --- a/transformer_engine/pytorch/tensor/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -2,10 +2,10 @@ # # See LICENSE for license information. -"""Tensor with quantized data""" +"""Pure Python base classes for quantization.""" from __future__ import annotations -from typing import Callable, Optional, Tuple, Iterable, Any, Dict, Union +from typing import Optional, Tuple, Iterable, Any, Dict, Union import abc import copy import warnings @@ -14,6 +14,11 @@ from torch.utils._pytree import tree_map from transformer_engine.common.recipe import Recipe +from transformer_engine.pytorch.tensor._quantization_helpers import ( + _QuantizeFunc, + _IdentityFunc, + _stride_from_shape, +) class QuantizedTensorStorage: @@ -310,73 +315,6 @@ def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-a return True -class _QuantizeFunc(torch.autograd.Function): - """Quantize tensor""" - - @staticmethod - def forward( - _ctx: Optional[torch.autograd.function.FunctionCtx], # unused - tensor: torch.Tensor, - quantize_impl: Callable, - ) -> QuantizedTensor: - # pylint: disable=missing-function-docstring - return quantize_impl(tensor) - - @staticmethod - def backward( - _ctx: torch.autograd.function.FunctionCtx, # unused - grad: torch.Tensor, - ) -> Tuple[Optional[torch.Tensor], ...]: - # pylint: disable=missing-function-docstring - # Assume that we want gradients in full precision - return grad, None - - -class _IdentityFunc(torch.autograd.Function): - """Identity function - - If constructor keyword-arguments are provided, then construct a - new Float8Tensor using the provided tensor's attributes. - - """ - - @staticmethod - def forward( - ctx, tensor: QuantizedTensor, init_kwargs: Optional[Dict[str, Any]] = None - ) -> QuantizedTensor: - # pylint: disable=missing-function-docstring - - # Return input tensor if constructor kwargs are not provided - if init_kwargs is None: - return tensor.detach() - - # Construct new tensor if constructor kwargs are provided - ctx.input_dtype = tensor.dtype - kwargs = tensor.get_metadata() - for key, val in init_kwargs.items(): - kwargs[key] = val - return type(tensor)(tensor.shape, tensor.dtype, **kwargs) - - @staticmethod - def backward(ctx, grad_output): - # pylint: disable=missing-function-docstring - grad_input = grad_output - if grad_input.dtype == ctx.input_dtype: - grad_input = grad_input.detach() - else: - grad_input = grad_input.to(ctx.input_dtype) - return grad_input, None - - -def _stride_from_shape(shape: list[int]): - if len(shape) == 0: - return [] - rstride = [1] - for d in reversed(shape[1:]): - rstride.append(rstride[-1] * d) - return list(reversed(rstride)) - - class QuantizedTensor(torch.Tensor): """Abstract base class for tensor with quantized data diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 7689e20194..ada624a902 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -6,7 +6,7 @@ import torch -from .quantized_tensor import ( +from ..quantized_tensor import ( QuantizedTensorStorage, QuantizedTensor, Quantizer, diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py new file mode 100644 index 0000000000..2214edbff2 --- /dev/null +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -0,0 +1,84 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Private helper functions and classes for quantized tensor implementations. + +This module contains internal autograd functions and utilities that support +the quantization machinery. +""" + +from __future__ import annotations +from typing import Callable, Optional, Tuple, Any, Dict, TYPE_CHECKING +import torch + +if TYPE_CHECKING: + from transformer_engine.pytorch.quantized_tensor import QuantizedTensor + + +class _QuantizeFunc(torch.autograd.Function): + """Quantize tensor""" + + @staticmethod + def forward( + _ctx: Optional[torch.autograd.function.FunctionCtx], # unused + tensor: torch.Tensor, + quantize_impl: Callable, + ) -> QuantizedTensor: + # pylint: disable=missing-function-docstring + return quantize_impl(tensor) + + @staticmethod + def backward( + _ctx: torch.autograd.function.FunctionCtx, # unused + grad: torch.Tensor, + ) -> Tuple[Optional[torch.Tensor], ...]: + # pylint: disable=missing-function-docstring + # Assume that we want gradients in full precision + return grad, None + + +class _IdentityFunc(torch.autograd.Function): + """Identity function + + If constructor keyword-arguments are provided, then construct a + new Float8Tensor using the provided tensor's attributes. + + """ + + @staticmethod + def forward( + ctx, tensor: QuantizedTensor, init_kwargs: Optional[Dict[str, Any]] = None + ) -> QuantizedTensor: + # pylint: disable=missing-function-docstring + + # Return input tensor if constructor kwargs are not provided + if init_kwargs is None: + return tensor.detach() + + # Construct new tensor if constructor kwargs are provided + ctx.input_dtype = tensor.dtype + kwargs = tensor.get_metadata() + for key, val in init_kwargs.items(): + kwargs[key] = val + return type(tensor)(tensor.shape, tensor.dtype, **kwargs) + + @staticmethod + def backward(ctx, grad_output): + # pylint: disable=missing-function-docstring + grad_input = grad_output + if grad_input.dtype == ctx.input_dtype: + grad_input = grad_input.detach() + else: + grad_input = grad_input.to(ctx.input_dtype) + return grad_input, None + + +def _stride_from_shape(shape: list[int]): + """Calculate stride from shape for contiguous tensors""" + if len(shape) == 0: + return [] + rstride = [1] + for d in reversed(shape[1:]): + rstride.append(rstride[-1] * d) + return list(reversed(rstride)) diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 48762499b9..8054374c81 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -14,11 +14,8 @@ from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage -from .quantized_tensor import ( - QuantizedTensor, - Quantizer, - _IdentityFunc, -) +from ..quantized_tensor import QuantizedTensor, Quantizer +from ._quantization_helpers import _IdentityFunc from ..utils import devices_match, round_up_to_nearest_multiple aten = torch.ops.aten diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index a4e68e53b0..de112bb3fd 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -14,11 +14,8 @@ from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling, Recipe from ..utils import canonicalize_process_group, devices_match from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func -from .quantized_tensor import ( - QuantizedTensor, - Quantizer, - _IdentityFunc, -) +from ..quantized_tensor import QuantizedTensor, Quantizer +from ._quantization_helpers import _IdentityFunc from ..constants import dist_group_type aten = torch.ops.aten diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 700de24c4e..5ef5708fdb 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -17,11 +17,8 @@ from ..utils import devices_match, round_up_to_nearest_multiple from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func -from .quantized_tensor import ( - QuantizedTensor, - Quantizer, - _IdentityFunc, -) +from ..quantized_tensor import QuantizedTensor, Quantizer +from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 5e2eeed726..7a5f8858f2 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -22,7 +22,8 @@ ) from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func -from .quantized_tensor import QuantizedTensor, Quantizer, _IdentityFunc +from ..quantized_tensor import QuantizedTensor, Quantizer +from ._quantization_helpers import _IdentityFunc aten = torch.ops.aten diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 9040ea3a43..c2d5e8b3fa 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -13,12 +13,10 @@ from transformer_engine_torch import DType as TE_DType from transformer_engine_torch import Float8BlockScaleTensorFormat -from ..quantized_tensor import QuantizedTensorStorage +from ...quantized_tensor import QuantizedTensorStorage, Quantizer from ...constants import TE_DType_To_Torch -from ..quantized_tensor import Quantizer - from ...utils import _empty_tensor diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index b9533edb6e..a31f6a3799 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -12,12 +12,10 @@ import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType -from ..quantized_tensor import QuantizedTensorStorage +from ...quantized_tensor import QuantizedTensorStorage, Quantizer from ...constants import TE_DType as torch_to_transformer_engine_dtype -from ..quantized_tensor import Quantizer - from ...utils import is_non_tn_fp8_gemm_supported, _empty_tensor diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index c1f30146c9..2cca0829db 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -13,12 +13,10 @@ import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType -from ..quantized_tensor import QuantizedTensorStorage +from ...quantized_tensor import QuantizedTensorStorage, Quantizer from ...constants import TE_DType as torch_to_transformer_engine_dtype -from ..quantized_tensor import Quantizer - from ...utils import _empty_tensor diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 350103f7ca..67543a8e2a 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -16,10 +16,9 @@ # import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType -from ..quantized_tensor import QuantizedTensorStorage +from ...quantized_tensor import QuantizedTensorStorage, Quantizer # from ...constants import TE_DType as torch_to_transformer_engine_dtype -from ..quantized_tensor import Quantizer from ...utils import _empty_tensor diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 72c465edb2..8354823b32 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -4,13 +4,13 @@ """Helper functions for using fp8 tensors as weights""" -import os -from typing import Optional, List, Union +from typing import Optional, Union, List import torch + import transformer_engine_torch as tex from transformer_engine_torch import multi_tensor_scale, multi_tensor_compute_scale_and_scale_inv -from .quantized_tensor import QuantizedTensor, Quantizer, QuantizedTensorStorage +from ..quantized_tensor import QuantizedTensor, Quantizer, QuantizedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer @@ -459,18 +459,13 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten raise ValueError(f"post_processing for {type(model_weight)} is not supported") -def is_experimental(x: Optional[Union[Quantizer, QuantizedTensorStorage]] = None) -> bool: - """Check if an environment or object is using experimental Kitchen middleware. +def is_custom(x: Optional[Union[Quantizer, QuantizedTensorStorage]] = None) -> bool: + """Check if an object is custom. Returns False if x is a torch.Tensor. """ - # Detect if the environment is experimental - if x is None: - return int(os.getenv("QAT_PARAMS", "0")) > 0 - - # Detect if the object is experimental - if isinstance(x, torch.Tensor): + if x is None or isinstance(x, torch.Tensor): return False if not isinstance(x, (Quantizer, QuantizedTensorStorage)): raise AssertionError("Object must be a Quantizer or QuantizedTensorStorage instance") - return hasattr(x, "experimental") and x.experimental + return hasattr(x, "custom") and x.custom diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 2be0aed4a8..90c6289963 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -12,7 +12,7 @@ import torch from . import torch_version -from .tensor.quantized_tensor import Quantizer +from .quantized_tensor import Quantizer from ..debug.pytorch.debug_quantization import DebugQuantizedTensor From 818b30cc4b07bcac955b17a6a12ca9708d7f0a7e Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 22 Oct 2025 08:51:36 -0700 Subject: [PATCH 016/521] [JAX] NVFP4 recipe with option to enable/disable SR, RHT, and 2D quantization (#2270) * [JAX] Support recipe flags for disabling SR, RHT, and 2D quantization Signed-off-by: Jeremy Berchtold * lint Signed-off-by: Jeremy Berchtold * Fix issue with SR state being erased due to pytree handling of NVFP4Quantizer Signed-off-by: Jeremy Berchtold * Add test for SR state preservation across VJP boundaries Signed-off-by: Jeremy Berchtold * Fix sharding of SR rng state Signed-off-by: Jeremy Berchtold * lint Signed-off-by: Jeremy Berchtold * update tolerances slightly now that SR is enabled Signed-off-by: Jeremy Berchtold * lint Signed-off-by: Jeremy Berchtold * Use hashlib for deterministic hashes across runs for SR Signed-off-by: Jeremy Berchtold * rename uses_rht on scaled tensors to has_applied_rht Signed-off-by: Jeremy Berchtold * add assert Signed-off-by: Jeremy Berchtold * Move decision of whether to use RHT into helper.py and add dedicated RHT tests Signed-off-by: Jeremy Berchtold * lint Signed-off-by: Jeremy Berchtold * fix use_rht attr usage Signed-off-by: Jeremy Berchtold * fix pure-jax rht usage criteria Signed-off-by: Jeremy Berchtold * Adjust tolerances after rebase Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- .../encoder/test_multiprocessing_encoder.py | 4 +- .../jax/encoder/test_single_gpu_encoder.py | 2 +- tests/jax/test_custom_call_compute.py | 155 ++++++++++++------ tests/jax/test_helper.py | 82 ++++++++- transformer_engine/jax/cpp_extensions/gemm.py | 16 +- .../jax/cpp_extensions/quantization.py | 40 +++-- .../jax/quantize/dequantizer.py | 11 +- transformer_engine/jax/quantize/hadamard.py | 26 --- transformer_engine/jax/quantize/helper.py | 90 +++++++--- transformer_engine/jax/quantize/metadata.py | 20 +++ transformer_engine/jax/quantize/quantizer.py | 34 +++- transformer_engine/jax/quantize/tensor.py | 25 +++ transformer_engine/jax/sharding.py | 13 ++ 13 files changed, 382 insertions(+), 136 deletions(-) diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index c2e97029b0..9605adf771 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -672,7 +672,7 @@ def test_te_mxfp8(self): def test_te_nvfp4(self): """Test Transformer Engine with NVFP4""" result = self.exec(True, "NVFP4BlockScaling") - assert result[0] < 0.451 and result[1] > 0.79 + assert result[0] < 0.451 and result[1] > 0.788 @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") def test_te_bf16_shardy(self): @@ -710,7 +710,7 @@ def test_te_mxfp8_shardy(self): def test_te_nvfp4_shardy(self): """Test Transformer Engine with NVFP4""" result = self.exec(True, "NVFP4BlockScaling", enable_shardy=True) - assert result[0] < 0.451 and result[1] > 0.79 + assert result[0] < 0.451 and result[1] > 0.788 if __name__ == "__main__": diff --git a/examples/jax/encoder/test_single_gpu_encoder.py b/examples/jax/encoder/test_single_gpu_encoder.py index 1c62de7fa4..81f2d6c744 100644 --- a/examples/jax/encoder/test_single_gpu_encoder.py +++ b/examples/jax/encoder/test_single_gpu_encoder.py @@ -390,7 +390,7 @@ def test_te_nvfp4(self): self.args.use_fp8 = True self.args.fp8_recipe = "NVFP4BlockScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.476 and actual[1] > 0.775 + assert actual[0] < 0.477 and actual[1] > 0.769 if __name__ == "__main__": diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 2934e48df1..1217ebf65f 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -40,7 +40,6 @@ QuantizerFactory, QuantizeLayout, noop_quantizer_set, - should_use_rht, ) from transformer_engine.jax.quantize import helper from transformer_engine.jax.activation import activation @@ -685,21 +684,14 @@ class TestQuantize: Purely quantization related tests that will always test on a wider set of types and shapes """ - def _skip_for_fp4(self, input_shape, q_dtype, scaling_mode, q_layout, flatten_axis): - """Temporary hack to skip unsupported FP4 cases until we implement them""" + def _skip_unsupported_dtypes(self, q_dtype, scaling_mode): + """Skip unsupported dtypes for given scaling mode. For example, NVFP4 only supports the float4_e2m1 dtype not float8 dtypes.""" if q_dtype not in scaling_mode.get_compatible_q_dtypes(): pytest.skip(f"Quantize dtype {q_dtype} is not supported by {scaling_mode}") return - # HACK: FIXME TODO(jberchtold) - row = reduce(operator.mul, input_shape[flatten_axis:], 1) - col = reduce(operator.mul, input_shape[:flatten_axis], 1) - will_use_rht = should_use_rht(scaling_mode, q_layout=q_layout) - if will_use_rht and (row % 64 != 0 or col % 128 != 0): - pytest.skip("Unfused RHT is not supported currently, skipping") - def test_qdq(self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatten_axis): - self._skip_for_fp4(input_shape, q_dtype, scaling_mode, q_layout, flatten_axis) + self._skip_unsupported_dtypes(q_dtype, scaling_mode) key = jax.random.PRNGKey(0) @@ -780,22 +772,8 @@ def test_qdq(self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatt assert_dequantized_scaled_tensor(scaled_tensor, x) def _should_use_precise_comparison( - self, in_dtype, scaling_mode, q_layout, input_shape, flatten_axis + self, in_dtype, scaling_mode, quantizer, input_shape, flatten_axis ): - # TODO(jberchtold): Remove this hack once we have a better solution to ensure bitwise identical results between TE and JAX RHT+quant implementations. Currently for certain shapes the quantized fp4 data differs by a small amount on <0.5% of the values. - RHT_SLIGHT_MISMATCH_SHAPES = [ - ((32, 256, 128), -1), - ((64, 32, 32, 256), -1), - ((8192, 2, 4096), -2), - ] - - if ( - should_use_rht(scaling_mode, q_layout=q_layout) - and (input_shape, flatten_axis) in RHT_SLIGHT_MISMATCH_SHAPES - ): - # TE fused RHT+quant and JAX RHT+quant have slight implementation differences which can lead to small numerical differences on certain shapes - return False - if scaling_mode.is_nvfp4_scaling and in_dtype != jnp.bfloat16: # With NVFP4 scaling, TE kernels internally use bfloat16 so using a different input dtype can lead to small numerical differences compared to the JAX implementation return False @@ -805,7 +783,7 @@ def _should_use_precise_comparison( def test_quantize_bitwise( self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatten_axis ): - self._skip_for_fp4(input_shape, q_dtype, scaling_mode, q_layout, flatten_axis) + self._skip_unsupported_dtypes(q_dtype, scaling_mode) key = jax.random.PRNGKey(0) input = jax.random.uniform(key, input_shape, in_dtype) @@ -816,28 +794,20 @@ def test_quantize_bitwise( jax_output = _jax_quantize(input, quantizer=jax_quantizer, flatten_axis=flatten_axis) - try: - te_output = tex.quantize(input, quantizer=te_quantizer, flatten_axis=flatten_axis) - except AssertionError as e: - if should_use_rht(scaling_mode, q_layout=q_layout) and in_dtype != jnp.bfloat16: - error_message = e.args[0] - if "RHT requires input to be bfloat16" in error_message: - # Successfully caught the expected error, early return from the test - return - raise e + te_output = tex.quantize(input, quantizer=te_quantizer, flatten_axis=flatten_axis) assert_bitwise_scaled_tensors( te_output, jax_output, precise_comparison=self._should_use_precise_comparison( - in_dtype, scaling_mode, q_layout, input_shape, flatten_axis + in_dtype, scaling_mode, te_quantizer, input_shape, flatten_axis ), ) def test_quantize_bitwise_jitted( self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatten_axis ): - self._skip_for_fp4(input_shape, q_dtype, scaling_mode, q_layout, flatten_axis) + self._skip_unsupported_dtypes(q_dtype, scaling_mode) key = jax.random.PRNGKey(0) input = jax.random.uniform(key, input_shape, in_dtype) @@ -851,21 +821,13 @@ def test_quantize_bitwise_jitted( jax_output = jax_impl_func_jit(input, quantizer=jax_quantizer, flatten_axis=flatten_axis) - try: - te_output = te_impl_func_jit(input, quantizer=te_quantizer, flatten_axis=flatten_axis) - except AssertionError as e: - if should_use_rht(scaling_mode, q_layout=q_layout) and in_dtype != jnp.bfloat16: - error_message = e.args[0] - if "RHT requires input to be bfloat16" in error_message: - # Successfully caught the expected error, early return from the test - return - raise e + te_output = te_impl_func_jit(input, quantizer=te_quantizer, flatten_axis=flatten_axis) assert_bitwise_scaled_tensors( te_output, jax_output, precise_comparison=self._should_use_precise_comparison( - in_dtype, scaling_mode, q_layout, input_shape, flatten_axis + in_dtype, scaling_mode, te_quantizer, input_shape, flatten_axis ), ) @@ -985,12 +947,6 @@ def _test_sr( def test_sr_nvfp4(self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatten_axis): """Tests that the mean absolute error of stochastic rounding is smaller than round nearest quantization over multiple samples for both TE and JAX implementations. Asserts that the MAE of both implementations is close to each other.""" - # HACK: FIXME TODO(jberchtold) - row = reduce(operator.mul, input_shape[flatten_axis:], 1) - col = reduce(operator.mul, input_shape[:flatten_axis], 1) - will_use_rht = should_use_rht(scaling_mode, q_layout=q_layout) - if will_use_rht and (row % 64 != 0 or col % 128 != 0): - pytest.skip("Unfused RHT is not supported currently, skipping") key = jax.random.PRNGKey(0) inputs = jax.random.uniform(key, input_shape, in_dtype) @@ -1007,6 +963,97 @@ def test_sr_nvfp4(self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, assert_allclose(te_mean_error, jax_mean_error, rtol=0.2, atol=1e-4) +@pytest_parametrize_wrapper("in_dtype", [jnp.bfloat16]) +@pytest_parametrize_wrapper("q_dtype", [jnp.float4_e2m1fn]) +@pytest_parametrize_wrapper( + "scaling_mode", [s for s in supported_scaling_modes if s == ScalingMode.NVFP4_1D_SCALING] +) +class TestRandomizedHadamardTransform: + + @pytest_parametrize_wrapper( + "q_layout", [QuantizeLayout.ROWWISE_COLWISE, QuantizeLayout.COLWISE] + ) + @pytest_parametrize_wrapper("input_shape,flatten_axis", [((64, 128), -1)]) + def test_rht_quantize_bitwise_jitted( + self, in_dtype, q_dtype, scaling_mode, q_layout, input_shape, flatten_axis + ): + key = jax.random.PRNGKey(0) + inputs = jax.random.uniform(key, input_shape, in_dtype) + + te_quantizer, jax_quantizer = QuantizerFactory.create( + n_quantizers=2, + q_dtype=q_dtype, + scaling_mode=scaling_mode, + q_layout=q_layout, + use_rht=True, + ) + + jax_impl_func_jit = jax.jit(_jax_quantize, static_argnums=(2, 3)) + te_impl_func_jit = jax.jit(tex.quantize, static_argnums=(2,)) + + jax_output = jax_impl_func_jit(inputs, quantizer=jax_quantizer, flatten_axis=flatten_axis) + + te_output = te_impl_func_jit(inputs, quantizer=te_quantizer, flatten_axis=flatten_axis) + + assert_bitwise_scaled_tensors(te_output, jax_output) + + def _ref_gemm_with_jnp_dot(self, a, b, data_layout): + if data_layout[0] == "T": + a = jnp.swapaxes(a, -1, -2) + if data_layout[1] == "T": + b = jnp.swapaxes(b, -1, -2) + return jnp.dot(a, b) + + def _generate_gemm_input(self, m, n, k, data_layout): + key = jax.random.PRNGKey(0) + subkeys = jax.random.split(key, 2) + x = jax.random.uniform( + subkeys[0], + (m if data_layout[0] == "N" else k, k if data_layout[0] == "N" else m), + dtype=jnp.bfloat16, + ) / jnp.sqrt(k) + w = jax.random.uniform( + subkeys[1], + (k if data_layout[1] == "N" else n, n if data_layout[1] == "N" else k), + dtype=jnp.bfloat16, + ) / jnp.sqrt(n) + lhs_contracting_dim = (1,) if data_layout[0] == "N" else (0,) + rhs_contracting_dim = (0,) if data_layout[1] == "N" else (1,) + contracting_dims = (lhs_contracting_dim, rhs_contracting_dim) + + return (x, w, contracting_dims) + + @pytest_parametrize_wrapper("m,n,k", [(64, 32, 64)]) + # We do not test NN and TT layouts here as they do not have both inputs using RHT due to RHT only supporting the colwise layout currently + @pytest_parametrize_wrapper("data_layout", ["TN", "NT"]) + @pytest_parametrize_wrapper("with_jax_gemm", [True, False]) + def test_rht_gemm(self, in_dtype, q_dtype, scaling_mode, m, n, k, data_layout, with_jax_gemm): + key = jax.random.PRNGKey(0) + + lhs_scaling_mode, rhs_scaling_mode = scaling_mode, scaling_mode + x, w, contracting_dims = self._generate_gemm_input(m, n, k, data_layout) + lhs_quantizer = QuantizerFactory.create( + scaling_mode=lhs_scaling_mode, + q_dtype=jnp.float4_e2m1fn, + use_rht=True, + ) + rhs_quantizer = QuantizerFactory.create( + scaling_mode=rhs_scaling_mode, + q_dtype=jnp.float4_e2m1fn, + use_rht=True, + ) + with use_jax_gemm(enabled=with_jax_gemm): + primitive_out = tex.gemm( + x, + w, + contracting_dims=contracting_dims, + lhs_quantizer=lhs_quantizer, + rhs_quantizer=rhs_quantizer, + ) + ref_out = self._ref_gemm_with_jnp_dot(x, w, data_layout) + assert_allclose(primitive_out, ref_out, dtype=jnp.float4_e2m1fn) + + @pytest.mark.skipif(not is_fp8_supported, reason=fp8_unsupported_reason) @pytest_parametrize_wrapper("in_dtype", QUANTIZATION_INPUT_DTYPE) @pytest_parametrize_wrapper("input_shape", [(8, 16, 32)]) diff --git a/tests/jax/test_helper.py b/tests/jax/test_helper.py index ca804625c6..fc88b7ef77 100644 --- a/tests/jax/test_helper.py +++ b/tests/jax/test_helper.py @@ -3,11 +3,13 @@ # See LICENSE for license information. import unittest +from functools import partial import flax import jax import jax.numpy as jnp import numpy as np +from flax import linen as nn from utils import assert_allclose from transformer_engine.common.recipe import ( @@ -24,15 +26,51 @@ ScalingMode, update_collections, TensorSource, + QuantizerFactory, + QuantizeLayout, ) from transformer_engine.jax.quantize.helper import _format2dtypes from transformer_engine.jax.sharding import MeshResource, global_mesh_resource +from transformer_engine.jax.flax.module import TransformerEngineBase is_fp8_supported, reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING) is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING) is_nvfp4_supported, nvfp4_reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING) +def quantizer_check_vjp(outer_quantizer_set, assertion_func, x): + """Check that the quantizers in the quantizer set are as expected and reconstructed correctly from flattened pytree representations across VJP boundaries.""" + + # Define a function with a custom VJP (vector-Jacobian product) + @partial(jax.custom_vjp, nondiff_argnums=(1,)) + def quantizer_check(inner_quantizer_set, assertion_func, x): + return quantizer_check_fwd(inner_quantizer_set, assertion_func, x) + + def quantizer_check_fwd(inner_quantizer_set, assertion_func, x): + assertion_func(inner_quantizer_set.x, TensorSource.X) + assertion_func(inner_quantizer_set.kernel, TensorSource.KERNEL) + assertion_func(inner_quantizer_set.dgrad, TensorSource.DGRAD) + return x + + def quantizer_check_bwd(ctx, g): + return (g,) + + quantizer_check.defvjp(quantizer_check_fwd, quantizer_check_bwd) + return quantizer_check(outer_quantizer_set, assertion_func, x) + + +class TestModule(TransformerEngineBase): + """A simple module to test quantizer creation and reconstruction across VJP boundaries.""" + + # Signature: (quantizer: Quantizer, tensor_source: TensorSource) -> None + assertion_func: callable + + @nn.compact + def __call__(self, x): + quantizer_set = self.generate_quantizer_set() + return quantizer_check_vjp(quantizer_set, self.assertion_func, x) + + class TestHelper(unittest.TestCase): @unittest.skipIf(not is_fp8_supported, reason=reason) @@ -89,12 +127,43 @@ def _compare_nvfp4_scaling(self, test): for tensor_source in TensorSource: target_scaling_mode = ( ScalingMode.NVFP4_2D_SCALING - if tensor_source == TensorSource.KERNEL + if (not test.disable_2d_quantization) and tensor_source == TensorSource.KERNEL else ScalingMode.NVFP4_1D_SCALING ) self.assertEqual( get_quantize_config().get_scaling_mode(tensor_source), target_scaling_mode ) + self.assertEqual( + get_quantize_config().DISABLE_STOCHASTIC_ROUNDING, test.disable_stochastic_rounding + ) + self.assertEqual(get_quantize_config().DISABLE_RHT, test.disable_rht) + self.assertEqual( + get_quantize_config().DISABLE_2D_QUANTIZATION, test.disable_2d_quantization + ) + + def _compare_nvfp4_scaling_quantizers(self, test): + """Check that the quantizers created have the expected stochastic rounding state and the state is preserved across VJP boundaries.""" + + def assertion_func(quantizer, tensor_source): + if test.disable_stochastic_rounding or tensor_source != TensorSource.DGRAD: + self.assertIsNone(quantizer.stochastic_rounding_rng_state) + else: + self.assertIsNotNone(quantizer.stochastic_rounding_rng_state) + + expected_rht = ( + quantizer.scaling_mode == ScalingMode.NVFP4_1D_SCALING + and quantizer.q_layout in {QuantizeLayout.ROWWISE_COLWISE, QuantizeLayout.COLWISE} + and not test.disable_rht + ) + self.assertEqual(quantizer.use_rht, expected_rht) + + x = jnp.ones((), dtype=jnp.float32) + test_module = TestModule(assertion_func=assertion_func) + param_key, sr_key = jax.random.split(jax.random.PRNGKey(0)) + rngs = {"params": param_key, "sr_rng": sr_key} + variables = test_module.init(rngs, x) + + jax.jit(jax.value_and_grad(test_module.apply), static_argnums=(2,))(variables, x, rngs=rngs) @unittest.skipIf(not is_fp8_supported, reason=reason) def test_autocast_delayed_scaling(self): @@ -171,5 +240,16 @@ def test_autocast_nvfp4_block_scaling(self): with autocast(enabled=True, recipe=bs, mesh_resource=MeshResource()): self.assertTrue(get_quantize_config().is_fp8_enabled()) self._compare_nvfp4_scaling(bs) + self._compare_nvfp4_scaling_quantizers(bs) + + bs = NVFP4BlockScaling( + disable_stochastic_rounding=True, + disable_rht=True, + disable_2d_quantization=True, + ) + with autocast(enabled=True, recipe=bs, mesh_resource=MeshResource()): + self.assertTrue(get_quantize_config().is_fp8_enabled()) + self._compare_nvfp4_scaling(bs) + self._compare_nvfp4_scaling_quantizers(bs) self._check_default_state() diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index b37c4bd848..778f77c0d5 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -44,7 +44,6 @@ noop_quantizer_set, is_fp8_gemm_with_all_layouts_supported, apply_padding_to_scale_inv, - should_use_rht, ) from .misc import get_padded_spec, is_all_reduce_in_float32 from ..sharding import ( @@ -169,16 +168,13 @@ def _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_ assert not isinstance(lhs_q, ScaledTensor2x) assert not isinstance(rhs_q, ScaledTensor2x) - def uses_rht(q: AbstractBaseTensor) -> bool: - return isinstance(q, ScaledTensor1x) and should_use_rht( - q.scaling_mode, is_colwise=q.is_colwise - ) + def has_rht_applied(q: AbstractBaseTensor) -> bool: + return isinstance(q, ScaledTensor1x) and q.has_rht_applied - # TODO(jberchtold): Move RHT usage check to a bool flag on the ScaledTensor class - assert uses_rht(lhs_q) == uses_rht(rhs_q), ( - "With NVFP4_1D_SCALING, if one operand is colwise quantized, the other must be colwise" - " quantized as well. This is to ensure the RHT is applied to both and will cancel out in" - " the GEMM." + assert has_rht_applied(lhs_q) == has_rht_applied(rhs_q), ( + "With NVFP4_1D_SCALING, if one operand is quantized with RHT, the other must be quantized" + " with RHT as well. This is to ensure the RHT is applied to both and will cancel out in the" + " GEMM." ) return lhs_q, rhs_q diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index b3f1e60f9a..67c505bc98 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -31,7 +31,7 @@ from ..sharding import ( all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp, - num_of_devices, + get_num_devices_in_mesh, ) from ..quantize import ( ScaledTensor2x, @@ -45,7 +45,6 @@ compute_scale_from_amax, NoScaleTensor, get_rht_matrix, - should_use_rht, ) @@ -108,17 +107,18 @@ def abstract( "sr_rng_state must be a uint32 array when stochastic_rounding is True but" f" received {sr_rng_state_aval}" ) - if is_outer: + if is_outer and get_num_devices_in_mesh() > 1: assert ( - sr_rng_state_aval.shape[0] == num_of_devices() + sr_rng_state_aval.shape[0] == get_num_devices_in_mesh() and sr_rng_state_aval.shape[1] == 4 ), ( "sr_rng_state must be of shape (num_devices, 4) when stochastic_rounding is" f" True and is_outer is True but received {sr_rng_state_aval.shape}" ) else: - assert sr_rng_state_aval.shape == (4,), ( - "Sharded sr_rng_state must be of shape (4,) per device when" + # We cannot assert the shape is exactly (4,) here because if the quantized data is not perfectly sharded across all devices then we will have extra rng state here. For example, this could occur when the weights are not sharded when using data parallelism. However, this is okay because the extra rng state will simply not be used and each device still has a unique rng state. + assert sr_rng_state_aval.size >= 4, ( + "Sharded sr_rng_state must have at least 4 elements per device when" f" stochastic_rounding is True but received {sr_rng_state_aval.shape}" ) @@ -552,8 +552,13 @@ def partition( desc="BaseDBiasQuantizePrimitive.colwise_scale_inv", ) - # TODO(jberchtold): Assert the sr_rng state is sharded along all mesh axes - arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + arg_shardings = list(arg_i.sharding for arg_i in arg_infos) + arg_shardings[3] = NamedSharding( + mesh, + PartitionSpec(tuple(x for x in x_spec if x is not None), None), + desc="BaseDBiasQuantizePrimitive.sr_rng_state", + ) + arg_shardings = tuple(arg_shardings) out_shardings = ( out_sharding, colwise_out_sharding, @@ -564,6 +569,9 @@ def partition( ) def sharded_impl(x, scale, amax, sr_rng_state, post_rht_amax, rht_matrix): + if sr_rng_state.size > 4: + # See comment in abstract method for explanation of why we cannot assert exact shape + sr_rng_state = sr_rng_state.flatten()[:4] ( local_x, local_colwise_x, @@ -754,9 +762,10 @@ def _quantize_dbias_impl( # If TE/common custom quantize op is disabled, or if quantizer layout is COLWISE, # fall back on the native-JAX quantize implementation PrimitiveClass = DBiasQuantizePrimitive if is_dbias else QuantizePrimitive - is_unsupported = ( - quantizer.q_layout == QuantizeLayout.COLWISE - and quantizer.scaling_mode != ScalingMode.NVFP4_1D_SCALING + is_unsupported = quantizer.q_layout == QuantizeLayout.COLWISE and not ( + quantizer.scaling_mode == ScalingMode.NVFP4_1D_SCALING + and hasattr(quantizer, "use_rht") + and quantizer.use_rht ) if is_unsupported or not PrimitiveClass.enabled(): if is_dbias: @@ -792,7 +801,7 @@ def _quantize_dbias_impl( rht_matrix = jnp.empty((1, 1), jnp.bfloat16) amax = x.amax - if should_use_rht(quantizer.scaling_mode, q_layout=quantizer.q_layout): + if hasattr(quantizer, "use_rht") and quantizer.use_rht: use_rht = True rht_matrix = get_rht_matrix() @@ -861,7 +870,11 @@ def _quantize_dbias_impl( x.data, scale, amax, - sr_rng_state if sr_rng_state is not None else jnp.empty((num_of_devices(), 1), jnp.uint32), + ( + sr_rng_state + if sr_rng_state is not None + else jnp.empty((get_num_devices_in_mesh(), 1), jnp.uint32) + ), post_rht_amax if post_rht_amax is not None else jnp.zeros((1,), jnp.float32), rht_matrix, out_dtype=quantizer.q_dtype, @@ -902,6 +915,7 @@ def _quantize_dbias_impl( q_layout=quantizer.q_layout, data_layout=quantizer.get_data_layout(), flatten_axis=flatten_axis, + colwise_has_rht_applied=use_rht, ) return out, dbias.astype(dq_dtype) diff --git a/transformer_engine/jax/quantize/dequantizer.py b/transformer_engine/jax/quantize/dequantizer.py index b4da6f3bed..80ebc6b875 100644 --- a/transformer_engine/jax/quantize/dequantizer.py +++ b/transformer_engine/jax/quantize/dequantizer.py @@ -15,7 +15,7 @@ import jax.numpy as jnp from .scaling_modes import ScalingMode -from .hadamard import apply_rht, should_use_rht +from .hadamard import apply_rht __all__ = ["ScalingModeToDequantizerMap"] @@ -171,7 +171,9 @@ class NVFP4Dequantizer(Dequantizer): """ @staticmethod - def _dequantize_func(data, scale_inv, amax, dq_dtype, scaling_mode, is_colwise, flatten_axis): + def _dequantize_func( + data, scale_inv, amax, dq_dtype, scaling_mode, is_colwise, flatten_axis, has_rht_applied + ): """Dequantize a tensor using block scaling. Args: @@ -182,6 +184,7 @@ def _dequantize_func(data, scale_inv, amax, dq_dtype, scaling_mode, is_colwise, scaling_mode: The scaling mode used for quantization is_colwise: Whether the scaling is column-wise flatten_axis: The axis along which the tensor could be flattened to 2D + has_rht_applied: Whether the quantization has RHT applied and we need to apply the inverse RHT to dequantize Returns: The dequantized tensor @@ -223,8 +226,7 @@ def _dequantize_func(data, scale_inv, amax, dq_dtype, scaling_mode, is_colwise, out = jnp.asarray(data * scale_inv, dq_dtype).reshape(data_shape) # Apply inverse of RHT if needed - use_rht = should_use_rht(scaling_mode, is_colwise=is_colwise) - if use_rht: + if has_rht_applied: out = apply_rht(out, inverse=True) return out @@ -247,6 +249,7 @@ def dequantize(scaled_tensor): scaled_tensor.scaling_mode, scaled_tensor.is_colwise, scaled_tensor.flatten_axis, + scaled_tensor.has_rht_applied, ) diff --git a/transformer_engine/jax/quantize/hadamard.py b/transformer_engine/jax/quantize/hadamard.py index c0b74ef75e..5f6f0ec2b5 100644 --- a/transformer_engine/jax/quantize/hadamard.py +++ b/transformer_engine/jax/quantize/hadamard.py @@ -4,32 +4,6 @@ """Randomized Hadamard Transform (RHT) utilities for JAX.""" import jax.numpy as jnp -from .scaling_modes import ScalingMode - - -def should_use_rht(scaling_mode, is_colwise=None, q_layout=None) -> bool: - """Determine if RHT (Randomized Hadamard Transform) should be used. - - Args: - scaling_mode: The scaling mode of the tensor. - is_colwise: Whether the tensor is column-wise. Only one of is_colwise or q_layout should be provided. - q_layout: The quantization layout of the tensor. Only one of is_colwise or q_layout should be provided. - - Returns: - bool: True if RHT should be used, False otherwise. - """ - # Delayed import to avoid circular dependencies - from .quantizer import QuantizeLayout - - assert (is_colwise is None) != ( - q_layout is None - ), "Exactly one of is_colwise or q_layout must be provided." - - if q_layout is not None: - is_colwise = q_layout in {QuantizeLayout.COLWISE, QuantizeLayout.ROWWISE_COLWISE} - - return scaling_mode == ScalingMode.NVFP4_1D_SCALING and is_colwise - def get_wgrad_sign_vector() -> list[int]: """Get a fixed sign vector for the RHT used in NVFP4 weight gradient quantization.""" diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index 06c67b62ee..e8b33c1d1c 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -12,6 +12,7 @@ from contextlib import contextmanager from dataclasses import dataclass from enum import Enum +import hashlib from typing import Optional, Tuple, Dict, Union, Sequence, Type, List from functools import reduce, lru_cache import operator @@ -35,7 +36,7 @@ from transformer_engine.jax.sharding import ( global_shard_guard, MeshResource, - num_of_devices, + get_num_devices_in_mesh, get_all_mesh_axes, with_sharding_constraint, ) @@ -561,29 +562,87 @@ def get_quantize_flax_meta( return QuantizeMeta() +@dataclass class NVFP4ScalingQuantizeConfig(BaseQuantizeConfig): """Configuration class for NVFP4 scaling recipe. This class provides specific initialization and finalization for NVFP4 scaling quantization mode. """ + DISABLE_STOCHASTIC_ROUNDING: bool = False + DISABLE_RHT: bool = False + DISABLE_2D_QUANTIZATION: bool = False + def initialize_from_recipe(self, fp8_recipe: Recipe) -> None: - """Initialize block scaling FP8 configuration. + """Initialize block scaling NVFP4 configuration. Args: - fp8_recipe: The FP8 recipe to use for initialization + fp8_recipe: The quantization recipe to use for initialization """ + assert isinstance(fp8_recipe, NVFP4BlockScaling) + self.INITIALIZED = True self.FWD_DTYPE, self.BWD_DTYPE = _format2dtypes(fp8_recipe.fp4_format) self.AMAX_HISTORY_LEN = 0 + self.DISABLE_STOCHASTIC_ROUNDING = fp8_recipe.disable_stochastic_rounding + self.DISABLE_RHT = fp8_recipe.disable_rht + self.DISABLE_2D_QUANTIZATION = fp8_recipe.disable_2d_quantization + def get_scaling_mode(self, tensor_source: TensorSource) -> ScalingMode: """Gets the scaling mode for a specific tensor's usage type.""" - if tensor_source == TensorSource.KERNEL: + if (not self.DISABLE_2D_QUANTIZATION) and tensor_source == TensorSource.KERNEL: return ScalingMode.NVFP4_2D_SCALING # for x and grad return ScalingMode.NVFP4_1D_SCALING + def _make_rht_quantize_meta(self, q_layout, tensor_source: TensorSource) -> QuantizeMeta: + """Create the quantization metadata for RHT if applicable.""" + # Imported here to prevent circular import + from transformer_engine.jax.quantize import QuantizeLayout + + use_rht = self.get_scaling_mode( + tensor_source + ) == ScalingMode.NVFP4_1D_SCALING and q_layout in { + QuantizeLayout.ROWWISE_COLWISE, + QuantizeLayout.COLWISE, + } + if self.DISABLE_RHT: + use_rht = False + return QuantizeMeta(use_rht=use_rht) + + def _make_stochastic_rounding_rng_state( + self, module, tensor_source: TensorSource, quantizer_name: str + ) -> jnp.ndarray: + """Create the stochastic rounding rng state if applicable.""" + if self.DISABLE_STOCHASTIC_ROUNDING: + return QuantizeMeta() + + if tensor_source != TensorSource.DGRAD: + # Only DGRAD uses stochastic rounding + return QuantizeMeta() + + sr_jax_rng = module.make_rng("sr_rng") + # Get a unique key for this quantizer + # Use hashlib to get a deterministic hash value for quantizer_name + quantizer_hash = ( + int(hashlib.sha256(quantizer_name.encode("utf-8")).hexdigest(), 16) + % jnp.iinfo(jnp.int32).max + ) + sr_jax_rng = jax.jit(jax.random.fold_in)(sr_jax_rng, quantizer_hash) + + # Generate 4 random uint32 values from the JAX PRNG key + shape = (4,) + if get_num_devices_in_mesh() > 1: + shape = (get_num_devices_in_mesh(), 4) + sr_jax_rng_state = jax.random.randint( + sr_jax_rng, shape, 0, jnp.iinfo(jnp.int32).max, dtype=jnp.int32 + ).view(jnp.uint32) + sr_jax_rng_state = with_sharding_constraint( + sr_jax_rng_state, jax.sharding.PartitionSpec(get_all_mesh_axes(), None) + ) + return QuantizeMeta(stochastic_rounding_rng_state=sr_jax_rng_state) + def get_quantize_flax_meta( self, module, @@ -603,27 +662,14 @@ def get_quantize_flax_meta( Returns: The quantization metadata for the specified module and tensor. It can be empty if no metadata is needed. """ - if tensor_source != TensorSource.DGRAD: - # Only DGRAD uses stochastic rounding - return QuantizeMeta() - - # TODO(jberchtold): This assumes SR is always enabled for NVFP4. Use flag from recipe to toggle it. - sr_jax_rng = module.make_rng("sr_rng") - # Get a unique key for this quantizer - sr_jax_rng = jax.jit(jax.random.fold_in)( - sr_jax_rng, hash(quantizer_name) % jnp.iinfo(jnp.int32).max - ) + # Imported here to prevent circular import + from transformer_engine.jax.quantize import QuantizeLayout - # Generate 4 random uint32 values from the JAX PRNG key - sr_jax_rng_state = jax.random.randint( - sr_jax_rng, (num_of_devices(), 4), 0, jnp.iinfo(jnp.int32).max, dtype=jnp.int32 - ).view(jnp.uint32) - sr_jax_rng_state = with_sharding_constraint( - sr_jax_rng_state, jax.sharding.PartitionSpec(get_all_mesh_axes(), None) + return QuantizeMeta.merge( + self._make_rht_quantize_meta(QuantizeLayout.ROWWISE_COLWISE, tensor_source), + self._make_stochastic_rounding_rng_state(module, tensor_source, quantizer_name), ) - return QuantizeMeta(stochastic_rounding_rng_state=sr_jax_rng_state) - _QUANTIZE_CONFIG = NoOpQuantizeConfig() diff --git a/transformer_engine/jax/quantize/metadata.py b/transformer_engine/jax/quantize/metadata.py index 11a349ed7d..a987643eb7 100644 --- a/transformer_engine/jax/quantize/metadata.py +++ b/transformer_engine/jax/quantize/metadata.py @@ -26,6 +26,26 @@ class QuantizeMeta: """ + @staticmethod + def merge(a: "QuantizeMeta", b: "QuantizeMeta") -> "QuantizeMeta": + """Merge two QuantizeMeta instances. + + Args: + a (QuantizeMeta): The first QuantizeMeta instance. + b (QuantizeMeta): The second QuantizeMeta instance. + + Returns: + QuantizeMeta: A new QuantizeMeta instance with merged metadata. + """ + assert isinstance(a, QuantizeMeta) + assert isinstance(b, QuantizeMeta) + for key in b.get_kwargs_dictionary().keys(): + if key in a.get_kwargs_dictionary(): + assert ( + a.get_kwargs_dictionary()[key] == b.get_kwargs_dictionary()[key] + ), f"Conflict in merging QuantizeMeta: {key} has different values." + return QuantizeMeta(**{**a.get_kwargs_dictionary(), **b.get_kwargs_dictionary()}) + def __init__(self, **kwargs): self._kwargs = kwargs diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index 7bc08f834f..d138b58dad 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -19,7 +19,7 @@ from transformer_engine.common import recipe from .scaling_modes import ScalingMode -from .hadamard import apply_rht, should_use_rht +from .hadamard import apply_rht from .tensor import ( ScaledTensor, ScaledTensor1x, @@ -590,11 +590,13 @@ class NVFP4Quantizer(Quantizer): q_layout: Quantization axis data_layout: Data layout string (default: "NT") stochastic_rounding_rng_state: RNG state for stochastic rounding, must be of shape (4,) and dtype uint32. If None, stochastic rounding is disabled. + use_rht: Whether to apply Randomized Hadamard Transform (RHT) before quantization. """ scaling_mode: ScalingMode = ScalingMode.NVFP4_1D_SCALING q_layout: QuantizeLayout = QuantizeLayout.ROWWISE_COLWISE data_layout: str = "NT" + use_rht: bool = False stochastic_rounding_rng_state: Optional[jnp.ndarray] = None def __post_init__(self): @@ -603,6 +605,30 @@ def __post_init__(self): ), "NVFP4 quantization must use a q_dtype of float4_e2m1fn" assert self.scaling_mode.is_nvfp4_scaling, "NVFP4Quantizer must use NVFP4 scaling modes" + def tree_flatten(self): + """Flatten the quantizer for JAX tree operations. + + Returns: + Tuple of (children, aux_data) for tree operations + """ + children = (self.stochastic_rounding_rng_state,) + aux_data = (self.q_dtype, self.scaling_mode, self.q_layout, self.data_layout, self.use_rht) + return (children, aux_data) + + @classmethod + def tree_unflatten(cls, aux_data, children): + """Reconstruct a quantizer from its flattened representation. + + Args: + aux_data: Auxiliary data containing quantizer parameters + children: Unused children data + + Returns: + A reconstructed Quantizer instance + """ + stochastic_rounding_rng_state = children[0] + return cls(*aux_data, stochastic_rounding_rng_state=stochastic_rounding_rng_state) + def _apply_stochastic_rounding(self, x): assert ( self.stochastic_rounding_rng_state is not None @@ -688,8 +714,9 @@ def _quantize_func(self, x, is_colwise=False, dq_dtype=None, flatten_axis=-1) -> flatten_axis = x.ndim - flatten_axis x_shape = x.shape - if should_use_rht(self.scaling_mode, is_colwise=is_colwise): - # We only apply RHT for 1D colwise nvfp4 + # We currently only have a single flag 'use_rht' on the quantizer. To avoid an unused rowwise flag, we assume RHT is only used for colwise quantization for now. + use_rht = self.use_rht and is_colwise and self.scaling_mode == ScalingMode.NVFP4_1D_SCALING + if use_rht: x = apply_rht(x) dq_dtype = dq_dtype if dq_dtype is not None else x.dtype @@ -790,6 +817,7 @@ def repeat_to_shape(x, target_shape): scaling_mode=self.scaling_mode, dq_dtype=dq_dtype, flatten_axis=rowwise_flatten_axis, + has_rht_applied=use_rht, ) diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index 2d2d78190f..6c358a044e 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -175,6 +175,7 @@ class ScaledTensor1x(AbstractBaseTensor1x, ScaledTensor): is_colwise: Whether the tensor uses column-wise quantization data_layout: The data_layout specification for the tensor flatten_axis: The quantization axis for the tensor + has_rht_applied: Whether the tensor had the Randomized Hadamard Transform (RHT) applied during quantization """ scale_inv: jnp.ndarray @@ -184,6 +185,7 @@ class ScaledTensor1x(AbstractBaseTensor1x, ScaledTensor): is_colwise: bool data_layout: str flatten_axis: int + has_rht_applied: bool def __post_init__(self): """Validates and adjusts the scale_inv shape after initialization. @@ -243,6 +245,7 @@ def tree_flatten(self): self.is_colwise, self.data_layout, self.flatten_axis, + self.has_rht_applied, ) return (children, aux_data) @@ -314,6 +317,7 @@ def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[st is_colwise=self.is_colwise, data_layout=self.data_layout, flatten_axis=self.flatten_axis, + has_rht_applied=self.has_rht_applied, ) @@ -354,6 +358,7 @@ def __init__( self.group_sizes = group_sizes self.original_shape = original_shape self.group_axis = group_axis + # TODO(Phuong):Handle RHT for grouped quantization once grouped quantization supports NVFP4 super().__init__( data=data, scale_inv=scale_inv, @@ -364,6 +369,7 @@ def __init__( is_colwise=is_colwise, data_layout=data_layout, flatten_axis=flatten_axis, + has_rht_applied=False, ) def __post_init__(self): @@ -515,6 +521,7 @@ def create_1x( group_sizes=None, original_shape=None, group_axis=0, + has_rht_applied=False, ): """Creates a single-scale quantized tensor. @@ -530,6 +537,7 @@ def create_1x( group_sizes: Array of ints containing the size of each group (default: None) original_shape: The original shape of the tensor before grouping (default: None) group_axis: The axis along which grouping is performed (default: 0) + has_rht_applied: Whether the tensor had the Randomized Hadamard Transform (RHT) applied during quantization (default: False) Returns: A ScaledTensor1x or GroupedScaledTensor1x instance depending on whether group_sizes is provided @@ -593,6 +601,7 @@ def create_1x( is_colwise=is_colwise, data_layout=data_layout, flatten_axis=flatten_axis, + has_rht_applied=has_rht_applied, ) @staticmethod @@ -610,6 +619,8 @@ def create_2x( group_sizes=None, original_shape=None, group_axis=0, + rowwise_has_rht_applied=False, + colwise_has_rht_applied=False, ): """Creates a double-scale quantized tensor. @@ -626,6 +637,8 @@ def create_2x( group_sizes: Array containing the size of each group (default: None) original_shape: The original shape of the tensor before grouping (default: None) group_axis: The axis along which grouping is performed (default: 0) + rowwise_has_rht_applied: Whether the row-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) + colwise_has_rht_applied: Whether the column-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) Returns: A ScaledTensor2x instance @@ -648,6 +661,7 @@ def create_2x( group_sizes=group_sizes, original_shape=original_shape, group_axis=group_axis, + has_rht_applied=rowwise_has_rht_applied, ) colwise_tensor = ScaledTensorFactory.create_1x( colwise_data, @@ -661,6 +675,7 @@ def create_2x( group_sizes=group_sizes, original_shape=original_shape, group_axis=group_axis, + has_rht_applied=colwise_has_rht_applied, ) return ScaledTensor2x(rowwise_tensor, colwise_tensor) @@ -680,6 +695,8 @@ def create( group_sizes: jnp.ndarray = None, original_shape: Tuple[int] = None, group_axis: int = 0, + rowwise_has_rht_applied: bool = False, + colwise_has_rht_applied: bool = False, ): """Creates a scaled tensor based on the quantization axis. @@ -696,10 +713,14 @@ def create( group_sizes: Array containing the size of each group (default: None) original_shape: The original shape of the tensor before grouping (default: None) group_axis: The axis along which grouping is performed (default: 0) + rowwise_has_rht_applied: Whether the row-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) + colwise_has_rht_applied: Whether the col-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) Returns: Either a ScaledTensor1x or ScaledTensor2x instance depending on q_layout """ + assert not rowwise_has_rht_applied, "RHT is not supported for rowwise quantization yet" + if q_layout == QuantizeLayout.ROWWISE_COLWISE: return ScaledTensorFactory.create_2x( data, @@ -715,6 +736,8 @@ def create( group_sizes=group_sizes, original_shape=original_shape, group_axis=group_axis, + rowwise_has_rht_applied=rowwise_has_rht_applied, + colwise_has_rht_applied=colwise_has_rht_applied, ) is_colwise = q_layout == QuantizeLayout.COLWISE @@ -731,6 +754,7 @@ def create( group_sizes=group_sizes, original_shape=original_shape, group_axis=group_axis, + has_rht_applied=colwise_has_rht_applied, ) return ScaledTensorFactory.create_1x( @@ -745,6 +769,7 @@ def create( group_sizes=group_sizes, original_shape=original_shape, group_axis=group_axis, + has_rht_applied=rowwise_has_rht_applied, ) diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 8eeaca4cc8..adb67e358f 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -238,6 +238,19 @@ def num_of_devices(): return len(jax.devices()) +def get_num_devices_in_mesh(mesh=None): + """ + Get the number of devices in the given mesh. + If the mesh is None, it would be replaced + by the global mesh. + """ + if mesh is None: + mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + if mesh.empty: + return 1 + return np.prod(list(mesh.shape.values())) + + def get_mesh_axis_size(axis, mesh=None): """ Get the axis size of the given mesh. From 2ac3c16876fe3bcd4866f8a62251802ea5530888 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 22 Oct 2025 13:11:31 -0700 Subject: [PATCH 017/521] [JAX] Defer TE/JAX cublas shape check on fp8 gemms until lowering (#2292) Defer cublas check on fp8 gemms until lowering Signed-off-by: Jeremy Berchtold --- transformer_engine/jax/cpp_extensions/gemm.py | 46 +++++++++---------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 778f77c0d5..72bee251c4 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -470,29 +470,6 @@ def _dims_are_consecutive(dims): f" LHS dtype != RHS dtype, lhs.dtype={lhs.dtype}, rhs.dtype={rhs.dtype}" ) - lhs_axis_boundary = get_lhs_axis_boundary(lhs_contracting_dims, lhs_is_transposed) - lhs_contracting_size = ( - reduce(operator.mul, lhs.shape[lhs_axis_boundary:]) - if lhs_is_transposed - else reduce(operator.mul, lhs.shape[:lhs_axis_boundary]) - ) - assert_cublas_requirements( - scaling_mode, - lhs_contracting_size, - "LHS", - ) - rhs_axis_boundary = get_rhs_axis_boundary(rhs_contracting_dims, rhs_is_transposed) - rhs_contracting_size = ( - reduce(operator.mul, rhs.shape[:rhs_axis_boundary]) - if rhs_is_transposed - else reduce(operator.mul, rhs.shape[rhs_axis_boundary:]) - ) - assert_cublas_requirements( - scaling_mode, - rhs_contracting_size, - "RHS", - ) - # Determine output shape and dtype assert ( dtypes.canonicalize_dtype(out_dtype).itemsize > 1 @@ -601,6 +578,29 @@ def lowering( (lhs_aval.ndim, rhs_aval.ndim), (lhs_cdims, rhs_cdims) ) + lhs_axis_boundary = get_lhs_axis_boundary(lhs_cdims, lhs_transposed) + lhs_contracting_size = ( + reduce(operator.mul, lhs_aval.shape[lhs_axis_boundary:]) + if lhs_transposed + else reduce(operator.mul, lhs_aval.shape[:lhs_axis_boundary]) + ) + assert_cublas_requirements( + scaling_mode, + lhs_contracting_size, + "LHS", + ) + rhs_axis_boundary = get_rhs_axis_boundary(rhs_cdims, rhs_transposed) + rhs_contracting_size = ( + reduce(operator.mul, rhs_aval.shape[:rhs_axis_boundary]) + if rhs_transposed + else reduce(operator.mul, rhs_aval.shape[rhs_axis_boundary:]) + ) + assert_cublas_requirements( + scaling_mode, + rhs_contracting_size, + "RHS", + ) + args = (lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alpha, beta) kwargs = { "scaling_mode": int(scaling_mode.value), From 66acb8e97baa095c8a6e0001bc27aca4f6a8574e Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 22 Oct 2025 20:33:49 -0400 Subject: [PATCH 018/521] Include TE core headers in final build (#2291) Include TE core headers in build Signed-off-by: Kirthi Shankar Sivamani --- MANIFEST.in | 1 + 1 file changed, 1 insertion(+) create mode 100644 MANIFEST.in diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000000..c34025772a --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +recursive-include transformer_engine/common/include *.* From eb34783cb774438a4367e45d478744d2799e1a7f Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 22 Oct 2025 22:31:08 -0700 Subject: [PATCH 019/521] Overhaul the compilation for the arch-specific features (#2279) * Added sm_120f to the build Signed-off-by: Przemek Tredak * Change the arch specific handling Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * Support for CUDA<12.9 Signed-off-by: Przemek Tredak * Moved through the rest of the files Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * Common cases Signed-off-by: Przemek Tredak * Remove pure 100 from the list Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * CMake changes, (not yet working) Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * Do not pass the arch-specific thing from build_tools Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * Moved some of the files to arch-specific compilation Signed-off-by: Przemek Tredak * Fix and also changing the order of compilation to hopefully get the compilation time lower Signed-off-by: Przemek Tredak * Fix for the files overwriting custom compile properties Signed-off-by: Przemek Tredak * Actually make this whole thing work Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add space to the error message Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak * Apply suggestions from code review Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak * Fixes from review Signed-off-by: Przemek Tredak * Changing the naming to be more intuitive Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add missing cassert include for device-side asserts 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: Przemyslaw Tredak Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> --- build_tools/utils.py | 6 +- transformer_engine/common/CMakeLists.txt | 206 +++++++++--- .../hadamard_transform_cast_fusion.cu | 27 +- ...quantize_transpose_vector_blockwise_fp4.cu | 76 ++--- .../common/util/nvfp4_transpose.cuh | 290 ++++++++-------- transformer_engine/common/util/ptx.cuh | 310 +++++++++++++++--- transformer_engine/common/utils.cuh | 1 + 7 files changed, 610 insertions(+), 306 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index 296f928b71..395b41261b 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -257,11 +257,9 @@ def cuda_archs() -> str: if archs is None: version = cuda_version() if version >= (13, 0): - archs = "75;80;89;90;100;100a;103a;120" - elif version >= (12, 9): - archs = "70;80;89;90;100;100a;103a;120" + archs = "75;80;89;90;100;120" elif version >= (12, 8): - archs = "70;80;89;90;100;100a;120" + archs = "70;80;89;90;100;120" else: archs = "70;80;89;90" return archs diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e6be47686a..175abd3530 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -5,15 +5,6 @@ cmake_minimum_required(VERSION 3.21) # Language options -if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) - if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.0) - set(CMAKE_CUDA_ARCHITECTURES 75 80 89 90 100 120) - elseif (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8) - set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90 100 120) - else () - set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90) - endif() -endif() set(CMAKE_CXX_STANDARD 17) set(CMAKE_CUDA_STANDARD 17) set(CMAKE_CUDA_STANDARD_REQUIRED ON) @@ -30,8 +21,62 @@ project(transformer_engine LANGUAGES CUDA CXX) # CUDA Toolkit find_package(CUDAToolkit REQUIRED) -if (CUDAToolkit_VERSION VERSION_LESS 12.0) - message(FATAL_ERROR "CUDA 12.0+ is required, but found CUDA ${CUDAToolkit_VERSION}") +if (CUDAToolkit_VERSION VERSION_LESS 12.1) + message(FATAL_ERROR "CUDA 12.1+ is required, but found CUDA ${CUDAToolkit_VERSION}") +endif() + +# Process GPU architectures +if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) + if (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 13.0) + set(CMAKE_CUDA_ARCHITECTURES 75 80 89 90 100 120) + elseif (CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.8) + set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90 100 120) + else () + set(CMAKE_CUDA_ARCHITECTURES 70 80 89 90) + endif() +endif() + +# Process CMAKE_CUDA_ARCHITECTURES to separate generic and specific architectures +set(NVTE_GENERIC_ARCHS) +set(NVTE_SPECIFIC_ARCHS) + +# Check for architecture 100 +list(FIND CMAKE_CUDA_ARCHITECTURES "100" arch_100_index) +if(NOT arch_100_index EQUAL -1) + list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "100") + list(APPEND NVTE_GENERIC_ARCHS "100") + list(APPEND NVTE_SPECIFIC_ARCHS "100a") + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9) + list(APPEND NVTE_SPECIFIC_ARCHS "103a") + endif() +endif() + +# Check for architecture 101 (if we see this we are in toolkit <= 12.9) +list(FIND CMAKE_CUDA_ARCHITECTURES "101" arch_101_index) +if(NOT arch_101_index EQUAL -1) + list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "101") + list(APPEND NVTE_GENERIC_ARCHS "101") + list(APPEND NVTE_SPECIFIC_ARCHS "101a") +endif() + +# Check for architecture 110 (if we see this we are in toolkit >= 13.0) +list(FIND CMAKE_CUDA_ARCHITECTURES "110" arch_110_index) +if(NOT arch_110_index EQUAL -1) + list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "110") + list(APPEND NVTE_GENERIC_ARCHS "110") + list(APPEND NVTE_SPECIFIC_ARCHS "110f") +endif() + +# Check for architecture 120 +list(FIND CMAKE_CUDA_ARCHITECTURES "120" arch_120_index) +if(NOT arch_120_index EQUAL -1) + list(REMOVE_ITEM CMAKE_CUDA_ARCHITECTURES "120") + list(APPEND NVTE_GENERIC_ARCHS "120") + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.9) + list(APPEND NVTE_SPECIFIC_ARCHS "120f") + else() + list(APPEND NVTE_SPECIFIC_ARCHS "120a") + endif() endif() # cuDNN frontend API @@ -78,9 +123,28 @@ endif() # Configure Transformer Engine library include_directories(${PROJECT_SOURCE_DIR}/..) set(transformer_engine_SOURCES) -list(APPEND transformer_engine_SOURCES +set(transformer_engine_cpp_sources) +set(transformer_engine_cuda_sources) +set(transformer_engine_cuda_arch_specific_sources) + +list(APPEND transformer_engine_cpp_sources cudnn_utils.cpp transformer_engine.cpp + fused_attn/fused_attn.cpp + gemm/config.cpp + normalization/common.cpp + normalization/layernorm/ln_api.cpp + normalization/rmsnorm/rmsnorm_api.cpp + util/cuda_driver.cpp + util/cuda_nvml.cpp + util/cuda_runtime.cpp + util/multi_stream.cpp + util/rtc.cpp + comm_gemm_overlap/userbuffers/ipcsocket.cc + comm_gemm_overlap/userbuffers/userbuffers-host.cpp + comm_gemm_overlap/comm_gemm_overlap.cpp) + +list(APPEND transformer_engine_cuda_sources common.cu multi_tensor/adam.cu multi_tensor/compute_scale.cu @@ -92,40 +156,23 @@ list(APPEND transformer_engine_SOURCES transpose/cast_transpose_fusion.cu transpose/transpose_fusion.cu transpose/multi_cast_transpose.cu - transpose/quantize_transpose_square_blockwise.cu transpose/quantize_transpose_vector_blockwise.cu transpose/swap_first_dims.cu - transpose/quantize_transpose_vector_blockwise_fp4.cu - activation/gelu.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 fused_attn/fused_attn_f16_arbitrary_seqlen.cu - activation/relu.cu - activation/swiglu.cu fused_attn/fused_attn_fp8.cu - fused_attn/fused_attn.cpp fused_attn/utils.cu - gemm/config.cpp gemm/cublaslt_gemm.cu - gemm/cutlass_grouped_gemm.cu - normalization/common.cpp - normalization/layernorm/ln_api.cpp normalization/layernorm/ln_bwd_semi_cuda_kernel.cu normalization/layernorm/ln_fwd_cuda_kernel.cu - normalization/rmsnorm/rmsnorm_api.cpp normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu permutation/permutation.cu - util/cast.cu util/padding.cu - util/cuda_driver.cpp - util/cuda_nvml.cpp - util/cuda_runtime.cpp - util/multi_stream.cpp - util/rtc.cpp swizzle/swizzle.cu swizzle/swizzle_block_scaling.cu fused_softmax/scaled_masked_softmax.cu @@ -139,12 +186,58 @@ list(APPEND transformer_engine_SOURCES recipe/delayed_scaling.cu recipe/fp8_block_scaling.cu recipe/nvfp4.cu + comm_gemm_overlap/userbuffers/userbuffers.cu) + +list(APPEND transformer_engine_cuda_arch_specific_sources + gemm/cutlass_grouped_gemm.cu + util/cast.cu + activation/gelu.cu + activation/relu.cu + activation/swiglu.cu + transpose/quantize_transpose_square_blockwise.cu + transpose/quantize_transpose_vector_blockwise_fp4.cu hadamard_transform/hadamard_transform.cu - hadamard_transform/hadamard_transform_cast_fusion.cu - comm_gemm_overlap/userbuffers/ipcsocket.cc - comm_gemm_overlap/userbuffers/userbuffers-host.cpp - comm_gemm_overlap/userbuffers/userbuffers.cu - comm_gemm_overlap/comm_gemm_overlap.cpp) + hadamard_transform/hadamard_transform_cast_fusion.cu) + +# Compiling the files with the worst compilation time first to hopefully overlap +# better with the faster-compiling cpp files +list(APPEND transformer_engine_SOURCES ${transformer_engine_cuda_arch_specific_sources} + ${transformer_engine_cuda_sources} + ${transformer_engine_cpp_sources}) + +# Set compile options for CUDA sources with generic architectures +foreach(cuda_source IN LISTS transformer_engine_cuda_sources) + set(arch_compile_options) + foreach(arch IN LISTS NVTE_GENERIC_ARCHS) + list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}") + endforeach() + + if(arch_compile_options) + set_property( + SOURCE ${cuda_source} + APPEND + PROPERTY + COMPILE_OPTIONS ${arch_compile_options} + ) + endif() +endforeach() + +# Set compile options for CUDA sources with specific architectures +foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) + set(arch_compile_options) + foreach(arch IN LISTS NVTE_SPECIFIC_ARCHS) + list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}") + endforeach() + + if(arch_compile_options) + set_property( + SOURCE ${cuda_source} + APPEND + PROPERTY + COMPILE_OPTIONS ${arch_compile_options} + ) + endif() +endforeach() if (NVTE_WITH_CUBLASMP) list(APPEND transformer_engine_SOURCES @@ -249,28 +342,35 @@ target_include_directories(transformer_engine PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/string_headers") # Compiler options -set_source_files_properties(fused_softmax/scaled_masked_softmax.cu - fused_softmax/scaled_upper_triang_masked_softmax.cu - fused_softmax/scaled_aligned_causal_masked_softmax.cu - multi_tensor/adam.cu - multi_tensor/compute_scale.cu - multi_tensor/l2norm.cu - multi_tensor/scale.cu - multi_tensor/sgd.cu - fused_attn/flash_attn.cu - fused_attn/context_parallel.cu - fused_attn/kv_cache.cu - PROPERTIES - COMPILE_OPTIONS "--use_fast_math") +set(nvte_sources_with_fast_math) +list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu + fused_softmax/scaled_upper_triang_masked_softmax.cu + fused_softmax/scaled_aligned_causal_masked_softmax.cu + multi_tensor/adam.cu + multi_tensor/compute_scale.cu + multi_tensor/l2norm.cu + multi_tensor/scale.cu + multi_tensor/sgd.cu + fused_attn/flash_attn.cu + fused_attn/context_parallel.cu + fused_attn/kv_cache.cu) + option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --use_fast_math option" OFF) if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) - set_source_files_properties(activation/gelu.cu - activation/relu.cu - activation/swiglu.cu - util/cast.cu - PROPERTIES - COMPILE_OPTIONS "--use_fast_math") + list(APPEND nvte_sources_with_fast_math activation/gelu.cu + activation/relu.cu + activation/swiglu.cu + util/cast.cu) endif() + +foreach(cuda_source IN LISTS nvte_sources_with_fast_math) + set_property( + SOURCE ${cuda_source} + APPEND + PROPERTY + COMPILE_OPTIONS "--use_fast_math") +endforeach() + set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr") set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -O3") 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 ce191b5ffd..263a32623e 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -97,22 +97,23 @@ cutlass::Array StochasticNumericConverterBase(cutlass::Array const &input, cutlass::Array const &rbits) { using result_type = cutlass::Array; result_type output; -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - auto output_ptr = reinterpret_cast(&output); - asm volatile( \ - "{\n" \ - "cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\n" \ - "cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\n" \ - "}" \ - : "=h"(output_ptr[0]), + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + auto output_ptr = reinterpret_cast(&output); + asm volatile( \ + "{\n" \ + "cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\n" \ + "cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\n" \ + "}" \ + : "=h"(output_ptr[0]), "=h"(output_ptr[1]) - : "f"(input[0]), "f"(input[1]), "f"(input[2]), "f"(input[3]), + : "f"(input[0]), "f"(input[1]), "f"(input[2]), "f"(input[3]), "f"(input[4]), "f"(input[5]), "f"(input[6]), "f"(input[7]), "r"(rbits[0]), "r"(rbits[1])); -#else - NVTE_DEVICE_ERROR("FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + } else { + NVTE_DEVICE_ERROR("FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } return output; } 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 eced2c4bb6..fed18c51f8 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -264,48 +264,50 @@ __device__ __forceinline__ size_t scale_factor_swizzled_offset(size_t row_idx, s __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_rounding( const float2 in01, const float2 in23, const uint32_t rbits) { -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - uint16_t out_4x; - asm volatile( - "{\n" - "cvt.rs.satfinite.e2m1x4.f32 %0, {%3, %4, %1, %2}, %5; \n\t" - "}" - : "=h"(out_4x) - : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x), "r"(rbits)); - return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x); -#else - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - uint16_t dummy = 0; - return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + uint16_t out_4x; + asm volatile( + "{\n" + "cvt.rs.satfinite.e2m1x4.f32 %0, {%3, %4, %1, %2}, %5; \n\t" + "}" + : "=h"(out_4x) + : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x), "r"(rbits)); + return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt.rs PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + uint16_t dummy = 0; + return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); + } } __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_rn(const float2 in01, const float2 in23, const uint32_t rbits) { -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - // NOTE: rbits unused for rn. - uint32_t out_4x; // Only need 16 bit. Using 32 bit container for packing. - asm volatile( - "{\n" - ".reg.b8 f0; \n\t" - ".reg.b8 f1; \n\t" - "cvt.rn.satfinite.e2m1x2.f32 f0, %1, %2;\n\t" - "cvt.rn.satfinite.e2m1x2.f32 f1, %3, %4;\n\t" - "mov.b32 %0, {f0, f1, f0, f1};\n\t" - "}" - : "=r"(out_4x) - : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x)); - return reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x)[0]; -#else - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - uint16_t dummy = 0; - return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + constexpr bool has_fp4 = ARCH_BLACKWELL_FAMILY; + if constexpr (has_fp4) { + // NOTE: rbits unused for rn. + uint32_t out_4x; // Only need 16 bit. Using 32 bit container for packing. + asm volatile( + "{\n" + ".reg.b8 f0; \n\t" + ".reg.b8 f1; \n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, %1, %2;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, %3, %4;\n\t" + "mov.b32 %0, {f0, f1, f0, f1};\n\t" + "}" + : "=r"(out_4x) + : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x)); + return reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x)[0]; + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + uint16_t dummy = 0; + return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy); + } } template diff --git a/transformer_engine/common/util/nvfp4_transpose.cuh b/transformer_engine/common/util/nvfp4_transpose.cuh index 712b557c5d..45fa29f0e9 100644 --- a/transformer_engine/common/util/nvfp4_transpose.cuh +++ b/transformer_engine/common/util/nvfp4_transpose.cuh @@ -15,10 +15,9 @@ #include #include -#if CUDA_VERSION > 12080 +#if FP4_TYPE_SUPPORTED #include -#endif // CUDA_VERSION > 12080 - +#endif // FP4_TYPE_SUPPORTED #include #include "../common.h" @@ -30,7 +29,7 @@ namespace transformer_engine { -#if CUDA_VERSION > 12080 +#if FP4_TYPE_SUPPORTED namespace nvfp4_transpose { using RNG = decltype(curanddx::Generator() + curanddx::PhiloxRounds<10>() + @@ -152,89 +151,89 @@ __device__ __forceinline__ uint32_t get_rbits(RNG &rng, uint4 &random_uint4, int return rbits; } -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding( const uint64_t in_4x, const float2 scale, const uint32_t rbits) { uint16_t out_4x = 0; -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b16 v0_bf16; \n\t" - ".reg.b16 v1_bf16; \n\t" - ".reg.b16 v2_bf16; \n\t" - ".reg.b16 v3_bf16; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" - "cvt.f32.bf16 v0, v0_bf16; \n\t" - "cvt.f32.bf16 v1, v1_bf16; \n\t" - "cvt.f32.bf16 v2, v2_bf16; \n\t" - "cvt.f32.bf16 v3, v3_bf16; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %3; \n\t" // mind the shuffled elements order - "}" - : "=h"(out_4x) - : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); -#else - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b16 v0_bf16; \n\t" + ".reg.b16 v1_bf16; \n\t" + ".reg.b16 v2_bf16; \n\t" + ".reg.b16 v3_bf16; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" + "cvt.f32.bf16 v0, v0_bf16; \n\t" + "cvt.f32.bf16 v1, v1_bf16; \n\t" + "cvt.f32.bf16 v2, v2_bf16; \n\t" + "cvt.f32.bf16 v3, v3_bf16; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %3; \n\t" // mind the shuffled elements order + "}" + : "=h"(out_4x) + : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } return *reinterpret_cast(&out_4x); } __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_rn(const uint64_t in_4x, const float2 scale, const uint32_t rbits) { - // NOTE: rbits unused for rn. + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; uint32_t out_4x = 0; // Only need 16 bit. Using 32 bit container for packing. -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b16 v0_bf16; \n\t" - ".reg.b16 v1_bf16; \n\t" - ".reg.b16 v2_bf16; \n\t" - ".reg.b16 v3_bf16; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - ".reg.b8 f0; \n\t" - ".reg.b8 f1; \n\t" - "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" - "cvt.f32.bf16 v0, v0_bf16; \n\t" - "cvt.f32.bf16 v1, v1_bf16; \n\t" - "cvt.f32.bf16 v2, v2_bf16; \n\t" - "cvt.f32.bf16 v3, v3_bf16; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" - "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" - "mov.b32 %0, {f0, f1, f0, f1};\n\t" - "}" - : "=r"(out_4x) - : "l"(in_4x), "l"(reinterpret_cast(scale))); -#else - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + if constexpr (is_blackwell) { + // NOTE: rbits unused for rn. + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b16 v0_bf16; \n\t" + ".reg.b16 v1_bf16; \n\t" + ".reg.b16 v2_bf16; \n\t" + ".reg.b16 v3_bf16; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + ".reg.b8 f0; \n\t" + ".reg.b8 f1; \n\t" + "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" + "cvt.f32.bf16 v0, v0_bf16; \n\t" + "cvt.f32.bf16 v1, v1_bf16; \n\t" + "cvt.f32.bf16 v2, v2_bf16; \n\t" + "cvt.f32.bf16 v3, v3_bf16; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" + "mov.b32 %0, {f0, f1, f0, f1};\n\t" + "}" + : "=r"(out_4x) + : "l"(in_4x), "l"(reinterpret_cast(scale))); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } return reinterpret_cast(&out_4x)[0]; } @@ -252,34 +251,35 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x(const uint64_t in_4x __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding( const float2 in01, const float2 in23, const float2 scale, const uint32_t rbits) { uint16_t out_4x = 0; -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - "mov.b64 {v0, v1} , %1; \n\t" - "mov.b64 {v2, v3} , %2; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %4; \n\t" // mind the shuffled elements order - "}" - : "=h"(out_4x) - : "l"(reinterpret_cast(in01)), - "l"(reinterpret_cast(in23)), - "l"(reinterpret_cast(scale)), "r"(rbits)); -#else - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + "mov.b64 {v0, v1} , %1; \n\t" + "mov.b64 {v2, v3} , %2; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %4; \n\t" // mind the shuffled elements order + "}" + : "=h"(out_4x) + : "l"(reinterpret_cast(in01)), + "l"(reinterpret_cast(in23)), + "l"(reinterpret_cast(scale)), "r"(rbits)); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } return *reinterpret_cast(&out_4x); } @@ -287,40 +287,41 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_rn(const float2 const float2 in23, const float2 scale, const uint32_t rbits) { - // NOTE: rbits unused for rn. + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; uint32_t out_4x = 0; // Only need 16 bit. Using 32 bit container for packing. -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - ".reg.b8 f0; \n\t" - ".reg.b8 f1; \n\t" - "mov.b64 {v0, v1} , %1; \n\t" - "mov.b64 {v2, v3} , %2; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" - "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" - "mov.b32 %0, {f0, f1, f0, f1};\n\t" - "}" - : "=r"(out_4x) - : "l"(reinterpret_cast(in01)), - "l"(reinterpret_cast(in23)), - "l"(reinterpret_cast(scale))); -#else - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); -#endif // CUDA_ARCH_HAS_FEATURE_SM10X_ALL + if constexpr (is_blackwell) { + // NOTE: rbits unused for rn. + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + ".reg.b8 f0; \n\t" + ".reg.b8 f1; \n\t" + "mov.b64 {v0, v1} , %1; \n\t" + "mov.b64 {v2, v3} , %2; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" + "mov.b32 %0, {f0, f1, f0, f1};\n\t" + "}" + : "=r"(out_4x) + : "l"(reinterpret_cast(in01)), + "l"(reinterpret_cast(in23)), + "l"(reinterpret_cast(scale))); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } return reinterpret_cast(&out_4x)[0]; } @@ -335,8 +336,6 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x(const float2 in01, c } } -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - template __global__ void __launch_bounds__(THREADS_NUM) @@ -1380,18 +1379,13 @@ __global__ void __launch_bounds__(THREADS_NUM) #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } } // namespace nvfp4_transpose -#endif // CUDA_VERSION > 12080 - -// Compile-time flag to choose kernel variant -#ifndef USE_2D_NVFP4_KERNEL -#define USE_2D_NVFP4_KERNEL 0 -#endif +#endif // FP4_TYPE_SUPPORTED template void nvfp4_quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, const QuantizationConfig *quant_config, cudaStream_t stream) { -#if CUDA_VERSION > 12080 +#if FP4_TYPE_SUPPORTED bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to @@ -1509,7 +1503,7 @@ void nvfp4_quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *o });); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); -#endif // CUDA_VERSION > 12080 +#endif // FP4_TYPE_SUPPORTED } } // namespace transformer_engine diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 85717afdf2..aeac2b4a2c 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -18,44 +18,165 @@ #include #endif // CUDA_VERSION >= 12080 +#include "common/utils.cuh" + namespace transformer_engine { + namespace ptx { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +template +struct ArchSpecific { + constexpr static int id = N * 10; + + template + constexpr static bool compatible() { + if constexpr (CurrentArch == id) { + static_assert(ArchSpecific == CurrentArch, + "Compiled for the generic architecture, while utilizing arch-specific " + "features. Please compile for smXXXa architecture instead of smXXX " + "architecture."); + return true; + } else { + return false; + } + } +}; + +template +struct FamilySpecific { + constexpr static int id = N * 10; + + template + constexpr static bool compatible() { + if constexpr ((CurrentArch / 100) == (id / 100)) { + static_assert(FamilySpecific == CurrentArch, + "Compiled for the generic architecture, while utilizing family-specific " + "features. Please compile for smXXXf architecture instead of smXXX " + "architecture."); + return true; + } else { + return false; + } + } +}; + +template +constexpr bool is_supported_arch() { + if constexpr (T::template compatible()) { + return true; + } else if constexpr (sizeof...(U) != 0) { + return is_supported_arch(); + } else { + return false; + } +} + +#if CUDA_VERSION < 12090 +#if __CUDA_ARCH_HAS_FEATURE__(SM90_ALL) +#define __CUDA_ARCH_SPECIFIC__ 900 +#define __CUDA_ARCH_FAMILY_SPECIFIC__ 900 +#endif +#if __CUDA_ARCH_HAS_FEATURE__(SM100_ALL) +#define __CUDA_ARCH_SPECIFIC__ 1000 +#define __CUDA_ARCH_FAMILY_SPECIFIC__ 1000 +#endif +#if __CUDA_ARCH_HAS_FEATURE__(SM101_ALL) +#define __CUDA_ARCH_SPECIFIC__ 1010 +#define __CUDA_ARCH_FAMILY_SPECIFIC__ 1010 +#endif +#if __CUDA_ARCH_HAS_FEATURE__(SM120_ALL) +#define __CUDA_ARCH_SPECIFIC__ 1200 +#define __CUDA_ARCH_FAMILY_SPECIFIC__ 1200 +#endif +#endif + +#ifdef __CUDA_ARCH__ +#define __NVTE_CURRENT_ARCH__ constexpr int current_arch = __CUDA_ARCH__; +#else +#define __NVTE_CURRENT_ARCH__ constexpr int current_arch = 0; +#endif + +#ifdef __CUDA_ARCH_SPECIFIC__ +#define __NVTE_ARCH_SPECIFIC__ constexpr int ArchSpecific = __CUDA_ARCH_SPECIFIC__; +#else +#define __NVTE_ARCH_SPECIFIC__ constexpr int ArchSpecific = 0; +#endif + +#ifdef __CUDA_ARCH_FAMILY_SPECIFIC__ +#define __NVTE_ARCH_FAMILY_SPECIFIC__ constexpr int FamilySpecific = __CUDA_ARCH_FAMILY_SPECIFIC__; +#else +#define __NVTE_ARCH_FAMILY_SPECIFIC__ constexpr int FamilySpecific = 0; +#endif + +#define NVTE_CUDA_ARCH_MATCHES(...) \ + [&] { \ + __NVTE_CURRENT_ARCH__ \ + __NVTE_ARCH_SPECIFIC__ \ + __NVTE_ARCH_FAMILY_SPECIFIC__ \ + return transformer_engine::ptx::is_supported_arch(); \ + }(); + +#define ARCH_BLACKWELL_FAMILY \ + NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>, ptx::FamilySpecific<110>, \ + ptx::FamilySpecific<120>) +#define ARCH_HAS_STOCHASTIC_ROUNDING \ + NVTE_CUDA_ARCH_MATCHES(ptx::ArchSpecific<100>, ptx::ArchSpecific<103>) // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-init __device__ __forceinline__ void mbarrier_init(uint64_t *mbar, const uint32_t count) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.init.shared.b64 [%0], %1;" ::"r"(mbar_ptr), "r"(count) : "memory"); +#else + NVTE_DEVICE_ERROR("mbarrier_init is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-inval __device__ __forceinline__ void mbarrier_invalid(uint64_t *mbar) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.inval.shared.b64 [%0];" ::"r"(mbar_ptr) : "memory"); +#else + NVTE_DEVICE_ERROR("mbarrier_invalid is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive __device__ __forceinline__ void mbarrier_arrive(uint64_t *mbar) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.arrive.shared.b64 _, [%0];" ::"r"(mbar_ptr) : "memory"); +#else + NVTE_DEVICE_ERROR("mbarrier_arrive is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-arrive __device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t *mbar, const uint32_t tx_count) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); asm volatile("mbarrier.arrive.expect_tx.shared.b64 _, [%0], %1;" ::"r"(mbar_ptr), "r"(tx_count) : "memory"); +#else + NVTE_DEVICE_ERROR("mbarrier_arrive_expect_tx is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile("fence.mbarrier_init.release.cluster;"); +#else + NVTE_DEVICE_ERROR("fence_mbarrier_init_release_cluster is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor // global -> shared::cluster __device__ __forceinline__ void cp_async_bulk_tensor_1d_global_to_shared( uint64_t *dst_shmem, const uint64_t *src_global_ptr, const uint32_t size, uint64_t *mbar) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t dst_shmem_ptr = __cvta_generic_to_shared(dst_shmem); uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); // triggers async copy, i.e. the thread continues until wait() on mbarrier @@ -67,6 +188,9 @@ __device__ __forceinline__ void cp_async_bulk_tensor_1d_global_to_shared( ".mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"(dst_shmem_ptr), "l"(src_global_ptr), "r"(size), "r"(mbar_ptr) : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_1d_global_to_shared is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor @@ -74,6 +198,7 @@ __device__ __forceinline__ void cp_async_bulk_tensor_1d_global_to_shared( __device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared( uint64_t *dst_shmem, const uint64_t *tensor_map_ptr, const uint32_t offset_x, const uint32_t offset_y, uint64_t *mbar) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t dst_shmem_ptr = __cvta_generic_to_shared(dst_shmem); uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); // triggers async copy, i.e. the thread continues until wait() on mbarrier @@ -85,9 +210,13 @@ __device__ __forceinline__ void cp_async_bulk_tensor_2d_global_to_shared( ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" ::"r"(dst_shmem_ptr), "l"(tensor_map_ptr), "r"(offset_x), "r"(offset_y), "r"(mbar_ptr) : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_2d_global_to_shared is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ bool mbarrier_try_wait_parity(uint32_t mbar_ptr, const uint32_t parity) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t waitComplete; asm volatile( "{\n\t .reg .pred P_OUT; \n\t" @@ -98,15 +227,21 @@ __device__ __forceinline__ bool mbarrier_try_wait_parity(uint32_t mbar_ptr, cons : "r"(mbar_ptr), "r"(parity) : "memory"); return static_cast(waitComplete); +#else + NVTE_DEVICE_ERROR("mbarrier_try_wait_parity is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + return true; } __device__ __forceinline__ void mbarrier_wait_parity(uint64_t *mbar, const uint32_t parity) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); while (!mbarrier_try_wait_parity(mbar_ptr, parity)) { } -} - +#else + NVTE_DEVICE_ERROR("mbarrier_wait_parity is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} constexpr uint32_t FP32_MANTISSA_BITS = 23; constexpr uint32_t FP32_EXPONENT_BIAS = 127; @@ -121,55 +256,53 @@ __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { return __int_as_float(biased_exp << FP32_MANTISSA_BITS); } -#define CUDA_ARCH_HAS_FEATURE_SM10X_ALL \ - ((__CUDA_ARCH_HAS_FEATURE__(SM100_ALL)) || (__CUDA_ARCH_HAS_FEATURE__(SM101_ALL)) || \ - (__CUDA_ARCH_HAS_FEATURE__(SM103_ALL))) - __device__ __forceinline__ e8m0_t float_to_e8m0(float val) { -#if CUDA_ARCH_HAS_FEATURE_SM10X_ALL - - uint16_t out; - asm volatile( - "{\n" - "cvt.rp.satfinite.ue8m0x2.f32 %0, 0.0, %1;\n" - "}" - : "=h"(out) - : "f"(val)); - return *reinterpret_cast(&out); -#else - // TODO: nan/inf needs to be set for any value - // of nan/inf in input not just amax. - if (isnan(val)) { - return 0xFF; - } - if (isinf(val)) { - return 0xFE; - } - if (val == 0.0f) { - return 0x00; - } - uint32_t val_u32 = *reinterpret_cast(&val); - e8m0_t exponent = (val_u32 >> FP32_MANTISSA_BITS); - uint32_t mantissa = val_u32 & 0x7FFFFF; - // Round up exponent and deal with satfinite. - if ((mantissa > 0 && exponent != 0xFE) && !(exponent == 0 && mantissa <= 0x400000)) { - ++exponent; + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + uint16_t out; + asm volatile( + "{\n" + "cvt.rp.satfinite.ue8m0x2.f32 %0, 0.0, %1;\n" + "}" + : "=h"(out) + : "f"(val)); + return *reinterpret_cast(&out); + } else { + // TODO: nan/inf needs to be set for any value + // of nan/inf in input not just amax. + if (isnan(val)) { + return 0xFF; + } + if (isinf(val)) { + return 0xFE; + } + if (val == 0.0f) { + return 0x00; + } + uint32_t val_u32 = *reinterpret_cast(&val); + e8m0_t exponent = (val_u32 >> FP32_MANTISSA_BITS); + uint32_t mantissa = val_u32 & 0x7FFFFF; + // Round up exponent and deal with satfinite. + if ((mantissa > 0 && exponent != 0xFE) && !(exponent == 0 && mantissa <= 0x400000)) { + ++exponent; + } + return exponent; } - return exponent; -#endif } -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor // shared::cta -> global __device__ __forceinline__ void cp_async_bulk_tensor_1d_shared_to_global(uint64_t *dst_global_ptr, const uint64_t *src_shmem, const uint32_t size) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t src_shmem_ptr = __cvta_generic_to_shared(src_shmem); asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst_global_ptr), "r"(src_shmem_ptr), "r"(size) : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_1d_shared_to_global is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-tensor @@ -177,51 +310,93 @@ __device__ __forceinline__ void cp_async_bulk_tensor_1d_shared_to_global(uint64_ __device__ __forceinline__ void cp_async_bulk_tensor_2d_shared_to_global( const uint64_t *tensor_map_ptr, const uint32_t offset_x, const uint32_t offset_y, uint64_t *src_shmem) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) uint32_t src_shmem_ptr = __cvta_generic_to_shared(src_shmem); asm volatile("cp.async.bulk.tensor.2d.global.shared::cta.bulk_group [%0, {%1, %2}], [%3];" ::"l"( tensor_map_ptr), "r"(offset_x), "r"(offset_y), "r"(src_shmem_ptr) : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_2d_shared_to_global is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group __device__ __forceinline__ void cp_async_bulk_wait_group() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.wait_group 0;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_wait_group is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-wait-group template __device__ __forceinline__ void cp_async_bulk_wait_group_read() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.wait_group.read 0;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } template <> __device__ __forceinline__ void cp_async_bulk_wait_group_read<0>() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.wait_group.read 0;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } template <> __device__ __forceinline__ void cp_async_bulk_wait_group_read<1>() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.wait_group.read 1;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } template <> __device__ __forceinline__ void cp_async_bulk_wait_group_read<2>() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.wait_group.read 2;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } template <> __device__ __forceinline__ void cp_async_bulk_wait_group_read<4>() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.wait_group.read 4;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_wait_group_read is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-cp-async-bulk-commit-group __device__ __forceinline__ void cp_async_bulk_commit_group() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("cp.async.bulk.commit_group;"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_commit_group is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } // Proxy fence (bi-directional): -__device__ __forceinline__ void fence_proxy_async() { asm volatile("fence.proxy.async;"); } +__device__ __forceinline__ void fence_proxy_async() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("fence.proxy.async;"); +#else + NVTE_DEVICE_ERROR("fence_proxy_async is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} __device__ __forceinline__ void fence_proxy_async_shared_cta() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) asm volatile("fence.proxy.async.shared::cta;"); +#else + NVTE_DEVICE_ERROR("fence_proxy_async_shared_cta is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) } template @@ -282,15 +457,6 @@ static_assert(sizeof(fp4e2m1x2) == 1); static_assert(sizeof(fp4e2m1x4) == 2); #endif // CUDA_VERSION >= 12080 -// cvt.rn.satfinite.e2m1x2.f32 d, a, b; // Convert two FP32 values to two packed e2m1 - -// cvt.rn.satfinite{.relu}.{e2m1x2/e2m3x2/e3m2x2/ue8m0x2}.f32 introduced in PTX ISA version 8.6. - -// vt.rn.satfinite{.relu}.{e2m1x2/e2m3x2/e3m2x2/ue8m0x2}.f32 is supported on following architectures: -// sm_100a -// sm_101a -// sm_120a - // When converting to .e2m1x2 data formats, the destination operand d has .b8 type. // When converting two .f32 inputs to .e2m1x2, each input is converted to the specified format, // and the converted values are packed in the destination operand d such that the value @@ -313,6 +479,7 @@ __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, cons // SIMD like "Fused" cast + multiplication (x2) __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const floatx2 &in, const floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile( "{\n" ".reg.b64 val_pair; \n\t" @@ -325,10 +492,14 @@ __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const floatx2 &in, : "=h"(reinterpret_cast(out)) : "l"(reinterpret_cast(in)), "l"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_2x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void mul_cvt_2x(fp8e5m2x2 &out, const floatx2 &in, const floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile( "{\n" ".reg.b64 val_pair; \n\t" @@ -341,9 +512,13 @@ __device__ __forceinline__ void mul_cvt_2x(fp8e5m2x2 &out, const floatx2 &in, : "=h"(reinterpret_cast(out)) : "l"(reinterpret_cast(in)), "l"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_2x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const bf16x2 &in, const floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile( "{\n" ".reg.b64 val_pair_before; \n\t" @@ -363,9 +538,13 @@ __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const bf16x2 &in, con : "=h"(reinterpret_cast(out)) : "r"(reinterpret_cast(in)), "l"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_2x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void mul_cvt_2x(fp8e5m2x2 &out, const bf16x2 &in, const floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile( "{\n" ".reg.b64 val_pair_before; \n\t" @@ -385,9 +564,13 @@ __device__ __forceinline__ void mul_cvt_2x(fp8e5m2x2 &out, const bf16x2 &in, con : "=h"(reinterpret_cast(out)) : "r"(reinterpret_cast(in)), "l"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_2x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const fp16x2 &in, const floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile( "{\n" ".reg.b64 val_pair_before; \n\t" @@ -407,9 +590,13 @@ __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const fp16x2 &in, con : "=h"(reinterpret_cast(out)) : "r"(reinterpret_cast(in)), "l"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_2x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void mul_cvt_2x(fp8e5m2x2 &out, const fp16x2 &in, const floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile( "{\n" ".reg.b64 val_pair_before; \n\t" @@ -429,24 +616,33 @@ __device__ __forceinline__ void mul_cvt_2x(fp8e5m2x2 &out, const fp16x2 &in, con : "=h"(reinterpret_cast(out)) : "r"(reinterpret_cast(in)), "l"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_2x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } __device__ __forceinline__ void abs_max_2x(bf16x2 &dst, const bf16x2 &p1, const bf16x2 &p2) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;" : "=r"(reinterpret_cast(dst)) : "r"(reinterpret_cast(p1)), "r"(reinterpret_cast(p2))); +#else + NVTE_DEVICE_ERROR("abs_max_2x is only supported on SM 8.9+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) } __device__ __forceinline__ void abs_max_2x(fp16x2 &dst, const fp16x2 &p1, const fp16x2 &p2) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) asm volatile("max.xorsign.abs.f16x2 %0, %1, %2;" : "=r"(reinterpret_cast(dst)) : "r"(reinterpret_cast(p1)), "r"(reinterpret_cast(p2))); +#else + NVTE_DEVICE_ERROR("abs_max_2x is only supported on SM 8.9+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) } -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - } // namespace ptx namespace { @@ -464,6 +660,8 @@ __forceinline__ __device__ void initialize_barriers(uint64_t *mbar, const bool i } // Syncthreads so initialized barrier is visible to all threads. __syncthreads(); +#else + NVTE_DEVICE_ERROR("initialize_barriers is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -479,6 +677,8 @@ __forceinline__ __device__ void destroy_barriers(uint64_t *mbar, const bool is_m ptx::mbarrier_invalid(&mbar[iter]); } } +#else + NVTE_DEVICE_ERROR("destroy_barriers is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -498,6 +698,8 @@ __forceinline__ __device__ void copy_1d_to_shared(void *dst, const void *src, // Other threads just arrive ptx::mbarrier_arrive(barrier); } +#else + NVTE_DEVICE_ERROR("copy_1d_to_shared is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -517,6 +719,8 @@ __forceinline__ __device__ void copy_2d_to_shared(void *dst, const void *src, co // Other threads just arrive ptx::mbarrier_arrive(barrier); } +#else + NVTE_DEVICE_ERROR("copy_2d_to_shared is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -543,6 +747,8 @@ __forceinline__ __device__ void copy_2d_to_sharedx2(void *dst, const void *src, // Other threads just arrive ptx::mbarrier_arrive(barrier); } +#else + NVTE_DEVICE_ERROR("copy_2d_to_sharedx2 is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -572,6 +778,8 @@ __forceinline__ __device__ void copy_2d_to_sharedx3( // Other threads just arrive ptx::mbarrier_arrive(barrier); } +#else + NVTE_DEVICE_ERROR("copy_2d_to_sharedx3 is only supported on SM 10.0+."); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index bc764ac746..2d37e9c85a 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -16,6 +16,7 @@ #endif #if !defined(__CUDACC_RTC__) +#include #include #else // Importing C++ standard headers is a pain with NVRTC From e2f2a0b4ef206af541c903262476db8cbfab3fb8 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 23 Oct 2025 12:34:50 -0700 Subject: [PATCH 020/521] [JAX] Make SR rng state always 2D (num_devices, 4) to fix partitioning issue (#2294) * Make SR rng state always 2D (num_devices, 4) Signed-off-by: Jeremy Berchtold * fix pure-jax impl Signed-off-by: Jeremy Berchtold * fix test shape Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- tests/jax/test_custom_call_compute.py | 2 +- transformer_engine/jax/quantize/helper.py | 6 ++---- transformer_engine/jax/quantize/quantizer.py | 18 +++++++++++------- 3 files changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 1217ebf65f..11ff9d061c 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -876,7 +876,7 @@ def _sample_sr_qdq( for i in range(num_samples): iter_key = jax.random.fold_in(key, i) sr_rng_state = jax.random.randint( - iter_key, (4,), minval=0, maxval=2**30 - 1, dtype=jnp.uint32 + iter_key, (1, 4), minval=0, maxval=2**30 - 1, dtype=jnp.uint32 ) quantizer = QuantizerFactory.create( q_dtype=q_dtype, diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index e8b33c1d1c..d5093e70e4 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -631,10 +631,8 @@ def _make_stochastic_rounding_rng_state( ) sr_jax_rng = jax.jit(jax.random.fold_in)(sr_jax_rng, quantizer_hash) - # Generate 4 random uint32 values from the JAX PRNG key - shape = (4,) - if get_num_devices_in_mesh() > 1: - shape = (get_num_devices_in_mesh(), 4) + # Generate 4 random uint32 values per device from the JAX PRNG key + shape = (get_num_devices_in_mesh(), 4) sr_jax_rng_state = jax.random.randint( sr_jax_rng, shape, 0, jnp.iinfo(jnp.int32).max, dtype=jnp.int32 ).view(jnp.uint32) diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index d138b58dad..eb2b7b5924 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -34,6 +34,7 @@ TensorSource, ) from .device_utils import is_fp8_gemm_with_all_layouts_supported +from ..sharding import get_num_devices_in_mesh __all__ = [ "QuantizeLayout", @@ -633,9 +634,11 @@ def _apply_stochastic_rounding(self, x): assert ( self.stochastic_rounding_rng_state is not None ), "Stochastic rounding RNG state is not initialized" - assert self.stochastic_rounding_rng_state.shape == ( - 4, - ), "Stochastic rounding RNG state must be of shape (4,)" + expected_sr_rng_state_shape = (get_num_devices_in_mesh(), 4) + assert self.stochastic_rounding_rng_state.shape == expected_sr_rng_state_shape, ( + "Stochastic rounding RNG state must be of shape (num_devices_in_mesh, 4). Expected" + f" {expected_sr_rng_state_shape}, but got {self.stochastic_rounding_rng_state.shape}" + ) assert ( self.stochastic_rounding_rng_state.dtype == jnp.uint32 ), "Stochastic rounding RNG state must be of dtype uint32" @@ -643,14 +646,15 @@ def _apply_stochastic_rounding(self, x): # Default RNG state in JAX expects 2x 32-bit integers, use first 2 uint32s for initial state and fold in the other 2 uint32s key_bits = jnp.array( [ - self.stochastic_rounding_rng_state[0], - self.stochastic_rounding_rng_state[1], + # only take the first device's RNG state as the pure-JAX stochastic rounding impl only uses a single-device + self.stochastic_rounding_rng_state[0][0], + self.stochastic_rounding_rng_state[0][1], ], dtype=jnp.uint32, ) key = jax.random.wrap_key_data(key_bits) - key = jax.jit(jax.random.fold_in)(key, self.stochastic_rounding_rng_state[2]) - key = jax.jit(jax.random.fold_in)(key, self.stochastic_rounding_rng_state[3]) + key = jax.jit(jax.random.fold_in)(key, self.stochastic_rounding_rng_state[0][2]) + key = jax.jit(jax.random.fold_in)(key, self.stochastic_rounding_rng_state[0][3]) abs_x = jnp.abs(x) sign_x = jnp.sign(x) From 021e1e6239a44c334390ba8baf1d759166dfcedc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 24 Oct 2025 01:46:52 +0200 Subject: [PATCH 021/521] [PyTorch Debug] Fix issue with microbatching + debug value caching (#2108) * fix perf issue Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski --- tests/pytorch/debug/test_perf.py | 11 +++++++---- transformer_engine/pytorch/module/base.py | 8 +++++++- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/debug/test_perf.py b/tests/pytorch/debug/test_perf.py index 2d4b62b23f..ad40c31c02 100644 --- a/tests/pytorch/debug/test_perf.py +++ b/tests/pytorch/debug/test_perf.py @@ -28,13 +28,15 @@ def _run_cpu_overhead(debug_tools_initialized, layer, configs_dir, feature_dirs) model = torch.nn.Sequential( te.Linear(1, 1, name="linear1"), te.Linear(1, 1, name="linear2") ).cuda() - NUM_ITERS = 18000 + NUM_ITERS = 1800 elif layer == "transformer": model = torch.nn.Sequential( te.TransformerLayer(1, 1, 1, name="transformer1"), te.TransformerLayer(1, 1, 1, name="transformer2"), ).cuda() - NUM_ITERS = 2000 + NUM_ITERS = 200 + + NUM_INVOCATIONS_PER_ITER = 10 x = torch.randn(1, 1, 1).cuda() @@ -45,8 +47,9 @@ def _run_cpu_overhead(debug_tools_initialized, layer, configs_dir, feature_dirs) time_start = time.time() for i in range(NUM_ITERS): - y = model(x) - y.sum().backward() + for _ in range(NUM_INVOCATIONS_PER_ITER): + y = model(x) + y.sum().backward() if debug_tools_initialized: debug_api.step() torch.cuda.synchronize() diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 7f571ce011..53b9920a6a 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1523,7 +1523,13 @@ def is_debug_iter(self) -> bool: debug = False else: debug = TEDebugState.get_iteration() >= self.next_iter_when_debug_should_be_run - self.debug_last_iteration = TEDebugState.get_iteration() + self.debug_last_iteration = TEDebugState.get_iteration() + self.debug_enabled_in_this_iteration = debug + else: + # If this is the same iteration as previous invocation of the module, + # we use the debug value from the first invocation in the iteration. + debug = self.debug_enabled_in_this_iteration + return debug def no_debug_features_active(self, quantizers): From 6273cede50f50f6e48314fddb9d22da2d16ef871 Mon Sep 17 00:00:00 2001 From: buptzyb Date: Fri, 24 Oct 2025 21:56:03 +0800 Subject: [PATCH 022/521] [PyTorch] Support delay_wgrad_compute cudagraph (#1948) * support cudagraph dw Signed-off-by: Robin Zhang * fix lint Signed-off-by: Robin Zhang * fix ci Signed-off-by: Robin Zhang --------- Signed-off-by: Robin Zhang Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/graph.py | 92 +++++++++++++++---- transformer_engine/pytorch/module/base.py | 12 ++- .../pytorch/module/grouped_linear.py | 2 +- .../pytorch/module/layernorm_mlp.py | 2 +- 4 files changed, 85 insertions(+), 23 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 798d3209a0..9af9fb8870 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -322,14 +322,16 @@ def _make_graphed_callables( fwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] bwd_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] + bwd_dw_graphs = [torch.cuda.CUDAGraph() for _ in range(len(flatten_sample_args))] graph_callables = [None for _ in range(len(flatten_sample_args))] # For cases with multiple active RNG states, e.g. TP. if graph_safe_rng_available(): for _, state in get_all_rng_states().items(): - for fwd_graph, bwd_graph in zip(fwd_graphs, bwd_graphs): + for fwd_graph, bwd_graph, bwd_dw_graph in zip(fwd_graphs, bwd_graphs, bwd_dw_graphs): fwd_graph.register_generator_state(state) bwd_graph.register_generator_state(state) + bwd_dw_graph.register_generator_state(state) mempool = graph_pool_handle() if pool is None else pool @@ -366,21 +368,8 @@ def _make_graphed_callables( ), f"Warmup runs {len(warmup_func)} but only {len(set(warmup_func_idx))} are unique." # Filter the TE modules that cudagraph can access. - visited_te_modules = set() - - def hook_fn(module, inputs, outputs): # pylint: disable=unused-argument - if isinstance(module, TransformerEngineBaseModule): - visited_te_modules.add(module) - # If forward is called on a BasicOperation directly the hook will run - elif isinstance(module, BasicOperation): - visited_te_modules.add(module) - # If forward is called on a te.ops.Sequential it is not called on its constituent ops - elif isinstance(module, Sequential): - assert module._module_groups is not None, "Should have been initialized by warmup" - for module_group in module._module_groups: - if isinstance(module_group, OperationFuser): - for basic_op in module_group._basic_ops: - visited_te_modules.add(basic_op) + visited_te_modules = {} + need_bwd_dw_graph = {} # Run warmup and do the above filtering. with torch.cuda.stream(torch.cuda.Stream()): @@ -388,6 +377,31 @@ def hook_fn(module, inputs, outputs): # pylint: disable=unused-argument args = sample_args[func_idx] kwargs = sample_kwargs[func_idx] static_input_surface = per_callable_static_input_surfaces[func_idx] + + def hook_fn( + module, inputs, outputs, func_idx=func_idx + ): # pylint: disable=unused-argument + modules = set() + if isinstance(module, TransformerEngineBaseModule): + modules.add(module) + # If forward is called on a BasicOperation directly the hook will run + elif isinstance(module, BasicOperation): + modules.add(module) + # If forward is called on a te.ops.Sequential it is not called on its constituent ops + elif isinstance(module, Sequential): + assert ( + module._module_groups is not None + ), "Should have been initialized by warmup" + for module_group in module._module_groups: + if isinstance(module_group, OperationFuser): + for basic_op in module_group._basic_ops: + modules.add(basic_op) + if modules: + if func_idx not in visited_te_modules: + visited_te_modules[func_idx] = modules + else: + visited_te_modules[func_idx].update(modules) + for warmup_iter in range(num_warmup_iters): hooks = [] for module in func.modules(): @@ -432,6 +446,15 @@ def hook_fn(module, inputs, outputs): # pylint: disable=unused-argument module_params_with_grad ) per_callable_static_input_surfaces[func_idx] = static_input_surface + + # Run wgrad. This is essential for some TE modules when they have + # delay_wgrad_compute enabled. + need_backward_dw = False + for module in visited_te_modules.get(func_idx, set()): + if hasattr(module, "need_backward_dw") and module.need_backward_dw(): + need_backward_dw = True + module.backward_dw() + need_bwd_dw_graph[func_idx] = need_backward_dw else: grad_inputs = None del outputs, grad_inputs @@ -514,6 +537,17 @@ def hook_fn(module, inputs, outputs): # pylint: disable=unused-argument allow_unused=allow_unused_input, retain_graph=retain_graph_in_backward, ) + # If no one module needs the backward_dw, the bwd_dw_graph will be empty. + # So skip capturing it. + if need_bwd_dw_graph[per_callable_bwd_idx]: + bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx] + with _graph_context_wrapper(bwd_dw_graph, pool=mempool): + for module in visited_te_modules[per_callable_bwd_idx]: + if ( + hasattr(module, "need_backward_dw") + and module.need_backward_dw() + ): + module.backward_dw() # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs # that don't require grad. I couldn't think of a one-liner for this pattern. @@ -582,10 +616,12 @@ def hook_fn(module, inputs, outputs): # pylint: disable=unused-argument # Capture backward graphs in reverse order per_callable_static_grad_outputs = [] per_callable_static_grad_inputs = [] - for static_input_surface, static_outputs, bwd_graph in zip( + for static_input_surface, static_outputs, bwd_graph, bwd_dw_graph, bwd_idx in zip( reversed(per_callable_static_input_surfaces), reversed(per_callable_static_outputs), reversed(bwd_graphs), + reversed(bwd_dw_graphs), + reversed(range(len(per_callable_static_input_surfaces))), ): # For now, assumes all static_outputs require grad static_grad_outputs = tuple( @@ -601,6 +637,11 @@ def hook_fn(module, inputs, outputs): # pylint: disable=unused-argument allow_unused=allow_unused_input, retain_graph=retain_graph_in_backward, ) + if need_bwd_dw_graph[bwd_idx]: + with _graph_context_wrapper(bwd_dw_graph, pool=mempool): + for module in visited_te_modules[bwd_idx]: + if hasattr(module, "need_backward_dw") and module.need_backward_dw(): + module.backward_dw() # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs that # don't require grad. I couldn't think of a slick one-liner for this pattern. @@ -732,9 +773,10 @@ def functionalized(*user_args, **user_kwargs): ) func = graph_callables[i] + te_modules = visited_te_modules.get(i, set()) if isinstance(func, torch.nn.Module): - def make_graphed_forward(func, graph_training_state, graphed, orig_fwd): + def make_graphed_forward(func, graph_training_state, graphed, orig_fwd, te_modules): def new_fwd(*user_args, **user_kwargs): # If the module's training-or-eval state matches what we graphed, # run the graph, otherwise run the original forward method @@ -743,7 +785,7 @@ def new_fwd(*user_args, **user_kwargs): if FP8GlobalStateManager.is_fp8_enabled(): fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() for m in func.modules(): - if m not in visited_te_modules: + if m not in te_modules: # Only Set the FP8 meta for the modules included by forward continue if isinstance(m, TransformerEngineBaseModule): @@ -780,7 +822,7 @@ def new_fwd(*user_args, **user_kwargs): return new_fwd - forward = make_graphed_forward(func, func.training, graphed, func.forward) + forward = make_graphed_forward(func, func.training, graphed, func.forward, te_modules) if _order is None: func.forward = forward ret.append(func) @@ -789,6 +831,16 @@ def new_fwd(*user_args, **user_kwargs): else: ret.append(graphed) + # Attach backward_dw as an attribute to the graphed callable. + def backward_dw( + need_backward_dw=need_bwd_dw_graph.get(i, False), + bwd_dw_graph=bwd_dw_graphs[i], + ): + if need_backward_dw: + bwd_dw_graph.replay() + + setattr(ret[-1], "backward_dw", backward_dw) + if just_one_callable: return ret[0] diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 53b9920a6a..9b6ca9d9cd 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -662,6 +662,7 @@ def __init__(self) -> None: self._fp8_workspaces: Dict[str, QuantizedTensor] = {} self.activation_dtype: Optional[torch.dtype] = None self.wgrad_accumulation_and_reduce_hooks = [] + self.wgrad_store = None if not TEDebugState.debug_enabled: TEDebugState.initialize() @@ -1481,12 +1482,21 @@ def register_wgrad_accumulation_and_reduce_hooks(self, wgrad_accumulation_and_re """ self.wgrad_accumulation_and_reduce_hooks.append(wgrad_accumulation_and_reduce_hook) + def need_backward_dw(self): + """ + Check if this module needs to execute the delayed weight gradient computation. + This method should be used at the beginning of self.backward_dw() to determine if it + should actually be executed or just return without doing anything. + User can also manually call this method to check that before calling into backward_dw(). + """ + return self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute() + def backward_dw(self): """ Execute the delayed weight gradient computation. This method is called after the main backward pass to compute weight gradients. """ - if self.wgrad_store is None or not self.wgrad_store.delay_wgrad_compute(): + if not self.need_backward_dw(): return with torch.cuda.nvtx.range(f"_{self.__class__.__name__}_wgrad"): (wgrad, bgrad), _ = self.wgrad_store.pop() diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index aae85e2cab..bba97554c5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -840,7 +840,7 @@ def backward_dw(self): Execute the delayed weight gradient computation. This method is called after the main backward pass to compute weight gradients. """ - if self.wgrad_store is None or not self.wgrad_store.delay_wgrad_compute(): + if not self.need_backward_dw(): return with torch.cuda.nvtx.range("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index bae0f28251..ccf5dc0953 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2211,7 +2211,7 @@ def backward_dw(self): Execute the delayed weight gradient computation. This method is called after the main backward pass to compute weight gradients. """ - if self.wgrad_store is None or not self.wgrad_store.delay_wgrad_compute(): + if not self.need_backward_dw(): return with torch.cuda.nvtx.range("_LayerNormMLP_wgrad"): (fc2_wgrad, fc2_bias_grad_, *_), tensor_list_fc2 = self.wgrad_store.pop() From 060811c93615c7f8f671bdd870e4fe292b997836 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 24 Oct 2025 08:02:59 -0700 Subject: [PATCH 023/521] [Common] Fix checks in quantize_transpose_vector_blockwise_fp4 (#2299) fix checks in unoptimized non-rht fp4 quantize kernel Signed-off-by: Jeremy Berchtold --- .../quantize_transpose_vector_blockwise_fp4.cu | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) 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 fed18c51f8..4735fdcbe0 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -718,13 +718,11 @@ void quantize_transpose_vector_blockwise_fp4( // raise error if pow2_scale is true NVTE_CHECK(!pow2_scale, "No support for pow2_scale for MXFP4 for now"); - if (!return_identity && !return_transpose) { - return; - } + NVTE_CHECK(return_identity || return_transpose, + "At least one of return_identity or return_transpose must be true."); - if (use_2d_quantization && !return_identity) { - return; - } + NVTE_CHECK(return_identity || !use_2d_quantization, + "2D block quantization is only supported when return_identity is true."); const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u; size_t num_elements = row_length; @@ -777,7 +775,7 @@ void quantize_transpose_vector_blockwise_fp4( input.dtype, InputType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP4x2_ONLY( - output.dtype, 2, OutputType, + return_identity ? output.dtype : output_t.dtype, 2, OutputType, dim3 grid(num_blocks_x, num_blocks_y, 1); From 87cb26c63c4dc240a77d1b526374631d28810018 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 24 Oct 2025 17:01:51 -0700 Subject: [PATCH 024/521] [PyTorch] Add max_logit support for MuonClip (#2195) * add max_score for fused/unfused F16 non-CP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * calculate max per head instead of max over all heads Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fused attn max_score shape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert FE to github 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 to 1.15.0-rc Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * 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 * reduce ew kernels; fix causal masks; add more 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 * minor fix to 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 * remove logic for flash-attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * WIP: add CP support for p2p/a2a/all_gather 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 improvements of implementation/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 * WIP: add thd support Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add thd to UnfusedDPA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * more fixes for lint 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 to FE 1.15 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unneeded changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable unfused for thd + pad_between_seqs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable thd for unfused until bug is fixed 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 all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * rename max_score to max_logit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix all_gather Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable fused attn + thd 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> --- 3rdparty/cudnn-frontend | 2 +- .../attention/run_attention_with_cp.py | 15 +- tests/pytorch/attention/test_attention.py | 68 ++- .../attention/test_attention_with_cp.py | 6 +- tests/pytorch/utils.py | 3 + .../common/fused_attn/fused_attn.cpp | 80 ++-- .../fused_attn_f16_arbitrary_seqlen.cu | 410 ++++++++++++------ .../fused_attn_f16_arbitrary_seqlen.h | 46 +- .../common/fused_attn/fused_attn_fp8.cu | 6 +- transformer_engine/common/fused_attn/utils.h | 5 +- .../include/transformer_engine/fused_attn.h | 79 ++-- .../jax/csrc/extensions/attention.cpp | 32 +- .../dot_product_attention/backends.py | 69 ++- .../dot_product_attention/context_parallel.py | 79 +++- .../dot_product_attention.py | 15 + .../attention/dot_product_attention/utils.py | 91 ++++ .../pytorch/cpp_extensions/fused_attn.py | 18 + transformer_engine/pytorch/csrc/extensions.h | 4 +- .../pytorch/csrc/extensions/attention.cpp | 25 +- 19 files changed, 748 insertions(+), 305 deletions(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 80a8e4af4d..0b1577c8c8 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 80a8e4af4d89d33a2c59d51fcf9fda1c9d368cd4 +Subproject commit 0b1577c8c83401237d601d0d0db5210506705396 diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 1edffaf486..5ed67c3d5e 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -248,6 +248,7 @@ def run_dpa_with_cp( attn_mask_type=config.attn_mask_type, window_size=config.window_size, softmax_type=config.softmax_type, + return_max_logit=config.return_max_logit, ).cuda() if config.softmax_type != "vanilla": core_attn.softmax_offset.requires_grad = True @@ -308,6 +309,7 @@ def run_dpa_with_cp( fp8_context = autocast(enabled=True, recipe=fp8_recipe, amax_reduction_group=cp_comm_group) else: fp8_context = nullcontext() + max_logit = None with fp8_context: # q, k, v, out in FP8; dout in F16 out = core_attn( @@ -322,6 +324,8 @@ def run_dpa_with_cp( cu_seqlens_kv_padded=cu_seqlens_kv_padded, fp8_output=fp8_mha, ) + if config.return_max_logit: + out, max_logit = out if fp8_bwd and fp8_mha: dout_fp8 = dout_quantizer(dout) out.backward(dout_fp8) @@ -400,6 +404,7 @@ def run_dpa_with_cp( fp8_context = nullcontext() # run attention + max_logit_ = None with fp8_context: # q, k, v, out in FP8; dout in F16 out_ = core_attn( @@ -414,6 +419,8 @@ def run_dpa_with_cp( cu_seqlens_kv_padded=cu_seqlens_kv_padded, fp8_output=fp8_mha, ) + if config.return_max_logit: + out_, max_logit_ = out_ if fp8_bwd and fp8_mha: dout_fp8_ = dout_quantizer(dout_) out_.backward(dout_fp8_) @@ -495,15 +502,15 @@ def run_dpa_with_cp( ) atol, rtol, rmse_tol = get_tols(config, dtype) - tensors_cp = [out_, dq_, dk_, dv_, d_softmax_offset_] - tensors_no_cp = [out, dq, dk, dv, d_softmax_offset] - names = ["out", "dq", "dk", "dv", "d_softmax_offset"] + tensors_cp = [out_, dq_, dk_, dv_, d_softmax_offset_, max_logit_] + tensors_no_cp = [out, dq, dk, dv, d_softmax_offset, max_logit] + names = ["out", "dq", "dk", "dv", "d_softmax_offset", "max_logit"] names_cp = [x + "_cp" for x in names] names_no_cp = [x + "_no_cp" for x in names] is_fp8 = dtype == "fp8" for i, t in enumerate(tensors_no_cp): if t is not None: - if "softmax_offset" not in names[i]: + if "softmax_offset" not in names[i] and "max_logit" not in names[i]: if qkv_format == "bshd": compare_and_assert( t[:, 0], diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 3150c06abb..b05a0447c5 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -131,6 +131,11 @@ def test_dot_product_attention( if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] 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 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" + ) # Get backends is_training = True @@ -172,7 +177,7 @@ def test_dot_product_attention( # UnfusedDotProductAttention backend if unfused_attn_supported: - unfused_attn_fwd, unfused_attn_bwd = _run_dot_product_attention( + unfused_attn_fwd, unfused_max_logit, unfused_attn_bwd = _run_dot_product_attention( dtype, config, "UnfusedDotProductAttention", @@ -186,7 +191,7 @@ def test_dot_product_attention( # FusedAttention backend if fused_attn_supported: if len(fused_attn_backends) == 1: - fused_attn_fwd, fused_attn_bwd = _run_dot_product_attention( + fused_attn_fwd, fused_max_logit, fused_attn_bwd = _run_dot_product_attention( dtype, config, "FusedAttention", @@ -198,7 +203,7 @@ def test_dot_product_attention( ) if len(fused_attn_backends) == 2: os.environ["NVTE_FUSED_ATTN_BACKEND"] = "0" - fused_attn_fwd, fused_attn_bwd = _run_dot_product_attention( + fused_attn_fwd, _, fused_attn_bwd = _run_dot_product_attention( dtype, config, "FusedAttention", @@ -209,7 +214,7 @@ def test_dot_product_attention( is_training, ) os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" - fused_attn_fwd_1, fused_attn_bwd_1 = _run_dot_product_attention( + fused_attn_fwd_1, _, fused_attn_bwd_1 = _run_dot_product_attention( dtype, config, "FusedAttention", @@ -222,7 +227,7 @@ def test_dot_product_attention( # FlashAttention backend if flash_attn_supported: - flash_attn_fwd, flash_attn_bwd = _run_dot_product_attention( + flash_attn_fwd, _, flash_attn_bwd = _run_dot_product_attention( dtype, config, "FlashAttention", @@ -243,6 +248,8 @@ def test_dot_product_attention( if unfused_attn_supported and fused_attn_supported: logging.info("[test_dot_product_attention]: unfused attn vs fused attn") torch.testing.assert_close(fused_attn_fwd, unfused_attn_fwd, **tols) + if config.return_max_logit: + torch.testing.assert_close(fused_max_logit, unfused_max_logit, **tols) for i, _ in enumerate(unfused_attn_bwd): torch.testing.assert_close(fused_attn_bwd[i], unfused_attn_bwd[i], **tols) if fused_attn_supported and flash_attn_supported: @@ -266,6 +273,33 @@ def test_dpa_checkpoint(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, True, None, False, False) +model_configs_max_logit = { + # test: ModelConfig(b, sq, hq, dqk) + "max_logit_1": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096), + "max_logit_2": ModelConfig(2, 2048, 24, 128, attn_mask_type="causal"), + "max_logit_3": ModelConfig(2, 1, 16, 128, max_seqlen_kv=2048, attn_mask_type="padding_causal"), + "max_logit_4": ModelConfig( + 8, 128, 16, 192, max_seqlen_kv=2048, attn_bias_type="post_scale_bias" + ), + "max_logit_5": ModelConfig( + 8, 128, 16, 512, max_seqlen_kv=2048, attn_mask_type="causal", window_size=(20, 0) + ), + "max_logit_6": ModelConfig(8, 1, 16, 1024, max_seqlen_kv=2048), +} + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_max_logit]) +@pytest.mark.parametrize("model", model_configs_max_logit.keys()) +@pytest.mark.parametrize("qkv_layout", ["sbhd_sbhd_sbhd", "thd_thd_thd"]) +def test_dpa_max_logit(dtype, model_configs, model, qkv_layout): + """Test DotProductAttention module with checkpointing""" + config = model_configs[model] + config.return_max_logit = True + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + + model_configs_softmax = { # test: ModelConfig(b, sq, hq, dqk) "softmax_1_0": ModelConfig(2, 2048, 64, 64, num_gqa_groups=8), @@ -962,6 +996,8 @@ def _run_dot_product_attention( layout = layout.replace("d", "dqk") tensor_shape = [dim_to_num[j] for j in layout.split("_")] tensor = 0.1 * torch.randn(tensor_shape, dtype=dtype, device="cuda") + # tensor: with padding tokens + # tensor_orig: without padding tokens tensor_orig = tensor if qkv_format == "thd" and pad_between_seqs: tensor_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1071,6 +1107,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: layer_number=1, attention_type=config.attn_type, softmax_type=config.softmax_type, + return_max_logit=config.return_max_logit, ).to(dtype=dtype, device="cuda") if not is_training: block = block.eval() @@ -1108,16 +1145,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: alibi_slopes=alibi_slopes, fast_zero_fill=True, ) + max_logit = None + if config.return_max_logit: + out, max_logit = out if is_training: out.backward(d_out) + d_softmax_offset = None if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad + if backend in ["FlashAttention", "UnfusedDotProductAttention"]: if is_training: - return out, (q.grad, k.grad, v.grad, d_softmax_offset) + return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) else: - return out, (None, None, None, d_softmax_offset) + return out, max_logit, (None, None, None, d_softmax_offset) if backend == "FusedAttention": if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) @@ -1146,14 +1188,18 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: [v_grad_orig, v.grad[valid_range_kv[0] : valid_range_kv[1]]], dim=0 ) if is_training: - return out_orig, (q_grad_orig, k_grad_orig, v_grad_orig, d_softmax_offset) + return ( + out_orig, + max_logit, + (q_grad_orig, k_grad_orig, v_grad_orig, d_softmax_offset), + ) else: - return out_orig, (None, None, None, d_softmax_offset) + return out_orig, max_logit, (None, None, None, d_softmax_offset) else: if is_training: - return out, (q.grad, k.grad, v.grad, d_softmax_offset) + return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) else: - return out, (None, None, None, d_softmax_offset) + return out, max_logit, (None, None, None, d_softmax_offset) model_configs_te_layer = { diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 2c7f9d8578..e5c856acd8 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -137,8 +137,8 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): model_configs_fused_attn = { # test: ModelConfig(b, sq, hq, dqk) - "cp_1_0": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal"), # MHA - "cp_1_1": ModelConfig(2, 4096, 12, 128), # MHA + "cp_1_0": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", return_max_logit=True), # MHA + "cp_1_1": ModelConfig(2, 4096, 12, 128, return_max_logit=True), # MHA "cp_1_2": ModelConfig( 2, 4096, 12, 128, attn_mask_type="causal", attn_bias_type="post_scale_bias" ), # MHA @@ -183,7 +183,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): qkv_formats = ["bshd", "sbhd", "thd"] cp_comm_types = ["p2p", "all_gather", "a2a", "a2a+p2p"] if test_essential: - configs = ["cp_1_0", "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"] + configs = ["cp_1_0", "cp_1_1", "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"] model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs} dtypes = ["bf16", "fp8"] qkv_formats = ["sbhd", "thd"] diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 72a1b3b534..485c739c03 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -205,6 +205,7 @@ def __init__( window_size: Tuple[int, int] = (-1, -1), context_parallel: bool = False, cp_comm_type: str = "p2p", + return_max_logit=False, total_requests: int = None, max_ctx_len: int = None, num_layers: int = 1, @@ -233,6 +234,7 @@ def __init__( self.window_size = check_set_window_size(self.attn_mask_type, window_size) self.context_parallel = context_parallel self.cp_comm_type = cp_comm_type + self.return_max_logit = return_max_logit self.total_requests = total_requests self.max_ctx_len = max_ctx_len self.num_layers = num_layers @@ -318,6 +320,7 @@ def test(): is_training=is_training, inference_params=inference_params, softmax_type=config.softmax_type, + return_max_logit=config.return_max_logit, ) ( use_flash_attention, diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 77cd8d235a..f6ee37d4c5 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -138,7 +138,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float dropout, 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, int64_t window_size_left, - int64_t window_size_right) { + int64_t window_size_right, bool return_max_logit) { using namespace transformer_engine; NVTE_Fused_Attn_Backend backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; const int device_id = cuda::current_device(); @@ -187,7 +187,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (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) && // 9.10.0: known bugs with SDPA FP8 - (cudnn_runtime_version != 91000)) { + (cudnn_runtime_version != 91000) && !return_max_logit) { if (cudnn_runtime_version >= 8900) { backend = NVTE_Fused_Attn_Backend::NVTE_FP8; } else { @@ -216,7 +216,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (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)) { + (softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) && !return_max_logit) { flag_m512 = true; } if ( @@ -418,8 +418,8 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, - size_t max_seqlen, bool is_training, float attn_scale, - float dropout, NVTE_QKV_Layout qkv_layout, + size_t max_seqlen, bool is_training, bool return_max_logit, + 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, NVTETensor workspace, @@ -460,7 +460,7 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h, h, max_seqlen, max_seqlen, d, d, window_size_left, window_size_right); + h, h, max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, return_max_logit); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -474,10 +474,10 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { #if (CUDNN_VERSION >= 8900) fused_attn_arbitrary_seqlen_fwd_qkvpacked( - b, h, max_seqlen, d, t, is_training, attn_scale, dropout, qkv_layout, bias_type, - attn_mask_type, softmax_type, window_size_left, window_size_right, input_QKV, input_Bias, - input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, input_cu_seqlens_padded, - input_rng_state, wkspace, stream, handle); + b, h, max_seqlen, d, t, is_training, return_max_logit, attn_scale, dropout, qkv_layout, + bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, input_QKV, + input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, + input_cu_seqlens_padded, 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"); @@ -544,7 +544,7 @@ void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, con NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h, h, - max_seqlen, max_seqlen, d, d, window_size_left, window_size_right); + max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -602,7 +602,7 @@ void nvte_fused_attn_fwd_kvpacked( const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, const NVTETensor page_table_v, const NVTETensor rng_state, size_t max_seqlen_q, - size_t max_seqlen_kv, bool is_training, float attn_scale, float dropout, + size_t max_seqlen_kv, bool is_training, bool return_max_logit, 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, NVTETensor workspace, cudaStream_t stream) { @@ -680,7 +680,8 @@ void nvte_fused_attn_fwd_kvpacked( NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right); + h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, + return_max_logit); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -695,12 +696,12 @@ void nvte_fused_attn_fwd_kvpacked( #if (CUDNN_VERSION >= 8903) fused_attn_arbitrary_seqlen_fwd_kvpacked( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, 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, attn_scale, - dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, - window_size_right, input_Q, input_KV, 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); + 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, input_Q, input_KV, 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.3 is required for BF16/FP16 fused attention with arbitrary sequence length. \n"); @@ -777,7 +778,7 @@ void nvte_fused_attn_bwd_kvpacked( NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right); + h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -832,18 +833,16 @@ void nvte_fused_attn_bwd_kvpacked( } } // NVTE fused attention FWD with separate Q, K and V -void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - const NVTETensor page_table_v, const NVTETensor rng_state, - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, - 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, NVTETensor workspace, cudaStream_t stream) { +void nvte_fused_attn_fwd( + const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, + const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, + const NVTETensor page_table_k, 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, + 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, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -913,7 +912,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right); + h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, + return_max_logit); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -928,12 +928,12 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso #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, attn_scale, - dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, - window_size_right, 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); + 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, 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"); @@ -1008,7 +1008,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right); + h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) 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 ba0f845789..950ced61bb 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 @@ -53,10 +53,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t num_pages_k, int64_t num_pages_v, 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, bool is_training, - float scaling_factor, float dropout_probability, NVTE_QKV_Layout layout, + 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, void *devPtrQ, void *devPtrK, - void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrSoftmaxStats, + 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, @@ -102,36 +102,40 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; + bool generate_stats = !return_max_logit; try { - FADescriptor_v1 descriptor{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, - scaling_factor, - is_training, - dropout_probability, - layout, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - true, - tensorType, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET}; + FADescriptor_v1 descriptor{ + 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, + scaling_factor, + is_training, + dropout_probability, + layout, + bias_type, + mask_type, + softmax_type, + window_size_left, + window_size_right, + true, + tensorType, + cudnn_frontend::DataType_t::NOT_SET, + cudnn_frontend::DataType_t::NOT_SET, + cudnn_frontend::DataType_t::NOT_SET, + return_max_logit, + }; namespace fe = cudnn_frontend; using graph_and_tensors = @@ -141,7 +145,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr, // V std::shared_ptr, // attn_scale std::shared_ptr, // O - std::shared_ptr, // Stats + std::shared_ptr, // S1 + std::shared_ptr, // S2 std::shared_ptr, // bias std::shared_ptr, // softmax_offset std::shared_ptr, // seq_q @@ -244,6 +249,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( sdpa_options = fe::graph::SDPA_attributes() .set_name("flash_attention") .set_is_inference(false) + .set_generate_stats(generate_stats) .set_causal_mask(is_causal) .set_causal_mask_bottom_right(is_bottom_right) .set_attn_scale(attn_scale); @@ -317,7 +323,36 @@ void fused_attn_arbitrary_seqlen_fwd_impl( sdpa_options.set_sink_token(softmax_offset); } - auto [O, Stats] = mha_graph->sdpa(Q, K, V, sdpa_options); + std::shared_ptr Max, Sum_Exp; + if (is_ragged_q && cudnn_runtime_version >= 90600) { + offset_stats = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("offset_stats") + .set_dim({b + 1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); + } + if (return_max_logit) { + Max = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Max") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + Sum_Exp = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Sum_Exp") + .set_dim({b, h, s_q, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + if (is_ragged_q && cudnn_runtime_version >= 90600) { + Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + Sum_Exp->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Max->set_stride({h * s_q, s_q, 1, 1}); + Sum_Exp->set_stride({h * s_q, s_q, 1, 1}); + } + sdpa_options.set_logit_max(Max); + sdpa_options.set_score_sum_exp(Sum_Exp); + } + + 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, @@ -332,17 +367,13 @@ void fused_attn_arbitrary_seqlen_fwd_impl( O->set_ragged_offset(offset_o); } - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (is_ragged_q && cudnn_runtime_version >= 90600) { - offset_stats = - mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("offset_stats") - .set_dim({b + 1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(get_cudnn_fe_dtype(ragged_offset_type))); - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); + if (!return_max_logit) { + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (is_ragged_q && cudnn_runtime_version >= 90600) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); + } } std::tuple, // Q @@ -351,7 +382,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr, // attn_scale std::shared_ptr> // O key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); - auto Stats_tuple = std::make_tuple(Stats); + auto Stats_tuple = + generate_stats ? std::make_tuple(Stats, nullptr) : std::make_tuple(Max, Sum_Exp); 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); @@ -384,7 +416,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( return return_tuple; }; - auto [mha_graph, Q, K, V, attn_scale, O, Stats, bias, softmax_offset, seq_q, seq_kv, + auto [mha_graph, Q, K, V, attn_scale, O, S1, S2, bias, softmax_offset, seq_q, seq_kv, page_table_k, page_table_v, offset_q, offset_o, offset_k, offset_v, offset_stats, dropout_seed, dropout_offset] = get_graph(sdpa_f16_fprop_cache, descriptor); @@ -417,9 +449,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( // Build variant pack std::unordered_map, void *> variant_pack = { - {Q, devPtrQ}, {K, devPtrK}, - {V, devPtrV}, {attn_scale, &scaling_factor}, - {O, devPtrO}, {Stats, devPtrSoftmaxStats}}; + {Q, devPtrQ}, {K, devPtrK}, {V, devPtrV}, {attn_scale, &scaling_factor}, + {O, devPtrO}, {S1, devPtrS1}}; + + if (return_max_logit) { + variant_pack[S2] = devPtrS2; + } if (is_bias) { variant_pack[bias] = devPtrBias; @@ -561,35 +596,38 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; try { - FADescriptor_v1 descriptor{b, - h, - hg, - s_q, - s_kv, - d_qk, - d_v, - 0, - 0, - 0, - 0, - 0, - 0, - bias_b, - bias_h, - scaling_factor, - true, - dropout_probability, - layout, - bias_type, - mask_type, - softmax_type, - window_size_left, - window_size_right, - deterministic, - tensorType, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET}; + FADescriptor_v1 descriptor{ + b, + h, + hg, + s_q, + s_kv, + d_qk, + d_v, + 0, + 0, + 0, + 0, + 0, + 0, + bias_b, + bias_h, + scaling_factor, + true, + dropout_probability, + layout, + bias_type, + mask_type, + softmax_type, + window_size_left, + window_size_right, + deterministic, + tensorType, + cudnn_frontend::DataType_t::NOT_SET, + cudnn_frontend::DataType_t::NOT_SET, + cudnn_frontend::DataType_t::NOT_SET, + false, + }; namespace fe = cudnn_frontend; using graph_and_tensors = @@ -1001,12 +1039,13 @@ void fused_attn_arbitrary_seqlen_bwd_impl( using namespace transformer_engine::fused_attn; void fused_attn_arbitrary_seqlen_fwd_qkvpacked( size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, - bool is_training, 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, const Tensor *input_QKV, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, const Tensor *cu_seqlens_padded, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { + 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, + const Tensor *input_QKV, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, + Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, + const Tensor *cu_seqlens_padded, const Tensor *rng_state, Tensor *workspace, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto QKV_type = input_QKV->data.dtype; @@ -1037,7 +1076,8 @@ void fused_attn_arbitrary_seqlen_fwd_qkvpacked( } void *devPtrO = output_O->data.dptr; - void *devPtrS = nullptr; + void *devPtrS1 = nullptr; + void *devPtrS2 = nullptr; void *devPtrCuSeqlens = cu_seqlens->data.dptr; void *devPtrSeqOffsets = cu_seqlens_padded->data.dptr; @@ -1051,14 +1091,34 @@ void fused_attn_arbitrary_seqlen_fwd_qkvpacked( size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_S->data.shape = {max_tokens, num_attn_heads, 1}; + if (return_max_logit) { + Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_Max->data.dptr = nullptr; + if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_Max->data.shape = {max_tokens, num_attn_heads, 1}; + } else { + output_Max->data.shape = {batch, num_attn_heads, max_seqlen, 1}; + } + output_Max->data.dtype = DType::kFloat32; + Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_Sum_Exp->data.dptr = nullptr; + if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_Sum_Exp->data.shape = {max_tokens, num_attn_heads, 1}; + } else { + output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen, 1}; + } + output_Sum_Exp->data.dtype = DType::kFloat32; } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen, 1}; + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_S->data.dptr = nullptr; + if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_S->data.shape = {max_tokens, num_attn_heads, 1}; + } else { + output_S->data.shape = {batch, num_attn_heads, max_seqlen, 1}; + } + output_S->data.dtype = DType::kFloat32; } - output_S->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}; @@ -1080,8 +1140,15 @@ void fused_attn_arbitrary_seqlen_fwd_qkvpacked( Aux_CTX_Tensors->size = i; } else if (Aux_CTX_Tensors->size >= 2) { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS = output_S->data.dptr; + if (return_max_logit) { + Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_Max->data.dptr; + Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS2 = output_Sum_Exp->data.dptr; + } else { + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_S->data.dptr; + } Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = rng_state->data.dptr; if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { @@ -1105,11 +1172,11 @@ void fused_attn_arbitrary_seqlen_fwd_qkvpacked( fused_attn_arbitrary_seqlen_fwd_impl( batch, num_attn_heads, num_attn_heads, max_seqlen, max_seqlen, head_dim, head_dim, max_batch_size, max_tokens, max_tokens, 0, 0, 0, 0, 0, 0, bias_b, bias_h, is_training, - attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS, - devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlens, devPtrCuSeqlens, nullptr, - nullptr, devPtrSeqOffsets, devPtrSeqOffsets, get_cudnn_fe_dtype(QKV_type), - workspace->data.dptr, &workspace_size, stream, handle); + return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, + window_size_left, window_size_right, devPtrQ, devPtrK, devPtrV, devPtrBias, + devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, + devPtrCuSeqlens, devPtrCuSeqlens, nullptr, nullptr, devPtrSeqOffsets, devPtrSeqOffsets, + get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1221,14 +1288,15 @@ void fused_attn_arbitrary_seqlen_fwd_kvpacked( 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, 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, 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, - const Tensor *input_Q, const Tensor *input_KV, 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) { + 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, const Tensor *input_Q, const Tensor *input_KV, + 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; @@ -1260,7 +1328,8 @@ void fused_attn_arbitrary_seqlen_fwd_kvpacked( } void *devPtrO = output_O->data.dptr; - void *devPtrS = nullptr; + void *devPtrS1 = nullptr; + void *devPtrS2 = nullptr; void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; @@ -1285,14 +1354,34 @@ void fused_attn_arbitrary_seqlen_fwd_kvpacked( size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_S->data.shape = {max_tokens_q, num_attn_heads, 1}; + if (return_max_logit) { + Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_Max->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_Max->data.shape = {max_tokens_q, num_attn_heads, 1}; + } else { + output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_Max->data.dtype = DType::kFloat32; + Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_Sum_Exp->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_Sum_Exp->data.shape = {max_tokens_q, num_attn_heads, 1}; + } else { + output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_Sum_Exp->data.dtype = DType::kFloat32; } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_S->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_S->data.shape = {max_tokens_q, num_attn_heads, 1}; + } else { + output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_S->data.dtype = DType::kFloat32; } - output_S->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}; @@ -1314,8 +1403,15 @@ void fused_attn_arbitrary_seqlen_fwd_kvpacked( Aux_CTX_Tensors->size = i; } else if (Aux_CTX_Tensors->size >= 2) { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS = output_S->data.dptr; + if (return_max_logit) { + Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_Max->data.dptr; + Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS2 = output_Sum_Exp->data.dptr; + } else { + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_S->data.dptr; + } Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = rng_state->data.dptr; if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { @@ -1340,11 +1436,12 @@ void fused_attn_arbitrary_seqlen_fwd_kvpacked( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, head_dim, 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, is_training, - attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS, - devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, - devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); + return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, + window_size_left, window_size_right, 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) { @@ -1471,14 +1568,14 @@ void fused_attn_arbitrary_seqlen_fwd( 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, - 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, 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) { + 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, 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; @@ -1488,7 +1585,8 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrK = input_K->data.dptr; void *devPtrV = input_V->data.dptr; void *devPtrO = output_O->data.dptr; - void *devPtrS = nullptr; + void *devPtrS1 = nullptr; + void *devPtrS2 = nullptr; void *devPtrBias = nullptr; size_t bias_b = 0; size_t bias_h = 0; @@ -1525,14 +1623,34 @@ void fused_attn_arbitrary_seqlen_fwd( size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_S->data.shape = {max_tokens_q, num_attn_heads, 1}; + if (return_max_logit) { + Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_Max->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_Max->data.shape = {max_tokens_q, num_attn_heads, 1}; + } else { + output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_Max->data.dtype = DType::kFloat32; + Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_Sum_Exp->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_Sum_Exp->data.shape = {max_tokens_q, num_attn_heads, 1}; + } else { + output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_Sum_Exp->data.dtype = DType::kFloat32; } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_S->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_S->data.shape = {max_tokens_q, num_attn_heads, 1}; + } else { + output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_S->data.dtype = DType::kFloat32; } - output_S->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}; @@ -1554,8 +1672,15 @@ void fused_attn_arbitrary_seqlen_fwd( Aux_CTX_Tensors->size = i; } else if (Aux_CTX_Tensors->size >= 2) { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS = output_S->data.dptr; + if (return_max_logit) { + Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_Max->data.dptr; + Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS2 = output_Sum_Exp->data.dptr; + } else { + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_S->data.dptr; + } Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = rng_state->data.dptr; if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { @@ -1580,11 +1705,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, is_training, - attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, devPtrQ, devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS, - devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, - devPtrPageTableK, devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); + return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, + window_size_left, window_size_right, 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) { 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 b9658b0530..a3181c6295 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 @@ -20,12 +20,13 @@ namespace transformer_engine { #if (CUDNN_VERSION >= 8900) void fused_attn_arbitrary_seqlen_fwd_qkvpacked( size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, - bool is_training, 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, const Tensor *input_QKV, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, const Tensor *cu_seqlens_padded, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + 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, + const Tensor *input_QKV, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, + Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, + const Tensor *cu_seqlens_padded, const Tensor *rng_state, Tensor *workspace, + cudaStream_t stream, cudnnHandle_t handle); void fused_attn_arbitrary_seqlen_bwd_qkvpacked( size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, @@ -41,14 +42,15 @@ void fused_attn_arbitrary_seqlen_fwd_kvpacked( 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, 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, 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, - const Tensor *input_Q, const Tensor *input_KV, 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); + 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, const Tensor *input_Q, const Tensor *input_KV, + 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_kvpacked( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, @@ -68,14 +70,14 @@ void fused_attn_arbitrary_seqlen_fwd( 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, - 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, 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); + 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, 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, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 21c544491a..7b85be972c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1710,7 +1710,8 @@ void fused_attn_fp8_fwd_impl_v1( qkv_tensor_type, o_tensor_type, cudnn_frontend::DataType_t::NOT_SET, - cudnn_frontend::DataType_t::NOT_SET}; + cudnn_frontend::DataType_t::NOT_SET, + false}; namespace fe = cudnn_frontend; using graph_and_tensors = @@ -2038,7 +2039,8 @@ void fused_attn_fp8_bwd_impl_v1( qkv_tensor_type, o_tensor_type, do_tensor_type, - dqkv_tensor_type}; + dqkv_tensor_type, + false}; namespace fe = cudnn_frontend; using graph_and_tensors = diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index f03774f8ed..72047a73f2 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -115,20 +115,21 @@ struct FADescriptor_v1 { cudnn_frontend::DataType_t o_tensor_type; cudnn_frontend::DataType_t do_tensor_type; cudnn_frontend::DataType_t dqkv_tensor_type; + bool generate_max_sum_exp; 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, attnScale, isTraining, dropoutProbability, layout, mask_type, softmax_type, window_size_left, window_size_right, deterministic, bias_type, qkv_tensor_type, - o_tensor_type, do_tensor_type, dqkv_tensor_type) < + o_tensor_type, do_tensor_type, dqkv_tensor_type, generate_max_sum_exp) < 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.attnScale, rhs.isTraining, rhs.dropoutProbability, rhs.layout, rhs.mask_type, rhs.softmax_type, rhs.window_size_left, rhs.window_size_right, rhs.deterministic, rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, - rhs.dqkv_tensor_type); + rhs.dqkv_tensor_type, rhs.generate_max_sum_exp); } }; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index a150978c4a..518fad20de 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -190,29 +190,30 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); /*! \brief Get fused attention backend based on input parameters. * - * \param[in] is_training Whether the model is in training mode. - * \param[in] q_dtype The data type of Tensor Q. - * \param[in] kv_dtype The data type of Tensors K, V. - * \param[in] qkv_layout The layout of Tensors Q, K, V. - * \param[in] bias_type The attention bias type. - * \param[in] attn_mask_type The attention mask type. - * \param[in] softmax_type The attention softmax type. - * \param[in] dropout The dropout probability. - * \param[in] num_attn_heads The number of heads in Q. - * \param[in] num_gqa_groups The number of heads in K, V. - * \param[in] max_seqlen_q The sequence length of Q. - * \param[in] max_seqlen_kv The sequence length of K, V. - * \param[in] head_dim_qk The head dimension of Q, K. - * \param[in] head_dim_v The head dimension of V. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). + * \param[in] is_training Whether the model is in training mode. + * \param[in] q_dtype The data type of Tensor Q. + * \param[in] kv_dtype The data type of Tensors K, V. + * \param[in] qkv_layout The layout of Tensors Q, K, V. + * \param[in] bias_type The attention bias type. + * \param[in] attn_mask_type The attention mask type. + * \param[in] softmax_type The attention softmax type. + * \param[in] dropout The dropout probability. + * \param[in] num_attn_heads The number of heads in Q. + * \param[in] num_gqa_groups The number of heads in K, V. + * \param[in] max_seqlen_q The sequence length of Q. + * \param[in] max_seqlen_kv The sequence length of K, V. + * \param[in] head_dim_qk The head dimension of Q, K. + * \param[in] head_dim_v The head dimension of V. + * \param[in] window_size_left Sliding window size (the left half). + * \param[in] window_size_right Sliding window size (the right half). + * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float dropout, 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, int64_t window_size_left, - int64_t window_size_right); + int64_t window_size_right, bool return_max_logit); /*! \brief Compute dot product attention with packed QKV input. * @@ -255,6 +256,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] max_seqlen Max sequence length used for computing, * it may be >= max(seqlen_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. + * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensor's layout. @@ -266,13 +268,16 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ -void nvte_fused_attn_fwd_qkvpacked( - const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, - const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, - bool is_training, 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, NVTETensor workspace, cudaStream_t stream); +void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, + const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, + NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, + const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, + size_t max_seqlen, bool is_training, bool return_max_logit, + 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, NVTETensor workspace, + cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with packed QKV input. * @@ -381,6 +386,7 @@ void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, con * \param[in] max_seqlen_kv Max sequence length used for computing for KV. * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. + * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensor's layout. @@ -399,7 +405,7 @@ void nvte_fused_attn_fwd_kvpacked( const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, const NVTETensor page_table_v, const NVTETensor rng_state, size_t max_seqlen_q, - size_t max_seqlen_kv, bool is_training, float attn_scale, float dropout, + size_t max_seqlen_kv, bool is_training, bool return_max_logit, 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, NVTETensor workspace, cudaStream_t stream); @@ -520,6 +526,7 @@ void nvte_fused_attn_bwd_kvpacked( * \param[in] max_seqlen_kv Max sequence length used for computing for K and V. * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. + * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensors' layout. @@ -531,18 +538,16 @@ void nvte_fused_attn_bwd_kvpacked( * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ -void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, - const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - const NVTETensor page_table_v, const NVTETensor rng_state, - size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, - 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, NVTETensor workspace, cudaStream_t stream); +void nvte_fused_attn_fwd( + const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, + const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, + const NVTETensor page_table_k, 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, + 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, NVTETensor workspace, cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 9277569e11..ffc0706fe7 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -22,7 +22,8 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend(bool is_training, DType q_dtype, DTy auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right); + q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, + false); return backend; } @@ -179,17 +180,18 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( qkv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), dummy_rng_state_tensor.data(), q_max_seqlen, is_training, - scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, query_workspace_tensor.data(), nullptr); + false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, + softmax_type, window_size_left, window_size_right, query_workspace_tensor.data(), + nullptr); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { nvte_fused_attn_fwd_kvpacked( q_tensor.data(), kv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, scaling_factor, - dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, query_workspace_tensor.data(), nullptr); + dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, + scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, + window_size_left, window_size_right, query_workspace_tensor.data(), nullptr); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { nvte_fused_attn_fwd( q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), @@ -197,8 +199,8 @@ 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, scaling_factor, dropout_probability, qkv_layout, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, + kv_max_seqlen, is_training, false, scaling_factor, dropout_probability, qkv_layout, + bias_type, mask_type, softmax_type, window_size_left, window_size_right, query_workspace_tensor.data(), nullptr); } else { NVTE_ERROR("Unsupported QKVLayout."); @@ -276,7 +278,8 @@ static void FusedAttnForwardImpl( auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right); + q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, + false); nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); /* Auxiliary tensors (to be propagated to the backward pass later) */ @@ -294,7 +297,7 @@ static void FusedAttnForwardImpl( nvte_fused_attn_fwd_qkvpacked( qkv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), - q_seq_offsets_tensor.data(), rng_state_tensor.data(), q_max_seqlen, is_training, + q_seq_offsets_tensor.data(), rng_state_tensor.data(), q_max_seqlen, is_training, false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, workspace_tensor.data(), stream); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { @@ -308,8 +311,8 @@ static void FusedAttnForwardImpl( s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, scaling_factor, dropout_probability, qkv_layout, - bias_type, mask_type, softmax_type, window_size_left, window_size_right, + q_max_seqlen, kv_max_seqlen, is_training, false, scaling_factor, dropout_probability, + qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, workspace_tensor.data(), stream); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; @@ -323,7 +326,7 @@ static void FusedAttnForwardImpl( dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, scaling_factor, + rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, workspace_tensor.data(), stream); } else { @@ -542,7 +545,8 @@ static void FusedAttnBackwardImpl( auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, - q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right); + q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, + false); PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, softmax_aux, rng_state, bias); diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6c19d868a1..d4903be902 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -58,6 +58,8 @@ combine_and_quantize, combine_and_dequantize, print_quantizers, + ConvertTHDtoBSHD, + ConvertBSHDtoTHD, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( AttentionLogging as attn_log, @@ -201,6 +203,7 @@ def __init__( attention_dropout_ctx: Optional[Callable] = nullcontext, layer_number: Optional[int] = None, softmax_type: str = "vanilla", + return_max_logit: Optional[bool] = False, ) -> None: super().__init__() @@ -209,6 +212,7 @@ def __init__( self.attention_dropout_ctx = attention_dropout_ctx self.layer_number = layer_number self.softmax_type = softmax_type + self.return_max_logit = return_max_logit def mask_func(x, y): return ( @@ -217,6 +221,7 @@ def mask_func(x, y): else attention_mask_func(x, y) ) + self.mask_func = mask_func self.scale_mask_softmax = FusedScaleMaskSoftmax(mask_func) # Dropout. Note that for a single iteration, this layer will generate @@ -238,6 +243,8 @@ def forward( qkv_layout: str = "sbh3d", cu_seqlens_q: Optional[torch.Tensor] = None, # pylint: disable=unused-argument cu_seqlens_kv: Optional[torch.Tensor] = None, # pylint: disable=unused-argument + max_seqlen_q: Optional[torch.Tensor] = None, # pylint: disable=unused-argument + max_seqlen_kv: Optional[torch.Tensor] = None, # pylint: disable=unused-argument attn_mask_type: str = "causal", attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, window_size: Optional[Tuple[int, int]] = None, @@ -261,6 +268,9 @@ def forward( if inference_params is not None and inference_params.is_paged: key_layer, value_layer = inference_params.convert_paged_to_nonpaged(self.layer_number) + # convert to sbhd + # training: bshd, thd + # inference: bshd, sbhd_2bshd, thd_2bshd if qkv_format == "bshd": # convert to sbhd and use sbhd implementation for now query_layer, key_layer, value_layer = [ @@ -269,9 +279,8 @@ def forward( if qkv_format == "sbhd_2bshd": key_layer, value_layer = [x.transpose(0, 1) for x in [key_layer, value_layer]] - total_tokens, batch_size = None, None if qkv_format == "thd_2bshd": - total_tokens, batch_size = query_layer.shape[0], key_layer.shape[0] + batch_size = key_layer.shape[0] query_layer = tex.convert_thd_to_bshd( query_layer, cu_seqlens_q, @@ -281,6 +290,26 @@ def forward( query_layer, key_layer, value_layer = [ x.transpose(0, 1) for x in [query_layer, key_layer, value_layer] ] + if qkv_format == "thd": + assert cu_seqlens_q is not None and cu_seqlens_kv is not None + assert max_seqlen_q is not None and max_seqlen_kv is not None + query_layer = ConvertTHDtoBSHD.apply( + query_layer, + cu_seqlens_q, + max_seqlen_q, + ) + key_layer, value_layer = [ + ConvertTHDtoBSHD.apply( + x, + cu_seqlens_kv, + max_seqlen_kv, + ) + for x in [key_layer, value_layer] + ] + query_layer, key_layer, value_layer = [ + x.transpose(0, 1).contiguous() for x in [query_layer, key_layer, value_layer] + ] + batch_size, max_seqlen_q, max_seqlen_kv = ( query_layer.shape[1], query_layer.shape[0], @@ -426,6 +455,15 @@ def forward( matmul_result, None, None, dP_quantizer, "dP_quantizer", None ) + # max attention score + max_logit = None + if self.return_max_logit: + # matmul_result [b, np, sq, dk], max_logit [np] + 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] if self.softmax_type != "vanilla": matmul_result = torch.cat( @@ -506,14 +544,13 @@ def forward( context_layer = context_layer.permute(0, 2, 1, 3).contiguous() # [b, sq, np, hn] --> [tq, np, hn] - context_layer = tex.convert_bshd_to_thd( + context_layer = ConvertBSHDtoTHD.apply( context_layer, cu_seqlens_q, - total_tokens, ) # [tq, np, hn] --> [tq, hp] - context_layer = context_layer.view(total_tokens, -1) + context_layer = context_layer.view(context_layer.shape[0], -1) if fp8: # quantize and dequantize O to emulate FP8 @@ -529,6 +566,9 @@ def forward( if fp8_output: context_layer = O_quantizer(context_layer) + if self.return_max_logit: + return context_layer, max_logit + return context_layer @@ -1067,6 +1107,7 @@ def forward( softmax_offset, fp8_output, layer_number, + return_max_logit, ): # pylint: disable=missing-function-docstring @@ -1102,6 +1143,7 @@ def forward( # FP8 attention: torch.float16 or torch.bfloat16 out_nominal_dtype = q.dtype + max_logit = None if fp8: fused_attention_backend = FusedAttnBackend["FP8"] @@ -1129,7 +1171,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 - out_, aux_ctx_tensors = fused_attn_fwd( + out_, aux_ctx_tensors, *_ = fused_attn_fwd( is_training, max_seqlen_q, max_seqlen_kv, @@ -1205,7 +1247,7 @@ def forward( qkvo_tensors = (q, k, v, out) else: # q, k, v, out_: torch.Tensor; dtype = torch.float16 or torch.bfloat16 - out_, aux_ctx_tensors = fused_attn_fwd( + out_, aux_ctx_tensors, *max_logit = fused_attn_fwd( is_training, max_seqlen_q, max_seqlen_kv, @@ -1233,6 +1275,7 @@ def forward( window_size, rng_gen, softmax_offset, + return_max_logit, ) out = out_ out_ret = out_ @@ -1327,10 +1370,12 @@ def forward( ctx.use_FAv2_bwd = use_FAv2_bwd ctx.deterministic = deterministic + if return_max_logit: + return out_ret, *max_logit return out_ret @staticmethod - def backward(ctx, d_out): + def backward(ctx, d_out, *_args): # pylint: disable=missing-function-docstring # d_out is expected to be in FP8 if is_output_fp8=True, @@ -1574,6 +1619,7 @@ def backward(ctx, d_out): d_softmax_offset, None, None, + None, ) @@ -1614,6 +1660,7 @@ def __init__( layer_number: Optional[int] = None, deterministic: bool = False, softmax_type: str = "vanilla", + return_max_logit: Optional[bool] = False, ) -> None: super().__init__() @@ -1627,6 +1674,7 @@ def __init__( self.layer_number = 1 if layer_number is None else layer_number self.deterministic = deterministic self.softmax_type = softmax_type + self.return_max_logit = return_max_logit def remove_extra_states_check(self, incompatible_keys): # pylint: disable=unused-argument """ @@ -1846,6 +1894,7 @@ def forward( softmax_offset=softmax_offset, fp8_output=fp8_output, layer_number=self.layer_number, + return_max_logit=self.return_max_logit, ) else: with self.attention_dropout_ctx(): @@ -1881,7 +1930,11 @@ def forward( softmax_offset, fp8_output, self.layer_number, + self.return_max_logit, ) + if self.return_max_logit: + # ...hd -> ...(hd) + return output[0].view(*output[0].shape[:-2], -1), output[1] # ...hd -> ...(hd) return output.view(*output.shape[:-2], -1) 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 e5ee8cc7db..f312cac798 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -617,6 +617,7 @@ def cp_p2p_fwd_fused_attn( rank, step, cp_size, + return_max_logit, q_part, k_part, v_part, @@ -693,7 +694,7 @@ def cp_p2p_fwd_fused_attn( fp8_meta_kwargs["s_quantizer"] = S_quantizer_per_step fp8_meta_kwargs["o_quantizer"] = O_quantizer_per_step - out_per_step, aux_ctx_tensors = fused_attn_fwd( + out_per_step, aux_ctx_tensors, *max_logit = fused_attn_fwd( is_training, max_seqlen_q_, max_seqlen_kv_, @@ -713,6 +714,7 @@ def cp_p2p_fwd_fused_attn( cu_seqlens_q_padded=cu_seqlens_q_padded_, cu_seqlens_kv_padded=cu_seqlens_kv_padded_, **fp8_meta_kwargs, + return_max_logit=return_max_logit, ) if fp8: @@ -721,7 +723,9 @@ def cp_p2p_fwd_fused_attn( softmax_lse_per_step, rng_states, *rest = aux_ctx_tensors attn_bias = rest[0] if len(rest) > 0 else None - return out_per_step, softmax_lse_per_step, rng_states, attn_bias + if return_max_logit: + return out_per_step, softmax_lse_per_step, rng_states, attn_bias, *max_logit + return out_per_step, softmax_lse_per_step, rng_states, attn_bias, None def cp_p2p_fwd_flash_attn( @@ -1086,6 +1090,7 @@ def forward( attn_bias, deterministic, use_fused_attention, + return_max_logit, fp8, fp8_meta, cp_group, @@ -1156,6 +1161,8 @@ def forward( amax_per_step = None S_quantizer_per_step = [None for _ in range(cp_size)] O_quantizer_per_step = [None for _ in range(cp_size)] + max_logit_per_step = [None for _ in range(cp_size)] + max_logit = None assert isinstance(k, q.__class__) and isinstance( v, q.__class__ @@ -1244,6 +1251,10 @@ def forward( q_f16 = q if use_fused_attention: fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"] + if return_max_logit: + max_logit_per_step = [ + torch.empty(q.shape[-2], dtype=q.dtype, device=q.device) for _ in range(cp_size) + ] # split qkv to two halves and prepare for load balancing assert qkv_format == "thd" or ( @@ -1418,6 +1429,7 @@ def forward( rank, i, cp_size, + return_max_logit, ] else: flash_attn_inputs = [ @@ -1462,6 +1474,7 @@ def forward( softmax_lse_per_step[i], rng_states[i], attn_biases[i], + max_logit_per_step[i], ) = cp_p2p_fwd_fused_attn( *fused_attn_inputs, *prepare_outputs, section ) @@ -1488,6 +1501,7 @@ def forward( softmax_lse_per_step[i], rng_states[i], attn_biases[i], + max_logit_per_step[i], ) = cp_p2p_fwd_fused_attn( *fused_attn_inputs, *prepare_outputs, section ) @@ -1514,6 +1528,7 @@ def forward( softmax_lse_per_step[i], rng_states[i], attn_biases[i], + max_logit_per_step[i], ) = cp_p2p_fwd_fused_attn( *fused_attn_inputs, *prepare_outputs, section ) @@ -1541,6 +1556,7 @@ def forward( softmax_lse_per_step[i], rng_states[i], attn_biases[i], + max_logit_per_step[i], ) = cp_p2p_fwd_fused_attn(*fused_attn_inputs, *prepare_outputs, section) else: out_per_step[i], softmax_lse_per_step[i], rng_states[i] = ( @@ -1600,11 +1616,20 @@ def forward( softmax_lse.view(*softmax_lse.shape[:-1], 2, -1), softmax_lse_per_step[i - 1], ) + if return_max_logit: + if i == 1: + max_logit = torch.clone(max_logit_per_step[0]) + else: + max_logit = torch.maximum(max_logit, max_logit_per_step[i - 1]) if i < cp_size: flash_attn_streams[(i - 1) % 2].record_event(fwd_results_correction_done) torch.cuda.current_stream().wait_stream(flash_attn_streams[1]) + if return_max_logit: + torch.distributed.all_reduce( + max_logit, op=torch.distributed.ReduceOp.MAX, group=cp_group + ) second_half_lse_seqlen = None if causal and rank < (cp_size - 1): @@ -1682,6 +1707,10 @@ def forward( elif qkv_format == "sbhd": # [s*b, h, d] -> [s, b, h, d] out = out.view(-1, ctx.batch_size, *out.shape[-2:]) + 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:]) @@ -1811,10 +1840,12 @@ def forward( nvtx_range_pop(f"{nvtx_label}") + if return_max_logit: + return out_ret, max_logit return out_ret @staticmethod - def backward(ctx, dout): + def backward(ctx, dout, *_args): # pylint: disable=missing-function-docstring # add NVTX range @@ -2522,6 +2553,7 @@ def backward(ctx, dout): None, None, None, + None, ) @@ -2577,6 +2609,7 @@ def forward( attn_bias, deterministic, use_fused_attention, + return_max_logit, window_size, cp_group, cp_stream, @@ -2682,6 +2715,8 @@ def forward( 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): @@ -2712,7 +2747,11 @@ def forward( # [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_]] if use_fused_attention: - out_per_step[i], [softmax_lse_per_step[i], rng_states[i]] = fused_attn_fwd( + ( + out_per_step[i], + [softmax_lse_per_step[i], rng_states[i]], + *max_logit_, + ) = fused_attn_fwd( is_training, max_seqlen_q, max_seqlen_kv_, @@ -2732,7 +2771,10 @@ def forward( cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_per_step[i], window_size=window_size_per_step[i], + return_max_logit=return_max_logit, ) + if return_max_logit: + max_logit_per_step[i] = max_logit_[0] else: fa_forward_args_thd = get_fa_args( True, @@ -2767,14 +2809,22 @@ def forward( if not use_flash_attn_3: rng_states[i] = fa_outputs[3] + 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 return_max_logit: + max_logit = torch.maximum(max_logit, max_logit_per_step[i - 1]) torch.cuda.current_stream().wait_stream(cp_stream) + 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": @@ -2811,10 +2861,12 @@ def forward( ctx.use_fused_attention = use_fused_attention ctx.use_flash_attn_3 = use_flash_attn_3 nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVAllGather.forward") + if return_max_logit: + return out, max_logit return out @staticmethod - def backward(ctx, dout): + def backward(ctx, dout, *_args): # pylint: disable=missing-function-docstring nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVAllGather.backward") cp_size = get_distributed_world_size(ctx.cp_group) @@ -3035,6 +3087,7 @@ def backward(ctx, dout): None, None, None, + None, ) @@ -3065,6 +3118,7 @@ def forward( attn_bias, deterministic, use_fused_attention, + return_max_logit, window_size, fp8, fp8_meta, @@ -3158,6 +3212,7 @@ def forward( 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) @@ -3203,7 +3258,7 @@ def forward( 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 = fused_attn_fwd( + out_, aux_ctx_tensors, *max_logit = fused_attn_fwd( is_training, max_seqlen_q, max_seqlen_kv, @@ -3226,6 +3281,7 @@ def forward( **fp8_meta_kwargs, softmax_type=softmax_type, softmax_offset=softmax_offset, + return_max_logit=return_max_logit, ) if isinstance(out_, Float8Tensor): out_fp8 = out_ @@ -3276,6 +3332,10 @@ def forward( out_ = flash_attn_a2a_communicate( out_, chunk_ids_for_a2a, seq_dim, cp_size, cp_group, cp_stream, False ) + 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": @@ -3362,10 +3422,12 @@ def forward( ctx.S_quantizer = S_quantizer.copy() ctx.S_quantizer.scale = S_quantizer.scale.clone() nvtx_range_pop("transformer_engine.AttnFuncWithCPAndQKVOA2A.forward") + if return_max_logit: + return out_ret, max_logit return out_ret @staticmethod - def backward(ctx, dout): + def backward(ctx, dout, *_args): # pylint: disable=missing-function-docstring nvtx_range_push("transformer_engine.AttnFuncWithCPAndQKVOA2A.backward") cp_size = get_distributed_world_size(ctx.cp_group) @@ -3599,6 +3661,7 @@ def backward(ctx, dout): None, None, None, + None, d_softmax_offset, None, ) @@ -3637,6 +3700,7 @@ def attn_forward_func_with_cp( softmax_offset=None, fp8_output=False, layer_number=1, + return_max_logit=False, ) -> torch.Tensor: """ Attention implementation with context parallelism (CP). CP partitions tensors along the sequence @@ -3784,6 +3848,7 @@ def attn_forward_func_with_cp( attn_bias, deterministic, use_fused_attention, + return_max_logit, ] if cp_comm_type in ["p2p", "a2a+p2p"]: 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 6d9ce9a522..0d1c0b0c05 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 @@ -255,6 +255,12 @@ class DotProductAttention(TransformerEngineBaseModule): where alpha is a learnable parameter in shape [h]. 'off-by-one' and 'learnable' softmax types are also called sink attention ('zero sink' and 'learnable sink'). + return_max_logit: Optional[bool], default = `False` + If true, returns the maximum attention score that can be used in a Muon optimizer to + rescale the Q and K projection weights (see `Muon is Scalable for LLM Training + `_). + max_logit = max(S), where S = mask(Q*K^T*softmax_scale + bias) in shape [b, h, s_q, s_kv], + and max_logit is in shape [h]. Parallelism parameters ---------------------- @@ -311,6 +317,7 @@ def __init__( cp_comm_type: str = "p2p", softmax_scale: Optional[float] = None, softmax_type: str = "vanilla", + return_max_logit: Optional[bool] = False, ) -> None: super().__init__() @@ -394,6 +401,7 @@ def __init__( self.attention_type = attention_type self.attention_dropout = attention_dropout + self.return_max_logit = return_max_logit self.softmax_type = softmax_type if self.softmax_type == "vanilla": @@ -431,6 +439,7 @@ def __init__( deterministic=self.deterministic, **attn_kwargs, softmax_type=self.softmax_type, + return_max_logit=self.return_max_logit, ) self.unfused_attention = UnfusedDotProductAttention( @@ -439,6 +448,7 @@ def __init__( **attn_kwargs, layer_number=layer_number, softmax_type=self.softmax_type, + return_max_logit=self.return_max_logit, ) def remove_extra_states_check(self, incompatible_keys): # pylint: disable=unused-argument @@ -1303,6 +1313,7 @@ def forward( fp8_meta=self.fp8_meta, inference_params=inference_params, softmax_type=self.softmax_type, + return_max_logit=self.return_max_logit, ) global _attention_backends if is_in_onnx_export_mode(): @@ -1502,6 +1513,8 @@ def forward( qkv_layout=qkv_layout, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, attn_mask_type=attn_mask_type, attention_mask=attention_mask, window_size=window_size, @@ -1523,6 +1536,8 @@ def forward( qkv_layout=qkv_layout, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, attn_mask_type=attn_mask_type, attention_mask=attention_mask, window_size=window_size, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 4cb39cda09..51279bd372 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -229,6 +229,8 @@ class AttentionParams: Inference-related parameters. See InferenceParams for details. softmax_type: str, default = "vanilla" The type of softmax operation. See DotProductAttention for details. + return_max_logit: bool, default = `False` + Whether to output max_logit. """ qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor @@ -257,6 +259,7 @@ class AttentionParams: fp8_meta: Union[Dict[str, Any], None] = None inference_params: Optional[InferenceParams] = None softmax_type: str = "vanilla" + return_max_logit: bool = False def __eq__(self, other): """ @@ -330,6 +333,7 @@ def get_attention_backend( fp8_meta = attention_params.fp8_meta inference_params = attention_params.inference_params softmax_type = attention_params.softmax_type + return_max_logit = attention_params.return_max_logit # Run config logger = logging.getLogger("DotProductAttention") @@ -477,6 +481,20 @@ def get_attention_backend( logger.debug("Disabling FusedAttention for FP8 current scaling with cuDNN < 9.14.0") use_fused_attention = False + # Filter: Return max_logit + if return_max_logit: + if use_flash_attention: + use_flash_attention = False + logger.debug("Disabling FlashAttention for max_logit") + if use_fused_attention and qkv_format == "thd": + use_fused_attention = False + logger.debug("Disabling FusedAttention for max_logit with qkv_format = thd") + if fp8 and fp8_meta["recipe"].fp8_dpa: + use_flash_attention = False + use_fused_attention = False + use_unfused_attention = False + logger.debug("Disabling all backends for max_logit with FP8 attention") + # Filter: KV cache # backend | precision | KV cache | architecture | qkv_format | page_size # --------------------------------------------------------------------------------------- @@ -913,6 +931,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt head_dim_v, window_size[0], window_size[1], + return_max_logit, ) if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug("Disabling FusedAttention as no backend supports the provided input") @@ -1649,6 +1668,78 @@ def backward(ctx, grad_output): return None, None, _pack_tensor(indices, grad_output) +class ConvertTHDtoBSHD(torch.autograd.Function): + """ + Convert a tensor from qkv_format = thd to qkv_format = bshd. + """ + + @staticmethod + def forward(ctx, thd_tensor, cu_seqlens, max_seqlen): + # pylint: disable=missing-function-docstring + batch_size = cu_seqlens.shape[0] - 1 + if not thd_tensor.is_contiguous(): + thd_tensor = thd_tensor.contiguous() + bshd_tensor = tex.convert_thd_to_bshd( + thd_tensor, + cu_seqlens, + batch_size, + max_seqlen, + ) + ctx.save_for_backward(cu_seqlens) + ctx.num_tokens = thd_tensor.shape[0] + return bshd_tensor + + @staticmethod + def backward(ctx, bshd_tensor): + # pylint: disable=missing-function-docstring + (cu_seqlens,) = ctx.saved_tensors + if not bshd_tensor.is_contiguous(): + bshd_tensor = bshd_tensor.contiguous() + thd_tensor = tex.convert_bshd_to_thd( + bshd_tensor, + cu_seqlens, + ctx.num_tokens, + ) + return thd_tensor, None, None + + +class ConvertBSHDtoTHD(torch.autograd.Function): + """ + Convert a tensor from qkv_format = bshd to qkv_format = thd. + """ + + @staticmethod + def forward(ctx, bshd_tensor, cu_seqlens): + # pylint: disable=missing-function-docstring + num_tokens = cu_seqlens[-1] + max_seqlen = bshd_tensor.shape[1] + if not bshd_tensor.is_contiguous(): + bshd_tensor = bshd_tensor.contiguous() + thd_tensor = tex.convert_bshd_to_thd( + bshd_tensor, + cu_seqlens, + num_tokens, + ) + ctx.save_for_backward(cu_seqlens) + ctx.max_seqlen = max_seqlen + return thd_tensor + + @staticmethod + def backward(ctx, thd_tensor): + # pylint: disable=missing-function-docstring + (cu_seqlens,) = ctx.saved_tensors + batch_size = cu_seqlens.shape[0] - 1 + if not thd_tensor.is_contiguous(): + thd_tensor = thd_tensor.contiguous() + bshd_tensor = tex.convert_thd_to_bshd( + thd_tensor, + cu_seqlens, + batch_size, + ctx.max_seqlen, + ) + return bshd_tensor, None + + def get_qkv_format( qkv_layout: str = "bshd_bshd_bshd", inference_params: InferenceParams = None, diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index f80c001a1d..eb43c75f6b 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -139,6 +139,7 @@ def fused_attn_fwd( window_size: Tuple[int, int] = (-1, -1), rng_gen: torch.Generator = None, softmax_offset: torch.Tensor = None, + return_max_logit: bool = False, ) -> Tuple[Union[torch.Tensor, None], ...]: """Fused Attention FWD for separate QKV input. @@ -216,6 +217,8 @@ def fused_attn_fwd( softmax_offset: torch.Tensor, default = None softmax offset tensor in shape [1, h_q, 1, 1]. See softmax_type in DotProductAttention for details. + return_max_logit: bool, default = False + whether to return the maximum attention score Returns ---------- @@ -246,6 +249,7 @@ def fused_attn_fwd( rng_state: torch.Tensor, optional, if backend is not F16_max512_seqlen 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 """ if attn_scale is None: @@ -315,8 +319,22 @@ def fused_attn_fwd( softmax_offset, rng_gen, rng_elts_per_thread, + return_max_logit, ) + if return_max_logit: + qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] + # thd: output_tensors: out [tq, h, d], Max [tq, h, 1], Sum_Exp [tq, h, 1] + # bshd: output_tensors: out [b, sq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] + # sbhd: output_tensors: out [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] + stats = output_tensors[1] + torch.log(output_tensors[2]) + amax_dims = (0, 2) if qkv_format == "thd" else (0, 2, 3) + # Max -> max_logit [h] + max_logit = torch.amax(output_tensors[1], dim=amax_dims).to(dtype=output_tensors[0].dtype) + aux_ctx_tensors = [stats] + aux_ctx_tensors.extend(output_tensors[3:]) + return output_tensors[0], aux_ctx_tensors, max_logit + # out, aux_ctx_tensors return output_tensors[0], output_tensors[1:] diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index d86a96959c..79fb798422 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -76,7 +76,7 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float p_dropout, 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, int64_t window_size_left, - int64_t window_size_right); + int64_t window_size_right, bool return_max_logit); std::pair quantizer_helper(py::handle quantizer, const std::vector &shape, DType dtype, @@ -94,7 +94,7 @@ std::vector fused_attn_fwd( const std::optional page_table_k, const std::optional page_table_v, py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread); + size_t rng_elts_per_thread, bool return_max_logit); std::vector fused_attn_bwd( size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 344bc4ab0b..f66c8aa619 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -45,11 +45,12 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float p_dropout, 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, int64_t window_size_left, - int64_t window_size_right) { + int64_t window_size_right, bool return_max_logit) { NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, attn_mask_type, softmax_type, p_dropout, num_attn_heads, num_gqa_groups, - max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right); + max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right, + return_max_logit); return fused_attention_backend; } @@ -106,7 +107,7 @@ std::vector fused_attn_fwd( const std::optional page_table_k, const std::optional page_table_v, py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread) { + size_t rng_elts_per_thread, bool return_max_logit) { auto none = py::none(); // create QKV tensor wrappers @@ -228,8 +229,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, - attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], workspace.data(), at::cuda::getCurrentCUDAStream()); + return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, + softmax_type, window_size[0], window_size[1], workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace and auxiliary output tensors @@ -249,7 +251,9 @@ std::vector fused_attn_fwd( }; // allocate memory for nvte_aux_tensor_pack.tensors // f16_max512 : S [b, h, sq, skv] - // f16_arbitrary: S [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] + // 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: Max [b, h, sq, 1], Sum_Exp [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] size_t i = 0; at::Tensor output_tensor; @@ -258,8 +262,8 @@ 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 - if (qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { + // fp8 has an additional softmax stats tensor, ZInv; return_max_logit=true has an additional Sum_Exp tensor + if (return_max_logit || qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { 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); @@ -285,8 +289,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, - attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size[0], - window_size[1], workspace.data(), at::cuda::getCurrentCUDAStream()); + return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, + softmax_type, window_size[0], window_size[1], workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // destroy tensor wrappers, but not allocated memory From d2945c6a571e3978677614d1fe08779966a5a4ef Mon Sep 17 00:00:00 2001 From: Tong Liu Date: Mon, 27 Oct 2025 16:10:58 +0800 Subject: [PATCH 025/521] [PyTorch] Use dummy wgrad in GroupedLinear (#2305) dummy wgrad Signed-off-by: tongliu Signed-off-by: Xin Yao Co-authored-by: Xin Yao --- .../pytorch/module/grouped_linear.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index bba97554c5..4d6b2f23b9 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -13,6 +13,7 @@ from transformer_engine.common.recipe import Recipe from .base import ( + get_dummy_wgrad, get_multi_stream_cublas_workspace, TransformerEngineBaseModule, _2X_ACC_FPROP, @@ -447,18 +448,15 @@ def handle_custom_ddp_from_mcore(weight, wgrad): ): weight.grad_added_to_main_grad = True if getattr(weight, "zero_out_wgrad", False): - wgrad = torch.zeros( - weight.main_grad.shape, - dtype=weight.dtype, - device=torch.cuda.current_device(), - requires_grad=False, + wgrad = get_dummy_wgrad( + list(weight.main_grad.shape), + weight.dtype, + zero=True, ) else: - wgrad = torch.empty( - weight.main_grad.shape, - dtype=weight.dtype, - device=torch.cuda.current_device(), - requires_grad=False, + wgrad = get_dummy_wgrad( + list(weight.main_grad.shape), + weight.dtype, ) elif ctx.fuse_wgrad_accumulation: wgrad = None From d7c9777e611a90337fffb0482a62ee2b60ef0353 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 27 Oct 2025 18:13:30 -0400 Subject: [PATCH 026/521] Remove `nvidia-mathdx` dependency (#2295) * Remove nvidia-mathdx dep Signed-off-by: Kirthi Shankar Sivamani * Fix SR Signed-off-by: Kirthi Shankar Sivamani * Add comment Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- .github/workflows/build.yml | 8 +- build_tools/wheel_utils/build_wheels.sh | 2 +- pyproject.toml | 3 +- transformer_engine/common/CMakeLists.txt | 23 ---- .../hadamard_transform_cast_fusion.cu | 18 +-- ...quantize_transpose_vector_blockwise_fp4.cu | 27 ++--- transformer_engine/common/util/curanddx.hpp | 106 ++++++++++++++++++ .../common/util/nvfp4_transpose.cuh | 28 +++-- 8 files changed, 145 insertions(+), 70 deletions(-) create mode 100644 transformer_engine/common/util/curanddx.hpp diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 506bc83f08..f40b281895 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -19,7 +19,7 @@ jobs: run: | apt-get update apt-get install -y git python3.9 pip cudnn9-cuda-12 - pip install cmake==3.21.0 pybind11[global] ninja nvidia-mathdx==25.1.1 + pip install cmake==3.21.0 pybind11[global] ninja - name: 'Checkout' uses: actions/checkout@v3 with: @@ -43,7 +43,7 @@ jobs: run: | apt-get update apt-get install -y git python3.9 pip cudnn9-cuda-12 - pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript nvidia-mathdx==25.1.1 + pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript - name: 'Checkout' uses: actions/checkout@v3 with: @@ -63,7 +63,7 @@ jobs: options: --user root steps: - name: 'Dependencies' - run: pip install pybind11[global] nvidia-mathdx==25.1.1 + run: pip install pybind11[global] - name: 'Checkout' uses: actions/checkout@v3 with: @@ -83,7 +83,7 @@ jobs: options: --user root steps: - name: 'Dependencies' - run: pip install torch pybind11[global] einops onnxscript nvidia-mathdx==25.1.1 + run: pip install torch pybind11[global] einops onnxscript - name: 'Checkout' uses: actions/checkout@v3 with: diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index 954a8f1c67..d0055b791d 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -23,7 +23,7 @@ git checkout $TARGET_BRANCH git submodule update --init --recursive # Install deps -/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel nvidia-mathdx==25.1.1 +/opt/python/cp310-cp310/bin/pip install cmake pybind11[global] ninja setuptools wheel if $BUILD_METAPACKAGE ; then cd /TransformerEngine diff --git a/pyproject.toml b/pyproject.toml index 8692ad9610..35a7c20727 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -3,8 +3,7 @@ # See LICENSE for license information. [build-system] -requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "nvidia-mathdx==25.1.1", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1"] +requires = ["setuptools>=61.0", "cmake>=3.21", "wheel", "pybind11[global]", "ninja", "pip", "torch>=2.1", "jax>=0.5.0", "flax>=0.7.1"] # Use legacy backend to import local packages in setup.py build-backend = "setuptools.build_meta:__legacy__" - diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 175abd3530..e388dd794b 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -98,28 +98,6 @@ set(CUTLASS_TOOLS_INCLUDE_DIR # Python find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) -# NVIDIA MathDX include directory (from Python package install location) -if(NOT DEFINED MATHDX_INCLUDE_DIR) - execute_process( - COMMAND ${Python_EXECUTABLE} -m pip show nvidia-mathdx - OUTPUT_VARIABLE _PIP_SHOW_MATHDX - ERROR_VARIABLE _PIP_SHOW_MATHDX_ERR - RESULT_VARIABLE _PIP_SHOW_MATHDX_RES - OUTPUT_STRIP_TRAILING_WHITESPACE) - if(NOT _PIP_SHOW_MATHDX_RES EQUAL 0) - message(FATAL_ERROR "Failed to query 'nvidia-mathdx' with pip (using ${Python_EXECUTABLE}): ${_PIP_SHOW_MATHDX_ERR}") - endif() - string(REGEX MATCH "Location: ([^\n\r]+)" _MATHDX_LOC_MATCH "${_PIP_SHOW_MATHDX}") - if(NOT _MATHDX_LOC_MATCH) - message(FATAL_ERROR "Could not parse installation location for 'nvidia-mathdx'. Output was:\n${_PIP_SHOW_MATHDX}") - endif() - set(MATHDX_LOCATION "${CMAKE_MATCH_1}") - set(MATHDX_INCLUDE_DIR "${MATHDX_LOCATION}/nvidia/mathdx/include") -endif() -if(NOT EXISTS "${MATHDX_INCLUDE_DIR}") - message(FATAL_ERROR "MATHDX include directory not found at ${MATHDX_INCLUDE_DIR}. Set MATHDX_INCLUDE_DIR or ensure 'nvidia-mathdx' is installed for ${Python_EXECUTABLE}.") -endif() - # Configure Transformer Engine library include_directories(${PROJECT_SOURCE_DIR}/..) set(transformer_engine_SOURCES) @@ -263,7 +241,6 @@ target_link_libraries(transformer_engine PUBLIC target_include_directories(transformer_engine PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) -target_include_directories(transformer_engine PRIVATE ${MATHDX_INCLUDE_DIR}) target_include_directories(transformer_engine SYSTEM PRIVATE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}/cccl) target_include_directories(transformer_engine PRIVATE "${CUDNN_FRONTEND_INCLUDE_DIR}") 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 263a32623e..12f02dba6b 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -19,9 +19,9 @@ #include "common/common.h" #include "common/util/cuda_runtime.h" +#include "common/util/curanddx.hpp" #include "common/util/ptx.cuh" #include "common/utils.cuh" -#include "curanddx.hpp" #include "cutlass/arch/barrier.h" #include "cutlass/cutlass.h" #include "cutlass/gemm/collective/builders/sm100_common.inl" @@ -38,15 +38,6 @@ namespace transformer_engine { namespace detail { namespace { -// Define a cuRANDDx descriptor -// Note curanddx::PhiloxRounds<4> means 4 rounds of philox4_32. If the operator is not specified, it will be default to 10. -// curanddx::SM<800>() does NOT mean the code can only run on SM 800. The operator is used for do some internal checks, e.g., -// if shared memory, if needed, is enough for the described problem, usually not applicable. - -// curanddx doc: https://docs.nvidia.com/cuda/curanddx/index.html -using RNG = decltype(curanddx::Generator() + curanddx::PhiloxRounds<10>() + curanddx::SM<800>() + curanddx::Thread()); - - using namespace cute; using cute::Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor @@ -502,8 +493,9 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, // Initialize RNG for tile const size_t rng_sequence = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256; - RNG rng(rng_seed, rng_sequence, rng_offset); - curanddx::uniform_bits dist; + + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = uint4{0, 0, 0, 0}; CUTLASS_PRAGMA_UNROLL @@ -511,7 +503,7 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, auto acc_scale = cutlass::minimum_with_nan_propagation{}(acc_scales[v], cutlass::platform::numeric_limits::max()); // auto acc_scale = acc_scales[v]; if constexpr (kEnableStochasticRounding) { - random_uint4 = dist.generate4(rng); + random_uint4 = rng.generate4(); output_frgs[v] = StochasticNumericConverter( cutlass::multiplies>{}( compute_frgs[v], 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 4735fdcbe0..b49a54fbdb 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -17,9 +17,9 @@ #include "common/common.h" #include "common/recipe/recipe_common.cuh" #include "common/transpose/cast_transpose.h" +#include "common/util/curanddx.hpp" #include "common/util/ptx.cuh" #include "common/utils.cuh" -#include "curanddx.hpp" namespace transformer_engine { @@ -33,14 +33,6 @@ using std::uint8_t; using transformer_engine::detail::TypeExtrema; -// Define a cuRANDDx descriptor -// Note curanddx::PhiloxRounds<4> means 4 rounds of philox4_32. If the operator is not specified, it will be default to 10. -// curanddx::SM<800>() does NOT mean the code can only run on SM 800. The operator is used for do some internal checks, e.g., -// if shared memory, if needed, is enough for the described problem, usually not applicable. -// curanddx doc: https://docs.nvidia.com/cuda/curanddx/index.html -using RNG = decltype(curanddx::Generator() + curanddx::PhiloxRounds<10>() + - curanddx::SM<800>() + curanddx::Thread()); - // clang-format off /* @@ -209,12 +201,15 @@ __device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_ return global_encode_scale; } -__device__ __forceinline__ uint32_t get_rbits(RNG& rng, uint4& random_uint4, int& rnd_idx) { +__device__ __forceinline__ uint32_t +get_rbits(transformer_engine::curanddx::detail::philox4x32_native_state<10>& + rng, // philox4x32_native_state<10>: 10 rounds of philox4_32 + uint4& random_uint4, int& rnd_idx) { if (rnd_idx == 4) { rnd_idx = 0; - curanddx::uniform_bits dist; - random_uint4 = dist.generate4(rng); + random_uint4 = rng.generate4(); } + // Treat uint4 as an array of 4x uint32_t elements for indexing const uint32_t* const rbits_arr = reinterpret_cast(&random_uint4); const uint32_t rbits = rbits_arr[rnd_idx++]; @@ -348,9 +343,11 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo threadIdx.x + block_idx_x * kThreadsPerBlock + block_idx_y * gridDim.x * kThreadsPerBlock; const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - RNG rng(rng_seed, rng_sequence, rng_offset); - curanddx::uniform_bits dist; - uint4 random_uint4 = kApplyStochasticRounding ? dist.generate4(rng) : uint4{0, 0, 0, 0}; + + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = kApplyStochasticRounding ? rng.generate4() : uint4{0, 0, 0, 0}; + int rnd_idx = 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x diff --git a/transformer_engine/common/util/curanddx.hpp b/transformer_engine/common/util/curanddx.hpp new file mode 100644 index 0000000000..4d7c90a019 --- /dev/null +++ b/transformer_engine/common/util/curanddx.hpp @@ -0,0 +1,106 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_UTIL_CURANDDX_HPP_ +#define TRANSFORMER_ENGINE_COMMON_UTIL_CURANDDX_HPP_ + +namespace transformer_engine { +namespace curanddx { +namespace detail { + +inline constexpr unsigned int philox4x32_w32_0 = 0x9E3779B9U; +inline constexpr unsigned int philox4x32_w32_1 = 0xBB67AE85U; +inline constexpr unsigned int philox4x32_m4x32_0 = 0xD2511F53U; +inline constexpr unsigned int philox4x32_m4x32_1 = 0xCD9E8D57U; + +__forceinline__ __device__ unsigned int mulhilo32(unsigned int a, unsigned int b, + unsigned int* hip) { + *hip = __umulhi(a, b); + return a * b; +} + +__forceinline__ __device__ uint4 single_round(uint4 ctr, uint2 key) { + unsigned int hi0; + unsigned int hi1; + unsigned int lo0 = mulhilo32(philox4x32_m4x32_0, ctr.x, &hi0); + unsigned int lo1 = mulhilo32(philox4x32_m4x32_1, ctr.z, &hi1); + + uint4 ret = {hi1 ^ ctr.y ^ key.x, lo1, hi0 ^ ctr.w ^ key.y, lo0}; + return ret; +} + +template +__forceinline__ __device__ uint4 multiple_rounds(uint4 c, uint2 k) { + for (unsigned int i = 0; i < Rounds - 1; i++) { + c = single_round(c, k); // 1 + k.x += philox4x32_w32_0; + k.y += philox4x32_w32_1; + } + return single_round(c, k); // Rounds +} + +template +struct philox4x32_native_state { + static constexpr unsigned int rounds = Rounds; + + uint4 ctr; + uint2 key; + + __forceinline__ __device__ void philox_state_incr() { + if (++ctr.x) return; + if (++ctr.y) return; + if (++ctr.z) return; + ++ctr.w; + } + + __forceinline__ __device__ void philox_state_incr(size_t n) { + unsigned int nlo = (unsigned int)(n); + unsigned int nhi = (unsigned int)(n >> 32); + + ctr.x += nlo; + if (ctr.x < nlo) nhi++; + + ctr.y += nhi; + if (nhi <= ctr.y) return; + if (++ctr.z) return; + ++ctr.w; + } + + __forceinline__ __device__ void philox_state_incr_hi(size_t n) { + unsigned int nlo = (unsigned int)(n); + unsigned int nhi = (unsigned int)(n >> 32); + + ctr.z += nlo; + if (ctr.z < nlo) nhi++; + + ctr.w += nhi; + } + + // offset is the total # of 128bits generated with a single generate4() call + __forceinline__ __device__ void skip_offset(size_t n) { philox_state_incr(n); } + + __forceinline__ __device__ void skip_subsequence(size_t n) { philox_state_incr_hi(n); } + + __forceinline__ __device__ void init(size_t seed, size_t subsequence, size_t offset) { + ctr = make_uint4(0, 0, 0, 0); + key.x = (unsigned int)seed; + key.y = (unsigned int)(seed >> 32); + + skip_subsequence(subsequence); + skip_offset(offset); + } + + __forceinline__ __device__ uint4 generate4() { + auto tmp = multiple_rounds(ctr, key); + philox_state_incr(); + return tmp; + } +}; +} // namespace detail +} // namespace curanddx +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_UTIL_CURANDDX_HPP_ diff --git a/transformer_engine/common/util/nvfp4_transpose.cuh b/transformer_engine/common/util/nvfp4_transpose.cuh index 45fa29f0e9..629520aeb7 100644 --- a/transformer_engine/common/util/nvfp4_transpose.cuh +++ b/transformer_engine/common/util/nvfp4_transpose.cuh @@ -32,9 +32,6 @@ namespace transformer_engine { #if FP4_TYPE_SUPPORTED namespace nvfp4_transpose { -using RNG = decltype(curanddx::Generator() + curanddx::PhiloxRounds<10>() + - curanddx::SM<800>() + curanddx::Thread()); - using namespace ptx; using nvfp4_scale_t = fp8e4m3; @@ -139,12 +136,15 @@ __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const return global_encode_scale; } -__device__ __forceinline__ uint32_t get_rbits(RNG &rng, uint4 &random_uint4, int &rnd_idx) { +__device__ __forceinline__ uint32_t +get_rbits(transformer_engine::curanddx::detail::philox4x32_native_state<10> + &rng, // philox4x32_native_state<10>: 10 rounds of philox4_32 + uint4 &random_uint4, int &rnd_idx) { if (rnd_idx == 4) { rnd_idx = 0; - curanddx::uniform_bits dist; - random_uint4 = dist.generate4(rng); + random_uint4 = rng.generate4(); } + // Treat uint4 as an array of 4x uint32_t elements for indexing const uint32_t *const rbits_arr = reinterpret_cast(&random_uint4); const uint32_t rbits = rbits_arr[rnd_idx++]; @@ -363,9 +363,11 @@ __global__ void __launch_bounds__(THREADS_NUM) threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - RNG rng(rng_seed, rng_sequence, rng_offset); - curanddx::uniform_bits dist; - uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? dist.generate4(rng) : uint4{0, 0, 0, 0}; + + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; + int rnd_idx = 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x @@ -874,9 +876,11 @@ __global__ void __launch_bounds__(THREADS_NUM) threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - RNG rng(rng_seed, rng_sequence, rng_offset); - curanddx::uniform_bits dist; - uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? dist.generate4(rng) : uint4{0, 0, 0, 0}; + + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; + int rnd_idx = 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x From a019c80a403fc0f596210237ccf6b954e392c260 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 27 Oct 2025 18:27:52 -0400 Subject: [PATCH 027/521] Submodule checkout during setup (#2293) * Add checks to submodule during setup and automatically checkout Signed-off-by: Kirthi Shankar Sivamani * fix import and formatting Signed-off-by: Kirthi Shankar Sivamani * provide envvar to skip submodule init in setup Signed-off-by: Kirthi Shankar Sivamani * Fix formatting Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- setup.py | 57 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/setup.py b/setup.py index a820265c30..ce3805d2eb 100644 --- a/setup.py +++ b/setup.py @@ -6,6 +6,8 @@ from importlib import metadata import os +import shutil +import subprocess import time from pathlib import Path from typing import List, Tuple @@ -126,9 +128,64 @@ def setup_requirements() -> Tuple[List[str], List[str]]: return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] +def git_check_submodules() -> None: + """ + Attempt to checkout git submodules automatically during setup. + + This runs successfully only if the submodules are + either in the correct or uninitialized state. + + Note to devs: With this, any updates to the submodules itself, e.g. moving to a newer + commit, must be commited before build. This also ensures that stale submodules aren't + being silently used by developers. + """ + + # Provide an option to skip these checks for development. + if bool(int(os.getenv("NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD", "0"))): + return + + # Require git executable. + if shutil.which("git") is None: + return + + # Require a .gitmodules file. + if not (current_file_path / ".gitmodules").exists(): + return + + try: + submodules = subprocess.check_output( + ["git", "submodule", "status", "--recursive"], + cwd=str(current_file_path), + text=True, + ).splitlines() + + for submodule in submodules: + # '-' start is for an uninitialized submodule. + # ' ' start is for a submodule on the correct commit. + assert submodule[0] in ( + " ", + "-", + ), ( + "Submodules are initialized incorrectly. If this is intended, set the " + "environment variable `NVTE_SKIP_SUBMODULE_CHECKS_DURING_BUILD` to a " + "non-zero value to skip these checks during development. Otherwise, " + "run `git submodule update --init --recursive` to checkout the correct" + " submodule commits." + ) + + subprocess.check_call( + ["git", "submodule", "update", "--init", "--recursive"], + cwd=str(current_file_path), + ) + except subprocess.CalledProcessError: + return + + if __name__ == "__main__": __version__ = te_version() + git_check_submodules() + with open("README.rst", encoding="utf-8") as f: long_description = f.read() From 4cf2f12b408540513f857eb6f516eebef757979b Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Mon, 27 Oct 2025 16:49:28 -0700 Subject: [PATCH 028/521] Change the pyTorch installation to CUDA 13 in Build All GitHub action (#2308) Change the pyTorch installation to CUDA 13 in Build All GitHub action to match the version in the JAX container Signed-off-by: Przemek Tredak --- .github/workflows/build.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index f40b281895..42c5f0342e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -83,7 +83,9 @@ jobs: options: --user root steps: - name: 'Dependencies' - run: pip install torch pybind11[global] einops onnxscript + run: | + pip install pybind11[global] einops onnxscript + pip install torch --index-url https://download.pytorch.org/whl/cu130 - name: 'Checkout' uses: actions/checkout@v3 with: From a8e4346ec6630bf808d8904a4fdb19cb6f54b48b Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 28 Oct 2025 09:19:50 -0400 Subject: [PATCH 029/521] [JAX] Use TE quantization when TE fused norm is disable (#2303) * jax norm + te quant Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- .../jax/cpp_extensions/normalization.py | 38 +++++++++++++------ 1 file changed, 26 insertions(+), 12 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/normalization.py b/transformer_engine/jax/cpp_extensions/normalization.py index 90ab5fb7fe..d09ce7ef74 100644 --- a/transformer_engine/jax/cpp_extensions/normalization.py +++ b/transformer_engine/jax/cpp_extensions/normalization.py @@ -27,7 +27,7 @@ NamedSharding, get_cudnn_version, ) -from .quantization import _quantize_dbias_impl, AmaxScope +from .quantization import quantize, AmaxScope from ..sharding import ( all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp_tpsp, @@ -945,7 +945,7 @@ def layernorm_fwd( beta: jnp.ndarray, zero_centered_gamma: bool, epsilon: float, - quantizer: Optional[Quantizer], + quantizer: Optional[Quantizer] = None, amax_scope: AmaxScope = AmaxScope.LOCAL, transpose_batch_sequence: bool = False, output_amax_when_no_scaling: bool = False, @@ -975,7 +975,16 @@ def layernorm_fwd( - Reciprocal of the standard deviation of the input tensor. Shape: (..., 1) """ if not NormFwdPrimitive.enabled(): - return _jax_layernorm(x, gamma, beta, zero_centered_gamma, epsilon, quantizer) + output, mu, rsigma = _jax_layernorm(x, gamma, beta, zero_centered_gamma, epsilon) + if quantizer is not None: + output = quantize( + output, + quantizer, + flatten_axis=-1, + amax_scope=amax_scope, + transpose_batch_sequence=transpose_batch_sequence, + ) + return (output, mu, rsigma) # TE/common does not support normalization with colwise only quantization yet if quantizer is not None and quantizer.q_layout == QuantizeLayout.COLWISE: @@ -1029,7 +1038,7 @@ def layernorm_fwd( transpose_batch_sequence=transpose_batch_sequence, output_amax_when_no_scaling=False, ) - out, _ = _quantize_dbias_impl( + out, _ = quantize( out, quantizer, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence ) return out, mu, rsigma @@ -1050,11 +1059,9 @@ def layernorm_fwd( transpose_batch_sequence=transpose_batch_sequence, output_amax_when_no_scaling=True, ) - out, _ = _quantize_dbias_impl( + out = quantize( out, - is_dbias=False, quantizer=quantizer, - dq_dtype=x.dtype, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, ) @@ -1219,7 +1226,16 @@ def rmsnorm_fwd( Shape: (..., 1) """ if not NormFwdPrimitive.enabled(): - return _jax_rmsnorm(x, gamma, zero_centered_gamma, epsilon, quantizer) + output, rsigma = _jax_rmsnorm(x, gamma, zero_centered_gamma, epsilon) + if quantizer is not None: + output = quantize( + output, + quantizer, + flatten_axis=-1, + amax_scope=amax_scope, + transpose_batch_sequence=transpose_batch_sequence, + ) + return (output, rsigma) # TE/common does not support normalization with colwise only quantization yet if quantizer is not None and quantizer.q_layout == QuantizeLayout.COLWISE: @@ -1274,7 +1290,7 @@ def rmsnorm_fwd( transpose_batch_sequence=transpose_batch_sequence, output_amax_when_no_scaling=False, ) - out, _ = _quantize_dbias_impl( + out = quantize( out.data, quantizer, amax_scope=amax_scope, @@ -1297,11 +1313,9 @@ def rmsnorm_fwd( transpose_batch_sequence=transpose_batch_sequence, output_amax_when_no_scaling=True, ) - out, _ = _quantize_dbias_impl( + out = quantize( out, - is_dbias=False, quantizer=quantizer, - dq_dtype=x.dtype, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, ) From c6cbcc85368adc03b23d4e55d1f258b4de19316a Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Wed, 29 Oct 2025 13:21:33 -0700 Subject: [PATCH 030/521] [Pytorch] Integrate GPT OSS Swiglu in TransformerLayer (#2312) * changes working Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add support for onnx, minor comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * greptile review comments Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/transformer.py Co-authored-by: Przemyslaw Tredak Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/module/layernorm_mlp.py Co-authored-by: Przemyslaw Tredak Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/transformer.py Co-authored-by: Przemyslaw Tredak Signed-off-by: vthumbe1503 * address review comments Signed-off-by: Varun Thumbe * address review comments Signed-off-by: Varun Thumbe * revert the name change 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> Co-authored-by: Przemyslaw Tredak --- tests/pytorch/test_onnx_export.py | 2 +- tests/pytorch/test_sanity.py | 4 +- .../pytorch/module/layernorm_mlp.py | 56 +++++++++++++++---- transformer_engine/pytorch/transformer.py | 9 ++- 4 files changed, 57 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index f8b4d7481d..2ce6eb82bb 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -68,7 +68,7 @@ fp8_recipes.append(recipe.Float8CurrentScaling()) fp8_recipes.append(None) -supported_activations = ["gelu", "relu", "reglu", "geglu", "swiglu"] +supported_activations = ["gelu", "relu", "reglu", "geglu", "swiglu", "clamped_swiglu"] all_normalizations = ["LayerNorm", "RMSNorm"] diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index e283842ec6..f12e80d4c3 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -122,6 +122,7 @@ def nvfp4_vanilla(): "sreglu", "silu", "swiglu", + "clamped_swiglu", ] all_normalizations = ["LayerNorm", "RMSNorm"] @@ -547,7 +548,7 @@ def test_sanity_layernorm_mlp( sigma = 0.023 init_method = init_method_normal(sigma) output_layer_init_method = scaled_init_method_normal(sigma, config.num_layers) - + activation_params = None if activation != "clamped_swiglu" else {"limit": 7.0, "alpha": 1.702} block = LayerNormMLP( config.hidden_size, 4 * config.hidden_size, @@ -555,6 +556,7 @@ def test_sanity_layernorm_mlp( output_layer_init_method=output_layer_init_method, zero_centered_gamma=zero_centered_gamma, activation=activation, + activation_params=activation_params, normalization=normalization, params_dtype=dtype, device="cuda", diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ccf5dc0953..889f545c1e 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -99,6 +99,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): "sreglu": (tex.sreglu, tex.dsreglu, None), "silu": (tex.silu, tex.dsilu, None), "swiglu": (tex.swiglu, tex.dswiglu, None), + "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), } if recipe.delayed() or recipe.mxfp8(): # Delayed scaling, fusion supported list: [tex.dbias_dgelu, tex.dbias_drelu, tex.dbias_dqgelu, tex.dbias_dsrelu] @@ -114,6 +115,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): "sreglu": (tex.sreglu, tex.dsreglu, None), "silu": (tex.silu, tex.dsilu, tex.dbias_dsilu), "swiglu": (tex.swiglu, tex.dswiglu, None), + "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), } # no activation fusion written yet # Per-tensor current scaling or fp8 blockwise scaling or custom quantization: [] @@ -135,6 +137,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): "sreglu": (tex.sreglu, tex.dsreglu, None), "silu": (tex.silu, tex.dsilu, None), "swiglu": (tex.swiglu, tex.dswiglu, None), + "clamped_swiglu": (tex.clamped_swiglu, tex.clamped_dswiglu, None), } raise NotImplementedError(f"Unhandled recipe type {recipe}") @@ -199,6 +202,7 @@ def forward( bwd_ln_sm_margin: int, zero_centered_gamma: bool, activation: str, + activation_params: Optional[dict], normalization: str, ub_overlap_ag: bool, ub_overlap_rs: bool, @@ -440,6 +444,7 @@ def forward( # ACTIVATION - sometimes activation is fused with the GEMM above. fc1_out_without_bias = None + act_params = activation_params or {} if bias_gelu_fusion: fc1_out = None @@ -449,7 +454,7 @@ def forward( act_out, _, fc1_out, _ = fc1_outputs elif debug: fc1_out, *_ = fc1_outputs - act_out = activation_func(fc1_out, None) + act_out = activation_func(fc1_out, None, **act_params) act_out = fc2_input_quantizer(act_out) else: fc1_out, *_ = fc1_outputs @@ -457,19 +462,19 @@ def forward( recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_block_scaling(): # tex.quantize does not support GELU fusion for blockwise - act_out = activation_func(fc1_out, None) + act_out = activation_func(fc1_out, None, **act_params) act_out = tex.quantize(act_out, fc2_input_quantizer) elif recipe.custom(): # tex.quantize does not support custom quantizers - act_out = activation_func(fc1_out, None) + act_out = activation_func(fc1_out, None, **act_params) act_out = fc2_input_quantizer(act_out) else: - act_out = activation_func(fc1_out, fc2_input_quantizer) + act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) else: if fp8_calibration: - act_out = activation_func(fc1_out, None) + act_out = activation_func(fc1_out, None, **act_params) else: - act_out = activation_func(fc1_out, fc2_input_quantizer) + act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) if not is_grad_enabled: clear_tensor_data(fc1_out) @@ -624,6 +629,7 @@ def forward( ctx.device = device ctx.activation_dtype = activation_dtype ctx.activation = activation + ctx.activation_params = activation_params ctx.fp8 = fp8 ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation @@ -1002,6 +1008,7 @@ def fc2_wgrad_gemm( # -------------------------------------------------- # bias computation + act_params = ctx.activation_params or {} fc1_bias_grad = None fuse_gemm_and_bias_fc1_wgrad = False if ctx.fc1_grad_output_quantizer is not None: @@ -1015,7 +1022,7 @@ def fc2_wgrad_gemm( dact = ctx.fc1_grad_output_quantizer(dact) elif ctx.debug: dact_func = _act_func(ctx.activation)[1] - dact = dact_func(fc2_dgrad, fc1_out.to(ctx.activation_dtype), None) + dact = dact_func(fc2_dgrad, fc1_out.to(ctx.activation_dtype), None, **act_params) fc1_bias_grad = dact.sum(dim=0) dact = ctx.fc1_grad_output_quantizer(dact) elif ( @@ -1027,7 +1034,10 @@ def fc2_wgrad_gemm( ctx.activation, ctx.fp8_recipe if ctx.fp8 else None )[2] fc1_bias_grad, dact = dbias_dact_quantize_func( - fc2_dgrad, fc1_out.to(ctx.activation_dtype), ctx.fc1_grad_output_quantizer + fc2_dgrad, + fc1_out.to(ctx.activation_dtype), + ctx.fc1_grad_output_quantizer, + **act_params, ) # quantize bgrad gelu fused else: # Fusion: gemm + gelu, @@ -1036,7 +1046,7 @@ def fc2_wgrad_gemm( ctx.activation, ctx.fp8_recipe if ctx.fp8 else None )[1] dact = activation_func_bwd( - fc2_dgrad, fc1_out.to(ctx.activation_dtype), None + fc2_dgrad, fc1_out.to(ctx.activation_dtype), None, **act_params ) # activation in high precision if ctx.fp8: @@ -1401,6 +1411,7 @@ def fc1_wgrad_gemm( None, # bwd_ln_sm_margin None, # zero_centered_gamma None, # activation + None, # activation_params None, # normalization None, # ub_overlap_ag None, # ub_overlap_rs @@ -1436,7 +1447,11 @@ class LayerNormMLP(TransformerEngineBaseModule): activation : str, default = 'gelu' activation function used. Options: 'gelu', 'geglu', 'qgelu', 'qgeglu', 'relu', 'reglu', 'srelu', 'sreglu', - 'silu', and 'swiglu'. + 'silu', 'swiglu', and 'clamped_swiglu'. + 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. 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)`. @@ -1537,6 +1552,7 @@ def __init__( bias: bool = True, normalization: str = "LayerNorm", activation: str = "gelu", + activation_params: Optional[dict] = None, output_layer_init_method: Optional[Callable] = None, fuse_wgrad_accumulation: bool = False, params_dtype: Optional[torch.dtype] = None, @@ -1564,6 +1580,7 @@ def __init__( assert normalization in ["LayerNorm", "RMSNorm"], "Unsupported normalization type!" self.use_bias = bias self.activation = activation + self.activation_params = activation_params self.return_bias = return_bias self.apply_bias = bias and not return_bias self.return_layernorm_output = return_layernorm_output @@ -1643,7 +1660,7 @@ def __init__( self.layer_norm_bias = None # FC1 init - if self.activation in ["geglu", "qgeglu", "reglu", "sreglu", "swiglu"]: + if self.activation in ["geglu", "qgeglu", "reglu", "sreglu", "swiglu", "clamped_swiglu"]: fc1_output_features = 2 * self.size_per_partition else: fc1_output_features = self.size_per_partition @@ -1897,6 +1914,7 @@ def forward( self.bwd_ln_sm_margin, self.zero_centered_gamma, self.activation, + self.activation_params, self.normalization, self.ub_overlap_ag, self.ub_overlap_rs, @@ -2026,6 +2044,19 @@ def onnx_forward(self, inp: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Ten fc1_out = onnx_gemm(fc1_weight, ln_out, fc1_bias) 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 + ) + + def _clamped_swiglu(x, limit, alpha): + 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) + return y activation_map = { "gelu": lambda x: torch.nn.functional.gelu(x, approximate="tanh"), @@ -2040,6 +2071,9 @@ def onnx_forward(self, inp: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Ten * x.chunk(2, -1)[1], "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 + ), } if self.activation not in activation_map: raise ValueError(f"Unsupported activation in onnx export: {self.activation}") diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 8a032b2f55..4c7599ad80 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -176,7 +176,12 @@ class TransformerLayer(torch.nn.Module): activation : str, default = 'gelu' Type of activation used in MLP block. Options are: 'gelu', 'geglu', 'qgelu', 'qgeglu', 'relu', 'reglu', 'srelu', 'sreglu', - 'silu', and 'swiglu'. + 'silu', 'swiglu', and 'clamped_swiglu'. + 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}`. 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 @@ -310,6 +315,7 @@ def __init__( ub_bulk_wgrad: bool = True, bias: bool = True, activation: str = "gelu", + activation_params: Optional[dict] = None, normalization: str = "LayerNorm", device: Union[torch.device, str] = "cuda", attn_input_format: str = "sbhd", @@ -475,6 +481,7 @@ def __init__( ub_overlap_rs=ub_overlap_rs, ub_overlap_ag=ub_overlap_ag, activation=activation, + activation_params=activation_params, normalization=normalization, device=device, name=name + ".layernorm_mlp" if name is not None else None, From f0295f9d9b0b6353dac0accd36a8030b55dbd733 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 30 Oct 2025 13:12:03 -0400 Subject: [PATCH 031/521] CMake to respect MAX_JOBS or NVTE_MAX_JOBS (#2319) fix max jobs Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e388dd794b..e0c42b2d9a 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -352,10 +352,10 @@ set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr") set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -O3") # Number of parallel build jobs -if(ENV{MAX_JOBS}) - set(BUILD_JOBS_STR "$ENV{MAX_JOBS}") -elseif(ENV{NVTE_BUILD_MAX_JOBS}) - set(BUILD_JOBS_STR "$ENV{NVTE_BUILD_MAX_JOBS}") +if($ENV{MAX_JOBS}) + set(BUILD_JOBS_STR $ENV{MAX_JOBS}) +elseif($ENV{NVTE_BUILD_MAX_JOBS}) + set(BUILD_JOBS_STR $ENV{NVTE_BUILD_MAX_JOBS}) else() set(BUILD_JOBS_STR "max") endif() From 5e8a9a961f5375cd7c098989f618194bbcf4e9cb Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 30 Oct 2025 10:27:42 -0700 Subject: [PATCH 032/521] [JAX] Fix: Skip determinism tests for bprop for all sm >=100 (#2315) * Fix: Skip determinism tests for bprop for all sm >=100 Signed-off-by: Kshitij Lakhani * Add username to TODO Signed-off-by: Kshitij Lakhani * Assert in fused attn bwd pass for sm100+ 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 | 6 +++--- transformer_engine/jax/cpp_extensions/attention.py | 7 +++++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 5b814cb99f..a5d73d9605 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -378,14 +378,14 @@ def _check_configs(self): pytest.skip( "seqlen_q > seqlen_kv is not supported with sliding window attention in cuDNN" ) - + # TODO(KshitijLakhani): Set the upper limit for skipping this test when cuDNN adds support if ( - get_device_compute_capability(0) == 100 + get_device_compute_capability(0) >= 100 and self.dropout_prob == 0.1 and self.attn_bias_type is not AttnBiasType.NO_BIAS ): pytest.skip( - "For sm100, bprop kernel support for dropout + determinism (bias) is not supported" + "For sm100+, bprop kernel support for dropout + determinism (bias) is not supported" ) # Test the MLA case where head dims for qk differ from head dims for v, only if the tensors # are provided in BSHD_BSHD_BSHD or THD_THD_THD formats diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index db2537c38f..c0cb6cda1f 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -2739,10 +2739,13 @@ def fused_attn_bwd( assert bias is None bias = jnp.zeros(0, dtype=qkv[0].dtype) - if 100 in get_all_device_compute_capability(): + # TODO(KshitijLakhani): Add a check for cuDNN version when determinism does get supported on + # sm100+ + compute_capabilities = get_all_device_compute_capability() + if any(x >= 100 for x in compute_capabilities): assert not ( attn_bias_type != AttnBiasType.NO_BIAS and dropout_probability != 0 - ), "For sm100, bprop kernel support for dropout + determinism (bias) is not supported" + ), "For sm100+, bprop kernel support for dropout + determinism (bias) is not supported" fused_config = _FusedAttnConfig( attn_bias_type=attn_bias_type, From 490a5f41ada5788bc6dd94ba54ab024e465e0ec6 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 30 Oct 2025 15:32:37 -0400 Subject: [PATCH 033/521] [PyTorch] Fix attention backend and tests for `sm120` (#2320) * Fix attention backend and tests for sm120 Signed-off-by: Kirthi Shankar Sivamani * Disable MLA only for backward Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/attention/test_attention.py | 22 +++++++----- .../attention/dot_product_attention/utils.py | 35 +++++++++++++++++++ 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index b05a0447c5..a671f1eec2 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -61,8 +61,16 @@ get_available_attention_backends, ) -# Check if hardware supports FP8 +# Check if hardware supports FP8 attention. fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) +fp8_attn_available, reason_for_no_fp8_attn = fp8_available, reason_for_no_fp8 +device_compute_capability = get_device_compute_capability() +if fp8_available and (device_compute_capability < (9, 0) or device_compute_capability >= (12, 0)): + fp8_attn_available = False + reason_for_no_fp8_attn = ( + "FP8 attention is not supported for compute capability =" + f" sm{device_compute_capability[0] * 10 + device_compute_capability[1]}" + ) # Reset RNG seed and states seed = 1234 @@ -1573,8 +1581,7 @@ def _run_transformer_layer( } -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.skipif(get_device_compute_capability() < (9, 0), reason="FP8 tests require Hopper.") +@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.skipif(get_cudnn_version() < (9, 3, 0), reason="cuDNN 9.3.0+ is required.") @pytest.mark.parametrize("model", ["large"]) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -1736,8 +1743,7 @@ def get_model(dtype, config): @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.skipif(get_device_compute_capability() < (9, 0), reason="FP8 tests require Hopper+.") +@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.parametrize("dtype", param_types_fp8_vs_f16) @pytest.mark.parametrize("model", model_configs_fp8_vs_f16.keys()) @pytest.mark.parametrize("qkv_format", qkv_format_fp8_vs_f16) @@ -1973,8 +1979,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: @pytest.mark.skipif(get_cudnn_version() < (9, 2, 1), reason="cuDNN 9.2.1+ is required.") -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.skipif(get_device_compute_capability() < (9, 0), reason="FP8 tests require Hopper+.") +@pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.parametrize("dtype", param_types_fp8_vs_f16) @pytest.mark.parametrize("model", model_configs_fp8_vs_f16.keys()) @pytest.mark.parametrize("qkv_layout", qkv_layout_fp8_vs_f16) @@ -2302,8 +2307,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: ), reason=f"""cuDNN {"8.9.3" if cudnn_frontend_version == 0 else "9.2.1"}+ is required.""", ) -@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.skipif(get_device_compute_capability() < (9, 0), reason="FP8 tests require Hopper+.") +@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) def test_custom_mha_fp8_vs_f16(dtype, model): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 51279bd372..25dc0e96c8 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -481,6 +481,20 @@ def get_attention_backend( logger.debug("Disabling FusedAttention for FP8 current scaling with cuDNN < 9.14.0") use_fused_attention = False + if device_compute_capability == (12, 0): + if use_flash_attention: + logger.debug( + "Disabling FlashAttention as FP8 is not supported" + " for compute capability = sm120" + ) + if use_fused_attention: + logger.debug( + "Disabling FusedAttention as FP8 is not supported" + " for compute capability = sm120" + ) + use_flash_attention = False + use_fused_attention = False + # Filter: Return max_logit if return_max_logit: if use_flash_attention: @@ -560,6 +574,20 @@ def get_attention_backend( qkv_layout, ) use_fused_attention = False + if ( + device_compute_capability == (12, 0) + and (head_dim_qk > 128 or head_dim_qk % 8 != 0) + and is_training + ): + if use_fused_attention: + logger.debug( + "Disabling FusedAttention as MLA for backward pass is not supported for compute" + " capability = sm120 for a head_dim_qk > 128 or head_dim_qk %%8 != 0. Found:" + " head_dim_qk = %s", + head_dim_qk, + ) + use_fused_attention = False + if use_flash_attention_2 and ( head_dim_qk > 256 or head_dim_qk % 8 != 0 @@ -629,6 +657,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]" ) use_flash_attention = False + if device_compute_capability == (12, 0): + if use_fused_attention: + logger.debug( + "Disabling FusedAttention as qkv_format = thd is" + " not supported for compute capability = sm120" + ) + use_fused_attention = False # Filter: Dropout if attention_dropout != 0.0 and use_flash_attention_3: From 0e80c847845be58dbb88f46a0975786ddc823798 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Thu, 30 Oct 2025 20:53:11 +0100 Subject: [PATCH 034/521] [Common] Split cast/gated kernels by scaling mode (#2248) * Separated gated and dequantize kernels Signed-off-by: Oleg Goncharov * Separated quantize, dequantize and gated functions Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed lint issues Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed persistent lint issues Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added missing compute capability 10.0 check for Quantize FP8 TMA kernels Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed the issue which was added again by autofix Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Changed files description. Completely removed non-identity activations from the NVFP4 transpose test suite Signed-off-by: Oleg Goncharov * Removed unsupported template arguments in NVFP4 quantize Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed undefined symbol error Signed-off-by: Oleg Goncharov * Fixed condition Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> * Fixed CUDA version check Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Changed arch conditions order 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 * Clean up Signed-off-by: Oleg Goncharov * Small fix Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Small fix Signed-off-by: Oleg Goncharov * Fixes per the PR review Signed-off-by: Oleg Goncharov * Fix Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Split quantize helper into two (FWD and BWD) functions Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Moved activation functions from cast.cu. Removed cast.cu from the fast-math compilation list Signed-off-by: Oleg Goncharov * Enabled fast math for activations by default Signed-off-by: Oleg Goncharov * Disabled fast math for activations by default Signed-off-by: Oleg Goncharov --------- Signed-off-by: Oleg Goncharov Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/common/CMakeLists.txt | 5 +- .../common/activation/activation_template.h | 27 +- transformer_engine/common/activation/gelu.cu | 26 + transformer_engine/common/activation/relu.cu | 26 + .../common/activation/swiglu.cu | 13 + transformer_engine/common/cast/cast.cu | 102 + .../common/cast/core/common.cuh | 97 + .../common/cast/dispatch/dequantize.cuh | 56 + .../common/cast/dispatch/gated.cuh | 161 ++ .../common/cast/dispatch/quantize.cuh | 326 +++ .../common/cast/fp8/dequantize_fp8.cuh | 54 + .../common/cast/fp8/gated_fp8.cuh | 394 +++ .../common/cast/fp8/quantize_fp8.cuh | 580 +++++ .../mxfp8/dequantize_mxfp8.cuh} | 171 +- .../mxfp8/gated_mxfp8.cuh} | 711 +----- .../common/cast/mxfp8/quantize_mxfp8.cuh | 722 ++++++ .../common/cast/nvfp4/core_nvfp4.cuh | 112 + .../common/cast/nvfp4/dequantize_nvfp4.cuh | 111 + .../common/cast/nvfp4/quantize_nvfp4.cuh | 688 ++++++ .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 1287 ++++++++++ transformer_engine/common/util/cast.cu | 201 -- .../common/util/cast_kernels.cuh | 2188 ----------------- transformer_engine/common/util/math.h | 2 + transformer_engine/common/util/ptx.cuh | 191 +- 24 files changed, 5073 insertions(+), 3178 deletions(-) create mode 100644 transformer_engine/common/cast/cast.cu create mode 100644 transformer_engine/common/cast/core/common.cuh create mode 100644 transformer_engine/common/cast/dispatch/dequantize.cuh create mode 100644 transformer_engine/common/cast/dispatch/gated.cuh create mode 100644 transformer_engine/common/cast/dispatch/quantize.cuh create mode 100644 transformer_engine/common/cast/fp8/dequantize_fp8.cuh create mode 100644 transformer_engine/common/cast/fp8/gated_fp8.cuh create mode 100644 transformer_engine/common/cast/fp8/quantize_fp8.cuh rename transformer_engine/common/{util/dequantize_kernels.cuh => cast/mxfp8/dequantize_mxfp8.cuh} (69%) rename transformer_engine/common/{util/cast_gated_kernels.cuh => cast/mxfp8/gated_mxfp8.cuh} (53%) create mode 100644 transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh create mode 100644 transformer_engine/common/cast/nvfp4/core_nvfp4.cuh create mode 100644 transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh create mode 100644 transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh create mode 100644 transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh delete mode 100644 transformer_engine/common/util/cast.cu delete mode 100644 transformer_engine/common/util/cast_kernels.cuh diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index e0c42b2d9a..62b769c77e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -168,7 +168,7 @@ list(APPEND transformer_engine_cuda_sources list(APPEND transformer_engine_cuda_arch_specific_sources gemm/cutlass_grouped_gemm.cu - util/cast.cu + cast/cast.cu activation/gelu.cu activation/relu.cu activation/swiglu.cu @@ -336,8 +336,7 @@ option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --u if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) list(APPEND nvte_sources_with_fast_math activation/gelu.cu activation/relu.cu - activation/swiglu.cu - util/cast.cu) + activation/swiglu.cu) endif() foreach(cuda_source IN LISTS nvte_sources_with_fast_math) diff --git a/transformer_engine/common/activation/activation_template.h b/transformer_engine/common/activation/activation_template.h index 1d9a3fb43c..7353c3e1d5 100644 --- a/transformer_engine/common/activation/activation_template.h +++ b/transformer_engine/common/activation/activation_template.h @@ -14,26 +14,17 @@ #include #include +#include "../cast/dispatch/gated.cuh" +#include "../cast/dispatch/quantize.cuh" #include "../common.h" -#include "../util/cast_gated_kernels.cuh" -#include "../util/cast_kernels.cuh" -#include "../util/math.h" -#include "../util/vectorized_pointwise.h" namespace transformer_engine { template void act_fn(const NVTETensor input, NVTETensor output, cudaStream_t stream) { using namespace detail; - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = false; constexpr bool IS_ACT = true; - constexpr NVTETensor dbias = nullptr; - constexpr NVTETensor workspace = nullptr; - constexpr const NVTETensor grad = nullptr; - - quantize_helper(input, grad, output, dbias, workspace, - nullptr, stream); + dispatch::quantize_fwd_helper(input, output, nullptr, stream); } template @@ -42,20 +33,17 @@ void dact_fn(const NVTETensor grad, const NVTETensor input, NVTETensor output, using namespace detail; constexpr bool IS_DBIAS = false; constexpr bool IS_DACT = true; - constexpr bool IS_ACT = false; constexpr NVTETensor dbias = nullptr; constexpr NVTETensor workspace = nullptr; - quantize_helper(input, grad, output, dbias, workspace, - nullptr, stream); + dispatch::quantize_bwd_helper(grad, input, output, dbias, workspace, + nullptr, stream); } template void gated_act_fn(const NVTETensor input, NVTETensor output, Param &p, cudaStream_t stream) { using namespace detail; - constexpr bool IS_DGATED = false; - constexpr NVTETensor grad = nullptr; - quantize_gated_helper(grad, input, output, p, stream); + dispatch::quantize_gated_fwd_helper(input, output, p, stream); } template (grad, input, output, p, stream); + dispatch::quantize_gated_bwd_helper(grad, input, output, p, stream); } } // namespace transformer_engine diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index 4949ba5906..4979023ef1 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -20,6 +20,19 @@ void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, 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_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_geglu); using namespace transformer_engine; @@ -48,6 +61,19 @@ void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu dact_fn>(grad, input, output, 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_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/relu.cu b/transformer_engine/common/activation/relu.cu index c74fc6eee9..c0ef9fd65a 100644 --- a/transformer_engine/common/activation/relu.cu +++ b/transformer_engine/common/activation/relu.cu @@ -20,6 +20,19 @@ void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, 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_reglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_reglu); using namespace transformer_engine; @@ -48,6 +61,19 @@ void nvte_dsrelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu dact_fn>(grad, input, output, 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_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/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index cafc48abba..6957a91e61 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -20,6 +20,19 @@ void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, 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_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_swiglu); using namespace transformer_engine; diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu new file mode 100644 index 0000000000..1ed46a3359 --- /dev/null +++ b/transformer_engine/common/cast/cast.cu @@ -0,0 +1,102 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include "../common.h" +#include "../transpose/cast_transpose.h" +#include "../util/multi_stream.h" +#include "../utils.cuh" +#include "dispatch/dequantize.cuh" +#include "dispatch/quantize.cuh" +#include "transformer_engine/transpose.h" + +void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize); + using namespace transformer_engine; + + constexpr bool IS_ACT = false; + dispatch::quantize_fwd_helper(input, output, nullptr, stream); +} + +void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor noop, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_noop); + using namespace transformer_engine; + + // Create config with noop tensor + QuantizationConfig quant_config; + quant_config.noop_tensor = noop; + + nvte_quantize_v2(input, output, reinterpret_cast(&quant_config), stream); +} + +void nvte_quantize_v2(const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_v2); + using namespace transformer_engine; + + constexpr bool IS_ACT = false; + 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_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_dequantize); + using namespace transformer_engine; + dispatch::dequantize_helper(*convertNVTETensorCheck(input), convertNVTETensorCheck(output), + stream); +} + +void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, + const NVTEQuantizationConfig quant_configs, + const size_t num_tensors, cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_quantize); + using namespace transformer_engine; + + constexpr bool IS_ACT = false; + + const size_t num_streams = nvte_get_num_compute_streams(); + + int num_stream_used = std::min(num_streams, num_tensors); + // wait for current stream to finish + NVTE_CHECK_CUDA(cudaEventRecord(detail::get_compute_stream_event(0), stream)); + for (int s = 0; s < num_stream_used; s++) { + NVTE_CHECK_CUDA( + cudaStreamWaitEvent(detail::get_compute_stream(s), detail::get_compute_stream_event(0))); + } + + for (int i = 0; i < num_tensors; i++) { + dispatch::quantize_fwd_helper( + inputs[i], outputs[i], quant_configs, detail::get_compute_stream(i % num_streams)); + } + + // record events on compute streams + for (int s = 0; s < num_stream_used; s++) { + NVTE_CHECK_CUDA( + cudaEventRecord(detail::get_compute_stream_event(s), detail::get_compute_stream(s))); + } + // wait for all compute streams to finish + for (int s = 0; s < num_stream_used; s++) { + NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream, detail::get_compute_stream_event(s))); + } +} diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh new file mode 100644 index 0000000000..b750142f5b --- /dev/null +++ b/transformer_engine/common/cast/core/common.cuh @@ -0,0 +1,97 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file common.cuh + * \brief Common functions in quantize. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../utils.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace common { +inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_block) { + const size_t N = product(t->data.shape); + const bool isFullTile = (N % elems_per_block == 0); + return isFullTile; +} + +inline bool dimensions_supported_by_TMA(const Tensor *const t) { + const size_t cols = t->flat_last_dim(); + constexpr size_t TMA_bytes = 16; + const size_t alignment_requirement = (TMA_bytes * 8) / typeToNumBits(t->dtype()); + return cols % alignment_requirement == 0; +} + +namespace kernel { + +constexpr size_t THREADS_PER_BLOCK = 256; +template +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, + const size_t rows, const size_t cols) { + using ComputeVec = Vec; + using OutputVec = Vec; + + const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; + + if (thread_id * nvec >= cols) { + return; + } + + const float *const thread_in_base = dbias_partial + thread_id * nvec; + OType *const thread_out_base = dbias_output + thread_id * nvec; + + ComputeVec ldg_vec; + ComputeVec acc_vec; + acc_vec.clear(); + for (int i = 0; i < rows; ++i) { + ldg_vec.load_from(thread_in_base + i * cols); +#pragma unroll + for (int e = 0; e < nvec; ++e) { + acc_vec.data.elt[e] += ldg_vec.data.elt[e]; + } + } + + OutputVec stg_vec; +#pragma unroll + for (int e = 0; e < nvec; ++e) { + stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); + } + stg_vec.store_to(thread_out_base); +} +} // namespace kernel + +template +void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, const size_t cols, + cudaStream_t stream) { + using namespace kernel; + constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 + constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); + + NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape."); + const size_t reduce_dbias_num_blocks = DIVUP(cols, THREADS_PER_BLOCK * reduce_dbias_nvec); + + reduce_dbias_kernel + <<>>( + reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace common +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_CORE_COMMON_CUH_ diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh new file mode 100644 index 0000000000..b8547915c7 --- /dev/null +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -0,0 +1,56 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file dequantize.cuh + * \brief Dequantize dispatcher. + */ + +#ifndef TRANSFORMER_ENGINE_DISPATCH_DEQUANTIZE_CUH_ +#define TRANSFORMER_ENGINE_DISPATCH_DEQUANTIZE_CUH_ + +#include + +#include "../../common.h" +#include "../fp8/dequantize_fp8.cuh" +#include "../mxfp8/dequantize_mxfp8.cuh" +#include "../nvfp4/dequantize_nvfp4.cuh" + +namespace transformer_engine { +namespace dispatch { + +inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t stream) { + CheckInputTensor(input, "cast_input"); + CheckOutputTensor(*output, "cast_output"); + + switch (input.scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: { + NVTE_CHECK(is_fp8_dtype(input.data.dtype), "Input must have FP8 type."); + NVTE_CHECK(!is_fp8_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."); + fp8::dequantize(input, output, stream); + break; + } + case NVTE_MXFP8_1D_SCALING: { + if (is_supported_by_CC_100()) { + mxfp8::dequantize(input, output, stream); + } else { + NVTE_ERROR("MXFP8 Dequantization is NOT supported by architectures < 10.0"); + } + break; + } + case NVTE_NVFP4_1D_SCALING: { + nvfp4::dequantize(input, output, stream); + break; + } + default: + NVTE_ERROR("Not implemented scaling mode: " + to_string(input.scaling_mode) + "."); + } +} + +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_DISPATCH_DEQUANTIZE_CUH_ diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh new file mode 100644 index 0000000000..4373090b72 --- /dev/null +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -0,0 +1,161 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated.cuh + * \brief Gated dispatcher. + */ + +#ifndef TRANSFORMER_ENGINE_DISPATCH_GATED_CUH_ +#define TRANSFORMER_ENGINE_DISPATCH_GATED_CUH_ + +#include + +#include "../../common.h" +#include "../../utils.cuh" +#include "../fp8/gated_fp8.cuh" +#include "../mxfp8/gated_mxfp8.cuh" + +namespace transformer_engine { +namespace dispatch { + +template +void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_output, ParamOP &p, + cudaStream_t stream) { + const Tensor input = *convertNVTETensorCheck(nvte_input); + Tensor *output = convertNVTETensorCheck(nvte_output); + + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output", /*allow_empty=*/false); + + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim() / 2; + + NVTE_CHECK(input.flat_last_dim() % 2 == 0, + "Wrong input shape. Expected (after flattening) last dimension to be even, ", "got [", + input.flat_first_dim(), ", ", input.flat_last_dim(), "]."); + NVTE_CHECK(output->flat_last_dim() == cols, + "Wrong output shape. Expected (after flattening) [*, ", cols, "], got [", + output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); + + NVTE_CHECK(output->has_data() || output->has_columnwise_data(), + "Either rowwise or columnwise output data need to be allocated."); + + switch (output->scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: { + const bool use_tma_kernels = (cols % 32 == 0) && is_supported_by_CC_100(); + if (use_tma_kernels) { + Tensor dummy_grad_tensor; + fp8::cast_gated_tma(input, dummy_grad_tensor, + output, p, stream); + } else { + fp8::cast_gated_fwd(input, output, p, stream); + } + break; + } + case NVTE_MXFP8_1D_SCALING: { + NVTE_CHECK(cols % 32 == 0, + "Invalid input shape. Expected the last dimension to be " + "divisible by 32, but got ", + cols, "."); + if (output->has_data()) { + NVTE_CHECK(is_fp8_dtype(output->data.dtype), + "The type of the output tensor should be FP8."); + } + if (output->has_columnwise_data()) { + NVTE_CHECK(is_fp8_dtype(output->columnwise_data.dtype), + "The type of the columnwise output tensor should be FP8."); + } + NVTE_CHECK(is_supported_by_CC_100(), + "Gated FWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); + Tensor dummy_grad_tensor; + mxfp8::quantize_gated(input, dummy_grad_tensor, + output, p, stream); + break; + } + default: + NVTE_ERROR("Not supported scaling mode: " + to_string(output->scaling_mode) + "."); + } +} + +template +void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte_gated_input, + NVTETensor nvte_output, ParamOP &p, cudaStream_t stream) { + const Tensor &grad = *(convertNVTETensorCheck(nvte_grad)); + const Tensor gated_input = *convertNVTETensorCheck(nvte_gated_input); + Tensor *output = convertNVTETensorCheck(nvte_output); + + CheckInputTensor(grad, "grad"); + CheckInputTensor(gated_input, "gated_input"); + CheckOutputTensor(*output, "output", /*allow_empty=*/false); + + NVTE_CHECK(gated_input.flat_last_dim() % 2 == 0, "Number of columns must be even, but got ", + gated_input.flat_last_dim(), "."); + + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + + NVTE_CHECK(!is_fp8_dtype(grad.data.dtype), "Grad input must be in higher precision."); + NVTE_CHECK(grad.data.dtype == gated_input.data.dtype, "Types of both inputs must match."); + + NVTE_CHECK(grad.flat_first_dim() == rows, + "Wrong Grad shape. Expected first dimension (after flattening) [", rows, ", *], got [", + grad.flat_first_dim(), ", ", grad.flat_last_dim(), "]."); + NVTE_CHECK(grad.flat_last_dim() == cols, + "Wrong Grad shape. Expected last dimension (after flattening) [", cols, ", *], got [", + grad.flat_first_dim(), ", ", grad.flat_last_dim(), "]."); + + NVTE_CHECK(output->has_data() || output->has_columnwise_data(), + "Either rowwise or columnwise output data need to be allocated."); + + NVTE_CHECK(output->flat_first_dim() == rows, "Wrong output shape. Expected (after flattening) [", + rows, ", *], got [", output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); + NVTE_CHECK(output->flat_last_dim() == cols * 2, + "Wrong output shape. Expected (after flattening) [*, ", cols * 2, "], got [", + output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); + NVTE_CHECK(gated_input.data.shape == output->data.shape, + "Gated input and output shapes must match. Input shape: ", gated_input.data.shape, + ", output shape: ", output->data.shape, "."); + + switch (output->scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: { + const bool use_tma_kernels = (cols % 32 == 0) && is_supported_by_CC_100(); + if (use_tma_kernels) { + fp8::cast_gated_tma(gated_input, grad, output, p, + stream); + } else { + fp8::cast_gated_bwd(gated_input, grad, output, p, stream); + } + break; + } + case NVTE_MXFP8_1D_SCALING: { + NVTE_CHECK(cols % 32 == 0, + "Invalid input shape. Expected the last dimension to be " + "divisible by 32, but got ", + cols, "."); + if (output->has_data()) { + NVTE_CHECK(is_fp8_dtype(output->data.dtype), + "The type of the output tensor should be FP8."); + } + if (output->has_columnwise_data()) { + NVTE_CHECK(is_fp8_dtype(output->columnwise_data.dtype), + "The type of the columnwise output tensor should be FP8."); + } + NVTE_CHECK(is_supported_by_CC_100(), + "Gated BWD NVTE_MXFP8_1D_SCALING is only supported on SM 10.0+"); + + mxfp8::quantize_gated(gated_input, grad, output, p, + stream); + break; + } + default: + NVTE_ERROR("Not supported scaling mode: " + to_string(output->scaling_mode) + "."); + } +} +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_DISPATCH_GATED_CUH_ diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh new file mode 100644 index 0000000000..9f7a4a9b01 --- /dev/null +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -0,0 +1,326 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize.cuh + * \brief Quantize dispatcher. + */ + +#ifndef TRANSFORMER_ENGINE_DISPATCH_QUANTIZE_CUH_ +#define TRANSFORMER_ENGINE_DISPATCH_QUANTIZE_CUH_ + +#include + +#include "../../common.h" +#include "../../transpose/cast_transpose.h" +#include "../../util/vectorized_pointwise.h" +#include "../core/common.cuh" +#include "../fp8/quantize_fp8.cuh" +#include "../mxfp8/quantize_mxfp8.cuh" +#include "../nvfp4/quantize_nvfp4.cuh" +#include "../nvfp4/quantize_transpose_nvfp4.cuh" + +namespace transformer_engine { +namespace dispatch { + +template +void quantize_fwd_helper(const NVTETensor input, NVTETensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + using namespace detail; + + const Tensor *input_tensor = convertNVTETensorCheck(input); + Tensor *output_tensor = convertNVTETensorCheck(output); + + // Quantization config + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Noop flag + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + // Check for unsupported options + if (quant_config_cpp.stochastic_rounding) { + NVTE_CHECK(output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING, + "Stochastic rounding is only supported for NVFP4 quantization."); + } + + NVTE_CHECK(output_tensor->has_data() || output_tensor->has_columnwise_data(), + "Either rowwise or columnwise output data need to be allocated."); + + // Dispatch to quantization kernel depending on data format + switch (output_tensor->scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: { + const Tensor *dummy_input_tensor = nullptr; + Tensor *dummy_dbias_tensor = nullptr; + Tensor *dummy_workspace_tensor = nullptr; + if (output_tensor->has_columnwise_data()) { + NVTE_CHECK(output_tensor->has_data(), + "Quantizing in only the columnwise direction not supported yet!"); + if constexpr (!IS_ACT) { + cast_transpose(*input_tensor, *noop_tensor, output_tensor, stream); + } else { + cast_transpose_fused( + *input_tensor, dummy_input_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, stream); + } + } else if (output_tensor->has_data()) { + fp8::quantize( + *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, stream); + } + break; + } + case NVTE_MXFP8_1D_SCALING: { + const Tensor *dummy_input_tensor = nullptr; + Tensor *dummy_dbias_tensor = nullptr; + Tensor *dummy_workspace_tensor = nullptr; + mxfp8::quantize( + *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, + dummy_workspace_tensor, stream); + break; + } + case NVTE_NVFP4_1D_SCALING: { + NVTE_CHECK(!IS_ACT, "IS_ACT is not supported by FWD NVTE_NVFP4_1D_SCALING"); + + // Check tensors + CheckNoopTensor(*noop_tensor, "cast_noop"); + CheckInputTensor(*input_tensor, "input"); + CheckOutputTensor(*output_tensor, "output", false); + + // Choose kernel + int32_t rows = input_tensor->flat_first_dim(); + int32_t cols = input_tensor->flat_last_dim(); + auto dtype = input_tensor->dtype(); + 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 (quant_config_cpp.nvfp4_2d_quantization) { + nvfp4::quantize_transpose( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } else { + nvfp4::quantize_transpose( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } + } else { + auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax + : output_tensor->columnwise_amax; + quantize_transpose_vector_blockwise_fp4( + /*input=*/input_tensor->data, /*global_amax=*/global_amax, + /*scale_inv=*/output_tensor->scale_inv, + /*scale_inv_t=*/output_tensor->columnwise_scale_inv, + /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data, + /*epsilon=*/0.0f, /*return_identity=*/output_tensor->has_data(), + /*return_transpose=*/output_tensor->has_columnwise_data(), /*pow2_scale=*/false, + /*swizzled_scale=*/false, + /*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); + } + break; + } + case NVTE_BLOCK_SCALING_2D: { + // TODO(kwyss): IS_ACT, ParamOP, OP parameters support. + NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for FWD NVTE_BLOCK_SCALING_2D"); + bool force_pow_2_scales = quant_config_cpp.force_pow_2_scales; + float epsilon = quant_config_cpp.amax_epsilon; + quantize_transpose_square_blockwise( + input_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, + output_tensor->data, output_tensor->columnwise_data, epsilon, + /*return_transpose=*/output_tensor->has_columnwise_data(), force_pow_2_scales, + /*noop_tensor=*/noop_tensor->data, stream); + break; + } + case NVTE_BLOCK_SCALING_1D: { + // TODO(kwyss): IS_ACT, ParamOP, OP parameters support. + NVTE_CHECK(!IS_ACT, "IS_ACT is not implemented for FWD NVTE_BLOCK_SCALING_1D"); + bool force_pow_2_scales = quant_config_cpp.force_pow_2_scales; + float epsilon = quant_config_cpp.amax_epsilon; + FP8BlockwiseRowwiseOption rowwise_option = FP8BlockwiseRowwiseOption::NONE; + FP8BlockwiseColumnwiseOption columnwise_option = FP8BlockwiseColumnwiseOption::NONE; + if (output_tensor->has_data()) { + bool rowwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == + Float8BlockScaleTensorFormat::COMPACT); + rowwise_option = rowwise_compact ? FP8BlockwiseRowwiseOption::ROWWISE_COMPACT + : FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; + } + if (output_tensor->has_columnwise_data()) { + bool columnwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == + Float8BlockScaleTensorFormat::COMPACT); + columnwise_option = columnwise_compact + ? FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT + : FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; + } + quantize_transpose_vector_blockwise( + input_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, + output_tensor->data, output_tensor->columnwise_data, epsilon, rowwise_option, + columnwise_option, force_pow_2_scales, noop_tensor->data, stream); + break; + } + default: + NVTE_ERROR("Not implemented scaling mode: " + to_string(output_tensor->scaling_mode) + "."); + } +} + +template +void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETensor output, + NVTETensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + using namespace detail; + + const Tensor *grad_tensor = convertNVTETensorCheck(grad); + const Tensor *input_tensor = convertNVTETensor(input); + + Tensor *output_tensor = convertNVTETensorCheck(output); + Tensor *dbias_tensor = convertNVTETensor(dbias); + Tensor *workspace_tensor = convertNVTETensor(workspace); + + // Quantization config + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Noop flag + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + // Check for unsupported options + if (quant_config_cpp.stochastic_rounding) { + NVTE_CHECK(output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING, + "Stochastic rounding is only supported for NVFP4 quantization."); + } + + NVTE_CHECK(output_tensor->has_data() || output_tensor->has_columnwise_data(), + "Either rowwise or columnwise output data need to be allocated."); + + // Dispatch to quantization kernel depending on data format + switch (output_tensor->scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: { + if (output_tensor->has_columnwise_data()) { + NVTE_CHECK(output_tensor->has_data(), + "Quantizing in only the columnwise direction not supported yet!"); + if constexpr (!IS_DBIAS && !IS_DACT) { + cast_transpose(*grad_tensor, *noop_tensor, output_tensor, stream); + } else { + cast_transpose_fused( + *grad_tensor, input_tensor, output_tensor, dbias_tensor, workspace_tensor, stream); + } + } else if (output_tensor->has_data()) { + fp8::quantize( + *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + stream); + } + break; + } + case NVTE_MXFP8_1D_SCALING: { + mxfp8::quantize( + *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + stream); + break; + } + case NVTE_NVFP4_1D_SCALING: { + NVTE_CHECK((!IS_DBIAS && !IS_DACT), + "IS_DBIAS and IS_DACT are not supported by BWD NVTE_NVFP4_1D_SCALING"); + + // Check tensors + CheckNoopTensor(*noop_tensor, "cast_noop"); + CheckInputTensor(*grad_tensor, "input"); + CheckOutputTensor(*output_tensor, "output", false); + + // Choose kernel + int32_t rows = grad_tensor->flat_first_dim(); + int32_t cols = grad_tensor->flat_last_dim(); + auto dtype = grad_tensor->dtype(); + 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 (quant_config_cpp.nvfp4_2d_quantization) { + nvfp4::quantize_transpose( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } else { + nvfp4::quantize_transpose( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } + } else { + auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax + : output_tensor->columnwise_amax; + quantize_transpose_vector_blockwise_fp4( + /*input=*/grad_tensor->data, /*global_amax=*/global_amax, + /*scale_inv=*/output_tensor->scale_inv, + /*scale_inv_t=*/output_tensor->columnwise_scale_inv, + /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data, + /*epsilon=*/0.0f, /*return_identity=*/output_tensor->has_data(), + /*return_transpose=*/output_tensor->has_columnwise_data(), /*pow2_scale=*/false, + /*swizzled_scale=*/false, + /*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); + } + break; + } + case NVTE_BLOCK_SCALING_2D: { + // TODO(kwyss): IS_BIAS, IS_DACT, ParamOP, OP parameters support. + NVTE_CHECK((!IS_DBIAS && !IS_DACT), + "IS_DBIAS and IS_DACT are not implemented for BWD NVTE_BLOCK_SCALING_2D"); + bool force_pow_2_scales = quant_config_cpp.force_pow_2_scales; + float epsilon = quant_config_cpp.amax_epsilon; + quantize_transpose_square_blockwise( + grad_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, + output_tensor->data, output_tensor->columnwise_data, epsilon, + /*return_transpose=*/output_tensor->has_columnwise_data(), force_pow_2_scales, + /*noop_tensor=*/noop_tensor->data, stream); + break; + } + case NVTE_BLOCK_SCALING_1D: { + // TODO(kwyss): IS_BIAS, IS_DACT, ParamOP, OP parameters support. + NVTE_CHECK((!IS_DBIAS && !IS_DACT), + "IS_DBIAS and IS_DACT are not implemented for BWD NVTE_BLOCK_SCALING_1D"); + bool force_pow_2_scales = quant_config_cpp.force_pow_2_scales; + float epsilon = quant_config_cpp.amax_epsilon; + FP8BlockwiseRowwiseOption rowwise_option = FP8BlockwiseRowwiseOption::NONE; + FP8BlockwiseColumnwiseOption columnwise_option = FP8BlockwiseColumnwiseOption::NONE; + if (output_tensor->has_data()) { + bool rowwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == + Float8BlockScaleTensorFormat::COMPACT); + rowwise_option = rowwise_compact ? FP8BlockwiseRowwiseOption::ROWWISE_COMPACT + : FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; + } + if (output_tensor->has_columnwise_data()) { + bool columnwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == + Float8BlockScaleTensorFormat::COMPACT); + columnwise_option = columnwise_compact + ? FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT + : FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; + } + quantize_transpose_vector_blockwise( + grad_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, + output_tensor->data, output_tensor->columnwise_data, epsilon, rowwise_option, + columnwise_option, force_pow_2_scales, noop_tensor->data, stream); + break; + } + default: + NVTE_ERROR("Not implemented scaling mode: " + to_string(output_tensor->scaling_mode) + "."); + } +} + +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_DISPATCH_QUANTIZE_CUH_ diff --git a/transformer_engine/common/cast/fp8/dequantize_fp8.cuh b/transformer_engine/common/cast/fp8/dequantize_fp8.cuh new file mode 100644 index 0000000000..2514758b5a --- /dev/null +++ b/transformer_engine/common/cast/fp8/dequantize_fp8.cuh @@ -0,0 +1,54 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file dequantize_fp8.cuh + * \brief CUDA kernels to dequantize from FP8. + */ + +#ifndef TRANSFORMER_ENGINE_DEQUANTIZE_FP8_CUH_ +#define TRANSFORMER_ENGINE_DEQUANTIZE_FP8_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/vectorized_pointwise.h" +#include "../../utils.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8 { +struct DequantizeParam { + const float *scale_inv; +}; + +__device__ inline float dequantize_func(float value, const DequantizeParam ¶m) { + return value * (*(param.scale_inv)); +} + +inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { + const size_t N = product(input.data.shape); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + input.data.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + output->data.dtype, OType, + + constexpr int nvec = 32 / sizeof(OType); + DequantizeParam p; p.scale_inv = reinterpret_cast(input.scale_inv.dptr); + VectorizedUnaryKernelLauncher( + reinterpret_cast(input.data.dptr), nullptr, + reinterpret_cast(output->data.dptr), nullptr, nullptr, nullptr, N, p, + stream);); // NOLINT(*) + ); // NOLINT(*) +} +} // namespace fp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_DEQUANTIZE_FP8_CUH_ diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh new file mode 100644 index 0000000000..225ef93ed9 --- /dev/null +++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh @@ -0,0 +1,394 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file gated_fp8.cuh + * \brief CUDA kernels to cast to FP8 with gated activations. + */ + +#ifndef TRANSFORMER_ENGINE_GATED_FP8_CUH_ +#define TRANSFORMER_ENGINE_GATED_FP8_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../util/vectorized_pointwise.h" +#include "../../utils.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8 { +namespace kernel { + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_PER_CHUNK = 512; +constexpr size_t THREADS_PER_CHUNK_X = CHUNK_DIM_X; +constexpr size_t THREADS_PER_CHUNK_Y = THREADS_PER_CHUNK / THREADS_PER_CHUNK_X; // 4 = 512 / 128 +constexpr size_t BUFFERS_NUM = 2; +constexpr size_t BUFFER_DIM_Y = 32; +constexpr size_t BUFFER_DIM_X = CHUNK_DIM_X; // 128 +constexpr size_t SHMEM_DIM_Y = BUFFER_DIM_Y; // 32 +constexpr size_t SHMEM_DIM_X = BUFFER_DIM_X; // 128 + +constexpr size_t BUFFER_STAGES_NUM = BUFFER_DIM_Y / THREADS_PER_CHUNK_Y; // 8 = 32 / 4 +constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y; // 4 = 128 / 32 +static_assert(ITERATIONS >= 1); + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + cast_fp8_gated_kernel(const __grid_constant__ CUtensorMap tensor_map_grad, + const __grid_constant__ CUtensorMap tensor_map_input_act, + const __grid_constant__ CUtensorMap tensor_map_input_gate, + const __grid_constant__ CUtensorMap tensor_map_output_act, + const __grid_constant__ CUtensorMap tensor_map_output_gate, + float *const amax_ptr, float *const scale_inv_ptr, + const float *const scale_ptr, const size_t rows, const size_t cols, + const ParamOP p) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + + const size_t chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; + + const size_t tid_Y = threadIdx.x / THREADS_PER_CHUNK_X; + const size_t tid_X = threadIdx.x % THREADS_PER_CHUNK_X; + + const size_t thread_offset_Y = tid_Y; + const size_t thread_offset_X = tid_X; + + float amax = 0; + const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; + + 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)); + + constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X; + constexpr size_t buff_elems_total = BUFFERS_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 = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + + constexpr size_t grad_mem = IS_BWD ? buff_size_aligned_in : 0; + + constexpr size_t in_act_mem = buff_size_aligned_in; + constexpr size_t in_gate_mem = buff_size_aligned_in; + constexpr size_t in_mem = in_act_mem + in_gate_mem; + + constexpr size_t out_act_mem = buff_size_aligned_out; + constexpr size_t in_transaction_size = buff_elems * sizeof(IType); + + // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned + IType *in_grad_sh = reinterpret_cast(dshmem); + IType *in_act_sh = reinterpret_cast(dshmem + grad_mem); + IType *in_gate_sh = reinterpret_cast(dshmem + grad_mem + in_act_mem); + OType *out_act_sh = reinterpret_cast(dshmem + grad_mem + in_mem); + OType *out_gate_sh = reinterpret_cast(dshmem + grad_mem + in_mem + out_act_mem); + + const uint64_t *TMAP_grad_in = reinterpret_cast(&tensor_map_grad); + const uint64_t *TMAP_in_act = reinterpret_cast(&tensor_map_input_act); + const uint64_t *TMAP_in_gate = reinterpret_cast(&tensor_map_input_gate); + const uint64_t *TMAP_output_act = reinterpret_cast(&tensor_map_output_act); + const uint64_t *TMAP_output_gate = reinterpret_cast(&tensor_map_output_gate); + + const bool is_master_thread = (threadIdx.x == 0); + +// 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[ITERATIONS]; + + initialize_barriers(mbar, is_master_thread); + + int parity = 0; + + // Prefetch data of the first stage + + if constexpr (IS_BWD) { + copy_2d_to_sharedx3(in_grad_sh, TMAP_grad_in, chunk_offset_X, chunk_offset_Y, in_act_sh, + TMAP_in_act, chunk_offset_X, chunk_offset_Y, in_gate_sh, TMAP_in_gate, + chunk_offset_X, chunk_offset_Y, in_transaction_size, &mbar[0], + is_master_thread); + } else { + copy_2d_to_sharedx2(in_act_sh, TMAP_in_act, chunk_offset_X, chunk_offset_Y, in_gate_sh, + TMAP_in_gate, chunk_offset_X, chunk_offset_Y, in_transaction_size, &mbar[0], + is_master_thread); + } + +#pragma unroll + for (int it = 0; it < ITERATIONS; ++it) { + const size_t buff = it % BUFFERS_NUM; + const size_t next_it = it + 1; + if (next_it < ITERATIONS) { + const size_t next_buff = next_it % BUFFERS_NUM; + const size_t chunk_it_offset_y = chunk_offset_Y + next_it * BUFFER_DIM_Y; + const size_t chunk_it_offset_x = chunk_offset_X; + if constexpr (IS_BWD) { + copy_2d_to_sharedx3( + &in_grad_sh[next_buff * buff_elems], TMAP_grad_in, chunk_it_offset_x, chunk_it_offset_y, + &in_act_sh[next_buff * buff_elems], TMAP_in_act, chunk_it_offset_x, chunk_it_offset_y, + &in_gate_sh[next_buff * buff_elems], TMAP_in_gate, chunk_it_offset_x, chunk_it_offset_y, + in_transaction_size, &mbar[next_it], is_master_thread); + } else { + copy_2d_to_sharedx2(&in_act_sh[next_buff * buff_elems], TMAP_in_act, chunk_it_offset_x, + chunk_it_offset_y, &in_gate_sh[next_buff * buff_elems], TMAP_in_gate, + chunk_it_offset_x, chunk_it_offset_y, in_transaction_size, + &mbar[next_it], is_master_thread); + } + } + + ptx::fence_proxy_async_shared_cta(); + + // Wait for the data to have arrived + ptx::mbarrier_wait_parity(&mbar[it], parity); + + IType *in_grad_sh_curr = in_grad_sh + buff * buff_elems; + IType *in_act_sh_curr = in_act_sh + buff * buff_elems; + IType *in_gate_sh_curr = in_gate_sh + buff * buff_elems; + OType *out_act_sh_curr = out_act_sh + buff * buff_elems; + OType *out_gate_sh_curr = out_gate_sh + buff * buff_elems; +#pragma unroll + for (int stage = 0; stage < BUFFER_STAGES_NUM; ++stage) { + const size_t stage_offset_Y = stage * THREADS_PER_CHUNK_Y; + const size_t shmem_offset_y = thread_offset_Y + stage_offset_Y; + const size_t shmem_offset_x = thread_offset_X; + const size_t shmem_idx = shmem_offset_y * SHMEM_DIM_X + shmem_offset_x; + + float act_elt = static_cast(in_act_sh_curr[shmem_idx]); + 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; + } + + if constexpr (IS_BWD) { + float grad_elt = static_cast(in_grad_sh_curr[shmem_idx]); + + const float x = act_elt; + float act_x; + float dact_x; + if constexpr (std::is_same::value) { + const float x = min(act_elt, p.limit); + const float s = sigmoidf(p.alpha * x); + act_x = x * s; + if (act_elt <= p.limit) { + dact_x = s + s * (1 - s) * p.alpha * x; + } else { + dact_x = 0.0f; + } + } else { + if constexpr ((ActOP == &silu) && (DActOP == &dsilu)) { + const float s = sigmoidf(x); + act_x = x * s; + dact_x = x * s * (1 - s) + s; + } else { + act_x = ActOP(x, p); + dact_x = DActOP(x, p); + } + } + float after_dact = dact_x * grad_elt * gate_elt; + float after_dgate = dgate_elt ? act_x * grad_elt : 0.0f; + + out_act_sh_curr[shmem_idx] = static_cast(scale * after_dact); + out_gate_sh_curr[shmem_idx] = static_cast(scale * after_dgate); + + amax = fmaxf(amax, fabsf(after_dact)); + amax = fmaxf(amax, fabsf(after_dgate)); + } else { + const float after_act = ActOP(act_elt, p) * gate_elt; + out_act_sh_curr[shmem_idx] = static_cast(scale * after_act); + amax = fmaxf(amax, fabsf(after_act)); + } + } + + // Wait for shared memory writes to be visible to TMA engine (cross-proxy fence) + 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 size_t chunk_it_offset_y = chunk_offset_Y + it * BUFFER_DIM_Y; + const size_t chunk_it_offset_x = chunk_offset_X; + + // dGeLU + ptx::cp_async_bulk_tensor_2d_shared_to_global(TMAP_output_act, chunk_it_offset_x, + chunk_it_offset_y, + reinterpret_cast(out_act_sh_curr)); + + if constexpr (IS_BWD) { + // dGate + ptx::cp_async_bulk_tensor_2d_shared_to_global( + TMAP_output_gate, chunk_it_offset_x, chunk_it_offset_y, + reinterpret_cast(out_gate_sh_curr)); + } + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + + // Wait for TMA transfer to have finished reading shared memory. + ptx::cp_async_bulk_wait_group_read(); + } + } + ptx::cp_async_bulk_wait_group_read<0>(); + __syncthreads(); + + if (amax_ptr != nullptr) { + const int warp_id = threadIdx.x / THREADS_PER_WARP; + // Reduce the amax over the block + amax = reduce_max(amax, warp_id); + // Update the global amax + if (is_master_thread) { + atomicMaxFloat(amax_ptr, amax); + } + } + + // Update scale-inverse + if (is_master_thread && blockIdx.x == 0 && (scale_inv_ptr != nullptr)) { + reciprocal(scale_inv_ptr, scale); + } + + // Destroy the barriers. This invalidates the memory region of the barrier. + // If further computations were to take place in the kernel, this allows the + // memory location of the shared memory barrier to be reused. + if (is_master_thread) { +#pragma unroll + for (int it = 0; it < ITERATIONS; ++it) { + ptx::mbarrier_invalid(&mbar[it]); + } + } +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +} // namespace kernel + +template +void cast_gated_tma(const Tensor &gated_input, const Tensor &grad, Tensor *output, ParamOP &p, + cudaStream_t stream) { + using namespace kernel; + checkCuDriverContext(stream); + + NVTE_CHECK(!output->has_columnwise_data(), "Only rowwise cast supported in this function."); + const size_t rows = gated_input.flat_first_dim(); + const size_t cols = gated_input.flat_last_dim() / 2; + const size_t output_cols = (IS_BWD ? 2 : 1) * cols; + + const size_t blocks_Y = DIVUP(rows, CHUNK_DIM_Y); + const size_t blocks_X = DIVUP(cols, CHUNK_DIM_X); + + float *const amax_ptr = reinterpret_cast(output->amax.dptr); + float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); + float *const scale_ptr = reinterpret_cast(output->scale.dptr); + + const dim3 block_dim(THREADS_PER_CHUNK); + const dim3 grid_dim(blocks_X, blocks_Y); + + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + gated_input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( + output->dtype(), OType, + + alignas(64) CUtensorMap tensor_map_grad{}; + alignas(64) CUtensorMap tensor_map_input_act{}; + alignas(64) CUtensorMap tensor_map_input_gate{}; + alignas(64) CUtensorMap tensor_map_output_act{}; + alignas(64) CUtensorMap tensor_map_output_gate{}; + + if constexpr (IS_BWD) { + create_2D_tensor_map(tensor_map_grad, grad.data, rows, cols, SHMEM_DIM_Y, SHMEM_DIM_X, + cols, 0, typeToNumBits(gated_input.dtype())); + } + + const uint32_t tensor_stride_elems = output_cols; + + create_2D_tensor_map(tensor_map_input_act, gated_input.data, rows, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, cols * 2, 0, typeToNumBits(gated_input.dtype())); + create_2D_tensor_map(tensor_map_input_gate, gated_input.data, rows, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, cols * 2, cols, typeToNumBits(gated_input.dtype())); + create_2D_tensor_map(tensor_map_output_act, output->data, rows, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, tensor_stride_elems, 0, typeToNumBits(output->dtype())); + create_2D_tensor_map(tensor_map_output_gate, output->data, rows, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, tensor_stride_elems, cols, + typeToNumBits(output->dtype())); + + const size_t buff_elems_total = BUFFERS_NUM * SHMEM_DIM_Y * SHMEM_DIM_X; + const size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); + const size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + const size_t grad_mem = (IS_BWD ? buff_size_aligned_in : 0); + const size_t in_act_mem = buff_size_aligned_in; + const size_t in_gate_mem = buff_size_aligned_in; + const size_t out_act_mem = buff_size_aligned_out; + const size_t out_gate_mem = buff_size_aligned_out; + + const size_t shmem_size = grad_mem + (in_act_mem + in_gate_mem) + + (out_act_mem + out_gate_mem) + TMA_SHMEM_ALIGNMENT; + + auto kernel = cast_fp8_gated_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + shmem_size)); + + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, tensor_map_output_act, + tensor_map_output_gate, amax_ptr, scale_inv_ptr, scale_ptr, rows, cols, p); + NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) + ); // NOLINT(*) +} + +template +void cast_gated_fwd(const Tensor &input, Tensor *output, ParamOP &p, cudaStream_t stream) { + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( + output->dtype(), OType, + + constexpr int nvec = 32 / sizeof(IType); + GatedActivationKernelLauncher( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->scale.dptr), + reinterpret_cast(output->amax.dptr), + reinterpret_cast(output->scale_inv.dptr), input.flat_first_dim(), + output->flat_last_dim(), p, stream);); // NOLINT(*) + ); // NOLINT(*) +} + +template +void cast_gated_bwd(const Tensor &input, const Tensor &grad, Tensor *output, ParamOP &p, + cudaStream_t stream) { + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( + output->dtype(), OType, + + constexpr int nvec = 32 / sizeof(IType); + DGatedActivationKernelLauncher( + reinterpret_cast(grad.data.dptr), + reinterpret_cast(input.data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->scale.dptr), + reinterpret_cast(output->amax.dptr), + reinterpret_cast(output->scale_inv.dptr), grad.flat_first_dim(), + grad.flat_last_dim(), p, stream);); // NOLINT(*) + ); // NOLINT(*) +} +} // namespace fp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GATED_FP8_CUH_ diff --git a/transformer_engine/common/cast/fp8/quantize_fp8.cuh b/transformer_engine/common/cast/fp8/quantize_fp8.cuh new file mode 100644 index 0000000000..efc5015b75 --- /dev/null +++ b/transformer_engine/common/cast/fp8/quantize_fp8.cuh @@ -0,0 +1,580 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_fp8.cuh + * \brief CUDA kernels to quantize to FP8. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_FP8_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_FP8_CUH_ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../transpose/cast_transpose.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../util/vectorized_pointwise.h" +#include "../../utils.cuh" +#include "../core/common.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace fp8 { +namespace quantize_2D_kernel { + +constexpr size_t FP8_CHUNK_DIM_Y = 128; +constexpr size_t FP8_CHUNK_DIM_X = 128; +constexpr size_t FP8_THREADS_PER_CHUNK = 128; +constexpr size_t FP8_BUFFERS_NUM = 2; +constexpr size_t FP8_PREFETCH_BUFFERS_NUM = 1; +static_assert(FP8_PREFETCH_BUFFERS_NUM < FP8_BUFFERS_NUM); + +constexpr size_t FP8_BUFFER_DIM_Y = 16; +constexpr size_t FP8_BUFFER_DIM_X = FP8_CHUNK_DIM_X; // 128 +constexpr size_t FP8_SHMEM_DIM_Y = FP8_BUFFER_DIM_Y; // 16 +constexpr size_t FP8_SHMEM_DIM_X = FP8_BUFFER_DIM_X; // 128 + +constexpr size_t FP8_BUFF_STAGES_NUM = FP8_BUFFER_DIM_Y; // 16 +constexpr size_t FP8_ITERATIONS = FP8_CHUNK_DIM_Y / FP8_BUFFER_DIM_Y; // 8 = 128 / 16 +static_assert(FP8_ITERATIONS >= FP8_PREFETCH_BUFFERS_NUM); + +template +__global__ void __launch_bounds__(FP8_THREADS_PER_CHUNK) + cast_fp8_2D_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_act_input, + const __grid_constant__ CUtensorMap tensor_map_output, + float *const dbias_workspace, float *const amax_ptr, + float *const scale_inv_ptr, const float *const scale_ptr, const size_t rows, + const size_t cols) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + + const size_t block_offset_Y = blockIdx.y * FP8_CHUNK_DIM_Y; + const size_t block_offset_X = blockIdx.x * FP8_CHUNK_DIM_X; + + const size_t tid_Y = threadIdx.x / FP8_THREADS_PER_CHUNK; + const size_t tid_X = threadIdx.x % FP8_THREADS_PER_CHUNK; + + const size_t thread_offset_Y = tid_Y; + const size_t thread_offset_X = tid_X; + + const size_t dbias_offset_Y = blockIdx.y + tid_Y; + const size_t my_column = blockIdx.x * FP8_CHUNK_DIM_X + thread_offset_X; + const bool col_out_of_bounds = my_column >= cols; + const size_t dbias_stride = cols; + + float partial_dbias = 0.f; + + float amax = 0; + const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; + + // The destination shared memory buffer of a bulk tensor operation should be 128-byte aligned + __shared__ alignas(TMA_SHMEM_ALIGNMENT) + IType in_sh[FP8_BUFFERS_NUM][FP8_SHMEM_DIM_Y][FP8_SHMEM_DIM_X]; + __shared__ alignas(TMA_SHMEM_ALIGNMENT) + IType act_in_sh[FP8_BUFFERS_NUM][FP8_SHMEM_DIM_Y][FP8_SHMEM_DIM_X]; + __shared__ alignas(TMA_SHMEM_ALIGNMENT) + OType out_sh[FP8_BUFFERS_NUM][FP8_SHMEM_DIM_Y][FP8_SHMEM_DIM_X]; + + constexpr size_t shmem_buff_size = sizeof(in_sh) / FP8_BUFFERS_NUM; + + const bool is_master_thread = (threadIdx.x == 0); + +// 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[FP8_ITERATIONS]; + + initialize_barriers(mbar, is_master_thread); + + int parity = 0; + + const size_t chunk_offset_Y = block_offset_Y; + const size_t chunk_offset_X = block_offset_X; + +#pragma unroll + for (int prefetch_buff = 0; prefetch_buff < FP8_PREFETCH_BUFFERS_NUM; ++prefetch_buff) { + const size_t chunk_stage_offset_Y = chunk_offset_Y + prefetch_buff * FP8_BUFFER_DIM_Y; + const size_t chunk_stage_offset_X = chunk_offset_X; + if constexpr (IS_DACT) { + copy_2d_to_sharedx2(&in_sh[prefetch_buff], &tensor_map_input, chunk_stage_offset_X, + chunk_stage_offset_Y, &act_in_sh[prefetch_buff], &tensor_map_act_input, + chunk_stage_offset_X, chunk_stage_offset_Y, shmem_buff_size, + &mbar[prefetch_buff], is_master_thread); + } else { + copy_2d_to_shared(&in_sh[prefetch_buff], &tensor_map_input, chunk_stage_offset_X, + chunk_stage_offset_Y, shmem_buff_size, &mbar[prefetch_buff], + is_master_thread); + } + } + +#pragma unroll + for (int iter = 0; iter < FP8_ITERATIONS; ++iter) { + const size_t buff = iter % FP8_BUFFERS_NUM; + const size_t next_iter = iter + FP8_PREFETCH_BUFFERS_NUM; + const size_t row_base = block_offset_Y + iter * FP8_BUFFER_DIM_Y; + if (next_iter < FP8_ITERATIONS) { + const size_t next_buff = next_iter % FP8_BUFFERS_NUM; + const size_t chunk_it_offset_y = chunk_offset_Y + next_iter * FP8_BUFFER_DIM_Y; + const size_t chunk_it_offset_x = chunk_offset_X; + if constexpr (IS_DACT) { + copy_2d_to_sharedx2(&in_sh[next_buff], &tensor_map_input, chunk_it_offset_x, + chunk_it_offset_y, &act_in_sh[next_buff], &tensor_map_act_input, + chunk_it_offset_x, chunk_it_offset_y, shmem_buff_size, &mbar[next_iter], + is_master_thread); + } else { + copy_2d_to_shared(&in_sh[next_buff], &tensor_map_input, chunk_it_offset_x, + chunk_it_offset_y, shmem_buff_size, &mbar[next_iter], is_master_thread); + } + } + + // Wait for the data to have arrived + ptx::mbarrier_wait_parity(&mbar[iter], parity); + +#pragma unroll + for (int stage = 0; stage < FP8_BUFF_STAGES_NUM; ++stage) { + const size_t stage_offset_Y = stage; + const size_t shmem_offset_y = thread_offset_Y + stage_offset_Y; + const size_t shmem_offset_x = thread_offset_X; + const size_t row = row_base + shmem_offset_y; + const bool row_out_of_bounds = row >= rows; + const bool out_of_bounds = col_out_of_bounds || row_out_of_bounds; + + float elt = static_cast(in_sh[buff][shmem_offset_y][shmem_offset_x]); + if constexpr (IS_DACT) { + float act_in_elt = static_cast(act_in_sh[buff][shmem_offset_y][shmem_offset_x]); + elt *= OP(act_in_elt, {}); + } + if constexpr (IS_DBIAS) { + if constexpr (IS_DACT) { + if (!out_of_bounds) { + partial_dbias += elt; + } + } else { + // If no activation, elt is 0 so we can safely do this + partial_dbias += elt; + } + } + __builtin_assume(amax >= 0); + if (IS_DACT) { + if (!out_of_bounds) { + amax = fmaxf(amax, fabsf(elt)); + } + } else { + // If no activation, elt is 0 so we can safely do this + amax = fmaxf(amax, fabsf(elt)); + } + out_sh[buff][shmem_offset_y][shmem_offset_x] = static_cast(elt * scale); + } + + // 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 size_t chunk_it_offset_y = chunk_offset_Y + iter * FP8_BUFFER_DIM_Y; + const size_t chunk_it_offset_x = chunk_offset_X; + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), chunk_it_offset_x, + chunk_it_offset_y, reinterpret_cast(&out_sh[buff])); + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + + // Wait for TMA transfer to have finished reading shared memory. + ptx::cp_async_bulk_wait_group_read(); + } + } + ptx::cp_async_bulk_wait_group_read<0>(); + __syncthreads(); + + parity ^= 1; + + if constexpr (IS_DBIAS) { + const size_t dbias_offset_X = my_column; + const size_t dbias_offset = dbias_offset_Y * dbias_stride + dbias_offset_X; + if (!col_out_of_bounds) { + dbias_workspace[dbias_offset] = partial_dbias; + } + } + + if (amax_ptr != nullptr) { + const int warp_id = threadIdx.x / THREADS_PER_WARP; + // Reduce the amax over the block + amax = reduce_max(amax, warp_id); + // Update the global amax + if (is_master_thread) { + atomicMaxFloat(amax_ptr, amax); + } + } + + // Update scale-inverse + if (is_master_thread && blockIdx.x == 0 && (scale_inv_ptr != nullptr)) { + reciprocal(scale_inv_ptr, scale); + } + + destroy_barriers(mbar, is_master_thread); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +} // namespace quantize_2D_kernel + +namespace quantize_1D_kernel { +using namespace quantize_2D_kernel; + +constexpr size_t CHUNKS_PER_BLOCK = 128; +constexpr size_t THREADS_PER_BLOCK = FP8_THREADS_PER_CHUNK; +constexpr size_t CHUNK_SIZE = THREADS_PER_BLOCK; +constexpr size_t ELEMS_PER_BLOCK = CHUNKS_PER_BLOCK * CHUNK_SIZE; +constexpr size_t CHUNKS_PER_ITERATION = 32; +constexpr size_t SHMEM_DIM = CHUNKS_PER_ITERATION * CHUNK_SIZE; +constexpr size_t ITERATIONS = CHUNKS_PER_BLOCK / CHUNKS_PER_ITERATION; +constexpr size_t SHMEM_BUFFERS = 2; +static_assert(CHUNKS_PER_BLOCK % CHUNKS_PER_ITERATION == 0); + +template +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + cast_fp8_1D_kernel(const IType *input_ptr, OType *output_ptr, float *const amax_ptr, + float *const scale_inv_ptr, const float *const scale_ptr, const size_t N) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + + const size_t block_offset = blockIdx.x * ELEMS_PER_BLOCK; + const IType *input = input_ptr + block_offset; + OType *output = output_ptr + block_offset; + + float amax = 0; + const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; + + // The destination shared memory buffer of a bulk tensor operation should be 128-byte aligned + __shared__ alignas(TMA_SHMEM_ALIGNMENT) IType in_sh[SHMEM_BUFFERS][SHMEM_DIM]; + __shared__ alignas(TMA_SHMEM_ALIGNMENT) OType out_sh[SHMEM_BUFFERS][SHMEM_DIM]; + + constexpr size_t transaction_size_IN = sizeof(in_sh) / SHMEM_BUFFERS; + constexpr size_t transaction_size_OUT = sizeof(out_sh) / SHMEM_BUFFERS; + + const bool is_master_thread = (threadIdx.x == 0); + +// 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[ITERATIONS]; + + initialize_barriers(mbar, is_master_thread); + + int parity = 0; + + copy_1d_to_shared(&(in_sh[0]), input, transaction_size_IN, &(mbar[0]), is_master_thread); + +#pragma unroll + for (int iter = 0; iter < ITERATIONS; ++iter) { + const size_t buff = iter % SHMEM_BUFFERS; + const size_t it_offset = iter * SHMEM_DIM; + + const size_t next_iter = iter + 1; + const size_t next_buff = next_iter % SHMEM_BUFFERS; + const size_t next_iter_offset = next_iter * SHMEM_DIM; + + if (next_iter < ITERATIONS) { + copy_1d_to_shared(&(in_sh[next_buff]), input + next_iter_offset, transaction_size_IN, + &(mbar[next_iter]), is_master_thread); + } + + ptx::fence_proxy_async_shared_cta(); + + // Wait for the data to have arrived + ptx::mbarrier_wait_parity(&mbar[iter], parity); + +#pragma unroll + for (int chunk = 0; chunk < CHUNKS_PER_ITERATION; ++chunk) { + const size_t shmem_offset = chunk * CHUNK_SIZE + threadIdx.x; + float elt = static_cast(in_sh[buff][shmem_offset]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + __builtin_assume(amax >= 0); + amax = fmaxf(amax, fabsf(elt)); + out_sh[buff][shmem_offset] = static_cast(elt * scale); + } + + // 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) { + ptx::cp_async_bulk_tensor_1d_shared_to_global( + reinterpret_cast(output + it_offset), + reinterpret_cast(&out_sh[buff]), transaction_size_OUT); + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + + // Wait for TMA transfer to have finished reading shared memory. + ptx::cp_async_bulk_wait_group_read<1>(); + } + } + ptx::cp_async_bulk_wait_group_read<0>(); + __syncthreads(); + + if (amax_ptr != nullptr) { + const int warp_id = threadIdx.x / THREADS_PER_WARP; + // Reduce the amax over the block + amax = reduce_max(amax, warp_id); + // Update the global amax + if (is_master_thread) { + atomicMaxFloat(amax_ptr, amax); + } + } + + // Update scale-inverse + if (is_master_thread && blockIdx.x == 0 && (scale_inv_ptr != nullptr)) { + reciprocal(scale_inv_ptr, scale); + } + + destroy_barriers(mbar, is_master_thread); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +} // namespace quantize_1D_kernel + +template +void quantize_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { + using namespace quantize_1D_kernel; + const size_t N = product(input.data.shape); + + const bool isFullTile = (N % ELEMS_PER_BLOCK == 0); + NVTE_CHECK(isFullTile, "Only full tiles are supported."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + + const size_t chunks = DIVUP(N, CHUNK_SIZE); + const size_t blocks = DIVUP(chunks, CHUNKS_PER_BLOCK); + + float *const amax_ptr = reinterpret_cast(output->amax.dptr); + float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); + const float *const scale_ptr = reinterpret_cast(output->scale.dptr); + + const dim3 block(THREADS_PER_BLOCK); + const dim3 grid(blocks); + + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + const IType *input_ptr = reinterpret_cast(input.data.dptr); + OType *output_ptr = reinterpret_cast(output->data.dptr); + + cast_fp8_1D_kernel<<>>( + input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N);); // NOLINT(*) + ); // NOLINT(*) + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +template +void quantize_2D(const Tensor &input, const Tensor *act_input, Tensor *output, Tensor *dbias, + Tensor *workspace, cudaStream_t stream) { + using namespace quantize_2D_kernel; + checkCuDriverContext(stream); + + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + 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; + const size_t blocks_X = chunks_X; + + const size_t dbias_rows = blocks_Y; + const size_t dbias_cols = cols; + + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + + 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(workspace != nullptr, "Workspace must be a tensor."); + + if (workspace->data.dptr == nullptr) { + workspace->data.shape = {dbias_rows, dbias_cols}; + workspace->data.dtype = DType::kFloat32; + return; + } + } + float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + float *const amax_ptr = reinterpret_cast(output->amax.dptr); + float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); + float *const scale_ptr = reinterpret_cast(output->scale.dptr); + + const dim3 block(FP8_THREADS_PER_CHUNK); + const dim3 grid(blocks_X, blocks_Y); + + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.data.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->data.dtype, OType, + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_act_input{}; + alignas(64) CUtensorMap tensor_map_output{}; + + create_2D_tensor_map(tensor_map_input, input.data, rows, cols, FP8_SHMEM_DIM_Y, + FP8_SHMEM_DIM_X, cols, 0, typeToNumBits(input.data.dtype)); + + if constexpr (IS_DACT) { + create_2D_tensor_map(tensor_map_act_input, act_input->data, rows, cols, FP8_SHMEM_DIM_Y, + FP8_SHMEM_DIM_X, cols, 0, typeToNumBits(input.data.dtype)); + } + + create_2D_tensor_map(tensor_map_output, output->data, rows, cols, FP8_SHMEM_DIM_Y, + FP8_SHMEM_DIM_X, cols, 0, typeToNumBits(output->data.dtype)); + + cast_fp8_2D_kernel + <<>>(tensor_map_input, tensor_map_act_input, tensor_map_output, + workspace_ptr, amax_ptr, scale_inv_ptr, scale_ptr, rows, + cols); + NVTE_CHECK_CUDA(cudaGetLastError()); + + if constexpr (IS_DBIAS) { + common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); + }); // NOLINT(*) + ); // NOLINT(*) +} + +namespace detail { +using Empty = transformer_engine::Empty; +__device__ inline float identity(float value, const Empty &) { return value; } +} // namespace detail + +template +void CastVectorizedUnaryKernelLauncher(const Tensor &input, const Tensor *noop, Tensor *output, + cudaStream_t stream) { + constexpr float (*UnaryOP)(float, const ParamOP &) = (OP == nullptr) ? detail::identity : OP; + const size_t N = product(input.data.shape); + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.data.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( + output->data.dtype, OType, + if (!is_fp8_dtype(output->data.dtype) || is_tensor_scaling(output->scaling_mode)) { + constexpr int nvec = 32 / sizeof(IType); + VectorizedUnaryKernelLauncher( + reinterpret_cast(input.data.dptr), + reinterpret_cast(noop->data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->scale.dptr), + reinterpret_cast(output->amax.dptr), + reinterpret_cast(output->scale_inv.dptr), N, {}, stream); + } else { + NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); + }); // NOLINT(*) + ); // NOLINT(*) +} + +template +void CastVectorizedUnaryGradKernelLauncher(const Tensor &grad, const Tensor *input, Tensor *output, + cudaStream_t stream) { + constexpr float (*UnaryOP)(float, const ParamOP &) = (OP == nullptr) ? detail::identity : OP; + const size_t N = product(input->data.shape); + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input->data.dtype, IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( + output->data.dtype, OType, + if (!is_fp8_dtype(output->data.dtype) || is_tensor_scaling(output->scaling_mode)) { + constexpr int nvec = 32 / sizeof(IType); + VectorizedUnaryGradKernelLauncher( + reinterpret_cast(grad.data.dptr), + reinterpret_cast(input->data.dptr), + reinterpret_cast(output->data.dptr), + reinterpret_cast(output->scale.dptr), + reinterpret_cast(output->amax.dptr), + reinterpret_cast(output->scale_inv.dptr), N, {}, stream); + } else { + NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); + }); // NOLINT(*) + ); // NOLINT(*) +} + +template +void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, Tensor *output, + Tensor *dbias, Tensor *workspace, cudaStream_t stream) { + using namespace quantize_1D_kernel; + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "cast_input"); + CheckOutputTensor(*output, "cast_output"); + + if constexpr (IS_DBIAS) { + NVTE_CHECK(dbias != nullptr); + CheckOutputTensor(*dbias, "dbias"); + } + if constexpr (IS_DACT) { + NVTE_CHECK(act_input != nullptr); + CheckInputTensor(*act_input, "activation_input"); + NVTE_CHECK(input.dtype() == act_input->dtype(), "Types of both inputs must match."); + NVTE_CHECK(input.data.shape == act_input->data.shape, "Shapes of both inputs must match."); + } + + NVTE_CHECK(!is_fp8_dtype(input.dtype()), "Input must be in higher precision."); + NVTE_CHECK(output->data.shape == input.data.shape, "Input and output shapes need to match."); + + // Supported by the Arch >= 10.0 + if (is_supported_by_CC_100()) { + if (!IS_DBIAS && !IS_DACT) { + if (common::full_tile_1D_tensor(output, ELEMS_PER_BLOCK) && is_fp8_dtype(output->dtype()) && + is_aligned_tensor_data(input, TMA_GMEM_ALIGNMENT) && + is_aligned_tensor_data(*output, TMA_GMEM_ALIGNMENT)) { + // Aligned AND FP8 + quantize_1D(input, output, stream); + } else { + // Unaligned + CastVectorizedUnaryKernelLauncher(input, noop, output, stream); + } + } else if (!IS_DBIAS && IS_DACT) { + if (common::dimensions_supported_by_TMA(output) && is_fp8_dtype(output->dtype()) && + is_aligned_tensor_data(input, TMA_GMEM_ALIGNMENT) && + is_aligned_tensor_data(*output, TMA_GMEM_ALIGNMENT) && + is_aligned_tensor_data(*act_input, TMA_GMEM_ALIGNMENT)) { + // Aligned AND FP8 (+dAct) + quantize_2D(input, act_input, output, dbias, workspace, + stream); + } else { + // Unaligned + CastVectorizedUnaryGradKernelLauncher(input, act_input, output, stream); + } + } else { + quantize_2D(input, act_input, output, dbias, workspace, + stream); + } + } else { + if (IS_DBIAS) { + // zhongboz: should we just ignore IS_ACT here? + NVTE_ERROR("Not implemented scaling mode or fusion: " + to_string(output->scaling_mode) + + " or IS_DBIAS=true" + " on GPU with compute capability < 10.0."); + } + if (!IS_DACT) { + CastVectorizedUnaryKernelLauncher(input, noop, output, stream); + } else { + CastVectorizedUnaryGradKernelLauncher(input, act_input, output, stream); + } + } +} + +} // namespace fp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_FP8_CUH_ diff --git a/transformer_engine/common/util/dequantize_kernels.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh similarity index 69% rename from transformer_engine/common/util/dequantize_kernels.cuh rename to transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index 9f70ce4cd4..fb43fce96b 100644 --- a/transformer_engine/common/util/dequantize_kernels.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -4,36 +4,27 @@ * See LICENSE for license information. ************************************************************************/ -/*! \file dequantize_kernels.cuh - * \brief CUDA kernels to cast from MXFP8. +/*! \file dequantize_mxfp8.cuh + * \brief CUDA kernels to dequantize from MXFP8. */ -#ifndef TRANSFORMER_ENGINE_DEQUANTIZE_KERNELS_CUH_ -#define TRANSFORMER_ENGINE_DEQUANTIZE_KERNELS_CUH_ +#ifndef TRANSFORMER_ENGINE_DEQUANTIZE_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_DEQUANTIZE_MXFP8_CUH_ #include #include #include -#include - -#include -#include -#include -#include - -#include "../common.h" -#include "../transpose/cast_transpose.h" -#include "../util/vectorized_pointwise.h" -#include "../utils.cuh" -#include "math.h" -#include "ptx.cuh" -#include "transformer_engine/activation.h" -#include "transformer_engine/transformer_engine.h" -#include "transformer_engine/transpose.h" +#include -namespace transformer_engine { +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" -namespace dequantization { +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace dequantize_kernel { constexpr size_t CHUNK_DIM_Y = 128; constexpr size_t CHUNK_DIM_X = 128; @@ -228,29 +219,10 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +} // namespace dequantize_kernel -void fp8_dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { - NVTE_CHECK(is_fp8_dtype(input.data.dtype), "Input must have FP8 type."); - NVTE_CHECK(!is_fp8_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 size_t N = product(input.data.shape); - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - input.data.dtype, IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - output->data.dtype, OType, - - constexpr int nvec = 32 / sizeof(OType); - detail::DequantizeParam p; - p.scale_inv = reinterpret_cast(input.scale_inv.dptr); - VectorizedUnaryKernelLauncher( - reinterpret_cast(input.data.dptr), nullptr, - reinterpret_cast(output->data.dptr), nullptr, nullptr, nullptr, N, p, - stream);); // NOLINT(*) - ); // NOLINT(*) -} - -void mxfp8_dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { +inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { + using namespace dequantize_kernel; bool use_rowwise_scaling = input.has_data(); bool use_colwise_scaling = input.has_columnwise_data(); checkCuDriverContext(stream); @@ -334,113 +306,8 @@ void mxfp8_dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); } - -#if CUDA_VERSION >= 12080 -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 thread_idx = blockIdx.x * blockDim.x + threadIdx.x; - const size_t x = thread_idx % M; - const size_t y = thread_idx / M; - - union fp4vec { - uint64_t vec; - fp4e2m1x4 small_vec[4]; - }; - using OVec = Vec; - const uint64_t *const input_vectorized = reinterpret_cast(input); - OVec *output_vec = reinterpret_cast(output); - - const size_t my_index = x + y * M; - const size_t 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]; - fp8e4m3 scale = scales[my_scale_index]; - float amax = *tensor_amax; - constexpr float factor_inv = 1.0 / (6.0 * 448.0); - float final_scale = static_cast(scale) * amax * factor_inv; -#pragma unroll - for (int i = 0; i < 4; i++) { - float4 current = static_cast(value.small_vec[i]); - OVec out; - out.data.elt[0] = static_cast(current.x * final_scale); - out.data.elt[1] = static_cast(current.y * final_scale); - out.data.elt[2] = static_cast(current.z * final_scale); - out.data.elt[3] = static_cast(current.w * final_scale); - output_vec[my_output_index + i] = out; - } -} -#endif // CUDA_VERSION - -void fp4_dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { -#if CUDA_VERSION >= 12080 - CheckInputTensor(input, "input"); - CheckOutputTensor(*output, "output"); - NVTE_CHECK(input.data.dtype == DType::kFloat4E2M1, "Input must have FP4 type."); - 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."); - - constexpr int FP4_BLOCK_SIZE = 16; - const size_t N = input.flat_first_dim(); - const size_t M = input.flat_last_dim(); - - 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, "."); - - const size_t Mread = M / FP4_BLOCK_SIZE; - const size_t total = N * Mread; - const size_t threads = 512; - const size_t blocks = DIVUP(total, threads); - - 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(*) - NVTE_CHECK_CUDA(cudaGetLastError()); -#else - NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); -#endif // CUDA_VERSION >= 12080 -} - -} // namespace dequantization - -namespace detail { - -void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t stream) { - CheckInputTensor(input, "cast_input"); - CheckOutputTensor(*output, "cast_output"); - - switch (input.scaling_mode) { - case NVTE_DELAYED_TENSOR_SCALING: { - dequantization::fp8_dequantize(input, output, stream); - break; - } - case NVTE_MXFP8_1D_SCALING: { - if (is_supported_by_CC_100()) { - dequantization::mxfp8_dequantize(input, output, stream); - } else { - NVTE_ERROR("MXFP8 Dequantization is NOT supported by architectures < 10.0"); - } - break; - } - case NVTE_NVFP4_1D_SCALING: { - dequantization::fp4_dequantize(input, output, stream); - break; - } - default: - NVTE_ERROR("Not implemented scaling mode: " + to_string(input.scaling_mode) + "."); - } -} - -} // namespace detail - +} // namespace mxfp8 +} // namespace dispatch } // namespace transformer_engine -#endif // TRANSFORMER_ENGINE_DEQUANTIZE_KERNELS_CUH_ +#endif // TRANSFORMER_ENGINE_DEQUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/util/cast_gated_kernels.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh similarity index 53% rename from transformer_engine/common/util/cast_gated_kernels.cuh rename to transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 93086bd827..4f0e1b80f7 100644 --- a/transformer_engine/common/util/cast_gated_kernels.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -4,280 +4,27 @@ * See LICENSE for license information. ************************************************************************/ -/*! \file cast_gated_kernels.cuh - * \brief CUDA gated activations kernels to cast to/from FP8/MXFP8. +/*! \file gated_mxfp8.cuh + * \brief CUDA kernels to cast to MXFP8 with gated activations. */ -#ifndef TRANSFORMER_ENGINE_CAST_GATED_KERNELS_CUH_ -#define TRANSFORMER_ENGINE_CAST_GATED_KERNELS_CUH_ +#ifndef TRANSFORMER_ENGINE_GATED_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_GATED_MXFP8_CUH_ #include #include #include -#include -#include +#include -#include - -#include "../common.h" -#include "../util/vectorized_pointwise.h" -#include "../utils.cuh" -#include "math.h" -#include "ptx.cuh" +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" namespace transformer_engine { - -namespace gated_kernels { - -constexpr size_t CHUNK_DIM_Y = 128; -constexpr size_t CHUNK_DIM_X = 128; -constexpr size_t THREADS_PER_CHUNK = 512; -constexpr size_t THREADS_PER_CHUNK_X = CHUNK_DIM_X; -constexpr size_t THREADS_PER_CHUNK_Y = THREADS_PER_CHUNK / THREADS_PER_CHUNK_X; // 4 = 512 / 128 -constexpr size_t BUFFERS_NUM = 2; -constexpr size_t BUFFER_DIM_Y = 32; -constexpr size_t BUFFER_DIM_X = CHUNK_DIM_X; // 128 -constexpr size_t SHMEM_DIM_Y = BUFFER_DIM_Y; // 32 -constexpr size_t SHMEM_DIM_X = BUFFER_DIM_X; // 128 - -constexpr size_t BUFFER_STAGES_NUM = BUFFER_DIM_Y / THREADS_PER_CHUNK_Y; // 8 = 32 / 4 -constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y; // 4 = 128 / 32 -static_assert(ITERATIONS >= 1); - -__device__ inline float sigmoidf(const float x) { return __frcp_rn(1.0f + __expf(-x)); } - -template -__global__ void __launch_bounds__(THREADS_PER_CHUNK) - cast_fp8_gated_kernel(const __grid_constant__ CUtensorMap tensor_map_grad, - const __grid_constant__ CUtensorMap tensor_map_input_act, - const __grid_constant__ CUtensorMap tensor_map_input_gate, - const __grid_constant__ CUtensorMap tensor_map_output_act, - const __grid_constant__ CUtensorMap tensor_map_output_gate, - float *const amax_ptr, float *const scale_inv_ptr, - const float *const scale_ptr, const size_t rows, const size_t cols, - const ParamOP p) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - - const size_t chunk_offset_Y = blockIdx.y * CHUNK_DIM_Y; - const size_t chunk_offset_X = blockIdx.x * CHUNK_DIM_X; - - const size_t tid_Y = threadIdx.x / THREADS_PER_CHUNK_X; - const size_t tid_X = threadIdx.x % THREADS_PER_CHUNK_X; - - const size_t thread_offset_Y = tid_Y; - const size_t thread_offset_X = tid_X; - - float amax = 0; - const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; - - 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)); - - constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X; - constexpr size_t buff_elems_total = BUFFERS_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 = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); - - constexpr size_t grad_mem = IS_DGATED ? buff_size_aligned_in : 0; - - constexpr size_t in_act_mem = buff_size_aligned_in; - constexpr size_t in_gate_mem = buff_size_aligned_in; - constexpr size_t in_mem = in_act_mem + in_gate_mem; - - constexpr size_t out_act_mem = buff_size_aligned_out; - constexpr size_t in_transaction_size = buff_elems * sizeof(IType); - - // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - IType *in_grad_sh = reinterpret_cast(dshmem); - IType *in_act_sh = reinterpret_cast(dshmem + grad_mem); - IType *in_gate_sh = reinterpret_cast(dshmem + grad_mem + in_act_mem); - OType *out_act_sh = reinterpret_cast(dshmem + grad_mem + in_mem); - OType *out_gate_sh = reinterpret_cast(dshmem + grad_mem + in_mem + out_act_mem); - - const uint64_t *TMAP_grad_in = reinterpret_cast(&tensor_map_grad); - const uint64_t *TMAP_in_act = reinterpret_cast(&tensor_map_input_act); - const uint64_t *TMAP_in_gate = reinterpret_cast(&tensor_map_input_gate); - const uint64_t *TMAP_output_act = reinterpret_cast(&tensor_map_output_act); - const uint64_t *TMAP_output_gate = reinterpret_cast(&tensor_map_output_gate); - - const bool is_master_thread = (threadIdx.x == 0); - -// 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[ITERATIONS]; - - initialize_barriers(mbar, is_master_thread); - - int parity = 0; - - // Prefetch data of the first stage - - if constexpr (IS_DGATED) { - copy_2d_to_sharedx3(in_grad_sh, TMAP_grad_in, chunk_offset_X, chunk_offset_Y, in_act_sh, - TMAP_in_act, chunk_offset_X, chunk_offset_Y, in_gate_sh, TMAP_in_gate, - chunk_offset_X, chunk_offset_Y, in_transaction_size, &mbar[0], - is_master_thread); - } else { - copy_2d_to_sharedx2(in_act_sh, TMAP_in_act, chunk_offset_X, chunk_offset_Y, in_gate_sh, - TMAP_in_gate, chunk_offset_X, chunk_offset_Y, in_transaction_size, &mbar[0], - is_master_thread); - } - -#pragma unroll - for (int it = 0; it < ITERATIONS; ++it) { - const size_t buff = it % BUFFERS_NUM; - const size_t next_it = it + 1; - if (next_it < ITERATIONS) { - const size_t next_buff = next_it % BUFFERS_NUM; - const size_t chunk_it_offset_y = chunk_offset_Y + next_it * BUFFER_DIM_Y; - const size_t chunk_it_offset_x = chunk_offset_X; - if constexpr (IS_DGATED) { - copy_2d_to_sharedx3( - &in_grad_sh[next_buff * buff_elems], TMAP_grad_in, chunk_it_offset_x, chunk_it_offset_y, - &in_act_sh[next_buff * buff_elems], TMAP_in_act, chunk_it_offset_x, chunk_it_offset_y, - &in_gate_sh[next_buff * buff_elems], TMAP_in_gate, chunk_it_offset_x, chunk_it_offset_y, - in_transaction_size, &mbar[next_it], is_master_thread); - } else { - copy_2d_to_sharedx2(&in_act_sh[next_buff * buff_elems], TMAP_in_act, chunk_it_offset_x, - chunk_it_offset_y, &in_gate_sh[next_buff * buff_elems], TMAP_in_gate, - chunk_it_offset_x, chunk_it_offset_y, in_transaction_size, - &mbar[next_it], is_master_thread); - } - } - - ptx::fence_proxy_async_shared_cta(); - - // Wait for the data to have arrived - ptx::mbarrier_wait_parity(&mbar[it], parity); - - IType *in_grad_sh_curr = in_grad_sh + buff * buff_elems; - IType *in_act_sh_curr = in_act_sh + buff * buff_elems; - IType *in_gate_sh_curr = in_gate_sh + buff * buff_elems; - OType *out_act_sh_curr = out_act_sh + buff * buff_elems; - OType *out_gate_sh_curr = out_gate_sh + buff * buff_elems; -#pragma unroll - for (int stage = 0; stage < BUFFER_STAGES_NUM; ++stage) { - const size_t stage_offset_Y = stage * THREADS_PER_CHUNK_Y; - const size_t shmem_offset_y = thread_offset_Y + stage_offset_Y; - const size_t shmem_offset_x = thread_offset_X; - const size_t shmem_idx = shmem_offset_y * SHMEM_DIM_X + shmem_offset_x; - - float act_elt = static_cast(in_act_sh_curr[shmem_idx]); - 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; - } - - if constexpr (IS_DGATED) { - float grad_elt = static_cast(in_grad_sh_curr[shmem_idx]); - - const float x = act_elt; - float act_x; - float dact_x; - if constexpr (std::is_same::value) { - const float x = min(act_elt, p.limit); - const float s = sigmoidf(p.alpha * x); - act_x = x * s; - if (act_elt <= p.limit) { - dact_x = s + s * (1 - s) * p.alpha * x; - } else { - dact_x = 0.0f; - } - } else { - if constexpr ((ActOP == &silu) && (DActOP == &dsilu)) { - const float s = sigmoidf(x); - act_x = x * s; - dact_x = x * s * (1 - s) + s; - } else { - act_x = ActOP(x, p); - dact_x = DActOP(x, p); - } - } - float after_dact = dact_x * grad_elt * gate_elt; - float after_dgate = dgate_elt ? act_x * grad_elt : 0.0f; - - out_act_sh_curr[shmem_idx] = static_cast(scale * after_dact); - out_gate_sh_curr[shmem_idx] = static_cast(scale * after_dgate); - - amax = fmaxf(amax, fabsf(after_dact)); - amax = fmaxf(amax, fabsf(after_dgate)); - } else { - const float after_act = ActOP(act_elt, p) * gate_elt; - out_act_sh_curr[shmem_idx] = static_cast(scale * after_act); - amax = fmaxf(amax, fabsf(after_act)); - } - } - - // Wait for shared memory writes to be visible to TMA engine (cross-proxy fence) - 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 size_t chunk_it_offset_y = chunk_offset_Y + it * BUFFER_DIM_Y; - const size_t chunk_it_offset_x = chunk_offset_X; - - // dGeLU - ptx::cp_async_bulk_tensor_2d_shared_to_global(TMAP_output_act, chunk_it_offset_x, - chunk_it_offset_y, - reinterpret_cast(out_act_sh_curr)); - - if constexpr (IS_DGATED) { - // dGate - ptx::cp_async_bulk_tensor_2d_shared_to_global( - TMAP_output_gate, chunk_it_offset_x, chunk_it_offset_y, - reinterpret_cast(out_gate_sh_curr)); - } - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - - // Wait for TMA transfer to have finished reading shared memory. - ptx::cp_async_bulk_wait_group_read(); - } - } - ptx::cp_async_bulk_wait_group_read<0>(); - __syncthreads(); - - if (amax_ptr != nullptr) { - const int warp_id = threadIdx.x / THREADS_PER_WARP; - // Reduce the amax over the block - amax = reduce_max(amax, warp_id); - // Update the global amax - if (is_master_thread) { - atomicMaxFloat(amax_ptr, amax); - } - } - - // Update scale-inverse - if (is_master_thread && blockIdx.x == 0 && (scale_inv_ptr != nullptr)) { - reciprocal(scale_inv_ptr, scale); - } - - // Destroy the barriers. This invalidates the memory region of the barrier. - // If further computations were to take place in the kernel, this allows the - // memory location of the shared memory barrier to be reused. - if (is_master_thread) { -#pragma unroll - for (int it = 0; it < ITERATIONS; ++it) { - ptx::mbarrier_invalid(&mbar[it]); - } - } -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} - -namespace mxfp8_kernel { +namespace dispatch { +namespace mxfp8 { +namespace gated_kernel { constexpr size_t CHUNK_DIM_Y = 64; constexpr size_t CHUNK_DIM_X = 64; @@ -302,20 +49,21 @@ constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4) / 1; // 128 // 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; // 4 = 128 / 32 -template __global__ void __launch_bounds__(THREADS_PER_CHUNK) - cast_mxfp8_gated_kernel(const __grid_constant__ CUtensorMap tensor_map_grad, - const __grid_constant__ CUtensorMap tensor_map_input_act, - const __grid_constant__ CUtensorMap tensor_map_input_gate, - const __grid_constant__ CUtensorMap tensor_map_output_act_rowwise, - const __grid_constant__ CUtensorMap tensor_map_output_gate_rowwise, - const __grid_constant__ CUtensorMap tensor_map_output_act_colwise, - const __grid_constant__ CUtensorMap tensor_map_output_gate_colwise, - e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, - const size_t rows, const size_t cols, const size_t scale_stride_rowwise, - const size_t scale_stride_colwise, const ParamOP p) { + quantize_gated_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_grad, + const __grid_constant__ CUtensorMap tensor_map_input_act, + const __grid_constant__ CUtensorMap tensor_map_input_gate, + const __grid_constant__ CUtensorMap tensor_map_output_act_rowwise, + const __grid_constant__ CUtensorMap tensor_map_output_gate_rowwise, + const __grid_constant__ CUtensorMap tensor_map_output_act_colwise, + const __grid_constant__ CUtensorMap tensor_map_output_gate_colwise, + e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, + const size_t rows, const size_t cols, + const size_t scale_stride_rowwise, + const size_t scale_stride_colwise, const ParamOP p) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) using IType2 = typename ptx::FPx2; using OType2 = typename ptx::FPx2; @@ -385,14 +133,14 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); - const size_t grad_mem = (IS_DGATED ? buff_size_aligned_in : 0); + const size_t grad_mem = (IS_BWD ? buff_size_aligned_in : 0); const size_t in_act_mem = buff_size_aligned_in; const size_t in_gate_mem = buff_size_aligned_in; const size_t in_mem = in_act_mem + in_gate_mem; const size_t out_act_mem = buff_size_aligned_out; - const size_t out_gate_mem = (IS_DGATED ? buff_size_aligned_out : 0); + const size_t out_gate_mem = (IS_BWD ? buff_size_aligned_out : 0); const size_t out_mem = out_act_mem + out_gate_mem; // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned @@ -427,7 +175,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) int parity = 0; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { copy_2d_to_sharedx3(&in_grad_sh[0], &tensor_map_grad, block_offset_X, block_offset_Y, &in_act_sh[0], &tensor_map_input_act, block_offset_X, block_offset_Y, &in_gate_sh[0], &tensor_map_input_gate, block_offset_X, block_offset_Y, @@ -454,7 +202,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; const size_t global_offset_X = block_offset_X; const size_t next_buff_offset = next_buff * BUFF_DIM; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { copy_2d_to_sharedx3(&in_grad_sh[next_buff_offset], &tensor_map_grad, global_offset_X, global_offset_Y, &in_act_sh[next_buff_offset], &tensor_map_input_act, global_offset_X, global_offset_Y, &in_gate_sh[next_buff_offset], @@ -497,7 +245,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) 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; } - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { float grad_elt = static_cast(in_grad_sh[shmem_offset_colwise]); const float x = act_elt; float act_x; @@ -526,20 +274,20 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // Numerical truncation: Downcast to IType (BF16/FP16), then upcast it back to FP32 if constexpr (!std::is_same_v) { after_act_elt = static_cast(static_cast(after_act_elt)); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { after_gate_elt = static_cast(static_cast(after_gate_elt)); } } after_act_colwise[i] = after_act_elt; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { after_gate_colwise[i] = after_gate_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(after_act_elt); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { cached_gate_sh[shmem_offset_colwise] = static_cast(after_gate_elt); } } @@ -549,7 +297,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if (!out_of_bounds) { thread_amax_act = fmaxf(thread_amax_act, fabsf(after_act_elt)); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { thread_amax_gate = fmaxf(thread_amax_gate, fabsf(after_gate_elt)); } } @@ -578,7 +326,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // All threads read the reduced amax (ACT) thread_amax_act = subamax_colwise_buff[0][tid_X_colwise]; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { // Make sure the previous read of the ACT values has been completed, // so the data are not rewritten __syncthreads(); @@ -622,7 +370,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float block_scale_inverse_act = ptx::exp2f_rcp(biased_exponent_act); float block_scale_inverse_gate; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { const e8m0_t biased_exponent_gate = ptx::float_to_e8m0(thread_amax_gate * Quantized_Limits::max_norm_rcp); @@ -639,7 +387,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) for (int i = 0; i < SCALE_DIM_Y / COLWISE_WAVEFRONT_SIZE; ++i) { const size_t shmem_offset_elt = shmem_offset_base_colwise + i * COLWISE_WAVEFRONT_SIZE * BUFF_DIM_X; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { OType2 out_pair; ptx::floatx2 in_pair = {after_act_colwise[i], after_gate_colwise[i]}; const ptx::floatx2 block_scale_inverse_2x_pair = {block_scale_inverse_act, @@ -685,7 +433,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // Load cached elements in_cached_act[w].load_from(&cached_act_sh[shmem_offset_rowwise]); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { in_cached_gate[w].load_from(&cached_gate_sh[shmem_offset_rowwise]); } // Since TMA requirement for the data alignment is 16B (i.e. cols % 8 == 0, in case of BF16 elements) @@ -695,7 +443,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) #pragma unroll for (int e = 0; e < PACK_SIZE; ++e) { thread_amax_act = fmaxf(thread_amax_act, fabsf(in_cached_act[w].data.elt[e])); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { thread_amax_gate = fmaxf(thread_amax_gate, fabsf(in_cached_gate[w].data.elt[e])); } } @@ -705,7 +453,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const IType2 in_cached_2x_act = {in_cached_act[w].data.elt[e], in_cached_act[w].data.elt[e + 1]}; ptx::abs_max_2x(thread_amax_2x_act, thread_amax_2x_act, in_cached_2x_act); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { const IType2 in_cached_2x_gate = {in_cached_gate[w].data.elt[e], in_cached_gate[w].data.elt[e + 1]}; ptx::abs_max_2x(thread_amax_2x_gate, thread_amax_2x_gate, in_cached_2x_gate); @@ -717,7 +465,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if constexpr (!std::is_same_v) { thread_amax_act = static_cast( __hmax(__habs(thread_amax_2x_act.x), __habs(thread_amax_2x_act.y))); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { thread_amax_gate = static_cast( __hmax(__habs(thread_amax_2x_gate.x), __habs(thread_amax_2x_gate.y))); } @@ -735,7 +483,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) in_act.load_from(&in_act_sh[shmem_offset_rowwise]); in_gate.load_from(&in_gate_sh[shmem_offset_rowwise]); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { in_grad.load_from(&in_grad_sh[shmem_offset_rowwise]); } @@ -753,7 +501,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) 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; } - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { float grad_elt = static_cast(in_grad.data.elt[e]); const float x = act_elt; float act_x; @@ -786,7 +534,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // Numerical truncation: Downcast to IType (BF16/FP16), then upcast it back to FP32 if constexpr (!std::is_same_v) { after_act_elt = static_cast(static_cast(after_act_elt)); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { after_gate_elt = static_cast(static_cast(after_gate_elt)); } } @@ -796,7 +544,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const bool out_of_bounds = (row_out_of_bounds_rowwise || swizzled_col_out_of_bounds); if (!out_of_bounds) { thread_amax_act = fmaxf(thread_amax_act, fabsf(after_act_elt)); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { thread_amax_gate = fmaxf(thread_amax_gate, fabsf(after_gate_elt)); } } @@ -822,7 +570,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float block_scale_inverse_gate; ptx::floatx2 block_scale_inverse_2x_gate; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { const e8m0_t biased_exponent_gate = ptx::float_to_e8m0(thread_amax_gate * Quantized_Limits::max_norm_rcp); const size_t scale_idx_gate = scale_idx + gate_scale_idx_offset_rowwise; @@ -853,7 +601,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } ptx::mul_cvt_2x(out_act_pair, in_act, block_scale_inverse_2x_act); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { IType2 in_gate; OType2 &out_gate_pair = reinterpret_cast(out_gate.data.elt[e]); @@ -873,7 +621,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_idx; out_act.store_to(&out_act_rowwise_sh[shmem_offset_rowwise]); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { out_gate.store_to(&out_gate_rowwise_sh[shmem_offset_rowwise]); } } @@ -894,7 +642,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output_act_rowwise), global_offset_X, global_offset_Y, reinterpret_cast(&out_act_rowwise_sh[buff_offset])); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output_gate_rowwise), global_offset_X, global_offset_Y, reinterpret_cast(&out_gate_rowwise_sh[buff_offset])); @@ -904,7 +652,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output_act_colwise), global_offset_X, global_offset_Y, reinterpret_cast(&out_act_colwise_sh[buff_offset])); - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { ptx::cp_async_bulk_tensor_2d_shared_to_global( reinterpret_cast(&tensor_map_output_gate_colwise), global_offset_X, global_offset_Y, reinterpret_cast(&out_gate_colwise_sh[buff_offset])); @@ -920,94 +668,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) destroy_barriers(mbar, is_master_thread); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } -} // namespace mxfp8_kernel +} // namespace gated_kernel -template -void cast_fp8_gated(const Tensor &grad, const Tensor &gated_input, Tensor *output, ParamOP p, +void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *output, ParamOP &p, cudaStream_t stream) { - checkCuDriverContext(stream); - - if (output->has_data()) { - NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); - } - if (output->has_columnwise_data()) { - NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); - } - - NVTE_CHECK(!output->has_columnwise_data(), "Only rowwise cast supported in this function."); - const size_t rows = gated_input.flat_first_dim(); - const size_t cols = gated_input.flat_last_dim() / 2; - const size_t output_cols = (IS_DGATED ? 2 : 1) * cols; - - const size_t blocks_Y = DIVUP(rows, CHUNK_DIM_Y); - const size_t blocks_X = DIVUP(cols, CHUNK_DIM_X); - - float *const amax_ptr = reinterpret_cast(output->amax.dptr); - float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); - float *const scale_ptr = reinterpret_cast(output->scale.dptr); - - const dim3 block_dim(THREADS_PER_CHUNK); - const dim3 grid_dim(blocks_X, blocks_Y); - - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - gated_input.dtype(), IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output->dtype(), OType, - - alignas(64) CUtensorMap tensor_map_grad{}; - alignas(64) CUtensorMap tensor_map_input_act{}; - alignas(64) CUtensorMap tensor_map_input_gate{}; - alignas(64) CUtensorMap tensor_map_output_act{}; - alignas(64) CUtensorMap tensor_map_output_gate{}; - - if constexpr (IS_DGATED) { - create_2D_tensor_map(tensor_map_grad, grad.data, rows, cols, SHMEM_DIM_Y, SHMEM_DIM_X, - cols, 0, typeToNumBits(gated_input.dtype())); - } - - const uint32_t tensor_stride_elems = output_cols; - - create_2D_tensor_map(tensor_map_input_act, gated_input.data, rows, cols, SHMEM_DIM_Y, - SHMEM_DIM_X, cols * 2, 0, typeToNumBits(gated_input.dtype())); - create_2D_tensor_map(tensor_map_input_gate, gated_input.data, rows, cols, SHMEM_DIM_Y, - SHMEM_DIM_X, cols * 2, cols, typeToNumBits(gated_input.dtype())); - create_2D_tensor_map(tensor_map_output_act, output->data, rows, cols, SHMEM_DIM_Y, - SHMEM_DIM_X, tensor_stride_elems, 0, typeToNumBits(output->dtype())); - create_2D_tensor_map(tensor_map_output_gate, output->data, rows, cols, SHMEM_DIM_Y, - SHMEM_DIM_X, tensor_stride_elems, cols, - typeToNumBits(output->dtype())); - - const size_t buff_elems_total = BUFFERS_NUM * SHMEM_DIM_Y * SHMEM_DIM_X; - const size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); - const size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); - const size_t grad_mem = (IS_DGATED ? buff_size_aligned_in : 0); - const size_t in_act_mem = buff_size_aligned_in; - const size_t in_gate_mem = buff_size_aligned_in; - const size_t out_act_mem = buff_size_aligned_out; - const size_t out_gate_mem = buff_size_aligned_out; - - const size_t shmem_size = grad_mem + (in_act_mem + in_gate_mem) + - (out_act_mem + out_gate_mem) + TMA_SHMEM_ALIGNMENT; - - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - cast_fp8_gated_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - - cast_fp8_gated_kernel - <<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, tensor_map_output_act, - tensor_map_output_gate, amax_ptr, scale_inv_ptr, scale_ptr, rows, cols, p); - NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) - ); // NOLINT(*) -} - -template -void cast_mxfp8_gated(const Tensor &grad, const Tensor &gated_input, Tensor *output, ParamOP p, - cudaStream_t stream) { + using namespace gated_kernel; checkCuDriverContext(stream); const bool USE_ROWWISE_SCALING = output->has_data(); @@ -1031,17 +698,11 @@ void cast_mxfp8_gated(const Tensor &grad, const Tensor &gated_input, Tensor *out const size_t rows = gated_input.flat_first_dim(); const size_t cols = gated_input.flat_last_dim() / 2; - const size_t output_cols = (IS_DGATED ? 2 : 1) * cols; - - constexpr size_t BUFF_DIM_Y = mxfp8_kernel::BUFF_DIM_Y; - constexpr size_t BUFF_DIM_X = mxfp8_kernel::BUFF_DIM_X; - constexpr size_t BUFFS_NUM = mxfp8_kernel::BUFFS_NUM; + const size_t output_cols = (IS_BWD ? 2 : 1) * cols; - const size_t blocks_Y = DIVUP(rows, mxfp8_kernel::CHUNK_DIM_Y); - const size_t blocks_X = DIVUP(cols, mxfp8_kernel::CHUNK_DIM_X); + const size_t blocks_Y = DIVUP(rows, CHUNK_DIM_Y); + const size_t blocks_X = DIVUP(cols, CHUNK_DIM_X); - constexpr size_t THREADS_PER_CHUNK_COLWISE = mxfp8_kernel::THREADS_PER_CHUNK_COLWISE; - constexpr size_t THREADS_PER_CHUNK_NON_COLWISE = mxfp8_kernel::THREADS_PER_CHUNK_NON_COLWISE; const size_t THREADS_PER_CHUNK = (scaling_type == ScalingType::COLWISE) ? THREADS_PER_CHUNK_COLWISE : THREADS_PER_CHUNK_NON_COLWISE; @@ -1073,7 +734,7 @@ void cast_mxfp8_gated(const Tensor &grad, const Tensor &gated_input, Tensor *out constexpr size_t input_type_bit_size = TypeInfo::size; constexpr size_t output_type_bit_size = TypeInfo::size; - if constexpr (IS_DGATED) { + if constexpr (IS_BWD) { create_2D_tensor_map(tensor_map_grad, grad.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, input_type_bit_size); } @@ -1110,238 +771,68 @@ void cast_mxfp8_gated(const Tensor &grad, const Tensor &gated_input, Tensor *out const size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); - const size_t grad_mem = (IS_DGATED ? buff_size_aligned_in : 0); + const size_t grad_mem = (IS_BWD ? buff_size_aligned_in : 0); const size_t in_act_mem = buff_size_aligned_in; const size_t in_gate_mem = buff_size_aligned_in; const size_t in_mem = grad_mem + in_act_mem + in_gate_mem; const size_t out_act_mem = buff_size_aligned_out; - const size_t out_gate_mem = (IS_DGATED ? buff_size_aligned_out : 0); + const size_t out_gate_mem = (IS_BWD ? buff_size_aligned_out : 0); size_t out_mem = out_act_mem + out_gate_mem; + if (USE_ROWWISE_SCALING && USE_COLWISE_SCALING) { out_mem *= 2; } const size_t shmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; switch (scaling_type) { - case ScalingType::ROWWISE: + case ScalingType::ROWWISE: { + auto kernel = + quantize_gated_mxfp8_kernel; NVTE_CHECK_CUDA(cudaFuncSetAttribute( - mxfp8_kernel::cast_mxfp8_gated_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - - mxfp8_kernel::cast_mxfp8_gated_kernel - <<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, - tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, - tensor_map_output_act_colwise, tensor_map_output_gate_colwise, - scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise, p); - NVTE_CHECK_CUDA(cudaGetLastError()); + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); + + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, + tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, + tensor_map_output_act_colwise, tensor_map_output_gate_colwise, scales_rowwise_ptr, + scales_colwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise, p); break; - case ScalingType::COLWISE: + } + case ScalingType::COLWISE: { + auto kernel = + quantize_gated_mxfp8_kernel; NVTE_CHECK_CUDA(cudaFuncSetAttribute( - mxfp8_kernel::cast_mxfp8_gated_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - - mxfp8_kernel::cast_mxfp8_gated_kernel - <<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, - tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, - tensor_map_output_act_colwise, tensor_map_output_gate_colwise, - scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise, p); - NVTE_CHECK_CUDA(cudaGetLastError()); + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); + + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, + tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, + tensor_map_output_act_colwise, tensor_map_output_gate_colwise, scales_rowwise_ptr, + scales_colwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise, p); break; - case ScalingType::BIDIMENSIONAL: + } + case ScalingType::BIDIMENSIONAL: { + auto kernel = + quantize_gated_mxfp8_kernel; NVTE_CHECK_CUDA(cudaFuncSetAttribute( - mxfp8_kernel::cast_mxfp8_gated_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - mxfp8_kernel::cast_mxfp8_gated_kernel - <<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, - tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, - tensor_map_output_act_colwise, tensor_map_output_gate_colwise, - scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise, p); - NVTE_CHECK_CUDA(cudaGetLastError()); - break; - }); // NOLINT(*) - ); // NOLINT(*) -} - -template -void cast_gated(const Tensor &input, Tensor *output, ParamOP p, cudaStream_t stream) { - CheckInputTensor(input, "gated_act_input"); - CheckOutputTensor(*output, "gated_act_output"); - NVTE_CHECK(input.flat_last_dim() % 2 == 0, - "Wrong input shape. Expected (after flattening) last dimension to be even, ", "got [", - input.flat_first_dim(), ", ", input.flat_last_dim(), "]."); - NVTE_CHECK(output->flat_last_dim() == input.flat_last_dim() / 2, - "Wrong output shape. Expected (after flattening) [*, ", input.flat_last_dim() / 2, - "], got [", output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); - - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.dtype(), IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( - output->dtype(), OType, - - if (!is_fp8_dtype(output->data.dtype) || - is_delayed_tensor_scaling(output->scaling_mode)) { - constexpr int nvec = 32 / sizeof(IType); - GatedActivationKernelLauncher( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output->data.dptr), - reinterpret_cast(output->scale.dptr), - reinterpret_cast(output->amax.dptr), - reinterpret_cast(output->scale_inv.dptr), input.flat_first_dim(), - output->flat_last_dim(), p, stream); - } else { - NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); - }); // NOLINT(*) - ); // NOLINT(*) -} - -template -void cast_dgated(const Tensor &grad, const Tensor &input, Tensor *output, ParamOP p, - cudaStream_t stream) { - CheckInputTensor(grad, "dgated_act_grad"); - CheckInputTensor(input, "dgated_act_input"); - CheckOutputTensor(*output, "dgated_act_output"); - NVTE_CHECK(output->flat_first_dim() == grad.flat_first_dim(), - "Wrong output shape. Expected (after flattening) [", grad.flat_first_dim(), - ", *], got [", output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); - NVTE_CHECK(output->flat_last_dim() == grad.flat_last_dim() * 2, - "Wrong output shape. Expected (after flattening) [*, ", grad.flat_last_dim() * 2, - "], got [", output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); - NVTE_CHECK(input.data.shape == output->data.shape, - "Input and output shapes must match. Input shape: ", input.data.shape, - ", output shape: ", output->data.shape, "."); - - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.dtype(), IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( - output->dtype(), OType, - - if (!is_fp8_dtype(output->data.dtype) || - is_delayed_tensor_scaling(output->scaling_mode)) { - constexpr int nvec = 32 / sizeof(IType); - DGatedActivationKernelLauncher( - reinterpret_cast(grad.data.dptr), - reinterpret_cast(input.data.dptr), - reinterpret_cast(output->data.dptr), - reinterpret_cast(output->scale.dptr), - reinterpret_cast(output->amax.dptr), - reinterpret_cast(output->scale_inv.dptr), grad.flat_first_dim(), - grad.flat_last_dim(), p, stream); - } else { - NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); - }); // NOLINT(*) - ); // NOLINT(*) -} - -template -void quantize_gated(const Tensor &grad, const Tensor &gated_input, Tensor *output, ParamOP p, - cudaStream_t stream) { - constexpr bool allow_empty = false; - CheckInputTensor(gated_input, "gated_input"); - CheckOutputTensor(*output, "output", allow_empty); - - NVTE_CHECK(gated_input.flat_last_dim() % 2 == 0, "Number of columns must be even."); + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - const size_t rows = gated_input.flat_first_dim(); - const size_t cols = gated_input.flat_last_dim() / 2; - const size_t output_cols = (IS_DGATED ? 2 : 1) * cols; - - if constexpr (IS_DGATED) { - CheckInputTensor(grad, "grad"); - NVTE_CHECK(!is_fp8_dtype(grad.data.dtype), "Grad input must be in higher precision."); - NVTE_CHECK(grad.data.dtype == gated_input.data.dtype, "Types of both inputs must match."); - NVTE_CHECK(grad.flat_first_dim() == rows, "Wrong dimension of the grad input."); - NVTE_CHECK(grad.flat_last_dim() == cols, "Wrong dimension of the grad input."); - } - - NVTE_CHECK(output->has_data() || output->has_columnwise_data(), - "Either rowwise or columnwise output data need to be allocated."); - - bool is_fp8_rowwise_output = true; - bool is_fp8_colwise_output = true; - if (output->has_data()) { - is_fp8_rowwise_output = is_fp8_dtype(output->data.dtype); - NVTE_CHECK(output->flat_first_dim() == rows, "Wrong dimension of the output."); - NVTE_CHECK(output->flat_last_dim() == output_cols, "Wrong dimension of the output."); - } - if (output->has_columnwise_data()) { - is_fp8_colwise_output = is_fp8_dtype(output->columnwise_data.dtype); - NVTE_CHECK(output->flat_first_dim() == rows, "Wrong dimension of the output."); - NVTE_CHECK(output->flat_last_dim() == output_cols, "Wrong dimension of the output."); - } - - const bool use_tma_kernels = is_fp8_rowwise_output && is_fp8_colwise_output && cols % 32 == 0; - - if (is_delayed_tensor_scaling(output->scaling_mode)) { - if (use_tma_kernels) { - cast_fp8_gated(grad, gated_input, output, p, stream); - } else { - if constexpr (IS_DGATED) { - cast_dgated(grad, gated_input, output, p, stream); - } else { - cast_gated(gated_input, output, p, stream); - } - } - } else if (is_mxfp8_scaling(output->scaling_mode)) { - if (use_tma_kernels) { - cast_mxfp8_gated(grad, gated_input, output, p, stream); - } else { - NVTE_ERROR("Invalid input shape. Expected the last dimension to be divisible ", - "by 32, got input of shape ", gated_input.data.shape); - } - } else { - NVTE_ERROR("Not supported scaling mode"); - } -} -} // namespace gated_kernels - -namespace detail { - -template -void quantize_gated_helper(const NVTETensor grad, const NVTETensor gated_input, NVTETensor output, - ParamOP p, cudaStream_t stream) { - using namespace gated_kernels; - Tensor grad_empty_tensor; - const Tensor &grad_tensor = IS_DGATED ? *(convertNVTETensorCheck(grad)) : grad_empty_tensor; - const Tensor gated_input_tensor = *convertNVTETensorCheck(gated_input); - Tensor *output_tensor = convertNVTETensorCheck(output); - - if (is_supported_by_CC_100()) { - quantize_gated(grad_tensor, gated_input_tensor, - output_tensor, p, stream); - } else { - if (is_delayed_tensor_scaling(output_tensor->scaling_mode)) { - if constexpr (IS_DGATED) { - cast_dgated(grad_tensor, gated_input_tensor, output_tensor, p, - stream); - } else { - cast_gated(gated_input_tensor, output_tensor, p, stream); - } - } else { - // MX scaling - NVTE_ERROR("Not supported by the Arch < 10.0"); - } - } + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, + tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, + tensor_map_output_act_colwise, tensor_map_output_gate_colwise, scales_rowwise_ptr, + scales_colwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise, p); + break; + } + } NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) + ); // NOLINT(*) } -} // namespace detail +} // namespace mxfp8 +} // namespace dispatch } // namespace transformer_engine -#endif // TRANSFORMER_ENGINE_CAST_GATED_KERNELS_CUH_ +#endif // TRANSFORMER_ENGINE_GATED_MXFP8_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh new file mode 100644 index 0000000000..5505de6050 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -0,0 +1,722 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_mxfp8.cuh + * \brief CUDA kernels to quantize to MXFP8. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_MXFP8_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "../core/common.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { + +constexpr size_t SCALE_DIM_Y = 32; +constexpr size_t SCALE_DIM_X = 32; + +constexpr size_t BUFFS_NUM = 2; +constexpr size_t PACK_SIZE = 4; +constexpr size_t WAVES = SCALE_DIM_X / PACK_SIZE; + +// Number of 1-byte elements that span 32 banks (4-byte each) of shared memory +constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4) / 1; // 128 + +// 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; // 4 = 128 / 32 + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) + quantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_act_input, + const __grid_constant__ CUtensorMap tensor_map_output_rowwise, + const __grid_constant__ CUtensorMap tensor_map_output_colwise, + e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, + const float *noop, float *const dbias_workspace, float *const amax_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 COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; + constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; + + using IType2 = typename ptx::FPx2; + using OType2 = typename ptx::FPx2; + + if constexpr (NO_ACTIVATIONS) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + } + constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; + constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; + + constexpr size_t BUFF_DIM_Y = THREADS_Y; + constexpr size_t BUFF_DIM_X = CHUNK_DIM_X; + constexpr size_t BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; + static_assert(BUFF_DIM_Y == 32); + + constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; + static_assert(STAGES >= 1); + + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; + + const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; + const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; + const size_t scales_block_offset_X_rowwise = blockIdx.x * CHUNK_DIM_X / SCALE_DIM_X; + const size_t scales_block_offset_Y_colwise = blockIdx.y * CHUNK_DIM_Y / SCALE_DIM_Y; + const size_t scales_block_offset_X_colwise = blockIdx.x * CHUNK_DIM_X; + + const size_t tid_Y_rowwise = threadIdx.x / THREADS_X; + const size_t tid_X_rowwise = threadIdx.x % THREADS_X; + const size_t tid_Y_colwise = 0; + const size_t tid_X_colwise = threadIdx.x; + + const size_t thread_offset_Y_rowwise = tid_Y_rowwise; + const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM_X; + const size_t thread_offset_Y_colwise = tid_Y_colwise; + const size_t thread_offset_X_colwise = tid_X_colwise; + + const size_t row_base_rowwise = block_offset_Y + thread_offset_Y_rowwise; + const size_t row_base_colwise = block_offset_Y + thread_offset_Y_colwise; + const size_t col_base_colwise = block_offset_X + thread_offset_X_colwise; + + const bool col_out_of_bounds_colwise = (col_base_colwise >= cols); + + const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; + const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; + const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; + const size_t scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; + + const bool rowwise_scale_is_within_bounds = scales_offset_X_rowwise < 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_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 = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 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); + IType *act_in_sh = reinterpret_cast(dshmem + elt_input_mem); + + OType *out_rowwise_data_sh = reinterpret_cast(dshmem + in_mem); + OType *out_colwise_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise); + IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer + + constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + const bool is_master_thread = (threadIdx.x == 0); + + float partial_dbias_colwise = 0.0f; + float thread_dbias_rowwise[SCALE_DIM_X]; + if constexpr (IS_DBIAS) { +#pragma unroll + for (int j = 0; j < SCALE_DIM_X; ++j) { + thread_dbias_rowwise[j] = 0.0f; + } + } + + float block_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); + + int parity = 0; + + if constexpr (IS_DACT) { + copy_2d_to_sharedx2(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, &act_in_sh[0], + &tensor_map_act_input, block_offset_X, block_offset_Y, shmem_buff_size, + &mbar[0], is_master_thread); + } else { + 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 size_t buff = stage % BUFFS_NUM; + const size_t next_stage = stage + 1; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + + 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 size_t next_buff = next_stage % BUFFS_NUM; + const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_buff_offset = next_buff * BUFF_DIM; + if constexpr (IS_DACT) { + copy_2d_to_sharedx2(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, + global_offset_Y, &act_in_sh[next_buff_offset], &tensor_map_act_input, + global_offset_X, global_offset_Y, shmem_buff_size, &mbar[next_stage], + is_master_thread); + } else { + 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], parity); + + float thread_amax = 0.0f; + if constexpr (COLWISE_SCALING) { + const size_t shmem_offset_base_colwise = buff * BUFF_DIM + tid_X_colwise; + thread_amax = 0.0f; + float in_compute_colwise[BUFF_DIM_Y]; + IType in_colwise_IType[BUFF_DIM_Y]; + + // 1. Read/Compute elements. Find MXFP8-block AMAX + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + IType thread_amax_f16 = static_cast(0.0f); +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; + in_colwise_IType[i] = in_sh[shmem_offset_colwise]; + thread_amax_f16 = __hmax(thread_amax_f16, __habs(in_colwise_IType[i])); + } + thread_amax = static_cast(thread_amax_f16); + } else { +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; + + float elt = static_cast(in_sh[shmem_offset_colwise]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); + elt *= OP(act_in_elt, {}); + } + if constexpr (IS_DBIAS) { + partial_dbias_colwise += 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) { + thread_amax = fmaxf(thread_amax, fabsf(elt)); + } + } else { + // If no activation, elt is 0 so we can safely do this + thread_amax = fmaxf(thread_amax, fabsf(elt)); + } + in_compute_colwise[i] = elt; + } + } + + // 2. Compute E8M0 scaling factor + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + + const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; + const size_t global_scales_offset_X = scales_offset_X_colwise; + const size_t scale_idx = + global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + scales_colwise[scale_idx] = biased_exponent; + + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; + +// 3. Scale elements +#pragma unroll + for (int i = 0; i < SCALE_DIM_Y; ++i) { + float in; + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + in = static_cast(in_colwise_IType[i]); + } else { + in = in_compute_colwise[i]; + } + const float scaled_out = in * block_scale_inverse; + + const size_t shmem_offset_elt = shmem_offset_base_colwise + i * BUFF_DIM_X; + out_colwise_data_sh[shmem_offset_elt] = static_cast(scaled_out); + } + } + + if constexpr (ROWWISE_SCALING) { + const size_t shmem_offset_base_rowwise = + buff * BUFF_DIM + thread_offset_Y_rowwise * BUFF_DIM_X; + thread_amax = 0.0f; + float in_compute_rowwise[SCALE_DIM_X]; + Vec in_cached[WAVES]; + + // used as an IType container for BF16/FP16 --> MXFP8 CAST ONLY + Vec in_IType[WAVES]; + + // 1. Read/Compute elements. Find MXFP8-block AMAX + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + 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]); + } + } + thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; + + const bool row_out_of_bounds_rowwise = (row_base_rowwise + stage_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) { + thread_amax = fmaxf(thread_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) { + thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; + + Vec in; + Vec act_in; + + in.load_from(&in_sh[shmem_offset_rowwise]); + if constexpr (IS_DACT) { + act_in.load_from(&act_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 (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + float act_in_elt = static_cast(act_in.data.elt[e]); + elt *= OP(act_in_elt, {}); + } + + // If DBIAS was computed in the 1st pass (COLWISE) then no need to compute it again + if constexpr (IS_DBIAS && (!COLWISE_SCALING)) { + thread_dbias_rowwise[j] += 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 + stage_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) { + thread_amax = fmaxf(thread_amax, fabsf(elt)); + } + } else { + // If no activation, elt is 0 so we can safely do this + thread_amax = fmaxf(thread_amax, fabsf(elt)); + } + in_compute_rowwise[j] = elt; + } + } + } + + // 2. Compute E8M0 scaling factor + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; + const int stage_scales_offset_X = scales_offset_X_rowwise; + const int 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; + } + + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; + + // 3. Scale elements +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + Vec out; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 2; ++e) { + IType2 in; + OType2 &out_pair = reinterpret_cast(out.data.elt[e]); + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + in = in_IType[w].data.elt[e]; + } else if constexpr (IS_CACHED_ACT_OP) { + in.x = in_cached[w].data.elt[2 * e]; + in.y = in_cached[w].data.elt[2 * e + 1]; + } else { + const int j = w * PACK_SIZE + 2 * e; + in.x = in_compute_rowwise[j]; + in.y = in_compute_rowwise[j + 1]; + } + ptx::mul_cvt_2x(out_pair, in, block_scale_inverse_2x); + } + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_idx; + out.store_to(&out_rowwise_data_sh[shmem_offset_rowwise]); + } + } + + __builtin_assume(block_amax >= 0); + __builtin_assume(thread_amax >= 0); + block_amax = fmaxf(block_amax, thread_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 = buff * BUFF_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])); + } + 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])); + } + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + } + } + + parity ^= 1; + + if constexpr (IS_DBIAS) { + float thread_partial_dbias = 0.0f; + if constexpr (COLWISE_SCALING) { + thread_partial_dbias = partial_dbias_colwise; + } else { + // Reusing dshmem (in_sh) as dbias buffer [HEIGHT x WIDTH] + // HEIGHT = THREADS_Y + // WIDTH = THREADS_X * (SCALE_DIM_X + 1) + // Added extra 1-element padding per thread_X to reduce bank conflicts + float *partial_dbias_rowwise = reinterpret_cast(dshmem); + + constexpr int DBIAS_BUFF_WIDTH = THREADS_X * (SCALE_DIM_X + 1); + + const int shmem_thread_offset = + tid_Y_rowwise * DBIAS_BUFF_WIDTH + tid_X_rowwise * (SCALE_DIM_X + 1); +#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_group_offset = shmem_thread_offset + swizzled_group_idx; +#pragma unroll + for (int e = 0; e < PACK_SIZE; ++e) { + const int j = w * PACK_SIZE + e; + const int shmem_elt_idx = swizzled_group_offset + e; + partial_dbias_rowwise[shmem_elt_idx] = thread_dbias_rowwise[j]; + } + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < THREADS_Y; ++i) { + // Add extra element offset per MXFP8 scaling block [1x32] + const int scaling_block = threadIdx.x / SCALE_DIM_X; + thread_partial_dbias += + partial_dbias_rowwise[i * DBIAS_BUFF_WIDTH + threadIdx.x + scaling_block]; + } + } + const int dbias_stride = cols; + const int dbias_offset_Y = blockIdx.y; + const int dbias_offset_X = blockIdx.x * CHUNK_DIM_X + threadIdx.x; + const int dbias_idx = dbias_offset_Y * dbias_stride + dbias_offset_X; + const bool col_out_of_bounds_dbias = (dbias_offset_X >= cols); + if (!col_out_of_bounds_dbias) { + dbias_workspace[dbias_idx] = thread_partial_dbias; + } + } + + if (amax_ptr != nullptr) { + const int warp_id = threadIdx.x / THREADS_PER_WARP; + // Reduce the amax over the block + block_amax = reduce_max(block_amax, warp_id); + } + + if (is_master_thread && amax_ptr != nullptr) { + atomicMaxFloat(amax_ptr, block_amax); + } + + destroy_barriers(mbar, is_master_thread); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +} // namespace quantize_kernel + +template +void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // TODO (ksivamani) + Tensor *output, Tensor *dbias, Tensor *workspace, cudaStream_t stream) { + using namespace quantize_kernel; + checkCuDriverContext(stream); + + bool use_rowwise_scaling = output->has_data(); + bool use_colwise_scaling = output->has_columnwise_data(); + NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); + + if (use_rowwise_scaling) { + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + } + 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 bool CAST_DBIAS_ONLY = IS_DBIAS && (!IS_DACT) && (!IS_ACT); + + constexpr size_t CHUNK_DIM_Y = CAST_DBIAS_ONLY ? 128 : 64; + constexpr size_t CHUNK_DIM_X = CAST_DBIAS_ONLY ? 128 : 64; + constexpr size_t THREADS_PER_CHUNK = CAST_DBIAS_ONLY ? 128 : 64; + + constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; + constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; + constexpr size_t BUFF_DIM_Y = THREADS_Y; + 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 = use_rowwise_scaling ? output->scale_inv.shape[1] : 1; + const size_t scale_stride_colwise = + use_colwise_scaling ? output->columnwise_scale_inv.shape[1] : 1; + + e8m0_t *const scales_rowwise_ptr = + use_rowwise_scaling ? reinterpret_cast(output->scale_inv.dptr) : nullptr; + e8m0_t *const scales_colwise_ptr = + use_colwise_scaling ? reinterpret_cast(output->columnwise_scale_inv.dptr) : nullptr; + const size_t dbias_rows = blocks_Y; + const size_t dbias_cols = cols; + + ScalingType scaling_type; + if (use_rowwise_scaling && (!use_colwise_scaling)) { + scaling_type = ScalingType::ROWWISE; + } else if ((!use_rowwise_scaling) && use_colwise_scaling) { + scaling_type = ScalingType::COLWISE; + } else if (use_rowwise_scaling && use_colwise_scaling) { + scaling_type = ScalingType::BIDIMENSIONAL; + } + + 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(workspace != nullptr, "Workspace must be a tensor."); + + if (workspace->data.dptr == nullptr) { + workspace->data.shape = {dbias_rows, dbias_cols}; + workspace->data.dtype = DType::kFloat32; + return; + } + } + + float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + float *const amax_ptr = reinterpret_cast(output->amax.dptr); + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_act_input{}; + alignas(64) CUtensorMap tensor_map_output_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_colwise{}; + + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; + + create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, + cols, 0, input_type_bit_size); + + if constexpr (IS_DACT) { + create_2D_tensor_map(tensor_map_act_input, act_input->data, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols, 0, input_type_bit_size); + } + + if (use_rowwise_scaling) { + create_2D_tensor_map(tensor_map_output_rowwise, output->data, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols, 0, output_type_bit_size); + } + + 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, output_type_bit_size); + } + + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; + constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); + const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); + const size_t out_mem = out_rowwise_mem + out_colwise_mem; + + const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; + + switch (scaling_type) { + case ScalingType::ROWWISE: { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + 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: { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + 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: { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + 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; + } + } + + if constexpr (IS_DBIAS) { + common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); + }); // NOLINT(*) + ); // NOLINT(*) +} + +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh new file mode 100644 index 0000000000..cff8464903 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -0,0 +1,112 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file core_nvfp4.cuh + * \brief Core functions used in NVFP4. + */ + +#ifndef TRANSFORMER_ENGINE_CORE_NVFP4_CUH_ +#define TRANSFORMER_ENGINE_CORE_NVFP4_CUH_ + +#include +#include +#include + +#include + +#include "../../common.h" +#include "../../util/curanddx.hpp" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" + +#if FP4_TYPE_SUPPORTED +#include +#endif // FP4_TYPE_SUPPORTED + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { + +using nvfp4_scale_t = fp8e4m3; + +namespace quantization_and_transposition_SF { +#if FP4_TYPE_SUPPORTED +// Used in transpose variant +// Compute per-block E4M3 encoding/decoding scaling factor +__device__ __forceinline__ nvfp4_scale_t compute_decoding_scaling_factor(const float block_amax, + const float S_enc) { + // constexpr float rcp_6f = 1.0f / 6.0f; + // const float S_dec_b = block_amax * rcp_6f; + // const nvfp4_scale_t S_dec_b_fp8 = static_cast(S_dec_b * S_enc); + // return S_dec_b_fp8; + // NOTE: Divide by 6.0f is not elegant and not efficient. + // However, this is part of the emulation code to ensure exact match. + using namespace detail; + constexpr float fp4_max = TypeExtrema::max; // 6.0f; + const float S_dec_b = block_amax / fp4_max * S_enc; + return static_cast(fminf(S_dec_b, TypeExtrema::max)); +} +#endif // FP4_TYPE_SUPPORTED +} // namespace quantization_and_transposition_SF + +namespace quantization_SF { +#if FP4_TYPE_SUPPORTED +// Used in non-transpose variant +// Compute per-block E4M3 encoding/decoding scaling factor +__device__ __forceinline__ fp8e4m3 compute_decoding_scaling_factor(const float block_amax, + const float S_enc) { + constexpr float rcp_6f = 1.0f / 6.0f; + // const float S_dec_b = block_amax * rcp_6f; + // const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); + // return S_dec_b_fp8; + return static_cast(block_amax * rcp_6f * S_enc); +} +#endif // FP4_TYPE_SUPPORTED +} // namespace quantization_SF + +namespace core { + +#if FP4_TYPE_SUPPORTED +using namespace ptx; + +// Compute the global encode scale factor for a given global amax +__device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const float global_amax) { + using namespace detail; + constexpr float fp8_max = TypeExtrema::max; // 448.0f; + 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 + global_encode_scale = fminf(global_encode_scale, TypeExtrema::max); + // If global amax is 0 or infinity, return 1 + if (global_amax == 0.0f || global_encode_scale == 0.0f) { + return 1.0f; + } + return global_encode_scale; +} + +__device__ __forceinline__ uint32_t +get_rbits(transformer_engine::curanddx::detail::philox4x32_native_state<10> &rng, + // philox4x32_native_state<10>: 10 rounds of philox4_32 + uint4 &random_uint4, int &rnd_idx) { + if (rnd_idx == 4) { + rnd_idx = 0; + random_uint4 = rng.generate4(); + } + // Treat uint4 as an array of 4x uint32_t elements for indexing + const uint32_t *const rbits_arr = reinterpret_cast(&random_uint4); + const uint32_t rbits = rbits_arr[rnd_idx++]; + return rbits; +} + +#endif // FP4_TYPE_SUPPORTED + +} // namespace core +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_CORE_NVFP4_CUH_ diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh new file mode 100644 index 0000000000..bf7b535be4 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -0,0 +1,111 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file dequantize_nvfp4.cuh + * \brief CUDA kernels to dequantize from NVFP4. + */ + +#ifndef TRANSFORMER_ENGINE_DEQUANTIZE_NVFP4_CUH_ +#define TRANSFORMER_ENGINE_DEQUANTIZE_NVFP4_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" + +#if FP4_TYPE_SUPPORTED +#include +#endif // FP4_TYPE_SUPPORTED + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { +namespace dequantize_kernel { +#if FP4_TYPE_SUPPORTED +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 thread_idx = blockIdx.x * blockDim.x + threadIdx.x; + const size_t x = thread_idx % M; + const size_t y = thread_idx / M; + + union fp4vec { + uint64_t vec; + fp4e2m1x4 small_vec[4]; + }; + using OVec = Vec; + const uint64_t *const input_vectorized = reinterpret_cast(input); + OVec *output_vec = reinterpret_cast(output); + + const size_t my_index = x + y * M; + const size_t 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]; + fp8e4m3 scale = scales[my_scale_index]; + float amax = *tensor_amax; + constexpr float factor_inv = 1.0 / (6.0 * 448.0); + float final_scale = static_cast(scale) * amax * factor_inv; +#pragma unroll + for (int i = 0; i < 4; i++) { + float4 current = static_cast(value.small_vec[i]); + OVec out; + out.data.elt[0] = static_cast(current.x * final_scale); + out.data.elt[1] = static_cast(current.y * final_scale); + out.data.elt[2] = static_cast(current.z * final_scale); + out.data.elt[3] = static_cast(current.w * final_scale); + output_vec[my_output_index + i] = out; + } +} +#endif // FP4_TYPE_SUPPORTED +} // namespace dequantize_kernel + +inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace dequantize_kernel; + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output"); + NVTE_CHECK(input.data.dtype == DType::kFloat4E2M1, "Input must have FP4 type."); + 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."); + + constexpr int FP4_BLOCK_SIZE = 16; + const size_t N = input.flat_first_dim(); + const size_t M = input.flat_last_dim(); + + 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, "."); + + const size_t Mread = M / FP4_BLOCK_SIZE; + const size_t total = N * Mread; + const size_t threads = 512; + const size_t blocks = DIVUP(total, threads); + + 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(*) + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); +#endif // FP4_TYPE_SUPPORTED +} +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_DEQUANTIZE_NVFP4_CUH_ diff --git a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh new file mode 100644 index 0000000000..83ad8fd40b --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh @@ -0,0 +1,688 @@ +/************************************************************************* + * Copyright (c) 2022-2025, 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 buff_size_nvfp4_scales = + CHUNK_DIM_Y * (CHUNK_DIM_X / SCALE_DIM_X) * sizeof(fp8e4m3); + constexpr size_t buff_size_mxfp8_scales = + (CHUNK_DIM_Y / SCALE_DIM_Y) * CHUNK_DIM_X * sizeof(fp8e8m0); + + 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); + constexpr size_t out_mem_rowwise_scales = (ROWWISE_SCALING ? buff_size_nvfp4_scales : 0); + constexpr size_t out_mem_colwise_scales = (COLWISE_SCALING ? buff_size_mxfp8_scales : 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); + e8m0_t *out_colwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + 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"); + + 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_ diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh new file mode 100644 index 0000000000..7322bf2655 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -0,0 +1,1287 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_transpose_nvfp4.cuh + * \brief CUDA kernels to cast to NVFP4 and transpose. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_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_transpose_kernel { + +using namespace quantization_and_transposition_SF; +using namespace core; +using namespace ptx; + +#if FP4_TYPE_SUPPORTED + +constexpr size_t SCALE_DIM = 16; // NVFP4 block (x16 elts) + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_NUM = 128; + +constexpr size_t SCALES_PER_CHUNK_Y = CHUNK_DIM_Y / SCALE_DIM; +constexpr size_t SCALES_PER_CHUNK_X = CHUNK_DIM_X / SCALE_DIM; + +constexpr size_t SCALES_PER_THREAD = 2 * (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM / THREADS_NUM; + +// Each call generates 4x uint32_t random numbers +constexpr size_t RNG_GENS_PER_THREAD = SCALES_PER_THREAD / 4; + +constexpr size_t TILE_DIM_Y = 32; +constexpr size_t TILE_DIM_X = 128; + +// SHould this be SCALE_DIM or BLOCK_DIM? Both are 16, should work for both 1D and 2D +constexpr size_t SCALES_PER_TILE_Y = TILE_DIM_Y / SCALE_DIM; +constexpr size_t SCALES_PER_TILE_X = TILE_DIM_X / SCALE_DIM; // 128 / 16 = 8 + +constexpr size_t TILES_Y = CHUNK_DIM_Y / TILE_DIM_Y; +constexpr size_t TILES_X = CHUNK_DIM_X / TILE_DIM_X; +constexpr size_t STAGES = TILES_Y * TILES_X; + +constexpr size_t BUFFS_NUM = 2; +constexpr size_t BUFF_DIM_Y = TILE_DIM_Y; +constexpr size_t BUFF_DIM_X = TILE_DIM_X; +constexpr size_t BUFF_SIZE = BUFF_DIM_Y * BUFF_DIM_X; +constexpr size_t BUFF_SIZE_TOTAL = BUFF_SIZE * BUFFS_NUM; + +// Input buffer (BF16) +constexpr size_t BUFF_IN_DIM_Y = BUFF_DIM_Y; +constexpr size_t BUFF_IN_DIM_X = BUFF_DIM_X; +constexpr size_t BUFF_IN_SIZE = BUFF_IN_DIM_Y * BUFF_IN_DIM_X; + +// Output buffer (NVFP4) +constexpr size_t BUFF_OUT_DIM_Y = BUFF_DIM_Y; +constexpr size_t BUFF_OUT_DIM_X = (BUFF_DIM_X * 4) / 8; +constexpr size_t BUFF_OUT_SIZE = BUFF_OUT_DIM_Y * BUFF_OUT_DIM_X; + +// Output transpose buffer (NVFP4) +constexpr size_t BUFF_OUT_T_DIM_Y = BUFF_DIM_X; +constexpr size_t BUFF_OUT_T_DIM_X = (BUFF_DIM_Y * 4) / 8; +constexpr size_t BUFF_OUT_T_SIZE = BUFF_OUT_T_DIM_Y * BUFF_OUT_T_DIM_X; + +// Manual swizzling parameters to reduce SHMEM bank conflicts +constexpr size_t PACK_SIZE = 8; +constexpr size_t WAVES = SCALE_DIM / PACK_SIZE; + +constexpr size_t SCALING_FACTORS_PER_TILE_X = TILE_DIM_X / SCALE_DIM; +constexpr size_t THREADS_X_ROWWISE = SCALING_FACTORS_PER_TILE_X; // 128 / 16 = 8 +constexpr size_t THREADS_Y_ROWWISE = THREADS_NUM / THREADS_X_ROWWISE; // 128 / 8 = 16 + +constexpr size_t ITERATIONS_NORMAL = BUFF_DIM_Y / THREADS_Y_ROWWISE; // 32/ 16 = 2 +constexpr size_t ITERATIONS_TRANSPOSE = BUFF_IN_DIM_Y / SCALE_DIM; +constexpr size_t BUFF_OUT_IT_OFFSET = BUFF_OUT_T_DIM_X / ITERATIONS_TRANSPOSE; + +static_assert(BUFF_DIM_Y >= SCALE_DIM && + "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"); + +// 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; // 8 = 128 / 16 + +template +__global__ void __launch_bounds__(THREADS_NUM) + quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_output, + const __grid_constant__ CUtensorMap tensor_map_output_t, + nvfp4_scale_t *const scales_ptr, + nvfp4_scale_t *const scales_t_ptr, const float *noop, + const float *const amax_rowwise_ptr, + const float *const amax_colwise_ptr, const size_t rows, + const size_t cols, const size_t scale_stride, + const size_t scale_stride_t, const size_t *rng_state) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + 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; + } + } + + const size_t rng_sequence = + threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; + // Index of the random number. It increments each time when used and resets to 0 if reaches 4x + int rnd_idx = 0; + + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS; + + const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; + + const size_t block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; + const size_t block_offset_X_t = blockIdx.y * CHUNK_DIM_Y; + + const size_t chunk_rows = rows - block_offset_Y; + + const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; + const size_t scales_block_offset_X_rowwise = blockIdx.x * SCALES_PER_CHUNK_X; + const size_t scales_block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; + const size_t scales_block_offset_X_t = blockIdx.y * SCALES_PER_CHUNK_Y; + + const size_t tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; + const size_t tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; + const size_t tid_X_colwise = threadIdx.x; + const size_t tid_Y_t = tid_X_colwise; + // const size_t tid_X_t = 0; + + const size_t thread_offset_Y_rowwise = tid_Y_rowwise; + const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM; + const size_t thread_offset_X_colwise = tid_X_colwise; + + const size_t row_base_rowwise = block_offset_Y + thread_offset_Y_rowwise; + const size_t row_base_colwise = block_offset_Y; + const size_t col_base_colwise = block_offset_X + thread_offset_X_colwise; + + const bool col_out_of_bounds_colwise = (col_base_colwise >= cols); + + const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; + const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; + const size_t scales_offset_Y_t = scales_block_offset_Y_t + tid_Y_t; + const size_t scales_offset_X_t = scales_block_offset_X_t; + + const size_t SFs_per_row = cols / SCALE_DIM; + + const bool rowwise_scale_is_within_bounds_X = scales_offset_X_rowwise < SFs_per_row; + const bool colwise_scale_is_within_bounds_Y = scales_offset_Y_t < 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 = + DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); + + constexpr size_t in_mem = buff_size_aligned_in; + + constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; + constexpr size_t out_mem_colwise_data = buff_size_aligned_out; + constexpr size_t out_mem_rowwise_scales = 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_data_sh = reinterpret_cast(dshmem + in_mem); + fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + + nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer + + constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + const bool is_master_thread = (threadIdx.x == 0); + + // Compute a global encoding/decoding scaling factors for all S_dec_b + const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) + ? 1.0f + : compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); + // NOTE: This is to match with how emulation code was written. + const float S_dec_rowwise = 1.0 / S_enc_rowwise; + + const float S_enc_colwise = (amax_colwise_ptr == nullptr) + ? S_enc_rowwise + : compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); + const float S_dec_colwise = 1.0 / S_enc_colwise; + + 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 (size_t stage = 0; stage < STAGES; ++stage) { + const size_t buff = stage % BUFFS_NUM; + const size_t next_stage = stage + 1; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + + const size_t buff_offset_in = buff * BUFF_IN_SIZE; + const size_t buff_offset_out = buff * BUFF_OUT_SIZE; + const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE; + + 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 size_t next_buff = next_stage % BUFFS_NUM; + const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_buff_offset = next_buff * BUFF_IN_SIZE; + + 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; + + // COLWISE scaling + if constexpr (RETURN_TRANSPOSE) { +#pragma unroll + for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { + const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; + const size_t in_thread_offset_X = thread_offset_X_colwise; + + const size_t out_t_thread_offset_Y = thread_offset_X_colwise; + const size_t out_t_thread_offset_X = 0 + it * BUFF_OUT_IT_OFFSET; + + const size_t shmem_offset_base_colwise_in = + buff_offset_in + in_thread_offset_Y * BUFF_IN_DIM_X + in_thread_offset_X; + const size_t shmem_offset_base_colwise_out_t = + buff_offset_out_t + out_t_thread_offset_Y * BUFF_OUT_T_DIM_X + out_t_thread_offset_X; + + block_amax = 0.0f; + float in_compute_colwise[SCALE_DIM]; + IType in_colwise_IType[SCALE_DIM]; + // 1. Read/Compute elements. Find NVFP4-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; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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 E4M3 scaling factor + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax, S_enc_colwise); + + // Store scaling factors through SHMEM + const size_t scale_idx_sh = + tid_Y_t * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + it; + out_colwise_scales_sh[scale_idx_sh] = S_dec_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_colwise), float_max); // S_enc_b_fp8 + const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + // 3. Scale elements + fp4e2m1x4 regs[SCALE_DIM / 4]; + +#pragma unroll + for (int e = 0; e < SCALE_DIM / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); + regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); + const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); + regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } + } + + const int group = thread_lane / 16; + uint32_t val[2]; + uint32_t *regs_4x = reinterpret_cast(regs); + + // Helps reducing bank conflicts + switch (group) { + case 0: + val[0] = regs_4x[0]; + val[1] = regs_4x[1]; + break; + case 1: + val[0] = regs_4x[1]; + val[1] = regs_4x[0]; + + break; + } + uint32_t *out_t_data_sh_as_uint32_t = + reinterpret_cast(&out_t_data_sh[shmem_offset_base_colwise_out_t]); + out_t_data_sh_as_uint32_t[group] = val[0]; // idx1 = (group + 0) % 2; + out_t_data_sh_as_uint32_t[(group + 1) & 1] = val[1]; // idx2 = (group + 1) % 2; + } + } + + // ROWWISE scaling + { + const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; +#pragma unroll + for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { + const size_t it_thread_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; + + const size_t shmem_offset_base_rowwise_in = + buff_offset_in + it_thread_offset_Y_rowwise * BUFF_IN_DIM_X; + const size_t shmem_offset_base_rowwise_out = + buff_offset_out + it_thread_offset_Y_rowwise * BUFF_OUT_DIM_X; + + const size_t it_offset_Y = stage_offset_Y + it * THREADS_Y_ROWWISE; + + block_amax = 0.0f; + float in_compute_rowwise[SCALE_DIM]; + 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t 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 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 = 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; + 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 +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + Vec out; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + IType2 in01; + IType2 in23; + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_IType[w].data.elt[2 * e]); + out.data.elt[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else if constexpr (IS_CACHED_ACT_OP) { + const uint64_t elts = *reinterpret_cast(&in_cached[w].data.elt[4 * e]); + out.data.elt[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const int j = w * PACK_SIZE + 4 * e; + const float2 in01 = make_float2(in_compute_rowwise[j], in_compute_rowwise[j + 1]); + const float2 in23 = make_float2(in_compute_rowwise[j + 2], in_compute_rowwise[j + 3]); + out.data.elt[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } + } + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; + out.store_to(&out_data_sh[shmem_offset_rowwise]); + } + } + } + + __builtin_assume(thread_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 size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + + const size_t global_offset_Y_t = block_offset_Y_t; + const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, + reinterpret_cast(&out_data_sh[buff_offset_out])); + + if constexpr (RETURN_TRANSPOSE) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_t), global_offset_X_t, + global_offset_Y_t, reinterpret_cast(&out_t_data_sh[buff_offset_out_t])); + } + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + } + } // end of stages + + // Vectorized store scaling factors through SHMEM + if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { + using ScalesVec = Vec; + const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; + ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); + const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; + const size_t count = // number of scales in Y dimension of this chunk + (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); + nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; + constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { + // Fast path: vectorized store when destination is properly aligned + scales_vec.store_to(dst); + } else { + // Safe path: element-wise store for tails or unaligned destinations + scales_vec.store_to_elts(dst, 0, count); + } + } + + destroy_barriers(mbar, is_master_thread); +#else + NVTE_DEVICE_ERROR("sm_100 or higher is required."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template +__global__ void __launch_bounds__(THREADS_NUM) + quantize_transpose_nvfp4_2D_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_output, + const __grid_constant__ CUtensorMap tensor_map_output_t, + nvfp4_scale_t *const scales_ptr, + nvfp4_scale_t *const scales_t_ptr, const float *noop, + const float *const amax_rowwise_ptr, + const float *const amax_colwise_ptr, const size_t rows, + const size_t cols, const size_t scale_stride, + const size_t scale_stride_t, const size_t *rng_state) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + 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; + } + } + const size_t rng_sequence = + threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; + int rnd_idx = + 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x + + // NEW: 2D Block-based scaling constants + constexpr size_t BLOCK_DIM = 16; + constexpr size_t BLOCKS_PER_TILE_Y = TILE_DIM_Y / BLOCK_DIM; // 32/16 = 2 + constexpr size_t BLOCKS_PER_TILE_X = TILE_DIM_X / BLOCK_DIM; // 128/16 = 8 + constexpr size_t ITERATIONS_BLOCK = 2; // iterations to calculate 2d block amaxes of 1 tile + constexpr size_t BLOCKS_PER_WARP = BLOCKS_PER_TILE_X / (THREADS_NUM / 32); // 8 / (128/32) = 2 + + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS; + + const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; + + const size_t block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; + const size_t block_offset_X_t = blockIdx.y * CHUNK_DIM_Y; + + const size_t chunk_rows = rows - block_offset_Y; + + const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; + const size_t scales_block_offset_X_rowwise = blockIdx.x * SCALES_PER_CHUNK_X; + const size_t scales_block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; + const size_t scales_block_offset_X_t = blockIdx.y * SCALES_PER_CHUNK_Y; + + const size_t tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; + const size_t tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; + const size_t tid_X_colwise = threadIdx.x; + const size_t tid_Y_t = tid_X_colwise; + + const size_t thread_offset_Y_rowwise = tid_Y_rowwise; + const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM; + const size_t thread_offset_X_colwise = tid_X_colwise; + + const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; + const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; + const size_t scales_offset_Y_t = scales_block_offset_Y_t + tid_Y_t; + const size_t scales_offset_X_t = scales_block_offset_X_t; + + const size_t SFs_per_row = cols / SCALE_DIM; + + const bool rowwise_scale_is_within_bounds_X = scales_offset_X_rowwise < SFs_per_row; + const bool colwise_scale_is_within_bounds_Y = scales_offset_Y_t < 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 = + DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); + + constexpr size_t in_mem = buff_size_aligned_in; + + constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; + constexpr size_t out_mem_colwise_data = buff_size_aligned_out; + constexpr size_t out_mem_rowwise_scales = 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_data_sh = reinterpret_cast(dshmem + in_mem); + fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + + nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer + + constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + const bool is_master_thread = (threadIdx.x == 0); + + // Compute a global encoding/decoding scaling factors for all S_dec_b + const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) + ? 1.0f + : compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); + // NOTE: This is to match with how emulation code was written. + const float S_dec_rowwise = 1.0 / S_enc_rowwise; + + const float S_enc_colwise = (amax_colwise_ptr == nullptr) + ? S_enc_rowwise + : compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); + const float S_dec_colwise = 1.0 / S_enc_colwise; + + const size_t warp_id = threadIdx.x / 32; + const size_t lane_id = threadIdx.x % 32; + float thread_amax = 0.0f; + const size_t block_in_warp = lane_id / BLOCKS_PER_WARP; + +// 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]; + + __shared__ __align__(16) float block_amax_matrix[BLOCKS_PER_TILE_Y][BLOCKS_PER_TILE_X + 1]; + + // Helper function for warp reduction + auto warp_reduce_amax = [](float thread_amax, int block_in_warp) -> float { +#pragma unroll + for (int delta = 8; delta >= 1; delta /= 2) { + float other_amax = __shfl_xor_sync(0xffffffff, thread_amax, delta); + thread_amax = fmaxf(thread_amax, other_amax); + } + return thread_amax; + }; + + 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 (size_t stage = 0; stage < STAGES; ++stage) { + const size_t buff = stage % BUFFS_NUM; + const size_t next_stage = stage + 1; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + + const size_t buff_offset_in = buff * BUFF_IN_SIZE; + const size_t buff_offset_out = buff * BUFF_OUT_SIZE; + const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE; + + 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 size_t next_buff = next_stage % BUFFS_NUM; + const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_buff_offset = next_buff * BUFF_IN_SIZE; + + 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; + +#pragma unroll + for (size_t block_iter = 0; block_iter < ITERATIONS_BLOCK; ++block_iter) { + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; + const size_t block_in_tile_y = block_iter; + const size_t block_in_tile_x = threadIdx.x / BLOCK_DIM; + + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + for (int elem = 0; elem < BLOCK_DIM; elem += 2) { + const size_t elem_0_row = block_iter * BLOCK_DIM + elem; + const size_t elem_1_row = elem_0_row + 1; + const size_t elem_0_col = warp_id * BLOCKS_PER_WARP * BLOCK_DIM + lane_id; + const size_t elem_1_col = elem_0_col; + + const size_t shmem_offset_0 = buff_offset_in + elem_0_row * BUFF_IN_DIM_X + elem_0_col; + const size_t shmem_offset_1 = buff_offset_in + elem_1_row * BUFF_IN_DIM_X + elem_1_col; + + IType2 val_2x; + val_2x.x = in_sh[shmem_offset_0]; + val_2x.y = in_sh[shmem_offset_1]; + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, val_2x); + } + + thread_amax = + static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); + } else { + for (int elem = 0; elem < BLOCK_DIM; ++elem) { + const size_t elem_row = block_iter * BLOCK_DIM + elem; + const size_t elem_col = warp_id * BLOCKS_PER_WARP * BLOCK_DIM + lane_id; + + // Bounds checking + const bool row_out_of_bounds = (block_offset_Y + stage_offset_Y + elem_row >= rows); + const bool col_out_of_bounds = (block_offset_X + elem_col >= cols); + if (!row_out_of_bounds && !col_out_of_bounds) { + const size_t shmem_offset = buff_offset_in + elem_row * BUFF_IN_DIM_X + elem_col; + float elt = static_cast(in_sh[shmem_offset]); + + if constexpr (COMPUTE_ACTIVATIONS) { + elt = OP(elt, {}); + } + if constexpr (!std::is_same_v) { + elt = static_cast(static_cast(elt)); + } + // Cache computed activations + if constexpr (IS_CACHED_ACT_OP) { + cached_act_sh[shmem_offset] = static_cast(elt); + } + + thread_amax = fmaxf(thread_amax, fabsf(elt)); + } + } + } + // Warp reduction to get block amax + block_amax = warp_reduce_amax(thread_amax, block_in_warp); + + if (lane_id == 0 || lane_id == 16) { + block_amax_matrix[block_in_tile_y][block_in_tile_x] = block_amax; + } + } + + // sync thread to ensure block_amax_matrix is done storing + __syncthreads(); + + // COLWISE scaling + if constexpr (RETURN_TRANSPOSE) { +#pragma unroll + for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { + const size_t block_in_tile_y = it; + const size_t block_in_tile_x = threadIdx.x / BLOCK_DIM; + + const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; + const size_t in_thread_offset_X = thread_offset_X_colwise; + + const size_t out_t_thread_offset_Y = thread_offset_X_colwise; + const size_t out_t_thread_offset_X = 0 + it * BUFF_OUT_IT_OFFSET; + + const size_t shmem_offset_base_colwise_in = + buff_offset_in + in_thread_offset_Y * BUFF_IN_DIM_X + in_thread_offset_X; + const size_t shmem_offset_base_colwise_out_t = + buff_offset_out_t + out_t_thread_offset_Y * BUFF_OUT_T_DIM_X + out_t_thread_offset_X; + + block_amax = block_amax_matrix[block_in_tile_y][block_in_tile_x]; + float in_compute_colwise[SCALE_DIM]; + IType in_colwise_IType[SCALE_DIM]; + // 3. Scale elements + + // Load data in + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + i * BUFF_IN_DIM_X; + in_colwise_IType[i] = in_sh[shmem_offset_colwise]; + } + } else { + for (int i = 0; i < SCALE_DIM; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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); + } + + in_compute_colwise[i] = elt; + } + } + + // 2. Compute E4M3 scaling factor + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax, S_enc_colwise); + + // // Store scaling factors through SHMEM + const size_t scale_idx_sh = + tid_Y_t * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + it; + out_colwise_scales_sh[scale_idx_sh] = S_dec_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_colwise), float_max); // S_enc_b_fp8 + const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + fp4e2m1x4 regs[SCALE_DIM / 4]; +#pragma unroll + for (int e = 0; e < SCALE_DIM / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); + regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); + const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); + regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } + } + + const int group = thread_lane / 16; + uint32_t val[2]; + uint32_t *regs_4x = reinterpret_cast(regs); + + // Helps reducing bank conflicts + switch (group) { + case 0: + val[0] = regs_4x[0]; + val[1] = regs_4x[1]; + break; + case 1: + val[0] = regs_4x[1]; + val[1] = regs_4x[0]; + break; + } + uint32_t *out_t_data_sh_as_uint32_t = + reinterpret_cast(&out_t_data_sh[shmem_offset_base_colwise_out_t]); + out_t_data_sh_as_uint32_t[group] = val[0]; // idx1 = (group + 0) % 2; + out_t_data_sh_as_uint32_t[(group + 1) & 1] = val[1]; // idx2 = (group + 1) % 2; + } + } + + // ROWWISE scaling + { + const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; +#pragma unroll + for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { + const size_t block_in_tile_y = it; + const size_t block_in_tile_x = tid_X_rowwise; + const size_t it_thread_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; + + const size_t shmem_offset_base_rowwise_in = + buff_offset_in + it_thread_offset_Y_rowwise * BUFF_IN_DIM_X; + const size_t shmem_offset_base_rowwise_out = + buff_offset_out + it_thread_offset_Y_rowwise * BUFF_OUT_DIM_X; + + block_amax = block_amax_matrix[block_in_tile_y][block_in_tile_x]; + float in_compute_rowwise[SCALE_DIM]; + 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; + // Load elements + in_IType[w].load_from(&in_sh[shmem_offset_rowwise]); + } + } else if constexpr (IS_CACHED_ACT_OP) { + // ensures that all writes to cache made in the section above are visible to all threads + __syncthreads(); +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; + + // Load cached elements + in_cached[w].load_from(&cached_act_sh[shmem_offset_rowwise]); + } + } else { +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t 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)); + } + in_compute_rowwise[j] = elt; + } + } + } + + // 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 = 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; + 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 +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + Vec out; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + IType2 in01; + IType2 in23; + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_IType[w].data.elt[2 * e]); + out.data.elt[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else if constexpr (IS_CACHED_ACT_OP) { + const uint64_t elts = *reinterpret_cast(&in_cached[w].data.elt[4 * e]); + out.data.elt[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const int j = w * PACK_SIZE + 4 * e; + const float2 in01 = make_float2(in_compute_rowwise[j], in_compute_rowwise[j + 1]); + const float2 in23 = make_float2(in_compute_rowwise[j + 2], in_compute_rowwise[j + 3]); + out.data.elt[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } + } + + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; + out.store_to(&out_data_sh[shmem_offset_rowwise]); + } + } + } + + __builtin_assume(thread_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 size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + + const size_t global_offset_Y_t = block_offset_Y_t; + const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, + reinterpret_cast(&out_data_sh[buff_offset_out])); + + if constexpr (RETURN_TRANSPOSE) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_t), global_offset_X_t, + global_offset_Y_t, reinterpret_cast(&out_t_data_sh[buff_offset_out_t])); + } + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + } + } // end of stages + + // Vectorized store scaling factors through SHMEM + if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { + using ScalesVec = Vec; + const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; + ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); + const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; + const size_t count = // number of scales in Y dimension of this chunk + (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); + nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; + constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { + // Fast path: vectorized store when destination is properly aligned + scales_vec.store_to(dst); + } else { + // Safe path: element-wise store for tails or unaligned destinations + scales_vec.store_to_elts(dst, 0, count); + } + } + + destroy_barriers(mbar, is_master_thread); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +#endif // FP4_TYPE_SUPPORTED +} // namespace quantize_transpose_kernel + +template +void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace quantize_transpose_kernel; + using namespace ptx; + bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; + + // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to + // return the transposed data. + // TODO(Frank): Is there a better way to do this? + bool return_transpose = output->has_columnwise_data(); + + constexpr bool COMPUTE_ACTIVATIONS = false; + using ParamOP = Empty; + constexpr float (*OP)(float, const ParamOP &) = nullptr; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output", false); + + NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); + 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"); + if (return_transpose) { + NVTE_CHECK(output->has_columnwise_data(), "NVFP4 transposed output tensor must be allocated."); + NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), + "Transposed output must have FP4 type."); + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, + "Transposed scaling tensor must be allocated"); + } + + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + + NVTE_CHECK(rows % 32 == 0, + "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA + NVTE_CHECK(cols % 32 == 0, + "Number of tensor cols must be a multiple of 32"); // 16B alignment for TMA + + 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_NUM; + + const size_t scale_stride = output->scale_inv.shape[1]; + const size_t scale_stride_transpose = + return_transpose ? output->columnwise_scale_inv.shape[1] : 0; + + nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); + nvfp4_scale_t *const scales_transpose_ptr = + reinterpret_cast(output->columnwise_scale_inv.dptr); + + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); + const float *const amax_colwise_ptr = + reinterpret_cast(output->columnwise_amax.dptr); + + const NVTETensor rng_state_tensor = (quant_config != nullptr) ? quant_config->rng_state : nullptr; + const size_t *rng_state = nullptr; + if (rng_state_tensor != nullptr) { + 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}, + "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); + } + + using IType = bf16; + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_output{}; + alignas(64) CUtensorMap tensor_map_output_transpose{}; + + 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, output->data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, + 4); + if (return_transpose) { + create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, + BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); + } + 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 = + DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); + + constexpr size_t in_mem = buff_size_aligned_in; + + constexpr size_t out_data_mem = buff_size_aligned_out; + constexpr size_t out_data_transpose_mem = buff_size_aligned_out; + constexpr size_t out_scales_transpose_mem = buff_size_scales; + + constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; + + constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, + + 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; + } + + 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 +} + +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_CUH_ diff --git a/transformer_engine/common/util/cast.cu b/transformer_engine/common/util/cast.cu deleted file mode 100644 index 107965d342..0000000000 --- a/transformer_engine/common/util/cast.cu +++ /dev/null @@ -1,201 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -#include "../common.h" -#include "../transpose/cast_transpose.h" -#include "../util/multi_stream.h" -#include "../util/vectorized_pointwise.h" -#include "../utils.cuh" -#include "cast_kernels.cuh" -#include "dequantize_kernels.cuh" -#include "math.h" -#include "ptx.cuh" -#include "transformer_engine/activation.h" -#include "transformer_engine/transpose.h" - -void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = false; - constexpr bool IS_ACT = false; - constexpr NVTETensor dbias = nullptr; - constexpr NVTETensor workspace = nullptr; - constexpr const NVTETensor grad = nullptr; - - detail::quantize_helper(input, grad, output, dbias, - workspace, nullptr, stream); -} - -void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor noop, - cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_noop); - using namespace transformer_engine; - - // Create config with noop tensor - QuantizationConfig quant_config; - quant_config.noop_tensor = noop; - - nvte_quantize_v2(input, output, reinterpret_cast(&quant_config), stream); -} - -void nvte_quantize_v2(const NVTETensor input, NVTETensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_v2); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = false; - constexpr bool IS_ACT = false; - constexpr NVTETensor dbias = nullptr; - constexpr NVTETensor workspace = nullptr; - constexpr const NVTETensor grad = nullptr; - - detail::quantize_helper( - input, grad, output, dbias, workspace, 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 bool IS_ACT = false; - constexpr const NVTETensor activation_input = nullptr; - - detail::quantize_helper( - activation_input, 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; - constexpr bool IS_ACT = false; - - detail::quantize_helper>( - activation_input, 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; - constexpr bool IS_ACT = false; - - detail::quantize_helper>( - activation_input, 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; - constexpr bool IS_ACT = false; - - detail::quantize_helper>( - activation_input, 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; - constexpr bool IS_ACT = false; - - detail::quantize_helper>( - activation_input, 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; - constexpr bool IS_ACT = false; - - detail::quantize_helper>( - activation_input, 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; - detail::dequantize_helper(*convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream); -} - -void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, - const NVTEQuantizationConfig quant_configs, - const size_t num_tensors, cudaStream_t stream) { - NVTE_API_CALL(nvte_multi_tensor_quantize); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = false; - constexpr bool IS_ACT = false; - constexpr NVTETensor dbias = nullptr; - constexpr NVTETensor workspace = nullptr; - constexpr const NVTETensor grad = nullptr; - - const size_t num_streams = nvte_get_num_compute_streams(); - - int num_stream_used = std::min(num_streams, num_tensors); - // wait for current stream to finish - NVTE_CHECK_CUDA(cudaEventRecord(detail::get_compute_stream_event(0), stream)); - for (int s = 0; s < num_stream_used; s++) { - NVTE_CHECK_CUDA( - cudaStreamWaitEvent(detail::get_compute_stream(s), detail::get_compute_stream_event(0))); - } - - for (int i = 0; i < num_tensors; i++) { - detail::quantize_helper( - inputs[i], grad, outputs[i], dbias, workspace, nullptr, - detail::get_compute_stream(i % num_streams)); - } - - // record events on compute streams - for (int s = 0; s < num_stream_used; s++) { - NVTE_CHECK_CUDA( - cudaEventRecord(detail::get_compute_stream_event(s), detail::get_compute_stream(s))); - } - // wait for all compute streams to finish - for (int s = 0; s < num_stream_used; s++) { - NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream, detail::get_compute_stream_event(s))); - } -} diff --git a/transformer_engine/common/util/cast_kernels.cuh b/transformer_engine/common/util/cast_kernels.cuh deleted file mode 100644 index b0498602b5..0000000000 --- a/transformer_engine/common/util/cast_kernels.cuh +++ /dev/null @@ -1,2188 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file cast_kernels.cuh - * \brief CUDA kernels to cast to/from FP8/MXFP8. - */ - -#ifndef TRANSFORMER_ENGINE_CAST_KERNELS_CUH_ -#define TRANSFORMER_ENGINE_CAST_KERNELS_CUH_ - -#include -#include -#include -#include - -#include - -#include "../common.h" -#include "../transpose/cast_transpose.h" -#include "../util/vectorized_pointwise.h" -#include "../utils.cuh" -#include "math.h" -#include "nvfp4_transpose.cuh" -#include "ptx.cuh" -#include "transformer_engine/transformer_engine.h" - -namespace transformer_engine { - -namespace mxfp8_kernel { - -constexpr size_t SCALE_DIM_Y = 32; -constexpr size_t SCALE_DIM_X = 32; - -constexpr size_t BUFFS_NUM = 2; -constexpr size_t PACK_SIZE = 4; -constexpr size_t WAVES = SCALE_DIM_X / PACK_SIZE; - -// Number of 1-byte elements that span 32 banks (4-byte each) of shared memory -constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4) / 1; // 128 - -// 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; // 4 = 128 / 32 - -template -__global__ void __launch_bounds__(THREADS_PER_CHUNK) - cast_mxfp8_2D_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ CUtensorMap tensor_map_act_input, - const __grid_constant__ CUtensorMap tensor_map_output_rowwise, - const __grid_constant__ CUtensorMap tensor_map_output_colwise, - e8m0_t *const scales_rowwise, e8m0_t *const scales_colwise, - const float *noop, float *const dbias_workspace, float *const amax_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 COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; - constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; - - using IType2 = typename ptx::FPx2; - using OType2 = typename ptx::FPx2; - - if constexpr (NO_ACTIVATIONS) { - if (noop != nullptr && noop[0] == 1.0f) { - return; - } - } - constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; - constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; - - constexpr size_t BUFF_DIM_Y = THREADS_Y; - constexpr size_t BUFF_DIM_X = CHUNK_DIM_X; - constexpr size_t BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; - static_assert(BUFF_DIM_Y == 32); - - constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; - static_assert(STAGES >= 1); - - constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; - - const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; - const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; - const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; - const size_t scales_block_offset_X_rowwise = blockIdx.x * CHUNK_DIM_X / SCALE_DIM_X; - const size_t scales_block_offset_Y_colwise = blockIdx.y * CHUNK_DIM_Y / SCALE_DIM_Y; - const size_t scales_block_offset_X_colwise = blockIdx.x * CHUNK_DIM_X; - - const size_t tid_Y_rowwise = threadIdx.x / THREADS_X; - const size_t tid_X_rowwise = threadIdx.x % THREADS_X; - const size_t tid_Y_colwise = 0; - const size_t tid_X_colwise = threadIdx.x; - - const size_t thread_offset_Y_rowwise = tid_Y_rowwise; - const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM_X; - const size_t thread_offset_Y_colwise = tid_Y_colwise; - const size_t thread_offset_X_colwise = tid_X_colwise; - - const size_t row_base_rowwise = block_offset_Y + thread_offset_Y_rowwise; - const size_t row_base_colwise = block_offset_Y + thread_offset_Y_colwise; - const size_t col_base_colwise = block_offset_X + thread_offset_X_colwise; - - const bool col_out_of_bounds_colwise = (col_base_colwise >= cols); - - const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; - const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; - const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; - const size_t scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; - - const bool rowwise_scale_is_within_bounds = scales_offset_X_rowwise < 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_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 = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); - - constexpr size_t elt_input_mem = buff_size_aligned_in; - constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); - constexpr size_t in_mem = elt_input_mem + act_input_mem; - - constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 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); - IType *act_in_sh = reinterpret_cast(dshmem + elt_input_mem); - - OType *out_rowwise_data_sh = reinterpret_cast(dshmem + in_mem); - OType *out_colwise_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise); - IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer - - constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; - - const bool is_master_thread = (threadIdx.x == 0); - - float partial_dbias_colwise = 0.0f; - float thread_dbias_rowwise[SCALE_DIM_X]; - if constexpr (IS_DBIAS) { -#pragma unroll - for (int j = 0; j < SCALE_DIM_X; ++j) { - thread_dbias_rowwise[j] = 0.0f; - } - } - - float block_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); - - int parity = 0; - - if constexpr (IS_DACT) { - copy_2d_to_sharedx2(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, &act_in_sh[0], - &tensor_map_act_input, block_offset_X, block_offset_Y, shmem_buff_size, - &mbar[0], is_master_thread); - } else { - 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 size_t buff = stage % BUFFS_NUM; - const size_t next_stage = stage + 1; - const size_t stage_offset_Y = stage * BUFF_DIM_Y; - - 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 size_t next_buff = next_stage % BUFFS_NUM; - const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; - const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; - const size_t global_offset_X = block_offset_X; - const size_t next_buff_offset = next_buff * BUFF_DIM; - if constexpr (IS_DACT) { - copy_2d_to_sharedx2(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, - global_offset_Y, &act_in_sh[next_buff_offset], &tensor_map_act_input, - global_offset_X, global_offset_Y, shmem_buff_size, &mbar[next_stage], - is_master_thread); - } else { - 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], parity); - - float thread_amax = 0.0f; - if constexpr (COLWISE_SCALING) { - const size_t shmem_offset_base_colwise = buff * BUFF_DIM + tid_X_colwise; - thread_amax = 0.0f; - float in_compute_colwise[BUFF_DIM_Y]; - IType in_colwise_IType[BUFF_DIM_Y]; - - // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - IType thread_amax_f16 = static_cast(0.0f); -#pragma unroll - for (int i = 0; i < BUFF_DIM_Y; ++i) { - const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; - in_colwise_IType[i] = in_sh[shmem_offset_colwise]; - thread_amax_f16 = __hmax(thread_amax_f16, __habs(in_colwise_IType[i])); - } - thread_amax = static_cast(thread_amax_f16); - } else { -#pragma unroll - for (int i = 0; i < BUFF_DIM_Y; ++i) { - const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; - - float elt = static_cast(in_sh[shmem_offset_colwise]); - if constexpr (IS_ACT) { - elt = OP(elt, {}); - } - if constexpr (IS_DACT) { - float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); - elt *= OP(act_in_elt, {}); - } - if constexpr (IS_DBIAS) { - partial_dbias_colwise += 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) { - thread_amax = fmaxf(thread_amax, fabsf(elt)); - } - } else { - // If no activation, elt is 0 so we can safely do this - thread_amax = fmaxf(thread_amax, fabsf(elt)); - } - in_compute_colwise[i] = elt; - } - } - - // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - - const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; - const size_t global_scales_offset_X = scales_offset_X_colwise; - const size_t scale_idx = - global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; - scales_colwise[scale_idx] = biased_exponent; - - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); - const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; - -// 3. Scale elements -#pragma unroll - for (int i = 0; i < SCALE_DIM_Y; ++i) { - float in; - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - in = static_cast(in_colwise_IType[i]); - } else { - in = in_compute_colwise[i]; - } - const float scaled_out = in * block_scale_inverse; - - const size_t shmem_offset_elt = shmem_offset_base_colwise + i * BUFF_DIM_X; - out_colwise_data_sh[shmem_offset_elt] = static_cast(scaled_out); - } - } - - if constexpr (ROWWISE_SCALING) { - const size_t shmem_offset_base_rowwise = - buff * BUFF_DIM + thread_offset_Y_rowwise * BUFF_DIM_X; - thread_amax = 0.0f; - float in_compute_rowwise[SCALE_DIM_X]; - Vec in_cached[WAVES]; - - // used as an IType container for BF16/FP16 --> MXFP8 CAST ONLY - Vec in_IType[WAVES]; - - // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + 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]); - } - } - thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; - - const bool row_out_of_bounds_rowwise = (row_base_rowwise + stage_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) { - thread_amax = fmaxf(thread_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) { - thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; - - Vec in; - Vec act_in; - - in.load_from(&in_sh[shmem_offset_rowwise]); - if constexpr (IS_DACT) { - act_in.load_from(&act_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 (IS_ACT) { - elt = OP(elt, {}); - } - if constexpr (IS_DACT) { - float act_in_elt = static_cast(act_in.data.elt[e]); - elt *= OP(act_in_elt, {}); - } - - // If DBIAS was computed in the 1st pass (COLWISE) then no need to compute it again - if constexpr (IS_DBIAS && (!COLWISE_SCALING)) { - thread_dbias_rowwise[j] += 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 + stage_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) { - thread_amax = fmaxf(thread_amax, fabsf(elt)); - } - } else { - // If no activation, elt is 0 so we can safely do this - thread_amax = fmaxf(thread_amax, fabsf(elt)); - } - in_compute_rowwise[j] = elt; - } - } - } - - // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; - const int stage_scales_offset_X = scales_offset_X_rowwise; - const int 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; - } - - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); - const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; - - // 3. Scale elements -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - Vec out; -#pragma unroll - for (int e = 0; e < PACK_SIZE / 2; ++e) { - IType2 in; - OType2 &out_pair = reinterpret_cast(out.data.elt[e]); - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - in = in_IType[w].data.elt[e]; - } else if constexpr (IS_CACHED_ACT_OP) { - in.x = in_cached[w].data.elt[2 * e]; - in.y = in_cached[w].data.elt[2 * e + 1]; - } else { - const int j = w * PACK_SIZE + 2 * e; - in.x = in_compute_rowwise[j]; - in.y = in_compute_rowwise[j + 1]; - } - ptx::mul_cvt_2x(out_pair, in, block_scale_inverse_2x); - } - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_idx; - out.store_to(&out_rowwise_data_sh[shmem_offset_rowwise]); - } - } - - __builtin_assume(block_amax >= 0); - __builtin_assume(thread_amax >= 0); - block_amax = fmaxf(block_amax, thread_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 = buff * BUFF_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])); - } - 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])); - } - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - } - } - - parity ^= 1; - - if constexpr (IS_DBIAS) { - float thread_partial_dbias = 0.0f; - if constexpr (COLWISE_SCALING) { - thread_partial_dbias = partial_dbias_colwise; - } else { - // Reusing dshmem (in_sh) as dbias buffer [HEIGHT x WIDTH] - // HEIGHT = THREADS_Y - // WIDTH = THREADS_X * (SCALE_DIM_X + 1) - // Added extra 1-element padding per thread_X to reduce bank conflicts - float *partial_dbias_rowwise = reinterpret_cast(dshmem); - - constexpr int DBIAS_BUFF_WIDTH = THREADS_X * (SCALE_DIM_X + 1); - - const int shmem_thread_offset = - tid_Y_rowwise * DBIAS_BUFF_WIDTH + tid_X_rowwise * (SCALE_DIM_X + 1); -#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_group_offset = shmem_thread_offset + swizzled_group_idx; -#pragma unroll - for (int e = 0; e < PACK_SIZE; ++e) { - const int j = w * PACK_SIZE + e; - const int shmem_elt_idx = swizzled_group_offset + e; - partial_dbias_rowwise[shmem_elt_idx] = thread_dbias_rowwise[j]; - } - } - __syncthreads(); -#pragma unroll - for (int i = 0; i < THREADS_Y; ++i) { - // Add extra element offset per MXFP8 scaling block [1x32] - const int scaling_block = threadIdx.x / SCALE_DIM_X; - thread_partial_dbias += - partial_dbias_rowwise[i * DBIAS_BUFF_WIDTH + threadIdx.x + scaling_block]; - } - } - const int dbias_stride = cols; - const int dbias_offset_Y = blockIdx.y; - const int dbias_offset_X = blockIdx.x * CHUNK_DIM_X + threadIdx.x; - const int dbias_idx = dbias_offset_Y * dbias_stride + dbias_offset_X; - const bool col_out_of_bounds_dbias = (dbias_offset_X >= cols); - if (!col_out_of_bounds_dbias) { - dbias_workspace[dbias_idx] = thread_partial_dbias; - } - } - - if (amax_ptr != nullptr) { - const int warp_id = threadIdx.x / THREADS_PER_WARP; - // Reduce the amax over the block - block_amax = reduce_max(block_amax, warp_id); - } - - if (is_master_thread && amax_ptr != nullptr) { - atomicMaxFloat(amax_ptr, block_amax); - } - - destroy_barriers(mbar, is_master_thread); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} -} // namespace mxfp8_kernel - -namespace nvfp4_kernel { - -using namespace ptx; - -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 - -// Compute per-block E4M3 encoding/decoding scaling factor -__device__ __forceinline__ fp8e4m3 compute_decoding_scaling_factor(const float block_amax, - const float S_enc) { - constexpr float rcp_6f = 1.0f / 6.0f; - // const float S_dec_b = block_amax * rcp_6f; - // const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); - // return S_dec_b_fp8; - return static_cast(block_amax * rcp_6f * S_enc); -} - -#define DIRECT_SCALING_FACTORS_STORE 1 - -template -__global__ void __launch_bounds__(THREADS_PER_CHUNK) - cast_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 buff_size_nvfp4_scales = - CHUNK_DIM_Y * (CHUNK_DIM_X / SCALE_DIM_X) * sizeof(fp8e4m3); - constexpr size_t buff_size_mxfp8_scales = - (CHUNK_DIM_Y / SCALE_DIM_Y) * CHUNK_DIM_X * sizeof(fp8e8m0); - - 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); - constexpr size_t out_mem_rowwise_scales = (ROWWISE_SCALING ? buff_size_nvfp4_scales : 0); - constexpr size_t out_mem_colwise_scales = (COLWISE_SCALING ? buff_size_mxfp8_scales : 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); - e8m0_t *out_colwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); - 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 nvfp4_kernel - -constexpr size_t FP8_CHUNK_DIM_Y = 128; -constexpr size_t FP8_CHUNK_DIM_X = 128; -constexpr size_t FP8_THREADS_PER_CHUNK = 128; -constexpr size_t FP8_BUFFERS_NUM = 2; -constexpr size_t FP8_PREFETCH_BUFFERS_NUM = 1; -static_assert(FP8_PREFETCH_BUFFERS_NUM < FP8_BUFFERS_NUM); - -constexpr size_t FP8_BUFFER_DIM_Y = 16; -constexpr size_t FP8_BUFFER_DIM_X = FP8_CHUNK_DIM_X; // 128 -constexpr size_t FP8_SHMEM_DIM_Y = FP8_BUFFER_DIM_Y; // 16 -constexpr size_t FP8_SHMEM_DIM_X = FP8_BUFFER_DIM_X; // 128 - -constexpr size_t FP8_BUFF_STAGES_NUM = FP8_BUFFER_DIM_Y; // 16 -constexpr size_t FP8_ITERATIONS = FP8_CHUNK_DIM_Y / FP8_BUFFER_DIM_Y; // 8 = 128 / 16 -static_assert(FP8_ITERATIONS >= FP8_PREFETCH_BUFFERS_NUM); - -template -__global__ void __launch_bounds__(FP8_THREADS_PER_CHUNK) - cast_fp8_2D_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ CUtensorMap tensor_map_act_input, - const __grid_constant__ CUtensorMap tensor_map_output, - float *const dbias_workspace, float *const amax_ptr, - float *const scale_inv_ptr, const float *const scale_ptr, const size_t rows, - const size_t cols) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - - const size_t block_offset_Y = blockIdx.y * FP8_CHUNK_DIM_Y; - const size_t block_offset_X = blockIdx.x * FP8_CHUNK_DIM_X; - - const size_t tid_Y = threadIdx.x / FP8_THREADS_PER_CHUNK; - const size_t tid_X = threadIdx.x % FP8_THREADS_PER_CHUNK; - - const size_t thread_offset_Y = tid_Y; - const size_t thread_offset_X = tid_X; - - const size_t dbias_offset_Y = blockIdx.y + tid_Y; - const size_t my_column = blockIdx.x * FP8_CHUNK_DIM_X + thread_offset_X; - const bool col_out_of_bounds = my_column >= cols; - const size_t dbias_stride = cols; - - float partial_dbias = 0.f; - - float amax = 0; - const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; - - // The destination shared memory buffer of a bulk tensor operation should be 128-byte aligned - __shared__ alignas(TMA_SHMEM_ALIGNMENT) - IType in_sh[FP8_BUFFERS_NUM][FP8_SHMEM_DIM_Y][FP8_SHMEM_DIM_X]; - __shared__ alignas(TMA_SHMEM_ALIGNMENT) - IType act_in_sh[FP8_BUFFERS_NUM][FP8_SHMEM_DIM_Y][FP8_SHMEM_DIM_X]; - __shared__ alignas(TMA_SHMEM_ALIGNMENT) - OType out_sh[FP8_BUFFERS_NUM][FP8_SHMEM_DIM_Y][FP8_SHMEM_DIM_X]; - - constexpr size_t shmem_buff_size = sizeof(in_sh) / FP8_BUFFERS_NUM; - - const bool is_master_thread = (threadIdx.x == 0); - -// 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[FP8_ITERATIONS]; - - initialize_barriers(mbar, is_master_thread); - - int parity = 0; - - const size_t chunk_offset_Y = block_offset_Y; - const size_t chunk_offset_X = block_offset_X; - -#pragma unroll - for (int prefetch_buff = 0; prefetch_buff < FP8_PREFETCH_BUFFERS_NUM; ++prefetch_buff) { - const size_t chunk_stage_offset_Y = chunk_offset_Y + prefetch_buff * FP8_BUFFER_DIM_Y; - const size_t chunk_stage_offset_X = chunk_offset_X; - if constexpr (IS_DACT) { - copy_2d_to_sharedx2(&in_sh[prefetch_buff], &tensor_map_input, chunk_stage_offset_X, - chunk_stage_offset_Y, &act_in_sh[prefetch_buff], &tensor_map_act_input, - chunk_stage_offset_X, chunk_stage_offset_Y, shmem_buff_size, - &mbar[prefetch_buff], is_master_thread); - } else { - copy_2d_to_shared(&in_sh[prefetch_buff], &tensor_map_input, chunk_stage_offset_X, - chunk_stage_offset_Y, shmem_buff_size, &mbar[prefetch_buff], - is_master_thread); - } - } - -#pragma unroll - for (int iter = 0; iter < FP8_ITERATIONS; ++iter) { - const size_t buff = iter % FP8_BUFFERS_NUM; - const size_t next_iter = iter + FP8_PREFETCH_BUFFERS_NUM; - const size_t row_base = block_offset_Y + iter * FP8_BUFFER_DIM_Y; - if (next_iter < FP8_ITERATIONS) { - const size_t next_buff = next_iter % FP8_BUFFERS_NUM; - const size_t chunk_it_offset_y = chunk_offset_Y + next_iter * FP8_BUFFER_DIM_Y; - const size_t chunk_it_offset_x = chunk_offset_X; - if constexpr (IS_DACT) { - copy_2d_to_sharedx2(&in_sh[next_buff], &tensor_map_input, chunk_it_offset_x, - chunk_it_offset_y, &act_in_sh[next_buff], &tensor_map_act_input, - chunk_it_offset_x, chunk_it_offset_y, shmem_buff_size, &mbar[next_iter], - is_master_thread); - } else { - copy_2d_to_shared(&in_sh[next_buff], &tensor_map_input, chunk_it_offset_x, - chunk_it_offset_y, shmem_buff_size, &mbar[next_iter], is_master_thread); - } - } - - // Wait for the data to have arrived - ptx::mbarrier_wait_parity(&mbar[iter], parity); - -#pragma unroll - for (int stage = 0; stage < FP8_BUFF_STAGES_NUM; ++stage) { - const size_t stage_offset_Y = stage; - const size_t shmem_offset_y = thread_offset_Y + stage_offset_Y; - const size_t shmem_offset_x = thread_offset_X; - const size_t row = row_base + shmem_offset_y; - const bool row_out_of_bounds = row >= rows; - const bool out_of_bounds = col_out_of_bounds || row_out_of_bounds; - - float elt = static_cast(in_sh[buff][shmem_offset_y][shmem_offset_x]); - if constexpr (IS_DACT) { - float act_in_elt = static_cast(act_in_sh[buff][shmem_offset_y][shmem_offset_x]); - elt *= OP(act_in_elt, {}); - } - if constexpr (IS_DBIAS) { - if constexpr (IS_DACT) { - if (!out_of_bounds) { - partial_dbias += elt; - } - } else { - // If no activation, elt is 0 so we can safely do this - partial_dbias += elt; - } - } - __builtin_assume(amax >= 0); - if (IS_DACT) { - if (!out_of_bounds) { - amax = fmaxf(amax, fabsf(elt)); - } - } else { - // If no activation, elt is 0 so we can safely do this - amax = fmaxf(amax, fabsf(elt)); - } - out_sh[buff][shmem_offset_y][shmem_offset_x] = static_cast(elt * scale); - } - - // 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 size_t chunk_it_offset_y = chunk_offset_Y + iter * FP8_BUFFER_DIM_Y; - const size_t chunk_it_offset_x = chunk_offset_X; - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output), chunk_it_offset_x, - chunk_it_offset_y, reinterpret_cast(&out_sh[buff])); - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - - // Wait for TMA transfer to have finished reading shared memory. - ptx::cp_async_bulk_wait_group_read(); - } - } - ptx::cp_async_bulk_wait_group_read<0>(); - __syncthreads(); - - parity ^= 1; - - if constexpr (IS_DBIAS) { - const size_t dbias_offset_X = my_column; - const size_t dbias_offset = dbias_offset_Y * dbias_stride + dbias_offset_X; - if (!col_out_of_bounds) { - dbias_workspace[dbias_offset] = partial_dbias; - } - } - - if (amax_ptr != nullptr) { - const int warp_id = threadIdx.x / THREADS_PER_WARP; - // Reduce the amax over the block - amax = reduce_max(amax, warp_id); - // Update the global amax - if (is_master_thread) { - atomicMaxFloat(amax_ptr, amax); - } - } - - // Update scale-inverse - if (is_master_thread && blockIdx.x == 0 && (scale_inv_ptr != nullptr)) { - reciprocal(scale_inv_ptr, scale); - } - - destroy_barriers(mbar, is_master_thread); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} - -constexpr size_t CHUNKS_PER_BLOCK = 128; -constexpr size_t THREADS_PER_BLOCK = FP8_THREADS_PER_CHUNK; -constexpr size_t CHUNK_SIZE = THREADS_PER_BLOCK; -constexpr size_t ELEMS_PER_BLOCK = CHUNKS_PER_BLOCK * CHUNK_SIZE; -constexpr size_t CHUNKS_PER_ITERATION = 32; -constexpr size_t SHMEM_DIM = CHUNKS_PER_ITERATION * CHUNK_SIZE; -constexpr size_t ITERATIONS = CHUNKS_PER_BLOCK / CHUNKS_PER_ITERATION; -constexpr size_t SHMEM_BUFFERS = 2; -static_assert(CHUNKS_PER_BLOCK % CHUNKS_PER_ITERATION == 0); - -template -__global__ void __launch_bounds__(THREADS_PER_BLOCK) - cast_fp8_1D_kernel(const IType *input_ptr, OType *output_ptr, float *const amax_ptr, - float *const scale_inv_ptr, const float *const scale_ptr, const size_t N) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - - const size_t block_offset = blockIdx.x * ELEMS_PER_BLOCK; - const IType *input = input_ptr + block_offset; - OType *output = output_ptr + block_offset; - - float amax = 0; - const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1; - - // The destination shared memory buffer of a bulk tensor operation should be 128-byte aligned - __shared__ alignas(TMA_SHMEM_ALIGNMENT) IType in_sh[SHMEM_BUFFERS][SHMEM_DIM]; - __shared__ alignas(TMA_SHMEM_ALIGNMENT) OType out_sh[SHMEM_BUFFERS][SHMEM_DIM]; - - constexpr size_t transaction_size_IN = sizeof(in_sh) / SHMEM_BUFFERS; - constexpr size_t transaction_size_OUT = sizeof(out_sh) / SHMEM_BUFFERS; - - const bool is_master_thread = (threadIdx.x == 0); - -// 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[ITERATIONS]; - - initialize_barriers(mbar, is_master_thread); - - int parity = 0; - - copy_1d_to_shared(&(in_sh[0]), input, transaction_size_IN, &(mbar[0]), is_master_thread); - -#pragma unroll - for (int iter = 0; iter < ITERATIONS; ++iter) { - const size_t buff = iter % SHMEM_BUFFERS; - const size_t it_offset = iter * SHMEM_DIM; - - const size_t next_iter = iter + 1; - const size_t next_buff = next_iter % SHMEM_BUFFERS; - const size_t next_iter_offset = next_iter * SHMEM_DIM; - - if (next_iter < ITERATIONS) { - copy_1d_to_shared(&(in_sh[next_buff]), input + next_iter_offset, transaction_size_IN, - &(mbar[next_iter]), is_master_thread); - } - - ptx::fence_proxy_async_shared_cta(); - - // Wait for the data to have arrived - ptx::mbarrier_wait_parity(&mbar[iter], parity); - -#pragma unroll - for (int chunk = 0; chunk < CHUNKS_PER_ITERATION; ++chunk) { - const size_t shmem_offset = chunk * CHUNK_SIZE + threadIdx.x; - float elt = static_cast(in_sh[buff][shmem_offset]); - if constexpr (IS_ACT) { - elt = OP(elt, {}); - } - __builtin_assume(amax >= 0); - amax = fmaxf(amax, fabsf(elt)); - out_sh[buff][shmem_offset] = static_cast(elt * scale); - } - - // 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) { - ptx::cp_async_bulk_tensor_1d_shared_to_global( - reinterpret_cast(output + it_offset), - reinterpret_cast(&out_sh[buff]), transaction_size_OUT); - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - - // Wait for TMA transfer to have finished reading shared memory. - ptx::cp_async_bulk_wait_group_read<1>(); - } - } - ptx::cp_async_bulk_wait_group_read<0>(); - __syncthreads(); - - if (amax_ptr != nullptr) { - const int warp_id = threadIdx.x / THREADS_PER_WARP; - // Reduce the amax over the block - amax = reduce_max(amax, warp_id); - // Update the global amax - if (is_master_thread) { - atomicMaxFloat(amax_ptr, amax); - } - } - - // Update scale-inverse - if (is_master_thread && blockIdx.x == 0 && (scale_inv_ptr != nullptr)) { - reciprocal(scale_inv_ptr, scale); - } - - destroy_barriers(mbar, is_master_thread); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} - -constexpr size_t DBIAS_THREADS_PER_BLOCK = 256; -template -__global__ void __launch_bounds__(DBIAS_THREADS_PER_BLOCK) - reduce_dbias_kernel(OType *const dbias_output, const float *const dbias_partial, - const size_t rows, const size_t cols) { - using ComputeVec = Vec; - using OutputVec = Vec; - - const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; - - if (thread_id * nvec >= cols) { - return; - } - - const float *const thread_in_base = dbias_partial + thread_id * nvec; - OType *const thread_out_base = dbias_output + thread_id * nvec; - - ComputeVec ldg_vec; - ComputeVec acc_vec; - acc_vec.clear(); - for (int i = 0; i < rows; ++i) { - ldg_vec.load_from(thread_in_base + i * cols); -#pragma unroll - for (int e = 0; e < nvec; ++e) { - acc_vec.data.elt[e] += ldg_vec.data.elt[e]; - } - } - - OutputVec stg_vec; -#pragma unroll - for (int e = 0; e < nvec; ++e) { - stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); - } - stg_vec.store_to(thread_out_base); -} - -template -void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, const size_t cols, - cudaStream_t stream) { - constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 - constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); - - NVTE_CHECK(cols % reduce_dbias_nvec == 0, "Unsupported shape."); - const size_t reduce_dbias_num_blocks = DIVUP(cols, DBIAS_THREADS_PER_BLOCK * reduce_dbias_nvec); - - reduce_dbias_kernel - <<>>( - reinterpret_cast(dbias->data.dptr), workspace_ptr, rows, cols); - NVTE_CHECK_CUDA(cudaGetLastError()); -} - -template -void cast_fp8_1D(const Tensor &input, Tensor *output, cudaStream_t stream) { - const size_t N = product(input.data.shape); - - const bool isFullTile = (N % ELEMS_PER_BLOCK == 0); - NVTE_CHECK(isFullTile, "Only full tiles are supported."); - NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); - NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); - - const size_t chunks = DIVUP(N, CHUNK_SIZE); - const size_t blocks = DIVUP(chunks, CHUNKS_PER_BLOCK); - - float *const amax_ptr = reinterpret_cast(output->amax.dptr); - float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); - const float *const scale_ptr = reinterpret_cast(output->scale.dptr); - - const dim3 block(THREADS_PER_BLOCK); - const dim3 grid(blocks); - - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.dtype(), IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output->dtype(), OType, - const IType *input_ptr = reinterpret_cast(input.data.dptr); - OType *output_ptr = reinterpret_cast(output->data.dptr); - - cast_fp8_1D_kernel<<>>( - input_ptr, output_ptr, amax_ptr, scale_inv_ptr, scale_ptr, N);); // NOLINT(*) - ); // NOLINT(*) - NVTE_CHECK_CUDA(cudaGetLastError()); -} - -template -void cast_fp8_2D(const Tensor &input, const Tensor *act_input, Tensor *output, Tensor *dbias, - Tensor *workspace, cudaStream_t stream) { - checkCuDriverContext(stream); - - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); - 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; - const size_t blocks_X = chunks_X; - - const size_t dbias_rows = blocks_Y; - const size_t dbias_cols = cols; - - NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); - NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); - - 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(workspace != nullptr, "Workspace must be a tensor."); - - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {dbias_rows, dbias_cols}; - workspace->data.dtype = DType::kFloat32; - return; - } - } - float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; - float *const amax_ptr = reinterpret_cast(output->amax.dptr); - float *const scale_inv_ptr = reinterpret_cast(output->scale_inv.dptr); - float *const scale_ptr = reinterpret_cast(output->scale.dptr); - - const dim3 block(FP8_THREADS_PER_CHUNK); - const dim3 grid(blocks_X, blocks_Y); - - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.data.dtype, IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output->data.dtype, OType, - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_act_input{}; - alignas(64) CUtensorMap tensor_map_output{}; - - create_2D_tensor_map(tensor_map_input, input.data, rows, cols, FP8_SHMEM_DIM_Y, - FP8_SHMEM_DIM_X, cols, 0, typeToNumBits(input.data.dtype)); - - if constexpr (IS_DACT) { - create_2D_tensor_map(tensor_map_act_input, act_input->data, rows, cols, FP8_SHMEM_DIM_Y, - FP8_SHMEM_DIM_X, cols, 0, typeToNumBits(input.data.dtype)); - } - - create_2D_tensor_map(tensor_map_output, output->data, rows, cols, FP8_SHMEM_DIM_Y, - FP8_SHMEM_DIM_X, cols, 0, typeToNumBits(output->data.dtype)); - - cast_fp8_2D_kernel - <<>>(tensor_map_input, tensor_map_act_input, tensor_map_output, - workspace_ptr, amax_ptr, scale_inv_ptr, scale_ptr, rows, - cols); - NVTE_CHECK_CUDA(cudaGetLastError()); - - if constexpr (IS_DBIAS) { - reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); - }); // NOLINT(*) - ); // NOLINT(*) -} - -template -void mxfp8_quantize(const Tensor &input, const Tensor *act_input, - const Tensor *noop, // TODO (ksivamani) - Tensor *output, Tensor *dbias, Tensor *workspace, cudaStream_t stream) { - using namespace mxfp8_kernel; - checkCuDriverContext(stream); - - bool use_rowwise_scaling = output->has_data(); - bool use_colwise_scaling = output->has_columnwise_data(); - NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); - - if (use_rowwise_scaling) { - NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); - } - 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 bool CAST_DBIAS_ONLY = IS_DBIAS && (!IS_DACT) && (!IS_ACT); - - constexpr size_t CHUNK_DIM_Y = CAST_DBIAS_ONLY ? 128 : 64; - constexpr size_t CHUNK_DIM_X = CAST_DBIAS_ONLY ? 128 : 64; - constexpr size_t THREADS_PER_CHUNK = CAST_DBIAS_ONLY ? 128 : 64; - - constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; - constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; - constexpr size_t BUFF_DIM_Y = THREADS_Y; - 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 = use_rowwise_scaling ? output->scale_inv.shape[1] : 1; - const size_t scale_stride_colwise = - use_colwise_scaling ? output->columnwise_scale_inv.shape[1] : 1; - - e8m0_t *const scales_rowwise_ptr = - use_rowwise_scaling ? reinterpret_cast(output->scale_inv.dptr) : nullptr; - e8m0_t *const scales_colwise_ptr = - use_colwise_scaling ? reinterpret_cast(output->columnwise_scale_inv.dptr) : nullptr; - const size_t dbias_rows = blocks_Y; - const size_t dbias_cols = cols; - - ScalingType scaling_type; - if (use_rowwise_scaling && (!use_colwise_scaling)) { - scaling_type = ScalingType::ROWWISE; - } else if ((!use_rowwise_scaling) && use_colwise_scaling) { - scaling_type = ScalingType::COLWISE; - } else if (use_rowwise_scaling && use_colwise_scaling) { - scaling_type = ScalingType::BIDIMENSIONAL; - } - - 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(workspace != nullptr, "Workspace must be a tensor."); - - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {dbias_rows, dbias_cols}; - workspace->data.dtype = DType::kFloat32; - return; - } - } - - float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; - float *const amax_ptr = reinterpret_cast(output->amax.dptr); - const float *noop_ptr = reinterpret_cast(noop->data.dptr); - - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - input.dtype(), IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output->dtype(), OType, - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_act_input{}; - alignas(64) CUtensorMap tensor_map_output_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_colwise{}; - - constexpr size_t input_type_bit_size = TypeInfo::size; - constexpr size_t output_type_bit_size = TypeInfo::size; - - create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, - cols, 0, input_type_bit_size); - - if constexpr (IS_DACT) { - create_2D_tensor_map(tensor_map_act_input, act_input->data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols, 0, input_type_bit_size); - } - - if (use_rowwise_scaling) { - create_2D_tensor_map(tensor_map_output_rowwise, output->data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols, 0, output_type_bit_size); - } - - 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, output_type_bit_size); - } - - constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; - constexpr size_t buff_elems_total = mxfp8_kernel::BUFFS_NUM * buff_elems; - constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; - constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; - constexpr size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); - - constexpr size_t elt_input_mem = buff_size_aligned_in; - constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); - constexpr size_t in_mem = elt_input_mem + act_input_mem; - - const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); - const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); - const size_t out_mem = out_rowwise_mem + out_colwise_mem; - - const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - - switch (scaling_type) { - case ScalingType::ROWWISE: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - cast_mxfp8_2D_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - cast_mxfp8_2D_kernel - <<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - 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: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - cast_mxfp8_2D_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - cast_mxfp8_2D_kernel - <<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - 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: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - cast_mxfp8_2D_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - cast_mxfp8_2D_kernel - <<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - 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; - } - - if constexpr (IS_DBIAS) { - reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); - }); // NOLINT(*) - ); // NOLINT(*) -} - -// This kernel supports only two scaling cases: -// 1. r16c0 - Rowwise NVFP4 -// 2. r16c32 - Rowwise NVFP4 AND Colwise MXFP8 -template -void nvfp4_quantize(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { - using namespace nvfp4_kernel; - using namespace ptx; - checkCuDriverContext(stream); - - 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"); - - 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, nvfp4_kernel::BUFF_DIM_Y, - BUFF_DIM_X, cols, 0, sizeof(IType) * 8); - - create_2D_tensor_map(tensor_map_output_rowwise, output->data, rows, cols, - nvfp4_kernel::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, - nvfp4_kernel::BUFF_DIM_Y, BUFF_DIM_X, cols, 0, sizeof(OType) * 8); - } - - constexpr size_t buff_elems = nvfp4_kernel::BUFF_DIM_Y * BUFF_DIM_X; - constexpr size_t buff_elems_total = nvfp4_kernel::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: - cudaFuncSetAttribute( - cast_nvfp4_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - - cast_nvfp4_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: - cudaFuncSetAttribute( - cast_nvfp4_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - - cast_nvfp4_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; - }); // NOLINT(*) - ); // NOLINT(*) -} - -namespace detail { - -using Empty = transformer_engine::Empty; - -__device__ inline float identity(float value, const Empty &) { return value; } - -struct DequantizeParam { - const float *scale_inv; -}; - -__device__ inline float dequantize_func(float value, const DequantizeParam ¶m) { - return value * (*(param.scale_inv)); -} - -} // namespace detail - -template -void CastVectorizedUnaryKernelLauncher(const Tensor &input, const Tensor *noop, Tensor *output, - cudaStream_t stream) { - constexpr float (*UnaryOP)(float, const ParamOP &) = (OP == nullptr) ? detail::identity : OP; - const size_t N = product(input.data.shape); - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.data.dtype, IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( - output->data.dtype, OType, - if (!is_fp8_dtype(output->data.dtype) || is_tensor_scaling(output->scaling_mode)) { - constexpr int nvec = 32 / sizeof(IType); - VectorizedUnaryKernelLauncher( - reinterpret_cast(input.data.dptr), - reinterpret_cast(noop->data.dptr), - reinterpret_cast(output->data.dptr), - reinterpret_cast(output->scale.dptr), - reinterpret_cast(output->amax.dptr), - reinterpret_cast(output->scale_inv.dptr), N, {}, stream); - } else { - NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); - }); // NOLINT(*) - ); // NOLINT(*) -} - -template -void CastVectorizedUnaryGradKernelLauncher(const Tensor &grad, const Tensor *input, Tensor *output, - cudaStream_t stream) { - constexpr float (*UnaryOP)(float, const ParamOP &) = (OP == nullptr) ? detail::identity : OP; - const size_t N = product(input->data.shape); - TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input->data.dtype, IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT( - output->data.dtype, OType, - if (!is_fp8_dtype(output->data.dtype) || is_tensor_scaling(output->scaling_mode)) { - constexpr int nvec = 32 / sizeof(IType); - VectorizedUnaryGradKernelLauncher( - reinterpret_cast(grad.data.dptr), - reinterpret_cast(input->data.dptr), - reinterpret_cast(output->data.dptr), - reinterpret_cast(output->scale.dptr), - reinterpret_cast(output->amax.dptr), - reinterpret_cast(output->scale_inv.dptr), N, {}, stream); - } else { - NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); - }); // NOLINT(*) - ); // NOLINT(*) -} - -namespace { - -static bool is_full_tile_1D_tensor(const Tensor *const t) { - const size_t N = product(t->data.shape); - const bool isFullTile = (N % ELEMS_PER_BLOCK == 0); - return isFullTile; -} - -bool dimensions_supported_by_TMA(const Tensor *const t) { - const size_t cols = t->flat_last_dim(); - constexpr size_t TMA_bytes = 16; - const size_t alignment_requirement = (TMA_bytes * 8) / typeToNumBits(t->dtype()); - return cols % alignment_requirement == 0; -} - -} // namespace - -// Supported by the Arch >= 10.0 -template -void fp8_quantize_arch_ge_100(const Tensor &input, const Tensor *act_input, const Tensor *noop, - Tensor *output, Tensor *dbias, Tensor *workspace, - cudaStream_t stream) { - switch (output->scaling_mode) { - case NVTE_DELAYED_TENSOR_SCALING: { - if (!IS_DBIAS && !IS_DACT) { - if (is_full_tile_1D_tensor(output) && is_fp8_dtype(output->dtype()) && - is_aligned_tensor_data(input, TMA_GMEM_ALIGNMENT) && - is_aligned_tensor_data(*output, TMA_GMEM_ALIGNMENT)) { - // Aligned AND FP8 - cast_fp8_1D(input, output, stream); - } else { - // Unaligned - CastVectorizedUnaryKernelLauncher(input, noop, output, stream); - } - } else if (!IS_DBIAS && IS_DACT) { - if (dimensions_supported_by_TMA(output) && is_fp8_dtype(output->dtype()) && - is_aligned_tensor_data(input, TMA_GMEM_ALIGNMENT) && - is_aligned_tensor_data(*output, TMA_GMEM_ALIGNMENT) && - is_aligned_tensor_data(*act_input, TMA_GMEM_ALIGNMENT)) { - // Aligned AND FP8 (+dAct) - cast_fp8_2D(input, act_input, output, dbias, workspace, - stream); - } else { - // Unaligned - CastVectorizedUnaryGradKernelLauncher(input, act_input, output, stream); - } - } else { - cast_fp8_2D(input, act_input, output, dbias, workspace, - stream); - } - break; - } - case NVTE_MXFP8_1D_SCALING: { - mxfp8_quantize(input, act_input, noop, output, dbias, - workspace, stream); - break; - } - default: - NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); - } -} - -// Supported by the Arch < 10.0 -template -void fp8_quantize_arch_l_100(const Tensor &input, const Tensor *act_input, const Tensor *noop, - Tensor *output, Tensor *dbias, Tensor *workspace, - cudaStream_t stream) { - if (!is_tensor_scaling(output->scaling_mode) || IS_DBIAS) { - // zhongboz: should we just ignore IS_ACT here? - NVTE_ERROR("Not implemented scaling mode or fusion: " + to_string(output->scaling_mode) + - " or IS_DBIAS=true" + " on GPU with compute capability < 10.0."); - } - switch (output->scaling_mode) { - case NVTE_DELAYED_TENSOR_SCALING: { - if (!IS_DACT) { - CastVectorizedUnaryKernelLauncher(input, noop, output, stream); - } else { - CastVectorizedUnaryGradKernelLauncher(input, act_input, output, stream); - } - break; - } - default: - NVTE_ERROR("Not implemented scaling mode: " + to_string(output->scaling_mode) + "."); - } -} - -template -void fp8_quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, Tensor *output, - Tensor *dbias, Tensor *workspace, cudaStream_t stream) { - CheckNoopTensor(*noop, "cast_noop"); - CheckInputTensor(input, "cast_input"); - CheckOutputTensor(*output, "cast_output"); - - if constexpr (IS_DBIAS) { - NVTE_CHECK(dbias != nullptr); - CheckOutputTensor(*dbias, "dbias"); - } - if constexpr (IS_DACT) { - NVTE_CHECK(act_input != nullptr); - CheckInputTensor(*act_input, "activation_input"); - NVTE_CHECK(input.dtype() == act_input->dtype(), "Types of both inputs must match."); - NVTE_CHECK(input.data.shape == act_input->data.shape, "Shapes of both inputs must match."); - } - - NVTE_CHECK(!is_fp8_dtype(input.dtype()), "Input must be in higher precision."); - NVTE_CHECK(output->data.shape == input.data.shape, "Input and output shapes need to match."); - - // Supported by the Arch >= 10.0 - if (is_supported_by_CC_100()) { - fp8_quantize_arch_ge_100(input, act_input, noop, output, - dbias, workspace, stream); - } else { - // Supported by the Arch < 10.0 - fp8_quantize_arch_l_100(input, act_input, noop, output, - dbias, workspace, stream); - } -} - -namespace detail { - -template -void quantize_helper(const NVTETensor input, const NVTETensor grad, NVTETensor output, - NVTETensor dbias, NVTETensor workspace, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - const Tensor *input_tensor; - const Tensor *activation_input_tensor; - if constexpr (IS_DBIAS || IS_DACT) { - // backward - input is incoming gradient - input_tensor = convertNVTETensorCheck(grad); - activation_input_tensor = convertNVTETensor(input); - } else { - // forward = input is activation input - input_tensor = convertNVTETensorCheck(input); - activation_input_tensor = nullptr; - } - auto output_tensor = convertNVTETensorCheck(output); - auto dbias_tensor = convertNVTETensor(dbias); - auto workspace_tensor = convertNVTETensor(workspace); - - // Quantization config - QuantizationConfig quant_config_cpp; - if (quant_config != nullptr) { - quant_config_cpp = *reinterpret_cast(quant_config); - } - - // Noop flag - Tensor dummy_tensor; - Tensor *noop_tensor = &dummy_tensor; - if (quant_config_cpp.noop_tensor != nullptr) { - noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); - } - - // Check for unsupported options - if (quant_config_cpp.stochastic_rounding) { - NVTE_CHECK(output_tensor->scaling_mode == NVTE_NVFP4_1D_SCALING, - "Stochastic rounding is only supported for NVFP4 quantization."); - } - - // Dispatch to quantization kernel depending on data format - switch (output_tensor->scaling_mode) { - case NVTE_DELAYED_TENSOR_SCALING: { - if (output_tensor->has_columnwise_data()) { - NVTE_CHECK(output_tensor->has_data(), - "Quantizing in only the columnwise direction not supported yet!"); - if constexpr (!IS_DBIAS && !IS_DACT && !IS_ACT) { - cast_transpose(*input_tensor, *noop_tensor, output_tensor, stream); - } else { - cast_transpose_fused( - *input_tensor, activation_input_tensor, output_tensor, dbias_tensor, workspace_tensor, - stream); - } - } else if (output_tensor->has_data()) { - fp8_quantize( - *input_tensor, activation_input_tensor, noop_tensor, output_tensor, dbias_tensor, - workspace_tensor, stream); - } - break; - } - case NVTE_MXFP8_1D_SCALING: { - mxfp8_quantize( - *input_tensor, activation_input_tensor, noop_tensor, output_tensor, dbias_tensor, - workspace_tensor, stream); - break; - } - case NVTE_NVFP4_1D_SCALING: { - // Check tensors - CheckNoopTensor(*noop_tensor, "cast_noop"); - CheckInputTensor(*input_tensor, "input"); - CheckOutputTensor(*output_tensor, "output", false); - - // Choose kernel - int32_t rows = input_tensor->flat_first_dim(); - int32_t cols = input_tensor->flat_last_dim(); - auto dtype = input_tensor->dtype(); - 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 (quant_config_cpp.nvfp4_2d_quantization) { - nvfp4_quantize_transpose( - *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); - } else { - nvfp4_quantize_transpose( - *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); - } - } else { - auto &global_amax = (output_tensor->amax.dptr != nullptr) ? output_tensor->amax - : output_tensor->columnwise_amax; - NVTE_CHECK((!IS_DBIAS && !IS_DACT && !IS_ACT), - "IS_DBIAS, IS_DACT, and IS_ACT not implemented for NVTE_NVFP4_1D_SCALING for " - "2D quantization"); - quantize_transpose_vector_blockwise_fp4( - /*input=*/input_tensor->data, /*global_amax=*/global_amax, - /*scale_inv=*/output_tensor->scale_inv, - /*scale_inv_t=*/output_tensor->columnwise_scale_inv, - /*output=*/output_tensor->data, /*output_t=*/output_tensor->columnwise_data, - /*epsilon=*/0.0f, /*return_identity=*/output_tensor->has_data(), - /*return_transpose=*/output_tensor->has_columnwise_data(), /*pow2_scale=*/false, - /*swizzled_scale=*/false, - /*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); - } - break; - } - case NVTE_BLOCK_SCALING_2D: { - // TODO(kwyss): IS_BIAS, IS_DACT, IS_ACT, ParamOP, OP parameters support. - NVTE_CHECK((!IS_DBIAS && !IS_DACT && !IS_ACT), - "IS_DBIAS, IS_DACT, and IS_ACT not implemented for NVTE_BLOCK_SCALING_2D"); - bool force_pow_2_scales = quant_config_cpp.force_pow_2_scales; - float epsilon = quant_config_cpp.amax_epsilon; - quantize_transpose_square_blockwise( - input_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, - output_tensor->data, output_tensor->columnwise_data, epsilon, - /*return_transpose=*/output_tensor->has_columnwise_data(), force_pow_2_scales, - /*noop_tensor=*/noop_tensor->data, stream); - break; - } - case NVTE_BLOCK_SCALING_1D: { - // TODO(kwyss): IS_BIAS, IS_DACT, IS_ACT, ParamOP, OP parameters support. - NVTE_CHECK((!IS_DBIAS && !IS_DACT && !IS_ACT), - "IS_DBIAS, IS_DACT, and IS_ACT not implemented for NVTE_BLOCK_SCALING_1D"); - bool force_pow_2_scales = quant_config_cpp.force_pow_2_scales; - float epsilon = quant_config_cpp.amax_epsilon; - FP8BlockwiseRowwiseOption rowwise_option = FP8BlockwiseRowwiseOption::NONE; - FP8BlockwiseColumnwiseOption columnwise_option = FP8BlockwiseColumnwiseOption::NONE; - if (output_tensor->has_data()) { - bool rowwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT); - rowwise_option = rowwise_compact ? FP8BlockwiseRowwiseOption::ROWWISE_COMPACT - : FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; - } - if (output_tensor->has_columnwise_data()) { - bool columnwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT); - columnwise_option = columnwise_compact - ? FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT - : FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; - } - quantize_transpose_vector_blockwise( - input_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, - output_tensor->data, output_tensor->columnwise_data, epsilon, rowwise_option, - columnwise_option, force_pow_2_scales, noop_tensor->data, stream); - break; - } - default: - NVTE_ERROR("Not implemented scaling mode: " + to_string(output_tensor->scaling_mode) + "."); - } -} - -} // namespace detail -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_CAST_KERNELS_CUH_ diff --git a/transformer_engine/common/util/math.h b/transformer_engine/common/util/math.h index 2f20817fb0..005a600670 100644 --- a/transformer_engine/common/util/math.h +++ b/transformer_engine/common/util/math.h @@ -36,6 +36,8 @@ __device__ inline OType sigmoid(const IType val, const Empty&) { return 1.f / (1.f + expf(-cval)); } +__device__ inline float sigmoidf(const float x) { return __frcp_rn(1.0f + __expf(-x)); } + template __device__ inline OType dsigmoid(const IType val, const Empty& e) { const float cval = val; diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index aeac2b4a2c..6605d9cad1 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -449,13 +449,12 @@ static_assert(sizeof(fp16x2) == 4); static_assert(sizeof(fp8e4m3x2) == 2); static_assert(sizeof(fp8e5m2x2) == 2); -#if CUDA_VERSION >= 12080 +#if FP4_TYPE_SUPPORTED using fp4e2m1 = __nv_fp4_e2m1; using fp4e2m1x2 = __nv_fp4x2_e2m1; using fp4e2m1x4 = __nv_fp4x4_e2m1; static_assert(sizeof(fp4e2m1x2) == 1); static_assert(sizeof(fp4e2m1x4) == 2); -#endif // CUDA_VERSION >= 12080 // When converting to .e2m1x2 data formats, the destination operand d has .b8 type. // When converting two .f32 inputs to .e2m1x2, each input is converted to the specified format, @@ -464,7 +463,6 @@ static_assert(sizeof(fp4e2m1x4) == 2); // from input b is stored in the lower 4 bits of d. // SIMD like "Fused" cast + multiplication (x4) -#if CUDA_VERSION >= 12080 template __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, const Tx2 &in23, const float scale) { @@ -474,7 +472,192 @@ __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, cons const float x3 = static_cast(in23.y) * scale; out = fp4e2m1x4(make_float4(x0, x1, x2, x3)); } -#endif // CUDA_VERSION >= 12080 + +__device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding( + const uint64_t in_4x, const float2 scale, const uint32_t rbits) { + uint16_t out_4x = 0; + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b16 v0_bf16; \n\t" + ".reg.b16 v1_bf16; \n\t" + ".reg.b16 v2_bf16; \n\t" + ".reg.b16 v3_bf16; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" + "cvt.f32.bf16 v0, v0_bf16; \n\t" + "cvt.f32.bf16 v1, v1_bf16; \n\t" + "cvt.f32.bf16 v2, v2_bf16; \n\t" + "cvt.f32.bf16 v3, v3_bf16; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %3; \n\t" // mind the shuffled elements order + "}" + : "=h"(out_4x) + : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return *reinterpret_cast(&out_4x); +} + +__device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_rn(const uint64_t in_4x, + const float2 scale, + const uint32_t rbits) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + uint32_t out_4x = 0; // Only need 16 bit. Using 32 bit container for packing. + if constexpr (is_blackwell) { + // NOTE: rbits unused for rn. + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b16 v0_bf16; \n\t" + ".reg.b16 v1_bf16; \n\t" + ".reg.b16 v2_bf16; \n\t" + ".reg.b16 v3_bf16; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + ".reg.b8 f0; \n\t" + ".reg.b8 f1; \n\t" + "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" + "cvt.f32.bf16 v0, v0_bf16; \n\t" + "cvt.f32.bf16 v1, v1_bf16; \n\t" + "cvt.f32.bf16 v2, v2_bf16; \n\t" + "cvt.f32.bf16 v3, v3_bf16; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" + "mov.b32 %0, {f0, f1, f0, f1};\n\t" + "}" + : "=r"(out_4x) + : "l"(in_4x), "l"(reinterpret_cast(scale))); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return reinterpret_cast(&out_4x)[0]; +} + +template +__device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x(const uint64_t in_4x, + const float2 scale, + const uint32_t rbits) { + if constexpr (USE_STOCHASTIC_ROUNDING) { + return mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding(in_4x, scale, rbits); + } else { + return mul_cvt_bf16_to_fp4_4x_with_rn(in_4x, scale, rbits); + } +} + +__device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding( + const float2 in01, const float2 in23, const float2 scale, const uint32_t rbits) { + uint16_t out_4x = 0; + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + "mov.b64 {v0, v1} , %1; \n\t" + "mov.b64 {v2, v3} , %2; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %4; \n\t" // mind the shuffled elements order + "}" + : "=h"(out_4x) + : "l"(reinterpret_cast(in01)), + "l"(reinterpret_cast(in23)), + "l"(reinterpret_cast(scale)), "r"(rbits)); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return *reinterpret_cast(&out_4x); +} + +__device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_rn(const float2 in01, + const float2 in23, + const float2 scale, + const uint32_t rbits) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + uint32_t out_4x = 0; // Only need 16 bit. Using 32 bit container for packing. + if constexpr (is_blackwell) { + // NOTE: rbits unused for rn. + asm volatile( + "{\n" + ".reg.b64 v01; \n\t" + ".reg.b64 v23; \n\t" + ".reg.b32 v0; \n\t" + ".reg.b32 v1; \n\t" + ".reg.b32 v2; \n\t" + ".reg.b32 v3; \n\t" + ".reg.b8 f0; \n\t" + ".reg.b8 f1; \n\t" + "mov.b64 {v0, v1} , %1; \n\t" + "mov.b64 {v2, v3} , %2; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order + "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" + "mov.b32 %0, {f0, f1, f0, f1};\n\t" + "}" + : "=r"(out_4x) + : "l"(reinterpret_cast(in01)), + "l"(reinterpret_cast(in23)), + "l"(reinterpret_cast(scale))); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return reinterpret_cast(&out_4x)[0]; +} + +template +__device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x(const float2 in01, const float2 in23, + const float2 scale, + const uint32_t rbits) { + if constexpr (USE_STOCHASTIC_ROUNDING) { + return mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding(in01, in23, scale, rbits); + } else { + return mul_cvt_fp32_to_fp4_4x_with_rn(in01, in23, scale, rbits); + } +} +#endif // FP4_TYPE_SUPPORTED // SIMD like "Fused" cast + multiplication (x2) __device__ __forceinline__ void mul_cvt_2x(fp8e4m3x2 &out, const floatx2 &in, From 26370b117169aec87df9e86f90814a4faabbcc09 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 30 Oct 2025 15:50:16 -0700 Subject: [PATCH 035/521] [PyT] Bump the min version expected to supported FP8 current scaling determinism on Blackwell (#2316) * Bump the min version expected to supported FP8 cs det on Blackwell Signed-off-by: Kshitij Lakhani * Disable fused attn for cudnn < 9.14 for FP8 CS. Disable fused attn for cudnn < 9.18 for FP8 deterministic CS 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> --- .../attention/dot_product_attention/utils.py | 18 +++++++++++++++--- 1 file changed, 15 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 25dc0e96c8..7d4a4f86d9 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -477,9 +477,21 @@ def get_attention_backend( if device_compute_capability < (10, 0): logger.debug("Disabling FusedAttention for FP8 current scaling on arch < sm100") use_fused_attention = False - elif cudnn_version < (9, 14, 0): - logger.debug("Disabling FusedAttention for FP8 current scaling with cuDNN < 9.14.0") - 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( + "Disabling FusedAttention for FP8 current scaling with cuDNN < 9.14.0" + ) + use_fused_attention = False + else: + if deterministic and cudnn_version < (9, 18, 0): + logger.debug( + "Disabling FusedAttention for FP8 current scaling requiring determinism" + " with cuDNN < 9.18.0" + ) + use_fused_attention = False if device_compute_capability == (12, 0): if use_flash_attention: From 1269b2e209c392d41d81f12391cdabc0d5a132fd Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 30 Oct 2025 16:45:44 -0700 Subject: [PATCH 036/521] [JAX] Ensure JAX reference impl uses an accurate backend in our tests (#2322) Ensure JAX reference impl uses an accurate backend Signed-off-by: Jeremy Berchtold --- qa/L1_jax_distributed_unittest/test.sh | 3 ++- qa/L2_jax_distributed_unittest/test.sh | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 270f0df15e..42b70a28e0 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -8,5 +8,6 @@ set -xe : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" -NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* +# Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate. +XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* SCRIPT_NAME=$TE_PATH/tests/jax/test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index 0b73726502..de5624a596 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -8,4 +8,5 @@ set -xe : ${XML_LOG_DIR:=/logs} 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.xml $TE_PATH/tests/jax/test_distributed_* +# Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate. +XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* From 006670de2f518022ff2a625857f626137a764266 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 31 Oct 2025 08:36:03 -0700 Subject: [PATCH 037/521] [JAX] Fix mesh resource requirement when no mesh (#2307) * Fix mesh resource requirement when no mesh Signed-off-by: Jeremy Berchtold * do not require meshresource if all axes are manual axes Signed-off-by: Jeremy Berchtold * remove abstract_mesh is None check Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- transformer_engine/jax/sharding.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index adb67e358f..7f204e768b 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -75,6 +75,16 @@ def get_sharding_map_logic_axis_to_mesh_axis(): """ Generate a dict to map logical axes to mesh axes. """ + mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + if mesh is None or mesh.empty: + # If no mesh is defined, return an empty dict and do not require a MeshResource context to be present + return {} + + abstract_mesh = get_abstract_mesh() + if sorted(abstract_mesh.manual_axes) == sorted(mesh.axis_names): + # If all mesh axes are manual axes, return an empty dict and do not require a MeshResource context to be present + return {} + gsr = global_mesh_resource() is_tpsp_enabled = gsr.tpsp_resource is not None and get_mesh_axis_size(gsr.tpsp_resource) > 1 From e7227af98070ebfcdb08b7f0a99bb87abe7b8532 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Fri, 31 Oct 2025 17:15:21 +0100 Subject: [PATCH 038/521] [Common] Deleted unused header (#2324) Deleted unused header Signed-off-by: Oleg Goncharov --- .../common/util/nvfp4_transpose.cuh | 1514 ----------------- 1 file changed, 1514 deletions(-) delete mode 100644 transformer_engine/common/util/nvfp4_transpose.cuh diff --git a/transformer_engine/common/util/nvfp4_transpose.cuh b/transformer_engine/common/util/nvfp4_transpose.cuh deleted file mode 100644 index 629520aeb7..0000000000 --- a/transformer_engine/common/util/nvfp4_transpose.cuh +++ /dev/null @@ -1,1514 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file nvfp4_transpose.cuh - * \brief CUDA kernels to cast to NVFP4 and transpose. - */ - -#ifndef TRANSFORMER_ENGINE_NVFP4_TRANSPOSE_CUH_ -#define TRANSFORMER_ENGINE_NVFP4_TRANSPOSE_CUH_ - -#include -#include -#include - -#if FP4_TYPE_SUPPORTED -#include -#endif // FP4_TYPE_SUPPORTED -#include - -#include "../common.h" -#include "../utils.cuh" -#include "curanddx.hpp" -#include "math.h" -#include "ptx.cuh" -#include "transformer_engine/transformer_engine.h" - -namespace transformer_engine { - -#if FP4_TYPE_SUPPORTED -namespace nvfp4_transpose { - -using namespace ptx; -using nvfp4_scale_t = fp8e4m3; - -constexpr size_t SCALE_DIM = 16; // NVFP4 block (x16 elts) - -constexpr size_t CHUNK_DIM_Y = 128; -constexpr size_t CHUNK_DIM_X = 128; -constexpr size_t THREADS_NUM = 128; - -constexpr size_t SCALES_PER_CHUNK_Y = CHUNK_DIM_Y / SCALE_DIM; -constexpr size_t SCALES_PER_CHUNK_X = CHUNK_DIM_X / SCALE_DIM; - -constexpr size_t SCALES_PER_THREAD = 2 * (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM / THREADS_NUM; -constexpr size_t RNG_GENS_PER_THREAD = - SCALES_PER_THREAD / 4; // Each call generates 4x uint32_t random numbers - -constexpr size_t TILE_DIM_Y = 32; -constexpr size_t TILE_DIM_X = 128; - -// SHould this be SCALE_DIM or BLOCK_DIM? Both are 16, should work for both 1D and 2D -constexpr size_t SCALES_PER_TILE_Y = TILE_DIM_Y / SCALE_DIM; -constexpr size_t SCALES_PER_TILE_X = TILE_DIM_X / SCALE_DIM; // 128 / 16 = 8 - -constexpr size_t TILES_Y = CHUNK_DIM_Y / TILE_DIM_Y; -constexpr size_t TILES_X = CHUNK_DIM_X / TILE_DIM_X; -constexpr size_t STAGES = TILES_Y * TILES_X; - -constexpr size_t BUFFS_NUM = 2; -constexpr size_t BUFF_DIM_Y = TILE_DIM_Y; -constexpr size_t BUFF_DIM_X = TILE_DIM_X; -constexpr size_t BUFF_SIZE = BUFF_DIM_Y * BUFF_DIM_X; -constexpr size_t BUFF_SIZE_TOTAL = BUFF_SIZE * BUFFS_NUM; - -// Input buffer (BF16) -constexpr size_t BUFF_IN_DIM_Y = BUFF_DIM_Y; -constexpr size_t BUFF_IN_DIM_X = BUFF_DIM_X; -constexpr size_t BUFF_IN_SIZE = BUFF_IN_DIM_Y * BUFF_IN_DIM_X; - -// Output buffer (NVFP4) -constexpr size_t BUFF_OUT_DIM_Y = BUFF_DIM_Y; -constexpr size_t BUFF_OUT_DIM_X = (BUFF_DIM_X * 4) / 8; -constexpr size_t BUFF_OUT_SIZE = BUFF_OUT_DIM_Y * BUFF_OUT_DIM_X; - -// Output transpose buffer (NVFP4) -constexpr size_t BUFF_OUT_T_DIM_Y = BUFF_DIM_X; -constexpr size_t BUFF_OUT_T_DIM_X = (BUFF_DIM_Y * 4) / 8; -constexpr size_t BUFF_OUT_T_SIZE = BUFF_OUT_T_DIM_Y * BUFF_OUT_T_DIM_X; - -// Manual swizzling parameters to reduce SHMEM bank conflicts -constexpr size_t PACK_SIZE = 8; -constexpr size_t WAVES = SCALE_DIM / PACK_SIZE; - -constexpr size_t SCALING_FACTORS_PER_TILE_X = TILE_DIM_X / SCALE_DIM; -constexpr size_t THREADS_X_ROWWISE = SCALING_FACTORS_PER_TILE_X; // 128 / 16 = 8 -constexpr size_t THREADS_Y_ROWWISE = THREADS_NUM / THREADS_X_ROWWISE; // 128 / 8 = 16 - -constexpr size_t ITERATIONS_NORMAL = BUFF_DIM_Y / THREADS_Y_ROWWISE; // 32/ 16 = 2 -constexpr size_t ITERATIONS_TRANSPOSE = BUFF_IN_DIM_Y / SCALE_DIM; -constexpr size_t BUFF_OUT_IT_OFFSET = BUFF_OUT_T_DIM_X / ITERATIONS_TRANSPOSE; - -static_assert(BUFF_DIM_Y >= SCALE_DIM && - "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"); - -// 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; // 8 = 128 / 16 - -// Compute per-block E4M3 encoding/decoding scaling factor -__device__ __forceinline__ nvfp4_scale_t compute_decoding_scaling_factor(const float block_amax, - const float S_enc) { - // constexpr float rcp_6f = 1.0f / 6.0f; - // const float S_dec_b = block_amax * rcp_6f; - // const nvfp4_scale_t S_dec_b_fp8 = static_cast(S_dec_b * S_enc); - // return S_dec_b_fp8; - // NOTE: Divide by 6.0f is not elegant and not efficient. - // However, this is part of the emulation code to ensure exact match. - using namespace detail; - constexpr float fp4_max = TypeExtrema::max; // 6.0f; - const float S_dec_b = block_amax / fp4_max * S_enc; - return static_cast(fminf(S_dec_b, TypeExtrema::max)); -} - -// Compute the global encode scale factor for a given global amax -__device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const float global_amax) { - using namespace detail; - constexpr float fp8_max = TypeExtrema::max; // 448.0f; - 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 - global_encode_scale = fminf(global_encode_scale, TypeExtrema::max); - // If global amax is 0 or infinity, return 1 - if (global_amax == 0.0f || global_encode_scale == 0.0f) { - return 1.0f; - } - return global_encode_scale; -} - -__device__ __forceinline__ uint32_t -get_rbits(transformer_engine::curanddx::detail::philox4x32_native_state<10> - &rng, // philox4x32_native_state<10>: 10 rounds of philox4_32 - uint4 &random_uint4, int &rnd_idx) { - if (rnd_idx == 4) { - rnd_idx = 0; - random_uint4 = rng.generate4(); - } - - // Treat uint4 as an array of 4x uint32_t elements for indexing - const uint32_t *const rbits_arr = reinterpret_cast(&random_uint4); - const uint32_t rbits = rbits_arr[rnd_idx++]; - return rbits; -} - -__device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding( - const uint64_t in_4x, const float2 scale, const uint32_t rbits) { - uint16_t out_4x = 0; - constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; - if constexpr (has_rs) { - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b16 v0_bf16; \n\t" - ".reg.b16 v1_bf16; \n\t" - ".reg.b16 v2_bf16; \n\t" - ".reg.b16 v3_bf16; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" - "cvt.f32.bf16 v0, v0_bf16; \n\t" - "cvt.f32.bf16 v1, v1_bf16; \n\t" - "cvt.f32.bf16 v2, v2_bf16; \n\t" - "cvt.f32.bf16 v3, v3_bf16; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %3; \n\t" // mind the shuffled elements order - "}" - : "=h"(out_4x) - : "l"(in_4x), "l"(reinterpret_cast(scale)), "r"(rbits)); - } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - } - return *reinterpret_cast(&out_4x); -} - -__device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_rn(const uint64_t in_4x, - const float2 scale, - const uint32_t rbits) { - constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; - uint32_t out_4x = 0; // Only need 16 bit. Using 32 bit container for packing. - if constexpr (is_blackwell) { - // NOTE: rbits unused for rn. - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b16 v0_bf16; \n\t" - ".reg.b16 v1_bf16; \n\t" - ".reg.b16 v2_bf16; \n\t" - ".reg.b16 v3_bf16; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - ".reg.b8 f0; \n\t" - ".reg.b8 f1; \n\t" - "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16} , %1; \n\t" - "cvt.f32.bf16 v0, v0_bf16; \n\t" - "cvt.f32.bf16 v1, v1_bf16; \n\t" - "cvt.f32.bf16 v2, v2_bf16; \n\t" - "cvt.f32.bf16 v3, v3_bf16; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %2; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %2; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" - "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" - "mov.b32 %0, {f0, f1, f0, f1};\n\t" - "}" - : "=r"(out_4x) - : "l"(in_4x), "l"(reinterpret_cast(scale))); - } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - } - return reinterpret_cast(&out_4x)[0]; -} - -template -__device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x(const uint64_t in_4x, - const float2 scale, - const uint32_t rbits) { - if constexpr (USE_STOCHASTIC_ROUNDING) { - return mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding(in_4x, scale, rbits); - } else { - return mul_cvt_bf16_to_fp4_4x_with_rn(in_4x, scale, rbits); - } -} - -__device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding( - const float2 in01, const float2 in23, const float2 scale, const uint32_t rbits) { - uint16_t out_4x = 0; - constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; - if constexpr (has_rs) { - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - "mov.b64 {v0, v1} , %1; \n\t" - "mov.b64 {v2, v3} , %2; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rs.satfinite.e2m1x4.f32 %0, {v2, v3, v0, v1}, %4; \n\t" // mind the shuffled elements order - "}" - : "=h"(out_4x) - : "l"(reinterpret_cast(in01)), - "l"(reinterpret_cast(in23)), - "l"(reinterpret_cast(scale)), "r"(rbits)); - } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - } - return *reinterpret_cast(&out_4x); -} - -__device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_rn(const float2 in01, - const float2 in23, - const float2 scale, - const uint32_t rbits) { - constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; - uint32_t out_4x = 0; // Only need 16 bit. Using 32 bit container for packing. - if constexpr (is_blackwell) { - // NOTE: rbits unused for rn. - asm volatile( - "{\n" - ".reg.b64 v01; \n\t" - ".reg.b64 v23; \n\t" - ".reg.b32 v0; \n\t" - ".reg.b32 v1; \n\t" - ".reg.b32 v2; \n\t" - ".reg.b32 v3; \n\t" - ".reg.b8 f0; \n\t" - ".reg.b8 f1; \n\t" - "mov.b64 {v0, v1} , %1; \n\t" - "mov.b64 {v2, v3} , %2; \n\t" - "mov.b64 v01, {v0, v1}; \n\t" - "mov.b64 v23, {v2, v3}; \n\t" - "mul.f32x2 v01, v01, %3; \n\t" // mind the shuffled elements order - "mul.f32x2 v23, v23, %3; \n\t" // mind the shuffled elements order - "mov.b64 {v1, v0}, v01; \n\t" - "mov.b64 {v3, v2}, v23; \n\t" - "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" - "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" - "mov.b32 %0, {f0, f1, f0, f1};\n\t" - "}" - : "=r"(out_4x) - : "l"(reinterpret_cast(in01)), - "l"(reinterpret_cast(in23)), - "l"(reinterpret_cast(scale))); - } else { - NVTE_DEVICE_ERROR( - "FP4 cvt PTX instructions are architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); - } - return reinterpret_cast(&out_4x)[0]; -} - -template -__device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x(const float2 in01, const float2 in23, - const float2 scale, - const uint32_t rbits) { - if constexpr (USE_STOCHASTIC_ROUNDING) { - return mul_cvt_fp32_to_fp4_4x_with_stochastic_rounding(in01, in23, scale, rbits); - } else { - return mul_cvt_fp32_to_fp4_4x_with_rn(in01, in23, scale, rbits); - } -} - -template -__global__ void __launch_bounds__(THREADS_NUM) - nvfp4_transpose_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ CUtensorMap tensor_map_output, - const __grid_constant__ CUtensorMap tensor_map_output_t, - nvfp4_scale_t *const scales_ptr, nvfp4_scale_t *const scales_t_ptr, - const float *noop, const float *const amax_rowwise_ptr, - const float *const amax_colwise_ptr, const size_t rows, - const size_t cols, const size_t scale_stride, - const size_t scale_stride_t, const size_t *rng_state) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - 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; - } - } - - const size_t rng_sequence = - threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; - const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; - const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; - rng.init(rng_seed, rng_sequence, rng_offset); - uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; - - int rnd_idx = - 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x - - constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS; - - const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; - const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; - - const size_t block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; - const size_t block_offset_X_t = blockIdx.y * CHUNK_DIM_Y; - - const size_t chunk_rows = rows - block_offset_Y; - - const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; - const size_t scales_block_offset_X_rowwise = blockIdx.x * SCALES_PER_CHUNK_X; - const size_t scales_block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; - const size_t scales_block_offset_X_t = blockIdx.y * SCALES_PER_CHUNK_Y; - - const size_t tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; - const size_t tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; - const size_t tid_X_colwise = threadIdx.x; - const size_t tid_Y_t = tid_X_colwise; - // const size_t tid_X_t = 0; - - const size_t thread_offset_Y_rowwise = tid_Y_rowwise; - const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM; - const size_t thread_offset_X_colwise = tid_X_colwise; - - const size_t row_base_rowwise = block_offset_Y + thread_offset_Y_rowwise; - const size_t row_base_colwise = block_offset_Y; - const size_t col_base_colwise = block_offset_X + thread_offset_X_colwise; - - const bool col_out_of_bounds_colwise = (col_base_colwise >= cols); - - const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; - const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; - const size_t scales_offset_Y_t = scales_block_offset_Y_t + tid_Y_t; - const size_t scales_offset_X_t = scales_block_offset_X_t; - - const size_t SFs_per_row = cols / SCALE_DIM; - - const bool rowwise_scale_is_within_bounds_X = scales_offset_X_rowwise < SFs_per_row; - const bool colwise_scale_is_within_bounds_Y = scales_offset_Y_t < 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 = - DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - - constexpr size_t in_mem = buff_size_aligned_in; - - constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; - constexpr size_t out_mem_colwise_data = buff_size_aligned_out; - constexpr size_t out_mem_rowwise_scales = 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_data_sh = reinterpret_cast(dshmem + in_mem); - fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - - nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); - IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer - - constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; - - const bool is_master_thread = (threadIdx.x == 0); - - // Compute a global encoding/decoding scaling factors for all S_dec_b - const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) - ? 1.0f - : compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); - // NOTE: This is to match with how emulation code was written. - const float S_dec_rowwise = 1.0 / S_enc_rowwise; - - const float S_enc_colwise = (amax_colwise_ptr == nullptr) - ? S_enc_rowwise - : compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); - const float S_dec_colwise = 1.0 / S_enc_colwise; - - 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 (size_t stage = 0; stage < STAGES; ++stage) { - const size_t buff = stage % BUFFS_NUM; - const size_t next_stage = stage + 1; - const size_t stage_offset_Y = stage * BUFF_DIM_Y; - - const size_t buff_offset_in = buff * BUFF_IN_SIZE; - const size_t buff_offset_out = buff * BUFF_OUT_SIZE; - const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE; - - 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 size_t next_buff = next_stage % BUFFS_NUM; - const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; - const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; - const size_t global_offset_X = block_offset_X; - const size_t next_buff_offset = next_buff * BUFF_IN_SIZE; - - 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; - - // COLWISE scaling - if constexpr (RETURN_TRANSPOSE) { -#pragma unroll - for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { - const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; - const size_t in_thread_offset_X = thread_offset_X_colwise; - - const size_t out_t_thread_offset_Y = thread_offset_X_colwise; - const size_t out_t_thread_offset_X = 0 + it * BUFF_OUT_IT_OFFSET; - - const size_t shmem_offset_base_colwise_in = - buff_offset_in + in_thread_offset_Y * BUFF_IN_DIM_X + in_thread_offset_X; - const size_t shmem_offset_base_colwise_out_t = - buff_offset_out_t + out_t_thread_offset_Y * BUFF_OUT_T_DIM_X + out_t_thread_offset_X; - - block_amax = 0.0f; - float in_compute_colwise[SCALE_DIM]; - IType in_colwise_IType[SCALE_DIM]; - // 1. Read/Compute elements. Find NVFP4-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; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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 E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_colwise); - - // Store scaling factors through SHMEM - const size_t scale_idx_sh = - tid_Y_t * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + it; - out_colwise_scales_sh[scale_idx_sh] = S_dec_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_colwise), float_max); // S_enc_b_fp8 - const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; - - // 3. Scale elements - fp4e2m1x4 regs[SCALE_DIM / 4]; - -#pragma unroll - for (int e = 0; e < SCALE_DIM / 4; ++e) { - const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); - regs[e] = mul_cvt_bf16_to_fp4_4x(elts, block_scale_inverse_2x, - rbits); - } else { - const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); - const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); - regs[e] = mul_cvt_fp32_to_fp4_4x( - in01, in23, block_scale_inverse_2x, rbits); - } - } - - const int group = thread_lane / 16; - uint32_t val[2]; - uint32_t *regs_4x = reinterpret_cast(regs); - - // Helps reducing bank conflicts - switch (group) { - case 0: - val[0] = regs_4x[0]; - val[1] = regs_4x[1]; - break; - case 1: - val[0] = regs_4x[1]; - val[1] = regs_4x[0]; - - break; - } - uint32_t *out_t_data_sh_as_uint32_t = - reinterpret_cast(&out_t_data_sh[shmem_offset_base_colwise_out_t]); - out_t_data_sh_as_uint32_t[group] = val[0]; // idx1 = (group + 0) % 2; - out_t_data_sh_as_uint32_t[(group + 1) & 1] = val[1]; // idx2 = (group + 1) % 2; - } - } - - // ROWWISE scaling - { - const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; -#pragma unroll - for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { - const size_t it_thread_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; - - const size_t shmem_offset_base_rowwise_in = - buff_offset_in + it_thread_offset_Y_rowwise * BUFF_IN_DIM_X; - const size_t shmem_offset_base_rowwise_out = - buff_offset_out + it_thread_offset_Y_rowwise * BUFF_OUT_DIM_X; - - const size_t it_offset_Y = stage_offset_Y + it * THREADS_Y_ROWWISE; - - block_amax = 0.0f; - float in_compute_rowwise[SCALE_DIM]; - 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t 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 size_t 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 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 = 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; - 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 -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - Vec out; -#pragma unroll - for (int e = 0; e < PACK_SIZE / 4; ++e) { - const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); - IType2 in01; - IType2 in23; - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - const uint64_t elts = *reinterpret_cast(&in_IType[w].data.elt[2 * e]); - out.data.elt[e] = mul_cvt_bf16_to_fp4_4x( - elts, block_scale_inverse_2x, rbits); - } else if constexpr (IS_CACHED_ACT_OP) { - const uint64_t elts = *reinterpret_cast(&in_cached[w].data.elt[4 * e]); - out.data.elt[e] = mul_cvt_bf16_to_fp4_4x( - elts, block_scale_inverse_2x, rbits); - } else { - const int j = w * PACK_SIZE + 4 * e; - const float2 in01 = make_float2(in_compute_rowwise[j], in_compute_rowwise[j + 1]); - const float2 in23 = make_float2(in_compute_rowwise[j + 2], in_compute_rowwise[j + 3]); - out.data.elt[e] = mul_cvt_fp32_to_fp4_4x( - in01, in23, block_scale_inverse_2x, rbits); - } - } - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; - out.store_to(&out_data_sh[shmem_offset_rowwise]); - } - } - } - - __builtin_assume(thread_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 size_t global_offset_Y = block_offset_Y + stage_offset_Y; - const size_t global_offset_X = block_offset_X; - - const size_t global_offset_Y_t = block_offset_Y_t; - const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; - - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, - reinterpret_cast(&out_data_sh[buff_offset_out])); - - if constexpr (RETURN_TRANSPOSE) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_t), global_offset_X_t, - global_offset_Y_t, reinterpret_cast(&out_t_data_sh[buff_offset_out_t])); - } - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - } - } // end of stages - - // Vectorized store scaling factors through SHMEM - if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { - using ScalesVec = Vec; - const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; - ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); - const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; - const size_t count = // number of scales in Y dimension of this chunk - (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); - nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); - if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { - // Fast path: vectorized store when destination is properly aligned - scales_vec.store_to(dst); - } else { - // Safe path: element-wise store for tails or unaligned destinations - scales_vec.store_to_elts(dst, 0, count); - } - } - - destroy_barriers(mbar, is_master_thread); -#else - NVTE_DEVICE_ERROR("sm_100 or higher is required."); -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} - -template -__global__ void __launch_bounds__(THREADS_NUM) - nvfp4_transpose_kernel_2D(const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ CUtensorMap tensor_map_output, - const __grid_constant__ CUtensorMap tensor_map_output_t, - nvfp4_scale_t *const scales_ptr, nvfp4_scale_t *const scales_t_ptr, - const float *noop, const float *const amax_rowwise_ptr, - const float *const amax_colwise_ptr, const size_t rows, - const size_t cols, const size_t scale_stride, - const size_t scale_stride_t, const size_t *rng_state) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - 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; - } - } - const size_t rng_sequence = - threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; - const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; - const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; - rng.init(rng_seed, rng_sequence, rng_offset); - uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; - - int rnd_idx = - 0; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x - - // NEW: 2D Block-based scaling constants - constexpr size_t BLOCK_DIM = 16; - constexpr size_t BLOCKS_PER_TILE_Y = TILE_DIM_Y / BLOCK_DIM; // 32/16 = 2 - constexpr size_t BLOCKS_PER_TILE_X = TILE_DIM_X / BLOCK_DIM; // 128/16 = 8 - constexpr size_t ITERATIONS_BLOCK = 2; // iterations to calculate 2d block amaxes of 1 tile - constexpr size_t BLOCKS_PER_WARP = BLOCKS_PER_TILE_X / (THREADS_NUM / 32); // 8 / (128/32) = 2 - - constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS; - - const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; - const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; - - const size_t block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; - const size_t block_offset_X_t = blockIdx.y * CHUNK_DIM_Y; - - const size_t chunk_rows = rows - block_offset_Y; - - const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; - const size_t scales_block_offset_X_rowwise = blockIdx.x * SCALES_PER_CHUNK_X; - const size_t scales_block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; - const size_t scales_block_offset_X_t = blockIdx.y * SCALES_PER_CHUNK_Y; - - const size_t tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; - const size_t tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; - const size_t tid_X_colwise = threadIdx.x; - const size_t tid_Y_t = tid_X_colwise; - - const size_t thread_offset_Y_rowwise = tid_Y_rowwise; - const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM; - const size_t thread_offset_X_colwise = tid_X_colwise; - - const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; - const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; - const size_t scales_offset_Y_t = scales_block_offset_Y_t + tid_Y_t; - const size_t scales_offset_X_t = scales_block_offset_X_t; - - const size_t SFs_per_row = cols / SCALE_DIM; - - const bool rowwise_scale_is_within_bounds_X = scales_offset_X_rowwise < SFs_per_row; - const bool colwise_scale_is_within_bounds_Y = scales_offset_Y_t < 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 = - DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - - constexpr size_t in_mem = buff_size_aligned_in; - - constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; - constexpr size_t out_mem_colwise_data = buff_size_aligned_out; - constexpr size_t out_mem_rowwise_scales = 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_data_sh = reinterpret_cast(dshmem + in_mem); - fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - - nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); - IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer - - constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; - - const bool is_master_thread = (threadIdx.x == 0); - - // Compute a global encoding/decoding scaling factors for all S_dec_b - const float S_enc_rowwise = (amax_rowwise_ptr == nullptr) - ? 1.0f - : compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); - // NOTE: This is to match with how emulation code was written. - const float S_dec_rowwise = 1.0 / S_enc_rowwise; - - const float S_enc_colwise = (amax_colwise_ptr == nullptr) - ? S_enc_rowwise - : compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); - const float S_dec_colwise = 1.0 / S_enc_colwise; - - const size_t warp_id = threadIdx.x / 32; - const size_t lane_id = threadIdx.x % 32; - float thread_amax = 0.0f; - const size_t block_in_warp = lane_id / BLOCKS_PER_WARP; - -// 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]; - - __shared__ __align__(16) float block_amax_matrix[BLOCKS_PER_TILE_Y][BLOCKS_PER_TILE_X + 1]; - - // Helper function for warp reduction - auto warp_reduce_amax = [](float thread_amax, int block_in_warp) -> float { -#pragma unroll - for (int delta = 8; delta >= 1; delta /= 2) { - float other_amax = __shfl_xor_sync(0xffffffff, thread_amax, delta); - thread_amax = fmaxf(thread_amax, other_amax); - } - return thread_amax; - }; - - 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 (size_t stage = 0; stage < STAGES; ++stage) { - const size_t buff = stage % BUFFS_NUM; - const size_t next_stage = stage + 1; - const size_t stage_offset_Y = stage * BUFF_DIM_Y; - - const size_t buff_offset_in = buff * BUFF_IN_SIZE; - const size_t buff_offset_out = buff * BUFF_OUT_SIZE; - const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE; - - 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 size_t next_buff = next_stage % BUFFS_NUM; - const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; - const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; - const size_t global_offset_X = block_offset_X; - const size_t next_buff_offset = next_buff * BUFF_IN_SIZE; - - 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; - -#pragma unroll - for (size_t block_iter = 0; block_iter < ITERATIONS_BLOCK; ++block_iter) { - IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; - const size_t block_in_tile_y = block_iter; - const size_t block_in_tile_x = threadIdx.x / BLOCK_DIM; - - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - for (int elem = 0; elem < BLOCK_DIM; elem += 2) { - const size_t elem_0_row = block_iter * BLOCK_DIM + elem; - const size_t elem_1_row = elem_0_row + 1; - const size_t elem_0_col = warp_id * BLOCKS_PER_WARP * BLOCK_DIM + lane_id; - const size_t elem_1_col = elem_0_col; - - const size_t shmem_offset_0 = buff_offset_in + elem_0_row * BUFF_IN_DIM_X + elem_0_col; - const size_t shmem_offset_1 = buff_offset_in + elem_1_row * BUFF_IN_DIM_X + elem_1_col; - - IType2 val_2x; - val_2x.x = in_sh[shmem_offset_0]; - val_2x.y = in_sh[shmem_offset_1]; - ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, val_2x); - } - - thread_amax = - static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); - } else { - for (int elem = 0; elem < BLOCK_DIM; ++elem) { - const size_t elem_row = block_iter * BLOCK_DIM + elem; - const size_t elem_col = warp_id * BLOCKS_PER_WARP * BLOCK_DIM + lane_id; - - // Bounds checking - const bool row_out_of_bounds = (block_offset_Y + stage_offset_Y + elem_row >= rows); - const bool col_out_of_bounds = (block_offset_X + elem_col >= cols); - if (!row_out_of_bounds && !col_out_of_bounds) { - const size_t shmem_offset = buff_offset_in + elem_row * BUFF_IN_DIM_X + elem_col; - float elt = static_cast(in_sh[shmem_offset]); - - if constexpr (COMPUTE_ACTIVATIONS) { - elt = OP(elt, {}); - } - if constexpr (!std::is_same_v) { - elt = static_cast(static_cast(elt)); - } - // Cache computed activations - if constexpr (IS_CACHED_ACT_OP) { - cached_act_sh[shmem_offset] = static_cast(elt); - } - - thread_amax = fmaxf(thread_amax, fabsf(elt)); - } - } - } - // Warp reduction to get block amax - block_amax = warp_reduce_amax(thread_amax, block_in_warp); - - if (lane_id == 0 || lane_id == 16) { - block_amax_matrix[block_in_tile_y][block_in_tile_x] = block_amax; - } - } - - // sync thread to ensure block_amax_matrix is done storing - __syncthreads(); - - // COLWISE scaling - if constexpr (RETURN_TRANSPOSE) { -#pragma unroll - for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { - const size_t block_in_tile_y = it; - const size_t block_in_tile_x = threadIdx.x / BLOCK_DIM; - - const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; - const size_t in_thread_offset_X = thread_offset_X_colwise; - - const size_t out_t_thread_offset_Y = thread_offset_X_colwise; - const size_t out_t_thread_offset_X = 0 + it * BUFF_OUT_IT_OFFSET; - - const size_t shmem_offset_base_colwise_in = - buff_offset_in + in_thread_offset_Y * BUFF_IN_DIM_X + in_thread_offset_X; - const size_t shmem_offset_base_colwise_out_t = - buff_offset_out_t + out_t_thread_offset_Y * BUFF_OUT_T_DIM_X + out_t_thread_offset_X; - - block_amax = block_amax_matrix[block_in_tile_y][block_in_tile_x]; - float in_compute_colwise[SCALE_DIM]; - IType in_colwise_IType[SCALE_DIM]; - // 3. Scale elements - - // Load data in - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { -#pragma unroll - for (int i = 0; i < SCALE_DIM; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise_in + i * BUFF_IN_DIM_X; - in_colwise_IType[i] = in_sh[shmem_offset_colwise]; - } - } else { - for (int i = 0; i < SCALE_DIM; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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); - } - - in_compute_colwise[i] = elt; - } - } - - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_colwise); - - // // Store scaling factors through SHMEM - const size_t scale_idx_sh = - tid_Y_t * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + it; - out_colwise_scales_sh[scale_idx_sh] = S_dec_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_colwise), float_max); // S_enc_b_fp8 - const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; - - fp4e2m1x4 regs[SCALE_DIM / 4]; -#pragma unroll - for (int e = 0; e < SCALE_DIM / 4; ++e) { - const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); - regs[e] = mul_cvt_bf16_to_fp4_4x(elts, block_scale_inverse_2x, - rbits); - } else { - const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); - const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); - regs[e] = mul_cvt_fp32_to_fp4_4x( - in01, in23, block_scale_inverse_2x, rbits); - } - } - - const int group = thread_lane / 16; - uint32_t val[2]; - uint32_t *regs_4x = reinterpret_cast(regs); - - // Helps reducing bank conflicts - switch (group) { - case 0: - val[0] = regs_4x[0]; - val[1] = regs_4x[1]; - break; - case 1: - val[0] = regs_4x[1]; - val[1] = regs_4x[0]; - break; - } - uint32_t *out_t_data_sh_as_uint32_t = - reinterpret_cast(&out_t_data_sh[shmem_offset_base_colwise_out_t]); - out_t_data_sh_as_uint32_t[group] = val[0]; // idx1 = (group + 0) % 2; - out_t_data_sh_as_uint32_t[(group + 1) & 1] = val[1]; // idx2 = (group + 1) % 2; - } - } - - // ROWWISE scaling - { - const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; -#pragma unroll - for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { - const size_t block_in_tile_y = it; - const size_t block_in_tile_x = tid_X_rowwise; - const size_t it_thread_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; - - const size_t shmem_offset_base_rowwise_in = - buff_offset_in + it_thread_offset_Y_rowwise * BUFF_IN_DIM_X; - const size_t shmem_offset_base_rowwise_out = - buff_offset_out + it_thread_offset_Y_rowwise * BUFF_OUT_DIM_X; - - block_amax = block_amax_matrix[block_in_tile_y][block_in_tile_x]; - float in_compute_rowwise[SCALE_DIM]; - 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; - // Load elements - in_IType[w].load_from(&in_sh[shmem_offset_rowwise]); - } - } else if constexpr (IS_CACHED_ACT_OP) { - // ensures that all writes to cache made in the section above are visible to all threads - __syncthreads(); -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; - - // Load cached elements - in_cached[w].load_from(&cached_act_sh[shmem_offset_rowwise]); - } - } else { -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t 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 size_t 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)); - } - in_compute_rowwise[j] = elt; - } - } - } - - // 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 = 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; - 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 -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - Vec out; -#pragma unroll - for (int e = 0; e < PACK_SIZE / 4; ++e) { - const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); - IType2 in01; - IType2 in23; - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - const uint64_t elts = *reinterpret_cast(&in_IType[w].data.elt[2 * e]); - out.data.elt[e] = mul_cvt_bf16_to_fp4_4x( - elts, block_scale_inverse_2x, rbits); - } else if constexpr (IS_CACHED_ACT_OP) { - const uint64_t elts = *reinterpret_cast(&in_cached[w].data.elt[4 * e]); - out.data.elt[e] = mul_cvt_bf16_to_fp4_4x( - elts, block_scale_inverse_2x, rbits); - } else { - const int j = w * PACK_SIZE + 4 * e; - const float2 in01 = make_float2(in_compute_rowwise[j], in_compute_rowwise[j + 1]); - const float2 in23 = make_float2(in_compute_rowwise[j + 2], in_compute_rowwise[j + 3]); - out.data.elt[e] = mul_cvt_fp32_to_fp4_4x( - in01, in23, block_scale_inverse_2x, rbits); - } - } - - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; - const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; - out.store_to(&out_data_sh[shmem_offset_rowwise]); - } - } - } - - __builtin_assume(thread_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 size_t global_offset_Y = block_offset_Y + stage_offset_Y; - const size_t global_offset_X = block_offset_X; - - const size_t global_offset_Y_t = block_offset_Y_t; - const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; - - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, - reinterpret_cast(&out_data_sh[buff_offset_out])); - - if constexpr (RETURN_TRANSPOSE) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_t), global_offset_X_t, - global_offset_Y_t, reinterpret_cast(&out_t_data_sh[buff_offset_out_t])); - } - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - } - } // end of stages - - // Vectorized store scaling factors through SHMEM - if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { - using ScalesVec = Vec; - const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; - ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); - const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; - const size_t count = // number of scales in Y dimension of this chunk - (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); - nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; - constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); - if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { - // Fast path: vectorized store when destination is properly aligned - scales_vec.store_to(dst); - } else { - // Safe path: element-wise store for tails or unaligned destinations - scales_vec.store_to_elts(dst, 0, count); - } - } - - destroy_barriers(mbar, is_master_thread); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} -} // namespace nvfp4_transpose -#endif // FP4_TYPE_SUPPORTED - -template -void nvfp4_quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED - bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; - - // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to - // return the transposed data. - // TODO(Frank): Is there a better way to do this? - bool return_transpose = output->has_columnwise_data(); - - using namespace nvfp4_transpose; - using namespace ptx; - - checkCuDriverContext(stream); - CheckNoopTensor(*noop, "cast_noop"); - CheckInputTensor(input, "input"); - CheckOutputTensor(*output, "output", false); - - NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - 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"); - if (return_transpose) { - NVTE_CHECK(output->has_columnwise_data(), "NVFP4 transposed output tensor must be allocated."); - NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), - "Transposed output must have FP4 type."); - NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, - "Transposed scaling tensor must be allocated"); - } - - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); - - NVTE_CHECK(rows % 32 == 0, - "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA - NVTE_CHECK(cols % 32 == 0, - "Number of tensor cols must be a multiple of 32"); // 16B alignment for TMA - - 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_NUM; - - const size_t scale_stride = output->scale_inv.shape[1]; - const size_t scale_stride_transpose = - return_transpose ? output->columnwise_scale_inv.shape[1] : 0; - - nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); - nvfp4_scale_t *const scales_transpose_ptr = - reinterpret_cast(output->columnwise_scale_inv.dptr); - - const float *noop_ptr = reinterpret_cast(noop->data.dptr); - const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); - const float *const amax_colwise_ptr = - reinterpret_cast(output->columnwise_amax.dptr); - - const NVTETensor rng_state_tensor = (quant_config != nullptr) ? quant_config->rng_state : nullptr; - const size_t *rng_state = nullptr; - if (rng_state_tensor != nullptr) { - 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}, - "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); - } - - using IType = bf16; - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_output{}; - alignas(64) CUtensorMap tensor_map_output_transpose{}; - - 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, output->data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, - 4); - if (return_transpose) { - create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, - BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); - } - 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 = - DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); - - constexpr size_t in_mem = buff_size_aligned_in; - - constexpr size_t out_data_mem = buff_size_aligned_out; - constexpr size_t out_data_transpose_mem = buff_size_aligned_out; - constexpr size_t out_scales_transpose_mem = buff_size_scales; - - constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; - - constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; - - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, - - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - auto kernel = nvfp4_transpose_kernel; - - if constexpr (use_2d_quantization) { - kernel = nvfp4_transpose_kernel_2D; - } - - 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 -} -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_NVFP4_TRANSPOSE_CUH_ From c57ffc51a83ce16f7961df5b7b65c08080eb6639 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 3 Nov 2025 11:12:49 -0500 Subject: [PATCH 039/521] [JAX] L1_jax_distributed_test suit with individual executions (#2321) * L1 rework Signed-off-by: Phuong Nguyen * comment out test_multi_process_grouped_gemm for now Signed-off-by: Phuong Nguyen * rm e5m2 from test norm + MXFP8 Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- qa/L1_jax_distributed_unittest/test.sh | 36 +++++++++++++++++++++++--- tests/jax/multi_process_launch.sh | 10 ++++++- tests/jax/test_custom_call_compute.py | 7 ++++- 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 42b70a28e0..f4ea2dd68e 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -2,12 +2,42 @@ # # See LICENSE for license information. -set -xe +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" +export NVTE_JAX_UNITTEST_LEVEL="L1" + # Use --xla_gpu_enable_triton_gemm=false to ensure the reference JAX implementation we are using is accurate. -XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" NVTE_JAX_UNITTEST_LEVEL="L1" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/jax/test_distributed_* -SCRIPT_NAME=$TE_PATH/tests/jax/test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh +export XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_triton_gemm=false" + +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_dense.xml $TE_PATH/tests/jax/test_distributed_dense.py || test_fail "test_distributed_dense.py" + +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_helper.xml $TE_PATH/tests/jax/test_distributed_helper.py || test_fail "test_distributed_helper.py" + +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_layernorm.xml $TE_PATH/tests/jax/test_distributed_layernorm.py || test_fail "test_distributed_layernorm.py" + +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_mlp.xml $TE_PATH/tests/jax/test_distributed_layernorm_mlp.py || test_fail "test_distributed_layernorm_mlp.py" + +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py" + +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" + +# TODO(Phuong): add this test back after it is verified +# SCRIPT_NAME=$TE_PATH/tests/jax/test_multi_process_distributed_grouped_gemm.py bash $TE_PATH/tests/jax/multi_process_launch.sh || test_fail "test_multi_process_distributed_grouped_gemm.py" + +if [ $RET -ne 0 ]; then + echo "Error: some sub-tests failed: $FAILED_CASES" + exit 1 +fi +echo "All tests passed" +exit 0 diff --git a/tests/jax/multi_process_launch.sh b/tests/jax/multi_process_launch.sh index fcb066de75..d430e0f413 100644 --- a/tests/jax/multi_process_launch.sh +++ b/tests/jax/multi_process_launch.sh @@ -18,6 +18,14 @@ do CUDA_VISIBLE_DEVICES=$i python $SCRIPT_NAME 127.0.0.1:12345 $i $NUM_RUNS > /dev/null 2>&1 & done -CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_RUNS +CUDA_VISIBLE_DEVICES=0 python $SCRIPT_NAME 127.0.0.1:12345 0 $NUM_RUNS | tee stdout_multi_process.txt wait + +RET=0 +if grep -q "FAILED" stdout_multi_process.txt; then + RET=1 +fi + +rm -f stdout_multi_process.txt +exit "$RET" diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 11ff9d061c..cecdb31218 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -605,7 +605,12 @@ def test_norm_forward_with_tensor_scaling_fp8( ) @pytest.mark.skipif(not is_mxfp8_supported, reason=mxfp8_unsupported_reason) - @pytest.mark.parametrize("out_dtype", [jnp.float8_e4m3fn, jnp.float8_e5m2]) + @pytest.mark.parametrize( + "out_dtype", + [ + jnp.float8_e4m3fn, + ], + ) def test_norm_forward_with_block_scaling_fp8( self, n, hidden, norm_type, zero_centered_gamma, epsilon, inp_dtype, out_dtype ): From 3d76218ee75626f7d749025d8d877468547f78c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 4 Nov 2025 01:26:39 +0100 Subject: [PATCH 040/521] [PyTorch debug] Fixes to debug tests failures (#2268) * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix: Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix 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> --- qa/L0_pytorch_debug_unittest/test.sh | 41 +++++-- tests/pytorch/debug/run_distributed.py | 7 +- .../test_switch_to_nondebug_mode.yaml | 11 ++ tests/pytorch/debug/test_perf.py | 114 +++++++++--------- .../debug/features/_test_dummy_feature.py | 46 ++++++- 5 files changed, 139 insertions(+), 80 deletions(-) create mode 100644 tests/pytorch/debug/test_configs/test_switch_to_nondebug_mode.yaml diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index 9980ccfb05..b6c42109b5 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -2,7 +2,19 @@ # # See LICENSE for license information. +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} : ${NVTE_TEST_NVINSPECT_FEATURE_DIRS:=$TE_PATH/transformer_engine/debug/features} @@ -14,24 +26,27 @@ mkdir -p "$XML_LOG_DIR" # Nvinspect will be disabled if no feature is active. : ${NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE:=$TE_PATH/tests/pytorch/debug/test_configs/dummy_feature.yaml} -FAIL=0 - # It is not installed as a requirement, # because it is not available on PyPI. pip uninstall -y nvdlfw-inspect pip install git+https://github.com/NVIDIA/nvidia-dlfw-inspect.git -pip install pytest==8.2.1 -pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity.xml $TE_PATH/tests/pytorch/debug/test_sanity.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || FAIL=1 -pytest -v -s --junitxml=$XML_LOG_DIR/test_config.xml $TE_PATH/tests/pytorch/debug/test_config.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || FAIL=1 -pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics.xml $TE_PATH/tests/pytorch/debug/test_numerics.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || FAIL=1 -pytest -v -s --junitxml=$XML_LOG_DIR/test_log.xml $TE_PATH/tests/pytorch/debug/test_log.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR || FAIL=1 -NVTE_TORCH_COMPILE=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_api_features.xml $TE_PATH/tests/pytorch/debug/test_api_features.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR || FAIL=1 -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 || FAIL=1 +pip install pytest==8.2.1 || error_exit "Failed to install pytest" +pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity.xml $TE_PATH/tests/pytorch/debug/test_sanity.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || test_fail "test_sanity.py" +pytest -v -s --junitxml=$XML_LOG_DIR/test_config.xml $TE_PATH/tests/pytorch/debug/test_config.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || test_fail "test_config.py" +pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics.xml $TE_PATH/tests/pytorch/debug/test_numerics.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || test_fail "test_numerics.py" +pytest -v -s --junitxml=$XML_LOG_DIR/test_log.xml $TE_PATH/tests/pytorch/debug/test_log.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR || test_fail "test_log.py" +NVTE_TORCH_COMPILE=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_api_features.xml $TE_PATH/tests/pytorch/debug/test_api_features.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR || test_fail "test_api_features.py" +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 || FAIL=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 NVTE_FUSED_ATTN=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics_2.xml $TE_PATH/tests/pytorch/test_numerics.py || FAIL=1 - -exit $FAIL +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 + echo "Error in the following test cases:$FAILED_CASES" + exit 1 +fi +echo "All tests passed" +exit 0 diff --git a/tests/pytorch/debug/run_distributed.py b/tests/pytorch/debug/run_distributed.py index fee2189fa6..358841943a 100644 --- a/tests/pytorch/debug/run_distributed.py +++ b/tests/pytorch/debug/run_distributed.py @@ -668,11 +668,12 @@ def _run_test_with_combinations( _init_distributed() test_log_expert_parallel() - for parallel_mode in ["column", "row"]: - for gather_weight in [True, False]: - test_log_distributed(parallel_mode, gather_weight) if fp8_available: + for parallel_mode in ["column", "row"]: + for gather_weight in [True, False]: + test_log_distributed(parallel_mode, gather_weight) + for parallel_mode in ["row", "column"]: test_disable_fp8_layer(parallel_mode) diff --git a/tests/pytorch/debug/test_configs/test_switch_to_nondebug_mode.yaml b/tests/pytorch/debug/test_configs/test_switch_to_nondebug_mode.yaml new file mode 100644 index 0000000000..224be46180 --- /dev/null +++ b/tests/pytorch/debug/test_configs/test_switch_to_nondebug_mode.yaml @@ -0,0 +1,11 @@ +test_switch_to_nondebug_mode: + enabled: True + layers: + layer_name_regex_pattern: .* + transformer_engine: + TestDummyFeature: + enabled: True + inspect_only_once: True + tensors: [weight, activation, gradient, output, wgrad, dgrad] + gemms: [wgrad, dgrad, fprop] + diff --git a/tests/pytorch/debug/test_perf.py b/tests/pytorch/debug/test_perf.py index ad40c31c02..c8c9ae3c1f 100644 --- a/tests/pytorch/debug/test_perf.py +++ b/tests/pytorch/debug/test_perf.py @@ -6,74 +6,70 @@ import pytest import torch import transformer_engine.pytorch as te -import time import nvdlfw_inspect.api as debug_api from transformer_engine.debug.pytorch.debug_state import TEDebugState -def _run_cpu_overhead(debug_tools_initialized, layer, configs_dir, feature_dirs): - debug_api.end_debug() - TEDebugState._reset() - if debug_tools_initialized: - # This config log stats starting from 0, every N iterations for huge N >> NUM_ITERS. - # So after 1 warm-up iteration, this layers should work in non-debug mode. - debug_api.initialize( - config_file=configs_dir + "/perf_config.yaml", feature_dirs=feature_dirs - ) - - try: - if layer == "linear": - model = torch.nn.Sequential( - te.Linear(1, 1, name="linear1"), te.Linear(1, 1, name="linear2") - ).cuda() - NUM_ITERS = 1800 - elif layer == "transformer": - model = torch.nn.Sequential( - te.TransformerLayer(1, 1, 1, name="transformer1"), - te.TransformerLayer(1, 1, 1, name="transformer2"), - ).cuda() - NUM_ITERS = 200 - - NUM_INVOCATIONS_PER_ITER = 10 +@pytest.mark.parametrize("use_microbatching", [False, True]) +def test_layer_switches_to_nondebug_mode(configs_dir, feature_dirs, use_microbatching): + """ + Test that layers switch to non-debug mode when no features are active. - x = torch.randn(1, 1, 1).cuda() + Uses TestDummyFeature with inspect_only_once=True, which makes inspect_tensor_enabled return (False, None). + The TE should: + 1. Call inspect_tensor_enabled to check if feature is needed + 2. Never call inspect_tensor + 3. Allow layers to switch to non-debug mode for optimal performance, + so that inspect_tensor_enabled is never called again. - y = model(x) - y.sum().backward() - debug_api.step() - torch.cuda.synchronize() + Tests both with and without microbatching to ensure proper behavior in both scenarios. + """ - time_start = time.time() - for i in range(NUM_ITERS): - for _ in range(NUM_INVOCATIONS_PER_ITER): + try: + debug_api.initialize( + config_file=configs_dir + "/test_switch_to_nondebug_mode.yaml", + feature_dirs=feature_dirs, + ) + import transformer_engine.debug.features._test_dummy_feature as dummy_feature + + # Reset counters + dummy_feature._inspect_tensor_enabled_call_count = 0 + dummy_feature._inspect_tensor_call_count = 0 + + model = te.Linear(256, 256, name="test_linear").cuda() + x = torch.randn(8, 256, 256).cuda() + + # Run multiple iterations + for i in range(20): + if use_microbatching: + # Alternate between first and non-first microbatch + is_first_microbatch = i % 2 == 0 + y = model(x, is_first_microbatch=is_first_microbatch) + else: + # Run without specifying is_first_microbatch y = model(x) - y.sum().backward() - if debug_tools_initialized: - debug_api.step() - torch.cuda.synchronize() - time_end = time.time() - - finally: - if debug_tools_initialized: - debug_api.end_debug() - - return time_end - time_start - - -@pytest.mark.parametrize("layer", ["linear", "transformer"]) -def test_cpu_overhead(layer, configs_dir, feature_dirs): - # runs one layer many times on very small tensor - # - gpu time should be negligible, so time should be dominated by cpu time. - # if layers does not invoke any feature in current iteration, - # then it changed into non-debug mode and should not have any non-negligible cpu overhead - # compared to layer without debug tools initialized. - - with_debug_tools = _run_cpu_overhead(True, layer, configs_dir, feature_dirs) - without_debug_tools = _run_cpu_overhead(False, layer, configs_dir, feature_dirs) + y.sum().backward() + debug_api.step() + + # Verify inspect_tensor_enabled was called only once per tensor + # (activation, weight, gradient, output, wgrad, dgrad) + enabled_call_count = dummy_feature._inspect_tensor_enabled_call_count + microbatch_info = "with microbatching" if use_microbatching else "without microbatching" + assert enabled_call_count == 6, ( + f"inspect_tensor_enabled was called {enabled_call_count} times ({microbatch_info}), " + "but should be called 6 times to check if feature is needed for each tensor " + "(activation, weight, gradient, output, wgrad, dgrad)" + ) - print(f"with_debug_tools: {with_debug_tools} s") - print(f"without_debug_tools: {without_debug_tools} s") + # Verify inspect_tensor was never called - it should not be called if inspect_tensor_enabled returns (False, None) + inspect_call_count = dummy_feature._inspect_tensor_call_count + assert inspect_call_count == 0, ( + f"inspect_tensor was called {inspect_call_count} times ({microbatch_info}), " + "but should never be called when inspect_tensor_enabled returns (False, None)" + ) - assert with_debug_tools < without_debug_tools * 1.25 # 25% overhead margin + finally: + debug_api.end_debug() + TEDebugState._reset() diff --git a/transformer_engine/debug/features/_test_dummy_feature.py b/transformer_engine/debug/features/_test_dummy_feature.py index c8a31a3436..4dee97b707 100644 --- a/transformer_engine/debug/features/_test_dummy_feature.py +++ b/transformer_engine/debug/features/_test_dummy_feature.py @@ -7,19 +7,55 @@ from nvdlfw_inspect.registry import Registry, api_method from transformer_engine.debug.features.api import TEConfigAPIMapper +# Module-level counters for tracking invocations +# NOTE: These must be accessed via the full module path +# (transformer_engine.debug.features._test_dummy_feature._inspect_tensor_enabled_call_count) +# to ensure the same module instance is used when the feature is loaded by the debug framework +# and when imported by tests. Using just the variable name would create separate instances +# in different import contexts. +_inspect_tensor_enabled_call_count = 0 +_inspect_tensor_call_count = 0 + @Registry.register_feature(namespace="transformer_engine") class TestDummyFeature(TEConfigAPIMapper): """ - This is feature used only in tests. It invokes look_at_tensor_before_process - and does nothing. + This is feature used only in tests. It invokes inspect_tensor and does nothing. If no features are used, then TE layer automatically switches to the non-debug mode. This feature is invoked for each GEMM to prevent this behavior. + + Config options: + - inspect_only_once: if True, return (False, None) from inspect_tensor_enabled to test caching behavior + + Note: This feature always tracks invocations for testing purposes. """ @api_method - def inspect_tensor_enabled(self, *_args, **_kwargs): - """API call used to determine whether to run look_at_tensor_before_process - in the forward pass.""" + def inspect_tensor_enabled(self, config, *_args, **_kwargs): + """API call used to determine whether to run inspect_tensor in the forward pass. + + Always tracks calls for testing purposes. + + Returns: + - If inspect_only_once=True in config: returns (False, None) - check once, never call inspect_tensor + - Otherwise: returns True - feature is always enabled + """ + # Access counter via full module path to ensure we're modifying the same module-level + # variable regardless of import context (debug framework vs test import) + import transformer_engine.debug.features._test_dummy_feature as dummy_feature # pylint: disable=import-self + + dummy_feature._inspect_tensor_enabled_call_count += 1 + + inspect_only_once = config.get("inspect_only_once", False) + if inspect_only_once: + return False, None return True + + @api_method + def inspect_tensor(self, _config, *_args, **_kwargs): + """This method does nothing but always tracks invocations for testing.""" + # Access counter via full module path to ensure shared state across import contexts + import transformer_engine.debug.features._test_dummy_feature as dummy_feature # pylint: disable=import-self + + dummy_feature._inspect_tensor_call_count += 1 From 77a006352bbe67b6e66fb5a2c5be8ff2d0dd9cde 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, 5 Nov 2025 14:39:41 +0100 Subject: [PATCH 041/521] [PyTorch Debug] Add max_blockwise_dynamic_range stats (#2137) * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix 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> --- tests/pytorch/debug/test_log.py | 198 +++++++++++++++++- .../debug/features/log_tensor_stats.py | 77 ++++++- .../debug/features/utils/stats_buffer.py | 6 +- .../debug/features/utils/stats_computation.py | 116 ++++++++++ 4 files changed, 389 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/debug/test_log.py b/tests/pytorch/debug/test_log.py index e9d074821d..0f833d41fb 100644 --- a/tests/pytorch/debug/test_log.py +++ b/tests/pytorch/debug/test_log.py @@ -18,7 +18,11 @@ ) from transformer_engine.pytorch.quantization import RecipeState from transformer_engine.debug.pytorch.debug_state import TEDebugState - +from transformer_engine.debug.features.utils.stats_computation import ( + compute_max_blockwise_dynamic_range, + BlockwiseDynamicRangeStat, +) +import math fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) @@ -154,7 +158,7 @@ def test_sanity(feature_dirs): @pytest.mark.parametrize("fp8_recipe", fp8_recipes) -def test_numerics(fp8_recipe, feature_dirs): +def test_log_quantized_stats_numerics(fp8_recipe, feature_dirs): if not fp8_available: pytest.skip(reason_for_no_fp8) if not mxfp8_available and fp8_recipe == recipe.MXFP8BlockScaling(): @@ -210,6 +214,107 @@ def test_numerics(fp8_recipe, feature_dirs): assert overflows == pytest.approx(expected.cpu(), abs=1e-4) +LOG_HIGH_PRECISION_CONFIG = """ +log: + layers: + layer_name_regex_pattern: .* + enabled: + True + transformer_engine: + LogTensorStats: + enabled: True + stats: + - dynamic_range + - max_blockwise_dynamic_range: + block_size: 4 + dims: 1 + - max_blockwise_dynamic_range: + block_size: 4 + dims: 2 + tensors: [activation, gradient, weight] + freq: 2 + start_step: 0 + end_step: 10 +""" + + +@pytest.mark.parametrize("tensor_name", ["activation", "weight", "gradient"]) +def test_log_stats_numerics(feature_dirs, tensor_name): + """Check correctness of dynamic range and max blockwise dynamic range stats. + + Tests different tensor types: + - activation/weight: use both orientations (rowwise + columnwise), takes max + - gradient/dgrad: use single orientation (rowwise only) + """ + log_only_bare_stats_config = LOG_HIGH_PRECISION_CONFIG + + with debug_session(log_only_bare_stats_config, feature_dirs) as log_dir: + # There is 1024 x 1024 tensor with very small epsilon values in almost all elements, + # one row of large value A and three rows of large value B. + epsilon = 1e-10 + A = 1000 + B = 50 + tensor = torch.zeros(1024, 1024).cuda() + epsilon + tensor[0, :] = A + tensor[1:4, :] = B + + debug_api.transformer_engine.inspect_tensor( + layer_name="layer_name", + tensor_name=tensor_name, + iteration=0, + tp_group=None, + tensor=tensor, + quantizer=None, + rowwise_quantized_tensor=None, + columnwise_quantized_tensor=None, + ) + debug_api.step() + + output = read_log(log_dir) + + max_over_orientations = tensor_name in ["activation", "weight"] + max_over_orientations_suffix = "_max_over_orientations" if max_over_orientations else "" + + # Track which stats were found to ensure all are present + found_dims_1 = False + found_dims_2 = False + found_dynamic_range = False + + for line in output.splitlines(): + if f"max_blockwise_dynamic_range_block_size_4_dims_1{max_over_orientations_suffix}" in line: + max_blockwise_dynamic_range_block_size_4_dims_1 = float(line.split("value=")[1]) + if max_over_orientations: + # Columnwise blocks have mixed values [A, B, B, B] -> dynamic_range = log2(A/B) + expected = math.log2(A) - math.log2(B) + else: + # Rowwise blocks have uniform values -> dynamic_range = 0 + expected = 0 + assert max_blockwise_dynamic_range_block_size_4_dims_1 == pytest.approx( + expected, abs=1e-4 + ) + found_dims_1 = True + elif ( + f"max_blockwise_dynamic_range_block_size_4_dims_2{max_over_orientations_suffix}" in line + ): + max_blockwise_dynamic_range_block_size_4_dims_2 = float(line.split("value=")[1]) + # For 2D blocks (4x4 tiles), blocks always contain mixed values from different rows + expected = math.log2(A) - math.log2(B) + assert max_blockwise_dynamic_range_block_size_4_dims_2 == pytest.approx( + expected, abs=1e-4 + ) + found_dims_2 = True + elif "_dynamic_range" in line and "max_blockwise_dynamic_range" not in line: + dynamic_range = float(line.split("value=")[1]) + expected = math.log2(A) - math.log2(epsilon) + assert dynamic_range == pytest.approx(expected, abs=1e-4) + found_dynamic_range = True + + # Ensure all expected stats were found in the output + assert found_dims_1, "max_blockwise_dynamic_range (dims=1) not found in output" + assert found_dims_2, "max_blockwise_dynamic_range (dims=2) not found in output" + assert found_dynamic_range, "dynamic_range not found in output" + + @pytest.mark.parametrize("layer", ["linear", "transformer"]) def test_log_every_3_or_5_layers(layer, configs_dir, feature_dirs): if not fp8_available: @@ -256,3 +361,92 @@ def test_log_every_3_or_5_layers(layer, configs_dir, feature_dirs): debug_api.end_debug() TEDebugState._reset() + + +def test_compute_max_blockwise_dynamic_range_direct(): + """Direct unit test for compute_max_blockwise_dynamic_range function. + + Tests the function with various configurations to ensure correct behavior + for different block sizes, dimensions, and orientation settings. + """ + # Create test tensor with uniform rows but mixed columns + # Row 0: all 1000, Row 1-3: all 50, remaining: all 0.01 + epsilon = 0.01 + A = 1000.0 + B = 50.0 + tensor = torch.zeros(1024, 1024).cuda() + epsilon + tensor[0, :] = A + tensor[1:4, :] = B + + # Test 1: dims=1, max_over_orientations=False (rowwise only) + # Rowwise blocks have uniform values -> dynamic_range should be 0 + stat_config = BlockwiseDynamicRangeStat(block_size=4, dims=1, max_over_orientations=False) + result = compute_max_blockwise_dynamic_range(tensor, stat_config) + assert result.item() == pytest.approx( + 0.0, abs=1e-4 + ), "Rowwise 1D blocks with uniform values should have dynamic_range=0" + + # Test 2: dims=1, max_over_orientations=True (max of rowwise and columnwise) + # Columnwise blocks have mixed values [A, B, B, B] -> dynamic_range = log2(A/B) + stat_config = BlockwiseDynamicRangeStat(block_size=4, dims=1, max_over_orientations=True) + result = compute_max_blockwise_dynamic_range(tensor, stat_config) + expected = math.log2(A) - math.log2(B) + assert result.item() == pytest.approx(expected, abs=1e-4), ( + f"Max over orientations should capture columnwise dynamic_range, expected {expected}, got" + f" {result.item()}" + ) + + # Test 3: dims=2, block_size=4 (4x4 tiles) + # 2D blocks span multiple rows -> always have mixed values + stat_config = BlockwiseDynamicRangeStat(block_size=4, dims=2, max_over_orientations=False) + result = compute_max_blockwise_dynamic_range(tensor, stat_config) + expected = math.log2(A) - math.log2(B) + assert result.item() == pytest.approx(expected, abs=1e-4), ( + f"2D blocks should capture mixed values from different rows, expected {expected}, got" + f" {result.item()}" + ) + + # Test 4: Different block size + # With block_size=8, columnwise blocks contain [A, B, B, B, epsilon, epsilon, epsilon, epsilon] + # So max=A, min=epsilon (not B anymore) + stat_config = BlockwiseDynamicRangeStat(block_size=8, dims=1, max_over_orientations=True) + result = compute_max_blockwise_dynamic_range(tensor, stat_config) + expected = math.log2(A) - math.log2(epsilon) # min is epsilon, not B + assert result.item() == pytest.approx( + expected, abs=1e-4 + ), f"Block size 8 should work correctly, expected {expected}, got {result.item()}" + + # Test 5: Tensor with all uniform values -> dynamic_range should be 0 + uniform_tensor = torch.ones(64, 64).cuda() * 42.0 + stat_config = BlockwiseDynamicRangeStat(block_size=4, dims=1, max_over_orientations=True) + result = compute_max_blockwise_dynamic_range(uniform_tensor, stat_config) + assert result.item() == pytest.approx( + 0.0, abs=1e-4 + ), "Uniform tensor should have dynamic_range=0" + + # Test 6: 3D tensor flattening validation using 2D/3D comparison + # Create a 4x4 tensor with distinct 2x2 blocks, compute with dims=2, block_size=2 + # Then reshape to 3D and compute again - results should match if flattening is correct + tensor_2d = torch.tensor( + [ + [1.0, 1.0, 10.0, 10.0], + [1.0, 1.0, 10.0, 10.0], + [100.0, 100.0, 1000.0, 1000.0], + [100.0, 100.0, 1000.0, 1000.0], + ] + ).cuda() + + # Compute on 2D tensor: 4 blocks of 2x2, max range is log2(1000/100) + stat_config = BlockwiseDynamicRangeStat(block_size=2, dims=2, max_over_orientations=False) + result_2d = compute_max_blockwise_dynamic_range(tensor_2d, stat_config) + + # Reshape to 3D [2, 2, 4] and compute - should give same result if flattening is correct + tensor_3d = tensor_2d.reshape(2, 2, 4) + result_3d = compute_max_blockwise_dynamic_range(tensor_3d, stat_config) + + assert result_2d.item() == pytest.approx(result_3d.item(), abs=1e-6), ( + "3D tensor [2,2,4] flattened to [4,4] must give same result as original 2D, got" + f" 2D={result_2d.item()}, 3D={result_3d.item()}" + ) + + print("All direct tests for compute_max_blockwise_dynamic_range passed!") diff --git a/transformer_engine/debug/features/log_tensor_stats.py b/transformer_engine/debug/features/log_tensor_stats.py index e917cf9a00..ff37e659a0 100644 --- a/transformer_engine/debug/features/log_tensor_stats.py +++ b/transformer_engine/debug/features/log_tensor_stats.py @@ -4,7 +4,7 @@ """LogTensorStats Feature support for nvidia-dlframework-inspect""" -from typing import Dict, Optional +from typing import Dict, Optional, List import torch @@ -19,6 +19,10 @@ from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS from transformer_engine.debug.features.utils import next_enabled_iter, get_reduction_params +from transformer_engine.debug.features.utils.stats_computation import ( + add_max_blockwise_dynamic_range_stats, + BlockwiseDynamicRangeStat, +) @Registry.register_feature(namespace="transformer_engine") @@ -44,7 +48,14 @@ class LogTensorStats(BaseLogTensorStats): - l1_norm - l2_norm - cur_amax – maximal absolute value of a tensor, - - dynamic_range – equal to `torch.log2(amax) - torch.log2(amin)` + - dynamic_range – equal to `torch.log2(amax) - torch.log2(nonzero_amin)` + - max_blockwise_dynamic_range – Computes the maximum dynamic range `log2(amax) - log2(nonzero_amin)` across all blocks of size block_size within the tensor. + If tensor and its transpose is needed in training, this stat is computed for both orientations and the maximum is returned. + For `dim=1` there are block_size consecutive elements in the block, for `dim=2` the block is block_size x block_size elements tile. + + - block_size: int, default = 32 + - dims: int, default = 1, allowed values are 1 and 2 + tensors/tensors_struct: List[str] list of tensors to log @@ -88,6 +99,60 @@ class LogTensorStats(BaseLogTensorStats): stats: [dynamic_range] """ + def _is_supported_stat(self, stat: str | Dict): + """Returns True if the stat is supported by this feature, False otherwise.""" + if isinstance(stat, dict): + stat_name = list(stat.keys())[0] + if stat_name == "max_blockwise_dynamic_range": + stat_dict = stat[stat_name] + if not isinstance(stat_dict, dict): + return False + # Ensure only supported keys are present + allowed_keys = {"block_size", "dims"} + if any(k not in allowed_keys for k in stat_dict.keys()): + return False + block_size = stat_dict.get("block_size", 32) + dims = stat_dict.get("dims", 1) + # Type and value validation + if not isinstance(block_size, int) or not isinstance(dims, int): + return False + if block_size > 0 and dims in [1, 2]: + return True + return False + return stat in BaseLogTensorStats._get_supported_stats_list(None) | { + "cur_amax", + "dynamic_range", + } + + def _parse_max_blockwise_dynamic_range_stats( + self, stats: List[str | Dict], tensor_name: str + ) -> List[str | BlockwiseDynamicRangeStat]: + """ + Adds all max_blockwise_dynamic_range stats to the stat computation logic. + Changes the types of the stats from Dict to BlockwiseDynamicRangeStat named tuple, + for other stats nothing is changed. + + For example, if the stats is [{"max_blockwise_dynamic_range": {"block_size": 32, "dims": 1}}], + it will be changed to [BlockwiseDynamicRangeStat(block_size=32, dims=1, max_over_orientations=True)] + or [BlockwiseDynamicRangeStat(block_size=32, dims=1, max_over_orientations=False)] depending on tensor_name. + + """ + max_over_orientations = tensor_name in ["activation", "weight"] + parsed_stats = [] + for stat in stats: + if isinstance(stat, dict): + block_size = stat["max_blockwise_dynamic_range"].get("block_size", 32) + dims = stat["max_blockwise_dynamic_range"].get("dims", 1) + + # Register stat and return the named tuple + parsed_stat = add_max_blockwise_dynamic_range_stats( + block_size, dims, max_over_orientations + ) + parsed_stats.append(parsed_stat) + else: + parsed_stats.append(stat) + return parsed_stats + def _get_supported_stats_list(self): """Returns stats this feature can log.""" return BaseLogTensorStats._get_supported_stats_list(None) | {"cur_amax", "dynamic_range"} @@ -147,14 +212,16 @@ def inspect_tensor( ) for stat in config["stats"]: - assert ( - stat in self._get_supported_stats_list() + assert self._is_supported_stat( + stat ), f"[NVTORCH INSPECT ERROR] Statistic {stat} is not supported." + stats = self._parse_max_blockwise_dynamic_range_stats(config["stats"], tensor_name) + STATS_BUFFERS.try_add_buffer( layer_name=layer_name, tensor_name=tensor_name, - stats=config["stats"], + stats=stats, options=options, reduction_group=reduction_group, reduce_within_microbatch=reduce_within_microbatch, diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py index 20236fb950..b5b462f5a2 100644 --- a/transformer_engine/debug/features/utils/stats_buffer.py +++ b/transformer_engine/debug/features/utils/stats_buffer.py @@ -130,8 +130,12 @@ def log(self): for stat_name in self.stats_to_log: combiner = STATS[stat_name][1] stat_value = combiner(gathered_helper_stats) + + # Convert stat key to string for logging (uses __str__ for named tuples) + stat_name_str = str(stat_name) + MetricLogger.log_scalar( - f"{self.layer_name}_{self.tensor_name}_{stat_name}", stat_value, self.iteration + f"{self.layer_name}_{self.tensor_name}_{stat_name_str}", stat_value, self.iteration ) output[(self.layer_name, self.tensor_name, stat_name, self.iteration)] = ( stat_value # for debugging purposes diff --git a/transformer_engine/debug/features/utils/stats_computation.py b/transformer_engine/debug/features/utils/stats_computation.py index 2fa6985acf..8c480441c5 100644 --- a/transformer_engine/debug/features/utils/stats_computation.py +++ b/transformer_engine/debug/features/utils/stats_computation.py @@ -7,12 +7,25 @@ """ import math +from collections import namedtuple + import torch import torch.nn.functional as F import transformer_engine_torch as tex from transformer_engine.common.recipe import Format +class BlockwiseDynamicRangeStat( + namedtuple("BlockwiseDynamicRangeStat", ["block_size", "dims", "max_over_orientations"]) +): + """Named tuple representing a blockwise dynamic range statistic configuration.""" + + def __str__(self) -> str: + """Convert to string representation for stat name. Used for logging.""" + suffix = "_max_over_orientations" if self.max_over_orientations else "" + return f"max_blockwise_dynamic_range_block_size_{self.block_size}_dims_{self.dims}{suffix}" + + @torch.compile def _compute_dynamic_range_top(tensor): """Computes the log2 of the amax of the tensor""" @@ -26,6 +39,7 @@ def _compute_dynamic_range_top(tensor): return torch.log2(amax) +@torch.compile def _compute_dynamic_range_bottom(tensor): """Computes the log2 of the amin of the tensor""" tensor_abs = tensor.abs() @@ -37,6 +51,76 @@ def _compute_dynamic_range_bottom(tensor): return torch.log2(amin) +def compute_max_blockwise_dynamic_range(tensor, stat_config): + """ + Computes maximum blockwise dynamic range (log2 max/min_nonzero) within blocks. + + Flattens tensor to 2D and computes maximum dynamic range within blocks. If max_over_orientations + is True, computes for both rowwise and columnwise orientations and returns the maximum, + capturing the worst-case scenario regardless of how the tensor is used in GEMM operations. + If False, computes only for rowwise orientation. + + Returns 0 if all blocks are zeros, otherwise computes dynamic range over non-zero blocks. + + Args: + tensor: Input tensor (will be flattened to 2D) + stat_config: BlockwiseDynamicRangeStat named tuple with: + - block_size: Size of blocks (int) + - dims: 1 for 1D blocks (consecutive elements), 2 for 2D blocks (tiles) + - max_over_orientations: If True, compute max over rowwise and columnwise orientations + """ + # Extract parameters from stat_config + block_size = stat_config.block_size + dims = stat_config.dims + max_over_orientations = stat_config.max_over_orientations + + def _compute_for_one_orientation(tensor): + total_numel = tensor.numel() + assert dims in [1, 2], f"dims must be 1 or 2, got {dims}" + + # torch.compile friendly code - standard ** power does not work with jit + total_block_size = block_size * block_size if dims == 2 else block_size + assert ( + total_numel % total_block_size == 0 + ), f"Tensor numel ({total_numel}) is not divisible by block_size ({block_size})." + + tensor = tensor.abs().float() + if dims == 1: + tensor = tensor.reshape(-1, block_size) + per_block_amax = tensor.amax(dim=1) + per_block_amin = tensor.masked_fill(tensor == 0, float("inf")).amin(dim=1) + else: + # We want to have tensor of shape [nr_blocks, block_size, block_size], + # where each block is a block_size x block_size tile of the original tensor. + dim_y = tensor.shape[-1] // block_size + tensor = ( + tensor.reshape(-1, block_size, dim_y, block_size) + .permute(0, 2, 1, 3) + .reshape(-1, block_size, block_size) + ) + per_block_amax = tensor.amax(dim=(1, 2)) + per_block_amin = tensor.masked_fill(tensor == 0, float("inf")).amin(dim=(1, 2)) + + # Identify blocks that contain any non-zero element + nonzero_blocks = per_block_amax != 0 + dynamic_range_per_block = torch.where( + nonzero_blocks, + torch.log2(per_block_amax) - torch.log2(per_block_amin), + torch.zeros_like(per_block_amax, dtype=torch.float32), + ) + return dynamic_range_per_block.max() + + # Flatten to 2D + tensor_2d = tensor.reshape(-1, tensor.shape[-1]) + if max_over_orientations: + return max( + _compute_for_one_orientation(tensor_2d), # Rowwise orientation + _compute_for_one_orientation(tensor_2d.transpose(-2, -1)), # Columnwise orientation + ) + return _compute_for_one_orientation(tensor_2d) + + +@torch.compile def compute_variance(variances, numels, sums): """Welford algorithm is used for numerically stable distributed variance computation.""" mean = torch.sum(sums) / torch.sum(numels) @@ -45,6 +129,7 @@ def compute_variance(variances, numels, sums): return var +@torch.compile def compute_std(variances, numels, sums): """Computates standard deviation.""" return torch.sqrt(compute_variance(variances, numels, sums)) @@ -316,6 +401,37 @@ def add_mse_stats(recipe_name: str, columnwise: bool = False): DEPENDENCIES[stat_mse] = {stat_mse, stat_err, "numel"} +def add_max_blockwise_dynamic_range_stats( + block_size: int, dims: int, max_over_orientations: bool = False +): + """Register max_blockwise_X_dynamic_range stats for the recipe. + + Args: + block_size: Size of blocks for computing blockwise dynamic range + dims: 1 for 1D blocks, 2 for 2D blocks + max_over_orientations: Whether to compute max over rowwise and columnwise orientations + + Returns: + BlockwiseDynamicRangeStat named tuple representing this stat (used as the stat key) + """ + # Use named tuple directly as the stat key - this is cleaner than string keys + stat_key = BlockwiseDynamicRangeStat(block_size, dims, max_over_orientations) + + if stat_key in stats_to_num: + return stat_key # already registered + + assert dims in [1, 2], f"dims must be 1 or 2, got {dims}" + stats_to_num[stat_key] = len(stats_to_num) + DEPENDENCIES[stat_key] = {stat_key} + + STATS[stat_key] = ( + lambda x, aux_dict, _stat_key=stat_key: compute_max_blockwise_dynamic_range(x, _stat_key), + lambda buffers, _stat_key=stat_key: max(_get(buffers, _stat_key)), + ) + + return stat_key + + for _columnwise in [True, False]: for _recipe_name in [ "", # default recipe From b6020e3bce7e0a22c6bf988ddee578943c60821f 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, 5 Nov 2025 23:49:11 +0100 Subject: [PATCH 042/521] [JAX] Fix bug with pre scale bias (#2300) * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski --- transformer_engine/jax/flax/transformer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 1eafed4131..42c9451245 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -197,6 +197,7 @@ def __call__( fused_scale_factor = scale_factor if self.attn_bias_type == AttnBiasType.PRE_SCALE_BIAS: attn_weights += bias + bias = None def apply_swa_mask(original_mask: Array) -> Array: """Apply the sliding window mask to a given mask""" From dcaca2a67ee0c390cb900a4c29168aef6ac198d5 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 5 Nov 2025 16:11:53 -0800 Subject: [PATCH 043/521] [JAX] Try to use pre-downloaded dataset artifacts first (#2345) * Try to use pre-downloaded dataset artifacts first Signed-off-by: Jeremy Berchtold * Set HF_HUB_OFFLINE to disable any network calls to HF when the pre-downloaded dataset is available Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- examples/jax/encoder/common.py | 55 ++++++++++++++++--- .../encoder/test_model_parallel_encoder.py | 4 +- examples/jax/encoder/test_multigpu_encoder.py | 4 +- .../encoder/test_multiprocessing_encoder.py | 4 +- .../jax/encoder/test_single_gpu_encoder.py | 4 +- examples/jax/mnist/test_single_gpu_mnist.py | 4 +- 6 files changed, 57 insertions(+), 18 deletions(-) diff --git a/examples/jax/encoder/common.py b/examples/jax/encoder/common.py index 9ffcfe57da..819fdf443d 100644 --- a/examples/jax/encoder/common.py +++ b/examples/jax/encoder/common.py @@ -3,6 +3,9 @@ # See LICENSE for license information. """Shared functions for the encoder tests""" from functools import lru_cache +import os +import pathlib +import zipfile import jax import jax.numpy @@ -120,12 +123,48 @@ def get_quantization_recipe_from_name_string(name: str): raise ValueError(f"Invalid quantization_recipe, got {name}") -def hf_login_if_available(): - """Login to HF hub if available""" - try: - from huggingface_hub import login +@lru_cache(maxsize=None) +def _get_example_artifacts_dir() -> pathlib.Path: + """Path to directory with pre-downloaded datasets""" - login() - except Exception as e: - print(e) - pass + # Check environment variable + path = os.getenv("NVTE_TEST_CHECKPOINT_ARTIFACT_PATH") + if path: + return pathlib.Path(path).resolve() + + # Fallback to path in root dir + root_dir = pathlib.Path(__file__).resolve().parent.parent.parent + return root_dir / "artifacts" / "examples" / "jax" + + +def _unpack_cached_dataset(artifacts_dir: pathlib.Path, folder_name: str) -> None: + """Unpack a cached dataset if available""" + dataset_dir = artifacts_dir / folder_name + if not dataset_dir.exists(): + print(f"Cached dataset {folder_name} not found at {dataset_dir}, skipping unpack") + return + + # Disable any HF network calls since the dataset is cached locally + os.environ["HF_HUB_OFFLINE"] = "1" + + for filename in os.listdir(dataset_dir): + filepath = dataset_dir / filename + if not filename.endswith(".zip"): + continue + print(f"Unpacking cached dataset {folder_name} from {filepath}") + + with zipfile.ZipFile(filepath, "r") as zip_ref: + zip_ref.extractall(pathlib.Path.home() / ".cache" / "huggingface") + print( + f"Unpacked cached dataset {folder_name} to" + f" {pathlib.Path.home() / '.cache' / 'huggingface'}" + ) + + +# This is cached so we don't have to unpack datasets multiple times +@lru_cache(maxsize=None) +def unpack_cached_datasets_if_available() -> None: + """Unpack cached datasets if available""" + artifacts_dir = _get_example_artifacts_dir() + _unpack_cached_dataset(artifacts_dir, "mnist") + _unpack_cached_dataset(artifacts_dir, "encoder") diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index c6d867ef98..a3935da9ff 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -23,14 +23,14 @@ is_bf16_supported, get_quantization_recipe_from_name_string, assert_params_sufficiently_sharded, - hf_login_if_available, + unpack_cached_datasets_if_available, ) import transformer_engine.jax as te import transformer_engine.jax.cpp_extensions as tex import transformer_engine.jax.flax as te_flax from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode -hf_login_if_available() +unpack_cached_datasets_if_available() DEVICE_DP_AXIS = "data" DEVICE_TP_AXIS = "model" diff --git a/examples/jax/encoder/test_multigpu_encoder.py b/examples/jax/encoder/test_multigpu_encoder.py index 1004dd2dd2..80a2b043cb 100644 --- a/examples/jax/encoder/test_multigpu_encoder.py +++ b/examples/jax/encoder/test_multigpu_encoder.py @@ -22,14 +22,14 @@ from common import ( is_bf16_supported, get_quantization_recipe_from_name_string, - hf_login_if_available, + unpack_cached_datasets_if_available, ) import transformer_engine.jax as te import transformer_engine.jax.cpp_extensions as tex import transformer_engine.jax.flax as te_flax from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode -hf_login_if_available() +unpack_cached_datasets_if_available() DEVICE_DP_AXIS = "data" PARAMS_KEY = "params" diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index 9605adf771..4d21411169 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -27,13 +27,13 @@ is_mxfp8_supported, is_nvfp4_supported, get_quantization_recipe_from_name_string, - hf_login_if_available, + unpack_cached_datasets_if_available, ) import transformer_engine.jax as te import transformer_engine.jax.cpp_extensions as tex import transformer_engine.jax.flax as te_flax -hf_login_if_available() +unpack_cached_datasets_if_available() os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" DEVICE_DP_AXIS = "data" diff --git a/examples/jax/encoder/test_single_gpu_encoder.py b/examples/jax/encoder/test_single_gpu_encoder.py index 81f2d6c744..7835b08b23 100644 --- a/examples/jax/encoder/test_single_gpu_encoder.py +++ b/examples/jax/encoder/test_single_gpu_encoder.py @@ -19,13 +19,13 @@ from common import ( is_bf16_supported, get_quantization_recipe_from_name_string, - hf_login_if_available, + unpack_cached_datasets_if_available, ) import transformer_engine.jax as te import transformer_engine.jax.flax as te_flax from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode -hf_login_if_available() +unpack_cached_datasets_if_available() PARAMS_KEY = "params" DROPOUT_KEY = "dropout" diff --git a/examples/jax/mnist/test_single_gpu_mnist.py b/examples/jax/mnist/test_single_gpu_mnist.py index 2e9d56e93f..62f7954e0d 100644 --- a/examples/jax/mnist/test_single_gpu_mnist.py +++ b/examples/jax/mnist/test_single_gpu_mnist.py @@ -25,10 +25,10 @@ from encoder.common import ( is_bf16_supported, get_quantization_recipe_from_name_string, - hf_login_if_available, + unpack_cached_datasets_if_available, ) -hf_login_if_available() +unpack_cached_datasets_if_available() IMAGE_H = 28 IMAGE_W = 28 From f3b97c26b58212b96b0580bc884ddc263a1ea1c2 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Thu, 6 Nov 2025 11:47:17 -0800 Subject: [PATCH 044/521] Fix out of bounds access in the FP4 dequantize kernel (#2346) Signed-off-by: Przemek Tredak --- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 4 +++ .../tensor/storage/nvfp4_tensor_storage.py | 33 ++----------------- 2 files changed, 7 insertions(+), 30 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index bf7b535be4..5307cad37f 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -39,6 +39,10 @@ __global__ void __launch_bounds__(512) const size_t x = thread_idx % M; const size_t y = thread_idx / M; + if (y >= N) { + return; + } + union fp4vec { uint64_t vec; fp4e2m1x4 small_vec[4]; diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 67543a8e2a..04ab092ee2 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -13,12 +13,12 @@ import torch -# import transformer_engine_torch as tex +import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType from ...quantized_tensor import QuantizedTensorStorage, Quantizer -# from ...constants import TE_DType as torch_to_transformer_engine_dtype +from ...constants import TE_DType as torch_to_transformer_engine_dtype from ...utils import _empty_tensor @@ -45,34 +45,7 @@ def forward( # Dequantize row-wise data if tensor._rowwise_data is not None: - ### TODO(tmoon): Debug dequantize kernel and remove unfused impl - # return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype]) - - # Tensor properties - shape = list(tensor._rowwise_data.size()) - shape[-1] *= 2 - device = tensor._rowwise_data.device - - # Convert FP4E2M1 values to FP32 - data = tensor._rowwise_data.view(torch.uint8).to(torch.int32) - data = torch.stack((data & 0x0F, data >> 4), dim=-1).reshape(shape) - data = _fp4_e2m1_vals(device, dtype=torch.float32)[data] - data = data.to(torch.float32).contiguous() - - # Convert FP8E4M3 block scales to FP32 - block_scales = tensor._rowwise_scale_inv - block_scales = block_scales.reshape(-1, block_scales.size(-1)) - block_scales = block_scales[: math.prod(shape[:-1]), : shape[-1] // 16] - block_scales = block_scales.view(torch.float8_e4m3fn).to(torch.float32) - - # Convert amax to FP32 tensor scale - tensor_scale = tensor._amax_rowwise / (6.0 * 448.0) # Scale by FP4E2M1 and FP8E4M3 max - - # Apply scales - block_data = data.view(-1, 16) - block_data *= tensor_scale.view(()) * block_scales.reshape(-1, 1) - - return data.to(dtype) + return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype]) if tensor._columnwise_data is not None: raise NotImplementedError("Dequantizing column-wise NVFP4 data is not implemented yet!") From b14a3b62cd58ab800d55daf42337b36514e943c7 Mon Sep 17 00:00:00 2001 From: Kunlun Li <94586211+kunlunl@users.noreply.github.com> Date: Fri, 7 Nov 2025 05:19:46 +0800 Subject: [PATCH 045/521] Make FP8 weights compatible with older MCore version (#2342) * Make cast_master_weights_to_fp8 compatible with older MCore version Signed-off-by: kunlunl * Rename keep_columnwise to manual_post_all_gather_processing & Optimize unit test Signed-off-by: kunlunl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove redundant _test_mini_optimizer() Signed-off-by: kunlunl --------- Signed-off-by: kunlunl 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> --- .../run_cast_master_weights_to_fp8.py | 690 ---------------- .../test_cast_master_weights_to_fp8.py | 751 +++++++++++++++++- transformer_engine/pytorch/tensor/utils.py | 56 +- 3 files changed, 771 insertions(+), 726 deletions(-) delete mode 100644 tests/pytorch/distributed/run_cast_master_weights_to_fp8.py diff --git a/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py deleted file mode 100644 index 2f11a24ee8..0000000000 --- a/tests/pytorch/distributed/run_cast_master_weights_to_fp8.py +++ /dev/null @@ -1,690 +0,0 @@ -#!/usr/bin/python3 - -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -import argparse -import datetime -import os -import sys - -import torch -from torch import nn -import torch.distributed as dist - -from transformer_engine.common.recipe import ( - DelayedScaling, - Float8CurrentScaling, - Float8BlockScaling, - Format, - Recipe, -) -import transformer_engine.pytorch as te -from transformer_engine.pytorch import ( - QuantizedTensor, - Float8Tensor, - Float8BlockwiseQTensor, -) -from transformer_engine.pytorch.tensor import cast_master_weights_to_fp8 -from transformer_engine.pytorch.tensor.utils import post_all_gather_processing, replace_raw_data - - -def _get_raw_data(quantized_tensor): - """Get the underlying data of a quantized tensor, used in zero-1 optimizer""" - if isinstance(quantized_tensor, Float8Tensor): - assert hasattr(quantized_tensor, "_data"), "Float8Tensor does not have _data attribute" - assert quantized_tensor._data.dtype == torch.uint8, "Float8Tensor _data must be uint8" - return quantized_tensor._data - elif isinstance(quantized_tensor, Float8BlockwiseQTensor): - assert hasattr( - quantized_tensor, "_rowwise_data" - ), "Float8BlockwiseQTensor does not have _rowwise_data attribute" - assert ( - quantized_tensor._rowwise_data.dtype == torch.uint8 - ), "Float8BlockwiseQTensor _rowwise_data must be uint8" - return quantized_tensor._rowwise_data - else: - raise ValueError(f"Unsupported quantized tensor type: {type(quantized_tensor)}") - - -class MiniZero_1: - """A mini zero-1 optimizer implementation, just used for this test""" - - def __init__(self, weights, lr, dp_group): - self.rank = dist.get_rank(dp_group) - self.world_size = dist.get_world_size(dp_group) - - self.weights = weights - self.lr = lr - self.dp_group = dp_group - - # [self.offsets[i], self.offsets[i+1]) is the range of weights[i] in the global buffer - self.offsets = [0] - for weight in self.weights: - self.offsets.append(self.offsets[-1] + weight.numel()) - - # Padding to avoid global buffer cannot be divided by world size, so the offsets[-1] may - # not be the end range of the last weight. - if self.offsets[-1] % self.world_size != 0: - self.offsets[-1] += self.world_size - self.offsets[-1] % self.world_size - - self.master_weights = [] - # The start offset of the master weight in the weight - self.start_offsets = [] - # The overlapping area of the weight and this rank's local buffer - self.overlapping_areas = [] - - # The start and end of this rank's local buffer in the global buffer - rank_start = self.offsets[-1] // self.world_size * self.rank - rank_end = rank_start + self.offsets[-1] // self.world_size - - for weight, offset in zip(self.weights, self.offsets[:-1]): - if offset >= rank_end or (offset + weight.numel()) <= rank_start: - # This weight is not in this rank's local buffer - master_weight = None - start_offset = None - overlapping_area = None - else: - overlapping_start = max(rank_start, offset) - overlapping_end = min(rank_end, offset + weight.numel()) - length = overlapping_end - overlapping_start - start_offset = overlapping_start - offset - if isinstance(weight, QuantizedTensor): - # If weight is a FP8 tensor, we need to use the original high precision version - # to initialize the master weight. - high_precision_init_val = weight.get_high_precision_init_val().view(-1) - master_weight = high_precision_init_val.to(weight.device).float()[ - start_offset : start_offset + length - ] - else: - master_weight = ( - weight.detach().view(-1).float()[start_offset : start_offset + length] - ) - overlapping_area = (overlapping_start, overlapping_end) - self.master_weights.append(master_weight) - self.start_offsets.append(start_offset) - self.overlapping_areas.append(overlapping_area) - - # Create global buffer for grads reduce-scatter - self.grad_buffer = torch.empty( - [self.offsets[-1]], dtype=torch.float32, device=weights[0].device - ) - self.grad_buffer_slice = self.grad_buffer[rank_start:rank_end] - - # Create global buffer for weights all-gather - if isinstance(self.weights[0], QuantizedTensor): - weight_buffer_dtype = torch.uint8 - else: - weight_buffer_dtype = weights[0].dtype - self.weight_buffer = torch.empty( - [self.offsets[-1]], dtype=weight_buffer_dtype, device=weights[0].device - ) - self.weight_buffer_slice = self.weight_buffer[rank_start:rank_end] - - def step(self): - # ----------------------------------------------------------------------------------------- - # Step 1: Copy grads to the grad buffer - # ----------------------------------------------------------------------------------------- - for weight, offset in zip(self.weights, self.offsets[:-1]): - start = offset - end = offset + weight.numel() - self.grad_buffer[start:end].copy_(weight.main_grad.view(-1)) - - # ----------------------------------------------------------------------------------------- - # Step 2: Grads reduce-scatter - # ----------------------------------------------------------------------------------------- - # Don't use reduce_scatter directly to explicitly control the reduce order. - # dist.reduce_scatter_tensor(self.grad_buffer_slice, self.grad_buffer, op=dist.ReduceOp.AVG, - # group=self.dp_group) - buffers = [torch.empty_like(self.grad_buffer) for _ in range(self.world_size)] - dist.all_gather(buffers, self.grad_buffer, group=self.dp_group) - for i in range(1, self.world_size): - buffers[0] += buffers[i] - rank_start = self.offsets[-1] // self.world_size * self.rank - rank_end = rank_start + self.offsets[-1] // self.world_size - self.grad_buffer_slice.copy_(buffers[0][rank_start:rank_end]) - self.grad_buffer_slice /= self.world_size - - # ----------------------------------------------------------------------------------------- - # Step 3: Update master weights - # ----------------------------------------------------------------------------------------- - for master_weight, overlapping_area in zip(self.master_weights, self.overlapping_areas): - if master_weight is None: - # This weight's master weight is in other rank. - continue - grad = self.grad_buffer[overlapping_area[0] : overlapping_area[1]] - master_weight -= grad * self.lr - - # ----------------------------------------------------------------------------------------- - # Step 4: Cast master weights to BF16 or FP8, depending on the type of the weight - # ----------------------------------------------------------------------------------------- - if isinstance(self.weights[0], QuantizedTensor): - # FP8 weights case - for i in range(1, len(self.weights)): - assert isinstance(self.weights[i], QuantizedTensor) - cast_master_weights_to_fp8( - self.weights, self.master_weights, self.start_offsets, self.dp_group - ) - else: - # BF16 weights case - for weight, master_weight, start_offset in zip( - self.weights, self.master_weights, self.start_offsets - ): - if master_weight is None: - continue - start = start_offset - end = start_offset + master_weight.numel() - weight.data.view(-1)[start:end].copy_(master_weight) - - # ----------------------------------------------------------------------------------------- - # Step 5: Copy the updated weights (not all weights) to the weight buffer - # ----------------------------------------------------------------------------------------- - for i in range(len(self.weights)): - master_weight = self.master_weights[i] - if master_weight is None: - continue - start_offset = self.start_offsets[i] - if isinstance(self.weights[i], QuantizedTensor): - weight = _get_raw_data(self.weights[i]) - else: - weight = self.weights[i] - weight_slice = weight.view(-1)[start_offset : start_offset + master_weight.numel()] - overlapping_start, overlapping_end = self.overlapping_areas[i] - self.weight_buffer[overlapping_start:overlapping_end].copy_(weight_slice) - - # ----------------------------------------------------------------------------------------- - # Step 6: Weight all-gather (FP8 or BF16) - # ----------------------------------------------------------------------------------------- - dist.all_gather_into_tensor( - self.weight_buffer, self.weight_buffer_slice, group=self.dp_group - ) - - # ----------------------------------------------------------------------------------------- - # Step 7: Copy the gathered weights from weight buffer to the actual weights - # ----------------------------------------------------------------------------------------- - quantized_weights = [] - for weight, offset in zip(self.weights, self.offsets[:-1]): - start = offset - end = offset + weight.numel() - if isinstance(weight, QuantizedTensor): - quantized_weights.append(weight) - weight = _get_raw_data(weight) - weight.view(-1).data.copy_(self.weight_buffer[start:end]) - post_all_gather_processing(quantized_weights) - - -class MiniOptimizer: - - def __init__(self, weights, lr, dp_group): - self.world_size = dist.get_world_size(dp_group) - - self.weights = weights - self.lr = lr - self.dp_group = dp_group - - master_weights = [] - for weight in self.weights: - master_weights.append(weight.detach().float()) - self.master_weights = master_weights - - def step(self): - for weight, master_weight in zip(self.weights, self.master_weights): - main_grad = weight.main_grad - - # Don't use all-reduce directly to explicitly control the reduce order. - # dist.all_reduce(main_grad, op=dist.ReduceOp.AVG, group=self.dp_group) - buffers = [torch.empty_like(main_grad) for _ in range(self.world_size)] - dist.all_gather(buffers, main_grad, group=self.dp_group) - for i in range(1, self.world_size): - buffers[0] += buffers[i] - main_grad.copy_(buffers[0]) - main_grad /= self.world_size - - master_weight -= main_grad * self.lr - weight.data.copy_(master_weight) - - -class MiniFSDP: - def __init__(self, weights, lr, dp_group): - rank = dist.get_rank(dp_group) - world_size = dist.get_world_size(dp_group) - - self.weights = weights - self.lr = lr - self.dp_group = dp_group - - # Flatten the weights and pad to align with world size - if isinstance(weights[0], QuantizedTensor): - raw_data_list = [_get_raw_data(w).view(-1) for w in weights] - else: - raw_data_list = [w.view(-1) for w in weights] - self.flatten_weight, original_length = self._flatten_tensors_with_pad(raw_data_list) - - # Split flattened weights into shards - self.local_weight_shard = torch.chunk(self.flatten_weight, world_size)[rank] - self.local_main_grad_shard = torch.zeros_like( - self.local_weight_shard, dtype=torch.float32, device="cuda" - ) - shard_size = self.flatten_weight.size(0) // world_size - - # Map original tensors to flattened indices - tensor_indices = [] - cumulative_length = 0 - for tensor in raw_data_list: - length = tensor.size(0) - tensor_indices.append((cumulative_length, cumulative_length + length)) - cumulative_length += length - - # Build shard index mappings - self.weight_indices = [] - self.shard_indices = [] - for idx, (start, end) in enumerate(tensor_indices): - shard_start = rank * shard_size - shard_end = shard_start + shard_size - adjusted_end = min(shard_end, original_length) - - if start <= adjusted_end and end >= shard_start: - start_idx = max(start, shard_start) - end_idx = min(end, adjusted_end) - self.weight_indices.append((start_idx - start, end_idx - start)) - self.shard_indices.append((start_idx - shard_start, end_idx - shard_start)) - else: - self.weight_indices.append((None, None)) - self.shard_indices.append((None, None)) - - if isinstance(weights[idx], QuantizedTensor): - replace_raw_data( - weights[idx], self.flatten_weight[start:end].view(weights[idx].shape) - ) - else: - weights[idx].data = self.flatten_weight[start:end].view(weights[idx].shape) - - # Initialize local model weights and high-precision master weights - self.local_weights = [] - self.master_weights = [] - for i, weight in enumerate(self.weights): - weight_start, weight_end = self.weight_indices[i] - shard_start, shard_end = self.shard_indices[i] - if shard_start is not None and shard_end is not None: - local_weight_shard = self.local_weight_shard[shard_start:shard_end] - self.local_weights.append(local_weight_shard) - - if isinstance(weight, QuantizedTensor): - high_precision_init_val = weight.get_high_precision_init_val().view(-1) - master_weight_shard = high_precision_init_val.to(weight.device).float()[ - weight_start:weight_end - ] - else: - master_weight_shard = weight.detach().view(-1).float()[weight_start:weight_end] - self.master_weights.append(master_weight_shard) - else: - self.local_weights.append(None) - self.master_weights.append(None) - setattr( - weight, "main_grad", torch.zeros_like(weight, dtype=torch.float32, device="cuda") - ) - - def _flatten_tensors_with_pad(self, tensors): - """ - Flatten the list of tensors and pad them to align with the world size. - - Args: - tensors (list): List of tensors to flatten. - - Returns: - tuple: Flattened tensor and its original length before padding. - """ - world_size = dist.get_world_size(self.dp_group) - - flatten_tensor = torch.cat(tensors) - original_length = flatten_tensor.size(0) - - padding_needed = (world_size - original_length % world_size) % world_size - if padding_needed > 0: - zeros = torch.zeros(padding_needed, dtype=flatten_tensor.dtype, device="cuda") - flatten_tensor = torch.cat([flatten_tensor, zeros]) - - return flatten_tensor, original_length - - def zero_grad(self): - for weight in self.weights: - weight.grad = None - weight.main_grad.zero_() - - def step(self): - """ - Perform an optimization step for the distributed sharded model. - - This method includes: - 1. Gradient reduce-scatter: Synchronize gradients across all processes. - 2. Master weight update: Update high-precision master weights using local gradients. - 3. Precision casting: Cast updated master weights to FP8 or BF16 precision. - 4. Weight synchronization: All-gather updated weights across all processes. - - Returns: - None - """ - # Step 1: Reduce-scatter the gradients - main_grad_buffer, _ = self._flatten_tensors_with_pad( - [weight.main_grad.view(-1) for weight in self.weights] - ) - dist.reduce_scatter_tensor( - self.local_main_grad_shard, main_grad_buffer, group=self.dp_group - ) - self.local_main_grad_shard /= dist.get_world_size(self.dp_group) - - # Step 2: Update the master weights - for weight, master_weight, (shard_start, shard_end) in zip( - self.weights, self.master_weights, self.shard_indices - ): - if master_weight is None: - continue - - # Extract the local gradient shard for this weight - grad = self.local_main_grad_shard[shard_start:shard_end] - - # Update the master weight using gradient descent - master_weight -= grad * self.lr - - # Step 3: Cast master weights to FP8 or BF16 precision - if isinstance(self.weights[0], QuantizedTensor): - local_weights = [] - for local_weight in self.local_weights: - if local_weight is None: - local_weights.append(None) - continue - - local_weights.append(local_weight) - - cast_master_weights_to_fp8( - self.weights, - self.master_weights, - [idx[0] for idx in self.weight_indices], - self.dp_group, - local_weights, - ) - else: - for weight, master_weight in zip(self.local_weights, self.master_weights): - if master_weight is None: - continue - - # Copy updated master weights to local weights - weight.data.copy_(master_weight) - - # Step 4: All-gather updated weights across processes - dist.all_gather_into_tensor( - self.flatten_weight, self.local_weight_shard, group=self.dp_group - ) - quantized_weights = [] - for weight in self.weights: - if isinstance(weight, QuantizedTensor): - quantized_weights.append(weight) - post_all_gather_processing(quantized_weights) - - -def _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group): - rank = dist.get_rank(dp_group) - world_size = dist.get_world_size(dp_group) - - # Configuration constants - NUM_STEPS = 100 - SEED = 12345 - - torch.manual_seed(SEED) - torch.cuda.manual_seed(SEED) - - mock_groups = [dist.new_group(ranks=[i]) for i in range(world_size)] - mock_group = mock_groups[rank] - - linear_kwargs = { - "params_dtype": torch.bfloat16, - "bias": False, - "fuse_wgrad_accumulation": True, - } - - # Create model with FP8 weights - with te.quantized_model_init( - enabled=quantization is not None, - recipe=quantization_recipe(quantization), - preserve_high_precision_init_val=True, - ): - model_fp8 = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), - te.Linear(256 * 3, 128, **linear_kwargs), - ) - - # Create model with BF16 weights - model = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), - te.Linear(256 * 3, 128, **linear_kwargs), - ) - - # Make sure the BF16 model and FP8 model have the same initial weights - for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): - high_precision_init_val = w_fp8.get_high_precision_init_val() - w.data.copy_(high_precision_init_val) - - optimizer_fp8 = MiniFSDP([w for w in model_fp8.parameters()], 10.0, dp_group) - optimizer = MiniFSDP([w for w in model.parameters()], 10.0, dp_group) - - for _ in range(100): - optimizer_fp8.zero_grad() - optimizer.zero_grad() - - inputs = [ - torch.randn(16, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) - ] - # Choose based on rank to make sure the inputs of different ranks are different. - x = inputs[rank] - - with te.autocast( - enabled=quantization is not None, - recipe=quantization_recipe(quantization), - amax_reduction_group=mock_group, - ): - y_fp8 = model_fp8(x) - - with te.autocast( - enabled=quantization is not None, - recipe=quantization_recipe(quantization), - amax_reduction_group=mock_group, - ): - y = model(x) - - targets = [torch.randn_like(y) for _ in range(world_size)] - # Choose based on rank to make sure the targets of different ranks are different. - target = targets[rank] - loss_fp8 = nn.MSELoss()(y_fp8, target) - loss = nn.MSELoss()(y, target) - - loss_fp8.backward() - loss.backward() - - optimizer_fp8.step() - optimizer.step() - - torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) - - -def _test_mini_optimizer(dp_group): - """Make sure the implementation of MiniZero_1 and MiniFSDP is correct""" - rank = dist.get_rank(dp_group) - world_size = dist.get_world_size(dp_group) - - torch.manual_seed(12345) - torch.cuda.manual_seed(12345) - - weights = [ - torch.randn(256 * 256, dtype=torch.bfloat16, device="cuda"), - torch.randn(256 * 256 * 3, dtype=torch.bfloat16, device="cuda"), - torch.randn(256 * 256 * 2 - 1, dtype=torch.bfloat16, device="cuda"), - ] - - weights_1 = weights - weights_2 = [weight.clone() for weight in weights] - weights_3 = [weight.clone() for weight in weights] - - lr = 1.0 - optimizer_1 = MiniZero_1(weights_1, lr, dp_group) - optimizer_2 = MiniOptimizer(weights_2, lr, dp_group) - optimizer_3 = MiniFSDP(weights_3, lr, dp_group) - - for _ in range(100): - for w1, w2, w3 in zip(weights_1, weights_2, weights_3): - main_grads = [ - torch.randn_like(w1, dtype=torch.float32, device="cuda") for _ in range(world_size) - ] - # Choose based on rank to make sure the grads of different ranks are different. - main_grad = main_grads[rank] - w1.main_grad = main_grad - w2.main_grad = main_grad - w3.main_grad = main_grad - - optimizer_1.step() - optimizer_2.step() - optimizer_3.step() - - for w1, w2 in zip(weights_1, weights_2): - torch.testing.assert_close(w1, w2, atol=0, rtol=0) - for w1, w3 in zip(weights_1, weights_3): - torch.testing.assert_close(w1, w3, atol=0, rtol=0) - - -def quantization_recipe(quantization) -> Recipe: - """Quantization recipe setup""" - fp8_format = Format.HYBRID - if quantization == "fp8": - return DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") - elif quantization == "fp8_cs": - return Float8CurrentScaling(fp8_format=fp8_format) - elif quantization == "fp8_block": - return Float8BlockScaling(fp8_format=fp8_format) - else: - raise ValueError(f"Unsupported quantization: {quantization}") - - -def _test_cast_master_weights_to_fp8(quantization, dp_group): - rank = dist.get_rank(dp_group) - world_size = dist.get_world_size(dp_group) - - torch.manual_seed(12345) - torch.cuda.manual_seed(12345) - - mock_groups = [dist.new_group(ranks=[i]) for i in range(world_size)] - mock_group = mock_groups[rank] - - linear_kwargs = {"params_dtype": torch.bfloat16, "bias": False, "fuse_wgrad_accumulation": True} - - # Create model with FP8 weights - with te.quantized_model_init( - enabled=quantization is not None, - recipe=quantization_recipe(quantization), - preserve_high_precision_init_val=True, - ): - model_fp8 = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), - te.Linear(256 * 3, 128, **linear_kwargs), - ) - - # Create model with BF16 weights - model = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), - te.Linear(256 * 3, 128, **linear_kwargs), - ) - - # Make sure the BF16 model and FP8 model have the same initial weights - for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): - high_precision_init_val = w_fp8.get_high_precision_init_val() - w.data.copy_(high_precision_init_val) - - # Allocate main_grads for each weight - for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): - w_fp8.main_grad = torch.zeros_like(w_fp8, dtype=torch.float32, device="cuda") - w.main_grad = torch.zeros_like(w, dtype=torch.float32, device="cuda") - - optimizer_fp8 = MiniZero_1([w for w in model_fp8.parameters()], 10.0, dp_group) - optimizer = MiniZero_1([w for w in model.parameters()], 10.0, dp_group) - - for i in range(100): - for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): - w_fp8.main_grad.zero_() - w.main_grad.zero_() - - inputs = [ - torch.randn(16, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) - ] - # Choose based on rank to make sure the inputs of different ranks are different. - x = inputs[rank] - - with te.autocast( - enabled=quantization is not None, - recipe=quantization_recipe(quantization), - amax_reduction_group=mock_group, - ): - y_fp8 = model_fp8(x) - - with te.autocast( - enabled=quantization is not None, - recipe=quantization_recipe(quantization), - amax_reduction_group=mock_group, - ): - y = model(x) - - targets = [torch.randn_like(y) for _ in range(world_size)] - # Choose based on rank to make sure the targets of different ranks are different. - target = targets[rank] - loss_fp8 = nn.MSELoss()(y_fp8, target) - loss = nn.MSELoss()(y, target) - - loss_fp8.backward() - loss.backward() - - optimizer_fp8.step() - optimizer.step() - - torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) - - -def main(argv=None, namespace=None): - WORLD_RANK = int(os.getenv("RANK", "0")) - WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) - LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) - LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) - - assert WORLD_SIZE == LOCAL_SIZE # this test supports only 1 node - assert LOCAL_SIZE <= torch.cuda.device_count() - dist_init_kwargs = { - "backend": "nccl", - "rank": WORLD_RANK, - "world_size": WORLD_SIZE, - "timeout": datetime.timedelta(seconds=30), - } - dist_init_kwargs["init_method"] = "env://" - dist_init_kwargs["device_id"] = torch.device(f"cuda:{LOCAL_RANK}") - assert dist.is_nccl_available() - torch.cuda.set_device(LOCAL_RANK) - dist.init_process_group(**dist_init_kwargs) - - parser = argparse.ArgumentParser() - parser.add_argument( - "--quantization", type=str, default=None, choices=["fp8", "fp8_cs", "fp8_block"] - ) - args = parser.parse_args(argv, namespace) - - dp_group = dist.new_group(backend="nccl") - _test_mini_optimizer(dp_group) - _test_cast_master_weights_to_fp8(args.quantization, dp_group) - _test_fsdp_cast_master_weights_to_fp8(args.quantization, dp_group) - - dist.destroy_process_group() - return 0 - - -if __name__ == "__main__": - - sys.exit(main()) diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 5bf46b8d5f..0ff98e6cb7 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -2,39 +2,744 @@ # # See LICENSE for license information. +import argparse +import datetime import os import subprocess -from pathlib import Path +import sys +import pathlib import pytest import torch -from transformer_engine.pytorch import is_fp8_available, is_fp8_block_scaling_available +from torch import nn +import torch.distributed as dist +from transformer_engine.common.recipe import ( + DelayedScaling, + Float8CurrentScaling, + Float8BlockScaling, + Format, + Recipe, +) +import transformer_engine.pytorch as te +from transformer_engine.pytorch import ( + is_fp8_available, + is_fp8_block_scaling_available, + QuantizedTensor, + Float8Tensor, + Float8BlockwiseQTensor, +) +from transformer_engine.pytorch.tensor import cast_master_weights_to_fp8 +from transformer_engine.pytorch.tensor.utils import post_all_gather_processing, replace_raw_data -if torch.cuda.device_count() < 2: - pytest.skip("cast_master_weights_to_fp8 test needs at least 2 GPUs.") -fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) -fp8_block_scaling_available, reason_for_no_fp8_block_scaling = is_fp8_block_scaling_available( - return_reason=True -) +def _get_quantization_recipe(quantization) -> Recipe: + """Quantization recipe setup""" + fp8_format = Format.HYBRID + if quantization == "fp8": + return DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") + elif quantization == "fp8_cs": + return Float8CurrentScaling(fp8_format=fp8_format) + elif quantization == "fp8_block": + return Float8BlockScaling(fp8_format=fp8_format) + else: + raise ValueError(f"Unsupported quantization: {quantization}") + + +def _get_raw_data(quantized_tensor): + """Get the underlying data of a quantized tensor, used in zero-1 optimizer""" + if isinstance(quantized_tensor, Float8Tensor): + assert hasattr(quantized_tensor, "_data"), "Float8Tensor does not have _data attribute" + assert quantized_tensor._data.dtype == torch.uint8, "Float8Tensor _data must be uint8" + return quantized_tensor._data + elif isinstance(quantized_tensor, Float8BlockwiseQTensor): + assert hasattr( + quantized_tensor, "_rowwise_data" + ), "Float8BlockwiseQTensor does not have _rowwise_data attribute" + assert ( + quantized_tensor._rowwise_data.dtype == torch.uint8 + ), "Float8BlockwiseQTensor _rowwise_data must be uint8" + return quantized_tensor._rowwise_data + else: + raise ValueError(f"Unsupported quantized tensor type: {type(quantized_tensor)}") + + +class MiniOptimizer: + + def __init__(self, weights, lr, dp_group): + self.world_size = dist.get_world_size(dp_group) + + self.weights = weights + self.lr = lr + self.dp_group = dp_group + + master_weights = [] + for weight in self.weights: + master_weights.append(weight.detach().float()) + self.master_weights = master_weights + + def step(self): + for weight, master_weight in zip(self.weights, self.master_weights): + main_grad = weight.main_grad + + # Don't use all-reduce directly to explicitly control the reduce order. + # dist.all_reduce(main_grad, op=dist.ReduceOp.AVG, group=self.dp_group) + buffers = [torch.empty_like(main_grad) for _ in range(self.world_size)] + dist.all_gather(buffers, main_grad, group=self.dp_group) + for i in range(1, self.world_size): + buffers[0] += buffers[i] + main_grad.copy_(buffers[0]) + main_grad /= self.world_size + + master_weight -= main_grad * self.lr + weight.data.copy_(master_weight) + + +class MiniZero_1: + """A mini zero-1 optimizer implementation, just used for this test""" + + def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=False): + self.rank = dist.get_rank(dp_group) + self.world_size = dist.get_world_size(dp_group) + + self.weights = weights + self.lr = lr + self.dp_group = dp_group + self.manual_post_all_gather_processing = manual_post_all_gather_processing + + # [self.offsets[i], self.offsets[i+1]) is the range of weights[i] in the global buffer + self.offsets = [0] + for weight in self.weights: + self.offsets.append(self.offsets[-1] + weight.numel()) + + # Padding to avoid global buffer cannot be divided by world size, so the offsets[-1] may + # not be the end range of the last weight. + if self.offsets[-1] % self.world_size != 0: + self.offsets[-1] += self.world_size - self.offsets[-1] % self.world_size + + self.master_weights = [] + # The start offset of the master weight in the weight + self.start_offsets = [] + # The overlapping area of the weight and this rank's local buffer + self.overlapping_areas = [] + + # The start and end of this rank's local buffer in the global buffer + rank_start = self.offsets[-1] // self.world_size * self.rank + rank_end = rank_start + self.offsets[-1] // self.world_size + + for weight, offset in zip(self.weights, self.offsets[:-1]): + if offset >= rank_end or (offset + weight.numel()) <= rank_start: + # This weight is not in this rank's local buffer + master_weight = None + start_offset = None + overlapping_area = None + else: + overlapping_start = max(rank_start, offset) + overlapping_end = min(rank_end, offset + weight.numel()) + length = overlapping_end - overlapping_start + start_offset = overlapping_start - offset + if isinstance(weight, QuantizedTensor): + # If weight is a FP8 tensor, we need to use the original high precision version + # to initialize the master weight. + high_precision_init_val = weight.get_high_precision_init_val().view(-1) + master_weight = high_precision_init_val.to(weight.device).float()[ + start_offset : start_offset + length + ] + else: + master_weight = ( + weight.detach().view(-1).float()[start_offset : start_offset + length] + ) + overlapping_area = (overlapping_start, overlapping_end) + self.master_weights.append(master_weight) + self.start_offsets.append(start_offset) + self.overlapping_areas.append(overlapping_area) + + # Create global buffer for grads reduce-scatter + self.grad_buffer = torch.empty( + [self.offsets[-1]], dtype=torch.float32, device=weights[0].device + ) + self.grad_buffer_slice = self.grad_buffer[rank_start:rank_end] + + # Create global buffer for weights all-gather + if isinstance(self.weights[0], QuantizedTensor): + weight_buffer_dtype = torch.uint8 + else: + weight_buffer_dtype = weights[0].dtype + self.weight_buffer = torch.empty( + [self.offsets[-1]], dtype=weight_buffer_dtype, device=weights[0].device + ) + self.weight_buffer_slice = self.weight_buffer[rank_start:rank_end] + + def step(self): + # ----------------------------------------------------------------------------------------- + # Step 1: Copy grads to the grad buffer + # ----------------------------------------------------------------------------------------- + for weight, offset in zip(self.weights, self.offsets[:-1]): + start = offset + end = offset + weight.numel() + self.grad_buffer[start:end].copy_(weight.main_grad.view(-1)) + + # ----------------------------------------------------------------------------------------- + # Step 2: Grads reduce-scatter + # ----------------------------------------------------------------------------------------- + # Don't use reduce_scatter directly to explicitly control the reduce order. + # dist.reduce_scatter_tensor(self.grad_buffer_slice, self.grad_buffer, op=dist.ReduceOp.AVG, + # group=self.dp_group) + buffers = [torch.empty_like(self.grad_buffer) for _ in range(self.world_size)] + dist.all_gather(buffers, self.grad_buffer, group=self.dp_group) + for i in range(1, self.world_size): + buffers[0] += buffers[i] + rank_start = self.offsets[-1] // self.world_size * self.rank + rank_end = rank_start + self.offsets[-1] // self.world_size + self.grad_buffer_slice.copy_(buffers[0][rank_start:rank_end]) + self.grad_buffer_slice /= self.world_size + + # ----------------------------------------------------------------------------------------- + # Step 3: Update master weights + # ----------------------------------------------------------------------------------------- + for master_weight, overlapping_area in zip(self.master_weights, self.overlapping_areas): + if master_weight is None: + # This weight's master weight is in other rank. + continue + grad = self.grad_buffer[overlapping_area[0] : overlapping_area[1]] + master_weight -= grad * self.lr + + # ----------------------------------------------------------------------------------------- + # Step 4: Cast master weights to BF16 or FP8, depending on the type of the weight + # ----------------------------------------------------------------------------------------- + if isinstance(self.weights[0], QuantizedTensor): + # FP8 weights case + for i in range(1, len(self.weights)): + assert isinstance(self.weights[i], QuantizedTensor) + cast_master_weights_to_fp8( + self.weights, + self.master_weights, + self.start_offsets, + self.dp_group, + manual_post_all_gather_processing=self.manual_post_all_gather_processing, + ) + else: + # BF16 weights case + for weight, master_weight, start_offset in zip( + self.weights, self.master_weights, self.start_offsets + ): + if master_weight is None: + continue + start = start_offset + end = start_offset + master_weight.numel() + weight.data.view(-1)[start:end].copy_(master_weight) + + # ----------------------------------------------------------------------------------------- + # Step 5: Copy the updated weights (not all weights) to the weight buffer + # ----------------------------------------------------------------------------------------- + for i in range(len(self.weights)): + master_weight = self.master_weights[i] + if master_weight is None: + continue + start_offset = self.start_offsets[i] + if isinstance(self.weights[i], QuantizedTensor): + weight = _get_raw_data(self.weights[i]) + else: + weight = self.weights[i] + weight_slice = weight.view(-1)[start_offset : start_offset + master_weight.numel()] + overlapping_start, overlapping_end = self.overlapping_areas[i] + self.weight_buffer[overlapping_start:overlapping_end].copy_(weight_slice) + + # ----------------------------------------------------------------------------------------- + # Step 6: Weight all-gather (FP8 or BF16) + # ----------------------------------------------------------------------------------------- + dist.all_gather_into_tensor( + self.weight_buffer, self.weight_buffer_slice, group=self.dp_group + ) + + # ----------------------------------------------------------------------------------------- + # Step 7: Copy the gathered weights from weight buffer to the actual weights + # ----------------------------------------------------------------------------------------- + for weight, offset in zip(self.weights, self.offsets[:-1]): + start = offset + end = offset + weight.numel() + if isinstance(weight, QuantizedTensor): + weight = _get_raw_data(weight) + weight.view(-1).data.copy_(self.weight_buffer[start:end]) + + if self.manual_post_all_gather_processing: + quantized_weights = [ + weight for weight in self.weights if isinstance(weight, QuantizedTensor) + ] + post_all_gather_processing(quantized_weights) + + +class MiniFSDP: + def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=False): + rank = dist.get_rank(dp_group) + world_size = dist.get_world_size(dp_group) + + self.weights = weights + self.lr = lr + self.dp_group = dp_group + self.manual_post_all_gather_processing = manual_post_all_gather_processing + + # Flatten the weights and pad to align with world size + if isinstance(weights[0], QuantizedTensor): + raw_data_list = [_get_raw_data(w).view(-1) for w in weights] + else: + raw_data_list = [w.view(-1) for w in weights] + self.flatten_weight, original_length = self._flatten_tensors_with_pad(raw_data_list) + + # Split flattened weights into shards + self.local_weight_shard = torch.chunk(self.flatten_weight, world_size)[rank] + self.local_main_grad_shard = torch.zeros_like( + self.local_weight_shard, dtype=torch.float32, device="cuda" + ) + shard_size = self.flatten_weight.size(0) // world_size + + # Map original tensors to flattened indices + tensor_indices = [] + cumulative_length = 0 + for tensor in raw_data_list: + length = tensor.size(0) + tensor_indices.append((cumulative_length, cumulative_length + length)) + cumulative_length += length + + # Build shard index mappings + self.weight_indices = [] + self.shard_indices = [] + for idx, (start, end) in enumerate(tensor_indices): + shard_start = rank * shard_size + shard_end = shard_start + shard_size + adjusted_end = min(shard_end, original_length) + + if start <= adjusted_end and end >= shard_start: + start_idx = max(start, shard_start) + end_idx = min(end, adjusted_end) + self.weight_indices.append((start_idx - start, end_idx - start)) + self.shard_indices.append((start_idx - shard_start, end_idx - shard_start)) + else: + self.weight_indices.append((None, None)) + self.shard_indices.append((None, None)) + + if isinstance(weights[idx], QuantizedTensor): + replace_raw_data( + weights[idx], self.flatten_weight[start:end].view(weights[idx].shape) + ) + else: + weights[idx].data = self.flatten_weight[start:end].view(weights[idx].shape) + + # Initialize local model weights and high-precision master weights + self.local_weights = [] + self.master_weights = [] + for i, weight in enumerate(self.weights): + weight_start, weight_end = self.weight_indices[i] + shard_start, shard_end = self.shard_indices[i] + if shard_start is not None and shard_end is not None: + local_weight_shard = self.local_weight_shard[shard_start:shard_end] + self.local_weights.append(local_weight_shard) + + if isinstance(weight, QuantizedTensor): + high_precision_init_val = weight.get_high_precision_init_val().view(-1) + master_weight_shard = high_precision_init_val.to(weight.device).float()[ + weight_start:weight_end + ] + else: + master_weight_shard = weight.detach().view(-1).float()[weight_start:weight_end] + self.master_weights.append(master_weight_shard) + else: + self.local_weights.append(None) + self.master_weights.append(None) + setattr( + weight, "main_grad", torch.zeros_like(weight, dtype=torch.float32, device="cuda") + ) + + def _flatten_tensors_with_pad(self, tensors): + """ + Flatten the list of tensors and pad them to align with the world size. + + Args: + tensors (list): List of tensors to flatten. + + Returns: + tuple: Flattened tensor and its original length before padding. + """ + world_size = dist.get_world_size(self.dp_group) + + flatten_tensor = torch.cat(tensors) + original_length = flatten_tensor.size(0) + + padding_needed = (world_size - original_length % world_size) % world_size + if padding_needed > 0: + zeros = torch.zeros(padding_needed, dtype=flatten_tensor.dtype, device="cuda") + flatten_tensor = torch.cat([flatten_tensor, zeros]) + + return flatten_tensor, original_length + + def zero_grad(self): + for weight in self.weights: + weight.grad = None + weight.main_grad.zero_() + + def step(self): + """ + Perform an optimization step for the distributed sharded model. + + This method includes: + 1. Gradient reduce-scatter: Synchronize gradients across all processes. + 2. Master weight update: Update high-precision master weights using local gradients. + 3. Precision casting: Cast updated master weights to FP8 or BF16 precision. + 4. Weight synchronization: All-gather updated weights across all processes. + + Returns: + None + """ + # Step 1: Reduce-scatter the gradients + main_grad_buffer, _ = self._flatten_tensors_with_pad( + [weight.main_grad.view(-1) for weight in self.weights] + ) + dist.reduce_scatter_tensor( + self.local_main_grad_shard, main_grad_buffer, group=self.dp_group + ) + self.local_main_grad_shard /= dist.get_world_size(self.dp_group) + + # Step 2: Update the master weights + for weight, master_weight, (shard_start, shard_end) in zip( + self.weights, self.master_weights, self.shard_indices + ): + if master_weight is None: + continue + + # Extract the local gradient shard for this weight + grad = self.local_main_grad_shard[shard_start:shard_end] + + # Update the master weight using gradient descent + master_weight -= grad * self.lr + + # Step 3: Cast master weights to FP8 or BF16 precision + if isinstance(self.weights[0], QuantizedTensor): + local_weights = [] + for local_weight in self.local_weights: + if local_weight is None: + local_weights.append(None) + continue + + local_weights.append(local_weight) + + cast_master_weights_to_fp8( + self.weights, + self.master_weights, + [idx[0] for idx in self.weight_indices], + self.dp_group, + local_weights, + manual_post_all_gather_processing=self.manual_post_all_gather_processing, + ) + else: + for weight, master_weight in zip(self.local_weights, self.master_weights): + if master_weight is None: + continue -TEST_ROOT = Path(__file__).parent.resolve() -NUM_PROCS: int = min(2, torch.cuda.device_count()) -LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] + # Copy updated master weights to local weights + weight.data.copy_(master_weight) + + # Step 4: All-gather updated weights across processes + dist.all_gather_into_tensor( + self.flatten_weight, self.local_weight_shard, group=self.dp_group + ) + + if self.manual_post_all_gather_processing: + quantized_weights = [ + weight for weight in self.weights if isinstance(weight, QuantizedTensor) + ] + post_all_gather_processing(quantized_weights) + + +def _test_mini_optimizer(dp_group): + """Make sure the implementation of MiniZero_1 and MiniFSDP is correct""" + rank = dist.get_rank(dp_group) + world_size = dist.get_world_size(dp_group) + + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + + weights = [ + torch.randn(256 * 256, dtype=torch.bfloat16, device="cuda"), + torch.randn(256 * 256 * 3, dtype=torch.bfloat16, device="cuda"), + torch.randn(256 * 256 * 2 - 1, dtype=torch.bfloat16, device="cuda"), + ] + + weights_1 = weights + weights_2 = [weight.clone() for weight in weights] + weights_3 = [weight.clone() for weight in weights] + + lr = 1.0 + optimizer_1 = MiniZero_1(weights_1, lr, dp_group) + optimizer_2 = MiniOptimizer(weights_2, lr, dp_group) + optimizer_3 = MiniFSDP(weights_3, lr, dp_group) + + for _ in range(100): + for w1, w2, w3 in zip(weights_1, weights_2, weights_3): + main_grads = [ + torch.randn_like(w1, dtype=torch.float32, device="cuda") for _ in range(world_size) + ] + # Choose based on rank to make sure the grads of different ranks are different. + main_grad = main_grads[rank] + w1.main_grad = main_grad + w2.main_grad = main_grad + w3.main_grad = main_grad + + optimizer_1.step() + optimizer_2.step() + optimizer_3.step() + + for w1, w2 in zip(weights_1, weights_2): + torch.testing.assert_close(w1, w2, atol=0, rtol=0) + for w1, w3 in zip(weights_1, weights_3): + torch.testing.assert_close(w1, w3, atol=0, rtol=0) + + +def _test_cast_master_weights_to_fp8(quantization, dp_group, manual_post_all_gather_processing): + rank = dist.get_rank(dp_group) + world_size = dist.get_world_size(dp_group) + + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + + mock_groups = [dist.new_group(ranks=[i]) for i in range(world_size)] + mock_group = mock_groups[rank] + + linear_kwargs = {"params_dtype": torch.bfloat16, "bias": False, "fuse_wgrad_accumulation": True} + + # Create model with FP8 weights + with te.quantized_model_init( + enabled=quantization is not None, + recipe=_get_quantization_recipe(quantization), + preserve_high_precision_init_val=True, + ): + model_fp8 = nn.Sequential( + te.Linear(128, 256 + 16, **linear_kwargs), + te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(256 * 3, 128, **linear_kwargs), + ) + + # Create model with BF16 weights + model = nn.Sequential( + te.Linear(128, 256 + 16, **linear_kwargs), + te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(256 * 3, 128, **linear_kwargs), + ) + + # Make sure the BF16 model and FP8 model have the same initial weights + for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): + high_precision_init_val = w_fp8.get_high_precision_init_val() + w.data.copy_(high_precision_init_val) + + # Allocate main_grads for each weight + for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): + w_fp8.main_grad = torch.zeros_like(w_fp8, dtype=torch.float32, device="cuda") + w.main_grad = torch.zeros_like(w, dtype=torch.float32, device="cuda") + + optimizer_fp8 = MiniZero_1( + [w for w in model_fp8.parameters()], 10.0, dp_group, manual_post_all_gather_processing + ) + optimizer = MiniZero_1([w for w in model.parameters()], 10.0, dp_group) + + for i in range(100): + for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): + w_fp8.main_grad.zero_() + w.main_grad.zero_() + + inputs = [ + torch.randn(16, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) + ] + # Choose based on rank to make sure the inputs of different ranks are different. + x = inputs[rank] + + with te.autocast( + enabled=quantization is not None, + recipe=_get_quantization_recipe(quantization), + amax_reduction_group=mock_group, + ): + y_fp8 = model_fp8(x) + + with te.autocast( + enabled=quantization is not None, + recipe=_get_quantization_recipe(quantization), + amax_reduction_group=mock_group, + ): + y = model(x) + + targets = [torch.randn_like(y) for _ in range(world_size)] + # Choose based on rank to make sure the targets of different ranks are different. + target = targets[rank] + loss_fp8 = nn.MSELoss()(y_fp8, target) + loss = nn.MSELoss()(y, target) + + loss_fp8.backward() + loss.backward() + + optimizer_fp8.step() + optimizer.step() + + torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) + + +def _test_fsdp_cast_master_weights_to_fp8( + quantization, dp_group, manual_post_all_gather_processing +): + rank = dist.get_rank(dp_group) + world_size = dist.get_world_size(dp_group) + + # Configuration constants + NUM_STEPS = 100 + SEED = 12345 + + torch.manual_seed(SEED) + torch.cuda.manual_seed(SEED) + + mock_groups = [dist.new_group(ranks=[i]) for i in range(world_size)] + mock_group = mock_groups[rank] + + linear_kwargs = { + "params_dtype": torch.bfloat16, + "bias": False, + "fuse_wgrad_accumulation": True, + } + + # Create model with FP8 weights + with te.quantized_model_init( + enabled=quantization is not None, + recipe=_get_quantization_recipe(quantization), + preserve_high_precision_init_val=True, + ): + model_fp8 = nn.Sequential( + te.Linear(128, 256 + 16, **linear_kwargs), + te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(256 * 3, 128, **linear_kwargs), + ) + + # Create model with BF16 weights + model = nn.Sequential( + te.Linear(128, 256 + 16, **linear_kwargs), + te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(256 * 3, 128, **linear_kwargs), + ) + + # Make sure the BF16 model and FP8 model have the same initial weights + for w_fp8, w in zip(model_fp8.parameters(), model.parameters()): + high_precision_init_val = w_fp8.get_high_precision_init_val() + w.data.copy_(high_precision_init_val) + + optimizer_fp8 = MiniFSDP( + [w for w in model_fp8.parameters()], 10.0, dp_group, manual_post_all_gather_processing + ) + optimizer = MiniFSDP([w for w in model.parameters()], 10.0, dp_group) + + for _ in range(100): + optimizer_fp8.zero_grad() + optimizer.zero_grad() + + inputs = [ + torch.randn(16, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) + ] + # Choose based on rank to make sure the inputs of different ranks are different. + x = inputs[rank] + + with te.autocast( + enabled=quantization is not None, + recipe=_get_quantization_recipe(quantization), + amax_reduction_group=mock_group, + ): + y_fp8 = model_fp8(x) + + with te.autocast( + enabled=quantization is not None, + recipe=_get_quantization_recipe(quantization), + amax_reduction_group=mock_group, + ): + y = model(x) + + targets = [torch.randn_like(y) for _ in range(world_size)] + # Choose based on rank to make sure the targets of different ranks are different. + target = targets[rank] + loss_fp8 = nn.MSELoss()(y_fp8, target) + loss = nn.MSELoss()(y, target) + + loss_fp8.backward() + loss.backward() + + optimizer_fp8.step() + optimizer.step() + + torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) + + +def run_parallel_tests() -> None: + """Run parallel tests""" + + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + + assert WORLD_SIZE == LOCAL_SIZE # this test supports only 1 node + assert LOCAL_SIZE <= torch.cuda.device_count() + dist_init_kwargs = { + "backend": "nccl", + "rank": WORLD_RANK, + "world_size": WORLD_SIZE, + "timeout": datetime.timedelta(seconds=30), + } + dist_init_kwargs["init_method"] = "env://" + dist_init_kwargs["device_id"] = torch.device(f"cuda:{LOCAL_RANK}") + assert dist.is_nccl_available() + torch.cuda.set_device(LOCAL_RANK) + dist.init_process_group(**dist_init_kwargs) + dp_group = dist.new_group(backend="nccl") + + quantizations = [] + if is_fp8_available(): + quantizations.extend(["fp8", "fp8_cs"]) + if is_fp8_block_scaling_available(): + quantizations.append("fp8_block") + + manual_post_all_gather_processings = [False, True] + + _test_mini_optimizer(dp_group) + + for quantization in quantizations: + for post_ag_processing in manual_post_all_gather_processings: + _test_cast_master_weights_to_fp8(quantization, dp_group, post_ag_processing) + _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group, post_ag_processing) + + dist.destroy_process_group() + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2, reason="cast_master_weights_to_fp8 test needs at least 2 GPUs." +) +@pytest.mark.parametrize("world_size", [2]) +def test_cast_master_weights_to_fp8(world_size: int) -> None: + """Launch parallel job that runs parallel tests""" + python_exe = pathlib.Path(sys.executable).resolve() + current_file = pathlib.Path(__file__).resolve() + command = [ + python_exe, + "-m", + "torch.distributed.run", + f"--nproc_per_node={world_size}", + current_file, + "--parallel", + ] + result = subprocess.run( + command, + check=True, + ) -def _run_test(quantization): - test_path = TEST_ROOT / "run_cast_master_weights_to_fp8.py" - test_cmd = LAUNCH_CMD + [str(test_path)] + ["--quantization", quantization] - result = subprocess.run(test_cmd, env=os.environ, check=False) - assert result.returncode == 0 +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--parallel", action="store_true", help="Run parallel tests") + args = parser.parse_args() + if args.parallel: + run_parallel_tests() -@pytest.mark.parametrize("quantization", ["fp8", "fp8_cs", "fp8_block"]) -def test_cast_master_weights_to_fp8(quantization): - if quantization in ("fp8", "fp8_cs") and not fp8_available: - pytest.skip(reason_for_no_fp8) - if quantization == "fp8_block" and not fp8_block_scaling_available: - pytest.skip(reason_for_no_fp8_block_scaling) - _run_test(quantization) +if __name__ == "__main__": + main() diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 8354823b32..20aba6c2bf 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -48,7 +48,12 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): def cast_master_weights_to_fp8( - model_weights, master_weights, start_offsets, group, fsdp_shard_model_weights=None + model_weights, + master_weights, + start_offsets, + group, + fsdp_shard_model_weights=None, + manual_post_all_gather_processing=False, ): r"""Helper function to cast master weights to FP8 primary weights. @@ -69,6 +74,11 @@ def cast_master_weights_to_fp8( fsdp_shard_model_weights : list of FSDP shard model weights. If None, it means that the model weights are not sharded. Otherwise, it means that the model weights are sharded and we get target model weights data storage using the FSDP shard model weights. + manual_post_all_gather_processing: bool, default = `False`. + If False, post processing will be automatically triggered during next forward. + If True, the timing of calling post_all_gather_processing is left to the user. + Note that users must call `post_all_gather_processing` if it's set to True, + otherwise the weights won't be updated correctly. """ @@ -129,21 +139,18 @@ def cast_master_weights_to_fp8( f"cast_master_weights_to_fp8 for {type(quantizer)} is not supported yet" ) + extra_args = [group, use_fsdp_shard_model_weights, manual_post_all_gather_processing] if len(delayed_scaling_params) > 0: - _cast_master_weights_to_fp8_delayed_scaling( - delayed_scaling_params, group, use_fsdp_shard_model_weights - ) + _cast_master_weights_to_fp8_delayed_scaling(delayed_scaling_params, *extra_args) if len(current_scaling_params) > 0: - _cast_master_weights_to_fp8_current_scaling( - current_scaling_params, group, use_fsdp_shard_model_weights - ) + _cast_master_weights_to_fp8_current_scaling(current_scaling_params, *extra_args) if len(blockwise_scaling_params) > 0: - _cast_master_weights_to_fp8_blockwise_scaling( - blockwise_scaling_params, group, use_fsdp_shard_model_weights - ) + _cast_master_weights_to_fp8_blockwise_scaling(blockwise_scaling_params, *extra_args) -def _cast_master_weights_to_fp8_delayed_scaling(params, group, use_fsdp_shard_model_weights=False): +def _cast_master_weights_to_fp8_delayed_scaling( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): r"""Helper function to cast master weights to FP8 primary weights for delayed scaling. Parameters @@ -160,6 +167,13 @@ def _cast_master_weights_to_fp8_delayed_scaling(params, group, use_fsdp_shard_mo amaxes, scales, scale_invs = [], [], [] for model_weight, master_weight, start_offset, shard_model_weight_raw in params: + if not manual_post_all_gather_processing: + # Reset transpose cache for all model weights. + # We cannot create transpose cache here because users (like megatron) may want to + # overlap the all-gather of model weights and forward process, so the model weight is + # not updated currently. + model_weight._reset_caches() + quantizer = model_weight._get_quantizer() amaxes.append(quantizer.amax.view(1)) @@ -219,7 +233,9 @@ def _cast_master_weights_to_fp8_delayed_scaling(params, group, use_fsdp_shard_mo ) -def _cast_master_weights_to_fp8_current_scaling(params, group, use_fsdp_shard_model_weights=False): +def _cast_master_weights_to_fp8_current_scaling( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): r"""Helper function to cast master weights to FP8 primary weights for current scaling. Parameters @@ -297,6 +313,13 @@ def _cast_master_weights_to_fp8_current_scaling(params, group, use_fsdp_shard_mo for (model_weight, master_weight, start_offset, model_weight_fragment), scale in zip( params, scales ): + if not manual_post_all_gather_processing: + # Reset transpose cache for all model weights. + # We cannot create transpose cache here because users (like megatron) may want to + # overlap the all-gather of model weights and forward process, so the model weight is + # not updated currently. + model_weight._reset_caches() + # If master weight is None, it means that the master weight of the current model weight # is in other DP ranks. if master_weight is None: @@ -322,7 +345,7 @@ def _cast_master_weights_to_fp8_current_scaling(params, group, use_fsdp_shard_mo def _cast_master_weights_to_fp8_blockwise_scaling( - params, group, use_fsdp_shard_model_weights=False + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False ): r"""Helper function to cast master weights to FP8 primary weights for blockwise scaling. @@ -421,6 +444,13 @@ def _cast_master_weights_to_fp8_blockwise_scaling( for (model_weight, master_weight, start_offset, model_weight_fragment), scale in zip( params, scales ): + if not manual_post_all_gather_processing: + # Clear columnwise data for all model weights. + # We cannot create columnwise data here because users (like megatron) may want to + # overlap the all-gather of model weights and forward process, so the model weight is + # not updated at this moment. + model_weight.update_usage(rowwise_usage=True, columnwise_usage=False) + # If master weight is None, it means that the master weight of the current model weight # is in other DP ranks. if master_weight is None: From 4ff3eed10acfa0ef88de3f80e6ab6349f9604523 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 6 Nov 2025 16:46:40 -0800 Subject: [PATCH 046/521] [JAX] Add test to check jaxpr that amax is reused for nvfp4 recipe (#2348) * Add test to check jaxpr that amax is reused for nvfp4 recipe Signed-off-by: Jeremy Berchtold * Move test to test_helper.py and rename file Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_custom_call_compute.py | 1 - ...lper.py => test_recipe_characteristics.py} | 67 ++++++++++++++++++- 2 files changed, 66 insertions(+), 2 deletions(-) rename tests/jax/{test_helper.py => test_recipe_characteristics.py} (78%) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index cecdb31218..3d4f179abd 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -45,7 +45,6 @@ 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.common import recipe GEMM_CASES = [ (256, 256, 512), diff --git a/tests/jax/test_helper.py b/tests/jax/test_recipe_characteristics.py similarity index 78% rename from tests/jax/test_helper.py rename to tests/jax/test_recipe_characteristics.py index fc88b7ef77..33fde7e231 100644 --- a/tests/jax/test_helper.py +++ b/tests/jax/test_recipe_characteristics.py @@ -11,7 +11,7 @@ import numpy as np from flax import linen as nn -from utils import assert_allclose +from utils import assert_allclose, pytest_parametrize_wrapper from transformer_engine.common.recipe import ( DelayedScaling, MXFP8BlockScaling, @@ -22,6 +22,7 @@ from transformer_engine.jax import autocast from transformer_engine.jax.quantize import ( get_quantize_config, + get_supported_quantization_recipes, is_scaling_mode_supported, ScalingMode, update_collections, @@ -32,11 +33,15 @@ from transformer_engine.jax.quantize.helper import _format2dtypes from transformer_engine.jax.sharding import MeshResource, global_mesh_resource from transformer_engine.jax.flax.module import TransformerEngineBase +from transformer_engine.jax import flax as te_flax +import transformer_engine.jax as te is_fp8_supported, reason = is_scaling_mode_supported(ScalingMode.DELAYED_TENSOR_SCALING) is_mxfp8_supported, mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING) is_nvfp4_supported, nvfp4_reason = is_scaling_mode_supported(ScalingMode.NVFP4_1D_SCALING) +SUPPORTED_RECIPES = get_supported_quantization_recipes() + def quantizer_check_vjp(outer_quantizer_set, assertion_func, x): """Check that the quantizers in the quantizer set are as expected and reconstructed correctly from flattened pytree representations across VJP boundaries.""" @@ -253,3 +258,63 @@ def test_autocast_nvfp4_block_scaling(self): self._compare_nvfp4_scaling_quantizers(bs) self._check_default_state() + + +class TestJaxprAndHlo: + """Tests to verify Jaxpr and/or HLO of compiled modules apply expected recipe functionality and optimizations.""" + + @pytest_parametrize_wrapper( + "quantization_recipe", + [ + quantization_recipe + for quantization_recipe in SUPPORTED_RECIPES + if isinstance(quantization_recipe, NVFP4BlockScaling) + ], + ) + def test_layernorm_mlp_reuses_amax_nvfp4(self, quantization_recipe): + """Tests that layernorm_mlp reuses the amax computed in layernorm and the activation and does not recompute it during quantizaton.""" + + with te.autocast(enabled=True, recipe=quantization_recipe, mesh_resource=te.MeshResource()): + model = te_flax.LayerNormMLP( + layernorm_type="rmsnorm", + return_layernorm_output=False, + intermediate_dropout_rate=0.0, + dtype=jnp.bfloat16, + ) + + var_collect = model.init( + jax.random.PRNGKey(0), + jnp.ones((128, 128), dtype=jnp.bfloat16), + ) + + def loss_fn(x, rngs): + return jnp.mean(model.apply(var_collect, x, rngs=rngs)[0]) + + x = jax.random.normal(jax.random.PRNGKey(0), (128, 128), dtype=jnp.bfloat16) + rngs = {"sr_rng": jax.random.PRNGKey(1), "dropout": jax.random.PRNGKey(2)} + jaxpr = jax.make_jaxpr(jax.value_and_grad(loss_fn))(x, rngs=rngs) + + rht_amax_eqns = [ + eqn for eqn in jaxpr.jaxpr.eqns if eqn.primitive.name == "te_rht_amax_ffi_wrapper" + ] + + assert len(rht_amax_eqns) == 4, f"Expected 4 rht_amax_eqns, got {len(rht_amax_eqns)}" + + def assert_param(index, tensor_name, expected_value: bool): + if expected_value: + assert rht_amax_eqns[index].params["produce_regular_amax"] == True, ( + f"Expected produce_regular_amax for {tensor_name} to be True, indicating no" + " reuse of amax as this tensor does not have a previous operation to fuse" + " with" + ) + else: + assert rht_amax_eqns[index].params["produce_regular_amax"] == False, ( + f"Expected produce_regular_amax for {tensor_name} to be False, indicating" + " reuse of amax" + ) + + assert_param(0, "fwd ln+q", False) + assert_param(1, "fwd act+q", False) + # No previous op before incoming dgrad in the backward so amax is not reused + assert_param(2, "bwd dgrad", True) + assert_param(3, "bwd dact+q", False) From f62cad90b1ed100b39bc01b2c6556b4033811714 Mon Sep 17 00:00:00 2001 From: Michael Goldfarb Date: Thu, 6 Nov 2025 20:04:40 -0600 Subject: [PATCH 047/521] Fix sharding of segment position to match id in ring attention. (#2349) --- .../jax/cpp_extensions/attention.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index c0cb6cda1f..6a21480d8d 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -1784,6 +1784,9 @@ def partition(config, mesh, arg_infos, result_infos): ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[4] = seed_sharding + # Ensure segment_pos gets same sharding as ID. + arg_shardings[-1] = arg_shardings[-3] + arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) @@ -1991,7 +1994,13 @@ def partition(config, mesh, arg_infos, result_infos): dk_sharding = NamedSharding(mesh, PartitionSpec(*k_spec)) dv_sharding = NamedSharding(mesh, PartitionSpec(*v_spec)) dbias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) - arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + + arg_shardings = [arg_i.sharding for arg_i in arg_infos] + # Ensure segment_pos gets same sharding as ID. + arg_shardings[-1] = arg_shardings[-3] + arg_shardings[-2] = arg_shardings[-4] + arg_shardings = tuple(arg_shardings) + out_shardings = (dq_sharding, dk_sharding, dv_sharding, dbias_sharding) helper = _FusedAttnCPWithP2PHelper(mesh, config) @@ -2265,6 +2274,9 @@ def partition(config, mesh, arg_infos, result_infos): ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[4] = seed_sharding + # Ensure segment_pos gets same sharding as ID. + arg_shardings[-1] = arg_shardings[-3] + arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) @@ -2403,7 +2415,11 @@ def partition(config, mesh, arg_infos, result_infos): if not is_context_parallel: return FusedAttnBwdPrimitive.partition(config, mesh, arg_infos, result_infos) - arg_shardings = tuple(arg.sharding for arg in arg_infos) + arg_shardings = [arg_i.sharding for arg_i in arg_infos] + # Ensure segment_pos gets same sharding as ID. + arg_shardings[-1] = arg_shardings[-3] + arg_shardings[-2] = arg_shardings[-4] + arg_shardings = tuple(arg_shardings) # dq, dk, dv, dbias sharding = q, k, v, bias sharding out_shardings = tuple(arg.sharding for arg in arg_infos[:4]) From 26aad6b0faae88f4865fe6ace357b0f56485267e Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 7 Nov 2025 12:00:31 -0500 Subject: [PATCH 048/521] Disable cuDNN attention for known IMA and NaNs (#2344) * Fix cuDNN backend selection for more case. Add CG as a option as well Signed-off-by: Kirthi Shankar Sivamani * fix logic Signed-off-by: Kirthi Shankar Sivamani * Fix cuDNN checks Signed-off-by: Kirthi Shankar Sivamani * Add more checks Signed-off-by: Kirthi Shankar Sivamani * Fix cuddn version Signed-off-by: Kirthi Shankar Sivamani * Fix error message Signed-off-by: Kirthi Shankar Sivamani * Add check for window size Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- .../common/fused_attn/fused_attn.cpp | 109 +++++++++++------- .../include/transformer_engine/fused_attn.h | 56 +++++---- .../jax/csrc/extensions/attention.cpp | 54 ++++----- .../dot_product_attention/backends.py | 5 + .../dot_product_attention/context_parallel.py | 8 ++ .../dot_product_attention.py | 1 + .../attention/dot_product_attention/utils.py | 5 + .../pytorch/cpp_extensions/fused_attn.py | 8 ++ transformer_engine/pytorch/csrc/extensions.h | 6 +- .../pytorch/csrc/extensions/attention.cpp | 42 +++---- 10 files changed, 178 insertions(+), 116 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index f6ee37d4c5..9c6e9b33d5 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -138,7 +138,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit) { + int64_t window_size_right, bool return_max_logit, bool cuda_graph) { using namespace transformer_engine; NVTE_Fused_Attn_Backend backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; const int device_id = cuda::current_device(); @@ -166,7 +166,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( 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: {bshd, sbhd}, any seqlen, d=128, {no_mask, causal} + // 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 || @@ -407,6 +407,28 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( " Please upgrade your cuDNN version if possible." << std::endl; } + if ((cudnn_runtime_version == 91400) && (max_seqlen_kv > 1024) && (window_size_left != -1) && + (attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_MASK) && + (attn_mask_type != NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK)) { + backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + std::cout << "Warning: Given combination of attention mask (non-causal) and " + "max_seqlen_kv (> 1024) does not support fused attention for cuDNN 9.14.0. " + " Please upgrade your cuDNN version if possible." + << std::endl; + } + if ((cudnn_runtime_version <= 91500) && is_training && + (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && + (max_seqlen_kv % 128 != 0) && cuda_graph && + (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)) { + backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + std::cout << "Warning: Given combination of attention mask (non-padding)," + " max_seqlen_kv (not divisible by 128), and qkv_format (BSHD/SBHD) for" + " backward fused attention with graph capture requires cuDNN 9.15.1+. " + "Please upgrade your cuDNN version if possible." + << std::endl; + } } else { backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; } @@ -419,11 +441,11 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, bool is_training, bool return_max_logit, - 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, NVTETensor workspace, - cudaStream_t stream) { + bool cuda_graph, 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, + NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd_qkvpacked); using namespace transformer_engine; @@ -460,7 +482,8 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h, h, max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, return_max_logit); + h, h, max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, return_max_logit, + cuda_graph); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -496,16 +519,14 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, } } // NVTE fused attention BWD with packed QKV -void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, - const NVTETensor S, NVTETensor dP, - const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, - NVTETensor dBias, NVTETensor dSoftmaxOffset, - const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, - size_t max_seqlen, 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 deterministic, NVTETensor workspace, cudaStream_t stream) { +void nvte_fused_attn_bwd_qkvpacked( + const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, + NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, + NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, + size_t max_seqlen, 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 deterministic, bool cuda_graph, + NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_bwd_qkvpacked); using namespace transformer_engine; @@ -544,7 +565,7 @@ void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, con NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h, h, - max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, false); + max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, false, cuda_graph); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -602,10 +623,10 @@ void nvte_fused_attn_fwd_kvpacked( const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, 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, 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, - NVTETensor workspace, cudaStream_t stream) { + 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd_kvpacked); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -681,7 +702,7 @@ void nvte_fused_attn_fwd_kvpacked( NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, - return_max_logit); + return_max_logit, cuda_graph); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -728,7 +749,8 @@ void nvte_fused_attn_bwd_kvpacked( 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 deterministic, NVTETensor workspace, cudaStream_t stream) { + int64_t window_size_right, bool deterministic, bool cuda_graph, NVTETensor workspace, + cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_bwd_kvpacked); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -776,9 +798,10 @@ void nvte_fused_attn_bwd_kvpacked( const NVTEDType Q_type = static_cast(input_Q->data.dtype); const NVTEDType KV_type = static_cast(input_KV->data.dtype); - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, false); + NVTE_Fused_Attn_Backend fused_attention_backend = + nvte_get_fused_attn_backend(true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, + softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, + d, window_size_left, window_size_right, false, cuda_graph); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -833,16 +856,19 @@ void nvte_fused_attn_bwd_kvpacked( } } // NVTE fused attention FWD with separate Q, K and V -void nvte_fused_attn_fwd( - const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, - const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, - const NVTETensor page_table_k, 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, - 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, NVTETensor workspace, cudaStream_t stream) { +void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, + const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, + NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, + const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, + 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, NVTETensor workspace, + cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -913,7 +939,7 @@ void nvte_fused_attn_fwd( NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, - return_max_logit); + return_max_logit, cuda_graph); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -963,7 +989,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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 deterministic, - NVTETensor workspace, cudaStream_t stream) { + 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); @@ -1008,7 +1034,8 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false); + h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false, + cuda_graph); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 518fad20de..40e6a0b4bc 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -207,13 +207,14 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit); + int64_t window_size_right, bool return_max_logit, bool cuda_graph); /*! \brief Compute dot product attention with packed QKV input. * @@ -257,6 +258,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * it may be >= max(seqlen_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensor's layout. @@ -273,11 +275,11 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, bool is_training, bool return_max_logit, - 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, NVTETensor workspace, - cudaStream_t stream); + bool cuda_graph, 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, + NVTETensor workspace, cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with packed QKV input. * @@ -324,19 +326,18 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). * \param[in] deterministic Whether to execute with deterministic behaviours. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ -void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, - const NVTETensor S, NVTETensor dP, - const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, - NVTETensor dBias, NVTETensor dSoftmaxOffset, - const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, - size_t max_seqlen, 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 deterministic, NVTETensor workspace, cudaStream_t stream); +void nvte_fused_attn_bwd_qkvpacked( + const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, + NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, + NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, + size_t max_seqlen, 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 deterministic, bool cuda_graph, + NVTETensor workspace, cudaStream_t stream); /*! \brief Compute dot product attention with packed KV input. * @@ -387,6 +388,7 @@ void nvte_fused_attn_bwd_qkvpacked(const NVTETensor QKV, const NVTETensor O, con * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensor's layout. @@ -405,10 +407,10 @@ void nvte_fused_attn_fwd_kvpacked( const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, 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, 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, - NVTETensor workspace, cudaStream_t stream); + 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, NVTETensor workspace, cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with packed KV input. * @@ -461,6 +463,7 @@ void nvte_fused_attn_fwd_kvpacked( * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). * \param[in] deterministic Whether to execute with deterministic behaviours. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ @@ -472,7 +475,8 @@ void nvte_fused_attn_bwd_kvpacked( 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 deterministic, NVTETensor workspace, cudaStream_t stream); + int64_t window_size_right, bool deterministic, bool cuda_graph, NVTETensor workspace, + cudaStream_t stream); /*! \brief Compute dot product attention with separate Q, K and V. * @@ -527,6 +531,7 @@ void nvte_fused_attn_bwd_kvpacked( * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensors' layout. @@ -545,9 +550,9 @@ void nvte_fused_attn_fwd( const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, 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, - 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, NVTETensor workspace, cudaStream_t stream); + bool cuda_graph, 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, NVTETensor workspace, cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * @@ -605,6 +610,7 @@ void nvte_fused_attn_fwd( * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). * \param[in] deterministic Whether to execute with deterministic behaviours. + * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ @@ -619,7 +625,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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 deterministic, - NVTETensor workspace, cudaStream_t stream); + bool cuda_graph, NVTETensor workspace, cudaStream_t stream); /*! \brief Update the RNG state with the seed and calculated offset. * diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index ffc0706fe7..a99f4fae90 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -23,7 +23,7 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend(bool is_training, DType q_dtype, DTy is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false); + false, false); return backend; } @@ -180,7 +180,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( qkv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), dummy_rng_state_tensor.data(), q_max_seqlen, is_training, - false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, + false, false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, query_workspace_tensor.data(), nullptr); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { @@ -189,7 +189,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, + 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, query_workspace_tensor.data(), nullptr); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { @@ -199,7 +199,7 @@ 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, scaling_factor, dropout_probability, qkv_layout, + 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, query_workspace_tensor.data(), nullptr); } else { @@ -279,7 +279,7 @@ static void FusedAttnForwardImpl( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false); + false, false); nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); /* Auxiliary tensors (to be propagated to the backward pass later) */ @@ -298,7 +298,7 @@ static void FusedAttnForwardImpl( qkv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(), rng_state_tensor.data(), q_max_seqlen, is_training, false, - scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, + false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, workspace_tensor.data(), stream); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; @@ -311,7 +311,7 @@ static void FusedAttnForwardImpl( s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, scaling_factor, dropout_probability, + 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, workspace_tensor.data(), stream); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { @@ -326,9 +326,9 @@ static void FusedAttnForwardImpl( dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, scaling_factor, - dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, workspace_tensor.data(), stream); + 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, workspace_tensor.data(), stream); } else { NVTE_ERROR("Unsupported qkv_layout."); } @@ -480,7 +480,7 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( dummy_d_softmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), dummy_ragged_offset_tensor.data(), q_max_seqlen, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, - deterministic, query_workspace_tensor.data(), nullptr); + deterministic, false, query_workspace_tensor.data(), nullptr); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { nvte_fused_attn_bwd_kvpacked( q_tensor.data(), kv_tensor.data(), output_tensor.data(), doutput_tensor.data(), @@ -491,19 +491,19 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( 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, deterministic, query_workspace_tensor.data(), nullptr); + window_size_right, deterministic, false, query_workspace_tensor.data(), nullptr); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { - 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, deterministic, 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, bias_type, mask_type, softmax_type, window_size_left, + window_size_right, deterministic, false, query_workspace_tensor.data(), nullptr); } else { NVTE_ERROR("Unsupported qkv_layout."); } @@ -546,7 +546,7 @@ static void FusedAttnBackwardImpl( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false); + false, false); PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, softmax_aux, rng_state, bias); @@ -568,7 +568,7 @@ static void FusedAttnBackwardImpl( q_seq_offsets_tensor.data(), q_max_seqlen, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, deterministic, - workspace_tensor.data(), stream); + false, workspace_tensor.data(), stream); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; auto kv_shape = @@ -590,7 +590,7 @@ static void FusedAttnBackwardImpl( dummy_d_softmax_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, deterministic, + mask_type, softmax_type, window_size_left, window_size_right, deterministic, false, workspace_tensor.data(), stream); } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; @@ -617,7 +617,7 @@ static void FusedAttnBackwardImpl( 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, deterministic, - workspace_tensor.data(), stream); + false, workspace_tensor.data(), stream); } else { NVTE_ERROR("Unsupported qkv_layout."); } diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index d4903be902..147a85fc2f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -66,6 +66,7 @@ ) from transformer_engine.pytorch import export from transformer_engine.pytorch.export import is_in_onnx_export_mode +from transformer_engine.pytorch.graph import is_graph_capturing # Global vars for flash attn v2 and v3 imports flash_attn_cuda_bwd = None @@ -1199,6 +1200,7 @@ def forward( window_size, rng_gen, softmax_offset, + cuda_graph=is_graph_capturing(), ) # out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16 @@ -1276,6 +1278,7 @@ def forward( rng_gen, softmax_offset, return_max_logit, + is_graph_capturing(), ) out = out_ out_ret = out_ @@ -1515,6 +1518,7 @@ def backward(ctx, d_out, *_args): ctx.softmax_type, ctx.window_size, ctx.deterministic, + is_graph_capturing(), ) # dq, dk, dv: torch.Tensor; dtype = torch.float16 or torch.bfloat16 @@ -1579,6 +1583,7 @@ def backward(ctx, d_out, *_args): ctx.softmax_type, ctx.window_size, ctx.deterministic, + is_graph_capturing(), ) d_bias = None 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 f312cac798..00d609ab9e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -23,6 +23,7 @@ from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor 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, @@ -33,6 +34,7 @@ gather_along_first_dim, reduce_scatter_along_first_dim, ) + from transformer_engine.pytorch.quantized_tensor import ( prepare_for_saving, restore_from_saved, @@ -715,6 +717,7 @@ def cp_p2p_fwd_fused_attn( cu_seqlens_kv_padded=cu_seqlens_kv_padded_, **fp8_meta_kwargs, return_max_logit=return_max_logit, + cuda_graph=is_graph_capturing(), ) if fp8: @@ -977,6 +980,7 @@ def cp_p2p_bwd_fused_attn( attn_mask_type=attn_mask_type_, attn_bias_type=attn_bias_type, deterministic=deterministic, + cuda_graph=is_graph_capturing(), **fp8_meta_kwargs, ) @@ -2772,6 +2776,7 @@ def forward( cu_seqlens_kv_padded=cu_seqlens_kv_per_step[i], window_size=window_size_per_step[i], return_max_logit=return_max_logit, + cuda_graph=is_graph_capturing(), ) if return_max_logit: max_logit_per_step[i] = max_logit_[0] @@ -2986,6 +2991,7 @@ def backward(ctx, dout, *_args): attn_bias_type=ctx.attn_bias_type, window_size=window_size_per_step[i], deterministic=ctx.deterministic, + cuda_graph=is_graph_capturing(), ) else: dq_per_step[i], dk_per_step[i], dv_per_step[i] = [ @@ -3282,6 +3288,7 @@ def forward( softmax_type=softmax_type, softmax_offset=softmax_offset, return_max_logit=return_max_logit, + cuda_graph=is_graph_capturing(), ) if isinstance(out_, Float8Tensor): out_fp8 = out_ @@ -3559,6 +3566,7 @@ def backward(ctx, dout, *_args): attn_bias_type=ctx.attn_bias_type, window_size=ctx.window_size, deterministic=ctx.deterministic, + cuda_graph=is_graph_capturing(), **fp8_meta_kwargs, softmax_type=ctx.softmax_type, ) 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 0d1c0b0c05..4278820e7a 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 @@ -1314,6 +1314,7 @@ def forward( inference_params=inference_params, softmax_type=self.softmax_type, return_max_logit=self.return_max_logit, + cuda_graph=is_graph_capturing(), ) global _attention_backends if is_in_onnx_export_mode(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7d4a4f86d9..a08ba14196 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -231,6 +231,8 @@ class AttentionParams: The type of softmax operation. See DotProductAttention for details. return_max_logit: bool, default = `False` Whether to output max_logit. + cuda_graph: bool, default = `False` + Whether support for cuda graph capture is needed or not. """ qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor @@ -260,6 +262,7 @@ class AttentionParams: inference_params: Optional[InferenceParams] = None softmax_type: str = "vanilla" return_max_logit: bool = False + cuda_graph: bool = False def __eq__(self, other): """ @@ -334,6 +337,7 @@ def get_attention_backend( inference_params = attention_params.inference_params softmax_type = attention_params.softmax_type return_max_logit = attention_params.return_max_logit + cuda_graph = attention_params.cuda_graph # Run config logger = logging.getLogger("DotProductAttention") @@ -979,6 +983,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt window_size[0], window_size[1], return_max_logit, + cuda_graph, ) if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug("Disabling FusedAttention as no backend supports the provided input") diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index eb43c75f6b..e55ea2a54a 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -140,6 +140,7 @@ def fused_attn_fwd( rng_gen: torch.Generator = None, softmax_offset: torch.Tensor = None, return_max_logit: bool = False, + cuda_graph: bool = False, ) -> Tuple[Union[torch.Tensor, None], ...]: """Fused Attention FWD for separate QKV input. @@ -219,6 +220,8 @@ def fused_attn_fwd( See softmax_type in DotProductAttention for details. return_max_logit: bool, default = False whether to return the maximum attention score + cuda_graph: bool, default = False + whether or not cuda graph capture is enabled. Returns ---------- @@ -320,6 +323,7 @@ def fused_attn_fwd( rng_gen, rng_elts_per_thread, return_max_logit, + cuda_graph, ) if return_max_logit: @@ -367,6 +371,7 @@ def fused_attn_bwd( softmax_type: str = "vanilla", window_size: Tuple[int, int] = (-1, -1), deterministic: bool = False, + cuda_graph: bool = False, ) -> Tuple[Union[torch.Tensor, None], ...]: """Fused Attention BWD for packed KV input. @@ -439,6 +444,8 @@ def fused_attn_bwd( window and causal mask specifically. deterministic: bool, default = False whether to execute the backward pass with deterministic behaviours. + cuda_graph: bool, default = False + whether or not cuda graph capture is enabled. Returns ---------- @@ -509,6 +516,7 @@ def fused_attn_bwd( s_quantizer, dp_quantizer, dqkv_quantizer, + cuda_graph, ) return output_tensors diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 79fb798422..43eab96544 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -76,7 +76,7 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float p_dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit); + int64_t window_size_right, bool return_max_logit, bool cuda_graph); std::pair quantizer_helper(py::handle quantizer, const std::vector &shape, DType dtype, @@ -94,7 +94,7 @@ std::vector fused_attn_fwd( const std::optional page_table_k, const std::optional page_table_v, py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread, bool return_max_logit); + size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph); std::vector fused_attn_bwd( size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, @@ -106,7 +106,7 @@ std::vector fused_attn_bwd( const std::vector Aux_CTX_Tensors, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, - py::handle dp_quantizer, py::handle dqkv_quantizer); + py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph); at::Tensor fa_prepare_fwd(at::Tensor qkvi); at::Tensor fa_prepare_bwd(at::Tensor q, at::Tensor k, at::Tensor v); diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index f66c8aa619..d51aef4065 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -45,12 +45,12 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float p_dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit) { + int64_t window_size_right, bool return_max_logit, bool cuda_graph) { NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, attn_mask_type, softmax_type, p_dropout, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right, - return_max_logit); + return_max_logit, cuda_graph); return fused_attention_backend; } @@ -107,7 +107,7 @@ std::vector fused_attn_fwd( const std::optional page_table_k, const std::optional page_table_v, py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, const std::optional SoftmaxOffset, const std::optional rng_gen, - size_t rng_elts_per_thread, bool return_max_logit) { + size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph) { auto none = py::none(); // create QKV tensor wrappers @@ -229,7 +229,7 @@ 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, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, + return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size[0], window_size[1], workspace.data(), at::cuda::getCurrentCUDAStream()); }); @@ -289,7 +289,7 @@ 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, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, + return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size[0], window_size[1], workspace.data(), at::cuda::getCurrentCUDAStream()); }); @@ -312,7 +312,7 @@ std::vector fused_attn_bwd( const std::vector Aux_CTX_Tensors, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, - py::handle dp_quantizer, py::handle dqkv_quantizer) { + py::handle dp_quantizer, py::handle dqkv_quantizer, bool cuda_graph) { auto none = py::none(); // create QKV, O, dO tensor wrappers @@ -527,13 +527,14 @@ std::vector fused_attn_bwd( // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( - te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), - &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], - window_size[1], deterministic, workspace.data(), at::cuda::getCurrentCUDAStream()); + nvte_fused_attn_bwd(te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), + te_S.data(), te_dP.data(), &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], window_size[1], deterministic, cuda_graph, + workspace.data(), at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace @@ -543,13 +544,14 @@ std::vector fused_attn_bwd( // execute kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd( - te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), - &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], - window_size[1], deterministic, workspace.data(), at::cuda::getCurrentCUDAStream()); + nvte_fused_attn_bwd(te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), + te_S.data(), te_dP.data(), &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], window_size[1], deterministic, cuda_graph, + workspace.data(), at::cuda::getCurrentCUDAStream()); }); // destroy tensor wrappers From 5978f1d7544fe0e4b06a036600c80c6a8c8203fe Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Fri, 7 Nov 2025 11:17:16 -0800 Subject: [PATCH 049/521] [JAX] Default to fused attention in JAX DPA (#2363) * Default to fused attention in JAX DPA Signed-off-by: Kshitij Lakhani * Consolidate documentation for DPA in JAX Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> * Correctly update the documentation for defaults in JAX DPA Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> --------- Signed-off-by: Kshitij Lakhani Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- transformer_engine/jax/flax/transformer.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 42c9451245..86af6cf499 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -407,10 +407,10 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods Users can select between these two backends via the :attr:`NVTE_FUSED_ATTN` environment variable: - * Set :attr:`NVTE_FUSED_ATTN=0` for unfused attention (default). - * Set :attr:`NVTE_FUSED_ATTN=1` for fused attention. If the required cuDNN fused attention - kernel is not available on the system, a warning will be issued, and the module will - automatically fall back to the unfused backend. + * Set :attr:`NVTE_FUSED_ATTN=0` for unfused attention. + * Set :attr:`NVTE_FUSED_ATTN=1` for fused attention (default). If the required cuDNN fused + attention kernel is not available on the system, a warning will be issued, and the module + will automatically fall back to the unfused backend. .. note:: The DotProductAttention default setting enables non-deterministic kernels for reduced @@ -602,7 +602,8 @@ def __call__( else: assert bias is not None - enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "0")) + # Use fused attn (if kernel check below passes) by default + enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) sequence_dim = 0 if self.transpose_batch_sequence else 1 seqlen_q = query.shape[sequence_dim] From d20311bd5ae6bae3cd83368736d482c25f52a1a3 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 7 Nov 2025 14:31:51 -0500 Subject: [PATCH 050/521] Update cudnn frontend to v1.16.0 (#2362) Signed-off-by: Kirthi Shankar Sivamani --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 0b1577c8c8..be6c079be8 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 0b1577c8c83401237d601d0d0db5210506705396 +Subproject commit be6c079be8aaffa0fc079fcf039887e637c289c7 From 3454f84da64cc117fde3c7c5286041d698576dea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 7 Nov 2025 23:08:22 +0100 Subject: [PATCH 051/521] [common] Remove kvpacked and qkvpacked attention functions for every kernel type. (#2287) * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * depracted compile time warning + \warning -> \deprecated Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski 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> Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../common/fused_attn/fused_attn.cpp | 333 +++++++++-- .../fused_attn_f16_arbitrary_seqlen.cu | 530 +----------------- .../fused_attn_f16_arbitrary_seqlen.h | 47 -- .../fused_attn_f16_max512_seqlen.cu | 264 --------- .../fused_attn/fused_attn_f16_max512_seqlen.h | 37 -- .../common/fused_attn/fused_attn_fp8.cu | 418 -------------- .../common/fused_attn/fused_attn_fp8.h | 41 -- .../include/transformer_engine/fused_attn.h | 20 + 8 files changed, 302 insertions(+), 1388 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 9c6e9b33d5..ac6fefdc6a 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -15,6 +15,74 @@ #include "fused_attn_fp8.h" #include "utils.h" +namespace { +// Helper function to create a tensor view with modified shape and optional pointer offset +transformer_engine::Tensor make_tensor_view(const transformer_engine::Tensor *source, + const std::vector &shape, + size_t offset_bytes = 0) { + transformer_engine::Tensor view = *source; + if (offset_bytes > 0) { + view.data.dptr = static_cast(static_cast(source->data.dptr) + offset_bytes); + } + view.data.shape = shape; + view.nvte_tensor = 0; // Mark as unmanaged/local tensor view + return view; +} + +// Helper function to calculate stride for packed QKV tensor unpacking +size_t calculate_qkv_stride(NVTE_QKV_Layout_Group layout_group, transformer_engine::DType dtype, + size_t h, size_t d) { + size_t stride = 0; + if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { + stride = (transformer_engine::typeToNumBits(dtype) * h * d) / 8; + } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { + stride = (transformer_engine::typeToNumBits(dtype) * d) / 8; + } + return stride; +} + +// Helper function to determine unpacked shape for QKV packed tensor +std::vector calculate_qkv_unpacked_shape(const transformer_engine::Tensor *qkv_tensor, + size_t h, size_t d) { + std::vector unpacked_shape; + if (qkv_tensor->data.shape.size() == 4) { + // T3HD or TH3D (4D) -> THD (3D): remove dimension "3" at position 1 + unpacked_shape = {qkv_tensor->data.shape[0], h, d}; + } else { + // BS3HD/SB3HD or BSH3D/SBH3D (5D) -> BSHD/SBHD (4D): remove dimension "3" at position 2 + unpacked_shape = {qkv_tensor->data.shape[0], qkv_tensor->data.shape[1], h, d}; + } + return unpacked_shape; +} + +// Helper function to calculate stride for packed KV tensor unpacking +size_t calculate_kv_stride(NVTE_QKV_Layout_Group layout_group, transformer_engine::DType dtype, + size_t h_kv, size_t d) { + size_t stride = 0; + if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { + stride = (transformer_engine::typeToNumBits(dtype) * h_kv * d) / 8; + } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { + stride = (transformer_engine::typeToNumBits(dtype) * d) / 8; + } + return stride; +} + +// Helper function to determine unpacked shape for KV packed tensor +std::vector calculate_kv_unpacked_shape(const transformer_engine::Tensor *kv_tensor, + NVTE_QKV_Layout_Group layout_group, + NVTE_QKV_Format kv_format, size_t t_kv, size_t h_kv, + size_t d) { + std::vector unpacked_kv_shape; + if (kv_format == NVTE_QKV_Format::NVTE_THD) { + unpacked_kv_shape = {t_kv, h_kv, d}; + } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD || + layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { + unpacked_kv_shape = {kv_tensor->data.shape[0], kv_tensor->data.shape[1], h_kv, d}; + } + return unpacked_kv_shape; +} +} // namespace + // map NVTE_QKV_Layout to NVTE_QKV_Layout_Group NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { switch (qkv_layout) { @@ -436,6 +504,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( } // NVTE fused attention FWD with packed QKV +// DEPRECATED: This API is deprecated. +// Please use nvte_fused_attn_fwd with separate Q, K, V tensors instead. void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, @@ -487,30 +557,62 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) - fused_attn_max_512_fwd_qkvpacked(b, h, max_seqlen, d, is_training, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, input_QKV, input_Bias, - output_O, Aux_CTX_Tensors, input_cu_seqlens, input_rng_state, - wkspace, stream, handle); + // Unpack QKV and call the non-packed function + const auto QKV_type = input_QKV->data.dtype; + size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); + std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); + + // Create tensor views for Q, K, V + Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); + Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); + Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); + + fused_attn_max_512_fwd(b, h, max_seqlen, max_seqlen, d, is_training, attn_scale, dropout, + qkv_layout, bias_type, attn_mask_type, &Q_view, &K_view, &V_view, + input_Bias, output_O, Aux_CTX_Tensors, input_cu_seqlens, + input_cu_seqlens, 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_qkvpacked( - b, h, max_seqlen, d, t, is_training, return_max_logit, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, input_QKV, - input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, - input_cu_seqlens_padded, input_rng_state, wkspace, stream, handle); + // Unpack QKV and call the non-packed function + const auto QKV_type = input_QKV->data.dtype; + size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); + std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); + + // Create tensor views for Q, K, V + Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); + Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); + Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); + + fused_attn_arbitrary_seqlen_fwd( + b, h, h, max_seqlen, max_seqlen, d, d, t, t, 0, 0, 0, 0, 0, 0, is_training, + return_max_logit, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, &Q_view, &K_view, &V_view, input_Bias, + input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, input_cu_seqlens, + input_cu_seqlens_padded, input_cu_seqlens_padded, nullptr, nullptr, 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_qkvpacked(b, h, max_seqlen, d, is_training, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, input_QKV, input_output_S, output_O, - Aux_CTX_Tensors, input_cu_seqlens, input_rng_state, wkspace, - stream, handle); + // Unpack QKV and call the non-packed function + const auto QKV_type = input_QKV->data.dtype; + size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); + std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); + + // Create tensor views for Q, K, V + Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); + Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); + Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); + + fused_attn_fp8_fwd(b, h, h, max_seqlen, max_seqlen, d, is_training, attn_scale, dropout, + qkv_layout, bias_type, attn_mask_type, &Q_view, &K_view, &V_view, + input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens, + input_cu_seqlens, input_rng_state, wkspace, stream, handle); #else NVTE_ERROR("cuDNN 8.9.0 is required for FP8 fused attention. \n"); #endif @@ -519,6 +621,8 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, } } // NVTE fused attention BWD with packed QKV +// DEPRECATED: This API is deprecated. +// Please use nvte_fused_attn_bwd with separate Q, K, V tensors instead. void nvte_fused_attn_bwd_qkvpacked( const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, @@ -570,9 +674,25 @@ void nvte_fused_attn_bwd_qkvpacked( 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_qkvpacked( - b, h, max_seqlen, d, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, input_QKV, - input_dO, output_S, output_dQKV, output_dBias, input_cu_seqlens, wkspace, stream, handle); + + // Unpack QKV and dQKV and call the non-packed function + const auto QKV_type = input_QKV->data.dtype; + size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); + std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); + + // Create tensor views for Q, K, V and dQ, dK, dV + Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); + Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); + Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); + + Tensor dQ_view = make_tensor_view(output_dQKV, unpacked_shape); + Tensor dK_view = make_tensor_view(output_dQKV, unpacked_shape, stride); + Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); + + fused_attn_max_512_bwd(b, h, max_seqlen, max_seqlen, d, attn_scale, dropout, qkv_layout, + bias_type, attn_mask_type, &Q_view, &K_view, &V_view, input_dO, output_S, + &dQ_view, &dK_view, &dV_view, output_dBias, input_cu_seqlens, + input_cu_seqlens, wkspace, stream, handle); #else NVTE_ERROR("cuDNN 8.9.1 is required for BF16/FP16 fused attention with max_seqlen<=512. \n"); #endif @@ -588,12 +708,27 @@ void nvte_fused_attn_bwd_qkvpacked( if (softmax_type != NVTE_VANILLA_SOFTMAX) { input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_arbitrary_seqlen_bwd_qkvpacked( - b, h, max_seqlen, d, t, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size_left, window_size_right, deterministic, input_QKV, input_O, - input_dO, input_Bias, input_SoftmaxOffset, output_S, output_dQKV, output_dBias, - output_dSoftmaxOffset, input_cu_seqlens, input_cu_seqlens_padded, input_rng_state, wkspace, - stream, handle); + + // Unpack QKV and dQKV and call the non-packed function + const auto QKV_type = input_QKV->data.dtype; + size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); + std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); + + // Create tensor views for Q, K, V and dQ, dK, dV + Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); + Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); + Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); + + Tensor dQ_view = make_tensor_view(output_dQKV, unpacked_shape); + Tensor dK_view = make_tensor_view(output_dQKV, unpacked_shape, stride); + Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); + + fused_attn_arbitrary_seqlen_bwd( + b, h, h, max_seqlen, max_seqlen, d, d, t, t, attn_scale, dropout, qkv_layout, bias_type, + attn_mask_type, softmax_type, window_size_left, window_size_right, deterministic, &Q_view, + &K_view, &V_view, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, &dQ_view, + &dK_view, &dV_view, output_dBias, output_dSoftmaxOffset, input_cu_seqlens, input_cu_seqlens, + input_cu_seqlens_padded, input_cu_seqlens_padded, input_rng_state, wkspace, stream, handle); #else const char *err_msg = "cuDNN 8.9.0 is required for BF16/FP16 fused attention " @@ -605,10 +740,26 @@ void nvte_fused_attn_bwd_qkvpacked( 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_qkvpacked(b, h, max_seqlen, d, attn_scale, dropout, qkv_layout, bias_type, - attn_mask_type, input_QKV, input_O, input_dO, input_M, input_ZInv, - input_S, input_output_dP, output_dQKV, input_cu_seqlens, - input_rng_state, wkspace, stream, handle); + + // Unpack QKV and dQKV and call the non-packed function + const auto QKV_type = input_QKV->data.dtype; + size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); + std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); + + // Create tensor views for Q, K, V and dQ, dK, dV + Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); + Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); + Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); + + Tensor dQ_view = make_tensor_view(output_dQKV, unpacked_shape); + Tensor dK_view = make_tensor_view(output_dQKV, unpacked_shape, stride); + Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); + + fused_attn_fp8_bwd(b, h, h, max_seqlen, max_seqlen, d, attn_scale, dropout, qkv_layout, + bias_type, attn_mask_type, &Q_view, &K_view, &V_view, input_O, input_dO, + input_M, input_ZInv, input_S, input_output_dP, &dQ_view, &dK_view, &dV_view, + input_cu_seqlens, input_cu_seqlens, input_rng_state, wkspace, stream, + handle); #else NVTE_ERROR("cuDNN 8.9.0 is required for FP8 fused attention. \n"); #endif @@ -617,6 +768,8 @@ void nvte_fused_attn_bwd_qkvpacked( } } // NVTE fused attention FWD with packed KV +// DEPRECATED: This API is deprecated. +// Please use nvte_fused_attn_fwd with separate Q, K, V tensors instead. void nvte_fused_attn_fwd_kvpacked( const NVTETensor Q, const NVTETensor KV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, @@ -706,21 +859,40 @@ void nvte_fused_attn_fwd_kvpacked( if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) - fused_attn_max_512_fwd_kvpacked( - b, h_q, max_seqlen_q, max_seqlen_kv, d, is_training, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, input_Q, input_KV, input_Bias, output_O, Aux_CTX_Tensors, - input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); + // Unpack KV and call the non-packed function + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + size_t stride = calculate_kv_stride(layout_group, input_Q->data.dtype, h_kv, d); + std::vector unpacked_kv_shape = + calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); + + Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); + Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); + + fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d, is_training, attn_scale, dropout, + qkv_layout, bias_type, attn_mask_type, input_Q, &K_view, &V_view, + 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 >= 8903) - fused_attn_arbitrary_seqlen_fwd_kvpacked( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, t_q, t_kv, num_pages_k, num_pages_v, + // Unpack KV and call the non-packed function + const auto Q_type = input_Q->data.dtype; + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); + std::vector unpacked_kv_shape = + calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); + + Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); + Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); + + fused_attn_arbitrary_seqlen_fwd( + b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, 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, input_Q, input_KV, input_Bias, input_SoftmaxOffset, - output_O, Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, + window_size_left, window_size_right, input_Q, &K_view, &V_view, 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 @@ -729,10 +901,20 @@ void nvte_fused_attn_fwd_kvpacked( #endif } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { #if (CUDNN_VERSION >= 8900) - fused_attn_fp8_fwd_kvpacked( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, is_training, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, input_Q, input_KV, input_output_S, output_O, Aux_CTX_Tensors, - input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); + // Unpack KV and call the non-packed function + const auto Q_type = input_Q->data.dtype; + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); + std::vector unpacked_kv_shape = + calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); + + Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); + Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); + + fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, is_training, attn_scale, + dropout, qkv_layout, bias_type, attn_mask_type, input_Q, &K_view, &V_view, + 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 @@ -741,6 +923,8 @@ void nvte_fused_attn_fwd_kvpacked( } } // NVTE fused attention BWD with packed KV +// DEPRECATED: This API is deprecated. +// Please use nvte_fused_attn_bwd with separate Q, K, V tensors instead. void nvte_fused_attn_bwd_kvpacked( const NVTETensor Q, const NVTETensor KV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, @@ -806,10 +990,23 @@ void nvte_fused_attn_bwd_kvpacked( 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_kvpacked( - b, h_q, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, qkv_layout, bias_type, - attn_mask_type, input_Q, input_KV, input_dO, output_S, output_dQ, output_dKV, output_dBias, - input_cu_seqlens_q, input_cu_seqlens_kv, wkspace, stream, handle); + + // Unpack KV and dKV and call the non-packed function + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + size_t stride = calculate_kv_stride(layout_group, input_Q->data.dtype, h_kv, d); + std::vector unpacked_kv_shape = + calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); + + Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); + Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); + + Tensor dK_view = make_tensor_view(output_dKV, unpacked_kv_shape); + Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); + + fused_attn_max_512_bwd(b, h_q, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, qkv_layout, + bias_type, attn_mask_type, input_Q, &K_view, &V_view, input_dO, output_S, + output_dQ, &dK_view, &dV_view, 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 @@ -825,13 +1022,29 @@ void nvte_fused_attn_bwd_kvpacked( if (softmax_type != NVTE_VANILLA_SOFTMAX) { input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); } - fused_attn_arbitrary_seqlen_bwd_kvpacked( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, t_q, t_kv, attn_scale, dropout, qkv_layout, + + // Unpack KV and dKV and call the non-packed function + const auto Q_type = input_Q->data.dtype; + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); + std::vector unpacked_kv_shape = + calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); + + Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); + Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); + + // Create tensor views for dK, dV + Tensor dK_view = make_tensor_view(output_dKV, unpacked_kv_shape); + Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); + + fused_attn_arbitrary_seqlen_bwd( + b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, t_q, t_kv, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, deterministic, - input_Q, input_KV, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, output_dQ, - output_dKV, 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); + input_Q, &K_view, &V_view, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, + output_dQ, &dK_view, &dV_view, 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.3 is required for BF16/FP16 fused attention " @@ -843,11 +1056,25 @@ void nvte_fused_attn_bwd_kvpacked( 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_kvpacked(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, input_Q, input_KV, input_O, - input_dO, input_M, input_ZInv, input_S, input_output_dP, output_dQ, - output_dKV, input_cu_seqlens_q, input_cu_seqlens_kv, - input_rng_state, wkspace, stream, handle); + + // Unpack KV and dKV and call the non-packed function + const auto Q_type = input_Q->data.dtype; + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); + std::vector unpacked_kv_shape = + calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); + + Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); + Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); + + Tensor dK_view = make_tensor_view(output_dKV, unpacked_kv_shape); + Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); + + fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, + qkv_layout, bias_type, attn_mask_type, input_Q, &K_view, &V_view, input_O, + input_dO, input_M, input_ZInv, input_S, input_output_dP, output_dQ, &dK_view, + &dV_view, 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 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 950ced61bb..14468b543a 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 @@ -1037,532 +1037,6 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } // namespace fused_attn using namespace transformer_engine::fused_attn; -void fused_attn_arbitrary_seqlen_fwd_qkvpacked( - size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, - 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, - const Tensor *input_QKV, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, - const Tensor *cu_seqlens_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - const auto QKV_type = input_QKV->data.dtype; - void *devPtrQKV = input_QKV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - stride = (typeToNumBits(QKV_type) * num_attn_heads * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void *devPtrQ = static_cast(devPtrQKV); - void *devPtrK = static_cast(static_cast(devPtrQKV) + stride); - void *devPtrV = static_cast(static_cast(devPtrQKV) + 2 * stride); - - void *devPtrBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { - devPtrBias = input_Bias->data.dptr; - bias_b = input_Bias->data.shape[0]; - bias_h = input_Bias->data.shape[1]; - } - void *devPtrSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - } - - void *devPtrO = output_O->data.dptr; - void *devPtrS1 = nullptr; - void *devPtrS2 = nullptr; - void *devPtrCuSeqlens = cu_seqlens->data.dptr; - void *devPtrSeqOffsets = cu_seqlens_padded->data.dptr; - - size_t max_batch_size = 0; - size_t max_tokens = 0; - if (qkv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - max_tokens = get_max_tokens(num_tokens); - } - - size_t i = 0; - if (Aux_CTX_Tensors->size == 0) { - const auto cudnn_runtime_version = cudnnGetVersion(); - if (return_max_logit) { - Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_Max->data.dptr = nullptr; - if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_Max->data.shape = {max_tokens, num_attn_heads, 1}; - } else { - output_Max->data.shape = {batch, num_attn_heads, max_seqlen, 1}; - } - output_Max->data.dtype = DType::kFloat32; - Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_Sum_Exp->data.dptr = nullptr; - if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_Sum_Exp->data.shape = {max_tokens, num_attn_heads, 1}; - } else { - output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen, 1}; - } - output_Sum_Exp->data.dtype = DType::kFloat32; - } else { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if (qkv_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_S->data.shape = {max_tokens, num_attn_heads, 1}; - } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen, 1}; - } - output_S->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; - - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_bias->data.dptr = nullptr; - output_bias->data.shape = {bias_b, bias_h, max_seqlen, max_seqlen}; - output_bias->data.dtype = QKV_type; - } - - 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) { - if (return_max_logit) { - Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_Max->data.dptr; - Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS2 = output_Sum_Exp->data.dptr; - } else { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_S->data.dptr; - } - Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_rng_state->data.dptr = rng_state->data.dptr; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_bias->data.dptr = devPtrBias; - } - 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 *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_fwd_impl( - batch, num_attn_heads, num_attn_heads, max_seqlen, max_seqlen, head_dim, head_dim, - max_batch_size, max_tokens, max_tokens, 0, 0, 0, 0, 0, 0, bias_b, bias_h, is_training, - return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, devPtrQ, devPtrK, devPtrV, devPtrBias, - devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, devPtrDropoutSeed, devPtrDropoutOffset, - devPtrCuSeqlens, devPtrCuSeqlens, nullptr, nullptr, devPtrSeqOffsets, devPtrSeqOffsets, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, 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_arbitrary_seqlen_bwd_qkvpacked( - size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, - 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 deterministic, const Tensor *input_QKV, const Tensor *input_O, - const Tensor *input_dO, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_S, Tensor *output_dQKV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens, const Tensor *cu_seqlens_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - const auto QKV_type = input_QKV->data.dtype; - void *devPtrQKV = input_QKV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - stride = (typeToNumBits(QKV_type) * num_attn_heads * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void *devPtrQ = devPtrQKV; - void *devPtrK = static_cast(static_cast(devPtrQKV) + stride); - void *devPtrV = static_cast(static_cast(devPtrQKV) + 2 * stride); - - void *devPtrO = input_O->data.dptr; - void *devPtrdO = input_dO->data.dptr; - void *devPtrBias = nullptr; - void *devPtrdBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { - devPtrBias = input_Bias->data.dptr; - devPtrdBias = output_dBias->data.dptr; - bias_b = output_dBias->data.shape[0]; - bias_h = output_dBias->data.shape[1]; - } - - size_t max_batch_size = 0; - size_t max_tokens = 0; - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if (qkv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - max_tokens = get_max_tokens(num_tokens); - } - - void *devPtrdQKV = output_dQKV->data.dptr; - void *devPtrdQ = devPtrdQKV; - void *devPtrdK = static_cast(static_cast(devPtrdQKV) + stride); - void *devPtrdV = static_cast(static_cast(devPtrdQKV) + 2 * stride); - - void *devPtrSoftmaxStats = nullptr; - devPtrSoftmaxStats = output_S->data.dptr; - void *devPtrSoftmaxOffset = nullptr; - void *devPtrdSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - devPtrdSoftmaxOffset = output_dSoftmaxOffset->data.dptr; - } - - void *devPtrCuSeqlens = cu_seqlens->data.dptr; - void *devPtrSeqOffsets = cu_seqlens_padded->data.dptr; - - void *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_bwd_impl( - batch, num_attn_heads, num_attn_heads, max_seqlen, max_seqlen, head_dim, head_dim, - max_batch_size, max_tokens, max_tokens, bias_b, bias_h, attn_scale, p_dropout, qkv_layout, - bias_type, mask_type, softmax_type, window_size_left, window_size_right, deterministic, - devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, - devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, - devPtrDropoutOffset, devPtrCuSeqlens, devPtrCuSeqlens, devPtrSeqOffsets, devPtrSeqOffsets, - get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, 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_arbitrary_seqlen_fwd_kvpacked( - 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, 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, const Tensor *input_Q, const Tensor *input_KV, - 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; - void *devPtrQ = input_Q->data.dptr; - void *devPtrKV = input_KV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - stride = (typeToNumBits(QKV_type) * num_gqa_groups * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void *devPtrK = devPtrKV; - void *devPtrV = static_cast(static_cast(devPtrKV) + stride); - - void *devPtrBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { - devPtrBias = input_Bias->data.dptr; - bias_b = input_Bias->data.shape[0]; - bias_h = input_Bias->data.shape[1]; - } - void *devPtrSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - } - - void *devPtrO = output_O->data.dptr; - void *devPtrS1 = nullptr; - void *devPtrS2 = nullptr; - - void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; - void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; - void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void *devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - void *devPtrPageTableK = page_table_k->data.dptr; - void *devPtrPageTableV = page_table_v->data.dptr; - - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = get_max_tokens(num_tokens_kv); - } - - size_t i = 0; - if (Aux_CTX_Tensors->size == 0) { - const auto cudnn_runtime_version = cudnnGetVersion(); - if (return_max_logit) { - Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_Max->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_Max->data.shape = {max_tokens_q, num_attn_heads, 1}; - } else { - output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_Max->data.dtype = DType::kFloat32; - Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_Sum_Exp->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_Sum_Exp->data.shape = {max_tokens_q, num_attn_heads, 1}; - } else { - output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_Sum_Exp->data.dtype = DType::kFloat32; - } else { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_S->data.shape = {max_tokens_q, num_attn_heads, 1}; - } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_S->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; - - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_bias->data.dptr = nullptr; - output_bias->data.shape = {bias_b, bias_h, max_seqlen_q, max_seqlen_kv}; - output_bias->data.dtype = QKV_type; - } - - 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) { - if (return_max_logit) { - Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_Max->data.dptr; - Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS2 = output_Sum_Exp->data.dptr; - } else { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_S->data.dptr; - } - Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_rng_state->data.dptr = rng_state->data.dptr; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_bias->data.dptr = devPtrBias; - } - 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 *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_fwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, head_dim, - 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, is_training, - return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, 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) { - 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_arbitrary_seqlen_bwd_kvpacked( - 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, size_t num_tokens_q, size_t num_tokens_kv, - 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 deterministic, const Tensor *input_Q, const Tensor *input_KV, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, - const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dKV, - Tensor *output_dBias, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - const auto QKV_type = input_Q->data.dtype; - void *devPtrQ = input_Q->data.dptr; - void *devPtrKV = input_KV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - stride = (typeToNumBits(QKV_type) * num_gqa_groups * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void *devPtrK = devPtrKV; - void *devPtrV = static_cast(static_cast(devPtrKV) + stride); - - void *devPtrO = input_O->data.dptr; - void *devPtrdO = input_dO->data.dptr; - void *devPtrBias = nullptr; - void *devPtrdBias = nullptr; - size_t bias_b = 0; - size_t bias_h = 0; - if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { - devPtrBias = input_Bias->data.dptr; - devPtrdBias = output_dBias->data.dptr; - bias_b = output_dBias->data.shape[0]; - bias_h = output_dBias->data.shape[1]; - } - - size_t max_batch_size = 0; - size_t max_tokens_q = 0; - size_t max_tokens_kv = 0; - NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - if (q_format == NVTE_QKV_Format::NVTE_THD || kv_format == NVTE_QKV_Format::NVTE_THD) { - max_batch_size = get_max_batch_size(batch); - } - if (q_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_q = get_max_tokens(num_tokens_q); - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - max_tokens_kv = get_max_tokens(num_tokens_kv); - } - - void *devPtrdQ = output_dQ->data.dptr; - void *devPtrdKV = output_dKV->data.dptr; - void *devPtrdK = devPtrdKV; - void *devPtrdV = static_cast(static_cast(devPtrdKV) + stride); - - void *devPtrSoftmaxStats = nullptr; - devPtrSoftmaxStats = output_S->data.dptr; - void *devPtrSoftmaxOffset = nullptr; - void *devPtrdSoftmaxOffset = nullptr; - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; - devPtrdSoftmaxOffset = output_dSoftmaxOffset->data.dptr; - } - - void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; - void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; - void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; - void *devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - - void *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_arbitrary_seqlen_bwd_impl( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, head_dim, - max_batch_size, max_tokens_q, max_tokens_kv, bias_b, bias_h, attn_scale, p_dropout, - qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, - 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) { - 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_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, @@ -1604,8 +1078,8 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; void *devPtrSeqOffsetsKV = cu_seqlens_kv_padded->data.dptr; - void *devPtrPageTableK = page_table_k->data.dptr; - void *devPtrPageTableV = page_table_v->data.dptr; + void *devPtrPageTableK = page_table_k ? page_table_k->data.dptr : nullptr; + void *devPtrPageTableV = page_table_v ? page_table_v->data.dptr : nullptr; size_t max_batch_size = 0; size_t max_tokens_q = 0; 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 a3181c6295..872b798bb4 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 @@ -18,53 +18,6 @@ namespace transformer_engine { #if (CUDNN_VERSION >= 8900) -void fused_attn_arbitrary_seqlen_fwd_qkvpacked( - size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, - 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, - const Tensor *input_QKV, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, - const Tensor *cu_seqlens_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_arbitrary_seqlen_bwd_qkvpacked( - size_t batch, size_t num_attn_heads, size_t max_seqlen, size_t head_dim, size_t num_tokens, - 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 deterministic, const Tensor *input_QKV, const Tensor *input_O, - const Tensor *input_dO, const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, - Tensor *output_S, Tensor *output_dQKV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, - const Tensor *cu_seqlens, const Tensor *cu_seqlens_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_arbitrary_seqlen_fwd_kvpacked( - 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, 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, const Tensor *input_Q, const Tensor *input_KV, - 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_kvpacked( - 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, size_t num_tokens_q, size_t num_tokens_kv, - 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 deterministic, const Tensor *input_Q, const Tensor *input_KV, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, - const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dKV, - Tensor *output_dBias, Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); - 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, 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 89528fa3c4..1028df6452 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 @@ -1215,150 +1215,6 @@ void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv } // namespace fused_attn using namespace transformer_engine::fused_attn; -void fused_attn_max_512_fwd_qkvpacked( - size_t batch, size_t num_head, size_t 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_QKV, const Tensor *input_Bias, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - // QKV shape is [b, s, 3, h, d] - void *devPtrQKV = input_QKV->data.dptr; - const auto stride = 2 * num_head * head_dim; - - void *devPtrQ = static_cast(devPtrQKV); - void *devPtrK = static_cast(static_cast(devPtrQKV) + stride); - void *devPtrV = static_cast(static_cast(devPtrQKV) + 2 * stride); - - void *devPtrBias = static_cast(input_Bias->data.dptr); - - void *devPtrO = output_O->data.dptr; - - void *devPtrS = nullptr; - - 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, max_seqlen, max_seqlen}; - output_S->data.dtype = input_QKV->data.dtype; - } 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 *devPtrCuSeqlen = 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); - - const DType QKV_type = input_QKV->data.dtype; - size_t workspace_size = 0; - - fused_attn_max_512_fwd_impl( - batch, num_head, max_seqlen, max_seqlen, head_dim, is_training, attn_scale, p_dropout, - qkv_layout, bias_type, mask_type, devPtrQ, devPtrK, devPtrV, devPtrS, devPtrO, devPtrBias, - devPtrCuSeqlen, devPtrCuSeqlen, devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, - &workspace_size, get_cudnn_dtype(QKV_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_fwd_kvpacked(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_KV, - 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; - - NVTE_CHECK(bias_type == NVTE_Bias_Type::NVTE_NO_BIAS || - bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS, - "NVTE_PRE_SCALE_BIAS is not implemented in fused_attn_max_512."); - - // Q shape is [b, s, h, d] - void *devPtrQ = input_Q->data.dptr; - - // KV shape is [b, s, 2, h, d] - const auto stride = 2 * num_head * head_dim; - void *devPtrK = input_KV->data.dptr; - void *devPtrV = static_cast(static_cast(devPtrK) + stride); - - 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_KV->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_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, @@ -1429,126 +1285,6 @@ void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen, } } -void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t num_head, size_t 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_QKV, - const Tensor *input_dO, Tensor *output_S, Tensor *output_dQKV, - Tensor *output_dBias, const Tensor *cu_seqlens, - Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle) { - using namespace transformer_engine; - - // QKV shape is [b, s, 3, h, d] - void *devPtrQKV = input_QKV->data.dptr; - - auto stride = 2 * num_head * head_dim; - void *devPtrQ = devPtrQKV; - void *devPtrK = static_cast(static_cast(devPtrQKV) + stride); - void *devPtrV = static_cast(static_cast(devPtrQKV) + 2 * stride); - - void *devPtrdO = input_dO->data.dptr; - - // dQKV shape is [b, s, 3, h, d] - void *devPtrdQKV = output_dQKV->data.dptr; - void *devPtrdQ = devPtrdQKV; - void *devPtrdK = static_cast(static_cast(devPtrdQKV) + stride); - void *devPtrdV = static_cast(static_cast(devPtrdQKV) + 2 * stride); - - void *devPtrdBias = output_dBias->data.dptr; - - void *devPtrS = output_S->data.dptr; - - // devPtrdS reuses the memory of devPtrS - void *devPtrdS = devPtrS; - - void *devPtrCuSeqlens = cu_seqlens->data.dptr; - - const auto qkv_type = input_QKV->data.dtype; - size_t workspace_size = 0; - - fused_attn_max_512_bwd_impl(batch, num_head, max_seqlen, max_seqlen, head_dim, attn_scale, - p_dropout, qkv_layout, mask_type, bias_type, devPtrQ, devPtrK, - devPtrV, devPtrS, devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdS, - devPtrdBias, devPtrCuSeqlens, devPtrCuSeqlens, workspace->data.dptr, - &workspace_size, get_cudnn_dtype(qkv_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_kvpacked(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_KV, - const Tensor *input_dO, Tensor *output_S, Tensor *output_dQ, - Tensor *output_dKV, 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; - - // Q shape is [b, s, h, d] - // KV shape is [b, s, 2, h, d] - auto stride = 2 * num_head * head_dim; - void *devPtrQ = input_Q->data.dptr; - void *devPtrK = input_KV->data.dptr; - void *devPtrV = static_cast(static_cast(devPtrK) + stride); - - void *devPtrdO = input_dO->data.dptr; - - // dQ shape is [b, s, h, d] - // dKV shape is [b, s, 2, h, d] - void *devPtrdQ = output_dQ->data.dptr; - void *devPtrdK = output_dKV->data.dptr; - void *devPtrdV = static_cast(static_cast(devPtrdK) + stride); - - 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_KV->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."); - } -} 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, 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 171fe846ce..57b7afcf43 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 @@ -18,25 +18,6 @@ namespace transformer_engine { #if (CUDNN_VERSION >= 8901) -void fused_attn_max_512_fwd_qkvpacked(size_t batch, size_t num_head, size_t max_seqlen, - size_t head_size, 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_QKV, const Tensor *input_Bias, - Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, - const Tensor *cu_seqlens, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_max_512_fwd_kvpacked(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_KV, - 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_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, @@ -47,24 +28,6 @@ void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen, const Tensor *kv_cu_seqlens, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); -void fused_attn_max_512_bwd_qkvpacked(size_t batch, size_t num_head, size_t 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_QKV, - const Tensor *input_dO, Tensor *output_S, Tensor *output_dQKV, - Tensor *output_dBias, const Tensor *cu_seqlens, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_max_512_bwd_kvpacked(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_KV, - const Tensor *input_dO, Tensor *output_S, Tensor *output_dQ, - Tensor *output_dKV, Tensor *output_dBias, - const Tensor *q_cu_seqlens, const Tensor *kv_cu_seqlens, - 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, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 7b85be972c..5d806290a9 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -2407,424 +2407,6 @@ void fused_attn_fp8_bwd_impl_v1( } // namespace fused_attn #if (CUDNN_VERSION >= 8900) -// fused attention FWD FP8 with packed QKV -void fused_attn_fp8_fwd_qkvpacked(size_t batch, size_t num_attn_heads, size_t 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_QKV, Tensor* input_output_S, Tensor* output_O, - NVTETensorPack* Aux_CTX_Tensors, const Tensor* cu_seqlens, - const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, - cudnnHandle_t handle) { - using namespace transformer_engine; - const DType QKV_type = input_QKV->data.dtype; - const DType O_type = output_O->data.dtype; - void* devPtrQKV = input_QKV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - stride = (typeToNumBits(QKV_type) * num_attn_heads * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void* devPtrQ = static_cast(devPtrQKV); - void* devPtrK = static_cast(static_cast(devPtrQKV) + stride); - void* devPtrV = static_cast(static_cast(devPtrQKV) + 2 * stride); - void* devPtrDescaleQ = input_QKV->scale_inv.dptr; - void* devPtrDescaleK = input_QKV->scale_inv.dptr; - void* devPtrDescaleV = input_QKV->scale_inv.dptr; - - void* devPtrO = output_O->data.dptr; - void* devPtrAmaxO = output_O->amax.dptr; - void* devPtrScaleO = output_O->scale.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]); - output_M->data.dptr = nullptr; - output_M->data.shape = {batch, num_attn_heads, max_seqlen, 1}; - output_M->data.dtype = DType::kFloat32; - output_ZInv->data.dptr = nullptr; - output_ZInv->data.shape = {batch, num_attn_heads, max_seqlen, 1}; - output_ZInv->data.dtype = DType::kFloat32; - 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]); - devPtrM = output_M->data.dptr; - devPtrZInv = output_ZInv->data.dptr; - output_rng_state->data.dptr = rng_state->data.dptr; - } 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* devPtrcuSeqlens = - reinterpret_cast(reinterpret_cast(cu_seqlens->data.dptr)); - void* devPtrDropoutSeed = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); - void* devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - 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)) { - fused_attn::fused_attn_fp8_fwd_impl_v1( - batch, num_attn_heads, num_attn_heads, max_seqlen, max_seqlen, 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, devPtrcuSeqlens, devPtrcuSeqlens, - devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), - get_cudnn_fe_dtype(O_type), 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, max_seqlen, head_dim, is_training, attn_scale, p_dropout, - qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrDescaleQ, - devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, - devPtrAmaxS, devPtrcuSeqlens, devPtrcuSeqlens, 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"); - } - - 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; - } -} -// fused attention BWD FP8 with packed QKV -void fused_attn_fp8_bwd_qkvpacked( - size_t batch, size_t num_attn_heads, size_t 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_QKV, 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_dQKV, const Tensor* cu_seqlens, const Tensor* rng_state, Tensor* workspace, - cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - const DType QKV_type = input_QKV->data.dtype; - const DType dO_type = input_dO->data.dtype; - const DType dQKV_type = output_dQKV->data.dtype; - void* devPtrQKV = input_QKV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - stride = (typeToNumBits(QKV_type) * num_attn_heads * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void* devPtrQ = devPtrQKV; - void* devPtrK = static_cast(static_cast(devPtrQKV) + stride); - void* devPtrV = static_cast(static_cast(devPtrQKV) + 2 * stride); - void* devPtrDescaleQ = input_QKV->scale_inv.dptr; - void* devPtrDescaleK = input_QKV->scale_inv.dptr; - void* devPtrDescaleV = input_QKV->scale_inv.dptr; - - void* devPtrO = input_O->data.dptr; - const DType O_type = input_O->data.dtype; - void* devPtrDescaleO = nullptr; - if (O_type == DType::kFloat8E4M3 || O_type == DType::kFloat8E5M2) { - devPtrDescaleO = input_O->scale_inv.dptr; - } - void* devPtrdO = input_dO->data.dptr; - void* devPtrDescaledO = input_dO->scale_inv.dptr; - - void* devPtrM = input_M->data.dptr; - void* devPtrZInv = input_ZInv->data.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* devPtrdQKV = output_dQKV->data.dptr; - void* devPtrdQ = devPtrdQKV; - void* devPtrdK = static_cast(static_cast(devPtrdQKV) + stride); - void* devPtrdV = static_cast(static_cast(devPtrdQKV) + 2 * stride); - void* devPtrAmaxdQ = output_dQKV->amax.dptr; - void* devPtrAmaxdK = output_dQKV->amax.dptr; - void* devPtrAmaxdV = output_dQKV->amax.dptr; - void* devPtrScaledQ = output_dQKV->scale.dptr; - void* devPtrScaledK = output_dQKV->scale.dptr; - void* devPtrScaledV = output_dQKV->scale.dptr; - - void* devPtrcuSeqlens = - reinterpret_cast(reinterpret_cast(cu_seqlens->data.dptr)); - void* devPtrDropoutSeed = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); - void* devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - 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)) { - fused_attn::fused_attn_fp8_bwd_impl_v1( - batch, num_attn_heads, num_attn_heads, max_seqlen, max_seqlen, head_dim, attn_scale, - p_dropout, qkv_layout, bias_type, mask_type, 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, devPtrcuSeqlens, devPtrcuSeqlens, - 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) { - fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, max_seqlen, max_seqlen, head_dim, 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, devPtrcuSeqlens, - devPtrcuSeqlens, 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"); - } - - 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; - } -} -// fused attention FWD FP8 with packed KV -void fused_attn_fp8_fwd_kvpacked(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_KV, 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; - const DType QKV_type = input_Q->data.dtype; - const DType O_type = output_O->data.dtype; - void* devPtrQ = input_Q->data.dptr; - void* devPtrKV = input_KV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - stride = (typeToNumBits(QKV_type) * num_gqa_groups * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void* devPtrK = devPtrKV; - void* devPtrV = static_cast(static_cast(devPtrKV) + stride); - void* devPtrDescaleQ = input_Q->scale_inv.dptr; - void* devPtrDescaleK = input_KV->scale_inv.dptr; - void* devPtrDescaleV = input_KV->scale_inv.dptr; - - void* devPtrO = output_O->data.dptr; - void* devPtrAmaxO = output_O->amax.dptr; - void* devPtrScaleO = output_O->scale.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]); - 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; - 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]); - devPtrM = output_M->data.dptr; - devPtrZInv = output_ZInv->data.dptr; - output_rng_state->data.dptr = rng_state->data.dptr; - } 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 = - reinterpret_cast(reinterpret_cast(cu_seqlens_kv->data.dptr)); - void* devPtrDropoutSeed = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); - void* devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - 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)) { - 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); - } 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, - 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"); - } - - 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; - } -} -// fused attention BWD FP8 with packed KV -void fused_attn_fp8_bwd_kvpacked( - 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, - const Tensor* input_Q, const Tensor* input_KV, 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_dKV, 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; - const DType QKV_type = input_Q->data.dtype; - const DType dO_type = input_dO->data.dtype; - const DType dQKV_type = output_dQ->data.dtype; - void* devPtrQ = input_Q->data.dptr; - void* devPtrKV = input_KV->data.dptr; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - stride = (typeToNumBits(QKV_type) * num_gqa_groups * head_dim) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - stride = (typeToNumBits(QKV_type) * head_dim) / 8; - } - void* devPtrK = devPtrKV; - void* devPtrV = static_cast(static_cast(devPtrKV) + stride); - void* devPtrDescaleQ = input_Q->scale_inv.dptr; - void* devPtrDescaleK = input_KV->scale_inv.dptr; - void* devPtrDescaleV = input_KV->scale_inv.dptr; - - void* devPtrO = input_O->data.dptr; - const DType O_type = input_O->data.dtype; - void* devPtrDescaleO = nullptr; - if (O_type == DType::kFloat8E4M3 || O_type == DType::kFloat8E5M2) { - devPtrDescaleO = input_O->scale_inv.dptr; - } - void* devPtrdO = input_dO->data.dptr; - void* devPtrDescaledO = input_dO->scale_inv.dptr; - - void* devPtrM = input_M->data.dptr; - void* devPtrZInv = input_ZInv->data.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* devPtrdQ = output_dQ->data.dptr; - void* devPtrdKV = output_dKV->data.dptr; - void* devPtrdK = devPtrdKV; - void* devPtrdV = static_cast(static_cast(devPtrdKV) + stride); - void* devPtrAmaxdQ = output_dQ->amax.dptr; - void* devPtrAmaxdK = output_dKV->amax.dptr; - void* devPtrAmaxdV = output_dKV->amax.dptr; - void* devPtrScaledQ = output_dQ->scale.dptr; - void* devPtrScaledK = output_dKV->scale.dptr; - void* devPtrScaledV = output_dKV->scale.dptr; - - void* devPtrcuSeqlensQ = - reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); - void* devPtrcuSeqlensKV = - reinterpret_cast(reinterpret_cast(cu_seqlens_kv->data.dptr)); - void* devPtrDropoutSeed = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr)); - void* devPtrDropoutOffset = - reinterpret_cast(reinterpret_cast(rng_state->data.dptr) + 1); - - 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)) { - 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, 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_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) { - fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim, 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"); - } - - 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; - } -} // 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, diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 3daf45d162..c2efa25829 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -13,47 +13,6 @@ namespace transformer_engine { #if (CUDNN_VERSION >= 8900) -// fused attention FWD FP8 with packed QKV -void fused_attn_fp8_fwd_qkvpacked(size_t batch, size_t num_attn_heads, size_t 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_QKV, Tensor *input_output_S, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle); - -// fused attention BWD FP8 with packed QKV -void fused_attn_fp8_bwd_qkvpacked( - size_t batch, size_t num_attn_heads, size_t 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_QKV, 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_dQKV, const Tensor *cu_seqlens, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); - -// fused attention FWD FP8 with packed KV -void fused_attn_fp8_fwd_kvpacked(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_KV, 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 packed KV -void fused_attn_fp8_bwd_kvpacked( - 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, - const Tensor *input_Q, const Tensor *input_KV, 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_dKV, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle); - // 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, diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 40e6a0b4bc..298dc63900 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -217,6 +217,8 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( int64_t window_size_right, bool return_max_logit, bool cuda_graph); /*! \brief Compute dot product attention with packed QKV input. + * + * \deprecated Please use `nvte_fused_attn_fwd` with separate Q, K, V tensors instead. * * Computes: * - P = Q * Transpose(K) + Bias @@ -270,6 +272,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ +[[deprecated( + "nvte_fused_attn_fwd_qkvpacked() is deprecated. Please use nvte_fused_attn_fwd() with separate " + "Q, K, V tensors instead.")]] void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, @@ -282,6 +287,8 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTETensor workspace, cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with packed QKV input. + * + * \deprecated Please use `nvte_fused_attn_bwd` with separate Q, K, V tensors instead. * * Support Matrix: \verbatim @@ -330,6 +337,9 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ +[[deprecated( + "nvte_fused_attn_bwd_qkvpacked() is deprecated. Please use nvte_fused_attn_bwd() with separate " + "Q, K, V tensors instead.")]] void nvte_fused_attn_bwd_qkvpacked( const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, @@ -340,6 +350,8 @@ void nvte_fused_attn_bwd_qkvpacked( NVTETensor workspace, cudaStream_t stream); /*! \brief Compute dot product attention with packed KV input. + * + * \deprecated Please use `nvte_fused_attn_fwd` with separate Q, K, V tensors instead. * * Computes: * - P = Q * Transpose(K) + Bias @@ -401,6 +413,9 @@ void nvte_fused_attn_bwd_qkvpacked( * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ +[[deprecated( + "nvte_fused_attn_fwd_kvpacked() is deprecated. Please use nvte_fused_attn_fwd() with separate " + "Q, K, V tensors instead.")]] void nvte_fused_attn_fwd_kvpacked( const NVTETensor Q, const NVTETensor KV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, @@ -413,6 +428,8 @@ void nvte_fused_attn_fwd_kvpacked( int64_t window_size_right, NVTETensor workspace, cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with packed KV input. + * + * \deprecated Please use `nvte_fused_attn_bwd` with separate Q, K, V tensors instead. * * Support Matrix: \verbatim @@ -467,6 +484,9 @@ void nvte_fused_attn_fwd_kvpacked( * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ +[[deprecated( + "nvte_fused_attn_bwd_kvpacked() is deprecated. Please use nvte_fused_attn_bwd() with separate " + "Q, K, V tensors instead.")]] void nvte_fused_attn_bwd_kvpacked( const NVTETensor Q, const NVTETensor KV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, From 5ea83432a400481b73e42de18bee7c206cb18fac Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 10 Nov 2025 10:42:57 -0800 Subject: [PATCH 052/521] Move Triton to common (#2359) * move triton to common and change paths Signed-off-by: tdophung * Formatting Signed-off-by: tdophung --------- Signed-off-by: tdophung --- transformer_engine/common/triton/__init__.py | 5 + .../common/triton/cross_entropy.py | 252 ++++++++ transformer_engine/common/triton/pad.py | 59 ++ .../common/triton/permutation.py | 605 +++++++++++++++++ transformer_engine/pytorch/distributed.py | 3 +- transformer_engine/pytorch/triton/__init__.py | 2 +- .../pytorch/triton/cross_entropy.py | 252 +------- transformer_engine/pytorch/triton/pad.py | 55 +- .../pytorch/triton/permutation.py | 609 +----------------- 9 files changed, 943 insertions(+), 899 deletions(-) create mode 100644 transformer_engine/common/triton/__init__.py create mode 100644 transformer_engine/common/triton/cross_entropy.py create mode 100644 transformer_engine/common/triton/pad.py create mode 100644 transformer_engine/common/triton/permutation.py diff --git a/transformer_engine/common/triton/__init__.py b/transformer_engine/common/triton/__init__.py new file mode 100644 index 0000000000..76c9b98d0e --- /dev/null +++ b/transformer_engine/common/triton/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Kernels written with OpenAI Triton.""" diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py new file mode 100644 index 0000000000..fc49ac20b7 --- /dev/null +++ b/transformer_engine/common/triton/cross_entropy.py @@ -0,0 +1,252 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Efficient Cross Entropy kernels written with OpenAI Triton.""" + +import triton +import triton.language as tl + + +@triton.jit +def online_softmax_kernel( + X_ptr, + X_stride, + Y_ptr, + Y_stride, + m_d_X_y_ptr, + m_d_X_y_stride, + rank, + n_cols, + BLOCK_SIZE: tl.constexpr, +): + """ + This kernel computes the m/d components on this TP rank for the online softmax. + + Parameters: + X_ptr: Pointer to input tensor. + X_stride (int): The stride of the input tensor. + Y_ptr: Pointer to target tensor. + Y_stride (int): The stride of the target tensor. + m_d_X_y_ptr: Pointer to m/d/X_y tensor. + m_d_X_y_stride (int): The stride of the m/d/X_y tensor. + rank (int): The rank of this device in the TP group. + n_cols (int): The number of columns in the input tensor. + BLOCK_SIZE (int): The block size for Triton operations. + """ + + program_id = tl.program_id(0).to(tl.int64) + + # locate the start index + X_ptr += program_id * X_stride + + # Load Y_ptr + Y_ptr += program_id * Y_stride + y = tl.load(Y_ptr) + + vocab_start_idx = rank * n_cols + vocab_end_idx = (rank + 1) * n_cols + if y >= vocab_start_idx: + if y < vocab_end_idx: + X_y = tl.load(X_ptr + y - vocab_start_idx).to(tl.float32) + else: + X_y = float("-inf") + else: + X_y = float("-inf") + + m_d_X_y_ptr += program_id * m_d_X_y_stride * 3 + + # 3. [Online softmax] first pass: find max + sum + m = float("-inf") # m is the max value. use the notation from the paper + d = 0.0 # d is the sum. use the notation from the paper + + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols, other=float("-inf")).to( + tl.float32 + ) + block_max = tl.max(X_block) + m_new = tl.maximum(m, block_max) + d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new)) + m = m_new + + tl.store(m_d_X_y_ptr, m) + tl.store(m_d_X_y_ptr + m_d_X_y_stride, d) + tl.store(m_d_X_y_ptr + (2 * m_d_X_y_stride), X_y) + + +@triton.jit +def cross_entropy_kernel( + X_ptr, + X_stride, + Y_ptr, + Y_stride, + loss_ptr, + loss_stride, + m_d_X_y_ptr, + m_d_X_y_stride, + rank, + world_size, + ignore_idx, + n_cols, + n_non_ignore, + reduce_loss: tl.constexpr, + label_smoothing: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + This kernel computes both cross entropy loss and the gradient of the input. + + Parameters: + X_ptr: Pointer to input tensor. + X_stride (int): The stride of the input tensor. + Y_ptr: Pointer to target tensor. + Y_stride (int): The stride of the target tensor. + loss_ptr: Pointer to tensor to store the loss. + loss_stride (int): The stride of the loss tensor. + m_d_X_y_ptr: Pointer to m/d/X_y tensor. + m_d_X_y_stride: The stride of m/d/X_y tensor. + rank (int): The rank of this device in the TP group. + world_size (int): The size of world involved in this distributed loss calculation. + ignore_idx (int): Tokens to be ignored for loss and gradient calculation. + n_cols (int): The number of columns in the input tensor. + n_non_ignore (int): The number of non-ignored elements in the batch. + label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing. + BLOCK_SIZE (int): The block size for Triton operations. + """ + + program_id = tl.program_id(0).to(tl.int64) + + # locate the start index + X_ptr += program_id * X_stride + + # Load Y_ptr + Y_ptr += program_id * Y_stride + y = tl.load(Y_ptr) + + if y == ignore_idx: + # set all X_ptr as 0 + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols) + return + + loss_ptr += program_id * loss_stride + m_d_X_y_ptr += program_id * 3 * m_d_X_y_stride + + # Need to reduce the m/d/X_y values from other TP ranks + m = tl.load(m_d_X_y_ptr) + d = tl.load(m_d_X_y_ptr + m_d_X_y_stride) + ori_X_y = tl.load(m_d_X_y_ptr + (2 * m_d_X_y_stride)) + + for i in range(1, world_size): + offset = i * 3 * n_non_ignore * m_d_X_y_stride + access_ptr = m_d_X_y_ptr + offset + m_new = tl.load(access_ptr) + d_new = tl.load(access_ptr + m_d_X_y_stride) + X_y_new = tl.load(access_ptr + (2 * m_d_X_y_stride)) + + d = d * tl.exp(m - tl.maximum(m, m_new)) + d_new * tl.exp(m_new - tl.maximum(m, m_new)) + m = tl.maximum(m, m_new) + ori_X_y = tl.maximum(ori_X_y, X_y_new) + + # Label smoothing is a general case of normal cross entropy + scaled_x_sum = 0.0 + eps = label_smoothing / (n_cols * world_size) + + # 4. [Online softmax] second pass: calculate the gradients + # dx_y = (softmax(x_y) - 1) / N + # dx_i = softmax(x_i) / N, i != y + # N is the number of non ignored elements in the batch + # For label smoothing: + # dx_i = (softmax(x_y) - label_smoothing / V) / N, V = n_cols, i != y + # dx_y = (softmax(x_y) - label_smoothing / V - (1 - label_smoothing)) / N + # = dx_i - (1 - label_smoothing) / N + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols, other=float("-inf")) + grad_dtype = X_block.dtype + X_block = X_block.to(tl.float32) + if label_smoothing > 0: + # scale X beforehand to avoid overflow + scaled_x_sum += tl.sum(tl.where(X_offsets < n_cols, -eps * X_block, 0.0)) + # Scale gradients based on reduction mode + # For reduce_loss=True: PyTorch will scale by 1/n_rows, so we need to scale by n_rows/n_non_ignore + # For reduce_loss=False: No additional scaling from PyTorch, so we don't scale here + if reduce_loss: + X_block = (tl.exp(X_block - m) / d - eps) / (n_non_ignore) + else: + X_block = tl.exp(X_block - m) / d - eps + tl.store(X_ptr + X_offsets, X_block.to(grad_dtype), mask=X_offsets < n_cols) + + # We need tl.debug_barrier() to ensure the new result of X_ptr is written + tl.debug_barrier() + + # 5. Calculate the loss + + # loss = log (softmax(X_y)) = log ((e ^ (X_y - max(X)) / sum(e ^ (X - max(X)))) + # = (X_y - max(X)) - log(sum(e ^ (X - max(X)))) + loss = -(ori_X_y - m - tl.log(d)) + + # Orginal loss = H(q, p), with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps + # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p) + # = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i)) + # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as: + # = (1 - label_smoothing) * H(q, p) + (-sum(x_i * eps) + label_smoothing * (m + logd)) + # Refer to H(q', p) in section 7 of the paper: https://arxiv.org/pdf/1512.00567 + if label_smoothing > 0: + smooth_loss = scaled_x_sum + label_smoothing * (m + tl.log(d)) + loss = loss * (1 - label_smoothing) + smooth_loss + + # 6. Specially handle the i==y case where `dx_y = (softmax(x_y) - (1 - label_smoothing) / N` + vocab_start_idx = rank * n_cols + vocab_end_idx = (rank + 1) * n_cols + if y >= vocab_start_idx: + if y < vocab_end_idx: + X_y = tl.load(X_ptr + y - vocab_start_idx) + # Apply the same conditional scaling logic for the target token + if reduce_loss: + X_y += -(1 - label_smoothing) / (n_non_ignore) + else: + X_y += -(1 - label_smoothing) + tl.store(X_ptr + y - vocab_start_idx, X_y) + + tl.store(loss_ptr, loss) + + +@triton.jit +def element_mul_kernel( + X_ptr, + X_stride, + grad_output_ptr, + grad_output_stride, + n_cols, + BLOCK_SIZE: tl.constexpr, +): + """ + This function multiplies each element of the tensor pointed by X_ptr with the value pointed by grad_output_ptr. + The multiplication is performed in-place on the tensor pointed by X_ptr. + + Parameters: + X_ptr: Pointer to the input tensor. + X_stride (int): The stride of the input tensor. + grad_output_ptr: Pointer to the gradient output value. + n_cols (int): The number of columns in the input tensor. + BLOCK_SIZE (int): The block size for Triton operations. + """ + + # Get the program ID and convert it to int64 to avoid overflow + program_id = tl.program_id(0).to(tl.int64) + + # Locate the start index + X_ptr += program_id * X_stride + + # Load the gradient output value + grad_output_ptr += program_id * grad_output_stride + grad_output = tl.load(grad_output_ptr) + + # Perform the element-wise multiplication + for i in range(0, n_cols, BLOCK_SIZE): + X_offsets = i + tl.arange(0, BLOCK_SIZE) + X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols) + tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols) diff --git a/transformer_engine/common/triton/pad.py b/transformer_engine/common/triton/pad.py new file mode 100644 index 0000000000..8f15e7dcba --- /dev/null +++ b/transformer_engine/common/triton/pad.py @@ -0,0 +1,59 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Efficient NVFP4 padding kernels written with OpenAI Triton . + +TODO(ksivamani): Documentation + +""" + +import triton +import triton.language as tl + + +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 128}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}, num_warps=8, num_stages=1), + ], + key=["out_dim0", "out_dim1"], +) +@triton.jit +def zero_pad_kernel( + inp_ptr, + out_ptr, + in_dim0: tl.constexpr, + in_dim1: tl.constexpr, + out_dim0: tl.constexpr, + out_dim1: tl.constexpr, + in_s0, + in_s1, + out_s0, + out_s1, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Pads a tensor assuming it's a columnwise scaling inverse.""" + + # tile over OUTPUT coordinates + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # output rows + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # output cols + om = offs_m[:, None] + on = offs_n[None, :] + + # edge masking for output + out_mask = (om < out_dim0) & (on < out_dim1) + + # valid input region is simply top-left (no offsets) + in_mask = (om < in_dim0) & (on < in_dim1) + + # load valid input, else zero (masked load touches memory only where True) + x = tl.load(inp_ptr + om * in_s0 + on * in_s1, mask=in_mask, other=0) + + # store to output (only within bounds of the output tile) + tl.store(out_ptr + om * out_s0 + on * out_s1, x, mask=out_mask) diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py new file mode 100644 index 0000000000..3a3a32014f --- /dev/null +++ b/transformer_engine/common/triton/permutation.py @@ -0,0 +1,605 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Efficient Permutation kernels written with OpenAI Triton.""" + +import triton +import triton.language as tl + +from triton.language import core +from triton.language.standard import _log2 +from packaging import version + + +# The following three argsort related kernels are adapted from +# the issue https://github.com/triton-lang/triton/issues/3698 + +get_int_dtype = core.get_int_dtype +if version.parse(triton.__version__) >= version.parse("3.5.0"): + get_int_dtype = triton.constexpr_function(get_int_dtype) + + +@triton.jit +def _compare_and_swap(x, indices, flip, i: tl.constexpr, n_dims: tl.constexpr): + n_outer: tl.constexpr = x.numel >> n_dims + shape: tl.constexpr = [n_outer * (2**i), 2, 2 ** (n_dims - i - 1)] + y = tl.reshape(x, shape) + z = tl.reshape(indices, shape) + + mask = tl.arange(0, 2)[None, :, None] + + l_value = tl.reshape(tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape), x.shape).to( + x.dtype + ) + r_value = tl.reshape(tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape), x.shape).to( + x.dtype + ) + + l_indice = tl.reshape(tl.broadcast_to(tl.sum(z * (1 - mask), 1)[:, None, :], shape), x.shape) + r_indice = tl.reshape(tl.broadcast_to(tl.sum(z * mask, 1)[:, None, :], shape), x.shape) + + idtype = get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) + + il_value = l_value.to(idtype, bitcast=True) + ir_value = r_value.to(idtype, bitcast=True) + ix = x.to(idtype, bitcast=True) + + flag1 = tl.where(((l_value > r_value) ^ flip) != 0, il_value ^ ir_value, tl.zeros_like(ix)) + ret = ix ^ flag1 + flag2 = tl.where(((l_value > r_value) ^ flip) != 0, l_indice ^ r_indice, tl.zeros_like(ix)) + ind = indices ^ flag2 + + return ret.to(x.dtype, bitcast=True), ind + + +@triton.jit +def _bitonic_merge(x, indices, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr): + n_outer: tl.constexpr = x.numel >> n_dims + tl.static_assert(stage <= n_dims) + """ + order_type 0 == ascending + order_type 1 == descending + order_type 2 == alternating + """ + if order == 2: + shape: tl.constexpr = [n_outer * (2 ** (n_dims - 1 - stage)), 2, 2**stage] + flip = tl.reshape(tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape) + else: + flip = tl.full(x.shape, value=order, dtype=tl.int32) + for i in tl.static_range(stage): + x, indices = _compare_and_swap(x, indices, flip, i + (n_dims - stage), n_dims) + return x, indices + + +@triton.jit +def _argsort(x, indices, n_dims: tl.constexpr): + for i in tl.static_range(1, n_dims + 1): + x, indices = _bitonic_merge(x, indices, i, 2 if i < n_dims else 1, n_dims) + return x, indices + + +@triton.jit +def _row_id_map_pass_1_kernel( + # pointers + routing_map_ptr, + row_id_map_ptr, + workspace_ptr, + # sizes + num_tokens, + # strides + stride_routing_map_token, + stride_routing_map_expert, + stride_row_id_map_token, + stride_row_id_map_expert, + # metas + BLOCK_SIZE: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offset = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + expert_token_mask = tl.load( + routing_map_ptr + pid_m * stride_routing_map_expert + offset * stride_routing_map_token, + mask=(offset < num_tokens), + other=0, + ).to(tl.int32) + row_id_within_token_block = tl.cumsum(expert_token_mask) * expert_token_mask + tl.store( + row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token, + row_id_within_token_block, + mask=offset < num_tokens, + ) + n_tokens_per_block = tl.sum(expert_token_mask) + tl.store(workspace_ptr + pid_m * tl.cdiv(num_tokens, BLOCK_SIZE) + pid_n, n_tokens_per_block) + + +@triton.jit +def _row_id_map_pass_2_kernel( + # pointers + row_id_map_ptr, + workspace_ptr, + # sizes + num_tokens, + # strides + stride_row_id_map_token, + stride_row_id_map_expert, + # metas + WORKSPACE_LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + chunk_idx = pid_m * tl.cdiv(num_tokens, BLOCK_SIZE) + pid_n + offset = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + row_id_within_token_block = tl.load( + row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token, + mask=(offset < num_tokens), + other=0, + ) + + workspace_off = tl.arange(0, WORKSPACE_LOAD_WIDTH) + n_tokens_per_chunk = tl.load(workspace_ptr + workspace_off, mask=workspace_off < chunk_idx) + row_id = tl.where( + row_id_within_token_block == 0, + -1, + row_id_within_token_block + tl.sum(n_tokens_per_chunk) - 1, + ) + tl.store( + row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token, + row_id, + mask=(offset < num_tokens), + ) + + +@triton.jit +def _row_id_map_pass_3_kernel( + # pointers + row_id_map_ptr, + # sizes + num_experts: tl.constexpr, + # strides + stride_row_id_map_token, + stride_row_id_map_expert, + # metas + LOAD_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + n_dims: tl.constexpr = _log2(LOAD_SIZE) + off = tl.arange(0, LOAD_SIZE) + row_id_map = tl.load( + row_id_map_ptr + pid * stride_row_id_map_token + stride_row_id_map_expert * off, + mask=off < num_experts, + other=-1, + ) + n_routed = tl.sum(tl.where(row_id_map != -1, 1, 0)) + indices = off + sorted_map, indices = _argsort(row_id_map, indices, n_dims=n_dims) + tl.store( + row_id_map_ptr + pid * stride_row_id_map_token + off * stride_row_id_map_expert, + sorted_map, + mask=off < n_routed, + ) + tl.store( + row_id_map_ptr + + pid * stride_row_id_map_token + + (num_experts + off) * stride_row_id_map_expert, + indices, + mask=off < n_routed, + ) + tl.store( + row_id_map_ptr + pid * stride_row_id_map_token + num_experts * 2 * stride_row_id_map_expert, + n_routed, + ) + + +@triton.jit +def _permute_kernel( + # pointers + input_ptr, + output_ptr, + row_id_map_ptr, + probs_ptr, + scale_ptr, + permuted_probs_ptr, + permuted_scale_ptr, + # sizes + num_experts: tl.constexpr, + hidden_size: tl.constexpr, + scale_hidden_dim, + # strides + stride_row_id_map_token, + stride_row_id_map_expert, + stride_input_token, + stride_input_hidden, + stride_output_token, + stride_output_hidden, + stride_probs_token, + stride_probs_expert, + stride_scale_token, + stride_scale_hidden, + stride_permuted_probs_token, + stride_permuted_scale_token, + stride_permuted_scale_hidden, + # metas + PERMUTE_PROBS: tl.constexpr, + PERMUTE_SCALE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + cur_off = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = cur_off < hidden_size + src_row = pid_t.to(tl.int64) + input_off = src_row * stride_input_token + cur_off * stride_input_hidden + inp = tl.load(input_ptr + input_off, mask=mask) + if PERMUTE_SCALE: + mask_scale = cur_off < scale_hidden_dim + scale_off = pid_t * stride_scale_token + cur_off * stride_scale_hidden + scale = tl.load(scale_ptr + scale_off, mask=mask_scale) + n_routed = tl.load( + row_id_map_ptr + + pid_t * stride_row_id_map_token + + num_experts * 2 * stride_row_id_map_expert + ) + for idx in tl.range(n_routed): + dst_row = tl.load( + row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert + ).to(tl.int64) + output_off = dst_row * stride_output_token + cur_off * stride_output_hidden + if PERMUTE_SCALE: + permuted_scale_off = ( + dst_row * stride_permuted_scale_token + cur_off * stride_permuted_scale_hidden + ) + tl.store(permuted_scale_ptr + permuted_scale_off, scale, mask=mask_scale) + if PERMUTE_PROBS: + expert_idx = tl.load( + row_id_map_ptr + + pid_t * stride_row_id_map_token + + (num_experts + idx) * stride_row_id_map_expert + ) + prob_off = pid_t * stride_probs_token + expert_idx * stride_probs_expert + prob = tl.load(probs_ptr + prob_off) + if pid_h == 0: + permuted_prob_off = dst_row * stride_permuted_probs_token + tl.store(permuted_probs_ptr + permuted_prob_off, prob) + if prob == 0.0: + # for routing_map padding + # dst_row != -1 and prob == 0.0 means that this slot is padded + tl.store(output_ptr + output_off, 0.0, mask=mask) + else: + tl.store(output_ptr + output_off, inp, mask=mask) + else: + tl.store(output_ptr + output_off, inp, mask=mask) + + +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}), + ], + key=["hidden_size"], + )(_permute_kernel) +except RuntimeError: + pass + + +@triton.jit +def _unpermute_kernel( + # pointers + input_ptr, + output_ptr, + row_id_map_ptr, + merging_probs_ptr, + permuted_probs_ptr, + unpermuted_probs_ptr, + # sizes + num_experts: tl.constexpr, + hidden_size: tl.constexpr, + # strides + stride_row_id_map_token, + stride_row_id_map_expert, + stride_input_token, + stride_input_hidden, + stride_output_token, + stride_output_hidden, + stride_merging_probs_token, + stride_merging_probs_expert, + stride_permuted_probs_token, + stride_unpermuted_probs_token, + stride_unpermuted_probs_expert, + # metas + PROBS_LOAD_WIDTH: tl.constexpr, + WITH_MERGING_PROBS: tl.constexpr, + PERMUTE_PROBS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + data_type = input_ptr.dtype.element_ty + compute_type = tl.float32 + + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + current_offset = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = current_offset < hidden_size + if PERMUTE_PROBS: + # write 0.0 to probs_grad that are not routed + if pid_h == 0: + map_load_off = tl.arange(0, PROBS_LOAD_WIDTH) + unpermuted_prob_off = ( + pid_t * stride_unpermuted_probs_token + + stride_unpermuted_probs_expert * map_load_off + ) + tl.store( + unpermuted_probs_ptr + unpermuted_prob_off, 0.0, mask=map_load_off < num_experts + ) + accumulator = tl.zeros((BLOCK_SIZE,), dtype=compute_type) + n_routed = tl.load( + row_id_map_ptr + + pid_t * stride_row_id_map_token + + num_experts * 2 * stride_row_id_map_expert + ) + for idx in tl.range(n_routed): + src_row = tl.load( + row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert + ).to(tl.int64) + input_off = src_row * stride_input_token + current_offset * stride_input_hidden + inp = tl.load(input_ptr + input_off, mask=mask) + inp = inp.to(compute_type) + if WITH_MERGING_PROBS: + expert_idx = tl.load( + row_id_map_ptr + + pid_t * stride_row_id_map_token + + (num_experts + idx) * stride_row_id_map_expert + ) + merging_prob_off = ( + pid_t * stride_merging_probs_token + expert_idx * stride_merging_probs_expert + ) + merging_prob = tl.load(merging_probs_ptr + merging_prob_off).to(compute_type) + inp *= merging_prob + accumulator += inp + if PERMUTE_PROBS: + if pid_h == 0: + expert_idx = tl.load( + row_id_map_ptr + + pid_t * stride_row_id_map_token + + (num_experts + idx) * stride_row_id_map_expert + ) + unpermuted_prob_off = ( + pid_t * stride_unpermuted_probs_token + + expert_idx * stride_unpermuted_probs_expert + ) + permuted_prob_off = src_row * stride_permuted_probs_token + prob = tl.load(permuted_probs_ptr + permuted_prob_off) + tl.store(unpermuted_probs_ptr + unpermuted_prob_off, prob) + accumulator = accumulator.to(data_type) + dst_row = pid_t.to(tl.int64) + output_off = dst_row * stride_output_token + current_offset * stride_output_hidden + tl.store(output_ptr + output_off, accumulator, mask=mask) + + +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}), + ], + key=["hidden_size"], + )(_unpermute_kernel) +except RuntimeError: + pass + + +@triton.jit +def _unpermute_bwd_with_merging_probs_kernel( + # pointers + fwd_output_grad_ptr, + fwd_input_grad_ptr, + fwd_input_ptr, + merging_probs_ptr, + merging_probs_grad_ptr, + row_id_map_ptr, + # sizes + num_experts: tl.constexpr, + hidden_size: tl.constexpr, + # strides + stride_row_id_map_token, + stride_row_id_map_expert, + stride_fwd_output_grad_token, + stride_fwd_output_grad_hidden, + stride_fwd_input_grad_token, + stride_fwd_input_grad_hidden, + stride_fwd_input_token, + stride_fwd_input_hidden, + stride_merging_probs_token, + stride_merging_probs_expert, + stride_merging_probs_grad_token, + stride_merging_probs_grad_expert, + # metas + PROBS_LOAD_WIDTH: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + data_type = fwd_output_grad_ptr.dtype.element_ty + compute_type = tl.float32 + + pid = tl.program_id(0) + map_load_off = tl.arange(0, PROBS_LOAD_WIDTH) + token_probs_grad_off = ( + pid * stride_merging_probs_grad_token + stride_merging_probs_grad_expert * map_load_off + ) + tl.store(merging_probs_grad_ptr + token_probs_grad_off, 0.0, mask=map_load_off < num_experts) + n_routed = tl.load( + row_id_map_ptr + pid * stride_row_id_map_token + num_experts * 2 * stride_row_id_map_expert + ) + for idx in tl.range(n_routed): + dst_row = tl.load( + row_id_map_ptr + pid * stride_row_id_map_token + idx * stride_row_id_map_expert + ).to(tl.int64) + expert_idx = tl.load( + row_id_map_ptr + + pid * stride_row_id_map_token + + (num_experts + idx) * stride_row_id_map_expert + ) + prob_grad_accum = tl.zeros((BLOCK_SIZE,), dtype=compute_type) + current_start = 0 + while current_start < hidden_size: + current_offset = current_start + tl.arange(0, BLOCK_SIZE) + mask = current_offset < hidden_size + src_row = pid.to(tl.int64) + input_off = ( + src_row * stride_fwd_output_grad_token + + current_offset * stride_fwd_output_grad_hidden + ) + inp = tl.load(fwd_output_grad_ptr + input_off, mask=mask) + inp = inp.to(compute_type) + merging_prob_off = ( + pid * stride_merging_probs_token + expert_idx * stride_merging_probs_expert + ) + merging_prob = tl.load(merging_probs_ptr + merging_prob_off).to(compute_type) + output = inp * merging_prob + output = output.to(data_type) + output_off = ( + dst_row * stride_fwd_input_grad_token + + current_offset * stride_fwd_input_grad_hidden + ) + tl.store(fwd_input_grad_ptr + output_off, output, mask=mask) + + fwd_input_off = ( + dst_row * stride_fwd_input_token + current_offset * stride_fwd_input_hidden + ) + fwd_input = tl.load(fwd_input_ptr + fwd_input_off, mask=mask) + prob_grad_accum += fwd_input.to(compute_type) * inp + current_start += BLOCK_SIZE + probs_grad = tl.sum(prob_grad_accum).to(merging_probs_grad_ptr.dtype.element_ty) + probs_grad_off = ( + pid * stride_merging_probs_grad_token + expert_idx * stride_merging_probs_grad_expert + ) + tl.store(merging_probs_grad_ptr + probs_grad_off, probs_grad) + + +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}), + ], + key=["hidden_size"], + )(_unpermute_bwd_with_merging_probs_kernel) +except RuntimeError: + pass + + +@triton.jit +def _make_chunk_sort_map_kernel( + # pointers + split_sizes_ptr, + sorted_indices_ptr, + dst_rows_ptr, + # sizes + num_splits: tl.constexpr, + # metas + IDX_LOAD_WIDTH: tl.constexpr, +): + pid = tl.program_id(0) + + load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) + sorted_indices = tl.load( + sorted_indices_ptr + load_split_offset, mask=load_split_offset < num_splits + ) + + # get chunk idx of the current token in the input tensor + input_split_sizes = tl.load( + split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0 + ).to(tl.int32) + input_split_sizes_cumsum = tl.cumsum(input_split_sizes) + input_split_sizes_mask = tl.where(input_split_sizes_cumsum <= pid, 1, 0) + input_chunk_idx = tl.sum(input_split_sizes_mask) + input_split_sizes_presum = tl.sum(input_split_sizes * input_split_sizes_mask) + in_chunk_offset = pid - input_split_sizes_presum + + # get chunk idx of the current token in the output tensor + output_chunk_mask = tl.where(sorted_indices == input_chunk_idx, 1, 0) + output_chunk_idx = tl.argmax(output_chunk_mask, axis=-1) + + # make row_id_map + output_split_sizes = tl.load( + split_sizes_ptr + sorted_indices, mask=load_split_offset < num_splits + ).to(tl.int32) + output_pre_split_sizes = tl.where(load_split_offset < output_chunk_idx, output_split_sizes, 0) + dst_row = tl.sum(output_pre_split_sizes) + in_chunk_offset + tl.store(dst_rows_ptr + pid, dst_row) + + +@triton.jit +def _sort_chunks_by_map_kernel( + # pointers + input_ptr, + output_ptr, + row_id_map_ptr, + probs_ptr, + permuted_probs_ptr, + # sizes + hidden_size: tl.constexpr, + # strides + stride_input_token, + stride_input_hidden, + stride_output_token, + stride_output_hidden, + stride_probs_token, + stride_permuted_probs_token, + # metas + PERMUTE_PROBS: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + FORWARD: tl.constexpr, +): + pid_t = tl.program_id(0) + pid_h = tl.program_id(1) + if FORWARD: + src_row = pid_t.to(tl.int64) + dst_row = tl.load(row_id_map_ptr + pid_t).to(tl.int64) + else: + src_row = tl.load(row_id_map_ptr + pid_t).to(tl.int64) + dst_row = pid_t.to(tl.int64) + current_offset = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = current_offset < hidden_size + input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden + output_offsets = dst_row * stride_output_token + current_offset * stride_output_hidden + inp = tl.load(input_ptr + input_offsets, mask=mask) + tl.store(output_ptr + output_offsets, inp, mask=mask) + if PERMUTE_PROBS: + if pid_h == 0: + prob_off = src_row * stride_probs_token + prob = tl.load(probs_ptr + prob_off) + permuted_prob_off = dst_row * stride_permuted_probs_token + tl.store(permuted_probs_ptr + permuted_prob_off, prob) + + +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}), + ], + key=["hidden_size"], + )(_sort_chunks_by_map_kernel) +except RuntimeError: + pass diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 8c14d5ab7f..e938509e58 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -29,12 +29,14 @@ import transformer_engine_torch as tex +from transformer_engine.pytorch.triton.pad import pad_columnwise_scale_inv from . import torch_version from .utils import ( is_non_tn_fp8_gemm_supported, safely_set_viewless_tensor_data, needs_quantized_gemm, ) + from .constants import dist_group_type from .quantization import FP8GlobalStateManager, autocast from .tensor.float8_tensor import Float8Quantizer, Float8Tensor, Float8CurrentScalingQuantizer @@ -46,7 +48,6 @@ from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from .tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage -from .triton.pad import pad_columnwise_scale_inv from ..debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer diff --git a/transformer_engine/pytorch/triton/__init__.py b/transformer_engine/pytorch/triton/__init__.py index 76c9b98d0e..80766864d9 100644 --- a/transformer_engine/pytorch/triton/__init__.py +++ b/transformer_engine/pytorch/triton/__init__.py @@ -2,4 +2,4 @@ # # See LICENSE for license information. -"""Kernels written with OpenAI Triton.""" +"""PyTorch wrappers for Triton kernels.""" diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index 7cfff1da9d..d7e2256e24 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Efficient Cross Entropy kernels written with OpenAI Triton.""" +"""PyTorch wrapper functions for Cross Entropy Triton kernels.""" from typing import Union from functools import reduce @@ -12,257 +12,17 @@ import torch.distributed as dist import triton -import triton.language as tl - - -@triton.jit -def online_softmax_kernel( - X_ptr, - X_stride, - Y_ptr, - Y_stride, - m_d_X_y_ptr, - m_d_X_y_stride, - rank, - n_cols, - BLOCK_SIZE: tl.constexpr, -): - """ - This kernel computes the m/d components on this TP rank for the online softmax. - - Parameters: - X_ptr: Pointer to input tensor. - X_stride (int): The stride of the input tensor. - Y_ptr: Pointer to target tensor. - Y_stride (int): The stride of the target tensor. - m_d_X_y_ptr: Pointer to m/d/X_y tensor. - m_d_X_y_stride (int): The stride of the m/d/X_y tensor. - rank (int): The rank of this device in the TP group. - n_cols (int): The number of columns in the input tensor. - BLOCK_SIZE (int): The block size for Triton operations. - """ - - program_id = tl.program_id(0).to(tl.int64) - - # locate the start index - X_ptr += program_id * X_stride - - # Load Y_ptr - Y_ptr += program_id * Y_stride - y = tl.load(Y_ptr) - - vocab_start_idx = rank * n_cols - vocab_end_idx = (rank + 1) * n_cols - if y >= vocab_start_idx: - if y < vocab_end_idx: - X_y = tl.load(X_ptr + y - vocab_start_idx).to(tl.float32) - else: - X_y = float("-inf") - else: - X_y = float("-inf") - - m_d_X_y_ptr += program_id * m_d_X_y_stride * 3 - - # 3. [Online softmax] first pass: find max + sum - m = float("-inf") # m is the max value. use the notation from the paper - d = 0.0 # d is the sum. use the notation from the paper - - for i in range(0, n_cols, BLOCK_SIZE): - X_offsets = i + tl.arange(0, BLOCK_SIZE) - X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols, other=float("-inf")).to( - tl.float32 - ) - block_max = tl.max(X_block) - m_new = tl.maximum(m, block_max) - d = d * tl.exp(m - m_new) + tl.sum(tl.exp(X_block - m_new)) - m = m_new - - tl.store(m_d_X_y_ptr, m) - tl.store(m_d_X_y_ptr + m_d_X_y_stride, d) - tl.store(m_d_X_y_ptr + (2 * m_d_X_y_stride), X_y) - - -@triton.jit -def cross_entropy_kernel( - X_ptr, - X_stride, - Y_ptr, - Y_stride, - loss_ptr, - loss_stride, - m_d_X_y_ptr, - m_d_X_y_stride, - rank, - world_size, - ignore_idx, - n_cols, - n_non_ignore, - reduce_loss: tl.constexpr, - label_smoothing: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - """ - This kernel computes both cross entropy loss and the gradient of the input. - - Parameters: - X_ptr: Pointer to input tensor. - X_stride (int): The stride of the input tensor. - Y_ptr: Pointer to target tensor. - Y_stride (int): The stride of the target tensor. - loss_ptr: Pointer to tensor to store the loss. - loss_stride (int): The stride of the loss tensor. - m_d_X_y_ptr: Pointer to m/d/X_y tensor. - m_d_X_y_stride: The stride of m/d/X_y tensor. - rank (int): The rank of this device in the TP group. - world_size (int): The size of world involved in this distributed loss calculation. - ignore_idx (int): Tokens to be ignored for loss and gradient calculation. - n_cols (int): The number of columns in the input tensor. - n_non_ignore (int): The number of non-ignored elements in the batch. - label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing. - BLOCK_SIZE (int): The block size for Triton operations. - """ - - program_id = tl.program_id(0).to(tl.int64) - - # locate the start index - X_ptr += program_id * X_stride - - # Load Y_ptr - Y_ptr += program_id * Y_stride - y = tl.load(Y_ptr) - - if y == ignore_idx: - # set all X_ptr as 0 - for i in range(0, n_cols, BLOCK_SIZE): - X_offsets = i + tl.arange(0, BLOCK_SIZE) - tl.store(X_ptr + X_offsets, 0.0, mask=X_offsets < n_cols) - return - - loss_ptr += program_id * loss_stride - m_d_X_y_ptr += program_id * 3 * m_d_X_y_stride - - # Need to reduce the m/d/X_y values from other TP ranks - m = tl.load(m_d_X_y_ptr) - d = tl.load(m_d_X_y_ptr + m_d_X_y_stride) - ori_X_y = tl.load(m_d_X_y_ptr + (2 * m_d_X_y_stride)) - - for i in range(1, world_size): - offset = i * 3 * n_non_ignore * m_d_X_y_stride - access_ptr = m_d_X_y_ptr + offset - m_new = tl.load(access_ptr) - d_new = tl.load(access_ptr + m_d_X_y_stride) - X_y_new = tl.load(access_ptr + (2 * m_d_X_y_stride)) - - d = d * tl.exp(m - tl.maximum(m, m_new)) + d_new * tl.exp(m_new - tl.maximum(m, m_new)) - m = tl.maximum(m, m_new) - ori_X_y = tl.maximum(ori_X_y, X_y_new) - - # Label smoothing is a general case of normal cross entropy - scaled_x_sum = 0.0 - eps = label_smoothing / (n_cols * world_size) - - # 4. [Online softmax] second pass: calculate the gradients - # dx_y = (softmax(x_y) - 1) / N - # dx_i = softmax(x_i) / N, i != y - # N is the number of non ignored elements in the batch - # For label smoothing: - # dx_i = (softmax(x_y) - label_smoothing / V) / N, V = n_cols, i != y - # dx_y = (softmax(x_y) - label_smoothing / V - (1 - label_smoothing)) / N - # = dx_i - (1 - label_smoothing) / N - for i in range(0, n_cols, BLOCK_SIZE): - X_offsets = i + tl.arange(0, BLOCK_SIZE) - X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols, other=float("-inf")) - grad_dtype = X_block.dtype - X_block = X_block.to(tl.float32) - if label_smoothing > 0: - # scale X beforehand to avoid overflow - scaled_x_sum += tl.sum(tl.where(X_offsets < n_cols, -eps * X_block, 0.0)) - # Scale gradients based on reduction mode - # For reduce_loss=True: PyTorch will scale by 1/n_rows, so we need to scale by n_rows/n_non_ignore - # For reduce_loss=False: No additional scaling from PyTorch, so we don't scale here - if reduce_loss: - X_block = (tl.exp(X_block - m) / d - eps) / (n_non_ignore) - else: - X_block = tl.exp(X_block - m) / d - eps - tl.store(X_ptr + X_offsets, X_block.to(grad_dtype), mask=X_offsets < n_cols) - - # We need tl.debug_barrier() to ensure the new result of X_ptr is written - tl.debug_barrier() - - # 5. Calculate the loss - - # loss = log (softmax(X_y)) = log ((e ^ (X_y - max(X)) / sum(e ^ (X - max(X)))) - # = (X_y - max(X)) - log(sum(e ^ (X - max(X)))) - loss = -(ori_X_y - m - tl.log(d)) - - # Orginal loss = H(q, p), with label smoothing regularization = H(q', p) and (label_smoothing / V) = eps - # H(q', p) = (1 - label_smoothing) * H(q, p) + label_smoothing * H(u, p) - # = (1 - label_smoothing) * H(q, p) + eps * sum(logsoftmax(x_i)) - # By using m (global max of xi) and d (sum of e^(xi-m)), we can simplify as: - # = (1 - label_smoothing) * H(q, p) + (-sum(x_i * eps) + label_smoothing * (m + logd)) - # Refer to H(q', p) in section 7 of the paper: https://arxiv.org/pdf/1512.00567 - if label_smoothing > 0: - smooth_loss = scaled_x_sum + label_smoothing * (m + tl.log(d)) - loss = loss * (1 - label_smoothing) + smooth_loss - - # 6. Specially handle the i==y case where `dx_y = (softmax(x_y) - (1 - label_smoothing) / N` - vocab_start_idx = rank * n_cols - vocab_end_idx = (rank + 1) * n_cols - if y >= vocab_start_idx: - if y < vocab_end_idx: - X_y = tl.load(X_ptr + y - vocab_start_idx) - # Apply the same conditional scaling logic for the target token - if reduce_loss: - X_y += -(1 - label_smoothing) / (n_non_ignore) - else: - X_y += -(1 - label_smoothing) - tl.store(X_ptr + y - vocab_start_idx, X_y) - - tl.store(loss_ptr, loss) +from transformer_engine.common.triton.cross_entropy import ( + online_softmax_kernel, + cross_entropy_kernel, + element_mul_kernel, +) # The optimal maximum block size depends on your hardware, your kernel, and your dtype MAX_FUSED_SIZE = 65536 // 2 -@triton.jit -def element_mul_kernel( - X_ptr, - X_stride, - grad_output_ptr, - grad_output_stride, - n_cols, - BLOCK_SIZE: tl.constexpr, -): - """ - This function multiplies each element of the tensor pointed by X_ptr with the value pointed by grad_output_ptr. - The multiplication is performed in-place on the tensor pointed by X_ptr. - - Parameters: - X_ptr: Pointer to the input tensor. - X_stride (int): The stride of the input tensor. - grad_output_ptr: Pointer to the gradient output value. - n_cols (int): The number of columns in the input tensor. - BLOCK_SIZE (int): The block size for Triton operations. - """ - - # Get the program ID and convert it to int64 to avoid overflow - program_id = tl.program_id(0).to(tl.int64) - - # Locate the start index - X_ptr += program_id * X_stride - - # Load the gradient output value - grad_output_ptr += program_id * grad_output_stride - grad_output = tl.load(grad_output_ptr) - - # Perform the element-wise multiplication - for i in range(0, n_cols, BLOCK_SIZE): - X_offsets = i + tl.arange(0, BLOCK_SIZE) - X_block = tl.load(X_ptr + X_offsets, mask=X_offsets < n_cols) - tl.store(X_ptr + X_offsets, X_block * grad_output, mask=X_offsets < n_cols) - - def cross_entropy_forward( _input: torch.Tensor, target: torch.Tensor, diff --git a/transformer_engine/pytorch/triton/pad.py b/transformer_engine/pytorch/triton/pad.py index 29b0daf310..790b8277b2 100644 --- a/transformer_engine/pytorch/triton/pad.py +++ b/transformer_engine/pytorch/triton/pad.py @@ -2,63 +2,12 @@ # # See LICENSE for license information. -"""NVFP4 padding kernels - -TODO(ksivamani): Documentation - -""" +"""PyTorch wrapper functions for padding Triton kernels.""" import torch - import triton -import triton.language as tl - - -@triton.autotune( - configs=[ - triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=2), - triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}, num_warps=4, num_stages=2), - triton.Config({"BLOCK_M": 256, "BLOCK_N": 128}, num_warps=8, num_stages=2), - triton.Config({"BLOCK_M": 128, "BLOCK_N": 256}, num_warps=8, num_stages=1), - ], - key=["out_dim0", "out_dim1"], -) -@triton.jit -def zero_pad_kernel( - inp_ptr, - out_ptr, - in_dim0: tl.constexpr, - in_dim1: tl.constexpr, - out_dim0: tl.constexpr, - out_dim1: tl.constexpr, - in_s0, - in_s1, - out_s0, - out_s1, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - """Pads a tensor assuming it's a columnwise scaling inverse.""" - - # tile over OUTPUT coordinates - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) # output rows - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) # output cols - om = offs_m[:, None] - on = offs_n[None, :] - - # edge masking for output - out_mask = (om < out_dim0) & (on < out_dim1) - - # valid input region is simply top-left (no offsets) - in_mask = (om < in_dim0) & (on < in_dim1) - - # load valid input, else zero (masked load touches memory only where True) - x = tl.load(inp_ptr + om * in_s0 + on * in_s1, mask=in_mask, other=0) - # store to output (only within bounds of the output tile) - tl.store(out_ptr + om * out_s0 + on * out_s1, x, mask=out_mask) +from transformer_engine.common.triton.pad import zero_pad_kernel def pad_columnwise_scale_inv(inp: torch.Tensor) -> torch.Tensor: diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 1474a664cf..da22299fe5 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -2,197 +2,23 @@ # # See LICENSE for license information. -"""Permutation kernels written with OpenAI Triton.""" +"""PyTorch wrapper functions for Permutation Triton kernels.""" from typing import Union import torch import triton -import triton.language as tl -from triton.language import core -from triton.language.standard import _log2 -from packaging import version - - -# The following three argsort related kernels are adapted from -# the issue https://github.com/triton-lang/triton/issues/3698 - -get_int_dtype = core.get_int_dtype -if version.parse(triton.__version__) >= version.parse("3.5.0"): - get_int_dtype = triton.constexpr_function(get_int_dtype) - - -@triton.jit -def _compare_and_swap(x, indices, flip, i: tl.constexpr, n_dims: tl.constexpr): - n_outer: tl.constexpr = x.numel >> n_dims - shape: tl.constexpr = [n_outer * (2**i), 2, 2 ** (n_dims - i - 1)] - y = tl.reshape(x, shape) - z = tl.reshape(indices, shape) - - mask = tl.arange(0, 2)[None, :, None] - - l_value = tl.reshape(tl.broadcast_to(tl.sum(y * (1 - mask), 1)[:, None, :], shape), x.shape).to( - x.dtype - ) - r_value = tl.reshape(tl.broadcast_to(tl.sum(y * mask, 1)[:, None, :], shape), x.shape).to( - x.dtype - ) - - l_indice = tl.reshape(tl.broadcast_to(tl.sum(z * (1 - mask), 1)[:, None, :], shape), x.shape) - r_indice = tl.reshape(tl.broadcast_to(tl.sum(z * mask, 1)[:, None, :], shape), x.shape) - - idtype = get_int_dtype(bitwidth=x.dtype.primitive_bitwidth, signed=True) - - il_value = l_value.to(idtype, bitcast=True) - ir_value = r_value.to(idtype, bitcast=True) - ix = x.to(idtype, bitcast=True) - - flag1 = tl.where(((l_value > r_value) ^ flip) != 0, il_value ^ ir_value, tl.zeros_like(ix)) - ret = ix ^ flag1 - flag2 = tl.where(((l_value > r_value) ^ flip) != 0, l_indice ^ r_indice, tl.zeros_like(ix)) - ind = indices ^ flag2 - - return ret.to(x.dtype, bitcast=True), ind - - -@triton.jit -def _bitonic_merge(x, indices, stage: tl.constexpr, order: tl.constexpr, n_dims: tl.constexpr): - n_outer: tl.constexpr = x.numel >> n_dims - tl.static_assert(stage <= n_dims) - """ - order_type 0 == ascending - order_type 1 == descending - order_type 2 == alternating - """ - if order == 2: - shape: tl.constexpr = [n_outer * (2 ** (n_dims - 1 - stage)), 2, 2**stage] - flip = tl.reshape(tl.broadcast_to(tl.arange(0, 2)[None, :, None], shape), x.shape) - else: - flip = tl.full(x.shape, value=order, dtype=tl.int32) - for i in tl.static_range(stage): - x, indices = _compare_and_swap(x, indices, flip, i + (n_dims - stage), n_dims) - return x, indices - - -@triton.jit -def _argsort(x, indices, n_dims: tl.constexpr): - for i in tl.static_range(1, n_dims + 1): - x, indices = _bitonic_merge(x, indices, i, 2 if i < n_dims else 1, n_dims) - return x, indices - - -@triton.jit -def _row_id_map_pass_1_kernel( - # pointers - routing_map_ptr, - row_id_map_ptr, - workspace_ptr, - # sizes - num_tokens, - # strides - stride_routing_map_token, - stride_routing_map_expert, - stride_row_id_map_token, - stride_row_id_map_expert, - # metas - BLOCK_SIZE: tl.constexpr, -): - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - offset = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - expert_token_mask = tl.load( - routing_map_ptr + pid_m * stride_routing_map_expert + offset * stride_routing_map_token, - mask=(offset < num_tokens), - other=0, - ).to(tl.int32) - row_id_within_token_block = tl.cumsum(expert_token_mask) * expert_token_mask - tl.store( - row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token, - row_id_within_token_block, - mask=offset < num_tokens, - ) - n_tokens_per_block = tl.sum(expert_token_mask) - tl.store(workspace_ptr + pid_m * tl.cdiv(num_tokens, BLOCK_SIZE) + pid_n, n_tokens_per_block) - - -@triton.jit -def _row_id_map_pass_2_kernel( - # pointers - row_id_map_ptr, - workspace_ptr, - # sizes - num_tokens, - # strides - stride_row_id_map_token, - stride_row_id_map_expert, - # metas - WORKSPACE_LOAD_WIDTH: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - chunk_idx = pid_m * tl.cdiv(num_tokens, BLOCK_SIZE) + pid_n - offset = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - row_id_within_token_block = tl.load( - row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token, - mask=(offset < num_tokens), - other=0, - ) - - workspace_off = tl.arange(0, WORKSPACE_LOAD_WIDTH) - n_tokens_per_chunk = tl.load(workspace_ptr + workspace_off, mask=workspace_off < chunk_idx) - row_id = tl.where( - row_id_within_token_block == 0, - -1, - row_id_within_token_block + tl.sum(n_tokens_per_chunk) - 1, - ) - tl.store( - row_id_map_ptr + pid_m * stride_row_id_map_expert + offset * stride_row_id_map_token, - row_id, - mask=(offset < num_tokens), - ) - - -@triton.jit -def _row_id_map_pass_3_kernel( - # pointers - row_id_map_ptr, - # sizes - num_experts: tl.constexpr, - # strides - stride_row_id_map_token, - stride_row_id_map_expert, - # metas - LOAD_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - n_dims: tl.constexpr = _log2(LOAD_SIZE) - off = tl.arange(0, LOAD_SIZE) - row_id_map = tl.load( - row_id_map_ptr + pid * stride_row_id_map_token + stride_row_id_map_expert * off, - mask=off < num_experts, - other=-1, - ) - n_routed = tl.sum(tl.where(row_id_map != -1, 1, 0)) - indices = off - sorted_map, indices = _argsort(row_id_map, indices, n_dims=n_dims) - tl.store( - row_id_map_ptr + pid * stride_row_id_map_token + off * stride_row_id_map_expert, - sorted_map, - mask=off < n_routed, - ) - tl.store( - row_id_map_ptr - + pid * stride_row_id_map_token - + (num_experts + off) * stride_row_id_map_expert, - indices, - mask=off < n_routed, - ) - tl.store( - row_id_map_ptr + pid * stride_row_id_map_token + num_experts * 2 * stride_row_id_map_expert, - n_routed, - ) +from transformer_engine.common.triton.permutation import ( + _row_id_map_pass_1_kernel, + _row_id_map_pass_2_kernel, + _row_id_map_pass_3_kernel, + _permute_kernel, + _unpermute_kernel, + _unpermute_bwd_with_merging_probs_kernel, + _make_chunk_sort_map_kernel, + _sort_chunks_by_map_kernel, +) def make_row_id_map( @@ -292,103 +118,6 @@ def make_row_id_map( return row_id_map -@triton.jit -def _permute_kernel( - # pointers - input_ptr, - output_ptr, - row_id_map_ptr, - probs_ptr, - scale_ptr, - permuted_probs_ptr, - permuted_scale_ptr, - # sizes - num_experts: tl.constexpr, - hidden_size: tl.constexpr, - scale_hidden_dim, - # strides - stride_row_id_map_token, - stride_row_id_map_expert, - stride_input_token, - stride_input_hidden, - stride_output_token, - stride_output_hidden, - stride_probs_token, - stride_probs_expert, - stride_scale_token, - stride_scale_hidden, - stride_permuted_probs_token, - stride_permuted_scale_token, - stride_permuted_scale_hidden, - # metas - PERMUTE_PROBS: tl.constexpr, - PERMUTE_SCALE: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - pid_t = tl.program_id(0) - pid_h = tl.program_id(1) - cur_off = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = cur_off < hidden_size - src_row = pid_t.to(tl.int64) - input_off = src_row * stride_input_token + cur_off * stride_input_hidden - inp = tl.load(input_ptr + input_off, mask=mask) - if PERMUTE_SCALE: - mask_scale = cur_off < scale_hidden_dim - scale_off = pid_t * stride_scale_token + cur_off * stride_scale_hidden - scale = tl.load(scale_ptr + scale_off, mask=mask_scale) - n_routed = tl.load( - row_id_map_ptr - + pid_t * stride_row_id_map_token - + num_experts * 2 * stride_row_id_map_expert - ) - for idx in tl.range(n_routed): - dst_row = tl.load( - row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert - ).to(tl.int64) - output_off = dst_row * stride_output_token + cur_off * stride_output_hidden - if PERMUTE_SCALE: - permuted_scale_off = ( - dst_row * stride_permuted_scale_token + cur_off * stride_permuted_scale_hidden - ) - tl.store(permuted_scale_ptr + permuted_scale_off, scale, mask=mask_scale) - if PERMUTE_PROBS: - expert_idx = tl.load( - row_id_map_ptr - + pid_t * stride_row_id_map_token - + (num_experts + idx) * stride_row_id_map_expert - ) - prob_off = pid_t * stride_probs_token + expert_idx * stride_probs_expert - prob = tl.load(probs_ptr + prob_off) - if pid_h == 0: - permuted_prob_off = dst_row * stride_permuted_probs_token - tl.store(permuted_probs_ptr + permuted_prob_off, prob) - if prob == 0.0: - # for routing_map padding - # dst_row != -1 and prob == 0.0 means that this slot is padded - tl.store(output_ptr + output_off, 0.0, mask=mask) - else: - tl.store(output_ptr + output_off, inp, mask=mask) - else: - tl.store(output_ptr + output_off, inp, mask=mask) - - -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}), - ], - key=["hidden_size"], - )(_permute_kernel) -except RuntimeError: - pass - - def permute_with_mask_map( inp: torch.Tensor, row_id_map: torch.Tensor, @@ -468,116 +197,6 @@ def permute_with_mask_map( return output, permuted_scale, permuted_probs -@triton.jit -def _unpermute_kernel( - # pointers - input_ptr, - output_ptr, - row_id_map_ptr, - merging_probs_ptr, - permuted_probs_ptr, - unpermuted_probs_ptr, - # sizes - num_experts: tl.constexpr, - hidden_size: tl.constexpr, - # strides - stride_row_id_map_token, - stride_row_id_map_expert, - stride_input_token, - stride_input_hidden, - stride_output_token, - stride_output_hidden, - stride_merging_probs_token, - stride_merging_probs_expert, - stride_permuted_probs_token, - stride_unpermuted_probs_token, - stride_unpermuted_probs_expert, - # metas - PROBS_LOAD_WIDTH: tl.constexpr, - WITH_MERGING_PROBS: tl.constexpr, - PERMUTE_PROBS: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - data_type = input_ptr.dtype.element_ty - compute_type = tl.float32 - - pid_t = tl.program_id(0) - pid_h = tl.program_id(1) - current_offset = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = current_offset < hidden_size - if PERMUTE_PROBS: - # write 0.0 to probs_grad that are not routed - if pid_h == 0: - map_load_off = tl.arange(0, PROBS_LOAD_WIDTH) - unpermuted_prob_off = ( - pid_t * stride_unpermuted_probs_token - + stride_unpermuted_probs_expert * map_load_off - ) - tl.store( - unpermuted_probs_ptr + unpermuted_prob_off, 0.0, mask=map_load_off < num_experts - ) - accumulator = tl.zeros((BLOCK_SIZE,), dtype=compute_type) - n_routed = tl.load( - row_id_map_ptr - + pid_t * stride_row_id_map_token - + num_experts * 2 * stride_row_id_map_expert - ) - for idx in tl.range(n_routed): - src_row = tl.load( - row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert - ).to(tl.int64) - input_off = src_row * stride_input_token + current_offset * stride_input_hidden - inp = tl.load(input_ptr + input_off, mask=mask) - inp = inp.to(compute_type) - if WITH_MERGING_PROBS: - expert_idx = tl.load( - row_id_map_ptr - + pid_t * stride_row_id_map_token - + (num_experts + idx) * stride_row_id_map_expert - ) - merging_prob_off = ( - pid_t * stride_merging_probs_token + expert_idx * stride_merging_probs_expert - ) - merging_prob = tl.load(merging_probs_ptr + merging_prob_off).to(compute_type) - inp *= merging_prob - accumulator += inp - if PERMUTE_PROBS: - if pid_h == 0: - expert_idx = tl.load( - row_id_map_ptr - + pid_t * stride_row_id_map_token - + (num_experts + idx) * stride_row_id_map_expert - ) - unpermuted_prob_off = ( - pid_t * stride_unpermuted_probs_token - + expert_idx * stride_unpermuted_probs_expert - ) - permuted_prob_off = src_row * stride_permuted_probs_token - prob = tl.load(permuted_probs_ptr + permuted_prob_off) - tl.store(unpermuted_probs_ptr + unpermuted_prob_off, prob) - accumulator = accumulator.to(data_type) - dst_row = pid_t.to(tl.int64) - output_off = dst_row * stride_output_token + current_offset * stride_output_hidden - tl.store(output_ptr + output_off, accumulator, mask=mask) - - -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}), - ], - key=["hidden_size"], - )(_unpermute_kernel) -except RuntimeError: - pass - - def unpermute_with_mask_map( inp: torch.Tensor, row_id_map: torch.Tensor, @@ -644,110 +263,6 @@ def unpermute_with_mask_map( return output, unpermuted_probs -@triton.jit -def _unpermute_bwd_with_merging_probs_kernel( - # pointers - fwd_output_grad_ptr, - fwd_input_grad_ptr, - fwd_input_ptr, - merging_probs_ptr, - merging_probs_grad_ptr, - row_id_map_ptr, - # sizes - num_experts: tl.constexpr, - hidden_size: tl.constexpr, - # strides - stride_row_id_map_token, - stride_row_id_map_expert, - stride_fwd_output_grad_token, - stride_fwd_output_grad_hidden, - stride_fwd_input_grad_token, - stride_fwd_input_grad_hidden, - stride_fwd_input_token, - stride_fwd_input_hidden, - stride_merging_probs_token, - stride_merging_probs_expert, - stride_merging_probs_grad_token, - stride_merging_probs_grad_expert, - # metas - PROBS_LOAD_WIDTH: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - data_type = fwd_output_grad_ptr.dtype.element_ty - compute_type = tl.float32 - - pid = tl.program_id(0) - map_load_off = tl.arange(0, PROBS_LOAD_WIDTH) - token_probs_grad_off = ( - pid * stride_merging_probs_grad_token + stride_merging_probs_grad_expert * map_load_off - ) - tl.store(merging_probs_grad_ptr + token_probs_grad_off, 0.0, mask=map_load_off < num_experts) - n_routed = tl.load( - row_id_map_ptr + pid * stride_row_id_map_token + num_experts * 2 * stride_row_id_map_expert - ) - for idx in tl.range(n_routed): - dst_row = tl.load( - row_id_map_ptr + pid * stride_row_id_map_token + idx * stride_row_id_map_expert - ).to(tl.int64) - expert_idx = tl.load( - row_id_map_ptr - + pid * stride_row_id_map_token - + (num_experts + idx) * stride_row_id_map_expert - ) - prob_grad_accum = tl.zeros((BLOCK_SIZE,), dtype=compute_type) - current_start = 0 - while current_start < hidden_size: - current_offset = current_start + tl.arange(0, BLOCK_SIZE) - mask = current_offset < hidden_size - src_row = pid.to(tl.int64) - input_off = ( - src_row * stride_fwd_output_grad_token - + current_offset * stride_fwd_output_grad_hidden - ) - inp = tl.load(fwd_output_grad_ptr + input_off, mask=mask) - inp = inp.to(compute_type) - merging_prob_off = ( - pid * stride_merging_probs_token + expert_idx * stride_merging_probs_expert - ) - merging_prob = tl.load(merging_probs_ptr + merging_prob_off).to(compute_type) - output = inp * merging_prob - output = output.to(data_type) - output_off = ( - dst_row * stride_fwd_input_grad_token - + current_offset * stride_fwd_input_grad_hidden - ) - tl.store(fwd_input_grad_ptr + output_off, output, mask=mask) - - fwd_input_off = ( - dst_row * stride_fwd_input_token + current_offset * stride_fwd_input_hidden - ) - fwd_input = tl.load(fwd_input_ptr + fwd_input_off, mask=mask) - prob_grad_accum += fwd_input.to(compute_type) * inp - current_start += BLOCK_SIZE - probs_grad = tl.sum(prob_grad_accum).to(merging_probs_grad_ptr.dtype.element_ty) - probs_grad_off = ( - pid * stride_merging_probs_grad_token + expert_idx * stride_merging_probs_grad_expert - ) - tl.store(merging_probs_grad_ptr + probs_grad_off, probs_grad) - - -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}), - ], - key=["hidden_size"], - )(_unpermute_bwd_with_merging_probs_kernel) -except RuntimeError: - pass - - def unpermute_with_mask_map_bwd_with_merging_probs( fwd_output_grad: torch.Tensor, row_id_map: torch.Tensor, @@ -813,47 +328,6 @@ def unpermute_with_mask_map_bwd_with_merging_probs( return act_grad, merging_probs_grad -@triton.jit -def _make_chunk_sort_map_kernel( - # pointers - split_sizes_ptr, - sorted_indices_ptr, - dst_rows_ptr, - # sizes - num_splits: tl.constexpr, - # metas - IDX_LOAD_WIDTH: tl.constexpr, -): - pid = tl.program_id(0) - - load_split_offset = tl.arange(0, IDX_LOAD_WIDTH) - sorted_indices = tl.load( - sorted_indices_ptr + load_split_offset, mask=load_split_offset < num_splits - ) - - # get chunk idx of the current token in the input tensor - input_split_sizes = tl.load( - split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0 - ).to(tl.int32) - input_split_sizes_cumsum = tl.cumsum(input_split_sizes) - input_split_sizes_mask = tl.where(input_split_sizes_cumsum <= pid, 1, 0) - input_chunk_idx = tl.sum(input_split_sizes_mask) - input_split_sizes_presum = tl.sum(input_split_sizes * input_split_sizes_mask) - in_chunk_offset = pid - input_split_sizes_presum - - # get chunk idx of the current token in the output tensor - output_chunk_mask = tl.where(sorted_indices == input_chunk_idx, 1, 0) - output_chunk_idx = tl.argmax(output_chunk_mask, axis=-1) - - # make row_id_map - output_split_sizes = tl.load( - split_sizes_ptr + sorted_indices, mask=load_split_offset < num_splits - ).to(tl.int32) - output_pre_split_sizes = tl.where(load_split_offset < output_chunk_idx, output_split_sizes, 0) - dst_row = tl.sum(output_pre_split_sizes) + in_chunk_offset - tl.store(dst_rows_ptr + pid, dst_row) - - def make_chunk_sort_map( split_sizes: torch.Tensor, sorted_indices: torch.Tensor, @@ -886,67 +360,6 @@ def make_chunk_sort_map( return row_id_map -@triton.jit -def _sort_chunks_by_map_kernel( - # pointers - input_ptr, - output_ptr, - row_id_map_ptr, - probs_ptr, - permuted_probs_ptr, - # sizes - hidden_size: tl.constexpr, - # strides - stride_input_token, - stride_input_hidden, - stride_output_token, - stride_output_hidden, - stride_probs_token, - stride_permuted_probs_token, - # metas - PERMUTE_PROBS: tl.constexpr, - BLOCK_SIZE: tl.constexpr, - FORWARD: tl.constexpr, -): - pid_t = tl.program_id(0) - pid_h = tl.program_id(1) - if FORWARD: - src_row = pid_t.to(tl.int64) - dst_row = tl.load(row_id_map_ptr + pid_t).to(tl.int64) - else: - src_row = tl.load(row_id_map_ptr + pid_t).to(tl.int64) - dst_row = pid_t.to(tl.int64) - current_offset = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = current_offset < hidden_size - input_offsets = src_row * stride_input_token + current_offset * stride_input_hidden - output_offsets = dst_row * stride_output_token + current_offset * stride_output_hidden - inp = tl.load(input_ptr + input_offsets, mask=mask) - tl.store(output_ptr + output_offsets, inp, mask=mask) - if PERMUTE_PROBS: - if pid_h == 0: - prob_off = src_row * stride_probs_token - prob = tl.load(probs_ptr + prob_off) - permuted_prob_off = dst_row * stride_permuted_probs_token - tl.store(permuted_probs_ptr + permuted_prob_off, prob) - - -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}), - ], - key=["hidden_size"], - )(_sort_chunks_by_map_kernel) -except RuntimeError: - pass - - def sort_chunks_by_map( inp: torch.Tensor, row_id_map: torch.Tensor, From 7a5859834084c8c9cc88bc1711270e0ee232e6bc Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 10 Nov 2025 10:59:22 -0800 Subject: [PATCH 053/521] [JAX] Fused layers argument default values changed (#2347) * Changing default activations in MLP, TransformerLayer, dropout rate after FC1 to 0, and return_layernorm_output to False Signed-off-by: tdophung * Fixing the failing tests by hard coding arguments to the previous values instead of relying on newer default values Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_distributed_layernorm_mlp.py | 2 ++ tests/jax/utils.py | 12 ++++++------ transformer_engine/jax/flax/module.py | 16 ++++++++-------- transformer_engine/jax/flax/transformer.py | 8 ++++---- 4 files changed, 20 insertions(+), 18 deletions(-) diff --git a/tests/jax/test_distributed_layernorm_mlp.py b/tests/jax/test_distributed_layernorm_mlp.py index 339097e9cc..667840da2c 100644 --- a/tests/jax/test_distributed_layernorm_mlp.py +++ b/tests/jax/test_distributed_layernorm_mlp.py @@ -389,6 +389,7 @@ def _test_layernorm_mlp( intermediate_dim=INTERMEDIATE, activations=activation_type, use_bias=use_bias, + return_layernorm_output=True, ) params_single = ln_mlp_single.init(init_rngs, x, deterministic=True) mlp_out_single, ln_out_single = ln_mlp_single.apply( @@ -417,6 +418,7 @@ def _test_layernorm_mlp( dot_1_input_axes=DOT_1_INPUT_AXES, dot_2_input_axes=DOT_2_INPUT_AXES, name="mlp", + return_layernorm_output=True, ) params_sharded = ln_mlp_sharded.init(init_rngs, x, deterministic=True) mlp_out_sharded, ln_out_sharded = ln_mlp_sharded.apply( diff --git a/tests/jax/utils.py b/tests/jax/utils.py index c28e68a15f..bbe8e65829 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -364,9 +364,9 @@ class MlpBlock(nn.Module): transpose_batch_sequence: bool intermediate_dim: int = 2048 - activations: Sequence[Union[str, Callable]] = ("relu",) + activations: Sequence[Union[str, Callable]] = ("gelu",) kernel_init: Initializer = None - intermediate_dropout_rate: float = 0.1 + intermediate_dropout_rate: float = 0.0 intermediate_dropout_dims: Sequence[int] = () use_bias: bool = False dtype: Any = jnp.float32 @@ -1035,14 +1035,14 @@ class EncoderLayer(nn.Module): hidden_dropout: float = 0.1 hidden_dropout_dims: Sequence[int] = () attention_dropout: float = 0.1 - intermediate_dropout: float = 0.1 + intermediate_dropout: float = 0.0 intermediate_dropout_dims: Sequence[int] = () transpose_batch_sequence: bool = True float32_attention_logits: bool = False scale_attn_logits: bool = False scaled_query_init: bool = True mlp_dim: int = 2048 - mlp_activations: Sequence[str] = ("relu",) + mlp_activations: Sequence[str] = ("gelu",) use_bias: bool = False dtype: Any = jnp.float32 apply_residual_connection_post_layernorm: bool = False @@ -1199,14 +1199,14 @@ class DecoderLayer(nn.Module): hidden_dropout: float = 0.1 hidden_dropout_dims: Sequence[int] = () attention_dropout: float = 0.1 - intermediate_dropout: float = 0.1 + intermediate_dropout: float = 0.0 intermediate_dropout_dims: Sequence[int] = () transpose_batch_sequence: bool = True float32_attention_logits: bool = False scale_attn_logits: bool = False scaled_query_init: bool = True mlp_dim: int = 2048 - mlp_activations: Sequence[str] = ("relu",) + mlp_activations: Sequence[str] = ("gelu",) use_bias: bool = False dtype: Any = jnp.float32 apply_residual_connection_post_layernorm: bool = False diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index c54ecb236f..33ea610985 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -597,7 +597,7 @@ class LayerNormDenseGeneral(TransformerEngineBase): bias_axes: Tuple[str, ...], default = () The name of axes used to shard bias with a corresponding mesh, only used when :attr:`use_bias=True`. - return_layernorm_output: bool, default = True + return_layernorm_output: bool, default = False Indicate whether to return the output of layer normalization. If set False, return None as the second tensor in outputs. enable_low_rank_adaptation: bool, default = False @@ -644,7 +644,7 @@ class LayerNormDenseGeneral(TransformerEngineBase): use_bias: bool = False bias_init: Initializer = nn.initializers.zeros bias_axes: Tuple[str, ...] = () - return_layernorm_output: bool = True + return_layernorm_output: bool = False enable_low_rank_adaptation: bool = False low_rank_adaptation_dim: int = 32 low_rank_adaptation_alpha: float = None @@ -891,10 +891,10 @@ class LayerNormMLP(TransformerEngineBase): The name of axes used to shard bias with a corresponding mesh for the weight of the second dense layer transformation. Only used when :attr:`use_bias=True`. - return_layernorm_output: bool, default = True + return_layernorm_output: bool, default = False Indicate whether to return the output of layer normalization. If set False, return None as the second tensor in outputs. - activations: Sequence[Union[str, Callable]], default = ('relu',) + activations: Sequence[Union[str, Callable]], default = ('gelu',) The sequence of activation functions to apply after the first dense layer transformation. Each activation has its own transformation layer. activation_params: dict, default = None @@ -903,7 +903,7 @@ class LayerNormMLP(TransformerEngineBase): need additional parameters. intermediate_dropout_rng_name: str, default = 'dropout' The key in given RNGs via flax.linen.Module.apply that for generating Dropout masks. - intermediate_dropout_rate: float, default = 0.1 + intermediate_dropout_rate: float, default = 0.0 Dropout probability for the dropout op after the :attr:`activations`. intermediate_hidden_dropout_dims: Sequence[int], default = () Dimensions that will share the same dropout mask for hidden @@ -959,11 +959,11 @@ class LayerNormMLP(TransformerEngineBase): bias_init: Initializer = nn.initializers.zeros bias_axes_1: Tuple[str, ...] = ("act", "mlp") bias_axes_2: Tuple[str, ...] = ("embed",) - return_layernorm_output: bool = True - activations: Sequence[Union[str, Callable]] = ("relu",) + return_layernorm_output: bool = False + activations: Sequence[Union[str, Callable]] = ("gelu",) activation_params: dict = None intermediate_dropout_rng_name: str = "dropout" - intermediate_dropout_rate: float = 0.1 + intermediate_dropout_rate: float = 0.0 intermediate_hidden_dropout_dims: Sequence[int] = () enable_low_rank_adaptation: bool = False low_rank_adaptation_dim: int = 32 diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 86af6cf499..d096e7997c 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -1620,7 +1620,7 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods Dimensions that will share the same dropout mask for hidden attention_dropout: float, default = 0.1 Dropout probability for the dropout op during multi-head attention. - intermediate_dropout: float, default = 0.1 + intermediate_dropout: float, default = 0.0 Dropout probability for the dropout op after FC1 layer. intermediate_dropout_dims: Sequence[int], default = () Dimensions that will share the same dropout mask for hidden after FC1 layer. @@ -1635,7 +1635,7 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal') Used for initializing weights of FC1 and FC2 layers. It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). - mlp_activations: Sequence[str], default = ('relu', ) + mlp_activations: Sequence[str], default = ('gelu', ) The sequence of activation functions to apply after the first linear transformation. Each activation has its own transformation layer. mlp_activation_params: dict = None @@ -1755,12 +1755,12 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods hidden_dropout: float = 0.1 hidden_dropout_dims: Sequence[int] = () attention_dropout: float = 0.1 - intermediate_dropout: float = 0.1 + intermediate_dropout: float = 0.0 intermediate_dropout_dims: Sequence[int] = () dropout_rng_name: str = "dropout" mha_kernel_init: Initializer = None mlp_kernel_init: Initializer = None - mlp_activations: Sequence[str] = ("relu",) + mlp_activations: Sequence[str] = ("gelu",) mlp_activation_params: dict = None use_bias: bool = False bias_init: Initializer = nn.initializers.zeros From 29537c96d06c4f3965fd3bcc668810dbb4245aaf Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 10 Nov 2025 19:31:55 -0800 Subject: [PATCH 054/521] [PyTorch] FSDP2 Support for TE (#2245) * fix for float8 tensor fsdp2 training Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * zeros_like should return fp32 for fsdp2 to work Signed-off-by: Varun Thumbe * minor cleanup Signed-off-by: Varun Thumbe * fix unsharded weights not releasing memory Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * implement using fsdp preallgather and postallgather functions Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * FSDP2 works on Hopper/L40 Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor comment Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * some fixes for fp8 + handwavy changes for mxfp8 Signed-off-by: Varun Thumbe * only transpose saved for backward pass allgather in case of L40/Hoppergst Signed-off-by: Varun Thumbe * missed minor change to hopper use-case Signed-off-by: Varun Thumbe * communicate only required data in mxfp8, fix for updating weight usages when required instead of doing upfront in fwd pass Signed-off-by: Varun Thumbe * changes for meta Dtensors for weights and better all gather data handling in fsdp hook functions Signed-off-by: Varun Thumbe * better solution to figure out forward pass in FSDP2 Signed-off-by: Varun Thumbe * adress review comments Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * everything functioning except hack for transformerlayer Signed-off-by: Varun Thumbe * fix merge conflict Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert change of commit id for cudnnt-frontend Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * unnecessary change Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor issues with linting, add some comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor stuff Signed-off-by: Varun Thumbe * revert space removal Add default usage handling for rowwise and columnwise data. Signed-off-by: vthumbe1503 * fix the fsdp state collection issue, and minor review comments addressing Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert change for dgrad redundant computation Signed-off-by: Varun Thumbe * bug: get fsdp param group's training state instead of root training state; 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 * address review comments Signed-off-by: Varun Thumbe * address review comments Signed-off-by: Varun Thumbe * address coderabbit 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 * address review comments Signed-off-by: Varun Thumbe * adress review comments; fix fp8 allgather test to do after fsdp lazy init 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 * remove detach Signed-off-by: Varun Thumbe * do what makes sense Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/tensor/float8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * address review comments Signed-off-by: Varun Thumbe * 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 * address review comments Signed-off-by: Varun Thumbe * adress 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 * have better dtype for fsdp_post_all_gather arguments Signed-off-by: Varun Thumbe * minor comment Signed-off-by: Varun Thumbe * improve comment Signed-off-by: Varun Thumbe * fix the error in CI Signed-off-by: Varun Thumbe * minor comment add Signed-off-by: Varun Thumbe * accidentally removed view function Signed-off-by: Varun Thumbe * fix minor bug for h100 Signed-off-by: Varun Thumbe * minor addition Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * implement padding removal/addition for allgather 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 * address review comments Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint error Signed-off-by: Varun Thumbe * adress 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 * improve the reset parameter logic for dtensors Signed-off-by: Varun Thumbe * other cosmetic changes Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cosmetic changes Signed-off-by: Varun Thumbe * cosmetic changes Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/module/layernorm_linear.py Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 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> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/pytorch/distributed/run_fsdp2_model.py | 321 +++++++++++++---- tests/pytorch/distributed/test_torch_fsdp2.py | 18 +- transformer_engine/pytorch/distributed.py | 37 ++ transformer_engine/pytorch/module/base.py | 43 ++- .../pytorch/module/grouped_linear.py | 24 +- .../pytorch/module/layernorm_linear.py | 19 +- .../pytorch/module/layernorm_mlp.py | 20 +- transformer_engine/pytorch/module/linear.py | 12 +- .../pytorch/quantized_tensor.py | 12 +- .../pytorch/tensor/float8_tensor.py | 230 +++++++++++- .../pytorch/tensor/mxfp8_tensor.py | 341 +++++++++++++++++- 11 files changed, 928 insertions(+), 149 deletions(-) diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/run_fsdp2_model.py index d3f8c82baa..c343299242 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/run_fsdp2_model.py @@ -9,57 +9,73 @@ import argparse import transformer_engine.pytorch as te -from transformer_engine.common.recipe import Format, DelayedScaling +from transformer_engine.common.recipe import ( + Format, + DelayedScaling, + Float8CurrentScaling, + MXFP8BlockScaling, +) import torch import torch.distributed as dist +from torch.distributed.tensor import DTensor import torch.nn.functional as F from torch import nn, optim from torch.distributed import DeviceMesh from torch.distributed._composable.fsdp import fully_shard from torch.distributed.device_mesh import init_device_mesh +from transformer_engine.pytorch import QuantizedTensor from contextlib import nullcontext +LOCAL_RANK = None -class SimpleNet(nn.Module): - def __init__(self, input_size, hidden_size, output_size): - super(SimpleNet, self).__init__() - self.fc1 = te.Linear(input_size, hidden_size) - self.fc2 = te.Linear(hidden_size, output_size) - def forward(self, x): - x = F.relu(self.fc1(x)) - x = self.fc2(x) - return x - - -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items()} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) +def dist_print(msg): + if LOCAL_RANK == 0: + print(msg) def _parse_args(argv=None, namespace=None): parser = argparse.ArgumentParser(description="Toy example for debugging fully_shard()") - parser.add_argument("--input-size", type=int, default=2048, help="Input size for the model") - parser.add_argument("--hidden-size", type=int, default=2048, help="Hidden layer size") - parser.add_argument("--output-size", type=int, default=2048, help="Output size for the model") - parser.add_argument("--batch-size", type=int, default=2048, help="Output size for the model") + parser.add_argument("--num-heads", type=int, default=8, help="Number of attn. heads") + parser.add_argument("--head-dim", type=int, default=64, help="Attention head size") + parser.add_argument("--batch-size", type=int, default=16, help="Batch size of input") + parser.add_argument("--seq-length", type=int, default=128, help="Sequence length of input") + parser.add_argument("--params-dtype", type=str, default="float32", help="Parameter dtype.") parser.add_argument( "--fp8-init", action="store_true", default=False, help="Initialize primary weights in FP8." ) + parser.add_argument( + "--recipe", + type=str, + default="mx_fp8_block_scaling", + help="Quantizer type.", + choices=["delayed_scaling", "current_scaling", "mx_fp8_block_scaling"], + ) + parser.add_argument( + "--layer-type", + type=str, + default="TransformerLayer", + choices=[ + "Linear", + "LayerNormLinear", + "LayerNormMLP", + "MultiheadAttention", + "TransformerLayer", + ], + help="Transformer Engine layer type", + ) + parser.add_argument("--num-layers", type=int, default=4, help="Number of layers in the model") parser.add_argument( "--iter", type=int, default=10, help="Number of iterations for forward pass" ) + parser.add_argument( + "--device", + type=str, + default="meta", + help="Device to run the model on.", + choices=["cuda", "meta"], + ) parser.add_argument("--seed", type=int, default=42, help="RNG seed.") # Adding hsdp_dim as a list argument, comma-separated parser.add_argument( @@ -74,10 +90,170 @@ def _parse_args(argv=None, namespace=None): return args -sub_modules_to_wrap = [te.Linear] +## Methods to help initialize the TE model in an FSDP2 setting +## with required configurations based on command line args +def get_te_layer_from_string(layer_name): + te_layer_types = [ + te.Linear, + te.LayerNormLinear, + te.LayerNormMLP, + te.MultiheadAttention, + te.TransformerLayer, + ] + te_layer_names = [layer.__name__ for layer in te_layer_types] + te_layer_map = dict(zip([name.lower() for name in te_layer_names], te_layer_types)) + if layer_name.lower() not in te_layer_map.keys(): + raise argparse.ArgumentTypeError( + f'"{layer_name}" is not a valid Transformer Engine layer, ' + f"please choose layer from {te_layer_names}." + ) + return te_layer_map[layer_name.lower()] + + +def get_recipe_from_string(recipe, fp8_format=Format.HYBRID): + if recipe == "delayed_scaling": + return DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max") + elif recipe == "current_scaling": + return Float8CurrentScaling(fp8_format=fp8_format) + elif recipe == "mx_fp8_block_scaling": + return MXFP8BlockScaling(fp8_format=fp8_format) + else: + raise ValueError(f"Unknown quantizer type: {recipe}") + + +def init_te_model(config): + hidden_size = config.num_heads * config.head_dim + args = [hidden_size, hidden_size] + inp_shape = [config.seq_length, config.batch_size, hidden_size] + out_shape = [config.seq_length, config.batch_size, hidden_size] + if config.params_dtype == "float16": + params_dtype = torch.float16 + elif config.params_dtype == "bfloat16": + params_dtype = torch.bfloat16 + else: + params_dtype = torch.float32 + kwargs = { + "params_dtype": params_dtype, + } + kwargs["device"] = config.device + + layer_type = get_te_layer_from_string(config.layer_type) + # We are creating model in a way so that we can test both reshard_after_forward=True/False cases. + # more details below. + if layer_type in [te.MultiheadAttention, te.TransformerLayer]: + # For this case, we are creating a model that resemebles production use-cases + # wherein there are mltiple TransformerLayers in the model. And we would need + # to shard each transformer layer. Since each transformer layer is not a root module, + # FSDP2's fully_shard assigns reshard_after_forward=False for all parameters of the model. + args[1] *= 4 # FFN hidden size + args.append(config.num_heads) + kwargs["fuse_qkv_params"] = True + if layer_type is te.MultiheadAttention: + kwargs["input_layernorm"] = True + model = nn.Sequential(*[layer_type(*args, **kwargs) for _ in range(config.num_layers)]) + elif layer_type == te.LayerNormLinear: + # For this case, we are creating a model with just one LayerNormLinear layer + # so that the model itself is a root module, and FSDP2's fully_shard assigns + # reshard_after_forward=True for the parameters of these model. + args[1] *= 3 # QKV projection + out_shape[-1] *= 3 + model = layer_type(*args, **kwargs) + else: + model = layer_type(*args, **kwargs) + + return model, inp_shape, out_shape + + +def get_device_mesh(world_size, sharding_dims): + dist_print(f"sharding-dims:{sharding_dims}") + device_ids = list(range(world_size)) + if sharding_dims is None: # FSDP + mesh = DeviceMesh("cuda", device_ids) + elif len(sharding_dims) == 1: + assert sharding_dims[0] == world_size + mesh = DeviceMesh("cuda", device_ids) + elif len(sharding_dims) == 2: # HSDP + assert sharding_dims[0] * sharding_dims[1] == world_size + mesh = init_device_mesh( + "cuda", + (sharding_dims[0], sharding_dims[1]), + mesh_dim_names=("replicate", "shard"), + ) + else: + assert False + return mesh + + +def shard_model_with_fsdp2(model, mesh): + for child in model.children(): + fully_shard(child, mesh=mesh) + fully_shard(model, mesh=mesh) + return model + + +#### Methods to save the custom attributes of QuantizedTensors before sharding +#### them with FSDP2, and restore them after sharding. +def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + if isinstance(param, QuantizedTensor): + # Ignore FP8 metadata attributes. Otherwise we will save duplicate copies + # for data/transpose FP8 tensors on top of FP8 tensors that FSDP2 will save. + ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] + else: + ignore_keys = [] + attrs = vars(param) + custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} + return custom_attrs + + +def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) + + +@torch.no_grad() +def test_fp8_fsdp2_allgather(model): + # Do manual allgather in fp32 and match against fp8 allgather done + # with fsdp2 + # FP32 manual weight allgather + fp32_allgathered_params = {} + for name, param in model.named_parameters(): + assert isinstance(param, DTensor) + local_tensor = param._local_tensor + device_mesh = param.device_mesh + dist_group = ( + device_mesh.get_group(mesh_dim="shard") + if device_mesh.ndim > 1 + else device_mesh.get_group() + ) + # Perform manual allgather on local_tensor. zeros_like will create hp tensor since torch_dispatch + # for local_tensor will go down the dequantization route. + gathered_tensor = [ + torch.zeros_like(local_tensor) for _ in range(dist.get_world_size(group=dist_group)) + ] + dist.all_gather(gathered_tensor, local_tensor.dequantize(), group=dist_group) + full_tensor = torch.cat(gathered_tensor, dim=0) + fp32_allgathered_params[name] = full_tensor + # FP8 allgather using FSDP2 + for module in model.modules(): + # Not all modules are wrapped/sharded with FSDP2. + if hasattr(module, "unshard"): + module.unshard() + # Make sure allgathered parameters match exactly + for name, param in model.named_parameters(): + assert torch.allclose(param.dequantize(), fp32_allgathered_params[name]) + # Revert model to original sharded state + for module in model.modules(): + # Not all modules are wrapped/sharded with FSDP2. + if hasattr(module, "reshard"): + module.reshard() def _train(args): + global LOCAL_RANK assert "TORCHELASTIC_RUN_ID" in os.environ WORLD_RANK = int(os.getenv("RANK", "0")) WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) @@ -103,74 +279,69 @@ def _train(args): # FP8 Configuration fp8_format = Format.HYBRID - fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max") - - # Create build context manager - if args.fp8_init: - from transformer_engine.pytorch import quantized_model_init + fp8_recipe = get_recipe_from_string(args.recipe, fp8_format) - build_model_context = quantized_model_init() + build_model_context_args = {} + if not args.fp8_init: + # Build model context (FP8 init) + build_model_context = nullcontext else: - build_model_context = nullcontext() + from transformer_engine.pytorch import fp8_model_init - # Build the model with the specified context - with build_model_context: - model = SimpleNet(args.input_size, args.hidden_size, args.output_size) + build_model_context = fp8_model_init + build_model_context_args["enabled"] = True + build_model_context_args["recipe"] = fp8_recipe - # Move the model to the correct device - model.to(device) + dist_print(f"Memory before model init: {torch.cuda.memory_allocated(device)/1e6} MB") + # Create the model on the meta/cuda device as per args + with build_model_context(**build_model_context_args): + model, inp_shape, out_shape = init_te_model(args) + dist_print( + f"Memory after model init on device {args.device}:" + f" {torch.cuda.memory_allocated(device)/1e6} MB" + ) - if LOCAL_RANK == 0: - print(f"Rank {LOCAL_RANK}: Applying FSDP fully_shard() to the model...") # Creating a DeviceMesh for fully_shard world_size = int(WORLD_SIZE) - device_ids = list(range(world_size)) - if LOCAL_RANK == 0: - print(f"sharding-dims:{args.sharding_dims}") # Setup the sharding mesh for FSDP/HSDP - if args.sharding_dims == None: # FSDP - mesh = DeviceMesh("cuda", device_ids) - elif len(args.sharding_dims) == 1: - assert args.sharding_dims[0] == device_ids[-1] + 1 - mesh = DeviceMesh("cuda", device_ids) - elif len(args.sharding_dims) == 2: # HSDP - assert args.sharding_dims[0] * args.sharding_dims[1] == device_ids[-1] + 1 - mesh = init_device_mesh( - "cuda", - (args.sharding_dims[0], args.sharding_dims[1]), - mesh_dim_names=("replicate", "shard"), - ) - else: - assert False - - # Apply FSDP/HSDP + mesh = get_device_mesh(world_size, args.sharding_dims) custom_attrs = save_custom_attrs(model) - for sub_module in model.modules(): - if any( - isinstance(sub_module, sub_module_to_wrap) for sub_module_to_wrap in sub_modules_to_wrap - ): - fully_shard(sub_module, mesh=mesh) - fully_shard(model, mesh=mesh) + model = shard_model_with_fsdp2(model, mesh) restore_custom_attrs(model, custom_attrs) + # model now has DTensors as its parameters + + if args.device == "meta": + # After FSDP2 has been applied, materialize and initialize the sharded parameters + # TE base.py's reset_parameters() handles DTensors with FP8 initialization + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + dist_print(f" Sharded parameters materialized and initialized on cuda device.") + + dist_print( + f"FSDP2 model in cuda, memory allocated: {torch.cuda.memory_allocated(device)/1e6} MB" + ) optimizer = optim.Adam(model.parameters(), lr=1e-3) for iteration in range(args.iter): # Zero the parameter gradients optimizer.zero_grad() - input_data = torch.randn(args.batch_size, args.input_size).to(device) + input_data = torch.randn(inp_shape).to(device) with te.autocast(enabled=True, recipe=fp8_recipe): output = model(input_data) - target = torch.randn(args.batch_size, args.output_size).to(device) + target = torch.randn(out_shape).to(device) loss = F.mse_loss(output, target) loss.backward() optimizer.step() - if LOCAL_RANK == 0: - print(f"Rank {LOCAL_RANK}: Iteration {iteration} completed.") + dist_print(f"Iteration {iteration} completed with loss {loss.item()}") + + # Some of the FSDP states are lazy initialized during FSDP forward pass + # so testing fp8 allgather at the end of the training loop. + if args.fp8_init: + test_fp8_fsdp2_allgather(model) dist.destroy_process_group() - if LOCAL_RANK == 0: - print(f"Rank {LOCAL_RANK}: Done...") return 0 diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 8fe4e8bc7c..91d6fc6ed1 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -12,22 +12,26 @@ fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) - +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) NUM_PROCS: int = torch.cuda.device_count() -def _run_test(fp_init, sharding_dims): +def _run_test(fp_init, sharding_dims, recipe, layer_type): test_path = Path(__file__).parent.resolve() / "run_fsdp2_model.py" test_cmd = ["torchrun", f"--nproc_per_node={NUM_PROCS}", str(test_path)] if fp_init: test_cmd += ["--fp8-init"] + if len(sharding_dims) == 1: test_cmd += ["--sharding-dims", str(sharding_dims[0])] elif len(sharding_dims) == 2: test_cmd += ["--sharding-dims", str(sharding_dims[0]), str(sharding_dims[1])] else: assert False + test_cmd += ["--recipe", recipe] + test_cmd += ["--layer-type", layer_type] + result = subprocess.run(test_cmd, env=os.environ, check=True) @@ -36,16 +40,20 @@ def _run_test(fp_init, sharding_dims): @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") @pytest.mark.parametrize("sharding_dims", ([NUM_PROCS], [2, NUM_PROCS // 2])) @pytest.mark.parametrize("fp8_init", (False, True)) -def test_distributed(fp8_init, sharding_dims): +@pytest.mark.parametrize("recipe", ("delayed_scaling", "current_scaling", "mx_fp8_block_scaling")) +@pytest.mark.parametrize("layer_type", ("LayerNormLinear", "TransformerLayer")) +def test_distributed(fp8_init, sharding_dims, recipe, layer_type): # Skip invalid configurations if torch.cuda.device_count() < 4: pytest.skip("FSDP2 test requires at least 4 GPUs") - if fp8_init and not fp8_available: + if recipe == "mx_fp8_block_scaling" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + elif not fp8_available: pytest.skip(reason_for_no_fp8) - _run_test(fp8_init, sharding_dims) + _run_test(fp8_init, sharding_dims, recipe, layer_type) def test_dummy() -> None: diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index e938509e58..620ea83013 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1886,6 +1886,43 @@ def allreduce( return inp, handle +def _get_module_fsdp_state(module): + """ + If module is an FSDP module, return its _FSDPState. + Otherwise, return the _FSDPState of the closest parent FSDP module + in the module hierarchy the module belongs to. + """ + + if hasattr(module, "_get_fsdp_state"): + # this will return correct fsdp state if module itself is an fsdp module + fsdp_state = module._get_fsdp_state() + elif getattr(module, "_te_cached_parent_fsdp_state", None) is not None: + # See if we have cached the parent fsdp state of the module + fsdp_state = module._te_cached_parent_fsdp_state + else: + from torch.distributed._composable_state import _module_state_mapping + + # Otherwise get the fsdp state of lca of module in the module hierarchy + min_nodes_in_parent = float("inf") + closest_parent_fsdp_mod = None + for fsdp_mod in _module_state_mapping.keys(): + all_submodules = list(fsdp_mod.modules()) + for submodule in all_submodules: + if submodule is module: + if min_nodes_in_parent > len(all_submodules): + closest_parent_fsdp_mod = fsdp_mod + min_nodes_in_parent = len(all_submodules) + if closest_parent_fsdp_mod is None: + raise RuntimeError( + "Module is not FSDP-wrapped and does not have any FSDP-wrapped parent modules." + ) + fsdp_state = closest_parent_fsdp_mod._get_fsdp_state() + # Cache the parent fsdp state of the module to avoid recomputing + # the closest parent fsdp module. + module._te_cached_parent_fsdp_state = fsdp_state + return fsdp_state + + def _fsdp_scatter_tensors( fsdp_group: dist_group_type, *tensors: torch.Tensor, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 9b6ca9d9cd..d2abe3a2de 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -17,6 +17,7 @@ import torch import torch.nn.functional as F +from torch.distributed.tensor import DTensor import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe @@ -1244,7 +1245,12 @@ def register_parameter(self, name, param, **kwargs): metedata used in deferred initialization. """ super().register_parameter(name, param) - self.param_init_meta[name] = _ParameterInitMeta(**kwargs) + # Initialize param_init_meta exactly once during the init. FSDP2 can call + # register parameter again to change parameters to DTensors. And it calls + # it without custom fp8 specific kwargs that we need. And so we dont want + # to reset/loose our fp8 init attributes. + if hasattr(self, "param_init_meta") and name not in self.param_init_meta: + self.param_init_meta[name] = _ParameterInitMeta(**kwargs) def reset_parameters(self, defer_init: Optional[bool] = False) -> None: """ @@ -1256,10 +1262,14 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: return for name, param in self.named_parameters(recurse=False): + # Check if parameter is a DTensor (FSDP2) or regular tensor + is_dtensor = isinstance(param, DTensor) + dtensor_param = param if is_dtensor else None + # Need to update/quantize local tensor in case of DTensor + param = param._local_tensor if is_dtensor else param # Ensure parameter is on a real device if param.device == torch.device("meta"): param = torch.empty_like(param, device="cuda") - # Initialize the parameter values on device init_fn = self.param_init_meta[name].init_fn get_rng_state_tracker = self.param_init_meta[name].get_rng_state_tracker @@ -1288,7 +1298,15 @@ 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): + device_mesh = dtensor_param.device_mesh + amax_reduction_group = ( + device_mesh.get_group(mesh_dim="shard") + if device_mesh.ndim > 1 + else device_mesh.get_group() + ) + quantizer.amax_reduction_group = amax_reduction_group + quantizer.with_amax_reduction = True # Quantize parameter param = quantizer(param) @@ -1296,7 +1314,18 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: # NOTE: Currently this can only be broken when primary weights are in Fp8 but # re-applying the nn.Parameter() wrap is a no-op when the input is already # a parameter so we always re-apply it just for extra safety. - param = torch.nn.Parameter(param) + if is_dtensor: + # recreate the DTensor from the parameter. + dtensor_param = DTensor.from_local( + param, + device_mesh=dtensor_param.device_mesh, + placements=dtensor_param.placements, + shape=dtensor_param.size(), + stride=dtensor_param.stride(), + ) + dtensor_param = torch.nn.Parameter(dtensor_param) + else: + param = torch.nn.Parameter(param) # Keep high-precision values on CPU if needed if high_precision_init_val is not None: @@ -1324,8 +1353,12 @@ def clear(self): param._high_precision_init_val = high_precision_init_val param.get_high_precision_init_val = MethodType(get, param) param.clear_high_precision_init_val = MethodType(clear, param) + # Update the parameter based on its type - setattr(self, name, param) + if not is_dtensor: + setattr(self, name, param) + else: + setattr(self, name, dtensor_param) @abstractmethod def forward(self): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 4d6b2f23b9..59dc2b2997 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -108,9 +108,15 @@ def forward( is_fp8_activation_recompute_enabled() and not in_fp8_activation_recompute_phase() ) - if weight_quantizers[0] is not None: + # No need to set the quantizer states if weight is already quantized + if weight_quantizers[0] is not None and not isinstance( + weights[0], QuantizedTensorStorage + ): for weight_quantizer in weight_quantizers: weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif isinstance(weights[0], QuantizedTensorStorage): + # If weights are already quantized, no need to set quantizer states + weight_quantizers = [weight._quantizer for weight in weights] if output_quantizers[0] is not None: for output_quantizer in output_quantizers: output_quantizer.set_usage(rowwise=True, columnwise=False) @@ -205,10 +211,6 @@ def forward( inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) else: inputmats = [None] * num_gemms - if inp.requires_grad: - for weight in weights_fp8: - if isinstance(weight, QuantizedTensorStorage): - weight.update_usage(columnwise_usage=True) if cpu_offloading: ctx.grad_added_to_main_grad = hasattr(weights[0], "grad_added_to_main_grad") @@ -354,13 +356,11 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], dtype=ctx.activation_dtype, device=ctx.device, ) - - for weight, quantizer in zip(weights, ctx.weight_quantizers): - if quantizer is not None and isinstance(weight, QuantizedTensorStorage): - weight.update_usage( - rowwise_usage=quantizer.rowwise_usage, - columnwise_usage=quantizer.columnwise_usage, - ) + # Make sure weights are available in column-wise format + # for dgrad computation. + for weight in weights: + if isinstance(weight, QuantizedTensorStorage): + weight.update_usage(columnwise_usage=True) general_grouped_gemm( weights, grad_output, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 933c7cde53..abe8c58298 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -276,12 +276,15 @@ def forward( # Prepare weight tensor # ------------------------------------------------------ weightmat = weight - quantized_weight = False + is_weight_param_quantized = False if fp8 or debug: - quantized_weight = not isinstance(weight, QuantizedTensorStorage) + is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) # Configure quantizer - if weight_quantizer is not None: + # If weight is already quantized, no need to set quantizer states + if is_weight_param_quantized: + weight_quantizer = weight._quantizer + elif weight_quantizer is not None: weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) # Get quantized weight @@ -413,10 +416,6 @@ def forward( ): ln_out.update_usage(rowwise_usage=False) - # Weight with column-wise usage is needed for dgrad GEMM. - if isinstance(weightmat, QuantizedTensorStorage): - weightmat.update_usage(columnwise_usage=True) - if cpu_offloading: mark_activation_offload(inputmat, mu, rsigma, ln_out) @@ -429,7 +428,7 @@ def forward( fsdp_group, mu, rsigma, - weightmat if quantized_weight else None, + weightmat if fp8 and not is_weight_param_quantized else None, ln_out if weight.requires_grad else None, ) nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") @@ -459,7 +458,7 @@ def forward( ctx.tensor_objects = tensor_objects ctx.requires_dgrad = inp_requires_grad ctx.requires_wgrad = weight.requires_grad - ctx.quantized_weight = quantized_weight + ctx.is_weight_param_quantized = is_weight_param_quantized if fuse_wgrad_accumulation and weight.requires_grad: # This check is needed to ensure that main_grad is not created # during the forward pass when using MCore FSDP as it creates @@ -563,7 +562,7 @@ def backward( ctx.fsdp_shapes, mu, rsigma, - weight if ctx.fp8 and ctx.quantized_weight else None, + weight if ctx.fp8 and not ctx.is_weight_param_quantized else None, ln_out, ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 889f545c1e..a358ae7ddf 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -351,8 +351,17 @@ def forward( # which handles weight caching etc. # FP8 cast to workspace buffer update_workspace = is_first_microbatch is None or is_first_microbatch - fc1_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) + # No need to set the quantizer states if weights are already quantized + if isinstance(fc1_weight, QuantizedTensorStorage): + fc1_weight_quantizer = fc1_weight._quantizer + elif fc1_weight_quantizer is not None: + fc1_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) + + if isinstance(fc2_weight, QuantizedTensorStorage): + fc2_weight_quantizer = fc2_weight._quantizer + elif fc2_weight_quantizer is not None: + fc2_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) + fc1_weight_final = module.get_weight_workspace( tensor=fc1_weight, quantizer=fc1_weight_quantizer, @@ -538,13 +547,6 @@ def forward( # Cache state for backward pass if is_grad_enabled: - - # Weight with column-wise usage is needed for dgrad GEMM. - if isinstance(fc1_weight_final, QuantizedTensorStorage): - fc1_weight_final.update_usage(columnwise_usage=True) - if isinstance(fc2_weight_final, QuantizedTensorStorage): - fc2_weight_final.update_usage(columnwise_usage=True) - if cpu_offloading: mark_activation_offload( inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ccb84e6642..0e2310a5a2 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -240,7 +240,8 @@ def forward( weightmat = weight if fp8 or debug: # Configure quantizer - if weight_quantizer is not None: + # No need to set the quantizer states if weight is already quantized + if weight_quantizer is not None and not isinstance(weight, QuantizedTensor): columnwise_usage = is_grad_enabled and inp.requires_grad if not columnwise_usage: columnwise_usage = ( @@ -248,7 +249,9 @@ def forward( and not in_fp8_activation_recompute_phase() ) weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - + elif isinstance(weight, QuantizedTensor): + # If weight is already quantized, no need to set quantizer states + weight_quantizer = weight._quantizer # Get quantized weight update_workspace = is_first_microbatch is None or is_first_microbatch weightmat = module.get_weight_workspace( @@ -389,11 +392,6 @@ def forward( if backward_needs_input: saved_inputmat = inputmat - # Weight with column-wise usage is needed for dgrad GEMM. - if inp.requires_grad: - if isinstance(weightmat, QuantizedTensorStorage): - weightmat.update_usage(columnwise_usage=True) - if cpu_offloading and saved_inputmat is not None: mark_activation_offload(saved_inputmat) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 15f5b6bd5e..7d49e3964f 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -433,6 +433,10 @@ def maybe_update_inplace(arg, new_arg, schema_arg): and schema_arg.alias_info.is_write ): arg.quantize_(new_arg) + elif isinstance(arg, list) and isinstance(new_arg, list): + # Recursively handle update for lists of tensors + for a, na in zip(arg, new_arg): + maybe_update_inplace(a, na, schema_arg) # In-place op: dequantize, perform op, and quantize if func._schema.is_mutable: @@ -489,20 +493,16 @@ def make_like( shape: Optional[Iterable[int]] = None, dtype: Optional[torch.dtype] = None, requires_grad: bool = False, - data: Optional[torch.Tensor] = None, ) -> QuantizedTensor: """Create new quantized tensor By default, new tensor has the same attributes and underlying - data. + data. This function is intended to create view of tensors. """ - if shape is None: - shape = data.shape if data is not None else tensor.shape + shape = shape if shape is not None else tensor.shape dtype = dtype if dtype is not None else tensor.dtype kwargs = tensor.get_metadata() - if data is not None: - kwargs["data"] = data return cls(shape=shape, dtype=dtype, requires_grad=requires_grad, **kwargs) def to_dtype(self, dtype: torch.dtype) -> QuantizedTensor: diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index de112bb3fd..eb2ac9a581 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -4,10 +4,10 @@ """Tensor class with FP8 data""" from __future__ import annotations -from typing import Optional, Tuple, Iterable, Union +from typing import Any, Optional, Tuple, Iterable, Union import warnings - import torch +from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType @@ -299,14 +299,12 @@ def make_empty( # Allocate FP8 data transpose if needed data_transpose = None if self.columnwise_usage: - inner_dim = data.size(-1) + transpose_shape = [data.size(-1)] + list(data.shape[:-1]) data_transpose = torch.empty( - inner_dim, - data.numel() // inner_dim, + transpose_shape, dtype=torch.uint8, device=device, ) - # Construct FP8 tensor return Float8Tensor( shape=shape, @@ -534,9 +532,36 @@ def remove_caches(self) -> None: self._transpose = None @classmethod - def __torch_dispatch__(cls, func, types, args, kwargs=None): + def make_like( + cls, + tensor: QuantizedTensor, + *, + shape: Optional[Iterable[int]] = None, + dtype: Optional[torch.dtype] = None, + requires_grad: bool = False, + data: Optional[torch.Tensor] = None, + data_transpose: Optional[torch.Tensor] = None, + ) -> QuantizedTensor: + """Create new quantized tensor + + By default, new tensor has the same attributes and underlying + data. - # View op + """ + if shape is None and data is not None: + shape = data.shape + new_tensor = super().make_like( + tensor, shape=shape, dtype=dtype, requires_grad=requires_grad + ) + if data is not None: + new_tensor._data = data + if data_transpose is not None: + new_tensor._transpose = data_transpose + new_tensor._transpose_invalid = False + return new_tensor + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == aten.view.default: tensor = args[0] data = tensor._data @@ -555,6 +580,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): or out_transpose_shape[1:] != out_shape[:-1] ): out_transpose = None + else: + view_shape_for_transpose = [out_shape[-1]] + list(out_shape[:-1]) + out_transpose = out_transpose.view(*view_shape_for_transpose) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, @@ -587,11 +615,37 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [data] + list(args[1:]), kwargs, ) - return [ - Float8Tensor.make_like(tensor, data=split_tensor, shape=split_tensor.shape) - for split_tensor in func_out + t_func_out = [None] * len(func_out) + # Compute corresponding split of the transpose cache if available + if tensor._transpose is not None and not tensor._transpose_invalid: + transpose = tensor._transpose + ndim = data.dim() + # Figure out the original split dim + if "dim" in kwargs: + dim_to_split = kwargs["dim"] + else: + dim_to_split = args[2] if len(args) > 2 else 0 + # Dimension along which transpose needs to be split + t_dim = 0 if dim_to_split == ndim - 1 else dim_to_split + 1 + t_func_out = transpose.__torch_dispatch__( + func, + types, + [transpose, args[1], t_dim], + kwargs, + ) + outs = [ + Float8Tensor.make_like( + tensor, + data=split_tensor, + data_transpose=split_transpose_tensor, + shape=split_tensor.shape, + ) + for split_tensor, split_transpose_tensor in zip(func_out, t_func_out) ] + return outs + if func == aten.new_zeros.default: + # create fresh new tensor with zeros. tensor = args[0] data = tensor._data func_out = data.__torch_dispatch__( @@ -600,17 +654,63 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [data] + list(args[1:]), kwargs, ) - return Float8Tensor.make_like(tensor, data=func_out, shape=func_out.shape) + func_transposed_out = None + if tensor._transpose is not None and not tensor._transpose_invalid: + transpose = tensor._transpose + size = args[1] + t_shape = [size[-1]] + list(size[:-1]) + func_transposed_out = transpose.__torch_dispatch__( + func, + types, + [transpose, t_shape] + list(args[2:]), + kwargs, + ) + # deep copy the scale inverse tensor and quantizer as well. + scale_inv = tensor._scale_inv.detach().clone() + quantizer = tensor._quantizer.copy() + out_tensor = Float8Tensor( + data=func_out, + shape=func_out.shape, + dtype=tensor.dtype, + fp8_dtype=tensor._fp8_dtype, + fp8_scale_inv=scale_inv, + data_transpose=func_transposed_out, + quantizer=quantizer, + ) + return out_tensor + if func == torch.ops.aten.as_strided.default: tensor = args[0] data = tensor._data + # Apply as_strided to the primary uint8 data func_out = data.__torch_dispatch__( func, types, [data] + list(args[1:]), kwargs, ) - return Float8Tensor.make_like(tensor, data=func_out, shape=func_out.shape) + func_transposed_out = None + if tensor._transpose is not None and not tensor._transpose_invalid: + transpose = tensor._transpose + size = args[1] + stride = args[2] + if "storage_offset" in kwargs: + storage_offset = kwargs["storage_offset"] + else: + storage_offset = args[3] if len(args) > 3 else 0 + # Shape and strided needed for transpose matrix + t_size = [size[-1]] + list(size[:-1]) + t_stride = [stride[-1]] + list(stride[:-1]) + func_transposed_out = transpose.__torch_dispatch__( + func, + types, + [transpose, t_size, t_stride, storage_offset] + list(args[4:]), + kwargs, + ) + return Float8Tensor.make_like( + tensor, data=func_out, data_transpose=func_transposed_out, shape=func_out.shape + ) + if func == torch.ops.aten.detach.default: return cls.detach(args[0]) if func == torch.ops.aten.clone.default: @@ -632,9 +732,105 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) else: pass - return super().__torch_dispatch__(func, types, args, kwargs) + def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + """Functions FSDP2 calls before all-gather of the + weights for both forward and backward passes. + Args: + mesh (torch.distributed.DeviceMesh): DeviceMesh used by FSDP2 + to shard the weights. + orig_size (torch.Size): Original size of the weight tensor.(For us same as self.shape) + contiguous_orig_stride (Tuple[int]): Original stride of the weight tensor + (For us same as self.stride()) + module (FSDPModule): FSDP module. FSDP wrapped module wrapped using fully_shard + that contains this FP8 tensor. + mp_policy (MixedPrecisionPolicy): Mixed precision policy used by FSDP2. + + Returns: + shareded_tensors: Tuple[torch.Tensor, ...]: Tuple of tensors + that need to be all-gathered.(In this case uint8 data tensor) + metadata: Tuple[Any]: Metadata needed for reconstructing the + Float8Tensor after all-gather. + """ + # pylint: disable=unused-argument + # Importing here to avoid circular imports + from transformer_engine.pytorch.distributed import _get_module_fsdp_state + + if isinstance(self._quantizer, Float8CurrentScalingQuantizer) and mesh is not None: + # When sharded weight is updated after reduce scattering the gradients in FSDP2, + # we need to do amax reduction across the mesh to make sure all weight shards are + # updated with same scale inverse. Setting the state below in the quantizer will make + # sure that updated Quantized weight tensor have same scale inverse across all shards. + self._quantizer.amax_reduction_group = mesh.get_group() + self._quantizer.with_amax_reduction = True + quantizer = self._quantizer.copy() # quantizer to be used for allgathered weights + fsdp_state = _get_module_fsdp_state(module) + reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward + # If weights are resharded after forward pass, then its enough to set the quantizer usages + # based on whether its forward or backward pass for the allgathered weights. + # If not resharded after forward pass, the same weights allgathered in forward + # are used again in backward and so we dont change the quantizer usages which might need + # both rowwise and columnwise usages. + if reshard_after_forward: + training_state = fsdp_state._fsdp_param_group._training_state + is_backward_pass = training_state == TrainingState.PRE_BACKWARD + # In case of hopper/L40, only one of data/transpose is needed + # based on forward or backward pass. So setting the quantizer usages appropriately. + quantizer.set_usage(rowwise=not is_backward_pass, columnwise=is_backward_pass) + sharded_tensors = (self._data,) + metadata = (self._scale_inv, self._fp8_dtype, quantizer) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[Float8Tensor] = None, + ): + """Functions FSDP2 calls after all-gather of the + weights for both forward and backward passes. + Args: + all_gather_outputs (Tuple[torch.Tensor, ...]): sharded_tensors sent out in fsdp_pre_all_gather from each rank + are all-gathered and received here as a tuple. + metadata (Any): metadata sent out in fsdp_pre_all_gather used for reconstructing the Float8Tensor. + param_dtype (torch.dtype): high precision dtype of the Float8Tensor. + out (Optional[torch.Tensor], optional): _description_. Defaults to None. + + Returns: + Tuple[Float8Tensor, Tuple[torch.Tensor, ...]]: Allgathered Float8Tensor and tuple of internal tensors + used by the Float8Tensor that was being computed after allgather. + """ + + (data,) = all_gather_outputs + (fp8_scale_inv, fp8_dtype, quantizer) = metadata + orig_shape = data.size() + # Quantizer has only columnwise usage set for backward pass + # In Blackwell+ architectures, transpose is not needed at all, + # even if columnwise usage is set. and is going to be handled + # internally in the update_usage method. + if out is not None: + out._data = data + else: + fp8_args = { + "shape": orig_shape, + "dtype": param_dtype, + "fp8_scale_inv": fp8_scale_inv, + "fp8_dtype": fp8_dtype, + "quantizer": quantizer, + "requires_grad": False, + "data": data, + } + out = Float8Tensor(**fp8_args) + + out.update_usage( + rowwise_usage=quantizer.rowwise_usage, + columnwise_usage=quantizer.columnwise_usage, + ) + return out, all_gather_outputs + @classmethod def _make_in_reduce_ex( cls, @@ -752,6 +948,9 @@ def forward( out_transpose_shape = out_transpose.size() if out_transpose_shape[0] != out_shape[-1] or out_transpose_shape[1:] != out_shape[:-1]: out_transpose = None + else: + view_shape_for_transpose = [shape[-1]] + list(shape[:-1]) + out_transpose = out_transpose.view(*view_shape_for_transpose) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, @@ -796,6 +995,9 @@ def forward( out_transpose_shape = out_transpose.size() if out_transpose_shape[0] != out_shape[-1] or out_transpose_shape[1:] != out_shape[:-1]: out_transpose = None + else: + reshape_shape_for_transpose = [shape[-1]] + list(shape[:-1]) + out_transpose = out_transpose.reshape(*reshape_shape_for_transpose) return Float8Tensor( shape=out_shape, dtype=tensor.dtype, diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 5ef5708fdb..d981f71579 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -6,16 +6,17 @@ from __future__ import annotations from collections.abc import Iterable import math -from typing import Optional, Tuple, Union +from typing import Optional, Tuple, Union, Any +import warnings import torch +from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe from ..constants import MXFP8_BLOCK_SCALING_SIZE from ..utils import devices_match, round_up_to_nearest_multiple - from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func from ..quantized_tensor import QuantizedTensor, Quantizer from ._quantization_helpers import _IdentityFunc @@ -298,7 +299,6 @@ def contiguous( memory_format: torch.memory_format = torch.contiguous_format, ) -> MXFP8Tensor: """Returns tensor with data in provided memory format - Returns `self` if data is already in correct memory format. """ @@ -314,7 +314,6 @@ def contiguous( @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): - # View op if func == aten.view.default: tensor = args[0] @@ -338,9 +337,335 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): fp8_dtype=tensor._fp8_dtype, ) + if func == torch.ops.aten.copy_.default: + dst, src = args[0], args[1] + # Booleans to check if src has all the usages that dst needs to respect dst quantizer usages. + # If not, default to base class behavior. + rowwise_matches = src._rowwise_data is not None or dst._rowwise_data is None + columnwise_matches = src._columnwise_data is not None or dst._columnwise_data is None + if ( + isinstance(src, MXFP8Tensor) + and isinstance(dst, MXFP8Tensor) + and rowwise_matches + and columnwise_matches + ): + if dst._rowwise_data is not None: + dst._rowwise_data.copy_(src._rowwise_data.detach()) + dst._rowwise_scale_inv.copy_(src._rowwise_scale_inv.detach()) + if dst._columnwise_data is not None: + dst._columnwise_data.copy_(src._columnwise_data.detach()) + dst._columnwise_scale_inv.copy_(src._columnwise_scale_inv.detach()) + return dst + + # FSDP2 related functions. + if func == aten.split.Tensor: + # This is called if entire model is initialized on CUDA device and + # then splitted. Finally the shard needed by the process is used + # and other splitted shards are discarded. + if "dim" in kwargs: + dim_to_split = kwargs["dim"] + else: + dim_to_split = args[2] if len(args) > 2 else 0 + tensor = args[0] + split_size = args[1] + dim0_size = tensor.size(0) + dimlast_size = math.prod(tensor.shape[1:]) + if ( + dim0_size % split_size != 0 + or dim_to_split != 0 + or split_size % MXFP8_BLOCK_SCALING_SIZE != 0 + or dimlast_size % MXFP8_BLOCK_SCALING_SIZE != 0 + ): + # Handle splitting by dequantizing and splitting the hp tensor + return super().__torch_dispatch__(func, types, args, kwargs) + + out_data = [] + for data in [tensor._rowwise_data, tensor._columnwise_data]: + func_out = ( + data.__torch_dispatch__( + func, + types, + [data] + list(args[1:]), + kwargs, + ) + if data is not None + else None + ) + out_data.append(func_out) + + scale_invs = [tensor._rowwise_scale_inv, tensor._columnwise_scale_inv] + split_sizes_for_scale = [split_size, split_size // MXFP8_BLOCK_SCALING_SIZE] + # Padding requirements: rowwise dim0 should be divisble by 128, columnwise dim0 should be divisble by 4 + padding_multiples = [128, 4] + for scale_inv, scale_split_size, pad_multiple in zip( + scale_invs, split_sizes_for_scale, padding_multiples + ): + scale_inv_out = ( + scale_inv.__torch_dispatch__( + func, + types, + [scale_inv, scale_split_size] + list(args[2:]), + kwargs, + ) + if scale_inv is not None + else None + ) + # Pad scale_inv_out to be a multiple of pad_multiple + if scale_inv_out is not None: + current_shape = scale_inv_out.shape + pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple + if pad_dim0 > 0: + scale_inv_out = torch.nn.functional.pad(scale_inv_out, (0, 0, 0, pad_dim0)) + + out_data.append(scale_inv_out) + return [ + MXFP8Tensor( + shape=( + splitted_tensor_data[0].size() + if splitted_tensor_data[0] is not None + else splitted_tensor_data[1].size() + ), + dtype=tensor.dtype, + rowwise_data=splitted_tensor_data[0], + rowwise_scale_inv=splitted_tensor_data[2], + columnwise_data=splitted_tensor_data[1], + columnwise_scale_inv=splitted_tensor_data[3], + quantizer=tensor._quantizer, + requires_grad=False, + fp8_dtype=tensor._fp8_dtype, + ) + for splitted_tensor_data in zip(*out_data) + ] + if func == torch.ops.aten.as_strided.default: + # Applied on unsharded param in FSDP2. In our case, this should be a no-op + # This is needed for the case where some MXFP8 shards need padding i.e dimension 0 + # of the unsharded param is not a multiple of the world size. If that is the case, + # we down the dequantization route and weights are allgathered in high precision. + # If weight doesnt need padding, this is just a no-op. + shape = args[1] + strides = args[2] + tensor = args[0] + if ( + len(shape) != 2 + or len(strides) != 2 + or strides[1] != 1 + or shape[0] != tensor.shape[0] + or shape[1] != tensor.shape[1] + ): + return super().__torch_dispatch__(func, types, args, kwargs) + + return MXFP8Tensor.make_like(tensor) + + if func == aten.slice.Tensor: + # FSDP2 needed function. + # We need slicing for the case where some MXFP8 weight shards need padding i.e dimension 0 + # of the unsharded param is not a multiple of the world size. If that is the case, + # we down the dequantization route and weights are allgathered in high precision instead. + # If sharded weight doesnt have padding, this is just a no-op. + dim = args[1] + start = args[2] + length = args[3] + tensor = args[0] + if ( + dim != 0 + or length != tensor.shape[0] + or start != 0 + or length % MXFP8_BLOCK_SCALING_SIZE != 0 + or start % MXFP8_BLOCK_SCALING_SIZE != 0 + ): + return super().__torch_dispatch__(func, types, args, kwargs) + return MXFP8Tensor.make_like(tensor) + + if func == aten.new_zeros.default: + rowwise_data = None + columnwise_data = None + rowwise_scale_inv = None + columnwise_scale_inv = None + tensor = args[0] + shape = args[1] + first_dim = math.prod(shape[:-1]) + last_dim = shape[-1] + if ( + first_dim % MXFP8_BLOCK_SCALING_SIZE != 0 + or last_dim % MXFP8_BLOCK_SCALING_SIZE != 0 + ): + return super().__torch_dispatch__(func, types, args, kwargs) + rowwise_scale_inv_shape = [first_dim, last_dim // MXFP8_BLOCK_SCALING_SIZE] + columnwise_scale_inv_shape = [ + first_dim // MXFP8_BLOCK_SCALING_SIZE, + last_dim, + ] + if tensor._rowwise_data is not None: + rowwise_data = tensor._rowwise_data.__torch_dispatch__( + func, + types, + [tensor._rowwise_data] + list(args[1:]), + kwargs, + ) + rowwise_scale_inv = tensor._rowwise_scale_inv.__torch_dispatch__( + func, + types, + [tensor._rowwise_scale_inv, rowwise_scale_inv_shape] + list(args[2:]), + kwargs, + ) + if tensor._columnwise_data is not None: + columnwise_data = tensor._columnwise_data.__torch_dispatch__( + func, + types, + [tensor._columnwise_data] + list(args[1:]), + kwargs, + ) + columnwise_scale_inv = tensor._columnwise_scale_inv.__torch_dispatch__( + func, + types, + [tensor._columnwise_scale_inv, columnwise_scale_inv_shape] + list(args[2:]), + kwargs, + ) + return MXFP8Tensor( + shape=args[1], + dtype=tensor.dtype, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + quantizer=tensor._quantizer.copy(), + requires_grad=False, + fp8_dtype=tensor._fp8_dtype, + ) # Default case return super().__torch_dispatch__(func, types, args, kwargs) + def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + """Functions FSDP2 calls before all-gather of the + weights for both forward and backward passes. + Args: + mesh (torch.distributed.DeviceMesh): DeviceMesh used by FSDP2 + to shard the weights. + orig_size (torch.Size): Original size of the weight tensor.(For us same as self.shape) + contiguous_orig_stride (Tuple[int]): Original stride of the weight tensor + (For us same as self.stride()). + module (FSDPModule): FSDP module. FSDP wrapped module wrapped using fully_shard + that contains this MXFP8 tensor. + mp_policy (MixedPrecisionPolicy): Mixed precision policy used by FSDP2. + + Returns: + sharded_tensors: Tuple[torch.Tensor, ...]: Tuple of tensors + that need to be all-gathered. + metadata: Tuple[Any]: Metadata needed for reconstructing the + MXFP8Tensor after all-gather. + """ + # pylint: disable=unused-argument + from transformer_engine.pytorch.distributed import _get_module_fsdp_state + + fsdp_state = _get_module_fsdp_state(module) + reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward + quantizer = self._quantizer.copy() + # Remove padding from scale inverses before allgather + # Rowwise scale_inv should be divisible by [128,4], columnwise by [4, 128] + rowwise_scale_inv = self._rowwise_scale_inv + columnwise_scale_inv = self._columnwise_scale_inv + shape = self.shape + if rowwise_scale_inv is not None: + # Remove padding from rowwise scale_inv + flattened_in_shape0 = math.prod(shape[:-1]) + if rowwise_scale_inv.size(0) != flattened_in_shape0: + rowwise_scale_inv = rowwise_scale_inv[:flattened_in_shape0] + + if columnwise_scale_inv is not None: + # Remove padding from columnwise scale_inv + flattened_in_shape0 = math.prod(shape[:-1]) // MXFP8_BLOCK_SCALING_SIZE + if columnwise_scale_inv.size(0) != flattened_in_shape0: + columnwise_scale_inv = columnwise_scale_inv[:flattened_in_shape0] + + sharded_tensors = (self._rowwise_data, rowwise_scale_inv) + # If weights are resharded after forward pass, then its enough to set the quantizer usages + # based on whether its forward or backward pass for the allgathered weights. + # If not resharded after forward pass, the same weights allgathered in forward + # are used again in backward. And hence if we need the columnwise data/scale_inv, + # we need to send them as well for allgather in forward pass itself. + if reshard_after_forward: + training_state = fsdp_state._fsdp_param_group._training_state + is_backward_pass = training_state == TrainingState.PRE_BACKWARD + # Allgather only the necessary tensors based on forward/backward pass + quantizer.set_usage(rowwise=not is_backward_pass, columnwise=is_backward_pass) + sharded_tensors = ( + (self._columnwise_data, columnwise_scale_inv) + if is_backward_pass + else sharded_tensors + ) + else: + if quantizer.columnwise_usage: + # If weights are not resharded after forward, then both + # rowwise and columnwise data/scale_inv need to be allgathered. + sharded_tensors += (self._columnwise_data, columnwise_scale_inv) + metadata = (self._fp8_dtype, quantizer) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[MXFP8Tensor] = None, + ): + """Functions FSDP2 calls after all-gather of the + weights for both forward and backward passes. + Args: + all_gather_outputs (Tuple[torch.Tensor, ...]): sharded_tensors sent out in fsdp_pre_all_gather from each rank + are all-gathered and received here as a tuple. + metadata (Any): metadata sent out in fsdp_pre_all_gather used for reconstructing the MXFP8Tensor. + param_dtype (torch.dtype): high precision dtype of the MXFP8Tensor. + out (Optional[torch.Tensor], optional): _description_. Defaults to None. + Returns: + Tuple[MXFP8Tensor, Tuple[torch.Tensor, ...]]: Allgathered MXFP8Tensor and tuple of internal tensors + used by the MXFP8Tensor that was being computed after allgather. + """ + fp8_dtype, quantizer = metadata + rowwise_data, rowwise_scale_inv = ( + all_gather_outputs[:2] if quantizer.rowwise_usage else (None, None) + ) + columnwise_data, columnwise_scale_inv = ( + all_gather_outputs[-2:] if quantizer.columnwise_usage else (None, None) + ) + + # Add padding to scale_inv tensors to be multiples of [128, 4]for rowwise and [4, 128] for columnwise + if rowwise_scale_inv is not None: + # Pad rowwise_scale_inv to be a multiple of [128, 4] + current_shape = rowwise_scale_inv.shape + pad_dim0 = (128 - current_shape[0] % 128) % 128 + if pad_dim0 > 0: + rowwise_scale_inv = torch.nn.functional.pad(rowwise_scale_inv, (0, 0, 0, pad_dim0)) + + if columnwise_scale_inv is not None: + # Pad columnwise_scale_inv to be a multiple of [4, 128] + current_shape = columnwise_scale_inv.shape + pad_dim0 = (4 - current_shape[0] % 4) % 4 + if pad_dim0 > 0: + columnwise_scale_inv = torch.nn.functional.pad( + columnwise_scale_inv, (0, 0, 0, pad_dim0) + ) + + if out is not None: + out._rowwise_data = rowwise_data + out._rowwise_scale_inv = rowwise_scale_inv + out._columnwise_data = columnwise_data + out._columnwise_scale_inv = columnwise_scale_inv + out._quantizer = quantizer + else: + out = MXFP8Tensor( + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=fp8_dtype, + dtype=param_dtype, + shape=rowwise_data.shape if rowwise_data is not None else columnwise_data.shape, + quantizer=quantizer, + ) + + return out, all_gather_outputs + @classmethod def _make_in_reduce_ex( cls, @@ -478,10 +803,14 @@ def forward( shape[i] = d_inferred break if shape[-1] != ctx.shape[-1]: - raise RuntimeError( - "MXFP8Tensor does not support reshaping inner dimension " + warnings.warn( + "MXFP8Tensor does not support reshaping inner dimension. " f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)})" + "If you are using this for FSDP2 without compiled_autograd_enabled," + "then ignore this warning. Since this view is not going to be used anywhere. ", + stacklevel=2, ) + return tensor.dequantize().view(*shape) # Construct new tensor if shape is provided new_rowwise_data = None From f8693d2b044ab83624e020ffa69f3412a916993d Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Wed, 12 Nov 2025 07:38:41 -0800 Subject: [PATCH 055/521] Fix CI failure related to bug in MXFP8 copy implementation (#2369) * fix ci issue Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert back testing changes 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: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../pytorch/tensor/mxfp8_tensor.py | 32 +++++++++---------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index d981f71579..15e0b86c90 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -339,23 +339,21 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == torch.ops.aten.copy_.default: dst, src = args[0], args[1] - # Booleans to check if src has all the usages that dst needs to respect dst quantizer usages. - # If not, default to base class behavior. - rowwise_matches = src._rowwise_data is not None or dst._rowwise_data is None - columnwise_matches = src._columnwise_data is not None or dst._columnwise_data is None - if ( - isinstance(src, MXFP8Tensor) - and isinstance(dst, MXFP8Tensor) - and rowwise_matches - and columnwise_matches - ): - if dst._rowwise_data is not None: - dst._rowwise_data.copy_(src._rowwise_data.detach()) - dst._rowwise_scale_inv.copy_(src._rowwise_scale_inv.detach()) - if dst._columnwise_data is not None: - dst._columnwise_data.copy_(src._columnwise_data.detach()) - dst._columnwise_scale_inv.copy_(src._columnwise_scale_inv.detach()) - return dst + if isinstance(src, MXFP8Tensor) and isinstance(dst, MXFP8Tensor): + # Booleans to check if src has all the usages that dst needs to respect dst quantizer usages. + # If not, default to base class behavior. + rowwise_matches = src._rowwise_data is not None or dst._rowwise_data is None + columnwise_matches = ( + src._columnwise_data is not None or dst._columnwise_data is None + ) + if rowwise_matches and columnwise_matches: + if dst._rowwise_data is not None: + dst._rowwise_data.copy_(src._rowwise_data.detach()) + dst._rowwise_scale_inv.copy_(src._rowwise_scale_inv.detach()) + if dst._columnwise_data is not None: + dst._columnwise_data.copy_(src._columnwise_data.detach()) + dst._columnwise_scale_inv.copy_(src._columnwise_scale_inv.detach()) + return dst # FSDP2 related functions. if func == aten.split.Tensor: From e4bfa628632e15ef8bc1fae9b2e89686f6a097ea Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Wed, 12 Nov 2025 23:47:41 +0530 Subject: [PATCH 056/521] [Feature] Enable rope application with offsets for training (#2188) * enable applying rope offsets in backwared Signed-off-by: Sudhakar Singh * add tests for rope offsets for thd/bshd/sbhd formats Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fixes 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> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fused_rope.py | 170 ++++++++++++++---- .../common/fused_rope/fused_rope.cu | 85 ++++----- .../include/transformer_engine/fused_rope.h | 13 +- transformer_engine/pytorch/attention/rope.py | 85 ++++----- transformer_engine/pytorch/csrc/extensions.h | 1 + .../pytorch/csrc/extensions/apply_rope.cpp | 17 +- 6 files changed, 236 insertions(+), 135 deletions(-) diff --git a/tests/pytorch/test_fused_rope.py b/tests/pytorch/test_fused_rope.py index aaf2eca2d3..9e4ddbdad1 100644 --- a/tests/pytorch/test_fused_rope.py +++ b/tests/pytorch/test_fused_rope.py @@ -58,10 +58,6 @@ def test_fused_rope( # are with the maximum length of the rope embeddings. pytest.skip("Skipping test with margin=0 and start_positions=True") - if start_positions == True and cp_size > 1: - # `start_positions` is only supported for `cp_size=1` and inference. - pytest.skip("Skipping test with cp_size>1 and start_positions=True") - device = torch.device("cuda:0") batch_size, head_num = 2, 64 t = torch.rand( @@ -102,11 +98,8 @@ def test_fused_rope( cp_rank=cp_rank, ).to(dtype) loss_unfused = loss_func(output_unfused) - - if not isinstance(start_positions, torch.Tensor): - loss_unfused.backward() - grad_unfused = t.grad.detach().clone() - + loss_unfused.backward() + grad_unfused = t.grad.detach().clone() t.grad = None # fused @@ -121,17 +114,12 @@ def test_fused_rope( cp_rank=cp_rank, ) loss_fused = loss_func(output_fused) - - if not isinstance(start_positions, torch.Tensor): - loss_fused.backward() - grad_fused = t.grad.detach().clone() + loss_fused.backward() + grad_fused = t.grad.detach().clone() t.grad = None torch.testing.assert_close(output_fused, output_unfused) - - if not isinstance(start_positions, torch.Tensor): - torch.testing.assert_close(grad_fused, grad_unfused) - + torch.testing.assert_close(grad_fused, grad_unfused) assert output_fused.is_contiguous() @@ -156,10 +144,6 @@ def test_fused_rope_thd( margin: int, ) -> None: - if start_positions == True and cp_size > 1: - # `start_positions` is only supported for `cp_size=1` and inference. - pytest.skip("Skipping test with cp_size>1 and start_positions=True") - device = torch.device("cuda:0") batch_size, head_num = 2, 64 cu_seqlens = [0, 400, 542, 711, 727, 752, 1270, 1426, 1450, 1954, 2044, 2048] @@ -214,10 +198,8 @@ def test_fused_rope_thd( cp_rank=cp_rank, ).to(dtype) loss_unfused = loss_func(output_unfused) - - if not isinstance(start_positions, torch.Tensor): - loss_unfused.backward() - grad_unfused = t.grad.detach().clone() + loss_unfused.backward() + grad_unfused = t.grad.detach().clone() t.grad = None # fused @@ -233,18 +215,142 @@ def test_fused_rope_thd( cp_rank=cp_rank, ) loss_fused = loss_func(output_fused) - - if not isinstance(start_positions, torch.Tensor): - loss_fused.backward() - grad_fused = t.grad.detach().clone() + loss_fused.backward() + grad_fused = t.grad.detach().clone() t.grad = None torch.testing.assert_close(output_fused, output_unfused) + torch.testing.assert_close(grad_fused, grad_unfused) + assert output_fused.is_contiguous() - if not isinstance(start_positions, torch.Tensor): - torch.testing.assert_close(grad_fused, grad_unfused) - assert output_fused.is_contiguous() +@pytest.mark.parametrize("start_positions", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("hidden_size", [128, 256]) +@pytest.mark.parametrize("rotary_percent", [1.0]) +@pytest.mark.parametrize("loss_func", [_overlapping_grad]) +@pytest.mark.parametrize("cp_size", [2]) +@pytest.mark.parametrize("interleaved", [False, True]) +def test_unfused_rope_thd_vs_bshd( + dtype: torch.dtype, + hidden_size: int, + rotary_percent: float, + loss_func: Callable, + cp_size: int, + interleaved: bool, + start_positions: bool, +) -> None: + """ + This is just a sanity check to ensure that the unfused RoPE in THD/SBHD/BSHD + formats are the same. + """ + device = torch.device("cuda:0") + seqlen, max_seqlen = 16, 2048 + batch_size, head_num = 4, 256 + + # NOTE: dtype=torch.int32 is important, otherwise the cumsum will be in int64 and + # that causes unexpected issues. + seq_lens = torch.tensor([seqlen for _ in range(batch_size)], dtype=torch.int32) + + cu_seqlens = torch.cumsum(torch.cat([torch.zeros(1, dtype=torch.int32), seq_lens]), dim=0).to( + device=device, dtype=torch.int32 + ) + + # Create a tensor in THD format + thd = torch.rand( + (cu_seqlens[-1] // cp_size, head_num, hidden_size), + dtype=dtype, + device=device, + ) + thd.requires_grad = True + + # Clone the tensor to create a tensor in BSHD format + bshd = thd.view(batch_size, -1, head_num, hidden_size).clone().detach() + bshd = bshd.to(dtype=dtype, device=device) + bshd.requires_grad = True + + # Clone the tensor to create a tensor in SBHD format + sbhd = bshd.transpose(1, 0).clone().detach() + sbhd = sbhd.to(dtype=dtype, device=device) + sbhd.requires_grad = True + + rotary_pos_emb = RotaryPositionEmbedding(hidden_size, rotary_percent, interleaved=interleaved) + emb = rotary_pos_emb(max_seqlen) + assert emb.is_contiguous() + + start_positions = cu_seqlens[:-1] if start_positions else None + + for cp_rank in range(cp_size): + # unfused bshd + output_unfused_bshd = apply_rotary_pos_emb( + bshd.float(), + emb, + start_positions=start_positions, + interleaved=interleaved, + fused=False, + tensor_format="bshd", + cu_seqlens=cu_seqlens, + cp_size=cp_size, + cp_rank=cp_rank, + ).to(dtype) + loss_unfused_bshd = loss_func(output_unfused_bshd) + loss_unfused_bshd.backward() + grad_unfused_bshd = bshd.grad.detach().clone() + bshd.grad = None + + # unfused sbhd + output_unfused_sbhd = apply_rotary_pos_emb( + sbhd.float(), + emb, + start_positions=start_positions, + interleaved=interleaved, + fused=False, + tensor_format="sbhd", + cu_seqlens=cu_seqlens, + cp_size=cp_size, + cp_rank=cp_rank, + ).to(dtype) + + loss_unfused_sbhd = loss_func(output_unfused_sbhd) + loss_unfused_sbhd.backward() + grad_unfused_sbhd = sbhd.grad.detach().clone() + sbhd.grad = None + + # unfused thd + output_unfused_thd = apply_rotary_pos_emb( + thd.float(), + emb, + start_positions=start_positions, + tensor_format="thd", + interleaved=interleaved, + fused=False, + cu_seqlens=cu_seqlens, + cp_size=cp_size, + cp_rank=cp_rank, + ).to(dtype) + + loss_unfused_thd = loss_func(output_unfused_thd) + loss_unfused_thd.backward() + grad_unfused_thd = thd.grad.detach().clone() + thd.grad = None + + torch.testing.assert_close( + output_unfused_bshd.reshape(*output_unfused_thd.shape), output_unfused_thd + ) + torch.testing.assert_close( + output_unfused_sbhd.transpose(1, 0).reshape(*output_unfused_thd.shape), + output_unfused_thd, + ) + torch.testing.assert_close( + grad_unfused_bshd.reshape(*grad_unfused_thd.shape), grad_unfused_thd + ) + torch.testing.assert_close( + grad_unfused_sbhd.transpose(1, 0).reshape(*grad_unfused_thd.shape), grad_unfused_thd + ) + + assert output_unfused_thd.is_contiguous() + assert output_unfused_bshd.is_contiguous() + assert output_unfused_sbhd.is_contiguous() @pytest.mark.parametrize("start_positions", [True, False]) diff --git a/transformer_engine/common/fused_rope/fused_rope.cu b/transformer_engine/common/fused_rope/fused_rope.cu index ccd0bc44c5..597a5d3c29 100644 --- a/transformer_engine/common/fused_rope/fused_rope.cu +++ b/transformer_engine/common/fused_rope/fused_rope.cu @@ -155,18 +155,18 @@ __global__ void fused_rope_forward_kernel(const scalar_t *src, const int *cu_seq cur_seqlens = s; } - int s_id_for_freqs; + // Offset the RoPE embedding by start_positions if provided. + int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id]; + int s_id_for_freqs = s_id + begin_offset; + + // If CP_SIZE > 1, offset the RoPE embedding by cp_rank based on the dual-chunk order. if (cp_size > 1) { assert(cur_seqlens % 2 == 0); if (s_id < cur_seqlens / 2) { - s_id_for_freqs = s_id + cp_rank * cur_seqlens / 2; + s_id_for_freqs += cp_rank * cur_seqlens / 2; } else { - s_id_for_freqs = - cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 + s_id - cur_seqlens / 2; + s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2; } - } else { - int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id]; - s_id_for_freqs = s_id + begin_offset; } fused_rope_block_forward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block, @@ -175,11 +175,11 @@ __global__ void fused_rope_forward_kernel(const scalar_t *src, const int *cu_seq template __global__ void fused_rope_backward_kernel( - const scalar_t *src, const int *cu_seqlens, const float *freqs, scalar_t *dst, - const bool interleaved, const int cp_size, const int cp_rank, const int s, const int h, - const int d, const int d2, const int stride_s_or_t, const int stride_b, const int stride_h, - const int stride_d, const int o_stride_s_or_t, const int o_stride_b, const int o_stride_h, - const int o_stride_d) { + const scalar_t *src, const int *cu_seqlens, const float *freqs, const int *start_positions, + scalar_t *dst, const bool interleaved, const int cp_size, const int cp_rank, const int s, + const int h, const int d, const int d2, const int stride_s_or_t, const int stride_b, + const int stride_h, const int stride_d, const int o_stride_s_or_t, const int o_stride_b, + const int o_stride_h, const int o_stride_d) { int s_id = blockIdx.x, b_id = blockIdx.y; int offset_block, offset_block_dst; int cur_seqlens; @@ -197,17 +197,18 @@ __global__ void fused_rope_backward_kernel( cur_seqlens = s; } - int s_id_for_freqs; + // Offset the RoPE embedding by start_positions if provided. + int begin_offset = (start_positions == nullptr) ? 0 : start_positions[b_id]; + int s_id_for_freqs = s_id + begin_offset; + + // If CP_SIZE > 1, offset the RoPE embedding by cp_rank based on the dual-chunk order. if (cp_size > 1) { assert(cur_seqlens % 2 == 0); if (s_id < cur_seqlens / 2) { - s_id_for_freqs = s_id + cp_rank * cur_seqlens / 2; + s_id_for_freqs += cp_rank * cur_seqlens / 2; } else { - s_id_for_freqs = - cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 + s_id - cur_seqlens / 2; + s_id_for_freqs += cur_seqlens * cp_size - (cp_rank + 1) * cur_seqlens / 2 - cur_seqlens / 2; } - } else { - s_id_for_freqs = s_id; } fused_rope_block_backward(src, freqs, dst, interleaved, s_id_for_freqs, offset_block, @@ -495,12 +496,12 @@ void fused_rope_forward_launcher(const scalar_t *input, const int *cu_seqlens, c template void fused_rope_backward_launcher(const scalar_t *output_grads, const int *cu_seqlens, - const float *freqs, scalar_t *input_grads, - const NVTE_QKV_Format qkv_format, const bool interleaved, - const int cp_size, const int cp_rank, const int s, const int b, - const int h, const int d, const int d2, const int stride_s_or_t, - const int stride_b, const int stride_h, const int stride_d, - cudaStream_t stream) { + const float *freqs, const int *start_positions, + scalar_t *input_grads, const NVTE_QKV_Format qkv_format, + const bool interleaved, const int cp_size, const int cp_rank, + const int s, const int b, const int h, const int d, const int d2, + const int stride_s_or_t, const int stride_b, const int stride_h, + const int stride_d, cudaStream_t stream) { int warps_per_block = h < 16 ? 4 : 8; dim3 blocks(s, b); dim3 threads(THREADS_PER_WARP, warps_per_block); @@ -521,9 +522,9 @@ void fused_rope_backward_launcher(const scalar_t *output_grads, const int *cu_se const int o_stride_d = 1; fused_rope_backward_kernel<<>>( - output_grads, cu_seqlens, freqs, input_grads, interleaved, cp_size, cp_rank, s, h, d, d2, - stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b, o_stride_h, - o_stride_d); + output_grads, cu_seqlens, freqs, start_positions, input_grads, interleaved, cp_size, cp_rank, + s, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, o_stride_s_or_t, o_stride_b, + o_stride_h, o_stride_d); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -590,16 +591,18 @@ void fused_rope_forward(const Tensor &input, const Tensor &cu_seqlens, const Ten } void fused_rope_backward(const Tensor &output_grads, const Tensor &cu_seqlens, const Tensor &freqs, - Tensor *input_grads, const NVTE_QKV_Format qkv_format, - const bool interleaved, const int cp_size, const int cp_rank, const int s, - const int b, const int h, const int d, const int d2, - const int stride_s_or_t, const int stride_b, const int stride_h, - const int stride_d, cudaStream_t stream) { + const Tensor &start_positions, Tensor *input_grads, + const NVTE_QKV_Format qkv_format, const bool interleaved, + const int cp_size, const int cp_rank, const int s, const int b, + const int h, const int d, const int d2, const int stride_s_or_t, + const int stride_b, const int stride_h, const int stride_d, + cudaStream_t stream) { TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( output_grads.data.dtype, scalar_t, fused_rope_backward_launcher(reinterpret_cast(output_grads.data.dptr), reinterpret_cast(cu_seqlens.data.dptr), reinterpret_cast(freqs.data.dptr), + reinterpret_cast(start_positions.data.dptr), reinterpret_cast(input_grads->data.dptr), qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, stream);); @@ -663,18 +666,18 @@ void nvte_fused_rope_forward(const NVTETensor input, const NVTETensor cu_seqlens } void nvte_fused_rope_backward(const NVTETensor output_grads, const NVTETensor cu_seqlens, - const NVTETensor freqs, NVTETensor input_grads, - const NVTE_QKV_Format qkv_format, const bool interleaved, - const int cp_size, const int cp_rank, const int s, const int b, - const int h, const int d, const int d2, const int stride_s_or_t, - const int stride_b, const int stride_h, const int stride_d, - cudaStream_t stream) { + const NVTETensor freqs, const NVTETensor start_positions, + NVTETensor input_grads, const NVTE_QKV_Format qkv_format, + const bool interleaved, const int cp_size, const int cp_rank, + const int s, const int b, const int h, const int d, const int d2, + const int stride_s_or_t, const int stride_b, const int stride_h, + const int stride_d, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_rope_backward); using namespace transformer_engine; fused_rope_backward(*convertNVTETensorCheck(output_grads), *convertNVTETensorCheck(cu_seqlens), - *convertNVTETensorCheck(freqs), convertNVTETensorCheck(input_grads), - qkv_format, interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s_or_t, - stride_b, stride_h, stride_d, stream); + *convertNVTETensorCheck(freqs), *convertNVTETensorCheck(start_positions), + convertNVTETensorCheck(input_grads), qkv_format, interleaved, cp_size, + cp_rank, s, b, h, d, d2, stride_s_or_t, stride_b, stride_h, stride_d, stream); } void nvte_fused_qkv_rope_forward(const NVTETensor qkv_input, const NVTETensor q_freqs, diff --git a/transformer_engine/common/include/transformer_engine/fused_rope.h b/transformer_engine/common/include/transformer_engine/fused_rope.h index 610868f932..19047f463b 100644 --- a/transformer_engine/common/include/transformer_engine/fused_rope.h +++ b/transformer_engine/common/include/transformer_engine/fused_rope.h @@ -51,6 +51,7 @@ void nvte_fused_rope_forward(const NVTETensor input, const NVTETensor cu_seqlens * \param[in] cu_seqlens The cumulative sum of sequence lengths tensor. * (Required for the thd format, empty tensor for other formats) * \param[in] freqs The freqs tensor. + * \param[in] start_positions The beginning offsets for applying RoPE embeddings. * \param[out] input_grads Input gradient tensor to calculate. * \param[in] qkv_format QKV format. * \param[in] interleaved Whether to use interleaved rotary position embedding. @@ -68,12 +69,12 @@ void nvte_fused_rope_forward(const NVTETensor input, const NVTETensor cu_seqlens * \param[in] stream CUDA stream used for the operation. */ void nvte_fused_rope_backward(const NVTETensor output_grads, const NVTETensor cu_seqlens, - const NVTETensor freqs, NVTETensor input_grads, - const NVTE_QKV_Format qkv_format, const bool interleaved, - const int cp_size, const int cp_rank, const int s, const int b, - const int h, const int d, const int d2, const int stride_s_or_t, - const int stride_b, const int stride_h, const int stride_d, - cudaStream_t stream); + const NVTETensor freqs, const NVTETensor start_positions, + NVTETensor input_grads, const NVTE_QKV_Format qkv_format, + const bool interleaved, const int cp_size, const int cp_rank, + const int s, const int b, const int h, const int d, const int d2, + const int stride_s_or_t, const int stride_b, const int stride_h, + const int stride_d, cudaStream_t stream); /*! \brief Apply rotary positional embedding to the combined QKV input tensor. * diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index cc23d65a3e..0e1222c22f 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -149,7 +149,7 @@ def forward( cp_size, cp_rank, ) - ctx.save_for_backward(freqs, cu_seqlens) + ctx.save_for_backward(freqs, cu_seqlens, start_positions) ctx.tensor_format = tensor_format ctx.cp_size = cp_size ctx.cp_rank = cp_rank @@ -160,10 +160,11 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: """Fused RoPE backward.""" - freqs, cu_seqlens = ctx.saved_tensors + freqs, cu_seqlens, start_positions = ctx.saved_tensors grad_input = tex.fused_rope_backward( grad_output, freqs, + start_positions, QKVFormat[ctx.tensor_format], ctx.interleaved, cu_seqlens, @@ -171,7 +172,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ctx.cp_rank, ) - return grad_input, None, None, None, None, None, None, None + return grad_input, None, None, None, None, None, None, None, None class FusedQKVRoPEFunc(torch.autograd.Function): @@ -278,7 +279,6 @@ def _rotate_half(x: torch.Tensor, interleaved: bool) -> torch.Tensor: def _apply_rotary_pos_emb_base( t: torch.Tensor, freqs: torch.Tensor, - start_positions: torch.Tensor = None, tensor_format: str = "sbhd", interleaved: bool = False, ) -> torch.Tensor: @@ -291,45 +291,19 @@ def _apply_rotary_pos_emb_base( Input tensor of shape `[s, b, h, d]` or `[b, s, h, d]`, on which rotary positional embedding will be applied. freqs: torch.Tensor - Rotary positional embedding tensor of shape `[s2, 1, 1, d2]` and dtype 'float', - with `s2 >= s` and `d2 <= d`. - start_positions: torch.Tensor, default = None. - Tokens in a sequence `i` should be applied with position encoding offset by - `start_positions[i]`. If `start_positions=None`, there's no offset. + Rotary positional embedding tensor of shape `[s2, 1, 1, d2]` or `[s2, b, 1, d2]` + and dtype 'float', with `s2 >= s` and `d2 <= d`. tensor_format: {'sbhd', 'bshd'}, default = 'sbhd' Should be `bshd` if `t` is of shape `[bs, seq, ...]`, or `sbhd` if `t` is of shape `[seq, bs, ...]`. interleaved: bool, default = False Whether to use interleaved rotary position embedding. """ - max_seq_len = freqs.shape[0] - cur_seq_len = t.shape[1] if tensor_format == "bshd" else t.shape[0] - - # In case `start_positions` are provided, create a staggered `freqs` tensor - # offset by the values in `start_positions`. - # `start_positions` is only supported for `cp_size=1` and inference. - if start_positions is not None: - max_offset = torch.max(start_positions) - assert ( - max_offset + cur_seq_len <= max_seq_len - ), f"Rotary Embeddings only suppported up to {max_seq_len} sequence length!" - - # Stack staggered rope embeddings along the batch dimension - freqs = torch.concatenate([freqs[i : i + cur_seq_len] for i in start_positions], dim=1) - - # Note that from this point, `freqs` has a shape `(s,b,1,d)`. - - # Only apply the rotary embeddings up to the sequence length of the running - # input. - assert ( - cur_seq_len <= max_seq_len - ), f"Rotary Embeddings only supported up to {max_seq_len} sequence length!" - freqs = freqs[:cur_seq_len] - # [seq, 1, 1, dim] -> [1, seq, 1, dim] or # [seq, b, 1, dim] -> [b, seq, 1, dim] if tensor_format == "bshd": freqs = freqs.transpose(0, 1) + # cos/sin first then dtype conversion for better precision cos_ = torch.cos(freqs).to(t.dtype) sin_ = torch.sin(freqs).to(t.dtype) @@ -366,7 +340,7 @@ def _get_freqs_on_this_cp_rank( ) # cp_size == 1 - return freqs + return freqs[:seqlen] def apply_rotary_pos_emb( @@ -388,13 +362,13 @@ def apply_rotary_pos_emb( Training: qkv_formats: "thd", "bshd", "sbhd" context parallel: yes - start_positions: no + start_positions: yes interleaving: yes Inference: qkv_formats: "thd", "bshd", "sbhd" context parallelism: no start_positions: yes - interleaving: yes + interleaving: yes Parameters ---------- @@ -423,22 +397,17 @@ def apply_rotary_pos_emb( cp_rank: int, default = 0. Context parallel rank. Only valid when `tensor_format` is 'thd' and `fused` is True. """ - - # `start_positions` is only supported for `cp_size=1` and inference. - assert not ( - cp_size > 1 and start_positions is not None - ), """start_positions != None with CP SIZE > 1 is not supported!""" - assert ( tensor_format != "thd" or cu_seqlens is not None ), "cu_seqlens must not be None when tensor_format is 'thd'." + # Fused apply rope logic for THD/BSHD/SBHD formats if fused: return FusedRoPEFunc.apply( t, freqs, start_positions, tensor_format, interleaved, cu_seqlens, cp_size, cp_rank ) - # Unfused THD format + # Unfused apply rope logic for THD format if tensor_format == "thd": cu_seqlens = cu_seqlens // cp_size seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() @@ -447,15 +416,18 @@ def apply_rotary_pos_emb( # `s1hd` tensors (for each sequence) and applies rotary embedding to # those sequences individually. # Note that if `start_positions` is not `None`, then for each sequence, - # it's corresponding rope offset is also supplied from `start_positions` - # individually. + # the freqs supplied are offset by the corresponding `start_positions` value. return torch.cat( [ _apply_rotary_pos_emb_base( x.unsqueeze(1), - _get_freqs_on_this_cp_rank(freqs, x.size(0), cp_size, cp_rank), - start_positions=( - start_positions[idx : idx + 1] if start_positions is not None else None + _get_freqs_on_this_cp_rank( + ( + freqs[start_positions[idx] :] if start_positions is not None else freqs + ), # offset the freqs + x.size(0), + cp_size, + cp_rank, ), interleaved=interleaved, ) @@ -463,17 +435,28 @@ def apply_rotary_pos_emb( ] ).squeeze(1) - # Unfused SBHD/BSHD format + # Unfused apply rope logic for SBHD/BSHD format follows ... + if tensor_format == "sbhd": seqlen = t.size(0) elif tensor_format == "bshd": seqlen = t.size(1) else: raise ValueError(f"Unsupported tensor_format: {tensor_format}.") + + if start_positions is not None: + max_offset = torch.max(start_positions) + assert ( + max_offset + seqlen * cp_size <= freqs.shape[0] + ), f"Rotary Embeddings only suppported up to {freqs.shape[0]} sequence length!" + + # Stack staggered rope embeddings along the batch dimension + freqs = torch.concatenate([freqs[i : i + seqlen * cp_size] for i in start_positions], dim=1) + # Note that from this point, `freqs` has a shape `(s,b,1,d)`. + return _apply_rotary_pos_emb_base( t, _get_freqs_on_this_cp_rank(freqs, seqlen, cp_size, cp_rank), - start_positions, tensor_format, interleaved=interleaved, ) @@ -505,7 +488,7 @@ def apply_fused_qkv_rotary_pos_emb( qkv_formats: "bshd", "sbhd" context parallelism: no start_positions: yes - interleaving: yes + interleaving: yes Parameters ---------- diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 43eab96544..77fb348589 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -346,6 +346,7 @@ at::Tensor fused_rope_forward(const at::Tensor &input, const at::Tensor &freqs, const int cp_rank); at::Tensor fused_rope_backward(const at::Tensor &output_grads, const at::Tensor &freqs, + const std::optional start_positions, const NVTE_QKV_Format qkv_format, const bool interleaved, const std::optional cu_seqlens, const int cp_size, const int cp_rank); diff --git a/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp b/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp index 064da8a670..d1dcf68c3d 100644 --- a/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp +++ b/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp @@ -163,6 +163,7 @@ std::tuple fused_qkv_rope_forward( } at::Tensor fused_rope_backward(const at::Tensor &output_grads, const at::Tensor &freqs, + const std::optional start_positions, const NVTE_QKV_Format qkv_format, const bool interleaved, const std::optional cu_seqlens, const int cp_size, const int cp_rank) { @@ -180,6 +181,12 @@ at::Tensor fused_rope_backward(const at::Tensor &output_grads, const at::Tensor auto freqs_cu = makeTransformerEngineTensor(freqs); auto input_grads_cu = makeTransformerEngineTensor(input_grads); + auto start_positions_cu = TensorWrapper(); // empty start_positions tensor + if (start_positions) { + start_positions_cu = makeTransformerEngineTensor(start_positions.value()); + TORCH_CHECK(start_positions_cu.ndim() == 1, "expected 1D tensor"); + } + if (qkv_format == NVTE_QKV_Format::NVTE_THD) { TORCH_CHECK(output_grads.dim() == 3, "expected 3D tensor"); TORCH_CHECK(cu_seqlens.has_value(), "expected cu_seqlens tensor"); @@ -208,8 +215,8 @@ at::Tensor fused_rope_backward(const at::Tensor &output_grads, const at::Tensor auto cu_seqlens_cu = makeTransformerEngineTensor(cu_seqlens.value()); nvte_fused_rope_backward(output_grads_cu.data(), cu_seqlens_cu.data(), freqs_cu.data(), - input_grads_cu.data(), qkv_format, interleaved, cp_size, cp_rank, - max_s, b, h, d, d2, stride_t, + start_positions_cu.data(), input_grads_cu.data(), qkv_format, + interleaved, cp_size, cp_rank, max_s, b, h, d, d2, stride_t, /*stride_b=*/0, stride_h, stride_d, at::cuda::getCurrentCUDAStream()); return input_grads; @@ -246,9 +253,9 @@ at::Tensor fused_rope_backward(const at::Tensor &output_grads, const at::Tensor auto cu_seqlens_cu = TensorWrapper(); // empty cu_seqlens tensor nvte_fused_rope_backward(output_grads_cu.data(), cu_seqlens_cu.data(), freqs_cu.data(), - input_grads_cu.data(), qkv_format, interleaved, cp_size, cp_rank, s, b, - h, d, d2, stride_s, stride_b, stride_h, stride_d, - at::cuda::getCurrentCUDAStream()); + start_positions_cu.data(), input_grads_cu.data(), qkv_format, + interleaved, cp_size, cp_rank, s, b, h, d, d2, stride_s, stride_b, + stride_h, stride_d, at::cuda::getCurrentCUDAStream()); return input_grads; } From c544ced2ea3c06950be9c33ad1802c831e44ff58 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 12 Nov 2025 18:55:14 -0500 Subject: [PATCH 057/521] [JAX] Relax tolerance for the test_multiprocessing_encoder.py with NVFP4 by 0.001 (#2375) relax tol Signed-off-by: Phuong Nguyen --- examples/jax/encoder/test_multiprocessing_encoder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index 4d21411169..f3092278e8 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -672,7 +672,7 @@ def test_te_mxfp8(self): def test_te_nvfp4(self): """Test Transformer Engine with NVFP4""" result = self.exec(True, "NVFP4BlockScaling") - assert result[0] < 0.451 and result[1] > 0.788 + assert result[0] < 0.451 and result[1] > 0.787 @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") def test_te_bf16_shardy(self): @@ -710,7 +710,7 @@ def test_te_mxfp8_shardy(self): def test_te_nvfp4_shardy(self): """Test Transformer Engine with NVFP4""" result = self.exec(True, "NVFP4BlockScaling", enable_shardy=True) - assert result[0] < 0.451 and result[1] > 0.788 + assert result[0] < 0.451 and result[1] > 0.787 if __name__ == "__main__": From d8f1e68f7c414f3e7985a8b41de4443b2f819af3 Mon Sep 17 00:00:00 2001 From: Lifu Zhang Date: Wed, 12 Nov 2025 16:34:25 -0800 Subject: [PATCH 058/521] fix gradient accumulation fusion for FSDP (#2371) Signed-off-by: Lifu Zhang Co-authored-by: Lifu Zhang Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/pytorch/module/grouped_linear.py | 6 +++--- transformer_engine/pytorch/module/layernorm_linear.py | 4 ++-- transformer_engine/pytorch/module/linear.py | 4 ++-- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 59dc2b2997..f336a743da 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -293,9 +293,9 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], origin_weights[i] = ctx.weight_objects[i] ctx.weight_objects[i] = None - if ctx.fuse_wgrad_accumulation: - for i in range(N): - origin_weights[i].main_grad = main_grads[i] + if ctx.fuse_wgrad_accumulation: + for i in range(N): + origin_weights[i].main_grad = main_grads[i] # Preprocess grad output grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index abe8c58298..20a67cba48 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -572,8 +572,8 @@ def backward( if ctx.cpu_offloading: if ctx.grad_added_to_main_grad: origin_weight = ctx.weight_object - if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: - origin_weight.main_grad = main_grad + if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: + origin_weight.main_grad = main_grad # Configure Userbuffers communication (comm+GEMM overlap) ctx.ub_obj_gradout = None diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 0e2310a5a2..46b9dbd85b 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -508,8 +508,8 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if ctx.cpu_offloading: if ctx.grad_added_to_main_grad: weight = ctx.weight_object - if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: - weight.main_grad = main_grad + if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: + weight.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already From d0d4063130e3ae8b40e557919eb04bc76b721c0c Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Thu, 13 Nov 2025 14:28:09 +0100 Subject: [PATCH 059/521] [PyTorch] Fix amax computation using output_t data in normalization (#2355) Fix amax computation using output_t data in normalization Signed-off-by: Evgeny --- tests/cpp/operator/test_normalization.h | 12 +++++++++++- .../normalization/layernorm/ln_fwd_kernels.cuh | 18 ++++++++++++++++-- .../rmsnorm/rmsnorm_fwd_kernels.cuh | 18 ++++++++++++++++-- 3 files changed, 43 insertions(+), 5 deletions(-) diff --git a/tests/cpp/operator/test_normalization.h b/tests/cpp/operator/test_normalization.h index fe69852d00..271345686e 100644 --- a/tests/cpp/operator/test_normalization.h +++ b/tests/cpp/operator/test_normalization.h @@ -114,8 +114,18 @@ void compute_ref_output(NormType norm_type, tmp = current * rsigma[i] * g; } + // Write output (scaled only for fp8 paths) output[i * H + j] = static_cast(tmp * scale); - current_max = fmaxf(current_max, fabsf(tmp)); + + // amax semantics: + // - fp8_out (scale != 1): amax on pre-scale compute value 'tmp' + // - non-fp8_out (scale == 1): amax on value converted to OutputType (e.g., bf16) + if (scale != 1.f) { + current_max = fmaxf(current_max, fabsf(tmp)); + } else { + OutputType out_t_val = static_cast(tmp); + current_max = fmaxf(current_max, fabsf(static_cast(out_t_val))); + } } } diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 6050b164d5..38c4096073 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -123,7 +123,14 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_tuned_kernel( if (requires_amax) { __builtin_assume(amax >= 0); - amax = fmaxf(amax, fabsf(temp_output)); + if (params.fp8_out) { + // For fp8_out, keep amax on pre-scale compute_t + amax = fmaxf(amax, fabsf(temp_output)); + } else { + // Otherwise compute amax on the value converted to output_t (e.g., bf16) + output_t out_t_val = output_t(temp_output); + amax = fmaxf(amax, fabsf(compute_t(out_t_val))); + } } if (params.fp8_out) { temp_output = temp_output * scale; @@ -290,7 +297,14 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void ln_fwd_general_kerne if (col + jt < params.cols) { compute_t z_ij = z.data.elt[jt]; __builtin_assume(amax >= 0); - amax = fmaxf(amax, fabsf(z_ij)); + if (params.fp8_out) { + // For fp8_out, keep amax on pre-scale compute_t + amax = fmaxf(amax, fabsf(z_ij)); + } else { + // Otherwise compute amax on the value converted to output_t (e.g., bf16) + output_t out_t_val = output_t(z_ij); + amax = fmaxf(amax, fabsf(compute_t(out_t_val))); + } if (params.fp8_out) { z.data.elt[jt] = z_ij * scale; } diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index fc093b73a7..7fed7f123a 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -115,7 +115,14 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_tuned_ke if (requires_amax) { __builtin_assume(amax >= 0); - amax = fmaxf(amax, fabsf(temp_output)); + if (params.fp8_out) { + // For fp8_out, keep amax on pre-scale compute_t + amax = fmaxf(amax, fabsf(temp_output)); + } else { + // Otherwise compute amax on the value converted to output_t (e.g., bf16) + output_t out_t_val = output_t(temp_output); + amax = fmaxf(amax, fabsf(compute_t(out_t_val))); + } } if (params.fp8_out) { temp_output = temp_output * scale; @@ -265,7 +272,14 @@ __global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) void rmsnorm_fwd_general_ if (col + jt < params.cols) { compute_t z_ij = z.data.elt[jt]; __builtin_assume(amax >= 0); - amax = fmaxf(amax, fabsf(z_ij)); + if (params.fp8_out) { + // For fp8_out, keep amax on pre-scale compute_t + amax = fmaxf(amax, fabsf(z_ij)); + } else { + // Otherwise compute amax on the value converted to output_t (e.g., bf16) + output_t out_t_val = output_t(z_ij); + amax = fmaxf(amax, fabsf(compute_t(out_t_val))); + } if (params.fp8_out) { z.data.elt[jt] = z_ij * scale; } From ef28c86582b5539655e7286366c699bec99284bc Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 13 Nov 2025 12:30:59 -0500 Subject: [PATCH 060/521] [JAX] NVFP4 scale swizzling via nvte kernel (#2350) * swizzle via nvte Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- transformer_engine/jax/cpp_extensions/gemm.py | 5 ++ .../jax/csrc/extensions/gemm.cpp | 62 ++++++++++++++----- 2 files changed, 50 insertions(+), 17 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 72bee251c4..9ffec2c6af 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -533,6 +533,9 @@ def _dims_are_consecutive(dims): # Declare cuBLAS workspace workspace_size = get_cublas_workspace_size_bytes() + # NVFP4 swizzling happen in via nvte kernel instead of JAX transposes + if scaling_mode.is_nvfp4_scaling: + workspace_size += lhs_scale_inv.size + rhs_scale_inv.size if not collective_op.is_none: workspace_size *= get_cgemm_num_max_streams() # cuBLAS workspace ptr must be 256 bytes aligned but JAX buffers are not @@ -662,6 +665,8 @@ def impl( rhs_scale_inv = apply_padding_to_scale_inv( rhs_scale_inv, scaling_mode, rhs.shape, not rhs_transposed, rhs_flatten_axis ) + # Only perform JAX-based swizzle for MXFP8, NVFP4 swizzle will go though nvte kernel + if scaling_mode.is_mxfp8_scaling: lhs_scale_inv = swizzled_scale(lhs_scale_inv, lhs_flatten_axis, lhs_transposed) rhs_scale_inv = swizzled_scale(rhs_scale_inv, rhs_flatten_axis, not rhs_transposed) diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 8a3658a0ba..6566ff1689 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -34,8 +34,8 @@ static uint8_t *move_ptr_to_next_256B_aligned(uint8_t *ptr) { } std::tuple> xla_buffer_to_nvte_gemm_operand( - cudaStream_t stream, Buffer_Type buffer, Buffer_Type scale_inv, JAXX_Scaling_Mode scaling_mode, - size_t axis_boundary, bool rowwise) { + cudaStream_t stream, Buffer_Type buffer, Buffer_Type scale_inv, uint8_t *swizzle_scale_ptr, + JAXX_Scaling_Mode scaling_mode, size_t axis_boundary, bool rowwise) { // Set tensor data with collapsed 2D shape auto buffer_dims = buffer.dimensions(); std::vector input_shape = {product(buffer_dims, 0, axis_boundary), @@ -56,17 +56,32 @@ std::tuple> xla_buffer_to_nvte_gemm_operand( NVTE_CHECK(scale_inv.element_count() > 0, "Missing inverse scaling factor for quantized GEMM."); std::vector scale_shape = {1}; - if (scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING) { + auto is_nvfp4 = is_nvfp4_scaling(scaling_mode); + auto scale_dtype = convert_ffi_datatype_to_te_dtype(scale_inv.element_type()); + if (scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING || is_nvfp4) { // Block scaling also needs to be collapsed to match 2D data scale_shape = {product(scale_inv.dimensions(), 0, axis_boundary), product(scale_inv.dimensions(), axis_boundary, scale_inv.dimensions().size())}; + NVTE_CHECK(typeToSize(scale_dtype) == 1, + "Inverse scale factors need to have an 8-bit data type."); } - - auto scale_dtype = convert_ffi_datatype_to_te_dtype(scale_inv.element_type()); - if (rowwise) { + if (!is_nvfp4) { + if (rowwise) { + input.set_rowwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); + } else { + input.set_columnwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); + } + } else { // Swizzle for NVFP4 + NVTE_CHECK(rowwise, "NVFP4 GEMM expects rowwise for both LHS and RHS"); input.set_rowwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); - } else { - input.set_columnwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); + // Create tensor to hold swizzled scale factor + TensorWrapper output(get_nvte_scaling_mode(scaling_mode)); + output.set_rowwise_data(buffer.untyped_data(), input_dtype, input_shape); + output.set_rowwise_scale_inv(swizzle_scale_ptr, scale_dtype, scale_shape); + // Launch swizzle kernel + nvte_swizzle_scaling_factors(input.data(), output.data(), stream); + // Set swizzled scales into the input tensor + input.set_rowwise_scale_inv(swizzle_scale_ptr, scale_dtype, scale_shape); } } @@ -145,16 +160,34 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i int64_t lhs_axis_boundary, int64_t rhs_axis_boundary, bool lhs_transposed, bool rhs_transposed, bool fuse_bias, bool fuse_gelu, bool grad, bool use_split_accumulator, JAXX_Collective_Op collective_op) { + // cuBLAS workspace + 256 alignment enforcement (+ swizzle scales) + uint8_t *lhs_swizzle_scale_ptr = nullptr, *rhs_swizzle_scale_ptr = nullptr; + auto workspace_ptr = reinterpret_cast(workspace->untyped_data()); + workspace_ptr = move_ptr_to_next_256B_aligned(workspace_ptr); + size_t workspace_size = static_cast(workspace->element_count()) - 256; + if (is_nvfp4_scaling(scaling_mode)) { + auto lhs_scale_size = product(lhs_scale_inv.dimensions()); + auto rhs_scale_size = product(rhs_scale_inv.dimensions()); + workspace_size = workspace_size - lhs_scale_size - rhs_scale_size; + lhs_swizzle_scale_ptr = workspace_ptr; + rhs_swizzle_scale_ptr = workspace_ptr + lhs_scale_size; + workspace_ptr = rhs_swizzle_scale_ptr + rhs_scale_size; + } + auto workspace_ = TensorWrapper(workspace_ptr, std::vector{workspace_size}, DType::kByte); + // NOTE: TensorWrapper operands are always rowwise for full-precision GEMM, or FP8 GEMM when // device supports non-TN layouts (compute capability >= 10.0, excluding 12.x) bool always_rowwise = (scaling_mode == JAXX_Scaling_Mode::NO_SCALING || (is_tensor_scaling(scaling_mode) && nvte_is_non_tn_fp8_gemm_supported())); bool make_lhs_rowwise = (always_rowwise) ? true : !lhs_transposed; bool make_rhs_rowwise = (always_rowwise) ? true : rhs_transposed; - auto [lhs_, lhs_shape] = xla_buffer_to_nvte_gemm_operand(stream, lhs, lhs_scale_inv, scaling_mode, - lhs_axis_boundary, make_lhs_rowwise); - auto [rhs_, rhs_shape] = xla_buffer_to_nvte_gemm_operand(stream, rhs, rhs_scale_inv, scaling_mode, - rhs_axis_boundary, make_rhs_rowwise); + + auto [lhs_, lhs_shape] = + xla_buffer_to_nvte_gemm_operand(stream, lhs, lhs_scale_inv, lhs_swizzle_scale_ptr, + scaling_mode, lhs_axis_boundary, make_lhs_rowwise); + auto [rhs_, rhs_shape] = + xla_buffer_to_nvte_gemm_operand(stream, rhs, rhs_scale_inv, rhs_swizzle_scale_ptr, + scaling_mode, rhs_axis_boundary, make_rhs_rowwise); std::vector out_shape = {(lhs_transposed) ? lhs_shape[1] : lhs_shape[0], (rhs_transposed) ? rhs_shape[0] : rhs_shape[1]}; @@ -191,11 +224,6 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i } auto pre_gelu_ = TensorWrapper(pre_gelu_ptr, pre_gelu_shape, pre_gelu_dtype); - // cuBLAS workspace + 256 alignment enforcement - auto workspace_ptr = reinterpret_cast(workspace->untyped_data()); - workspace_ptr = move_ptr_to_next_256B_aligned(workspace_ptr); - std::vector workspace_shape = {static_cast(workspace->element_count()) - 256}; - auto workspace_ = TensorWrapper(workspace_ptr, workspace_shape, DType::kByte); auto num_math_sm = cuda::sm_count() - getenv("NVTE_EXT_MARGIN_SM", 0); float one = 1.; From 9440b76aa8d4333d1347c44f415c5789f1e038be Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 13 Nov 2025 14:05:18 -0500 Subject: [PATCH 061/521] [JAX] Shardy rule + QuantizeLayout Rework (#2364) * shardy + quantize_layout rework Signed-off-by: Phuong Nguyen * add assertion for NVFP4 in fused act and fused norm primitive Signed-off-by: Phuong Nguyen * add assertions Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- .../jax/cpp_extensions/activation.py | 170 +++++++++--------- transformer_engine/jax/cpp_extensions/misc.py | 8 +- .../jax/cpp_extensions/normalization.py | 109 ++++++----- .../jax/cpp_extensions/quantization.py | 90 +++++----- transformer_engine/jax/csrc/extensions.h | 6 +- .../jax/csrc/extensions/activation.cpp | 45 ++--- transformer_engine/jax/csrc/extensions/misc.h | 14 +- .../jax/csrc/extensions/normalization.cpp | 29 ++- .../jax/csrc/extensions/pybind.cpp | 9 +- .../jax/csrc/extensions/quantization.cpp | 40 ++--- transformer_engine/jax/quantize/__init__.py | 1 + transformer_engine/jax/quantize/misc.py | 61 +++++++ transformer_engine/jax/quantize/quantizer.py | 55 ++---- .../jax/quantize/scaling_modes.py | 130 ++++++++++---- transformer_engine/jax/quantize/tensor.py | 36 ++-- 15 files changed, 456 insertions(+), 347 deletions(-) create mode 100644 transformer_engine/jax/quantize/misc.py diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index bb3c56bcf1..aa84fafd3e 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -10,7 +10,7 @@ import jax import jax.numpy as jnp from jax import dtypes, ffi -from jax.experimental.custom_partitioning import SdyShardingRule +from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING from jax.sharding import PartitionSpec import numpy as np @@ -159,7 +159,7 @@ class ActLuPrimitive(BasePrimitive): 11, 12, 13, - ) # out_dtype, act_enum, act_len, scaling_mode, is_2x, scale_dtype, act_params, amax_scope, transpose_batch_sequence, output_amax_when_no_scaling, is_outer + ) # out_dtype, act_enum, act_len, scaling_mode, quantize_layout, scale_dtype, act_params, amax_scope, transpose_batch_sequence, output_amax_when_no_scaling, is_outer inner_primitive = None outer_primitive = None @@ -173,7 +173,7 @@ def abstract( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -201,6 +201,13 @@ def abstract( "Current tensor scaling is not yet supported for fused activation and quantization." " Please do activation in higher-precision then quantize with current tensor scaling." ) + assert not ScalingMode(scaling_mode).is_nvfp4_scaling, ( + "NVFP4 block scaling is not yet supported for fused activation and quantization." + " Please do activation in higher-precision then quantize with current tensor scaling." + ) + assert ( + not quantize_layout.is_colwise_only + ), "Fused activation with colwise-only quantization is not supported." out_shape = (*x_aval.shape[:-2], x_aval.shape[-1]) # Exclude act dim out_aval = x_aval.update(shape=out_shape, dtype=out_dtype) @@ -210,7 +217,7 @@ def abstract( rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( scaling_mode ).get_scale_shape_2x(out_shape, is_padded=not is_outer, flatten_axis=-1) - if not is_2x: + if quantize_layout.is_rowwise_only: out_shape = (1,) colwise_scale_inv_shape = (1,) colwise_out_aval = jax.core.ShapedArray(shape=out_shape, dtype=out_dtype) @@ -232,7 +239,7 @@ def lowering( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -259,7 +266,7 @@ def lowering( amax, act_enum=act_enum, scaling_mode=scaling_mode.value, - is_2x=is_2x, + quantize_layout=quantize_layout.value.value, act_params=act_params.to_ffi_lowering_dict(), output_amax_when_no_scaling=output_amax_when_no_scaling, ) @@ -274,7 +281,7 @@ def impl( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -297,7 +304,7 @@ def impl( act_enum=act_enum, act_len=act_len, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, act_params=act_params, amax_scope=amax_scope, @@ -313,7 +320,7 @@ def impl( scale_inv = jax.lax.slice( scale_inv, [0] * len(rowwise_scale_inv_shape), rowwise_scale_inv_shape ) - if is_2x: + if quantize_layout.is_rowwise_colwise: colwise_scale_inv = jax.lax.slice( colwise_scale_inv, [0] * len(colwise_scale_inv_shape), colwise_scale_inv_shape ) @@ -329,7 +336,7 @@ def batcher( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -356,7 +363,7 @@ def batcher( act_enum=act_enum, act_len=act_len, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, act_params=act_params, amax_scope=amax_scope, @@ -373,7 +380,7 @@ def infer_sharding_from_operands( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -402,7 +409,7 @@ def infer_sharding_from_operands( out_spec = (*x_spec[:-2], x_spec[-1]) out_sharding = NamedSharding(mesh, PartitionSpec(*out_spec), desc="ActLuPrimitive.out") - if is_2x: + if quantize_layout.is_rowwise_colwise: if scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING.value: colwise_out_spec = multidim_transpose(out_spec, transpose_axis=-1) else: @@ -419,7 +426,7 @@ def infer_sharding_from_operands( elif scaling_mode == ScalingMode.MXFP8_1D_SCALING.value: scale_inv_spec = out_spec - if is_2x: + if quantize_layout.is_rowwise_colwise: colwise_scale_inv_spec = scale_inv_spec scale_inv_sharding = NamedSharding( @@ -444,7 +451,7 @@ def partition( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -462,7 +469,7 @@ def partition( out_spec = (*x_spec[:-2], x_spec[-1]) out_sharding = NamedSharding(mesh, PartitionSpec(*out_spec), desc="ActLuPrimitive.out") - if is_2x: + if quantize_layout.is_rowwise_colwise: if scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING.value: colwise_out_spec = multidim_transpose(out_spec, transpose_axis=-1) else: @@ -479,7 +486,10 @@ def partition( elif scaling_mode == ScalingMode.MXFP8_1D_SCALING.value: scale_inv_spec = out_spec - if is_2x: + if quantize_layout.is_rowwise_colwise: + assert not ScalingMode( + scaling_mode + ).is_colwise_transposed, "Transpose layout scaling modes are not supported here yet" colwise_scale_inv_spec = scale_inv_spec scale_inv_sharding = NamedSharding( @@ -514,7 +524,7 @@ def sharded_impl(x, scale, amax): act_enum=act_enum, act_len=act_len, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, act_params=act_params, amax_scope=amax_scope, @@ -550,7 +560,7 @@ def shardy_sharding_rule( act_enum, act_len, scaling_mode, - is_2x, + quantize_layout, scale_dtype, act_params, amax_scope, @@ -574,37 +584,28 @@ def shardy_sharding_rule( mesh, result_types, ) - prefix = "ActLu_" + prefix = "ActLu" input_shape = value_types[0].shape output_shape = input_shape[:-2] + input_shape[-1:] # Here we pass len of output so that the scales are propagated correctly scale_rules = ScalingMode(scaling_mode).get_shardy_sharding_rules( - output_shape, unique_var=prefix + "x", flatten_axis=-1 + output_shape, unique_var=prefix, flatten_axis=-1, q_layout=quantize_layout ) - x_axes = scale_rules.input_spec - # Correct input spec with act dim - x_axes = x_axes[:-1] + (prefix + "_act_dim",) + x_axes[-1:] - out = scale_rules.input_spec - - colwise_out = (prefix + "out_colwise",) - colwise_scale_inv = (prefix + "scale_inv_colwise",) - if is_2x: - colwise_scale_inv = scale_rules.colwise_rule - if scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING.value: - colwise_out = multidim_transpose(out, transpose_axis=-1) - else: - colwise_out = out - colwise_scale_inv = scale_rules.colwise_rule - - amax = (prefix + "amax",) + # Correct the input spec with act dim + input_spec = scale_rules.input_spec + input_spec = input_spec[:-1] + (prefix + "_act_dim",) + input_spec[-1:] + amax = (BATCHING + prefix + "_amax",) + scale = (BATCHING + prefix + "_scale",) return SdyShardingRule( + (tuple(input_spec), scale, amax), ( - x_axes, - ("…1",), + scale_rules.rowwise_out_spec, + scale_rules.colwise_out_spec, + scale_rules.rowwise_scale_spec, + scale_rules.colwise_scale_spec, amax, ), - (out, colwise_out, scale_rules.rowwise_rule, colwise_scale_inv, amax), **scale_rules.factor_sizes, ) @@ -612,7 +613,6 @@ def shardy_sharding_rule( register_primitive(ActLuPrimitive) -# TODO(Jeremy): replace is_2x with q_layout class BaseDActLuDBiasQuantizePrimitive(BasePrimitive): """ DActLu DBias Cast Transpose Primitive @@ -620,7 +620,7 @@ class BaseDActLuDBiasQuantizePrimitive(BasePrimitive): name = "te_dact_dbias_quantize_ffi" multiple_results = True - # out_dtype, scaling_mode, is_2x, scale_dtype, is_dbias, act_enum, act_len, act_params, amax_scope, transpose_batch_sequence, output_amax_when_no_scaling, is_outer + # out_dtype, scaling_mode, quantize_layout, scale_dtype, is_dbias, act_enum, act_len, act_params, amax_scope, transpose_batch_sequence, output_amax_when_no_scaling, is_outer impl_static_args = (4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15) inner_primitive = None outer_primitive = None @@ -634,7 +634,7 @@ def abstract( *, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -678,7 +678,7 @@ def abstract( rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( scaling_mode ).get_scale_shape_2x(x_aval.shape, is_padded=not is_outer, flatten_axis=-2) - if is_2x: + if quantize_layout.is_rowwise_colwise: if ScalingMode(scaling_mode).is_tensor_scaling(): colwise_out_shape = multidim_transpose(out_shape, transpose_axis=-2) else: @@ -700,7 +700,7 @@ def abstract( jax_dtype_to_te_dtype(x_aval.dtype), jax_dtype_to_te_dtype(out_dtype), scaling_mode, - is_2x, + quantize_layout.value, ) wkspace_shape = wkspace_info[0] wkspace_dtype = te_dtype_to_jax_dtype(wkspace_info[1]) @@ -741,7 +741,7 @@ def lowering( *, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -777,7 +777,7 @@ def lowering( scale, amax, scaling_mode=scaling_mode.value, - is_2x=is_2x, + quantize_layout=quantize_layout.value.value, is_dbias=is_dbias, act_enum=int(act_enum), act_params=act_params.to_ffi_lowering_dict(), @@ -792,7 +792,7 @@ def impl( amax, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -816,7 +816,7 @@ def impl( amax, out_dtype=out_dtype, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, is_dbias=is_dbias, act_enum=act_enum, @@ -835,7 +835,7 @@ def impl( scale_inv = jax.lax.slice( scale_inv, [0] * len(rowwise_scale_inv_shape), rowwise_scale_inv_shape ) - if is_2x: + if quantize_layout.is_rowwise_colwise: colwise_scale_inv = jax.lax.slice( colwise_scale_inv, [0] * len(colwise_scale_inv_shape), colwise_scale_inv_shape ) @@ -848,7 +848,7 @@ def batcher( *, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -883,7 +883,7 @@ def batcher( amax, out_dtype=out_dtype, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, is_dbias=is_dbias, act_enum=act_enum, @@ -901,7 +901,7 @@ def batcher( def infer_sharding_from_operands( out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -928,7 +928,7 @@ def infer_sharding_from_operands( out_sharding = NamedSharding( mesh, PartitionSpec(*x_spec), desc="BaseDActLuDBiasQuantizePrimitive.out" ) - if is_2x: + if quantize_layout.is_rowwise_colwise: if scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING.value: colwise_x_spec = multidim_transpose(x_spec, transpose_axis=-2) else: @@ -954,7 +954,7 @@ def infer_sharding_from_operands( elif scaling_mode == ScalingMode.MXFP8_1D_SCALING.value: scale_inv_spec = x_spec - if is_2x: + if quantize_layout.is_rowwise_colwise: colwise_scale_inv_spec = scale_inv_spec scale_inv_sharding = NamedSharding( @@ -981,7 +981,7 @@ def infer_sharding_from_operands( def partition( out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -1003,7 +1003,7 @@ def partition( mesh, PartitionSpec(*x_spec), desc="BaseDActLuDBiasQuantizePrimitive.out" ) - if is_2x: + if quantize_layout.is_rowwise_colwise: if scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING.value: colwise_x_spec = multidim_transpose(x_spec, transpose_axis=-2) else: @@ -1029,7 +1029,7 @@ def partition( elif scaling_mode == ScalingMode.MXFP8_1D_SCALING.value: scale_inv_spec = x_spec - if is_2x: + if quantize_layout.is_rowwise_colwise: colwise_scale_inv_spec = scale_inv_spec scale_inv_sharding = NamedSharding( @@ -1066,7 +1066,7 @@ def sharded_impl(dz, x, scale, amax): amax, out_dtype=out_dtype, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, is_dbias=is_dbias, act_enum=act_enum, @@ -1102,7 +1102,7 @@ def sharded_impl(dz, x, scale, amax): def shardy_sharding_rule( out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, is_dbias, act_enum, @@ -1132,28 +1132,30 @@ def shardy_sharding_rule( ) prefix = "DActLuDBias_" + # get sharding rules base on the input shape scale_rules = ScalingMode(scaling_mode).get_shardy_sharding_rules( - value_types[1].shape, unique_var=prefix + "x", flatten_axis=-2 + value_types[1].shape, + unique_var=prefix, + flatten_axis=-2, + q_layout=quantize_layout, ) - x_axes = scale_rules.input_spec - dz_axes = (*x_axes[:-2], x_axes[-1]) - out = x_axes - colwise_out = (prefix + "out_colwise",) - colwise_scale_inv = (prefix + "scale_inv_colwise",) - if is_2x: - if scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING.value: - colwise_out = tuple(multidim_transpose(x_axes, transpose_axis=-2)) - else: - colwise_out = out - colwise_scale_inv = scale_rules.colwise_rule - - dbias = x_axes[-2:] if is_dbias else (prefix + "dbias",) - amax = (prefix + "amax",) + input_spec = scale_rules.input_spec + dz_spec = (*input_spec[:-2], input_spec[-1]) + dbias = input_spec[-2:] if is_dbias else (prefix + "_dbias",) + amax = (prefix + "_amax",) + scale = (prefix + "_scale",) return SdyShardingRule( - (dz_axes, x_axes, ("…2",), amax), - (out, colwise_out, scale_rules.rowwise_rule, colwise_scale_inv, amax, dbias), + (tuple(dz_spec), tuple(input_spec), scale, amax), + ( + scale_rules.rowwise_out_spec, + scale_rules.colwise_out_spec, + scale_rules.rowwise_scale_spec, + scale_rules.colwise_scale_spec, + amax, + dbias, + ), **scale_rules.factor_sizes, ) @@ -1269,7 +1271,7 @@ def act_lu( return _jax_act_lu(x, activation_type, quantizer, act_params) # TE/common does not support colwise-only quantization yet - if quantizer is not None and quantizer.q_layout == QuantizeLayout.COLWISE: + if quantizer is not None and quantizer.q_layout.is_colwise_only: return _jax_act_lu(x, activation_type, quantizer, act_params) # TE/common does not support 2x quantization for DelayedScaling yet war_output = try_apply_delayed_scaling_2x_war( @@ -1298,7 +1300,7 @@ def act_lu( act_enum=act_type_id, act_len=act_len, scaling_mode=ScalingMode.NO_SCALING.value, - is_2x=False, + quantize_layout=QuantizeLayout.ROWWISE, scale_dtype=jnp.float32, act_params=act_params, amax_scope=amax_scope, @@ -1354,7 +1356,7 @@ def act_lu( act_enum=act_type_id, act_len=act_len, scaling_mode=quantizer.scaling_mode.value, - is_2x=quantizer.is_2x2x(), + quantize_layout=quantizer.q_layout, scale_dtype=quantizer.get_scale_dtype(), act_params=act_params, amax_scope=amax_scope, @@ -1415,7 +1417,7 @@ def quantize_dact_dbias( act_type_id = ActivationEnum[activation_type] PrimitiveClass = DActLuDBiasQuantizePrimitive if is_dbias else DActLuQuantizePrimitive if not PrimitiveClass.enabled() or ( - quantizer is not None and quantizer.q_layout == QuantizeLayout.COLWISE + quantizer is not None and quantizer.q_layout.is_colwise_only ): return _jax_quantize_dact_dbias(dz, x, activation_type, is_dbias, quantizer, act_params) if quantizer is None: @@ -1428,7 +1430,7 @@ def quantize_dact_dbias( out_dtype=(jnp.float32 if is_dbias else x.dtype), # default value for no scaling, TE/common ignore this value when scale is unset scaling_mode=ScalingMode.NO_SCALING.value, - is_2x=False, # unused + quantize_layout=QuantizeLayout.ROWWISE, # unused scale_dtype=jnp.float32, # unused is_dbias=False, act_enum=act_type_id, @@ -1555,7 +1557,7 @@ def quantize_dact_dbias( amax, out_dtype=quantizer.q_dtype, scaling_mode=quantizer.scaling_mode.value, - is_2x=quantizer.is_2x2x(), + quantize_layout=quantizer.q_layout, scale_dtype=quantizer.get_scale_dtype(), is_dbias=is_dbias, act_enum=act_type_id, @@ -1568,7 +1570,7 @@ def quantize_dact_dbias( ) # For DelayedScaling transpose, the scale buffer is shared for both rowwise and colwise - if quantizer.scaling_mode.is_tensor_scaling() and quantizer.is_2x2x(): + if quantizer.scaling_mode.is_tensor_scaling() and quantizer.q_layout.is_rowwise_colwise: colwise_scale_inv = rowwise_scale_inv quantizer.update(updated_amax) diff --git a/transformer_engine/jax/cpp_extensions/misc.py b/transformer_engine/jax/cpp_extensions/misc.py index 572d82f18d..f15fe72bad 100644 --- a/transformer_engine/jax/cpp_extensions/misc.py +++ b/transformer_engine/jax/cpp_extensions/misc.py @@ -207,7 +207,9 @@ def should_apply_1x_fused_dbias_war_for_arch_l_100(is_dbias: bool = False, quant break # _quantize_dbias_impl forcing 1x quantization for tensor scaling switches q_layout to ROWWISE, # but this fails when bias fusion is turned on with arch < 100. - force_1x_quantization = quantizer.scaling_mode.is_tensor_scaling() and quantizer.is_2x2x() + force_1x_quantization = ( + quantizer.scaling_mode.is_tensor_scaling() and quantizer.q_layout.is_rowwise_colwise + ) return ( (force_1x_quantization or quantizer.q_layout == QuantizeLayout.ROWWISE) and arch_l_100 @@ -229,7 +231,9 @@ def try_apply_delayed_scaling_2x_war(f, *args, quantizer=None, flatten_axis=-1, @return: the output of 'f' with the colwise output calculated """ should_apply_war = ( - quantizer is not None and quantizer.scaling_mode.is_tensor_scaling() and quantizer.is_2x2x() + quantizer is not None + and quantizer.scaling_mode.is_tensor_scaling() + and quantizer.q_layout.is_rowwise_colwise ) if not should_apply_war: return None diff --git a/transformer_engine/jax/cpp_extensions/normalization.py b/transformer_engine/jax/cpp_extensions/normalization.py index d09ce7ef74..92efb91a76 100644 --- a/transformer_engine/jax/cpp_extensions/normalization.py +++ b/transformer_engine/jax/cpp_extensions/normalization.py @@ -11,7 +11,7 @@ import jax import jax.numpy as jnp from jax import dtypes, ffi -from jax.experimental.custom_partitioning import SdyShardingRule +from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING from jax.interpreters.mlir import ir from jax.sharding import PartitionSpec @@ -112,7 +112,7 @@ def abstract( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -148,6 +148,13 @@ def abstract( "Current tensor scaling is not supported for fused norm and quantization. Please do" " norm in higher-precision then quantize with current tensor scaling." ) + assert not ScalingMode(scaling_mode).is_nvfp4_scaling, ( + "NVFP4 block scaling is not yet supported for fused norm and quantization." + " Please do norm in higher-precision then quantize with current tensor scaling." + ) + assert ( + not quantize_layout.is_colwise_only + ), "Fused norm with colwise-only quantization is not supported." mu_rsigama_dtype = jnp.float32 @@ -165,7 +172,7 @@ def abstract( updated_amax_aval = jax.core.ShapedArray(shape=(1,), dtype=jnp.float32) - colwise_out_shape = x_aval.shape if is_2x else (1,) + colwise_out_shape = x_aval.shape if quantize_layout.has_colwise else (1,) colwise_out_aval = jax.core.ShapedArray(shape=colwise_out_shape, dtype=out_dtype) rowwise_scale_inv_shape, colwise_scale_inv_shape = ScalingMode( @@ -173,7 +180,7 @@ def abstract( ).get_scale_shape_2x(x_aval.shape, is_padded=not is_outer) scale_inv_aval = jax.core.ShapedArray(shape=rowwise_scale_inv_shape, dtype=scale_dtype) - colwise_scale_inv_shape = colwise_scale_inv_shape if is_2x else (1,) + colwise_scale_inv_shape = colwise_scale_inv_shape if quantize_layout.has_colwise else (1,) colwise_scale_inv_aval = jax.core.ShapedArray( shape=colwise_scale_inv_shape, dtype=scale_dtype ) @@ -189,7 +196,7 @@ def abstract( zero_centered_gamma, epsilon, get_forward_sm_margin(), - is_2x, + True, # is_training ) wkspace_aval = jax.core.ShapedArray( shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) @@ -245,7 +252,7 @@ def lowering( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -287,7 +294,7 @@ def lowering( epsilon=epsilon, sm_margin=sm_margin, scaling_mode=scaling_mode.value, - is_2x=is_2x, + quantize_layout=quantize_layout.value.value, output_amax_when_no_scaling=output_amax_when_no_scaling, ) @@ -303,7 +310,7 @@ def impl( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -335,7 +342,7 @@ def impl( epsilon=epsilon, out_dtype=out_dtype, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, @@ -349,7 +356,7 @@ def impl( scale_inv = scale_inv.flatten()[: reduce(operator.mul, rowwise_scale_inv_shape, 1)].reshape( rowwise_scale_inv_shape ) - if is_2x: + if quantize_layout.has_colwise: colwise_scale_inv = colwise_scale_inv.flatten()[ : reduce(operator.mul, colwise_scale_inv_shape, 1) ].reshape(colwise_scale_inv_shape) @@ -373,7 +380,7 @@ def batcher( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -409,7 +416,7 @@ def batcher( epsilon=epsilon, out_dtype=out_dtype, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, @@ -426,7 +433,7 @@ def infer_sharding_from_operands( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -450,7 +457,7 @@ def infer_sharding_from_operands( ) out_sharding = NamedSharding(mesh, PartitionSpec(*out_spec), desc="NormFwdPrimitive.out") - colwise_out_spec = out_spec if is_2x else (None,) + colwise_out_spec = out_spec if quantize_layout.has_colwise else (None,) colwise_out_sharding = NamedSharding( mesh, PartitionSpec(*colwise_out_spec), desc="NormFwdPrimitive.colwise_out" ) @@ -488,7 +495,7 @@ def partition( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -524,7 +531,7 @@ def partition( ) out_sharding = NamedSharding(mesh, PartitionSpec(*out_spec), desc="NormFwdPrimitive.out") - colwise_out_spec = out_spec if is_2x else (None,) + colwise_out_spec = out_spec if quantize_layout.has_colwise else (None,) colwise_out_sharding = NamedSharding( mesh, PartitionSpec(*colwise_out_spec), desc="NormFwdPrimitive.colwise_out" ) @@ -586,7 +593,7 @@ def sharded_impl(x, scale, amax, gamma, beta): epsilon=epsilon, out_dtype=out_dtype, scaling_mode=scaling_mode, - is_2x=is_2x, + quantize_layout=quantize_layout, scale_dtype=scale_dtype, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, @@ -623,7 +630,7 @@ def shardy_sharding_rule( epsilon, out_dtype, scaling_mode, - is_2x, + quantize_layout, scale_dtype, amax_scope, transpose_batch_sequence, @@ -646,25 +653,29 @@ def shardy_sharding_rule( result_types, ) - prefix = "NormFwd_" + prefix = "NormFwd" scale_rules = ScalingMode(scaling_mode).get_shardy_sharding_rules( - value_types[0].shape, unique_var=prefix + "x", flatten_axis=-1 + value_types[0].shape, + unique_var=prefix, + flatten_axis=-1, + q_layout=quantize_layout, ) - x_axes = scale_rules.input_spec + input_spec = scale_rules.input_spec - out = x_axes - colwise_out = out if is_2x else (prefix + "out_colwise",) - rsigma = x_axes[:-1] - mu = (prefix + "mu",) if norm_type == NVTE_Norm_Type.RMSNorm else rsigma - amax = (prefix + "amax",) + rsigma = input_spec[:-1] + mu = (BATCHING + prefix + "_mu",) if norm_type == NVTE_Norm_Type.RMSNorm else rsigma + amax = (BATCHING + prefix + "_amax",) + scale = (BATCHING + prefix + "_scale",) + gamma = (BATCHING + prefix + "_gamma",) + beta = (BATCHING + prefix + "_beta",) return SdyShardingRule( - (x_axes, ("…1",), amax, ("…2",), ("…3",)), + (input_spec, scale, amax, gamma, beta), ( - out, - colwise_out, - scale_rules.rowwise_rule, - scale_rules.colwise_rule, + scale_rules.rowwise_out_spec, + scale_rules.colwise_out_spec, + scale_rules.rowwise_scale_spec, + scale_rules.colwise_scale_spec, amax, mu, rsigma, @@ -987,7 +998,7 @@ def layernorm_fwd( return (output, mu, rsigma) # TE/common does not support normalization with colwise only quantization yet - if quantizer is not None and quantizer.q_layout == QuantizeLayout.COLWISE: + if quantizer is not None and quantizer.q_layout.is_colwise_only: return _jax_layernorm(x, gamma, beta, zero_centered_gamma, epsilon, quantizer) scale = ( @@ -1008,7 +1019,7 @@ def layernorm_fwd( epsilon=epsilon, out_dtype=x.dtype, scaling_mode=ScalingMode.NO_SCALING.value, - is_2x=False, + quantize_layout=QuantizeLayout.ROWWISE, scale_dtype=jnp.float32, amax_scope=amax_scope, transpose_batch_sequence=False, @@ -1067,10 +1078,11 @@ def layernorm_fwd( ) return out, mu, rsigma - is_2x2x = quantizer.is_2x2x() - # TE/common normalization doesn't support 2x delayed scaling - if quantizer.is_2x2x() and quantizer.scaling_mode.is_tensor_scaling(): - is_2x2x = False + # TE/common Norm doesn't support 2x delayed scaling so do 1x then JAX transpose + q_layout = quantizer.q_layout + if quantizer.q_layout.is_rowwise_colwise and quantizer.scaling_mode.is_tensor_scaling(): + q_layout = QuantizeLayout.ROWWISE + ( rowwise_casted_output, colwise_casted_output, @@ -1090,7 +1102,7 @@ def layernorm_fwd( epsilon=epsilon, out_dtype=quantizer.q_dtype, scaling_mode=quantizer.scaling_mode.value, - is_2x=is_2x2x, + quantize_layout=q_layout, scale_dtype=quantizer.get_scale_dtype(), amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, @@ -1099,8 +1111,7 @@ def layernorm_fwd( ) quantizer.update(updated_amax) - # TE/common Norm doesn't support 2x delayed scaling so do 1x then JAX transpose - if quantizer.is_2x2x() and quantizer.scaling_mode.is_tensor_scaling(): + if quantizer.q_layout.is_rowwise_colwise and quantizer.scaling_mode.is_tensor_scaling(): colwise_casted_output = jnp.transpose( rowwise_casted_output, (-1, *range(rowwise_casted_output.ndim - 1)) ) @@ -1238,7 +1249,7 @@ def rmsnorm_fwd( return (output, rsigma) # TE/common does not support normalization with colwise only quantization yet - if quantizer is not None and quantizer.q_layout == QuantizeLayout.COLWISE: + if quantizer is not None and quantizer.q_layout.is_colwise_only: return _jax_rmsnorm(x, gamma, zero_centered_gamma, epsilon, quantizer) scale = ( @@ -1261,7 +1272,7 @@ def rmsnorm_fwd( epsilon=epsilon, out_dtype=x.dtype, scaling_mode=ScalingMode.NO_SCALING.value, - is_2x=False, + quantize_layout=QuantizeLayout.ROWWISE, scale_dtype=jnp.float32, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, @@ -1321,10 +1332,11 @@ def rmsnorm_fwd( ) return out, rsigma - is_2x2x = quantizer.is_2x2x() - # TE/common normalization doesn't support 2x delayed scaling - if quantizer.is_2x2x() and quantizer.scaling_mode.is_tensor_scaling(): - is_2x2x = False + # TE/common Norm doesn't support 2x delayed scaling so do 1x then JAX transpose + q_layout = quantizer.q_layout + if quantizer.q_layout.is_rowwise_colwise and quantizer.scaling_mode.is_tensor_scaling(): + q_layout = QuantizeLayout.ROWWISE + ( rowwise_casted_output, colwise_casted_output, @@ -1344,7 +1356,7 @@ def rmsnorm_fwd( epsilon=epsilon, out_dtype=quantizer.q_dtype, scaling_mode=quantizer.scaling_mode.value, - is_2x=is_2x2x, + quantize_layout=q_layout, scale_dtype=quantizer.get_scale_dtype(), amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, @@ -1353,8 +1365,7 @@ def rmsnorm_fwd( ) quantizer.update(updated_amax) - # TE/common Norm doesn't support 2x delayed scaling so do 1x then JAX transpose - if quantizer.is_2x2x() and quantizer.scaling_mode.is_tensor_scaling(): + if quantizer.q_layout.is_rowwise_colwise and quantizer.scaling_mode.is_tensor_scaling(): colwise_casted_output = jnp.transpose( rowwise_casted_output, (-1, *range(rowwise_casted_output.ndim - 1)) ) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 67c505bc98..a0e1a6406f 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -11,7 +11,7 @@ import jax import jax.numpy as jnp from jax import dtypes, ffi -from jax.experimental.custom_partitioning import SdyShardingRule +from jax.experimental.custom_partitioning import SdyShardingRule, BATCHING from jax.sharding import PartitionSpec import transformer_engine_jax @@ -122,7 +122,7 @@ def abstract( f" stochastic_rounding is True but received {sr_rng_state_aval.shape}" ) - if q_layout in (QuantizeLayout.ROWWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if QuantizeLayout(q_layout).has_rowwise: rowwise_out_shape = out_shape else: rowwise_out_shape = (1,) @@ -170,7 +170,7 @@ def abstract( broadcast_2d_scale_shape_to_1d=True, ) - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if QuantizeLayout(q_layout).has_colwise: if ScalingMode(scaling_mode).is_colwise_transposed: colwise_out_shape = multidim_transpose(out_shape, transpose_axis=flatten_axis) else: @@ -194,9 +194,7 @@ def abstract( jax_dtype_to_te_dtype(out_dtype), jax_dtype_to_te_dtype(scale_dtype), scaling_mode, - QuantizeLayout( - q_layout - ), # For now until we have auto-decoding for QuantizeLayout enum + q_layout.value, ) wkspace_shape = wkspace_info[0] wkspace_dtype = te_dtype_to_jax_dtype(wkspace_info[1]) @@ -272,7 +270,7 @@ def lowering( post_rht_amax, rht_matrix, scaling_mode=scaling_mode.value, - q_layout=q_layout, + q_layout=q_layout.value.value, flatten_axis=flatten_axis, is_dbias=is_dbias, stochastic_rounding=stochastic_rounding, @@ -335,7 +333,7 @@ def impl( scale_inv = jax.lax.slice( scale_inv, [0] * len(rowwise_scale_inv_shape), rowwise_scale_inv_shape ) - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_colwise: colwise_scale_inv = jax.lax.slice( colwise_scale_inv, [0] * len(colwise_scale_inv_shape), colwise_scale_inv_shape ) @@ -424,7 +422,7 @@ def infer_sharding_from_operands( PartitionSpec(*x_spec), desc="BaseDBiasQuantizePrimitive.out_sharding", ) - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_colwise: if ScalingMode(scaling_mode).is_colwise_transposed: colwise_out_spec = multidim_transpose(x_spec, transpose_axis=flatten_axis) else: @@ -448,7 +446,7 @@ def infer_sharding_from_operands( if ScalingMode(scaling_mode).is_block_scaling: scale_inv_spec = x_spec - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_colwise: if ( ScalingMode(scaling_mode).is_block_scaling and ScalingMode(scaling_mode).is_colwise_transposed @@ -505,7 +503,7 @@ def partition( desc="BaseDBiasQuantizePrimitive.out_sharding", ) - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_colwise: if ScalingMode(scaling_mode).is_colwise_transposed: colwise_out_spec = multidim_transpose(x_spec, transpose_axis=flatten_axis) else: @@ -529,7 +527,7 @@ def partition( if ScalingMode(scaling_mode).is_block_scaling: scale_inv_spec = x_spec - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_colwise: if ( ScalingMode(scaling_mode).is_block_scaling and ScalingMode(scaling_mode).is_colwise_transposed @@ -643,39 +641,37 @@ def shardy_sharding_rule( result_types, ) - prefix = "DBiasQuantize_" + prefix = "DBiasQuantize" scale_rules = ScalingMode(scaling_mode).get_shardy_sharding_rules( value_types[0].shape, - unique_var=prefix + "x", + unique_var=prefix, flatten_axis=flatten_axis, + q_layout=q_layout, broadcast_2d_scale_shape_to_1d=True, ) - x_axes = scale_rules.input_spec - - out = x_axes - colwise_out = (prefix + "out_colwise",) - colwise_scale_inv = (prefix + "colwise_scale_inv",) - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): - colwise_scale_inv = scale_rules.colwise_rule - if ScalingMode(scaling_mode).is_colwise_transposed: - colwise_out = tuple(multidim_transpose(x_axes, transpose_axis=flatten_axis)) - colwise_scale_inv = tuple( - multidim_transpose(colwise_scale_inv, transpose_axis=flatten_axis) - ) - else: - colwise_out = x_axes - - dbias = x_axes[flatten_axis:] if is_dbias else (prefix + "dbias",) - amax = (prefix + "amax",) - sr_rng_state = (prefix + "sr_rng_state_partition_axis", prefix + "sr_rng_state_data_axis") + input_spec = scale_rules.input_spec + dbias = input_spec[flatten_axis:] if is_dbias else (prefix + "_dbias",) + amax = (BATCHING + prefix + "_amax",) + scale = (BATCHING + prefix + "_scale",) + sr_rng_state = ( + BATCHING + prefix + "_sr_rng_state_partition_axis", + BATCHING + prefix + "sr_rng_state_data_axis", + ) - post_rht_amax = (prefix + "post_rht_amax",) - rht_matrix = (prefix + "rht_matrix_1", prefix + "rht_matrix_2") + post_rht_amax = (BATCHING + prefix + "_post_rht_amax",) + rht_matrix = (BATCHING + prefix + "_rht_matrix_1", BATCHING + prefix + "_rht_matrix_2") return SdyShardingRule( - (x_axes, ("…1",), amax, sr_rng_state, post_rht_amax, rht_matrix), - (out, colwise_out, scale_rules.rowwise_rule, colwise_scale_inv, amax, dbias), + (input_spec, scale, amax, sr_rng_state, post_rht_amax, rht_matrix), + ( + scale_rules.rowwise_out_spec, + scale_rules.colwise_out_spec, + scale_rules.rowwise_scale_spec, + scale_rules.colwise_scale_spec, + amax, + dbias, + ), **scale_rules.factor_sizes, ) @@ -762,7 +758,7 @@ def _quantize_dbias_impl( # If TE/common custom quantize op is disabled, or if quantizer layout is COLWISE, # fall back on the native-JAX quantize implementation PrimitiveClass = DBiasQuantizePrimitive if is_dbias else QuantizePrimitive - is_unsupported = quantizer.q_layout == QuantizeLayout.COLWISE and not ( + is_unsupported = quantizer.q_layout.is_colwise_only and not ( quantizer.scaling_mode == ScalingMode.NVFP4_1D_SCALING and hasattr(quantizer, "use_rht") and quantizer.use_rht @@ -845,7 +841,7 @@ def _quantize_dbias_impl( is_1x_kernel_supported = not (is_dbias and get_min_device_compute_capability() < 100) force_1x_quantization = ( quantizer.scaling_mode.is_tensor_scaling() - and quantizer.is_2x2x() + and quantizer.q_layout.is_rowwise_colwise and is_1x_kernel_supported ) q_layout = quantizer.q_layout @@ -879,7 +875,7 @@ def _quantize_dbias_impl( rht_matrix, out_dtype=quantizer.q_dtype, scaling_mode=quantizer.scaling_mode.value, - q_layout=q_layout.value, + q_layout=q_layout, flatten_axis=flatten_axis, scale_dtype=quantizer.get_scale_dtype(), is_dbias=is_dbias if not quantizer.scaling_mode.is_nvfp4_scaling else False, @@ -888,10 +884,10 @@ def _quantize_dbias_impl( use_rht=use_rht, ) # For DelayedScaling2x, the scale buffer is shared between rowwise and colwise - if quantizer.scaling_mode.is_tensor_scaling() and quantizer.is_2x2x(): + if quantizer.scaling_mode.is_tensor_scaling() and quantizer.q_layout.is_rowwise_colwise: colwise_scale_inv = rowwise_scale_inv - if q_layout == QuantizeLayout.ROWWISE: + if q_layout.is_rowwise_only: # Quantizer requires 2x quantization, but we are using 1x quantization # for performance reasons, so we need to generate the colwise data in JAX if flatten_axis < 0: @@ -1043,7 +1039,7 @@ def abstract( flatten_axis=flatten_axis, ) - if q_layout in (QuantizeLayout.ROWWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_rowwise: rowwise_out_shape = out_shape else: rowwise_out_shape = (1,) @@ -1052,7 +1048,7 @@ def abstract( amax_aval = jax.core.ShapedArray(shape=(group_sizes_aval.size,), dtype=jnp.float32) - if q_layout in (QuantizeLayout.COLWISE.value, QuantizeLayout.ROWWISE_COLWISE.value): + if q_layout.has_colwise: colwise_out_shape = out_shape else: colwise_out_shape = (1,) @@ -1117,7 +1113,7 @@ def lowering( scale, group_sizes, scaling_mode=scaling_mode.value, - q_layout=q_layout, + q_layout=q_layout.value.value, flatten_axis=flatten_axis, ) @@ -1240,7 +1236,7 @@ def grouped_quantize( ) # WAR for tensor_scaling as TE/Common does not support q_layout = COLWISE yet # So we performance ROWWISE_COLWISE and use the colwise_tensor_output - apply_colwise_war = is_tensor_scaling and quantizer.q_layout == QuantizeLayout.COLWISE + apply_colwise_war = is_tensor_scaling and quantizer.q_layout.is_colwise_only q_layout = QuantizeLayout.ROWWISE_COLWISE if apply_colwise_war else quantizer.q_layout ( rowwise_casted_output, @@ -1254,7 +1250,7 @@ def grouped_quantize( group_sizes, out_dtype=quantizer.q_dtype, scaling_mode=quantizer.scaling_mode.value, - q_layout=q_layout.value, + q_layout=q_layout, flatten_axis=flatten_axis, group_axis=group_axis, scale_dtype=quantizer.get_scale_dtype(), @@ -1262,7 +1258,7 @@ def grouped_quantize( # For DelayedScaling2x and CurrentScaling2x, the scale buffer # is shared between rowwise and colwise - if is_tensor_scaling and quantizer.is_2x2x() or apply_colwise_war: + if is_tensor_scaling and quantizer.q_layout.is_rowwise_colwise or apply_colwise_war: colwise_scale_inv = rowwise_scale_inv # TODO(Phuong): store the whole updated_amax in the grouped_quantize instead? diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 87c6fa91cd..c1c7e0d665 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -57,7 +57,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(DActLuDBiasQuantizeInitializeHandler); pybind11::tuple GetDActDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype, DType out_dtype, - JAXX_Scaling_Mode scaling_mode, bool is_2x); + JAXX_Scaling_Mode scaling_mode, + JAXX_Quantize_Layout quantize_layout); // Normalization XLA_FFI_DECLARE_HANDLER_SYMBOL(NormForwardInitializeHandler); @@ -87,7 +88,7 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(DequantizeHandler); pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype, DType out_dtype, DType scale_dtype, JAXX_Scaling_Mode scaling_mode, - QuantizeLayout q_layout); + JAXX_Quantize_Layout quantize_layout); // Softmax XLA_FFI_DECLARE_HANDLER_SYMBOL(ScaledSoftmaxForwardHandler); @@ -162,5 +163,6 @@ XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( // ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Collective_Op); +XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Quantize_Layout); #endif // TRANSFORMER_ENGINE_JAX_CSRC_FP8_MODULES_H_ diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index f512321c38..34ce29ae13 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -18,7 +18,8 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal Buffer_Type amax_buf, Result_Type output_buf, Result_Type colwise_output_buf, Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf, Result_Type updated_amax_buf, int64_t act_enum, JAXX_Scaling_Mode scaling_mode, - bool is_2x_int, ActivationConfig act_params, bool output_amax_when_no_scaling) { + JAXX_Quantize_Layout quantize_layout, ActivationConfig act_params, + bool output_amax_when_no_scaling) { // parameters for clamped swiglu used in GPT OSS auto swiglu_limit = act_params.clamped_swiglu.limit; auto swiglu_alpha = act_params.clamped_swiglu.alpha; @@ -40,7 +41,6 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal auto n = input_dims.back(); auto act_type = static_cast(act_enum); auto act_len = input_dims[input_dims.size() - 2]; - auto is_2x = static_cast(is_2x_int); auto flatten_axis = output_buf->dimensions().size() - 1; // output does not have act axis auto input_shape = std::vector{m, static_cast(act_len * n)}; @@ -77,7 +77,7 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal } } - if (is_2x) { + if (is_quantize_2x2x(quantize_layout)) { auto &tmp_shape = (scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING) ? output_trans_shape : output_shape; @@ -158,7 +158,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ActLuHandler, ActLuFFI, .Ret() // updated_amax .Attr("act_enum") .Attr("scaling_mode") - .Attr("is_2x") + .Attr("quantize_layout") .Attr("act_params") .Attr("output_amax_when_no_scaling"), FFI_CudaGraph_Traits); @@ -167,11 +167,12 @@ Error_Type ActLuInitializeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer Buffer_Type amax_buf, Result_Type output_buf, Result_Type colwise_output_buf, Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf, Result_Type updated_amax_buf, - int64_t act_enum, JAXX_Scaling_Mode scaling_mode, bool is_2x_int, - ActivationConfig act_params, bool output_amax_when_no_scaling) { + int64_t act_enum, JAXX_Scaling_Mode scaling_mode, + JAXX_Quantize_Layout quantize_layout, ActivationConfig act_params, + bool output_amax_when_no_scaling) { return wrapInStreamCapture(std::function(ActLuFFI), stream, input_buf, scale_buf, amax_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, - updated_amax_buf, act_enum, scaling_mode, is_2x_int, act_params, + updated_amax_buf, act_enum, scaling_mode, quantize_layout, act_params, output_amax_when_no_scaling); } @@ -188,13 +189,14 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(ActLuInitializeHandler, ActLuInitializeFFI, .Ret() // updated_amax .Attr("act_enum") .Attr("scaling_mode") - .Attr("is_2x") + .Attr("quantize_layout") .Attr("act_params") .Attr("output_amax_when_no_scaling")); pybind11::tuple GetDActDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype, DType out_dtype, - JAXX_Scaling_Mode scaling_mode, bool is_2x) { + JAXX_Scaling_Mode scaling_mode, + JAXX_Quantize_Layout quantize_layout) { auto input_shape = std::vector{batch_size, hidden_size}; auto dact_input_shape = std::vector{batch_size, hidden_size}; auto output_shape = std::vector{batch_size, hidden_size}; @@ -226,7 +228,7 @@ pybind11::tuple GetDActDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hid std::vector{1}); } - if (is_2x) { + if (is_quantize_2x2x(quantize_layout)) { auto &tmp_shape = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING ? output_trans_shape : output_shape; output_tensor.set_columnwise_data(reinterpret_cast(&temp), out_dtype, tmp_shape); @@ -260,9 +262,9 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, Result_Type colwise_output_buf, Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf, Result_Type updated_amax_buf, Result_Type dbias_buf, Result_Type workspace_buf, - JAXX_Scaling_Mode scaling_mode, int64_t act_enum, bool is_2x, - bool is_dbias, ActivationConfig act_params, - bool output_amax_when_no_scaling) { + JAXX_Scaling_Mode scaling_mode, int64_t act_enum, + JAXX_Quantize_Layout quantize_layout, bool is_dbias, + ActivationConfig act_params, bool output_amax_when_no_scaling) { // parameters for clamped swiglu used in GPT OSS auto swiglu_limit = act_params.clamped_swiglu.limit; auto swiglu_alpha = act_params.clamped_swiglu.alpha; @@ -340,7 +342,7 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, } } - if (is_2x) { + if (is_quantize_2x2x(quantize_layout)) { auto &tmp_shape = (scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING) ? output_trans_shape : output_shape; @@ -370,7 +372,8 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, // fused_dgated_dbias is not available, so we use dact_lu + quantize_dbias in Python instead NVTE_CHECK(!(act_len == 2 && is_dbias), "Unsupported DGatedActedDBias Fusion!"); - NVTE_CHECK(!(scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING && is_2x && act_len == 2), + NVTE_CHECK(!(scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING && + is_quantize_2x2x(quantize_layout) && act_len == 2), "TE/common does not support delayed scaling for 2x with gated activations."); if (is_dbias) { @@ -465,7 +468,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(DActLuDBiasQuantizeHandler, DActLuDBiasQuantizeFFI .Ret() // wkspace .Attr("scaling_mode") .Attr("act_enum") - .Attr("is_2x") + .Attr("quantize_layout") .Attr("is_dbias") .Attr("act_params") .Attr("output_amax_when_no_scaling"), @@ -476,13 +479,13 @@ Error_Type DActLuDBiasQuantizeInitializeFFI( Buffer_Type amax_buf, Result_Type output_buf, Result_Type colwise_output_buf, Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf, Result_Type updated_amax_buf, Result_Type dbias_buf, Result_Type workspace_buf, JAXX_Scaling_Mode scaling_mode, - int64_t act_enum, bool is_2x, bool is_dbias, ActivationConfig act_params, - bool output_amax_when_no_scaling) { + int64_t act_enum, JAXX_Quantize_Layout quantize_layout, bool is_dbias, + ActivationConfig act_params, bool output_amax_when_no_scaling) { return wrapInStreamCapture(std::function(DActLuDBiasQuantizeFFI), stream, input_buf, act_input_buf, scale_buf, amax_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, updated_amax_buf, dbias_buf, - workspace_buf, scaling_mode, act_enum, is_2x, is_dbias, act_params, - output_amax_when_no_scaling); + workspace_buf, scaling_mode, act_enum, quantize_layout, is_dbias, + act_params, output_amax_when_no_scaling); } XLA_FFI_DEFINE_HANDLER_SYMBOL(DActLuDBiasQuantizeInitializeHandler, @@ -502,7 +505,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(DActLuDBiasQuantizeInitializeHandler, .Ret() // wkspace .Attr("scaling_mode") .Attr("act_enum") - .Attr("is_2x") + .Attr("quantize_layout") .Attr("is_dbias") .Attr("act_params") .Attr("output_amax_when_no_scaling")); diff --git a/transformer_engine/jax/csrc/extensions/misc.h b/transformer_engine/jax/csrc/extensions/misc.h index 07e9aec7e9..21b50c1af4 100644 --- a/transformer_engine/jax/csrc/extensions/misc.h +++ b/transformer_engine/jax/csrc/extensions/misc.h @@ -34,12 +34,24 @@ inline size_t product(const std::vector &shape) { return ret; } -enum class QuantizeLayout { +enum class JAXX_Quantize_Layout : int64_t { ROWWISE, COLWISE, ROWWISE_COLWISE, }; +inline bool is_quantize_rowwise(const JAXX_Quantize_Layout &layout) { + return layout == JAXX_Quantize_Layout::ROWWISE || layout == JAXX_Quantize_Layout::ROWWISE_COLWISE; +} + +inline bool is_quantize_colwise(const JAXX_Quantize_Layout &layout) { + return layout == JAXX_Quantize_Layout::COLWISE || layout == JAXX_Quantize_Layout::ROWWISE_COLWISE; +} + +inline bool is_quantize_2x2x(const JAXX_Quantize_Layout &layout) { + return layout == JAXX_Quantize_Layout::ROWWISE_COLWISE; +} + enum class JAXX_Scaling_Mode : int64_t { NO_SCALING = 0, DELAYED_TENSOR_SCALING = 1, diff --git a/transformer_engine/jax/csrc/extensions/normalization.cpp b/transformer_engine/jax/csrc/extensions/normalization.cpp index 378e009c83..b01e23c128 100644 --- a/transformer_engine/jax/csrc/extensions/normalization.cpp +++ b/transformer_engine/jax/csrc/extensions/normalization.cpp @@ -66,7 +66,7 @@ Error_Type NormForwardFFI(cudaStream_t stream, Buffer_Type x_buf, Buffer_Type sc Result_Type updated_amax_buf, Result_Type mu_buf, Result_Type rsigma_buf, Result_Type wkspace_buf, int norm_type, bool zero_centered_gamma, double epsilon, int64_t sm_margin, JAXX_Scaling_Mode scaling_mode, - bool is_2x, bool output_amax_when_no_scaling) { + JAXX_Quantize_Layout quantize_layout, bool output_amax_when_no_scaling) { auto in_dtype = convert_ffi_datatype_to_te_dtype(x_buf.element_type()); auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type()); auto w_dtype = convert_ffi_datatype_to_te_dtype(gamma_buf.element_type()); @@ -86,7 +86,6 @@ Error_Type NormForwardFFI(cudaStream_t stream, Buffer_Type x_buf, Buffer_Type sc NVTE_CHECK(amax == updated_amax && amax != nullptr, "amax and updated_amax should be aliased"); auto _norm_type = static_cast(norm_type); - auto _is_2x = static_cast(is_2x); auto x_size = product(x_buf.dimensions()); auto gamma_size = product(gamma_buf.dimensions()); @@ -134,7 +133,7 @@ Error_Type NormForwardFFI(cudaStream_t stream, Buffer_Type x_buf, Buffer_Type sc output_tensor.set_scale(scale, DType::kFloat32, std::vector{1}); } - if (_is_2x) { + if (is_quantize_2x2x(quantize_layout)) { output_tensor.set_columnwise_data(colwise_output_buf->untyped_data(), static_cast(out_dtype), input_shape); output_tensor.set_columnwise_scale_inv( @@ -185,25 +184,23 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(NormForwardHandler, NormForwardFFI, .Attr("epsilon") .Attr("sm_margin") .Attr("scaling_mode") - .Attr("is_2x") + .Attr("quantize_layout") .Attr("output_amax_when_no_scaling"), FFI_CudaGraph_Traits); -Error_Type NormForwardInitializeFFI(cudaStream_t stream, Buffer_Type x_buf, Buffer_Type scale_buf, - Buffer_Type amax_buf, Buffer_Type gamma_buf, - Buffer_Type beta_buf, Result_Type output_buf, - Result_Type colwise_output_buf, Result_Type scale_inv_buf, - Result_Type colwise_scale_inv_buf, Result_Type updated_amax_buf, - Result_Type mu_buf, Result_Type rsigma_buf, - Result_Type wkspace_buf, int norm_type, - bool zero_centered_gamma, double epsilon, int64_t sm_margin, - JAXX_Scaling_Mode scaling_mode, bool is_2x, - bool output_amax_when_no_scaling) { +Error_Type NormForwardInitializeFFI( + cudaStream_t stream, Buffer_Type x_buf, Buffer_Type scale_buf, Buffer_Type amax_buf, + Buffer_Type gamma_buf, Buffer_Type beta_buf, Result_Type output_buf, + Result_Type colwise_output_buf, Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf, + Result_Type updated_amax_buf, Result_Type mu_buf, Result_Type rsigma_buf, + Result_Type wkspace_buf, int norm_type, bool zero_centered_gamma, double epsilon, + int64_t sm_margin, JAXX_Scaling_Mode scaling_mode, JAXX_Quantize_Layout quantize_layout, + bool output_amax_when_no_scaling) { return wrapInStreamCapture(std::function(NormForwardFFI), stream, x_buf, scale_buf, amax_buf, gamma_buf, beta_buf, output_buf, colwise_output_buf, scale_inv_buf, colwise_scale_inv_buf, updated_amax_buf, mu_buf, rsigma_buf, wkspace_buf, norm_type, zero_centered_gamma, epsilon, sm_margin, - scaling_mode, is_2x, output_amax_when_no_scaling); + scaling_mode, quantize_layout, output_amax_when_no_scaling); } XLA_FFI_DEFINE_HANDLER_SYMBOL(NormForwardInitializeHandler, NormForwardInitializeFFI, @@ -227,7 +224,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(NormForwardInitializeHandler, NormForwardInitializ .Attr("epsilon") .Attr("sm_margin") .Attr("scaling_mode") - .Attr("is_2x") + .Attr("quantize_layout") .Attr("output_amax_when_no_scaling")); pybind11::tuple GetNormBackwardWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype, diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index d740df0e2a..e57d07872e 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -176,11 +176,10 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVFP4_2D_SCALING", JAXX_Scaling_Mode::NVFP4_2D_SCALING) .export_values(); - pybind11::enum_(m, "QuantizeLayout", - pybind11::module_local()) - .value("ROWWISE", transformer_engine::jax::QuantizeLayout::ROWWISE) - .value("COLWISE", transformer_engine::jax::QuantizeLayout::COLWISE) - .value("ROWWISE_COLWISE", transformer_engine::jax::QuantizeLayout::ROWWISE_COLWISE) + pybind11::enum_(m, "JAXX_Quantize_Layout", pybind11::module_local()) + .value("ROWWISE", JAXX_Quantize_Layout::ROWWISE) + .value("COLWISE", JAXX_Quantize_Layout::COLWISE) + .value("ROWWISE_COLWISE", JAXX_Quantize_Layout::ROWWISE_COLWISE) .export_values(); pybind11::enum_(m, "JAXX_Collective_Op", pybind11::module_local()) diff --git a/transformer_engine/jax/csrc/extensions/quantization.cpp b/transformer_engine/jax/csrc/extensions/quantization.cpp index a45a698822..1f7db84383 100644 --- a/transformer_engine/jax/csrc/extensions/quantization.cpp +++ b/transformer_engine/jax/csrc/extensions/quantization.cpp @@ -20,7 +20,7 @@ namespace jax { pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size, DType in_dtype, DType out_dtype, DType scale_dtype, JAXX_Scaling_Mode scaling_mode, - QuantizeLayout q_layout) { + JAXX_Quantize_Layout q_layout) { auto input_shape = std::vector{batch_size, hidden_size}; auto output_shape = std::vector{batch_size, hidden_size}; auto output_trans_shape = std::vector{hidden_size, batch_size}; @@ -42,7 +42,7 @@ pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_ auto output_tensor = TensorWrapper(get_nvte_scaling_mode(scaling_mode)); auto scale_shape = std::vector{1}; // Only the pointers will be checked for scale_inv, thus the shapes do not matter - if (q_layout == QuantizeLayout::ROWWISE_COLWISE || q_layout == QuantizeLayout::ROWWISE) { + if (is_quantize_rowwise(q_layout)) { output_tensor.set_rowwise_data(reinterpret_cast(&temp), out_dtype, output_shape); if (scaling_mode != JAXX_Scaling_Mode::NO_SCALING) { if (is_nvfp4) @@ -52,7 +52,7 @@ pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_ } } - if (q_layout == QuantizeLayout::ROWWISE_COLWISE || q_layout == QuantizeLayout::COLWISE) { + if (is_quantize_colwise(q_layout)) { auto &tmp_shape = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING ? output_trans_shape : output_shape; output_tensor.set_columnwise_data(reinterpret_cast(&temp), out_dtype, tmp_shape); @@ -90,8 +90,8 @@ Error_Type DBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_T Result_Type scale_inv_buf, Result_Type colwise_scale_inv_buf, Result_Type updated_amax_buf, Result_Type dbias_buf, Result_Type workspace_buf, JAXX_Scaling_Mode scaling_mode, - int64_t quantize_layout_enum, bool is_dbias, int64_t flatten_axis, - bool stochastic_rounding, bool use_rht) { + JAXX_Quantize_Layout quantize_layout, bool is_dbias, + int64_t flatten_axis, bool stochastic_rounding, bool use_rht) { 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()); auto workspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type()); @@ -101,8 +101,6 @@ Error_Type DBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_T auto *input = input_buf.untyped_data(); - auto const quantize_layout = static_cast(quantize_layout_enum); - auto *output = output_buf->untyped_data(); auto *output_trans = output_trans_buf->untyped_data(); auto *dbias = dbias_buf->untyped_data(); @@ -127,15 +125,13 @@ Error_Type DBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_T bool const is_tensor_scaling = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING || scaling_mode == JAXX_Scaling_Mode::CURRENT_TENSOR_SCALING; - bool const is_mxfp8 = scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING; bool const is_nvfp4 = scaling_mode == JAXX_Scaling_Mode::NVFP4_1D_SCALING || scaling_mode == JAXX_Scaling_Mode::NVFP4_2D_SCALING; NVTE_CHECK(!stochastic_rounding || is_nvfp4, "Stochastic rounding is only supported for NVFP4."); NVTE_CHECK(!use_rht || is_nvfp4, "RHT is only supported for NVFP4 scaling"); - if (quantize_layout == QuantizeLayout::ROWWISE || - quantize_layout == QuantizeLayout::ROWWISE_COLWISE) { + if (is_quantize_rowwise(quantize_layout)) { output_tensor.set_rowwise_data(output, out_dtype, output_shape); if (is_tensor_scaling) { @@ -180,10 +176,9 @@ Error_Type DBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_T quant_config.set_rng_state(sr_rng_state_tensor.data()); } - if (quantize_layout == QuantizeLayout::COLWISE || - quantize_layout == QuantizeLayout::ROWWISE_COLWISE) { + if (is_quantize_colwise(quantize_layout)) { if (is_nvfp4 && use_rht) { - if (quantize_layout == QuantizeLayout::ROWWISE_COLWISE) { + if (is_quantize_2x2x(quantize_layout)) { // Do regular rowwise quantization without RHT nvte_quantize_v2(input_tensor.data(), output_tensor.data(), quant_config, stream); } @@ -281,7 +276,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(DBiasQuantizeHandler, DBiasQuantizeFFI, .Ret() // dbias .Ret() // wkspace .Attr("scaling_mode") - .Attr("q_layout") + .Attr("q_layout") .Attr("is_dbias") .Attr("flatten_axis") .Attr("stochastic_rounding") @@ -323,7 +318,7 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty Buffer_Type group_sizes, Result_Type outputs, Result_Type colwise_outputs, Result_Type scale_invs, Result_Type colwise_scale_invs, Result_Type amaxs, - JAXX_Scaling_Mode scaling_mode, int64_t quantize_layout_enum, + JAXX_Scaling_Mode scaling_mode, JAXX_Quantize_Layout quantize_layout, int64_t flatten_axis) { NVTE_CHECK(scaling_mode != JAXX_Scaling_Mode::NO_SCALING, "Unsupported scaling mode: ", static_cast(scaling_mode)); @@ -336,7 +331,6 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty auto group_size_dtype = convert_ffi_datatype_to_te_dtype(group_sizes.element_type()); auto sinv_dtype = convert_ffi_datatype_to_te_dtype(scale_invs->element_type()); auto amax_dtype = convert_ffi_datatype_to_te_dtype(amaxs->element_type()); - auto const quantize_layout = static_cast(quantize_layout_enum); auto *input_ptr = reinterpret_cast(inputs.untyped_data()); auto *scale_ptr = reinterpret_cast(scales.untyped_data()); @@ -346,10 +340,6 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty auto *colwise_sinv_ptr = reinterpret_cast(colwise_scale_invs->untyped_data()); auto *amax_ptr = reinterpret_cast(amaxs->untyped_data()); - bool has_rowwise = quantize_layout == QuantizeLayout::ROWWISE || - quantize_layout == QuantizeLayout::ROWWISE_COLWISE; - bool has_colwise = quantize_layout == QuantizeLayout::COLWISE || - quantize_layout == QuantizeLayout::ROWWISE_COLWISE; bool is_delayed_scaling = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING; bool const is_tensor_scaling = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING || scaling_mode == JAXX_Scaling_Mode::CURRENT_TENSOR_SCALING; @@ -359,8 +349,8 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty size_t output_dtype_bytes = te_dtype_bytes(out_dtype); size_t sinv_dtype_bytes = te_dtype_bytes(sinv_dtype); size_t group_size_dtype_bytes = te_dtype_bytes(group_size_dtype); - size_t colwise_output_dtype_bytes = has_colwise ? output_dtype_bytes : 0; - size_t colwise_sinv_dtype_bytes = has_colwise ? sinv_dtype_bytes : 0; + size_t colwise_output_dtype_bytes = is_quantize_colwise(quantize_layout) ? output_dtype_bytes : 0; + size_t colwise_sinv_dtype_bytes = is_quantize_colwise(quantize_layout) ? sinv_dtype_bytes : 0; size_t scale_dtype_bytes = is_tensor_scaling ? te_dtype_bytes(scale_dtype) : 0; size_t amax_dtype_bytes = is_tensor_scaling ? te_dtype_bytes(amax_dtype) : 0; @@ -423,7 +413,7 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty auto inp_i = TensorWrapper(static_cast(input_ptr), shape_i, in_dtype); auto out_i = TensorWrapper(get_nvte_scaling_mode(scaling_mode)); - if (has_rowwise) { + if (is_quantize_rowwise(quantize_layout)) { out_i.set_rowwise_data(static_cast(output_ptr), out_dtype, shape_i); if (is_fp8_dtype(out_dtype)) { @@ -442,7 +432,7 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty } } - if (has_colwise) { + if (is_quantize_colwise(quantize_layout)) { auto &tmp_shape = is_tensor_scaling ? shape_trans_i : shape_i; out_i.set_columnwise_data(static_cast(colwise_output_ptr), out_dtype, tmp_shape); // For 2x delayed scaling, the scale buffer is shared between rowwise and columnwise scaling @@ -501,7 +491,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedQuantizeHandler, GroupedQuantizeFFI, .Ret() // scale_inv colwise .Ret() // amax .Attr("scaling_mode") - .Attr("q_layout") + .Attr("q_layout") .Attr("flatten_axis")); } // namespace jax diff --git a/transformer_engine/jax/quantize/__init__.py b/transformer_engine/jax/quantize/__init__.py index 9616965c75..878067a783 100644 --- a/transformer_engine/jax/quantize/__init__.py +++ b/transformer_engine/jax/quantize/__init__.py @@ -17,3 +17,4 @@ from .hadamard import * from .helper import * from .device_utils import * +from .misc import * diff --git a/transformer_engine/jax/quantize/misc.py b/transformer_engine/jax/quantize/misc.py new file mode 100644 index 0000000000..c1e169d005 --- /dev/null +++ b/transformer_engine/jax/quantize/misc.py @@ -0,0 +1,61 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +""" +This module provides additional enum and utilities for quantizing tensors in JAX. +""" +from dataclasses import dataclass +from enum import Enum + +from transformer_engine_jax import JAXX_Quantize_Layout + +__all__ = [ + "QuantizeLayout", +] + + +@dataclass(frozen=True) +class QuantizeLayout(Enum): + "Wrapper for JAXX_Quantize_Layout" + + ROWWISE = JAXX_Quantize_Layout.ROWWISE + COLWISE = JAXX_Quantize_Layout.COLWISE + ROWWISE_COLWISE = JAXX_Quantize_Layout.ROWWISE_COLWISE + + @property + def has_rowwise(self) -> bool: + """If the layout has the rowwise component""" + return self.value in (JAXX_Quantize_Layout.ROWWISE, JAXX_Quantize_Layout.ROWWISE_COLWISE) + + @property + def has_colwise(self) -> bool: + """If the layout has the colwise component""" + return self.value in (JAXX_Quantize_Layout.COLWISE, JAXX_Quantize_Layout.ROWWISE_COLWISE) + + @property + def is_rowwise_colwise(self) -> bool: + """If layout is both rowwise and colwise""" + return self.value == JAXX_Quantize_Layout.ROWWISE_COLWISE + + @property + def is_rowwise_only(self) -> bool: + """If layout is rowwise only""" + return self.value == JAXX_Quantize_Layout.ROWWISE + + @property + def is_colwise_only(self) -> bool: + """If layout is colwise only""" + return self.value == JAXX_Quantize_Layout.COLWISE + + def __eq__(self, other): + """Compare this quantize layout with another. + + Args: + other: The other quantize layout to compare with + + Returns: + True if the modes are equal, False otherwise + """ + if not isinstance(other, QuantizeLayout): + return False + return self.value == other.value diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index eb2b7b5924..8a54f0b1db 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -15,10 +15,10 @@ import jax import jax.numpy as jnp from jax.tree_util import register_pytree_node_class -from transformer_engine_jax import QuantizeLayout from transformer_engine.common import recipe from .scaling_modes import ScalingMode +from .misc import QuantizeLayout from .hadamard import apply_rht from .tensor import ( ScaledTensor, @@ -37,7 +37,6 @@ from ..sharding import get_num_devices_in_mesh __all__ = [ - "QuantizeLayout", "Quantizer", "QuantizerSet", "CurrentScaleQuantizer", @@ -118,14 +117,6 @@ def update(self, *args, **kwargs): """Update quantizer state (no-op in base class).""" del args, kwargs - def is_2x2x(self) -> bool: - """Check if quantizer uses both row-wise and column-wise quantization. - - Returns: - True if using both row-wise and column-wise quantization - """ - return self.q_layout == QuantizeLayout.ROWWISE_COLWISE - def get_data_layout(self) -> str: """Get the data data_layout string. @@ -135,11 +126,11 @@ def get_data_layout(self) -> str: Raises: ValueError: If quantization axis is invalid """ - if self.q_layout == QuantizeLayout.ROWWISE_COLWISE: + if self.q_layout.is_rowwise_colwise: return self.data_layout - if self.q_layout == QuantizeLayout.ROWWISE: + if self.q_layout.is_rowwise_only: return self.data_layout[0] - if self.q_layout == QuantizeLayout.COLWISE: + if self.q_layout.is_colwise_only: return self.data_layout[1] raise ValueError(f"Invalid q_layout: {self.q_layout}") @@ -174,18 +165,10 @@ def quantize( """ del kwargs - is_rowwise = ( - is_rowwise - if is_rowwise is not None - else (self.q_layout == QuantizeLayout.ROWWISE or self.is_2x2x()) - ) - is_colwise = ( - is_colwise - if is_colwise is not None - else (self.q_layout == QuantizeLayout.COLWISE or self.is_2x2x()) - ) + is_rowwise = is_rowwise if is_rowwise is not None else self.q_layout.has_rowwise + is_colwise = is_colwise if is_colwise is not None else self.q_layout.has_colwise - if (is_rowwise and is_colwise) or self.is_2x2x(): + if is_rowwise and is_colwise: rowwise_tensor = self._quantize_func(x, dq_dtype=dq_dtype, flatten_axis=flatten_axis) colwise_tensor = self._quantize_func( x, is_colwise=True, dq_dtype=dq_dtype, flatten_axis=flatten_axis @@ -299,16 +282,8 @@ def quantize( flatten_axis += x.ndim assert 0 < flatten_axis < x.ndim, "flatten_axis is out of bounds!" - is_rowwise = ( - is_rowwise - if is_rowwise is not None - else (self.q_layout == QuantizeLayout.ROWWISE or self.is_2x2x()) - ) - is_colwise = ( - is_colwise - if is_colwise is not None - else (self.q_layout == QuantizeLayout.COLWISE or self.is_2x2x()) - ) + is_rowwise = is_rowwise if is_rowwise is not None else self.q_layout.has_rowwise + is_colwise = is_colwise if is_colwise is not None else self.q_layout.has_colwise rowwise_tensor = self._quantize_func(x, dq_dtype=dq_dtype, flatten_axis=flatten_axis) colwise_tensor = None @@ -974,16 +949,8 @@ def quantize( flatten_axis += x.ndim assert 0 < flatten_axis < x.ndim, "flatten_axis is out of bounds!" - is_rowwise = ( - is_rowwise - if is_rowwise is not None - else (self.q_layout == QuantizeLayout.ROWWISE or self.is_2x2x()) - ) - is_colwise = ( - is_colwise - if is_colwise is not None - else (self.q_layout == QuantizeLayout.COLWISE or self.is_2x2x()) - ) + is_rowwise = is_rowwise if is_rowwise is not None else self.q_layout.has_rowwise + is_colwise = is_colwise if is_colwise is not None else self.q_layout.has_colwise assert is_rowwise or is_colwise, "No quantization layout is specified" original_shape = x.shape diff --git a/transformer_engine/jax/quantize/scaling_modes.py b/transformer_engine/jax/quantize/scaling_modes.py index d490e02752..eea27a35d8 100644 --- a/transformer_engine/jax/quantize/scaling_modes.py +++ b/transformer_engine/jax/quantize/scaling_modes.py @@ -21,7 +21,8 @@ from jax.tree_util import register_pytree_node_class import jax.numpy as jnp -from transformer_engine_jax import JAXX_Scaling_Mode, QuantizeLayout +from transformer_engine_jax import JAXX_Scaling_Mode +from .misc import QuantizeLayout from .device_utils import is_fp8_gemm_with_all_layouts_supported @@ -72,16 +73,18 @@ class QuantizeShardyRules: Attributes: input_spec: Specification for the input axes - rowwise_rule: Sharding rule for the row-wise scale tensor, depends on - the axes in `input_spec` - colwise_rule: Likewise for the column-wise scale tensor. - factor_sizes: For block scaling, contains the block size factor, which is - used in `input_spec`. + rowwise_out_spec: Sharding spec for the rowwise quantized data + rowwise_scale_spec: Sharding spec for the rowwise scale + colwise_out_spec: Sharding spec for the colwise quantized data + colwise_scale_spec: Sharding spec for the colwise scale + factor_sizes: For block scaling, contains the block size factor """ input_spec: Tuple[str] - rowwise_rule: Tuple[str] - colwise_rule: Tuple[str] + rowwise_out_spec: Tuple[str] + rowwise_scale_spec: Tuple[str] + colwise_out_spec: Tuple[str] + colwise_scale_spec: Tuple[str] factor_sizes: Dict[str, int] @@ -166,7 +169,9 @@ def get_shardy_sharding_rules( input_shape, unique_var, flatten_axis, + q_layout, broadcast_2d_scale_shape_to_1d, + is_colwise_transposed, ) -> QuantizeShardyRules: """Sharding rules for the input and (row, col)wise scale tensors. @@ -174,7 +179,9 @@ def get_shardy_sharding_rules( input_shape: The shape of the input tensor (for which we produce the scale tensor) unique_var: An otherwise unused Shardy variable name prefix flatten_axis: Axis along which data can be flattened to 2D for quantization + q_layout: The layout of the quantized tensor broadcast_2d_scale_shape_to_1d: Whether to broadcast the 2D scale shape to 1D. + is_colwise_transposed: Whether the column-wise tensors are transposed. Returns: The Shardy rules for the scaling mode @@ -268,7 +275,9 @@ def get_shardy_sharding_rules( input_shape, unique_var, flatten_axis, + q_layout, broadcast_2d_scale_shape_to_1d, + is_colwise_transposed, ) -> QuantizeShardyRules: """Sharding rules for the input and (row, col)wise scale tensors. @@ -281,10 +290,17 @@ def get_shardy_sharding_rules( Returns: The Shardy rules for the scaling mode """ - del flatten_axis, broadcast_2d_scale_shape_to_1d - input_spec = tuple(f"{unique_var}{i}" for i in range(len(input_shape))) - scale_var = BATCHING + unique_var + "_scale_inv" - return QuantizeShardyRules(input_spec, (scale_var,), (scale_var,), {}) + del broadcast_2d_scale_shape_to_1d + input_spec = tuple(f"{unique_var}_x_{i}" for i in range(len(input_shape))) + output_spec = tuple(input_spec) + return QuantizeShardyRules( + input_spec, + output_spec, + (BATCHING + f"{unique_var}_scale",), + (BATCHING + f"{unique_var}_colwise_output",), + (BATCHING + f"{unique_var}_colwise_scale",), + {}, + ) class CurrentScalingModeMetadataImpl(ScalingModeMetadataImpl): @@ -376,7 +392,9 @@ def get_shardy_sharding_rules( input_shape, unique_var, flatten_axis, + q_layout, broadcast_2d_scale_shape_to_1d, + is_colwise_transposed, ) -> QuantizeShardyRules: """Sharding rules for the input and (row, col)wise scale tensors. @@ -385,14 +403,26 @@ def get_shardy_sharding_rules( unique_var: An otherwise unused Shardy variable name prefix flatten_axis: Axis along which data can be flattened to 2D for quantization broadcast_2d_scale_shape_to_1d: Whether to broadcast the 2D scale shape to 1D. - + q_layout: The layout of the quantized tensor + is_colwise_transposed: Whether the colwise scaling is transposed Returns: The Shardy rules for the scaling mode """ - del flatten_axis, broadcast_2d_scale_shape_to_1d - input_spec = tuple(f"{unique_var}{i}" for i in range(len(input_shape))) - scale_var = BATCHING + unique_var + "_scale_inv" - return QuantizeShardyRules(input_spec, (scale_var,), (scale_var,), {}) + del broadcast_2d_scale_shape_to_1d + input_spec = tuple(f"{unique_var}x_{i}" for i in range(len(input_shape))) + output_spec = input_spec + colwise_output_spec = (BATCHING + f"{unique_var}_colwise_output",) + + if q_layout.has_colwise: + from ..cpp_extensions.misc import multidim_transpose + + colwise_output_spec = input_spec + if is_colwise_transposed: + colwise_output_spec = multidim_transpose( + colwise_output_spec, transpose_axis=flatten_axis + ) + scale = (BATCHING + unique_var + "_scale_inv",) + return QuantizeShardyRules(input_spec, output_spec, scale, colwise_output_spec, scale, {}) class DelayedScalingModeMetadataImpl(CurrentScalingModeMetadataImpl): @@ -658,7 +688,9 @@ def get_shardy_sharding_rules( input_shape, unique_var, flatten_axis, + q_layout, broadcast_2d_scale_shape_to_1d, + is_colwise_transposed, ) -> QuantizeShardyRules: """Sharding rules for the input and (row, col)wise scale tensors. @@ -666,15 +698,18 @@ def get_shardy_sharding_rules( input_shape: The shape of the input tensor (for which we produce the scale tensor) unique_var: An otherwise unused Shardy variable name prefix flatten_axis: Axis along which data can be flattened to 2D for quantization + q_layout: The layout of the quantized tensor broadcast_2d_scale_shape_to_1d: Whether to broadcast the 2D scale shape to 1D. - + is_colwise_transposed: Whether the column-wise tensors are transposed. Returns: The Shardy rules for the scaling mode """ - # TODO(Phuong): to rework the shardy rule to handle transposes after NVFP4 is upstreamed + is_rowwise = q_layout.has_rowwise + is_colwise = q_layout.has_colwise + input_rank = len(input_shape) - input_spec = [f"{unique_var}_{i}" for i in range(input_rank)] flatten_axis = (flatten_axis + input_rank) % input_rank + input_spec = [f"{unique_var}_x_{i}" for i in range(input_rank)] assert ( self._block_dims[1] != 1 @@ -690,30 +725,56 @@ def get_shardy_sharding_rules( # We have to use two different factors in the two CompoundFactors because of Shardy # verifier requirements, even though they are the same. + # No CompoundFactor is needed if the dim has the same size as the blocksize blocksizes = {} - colwise_var = f"{unique_var}_None" rowwise_var = f"{unique_var}_None" - if not input_shape[-1] == block_size_1d: + colwise_var = f"{unique_var}_None" + if is_rowwise and not input_shape[-1] == block_size_1d: rowwise_var = input_spec[-1] + "_compound" input_spec[-1] = CompoundFactor(rowwise_var, "blocksize_x") blocksizes["blocksize_x"] = block_size_1d - if not input_shape[flatten_axis - 1] == block_size_1d: + if is_colwise and not input_shape[flatten_axis - 1] == block_size_1d: colwise_var = input_spec[flatten_axis - 1] + "_compound" input_spec[flatten_axis - 1] = CompoundFactor(colwise_var, "blocksize_y") blocksizes["blocksize_y"] = block_size_1d # The rowwise and colwise scale tensors should be sharded the same way as the input. # However, we need to adjust the dimensions where the block scaling factor applies. - rowwise = input_spec.copy() - rowwise[-1] = rowwise_var + if is_rowwise: + rowwise_out = input_spec.copy() + rowwise_scale = input_spec.copy() + rowwise_scale[-1] = rowwise_var + else: + rowwise_out = [ + BATCHING + f"{unique_var}_rowwise_output", + ] + rowwise_scale = [ + BATCHING + f"{unique_var}_rowwise_scale_inv", + ] - colwise = input_spec.copy() - colwise[flatten_axis - 1] = colwise_var + if is_colwise: + colwise_out = input_spec.copy() + colwise_scale = input_spec.copy() + colwise_scale[flatten_axis - 1] = colwise_var + if is_colwise_transposed: + from ..cpp_extensions.misc import multidim_transpose + + colwise_out = multidim_transpose(colwise_out, transpose_axis=flatten_axis) + colwise_scale = multidim_transpose(colwise_scale, transpose_axis=flatten_axis) + else: + colwise_out = [ + BATCHING + f"{unique_var}_colwise_output", + ] + colwise_scale = [ + BATCHING + f"{unique_var}_colwise_scale_inv", + ] return QuantizeShardyRules( tuple(input_spec), - tuple(rowwise), - tuple(colwise), + tuple(rowwise_out), + tuple(rowwise_scale), + tuple(colwise_out), + tuple(colwise_scale), blocksizes, ) @@ -850,7 +911,8 @@ def get_shardy_sharding_rules( self, input_shape, unique_var, - flatten_axis=-1, + flatten_axis, + q_layout, broadcast_2d_scale_shape_to_1d=False, ) -> Tuple[Tuple[str]]: """Sharding rules for the input and (row, col)wise scale tensors. @@ -859,13 +921,19 @@ def get_shardy_sharding_rules( input_shape: The shape of the input tensor (for which we produce the scale tensor) unique_var: An otherwise unused Shardy variable name prefix flatten_axis: Axis along which data can be flattened to 2D for quantization. + q_layout: The layout of the quantized tensor broadcast_2d_scale_shape_to_1d: Whether to broadcast the 2D scale shape to 1D. Defaults to False. Returns: The Shardy rules for the scaling mode """ return self._get_impl().get_shardy_sharding_rules( - input_shape, unique_var, flatten_axis, broadcast_2d_scale_shape_to_1d + input_shape, + unique_var, + flatten_axis, + q_layout, + broadcast_2d_scale_shape_to_1d, + self.is_colwise_transposed, ) def get_grouped_scale_shape_2x( diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index 6c358a044e..25db844098 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -15,10 +15,10 @@ import jax.numpy as jnp from jax.tree_util import register_pytree_node_class -from transformer_engine_jax import QuantizeLayout from .scaling_modes import ScalingMode, TensorUsage from .dequantizer import ScalingModeToDequantizerMap +from .misc import QuantizeLayout from ..sharding import ( with_sharding_constraint_by_logical_axes as original_with_sharding_constraint_by_logical_axes, ) @@ -128,9 +128,7 @@ def dequantize(self): def get_tensor(self, usage: TensorUsage): """Returns the tensor based on the tensor usage.""" q_layout = ScalingMode.NO_SCALING.get_quantize_layout(usage) - assert ( - q_layout == QuantizeLayout.ROWWISE - ), "Only ROWWISE layout is supported for NoScaleTensor" + assert q_layout.is_rowwise_only, "Only ROWWISE layout is supported for NoScaleTensor" return self def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[str, ...]): @@ -264,8 +262,8 @@ def dequantize(self): def get_tensor(self, usage: TensorUsage): """Returns the tensor based on the tensor usage.""" q_layout = self.scaling_mode.get_quantize_layout(usage) - colwise_usage_valid = q_layout == QuantizeLayout.COLWISE and self.is_colwise - rowwise_usage_valid = q_layout == QuantizeLayout.ROWWISE and not self.is_colwise + colwise_usage_valid = q_layout.is_colwise_only and self.is_colwise + rowwise_usage_valid = q_layout.is_rowwise_only and not self.is_colwise if colwise_usage_valid or rowwise_usage_valid: return self @@ -301,16 +299,15 @@ def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[st data = with_sharding_constraint_by_logical_axes(self.data, axis_names) - if self.scaling_mode == ScalingMode.MXFP8_1D_SCALING: - # TODO(Phuong): Handle padding !? + if self.scaling_mode.is_block_scaling: # Both MXFP8 and NVFP4 scale_inv = with_sharding_constraint_by_logical_axes(self.scale_inv, axis_names) else: scale_inv = self.scale_inv return ScaledTensor1x( data=data, - scale_inv=scale_inv, amax=self.amax, + scale_inv=scale_inv, scaling_mode=self.scaling_mode, dq_dtype=self.dq_dtype, _dq_func=self._dq_func, @@ -467,10 +464,10 @@ def get_tensor(self, usage: TensorUsage): q_layout_rowwise = self.rowwise_tensor.scaling_mode.get_quantize_layout(usage) q_layout_colwise = self.colwise_tensor.scaling_mode.get_quantize_layout(usage) - if q_layout_rowwise == QuantizeLayout.ROWWISE: + if q_layout_rowwise.is_rowwise_only: return self.rowwise_tensor - if q_layout_colwise == QuantizeLayout.COLWISE: + if q_layout_colwise.is_colwise_only: return self.colwise_tensor raise ValueError( @@ -548,13 +545,13 @@ def create_1x( dequantizer = ScalingModeToDequantizerMap.get(scaling_mode) if group_sizes is not None: - flatten_axis = len(original_shape) + flatten_axis if flatten_axis < 0 else flatten_axis + flatten_axis = (len(original_shape) + flatten_axis) % len(original_shape) assert ( original_shape is not None ), "original_shape is not given for GroupedScaledTensor1x" # Handling attrs of transposed tensors - group_axis = len(original_shape) + group_axis if group_axis < 0 else group_axis + group_axis = (len(original_shape) + group_axis) % len(original_shape) if data_layout == "T": if original_shape[0] == group_sizes.size: original_shape = ( @@ -587,7 +584,7 @@ def create_1x( ) # Handling attrs of transposed tensors - flatten_axis = data.ndim + flatten_axis if flatten_axis < 0 else flatten_axis + flatten_axis = (data.ndim + flatten_axis) % data.ndim if data_layout == "T": flatten_axis = data.ndim - flatten_axis @@ -669,7 +666,7 @@ def create_2x( colwise_amax, scaling_mode, dq_dtype, - is_colwise=True, # TODO(Phuong): set this correctly + is_colwise=True, data_layout=data_layout[1], flatten_axis=flatten_axis, group_sizes=group_sizes, @@ -721,7 +718,7 @@ def create( """ assert not rowwise_has_rht_applied, "RHT is not supported for rowwise quantization yet" - if q_layout == QuantizeLayout.ROWWISE_COLWISE: + if q_layout.is_rowwise_colwise: return ScaledTensorFactory.create_2x( data, scale_inv, @@ -740,15 +737,14 @@ def create( colwise_has_rht_applied=colwise_has_rht_applied, ) - is_colwise = q_layout == QuantizeLayout.COLWISE - if is_colwise: + if q_layout.is_colwise_only: return ScaledTensorFactory.create_1x( colwise_data, colwise_scale_inv, colwise_amax if colwise_amax is not None else amax, scaling_mode, dq_dtype, - is_colwise=is_colwise, + is_colwise=True, data_layout=data_layout[0], flatten_axis=flatten_axis, group_sizes=group_sizes, @@ -763,7 +759,7 @@ def create( amax, scaling_mode, dq_dtype, - is_colwise=is_colwise, + is_colwise=False, data_layout=data_layout[0], flatten_axis=flatten_axis, group_sizes=group_sizes, From 67d63d02f3efe1b8e0984788cc4e9ebf93bfd703 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 13 Nov 2025 13:58:32 -0800 Subject: [PATCH 062/521] [JAX] Support for checkpointing quantizations (#2356) * Support for checkpointing quantizations Signed-off-by: Jeremy Berchtold * Add jaxpr test for quant checkpoint name Signed-off-by: Jeremy Berchtold * Revert "Support for checkpointing quantizations" This reverts commit f7b784940369d0da2a77c57fa6ea744e883c5832. Signed-off-by: JAX Toolbox * Checkpoint quantizations Signed-off-by: Jeremy Berchtold * lint Signed-off-by: Jeremy Berchtold * revert other files Signed-off-by: Jeremy Berchtold * move checkpointing to VJPs Signed-off-by: Jeremy Berchtold * fix ci failure Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold Signed-off-by: JAX Toolbox Co-authored-by: JAX Toolbox --- tests/jax/test_recipe_characteristics.py | 121 +++++++++++++------ transformer_engine/jax/dense.py | 13 +- transformer_engine/jax/flax/module.py | 37 ++++-- transformer_engine/jax/layernorm_dense.py | 4 +- transformer_engine/jax/layernorm_mlp.py | 8 +- transformer_engine/jax/quantize/quantizer.py | 76 ++++++++++-- transformer_engine/jax/quantize/tensor.py | 55 +++++++++ 7 files changed, 252 insertions(+), 62 deletions(-) diff --git a/tests/jax/test_recipe_characteristics.py b/tests/jax/test_recipe_characteristics.py index 33fde7e231..b9c8fd7836 100644 --- a/tests/jax/test_recipe_characteristics.py +++ b/tests/jax/test_recipe_characteristics.py @@ -263,23 +263,16 @@ def test_autocast_nvfp4_block_scaling(self): class TestJaxprAndHlo: """Tests to verify Jaxpr and/or HLO of compiled modules apply expected recipe functionality and optimizations.""" - @pytest_parametrize_wrapper( - "quantization_recipe", - [ - quantization_recipe - for quantization_recipe in SUPPORTED_RECIPES - if isinstance(quantization_recipe, NVFP4BlockScaling) - ], - ) - def test_layernorm_mlp_reuses_amax_nvfp4(self, quantization_recipe): - """Tests that layernorm_mlp reuses the amax computed in layernorm and the activation and does not recompute it during quantizaton.""" - + def _generate_jaxpr_for_layernorm_mlp_fwd_bwd(self, quantization_recipe, ln_mlp_kwargs=None): + """Generates the jaxpr for a forward and backward pass of LayerNormMLP under the given quantization recipe.""" + ln_mlp_kwargs = ln_mlp_kwargs or {} with te.autocast(enabled=True, recipe=quantization_recipe, mesh_resource=te.MeshResource()): model = te_flax.LayerNormMLP( layernorm_type="rmsnorm", return_layernorm_output=False, intermediate_dropout_rate=0.0, dtype=jnp.bfloat16, + **ln_mlp_kwargs, ) var_collect = model.init( @@ -292,29 +285,83 @@ def loss_fn(x, rngs): x = jax.random.normal(jax.random.PRNGKey(0), (128, 128), dtype=jnp.bfloat16) rngs = {"sr_rng": jax.random.PRNGKey(1), "dropout": jax.random.PRNGKey(2)} - jaxpr = jax.make_jaxpr(jax.value_and_grad(loss_fn))(x, rngs=rngs) - - rht_amax_eqns = [ - eqn for eqn in jaxpr.jaxpr.eqns if eqn.primitive.name == "te_rht_amax_ffi_wrapper" - ] - - assert len(rht_amax_eqns) == 4, f"Expected 4 rht_amax_eqns, got {len(rht_amax_eqns)}" - - def assert_param(index, tensor_name, expected_value: bool): - if expected_value: - assert rht_amax_eqns[index].params["produce_regular_amax"] == True, ( - f"Expected produce_regular_amax for {tensor_name} to be True, indicating no" - " reuse of amax as this tensor does not have a previous operation to fuse" - " with" - ) - else: - assert rht_amax_eqns[index].params["produce_regular_amax"] == False, ( - f"Expected produce_regular_amax for {tensor_name} to be False, indicating" - " reuse of amax" - ) - - assert_param(0, "fwd ln+q", False) - assert_param(1, "fwd act+q", False) - # No previous op before incoming dgrad in the backward so amax is not reused - assert_param(2, "bwd dgrad", True) - assert_param(3, "bwd dact+q", False) + return jax.make_jaxpr(jax.value_and_grad(loss_fn))(x, rngs=rngs) + + @pytest_parametrize_wrapper( + "quantization_recipe", + [ + quantization_recipe + for quantization_recipe in SUPPORTED_RECIPES + if isinstance(quantization_recipe, NVFP4BlockScaling) + ], + ) + def test_layernorm_mlp_reuses_amax_nvfp4(self, quantization_recipe): + """Tests that layernorm_mlp reuses the amax computed in layernorm and the activation and does not recompute it during quantizaton.""" + + jaxpr = self._generate_jaxpr_for_layernorm_mlp_fwd_bwd(quantization_recipe) + + rht_amax_eqns = [ + eqn for eqn in jaxpr.jaxpr.eqns if eqn.primitive.name == "te_rht_amax_ffi_wrapper" + ] + + assert len(rht_amax_eqns) == 4, f"Expected 4 rht_amax_eqns, got {len(rht_amax_eqns)}" + + def assert_param(index, tensor_name, expected_value: bool): + if expected_value: + assert rht_amax_eqns[index].params["produce_regular_amax"] == True, ( + f"Expected produce_regular_amax for {tensor_name} to be True, indicating no" + " reuse of amax as this tensor does not have a previous operation to fuse" + " with" + ) + else: + assert rht_amax_eqns[index].params["produce_regular_amax"] == False, ( + f"Expected produce_regular_amax for {tensor_name} to be False, indicating" + " reuse of amax" + ) + + assert_param(0, "fwd ln+q", False) + assert_param(1, "fwd act+q", False) + # No previous op before incoming dgrad in the backward so amax is not reused + assert_param(2, "bwd dgrad", True) + assert_param(3, "bwd dact+q", False) + + @pytest_parametrize_wrapper("quantization_recipe", SUPPORTED_RECIPES) + @pytest_parametrize_wrapper( + "quantization_checkpoint_name", + [None, "quantization", "some_arbitrary_user_checkpoint_name"], + ) + def test_recipe_supports_quantization_checkpointing( + self, quantization_recipe, quantization_checkpoint_name + ): + """Tests that all supported quantization recipes correctly use checkpoint_name.""" + + kwargs = { + "quantization_checkpoint_name": quantization_checkpoint_name, + } + jaxpr = self._generate_jaxpr_for_layernorm_mlp_fwd_bwd(quantization_recipe, kwargs) + + checkpoint_name_eqns = [ + eqn + for eqn in jaxpr.jaxpr.eqns + if eqn.primitive.name == "name" and eqn.params["name"] == quantization_checkpoint_name + ] + + if quantization_checkpoint_name is None: + assert len(checkpoint_name_eqns) == 0, ( + "Expected 0 checkpoint_name eqns when quantization_checkpoint_name is None, got" + f" {len(checkpoint_name_eqns)}" + ) + return + + # 12 checkpointed values: + # - Fwd pass: + # - Input RMSNorm+Q -> 3 possible output tensors that will be used in the backward + # - Kernel Q -> 3 possible output tensors that will be used in the backward + # - Input Activation+Q -> 3 possible output tensors that will be used in the backward + # - Kernel Q -> 3 possible output tensors that will be used in the backward + expected_checkpoint_eqn_count = 12 + + assert len(checkpoint_name_eqns) == expected_checkpoint_eqn_count, ( + f"Expected {expected_checkpoint_eqn_count} checkpoint_name eqns when" + f" quantization_checkpoint_name is set, got {len(checkpoint_name_eqns)}" + ) diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 44c73a5b1e..c497775e01 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -19,6 +19,7 @@ from .cpp_extensions.amax import AmaxScope from .quantize import ( ScaledTensorFactory, + ScaledTensor, ScalingMode, QuantizeLayout, QuantizerSet, @@ -227,8 +228,8 @@ def _dense_fwd_rule( output += jnp.reshape(bias, bias_new_shape) ctx = ( - casted_x.get_tensor(usage=TensorUsage.LHS_TRANS), - casted_kernel.get_tensor(usage=TensorUsage.RHS_TRANS), + casted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(quantizer_set.x), + casted_kernel.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(quantizer_set.kernel), x.shape, kernel.shape, use_bias, @@ -529,8 +530,12 @@ def _grouped_dense_fwd_rule( ctx = ( group_sizes, - ctx_x, - ctx_kernel, + ctx_x.checkpoint(quantizer_set.x) if isinstance(ctx_x, ScaledTensor) else ctx_x, + ( + ctx_kernel.checkpoint(quantizer_set.kernel) + if isinstance(ctx_kernel, ScaledTensor) + else ctx_kernel + ), x.shape, kernel.shape, use_bias, diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 33ea610985..934af3d181 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -6,7 +6,7 @@ """ from functools import reduce import operator -from typing import Any, Callable, Iterable, List, Sequence, Tuple, Union, NewType +from typing import Any, Callable, Iterable, List, Sequence, Tuple, Union, NewType, Optional import numpy as np import jax.numpy as jnp @@ -345,7 +345,11 @@ class TransformerEngineBase(nn.Module): # pylint: disable=too-few-public-method """ def generate_quantizer_set( - self, postfix: str = "", variable_collection: str = None, fp8_recipe=None + self, + postfix: str = "", + variable_collection: str = None, + quantization_checkpoint_name: Optional[str] = None, + fp8_recipe=None, ): """ Generate a set of FP8 meta for a GEMM. @@ -375,7 +379,9 @@ def generate_quantizer_set( quantize_meta_set = QuantizeMetaSet(x=x_meta, kernel=kernel_meta, grad=grad_meta) quantizer_set = QuantizerFactory.create_set( - fp8_recipe=fp8_recipe, quantize_meta_set=quantize_meta_set + fp8_recipe=fp8_recipe, + quantize_meta_set=quantize_meta_set, + checkpoint_name=quantization_checkpoint_name, ) return quantizer_set @@ -424,6 +430,8 @@ class DenseGeneral(TransformerEngineBase): The data type used to allocate the initial parameters. transpose_batch_sequence: bool, default = False Indicate whether to transpose the batch and sequence dimensions of the input tensor. + quantization_checkpoint_name: Optional[str], default = None + The name for checkpointing quantizations. """ features: Union[Iterable[int], int] @@ -439,6 +447,7 @@ class DenseGeneral(TransformerEngineBase): dtype: DType = jnp.float32 input_axes: Tuple[str, ...] = () transpose_batch_sequence: bool = False + quantization_checkpoint_name: Optional[str] = None def __post_init__(self): if self.kernel_init is None: @@ -496,7 +505,9 @@ def __call__(self, inputs: Array) -> Array: else: bias = None - quantizer_set = self.generate_quantizer_set() + quantizer_set = self.generate_quantizer_set( + quantization_checkpoint_name=self.quantization_checkpoint_name + ) contract_ind = tuple(range(0, len(axis))) y = dense( inputs, @@ -628,6 +639,8 @@ class LayerNormDenseGeneral(TransformerEngineBase): value or None. When None is set, then no scaling is applied. transpose_batch_sequence: bool, default = False Indicate whether to transpose the batch and sequence dimensions of the input tensor. + quantization_checkpoint_name: Optional[str], default = None + The name for checkpointing quantizations. """ features: Union[Iterable[int], int] @@ -654,6 +667,7 @@ class LayerNormDenseGeneral(TransformerEngineBase): dot_input_axes: Tuple[str, ...] = None depth_scaling: float = None transpose_batch_sequence: bool = False + quantization_checkpoint_name: Optional[str] = None def __post_init__(self): if self.kernel_init is None: @@ -693,7 +707,9 @@ def __call__(self, inputs: Array) -> Array: input_dtype = inputs.dtype ln_output = None - quantizer_set = self.generate_quantizer_set() + quantizer_set = self.generate_quantizer_set( + quantization_checkpoint_name=self.quantization_checkpoint_name + ) fuse_layernorm = ( get_quantize_config().is_fp8_enabled() @@ -941,6 +957,8 @@ class LayerNormMLP(TransformerEngineBase): The data type used to allocate the initial parameters. transpose_batch_sequence: bool, default = False Indicate whether to transpose the batch and sequence dimensions of the input tensor. + quantization_checkpoint_name: Optional[str], default = None + The name for checkpointing quantizations. """ intermediate_dim: int = 2048 @@ -976,6 +994,7 @@ class LayerNormMLP(TransformerEngineBase): ffn1_ckpt_name: str = "ffn1" ffn2_ckpt_name: str = "ffn2" transpose_batch_sequence: bool = False + quantization_checkpoint_name: Optional[str] = None def __post_init__(self): if self.kernel_init is None: @@ -1010,8 +1029,12 @@ def __call__(self, inputs: Array, deterministic: bool = False) -> Array: """ assert self.axis == -1, "Only support axis == -1 at this moment" - ffn1_quantizer_set = self.generate_quantizer_set("_0") - ffn2_quantizer_set = self.generate_quantizer_set("_1") + ffn1_quantizer_set = self.generate_quantizer_set( + "_0", quantization_checkpoint_name=self.quantization_checkpoint_name + ) + ffn2_quantizer_set = self.generate_quantizer_set( + "_1", quantization_checkpoint_name=self.quantization_checkpoint_name + ) input_dtype = inputs.dtype ln_output = None diff --git a/transformer_engine/jax/layernorm_dense.py b/transformer_engine/jax/layernorm_dense.py index 705c742326..b9482b7bde 100644 --- a/transformer_engine/jax/layernorm_dense.py +++ b/transformer_engine/jax/layernorm_dense.py @@ -236,8 +236,8 @@ def _layernorm_dense_fwd_rule( output += jnp.reshape(bias, bias_new_shape) ctx = ( - casted_ln_out.get_tensor(TensorUsage.LHS_TRANS), - casted_kernel.get_tensor(TensorUsage.RHS_TRANS), + casted_ln_out.get_tensor(TensorUsage.LHS_TRANS).checkpoint(quantizer_set.x), + casted_kernel.get_tensor(TensorUsage.RHS_TRANS).checkpoint(quantizer_set.kernel), x.shape, kernel.shape, mu, diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py index 100848fdd5..2fd0f07d62 100644 --- a/transformer_engine/jax/layernorm_mlp.py +++ b/transformer_engine/jax/layernorm_mlp.py @@ -390,11 +390,11 @@ def _layernorm_mlp_fwd_rule( rsigma, gamma, beta, - casted_ln_out.get_tensor(TensorUsage.LHS_TRANS), - casted_kernel_1.get_tensor(TensorUsage.RHS_TRANS), + casted_ln_out.get_tensor(TensorUsage.LHS_TRANS).checkpoint(ffn1_quantizer_set.x), + casted_kernel_1.get_tensor(TensorUsage.RHS_TRANS).checkpoint(ffn1_quantizer_set.kernel), dot_1_output, - casted_act_out.get_tensor(TensorUsage.LHS_TRANS), - casted_kernel_2.get_tensor(TensorUsage.RHS_TRANS), + casted_act_out.get_tensor(TensorUsage.LHS_TRANS).checkpoint(ffn2_quantizer_set.x), + casted_kernel_2.get_tensor(TensorUsage.RHS_TRANS).checkpoint(ffn2_quantizer_set.kernel), x_contracting_dims, k_contracting_dims, kernel_1.shape, diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index 8a54f0b1db..adff317482 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -83,12 +83,15 @@ class Quantizer(ABC): q_dtype: The data type for quantized values scaling_mode: The scaling mode to use for quantization q_layout: The quantization axis (row-wise, column-wise, or both) + data_layout: The data layout string (e.g., "NT") + checkpoint_name: Optional name for checkpointing quantization state """ q_dtype: jnp.dtype scaling_mode: ScalingMode q_layout: QuantizeLayout data_layout: str + checkpoint_name: Optional[str] = None def tree_flatten(self): """Flatten the quantizer for JAX tree operations. @@ -97,7 +100,13 @@ def tree_flatten(self): Tuple of (children, aux_data) for tree operations """ children = () - aux_data = (self.q_dtype, self.scaling_mode, self.q_layout, self.data_layout) + aux_data = ( + self.q_dtype, + self.scaling_mode, + self.q_layout, + self.data_layout, + self.checkpoint_name, + ) return (children, aux_data) @classmethod @@ -337,7 +346,13 @@ def tree_flatten(self): Tuple of (children, aux_data) for tree operations """ children = (self.scale, self.amax_history) - aux_data = (self.q_dtype, self.scaling_mode, self.q_layout, self.data_layout) + aux_data = ( + self.q_dtype, + self.scaling_mode, + self.q_layout, + self.data_layout, + self.checkpoint_name, + ) return (children, aux_data) def _quantize_func( @@ -588,7 +603,14 @@ def tree_flatten(self): Tuple of (children, aux_data) for tree operations """ children = (self.stochastic_rounding_rng_state,) - aux_data = (self.q_dtype, self.scaling_mode, self.q_layout, self.data_layout, self.use_rht) + aux_data = ( + self.q_dtype, + self.scaling_mode, + self.q_layout, + self.data_layout, + self.checkpoint_name, + self.use_rht, + ) return (children, aux_data) @classmethod @@ -867,7 +889,14 @@ def tree_flatten(self): Tuple of (children, aux_data) for tree operations """ children = (self.quantizers,) - aux_data = (self.q_dtype, self.scaling_mode, self.q_layout, self.data_layout, self.n_groups) + aux_data = ( + self.q_dtype, + self.scaling_mode, + self.q_layout, + self.data_layout, + self.checkpoint_name, + self.n_groups, + ) return (children, aux_data) def __post_init__(self): @@ -1041,6 +1070,7 @@ def create( q_dtype: jnp.dtype = None, q_layout: QuantizeLayout = None, n_groups: int = None, + checkpoint_name: Optional[str] = None, **kwargs, ) -> Quantizer: """Create one or more quantizers with specified parameters. @@ -1052,6 +1082,7 @@ def create( q_layout: Quantization axis flatten_axis: The quantization axis for the tensor n_groups: Number of quantizers if GroupedQuantizer + checkpoint_name: Optional name for checkpointing quantizations **kwargs: Additional arguments for quantizer initialization Returns: @@ -1075,7 +1106,11 @@ def create( for _ in range(n_quantizers): quantizers.append( quantizer_type( - q_dtype=q_dtype, scaling_mode=scaling_mode, q_layout=q_layout, **kwargs + q_dtype=q_dtype, + scaling_mode=scaling_mode, + q_layout=q_layout, + checkpoint_name=checkpoint_name, + **kwargs, ) ) return quantizers[0] if len(quantizers) == 1 else tuple(quantizers) @@ -1089,6 +1124,7 @@ def _create_set( bwd_dtype, is_2x2x, n_groups, + checkpoint_name: Optional[str] = None, **kwargs, ) -> QuantizerSet: """Create a set of quantizers for forward and backward passes. @@ -1101,6 +1137,7 @@ def _create_set( bwd_dtype: Data type for backward pass is_2x2x: Whether to use 2x2x quantization n_groups + checkpoint_name: Optional name for checkpointing quantizations **kwargs: Additional arguments for quantizer initialization Returns: @@ -1123,12 +1160,32 @@ def _create_set( else: args_x = args_kernel = args_grad = {} - q_x = QuantizerFactory.create(1, x_scaling_mode, fwd_dtype, q_layout_x, n_groups, **args_x) + q_x = QuantizerFactory.create( + 1, + x_scaling_mode, + fwd_dtype, + q_layout_x, + n_groups, + checkpoint_name=checkpoint_name, + **args_x, + ) q_kernel = QuantizerFactory.create( - 1, kernel_scaling_mode, fwd_dtype, q_layout_kernel, n_groups, **args_kernel + 1, + kernel_scaling_mode, + fwd_dtype, + q_layout_kernel, + n_groups, + checkpoint_name=checkpoint_name, + **args_kernel, ) q_dgrad = QuantizerFactory.create( - 1, grad_scaling_mode, bwd_dtype, q_layout_dgrad, n_groups, **args_grad + 1, + grad_scaling_mode, + bwd_dtype, + q_layout_dgrad, + n_groups, + checkpoint_name=checkpoint_name, + **args_grad, ) return QuantizerSet(x=q_x, kernel=q_kernel, dgrad=q_dgrad) @@ -1140,6 +1197,7 @@ def create_set( bwd_dtype: jnp.dtype = None, is_2x2x: bool = None, n_groups: int = None, + checkpoint_name: Optional[str] = None, # TODO(jberchtold): rename fp8_recipe to quantization_recipe fp8_recipe: Optional[recipe.Recipe] = None, **kwargs, @@ -1153,6 +1211,7 @@ def create_set( bwd_dtype: Data type for backward pass, default is get_quantize_config().BWD_DTYPE is_2x2x: Whether to use 2x2x quantization, default is get_quantize_config().IF_QUANTIZE_2X n_groups: + checkpoint_name: Optional name for checkpointing quantizations fp8_recipe: Recipe to use for quantization. Scaling mode can be specified directly via the scaling_mode parameter or indirectly via recipe. Recipe is preferred as it will support additional recipes in future where scaling mode differs between x, kernel, and grad in the quantizer set. **kwargs: Additional arguments for quantizer initialization @@ -1208,6 +1267,7 @@ def create_set( bwd_dtype=bwd_dtype, is_2x2x=is_2x2x, n_groups=n_groups, + checkpoint_name=checkpoint_name, **kwargs, ) ) diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index 25db844098..90f139c3da 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -14,6 +14,7 @@ import jax.numpy as jnp from jax.tree_util import register_pytree_node_class +from jax.ad_checkpoint import checkpoint_name as jax_checkpoint_name from .scaling_modes import ScalingMode, TensorUsage @@ -89,6 +90,17 @@ def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[st The tensor with applied sharding constraints """ + @abstractmethod + def checkpoint(self, quantizer): + """Checkpoints the tensor with the given quantizer's checkpoint name if available. + + Args: + quantizer: The quantizer to use for checkpointing. If None, no checkpointing is applied. + + Returns: + The checkpointed tensor + """ + @dataclass class AbstractBaseTensor1x(AbstractBaseTensor): @@ -150,6 +162,18 @@ def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[st amax=self.amax, ) + def checkpoint(self, quantizer): + """Checkpoints the tensor with the given quantizer's checkpoint name if available. + + Args: + quantizer: The quantizer to use for checkpointing. If None, no checkpointing is applied. + + Returns: + The checkpointed tensor + """ + assert quantizer is None, "NoScaleTensor does not support quantization." + return self + class ScaledTensor(ABC): """Abstract base class for scaled tensors.""" @@ -317,6 +341,20 @@ def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[st has_rht_applied=self.has_rht_applied, ) + def checkpoint(self, quantizer): + """Checkpoints the tensor with the given quantizer's checkpoint name if available. + + Args: + quantizer: The quantizer to use for checkpointing. If None, no checkpointing is applied. + + Returns: + The checkpointed tensor + """ + if quantizer is None or quantizer.checkpoint_name is None: + return self + + return jax_checkpoint_name(self, name=quantizer.checkpoint_name) + @register_pytree_node_class @dataclass @@ -420,6 +458,20 @@ def tree_flatten(self): def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[str, ...]): raise NotImplementedError + def checkpoint(self, quantizer): + """Checkpoints the tensor with the given quantizer's checkpoint name if available. + + Args: + quantizer: The quantizer to use for checkpointing. If None, no checkpointing is applied. + + Returns: + The checkpointed tensor + """ + if quantizer is None or quantizer.checkpoint_name is None: + return self + + return jax_checkpoint_name(self, name=quantizer.checkpoint_name) + @register_pytree_node_class @dataclass @@ -496,6 +548,9 @@ def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[st return ScaledTensor2x(rowwise_tensor, colwise_tensor) + def checkpoint(self, quantizer): + raise NotImplementedError + @dataclass class ScaledTensorFactory: From 0ded11340ba28267d0826fca550e0666ea8a00aa Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 13 Nov 2025 18:40:37 -0500 Subject: [PATCH 063/521] [JAX] XLA_FLAG to WAR the current NCCL issue with test_distributed_softmax.py (#2378) * add war for test_distributed_softmax.py Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- qa/L1_jax_distributed_unittest/test.sh | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index f4ea2dd68e..886f27747e 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -28,7 +28,9 @@ python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/py python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_mlp.xml $TE_PATH/tests/jax/test_distributed_layernorm_mlp.py || test_fail "test_distributed_layernorm_mlp.py" -python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py" +# XLA_FLAGS to WAR for test_distributed_softmax issue with NCCL +# TODO(Kshitij): remove when NCCL issue is fixed +XLA_FLAGS="$XLA_FLAGS --xla_gpu_enable_nccl_comm_splitting=false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_softmax.xml $TE_PATH/tests/jax/test_distributed_softmax.py || test_fail "test_distributed_softmax.py" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_dist_fused_attn.xml $TE_PATH/tests/jax/test_distributed_fused_attn.py || test_fail "test_distributed_fused_attn.py" From 262c184eb8331dcb50037477881900e46bd5c5f2 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Fri, 14 Nov 2025 23:53:48 +0800 Subject: [PATCH 064/521] [PyTorch] Add reset cudagraph interface (#2367) * reset cudagraph Signed-off-by: Robin Zhang * use closure instead of mutable default values Signed-off-by: Robin Zhang * add test Signed-off-by: Robin Zhang * fix test Signed-off-by: Robin Zhang --------- Signed-off-by: Robin Zhang Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_cuda_graphs.py | 38 +++++++++++++++++++++++++---- transformer_engine/pytorch/graph.py | 27 +++++++++++++------- 2 files changed, 51 insertions(+), 14 deletions(-) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index fa8754d601..eacbf5168e 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -from typing import Iterable, List, Union +from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest import torch @@ -160,6 +160,20 @@ def get_outputs( return values +def reset_graphs( + graphed_callables: Union[Callable, Tuple[Callable, ...], Dict[Tuple[int, int], Callable]], +) -> None: + """Reset CUDA graphs.""" + if isinstance(graphed_callables, tuple) or isinstance(graphed_callables, list): + for callable in graphed_callables: + callable.reset() + elif isinstance(graphed_callables, dict): + for callable in graphed_callables.values(): + callable.reset() + else: + graphed_callables.reset() + + class _Sequential(torch.nn.Sequential): """Sequential model that forwards keyword arguments to modules""" @@ -322,7 +336,12 @@ def _test_cuda_graphs( output.backward(grad_output) optimizer.step() - return get_outputs(model, output) + outputs = get_outputs(model, output) + if graph_mode == "full": + reset_graphs(model) + elif graph_mode == "individual": + reset_graphs(modules) + return outputs @pytest.mark.parametrize("module", _test_cuda_graphs_modules) @@ -468,7 +487,10 @@ def _test_cuda_graphs_with_dot_product_attention( output = model(*inputs) output.backward(grad_output) - return get_outputs(model, output) + outputs = get_outputs(model, output) + if with_graph: + reset_graphs(model) + return outputs @pytest.mark.parametrize("dtype", dtypes) @@ -553,7 +575,10 @@ def _test_cuda_graphs_with_kwargs( output.backward(grad_output) optimizer.step() - return get_outputs(model, output) + outputs = get_outputs(model, output) + if with_graph: + reset_graphs(model) + return outputs def test_make_graphed_callables_with_kwargs( @@ -668,7 +693,10 @@ def backward(layer_idx: int, microbatch_idx: int): optimizer.step() outputs = [y for _, y in sorted(outputs.items())] - return get_outputs(model, outputs) + outputs = get_outputs(model, outputs) + if with_graph: + reset_graphs(layer_forwards) + return outputs def test_make_graphed_callables_with_interleaved_pipeline_parallelism( diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 9af9fb8870..f55f1dd128 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -756,6 +756,21 @@ def functionalized(*user_args, **user_kwargs): return functionalized + def make_graphed_attribute_functions(graph_idx): + + # Attach backward_dw as an attribute to the graphed callable. + def backward_dw(): + if need_bwd_dw_graph.get(graph_idx, False): + bwd_dw_graphs[graph_idx].replay() + + # Attach reset as an attribute to the graphed callable. + def reset(): + fwd_graphs[graph_idx].reset() + bwd_graphs[graph_idx].reset() + bwd_dw_graphs[graph_idx].reset() + + return backward_dw, reset + # Put together the final graphed callables ret = [] for i in range(len(sample_args)): @@ -831,15 +846,9 @@ def new_fwd(*user_args, **user_kwargs): else: ret.append(graphed) - # Attach backward_dw as an attribute to the graphed callable. - def backward_dw( - need_backward_dw=need_bwd_dw_graph.get(i, False), - bwd_dw_graph=bwd_dw_graphs[i], - ): - if need_backward_dw: - bwd_dw_graph.replay() - - setattr(ret[-1], "backward_dw", backward_dw) + backward_dw_func, reset_func = make_graphed_attribute_functions(i) + setattr(ret[-1], "backward_dw", backward_dw_func) + setattr(ret[-1], "reset", reset_func) if just_one_callable: return ret[0] From b88f727b44d7779200a7f57c279805930a3883ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 14 Nov 2025 18:26:42 +0100 Subject: [PATCH 065/521] [JAX] Make all jax attention calls use non-packed common calls (#2358) * fix Signed-off-by: Pawel Gadzinski * add notes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * small fixes 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> --- .../common/fused_attn/fused_attn.cpp | 2 +- .../jax/csrc/extensions/attention.cpp | 331 ++++++++---------- 2 files changed, 140 insertions(+), 193 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index ac6fefdc6a..611beb7b84 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -29,7 +29,7 @@ transformer_engine::Tensor make_tensor_view(const transformer_engine::Tensor *so return view; } -// Helper function to calculate stride for packed QKV tensor unpacking +// Helper function to calculate stride in bytes for packed QKV tensor unpacking size_t calculate_qkv_stride(NVTE_QKV_Layout_Group layout_group, transformer_engine::DType dtype, size_t h, size_t d) { size_t stride = 0; diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index a99f4fae90..ac7eba5c87 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -123,17 +123,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_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) { - // For qkv_packed - auto qkv_shape = std::vector{input_batch * q_max_seqlen, 3, attn_heads, qk_head_dim}; - auto qkv_tensor = TensorWrapper(nullptr, qkv_shape, dtype); - - // For kv_packed auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto kv_shape = std::vector{input_batch * kv_max_seqlen, 2, num_gqa_groups, v_head_dim}; - auto kv_tensor = TensorWrapper(nullptr, kv_shape, dtype); - - // For separate q, k, v auto k_shape = 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}; @@ -156,7 +147,6 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( nvte_tensor_pack_create(&aux_output_tensors); TensorWrapper query_workspace_tensor; - auto layout_group = nvte_get_qkv_layout_group(qkv_layout); 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; @@ -174,37 +164,14 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); auto ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal to kv_max_seqlen"); - nvte_fused_attn_fwd_qkvpacked( - qkv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), - s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), - ragged_offset_tensor.data(), dummy_rng_state_tensor.data(), q_max_seqlen, is_training, - false, false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, query_workspace_tensor.data(), - nullptr); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - nvte_fused_attn_fwd_kvpacked( - q_tensor.data(), kv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), - s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, query_workspace_tensor.data(), nullptr); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { - nvte_fused_attn_fwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, - 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, - query_workspace_tensor.data(), nullptr); - } else { - NVTE_ERROR("Unsupported QKVLayout."); - } + nvte_fused_attn_fwd( + q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), + dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, + 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, query_workspace_tensor.data(), nullptr); } nvte_tensor_pack_destroy(&aux_output_tensors); @@ -291,47 +258,57 @@ static void FusedAttnForwardImpl( /* Call the underlying NVTE API */ auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); + + // Prepare Q, K, V pointers and shapes based on layout + // Python passes dummy tensors for unused slots, so we extract from the actual packed data + 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}; + if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - auto qkv_shape = std::vector{input_batch * q_max_seqlen, 3, attn_heads, qk_head_dim}; - auto qkv_tensor = TensorWrapper(q, qkv_shape, dtype); - nvte_fused_attn_fwd_qkvpacked( - qkv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), s_tensor.data(), - o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), - q_seq_offsets_tensor.data(), rng_state_tensor.data(), q_max_seqlen, is_training, false, - false, scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, workspace_tensor.data(), stream); + // QKV packed in q: [batch*seqlen, 3, heads, dim] + // Python passes: q=packed_qkv, k=dummy, v=dummy + // Extract K and V pointers from the packed q data + NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal kv_max_seqlen"); + NVTE_CHECK(qk_head_dim == v_head_dim, + "For QKV packed layout, qk_head_dim must equal v_head_dim"); + size_t stride = (typeToSize(dtype) * attn_heads * qk_head_dim); + q_ptr = q; + k_ptr = static_cast(static_cast(q) + stride); + v_ptr = static_cast(static_cast(q) + 2 * stride); + // For packed QKV, all have same shape since they're views into the same packed tensor + k_shape = q_shape; + v_shape = q_shape; } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; - auto kv_shape = - std::vector{input_batch * kv_max_seqlen, 2, num_gqa_groups, qk_head_dim}; - auto q_tensor = TensorWrapper(q, q_shape, dtype); - auto kv_tensor = TensorWrapper(k, kv_shape, dtype); - nvte_fused_attn_fwd_kvpacked( - q_tensor.data(), kv_tensor.data(), bias_tensor.data(), dummy_softmax_offset_tensor.data(), - s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, - workspace_tensor.data(), stream); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { - 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_tensor = TensorWrapper(q, q_shape, dtype); - auto k_tensor = TensorWrapper(k, k_shape, dtype); - auto v_tensor = TensorWrapper(v, v_shape, dtype); - nvte_fused_attn_fwd( - q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, - 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, workspace_tensor.data(), stream); - } else { - NVTE_ERROR("Unsupported qkv_layout."); + // Q separate, KV packed in k: [batch*seqlen, 2, num_gqa_groups, dim] + // Python passes: q=query, k=packed_kv, v=dummy + // Extract V pointer from the packed k data + NVTE_CHECK(qk_head_dim == v_head_dim, + "For KV packed layout, qk_head_dim must equal v_head_dim"); + size_t stride = (typeToSize(dtype) * num_gqa_groups * qk_head_dim); + q_ptr = q; + k_ptr = k; + v_ptr = static_cast(static_cast(k) + stride); + // V has same shape as K since they're packed together + v_shape = k_shape; } + // else NVTE_HD_HD_HD: pointers and shapes already correct + + auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype); + auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); + auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); + + nvte_fused_attn_fwd( + q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), + dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, + 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, workspace_tensor.data(), stream); nvte_tensor_pack_destroy(&aux_output_tensors); } @@ -414,20 +391,9 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( NVTE_Mask_Type mask_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) { - // For qkv_packed - auto qkv_shape = std::vector{input_batch * q_max_seqlen, 3, attn_heads, qk_head_dim}; - auto qkv_tensor = TensorWrapper(nullptr, qkv_shape, dtype); - auto dqkv_tensor = TensorWrapper(nullptr, qkv_shape, dtype); - - // For kv_packed auto q_shape = 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 kv_shape = std::vector{input_batch * kv_max_seqlen, 2, num_gqa_groups, v_head_dim}; - auto kv_tensor = TensorWrapper(nullptr, kv_shape, dtype); - auto dkv_tensor = TensorWrapper(nullptr, kv_shape, dtype); - - // For separate q, k, v auto k_shape = 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); @@ -450,7 +416,6 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( TensorWrapper query_workspace_tensor; - auto layout_group = nvte_get_qkv_layout_group(qkv_layout); 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; @@ -471,42 +436,18 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); auto dummy_ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - nvte_fused_attn_bwd_qkvpacked( - qkv_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, dqkv_tensor.data(), dbias_tensor.data(), - dummy_d_softmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), - dummy_ragged_offset_tensor.data(), q_max_seqlen, scaling_factor, dropout_probability, - qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, - deterministic, false, query_workspace_tensor.data(), nullptr); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - nvte_fused_attn_bwd_kvpacked( - q_tensor.data(), kv_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(), dkv_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, deterministic, false, query_workspace_tensor.data(), nullptr); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { - 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, deterministic, false, query_workspace_tensor.data(), nullptr); - } else { - NVTE_ERROR("Unsupported qkv_layout."); - } + + 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, deterministic, false, query_workspace_tensor.data(), nullptr); } nvte_tensor_pack_destroy(&aux_input_tensors); @@ -552,76 +493,82 @@ static void FusedAttnBackwardImpl( softmax_aux, rng_state, bias); /* Call the underly NVTE API */ + // Prepare Q, K, V pointers and shapes based on layout + void *q_ptr = q; + void *k_ptr = k; + void *v_ptr = v; + 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}; + if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - auto qkv_shape = std::vector{input_batch * q_max_seqlen, 3, attn_heads, qk_head_dim}; - auto qkv_tensor = TensorWrapper(q, qkv_shape, dtype); - auto dqkv_tensor = TensorWrapper(dq, qkv_shape, dtype); - if (is_ragged) { - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(qkv_shape) * typeToSize(dtype), - stream); - } - nvte_fused_attn_bwd_qkvpacked(qkv_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, dqkv_tensor.data(), dbias_tensor.data(), - dummy_d_softmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), - q_seq_offsets_tensor.data(), q_max_seqlen, scaling_factor, - dropout_probability, qkv_layout, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, deterministic, - false, workspace_tensor.data(), stream); + // QKV packed in q: [batch*seqlen, 3, heads, dim] + NVTE_CHECK(q_max_seqlen == kv_max_seqlen, "q_max_seqlen must equal kv_max_seqlen"); + NVTE_CHECK(qk_head_dim == v_head_dim, + "For QKV packed layout, qk_head_dim must equal v_head_dim"); + size_t stride = (typeToSize(dtype) * attn_heads * qk_head_dim); + q_ptr = q; + k_ptr = static_cast(static_cast(q) + stride); + v_ptr = static_cast(static_cast(q) + 2 * stride); + dq_ptr = dq; + dk_ptr = static_cast(static_cast(dq) + stride); + dv_ptr = static_cast(static_cast(dq) + 2 * stride); + k_shape = q_shape; + v_shape = q_shape; } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; - auto kv_shape = - std::vector{input_batch * kv_max_seqlen, 2, num_gqa_groups, qk_head_dim}; - auto q_tensor = TensorWrapper(q, q_shape, dtype); - auto kv_tensor = TensorWrapper(k, kv_shape, dtype); - auto dq_tensor = TensorWrapper(dq, q_shape, dtype); - auto dkv_tensor = TensorWrapper(dk, kv_shape, dtype); - if (is_ragged) { - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * typeToSize(dtype), stream); - cudaMemsetAsync(dk, 0, transformer_engine::jax::product(kv_shape) * typeToSize(dtype), - stream); - } - nvte_fused_attn_bwd_kvpacked( - q_tensor.data(), kv_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(), dkv_tensor.data(), dbias_tensor.data(), - dummy_d_softmax_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, deterministic, false, - workspace_tensor.data(), stream); - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_HD_HD) { - 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_tensor = TensorWrapper(q, q_shape, dtype); - auto k_tensor = TensorWrapper(k, k_shape, dtype); - auto v_tensor = TensorWrapper(v, v_shape, dtype); - auto dq_tensor = TensorWrapper(dq, q_shape, dtype); - auto dk_tensor = TensorWrapper(dk, k_shape, dtype); - auto dv_tensor = TensorWrapper(dv, v_shape, dtype); - if (is_ragged) { - cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * typeToSize(dtype), stream); - cudaMemsetAsync(dk, 0, transformer_engine::jax::product(k_shape) * typeToSize(dtype), stream); - cudaMemsetAsync(dv, 0, transformer_engine::jax::product(v_shape) * typeToSize(dtype), stream); + // Q separate, KV packed in k: [batch*seqlen, 2, num_gqa_groups, dim] + NVTE_CHECK(qk_head_dim == v_head_dim, + "For KV packed layout, qk_head_dim must equal v_head_dim"); + size_t stride = (typeToSize(dtype) * num_gqa_groups * qk_head_dim); + q_ptr = q; + k_ptr = k; + v_ptr = static_cast(static_cast(k) + stride); + dq_ptr = dq; + dk_ptr = dk; + dv_ptr = static_cast(static_cast(dk) + stride); + // V has same shape as K since they're packed together + v_shape = k_shape; + } + + auto q_tensor = TensorWrapper(q_ptr, q_shape, dtype); + auto k_tensor = TensorWrapper(k_ptr, k_shape, dtype); + auto v_tensor = TensorWrapper(v_ptr, v_shape, dtype); + auto dq_tensor = TensorWrapper(dq_ptr, q_shape, dtype); + auto dk_tensor = TensorWrapper(dk_ptr, k_shape, dtype); + auto dv_tensor = TensorWrapper(dv_ptr, v_shape, dtype); + + if (is_ragged) { + size_t dtype_size = typeToSize(dtype); + if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { + // For packed QKV, dq contains all gradients (dq, dk, dv) - clear all at once + cudaMemsetAsync(dq, 0, 3 * transformer_engine::jax::product(q_shape) * dtype_size, stream); + } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { + // Clear dq + cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * dtype_size, stream); + // For packed KV, dk contains both dk and dv - clear all at once + cudaMemsetAsync(dk, 0, 2 * transformer_engine::jax::product(k_shape) * dtype_size, stream); + } else { + // All separate - clear each individually + cudaMemsetAsync(dq, 0, transformer_engine::jax::product(q_shape) * dtype_size, stream); + cudaMemsetAsync(dk, 0, transformer_engine::jax::product(k_shape) * dtype_size, stream); + cudaMemsetAsync(dv, 0, transformer_engine::jax::product(v_shape) * dtype_size, 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(), dummy_d_softmax_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, deterministic, - false, workspace_tensor.data(), stream); - } else { - NVTE_ERROR("Unsupported qkv_layout."); } + 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(), + 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, deterministic, false, workspace_tensor.data(), stream); + nvte_tensor_pack_destroy(&aux_input_tensors); } From a0754757660ac5a747b0be54c5398fca032161aa Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 14 Nov 2025 10:21:03 -0800 Subject: [PATCH 066/521] [JAX] Improve support and testing for direct recipe usage without autocast contexts (#2366) * Refactor to avoid storing a global quantization config so direct recipe passing works as intended Signed-off-by: Jeremy Berchtold * fix use_split_accumulator for current scaling recipe Signed-off-by: Jeremy Berchtold * fix tests that pass direct recipe and were missing quantize meta set Signed-off-by: Jeremy Berchtold * Revert "fix use_split_accumulator for current scaling recipe" This reverts commit a74ab7df812ec0a069b1bdd208debb93ec25a900. Signed-off-by: Jeremy Berchtold * fix ci failures Signed-off-by: Jeremy Berchtold * Fix amax_history post_init Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update transformer_engine/jax/quantize/quantizer.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ci failures Signed-off-by: Jeremy Berchtold * fix ci issue Signed-off-by: Jeremy Berchtold * address comments Signed-off-by: Jeremy Berchtold * make recipe assertion classes in test_recipe_characteristics not inherit from unittest.TestCase Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@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/jax/test_custom_call_compute.py | 19 +- tests/jax/test_layer.py | 21 +- tests/jax/test_recipe_characteristics.py | 338 +++++++++++------- transformer_engine/jax/cpp_extensions/base.py | 2 +- transformer_engine/jax/cpp_extensions/gemm.py | 23 +- .../jax/cpp_extensions/quantization.py | 4 +- transformer_engine/jax/dense.py | 3 +- transformer_engine/jax/flax/module.py | 34 +- transformer_engine/jax/layernorm_dense.py | 3 +- transformer_engine/jax/layernorm_mlp.py | 3 +- transformer_engine/jax/quantize/helper.py | 47 ++- transformer_engine/jax/quantize/quantizer.py | 89 +++-- 12 files changed, 373 insertions(+), 213 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 3d4f179abd..c8bd9d47c3 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -40,6 +40,8 @@ QuantizerFactory, QuantizeLayout, noop_quantizer_set, + QuantizeMetaSet, + QuantizeMeta, ) from transformer_engine.jax.quantize import helper from transformer_engine.jax.activation import activation @@ -1457,7 +1459,12 @@ def ref_func(x, w, bias, data_layout): value_n_grad_primitive_func = value_and_grad(primitive_func, (0, 1, 2)) value_n_grad_ref_func = value_and_grad(ref_func, (0, 1, 2)) - quantizer_set = QuantizerFactory.create_set(fp8_recipe=recipe) + quantizer_set = QuantizerFactory.create_set( + fp8_recipe=recipe, + quantize_meta_set=QuantizeMetaSet( + x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta() + ), + ) n_iterations = 3 if recipe.delayed() else 1 with use_jax_gemm(enabled=with_jax_gemm): @@ -1516,7 +1523,12 @@ def test_layernorm_dense_grad(self, m, n, k, recipe, norm_type, with_jax_gemm): gamma = jax.random.normal(subkeys[2], (k,)).astype(jnp.bfloat16) - quantizer_set = QuantizerFactory.create_set(fp8_recipe=recipe) + quantizer_set = QuantizerFactory.create_set( + fp8_recipe=recipe, + quantize_meta_set=QuantizeMetaSet( + x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta() + ), + ) if norm_type == "layernorm": beta = jax.random.normal(subkeys[3], (k,)).astype(jnp.bfloat16) @@ -1605,6 +1617,9 @@ def test_layernorm_mlp_grad( quantizer_sets = QuantizerFactory.create_set( n_quantizer_sets=2, fp8_recipe=recipe, + quantize_meta_set=QuantizeMetaSet( + x=QuantizeMeta(), kernel=QuantizeMeta(), grad=QuantizeMeta() + ), ) if norm_type == "layernorm": diff --git a/tests/jax/test_layer.py b/tests/jax/test_layer.py index d1b2535c4c..b51d6b2136 100644 --- a/tests/jax/test_layer.py +++ b/tests/jax/test_layer.py @@ -23,7 +23,8 @@ from transformer_engine.common import recipe from transformer_engine.jax.flax import TransformerLayer, TransformerLayerType from transformer_engine.jax.quantize import ( - get_quantize_config, + get_global_quantize_recipe, + get_quantize_config_with_recipe, ScalingMode, is_fp8_available, update_collections, @@ -358,7 +359,7 @@ def test_backward( ref_params, test_params = self._sync_params(ref_params, test_params) - if get_quantize_config().is_fp8_enabled(): + if get_quantize_config_with_recipe(get_global_quantize_recipe()).is_fp8_enabled(): for _ in range(4): _, updated_state = jax.value_and_grad(self._loss_fn, argnums=(3,), has_aux=False)( inputs, @@ -368,14 +369,24 @@ def test_backward( test_layer, ) if ( - get_quantize_config().get_scaling_mode(TensorSource.X) + get_quantize_config_with_recipe(get_global_quantize_recipe()).get_scaling_mode( + TensorSource.X + ) == ScalingMode.DELAYED_TENSOR_SCALING ): _, updated_quantize_meta = flax.core.pop( - updated_state[0], get_quantize_config().COLLECTION_NAME + updated_state[0], + get_quantize_config_with_recipe( + get_global_quantize_recipe() + ).COLLECTION_NAME, ) test_others = update_collections( - {get_quantize_config().COLLECTION_NAME: updated_quantize_meta}, test_others + { + get_quantize_config_with_recipe( + get_global_quantize_recipe() + ).COLLECTION_NAME: updated_quantize_meta + }, + test_others, ) del updated_quantize_meta del updated_state diff --git a/tests/jax/test_recipe_characteristics.py b/tests/jax/test_recipe_characteristics.py index b9c8fd7836..5171a6c622 100644 --- a/tests/jax/test_recipe_characteristics.py +++ b/tests/jax/test_recipe_characteristics.py @@ -4,6 +4,7 @@ import unittest from functools import partial +from abc import ABC, abstractmethod import flax import jax @@ -13,6 +14,7 @@ from utils import assert_allclose, pytest_parametrize_wrapper from transformer_engine.common.recipe import ( + Recipe, DelayedScaling, MXFP8BlockScaling, Float8CurrentScaling, @@ -21,13 +23,13 @@ from transformer_engine.common.recipe import Format as FP8Format from transformer_engine.jax import autocast from transformer_engine.jax.quantize import ( - get_quantize_config, + get_global_quantize_recipe, + get_quantize_config_with_recipe, get_supported_quantization_recipes, is_scaling_mode_supported, ScalingMode, update_collections, TensorSource, - QuantizerFactory, QuantizeLayout, ) from transformer_engine.jax.quantize.helper import _format2dtypes @@ -49,16 +51,17 @@ def quantizer_check_vjp(outer_quantizer_set, assertion_func, x): # Define a function with a custom VJP (vector-Jacobian product) @partial(jax.custom_vjp, nondiff_argnums=(1,)) def quantizer_check(inner_quantizer_set, assertion_func, x): - return quantizer_check_fwd(inner_quantizer_set, assertion_func, x) + return quantizer_check_fwd(inner_quantizer_set, assertion_func, x)[0] def quantizer_check_fwd(inner_quantizer_set, assertion_func, x): assertion_func(inner_quantizer_set.x, TensorSource.X) assertion_func(inner_quantizer_set.kernel, TensorSource.KERNEL) assertion_func(inner_quantizer_set.dgrad, TensorSource.DGRAD) - return x + return x, (inner_quantizer_set,) - def quantizer_check_bwd(ctx, g): - return (g,) + def quantizer_check_bwd(assertion_func, ctx, g): + (inner_quantizer_set,) = ctx + return (inner_quantizer_set, g) quantizer_check.defvjp(quantizer_check_fwd, quantizer_check_bwd) return quantizer_check(outer_quantizer_set, assertion_func, x) @@ -69,10 +72,11 @@ class TestModule(TransformerEngineBase): # Signature: (quantizer: Quantizer, tensor_source: TensorSource) -> None assertion_func: callable + direct_recipe: Recipe @nn.compact def __call__(self, x): - quantizer_set = self.generate_quantizer_set() + quantizer_set = self.generate_quantizer_set(fp8_recipe=self.direct_recipe) return quantizer_check_vjp(quantizer_set, self.assertion_func, x) @@ -97,167 +101,239 @@ def test_update_collections(self): self.assertEqual(updated_state["test2"], original_val) -class TestFP8Functions(unittest.TestCase): +def assert_fp8_format(quantizer, tensor_source, fp8_format): + if fp8_format == FP8Format.HYBRID: + if tensor_source == TensorSource.DGRAD: + assert quantizer.q_dtype == jnp.float8_e5m2 + else: + assert quantizer.q_dtype == jnp.float8_e4m3fn + elif fp8_format == FP8Format.E4M3: + assert quantizer.q_dtype == jnp.float8_e4m3fn + else: + raise ValueError(f"Unsupported FP8 format: {fp8_format}") - def _check_default_state(self): - self.assertFalse(get_quantize_config().is_fp8_enabled()) - - def _compare_delay_scaling(self, test): - self.assertEqual(get_quantize_config().MARGIN, test.margin) - self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp8_format)[0]) - self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp8_format)[1]) - self.assertEqual(get_quantize_config().AMAX_HISTORY_LEN, test.amax_history_len) - self.assertEqual(get_quantize_config().AMAX_COMPUTE_ALGO.value, test.amax_compute_algo) - - def _compare_current_scaling(self, test): - self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp8_format)[0]) - self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp8_format)[1]) - for tensor_source in TensorSource: - self.assertEqual( - get_quantize_config().get_scaling_mode(tensor_source), - ScalingMode.CURRENT_TENSOR_SCALING, - ) - def _compare_mxfp8_scaling(self, test): - self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp8_format)[0]) - self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp8_format)[1]) - for tensor_source in TensorSource: - self.assertEqual( - get_quantize_config().get_scaling_mode(tensor_source), ScalingMode.MXFP8_1D_SCALING - ) +class RecipeAssertionBase(ABC): + """Base class for defining recipe assertions.""" - def _compare_nvfp4_scaling(self, test): - self.assertEqual(get_quantize_config().FWD_DTYPE, _format2dtypes(test.fp4_format)[0]) - self.assertEqual(get_quantize_config().BWD_DTYPE, _format2dtypes(test.fp4_format)[1]) - for tensor_source in TensorSource: - target_scaling_mode = ( - ScalingMode.NVFP4_2D_SCALING - if (not test.disable_2d_quantization) and tensor_source == TensorSource.KERNEL - else ScalingMode.NVFP4_1D_SCALING - ) - self.assertEqual( - get_quantize_config().get_scaling_mode(tensor_source), target_scaling_mode - ) - self.assertEqual( - get_quantize_config().DISABLE_STOCHASTIC_ROUNDING, test.disable_stochastic_rounding - ) - self.assertEqual(get_quantize_config().DISABLE_RHT, test.disable_rht) - self.assertEqual( - get_quantize_config().DISABLE_2D_QUANTIZATION, test.disable_2d_quantization - ) + @abstractmethod + def assert_context(self, ref_recipe, quantize_config): + """Asserts that the quantize_config matches the expected properties from the reference recipe when the recipe is used with an autocast context. - def _compare_nvfp4_scaling_quantizers(self, test): - """Check that the quantizers created have the expected stochastic rounding state and the state is preserved across VJP boundaries.""" + Args: + ref_recipe: The reference quantization recipe. + quantize_config: The quantization configuration to be checked. + """ + pass - def assertion_func(quantizer, tensor_source): - if test.disable_stochastic_rounding or tensor_source != TensorSource.DGRAD: - self.assertIsNone(quantizer.stochastic_rounding_rng_state) - else: - self.assertIsNotNone(quantizer.stochastic_rounding_rng_state) + @abstractmethod + def assert_quantizers(self, ref_recipe, quantizer, tensor_source): + """Asserts that the quantizer matches the expected properties from the reference recipe. The quantizers are created in a small test Flax module TestModule and passed through a VJP boundary to ensure correct reconstruction. + + Args: + ref_recipe: The reference quantization recipe. + quantizer: The quantizer to be checked. + tensor_source: The source of the tensor (e.g., KERNEL, X, DGRAD). + """ + pass - expected_rht = ( - quantizer.scaling_mode == ScalingMode.NVFP4_1D_SCALING - and quantizer.q_layout in {QuantizeLayout.ROWWISE_COLWISE, QuantizeLayout.COLWISE} - and not test.disable_rht + +class DelayedScalingRecipeAssertion(RecipeAssertionBase): + + def assert_context(self, ref_recipe, quantize_config): + assert quantize_config.MARGIN == ref_recipe.margin + assert quantize_config.FWD_DTYPE == _format2dtypes(ref_recipe.fp8_format)[0] + assert quantize_config.BWD_DTYPE == _format2dtypes(ref_recipe.fp8_format)[1] + assert quantize_config.AMAX_HISTORY_LEN == ref_recipe.amax_history_len + assert quantize_config.AMAX_COMPUTE_ALGO.value == ref_recipe.amax_compute_algo + for tensor_source in TensorSource: + assert ( + quantize_config.get_scaling_mode(tensor_source) + == ScalingMode.DELAYED_TENSOR_SCALING ) - self.assertEqual(quantizer.use_rht, expected_rht) - x = jnp.ones((), dtype=jnp.float32) - test_module = TestModule(assertion_func=assertion_func) - param_key, sr_key = jax.random.split(jax.random.PRNGKey(0)) - rngs = {"params": param_key, "sr_rng": sr_key} - variables = test_module.init(rngs, x) + def assert_quantizers(self, ref_recipe: DelayedScaling, quantizer, tensor_source): + assert quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING + assert quantizer.margin == ref_recipe.margin + assert quantizer.amax_compute_algo.value == ref_recipe.amax_compute_algo + assert quantizer.amax_history.shape == (ref_recipe.amax_history_len,) + assert_fp8_format(quantizer, tensor_source, ref_recipe.fp8_format) - jax.jit(jax.value_and_grad(test_module.apply), static_argnums=(2,))(variables, x, rngs=rngs) - @unittest.skipIf(not is_fp8_supported, reason=reason) - def test_autocast_delayed_scaling(self): - self._check_default_state() +class CurrentScalingRecipeAssertion(RecipeAssertionBase): - with autocast(enabled=False, recipe=DelayedScaling(), mesh_resource=MeshResource()): - self._check_default_state() + def assert_context(self, ref_recipe, quantize_config): + assert quantize_config.FWD_DTYPE == _format2dtypes(ref_recipe.fp8_format)[0] + assert quantize_config.BWD_DTYPE == _format2dtypes(ref_recipe.fp8_format)[1] + for tensor_source in TensorSource: + assert ( + quantize_config.get_scaling_mode(tensor_source) + == ScalingMode.CURRENT_TENSOR_SCALING + ) - self._check_default_state() + def assert_quantizers(self, ref_recipe: Float8CurrentScaling, quantizer, tensor_source): + assert quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING + assert_fp8_format(quantizer, tensor_source, ref_recipe.fp8_format) - ds = DelayedScaling(margin=5.0, fp8_format=FP8Format.E4M3, amax_history_len=1) - with autocast(enabled=True, recipe=ds, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_delay_scaling(ds) - self._check_default_state() +class MXFP8RecipeAssertion(RecipeAssertionBase): - ds = DelayedScaling(margin=3.0, fp8_format=FP8Format.HYBRID, amax_history_len=1) - with autocast(enabled=True, recipe=ds, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_delay_scaling(ds) + def assert_context(self, ref_recipe, quantize_config): + assert quantize_config.FWD_DTYPE == _format2dtypes(ref_recipe.fp8_format)[0] + assert quantize_config.BWD_DTYPE == _format2dtypes(ref_recipe.fp8_format)[1] + for tensor_source in TensorSource: + assert quantize_config.get_scaling_mode(tensor_source) == ScalingMode.MXFP8_1D_SCALING - self._check_default_state() + def assert_quantizers(self, ref_recipe: MXFP8BlockScaling, quantizer, tensor_source): + assert quantizer.scaling_mode == ScalingMode.MXFP8_1D_SCALING + assert_fp8_format(quantizer, tensor_source, ref_recipe.fp8_format) - @unittest.skipIf(not is_fp8_supported, reason=reason) - def test_autocast_current_scaling(self): - self._check_default_state() - with autocast(enabled=False, recipe=Float8CurrentScaling(), mesh_resource=MeshResource()): - self._check_default_state() +class NVFP4RecipeAssertion(RecipeAssertionBase): - self._check_default_state() + def assert_context(self, ref_recipe, quantize_config): + assert quantize_config.FWD_DTYPE == _format2dtypes(ref_recipe.fp4_format)[0] + assert quantize_config.BWD_DTYPE == _format2dtypes(ref_recipe.fp4_format)[1] + for tensor_source in TensorSource: + target_scaling_mode = ( + ScalingMode.NVFP4_2D_SCALING + if (not ref_recipe.disable_2d_quantization) and tensor_source == TensorSource.KERNEL + else ScalingMode.NVFP4_1D_SCALING + ) + assert quantize_config.get_scaling_mode(tensor_source) == target_scaling_mode + assert quantize_config.DISABLE_STOCHASTIC_ROUNDING == ref_recipe.disable_stochastic_rounding + assert quantize_config.DISABLE_RHT == ref_recipe.disable_rht + assert quantize_config.DISABLE_2D_QUANTIZATION == ref_recipe.disable_2d_quantization + + def assert_quantizers(self, ref_recipe: NVFP4BlockScaling, quantizer, tensor_source): + if tensor_source == TensorSource.KERNEL and not ref_recipe.disable_2d_quantization: + assert quantizer.scaling_mode == ScalingMode.NVFP4_2D_SCALING + else: + assert quantizer.scaling_mode == ScalingMode.NVFP4_1D_SCALING + + if ref_recipe.disable_stochastic_rounding or tensor_source != TensorSource.DGRAD: + assert quantizer.stochastic_rounding_rng_state is None + else: + assert quantizer.stochastic_rounding_rng_state is not None + + expected_rht = ( + quantizer.scaling_mode == ScalingMode.NVFP4_1D_SCALING + and quantizer.q_layout in {QuantizeLayout.ROWWISE_COLWISE, QuantizeLayout.COLWISE} + and not ref_recipe.disable_rht + ) + assert quantizer.use_rht == expected_rht - cs = Float8CurrentScaling(fp8_format=FP8Format.E4M3) - with autocast(enabled=True, recipe=cs, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_current_scaling(cs) - self._check_default_state() +class TestFP8Functions(unittest.TestCase): - cs = Float8CurrentScaling(fp8_format=FP8Format.HYBRID) - with autocast(enabled=True, recipe=cs, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_current_scaling(cs) + def _check_default_state(self): + self.assertEqual(get_global_quantize_recipe(), None) - self._check_default_state() + def _test_recipe(self, quantization_recipe: Recipe, cls: RecipeAssertionBase): + """Tests a quantization recipe by verifying its behavior in both autocast and direct application contexts.""" + assert_context_func = cls().assert_context + assert_quantizer_func = partial(cls().assert_quantizers, quantization_recipe) + self._test_recipe_autocast(quantization_recipe, assert_context_func, assert_quantizer_func) + self._test_recipe_direct(quantization_recipe, assert_quantizer_func) - @unittest.skipIf(not is_mxfp8_supported, reason=mxfp8_reason) - def test_autocast_mxfp8_block_scaling(self): + def _test_recipe_autocast( + self, quantization_recipe, assert_context_func, assert_quantizer_func + ): + """Tests a quantization recipe within an autocast context by verifying the quantize config and quantizers in a test module.""" self._check_default_state() - - with autocast(enabled=False, recipe=MXFP8BlockScaling(), mesh_resource=MeshResource()): + with autocast(enabled=False, recipe=quantization_recipe, mesh_resource=MeshResource()): self._check_default_state() + with autocast(enabled=True, recipe=quantization_recipe, mesh_resource=MeshResource()): + quantize_config = self._get_global_quantize_config() + assert_context_func(quantization_recipe, quantize_config) + self._test_quantizer_in_model(assert_quantizer_func) + self._check_default_state() + def _test_recipe_direct(self, quantization_recipe, assert_quantizer_func): + """Tests a quantization recipe by directly passing it to a test module and verifying the quantizers.""" + self._check_default_state() + self._test_quantizer_in_model(assert_quantizer_func, direct_recipe=quantization_recipe) self._check_default_state() - bs = MXFP8BlockScaling() - with autocast(enabled=True, recipe=bs, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_mxfp8_scaling(bs) + def _test_quantizer_in_model(self, assert_quantizer_func, direct_recipe=None): + """Tests that the quantizers created in a test module match the expected properties by passing them through a VJP boundary. - self._check_default_state() + Args: + assert_quantizer_func: A function that asserts the properties of the quantizers. The function signature is (quantizer: Quantizer, tensor_source: TensorSource) -> None. + direct_recipe: An optional quantization recipe to be passed directly to the test module. This is an alternative API to using autocast contexts. + """ + x = jnp.ones((), dtype=jnp.float32) + test_module = TestModule(assertion_func=assert_quantizer_func, direct_recipe=direct_recipe) + param_key, sr_key = jax.random.split(jax.random.PRNGKey(0)) + rngs = {"params": param_key, "sr_rng": sr_key} + variables = test_module.init(rngs, x) - @unittest.skipIf(not is_nvfp4_supported, reason=nvfp4_reason) - def test_autocast_nvfp4_block_scaling(self): - self._check_default_state() + jax.jit(jax.value_and_grad(test_module.apply), static_argnums=(2,))(variables, x, rngs=rngs) - with autocast(enabled=False, recipe=NVFP4BlockScaling(), mesh_resource=MeshResource()): - self._check_default_state() + def _get_global_quantize_config(self): + quantization_recipe = get_global_quantize_recipe() + assert quantization_recipe is not None, "No global quantization recipe set" + quantize_config = get_quantize_config_with_recipe(quantization_recipe) + assert ( + quantize_config.is_fp8_enabled() + ), "Quantization not enabled in global quantize config" + return quantize_config - self._check_default_state() + @unittest.skipIf(not is_fp8_supported, reason=reason) + def test_autocast_delayed_scaling(self): + self._test_recipe( + quantization_recipe=DelayedScaling(), + cls=DelayedScalingRecipeAssertion, + ) + self._test_recipe( + quantization_recipe=DelayedScaling( + margin=5.0, fp8_format=FP8Format.E4M3, amax_history_len=1 + ), + cls=DelayedScalingRecipeAssertion, + ) + self._test_recipe( + quantization_recipe=DelayedScaling( + margin=3.0, fp8_format=FP8Format.HYBRID, amax_history_len=1 + ), + cls=DelayedScalingRecipeAssertion, + ) - bs = NVFP4BlockScaling() - with autocast(enabled=True, recipe=bs, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_nvfp4_scaling(bs) - self._compare_nvfp4_scaling_quantizers(bs) + @unittest.skipIf(not is_fp8_supported, reason=reason) + def test_autocast_current_scaling(self): + self._test_recipe( + quantization_recipe=Float8CurrentScaling(), + cls=CurrentScalingRecipeAssertion, + ) + self._test_recipe( + quantization_recipe=Float8CurrentScaling(margin=5.0, fp8_format=FP8Format.E4M3), + cls=CurrentScalingRecipeAssertion, + ) + self._test_recipe( + quantization_recipe=Float8CurrentScaling(margin=3.0, fp8_format=FP8Format.HYBRID), + cls=CurrentScalingRecipeAssertion, + ) - bs = NVFP4BlockScaling( - disable_stochastic_rounding=True, - disable_rht=True, - disable_2d_quantization=True, + @unittest.skipIf(not is_mxfp8_supported, reason=mxfp8_reason) + def test_autocast_mxfp8_block_scaling(self): + self._test_recipe( + quantization_recipe=MXFP8BlockScaling(), + cls=MXFP8RecipeAssertion, ) - with autocast(enabled=True, recipe=bs, mesh_resource=MeshResource()): - self.assertTrue(get_quantize_config().is_fp8_enabled()) - self._compare_nvfp4_scaling(bs) - self._compare_nvfp4_scaling_quantizers(bs) - self._check_default_state() + @unittest.skipIf(not is_nvfp4_supported, reason=nvfp4_reason) + def test_autocast_nvfp4_block_scaling(self): + self._test_recipe( + quantization_recipe=NVFP4BlockScaling(), + cls=NVFP4RecipeAssertion, + ) + self._test_recipe( + quantization_recipe=NVFP4BlockScaling( + disable_stochastic_rounding=True, + disable_rht=True, + disable_2d_quantization=True, + ), + cls=NVFP4RecipeAssertion, + ) class TestJaxprAndHlo: diff --git a/transformer_engine/jax/cpp_extensions/base.py b/transformer_engine/jax/cpp_extensions/base.py index 96b73909e1..556b587191 100644 --- a/transformer_engine/jax/cpp_extensions/base.py +++ b/transformer_engine/jax/cpp_extensions/base.py @@ -221,7 +221,7 @@ def manage_primitives(enable_names=None, disable_names=None, disable_all_first=F """ Helper function to manage primitive states by name without modifying environment variables. Allows enabling specific primitives, disabling specific primitives, or disabling all primitives. - This helper is used in the get_quantize_config().initialize() methods. + This helper is used in the get_quantize_config_with_recipe().initialize() methods. Args: enable_names: List of strings, each representing the name of a primitive class to enable. Defaults to None. diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 9ffec2c6af..c00b816f2e 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -38,12 +38,13 @@ ScalingMode, Quantizer, GroupedQuantizer, - get_quantize_config, QuantizerSet, QuantizeLayout, noop_quantizer_set, is_fp8_gemm_with_all_layouts_supported, apply_padding_to_scale_inv, + get_quantize_config_with_recipe, + get_global_quantize_recipe, ) from .misc import get_padded_spec, is_all_reduce_in_float32 from ..sharding import ( @@ -1246,7 +1247,7 @@ def _te_gemm( fuse_bias: bool = False, fuse_gelu: bool = False, grad: bool = False, - use_split_accumulator: bool = get_quantize_config().FP8_2X_ACC_FPROP, + use_split_accumulator: bool = None, transpose_batch_sequence: bool = False, collective_op: CollectiveOp = CollectiveOp.NONE, ) -> Tuple[jax.Array, ...]: @@ -1258,6 +1259,13 @@ def _te_gemm( DeprecationWarning, ) + if use_split_accumulator is None: + # TODO(jberchtold): Rework GEMM API to provide the context here instead of relying on global state and also + # use context of the GEMM type so we can decide between fprop, dgrad, and wgrad + use_split_accumulator = get_quantize_config_with_recipe( + get_global_quantize_recipe() + ).FP8_2X_ACC_FPROP + # Prepare non-quantized GEMM operands lhs_data = lhs rhs_data = rhs @@ -1720,10 +1728,15 @@ def _jax_gemm_impl(lhs, rhs): assert ( rhs.scaling_mode == lhs.scaling_mode ), f"rhs.scaling_mode={rhs.scaling_mode} != lhs.scaling_mode={lhs.scaling_mode}" + + # TODO(jberchtold): Rework GEMM API to provide the context here instead of relying on global state and also + # use context of the GEMM type so we can decide between fprop, dgrad, and wgrad + use_split_accumulator = get_quantize_config_with_recipe( + get_global_quantize_recipe() + ).FP8_2X_ACC_FPROP + precision = ( - jax.lax.Precision.HIGHEST - if get_quantize_config().FP8_2X_ACC_FPROP - else jax.lax.Precision.DEFAULT + jax.lax.Precision.HIGHEST if use_split_accumulator else jax.lax.Precision.DEFAULT ) return _jax_gemm_tensor_scaling_fp8(lhs, rhs, dim_nums, precision) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index a0e1a6406f..d16dab6d6c 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -820,7 +820,7 @@ def _quantize_dbias_impl( amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, ) - scale = compute_scale_from_amax(amax, quantizer.q_dtype) + scale = compute_scale_from_amax(amax, quantizer.q_dtype, margin=0.0) elif quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: scale = quantizer.scale # Make sure to reset amax to zeros for DelayedScaling @@ -1227,7 +1227,7 @@ def grouped_quantize( ) grouped_amax = jax.ops.segment_max(row_amax, segment_ids, num_segments=n_groups) for i in range(n_groups): - tmp_scale = compute_scale_from_amax(grouped_amax[i], quantizer.q_dtype) + tmp_scale = compute_scale_from_amax(grouped_amax[i], quantizer.q_dtype, margin=0.0) scale = scale.at[i].set(tmp_scale[0]) is_tensor_scaling = quantizer.scaling_mode in ( diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index c497775e01..613455b6c3 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -27,7 +27,6 @@ with_sharding_constraint_by_logical_axes, is_fp8_gemm_with_all_layouts_supported, TensorUsage, - get_quantize_config, ) @@ -95,7 +94,7 @@ def dense( if transpose_batch_sequence: warnings.warn("transpose_batch_sequence is not well tested, use with caution!") - if not get_quantize_config().is_fp8_enabled(): + if quantizer_set == noop_quantizer_set: input_dtype = x.dtype kernel = kernel.astype(input_dtype) diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 934af3d181..b5f1590229 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -33,10 +33,11 @@ ) from ..quantize import ( QuantizerFactory, - get_quantize_config, + get_global_quantize_recipe, QuantizeMetaSet, TensorSource, get_quantize_config_with_recipe, + noop_quantizer_set, ) PRNGKey = Any @@ -355,17 +356,17 @@ def generate_quantizer_set( Generate a set of FP8 meta for a GEMM. """ + if fp8_recipe is None: + fp8_recipe = get_global_quantize_recipe() + + quantize_config = get_quantize_config_with_recipe(fp8_recipe) + collection_name = ( variable_collection if variable_collection is not None - else get_quantize_config().COLLECTION_NAME + else quantize_config.COLLECTION_NAME ) - if fp8_recipe is None: - quantize_config = get_quantize_config() - else: - quantize_config = get_quantize_config_with_recipe(fp8_recipe) - x_meta = quantize_config.get_quantize_flax_meta( self, collection_name, postfix, TensorSource.X, "x" ) @@ -492,7 +493,11 @@ def __call__(self, inputs: Array) -> Array: self.dtype, ) - if not get_quantize_config().is_fp8_enabled(): + quantizer_set = self.generate_quantizer_set( + quantization_checkpoint_name=self.quantization_checkpoint_name + ) + + if quantizer_set == noop_quantizer_set: kernel = kernel.astype(input_dtype) if self.use_bias: @@ -505,9 +510,6 @@ def __call__(self, inputs: Array) -> Array: else: bias = None - quantizer_set = self.generate_quantizer_set( - quantization_checkpoint_name=self.quantization_checkpoint_name - ) contract_ind = tuple(range(0, len(axis))) y = dense( inputs, @@ -712,7 +714,7 @@ def __call__(self, inputs: Array) -> Array: ) fuse_layernorm = ( - get_quantize_config().is_fp8_enabled() + quantizer_set != noop_quantizer_set and not self.return_layernorm_output and self.enable_layernorm ) @@ -763,7 +765,7 @@ def __call__(self, inputs: Array) -> Array: kernel_shape, self.dtype, ) - if not get_quantize_config().is_fp8_enabled(): + if quantizer_set == noop_quantizer_set: kernel = kernel.astype(input_dtype) contract_ind = tuple(range(0, len(axis))) @@ -1042,7 +1044,7 @@ def __call__(self, inputs: Array, deterministic: bool = False) -> Array: # TODO(Phuong): use fuse_layernorm for high-precision # when NoOpQuantizer and Tensor are implemented fuse_layernorm = ( - get_quantize_config().is_fp8_enabled() + ffn1_quantizer_set != noop_quantizer_set and not self.return_layernorm_output and self.enable_layernorm ) @@ -1128,7 +1130,7 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args): self.dtype, ) - if not get_quantize_config().is_fp8_enabled(): + if ffn1_quantizer_set == noop_quantizer_set: kernel_1 = kernel_1.astype(input_dtype) hidden_size = inputs.shape[-1] @@ -1140,7 +1142,7 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args): kernel_2_shape, self.dtype, ) - if not get_quantize_config().is_fp8_enabled(): + if ffn2_quantizer_set == noop_quantizer_set: kernel_2 = kernel_2.astype(input_dtype) contract_ind = tuple(range(0, len(axis))) diff --git a/transformer_engine/jax/layernorm_dense.py b/transformer_engine/jax/layernorm_dense.py index b9482b7bde..14726553f2 100644 --- a/transformer_engine/jax/layernorm_dense.py +++ b/transformer_engine/jax/layernorm_dense.py @@ -23,7 +23,6 @@ noop_quantizer_set, with_sharding_constraint_by_logical_axes, TensorUsage, - get_quantize_config, ) @@ -73,7 +72,7 @@ def layernorm_dense( - Quantization is applied to both the normalized input and kernel """ - if not get_quantize_config().is_fp8_enabled(): + if quantizer_set == noop_quantizer_set: input_dtype = x.dtype kernel = kernel.astype(input_dtype) diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py index 2fd0f07d62..47fed6c3a7 100644 --- a/transformer_engine/jax/layernorm_mlp.py +++ b/transformer_engine/jax/layernorm_mlp.py @@ -28,7 +28,6 @@ QuantizerSet, noop_quantizer_set, TensorUsage, - get_quantize_config, ) @@ -114,7 +113,7 @@ def layernorm_mlp( not zero_centered_gamma ), "zero_centered_gamma is not supported if norm_type is 'rmsnorm'" - if not get_quantize_config().is_fp8_enabled(): + if quantizer_sets == (noop_quantizer_set, noop_quantizer_set): input_dtype = x.dtype kernel_1 = kernel_1.astype(input_dtype) kernel_2 = kernel_2.astype(input_dtype) diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index d5093e70e4..6358edf468 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -46,7 +46,7 @@ from .device_utils import get_device_compute_capability __all__ = [ - "get_quantize_config", + "get_global_quantize_recipe", "get_quantize_config_with_recipe", "autocast", "fp8_autocast", @@ -475,7 +475,12 @@ def get_quantize_flax_meta( (self.AMAX_HISTORY_LEN,), jnp.float32, ).value - return QuantizeMeta(scale=scale, amax_history=amax_history) + return QuantizeMeta( + margin=self.MARGIN, + amax_compute_algo=self.AMAX_COMPUTE_ALGO, + scale=scale, + amax_history=amax_history, + ) class CurrentScalingQuantizeConfig(BaseQuantizeConfig): @@ -669,14 +674,6 @@ def get_quantize_flax_meta( ) -_QUANTIZE_CONFIG = NoOpQuantizeConfig() - - -def get_quantize_config(): - """Global instance of BaseQuantizeConfig set by autocast context.""" - return _QUANTIZE_CONFIG - - def get_quantize_config_class( fp8_recipe: Recipe, ) -> Type[BaseQuantizeConfig]: @@ -687,6 +684,8 @@ def get_quantize_config_class( Returns: The quantization config class corresponding to the given recipe. """ + if fp8_recipe is None: + return NoOpQuantizeConfig if isinstance(fp8_recipe, DelayedScaling): return DelayedScalingQuantizeConfig if isinstance(fp8_recipe, MXFP8BlockScaling): @@ -701,10 +700,23 @@ def get_quantize_config_class( def get_quantize_config_with_recipe(fp8_recipe: Recipe): """Get the quantization configuration object based on the FP8 recipe.""" config = get_quantize_config_class(fp8_recipe)() - config.initialize_from_recipe(fp8_recipe) + if fp8_recipe is not None: + config.initialize_from_recipe(fp8_recipe) return config +_GLOBAL_RECIPE: Optional[Recipe] = None + + +def get_global_quantize_recipe() -> Optional[Recipe]: + """Get the global quantization recipe if set. + + Returns: + The global quantization recipe or None if not set. + """ + return _GLOBAL_RECIPE + + @contextmanager def autocast( enabled: bool = False, @@ -751,22 +763,21 @@ def autocast( if recipe is None: recipe = DelayedScaling() - global _QUANTIZE_CONFIG + global _GLOBAL_RECIPE - old_quantize_config = _QUANTIZE_CONFIG + old_global_recipe = _GLOBAL_RECIPE - _QUANTIZE_CONFIG = NoOpQuantizeConfig() + _GLOBAL_RECIPE = None try: with global_shard_guard(mesh_resource): if enabled: - _QUANTIZE_CONFIG = get_quantize_config_class(recipe)() - is_supported, reason = _QUANTIZE_CONFIG.is_supported() + _GLOBAL_RECIPE = recipe + is_supported, reason = get_quantize_config_class(_GLOBAL_RECIPE)().is_supported() assert is_supported, reason - _QUANTIZE_CONFIG.initialize_from_recipe(recipe) yield finally: - _QUANTIZE_CONFIG = old_quantize_config + _GLOBAL_RECIPE = old_global_recipe @contextmanager diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index adff317482..4edc187795 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -28,7 +28,7 @@ NoScaleTensor, ) from .helper import ( - get_quantize_config, + get_global_quantize_recipe, get_quantize_config_with_recipe, AmaxComputeAlgo, TensorSource, @@ -50,7 +50,7 @@ def compute_scale_from_amax( - amax: jnp.ndarray, q_dtype: jnp.dtype, scale: Optional[jnp.ndarray] = None + amax: jnp.ndarray, q_dtype: jnp.dtype, margin: float, scale: Optional[jnp.ndarray] = None ) -> jnp.ndarray: """Compute scale from amax value. @@ -64,7 +64,7 @@ def compute_scale_from_amax( fp8_max = jnp.astype(jnp.finfo(q_dtype).max, jnp.float32) if scale is None: scale = jnp.ones((1,)) - sf = (fp8_max / amax) / (2 ** get_quantize_config().MARGIN) + sf = (fp8_max / amax) / (2**margin) sf = jnp.where(amax > 0.0, sf, scale) sf = jnp.where(jnp.isfinite(amax), sf, scale) assert sf.shape == (1,), f"Expected sf.shape == (1,), but got {sf.shape}" @@ -223,6 +223,7 @@ class CurrentScaleQuantizer(Quantizer): Attributes: scaling_mode: Set to NVTE_DELAYED_TENSOR_SCALING q_layout: Quantization axis (default: ROWWISE_COLWISE) + data_layout: Data layout string (default: "NT") """ scaling_mode: ScalingMode = ScalingMode.CURRENT_TENSOR_SCALING @@ -254,8 +255,7 @@ def _quantize_func( compute_dtype = jnp.float32 dtype_max = (jnp.finfo(self.q_dtype).max).astype(compute_dtype) amax = x.amax or jnp.max(jnp.abs(x.data)).reshape((1,)) - fp8_max = jnp.astype(jnp.finfo(self.q_dtype).max, jnp.float32) - scale = (fp8_max / amax) / (2 ** get_quantize_config().MARGIN) + scale = compute_scale_from_amax(amax, self.q_dtype, margin=0.0) scaled_x = x.data.astype(compute_dtype) * scale clipped_scaled_x = jnp.clip(scaled_x, -dtype_max, dtype_max).astype(self.q_dtype) @@ -327,17 +327,23 @@ class DelayedScaleQuantizer(CurrentScaleQuantizer): Attributes: scaling_mode: Set to NVTE_DELAYED_TENSOR_SCALING q_layout: Quantization axis (default: ROWWISE_COLWISE) + data_layout: Data layout string (default: "NT") + margin: Margin value for scale computation + amax_compute_algo: Algorithm for computing amax scale: Current scaling factor amax_history: History of maximum absolute values """ - scaling_mode: ScalingMode = ScalingMode.DELAYED_TENSOR_SCALING - q_layout: QuantizeLayout = QuantizeLayout.ROWWISE_COLWISE + margin: float = 0.0 + amax_compute_algo: AmaxComputeAlgo = AmaxComputeAlgo.MAX scale: jnp.ndarray = field(default_factory=lambda: jnp.ones((1,), jnp.float32)) - amax_history: jnp.ndarray = field( - default_factory=lambda: jnp.zeros((get_quantize_config().AMAX_HISTORY_LEN,), jnp.float32) - ) + amax_history: jnp.ndarray = field(default_factory=lambda: jnp.zeros((1024,), jnp.float32)) + + def __post_init__(self): + assert self.margin is not None, "margin must be specified" + assert self.amax_compute_algo is not None, "amax_compute_algo must be specified" + assert self.amax_history is not None, "amax_history must be specified" def tree_flatten(self): """Flatten the quantizer for JAX tree operations. @@ -352,6 +358,8 @@ def tree_flatten(self): self.q_layout, self.data_layout, self.checkpoint_name, + self.margin, + self.amax_compute_algo, ) return (children, aux_data) @@ -407,12 +415,14 @@ def _update_amax_history(amax_history, new_amax): Returns: Updated AMAX history """ - amax_history = amax_history.at[0].set(new_amax[0]) + amax_history = amax_history.at[0].set(new_amax.reshape((1,))[0]) return amax_history @staticmethod - @partial(jax.jit, static_argnums=(2,)) - def _compute_scale(amax_history, scale, q_dtype): + @partial(jax.jit, static_argnums=(2, 3, 4)) + def _compute_scale( + amax_history, scale, q_dtype, amax_compute_algo: AmaxComputeAlgo, margin: float + ): """Compute new scale based on AMAX history. Args: @@ -424,12 +434,12 @@ def _compute_scale(amax_history, scale, q_dtype): Updated scale value """ # 2. Calculate the current scale - if get_quantize_config().AMAX_COMPUTE_ALGO is AmaxComputeAlgo.MAX: + if amax_compute_algo is AmaxComputeAlgo.MAX: amax = jnp.max(amax_history, axis=-1, keepdims=True) else: amax = amax_history[0:1] - return compute_scale_from_amax(amax, q_dtype, scale=scale) + return compute_scale_from_amax(amax, q_dtype, margin=margin, scale=scale) @staticmethod @jax.jit @@ -453,7 +463,9 @@ def update(self, new_amax: jnp.ndarray): new_amax: New maximum absolute value to add to history """ amax_history = self._update_amax_history(self.amax_history, new_amax) - self.scale = self._compute_scale(amax_history, self.scale, self.q_dtype) + self.scale = self._compute_scale( + amax_history, self.scale, self.q_dtype, self.amax_compute_algo, self.margin + ) self.amax_history = self._roll_and_reset_amax_history(amax_history) @@ -1124,6 +1136,7 @@ def _create_set( bwd_dtype, is_2x2x, n_groups, + is_inference_mode=False, checkpoint_name: Optional[str] = None, **kwargs, ) -> QuantizerSet: @@ -1137,6 +1150,7 @@ def _create_set( bwd_dtype: Data type for backward pass is_2x2x: Whether to use 2x2x quantization n_groups + is_inference_mode: Whether to create quantizers for inference mode. This option is not fully supported yet checkpoint_name: Optional name for checkpointing quantizations **kwargs: Additional arguments for quantizer initialization @@ -1149,7 +1163,7 @@ def _create_set( q_layout_x = q_layout_kernel = q_layout_dgrad = QuantizeLayout.ROWWISE if kernel_scaling_mode.is_1d_block_scaling(): q_layout_kernel = QuantizeLayout.COLWISE - if get_quantize_config().INFERENCE_MODE: + if is_inference_mode: q_layout_dgrad = None if "quantize_meta_set" in kwargs: @@ -1206,10 +1220,10 @@ def create_set( Args: n_quantizer_sets: Number of quantizer sets to create - scaling_mode: Scaling mode to use, default is get_quantize_config().get_scaling_mode - fwd_dtype: Data type for forward pass, default is get_quantize_config().FWD_DTYPE - bwd_dtype: Data type for backward pass, default is get_quantize_config().BWD_DTYPE - is_2x2x: Whether to use 2x2x quantization, default is get_quantize_config().IF_QUANTIZE_2X + scaling_mode: Scaling mode to use, default is get the scaling mode from the specified or global recipe + fwd_dtype: Data type for forward pass, default is get the fwd dtype from the specified or global recipe + bwd_dtype: Data type for backward pass, default is get the bwd dtype from the specified or global recipe + is_2x2x: Whether to use 2x2x quantization, default is determined based on the specified or global recipe n_groups: checkpoint_name: Optional name for checkpointing quantizations fp8_recipe: Recipe to use for quantization. Scaling mode can be specified directly via the scaling_mode parameter or indirectly via recipe. Recipe is preferred as it will support additional recipes in future where scaling mode differs between x, kernel, and grad in the quantizer set. @@ -1226,25 +1240,46 @@ def create_set( " scaling mode differs between x, kernel, and grad in the quantizer set." ) + # TODO(jberchtold): Currently this is a limitation because we only support automatically populating quantizer fields based on a given recipe when using Flax. In the generic quantizer logic, we cannot assume Flax is being used, so we require the user to provide the quantize_meta_set created by quantize_config.get_quantize_flax_meta() or the same data created by themselves if they are passing a recipe here directly. + assert ( + fp8_recipe is None or "quantize_meta_set" in kwargs + ), "When fp8_recipe is specified, quantize_meta_set must be provided in kwargs." + + if fp8_recipe is None: + fp8_recipe = get_global_quantize_recipe() + if fp8_recipe is not None: + assert scaling_mode is None, ( + "scaling_mode should not be specified when fp8_recipe is provided either directly" + " or through an autocast context." + ) + assert fwd_dtype is None, ( + "fwd_dtype should not be specified when fp8_recipe is provided either directly or" + " through an autocast context." + ) + assert bwd_dtype is None, ( + "bwd_dtype should not be specified when fp8_recipe is provided either directly or" + " through an autocast context." + ) quantize_config = get_quantize_config_with_recipe(fp8_recipe) x_scaling_mode = quantize_config.get_scaling_mode(TensorSource.X) kernel_scaling_mode = quantize_config.get_scaling_mode(TensorSource.KERNEL) grad_scaling_mode = quantize_config.get_scaling_mode(TensorSource.DGRAD) fwd_dtype = quantize_config.FWD_DTYPE bwd_dtype = quantize_config.BWD_DTYPE + is_inference_mode = quantize_config.INFERENCE_MODE else: if scaling_mode is not None: x_scaling_mode = scaling_mode kernel_scaling_mode = scaling_mode grad_scaling_mode = scaling_mode else: - x_scaling_mode = get_quantize_config().get_scaling_mode(TensorSource.X) - kernel_scaling_mode = get_quantize_config().get_scaling_mode(TensorSource.KERNEL) - grad_scaling_mode = get_quantize_config().get_scaling_mode(TensorSource.DGRAD) + # TODO(jberchtold): make a way to explicitly pass a no scaling recipe here if we need other quantization config attributes in the future since NoOpQuantizeConfig already exists, we just can't use it here with direct recipe passing because we cannot differentiate between fp8_recipe=None meaning no recipe specified vs explicitly no quantization desired. + x_scaling_mode = ScalingMode.NO_SCALING + kernel_scaling_mode = ScalingMode.NO_SCALING + grad_scaling_mode = ScalingMode.NO_SCALING + is_inference_mode = False - fwd_dtype = fwd_dtype or get_quantize_config().FWD_DTYPE - bwd_dtype = bwd_dtype or get_quantize_config().BWD_DTYPE if is_2x2x is None: # TODO(Jeremy): check x, kernel, grad separately for 2x if x_scaling_mode.is_1d_block_scaling(): @@ -1253,7 +1288,6 @@ def create_set( is_2x2x = not is_fp8_gemm_with_all_layouts_supported() else: # NO_SCALING ignores is_2x2x for now is_2x2x = False - is_inference_mode = get_quantize_config().INFERENCE_MODE assert not is_inference_mode, "Inference mode is not supported yet!" q_set = [] @@ -1267,6 +1301,7 @@ def create_set( bwd_dtype=bwd_dtype, is_2x2x=is_2x2x, n_groups=n_groups, + is_inference_mode=is_inference_mode, checkpoint_name=checkpoint_name, **kwargs, ) From c525760538b5cb1b77f3d93ab2c98d75b9453f43 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 14 Nov 2025 20:28:39 +0100 Subject: [PATCH 067/521] [PyTorch] Activation offloading refactor (#1762) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * init Signed-off-by: Pawel Gadzinski * offloading Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * all types Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * typo Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * init Signed-off-by: Pawel Gadzinski * api change Signed-off-by: Pawel Gadzinski * code drop Signed-off-by: Pawel Gadzinski * refactor Signed-off-by: Pawel Gadzinski * tests Signed-off-by: Pawel Gadzinski * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * example Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * cpu offload + debug warning Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * change empty_like implementation to use make_like Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * main_grad fix Signed-off-by: Pawel Gadzinski * manual synchornization Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * old path Signed-off-by: Pawel Gadzinski * remove example Signed-off-by: Pawel Gadzinski * api changes Signed-off-by: Pawel Gadzinski * reverted grouped linear Signed-off-by: Pawel Gadzinski * make odl code path work for modules Signed-off-by: Pawel Gadzinski * attention old code path Signed-off-by: Pawel Gadzinski * legacy tests Signed-off-by: Pawel Gadzinski * legacy tests Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * updated code path Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/pytorch/tensor/quantized_tensor.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nvfp4 support Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update tests/pytorch/test_cpu_offloading.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * small fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docs change Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- qa/L0_pytorch_unittest/test.sh | 3 +- tests/pytorch/test_cpu_offloading.py | 884 ++++++++--- tests/pytorch/test_cpu_offloading_v1.py | 215 +++ .../dot_product_attention/backends.py | 69 +- .../dot_product_attention.py | 8 - .../pytorch/attention/multi_head_attention.py | 5 +- transformer_engine/pytorch/cpu_offload.py | 1354 +++++++++-------- transformer_engine/pytorch/cpu_offload_v1.py | 743 +++++++++ .../pytorch/module/grouped_linear.py | 8 +- .../pytorch/module/layernorm_linear.py | 19 +- .../pytorch/module/layernorm_mlp.py | 21 +- transformer_engine/pytorch/module/linear.py | 11 +- .../pytorch/optimizers/fused_adam.py | 4 +- .../pytorch/quantized_tensor.py | 96 +- .../pytorch/tensor/float8_blockwise_tensor.py | 10 +- .../pytorch/tensor/float8_tensor.py | 36 +- .../pytorch/tensor/mxfp8_tensor.py | 36 +- .../pytorch/tensor/nvfp4_tensor.py | 45 +- .../float8_blockwise_tensor_storage.py | 7 + .../tensor/storage/float8_tensor_storage.py | 9 + .../tensor/storage/mxfp8_tensor_storage.py | 7 + 21 files changed, 2714 insertions(+), 876 deletions(-) create mode 100644 tests/pytorch/test_cpu_offloading_v1.py create mode 100644 transformer_engine/pytorch/cpu_offload_v1.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index b23ce3b6cf..e1ce680094 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -42,7 +42,8 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml 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_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" -NVTE_FLASH_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" +NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 64da83a210..c5b4b48b67 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -2,27 +2,41 @@ # # See LICENSE for license information. +import random import contextlib -import gc -import os -from typing import Iterable, Optional - import pytest +import os import torch - +from typing import Optional, List +from transformer_engine.pytorch.cpu_offload import ( + get_cpu_offload_context, + OffloadableLayerState, + DefaultOffloadSynchronizer, + start_offload, + mark_not_offload, +) +from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends -from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported -from utils import ModelConfig, get_available_attention_backends +from utils import ModelConfig +import transformer_engine_torch as tex # Check supported quantization schemes -fp8_available = te.is_fp8_available() -mxfp8_available = te.is_mxfp8_available() +fp8_available, _ = FP8GlobalStateManager.is_fp8_available() +fp8_block_scaling_available, _ = FP8GlobalStateManager.is_fp8_block_scaling_available() +mxfp8_available, _ = FP8GlobalStateManager.is_mxfp8_available() +nvfp4_available, _ = FP8GlobalStateManager.is_nvfp4_available() -quantization_recipes: Optional[recipe.Recipe] = [None] +quantization_recipes: List[Optional[recipe.Recipe]] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) +if fp8_block_scaling_available: + quantization_recipes.append(recipe.Float8BlockScaling()) +if mxfp8_available: + quantization_recipes.append(recipe.MXFP8BlockScaling()) +if nvfp4_available: + quantization_recipes.append(recipe.NVFP4BlockScaling()) + model_config = { "small": ModelConfig(8, 512, 8, 64, num_layers=5, eps=0.1), @@ -32,181 +46,709 @@ NUM_LAYERS = model_config["small"].num_layers EPSILON = model_config["small"].eps -# Flash attention saves some internal tensor for the backward pass -# that cannot be offloaded to CPU. -assert os.getenv("NVTE_FLASH_ATTN") == "0" +# Disable garbage collection to tests if there are reference cycles. +# We do not want them, because they can result in CUDA out of memory errors. +import gc -# Offloading is supported for attention only for fused and flash attention backends, -# so the use of bfloat16 is required. -# -# For the TransformerLayer, activation offloading with dropout is not supported, -# so we set hidden_dropout to 0.0. -model_types = { - "linear": lambda: te.Linear(SIZE, SIZE, params_dtype=torch.bfloat16), - "layernorm_mlp": lambda: te.LayerNormMLP(SIZE, SIZE, params_dtype=torch.bfloat16), - "layernorm_linear": lambda: te.LayerNormLinear(SIZE, SIZE, params_dtype=torch.bfloat16), - "multihead_attention": lambda: te.MultiheadAttention( - SIZE, NUM_HEADS, params_dtype=torch.bfloat16 - ), - "transformer_layer": lambda: te.TransformerLayer( - SIZE, SIZE, NUM_HEADS, params_dtype=torch.bfloat16, hidden_dropout=0.0 - ), - "linear_op": lambda: te.ops.Linear(SIZE, SIZE, dtype=torch.bfloat16), - "layernorm_mlp_ops": lambda: te.ops.Sequential( - te.ops.LayerNorm(SIZE, dtype=torch.bfloat16), - te.ops.Linear(SIZE, SIZE, dtype=torch.bfloat16), - te.ops.GELU(), - te.ops.Linear(SIZE, SIZE, dtype=torch.bfloat16), - ), -} +gc.disable() + + +class Utils: + tensor1 = torch.randn((1024, 1024), device="cuda", dtype=torch.bfloat16) + _B = 64 + _S = 256 + _H = 4 + _D = 256 + + @staticmethod + def long_job(stream: Optional[torch.cuda.Stream] = None): + NUM_ITERS = 6000 + if stream is None: + stream = torch.cuda.current_stream() + + with torch.cuda.stream(stream): + for i in range(NUM_ITERS): + Utils.tensor1.normal_() + + @staticmethod + def measure_time(func): + import time + + torch.cuda.synchronize() + start = time.time() + func() + torch.cuda.synchronize() + end = time.time() + return (end - start) * 1000 + + @staticmethod + def get_cuda_memory_mb(): + return torch.cuda.memory_allocated() / (1024**2) + + @staticmethod + def get_max_cuda_memory_mb(): + return torch.cuda.max_memory_allocated() / (1024**2) + + @staticmethod + def get_cpu_memory_mb() -> float: + import psutil, os + + return psutil.Process(os.getpid()).memory_info().rss / (1024**2) + + @staticmethod + def get_layer_names(): + return [ + "linear", + "layernorm_linear", + "layernorm_mlp", + "grouped_linear", + "multihead_attention", + "transformer_layer", + "linear_op", + "layernorm_mlp_ops", + ] + + @staticmethod + def create_layer(layer_type: str): + if layer_type == "linear": + return te.Linear(Utils._D, Utils._D, params_dtype=torch.bfloat16) + elif layer_type == "layernorm_linear": + return te.LayerNormLinear(Utils._D, Utils._D, params_dtype=torch.bfloat16) + elif layer_type == "layernorm_mlp": + return te.LayerNormMLP(Utils._D, Utils._D, params_dtype=torch.bfloat16) + elif layer_type == "multihead_attention": + return te.MultiheadAttention( + Utils._D, Utils._H, attention_dropout=0.0, params_dtype=torch.bfloat16 + ) + elif layer_type == "grouped_linear": + return te.GroupedLinear(Utils._H, Utils._D, Utils._D, params_dtype=torch.bfloat16) + elif layer_type == "transformer_layer": + return te.TransformerLayer( + Utils._D, + Utils._D, + Utils._H, + attention_dropout=0.0, + hidden_dropout=0.0, + params_dtype=torch.bfloat16, + ) + elif layer_type == "linear_op": + return te.ops.Linear(Utils._D, Utils._D, dtype=torch.bfloat16) + elif layer_type == "layernorm_mlp_ops": + return te.ops.Sequential( + te.ops.LayerNorm(Utils._D, dtype=torch.bfloat16), + te.ops.Linear(Utils._D, Utils._D, dtype=torch.bfloat16), + te.ops.GELU(), + te.ops.Linear(Utils._D, Utils._D, dtype=torch.bfloat16), + ) + else: + raise ValueError(f"Unknown layer type: {layer_type}") + + @staticmethod + def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) -> torch.Tensor: + shape = (Utils._B, Utils._S, Utils._D) + tensor = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + if recipe is None: + tensor = tensor.requires_grad_() if requires_grad else tensor + return tensor + elif recipe.delayed(): + quantizer = te.tensor.float8_tensor.Float8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + scale=torch.tensor([1.0], device="cuda"), + amax=torch.tensor([1.0], device="cuda"), + ) + return quantizer(tensor) + elif recipe.float8_current_scaling(): + quantizer = te.tensor.float8_tensor.Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, device="cuda" + ) + return quantizer(tensor) + elif recipe.float8_block_scaling(): + quantizer = te.tensor.float8_blockwise_tensor.Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True + ) + return quantizer(tensor) + elif recipe.mxfp8(): + quantizer = te.tensor.mxfp8_tensor.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + return quantizer(tensor) + elif recipe.nvfp4(): + quantizer = te.tensor.nvfp4_tensor.NVFP4Quantizer() + return quantizer(tensor) + + @staticmethod + def create_recipe_ctx(recipe: Optional[recipe.Recipe]): + if recipe is None: + return lambda: contextlib.nullcontext() + else: + return lambda: te.fp8_autocast(fp8_recipe=recipe) + + @staticmethod + 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()) + else: + return tensor.numel() * tensor.element_size() / (1024**2) + + @staticmethod + def memory_leak_check(): + # Should be called before each test. + # Only cublas workspaces and some global tensors are allowed to be allocated. + # All other allocations should be released. + # This is a simple check to catch memory leaks. + if Utils.get_cuda_memory_mb() > 1000: + memory_num = Utils.get_cuda_memory_mb() + import gc + + gc.collect() # We want next test to be run with clean state. + gc.disable() + raise RuntimeError(f"Memory leak: {memory_num} MB") + + +class TestsOffloadableLayerState: + @pytest.mark.parametrize("random_num_tensors", [True, False]) + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_general(self, random_num_tensors, recipe): + """ + Test general functionality of DefaultOffloadSynchronizer - offload NUM_LAYERS-1 out of NUM_LAYERS layers, + for each layer offload random number of random tensors. + Then do backward pass for each layer, and check if reloaded tensors are equal to original tensors. + """ + Utils.memory_leak_check() + NUM_ITERATIONS = 10 + + stream = torch.cuda.Stream() + + offload_layer_state = OffloadableLayerState( + offload_stream=stream, + ) + for _ in range(NUM_ITERATIONS): + original_tensors = [] + tensors_ids = [] + NUM_TENSORS = random.choice([1, 20]) if random_num_tensors else 1 + for _ in range(NUM_TENSORS): + tensor = Utils.create_tensor(recipe) + original_tensors.append(tensor) + tensor_id = offload_layer_state.push_tensor(tensor) + assert tensor.device.type == "cuda" + tensors_ids.append(tensor_id) + + offload_layer_state.start_offload() + offload_layer_state.release_activation_forward_gpu_memory() + offload_layer_state.start_reload() + + for j in range(len(tensors_ids)): + tensor_gpu = offload_layer_state.pop_tensor(tensors_ids[j]) + assert tensor_gpu.device.type == "cuda" + assert tensor_gpu.shape == original_tensors[j].shape + assert tensor_gpu.dtype == original_tensors[j].dtype + torch.testing.assert_close(tensor_gpu, original_tensors[j]) + offload_layer_state.release_all_memory() + torch.cuda.synchronize() + + def test_offload_base_tensor(self): + Utils.memory_leak_check() + stream = torch.cuda.Stream() + offload_layer_state = OffloadableLayerState( + offload_stream=stream, + ) + init_cuda_memory = Utils.get_cuda_memory_mb() + x = Utils.create_tensor(None) + x_size = Utils.get_tensor_size_mb(x) + x_1 = x[::2] + x_2 = x[1::2] + + start_offload(x_1, offload_base_tensor=True) + start_offload(x_2, offload_base_tensor=True) + x1_id = offload_layer_state.push_tensor(x_1) + x2_id = offload_layer_state.push_tensor(x_2) + del x_1, x_2 + offload_layer_state.start_offload() + offload_layer_state.release_activation_forward_gpu_memory() + + assert offload_layer_state.get_offloaded_total_size_mb() == pytest.approx(x_size, 0.1) + + offload_layer_state.start_reload() + x_1 = offload_layer_state.pop_tensor(x1_id) + x_2 = offload_layer_state.pop_tensor(x2_id) + assert x_1.device.type == "cuda" + assert x_2.device.type == "cuda" + + assert torch.allclose(x_1, x[::2]) + assert torch.allclose(x_2, x[1::2]) + del x + + assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory + x_size, 0.1) + + +class TestsDefaultOffloadSynchronizer: + @pytest.mark.parametrize("random_num_tensors", [True, False]) + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_general(self, random_num_tensors, recipe): + """ + Test general functionality of DefaultOffloadSynchronizer - offload NUM_LAYERS-1 out of NUM_LAYERS layers, + for each layer offload random number of random tensors. + Then do backward pass for each layer, and check if reloaded tensors are equal to original tensors. + """ + Utils.memory_leak_check() + NUM_LAYERS = 10 + NUM_ITERATIONS = 10 + + offload_synchronizer = DefaultOffloadSynchronizer( + num_layers=NUM_LAYERS, + num_offloaded_layers=NUM_LAYERS - 1, + ) + + for _ in range(NUM_ITERATIONS): + original_tensors = [] + tensors_ids = [] + layer_ids = [] + + for i in range(NUM_LAYERS): + NUM_LAYER_TENSORS = random.randint(1, 10) if random_num_tensors else 1 + layer_tensors = [] + layer_tensors_ids = [] + layer_id = offload_synchronizer.fwd_step() + for _ in range(NUM_LAYER_TENSORS): + tensor = Utils.create_tensor(recipe) + layer_tensors.append(tensor) + tensor_id = offload_synchronizer.push_tensor(tensor) + assert tensor.device.type == "cuda" + layer_tensors_ids.append(tensor_id) + layer_ids.append(layer_id) + tensors_ids.append(layer_tensors_ids) + original_tensors.append(layer_tensors) + for i in range(NUM_LAYERS - 1, -1, -1): + offload_synchronizer.bwd_step(layer_ids[i]) + for j in range(len(tensors_ids[i])): + tensor_gpu = offload_synchronizer.pop_tensor(tensors_ids[i][j]) + assert tensor_gpu.device.type == "cuda" + assert tensor_gpu.shape == original_tensors[i][j].shape + assert tensor_gpu.dtype == original_tensors[i][j].dtype + torch.testing.assert_close(tensor_gpu, original_tensors[i][j]) + offload_synchronizer.finish_part_of_bwd() + torch.cuda.synchronize() + + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_memory(self, recipe): + torch.cuda.synchronize() + Utils.memory_leak_check() + NUM_LAYERS = 10 + + torch.cuda.reset_peak_memory_stats() + + offload_synchronizer = DefaultOffloadSynchronizer( + num_layers=NUM_LAYERS, + num_offloaded_layers=NUM_LAYERS - 1, + ) -def _make_input() -> torch.Tensor: - """Generate random input tensor.""" - return torch.randn( - (128, SIZE, SIZE), - dtype=torch.bfloat16, - device="cuda", - requires_grad=True, - ) - - -def _warmup_model( - modules: Iterable[torch.nn.Module], - quantization_recipe: Optional[recipe.Recipe], -) -> None: - """Perform forward and backward pass""" - tensor = _make_input() - for module in modules: - with te.autocast( - enabled=quantization_recipe is not None, - recipe=quantization_recipe, + init_cuda_memory = Utils.get_cuda_memory_mb() + + tensor_ids = [] + + torch.cuda.synchronize() + for _ in range(NUM_LAYERS): + offload_synchronizer.fwd_step() + tensor = Utils.create_tensor(recipe) + tensor_size = Utils.get_tensor_size_mb(tensor) + tensor_id = offload_synchronizer.push_tensor(tensor) + assert tensor.device.type == "cuda" + tensor_ids.append(tensor_id) + del tensor, tensor_id + torch.cuda.synchronize() + + if recipe is None: + assert Utils.get_max_cuda_memory_mb() == pytest.approx( + init_cuda_memory + tensor_size, 0.1 + ) + assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory + tensor_size, 0.1) + + for i in range(NUM_LAYERS - 1, -1, -1): + offload_synchronizer.bwd_step(i) + tensor_gpu = offload_synchronizer.pop_tensor(tensor_ids[i]) + assert tensor_gpu.device.type == "cuda" + del tensor_gpu, tensor_ids[i] + offload_synchronizer.finish_part_of_bwd() + + del tensor_ids + torch.cuda.synchronize() + + if recipe is None: + assert Utils.get_max_cuda_memory_mb() == pytest.approx( + init_cuda_memory + tensor_size, 0.1 + ) + assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) + + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_multiple_tensor_offload(self, recipe): + Utils.memory_leak_check() + init_cpu_memory = Utils.get_cpu_memory_mb() + init_cuda_memory = Utils.get_cuda_memory_mb() + offload_synchronizer = DefaultOffloadSynchronizer( + num_layers=2, + num_offloaded_layers=1, + ) + x1 = Utils.create_tensor(recipe) + x_size = Utils.get_tensor_size_mb(x1) + offload_synchronizer.fwd_step() + offload_synchronizer.push_tensor(x1) + offload_synchronizer.push_tensor(x1) + offload_synchronizer.push_tensor(x1) + offload_synchronizer.fwd_step() + # Only one copy of tensor on cpu is allocated. + assert Utils.get_cpu_memory_mb() == pytest.approx(init_cpu_memory + 1 * x_size, 0.1) + del x1 + offload_synchronizer.bwd_step(1) + offload_synchronizer.bwd_step(0) + offload_synchronizer.finish_part_of_bwd() + + assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) + + +class TestTELayers: + @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_sanity(self, layer_type, recipe): + Utils.memory_leak_check() + + # Skip ops-based layers with Float8BlockScaling recipe + if ( + layer_type in ["linear_op", "layernorm_mlp_ops"] + and recipe is not None + and recipe.float8_block_scaling() ): - tensor = module(tensor) - tensor.sum().backward() - - -def _estimate_cached_weight_size( - model_name: str, - modules: Iterable[torch.nn.Module], - quantization_recipe: Optional[recipe.Recipe], -) -> float: - """Calculate the memory (in MiB) needed for weight caching.""" - - # The weight params are cached directly for unquantized compute - if quantization_recipe is None: - return 0 - - # Count number of weight param elements - param_elements = 0 - for module in modules: - for param in module.parameters(): - if param.dim() == 2: - param_elements += param.numel() - - # FP8 tensor-scaling caches one byte per element - if quantization_recipe.delayed() or quantization_recipe.float8_current_scaling(): - if not is_non_tn_fp8_gemm_supported() and model_name not in ( - "linear_op", - "layernorm_mlp_ops", + pytest.skip("Fusible operations do not support FP8 block scaling recipe") + + recipe_ctx = Utils.create_recipe_ctx(recipe) + init_cuda_memory = Utils.get_cuda_memory_mb() + OFFLOAD_LAYERS = 6 + NUM_LAYERS = 10 + offload_ctx, sync_function = get_cpu_offload_context( + enabled=True, + num_layers=OFFLOAD_LAYERS, + model_layers=NUM_LAYERS, + ) + layers = [Utils.create_layer(layer_type) for _ in range(NUM_LAYERS)] + inp = Utils.create_tensor(None) + m_splits = ( + {"m_splits": [Utils._B * Utils._S // Utils._H] * Utils._H} + if layer_type == "grouped_linear" + else {} + ) + out = inp + for i in range(NUM_LAYERS): + with offload_ctx, recipe_ctx(): + # Ops-based layers don't support is_first_microbatch parameter + if layer_type in ["linear_op", "layernorm_mlp_ops"]: + out = layers[i](out, **m_splits) + else: + out = layers[i](out, is_first_microbatch=False, **m_splits) + out = sync_function(out) + out.sum().backward() + torch.cuda.synchronize() + del out, inp, layers + + @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_memory(self, layer_type, recipe): + Utils.memory_leak_check() + + # Skip ops-based layers with Float8BlockScaling recipe + if ( + layer_type in ["linear_op", "layernorm_mlp_ops"] + and recipe is not None + and recipe.float8_block_scaling() + ): + pytest.skip("Fusible operations do not support FP8 block scaling recipe") + + offload_ctx, sync_function = get_cpu_offload_context( + enabled=True, + num_layers=1, + model_layers=2, + offload_activations=True, + offload_weights=False, + ) + recipe_ctx = Utils.create_recipe_ctx(recipe) + layer = Utils.create_layer(layer_type) + inp = Utils.create_tensor(None) + + m_splits = ( + {"m_splits": [Utils._B * Utils._S // Utils._H] * Utils._H} + if layer_type == "grouped_linear" + else {} + ) + + # Ops-based layers don't support is_first_microbatch parameter + is_ops_layer = layer_type in ["linear_op", "layernorm_mlp_ops"] + + with recipe_ctx(): + if is_ops_layer: + out = layer(inp, **m_splits) + else: + out = layer(inp, is_first_microbatch=True, **m_splits) + out.sum().backward() + + del inp + init_cuda_memory = Utils.get_cuda_memory_mb() + + # run layer without offload + inp = Utils.create_tensor(None) + with recipe_ctx(): + if is_ops_layer: + out = layer(inp, **m_splits) + else: + out = layer(inp, is_first_microbatch=False, **m_splits) + with recipe_ctx(): + out = out + 1 + del inp + cuda_memory_no_offload = Utils.get_cuda_memory_mb() + + out.sum().backward() + # run layer with offload + inp = Utils.create_tensor(None) + with offload_ctx, recipe_ctx(): + if is_ops_layer: + out = layer(inp, **m_splits) + else: + out = layer(inp, is_first_microbatch=False, **m_splits) + out = sync_function(out) + with offload_ctx, recipe_ctx(): + out = out + 1 + out = sync_function(out) + del inp + assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) + offloaded_memory_cpu = offload_ctx.offload_synchronizer.get_offloaded_total_size_mb() + + # This assertion verifies that the memory used by tensors on the CPU matches the memory saved from a layer. + # It helps catch cases where an offloaded tensor still has a live pointer, which would + # cause an unnecessary copy to the CPU and prevent GPU memory from being released. + assert Utils.get_cuda_memory_mb() + offloaded_memory_cpu == pytest.approx( + cuda_memory_no_offload, 0.1 + ) + out.sum().backward() + + @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) + @pytest.mark.parametrize("recipe", quantization_recipes) + def test_manual_synchronization(self, recipe, layer_type): + Utils.memory_leak_check() + + # Skip ops-based layers with Float8BlockScaling recipe + if ( + layer_type in ["linear_op", "layernorm_mlp_ops"] + and recipe is not None + and recipe.float8_block_scaling() ): - # Modules do not deallocate FP8 transpose for weights - return 2 * param_elements / 1024**2 - return param_elements / 1024**2 + pytest.skip("Fusible operations do not support FP8 block scaling recipe") + + offload_ctx, sync_function, manual_controller = get_cpu_offload_context( + enabled=True, + model_layers=6, + offload_activations=True, + manual_synchronization=True, + ) + layer_1 = Utils.create_layer(layer_type) + layer_2 = Utils.create_layer(layer_type) + inp1 = Utils.create_tensor(None) + inp2 = Utils.create_tensor(None) - # MXFP8 caches one data byte per element and one scale byte per 32 - # elements - if quantization_recipe.mxfp8(): - if model_name not in ("linear_op", "layernorm_mlp_ops"): - # Modules do not deallocate column-wise MXFP8 data for weights - return 2 * param_elements * (1 + 1 / 32) / 1024**2 - return param_elements * (1 + 1 / 32) / 1024**2 + recipe_ctx = Utils.create_recipe_ctx(recipe) - raise NotImplementedError(f"Unrecognized recipe ({quantization_recipe})") + m_splits = ( + {"m_splits": [Utils._B * Utils._S // Utils._H] * Utils._H} + if layer_type == "grouped_linear" + else {} + ) + + init_cuda_memory = Utils.get_cuda_memory_mb() + + # 1 fwd + with offload_ctx, recipe_ctx(): + out_1 = layer_1(inp1, **m_splits) + out_1 = sync_function(out_1) + + with offload_ctx, recipe_ctx(): + out_2 = layer_2(inp2, **m_splits) + out_2 = sync_function(out_2) + + mark_not_offload(out_1, out_2) + + del inp1, inp2 + + memory_before_offload = Utils.get_cuda_memory_mb() + manual_controller.start_offload_layer(0) + manual_controller.release_activation_forward_gpu_memory(0) + manual_controller.start_offload_layer(1) + manual_controller.release_activation_forward_gpu_memory(1) + memory_after_offload = Utils.get_cuda_memory_mb() + assert memory_after_offload + EPSILON < memory_before_offload + + manual_controller.start_reload_layer(0) + manual_controller.start_reload_layer(1) + + memory_after_reload = Utils.get_cuda_memory_mb() + assert memory_after_reload == pytest.approx(memory_before_offload, 0.1) + + out_1.sum().backward() + out_2.sum().backward() + + @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) + @pytest.mark.parametrize("use_cuda_graphs", [True, False]) + @pytest.mark.parametrize("retain_pinned_cpu_buffers", [True, False]) + @pytest.mark.parametrize("backend", ["FlashAttention", "FusedAttention", "UnfusedAttention"]) + def test_numerics( + self, + recipe, + layer_type, + use_cuda_graphs, + backend, + retain_pinned_cpu_buffers, + ): + # Skip ops-based layers with Float8BlockScaling recipe + if ( + layer_type in ["linear_op", "layernorm_mlp_ops"] + and recipe is not None + and recipe.float8_block_scaling() + ): + pytest.skip("Fusible operations do not support FP8 block scaling recipe") + recipe_ctx = Utils.create_recipe_ctx(recipe) -def _measure_cached_memory( - modules: Iterable[torch.nn.Module], - quantization_recipe: Optional[recipe.Recipe], - cpu_offload: bool, -) -> float: - """Measure the growth in allocated GPU memory in MiB after a model forward pass. + if use_cuda_graphs and not retain_pinned_cpu_buffers: + pytest.skip( + "Cuda graphs are not yet supported with cpu offloading when" + " retain_pinned_cpu_buffers is False." + ) - Memory measurement excludes the input and output tensors. + if backend == "FusedAttention" and use_cuda_graphs: + pytest.skip( + "Fused attention + cuda graphs is temporarily broken, not because of cpu offloading" + ) - """ + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" - # Reset memory - gc.collect() - torch.cuda.empty_cache() + if backend == "FlashAttention": + os.environ["NVTE_FLASH_ATTN"] = "1" + elif backend == "FusedAttention": + os.environ["NVTE_FUSED_ATTN"] = "1" + elif backend == "UnfusedAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" - # Context and sync function for CPU offloading - if cpu_offload: - offload_context, sync_function = te.get_cpu_offload_context( + offload_ctx, sync_function = get_cpu_offload_context( enabled=True, - num_layers=len(modules), - model_layers=len(modules) + 1, + num_layers=1, + model_layers=2, offload_activations=True, offload_weights=False, + retain_pinned_cpu_buffers=retain_pinned_cpu_buffers, ) - else: - offload_context = contextlib.nullcontext() - sync_function = lambda x: x - - # Forward pass, with dummy step to trigger offload for last module - inp = _make_input() - tensor = inp - memory_before_forward = torch.cuda.memory_allocated() / (1024**2) - for module in modules: - with te.autocast( - enabled=quantization_recipe is not None, recipe=quantization_recipe - ), offload_context: - tensor = module(tensor) - tensor = sync_function(tensor) - with offload_context: - tensor = tensor.clone() - tensor = sync_function(tensor) - memory_after_forward = (torch.cuda.memory_allocated() - tensor.nbytes) / (1024**2) - - # Backward pass - tensor.sum().backward() - torch.cuda.synchronize() - - # Memory usage in MiB - return memory_after_forward - memory_before_forward - - -@pytest.mark.parametrize("quantization_recipe", quantization_recipes) -@pytest.mark.parametrize("model_name", model_types.keys()) -def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: str) -> None: - """Check that CPU offloading runs and has expected memory usage.""" - - # Construct model - modules_list = [model_types[model_name]() for _ in range(NUM_LAYERS)] - if model_name in ["multihead_attention", "transformer_layer"]: - available_backends, *_ = get_available_attention_backends( - model_config["small"], - qkv_dtype=torch.bfloat16, - qkv_layout="sbhd_sbhd_sbhd", + + class Callable(torch.nn.Module): + def __init__(self, offload_ctx=None, sync_function=None): + super().__init__() + self.layers = torch.nn.ModuleList( + [Utils.create_layer(layer_type) for _ in range(2)] + ) + self.offload_ctx = offload_ctx + self.sync_function = sync_function + + def forward(self, x): + m_splits = ( + {"m_splits": [Utils._B * Utils._S // Utils._H] * Utils._H} + if layer_type == "grouped_linear" + else {} + ) + is_ops_layer = layer_type in ["linear_op", "layernorm_mlp_ops"] + for layer in self.layers: + with self.offload_ctx, recipe_ctx(): + if is_ops_layer: + x = layer(x, **m_splits) + else: + x = layer(x, is_first_microbatch=False, **m_splits) + if self.sync_function is not None: + x = self.sync_function(x) + return x + + callable_offload = Callable(offload_ctx=offload_ctx, sync_function=sync_function) + callable_no_offload = Callable(offload_ctx=contextlib.nullcontext(), sync_function=None) + + # copy parameters + for param_offload, param_no_offload in zip( + callable_offload.parameters(), callable_no_offload.parameters() + ): + param_offload.data.copy_(param_no_offload.data) + + x = Utils.create_tensor(None) + + if use_cuda_graphs: + callable_offload = te.make_graphed_callables( + callable_offload, + (x,), + enabled=recipe is not None, + recipe=(Utils.create_recipe_ctx(recipe) if recipe is not None else None), + ) + + # warm up (for example to compute sf for delayed scaling) + for _ in range(4): + out = callable_offload(x) + out.sum().backward() + out = callable_no_offload(x) + out.sum().backward() + + callable_offload.zero_grad(set_to_none=True) + out_offload = callable_offload(x) + out_offload.sum().backward() + + # save out and gradients + offload_outs = [out_offload] + for param in callable_offload.parameters(): + offload_outs.append(param.detach().clone()) + + torch.cuda.reset_peak_memory_stats() + out_no_offload = callable_no_offload(x) + out_no_offload.sum().backward() + + # collect gradients + no_offload_outs = [out_no_offload] + for param in callable_no_offload.parameters(): + no_offload_outs.append(param.detach().clone()) + + # check if tensors are the same + for i in range(len(offload_outs)): + assert torch.allclose(offload_outs[i], no_offload_outs[i]), f"Error in tensor {i}." + + torch.cuda.synchronize() + + def test_example_from_doc(self): + offload_stream = torch.cuda.Stream() + num_layers = 10 + layers = [Utils.create_layer("transformer_layer") for _ in range(num_layers)] + inp = [Utils.create_tensor(None) for _ in range(num_layers)] + out = [None] * num_layers + cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context( + enabled=True, + model_layers=num_layers, + manual_synchronization=True, + offload_stream=offload_stream, ) - _, fused_attn_supported, _ = available_backends - if not fused_attn_supported: - pytest.skip("Fused attention backend not available.") - os.environ["NVTE_FLASH_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True - - # Warmup - _warmup_model(modules_list, quantization_recipe) - - # Measure cached memory after forward pass - memory_without_offload = _measure_cached_memory(modules_list, quantization_recipe, False) - memory_with_offload = _measure_cached_memory(modules_list, quantization_recipe, True) - - # Check for expected memory usage - assert memory_with_offload < memory_without_offload - memory_from_cached_weights = _estimate_cached_weight_size( - model_name, - modules_list, - quantization_recipe, - ) - assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON + + for i in range(num_layers): + with cpu_offload_context: + out[i] = layers[i].forward(inp[i]) + out[i] = sync_function(out[i]) + manual_controller.start_offload_layer(i) + + offload_stream.synchronize() + for i in range(num_layers): + manual_controller.release_activation_forward_gpu_memory(i) + + for i in range(num_layers - 1, -1, -1): + # these calls are intended to be done in the backward pass + manual_controller.start_reload_layer(i) + + offload_stream.synchronize() + for i in range(num_layers): + out[i].sum().backward() diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py new file mode 100644 index 0000000000..8a8e036304 --- /dev/null +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -0,0 +1,215 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import contextlib +import gc +import os +from typing import Iterable, Optional + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported +from utils import ModelConfig, get_available_attention_backends + +# Check supported quantization schemes +fp8_available = te.is_fp8_available() +mxfp8_available = te.is_mxfp8_available() + +quantization_recipes: Optional[recipe.Recipe] = [None] +if fp8_available: + quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) + +model_config = { + "small": ModelConfig(8, 512, 8, 64, num_layers=5, eps=0.1), +} +SIZE = model_config["small"].hidden_size +NUM_HEADS = model_config["small"].num_heads +NUM_LAYERS = model_config["small"].num_layers +EPSILON = model_config["small"].eps + +# Flash attention saves some internal tensor for the backward pass +# that cannot be offloaded to CPU. +assert os.getenv("NVTE_FLASH_ATTN") == "0" + +# CPU offload v1 code path is enabled +assert os.environ.get("NVTE_CPU_OFFLOAD_V1", "0") == "1" + +# Offloading is supported for attention only for fused and flash attention backends, +# so the use of bfloat16 is required. +# +# For the TransformerLayer, activation offloading with dropout is not supported, +# so we set hidden_dropout to 0.0. +model_types = { + "linear": lambda: te.Linear(SIZE, SIZE, params_dtype=torch.bfloat16), + "layernorm_mlp": lambda: te.LayerNormMLP(SIZE, SIZE, params_dtype=torch.bfloat16), + "layernorm_linear": lambda: te.LayerNormLinear(SIZE, SIZE, params_dtype=torch.bfloat16), + "multihead_attention": lambda: te.MultiheadAttention( + SIZE, NUM_HEADS, params_dtype=torch.bfloat16 + ), + "transformer_layer": lambda: te.TransformerLayer( + SIZE, SIZE, NUM_HEADS, params_dtype=torch.bfloat16, hidden_dropout=0.0 + ), + "linear_op": lambda: te.ops.Linear(SIZE, SIZE, dtype=torch.bfloat16), + "layernorm_mlp_ops": lambda: te.ops.Sequential( + te.ops.LayerNorm(SIZE, dtype=torch.bfloat16), + te.ops.Linear(SIZE, SIZE, dtype=torch.bfloat16), + te.ops.GELU(), + te.ops.Linear(SIZE, SIZE, dtype=torch.bfloat16), + ), +} + + +def _make_input() -> torch.Tensor: + """Generate random input tensor.""" + return torch.randn( + (128, SIZE, SIZE), + dtype=torch.bfloat16, + device="cuda", + requires_grad=True, + ) + + +def _warmup_model( + modules: Iterable[torch.nn.Module], + quantization_recipe: Optional[recipe.Recipe], +) -> None: + """Perform forward and backward pass""" + tensor = _make_input() + for module in modules: + with te.autocast( + enabled=quantization_recipe is not None, + recipe=quantization_recipe, + ): + tensor = module(tensor) + tensor.sum().backward() + + +def _estimate_cached_weight_size( + model_name: str, + modules: Iterable[torch.nn.Module], + quantization_recipe: Optional[recipe.Recipe], +) -> float: + """Calculate the memory (in MiB) needed for weight caching.""" + + # The weight params are cached directly for unquantized compute + if quantization_recipe is None: + return 0 + + # Count number of weight param elements + param_elements = 0 + for module in modules: + for param in module.parameters(): + if param.dim() == 2: + param_elements += param.numel() + + # FP8 tensor-scaling caches one byte per element + if quantization_recipe.delayed() or quantization_recipe.float8_current_scaling(): + if not is_non_tn_fp8_gemm_supported() and model_name not in ( + "linear_op", + "layernorm_mlp_ops", + ): + # Modules do not deallocate FP8 transpose for weights + return 2 * param_elements / 1024**2 + return param_elements / 1024**2 + + # MXFP8 caches one data byte per element and one scale byte per 32 + # elements + if quantization_recipe.mxfp8(): + if model_name not in ("linear_op", "layernorm_mlp_ops"): + # Modules do not deallocate column-wise MXFP8 data for weights + return 2 * param_elements * (1 + 1 / 32) / 1024**2 + return param_elements * (1 + 1 / 32) / 1024**2 + + raise NotImplementedError(f"Unrecognized recipe ({quantization_recipe})") + + +def _measure_cached_memory( + modules: Iterable[torch.nn.Module], + quantization_recipe: Optional[recipe.Recipe], + cpu_offload: bool, +) -> float: + """Measure the growth in allocated GPU memory in MiB after a model forward pass. + + Memory measurement excludes the input and output tensors. + + """ + + # Reset memory + gc.collect() + torch.cuda.empty_cache() + + # Context and sync function for CPU offloading + if cpu_offload: + offload_context, sync_function = te.get_cpu_offload_context( + enabled=True, + num_layers=len(modules), + model_layers=len(modules) + 1, + offload_activations=True, + offload_weights=False, + ) + else: + offload_context = contextlib.nullcontext() + sync_function = lambda x: x + + # Forward pass, with dummy step to trigger offload for last module + inp = _make_input() + tensor = inp + memory_before_forward = torch.cuda.memory_allocated() / (1024**2) + for module in modules: + with te.autocast( + enabled=quantization_recipe is not None, recipe=quantization_recipe + ), offload_context: + tensor = module(tensor) + tensor = sync_function(tensor) + with offload_context: + tensor = tensor.clone() + tensor = sync_function(tensor) + memory_after_forward = (torch.cuda.memory_allocated() - tensor.nbytes) / (1024**2) + + # Backward pass + tensor.sum().backward() + torch.cuda.synchronize() + + # Memory usage in MiB + return memory_after_forward - memory_before_forward + + +@pytest.mark.parametrize("quantization_recipe", quantization_recipes) +@pytest.mark.parametrize("model_name", model_types.keys()) +def test_cpu_offload(quantization_recipe: Optional[recipe.Recipe], model_name: str) -> None: + """Check that CPU offloading runs and has expected memory usage.""" + + # Construct model + modules_list = [model_types[model_name]() for _ in range(NUM_LAYERS)] + if model_name in ["multihead_attention", "transformer_layer"]: + available_backends, *_ = get_available_attention_backends( + model_config["small"], + qkv_dtype=torch.bfloat16, + qkv_layout="sbhd_sbhd_sbhd", + ) + _, fused_attn_supported, _ = available_backends + if not fused_attn_supported: + pytest.skip("Fused attention backend not available.") + os.environ["NVTE_FLASH_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True + + # Warmup + _warmup_model(modules_list, quantization_recipe) + + # Measure cached memory after forward pass + memory_without_offload = _measure_cached_memory(modules_list, quantization_recipe, False) + memory_with_offload = _measure_cached_memory(modules_list, quantization_recipe, True) + + # Check for expected memory usage + assert memory_with_offload < memory_without_offload + memory_from_cached_weights = _estimate_cached_weight_size( + model_name, + modules_list, + quantization_recipe, + ) + assert abs(memory_with_offload - memory_from_cached_weights) < EPSILON diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 147a85fc2f..543055061b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -50,6 +50,13 @@ ) from transformer_engine.pytorch.attention.dot_product_attention.softmax import FusedScaleMaskSoftmax from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.cpu_offload import ( + is_cpu_offload_enabled, + start_offload, + mark_activation_offload, + NVTE_CPU_OFFLOAD_V1, +) +from transformer_engine.pytorch.cpu_offload_v1 import is_current_layer_offloaded # Import attention utils import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils @@ -737,6 +744,9 @@ def forward( x.contiguous() for x in (query_layer._data, key_layer._data, value_layer._data) ] + if is_cpu_offload_enabled(): + start_offload(query_layer, key_layer, value_layer, offload_base_tensor=True) + # get batch_size, max_seqlen and cu_seqlens batch_size, context_len = None, None if inference_params is None: @@ -877,12 +887,7 @@ def forward( fp8_output=fp8_output, ) else: - from transformer_engine.pytorch.cpu_offload import ( - CPUOffloadEnabled, - mark_activation_offload, - ) - - if CPUOffloadEnabled: + if is_cpu_offload_enabled(): mark_activation_offload( query_layer, key_layer, value_layer, cu_seqlens_q, cu_seqlens_kv ) @@ -1116,6 +1121,9 @@ def forward( nvtx_label = "transformer_engine.FusedAttnFunc.forward" nvtx_range_push(f"{nvtx_label}") + if is_cpu_offload_enabled(): + start_offload(q, k, v, offload_base_tensor=True) + # 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() @@ -1293,12 +1301,7 @@ def forward( # used when some tensors are base tensors and loose the "dtype" attribute ctx.nominal_dtype = out_nominal_dtype - from transformer_engine.pytorch.cpu_offload import ( - CPUOffloadEnabled, - mark_activation_offload, - ) - - if CPUOffloadEnabled: + if is_cpu_offload_enabled() and NVTE_CPU_OFFLOAD_V1: if ctx.fp8: tensor_list = fp8_tensors else: @@ -1309,6 +1312,7 @@ def forward( ctx.is_input_fp8 = is_input_fp8 ctx.is_output_fp8 = is_output_fp8 + tensors_to_save, tensor_objects = prepare_for_saving( *fp8_tensors, *qkvo_tensors, @@ -1339,27 +1343,26 @@ def forward( ctx.dropout_p = dropout_p ctx.fast_zero_fill = fast_zero_fill - from transformer_engine.pytorch.cpu_offload import ( - CPUOffloadedLayer, - ) - - # If interleaved tensor is offloaded, reloaded tensor will be - # non-interleaved, so we need to modify the QKV layout - # for backward - if CPUOffloadedLayer and CPUOffloadEnabled: - reload_layout = "" - split_list = qkv_layout.split("_") - for split in split_list: - temp_layout = "" - rep_count = 1 - for s in split: - if s.isalpha(): - temp_layout = temp_layout + s - else: - rep_count = int(s) - for _ in range(rep_count): - reload_layout = reload_layout + temp_layout + "_" - ctx.qkv_layout = reload_layout[:-1] + if NVTE_CPU_OFFLOAD_V1: + # If interleaved tensor is offloaded, reloaded tensor will be + # non-interleaved, so we need to modify the QKV layout + # for backward + if is_current_layer_offloaded() and is_cpu_offload_enabled(): + reload_layout = "" + split_list = qkv_layout.split("_") + for split in split_list: + temp_layout = "" + rep_count = 1 + for s in split: + if s.isalpha(): + temp_layout = temp_layout + s + else: + rep_count = int(s) + for _ in range(rep_count): + reload_layout = reload_layout + temp_layout + "_" + ctx.qkv_layout = reload_layout[:-1] + else: + ctx.qkv_layout = qkv_layout else: ctx.qkv_layout = qkv_layout 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 4278820e7a..4157e8d3a4 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 @@ -1494,14 +1494,6 @@ def forward( fp8_output=fp8_output, ) - from transformer_engine.pytorch.cpu_offload import CPUOffloadEnabled - - if CPUOffloadEnabled: - warnings.warn( - "Attention activation Offloading is only implemented" - "with Flash Attention and Fused Attention!" - ) - if use_unfused_attention: allow_emulation = os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" if checkpoint_core_attention: diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index b3bda677bb..2440693df4 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -33,6 +33,8 @@ from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb +from transformer_engine.pytorch.cpu_offload import start_offload, is_cpu_offload_enabled + # Force DotProductAttention to use a different recipe than the fp8_recipe set in autocast(). # Useful when GEMMs and attention use different recipes. Supported values are "DelayedScaling" # and "Float8CurrentScaling". Use other relevant variables here to define the recipe, e.g. fp8_dpa. @@ -971,7 +973,8 @@ def forward( # =========================== # Core attention computation # =========================== - + if is_cpu_offload_enabled(): + start_offload(query_layer, key_layer, value_layer, offload_base_tensor=True) context_layer = self.core_attention( query_layer, key_layer, diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 6edc126200..bfdee34752 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -3,698 +3,748 @@ # See LICENSE for license information. """Functionality for CPU offloading of tensors saved for backward pass.""" -from __future__ import annotations -from contextlib import nullcontext -from typing import Any, Dict, Optional +from __future__ import annotations +import contextlib +from collections import defaultdict +from dataclasses import dataclass, field +import os +import warnings +from typing import Any, Optional import torch - +from torch.autograd.graph import saved_tensors_hooks from transformer_engine.debug.pytorch.debug_state import TEDebugState -from .quantized_tensor import QuantizedTensorStorage -from .tensor.float8_tensor import Float8Tensor - -__all__ = ["get_cpu_offload_context"] +import transformer_engine.pytorch as te +import transformer_engine.pytorch.cpu_offload_v1 as v1_code_path +from .quantized_tensor import ( + restore_from_saved, + prepare_for_saving, +) -CPUOffloadEnabled = False -CPUOffloadedLayer = False - - -def mark_activation_offload(*tensors): - """Set the type of the offloading needed for a tensor.""" - if TEDebugState.debug_enabled: - raise RuntimeError("CPU offload is not supported in debug mode.") - for tensor in tensors: - if tensor is None: - continue - if type(tensor) in [torch.Tensor, torch.nn.Parameter]: - tensor.activation_offloading = True - else: - data_tensors = tensor.get_data_tensors() - for tensor in data_tensors: - if tensor is not None: - tensor.activation_offloading = True - # This is a hack to force clear the tensor after it is offloaded. - # It is needed, because .*TensorStorage classes are saved in the ctx, - # and they contain the reference to their data tensors. - tensor.needs_force_clear = True +__all__ = ["get_cpu_offload_context", "mark_not_offload", "start_offload"] +NVTE_CPU_OFFLOAD_V1 = os.environ.get("NVTE_CPU_OFFLOAD_V1", "0") == "1" -def is_cpu_offload_enabled() -> bool: - """Check if CPU offloading is currently enabled.""" - return CPUOffloadEnabled +OFFLOAD_SYNCHRONIZER = None -class CpuOffloadSavedTensorHook: - """Contex-manager that executes a pair of pack/unpack hooks for saved tensors. +def is_cpu_offload_enabled(): + """Returns True if CPU offload is enabled.""" + if NVTE_CPU_OFFLOAD_V1: + return v1_code_path.is_cpu_offload_enabled() + return OFFLOAD_SYNCHRONIZER is not None - In this context, the ``on_save_for_backward`` method will be called every time - a tensor is saved for backward (this includes intermediary results saved using - :func:`~torch.autograd.function._ContextMethodMixin.save_for_backward` but - also those recorded by a PyTorch-defined operation). - The ``on_get_saved_tensors`` method will be called when the backward function - of this op attempts to retrieve the saved tensor from context (this includes - :func: `torch.Tensor.backward()` or :func: `torch.autograd.grad()`. It takes the - as input the return value of the ``on_save_for_backward``, and is meant to return - an identical copy of the tensor being saved by ``on_save_for_backward`` in terms of - size, device and element values. +def mark_activation_offload(*tensors): + """Set the type of the offloading needed for a tensor.""" + if NVTE_CPU_OFFLOAD_V1: + v1_code_path.mark_activation_offload(*tensors) - Example: - >>> import torch - >>> from typing import Any - >>> - >>> class DummyHook(CpuOffloadSavedTensorHook): - ... - ... def on_save_for_backward(self, tensor: torch.Tensor) -> Any: - ... logging.info("On save", tensor) - ... return (tensor,) - ... - ... def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: - ... logging.info("On get", saved_state) - ... tensor, = saved_state - ... return tensor - ... - >>> a = torch.ones(5, requires_grad=True) - >>> b = torch.ones(5, requires_grad=True) * 2 - >>> with DummyHook(): - ... y = a * b - ... - On save tensor([1., 1., 1., 1., 1.], requires_grad=True) - On save tensor([2., 2., 2., 2., 2.], grad_fn=) - >>> y.sum().backward() - On get (tensor([1., 1., 1., 1., 1.], requires_grad=True),) - On get (tensor([2., 2., 2., 2., 2.], grad_fn=),) +def mark_not_offload(*tensors: torch.Tensor): + """Marks tensors to prevent them from being offloaded.""" + if NVTE_CPU_OFFLOAD_V1: + return - """ + tensors, tensor_obj = prepare_for_saving(*tensors) - def __init__(self) -> None: - self.inside_context = False + for tensor in tensors: + if tensor is not None: + setattr(tensor, "_TE_do_not_offload", True) - def __enter__(self): - global CPUOffloadEnabled - CPUOffloadEnabled = True + restore_from_saved(tensor_obj, tensors) - self.inside_context = True - torch._C._autograd._push_saved_tensors_default_hooks( - self.on_save_for_backward, self.on_get_saved_tensor - ) - def __exit__(self, *args: Any): - global CPUOffloadEnabled - CPUOffloadEnabled = False +def start_offload(*tensors: torch.Tensor, offload_base_tensor: bool = False): + """ + Marks point in on main stream where tensors are fully computed and ready to be offloaded. + If offload_base_tensor is True and the tensor is a view, the base tensor is offloaded + and reloaded - the stride and storage offset of the view are saved and restored after reload. + It is useful when multiple tensors are views of the same base tensor, + for example in MultiHeadAttention for interleaved q, k, v tensors. + """ + if NVTE_CPU_OFFLOAD_V1: + return - self.inside_context = False - torch._C._autograd._pop_saved_tensors_default_hooks() + def _mark_tensor_for_offload(t): + if t is None: + return + # Attach an event to mark when the tensor is ready for reload. + t.start_reload_event = torch.cuda.Event() + t.start_reload_event.record(torch.cuda.current_stream()) + if offload_base_tensor and t._base is not None: + setattr(t, "offload_base_tensor", True) - def on_save_for_backward(self, tensor: torch.Tensor) -> Any: - """On save for backward.""" - raise NotImplementedError( - "`on_save_for_backward: Callable[[torch.Tensor], Any]`" - "is not implemented in CpuOffloadHook class. Inherit " - "this class and implement your custom hooks" - ) + tensors, tensor_obj = prepare_for_saving(*tensors) - def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: - """On get saved tensor.""" - raise NotImplementedError( - "`on_get_saved_tensors: Callable[[Any], torch.Tensor]`" - "is not implemented in CpuOffloadHook class. Inherit " - "this class and implement your custom hooks" - ) + for tensor in tensors: + _mark_tensor_for_offload(tensor) + restore_from_saved(tensor_obj, tensors) -class CpuOffloadHookWithOffloadHandler(CpuOffloadSavedTensorHook): - """Context-manager that offloads/recovers tensors through an offload hander. - The hook just offloads/recovers the tensor object to the handler through `tensor_push` - and `tensor_pop` interface. How the offload-handler manages the offloading, recovering - or prefetching timing is transparent to this hook. +@dataclass +class TensorGroup: + """ + TensorGroup is a collection of tensors, events and auxiliary data. + It is used multiple times in the CPU offload code. """ - def __init__( - self, - offload_handler: OffloadHandler, - handler_extra_kwargs: Optional[Dict[str, Any]] = None, - debug: bool = False, - ) -> None: - if handler_extra_kwargs is None: - handler_extra_kwargs = {} - self.debug: bool = debug - self.offload_handler: OffloadHandler = offload_handler - self.handler_extra_kwargs: Dict[str, Any] = handler_extra_kwargs - super().__init__() - - def on_save_for_backward(self, tensor: torch.Tensor) -> Any: - retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs) - return retrieve_identifier - - def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: - tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs) - return tensor - + tensor_list: list[torch.Tensor] = field(default_factory=list) + events: list[torch.cuda.Event] = field(default_factory=list) + aux: Any = None -class OffloadHandler: - """A base class for CPU offload-handler.""" - def __init__(self) -> None: - pass +class TensorGroupProcessor: + """ + Suppose there is a tensor group T that needs to be offloaded. + Possibly we can switch T into (T_opt, aux), where T_opt is smaller and easier to offload, + offload T_opt, reload it and then restore T from (T_opt_reloaded, aux). - def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: - """Tensor push.""" - raise NotImplementedError( - "`tensor_push is not implented in OffloadHandler class. " - "Inherit this class and implement your custom tensor_push." - ) + This class contains static methods that perform these optimizations - for example + deduplication of tensors and restoring duplicates after reload. + """ - def tensor_pop(self, tensor_tag: Any, **kwargs): - """Tensor pop.""" - raise NotImplementedError( - "`tensor_pop is not implented in OffloadHandler class. " - "Inherit this class and implement your custom tensor_pop." - ) + @staticmethod + def tensor_group_process_before_offload(tensor_group: TensorGroup) -> tuple[TensorGroup, Any]: + """ + Call for a tensor group, just before offloading logic. + aux is a dictionary that contains auxiliary data, needed to restore pre-offload state. + """ + aux = {} + tensor_group = TensorGroupProcessor._switch_to_base_tensors(aux, tensor_group) + tensor_group = TensorGroupProcessor._deduplicate_tensors(aux, tensor_group) + return tensor_group, aux -class GroupCommitFunction(torch.autograd.Function): - """this is a dummy op with output identical to input. - However, it is necessary for marking a timepoint for offload handler to - accomplish all synchronizations. Implementing it as a function is necessary - because we need to actions in both forward and backward. - """ + @staticmethod + def tensor_group_process_after_reload(tensor_group: TensorGroup): + """ + Call for a tensor group, just after reload logic. + """ + assert tensor_group.aux is not None + tensor_group = TensorGroupProcessor._restore_tensor_duplicates(tensor_group) + tensor_group = TensorGroupProcessor._switch_to_views(tensor_group) + return tensor_group @staticmethod - def forward(ctx, tensor, cpu_offload_handler): - # pylint: disable=missing-function-docstring - cpu_offload_handler.on_group_commit_forward() - ctx.cpu_offload_handler = cpu_offload_handler - # return the identical tensor - return tensor + def _switch_to_base_tensors(aux, tensor_group: TensorGroup) -> TensorGroup: + """ + Changes tensors to base tensors and saves view options in aux. + + It we save multiple tensors which in fact are views of the same base tensor, + this will offload only this one base tensor. It is used for example in + MultiHeadAttention for interleaved q, k, v tensors. + """ + + def _check_if_offload_base_tensor(tensor: torch.Tensor) -> bool: + if getattr(tensor, "offload_base_tensor", False): + return True + if tensor._base is not None: + # If tensor is a view of a tensor and has the same elements, + # but with different strides, we can safely offload the base tensor. + # If tensor is a view on some part of a bigger tensor, + # the decision to offload the base tensor is non-trivial and we do not do it by default. + return tensor._base.numel() == tensor.numel() + return False + + aux["views"] = [] + for tensor_id in range( # pylint: disable=consider-using-enumerate + len(tensor_group.tensor_list) + ): + tensor = tensor_group.tensor_list[tensor_id] + if _check_if_offload_base_tensor(tensor): + aux["views"].append((tensor.shape, tensor.stride(), tensor.storage_offset())) + tensor = tensor._base + assert ( + tensor is not None + ), "Cannot offload base tensor, if the tensor is not a view." + tensor_group.tensor_list[tensor_id] = tensor + else: + aux["views"].append(None) + return tensor_group @staticmethod - def backward(ctx, grad_output): - # pylint: disable=missing-function-docstring - cpu_offload_handler = ctx.cpu_offload_handler - cpu_offload_handler.on_group_commit_backward() - return grad_output, None + def _deduplicate_tensors(aux, tensor_group: TensorGroup) -> TensorGroup: + """ + Deduplicate tensors. + """ + dedup_tensors: list[torch.Tensor] = [] + dedup_events: list[torch.cuda.Event] = [] + tensor_to_index: dict[int, int] = {} + aux["original_tensor_ids"] = [] + # If there are several duplicates of the same tensor, with different events, + # we keep only first event - every event is recorded when the tensor is ready to be offloaded, + # so it is the most optimal to use the first event. + for tensor_id, tensor in enumerate(tensor_group.tensor_list): + if id(tensor) in tensor_to_index: + aux["original_tensor_ids"].append(tensor_to_index[id(tensor)]) + else: + tensor_to_index[id(tensor)] = len(dedup_tensors) + dedup_tensors.append(tensor) + dedup_events.append(tensor_group.events[tensor_id]) + aux["original_tensor_ids"].append(tensor_to_index[id(tensor)]) -group_prefetch_offload_commit = GroupCommitFunction.apply + tensor_group.tensor_list = dedup_tensors + tensor_group.events = dedup_events + return tensor_group + @staticmethod + def _restore_tensor_duplicates(tensor_group: TensorGroup) -> TensorGroup: + """ + Restore tensor duplicates. + """ + new_tensor_list = [] + new_events_list = [] + for tensor_id in range(len(tensor_group.aux["original_tensor_ids"])): + original_tensor_id = tensor_group.aux["original_tensor_ids"][tensor_id] + new_tensor_list.append(tensor_group.tensor_list[original_tensor_id]) + new_events_list.append(tensor_group.events[original_tensor_id]) + + tensor_group.tensor_list = new_tensor_list + tensor_group.events = new_events_list + return tensor_group -class SynchronizedGroupOffloadHandler(OffloadHandler): - """Offload Handler that offloads/reloads in a synchronized way. - The device-to-host and host-to-device copying happen in the same stream - as the computation kernels, thus the copying will block computation. + @staticmethod + def _switch_to_views(tensor_group: TensorGroup) -> TensorGroup: + """ + Switch to views - reverse of _switch_to_base_tensors. + """ + for tensor_id, tensor in enumerate(tensor_group.tensor_list): + if tensor_group.aux["views"][tensor_id] is not None: + tensor_group.tensor_list[tensor_id] = tensor.as_strided( + *tensor_group.aux["views"][tensor_id] + ) + return tensor_group + + +class OffloadableLayerState: + """ + Class that manages offloading and reloading of tensors for a single layer. """ def __init__( - self, num_offload_group, tensor_need_offloading_checker=(lambda _: True), debug=False - ) -> None: - super().__init__() - - self.num_offload_group = num_offload_group - self.tensor_need_offloading_checker = tensor_need_offloading_checker - self.debug = debug - - self.groupid_reset() - - def groupid_reset(self): - """Groupid reset.""" - # Data structures to label saved tensors and book-keep their cpu copies. - # Currently, on push, create a new cpu tensor and copies; on pop, copies - # the tensor back to gpu and deletes the cpu tensor. - # These will increment whenever `group_commit()` is invoked - self.current_group, self.tensor_count_current_group = (0, 0) - self.torch_tensor_count = 0 - self.tensor_tag_to_state = {} - - def on_group_commit_forward(self): - """On group commit forward.""" - # finishing up with updating current group and tensor count - self.current_group += 1 # increment - self.tensor_count_current_group = 0 # reset - - def on_group_commit_backward(self): - """On group commit backward.""" - self.current_group -= 1 - assert self.current_group >= 0 - - @staticmethod - def offload(src_tensor, pin_memory=True): - """Offload.""" - - cpu_backup = torch.empty( - src_tensor.size(), - dtype=src_tensor.dtype, - layout=src_tensor.layout, - device="cpu", - pin_memory=pin_memory, + self, + offload_stream: torch.cuda.Stream, + retain_pinned_cpu_buffers: bool = False, + ): + self.offload_stream = offload_stream + self.retain_pinned_cpu_buffers = retain_pinned_cpu_buffers + + # There are 3 tensor groups: tensors on gpu before offload, + # tensors on cpu after offload, tensors on gpu after reload. + self.fwd_gpu_tensor_group = TensorGroup() + self.cpu_tensor_group = TensorGroup() + self.bwd_gpu_tensor_group = TensorGroup() + + self.aux: dict[str, Any] = {} + + # State can be one of: not_offloaded, offload_started, + # offload_finished, reload_started. + self.state = "not_offloaded" + + def _validate_state(self, func_name: str, allowed_states: list[str]): + assert ( + self.state in allowed_states + ), f"Invalid state: {self.state} for {func_name}, must be one of {allowed_states}" + + def start_offload(self): + """ + Start offloading of tensors. Puts copy from GPU to CPU tasks on offload stream. + Before each copy event, the offload stream waits for the event signalling that the tensor is ready to be offloaded. + This event is recorded in the start_offload or push_tensor call. + """ + self._validate_state(func_name="start_offload", allowed_states=["not_offloaded"]) + self.state = "offload_started" + + self.fwd_gpu_tensor_group, aux = TensorGroupProcessor.tensor_group_process_before_offload( + self.fwd_gpu_tensor_group ) - cpu_backup.copy_(src_tensor, non_blocking=pin_memory) - state = (src_tensor.device, cpu_backup) - return state - - @staticmethod - def reload(state, non_blocking=None, copy_buffer=None): - """Reload.""" - dev, cpu_backup = state - if non_blocking is None: - non_blocking = cpu_backup.is_pinned() - - if copy_buffer is None: - return cpu_backup.to(dev, non_blocking=non_blocking) - - assert cpu_backup.size() == copy_buffer.size(), "Can't copy two buffers of different sizes!" - - copy_buffer.copy_(cpu_backup, non_blocking=non_blocking) + allocate_cpu_buffers = ( + not self.retain_pinned_cpu_buffers or len(self.cpu_tensor_group.tensor_list) == 0 + ) - return copy_buffer + for tensor_id, tensor in enumerate(self.fwd_gpu_tensor_group.tensor_list): + assert tensor.is_contiguous() - def tensor_push(self, tensor: torch.Tensor, **kwargs): - """Tensor push.""" - # obtain a unique tensor tag - tensor_tag = (self.current_group, self.tensor_count_current_group) - self.tensor_count_current_group += 1 - assert tensor_tag not in self.tensor_tag_to_state - if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker( - tensor - ): - state = SynchronizedGroupOffloadHandler.offload(tensor) - self.tensor_tag_to_state[tensor_tag] = state - else: - # will be offloaded together after group commit - self.tensor_tag_to_state[tensor_tag] = tensor + # Wait for the moment the tensor is ready to be offloaded. + self.offload_stream.wait_event(self.fwd_gpu_tensor_group.events[tensor_id]) # type: ignore[arg-type] - return tensor_tag + with torch.cuda.stream(self.offload_stream): + if allocate_cpu_buffers: + # empty_like is defined also for QuantizedTensors + offloaded_tensor = torch.empty_like( + tensor, device=torch.device("cpu"), pin_memory=True + ) + self.cpu_tensor_group.tensor_list.append(offloaded_tensor) + else: + assert self.cpu_tensor_group.tensor_list[tensor_id].shape == tensor.shape, ( + "CPU buffer shape does not match the offloaded tensor shape:" + f" {self.cpu_tensor_group.tensor_list[tensor_id].shape} != {tensor.shape} " + " Make sure that tensor shaped do not change between" + " iterations if retain_pinned_cpu_buffers is True." + ) + offloaded_tensor = self.cpu_tensor_group.tensor_list[tensor_id] + offloaded_tensor.copy_(tensor, non_blocking=True) + + # aux is a dictionary that contains auxiliary data like information which tensors were deduplicated, + # 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) + + def release_activation_forward_gpu_memory(self): + """ + Release GPU memory of the activations. + Waits for offload to finish - memory needs to be kept alive when GPU->CPU copy is performed. + """ + self._validate_state( + func_name="release_activation_forward_gpu_memory", allowed_states=["offload_started"] + ) + self.state = "offload_finished" + + 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 + + def start_reload(self): + """ + Start reloading of tensors. + It allocates new tensors on GPU and puts copy from CPU tasks on offload stream. + """ + self._validate_state(func_name="start_reload", allowed_states=["offload_finished"]) + self.state = "reload_started" + + self.bwd_gpu_tensor_group = TensorGroup() + for tensor in self.cpu_tensor_group.tensor_list: + + # Notice that reloaded tensor is allocated on main stream, + # not offloaded stream. It is because PyTorch memory allocator + # cannot move tensors from pool of one stream to another without + # calling cudaFree and cudaMalloc again. + + # empty_like is defined also for QuantizedTensors. + reloaded_tensor = torch.empty_like(tensor, device=torch.device("cuda")) + self.offload_stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(self.offload_stream): + reloaded_tensor.copy_(tensor, non_blocking=True) + + reload_tensor_event = torch.cuda.Event() + reload_tensor_event.record(self.offload_stream) + self.bwd_gpu_tensor_group.events.append(reload_tensor_event) + self.bwd_gpu_tensor_group.tensor_list.append(reloaded_tensor) + + self.bwd_gpu_tensor_group.aux = self.aux + self.bwd_gpu_tensor_group = TensorGroupProcessor.tensor_group_process_after_reload( + self.bwd_gpu_tensor_group + ) - def tensor_pop(self, tensor_tag, **kwargs): - """Tensor pop.""" - assert tensor_tag in self.tensor_tag_to_state - state = self.tensor_tag_to_state.pop(tensor_tag) - if isinstance(state, tuple): - tensor = SynchronizedGroupOffloadHandler.reload(state) - else: - tensor = state + def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor: + """ + It is called when a tensor is saved for backward pass. + + If tensor is offloaded, returns int representing the index of the tensor in the offloaded tensor group. + If tensor is not offloaded, returns the tensor itself. + """ + self._validate_state(func_name="push_tensor", allowed_states=["not_offloaded"]) + + if self._check_if_offload(tensor): + self.fwd_gpu_tensor_group.tensor_list.append(tensor) + # The group is processed and offloaded at the end of the forward pass of current layer. + # To enable offloading of tensors faster we use self.offload_stream and record + # the events when the tensors are ready to be offloaded. + # It means that we do not need to wait to the end of current layer to start offloading. + if hasattr(tensor, "start_reload_event"): + self.fwd_gpu_tensor_group.events.append(tensor.start_reload_event) + else: + self.fwd_gpu_tensor_group.events.append(torch.cuda.Event()) + self.fwd_gpu_tensor_group.events[-1].record(torch.cuda.current_stream()) + return len(self.fwd_gpu_tensor_group.tensor_list) - 1 return tensor + def pop_tensor(self, tensor_or_tensor_id: torch.Tensor | int) -> torch.Tensor: + """ + It is called when a tensor is used in backward pass. + Returns the tensor. If tensor was offloaded/reloaded, wait for the reload of a tensor to finish. + """ + self._validate_state( + func_name="pop_tensor", allowed_states=["not_offloaded", "reload_started"] + ) -class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler): - """Compared to synchronize, this uses more memory because of the buffer but - achieves better performance due to the overlapping. D2h and h2d copying are - completely hidden behind computation if computation time of a layer is longer - than host-device communication time. Bulk offloading with delay and bulk reloading - with prefetch are implemented.""" + # 1. tensor not offloaded + if isinstance(tensor_or_tensor_id, torch.Tensor): + return tensor_or_tensor_id + # 2. the layer was not offloaded at all + if self.state == "not_offloaded": + return self.fwd_gpu_tensor_group.tensor_list[tensor_or_tensor_id] + + # 3. the layer was offloaded + assert self.state == "reload_started" + # wait for the tensor to be reloaded + torch.cuda.current_stream().wait_event( + self.bwd_gpu_tensor_group.events[tensor_or_tensor_id] + ) + return self.bwd_gpu_tensor_group.tensor_list[tensor_or_tensor_id] + + def release_all_memory(self): + """Release all gpu and cpu memory the state stored. Is called after the backward pass.""" + self.fwd_gpu_tensor_group = TensorGroup() + if not self.retain_pinned_cpu_buffers: + self.cpu_tensor_group = TensorGroup() + self.bwd_gpu_tensor_group = TensorGroup() + self.state = "not_offloaded" + + def _check_if_offload(self, t: torch.Tensor) -> bool: + """ + Check if tensor needs to be offloaded. + """ + if ( + not isinstance(t, torch.nn.Parameter) + and not getattr(t, "_TE_do_not_offload", False) + and not isinstance(t, torch._subclasses.FakeTensor) + and t.device.type == "cuda" + ): + if not t.is_contiguous() and not getattr(t, "offload_base_tensor", False): + warnings.warn( + "Tried to offload non-contiguous tensor, which is not supported. Offload of" + " this tensor will be skipped." + ) + return False + + return True + return False + + def get_offloaded_total_size_mb(self) -> float: + """ + Get total size of offloaded tensors in MB, used only for testing. + """ + + def get_tensor_size_mb(tensor): + if tensor is None: + return 0 + if isinstance(tensor, te.quantized_tensor.QuantizedTensorStorage): + return sum(get_tensor_size_mb(t) for t in tensor.get_data_tensors()) + return tensor.numel() * tensor.element_size() / (1024**2) + + total_size = 0 + for tensor in self.cpu_tensor_group.tensor_list: + total_size += get_tensor_size_mb(tensor) + return total_size + + +class OffloadSynchronizer: + """ + Base class responsible for synchronizing offloading and reloading of tensors for multiple layers. + In base class we only track layer number and + create OffloadableLayerState instances for all layers, but do not start offloading or reloading. + """ def __init__( self, - num_offload_group, # must be <= actual number of groups (number of commits) - num_model_group, - tensor_need_offloading_checker=(lambda t: True), - double_buffering=False, - debug=False, - ) -> None: - super().__init__( - num_offload_group=num_offload_group, - tensor_need_offloading_checker=tensor_need_offloading_checker, - debug=debug, - ) - # Number of layers in the model - self.num_layers = num_model_group - # Data Structure to maintain reference to activation tensors - self.tensor_tag_to_buf = {} - # Data structure to hold the FP8/MXFP8 tensor objects - self.fp8_tensor_object_map = {} - self.float8_transpose_cache_valid = {} - self.dereferencing_list = [] - # Tracking the number of layers offloaded - self.offloaded_group_count = 0 - # Core data structure that decides the window for offloading - self.layer_window_map = {} - - # Data structures fo double buffered reloading - self.double_buffering = double_buffering - self.reload_double_buffer = [[], []] - self.double_buffer_created = False - - # Logic to make offloading load balance across computation - # for optimal CPU/GPU interconnect usage - constant = 0 - for i in range(self.num_offload_group): - self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1 - if i < (self.num_layers % self.num_offload_group): - self.layer_window_map[i] += i + 1 - constant = i + 1 - else: - self.layer_window_map[i] += constant - - # allocate streams and events for synchronization - self.d2h_stream = torch.cuda.Stream() - self.h2d_stream = torch.cuda.Stream() - - def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: - global CPUOffloadedLayer - - torch_stray_tensor = isinstance( - tensor, - ( - torch._subclasses.fake_tensor.FakeTensor, - torch._subclasses.functional_tensor.FunctionalTensor, - ), + num_layers: int, + retain_pinned_cpu_buffers: bool = False, + offload_stream: Optional[torch.cuda.Stream] = None, + ): + self.num_layers = num_layers + self.offload_stream = offload_stream if offload_stream is not None else torch.cuda.Stream() + + self.layer_states = { + i: OffloadableLayerState(self.offload_stream, retain_pinned_cpu_buffers) + for i in range(num_layers) + } + + self.num_of_fwds = None + self.previous_bwd_layer_id = None + self.current_layer_id = None + + def fwd_step(self) -> int: + """ + Invoked before each layer forward. + """ + if self.num_of_fwds in [None, self.num_layers - 1]: + # reset the offload synchronizer + self.num_of_fwds = 0 + else: + self.num_of_fwds += 1 + self.current_layer_id = self.num_of_fwds + return self.current_layer_id + + def bwd_step(self, layer_num: int): + """ + Invoked before each layer backward. + """ + if self.previous_bwd_layer_id is not None: + self.layer_states[self.previous_bwd_layer_id].release_all_memory() + self.previous_bwd_layer_id = layer_num + self.current_layer_id = layer_num + + def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor: + """Default push tensor method""" + return self.layer_states[self.num_of_fwds].push_tensor(tensor) + + def pop_tensor(self, tensor_or_tensor_id: torch.Tensor | int) -> torch.Tensor: + """Default pop tensor method""" + return self.layer_states[self.current_layer_id].pop_tensor(tensor_or_tensor_id) + + def finish_part_of_bwd(self): + """ + We need to release memory of backward - this call does that. + It needs to be invoked after every backward pass - there may be + more than one in pipeline parallelism. + + It is needed, because call bwd_step is invoked before each layer backward, + but we need to release memory after the backward pass is finished. + """ + if self.previous_bwd_layer_id is not None: + self.layer_states[self.previous_bwd_layer_id].release_all_memory() + self.previous_bwd_layer_id = None + + def get_offloaded_total_size_mb(self) -> float: + """ + Get total size of offloaded tensors in MB, used only for testing. + """ + return sum( + self.layer_states[layer_id].get_offloaded_total_size_mb() + for layer_id in self.layer_states ) - is_quantized_tensor = isinstance(tensor, QuantizedTensorStorage) - - if not torch_stray_tensor: - - # obtain a unique tensor tag - tensor_tag = (self.current_group, self.tensor_count_current_group) - self.tensor_count_current_group += 1 - - assert tensor_tag not in self.tensor_tag_to_state - if is_quantized_tensor: - tensor_list, _ = tensor.prepare_for_saving() - - self.tensor_tag_to_state[tensor_tag] = [] - self.tensor_tag_to_buf[tensor_tag] = [] - - # Added support for de-duplicating FP8 param tensors - for _, value in self.fp8_tensor_object_map.items(): - if tensor is value: - self.dereferencing_list.append(tensor_tag) - break +class DefaultOffloadSynchronizer(OffloadSynchronizer): + """ + Default implementation of OffloadSynchronizer, + intended to be used in standard training workloads - with multiple forwards + and multiple backwards. + """ - self.fp8_tensor_object_map[tensor_tag] = tensor - if isinstance(tensor, Float8Tensor): - self.float8_transpose_cache_valid[tensor_tag] = getattr( - tensor, "_transpose_invalid" - ) + def __init__( + self, + num_layers: int, + num_offloaded_layers: int | None = None, + retain_pinned_cpu_buffers: bool = False, + offload_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__(num_layers, retain_pinned_cpu_buffers, offload_stream) + + # map of layers to bool meaning if layer needs to be offloaded + self.offload_layer_map: dict[int, bool] = {} + + # num_layer: int -> list of layers that need to finish offload by this moment + self.finish_offload_map: defaultdict[int, list[int]] = defaultdict(list) + # num_layer: int -> list of layers that need to start reload in this moment + self.start_reload_map: defaultdict[int, list[int]] = defaultdict(list) + + self._init_offload_synchronization_dicts(num_offloaded_layers) + + def _init_offload_synchronization_dicts(self, num_offloaded_layers: int): + """ + If synchronization dictionary is not provided, the number of offloaded layers is used to initialize + offload_layer_map, finish_offload_map and start_reload_map. + + The aim is to minimize memory usage by the end of the forward pass. + + The optimal strategy for that is to offload layers 0, ..., num_offloaded_layers - 1. + For layer i offload needs to finish before num_layers - num_offloaded_layers + i. + For layer i reload needs to start after num_layers - num_offloaded_layers + i. + + This ensures that - if all layers have memory footprint of T - then peak memory usage of saving activations is + (num_layers - num_offloaded_layers) * T. + """ + for layer_id in range(self.num_layers): + if layer_id < num_offloaded_layers: + self.offload_layer_map[layer_id] = True + self.finish_offload_map[self.num_layers - num_offloaded_layers + layer_id].append( + layer_id + ) + self.start_reload_map[self.num_layers - 1 - num_offloaded_layers + layer_id].append( + layer_id + ) else: - tensor_list = [tensor] - - for t in tensor_list: - if is_quantized_tensor: - self.tensor_tag_to_state[tensor_tag].append(t) - else: - self.tensor_tag_to_state[tensor_tag] = t - - if ( - self.current_group < self.num_offload_group - and self.tensor_need_offloading_checker(t) - ): - if is_quantized_tensor: - self.tensor_tag_to_buf[tensor_tag].append(t) - # Need to clear the internal data reference for the quantized tensors - tensor.clear() - else: - self.tensor_tag_to_buf[tensor_tag] = t - - # Needed to differentiate non offloaded layer's attention - # QKV layout of attention of non-offloaded layer needs - # to be modified while reloading - CPUOffloadedLayer = True - else: - tensor_tag = (-1, self.torch_tensor_count) - self.torch_tensor_count += 1 - self.tensor_tag_to_state[tensor_tag] = tensor + self.offload_layer_map[layer_id] = False - return tensor_tag + def fwd_step(self) -> int: + """ + Invoked before each layer forward. + """ + super().fwd_step() + if self.offload_layer_map.get(self.current_layer_id - 1, False): + self.layer_states[self.current_layer_id - 1].start_offload() - def tensor_pop(self, tensor_tag, **kwargs): - """Tensor pop.""" - global CPUOffloadedLayer + for layer in self.finish_offload_map[self.current_layer_id]: + self.layer_states[layer].release_activation_forward_gpu_memory() + return self.current_layer_id - assert tensor_tag in self.tensor_tag_to_state - tensor = self.tensor_tag_to_state.pop(tensor_tag) + def bwd_step(self, layer_num: int): + """ + Invoked before each layer backward. + """ + super().bwd_step(layer_num) - # Handling the quantized tensor case specially here - if isinstance(tensor, list): - # If it's a duplicated tensor, we don't need to locally - # write back a tensor as it would already be written - if tensor_tag in self.dereferencing_list: - self.dereferencing_list.remove(tensor_tag) - else: - self.fp8_tensor_object_map[tensor_tag].restore_from_saved(tensor) - tensor = self.fp8_tensor_object_map.pop(tensor_tag) + for layer in self.start_reload_map[layer_num]: + self.layer_states[layer].start_reload() - if self.double_buffering: - tensor._do_not_clear = True - self.tensor_tag_to_buf.pop(tensor_tag, None) - # the tensor should have been copied back in on_group_commit_backward() - # which invokes bulk_reload_group. - assert not isinstance(tensor, tuple) - return tensor +class ManualOffloadSynchronizer(OffloadSynchronizer): + """ + Manual implementation of OffloadSynchronizer, + all synchronization is done manually by the user by using + one of the following methods: + - start_offload_layer + - release_activation_forward_gpu_memory + - start_reload_layer + + This implementation is intended to be used in more complex trainigs workflows. + It is useful for example in pipeline parallelism. + """ - def bulk_offload_group(self, group_to_offload): - """Bulk offload group.""" - with torch.cuda.stream(self.d2h_stream): - for tensor_tag, state in self.tensor_tag_to_state.items(): - group_id, _ = tensor_tag - if group_id == group_to_offload: - assert not isinstance(state, tuple) - - is_quantized_tensor = isinstance(state, list) - - if is_quantized_tensor: - tensor_list = state - self.tensor_tag_to_state[tensor_tag] = [] - else: - tensor_list = [state] - - for tensor_on_device in tensor_list: - # `tensor_offloaded` is a hacky way of dealing with columnwise-only - # quantized tensors for CPU offloading. The complication is due to - # the `rowwise_data` being `None`. The offloading checker incorrectly - # returns `False` and the entire `state` ([None, columnwise_tensor]) - # is added to the tensor tag state dict. A better design would change - # how quantized tensors are kept track of in the offload handler. - # Currently at every stage it is ensured that a quantized tensor is a - # list whereas a non-quantized tensor is standalone object, which is - # not good! TODO(@sanandaraj5597) - tensor_offloaded = False - # if offload, return the reference to cpu copy - if self.tensor_need_offloading_checker(tensor_on_device): - tensor_offloaded = True - state = SynchronizedGroupOffloadHandler.offload(tensor_on_device) - if is_quantized_tensor: - if tensor_offloaded: - self.tensor_tag_to_state[tensor_tag].append(state) - else: - self.tensor_tag_to_state[tensor_tag].append(tensor_on_device) - else: - self.tensor_tag_to_state[tensor_tag] = state - - def synchronize_on_group_commit_forward(self, current_group): - """Synchronize on group commit forward.""" - global CPUOffloadedLayer - - # For the first group, kickstart the offload after we have - # the first compute completion - if current_group == 0: - self.d2h_stream.wait_stream(torch.cuda.current_stream()) - - if not self.double_buffer_created: - # Creating the first copy of double buffer for tensors that are offloaded - for tensor_tag, buf in self.tensor_tag_to_buf.items(): - if isinstance(buf, list): - for b in buf: - self.reload_double_buffer[0].append( - torch.empty_like(b) if self.double_buffering else None - ) - else: - self.reload_double_buffer[0].append( - torch.empty_like(buf) if self.double_buffering else None - ) - - self.bulk_offload_group(current_group) - - # Window map data structure helps us synchronize based on number - # of layers offloaded - if self.layer_window_map[self.offloaded_group_count] == current_group: - - # Stream synchronization both ways - self.d2h_stream.wait_stream(torch.cuda.current_stream()) - torch.cuda.current_stream().wait_stream(self.d2h_stream) - - # Time to free the activation memory after usage - for tensor_tag, tensor_buf in self.tensor_tag_to_buf.items(): - if tensor_tag[0] == self.offloaded_group_count: - if hasattr(tensor_buf, "needs_force_clear"): - # Need to clear activation tensor - sometimes references persist in the code. - # This is the case for example with the Float8TensorStorage class, - # which is saved directly inside the ctx while its internal tensors are - # saved inside save_for_backward. - tensor_buf.data = torch.Tensor() - # Release the pointer to the tensor - self.tensor_tag_to_buf[tensor_tag] = None - - # Time to offload the next group - if self.offloaded_group_count < (self.num_offload_group - 1): - self.bulk_offload_group(self.offloaded_group_count + 1) - - # Increment the offload group count to keep track - self.offloaded_group_count += 1 - - if current_group == (self.num_offload_group - 1): - CPUOffloadedLayer = False - - if not self.double_buffer_created: - # Creating second copy of double buffer for tensors that are offloaded - if current_group == (self.num_layers - 1): - for buf in self.reload_double_buffer[0]: - self.reload_double_buffer[1].append( - torch.empty_like(buf) if self.double_buffering else None - ) - self.double_buffer_created = True - - def on_group_commit_forward(self): - """This function will cause host device synchronization""" - # handle synchronization events - self.synchronize_on_group_commit_forward(self.current_group) - - super().on_group_commit_forward() - - def bulk_reload_group(self, group_to_reload): - """Bulk reload group.""" - assert group_to_reload < self.num_offload_group - - buffer_idx = 0 - double_buffer_idx = group_to_reload % 2 - - main_stream = torch.cuda.current_stream() - - with torch.cuda.stream(self.h2d_stream): - # move back tensors - for tensor_label, state in self.tensor_tag_to_state.items(): - group_id, _ = tensor_label - if group_id == group_to_reload: - - if isinstance(state, tuple): - if self.double_buffering: - reload_buffer = self.reload_double_buffer[double_buffer_idx][buffer_idx] - else: - with torch.cuda.stream(main_stream): - reload_buffer = torch.empty_like( - state[1], device=torch.cuda.current_device() - ) - - recovered_tensor = SynchronizedGroupOffloadHandler.reload( - state, True, reload_buffer - ) - buffer_idx = buffer_idx + 1 - self.tensor_tag_to_state[tensor_label] = recovered_tensor - elif isinstance(state, list): - tensor_list = [] - for state_tuple in state: - - if isinstance(state_tuple, tuple): - if self.double_buffering: - reload_buffer = self.reload_double_buffer[double_buffer_idx][ - buffer_idx - ] - else: - with torch.cuda.stream(main_stream): - reload_buffer = torch.empty_like( - state_tuple[1], device=torch.cuda.current_device() - ) - - tensor_list.append( - SynchronizedGroupOffloadHandler.reload( - state_tuple, - True, - reload_buffer, - ) - ) - buffer_idx = buffer_idx + 1 - else: - tensor_list.append(state_tuple) - - # No need to write back the duplicated tensor againn - # to the same location, this check ensures that - if tensor_label in self.dereferencing_list: - self.dereferencing_list.remove(tensor_label) - else: - _ = self.fp8_tensor_object_map[tensor_label].restore_from_saved( - tensor_list - ) - - if isinstance(self.fp8_tensor_object_map[tensor_label], Float8Tensor): - self.fp8_tensor_object_map[tensor_label]._transpose_invalid = ( - self.float8_transpose_cache_valid.pop(tensor_label) - ) - - self.tensor_tag_to_state[tensor_label] = self.fp8_tensor_object_map.pop( - tensor_label - ) - - def on_group_commit_backward(self): - # first decrement the current group. - # after last commit in forward, the group will +1; in backward it -1. - # Finally it should be decremented to 0. - self.current_group -= 1 - assert self.current_group >= 0 - - # Layer window data structure helps us to reload at right times - if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group: - - # Stream synchronization both ways - self.h2d_stream.wait_stream(torch.cuda.current_stream()) - torch.cuda.current_stream().wait_stream(self.h2d_stream) - - # Time to reload the next group - self.bulk_reload_group(self.offloaded_group_count - 1) - - # Decrease the offloading group counter - self.offloaded_group_count -= 1 if self.offloaded_group_count > 1 else 0 - - # Last group computation needs to wait till all the reloads complete - if self.current_group == 0: - torch.cuda.current_stream().wait_stream(self.h2d_stream) - self.offloaded_group_count = 0 + def start_offload_layer(self, layer_id: int): + """ + Start offloading of the layer. + Each tensor GPU->CPU copy is done asynchronously on the offload stream. + Start of each copy is started after tensor_push() is called on the current stream. + """ + self.layer_states[layer_id].start_offload() + + def release_activation_forward_gpu_memory(self, layer_id: int): + """ + Release memory of the activations of the layer. + It waits for the offload of the layer to finish. + """ + self.layer_states[layer_id].release_activation_forward_gpu_memory() + + def start_reload_layer(self, layer_id: int): + """ + Start reloading of the layer. + Each tensor reload is awaited to finish before tensor_pop() for that tensor is called on the current stream. + """ + self.layer_states[layer_id].start_reload() def get_cpu_offload_context( enabled: bool = False, - num_layers: int = 1, + num_layers: Optional[int] = 1, model_layers: int = 1, offload_activations: bool = True, offload_weights: bool = False, - double_buffering: bool = False, + double_buffering: bool = False, # pylint: disable=unused-argument + manual_synchronization: bool = False, + retain_pinned_cpu_buffers: bool = False, + offload_stream: Optional[torch.cuda.Stream] = None, ): """ - This function returns the CPU Offload context and the synchronizer function that needs to be - used after every transformer layer. Returns `nullcontext()` if offloading is not enabled. + CPU Offloading feature for seqeuences of layers. Can be used for arbitrary layers, not necessarily + for these provided by the TE. Usage: .. code-block:: python - cpu_offload_context, cpu_offload_synchronizer = get_cpu_offload_context(enabled=True) + cpu_offload_context, sync_function = get_cpu_offload_context(...) - with cpu_offload_context: - te_layer.forward(inp_tensor) - cpu_offload_synchronizer() + for _ in range(num_layers): + with cpu_offload_context: + x = layers[i].forward(x) + x = sync_function(x) Parameters ---------- enabled: bool, default = `False` When set to True, CPU Offloading functionality is enabled. num_layers: int, default = 1 - Determines the number of transformer layers - you want to offload activations/weights for. + Determines the number of layers + you want to offload activations/weights for. model_layers: int, default = 1 - Number of layers in the model that will be used under this context. + Number of layers in the model that will be used under this context. offload_activations: bool, default = `True` - When set to `True`, offloads the activations for the TE layer. + Deprecated. offload_weights: bool, default = `True` - When set to `True`, offloads the weights for the TE layer. + Deprecated. double_buffering: bool, default = `False` - When set to `True`, uses double buffering for offloading. + Deprecated. + retain_pinned_cpu_buffers: bool, default = `False` + If True, the pinned CPU buffers are retained after offloading + and reused for the next iteration. It is useful for cuda graphs capture. + manual_synchronization: bool, default = `False` + If True, the synchronization is done manually by the user. + Additional argument manual_controller is returned. See more in manual control section. + offload_stream: torch.cuda.Stream, default = `None` + If provided, the offload stream is used for offloading and reloading. + Otherwise, a new stream is allocated internally. It can be other than None + only if manual_synchronization is True. + + Manual synchronization + ---------- + By default, layers are offloaded/reloaded asynchronously + with respect to the current forward/backward stream with predefined synchronization, + to ensure that activation memory usage is equal to + `(num_layers - num_offloaded_layers) * T`, where `T` is the memory footprint of a layer. + + For more control over the offloading and reloading process, you can set `manual_synchronization=True`. + In this case, an additional argument, `manual_controller`, is returned. + + The `manual_controller` provides the following methods: + - `start_offload_layer(layer_id: int)` + - `release_activation_forward_gpu_memory(layer_id: int)` + - `start_reload_layer(layer_id: int)` + + If none of these methods are invoked for a given layer, that layer will not be offloaded or reloaded. + If `start_offload_layer()` is called for a layer, offload copies for that layer begin asynchronously on the offload stream. + + Since GPU activations must be kept in memory until the copy is finished, pointers to all activations are stored. + To release this memory, you need to call `release_activation_forward_gpu_memory(layer_id)`. + This method makes the current stream wait for an event recorded on the offload stream after all tensors from the layer have been offloaded. + + The `start_reload_layer()` method is used to start reloading a layer. + Each tensor reload is awaited to finish before `tensor_pop()` for that tensor is called on the current stream. + + You can provide an `offload_stream` to be used for offload and reload operations. + This allows for more detailed synchronization, such as delaying the start of offloading. + + Example: + .. code-block:: python + offload_stream = torch.cuda.Stream() + cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context( + enabled=True, model_layers=num_layers, manual_synchronization=True, offload_stream=offload_stream) + + for i in range(num_layers): + with cpu_offload_context: + out[i] = layers[i].forward(inp[i]) + out[i] = sync_function(out[i]) + manual_controller.start_offload_layer(i) + + offload_stream.synchronize() + for i in range(num_layers): + manual_controller.release_activation_forward_gpu_memory(i) + + for i in range(num_layers - 1, -1, -1): + manual_controller.start_reload_layer(i) + + offload_stream.synchronize() + for i in range(num_layers): + out[i].sum().backward() + + V1 code path + ---------- + If you want to use the v1 code path for offloading, + please set the environment variable NVTE_CPU_OFFLOAD_V1 to 1. """ + if NVTE_CPU_OFFLOAD_V1: + return v1_code_path.get_cpu_offload_context( + enabled=enabled, + num_layers=num_layers, + model_layers=model_layers, + offload_activations=offload_activations, + offload_weights=offload_weights, + double_buffering=double_buffering, + ) if not offload_weights and not offload_activations: raise ValueError( @@ -703,8 +753,6 @@ def get_cpu_offload_context( ) if offload_weights: - import warnings - warnings.warn( "Offloading weights is deprecated. Using offload_weights=True does not have any" " effect.", @@ -713,26 +761,100 @@ def get_cpu_offload_context( # Weights offloading is deprecated but we maintain backward compatibility by doing nothing. if not offload_activations: - return nullcontext(), lambda x: x + return contextlib.nullcontext(), lambda x: x - def tensor_need_offloading_checker_activations(tensor): - return hasattr(tensor, "activation_offloading") - - tensor_need_offloading_checker = tensor_need_offloading_checker_activations + if TEDebugState.debug_enabled: + raise RuntimeError("CPU offload is not supported in debug mode.") - cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler( - num_offload_group=num_layers, - num_model_group=model_layers, - tensor_need_offloading_checker=tensor_need_offloading_checker, - double_buffering=double_buffering, - ) + if not manual_synchronization: + assert ( + num_layers <= model_layers - 1 + ), "Cannot offload all layers without manual synchronization - last layer is not offloaded." + if num_layers == model_layers - 1: + warnings.warn( + "Offloading num_layers == model_layers - 1 is not recommended, it prevents" + " overlapping of computation and offload/reload." + ) + + assert ( + offload_stream is None or manual_synchronization + ), "offload_stream can be provided only if manual_synchronization is True" + + if manual_synchronization: + offload_synchronizer = ManualOffloadSynchronizer( + model_layers, retain_pinned_cpu_buffers, offload_stream + ) + else: + offload_synchronizer = DefaultOffloadSynchronizer( + model_layers, + num_layers, + retain_pinned_cpu_buffers, + offload_stream, + ) - def group_prefetch_offload_commit_async(tensor): - return group_prefetch_offload_commit(tensor, cpu_offload_handler) + class _CpuOffloadContext(contextlib.ContextDecorator): + def __init__(self): + self.current_layer = None + self.previous_offload_synchronizer = None + self.offload_synchronizer = offload_synchronizer + + self.inside_context = False + + def __enter__(self): + assert ( + self.inside_context is False + ), "Offloading context was entered without synchronization function being called." + self.inside_context = True + self._hooks_ctx = saved_tensors_hooks( + offload_synchronizer.push_tensor, offload_synchronizer.pop_tensor + ) + self._hooks_ctx.__enter__() + global OFFLOAD_SYNCHRONIZER + self.previous_offload_synchronizer = OFFLOAD_SYNCHRONIZER + OFFLOAD_SYNCHRONIZER = offload_synchronizer + self.current_layer = offload_synchronizer.fwd_step() + return self + + def __exit__(self, *args): + self._hooks_ctx.__exit__(*args) + global OFFLOAD_SYNCHRONIZER + OFFLOAD_SYNCHRONIZER = self.previous_offload_synchronizer + self.inside_context = False + + def synchronization_function(self, tensor): + """ + This function is used to catch the backward pass of the model. + """ + assert tensor.requires_grad is True + assert self.current_layer is not None + cur_layer = self.current_layer + assert ( + self.inside_context is False + ), "Synchronization function was called without offloading context being entered." + + def hook(_): + # offload_synchronizer.finish_part_of_bwd needs + # to be called after every backward pass - there may be + # more than one in pipeline parallelism. + torch.autograd.variable.Variable._execution_engine.queue_callback( + offload_synchronizer.finish_part_of_bwd + ) + offload_synchronizer.bwd_step(cur_layer) + + tensor.grad_fn.register_prehook(hook) + return tensor + + cpu_offload_context = _CpuOffloadContext() if enabled: + if manual_synchronization: + return ( + cpu_offload_context, + cpu_offload_context.synchronization_function, + offload_synchronizer, + ) return ( - CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler), - group_prefetch_offload_commit_async, + cpu_offload_context, + cpu_offload_context.synchronization_function, ) - return nullcontext(), group_prefetch_offload_commit_async + return contextlib.nullcontext(), lambda x: x diff --git a/transformer_engine/pytorch/cpu_offload_v1.py b/transformer_engine/pytorch/cpu_offload_v1.py new file mode 100644 index 0000000000..9f904864ab --- /dev/null +++ b/transformer_engine/pytorch/cpu_offload_v1.py @@ -0,0 +1,743 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Functionality for CPU offloading of tensors saved for backward pass.""" +from __future__ import annotations +from contextlib import nullcontext +from typing import Any, Dict, Optional + +import torch + +from transformer_engine.debug.pytorch.debug_state import TEDebugState +from .quantized_tensor import QuantizedTensorStorage +from .tensor.float8_tensor import Float8Tensor + +__all__ = ["get_cpu_offload_context"] + +CPUOffloadEnabled = False +CPUOffloadedLayer = False + + +def mark_activation_offload(*tensors): + """Set the type of the offloading needed for a tensor.""" + if TEDebugState.debug_enabled: + raise RuntimeError("CPU offload is not supported in debug mode.") + + for tensor in tensors: + if tensor is None: + continue + if type(tensor) in [torch.Tensor, torch.nn.Parameter]: + tensor.activation_offloading = True + else: + data_tensors = tensor.get_data_tensors() + for tensor in data_tensors: + if tensor is not None: + tensor.activation_offloading = True + # This is a hack to force clear the tensor after it is offloaded. + # It is needed, because .*TensorStorage classes are saved in the ctx, + # and they contain the reference to their data tensors. + tensor.needs_force_clear = True + + +def is_cpu_offload_enabled() -> bool: + """Check if CPU offloading is currently enabled.""" + return CPUOffloadEnabled + + +def is_current_layer_offloaded() -> bool: + """Check if current layers is being offloaded.""" + return CPUOffloadedLayer + + +class CpuOffloadSavedTensorHook: + """Contex-manager that executes a pair of pack/unpack hooks for saved tensors. + + In this context, the ``on_save_for_backward`` method will be called every time + a tensor is saved for backward (this includes intermediary results saved using + :func:`~torch.autograd.function._ContextMethodMixin.save_for_backward` but + also those recorded by a PyTorch-defined operation). + + The ``on_get_saved_tensors`` method will be called when the backward function + of this op attempts to retrieve the saved tensor from context (this includes + :func: `torch.Tensor.backward()` or :func: `torch.autograd.grad()`. It takes the + as input the return value of the ``on_save_for_backward``, and is meant to return + an identical copy of the tensor being saved by ``on_save_for_backward`` in terms of + size, device and element values. + + Example: + + >>> import torch + >>> from typing import Any + >>> + >>> class DummyHook(CpuOffloadSavedTensorHook): + ... + ... def on_save_for_backward(self, tensor: torch.Tensor) -> Any: + ... logging.info("On save", tensor) + ... return (tensor,) + ... + ... def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: + ... logging.info("On get", saved_state) + ... tensor, = saved_state + ... return tensor + ... + >>> a = torch.ones(5, requires_grad=True) + >>> b = torch.ones(5, requires_grad=True) * 2 + >>> with DummyHook(): + ... y = a * b + ... + On save tensor([1., 1., 1., 1., 1.], requires_grad=True) + On save tensor([2., 2., 2., 2., 2.], grad_fn=) + >>> y.sum().backward() + On get (tensor([1., 1., 1., 1., 1.], requires_grad=True),) + On get (tensor([2., 2., 2., 2., 2.], grad_fn=),) + + """ + + def __init__(self) -> None: + self.inside_context = False + + def __enter__(self): + global CPUOffloadEnabled + CPUOffloadEnabled = True + + self.inside_context = True + torch._C._autograd._push_saved_tensors_default_hooks( + self.on_save_for_backward, self.on_get_saved_tensor + ) + + def __exit__(self, *args: Any): + global CPUOffloadEnabled + CPUOffloadEnabled = False + + self.inside_context = False + torch._C._autograd._pop_saved_tensors_default_hooks() + + def on_save_for_backward(self, tensor: torch.Tensor) -> Any: + """On save for backward.""" + raise NotImplementedError( + "`on_save_for_backward: Callable[[torch.Tensor], Any]`" + "is not implemented in CpuOffloadHook class. Inherit " + "this class and implement your custom hooks" + ) + + def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: + """On get saved tensor.""" + raise NotImplementedError( + "`on_get_saved_tensors: Callable[[Any], torch.Tensor]`" + "is not implemented in CpuOffloadHook class. Inherit " + "this class and implement your custom hooks" + ) + + +class CpuOffloadHookWithOffloadHandler(CpuOffloadSavedTensorHook): + """Context-manager that offloads/recovers tensors through an offload hander. + + The hook just offloads/recovers the tensor object to the handler through `tensor_push` + and `tensor_pop` interface. How the offload-handler manages the offloading, recovering + or prefetching timing is transparent to this hook. + """ + + def __init__( + self, + offload_handler: OffloadHandler, + handler_extra_kwargs: Optional[Dict[str, Any]] = None, + debug: bool = False, + ) -> None: + if handler_extra_kwargs is None: + handler_extra_kwargs = {} + self.debug: bool = debug + self.offload_handler: OffloadHandler = offload_handler + self.handler_extra_kwargs: Dict[str, Any] = handler_extra_kwargs + super().__init__() + + def on_save_for_backward(self, tensor: torch.Tensor) -> Any: + retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs) + return retrieve_identifier + + def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: + tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs) + return tensor + + +class OffloadHandler: + """A base class for CPU offload-handler.""" + + def __init__(self) -> None: + pass + + def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: + """Tensor push.""" + raise NotImplementedError( + "`tensor_push is not implented in OffloadHandler class. " + "Inherit this class and implement your custom tensor_push." + ) + + def tensor_pop(self, tensor_tag: Any, **kwargs): + """Tensor pop.""" + raise NotImplementedError( + "`tensor_pop is not implented in OffloadHandler class. " + "Inherit this class and implement your custom tensor_pop." + ) + + +class GroupCommitFunction(torch.autograd.Function): + """this is a dummy op with output identical to input. + However, it is necessary for marking a timepoint for offload handler to + accomplish all synchronizations. Implementing it as a function is necessary + because we need to actions in both forward and backward. + """ + + @staticmethod + def forward(ctx, tensor, cpu_offload_handler): + # pylint: disable=missing-function-docstring + cpu_offload_handler.on_group_commit_forward() + ctx.cpu_offload_handler = cpu_offload_handler + # return the identical tensor + return tensor + + @staticmethod + def backward(ctx, grad_output): + # pylint: disable=missing-function-docstring + cpu_offload_handler = ctx.cpu_offload_handler + cpu_offload_handler.on_group_commit_backward() + return grad_output, None + + +group_prefetch_offload_commit = GroupCommitFunction.apply + + +class SynchronizedGroupOffloadHandler(OffloadHandler): + """Offload Handler that offloads/reloads in a synchronized way. + The device-to-host and host-to-device copying happen in the same stream + as the computation kernels, thus the copying will block computation. + """ + + def __init__( + self, num_offload_group, tensor_need_offloading_checker=(lambda _: True), debug=False + ) -> None: + super().__init__() + + self.num_offload_group = num_offload_group + self.tensor_need_offloading_checker = tensor_need_offloading_checker + self.debug = debug + + self.groupid_reset() + + def groupid_reset(self): + """Groupid reset.""" + # Data structures to label saved tensors and book-keep their cpu copies. + # Currently, on push, create a new cpu tensor and copies; on pop, copies + # the tensor back to gpu and deletes the cpu tensor. + # These will increment whenever `group_commit()` is invoked + self.current_group, self.tensor_count_current_group = (0, 0) + self.torch_tensor_count = 0 + self.tensor_tag_to_state = {} + + def on_group_commit_forward(self): + """On group commit forward.""" + # finishing up with updating current group and tensor count + self.current_group += 1 # increment + self.tensor_count_current_group = 0 # reset + + def on_group_commit_backward(self): + """On group commit backward.""" + self.current_group -= 1 + assert self.current_group >= 0 + + @staticmethod + def offload(src_tensor, pin_memory=True): + """Offload.""" + + cpu_backup = torch.empty( + src_tensor.size(), + dtype=src_tensor.dtype, + layout=src_tensor.layout, + device="cpu", + pin_memory=pin_memory, + ) + + cpu_backup.copy_(src_tensor, non_blocking=pin_memory) + state = (src_tensor.device, cpu_backup) + return state + + @staticmethod + def reload(state, non_blocking=None, copy_buffer=None): + """Reload.""" + dev, cpu_backup = state + if non_blocking is None: + non_blocking = cpu_backup.is_pinned() + + if copy_buffer is None: + return cpu_backup.to(dev, non_blocking=non_blocking) + + assert cpu_backup.size() == copy_buffer.size(), "Can't copy two buffers of different sizes!" + + copy_buffer.copy_(cpu_backup, non_blocking=non_blocking) + + return copy_buffer + + def tensor_push(self, tensor: torch.Tensor, **kwargs): + """Tensor push.""" + # obtain a unique tensor tag + tensor_tag = (self.current_group, self.tensor_count_current_group) + self.tensor_count_current_group += 1 + assert tensor_tag not in self.tensor_tag_to_state + if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker( + tensor + ): + state = SynchronizedGroupOffloadHandler.offload(tensor) + self.tensor_tag_to_state[tensor_tag] = state + else: + # will be offloaded together after group commit + self.tensor_tag_to_state[tensor_tag] = tensor + + return tensor_tag + + def tensor_pop(self, tensor_tag, **kwargs): + """Tensor pop.""" + assert tensor_tag in self.tensor_tag_to_state + state = self.tensor_tag_to_state.pop(tensor_tag) + if isinstance(state, tuple): + tensor = SynchronizedGroupOffloadHandler.reload(state) + else: + tensor = state + return tensor + + +class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler): + """Compared to synchronize, this uses more memory because of the buffer but + achieves better performance due to the overlapping. D2h and h2d copying are + completely hidden behind computation if computation time of a layer is longer + than host-device communication time. Bulk offloading with delay and bulk reloading + with prefetch are implemented.""" + + def __init__( + self, + num_offload_group, # must be <= actual number of groups (number of commits) + num_model_group, + tensor_need_offloading_checker=(lambda t: True), + double_buffering=False, + debug=False, + ) -> None: + super().__init__( + num_offload_group=num_offload_group, + tensor_need_offloading_checker=tensor_need_offloading_checker, + debug=debug, + ) + # Number of layers in the model + self.num_layers = num_model_group + # Data Structure to maintain reference to activation tensors + self.tensor_tag_to_buf = {} + # Data structure to hold the FP8/MXFP8 tensor objects + self.fp8_tensor_object_map = {} + self.float8_transpose_cache_valid = {} + self.dereferencing_list = [] + # Tracking the number of layers offloaded + self.offloaded_group_count = 0 + # Core data structure that decides the window for offloading + self.layer_window_map = {} + + # Data structures fo double buffered reloading + self.double_buffering = double_buffering + self.reload_double_buffer = [[], []] + self.double_buffer_created = False + + # Logic to make offloading load balance across computation + # for optimal CPU/GPU interconnect usage + constant = 0 + for i in range(self.num_offload_group): + self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1 + if i < (self.num_layers % self.num_offload_group): + self.layer_window_map[i] += i + 1 + constant = i + 1 + else: + self.layer_window_map[i] += constant + + # allocate streams and events for synchronization + self.d2h_stream = torch.cuda.Stream() + self.h2d_stream = torch.cuda.Stream() + + def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: + global CPUOffloadedLayer + + torch_stray_tensor = isinstance( + tensor, + ( + torch._subclasses.fake_tensor.FakeTensor, + torch._subclasses.functional_tensor.FunctionalTensor, + ), + ) + + is_quantized_tensor = isinstance(tensor, QuantizedTensorStorage) + + if not torch_stray_tensor: + + # obtain a unique tensor tag + tensor_tag = (self.current_group, self.tensor_count_current_group) + self.tensor_count_current_group += 1 + + assert tensor_tag not in self.tensor_tag_to_state + + if is_quantized_tensor: + tensor_list, _ = tensor.prepare_for_saving() + + self.tensor_tag_to_state[tensor_tag] = [] + self.tensor_tag_to_buf[tensor_tag] = [] + + # Added support for de-duplicating FP8 param tensors + for _, value in self.fp8_tensor_object_map.items(): + if tensor is value: + self.dereferencing_list.append(tensor_tag) + break + + self.fp8_tensor_object_map[tensor_tag] = tensor + if isinstance(tensor, Float8Tensor): + self.float8_transpose_cache_valid[tensor_tag] = getattr( + tensor, "_transpose_invalid" + ) + else: + tensor_list = [tensor] + + for t in tensor_list: + if is_quantized_tensor: + self.tensor_tag_to_state[tensor_tag].append(t) + else: + self.tensor_tag_to_state[tensor_tag] = t + + if ( + self.current_group < self.num_offload_group + and self.tensor_need_offloading_checker(t) + ): + if is_quantized_tensor: + self.tensor_tag_to_buf[tensor_tag].append(t) + # Need to clear the internal data reference for the quantized tensors + tensor.clear() + else: + self.tensor_tag_to_buf[tensor_tag] = t + + # Needed to differentiate non offloaded layer's attention + # QKV layout of attention of non-offloaded layer needs + # to be modified while reloading + CPUOffloadedLayer = True + else: + tensor_tag = (-1, self.torch_tensor_count) + self.torch_tensor_count += 1 + self.tensor_tag_to_state[tensor_tag] = tensor + + return tensor_tag + + def tensor_pop(self, tensor_tag, **kwargs): + """Tensor pop.""" + global CPUOffloadedLayer + + assert tensor_tag in self.tensor_tag_to_state + tensor = self.tensor_tag_to_state.pop(tensor_tag) + + # Handling the quantized tensor case specially here + if isinstance(tensor, list): + # If it's a duplicated tensor, we don't need to locally + # write back a tensor as it would already be written + if tensor_tag in self.dereferencing_list: + self.dereferencing_list.remove(tensor_tag) + else: + self.fp8_tensor_object_map[tensor_tag].restore_from_saved(tensor) + tensor = self.fp8_tensor_object_map.pop(tensor_tag) + + if self.double_buffering: + tensor._do_not_clear = True + + self.tensor_tag_to_buf.pop(tensor_tag, None) + # the tensor should have been copied back in on_group_commit_backward() + # which invokes bulk_reload_group. + assert not isinstance(tensor, tuple) + return tensor + + def bulk_offload_group(self, group_to_offload): + """Bulk offload group.""" + with torch.cuda.stream(self.d2h_stream): + for tensor_tag, state in self.tensor_tag_to_state.items(): + group_id, _ = tensor_tag + if group_id == group_to_offload: + assert not isinstance(state, tuple) + + is_quantized_tensor = isinstance(state, list) + + if is_quantized_tensor: + tensor_list = state + self.tensor_tag_to_state[tensor_tag] = [] + else: + tensor_list = [state] + + for tensor_on_device in tensor_list: + # `tensor_offloaded` is a hacky way of dealing with columnwise-only + # quantized tensors for CPU offloading. The complication is due to + # the `rowwise_data` being `None`. The offloading checker incorrectly + # returns `False` and the entire `state` ([None, columnwise_tensor]) + # is added to the tensor tag state dict. A better design would change + # how quantized tensors are kept track of in the offload handler. + # Currently at every stage it is ensured that a quantized tensor is a + # list whereas a non-quantized tensor is standalone object, which is + # not good! TODO(@sanandaraj5597) + tensor_offloaded = False + # if offload, return the reference to cpu copy + if self.tensor_need_offloading_checker(tensor_on_device): + tensor_offloaded = True + state = SynchronizedGroupOffloadHandler.offload(tensor_on_device) + if is_quantized_tensor: + if tensor_offloaded: + self.tensor_tag_to_state[tensor_tag].append(state) + else: + self.tensor_tag_to_state[tensor_tag].append(tensor_on_device) + else: + self.tensor_tag_to_state[tensor_tag] = state + + def synchronize_on_group_commit_forward(self, current_group): + """Synchronize on group commit forward.""" + global CPUOffloadedLayer + + # For the first group, kickstart the offload after we have + # the first compute completion + if current_group == 0: + self.d2h_stream.wait_stream(torch.cuda.current_stream()) + + if not self.double_buffer_created: + # Creating the first copy of double buffer for tensors that are offloaded + for tensor_tag, buf in self.tensor_tag_to_buf.items(): + if isinstance(buf, list): + for b in buf: + self.reload_double_buffer[0].append( + torch.empty_like(b) if self.double_buffering else None + ) + else: + self.reload_double_buffer[0].append( + torch.empty_like(buf) if self.double_buffering else None + ) + + self.bulk_offload_group(current_group) + + # Window map data structure helps us synchronize based on number + # of layers offloaded + if self.layer_window_map[self.offloaded_group_count] == current_group: + + # Stream synchronization both ways + self.d2h_stream.wait_stream(torch.cuda.current_stream()) + torch.cuda.current_stream().wait_stream(self.d2h_stream) + + # Time to free the activation memory after usage + for tensor_tag, tensor_buf in self.tensor_tag_to_buf.items(): + if tensor_tag[0] == self.offloaded_group_count: + if hasattr(tensor_buf, "needs_force_clear"): + # Need to clear activation tensor - sometimes references persist in the code. + # This is the case for example with the Float8TensorStorage class, + # which is saved directly inside the ctx while its internal tensors are + # saved inside save_for_backward. + tensor_buf.data = torch.Tensor() + # Release the pointer to the tensor + self.tensor_tag_to_buf[tensor_tag] = None + + # Time to offload the next group + if self.offloaded_group_count < (self.num_offload_group - 1): + self.bulk_offload_group(self.offloaded_group_count + 1) + + # Increment the offload group count to keep track + self.offloaded_group_count += 1 + + if current_group == (self.num_offload_group - 1): + CPUOffloadedLayer = False + + if not self.double_buffer_created: + # Creating second copy of double buffer for tensors that are offloaded + if current_group == (self.num_layers - 1): + for buf in self.reload_double_buffer[0]: + self.reload_double_buffer[1].append( + torch.empty_like(buf) if self.double_buffering else None + ) + self.double_buffer_created = True + + def on_group_commit_forward(self): + """This function will cause host device synchronization""" + # handle synchronization events + self.synchronize_on_group_commit_forward(self.current_group) + + super().on_group_commit_forward() + + def bulk_reload_group(self, group_to_reload): + """Bulk reload group.""" + assert group_to_reload < self.num_offload_group + + buffer_idx = 0 + double_buffer_idx = group_to_reload % 2 + + main_stream = torch.cuda.current_stream() + + with torch.cuda.stream(self.h2d_stream): + # move back tensors + for tensor_label, state in self.tensor_tag_to_state.items(): + group_id, _ = tensor_label + if group_id == group_to_reload: + + if isinstance(state, tuple): + if self.double_buffering: + reload_buffer = self.reload_double_buffer[double_buffer_idx][buffer_idx] + else: + with torch.cuda.stream(main_stream): + reload_buffer = torch.empty_like( + state[1], device=torch.cuda.current_device() + ) + + recovered_tensor = SynchronizedGroupOffloadHandler.reload( + state, True, reload_buffer + ) + buffer_idx = buffer_idx + 1 + self.tensor_tag_to_state[tensor_label] = recovered_tensor + elif isinstance(state, list): + tensor_list = [] + for state_tuple in state: + + if isinstance(state_tuple, tuple): + if self.double_buffering: + reload_buffer = self.reload_double_buffer[double_buffer_idx][ + buffer_idx + ] + else: + with torch.cuda.stream(main_stream): + reload_buffer = torch.empty_like( + state_tuple[1], device=torch.cuda.current_device() + ) + + tensor_list.append( + SynchronizedGroupOffloadHandler.reload( + state_tuple, + True, + reload_buffer, + ) + ) + buffer_idx = buffer_idx + 1 + else: + tensor_list.append(state_tuple) + + # No need to write back the duplicated tensor againn + # to the same location, this check ensures that + if tensor_label in self.dereferencing_list: + self.dereferencing_list.remove(tensor_label) + else: + _ = self.fp8_tensor_object_map[tensor_label].restore_from_saved( + tensor_list + ) + + if isinstance(self.fp8_tensor_object_map[tensor_label], Float8Tensor): + self.fp8_tensor_object_map[tensor_label]._transpose_invalid = ( + self.float8_transpose_cache_valid.pop(tensor_label) + ) + + self.tensor_tag_to_state[tensor_label] = self.fp8_tensor_object_map.pop( + tensor_label + ) + + def on_group_commit_backward(self): + # first decrement the current group. + # after last commit in forward, the group will +1; in backward it -1. + # Finally it should be decremented to 0. + self.current_group -= 1 + assert self.current_group >= 0 + + # Layer window data structure helps us to reload at right times + if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group: + + # Stream synchronization both ways + self.h2d_stream.wait_stream(torch.cuda.current_stream()) + torch.cuda.current_stream().wait_stream(self.h2d_stream) + + # Time to reload the next group + self.bulk_reload_group(self.offloaded_group_count - 1) + + # Decrease the offloading group counter + self.offloaded_group_count -= 1 if self.offloaded_group_count > 1 else 0 + + # Last group computation needs to wait till all the reloads complete + if self.current_group == 0: + torch.cuda.current_stream().wait_stream(self.h2d_stream) + self.offloaded_group_count = 0 + + +def get_cpu_offload_context( + enabled: bool = False, + num_layers: int = 1, + model_layers: int = 1, + offload_activations: bool = True, + offload_weights: bool = False, + double_buffering: bool = False, +): + """ + This function returns the CPU Offload context and the synchronizer function that needs to be + used after every transformer layer. Returns `nullcontext()` if offloading is not enabled. + + Usage: + + .. code-block:: python + + cpu_offload_context, cpu_offload_synchronizer = get_cpu_offload_context(enabled=True) + + with cpu_offload_context: + te_layer.forward(inp_tensor) + cpu_offload_synchronizer() + + Parameters + ---------- + enabled: bool, default = `False` + When set to True, CPU Offloading functionality is enabled. + num_layers: int, default = 1 + Determines the number of transformer layers + you want to offload activations/weights for. + model_layers: int, default = 1 + Number of layers in the model that will be used under this context. + offload_activations: bool, default = `True` + When set to `True`, offloads the activations for the TE layer. + offload_weights: bool, default = `True` + When set to `True`, offloads the weights for the TE layer. + double_buffering: bool, default = `False` + When set to `True`, uses double buffering for offloading. + + """ + + if not offload_weights and not offload_activations: + raise ValueError( + "CPU Offloading is enabled while it is not " + "mentioned what to offload (weights/activations)" + ) + + if offload_weights: + import warnings + + warnings.warn( + "Offloading weights is deprecated. Using offload_weights=True does not have any" + " effect.", + DeprecationWarning, + ) + + # Weights offloading is deprecated but we maintain backward compatibility by doing nothing. + if not offload_activations: + return nullcontext(), lambda x: x + + def tensor_need_offloading_checker_activations(tensor): + return hasattr(tensor, "activation_offloading") + + tensor_need_offloading_checker = tensor_need_offloading_checker_activations + + cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler( + num_offload_group=num_layers, + num_model_group=model_layers, + tensor_need_offloading_checker=tensor_need_offloading_checker, + double_buffering=double_buffering, + ) + + def group_prefetch_offload_commit_async(tensor): + return group_prefetch_offload_commit(tensor, cpu_offload_handler) + + if enabled: + return ( + CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler), + group_prefetch_offload_commit_async, + ) + return nullcontext(), group_prefetch_offload_commit_async diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f336a743da..1a56a06da3 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -41,7 +41,7 @@ from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing -from ..cpu_offload import is_cpu_offload_enabled +from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..quantized_tensor import ( @@ -135,6 +135,9 @@ def forward( else: inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) + if cpu_offloading: + start_offload(*inputmats) + # Initialize weights weights_fp8: list if fp8: @@ -196,6 +199,9 @@ def forward( for i in range(num_gemms): weight_quantizers[i].calibrate(weights[i]) + if cpu_offloading: + mark_not_offload(*weights_fp8, *weights) + if is_grad_enabled: ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 20a67cba48..4ed3ebb73f 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -66,10 +66,15 @@ from ...debug.pytorch.debug_state import TEDebugState from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..cpu_offload import ( + is_cpu_offload_enabled, + start_offload, + mark_not_offload, + mark_activation_offload, +) from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from ..export import is_in_onnx_export_mode, assert_warmed_up -from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ..cpp_extensions import ( general_gemm, @@ -158,6 +163,9 @@ def forward( ln_bias = cast_if_needed(ln_bias, activation_dtype) nvtx_range_pop(f"{nvtx_label}.norm_input_cast") + if is_cpu_offload_enabled(): + start_offload(inputmat) + tp_world_size = get_distributed_world_size(tp_group) weight_requires_grad = weight.requires_grad @@ -434,8 +442,14 @@ def forward( nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") if cpu_offloading: + mark_not_offload( + weightmat, + weight, + bias, + ln_weight, + ln_bias, + ) ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad") - if ctx.grad_added_to_main_grad: # If you are passing torch.nn.Parameter through the Torch hooks, you will # get back torch.Tensor. Torch rips off the Parameter wrapper. @@ -542,6 +556,7 @@ def backward( mu, rsigma, ) = restore_from_saved(ctx.tensor_objects, saved_tensors) + # Delete the references to tensor objects once they've been consumed # by the `restore_from_saved` method to construct back the actual tensors. ctx.tensor_objects = None diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index a358ae7ddf..c29775c926 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -69,7 +69,12 @@ from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ._common import apply_normalization, WeightGradStore -from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ..cpu_offload import ( + is_cpu_offload_enabled, + start_offload, + mark_not_offload, + mark_activation_offload, +) from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -235,6 +240,8 @@ def forward( ln_weight = cast_if_needed(ln_weight, activation_dtype) if ln_bias is not None: ln_bias = cast_if_needed(ln_bias, activation_dtype) + if is_cpu_offload_enabled(): + start_offload(inputmat) tp_world_size = get_distributed_world_size(tp_group) backwards_needs_fc1_input = is_grad_enabled and fc1_weight.requires_grad @@ -577,6 +584,18 @@ def forward( clear_tensor_data(act_out) act_out = None + if cpu_offloading: + mark_not_offload( + ln_weight, + ln_bias, + fc1_weight_final, + fc1_weight, + fc1_bias, + fc2_weight_final, + fc2_weight, + fc2_bias, + ) + tensors_to_save, tensor_objects = prepare_for_saving( inputmat, ln_weight, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 46b9dbd85b..00b78995fe 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -68,7 +68,12 @@ from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.utils import is_custom from ..export import is_in_onnx_export_mode, assert_warmed_up -from ..cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ..cpu_offload import ( + is_cpu_offload_enabled, + start_offload, + mark_not_offload, + mark_activation_offload, +) from ...debug.pytorch.debug_state import TEDebugState __all__ = ["Linear"] @@ -229,6 +234,9 @@ def forward( else: inputmat = cast_if_needed(inp, activation_dtype) # Cast for AMP inputmat_total = inputmat + + if is_cpu_offload_enabled(): + start_offload(inputmat) nvtx_range_pop(f"{nvtx_label}.input_cast_comm") # ------------------------------------------------------ # Input tensor is ready for GEMM... @@ -417,6 +425,7 @@ def forward( # weights if weights are externally touched outside this module ctx.weight_object = weight + mark_not_offload(weight, weightmat, bias) # TODO(ksivamani): Check memory usage tensors_to_save, tensor_objects = prepare_for_saving( saved_inputmat, diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 18f7e2031a..73b312ec28 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -372,9 +372,9 @@ def _initialize_state( """ dtype = self.name_to_dtype_map[state_name] if store_param_remainders: - data = torch.zeros_like(param, dtype=torch.int16) + data = torch.zeros(param.shape, dtype=torch.int16, device=param.device) else: - data = torch.empty_like(param, dtype=dtype) + data = torch.empty(param.shape, dtype=dtype, device=param.device) if zero_buffer: data.zero_() diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 7d49e3964f..c830b19e9f 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -9,6 +9,7 @@ import abc import copy import warnings +import math import torch from torch.utils._pytree import tree_map @@ -20,6 +21,11 @@ _stride_from_shape, ) +_quantized_tensor_cpu_supported_ops = ( + torch.ops.aten.empty_like.default, + torch.ops.aten.copy_.default, +) + class QuantizedTensorStorage: r"""Base class for all *TensorStorage classes. @@ -35,7 +41,7 @@ class QuantizedTensorStorage: XTensorStorage should contain all data members needed to implement the functionality of the tensor, while XTensor should only implement the functionality needed - to behave like regular torch.Tensor (liek __torch_dispatch__).""" + to behave like regular torch.Tensor (like __torch_dispatch__).""" _quantizer: Optional[Quantizer] @@ -63,6 +69,12 @@ def update_usage( f"{self.__class__.__name__} class does not implement update_usage function" ) + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the tensor""" + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement get_usages function" + ) + def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]: """Prepare the tensor base for saving for backward""" raise NotImplementedError( @@ -128,6 +140,7 @@ def prepare_for_saving( t, t_obj = tensor.prepare_for_saving() tensor_list.extend(t) tensor_objects_list.append(t_obj) + return tensor_list, tensor_objects_list @@ -314,6 +327,13 @@ def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-a """Returns whether or not given tensor can be quantized""" return True + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the quantizer""" + return { + "rowwise": self.rowwise_usage, + "columnwise": self.columnwise_usage, + } + class QuantizedTensor(torch.Tensor): """Abstract base class for tensor with quantized data @@ -325,7 +345,14 @@ class QuantizedTensor(torch.Tensor): """ - def __new__(cls, shape: Iterable[int], dtype: torch.dtype, *, requires_grad: bool = False): + def __new__( + cls, + shape: Iterable[int], + dtype: torch.dtype, + *, + requires_grad: bool = False, + device: Optional[torch.device] = None, + ): # We are assuming only contiguous tensors stride = _stride_from_shape(shape) instance = torch.Tensor._make_wrapper_subclass( @@ -336,7 +363,7 @@ def __new__(cls, shape: Iterable[int], dtype: torch.dtype, *, requires_grad: boo dtype=dtype, layout=torch.strided, requires_grad=requires_grad, - device=torch.cuda.current_device(), + device=torch.cuda.current_device() if device is None else device, ) return instance @@ -366,6 +393,9 @@ def detach(self) -> QuantizedTensor: def clear(self): """Deallocate this tensor's memory. Typically not needed and must be used carefully""" + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement clear function" + ) def __repr__(self, *, tensor_contents=None) -> str: return f"{self.__class__.__name__}(data={self.dequantize(dtype=self.dtype)})" @@ -407,6 +437,26 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == torch.ops.aten.copy_.default: dst = args[0] src = args[1] + if ( + isinstance(dst, QuantizedTensor) + and isinstance(src, QuantizedTensor) + and type(dst._quantizer) is type(src._quantizer) + and set(src.get_usages().keys()) == set(dst.get_usages().keys()) + and all( + src.get_usages()[usage] == dst.get_usages()[usage] + for usage in src.get_usages().keys() + ) + ): + + dst_tensors, dst_tensor_obj = dst.prepare_for_saving() + src_tensors, src_tensor_obj = src.prepare_for_saving() + for dst_tensor, src_tensor in zip(dst_tensors, src_tensors): + if dst_tensor is not None: + dst_tensor.copy_(src_tensor, *args[2:], **kwargs) + dst_tensor_obj.restore_from_saved(dst_tensors) + src_tensor_obj.restore_from_saved(src_tensors) + return None + if isinstance(dst, QuantizedTensor): dst.quantize_(src) else: @@ -419,6 +469,36 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == torch.ops.aten.view.default: raise NotImplementedError("{cls.__name__} class does not support tensor views") + # Empty like op + if func == torch.ops.aten.empty_like.default: + tensor = args[0] + device = kwargs.get("device", tensor.device) + requires_grad = kwargs.get("requires_grad", tensor.requires_grad) + pin_memory = kwargs.get("pin_memory", False) + usage = tensor.get_usages() + quantizer_usage = tensor._quantizer.get_usages() + tensor._quantizer.set_usage(**usage) + out = tensor._quantizer.make_empty( + shape=tensor.shape, + dtype=tensor.dtype, + device=device, + requires_grad=requires_grad, + pin_memory=pin_memory, + ) + tensor._quantizer.set_usage(**quantizer_usage) + return out + + if func == torch.ops.aten.numel.default: + tensor = args[0] + return math.prod(tensor.size()) + + if func == torch.ops.aten.is_pinned.default: + tensor = args[0] + for t in tensor.get_data_tensors(): + if t is not None: + return func(t) + return False # Or error out? + def maybe_unwrap(arg): if isinstance(arg, QuantizedTensor): return arg.dequantize(dtype=arg.dtype) @@ -463,6 +543,16 @@ def maybe_update_inplace(arg, new_arg, schema_arg): def __torch_function__(cls, func, types, args=(), kwargs=None): if kwargs is None: kwargs = {} + + def check_if_cpu(arg): + if isinstance(cls, QuantizedTensor) and arg.device.type == "cpu": + assert ( + func in _quantized_tensor_cpu_supported_ops + ), f"QuantizedTensor on CPU does not support this operation: {func}" + return arg + + args = tree_map(check_if_cpu, args) + # Do not force the QuantizedTensor type on the returned tensor return torch._C._disabled_torch_function_impl(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 8054374c81..8440c14b74 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -214,6 +214,7 @@ def make_empty( 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""" if device is None: @@ -229,12 +230,13 @@ def make_empty( data = None scale_inv = None if self.rowwise_usage: - data = torch.empty(shape, dtype=torch.uint8, device=device) + data = torch.empty(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.float32, device=device, + pin_memory=pin_memory, ) # Allocate FP8 data transpose if needed @@ -242,13 +244,17 @@ def make_empty( columnwise_scale_inv = None if self.columnwise_usage: columnwise_data = torch.empty( - self.get_columnwise_shape(shape), dtype=torch.uint8, device=device + self.get_columnwise_shape(shape), + 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.float32, device=device, + pin_memory=pin_memory, ) # Construct FP8 tensor diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index eb2ac9a581..7f7195a17f 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -101,6 +101,7 @@ def make_empty( dtype: torch.dtype = torch.float32, device: Optional[torch.device] = None, requires_grad: bool = False, + pin_memory: bool = False, ) -> Float8Tensor: # Canonicalize tensor attributes @@ -108,16 +109,19 @@ def make_empty( device = torch.device("cuda") # Allocate FP8 data - data = torch.empty(shape, dtype=torch.uint8, device=device) + 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 = [data.size(-1)] + list(data.shape[:-1]) + 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 @@ -125,7 +129,7 @@ def make_empty( shape=shape, dtype=dtype, data=data, - fp8_scale_inv=torch.empty(1, dtype=torch.float32, device=device), + 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, @@ -287,6 +291,7 @@ def make_empty( dtype: torch.dtype = torch.float32, device: Optional[torch.device] = None, requires_grad: bool = False, + pin_memory: bool = False, ) -> Float8Tensor: # Canonicalize tensor attributes @@ -294,23 +299,26 @@ def make_empty( device = torch.device("cuda") # Allocate FP8 data - data = torch.empty(shape, dtype=torch.uint8, device=device) + 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 = [data.size(-1)] + list(data.shape[:-1]) + 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), + 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, @@ -715,14 +723,22 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): return cls.detach(args[0]) if func == torch.ops.aten.clone.default: return cls.clone(args[0]) + if func == torch.ops.aten.copy_.default: dst, src = args[0], args[1] # Just copy FP8 attrs if copying between Float8Tensors if isinstance(src, Float8Tensor) and isinstance(dst, Float8Tensor): - dst._data.copy_(src._data.detach()) - dst._scale_inv.copy_(src._scale_inv.view(dst._scale_inv.size())) - if src._transpose is not None or dst._transpose is not None: - dst._create_transpose() + if dst._data is not None: + dst._data.copy_(src._data.detach(), *args[2:], **kwargs) + if dst._scale_inv is not None: + dst._scale_inv.copy_( + src._scale_inv.view(dst._scale_inv.size()), *args[2:], **kwargs + ) + if dst._transpose is not None and not dst._transpose_invalid: + if not src._transpose_invalid: + dst._transpose.copy_(src._transpose, *args[2:], **kwargs) + else: + dst._create_transpose() return dst elif func in _ops_to_preserve_subclass_in_fsdp2: # Ops in the _ops_to_preserve_subclass_in_fsdp2 are recommened to return the same class instance to work fine with the torch fsdp2 diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 15e0b86c90..7ca6e3b0dd 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -90,6 +90,7 @@ def make_empty( dtype: torch.dtype = torch.float32, device: Optional[torch.device] = None, requires_grad: bool = False, + pin_memory: bool = False, ) -> MXFP8Tensor: # Canonicalize tensor attributes @@ -105,24 +106,29 @@ def make_empty( ) # Allocate FP8 data - data = torch.empty(shape, dtype=torch.uint8, device=device) - 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, - ) + 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_like(data) + columnwise_data = torch.empty_like(data, 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 @@ -348,11 +354,17 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) if rowwise_matches and columnwise_matches: if dst._rowwise_data is not None: - dst._rowwise_data.copy_(src._rowwise_data.detach()) - dst._rowwise_scale_inv.copy_(src._rowwise_scale_inv.detach()) + dst._rowwise_data.copy_(src._rowwise_data.detach(), *args[2:], **kwargs) + dst._rowwise_scale_inv.copy_( + src._rowwise_scale_inv.detach(), *args[2:], **kwargs + ) if dst._columnwise_data is not None: - dst._columnwise_data.copy_(src._columnwise_data.detach()) - dst._columnwise_scale_inv.copy_(src._columnwise_scale_inv.detach()) + dst._columnwise_data.copy_( + src._columnwise_data.detach(), *args[2:], **kwargs + ) + dst._columnwise_scale_inv.copy_( + src._columnwise_scale_inv.detach(), *args[2:], **kwargs + ) return dst # FSDP2 related functions. diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 7a5f8858f2..31dbcf00a9 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -6,7 +6,7 @@ from __future__ import annotations from collections.abc import Iterable import math -from typing import Optional, Tuple, Union +from typing import Dict, Optional, Tuple, Union import functools import torch @@ -265,6 +265,7 @@ def make_empty( *, dtype: torch.dtype = torch.float32, device: Optional[torch.device] = None, + pin_memory: bool = False, requires_grad: bool = False, ) -> NVFP4Tensor: @@ -288,11 +289,18 @@ def make_empty( scale_inv = None amax_rowwise = None if self.rowwise_usage: - data = torch.empty(self.convert_shape_for_fp4(shape), dtype=torch.uint8, device=device) + 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) + 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) + amax_rowwise = torch.zeros(1, dtype=torch.float32, device=device, pin_memory=pin_memory) # Allocate FP8 data transpose if needed columnwise_data = None @@ -306,12 +314,15 @@ def make_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 + 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 ) - amax_columnwise = torch.zeros(1, dtype=torch.float32, device=device) # Construct FP8 tensor return NVFP4Tensor( @@ -498,6 +509,12 @@ def contiguous( return self raise ValueError("NVFP4Tensor does not support different memory formats!") + def get_usages(self) -> Dict[str, bool]: + return { + "rowwise": self._rowwise_data is not None, + "columnwise": self._columnwise_data is not None, + } + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): @@ -520,16 +537,20 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) if tensor._rowwise_data is not None: - rowwise_data = data_init_func(tensor._rowwise_data) - rowwise_scale_inv = scale_inv_init_func(tensor._rowwise_scale_inv) - amax_rowwise = torch.zeros_like(tensor._amax_rowwise) + rowwise_data = data_init_func(tensor._rowwise_data, *args[1:], **kwargs) + rowwise_scale_inv = scale_inv_init_func( + tensor._rowwise_scale_inv, *args[1:], **kwargs + ) + amax_rowwise = torch.zeros_like(tensor._amax_rowwise, *args[1:], **kwargs) else: rowwise_data, rowwise_scale_inv, amax_rowwise = None, None, None if tensor._columnwise_data is not None: - columnwise_data = data_init_func(tensor._columnwise_data) - columnwise_scale_inv = scale_inv_init_func(tensor._columnwise_scale_inv) - amax_columnwise = torch.zeros_like(tensor._amax_columnwise) + columnwise_data = data_init_func(tensor._columnwise_data, *args[1:], **kwargs) + columnwise_scale_inv = scale_inv_init_func( + tensor._columnwise_scale_inv, *args[1:], **kwargs + ) + amax_columnwise = torch.zeros_like(tensor._amax_columnwise, *args[1:], **kwargs) else: columnwise_data, columnwise_scale_inv, amax_columnwise = ( None, diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index c2d5e8b3fa..38d117b2a0 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -420,3 +420,10 @@ def update_usage( return return + + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the tensor""" + return { + "rowwise": self._rowwise_data is not None, + "columnwise": self._columnwise_data is not None, + } diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index a31f6a3799..8d12c30700 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -225,3 +225,12 @@ def update_usage( if not needs_data_transpose: self._transpose = None self._transpose_invalid = True + + def get_usages(self) -> Dict[str, bool]: + """Get the usage of the tensor""" + usages = {"rowwise": self._data is not None} + if is_non_tn_fp8_gemm_supported(): + usages["columnwise"] = self._data is not None + else: + usages["columnwise"] = self._transpose is not None and not self._transpose_invalid + return usages diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 2cca0829db..e7840d2c43 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -254,3 +254,10 @@ def update_usage( else: self._columnwise_data = None self._columnwise_scale_inv = None + + def get_usages(self) -> Tuple[bool, bool]: + """Get the usage of the tensor""" + return { + "rowwise": self._rowwise_data is not None, + "columnwise": self._columnwise_data is not None, + } From 389a6ba4f7be41e21b3db437d9ba23a01a44db1a Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 14 Nov 2025 12:57:07 -0800 Subject: [PATCH 068/521] [JAX] Use TE quant if TE fused act is disabled (#2374) * Use TE quant if TE fused act is disabled Signed-off-by: Jeremy Berchtold * Keep existing precision Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- .../jax/cpp_extensions/activation.py | 43 ++++++++++++++++--- 1 file changed, 36 insertions(+), 7 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index aa84fafd3e..e8249de170 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -27,7 +27,7 @@ should_apply_1x_fused_dbias_war_for_arch_l_100, NamedSharding, ) -from .quantization import _jax_dbias, _quantize_dbias_impl, AmaxScope +from .quantization import _jax_dbias, quantize, quantize_dbias, _quantize_dbias_impl, AmaxScope from ..sharding import all_reduce_max_along_all_axes_except_PP, all_reduce_sum_along_dp_fsdp from ..quantize import ScaledTensor, ScaledTensorFactory, NoScaleTensor from ..quantize import ( @@ -1268,7 +1268,19 @@ def act_lu( ) act_params = act_params if act_params is not None else ActivationParams() if not ActLuPrimitive.enabled(): - return _jax_act_lu(x, activation_type, quantizer, act_params) + act_out = _jax_act_lu(x, activation_type, act_params=act_params) + assert ( + act_out.data.dtype == x.dtype + ), f"JAX activation output dtype {act_out.data.dtype} must match input dtype {x.dtype}" + if quantizer is None: + return act_out + + return quantize( + act_out, + quantizer=quantizer, + amax_scope=amax_scope, + transpose_batch_sequence=transpose_batch_sequence, + ) # TE/common does not support colwise-only quantization yet if quantizer is not None and quantizer.q_layout.is_colwise_only: @@ -1330,11 +1342,12 @@ def act_lu( transpose_batch_sequence=transpose_batch_sequence, output_amax_when_no_scaling=True, ) - out, _ = _quantize_dbias_impl( + assert ( + out.data.dtype == x.dtype + ), f"Activation output dtype {out.data.dtype} must match input dtype {x.dtype}" + out = quantize( out, - is_dbias=False, quantizer=quantizer, - dq_dtype=x.dtype, amax_scope=amax_scope, transpose_batch_sequence=transpose_batch_sequence, ) @@ -1419,7 +1432,23 @@ def quantize_dact_dbias( if not PrimitiveClass.enabled() or ( quantizer is not None and quantizer.q_layout.is_colwise_only ): - return _jax_quantize_dact_dbias(dz, x, activation_type, is_dbias, quantizer, act_params) + if quantizer is None: + return _jax_quantize_dact_dbias(dz, x, activation_type, is_dbias, act_params=act_params) + dact_out, _ = _jax_quantize_dact_dbias( + dz, x, activation_type, is_dbias=False, act_params=act_params + ) + assert ( + dact_out.data.dtype == x.dtype + ), f"JAX dact output dtype {dact_out.data.dtype} must match input dtype {x.dtype}" + return quantize_dbias( + dact_out, + quantizer, + is_dbias=is_dbias, + flatten_axis=-2, + amax_scope=amax_scope, + transpose_batch_sequence=transpose_batch_sequence, + ) + if quantizer is None: output, _, _, _, updated_amax, _ = PrimitiveClass.outer_primitive.bind( dz, @@ -1465,7 +1494,7 @@ def quantize_dact_dbias( output_amax_when_no_scaling=output_amax_when_no_scaling, ) return _quantize_dbias_impl( - out.data, + out, quantizer, is_dbias=True, dq_dtype=x.dtype, From 66aed3ae5f6b900f4e85d2ed3eb1040c97891209 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Fri, 14 Nov 2025 17:05:02 -0800 Subject: [PATCH 069/521] Updated VERSION to 2.11.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 c7f2fd9b8e..5b70b33bd8 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.10.0.dev0 +2.11.0.dev0 From 42d227401e7eab1735a942f026bd2435167047e9 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 14 Nov 2025 17:09:51 -0800 Subject: [PATCH 070/521] [JAX] Quickstart documentation (#2310) * jax quickstart guide first commit Signed-off-by: tdophung * edit the syntax errors and remove unnecessary comments in utils. Add some footnotes in the quick start notebook Signed-off-by: tdophung * Fix greptiles comments on spelling, deepcopy, vjp function signature comaptibility with speedometer Signed-off-by: tdophung * Add Copyright to utils and fix some more greptiles complaints Signed-off-by: tdophung * Add comments to alternative of layers Signed-off-by: tdophung * Remove weight sharing between different iterations of the transformerLayer Signed-off-by: tdophung [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: tdophung * Add enum for attention implementations. Fix inconsistency between fuse and unfused TE impls to achieve same performance (removing extra dropout layer in fused layers. Also some minor wording changes Signed-off-by: tdophung [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix bug in TransformerLayer expected input shape being [sequence, batch, ...] instead of [batch, sequence,...] Signed-off-by: tdophung * Changing structure of notebook to bring fp8 ahead of fuse, to allow for fuse to take effect because quantization exist as suggested. Also make TransformerLayer perf get closer to Fused by setting hidden_dropout=0 Signed-off-by: tdophung * add option to choose between different attention implementation in call of BasicTETransformerLayer and demonstrated difference in runtime between using flax and using te's attetion implementation Signed-off-by: tdophung * Fix mistake in lacking attention_implementation in FuseTETransformerLayer Signed-off-by: tdophung * Removing AttentionWrapper and custom built DPA, using flax and TE's impl only, removing last mention of Pytorch Signed-off-by: tdophung * More changing to markdowns to remove pytorch Signed-off-by: tdophung * cosmetics fixes Signed-off-by: tdophung * changing names of all implementations Signed-off-by: tdophung * change fp8_autocast to autocast, make causal mask, and some wording changes Signed-off-by: tdophung --------- Signed-off-by: tdophung Co-authored-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- docs/examples/quickstart_jax.ipynb | 783 ++++++++++++++++++++++++++ docs/examples/quickstart_jax_utils.py | 86 +++ 2 files changed, 869 insertions(+) create mode 100644 docs/examples/quickstart_jax.ipynb create mode 100644 docs/examples/quickstart_jax_utils.py diff --git a/docs/examples/quickstart_jax.ipynb b/docs/examples/quickstart_jax.ipynb new file mode 100644 index 0000000000..0bf928d6ee --- /dev/null +++ b/docs/examples/quickstart_jax.ipynb @@ -0,0 +1,783 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "962d87bb", + "metadata": {}, + "source": [ + "\n", + "\n", + "# Getting Started\n", + "\n", + "## Overview\n", + "\n", + "Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, providing better performance with lower memory utilization in both training and inference. It provides support for 8-bit floating point (FP8) precision on Hopper, Ada, as well as 8-bit and 4-bit floating point (NVFP4) precision on Blackwell GPUs, implements a collection of highly optimized building blocks for popular Transformer architectures, and exposes an automatic-mixed-precision-like API that can be used seamlessly with your JAX code. It also includes a framework-agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers.\n", + "\n", + "This guide shows how to start using Transformer Engine with JAX. Similar tutorial for pyTorch is available [here](quickstart.ipynb).\n", + "We recommend you to try understanding the basics of JAX first, using these resources:\n", + "\n", + "- Thinking in JAX: https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html\n", + "- JAX 101: https://docs.jax.dev/en/latest/jax-101.html\n", + "- Key concepts in JAX: https://docs.jax.dev/en/latest/key-concepts.html#jax-arrays-jax-array\n", + "- Flax 101: https://flax-linen.readthedocs.io/en/latest/guides/flax_fundamentals/index.html\n", + "\n", + "## Let's build a Transformer decoder layer!\n", + "_This is based upon the GPT decoder layer with causal masking, which prevents each position from attending to future positions._\n", + "\n", + "
\n", + "\n", + "Summary\n", + " \n", + "We build a basic Transformer layer using regular Flax modules. This will be our baseline for later comparisons with Transformer Engine.\n", + "\n", + "
\n", + "\n", + "Let's start with creating the transformer layer using plain [FLAX Linen](https://flax.readthedocs.io/en/stable/) . Figure 1 shows the overall structure.\n", + "\n", + "
\n", + "\n", + "
Figure 1: Structure of a GPT decoder layer.
\n", + "
\n", + "\n", + "We construct the components as follows:\n", + "\n", + "- `LayerNorm`: `nn.LayerNorm` (Flax)\n", + "- `QKV Projection`: `nn.Dense` (conceptually there are three seperate `Dense` layers for Q, K, and V separately, but we fuse them together into a single `Dense` layer that is three times larger)\n", + "- `DotProductAttention`: `nn.MuliheadDotProductAttention` (Flax)\n", + "- `Projection`: `nn.Dense` (Flax)\n", + "- `Dropout`: `nn.Dropout` (Flax)\n", + "- `MLP`: `FlaxMLP` implemented using `nn.Dense` and `nn.gelu`\n", + "\n", + "Over the course of this tutorial we will use a few modules and helper functions defined in [quickstart_jax_utils.py](quickstart_jax_utils.py). Putting it all together: \n" + ] + }, + { + "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": 2, + "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", + "\n", + " @nn.compact\n", + " def __call__(self, x: jnp.ndarray) -> jnp.ndarray:\n", + " x = nn.Dense(features=self.ffn_hidden_size, use_bias=True)(x)\n", + " x = nn.gelu(x, approximate=True) # equivalent to tanh approximation\n", + " x = nn.Dense(features=self.hidden_size, use_bias=True)(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", + " \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)(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", + " # Reshape to [batch, seq_len, num_heads * head_dim] for Flax MultiHeadDotProductAttention\n", + " q_reshaped = q.reshape(q.shape[0], q.shape[1], self.hidden_size)\n", + " k_reshaped = k.reshape(k.shape[0], k.shape[1], self.hidden_size)\n", + " v_reshaped = v.reshape(v.shape[0], v.shape[1], self.hidden_size)\n", + " \n", + " # Attention using Flax's MultiHeadDotProductAttention\n", + " attention = nn.MultiHeadDotProductAttention(\n", + " num_heads=self.num_attention_heads,\n", + " qkv_features=self.kv_channels,\n", + " dropout_rate=self.attention_dropout,\n", + " )\n", + " x = attention(q_reshaped, k_reshaped, v_reshaped, mask=attention_mask, deterministic=deterministic)\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", + " )\n", + " x = mlp(x)\n", + " \n", + " return x + res\n" + ] + }, + { + "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)}, '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,)}, 'MultiHeadDotProductAttention_0': {'key': {'bias': (32, 4), 'kernel': (4096, 32, 4)}, 'out': {'bias': (4096,), 'kernel': (32, 4, 4096)}, 'query': {'bias': (32, 4), 'kernel': (4096, 32, 4)}, 'value': {'bias': (32, 4), 'kernel': (4096, 32, 4)}}}}\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: 17.708301544189453 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", + " dropout_key=dropout_key,\n", + " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "ccb16f31", + "metadata": {}, + "source": [ + "## Meet Transformer Engine\n", + "\n", + "
\n", + "\n", + "Summary\n", + " \n", + "Now that we have a basic Transformer layer in Flax, let's use Transformer Engine to speed up the training. The following examples show how to use TE modules.\n", + "\n", + "
\n", + "\n", + "As a reminder, the FlaxTransformerLayer above used:\n", + "\n", + "- `nn.LayerNorm`: Flax LayerNorm\n", + "- `nn.Dense`: Flax Dense layer for QKV projection \n", + "- `nn.MultiheadDotProductAttention`: Flax MultiheadDotProductAttention\n", + "- `nn.Dense`: Flax Dense layer for projection\n", + "- `nn.Dropout`: Flax Dropout\n", + "- `FlaxMLP`: Custom MLP implemented from `nn.Dense`\n", + "\n", + "Below we show how to use Transformer Engine Flax modules for better performance:\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "bed20d6b", + "metadata": {}, + "outputs": [], + "source": [ + "import transformer_engine.jax as te\n", + "import transformer_engine.jax.flax as te_flax" + ] + }, + { + "cell_type": "markdown", + "id": "f28cb444", + "metadata": {}, + "source": [ + "TE provides a set of Flax Linen modules that can be used to build Transformer layers. The simplest of the provided modules are the `DenseGeneral ` and `LayerNorm` layers, which we can use instead of `flax.linen.Dense` and ` flax.linen.LayerNorm`. Let's modify our `FlaxTransformerLayer`:" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "56105579", + "metadata": {}, + "outputs": [], + "source": [ + "from transformer_engine.jax.flax.transformer import DotProductAttention as TEDotProductAttention\n", + "\n", + "\n", + "class TEUnfusedMLP(nn.Module):\n", + " hidden_size : int\n", + " ffn_hidden_size: int\n", + "\n", + " @nn.compact\n", + " def __call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray:\n", + " x = te_flax.DenseGeneral(features=self.ffn_hidden_size, use_bias=True) (x)\n", + " x = x.reshape(*x.shape[:-1], 1, x.shape[-1])\n", + " x = te.activation.activation(x, activation_type=('gelu',))\n", + " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True) (x)\n", + " return x\n", + "\n", + "class TEUnfusedTransformerLayer(nn.Module):\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", + " use_te_attention: bool = True # True for TE attention, False for Flax attention\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 = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x)\n", + "\n", + " # Fused QKV projection\n", + " qkv = te_flax.DenseGeneral(features=3 * self.hidden_size, use_bias=True)(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", + " # Attention - either TE or Flax implementation\n", + " if self.use_te_attention:\n", + " # Use TE's DotProductAttention\n", + " attention = TEDotProductAttention(\n", + " head_dim=self.kv_channels,\n", + " num_attention_heads=self.num_attention_heads,\n", + " num_gqa_groups=self.num_attention_heads, # No GQA\n", + " attention_dropout=self.attention_dropout,\n", + " attn_mask_type='causal',\n", + " transpose_batch_sequence=False, # Input format is [batch, seq_len, ...]\n", + " )\n", + " x = attention(q, k, v, sequence_descriptor=None, deterministic=deterministic)\n", + " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", + " x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))\n", + " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x)\n", + " x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic)\n", + " else:\n", + " # Use Flax's MultiHeadDotProductAttention\n", + " q_reshaped = q.reshape(q.shape[0], q.shape[1], self.hidden_size)\n", + " k_reshaped = k.reshape(k.shape[0], k.shape[1], self.hidden_size)\n", + " v_reshaped = v.reshape(v.shape[0], v.shape[1], self.hidden_size)\n", + " \n", + " attention = nn.MultiHeadDotProductAttention(\n", + " num_heads=self.num_attention_heads,\n", + " qkv_features=self.kv_channels,\n", + " dropout_rate=self.attention_dropout,\n", + " )\n", + " x = attention(q_reshaped, k_reshaped, v_reshaped, mask=attention_mask, deterministic=deterministic)\n", + "\n", + " x = res + x\n", + "\n", + " # Second residual connection\n", + " res = x\n", + " x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x)\n", + "\n", + " # MLP\n", + " mlp = TEUnfusedMLP(\n", + " hidden_size=self.hidden_size,\n", + " ffn_hidden_size=self.ffn_hidden_size\n", + " )\n", + "\n", + " x = mlp(x, deterministic=deterministic)\n", + "\n", + " return x + res" + ] + }, + { + "cell_type": "markdown", + "id": "a76911ac", + "metadata": {}, + "source": [ + "Testing performance of the model, using `DenseGeneral`, `LayerNorm` and activation from TE, while keeping Flax's `MultiHeadDotProductAttention` the same as the first simple Transformer in JAX implementation. To read more about this implementation from Flax, you can refer to this documentation: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/attention.html" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "id": "4b67511f", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean time: 16.505107879638672 ms\n" + ] + } + ], + "source": [ + "te_unfused_transformer_with_flax_MHA = TEUnfusedTransformerLayer(\n", + " hidden_size, \n", + " ffn_hidden_size, \n", + " num_attention_heads,\n", + " use_te_attention=False\n", + ")\n", + "\n", + "te_params = te_unfused_transformer_with_flax_MHA.init(key, x, attention_mask=None, deterministic=False)\n", + "\n", + "utils.speedometer(\n", + " model_apply_fn=te_unfused_transformer_with_flax_MHA.apply,\n", + " variables=te_params, # Ensure the correct `params` is passed\n", + " input=x,\n", + " output_grad=dy,\n", + " dropout_key=dropout_key,\n", + " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + ")\n" + ] + }, + { + "cell_type": "markdown", + "id": "0b230058", + "metadata": {}, + "source": [ + "Now, we move on to also replace the attention sub-layer with TE's `DotProductAttention` implementation" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "5146cd99", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean time: 12.80329704284668 ms\n" + ] + } + ], + "source": [ + "te_unfused_transformer = TEUnfusedTransformerLayer(\n", + " hidden_size, \n", + " ffn_hidden_size, \n", + " num_attention_heads,\n", + ")\n", + "\n", + "te_params = te_unfused_transformer.init(key, x, attention_mask=None, deterministic=False)\n", + "\n", + "utils.speedometer(\n", + " model_apply_fn=te_unfused_transformer.apply,\n", + " variables=te_params, # Ensure the correct `params` is passed\n", + " input=x,\n", + " output_grad=dy,\n", + " dropout_key=dropout_key,\n", + " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "c9a101d3", + "metadata": {}, + "source": [ + "## Enabling Quantization (FP8 or FP4)\n", + "\n", + "
\n", + "\n", + "Summary\n", + " \n", + "We configure a TE module to perform compute in FP8.\n", + "\n", + "
\n", + "\n", + "Enabling FP8 support is very simple in Transformer Engine. We just need to wrap the modules within an [autocast](.../api/jax.rst#transformer_engine.jax.fp8_autocast) context manager. See the [FP8 tutorial](fp8_primer.ipynb) for a detailed explanation of FP8 recipes and the supported options.\n", + "\n", + "
\n", + "\n", + "Important: FP8 Metadata Initialization\n", + "\n", + "When using FP8, the model **must be initialized within the `autocast` context**. This creates a special collection called `fp8_metas` that contains scaling factors and other metadata required for FP8 computation. If you initialize a model outside of `autocast` and then try to use it with FP8, you will get a `ScopeCollectionNotFound` error because the `fp8_metas` collection was never created.\n", + "\n", + "
" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "c2eee376", + "metadata": {}, + "outputs": [], + "source": [ + "from transformer_engine.common.recipe import Format, DelayedScaling\n", + "fp8_format = Format.HYBRID\n", + "fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo=\"max\")" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "de96827c", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean time: 9.615030288696289 ms\n" + ] + } + ], + "source": [ + "with te.autocast(enabled=True, recipe=fp8_recipe):\n", + " te_unfused_params = te_unfused_transformer.init(key, x, attention_mask=None, deterministic=False)\n", + "\n", + " # Example usage of forward \n", + " y = te_unfused_transformer.apply(te_unfused_params, x, attention_mask=None, deterministic=True)\n", + "\n", + "utils.speedometer(\n", + " model_apply_fn=te_unfused_transformer.apply,\n", + " variables=te_unfused_params, # Ensure the correct `params` is passed\n", + " input=x,\n", + " output_grad=dy,\n", + " dropout_key=dropout_key,\n", + " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe}\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "3801b201", + "metadata": {}, + "source": [ + "\n", + "## Fused TE Modules\n", + "\n", + "
\n", + "\n", + "Summary\n", + " \n", + "We optimize the example Transformer layer with TE modules for fused operations.\n", + "\n", + "
\n", + "\n", + "The `DenseGeneral` layer is enough to build any Transformer model and it enables usage of the Transformer Engine even for very custom Transformers. However, having more knowledge about the model allows for additional optimizations such as kernel fusions in mixed-precision recipes, increasing the achievable speedup.\n", + "\n", + "Transformer Engine therefore provides coarser modules that span multiple layers:\n", + "\n", + "* `LayerNormDenseGeneral`\n", + "* `LayerNormMLP`\n", + "* `TransformerLayer`\n", + "\n", + "To see a complete list of all the functions TE Flax support, you can view it here: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/jax.html#modules\n", + "\n", + "Building a third iteration of our Transformer layer with `LayerNormDenseGeneral` and `LayerNormMLP`:" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "11203785", + "metadata": {}, + "outputs": [], + "source": [ + "class TEFusedTransformerLayer(nn.Module):\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", + "\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", + " res = x\n", + "\n", + " # Fused QKV projection\n", + " qkv,_ = te_flax.LayerNormDenseGeneral(features=3 * self.hidden_size, \n", + " epsilon=self.layernorm_eps, \n", + " use_bias=True, \n", + " return_layernorm_output=False)(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", + " # Attention using TE's DotProductAttention\n", + " attention = TEDotProductAttention(\n", + " head_dim=self.kv_channels,\n", + " num_attention_heads=self.num_attention_heads,\n", + " num_gqa_groups=self.num_attention_heads, \n", + " attention_dropout=self.attention_dropout,\n", + " attn_mask_type='causal',\n", + " transpose_batch_sequence=False, # Input format is [batch, seq_len, ...]\n", + " )\n", + " x = attention(q, k, v, sequence_descriptor=None, deterministic=deterministic)\n", + " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", + " x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))\n", + " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x)\n", + " x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic)\n", + "\n", + " x = res + x\n", + "\n", + " # Second residual connection\n", + " res = x\n", + " x,_ = te_flax.LayerNormMLP(intermediate_dim=self.ffn_hidden_size, \n", + " epsilon=self.layernorm_eps,\n", + " use_bias=True,\n", + " activations=('gelu',),\n", + " intermediate_dropout_rate=0.0,\n", + " return_layernorm_output=False\n", + " )(x, deterministic=deterministic)\n", + "\n", + " return x + res" + ] + }, + { + "cell_type": "markdown", + "id": "334cff59", + "metadata": {}, + "source": [ + "Similar to the unnfused model, we also compare the performance of fused model when using Flax's MultiheadDotProductAttention implementation and TE's." + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "6b0c705e", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean time: 9.331779479980469 ms\n" + ] + } + ], + "source": [ + "te_fused_transformer = TEFusedTransformerLayer(\n", + " hidden_size, \n", + " ffn_hidden_size, \n", + " num_attention_heads\n", + ")\n", + "\n", + "with te.autocast(enabled=True, recipe=fp8_recipe):\n", + " te_fused_params = te_fused_transformer.init(key, x, attention_mask=None, deterministic=False)\n", + " # Example usage of forward \n", + " y = te_fused_transformer.apply(te_fused_params, x, attention_mask=None, deterministic=True)\n", + "\n", + "utils.speedometer(\n", + " model_apply_fn=te_fused_transformer.apply,\n", + " variables=te_fused_params,\n", + " input=x,\n", + " output_grad=dy,\n", + " dropout_key=dropout_key,\n", + " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe}\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "a45c12c8", + "metadata": {}, + "source": [ + "Finally, the `TransformerLayer` module is convenient for creating standard Transformer architectures." + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "b2aaa8ef", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "te_transformer = te_flax.TransformerLayer(\n", + " hidden_size=hidden_size,\n", + " mlp_hidden_size=ffn_hidden_size, \n", + " num_attention_heads=num_attention_heads,\n", + " mlp_activations=(\"gelu\",),\n", + " self_attn_mask_type='causal',\n", + " layernorm_epsilon=1e-5,\n", + " use_bias=True,\n", + " intermediate_dropout=0.0,\n", + " enable_relative_embedding=False,\n", + " self_attn_bias_type='no_bias',\n", + " hidden_dropout=0.0\n", + ")\n", + "\n", + "with te.autocast(enabled=True, recipe=fp8_recipe):\n", + " te_transformer_params = te_transformer.init(key, x, deterministic=False)\n", + " y = te_transformer.apply(te_transformer_params, x, attention_mask=None, deterministic=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 17, + "id": "b9cdbf22", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Mean time: 9.23741340637207 ms\n" + ] + } + ], + "source": [ + "utils.speedometer(\n", + " model_apply_fn=te_transformer.apply,\n", + " model_init_fn=te_transformer.init,\n", + " variables=te_transformer_params,\n", + " input=x,\n", + " output_grad=dy,\n", + " dropout_key=dropout_key,\n", + " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe }\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/examples/quickstart_jax_utils.py b/docs/examples/quickstart_jax_utils.py new file mode 100644 index 0000000000..138427338d --- /dev/null +++ b/docs/examples/quickstart_jax_utils.py @@ -0,0 +1,86 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import jax +import jax.numpy as jnp +import time +import math + +from typing import Callable, Any, Dict, Optional, Tuple +from flax import linen as nn +import transformer_engine.jax as te +import transformer_engine.jax.flax as te_flax +from transformer_engine.jax.flax.transformer import DotProductAttention as TEDotProductAttention + + +def speedometer( + model_apply_fn: Callable, + variables: Any, + input: jnp.ndarray, + output_grad: jnp.ndarray, + dropout_key: jax.random.PRNGKey, + model_init_fn: Callable = None, + forward_kwargs: dict = {}, + autocast_kwargs: Optional[dict] = None, + timing_iters: int = 50, + warmup_iters: int = 50, +) -> None: + """Measure average runtime for a JAX module + Perform forward and backward passes . + """ + if autocast_kwargs is None: + autocast_kwargs = {"enabled": False} + model_init_fn = None + + train_step_fn = create_train_step_fn(model_apply_fn, autocast_kwargs, forward_kwargs) + + # Warm up runs + key = dropout_key + for _ in range(warmup_iters): + key, step_key = jax.random.split(key) + loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_key) + + # Timing runs + start = time.time() + for _ in range(timing_iters): + key, step_key = jax.random.split(key) + loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_key) + end = time.time() + + print(f"Mean time: {(end - start) * 1000 / timing_iters} ms") + + +def create_train_step_fn( + model_apply_fn: Callable, + autocast_kwargs: Dict[str, Any], + forward_kwargs: Dict[str, Any] = None, +) -> Callable: + """ + Creates a JIT-compiled function that performs one forward/backward pass. + """ + + if forward_kwargs is None: + forward_kwargs = {} + + def loss_fn(variables: Any, inp: jnp.ndarray, grad_target: jnp.ndarray, dropout_key): + rngs = {"dropout": dropout_key} + with te.autocast(**autocast_kwargs): + # Forward Pass: Apply the model using current parameters and variables + call_kwargs = {**forward_kwargs, "rngs": rngs} + out = model_apply_fn(variables, inp, **call_kwargs) + + # grad_target = derivative of L (loss fn) over y (output) = signma(L)/sigma(y) + # where grad_w(L) = gradient of loss over params = sigma(L)/sigma(y) * sigma(y)/sigma(w) --> chain rule + # sigma(y)/sigma(w) = J_model(w) + return jnp.vdot(out, grad_target) + + def fwd_bwd_fn(*args, **kwargs): + return jax.value_and_grad(loss_fn, argnums=(0, 1))(*args, **kwargs) + + # Use jax.value_and_grad to get the loss value and gradients simultaneously. (forward + backward pass) + # ∇_params[output^T · grad_target] = grad_target^T · J_output(params) = VJP + # fwd_bwd_fn = jax.value_and_grad(loss_fn, argnums=(0, 1)) + + # JIT-compile the fwd_bwd_fn + return jax.jit(fwd_bwd_fn) From e1edaaec2bb1e6542e0e2dff81d5217ff5e1eb89 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 17 Nov 2025 08:55:29 -0500 Subject: [PATCH 071/521] [PyTorch] Reduce CPU overheads (#2377) Initial changes to remove pytorch overheads Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/attention/test_attention.py | 14 +- tests/pytorch/debug/test_numerics.py | 2 - tests/pytorch/distributed/test_sanity.py | 26 ++- tests/pytorch/test_numerics.py | 7 - tests/pytorch/test_sanity.py | 4 +- .../dot_product_attention/backends.py | 9 +- .../dot_product_attention.py | 2 +- .../pytorch/cpp_extensions/gemm.py | 69 +++++- .../pytorch/csrc/extensions/attention.cpp | 5 + .../pytorch/csrc/extensions/gemm.cpp | 15 ++ .../pytorch/csrc/extensions/normalization.cpp | 10 + transformer_engine/pytorch/module/base.py | 81 ++----- .../pytorch/module/fp8_padding.py | 26 ++- .../pytorch/module/fp8_unpadding.py | 28 ++- .../pytorch/module/grouped_linear.py | 102 ++++---- .../pytorch/module/layernorm_linear.py | 180 ++++++--------- .../pytorch/module/layernorm_mlp.py | 217 +++++++----------- transformer_engine/pytorch/module/linear.py | 158 ++++++------- .../pytorch/ops/basic/basic_linear.py | 4 - .../ops/fused/userbuffers_backward_linear.py | 3 - .../ops/fused/userbuffers_forward_linear.py | 2 - .../pytorch/quantized_tensor.py | 5 - .../pytorch/tensor/float8_blockwise_tensor.py | 16 ++ .../pytorch/tensor/float8_tensor.py | 44 +++- .../pytorch/tensor/mxfp8_tensor.py | 12 + .../pytorch/tensor/nvfp4_tensor.py | 20 ++ transformer_engine/pytorch/utils.py | 19 ++ 27 files changed, 564 insertions(+), 516 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a671f1eec2..4f4cad97db 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2489,7 +2489,6 @@ def forward( max_s: int, fast_zero_fill: bool, fp8_meta: Dict[str, Any], - workspace: torch.Tensor, is_training: bool, mask_type: str, quantizers: list[Quantizer], @@ -2518,7 +2517,6 @@ def forward( qkv, *_ = ext.general_gemm( qkv_weight_fp8, inp_fp8, - workspace, bias=qkv_bias, out_dtype=qkv_weight_fp8.dtype, quantization_params=qkv_quantizer, @@ -2560,9 +2558,7 @@ def forward( s_quantizer=s_quantizer, ) - tensors_to_save, tensor_objects = prepare_for_saving( - q, k, v, inp_fp8, qkv_weight_fp8, workspace, out - ) + tensors_to_save, tensor_objects = prepare_for_saving(q, k, v, inp_fp8, qkv_weight_fp8, out) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects @@ -2592,7 +2588,7 @@ def forward( def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: with torch.cuda.nvtx.range("_DPA"): saved_tensors = ctx.saved_tensors - (q, k, v, inp_fp8, qkv_weight_fp8, workspace, out) = restore_from_saved( + (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_saved( ctx.tensor_objects, saved_tensors ) @@ -2648,7 +2644,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], qkv_dgrad, *_ = ext.general_gemm( qkv_weight_fp8, dqkv_c, - workspace, ctx.dtype, use_split_accumulator=_2X_ACC_DGRAD, layout="NN", @@ -2658,7 +2653,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], qkv_wgrad, *_ = ext.general_gemm( inp_fp8, dqkv, - workspace, ctx.dtype, use_split_accumulator=_2X_ACC_WGRAD, layout="NT", @@ -2709,9 +2703,6 @@ def __init__(self, config, params_dtype: torch.dtype = torch.float32): with torch.no_grad(): self.qkv_bias.zero_() self.qkv_weight.fill_(1.0) - self.workspace = torch.empty( - _CUBLASLT_WORKSPACE_SIZE_BYTES, dtype=torch.int8, device="cuda" - ) def forward( self, @@ -2730,7 +2721,6 @@ def forward( max_s, self.fast_zero_fill, self.fp8_meta, - self.workspace, self.training, self.mask_type, self.quantizers, diff --git a/tests/pytorch/debug/test_numerics.py b/tests/pytorch/debug/test_numerics.py index 2ad2c8fb8f..ed8cdc1773 100644 --- a/tests/pytorch/debug/test_numerics.py +++ b/tests/pytorch/debug/test_numerics.py @@ -82,7 +82,6 @@ def _fp8_gemm_kernel(tensor1, scale1, dtype1, tensor2, scale2, dtype2, use_split out, *_ = tepytorch.cpp_extensions.general_gemm( fp8_tensor1, fp8_tensor2, - tepytorch.module.base.get_workspace(), torch.float32, use_split_accumulator=use_split_accumulator, ) @@ -199,7 +198,6 @@ def _emulate_linear( wgrad, *_ = tepytorch.cpp_extensions.general_gemm( wgrad_input, wgrad_gradient, - tepytorch.module.base.get_workspace(), torch.float32, layout="NT", grad=True, diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index fbbbe29972..81e49f3c54 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -7,7 +7,7 @@ import pytest import torch import transformer_engine -from transformer_engine.pytorch import DotProductAttention, TransformerLayer, Linear +from transformer_engine.pytorch import DotProductAttention, TransformerLayer, Linear, GroupedLinear _current_file = pathlib.Path(__file__).resolve() sys.path.append(str(_current_file.parent.parent)) @@ -19,7 +19,9 @@ @pytest.mark.parametrize("model", ["small"]) -@pytest.mark.parametrize("module", ["TransformerLayer", "DotProductAttention", "Linear"]) +@pytest.mark.parametrize( + "module", ["TransformerLayer", "DotProductAttention", "Linear", "GroupedLinear"] +) def test_current_device(model, module): """Test cases where current device is different from tensor device""" @@ -58,7 +60,7 @@ def test_current_device(model, module): kwargs["cu_seqlens_kv"] = cu_seqlens_kv kwargs["max_seqlen_q"] = config.max_seqlen_q kwargs["max_seqlen_kv"] = config.max_seqlen_kv - if module == "DotProductAttention": + elif module == "DotProductAttention": model = DotProductAttention( config.num_heads, config.head_dim_qk, qkv_format="thd", attn_mask_type="padding" ) @@ -97,6 +99,24 @@ def test_current_device(model, module): requires_grad=True, ) ] + elif module == "GroupedLinear": + num_gemms = 4 + model = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + params_dtype=dtype, + device=f"cuda:{tensor_device}", + ) + args = [ + torch.randn( + (config.max_seqlen_q * config.batch_size * (num_gemms - 1), config.hidden_size), + dtype=dtype, + device=f"cuda:{tensor_device}", + requires_grad=True, + ), + [0] + [config.max_seqlen_q * config.batch_size] * (num_gemms - 1), # Empty first split. + ] current_device_before = torch.cuda.current_device() out = model(*args, **kwargs) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 01f1deb983..1925f2e2e9 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -44,7 +44,6 @@ ) from transformer_engine.pytorch import checkpoint as te_checkpoint from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm -from transformer_engine.pytorch.module.base import get_multi_stream_cublas_workspace, get_workspace from transformer_engine.common import recipe import transformer_engine_torch as tex from utils import ModelConfig, reset_rng_states @@ -2690,7 +2689,6 @@ def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass): general_gemm( A[i], B[i], - get_workspace(), dtype, grad=grad, accumulate=accumulate, @@ -2705,7 +2703,6 @@ def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass): B, out, dtype, - get_multi_stream_cublas_workspace(), m_splits=m_splits, grad=grad, accumulate=accumulate, @@ -2760,7 +2757,6 @@ def test_fp8gemm_with_unfused_quantization(N, datatype, input_quantizer, out_qua quantized_out, *_ = general_gemm( weight_fp8, inp_fp8, - get_workspace(), outp_type, quantization_params=out_quantizer, bias=None, @@ -2770,7 +2766,6 @@ def test_fp8gemm_with_unfused_quantization(N, datatype, input_quantizer, out_qua out, *_ = general_gemm( weight_fp8, inp_fp8, - get_workspace(), outp_type, quantization_params=None, bias=None, @@ -2846,7 +2841,6 @@ def test_fp8_grouped_gemm(shape, accumulate): general_gemm( A_fp8[i], B_fp8[i], - get_workspace(), dtype, out=out_ref[i], accumulate=accumulate, @@ -2856,7 +2850,6 @@ def test_fp8_grouped_gemm(shape, accumulate): B_fp8, out, dtype, - get_multi_stream_cublas_workspace(), m_splits=m_splits, accumulate=accumulate, ) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index f12e80d4c3..5c116496ef 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -36,7 +36,6 @@ from transformer_engine.common import recipe import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions import general_gemm -from transformer_engine.pytorch.module.base import get_workspace from transformer_engine.pytorch.tensor.utils import replace_raw_data from utils import ModelConfig @@ -912,7 +911,7 @@ def test_sanity_gemm_with_unalignment(N, offset, datatype): inp = torch.reshape(scratchpad[offset:-offset], (N, N)) weight = torch.reshape(scratchpad[offset * 2 :], (N, N)) - _ = general_gemm(A=weight, B=inp, workspace=get_workspace()) + _ = general_gemm(A=weight, B=inp) torch.cuda.synchronize() @@ -936,7 +935,6 @@ def test_sanity_fp8_gemm_with_unalignment(N, datatype): general_gemm( weight_fp8, inp_fp8, - get_workspace(), outp_type, bias=None, use_split_accumulator=False, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 543055061b..1480f900fd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -19,7 +19,12 @@ get_device_compute_capability, split_tensor_along_dim, ) -from transformer_engine.pytorch.utils import attention_mask_func, nvtx_range_push, nvtx_range_pop +from transformer_engine.pytorch.utils import ( + attention_mask_func, + nvtx_range_push, + nvtx_range_pop, + get_nvtx_range_context, +) from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Quantizer, Float8CurrentScalingQuantizer, @@ -1445,7 +1450,7 @@ def backward(ctx, d_out, *_args): dk = dk[..., : d_out.shape[-1]] dv = dv[..., : d_out.shape[-1]] else: - with torch.cuda.nvtx.range("FusedAttnFunc.backward"): + with get_nvtx_range_context("FusedAttnFunc.backward"): # get nominal data type of dq, dk, dv # FP16/BF16 attention: torch.float16 or torch.bfloat16 # FP8 attention: torch.float16 or torch.bfloat16 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 4157e8d3a4..83528330eb 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 @@ -975,7 +975,7 @@ def forward( Whether to enforce output to be in FP8 or not. """ - with torch.cuda.device(query_layer.device), self.prepare_forward( + with self.prepare_forward( query_layer, num_gemms=3, allow_non_contiguous=True, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index dd04112982..76a0e449c0 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -6,23 +6,59 @@ from typing import Iterable, Optional, Tuple, Union, List import os +import functools import torch import transformer_engine_torch as tex from ..constants import TE_DType from ..utils import get_sm_count, _empty_tensor -from ..quantized_tensor import Quantizer +from ..quantized_tensor import Quantizer, QuantizedTensor, QuantizedTensorStorage +from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage +from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer + __all__ = [ "general_gemm", "general_grouped_gemm", ] +_NUM_MAX_UB_STREAMS = 3 + + +def get_cublas_workspace_size_bytes() -> None: + """Return 32 MiB if using hopper, 4 MiB for all other architectures.""" + if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 9: + # 32 MiB for NVFP4 GEMM, plus additional 1024 B for alignment and misc scales + return 32 * 1024 * 1024 + 1024 + return 4_194_304 + + +@functools.lru_cache(maxsize=None) +def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Tensor: + """Returns workspace for cublas GEMM.""" + assert not (ub and grouped_gemm), "UB is unsupported for grouped GEMM." + + if ub: + return torch.empty( + get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device + ).repeat(_NUM_MAX_UB_STREAMS) + if grouped_gemm: + _multi_stream_cublas_workspace = [] + for _ in range(tex.get_num_cublas_streams()): + _multi_stream_cublas_workspace.append( + torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + ) + return _multi_stream_cublas_workspace + + return torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device) + + def validate_gemm_scale(scale: Optional[float], required: bool) -> float: """Validate whether a GEMM scaling factor is consistent with its usage""" if required: @@ -32,10 +68,35 @@ def validate_gemm_scale(scale: Optional[float], required: bool) -> float: return 0.0 +def get_tensor_device(tensor: torch.Tensor) -> int: + """Returns tensor device as an integer""" + if not isinstance(tensor, QuantizedTensorStorage): + return tensor.device.index + if isinstance(tensor, QuantizedTensor): + return tensor.device.index + if isinstance(tensor, (Float8BlockwiseQTensorStorage, MXFP8TensorStorage, NVFP4TensorStorage)): + return ( + tensor._rowwise_data.device.index + if tensor._rowwise_data is not None + else tensor._columnwise_data.device.index + ) + if isinstance(tensor, Float8TensorStorage): + return ( + tensor._data.device.index + if tensor._data is not None + else tensor._transpose.device.index + ) + try: + return ( + tensor._data.device.index if tensor._data is not None else tensor._data_t.device.index + ) + except AttributeError: + return torch.cuda.current_device() + + def general_gemm( A: torch.Tensor, B: torch.Tensor, - workspace: torch.Tensor, out_dtype: Optional[torch.dtype] = None, quantization_params: Optional[Quantizer] = None, gelu: bool = False, @@ -62,6 +123,7 @@ def general_gemm( alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) + workspace = get_cublas_workspace(get_tensor_device(A), ub is not None, False) if ub_type is not None: assert ub is not None, ( @@ -159,7 +221,6 @@ def general_grouped_gemm( B: List[torch.Tensor], out: List[torch.Tensor], out_dtype: torch.dtype, - workspaces: List[torch.Tensor], layout: str = "TN", m_splits: Optional[List[int]] = None, gelu: bool = False, @@ -187,6 +248,8 @@ def general_grouped_gemm( out_dtype = TE_DType[out[0].dtype] if D_dtype is None else D_dtype sm_count = get_sm_count() + workspaces = get_cublas_workspace(get_tensor_device(A[0]), False, True) + if grad and use_bias: grad_bias = [ torch.empty(B[i].shape[1], dtype=out[0].dtype, device="cuda") for i in range(num_gemms) diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index d51aef4065..2480d9aba9 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -108,6 +108,11 @@ std::vector fused_attn_fwd( py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, const std::optional SoftmaxOffset, const std::optional rng_gen, size_t rng_elts_per_thread, bool return_max_logit, bool cuda_graph) { + // Ensure that cuDNN handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(cu_seqlens_q.device()); + auto none = py::none(); // create QKV tensor wrappers diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 15404ad9a6..13e8bfb6e5 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -95,6 +95,11 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans bool bulk_overlap, float alpha, std::optional beta) { using namespace transformer_engine::pytorch::detail; + // Ensure that cublasLt handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(workspace.device()); + // Input tensors NVTE_CHECK(!A.is_none(), "Tensor A has not been provided"); NVTE_CHECK(!B.is_none(), "Tensor B has not been provided"); @@ -351,6 +356,11 @@ void te_atomic_gemm(at::Tensor A, at::Tensor A_scale_inverse, DType A_type, at::Tensor workspace, size_t workspaceSize, bool accumulate, bool use_split_accumulator, int math_sm_count, int m_split, int n_split, bool gemm_producer, at::Tensor counter) { + // Ensure that cublasLt handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(workspace.device()); + // TODO: Handle scaling modes NVTEScalingMode nvte_scaling_modeA = NVTE_DELAYED_TENSOR_SCALING; NVTEScalingMode nvte_scaling_modeB = NVTE_DELAYED_TENSOR_SCALING; @@ -400,6 +410,11 @@ std::optional> te_general_grouped_gemm( NVTE_ERROR("not implemented, D should be allocated for single output case."); } + // Ensure that cublasLt handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(workspace[0].device()); + void* output_data_ptr = nullptr; if (single_output) { output_data_ptr = (*D)[0].data_ptr(); diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index 3fa0fb0aa3..3c5c17fc6f 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -64,6 +64,11 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe const bool zero_centered_gamma) { using namespace transformer_engine::pytorch::detail; + // Ensure that cuDNN handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(input.cast().device()); + // Input and param tensors auto none = py::none(); const TensorWrapper &input_nvte = makeTransformerEngineTensor(input, none); @@ -294,6 +299,11 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w const int sm_margin, const bool zero_centered_gamma) { using namespace transformer_engine::pytorch::detail; + // Ensure that cuDNN handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(input.cast().device()); + // Input and param tensors auto none = py::none(); const TensorWrapper &input_nvte = makeTransformerEngineTensor(input, none); diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index d2abe3a2de..6d1d8c3540 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -39,13 +39,18 @@ _fsdp_gather_tensors, ) from ..constants import dist_group_type +from ..cpp_extensions.gemm import _NUM_MAX_UB_STREAMS from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage -from ..utils import is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype +from ..utils import ( + is_non_tn_fp8_gemm_supported, + torch_get_autocast_gpu_dtype, + get_nvtx_range_context, +) from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ...common.recipe import DelayedScaling, Recipe from ...debug.pytorch.debug_state import TEDebugState @@ -57,11 +62,8 @@ _2X_ACC_FPROP = False _2X_ACC_DGRAD = True _2X_ACC_WGRAD = True -_multi_stream_cublas_workspace = [] _dummy_wgrads = {} -_cublas_workspace = None _ub_communicators = None -_NUM_MAX_UB_STREAMS = 3 _MIN_STREAM_PRIORITY, _MAX_STREAM_PRIORITY = None, None layers_atomic_ring_exchange = [] @@ -75,35 +77,6 @@ class UserBufferQuantizationMode(Enum): FP8 = "fp8" -def get_cublas_workspace_size_bytes() -> None: - """Return 32 MiB if using hopper, 4 MiB for all other architectures.""" - if torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 9: - # 32 MiB for NVFP4 GEMM, plus additional 1024 B for alignment and misc scales - return 32 * 1024 * 1024 + 1024 - return 4_194_304 - - -def get_workspace() -> torch.Tensor: - """Returns workspace for cublas.""" - global _cublas_workspace - if _cublas_workspace is None: - _cublas_workspace = torch.empty( - get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda" - ) - return _cublas_workspace - - -def get_multi_stream_cublas_workspace() -> List[torch.Tensor]: - """Returns workspace for multi-stream cublas.""" - global _multi_stream_cublas_workspace - if not _multi_stream_cublas_workspace: - for _ in range(tex.get_num_cublas_streams()): - _multi_stream_cublas_workspace.append( - torch.empty(get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda") - ) - return _multi_stream_cublas_workspace - - def get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor: """Returns a dummy tensor of given shape.""" assert len(shape) == 2 @@ -276,16 +249,6 @@ def initialize_ub( flush=True, ) - # Allocate cuBLAS workspace with expanded size for chunking in overlapping GEMM calls - global _cublas_workspace - if _cublas_workspace is None: - _cublas_workspace = get_workspace().repeat(_NUM_MAX_UB_STREAMS) - elif _cublas_workspace.numel() != get_cublas_workspace_size_bytes() * _NUM_MAX_UB_STREAMS: - # This ensures we don't do `.repeat()` on an already expanded workspace - _cublas_workspace = torch.empty( - get_cublas_workspace_size_bytes(), dtype=torch.uint8, device="cuda" - ).repeat(_NUM_MAX_UB_STREAMS) - # Default buffer precision: AllGather buffers use fp8 when using fp8 recipe layers_all_gather_overlap = [ "qkv_fprop", @@ -1078,8 +1041,10 @@ def prepare_forward( """ self.allow_different_data_and_param_types = allow_different_data_and_param_types self.forwarded_at_least_once = True + # 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() FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) else: assert inp.is_cuda, "TransformerEngine needs CUDA." @@ -1091,25 +1056,27 @@ def prepare_forward( self.init_fp8_metadata(num_gemms=num_gemms) self._check_weight_tensor_recipe_correspondence() - if self.fp8 and self.sequence_parallel and self.fp8_meta["recipe"].delayed(): - assert self.fp8_meta["recipe"].reduce_amax, ( - "Amax reduction across tensor parallel group is " - "necessary when using sequence parallelism with FP8." - ) + delayed_scaling_recipe = self.fp8 and self.fp8_meta["recipe"].delayed() + if delayed_scaling_recipe: + if self.sequence_parallel: + assert self.fp8_meta["recipe"].reduce_amax, ( + "Amax reduction across tensor parallel group is " + "necessary when using sequence parallelism with FP8." + ) - if self.fp8 and not FP8GlobalStateManager.fp8_graph_capturing(): - FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta) + if not FP8GlobalStateManager.fp8_graph_capturing(): + FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta) - # Activation recomputation is used and this is the first forward phase. - if self.fp8 and self.training and is_fp8_activation_recompute_enabled(): - FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) + # Activation recomputation is used and this is the first forward phase. + if self.training and is_fp8_activation_recompute_enabled(): + FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) - with torch.cuda.nvtx.range(self.__class__.__name__ + " forward"): + with get_nvtx_range_context(self.__class__.__name__ + " forward"): if not allow_non_contiguous and not inp.is_contiguous(): inp = inp.contiguous() yield inp - if self.fp8 and in_fp8_activation_recompute_phase(): + if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase(): FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta) def set_nccl_overlap_warning_if_tp(self) -> None: @@ -1531,7 +1498,7 @@ def backward_dw(self): """ if not self.need_backward_dw(): return - with torch.cuda.nvtx.range(f"_{self.__class__.__name__}_wgrad"): + with get_nvtx_range_context(f"_{self.__class__.__name__}_wgrad"): (wgrad, bgrad), _ = self.wgrad_store.pop() if not self.fuse_wgrad_accumulation: weight_tensor = noop_cat(self._get_weight_tensors()) @@ -1628,6 +1595,8 @@ def _check_weight_tensor_recipe_correspondence(self) -> None: """ if not self.fp8 and not self.fp8_calibration: return + if not self.primary_weights_in_fp8: + return if not hasattr(self, "weight_names") or not self.weight_names: return diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index fca89fbaa9..fd9b9b4377 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -24,11 +24,14 @@ class _Fp8Padding(torch.autograd.Function): def forward( ctx, inp: torch.Tensor, - m_splits: List[int], - padded_m_splits: List[int], - is_grad_enabled: bool, + non_tensor_args: Tuple, ) -> torch.Tensor: # pylint: disable=missing-function-docstring + + # Reduce number of arguments to autograd function in order + # to reduce CPU overhead due to pytorch arg checking. + (m_splits, padded_m_splits, is_grad_enabled) = non_tensor_args + # Make sure input dimensions are compatible in_features = inp.shape[-1] @@ -65,7 +68,7 @@ def backward(ctx, grad_output: torch.Tensor): grad_output.view(-1, in_features), grad_input, ctx.padded_m_splits, ctx.m_splits ) - return (grad_input, None, None, None) + return grad_input, None class Fp8Padding(torch.nn.Module): @@ -128,19 +131,20 @@ def forward( if m_splits == padded_m_splits: return inp, m_splits - if torch.is_grad_enabled(): + is_grad_enabled = torch.is_grad_enabled() + + if is_grad_enabled: fn = _Fp8Padding.apply - args = [] + autograd_ctx = [] else: fn = _Fp8Padding.forward - args = [None] + autograd_ctx = [None] - args += ( - inp, + non_tensor_args = ( m_splits, padded_m_splits, - torch.is_grad_enabled(), + is_grad_enabled, ) - out = fn(*args) + out = fn(*autograd_ctx, inp, non_tensor_args) return out, padded_m_splits diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index 7a01f15729..58187c20ea 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -4,7 +4,7 @@ """FP8 Padding API""" -from typing import List, Optional +from typing import List, Optional, Tuple import torch @@ -24,11 +24,14 @@ class _Fp8Unpadding(torch.autograd.Function): def forward( ctx, inp: torch.Tensor, - m_splits: List[int], - padded_m_splits: List[int], - is_grad_enabled: bool, + non_tensor_args: Tuple, ) -> torch.Tensor: # pylint: disable=missing-function-docstring + + # Reduce number of arguments to autograd function in order + # to reduce CPU overhead due to pytorch arg checking. + (m_splits, padded_m_splits, is_grad_enabled) = non_tensor_args + in_features = inp.shape[-1] # Allocate cast and transpose output tensor @@ -63,7 +66,7 @@ def backward(ctx, grad_output: torch.Tensor): grad_output.view(-1, in_features), grad_input, ctx.m_splits, ctx.padded_m_splits ) - return (grad_input, None, None, None) + return grad_input, None class Fp8Unpadding(torch.nn.Module): @@ -126,19 +129,20 @@ def forward( if m_splits == padded_m_splits: return inp - if torch.is_grad_enabled(): + is_grad_enabled = torch.is_grad_enabled() + + if is_grad_enabled: fn = _Fp8Unpadding.apply - args = [] + autograd_ctx = [] else: fn = _Fp8Unpadding.forward - args = [None] + autograd_ctx = [None] - args += ( - inp, + non_tensor_args = ( m_splits, padded_m_splits, - torch.is_grad_enabled(), + is_grad_enabled, ) - out = fn(*args) + out = fn(*autograd_ctx, inp, non_tensor_args) return out diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1a56a06da3..b3a96df399 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -14,7 +14,6 @@ from transformer_engine.common.recipe import Recipe from .base import ( get_dummy_wgrad, - get_multi_stream_cublas_workspace, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -28,6 +27,7 @@ clear_tensor_data, init_method_constant, requires_grad, + get_nvtx_range_context, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -40,7 +40,6 @@ ) from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo -from ..graph import is_graph_capturing from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer @@ -63,28 +62,34 @@ class _GroupedLinear(torch.autograd.Function): def forward( ctx, inp: torch.Tensor, - m_splits: List[int], - use_bias: bool, - is_first_microbatch: Union[bool, None], - fp8: bool, - fp8_calibration: bool, - wgrad_store: WeightGradStore, - input_quantizers: List[Quantizer], - weight_quantizers: List[Quantizer], - output_quantizers: List[Quantizer], - grad_output_quantizers: List[Quantizer], - fuse_wgrad_accumulation: bool, - cpu_offloading: bool, - sequence_parallel: bool, - activation_dtype: torch.dtype, - is_grad_enabled: bool, - module, - skip_fp8_weight_update, - save_original_input, + non_tensor_args: Tuple, *weights_and_biases, ) -> torch.Tensor: # pylint: disable=missing-function-docstring + # 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, + fp8_calibration, + wgrad_store, + input_quantizers, + weight_quantizers, + output_quantizers, + grad_output_quantizers, + fuse_wgrad_accumulation, + cpu_offloading, + sequence_parallel, + activation_dtype, + is_grad_enabled, + module, + skip_fp8_weight_update, + save_original_input, + ) = non_tensor_args + num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] biases = weights_and_biases[num_gemms:] @@ -183,7 +188,6 @@ def forward( inputmats, [out], activation_dtype, - get_multi_stream_cublas_workspace(), single_output=True, m_splits=m_splits, bias=biases, @@ -284,7 +288,7 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring - with torch.cuda.nvtx.range("_GroupedLinear_backward"): + with get_nvtx_range_context("_GroupedLinear_backward"): saved_tensors = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) N = ctx.num_gemms inputmats = saved_tensors[:N] @@ -372,7 +376,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], grad_output, [dgrad], ctx.activation_dtype, - get_multi_stream_cublas_workspace(), single_output=True, layout="NN", m_splits=ctx.m_splits, @@ -419,7 +422,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], grouped_gemm_wgrad = functools.partial( general_grouped_gemm, out_dtype=ctx.activation_dtype, - workspaces=get_multi_stream_cublas_workspace(), layout="NT", grad=True, m_splits=ctx.m_splits, @@ -484,28 +486,11 @@ def handle_custom_ddp_from_mcore(weight, wgrad): ): grad_biases = [None] * ctx.num_gemms - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): + 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, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, - None, *wgrad_list, *grad_biases, ) @@ -765,16 +750,9 @@ def forward( ), "GroupedLinear doesn't support input tensor in FP8." assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." - if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor() - else: - skip_fp8_weight_update = None - if skip_fp8_weight_update is not None: - is_first_microbatch = False + is_grad_enabled = torch.is_grad_enabled() - with torch.cuda.device( - getattr(self, list(self.named_parameters())[0][0]).device - ), self.prepare_forward(inp, num_gemms=self.num_gemms) as inp: + with self.prepare_forward(inp, num_gemms=self.num_gemms) as inp: weight_tensors = self._get_weight_tensors() bias_tensors = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] @@ -794,7 +772,7 @@ def forward( # TODO: use internal after #1638 is merged. # pylint: disable=fixme for i in range(self.num_gemms): input_quantizers[i].internal = False - if torch.is_grad_enabled(): + if is_grad_enabled: grad_output_quantizers = [ self.quantizers["scaling_bwd"][ self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"] @@ -804,14 +782,14 @@ def forward( for i in range(self.num_gemms): grad_output_quantizers[i].internal = True - if torch.is_grad_enabled(): + if is_grad_enabled: linear_fn = _GroupedLinear.apply - args = [] + autograd_ctx = [] else: linear_fn = _GroupedLinear.forward - args = [None] - args += ( - inp, + autograd_ctx = [None] + + non_tensor_args = ( m_splits, self.apply_bias, is_first_microbatch, @@ -826,14 +804,12 @@ def forward( is_cpu_offload_enabled(), self.sequence_parallel, self.activation_dtype, - torch.is_grad_enabled(), + is_grad_enabled, self, - skip_fp8_weight_update, + None, # skip_fp8_weight_update self.save_original_input, - *weight_tensors, - *bias_tensors, ) - out = linear_fn(*args) + out = linear_fn(*autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors) if self.return_bias: return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors] @@ -846,7 +822,7 @@ def backward_dw(self): """ if not self.need_backward_dw(): return - with torch.cuda.nvtx.range("_GroupedLinear_wgrad"): + with get_nvtx_range_context("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() wgrad_list = tensor_list[2] weight_params = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 4ed3ebb73f..3adbbc22e9 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -19,7 +19,6 @@ from transformer_engine.pytorch.tensor.utils import is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, - get_workspace, get_ub, TransformerEngineBaseModule, get_dummy_wgrad, @@ -40,6 +39,7 @@ nvtx_range_push, requires_grad, needs_quantized_gemm, + get_nvtx_range_context, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -96,47 +96,53 @@ def forward( ln_bias: Union[torch.Tensor, None], weight: torch.Tensor, bias: torch.Tensor, - eps: float, - is_first_microbatch: Union[bool, None], - fp8: bool, - fp8_calibration: bool, - wgrad_store: WeightGradStore, - fuse_wgrad_accumulation: bool, - 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], - cpu_offloading: bool, - tp_group: Union[dist_group_type, None], - tp_size: int, - sequence_parallel: bool, - tensor_parallel: bool, - activation_dtype: torch.dtype, - parallel_mode: Union[str, None], - return_layernorm_output: bool, - return_layernorm_output_gathered: bool, - is_grad_enabled: bool, - fwd_ln_sm_margin: int, - bwd_ln_sm_margin: int, - zero_centered_gamma: bool, - normalization: str, - ub_overlap_ag_fprop: bool, - ub_overlap_rs_fprop: bool, - ub_overlap_ag_dgrad: bool, - ub_overlap_rs_dgrad: bool, - ub_bulk_wgrad: bool, - ub_bulk_dgrad: bool, - ub_name: str, - fsdp_group: Union[dist_group_type, None], - module: torch.nn.Module, - skip_fp8_weight_update: bool, - symmetric_ar_type: str, - debug: Optional[bool] = False, + non_tensor_args: Tuple, ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: # pylint: disable=missing-function-docstring + # Reduce number of arguments to autograd function in order + # to reduce CPU overhead due to pytorch arg checking. + ( + eps, + is_first_microbatch, + fp8, + fp8_calibration, + wgrad_store, + fuse_wgrad_accumulation, + input_quantizer, + weight_quantizer, + output_quantizer, + grad_input_quantizer, + grad_weight_quantizer, + grad_output_quantizer, + cpu_offloading, + tp_group, + tp_size, + sequence_parallel, + tensor_parallel, + activation_dtype, + parallel_mode, + return_layernorm_output, + return_layernorm_output_gathered, + is_grad_enabled, + fwd_ln_sm_margin, + bwd_ln_sm_margin, + zero_centered_gamma, + normalization, + ub_overlap_ag_fprop, + ub_overlap_rs_fprop, + ub_overlap_ag_dgrad, + ub_overlap_rs_dgrad, + ub_bulk_wgrad, + ub_bulk_dgrad, + ub_name, + fsdp_group, + module, + skip_fp8_weight_update, + symmetric_ar_type, + debug, + ) = non_tensor_args + # NVTX label for profiling nvtx_label = "transformer_engine._LayerNormLinear.forward" if ub_name is not None: @@ -355,7 +361,6 @@ def forward( gemm_out, *_, reduce_scatter_out = general_gemm( weightmat, ln_out_total, - get_workspace(), quantization_params=output_quantizer, out_dtype=activation_dtype, bias=bias, @@ -544,7 +549,7 @@ def backward( if ctx.ub_name is not None: nvtx_label = f"{nvtx_label}.{ctx.ub_name}" - with torch.cuda.nvtx.range("_LayerNormLinear_backward"): + with get_nvtx_range_context("_LayerNormLinear_backward"): saved_tensors = ctx.saved_tensors ( # pylint: disable=unbalanced-tuple-unpacking inputmat, @@ -731,7 +736,6 @@ def backward( gemm_out, *_, reduce_scatter_out = general_gemm( weight, grad_output, - get_workspace(), layout="NN", grad=True, quantization_params=ctx.grad_input_quantizer, @@ -858,7 +862,6 @@ def backward( # Arguments to include in wgrad GEMM closure wgrad_gemm_kwargs = { - "workspace": get_workspace(), "out_dtype": ( main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype ), @@ -1026,44 +1029,7 @@ def wgrad_gemm( dbeta, wgrad, grad_bias, - None, # eps - None, # is_first_microbatch - None, # fp8 - None, # fp8_calibration - None, # wgrad_store - None, # fuse_wgrad_accumulation - None, # input_quantizer - None, # weight_quantizer - None, # output_quantizer - None, # grad_input_quantizer - None, # grad_weight_quantizer - None, # grad_output_quantizer - None, # cpu_offloading - None, # tp_group - None, # tp_size - None, # sequence_parallel - None, # tensor_parallel - None, # activation_dtype - None, # parallel_mode - None, # return_layernorm_output - None, # return_layernorm_output_gathered - None, # is_grad_enabled - None, # fwd_ln_sm_margin - None, # bwd_ln_sm_margin - None, # zero_centered_gamma - None, # normalization - None, # ub_overlap_ag_fprop - None, # ub_overlap_rs_fprop - None, # ub_overlap_ag_dgrad - None, # ub_overlap_rs_dgrad - None, # ub_bulk_dgrad - None, # ub_bulk_wgrad - None, # ub_name - None, # fsdp_group - None, # debug - None, # module - None, # skip_fp8_weight_update - None, # symmetric_ar_type + None, ) @@ -1523,8 +1489,10 @@ def forward( first microbatch (since it is the first gradient being produced) """ + is_grad_enabled = torch.is_grad_enabled() + if is_in_onnx_export_mode(): - return self.onnx_forward(inp, fp8_output) + return self.onnx_forward(inp, fp8_output, is_grad_enabled) debug = self.is_debug_iter() @@ -1546,9 +1514,7 @@ def forward( ).is_fp8_ubuf(): fp8_grad = True - with torch.cuda.device( - getattr(self, list(self.named_parameters())[0][0]).device - ), self.prepare_forward( + with self.prepare_forward( inp, allow_non_contiguous=False # removed .contiguous from inside the layer ) as inp: @@ -1556,14 +1522,14 @@ def forward( weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() quantizers = ( - self._get_quantizers(fp8_output, fp8_grad) + self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) if not debug - else self._get_debug_quantizers(fp8_output, fp8_grad) + else self._get_debug_quantizers(fp8_output, fp8_grad, is_grad_enabled) ) if debug: if self.no_debug_features_active(quantizers): debug = False - quantizers = self._get_quantizers(fp8_output, fp8_grad) + quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) ( input_quantizer, @@ -1574,18 +1540,13 @@ def forward( grad_output_quantizer, ) = quantizers - if torch.is_grad_enabled(): + if is_grad_enabled: fwd_fn = _LayerNormLinear.apply - args = [] + autograd_ctx = [] else: fwd_fn = _LayerNormLinear.forward - args = [None] - args += ( - inp, - self.layer_norm_weight, - self.layer_norm_bias, - weight_tensor, - bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None, + autograd_ctx = [None] + non_tensor_args = ( self.eps, is_first_microbatch, self.fp8, @@ -1607,8 +1568,8 @@ def forward( self.parallel_mode, self.return_layernorm_output, self.return_layernorm_output_gathered, - torch.is_grad_enabled(), - self.fwd_ln_sm_margin if torch.is_grad_enabled() else self.inf_ln_sm_margin, + is_grad_enabled, + self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin, self.bwd_ln_sm_margin, self.zero_centered_gamma, self.normalization, @@ -1625,7 +1586,15 @@ def forward( self.symmetric_ar_type, debug, ) - out = fwd_fn(*args) + out = fwd_fn( + *autograd_ctx, + inp, + self.layer_norm_weight, + self.layer_norm_bias, + weight_tensor, + bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None, + non_tensor_args, + ) if self.return_layernorm_output: out, ln_out = out @@ -1641,7 +1610,7 @@ def forward( return out, ln_out return out - def _get_quantizers(self, fp8_output, fp8_grad): + def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): if not self.fp8: return [None] * 6 grad_input_quantizer = None @@ -1653,7 +1622,7 @@ def _get_quantizers(self, fp8_output, fp8_grad): (weight_quantizer,) = self._get_weight_quantizers() if fp8_output: output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] - if torch.is_grad_enabled(): + if is_grad_enabled: grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] grad_output_quantizer.internal = True if fp8_grad: @@ -1668,8 +1637,8 @@ def _get_quantizers(self, fp8_output, fp8_grad): grad_output_quantizer, ) - def _get_debug_quantizers(self, fp8_output, fp8_grad): - original_quantizers = self._get_quantizers(fp8_output, fp8_grad) + def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): + original_quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) assert TEDebugState.debug_enabled from ...debug.pytorch.debug_quantization import DebugQuantizer @@ -1694,6 +1663,7 @@ def onnx_forward( self, inp: torch.Tensor, fp8_output: bool, + is_grad_enabled: bool, ) -> torch.Tensor: """ ONNX-compatible version of the forward function that provides numerical equivalence @@ -1709,7 +1679,7 @@ def onnx_forward( weight_quantizer, output_quantizer, *_, - ) = self._get_quantizers(fp8_output, fp8_grad=False) + ) = self._get_quantizers(fp8_output, False, is_grad_enabled) inp_dtype = inp.dtype weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index c29775c926..35dcb10f34 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -20,7 +20,6 @@ from transformer_engine.pytorch.tensor.utils import is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, - get_workspace, _ub_communicators, get_ub, TransformerEngineBaseModule, @@ -45,6 +44,7 @@ clear_tensor_data, requires_grad, needs_quantized_gemm, + get_nvtx_range_context, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -174,55 +174,61 @@ def forward( fc1_bias: torch.Tensor, fc2_weight: torch.Tensor, fc2_bias: torch.Tensor, - eps: float, - is_first_microbatch: Union[bool, None], - fp8: bool, - fp8_calibration: bool, - wgrad_store: WeightGradStore, - fuse_wgrad_accumulation: bool, - fc1_input_quantizer: Optional[Quantizer], - fc1_weight_quantizer: Optional[Quantizer], - fc1_output_quantizer: Optional[Quantizer], - fc1_grad_input_quantizer: Optional[Quantizer], - fc1_grad_weight_quantizer: Optional[Quantizer], - fc1_grad_output_quantizer: Optional[Quantizer], - fc2_input_quantizer: Optional[Quantizer], - fc2_weight_quantizer: Optional[Quantizer], - fc2_output_quantizer: Optional[Quantizer], - fc2_grad_input_quantizer: Optional[Quantizer], - fc2_grad_weight_quantizer: Optional[Quantizer], - fc2_grad_output_quantizer: Optional[Quantizer], - cpu_offloading: bool, - tp_group: Union[dist_group_type, None], - tp_size: int, - sequence_parallel: bool, - tensor_parallel: bool, - activation_dtype: torch.dtype, - return_layernorm_output: bool, - return_layernorm_output_gathered: bool, - bias_gelu_fusion: bool, - set_parallel_mode: bool, - is_grad_enabled: bool, - fwd_ln_sm_margin: int, - bwd_ln_sm_margin: int, - zero_centered_gamma: bool, - activation: str, - activation_params: Optional[dict], - normalization: str, - ub_overlap_ag: bool, - ub_overlap_rs: bool, - ub_overlap_rs_dgrad: bool, - ub_bulk_wgrad: bool, - ub_bulk_dgrad: bool, - gemm_gelu_fusion: bool, - fsdp_group: Union[dist_group_type, None], - module: torch.nn.Module, - skip_fp8_weight_update: bool, - symmetric_ar_type: str, - debug: Optional[bool] = False, + non_tensor_args: Tuple, ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: # pylint: disable=missing-function-docstring + # Reduce number of arguments to autograd function in order + # to reduce CPU overhead due to pytorch arg checking. + ( + eps, + is_first_microbatch, + fp8, + fp8_calibration, + wgrad_store, + fuse_wgrad_accumulation, + fc1_input_quantizer, + fc1_weight_quantizer, + fc1_output_quantizer, + fc1_grad_input_quantizer, + fc1_grad_weight_quantizer, + fc1_grad_output_quantizer, + fc2_input_quantizer, + fc2_weight_quantizer, + fc2_output_quantizer, + fc2_grad_input_quantizer, + fc2_grad_weight_quantizer, + fc2_grad_output_quantizer, + cpu_offloading, + tp_group, + tp_size, + sequence_parallel, + tensor_parallel, + activation_dtype, + return_layernorm_output, + return_layernorm_output_gathered, + bias_gelu_fusion, + set_parallel_mode, + is_grad_enabled, + fwd_ln_sm_margin, + bwd_ln_sm_margin, + zero_centered_gamma, + activation, + activation_params, + normalization, + ub_overlap_ag, + ub_overlap_rs, + ub_overlap_rs_dgrad, + ub_bulk_wgrad, + ub_bulk_dgrad, + gemm_gelu_fusion, + fsdp_group, + module, + skip_fp8_weight_update, + symmetric_ar_type, + debug, + ) = non_tensor_args + # Make sure input dimensions are compatible in_features, inp_shape = ln_weight.numel(), inp.shape assert inp_shape[-1] == in_features, "GEMM not possible" @@ -433,7 +439,6 @@ def forward( fc1_outputs = general_gemm( fc1_weight_final, ln_out_total, - get_workspace(), quantization_params=( fc2_input_quantizer if gemm_gelu_fusion @@ -517,7 +522,6 @@ def forward( gemm_out, *_, reduce_scatter_out = general_gemm( fc2_weight_final, act_out, - get_workspace(), out_dtype=activation_dtype, bias=fc2_bias, quantization_params=fc2_output_quantizer, @@ -704,7 +708,7 @@ def backward( ctx, *grad_outputs: Tuple[torch.Tensor, ...] ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring - with torch.cuda.nvtx.range("_LayerNormMLP_backward"): + with get_nvtx_range_context("_LayerNormMLP_backward"): saved_tensors = ctx.saved_tensors ( # pylint: disable=unbalanced-tuple-unpacking inputmat, @@ -874,7 +878,6 @@ def backward( gemm_output, *_ = general_gemm( fc2_weight, grad_output, - get_workspace(), layout="NN", grad=True, quantization_params=( @@ -968,7 +971,6 @@ def backward( # Arguments to include in wgrad GEMM closure fc2_wgrad_gemm_kwargs = { - "workspace": get_workspace(), "out_dtype": ( origin_fc2_weight.main_grad.dtype if ctx.fuse_wgrad_accumulation @@ -1138,7 +1140,6 @@ def fc2_wgrad_gemm( gemm_out, *_, reduce_scatter_out = general_gemm( fc1_weight, dact, - get_workspace(), out=gemm_out, out_dtype=ctx.activation_dtype, quantization_params=ctx.fc1_grad_input_quantizer, @@ -1217,7 +1218,6 @@ def fc2_wgrad_gemm( # Arguments to include in wgrad GEMM closure fc1_wgrad_gemm_kwargs = { - "workspace": get_workspace(), "out_dtype": ( origin_fc1_weight.main_grad.dtype if ctx.fuse_wgrad_accumulation @@ -1399,52 +1399,7 @@ def fc1_wgrad_gemm( fc1_bias_grad if fc1_bias is not None else None, fc2_wgrad, # pylint: disable=possibly-used-before-assignment fc2_bias_grad, - None, # eps - None, # is_first_microbatch - None, # fp8 - None, # fp8_calibration - None, # wgrad_store - None, # fuse_wgrad_accumulation - None, # fc1_input_quantizer, - None, # fc1_weight_quantizer, - None, # fc1_output_quantizer, - None, # fc1_grad_input_quantizer, - None, # fc1_grad_weight_quantizer, - None, # fc1_grad_output_quantizer, - None, # fc2_input_quantizer, - None, # fc2_weight_quantizer, - None, # fc2_output_quantizer, - None, # fc2_grad_input_quantizer, - None, # fc2_grad_weight_quantizer, - None, # fc2_grad_output_quantizer, - None, # cpu_offloading - None, # tp_group - None, # tp_size - None, # sequence_parallel - None, # tensor_parallel - None, # activation_dtype - None, # return_layernorm_output - None, # return_layernorm_output_gathered - None, # bias_gelu_fusion - None, # set_parallel_mode - None, # is_grad_enabled - None, # fwd_ln_sm_margin - None, # bwd_ln_sm_margin - None, # zero_centered_gamma - None, # activation - None, # activation_params - None, # normalization - None, # ub_overlap_ag - None, # ub_overlap_rs - None, # ub_overlap_rs_dgrad - None, # ub_bulk_dgrad - None, # ub_bulk_wgrad - None, # gemm_gelu_fusion - None, # fsdp_group - None, # module - None, # skip_fp8_weight_update - None, # symmetric_ar_type - None, # debug + None, ) @@ -1827,8 +1782,10 @@ def forward( first microbatch (since it is the first gradient being produced) """ + is_grad_enabled = torch.is_grad_enabled() + if is_in_onnx_export_mode(): - return self.onnx_forward(inp) + return self.onnx_forward(inp, is_grad_enabled) debug = self.is_debug_iter() @@ -1844,19 +1801,17 @@ def forward( if get_ub("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()).is_fp8_ubuf(): fp8_output = True - with torch.cuda.device( - getattr(self, list(self.named_parameters())[0][0]).device - ), self.prepare_forward(inp, num_gemms=2) as inp: + with self.prepare_forward(inp, num_gemms=2) as inp: quantizers = ( - self._get_quantizers(fp8_output) + self._get_quantizers(fp8_output, is_grad_enabled) if not debug - else self._get_debug_quantizers(fp8_output) + else self._get_debug_quantizers(fp8_output, is_grad_enabled) ) if debug: if self.no_debug_features_active(quantizers): debug = False - quantizers = self._get_quantizers(fp8_output) + quantizers = self._get_quantizers(fp8_output, is_grad_enabled) # Get quantizers ( @@ -1888,20 +1843,14 @@ def forward( if self.bias_gelu_nvfusion and not use_reentrant_activation_recompute(): self.bias_gelu_nvfusion = False - if torch.is_grad_enabled(): + if is_grad_enabled: fwd_fn = _LayerNormMLP.apply - args = [] + autograd_ctx = [] else: fwd_fn = _LayerNormMLP.forward - args = [None] - args += ( - inp, - self.layer_norm_weight, - self.layer_norm_bias, - fc1_weight, - fc1_bias, - fc2_weight, - fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None, + autograd_ctx = [None] + + non_tensor_args = ( self.eps, is_first_microbatch, self.fp8, @@ -1930,8 +1879,8 @@ def forward( self.return_layernorm_output_gathered, self.bias_gelu_nvfusion and not self.fp8 and not debug, self.set_parallel_mode, - torch.is_grad_enabled(), - self.fwd_ln_sm_margin if torch.is_grad_enabled() else self.inf_ln_sm_margin, + is_grad_enabled, + self.fwd_ln_sm_margin if is_grad_enabled else self.inf_ln_sm_margin, self.bwd_ln_sm_margin, self.zero_centered_gamma, self.activation, @@ -1949,7 +1898,17 @@ def forward( self.symmetric_ar_type, debug, ) - out = fwd_fn(*args) + out = fwd_fn( + *autograd_ctx, + inp, + self.layer_norm_weight, + self.layer_norm_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None, + non_tensor_args, + ) if self.return_layernorm_output: out, ln_out = out @@ -1965,7 +1924,7 @@ def forward( return out, ln_out return out - def _get_quantizers(self, fp8_output): + def _get_quantizers(self, fp8_output, is_grad_enabled): ( fc1_input_quantizer, fc1_output_quantizer, @@ -1995,7 +1954,7 @@ def _get_quantizers(self, fp8_output): fc2_output_quantizer = self.quantizers["scaling_fwd"][ tex.FP8FwdTensors.GEMM2_OUTPUT ] - if torch.is_grad_enabled(): + if is_grad_enabled: fc2_grad_output_quantizer = self.quantizers["scaling_bwd"][ tex.FP8BwdTensors.GRAD_OUTPUT2 ] @@ -2020,7 +1979,9 @@ def _get_quantizers(self, fp8_output): fc2_grad_output_quantizer, ) - def onnx_forward(self, inp: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + def onnx_forward( + self, inp: torch.Tensor, is_grad_enabled: bool + ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ ONNX-compatible version of the forward function that provides numerical equivalence while only using operations that have defined ONNX symbolic translations. @@ -2037,7 +1998,7 @@ def onnx_forward(self, inp: torch.Tensor) -> Union[torch.Tensor, Tuple[torch.Ten fc2_weight_quantizer, output_quantizer, *_, - ) = self._get_quantizers(False) + ) = self._get_quantizers(False, is_grad_enabled) inp_dtype = inp.dtype fc1_weight, fc2_weight = self._get_weight_tensors() @@ -2122,10 +2083,10 @@ def _clamped_swiglu(x, limit, alpha): return fc2_out, fc2_bias.to(inp_dtype) return fc2_out - def _get_debug_quantizers(self, fp8_output): + def _get_debug_quantizers(self, fp8_output, is_grad_enabled): from ...debug.pytorch.debug_quantization import DebugQuantizer - base_quantizers = list(self._get_quantizers(fp8_output)) + base_quantizers = list(self._get_quantizers(fp8_output, is_grad_enabled)) assert TEDebugState.debug_enabled def make_debug(prefix, offset): @@ -2268,7 +2229,7 @@ def backward_dw(self): """ if not self.need_backward_dw(): return - with torch.cuda.nvtx.range("_LayerNormMLP_wgrad"): + with get_nvtx_range_context("_LayerNormMLP_wgrad"): (fc2_wgrad, fc2_bias_grad_, *_), tensor_list_fc2 = self.wgrad_store.pop() if self.use_bias and self.fc1_bias.grad is None: (fc1_wgrad, fc1_bias_grad, *_), _ = self.wgrad_store.pop() diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 00b78995fe..b3f8165a77 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -19,7 +19,6 @@ fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, - get_workspace, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -38,6 +37,7 @@ assert_dim_for_all_gather, nvtx_range_pop, nvtx_range_push, + get_nvtx_range_context, ) from ..distributed import ( set_tensor_model_parallel_attributes, @@ -90,42 +90,46 @@ def forward( weight: torch.Tensor, inp: torch.Tensor, bias: Optional[torch.Tensor], - is_first_microbatch: Union[bool, None], - fp8: bool, - fp8_calibration: bool, - wgrad_store: WeightGradStore, - 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], - fuse_wgrad_accumulation: bool, - cpu_offloading: bool, - tp_group: Union[dist_group_type, None], - tp_size: int, - sequence_parallel: bool, - tensor_parallel: bool, - activation_dtype: torch.dtype, - parallel_mode: Union[str, None], - is_grad_enabled: bool, - ub_overlap_rs_fprop: bool, - ub_overlap_ag_dgrad: bool, - ub_overlap_ag_fprop: bool, - ub_overlap_rs_dgrad: bool, - ub_bulk_dgrad: bool, - ub_bulk_wgrad: bool, - ub_name: str, - fp8_output: bool, # pylint: disable=unused-argument - fsdp_group: Union[dist_group_type, None], - module: torch.nn.Module, - skip_fp8_weight_update: bool, - symmetric_ar_type: str, - save_original_input: bool = False, - debug: Optional[bool] = False, + non_tensor_args: Tuple, ) -> torch.Tensor: # pylint: disable=missing-function-docstring + ( + is_first_microbatch, + fp8, + fp8_calibration, + wgrad_store, + input_quantizer, + weight_quantizer, + output_quantizer, + grad_input_quantizer, + grad_weight_quantizer, + grad_output_quantizer, + 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, # pylint: disable=unused-variable + fsdp_group, + module, + skip_fp8_weight_update, + symmetric_ar_type, + save_original_input, + debug, + ) = non_tensor_args + # NVTX label for profiling nvtx_label = "transformer_engine._Linear.forward" if ub_name is not None: @@ -320,7 +324,6 @@ def forward( gemm_out, *_, reduce_scatter_out = general_gemm( weightmat, inputmat_total, - get_workspace(), quantization_params=output_quantizer, out_dtype=activation_dtype, bias=bias, @@ -497,7 +500,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if ctx.ub_name is not None: nvtx_label = f"{nvtx_label}.{ctx.ub_name}" - with torch.cuda.nvtx.range("_Linear_backward"): + with get_nvtx_range_context("_Linear_backward"): saved_tensors = ctx.saved_tensors inputmat, weight_fp8, weight, bias = ( # pylint: disable=unbalanced-tuple-unpacking restore_from_saved(ctx.tensor_objects, saved_tensors) @@ -719,7 +722,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], gemm_out, *_, reduce_scatter_out = general_gemm( weight_fp8, grad_output, - get_workspace(), layout="NN", grad=True, quantization_params=ctx.grad_input_quantizer, @@ -845,7 +847,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], # Arguments to include in wgrad GEMM closure wgrad_gemm_kwargs = { - "workspace": get_workspace(), "out_dtype": ( main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype ), @@ -977,39 +978,7 @@ def wgrad_gemm( wgrad, dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, grad_bias, - None, # is_first_microbatch - None, # fp8 - None, # fp8_calibration - None, # wgrad_store - None, # input_quantizer - None, # weight_quantizer - None, # output_quantizer - None, # grad_input_quantizer - None, # grad_weight_quantizer - None, # grad_output_quantizer - None, # fuse_wgrad_accumulation - None, # cpu_offloading - None, # tp_group - None, # tp_size - None, # sequence_parallel - None, # tensor_parallel - None, # activation_dtype - None, # parallel_mode - None, # is_grad_enabled - None, # ub_overlap_rs_fprop - None, # ub_overlap_ag_dgrad - None, # ub_overlap_ag_fprop - None, # ub_overlap_rs_dgrad - None, # ub_bulk_dgrad - None, # ub_bulk_wgrad - None, # ub_name - None, # fp8_output - None, # fsdp_group - None, # module - None, # skip_fp8_weight_update - None, # symmetric_ar_type - None, # save_original_input - None, # debug + None, ) @@ -1403,8 +1372,10 @@ def forward( first microbatch (since it is the first gradient being produced) """ + is_grad_enabled = torch.is_grad_enabled() + if is_in_onnx_export_mode(): - return self.onnx_forward(inp, fp8_output) + return self.onnx_forward(inp, fp8_output, is_grad_enabled) debug = self.is_debug_iter() @@ -1426,9 +1397,7 @@ def forward( ).is_fp8_ubuf(): fp8_grad = True - with torch.cuda.device( - getattr(self, list(self.named_parameters())[0][0]).device - ), self.prepare_forward( + with self.prepare_forward( inp, allow_non_contiguous=isinstance(inp, QuantizedTensor), ) as inp: @@ -1436,14 +1405,14 @@ def forward( weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() quantizers = ( - self._get_quantizers(fp8_output, fp8_grad) + self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) if not debug - else self._get_debug_quantizers(fp8_output, fp8_grad) + else self._get_debug_quantizers(fp8_output, fp8_grad, is_grad_enabled) ) if debug: if self.no_debug_features_active(quantizers): debug = False - quantizers = self._get_quantizers(fp8_output, fp8_grad) + quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) ( input_quantizer, @@ -1454,16 +1423,14 @@ def forward( grad_output_quantizer, ) = quantizers - if torch.is_grad_enabled(): + if is_grad_enabled: linear_fn = _Linear.apply - args = [] + autograd_ctx = [] else: linear_fn = _Linear.forward - args = [None] - args += ( - weight_tensor, - inp, - bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None, + autograd_ctx = [None] + + non_tensor_args = ( is_first_microbatch, self.fp8, self.fp8_calibration, @@ -1482,7 +1449,7 @@ def forward( self.tp_size > 1, self.activation_dtype, self.parallel_mode, - torch.is_grad_enabled(), + is_grad_enabled, self.ub_overlap_rs_fprop, self.ub_overlap_ag_dgrad, self.ub_overlap_ag_fprop, @@ -1498,7 +1465,13 @@ def forward( self.save_original_input, debug, ) - out = linear_fn(*args) + out = linear_fn( + *autograd_ctx, + weight_tensor, + inp, + bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None, + non_tensor_args, + ) if self.gemm_bias_unfused_add: out = out + cast_if_needed(bias_tensor, self.activation_dtype) @@ -1506,7 +1479,7 @@ def forward( return out, cast_if_needed(bias_tensor, self.activation_dtype) return out - def _get_quantizers(self, fp8_output, fp8_grad): + def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): if not self.fp8: return [None] * 6 grad_input_quantizer = None @@ -1518,7 +1491,7 @@ def _get_quantizers(self, fp8_output, fp8_grad): (weight_quantizer,) = self._get_weight_quantizers() if fp8_output: output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] - if torch.is_grad_enabled(): + if is_grad_enabled: grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] grad_output_quantizer.internal = True if fp8_grad: @@ -1532,8 +1505,8 @@ def _get_quantizers(self, fp8_output, fp8_grad): grad_output_quantizer, ) - def _get_debug_quantizers(self, fp8_output, fp8_grad): - original_quantizers = self._get_quantizers(fp8_output, fp8_grad) + def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): + original_quantizers = self._get_quantizers(fp8_output, fp8_grad, is_grad_enabled) assert TEDebugState.debug_enabled from ...debug.pytorch.debug_quantization import DebugQuantizer @@ -1588,6 +1561,7 @@ def onnx_forward( self, inp: torch.Tensor, fp8_output: bool, + is_grad_enabled: bool, ) -> torch.Tensor: """ ONNX-compatible version of the forward function that provides numerical equivalence @@ -1604,7 +1578,7 @@ def onnx_forward( weight_quantizer, output_quantizer, *_, - ) = self._get_quantizers(fp8_output, False) + ) = self._get_quantizers(fp8_output, False, is_grad_enabled) inp_dtype = inp.dtype if input_quantizer is not None: diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 432d8c134b..749ab7a650 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -25,7 +25,6 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, get_dummy_wgrad, - get_workspace, ) from ...tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer @@ -585,7 +584,6 @@ def _functional_forward( y, *_ = general_gemm( w, x, - get_workspace(), out_dtype=dtype, quantization_params=output_quantizer, alpha=alpha, @@ -875,7 +873,6 @@ def _functional_backward( dx, *_ = general_gemm( w, dy, - get_workspace(), out_dtype=dtype, quantization_params=grad_input_quantizer, alpha=grad_input_alpha, @@ -928,7 +925,6 @@ def _functional_backward( dw, *_ = general_gemm( x, dy, - get_workspace(), out_dtype=dw_dtype, alpha=grad_weight_alpha, beta=grad_weight_beta, diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index fd1820d15d..32e4ee3657 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -19,7 +19,6 @@ fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, - get_workspace, ) from ...quantized_tensor import Quantizer from ...tensor.mxfp8_tensor import MXFP8Quantizer @@ -378,7 +377,6 @@ def _functional_backward( dx, *_ = general_gemm( w, dy, - get_workspace(), out_dtype=dtype, quantization_params=grad_input_quantizer, layout="NN", @@ -464,7 +462,6 @@ def _functional_backward( dw, *_ = general_gemm( x, dy, - get_workspace(), out_dtype=dw_dtype, accumulate=accumulate_into_grad_weight, layout="NT", diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 057eb576d7..d50d031ba7 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -18,7 +18,6 @@ from ...module.base import ( fill_userbuffers_buffer_for_all_gather, get_ub, - get_workspace, _2X_ACC_FPROP, ) from ...quantized_tensor import Quantizer @@ -243,7 +242,6 @@ def _functional_forward( gemm_output, *_, reduce_scatter_output = general_gemm( w, x, - get_workspace(), out_dtype=dtype, quantization_params=output_quantizer, bias=bias, diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index c830b19e9f..3e3f460b41 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -7,7 +7,6 @@ from __future__ import annotations from typing import Optional, Tuple, Iterable, Any, Dict, Union import abc -import copy import warnings import math @@ -297,10 +296,6 @@ def set_usage( if columnwise is not None: self.columnwise_usage = columnwise - def copy(self) -> Quantizer: - """Create shallow copy""" - return copy.copy(self) - def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: """Symbolic function for ONNX export""" raise NotImplementedError( diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 8440c14b74..069565f388 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -57,6 +57,22 @@ def __init__( self.block_scaling_dim = block_scaling_dim self.all_gather_usage = all_gather_usage + def copy(self) -> Float8BlockQuantizer: + """Create shallow copy""" + + quantizer = Float8BlockQuantizer( + fp8_dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + block_scaling_dim=self.block_scaling_dim, + all_gather_usage=self.all_gather_usage, + amax_epsilon=self.amax_epsilon, + force_pow_2_scales=self.force_pow_2_scales, + ) + quantizer.internal = self.internal + + return quantizer + def update_quantized( self, src: torch.Tensor, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 7f7195a17f..80e7ed4674 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -66,6 +66,20 @@ def __init__( self.amax = amax self.dtype = fp8_dtype + def copy(self) -> Float8Quantizer: + """Create shallow copy""" + + quantizer = Float8Quantizer( + scale=self.scale, + amax=self.amax, + fp8_dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + + return quantizer + def update_quantized( self, src: torch.Tensor, @@ -245,10 +259,16 @@ def __init__( amax_reduction_group: Optional[dist_group_type] = None, force_pow_2_scales: bool = False, amax_epsilon: float = 0.0, + scale: Optional[torch.Tensor] = None, + amax: Optional[torch.Tensor] = None, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) - self.scale = torch.empty(1, dtype=torch.float32, device=device) - self.amax = torch.empty(1, dtype=torch.float32, device=device) + 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 self.dtype = fp8_dtype self.use_existing_amax = use_existing_amax self.with_amax_reduction = with_amax_reduction @@ -256,6 +276,26 @@ def __init__( self.force_pow_2_scales = force_pow_2_scales self.amax_epsilon = amax_epsilon + def copy(self) -> Float8CurrentScalingQuantizer: + """Create shallow copy""" + + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=self.dtype, + device=0, + rowwise=self.rowwise_usage, + 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 + + return quantizer + def update_quantized( self, src: torch.Tensor, diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 7ca6e3b0dd..b42b328091 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -45,6 +45,18 @@ def __init__( super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = fp8_dtype + def copy(self) -> MXFP8Quantizer: + """Create shallow copy""" + + quantizer = MXFP8Quantizer( + fp8_dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + ) + quantizer.internal = self.internal + + return quantizer + def update_quantized( self, src: torch.Tensor, diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 31dbcf00a9..652163295c 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -176,6 +176,26 @@ def update_quantized( return dst + def copy(self) -> NVFP4Quantizer: + """Create shallow copy""" + + quantizer = NVFP4Quantizer( + fp4_dtype=self.dtype, + rowwise=self.rowwise_usage, + columnwise=self.columnwise_usage, + with_amax_reduction=self.with_amax_reduction, + amax_reduction_group=self.amax_reduction_group, + with_rht=self.with_rht, + with_post_rht_amax=self.with_post_rht_amax, + with_2d_quantization=self.with_2d_quantization, + stochastic_rounding=self.stochastic_rounding, + ) + quantizer.internal = self.internal + quantizer.rht_matrix = self.rht_matrix + quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t + + return quantizer + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 90c6289963..083117b7b4 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -8,6 +8,7 @@ import math import os from typing import Any, Callable, List, Optional, Sequence, Tuple, Union +from contextlib import nullcontext import numpy as np import torch @@ -592,6 +593,24 @@ def _nvtx_enabled() -> bool: _nvtx_range_messages: list[str] = [] +def get_nvtx_range_context(msg: str): + """Get NVTX context manager to tag module forward and backward passes. + + Set `NVTE_NVTX_ENABLED=1` in the environment to enable NVTX + context manager for module level profiling tags. + + Parameters + ---------- + msg: str + Message to associate with profiling context. + + """ + + if _nvtx_enabled(): + return torch.cuda.nvtx.range(msg) + return nullcontext() + + def nvtx_range_push(msg: str) -> None: """Push NVTX range onto stack, if NVTX range profiling is enabled From 1df4a69f761672f633d40ea3605327087d1ea737 Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Mon, 17 Nov 2025 20:48:49 +0100 Subject: [PATCH 072/521] [PyTorch] Enable reference Current Scaling recipe (#2368) * Enable reference current scaling recipe Signed-off-by: Evgeny * minor 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 * Test ref vs native Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Evgeny Co-authored-by: Kirthi Shankar Sivamani --- .../test_float8_current_scaling_exact.py | 132 +++++ .../quantization_current_scaling.py | 525 ++++++++++++++++++ .../custom_recipes/quantization_nvfp4.py | 6 +- 3 files changed, 660 insertions(+), 3 deletions(-) create mode 100644 transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index e4d6ce3651..fd47b66c7e 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -8,9 +8,15 @@ import pytest import transformer_engine.pytorch as te +import transformer_engine_torch as tex from transformer_engine.common.recipe import Float8CurrentScaling 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 ( + CurrentScalingQuantizerRef, +) # read env variable NVTE_TEST_FLOAT8_CURRENT_SCALING_EXACT_TENSOR_DUMP_DIR to override the default tensor dump directory @@ -749,6 +755,132 @@ def test_fp8_current_scaling_with_linear_module( ) +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestFP8CurrentScalingNativeVsRef: + @staticmethod + def _make_quantizers(rowwise=True, columnwise=True): + # TE native FP8 current scaling quantizer + te_quant = te.Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=torch.device("cuda"), + rowwise=rowwise, + columnwise=columnwise, + ) + # Reference quantizer + ref_quant = CurrentScalingQuantizerRef( + dtype=torch.float8_e4m3fn, + rowwise=rowwise, + columnwise=columnwise, + pow_2_scales=False, + eps=0.0, + ) + return te_quant, ref_quant + + @pytest.mark.parametrize( + "M, N, dtype", + [ + (128, 256, torch.bfloat16), + ], + ids=["rowwise"], + ) + def test_current_scaling_quantization_versus_reference(self, M, N, dtype): + device = "cuda" + seed = 123 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + x = torch.randn((M, N), dtype=dtype, device=device) + + te_quant, ref_quant = self._make_quantizers(rowwise=True, columnwise=False) + + # Native TE quantization + x_te = te_quant(x) + assert x_te._data is not None + qx_native = x_te._data.view(dtype=torch.float8_e4m3fn) + sx_native = x_te._scale_inv + + # Reference quantization + x_ref = ref_quant.quantize(x) + qx_ref = x_ref.data + sx_ref = x_ref.scale + + # Byte-for-byte equality on data and exact scale_inv match + torch.testing.assert_close(qx_native, qx_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(sx_native, sx_ref, atol=0.0, rtol=0.0) + + @pytest.mark.parametrize( + "M, K, N, out_dtype, accumulate", + [ + (128, 256, 96, torch.bfloat16, False), + (64, 128, 64, torch.float32, True), + ], + ids=["bf16_no_acc", "fp32_acc"], + ) + def test_current_scaling_gemm_versus_reference(self, M, K, N, out_dtype, accumulate): + device = "cuda" + seed = 42 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + x = torch.randn((M, K), dtype=torch.bfloat16, device=device) + w = torch.randn((N, K), dtype=torch.bfloat16, device=device) + out = torch.randn((M, N), dtype=out_dtype, device=device) if accumulate else None + + te_quant_x, ref_quant = self._make_quantizers(rowwise=True, columnwise=True) + te_quant_w, _ = self._make_quantizers(rowwise=True, columnwise=True) + + # Native TE quantization (direct) + qx_native = te_quant_x(x) + qw_native = te_quant_w(w) + + # Prepare inputs for reference qgemm + assert qx_native._data is not None and qw_native._data is not None + qx_data = qx_native._data.view(dtype=torch.float8_e4m3fn) + qw_data = qw_native._data.view(dtype=torch.float8_e4m3fn) + sx = qx_native._scale_inv + sw = qw_native._scale_inv + + # Reference GEMM + m_params = MMParams(out_dtype=out_dtype, use_split_accumulator=False) + y_ref = ref_quant.qgemm( + qx=qx_data, + qw=qw_data, + m_params=m_params, + out_dtype=out_dtype, + sx=sx, + sw=sw, + bias=None, + out=out.clone() if accumulate else None, + accumulate=accumulate, + gemm_type=None, + qresult_x=None, + qresult_w=None, + ) + + # Native TE GEMM + # return type is out, bias_grad, gelu_input, extra_output + y_native = tex.generic_gemm( + qw_native, # A + True, # transa (treat (N,K) as (K,N)) + qx_native, # B + False, # transb + out.clone() if accumulate else None, + None, # out quantizer + TE_DType[out_dtype], + None, # bias + TE_DType[torch.bfloat16], + False, # use_gelu + None, # gelu_input + False, # use_grad + torch.empty(0, dtype=torch.uint8, device=device), + 0, + accumulate, + False, # use_split_accumulator + )[0] + + torch.testing.assert_close(y_native, y_ref, atol=0.0, rtol=0.0) + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) class TestFP8CurrentScalingRecipeLayerNormLinear(TestFP8RecipeLayerNormLinearBase): diff --git a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py b/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py new file mode 100644 index 0000000000..96cbca772c --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py @@ -0,0 +1,525 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Current scaling recipe reference implementation.""" + +import dataclasses +import math +from typing import Optional, Tuple, Iterable + +import torch + +from transformer_engine.pytorch.custom_recipes import quantization +from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage, Quantizer + + +def current_scaling_ref_quantizer_factory(role): + """Factory function for current scaling reference quantizer. + + 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 + return CurrentScalingQuantizerRef( + dtype=dtype, + rowwise=True, + columnwise=True, + pow_2_scales=False, + eps=0.0, + ) + + +@dataclasses.dataclass +class CurrentScalingTensorRef(QuantizedTensorStorage): + """Reference implementation of current scaling quantized tensor""" + + data: Optional[torch.Tensor] = None + scale: Optional[torch.Tensor] = None + data_t: Optional[torch.Tensor] = None + scale_t: Optional[torch.Tensor] = None + + dtype: Optional[torch.dtype] = None + device: Optional[torch.device] = None + quant_dtype: Optional[torch.dtype] = None + original_shape: Optional[Tuple[int, ...]] = None + _quantizer: Optional[Quantizer] = None + + @property + def custom(self) -> bool: + """Flag to indicate this quantized tensor is custom.""" + return True + + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]: + """Prepare the quantization result for saving for backward""" + tensors = [self.data, self.data_t, self.scale, self.scale_t] + self.data = None + self.data_t = None + self.scale = None + self.scale_t = None + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + """Restore the quantization result from the saved tensors""" + self.data = tensors[0] + self.data_t = tensors[1] + self.scale = tensors[2] + self.scale_t = tensors[3] + return tensors[4:] + + # Compatibility + @property + def _data(self): + return self.data + + @_data.setter + def _data(self, value): + self.data = value + + @property + def _scale_inv(self): + return self.scale + + @_scale_inv.setter + def _scale_inv(self, value): + self.scale = value + + def __repr__(self): + return ( + f"{self.__class__.__name__}(" + f"dtype={self.dtype}, " + f"device={self.device}, " + f"quant_dtype={self.quant_dtype}, " + f"original_shape={self.original_shape}" + ")" + ) + + def update_usage( + self, + rowwise_usage: Optional[bool] = None, + columnwise_usage: Optional[bool] = None, + ): + """Generate or remove quantized data based on provided usage.""" + has_data = self.data is not None + has_data_transpose = self.data_t is not None + needs_data = has_data + needs_data_transpose = has_data_transpose + + if rowwise_usage is not None: + needs_data = rowwise_usage + if columnwise_usage is not None: + needs_data_transpose = columnwise_usage + + # Generate data that is required + if needs_data and not has_data: + raise RuntimeError("Cannot generate FP8 data, even from FP8 data transpose") + if needs_data_transpose and not has_data_transpose: + if not has_data: + raise RuntimeError("FP8 data is required to generate FP8 data transpose") + self._create_transpose() + + # Delete data that is not required + if not needs_data: + self.data = None + if not needs_data_transpose: + self.data_t = None + + def _create_transpose(self): + """Create transposed quantized tensor""" + if not self.data.is_contiguous(): + self.data = self.data.contiguous() + self.data_t = self.data.t().contiguous() + self.scale_t = self.scale + + def size(self, *args, **kwargs): + """Get the size of the quantized tensor""" + if self.data is not None: + return self.data.size(*args, **kwargs) + size = self.data_t.size(*args, **kwargs) + return torch.Size([size[-1], math.prod(size[:-1])]) + + +def _scale_from_amax_tensor( + x_dtype: torch.dtype, + amax: torch.Tensor, + quant_dtype: torch.dtype, + *, + eps: float, + pow_2_scales: bool, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Derives quantization and dequantization from amax and options. + + Reference implementation for scale calculation. + + Returns: + - scale: quantization scales + - scale_inv: dequantization scales + - amax: Amax tensor with updates made for extrema values. + """ + assert amax.dtype == torch.float, "amax must be a float tensor." + fp8_max = torch.finfo(quant_dtype).max + + # Clamping amax to avoid division by small numbers + amax = torch.max(amax, torch.tensor(eps)) + + # Compute scale factor + scale = torch.div(fp8_max, amax) + + # Take care of inf before pow_2_scales + scale = torch.where(scale == torch.inf, torch.finfo(x_dtype).max, scale) + + if pow_2_scales: + _, exp = torch.frexp(scale) + exp = exp - 1 + assert (exp > -127).all() + unity = torch.tensor([1.0], device=exp.device) + torch.ldexp(unity, exp, out=scale) + scale = torch.where(amax == float("inf"), 0.0, scale) + + # Handle overflow cases for amax zero causing NaN + scale = torch.where(amax == 0, 1.0, scale) + + # Compute scale_inv + scale_inv = torch.reciprocal(scale) + + return scale, scale_inv, amax + + +class CurrentScalingQuantizerRef(Quantizer): + """Reference implementation of current scaling quantizer""" + + def __init__( + self, + dtype: torch.dtype, + rowwise: bool = True, + columnwise: bool = True, + pow_2_scales: bool = False, + eps: float = 0.0, + ): + super().__init__(rowwise=rowwise, columnwise=columnwise) + self.internal = True + + self.dtype = dtype + self.pow_2_scales = pow_2_scales + self.eps = eps + + self.with_amax_reduction = False + self.amax_reduction_group = None + + @property + def custom(self) -> bool: + """Flag to indicate this quantizer is custom.""" + return True + + @property + def supports_allgather_fp8(self) -> bool: + """Flag to indicate this quantizer supports allgather fp8""" + return True + + @classmethod + def compute_scale( + cls, + x: torch.Tensor, + quant_dtype: torch.dtype, + eps=0.0, + pow_2_scales: bool = False, + ): + """Compute the scale from the amax tensor""" + # Use float32 for computation + x_fp32 = x.to(torch.float32) + + if x_fp32.numel() == 0: + amax = torch.empty(1, dtype=torch.float32, device=x.device) + else: + amax = torch.amax(torch.abs(x_fp32)).view(1) + + return _scale_from_amax_tensor( + x.dtype, + amax=amax, + quant_dtype=quant_dtype, + eps=eps, + pow_2_scales=pow_2_scales, + ) + + def _quantize(self, tensor: torch.Tensor) -> Tuple[ + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + ]: + """ + Python implementation of quantization (c++ kernel can be used as an option instead). + + Parameters + ---------- + tensor : torch.Tensor + Input tensor to quantize (should be 2D) + + Returns + ------- + Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]] + (qx, sx, qx_t, sx_t) where: + - qx: quantized data in row-major order (if rowwise_usage), None otherwise + - sx: empty scale tensor for qx (if rowwise_usage), None otherwise + - qx_t: quantized data in column-major order (if columnwise_usage), None otherwise + - sx_t: empty scale tensor for qx_t (if columnwise_usage), None otherwise + """ + # Handle amax reduction if enabled + if self.with_amax_reduction: + assert ( + self.amax_reduction_group is not None + ), "amax_reduction_group must be set when with_amax_reduction is True" + + # Compute local amax + if tensor.numel() == 0: + amax = torch.empty(1, dtype=torch.float32, device=tensor.device) + else: + amax = torch.amax(torch.abs(tensor)).view(1).to(torch.float32) + + # Reduce amax across all ranks + torch.distributed.all_reduce( + amax, group=self.amax_reduction_group, op=torch.distributed.ReduceOp.MAX + ) + + # Compute scale using the global amax + scale, scale_inv, _ = _scale_from_amax_tensor( + tensor.dtype, + amax=amax, + quant_dtype=self.dtype, + eps=self.eps, + pow_2_scales=self.pow_2_scales, + ) + else: + # compute scale factor using local amax + scale, scale_inv, _ = self.compute_scale( + tensor, + self.dtype, + eps=self.eps, + pow_2_scales=self.pow_2_scales, + ) + + qx: Optional[torch.Tensor] = (tensor.float() * scale).to(self.dtype) + sx: Optional[torch.Tensor] = scale_inv + + # transpose if needed + if self.columnwise_usage: + assert qx is not None + qx_t = qx.t().contiguous() + sx_t = sx + else: + qx_t, sx_t = None, None + + if not self.rowwise_usage: + qx = None + sx = None + + return qx, sx, qx_t, sx_t + + def quantize( + self, + tensor: torch.Tensor, + **kwargs, # pylint: disable=unused-argument + ) -> CurrentScalingTensorRef: + # sanity checks + assert tensor.dtype in utils.HIGH_PRECISION_FLOAT_DTYPES, "Unsupported input dtype." + + # Make it work with 3D tensors + original_shape = tensor.shape + if tensor.ndim > 2: + tensor = tensor.view(-1, tensor.shape[-1]) + + qx, sx, qx_t, sx_t = self._quantize(tensor) + + return CurrentScalingTensorRef( + data=qx, + scale=sx, + data_t=qx_t, + scale_t=sx_t, + dtype=tensor.dtype, + device=tensor.device, + quant_dtype=self.dtype, + _quantizer=self, + original_shape=original_shape, + ) + + def dequantize( + self, tensor: torch.Tensor, scale: torch.Tensor, dtype: Optional[torch.dtype] = None + ) -> torch.Tensor: + """Dequantize the quantized tensor""" + tensor = tensor.to(torch.float32) * scale + if dtype is None: + return tensor + return tensor.to(dtype) + + def qgemm( + self, + qx: torch.Tensor, + qw: torch.Tensor, + m_params: quantization.MMParams, + out_dtype: torch.dtype, + sx: torch.Tensor, + sw: torch.Tensor, + bias: torch.Tensor | None = None, + out: torch.Tensor | None = None, + accumulate: bool = False, + gemm_type: quantization.GEMMType = quantization.GEMMType.FPROP, # pylint: disable=unused-argument + qresult_x: QuantizedTensorStorage | None = None, # pylint: disable=unused-argument + qresult_w: QuantizedTensorStorage | None = None, # pylint: disable=unused-argument + ) -> torch.Tensor: + """Python implementation of quantized gemm.""" + M, K = qx.shape + N, _ = qw.shape + + if M == 0 or K == 0 or N == 0: + if accumulate: + assert out is not None + y = out + else: + y = torch.zeros((M, N), dtype=out_dtype, device=qx.device) + if bias is not None: + y += bias + return y + + # cublas fp8 gemm does not support fp32 bias + use_bias_in_gemm = ( + bias is not None and out_dtype != torch.float32 and bias.dtype != torch.float32 + ) + + # Run quantized gemm: y = qw * qx + scaled_mm_res = torch._scaled_mm( + qx, + qw.transpose(-1, -2), + scale_a=sx, + scale_b=sw, + out_dtype=out_dtype, + use_fast_accum=not m_params.use_split_accumulator, + bias=bias if use_bias_in_gemm else None, + ) + y = scaled_mm_res[0] if isinstance(scaled_mm_res, tuple) else scaled_mm_res + + if bias is not None and not use_bias_in_gemm: + # Check number of elements in bias tensor because it can be an empty tensor + if bias.numel(): + y += bias + + if accumulate: + assert out is not None, "Output tensor must be provided for accumulation." + out.add_(y) + y = out + else: + assert out is None, "Output tensor should be None when accumulate is False." + + return y + + def transpose_qresult(self, qresult: CurrentScalingTensorRef) -> CurrentScalingTensorRef: + """Python implementation of transpose qresult.""" + qx = qresult.data + scale = qresult.scale + assert qresult.data_t is None + assert qresult.scale_t is None + assert qx is not None + qx_t = qx.transpose(-2, -1).contiguous() + scale_t = scale + qresult.data_t = qx_t + qresult.scale_t = scale_t + return qresult + + def update_quantized( + self, + src: torch.Tensor, + dst: QuantizedTensorStorage, + *, + noop_flag: Optional[torch.Tensor] = None, + ) -> QuantizedTensorStorage: + """Update the quantized tensor with the given tensor in-place + + Parameters + ---------- + src: torch.Tensor + Source tensor to copy from + dst: ExperimentalQuantizedTensor + Destination ExperimentalQuantizedTensor to update + noop_flag: torch.Tensor, optional + float32 flag indicating whether to avoid performing update + """ + # Handle noop flag + if noop_flag is not None and noop_flag.item() != 0: + return dst + + # Make sure input is in expected format + if not src.is_contiguous(): + src = src.contiguous() + + # Store the original shape and reshape for processing + original_shape = src.shape + if src.ndim > 2: + src = src.view(-1, src.shape[-1]) + + qx, sx, qx_t, sx_t = self._quantize(src) + + # Update the destination with new data + dst.data = qx + dst.scale = sx + dst.data_t = qx_t + dst.scale_t = sx_t + dst.dtype = src.dtype + dst.quant_dtype = self.dtype + dst.original_shape = original_shape + + return dst + + def make_empty( + self, + shape: Iterable[int], + *, + dtype: torch.dtype = torch.float32, + device: Optional[torch.device] = None, + requires_grad: bool = False, # pylint: disable=unused-argument + ) -> CurrentScalingTensorRef: + assert len(shape) == 2, "shape is not 2d" + + # Canonicalize tensor attributes + if device is None: + device = torch.device("cuda") + + # Allocate quantized data + qx = torch.empty(shape, dtype=self.dtype, device=device) + sx = torch.empty(1, dtype=torch.float32, device=device) + + # Allocate quantized data transpose if needed + qx_t = None + sx_t = None + if self.columnwise_usage: + inner_dim = qx.size(-1) + qx_t = torch.empty( + inner_dim, + qx.numel() // inner_dim, + dtype=self.dtype, + device=device, + ) + sx_t = torch.empty(1, dtype=torch.float32, device=device) + + # Construct quantized tensor + return CurrentScalingTensorRef( + data=qx, + scale=sx, + data_t=qx_t, + scale_t=sx_t, + dtype=dtype, + device=device, + quant_dtype=self.dtype, + _quantizer=self, + original_shape=shape, + ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py index 1ce9079eb1..b371ca4842 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py @@ -18,9 +18,9 @@ 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 fp8_autocast: + Usage with CustomRecipe and autocast: custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) - with fp8_autocast(fp8_recipe=custom_recipe): + with autocast(fp8_recipe=custom_recipe): output = model(input) """ if role == "linear_input": @@ -338,7 +338,7 @@ def get_wgrad_sign_vector() -> torch.Tensor: class NVFP4QuantizerRef(Quantizer): - """NVFP4 quantizer for middleware between Transformer Engine and Kitchen""" + """Reference implementation of NVFP4 quantizer""" def __init__( self, From 7e593c3be96b3eebc384da1a2ab307727065c9ab Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 17 Nov 2025 13:26:52 -0800 Subject: [PATCH 073/521] Add num_splits support for FA3 backend (#2380) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [Common] Deleted unused header (#2324) Deleted unused header Signed-off-by: Oleg Goncharov Signed-off-by: Peter Dykas * [JAX] L1_jax_distributed_test suit with individual executions (#2321) * L1 rework Signed-off-by: Phuong Nguyen * comment out test_multi_process_grouped_gemm for now Signed-off-by: Phuong Nguyen * rm e5m2 from test norm + MXFP8 Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen Signed-off-by: Peter Dykas * for branch Signed-off-by: Peter Dykas * clean up and tests Signed-off-by: Peter Dykas * change tests Signed-off-by: Peter Dykas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Peter Dykas * [PyTorch debug] Fixes to debug tests failures (#2268) * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix: Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix 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> Signed-off-by: Peter Dykas * [PyTorch Debug] Add max_blockwise_dynamic_range stats (#2137) * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix 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> Signed-off-by: Peter Dykas * [JAX] Fix bug with pre scale bias (#2300) * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Peter Dykas * [JAX] Try to use pre-downloaded dataset artifacts first (#2345) * Try to use pre-downloaded dataset artifacts first Signed-off-by: Jeremy Berchtold * Set HF_HUB_OFFLINE to disable any network calls to HF when the pre-downloaded dataset is available Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold Signed-off-by: Peter Dykas * Fix out of bounds access in the FP4 dequantize kernel (#2346) Signed-off-by: Przemek Tredak Signed-off-by: Peter Dykas * Make FP8 weights compatible with older MCore version (#2342) * Make cast_master_weights_to_fp8 compatible with older MCore version Signed-off-by: kunlunl * Rename keep_columnwise to manual_post_all_gather_processing & Optimize unit test Signed-off-by: kunlunl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove redundant _test_mini_optimizer() Signed-off-by: kunlunl --------- Signed-off-by: kunlunl 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> Signed-off-by: Peter Dykas * [JAX] Add test to check jaxpr that amax is reused for nvfp4 recipe (#2348) * Add test to check jaxpr that amax is reused for nvfp4 recipe Signed-off-by: Jeremy Berchtold * Move test to test_helper.py and rename file Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Peter Dykas * Fix sharding of segment position to match id in ring attention. (#2349) Signed-off-by: Peter Dykas * Disable cuDNN attention for known IMA and NaNs (#2344) * Fix cuDNN backend selection for more case. Add CG as a option as well Signed-off-by: Kirthi Shankar Sivamani * fix logic Signed-off-by: Kirthi Shankar Sivamani * Fix cuDNN checks Signed-off-by: Kirthi Shankar Sivamani * Add more checks Signed-off-by: Kirthi Shankar Sivamani * Fix cuddn version Signed-off-by: Kirthi Shankar Sivamani * Fix error message Signed-off-by: Kirthi Shankar Sivamani * Add check for window size Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Peter Dykas * [JAX] Default to fused attention in JAX DPA (#2363) * Default to fused attention in JAX DPA Signed-off-by: Kshitij Lakhani * Consolidate documentation for DPA in JAX Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> * Correctly update the documentation for defaults in JAX DPA Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> --------- Signed-off-by: Kshitij Lakhani Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Peter Dykas * Update cudnn frontend to v1.16.0 (#2362) Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Peter Dykas * [common] Remove kvpacked and qkvpacked attention functions for every kernel type. (#2287) * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * depracted compile time warning + \warning -> \deprecated Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski 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> Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: Peter Dykas * Move Triton to common (#2359) * move triton to common and change paths Signed-off-by: tdophung * Formatting Signed-off-by: tdophung --------- Signed-off-by: tdophung Signed-off-by: Peter Dykas * [JAX] Fused layers argument default values changed (#2347) * Changing default activations in MLP, TransformerLayer, dropout rate after FC1 to 0, and return_layernorm_output to False Signed-off-by: tdophung * Fixing the failing tests by hard coding arguments to the previous values instead of relying on newer default values Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Signed-off-by: Peter Dykas * remove comment from gpt Signed-off-by: Peter Dykas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor changes for num_splits logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace None with 1 as default 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 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix dtype in pack/unpack when FP8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add fused_attn_supported constraint for some tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FA3 installation commands Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FA3 installation commands in DPA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * separate fused fp8 and f16 flags in tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * initialize fused_attn_supported_f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix FA installation in L3 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Oleg Goncharov Signed-off-by: Peter Dykas Signed-off-by: Phuong Nguyen Signed-off-by: Pawel Gadzinski Signed-off-by: Jeremy Berchtold Signed-off-by: Przemek Tredak Signed-off-by: kunlunl Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Kshitij Lakhani Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Signed-off-by: tdophung Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Co-authored-by: Phuong Nguyen Co-authored-by: root Co-authored-by: Peter Dykas Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak Co-authored-by: Kunlun Li <94586211+kunlunl@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Michael Goldfarb Co-authored-by: Kirthi Shankar Sivamani Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Teddy Do Co-authored-by: wdykas <73254672+wdykas@users.noreply.github.com> --- qa/L3_pytorch_FA_versions_test/test.sh | 6 +- tests/pytorch/attention/test_attention.py | 125 ++++++++++++------ tests/pytorch/utils.py | 10 ++ .../dot_product_attention/backends.py | 2 + .../dot_product_attention.py | 7 + .../attention/dot_product_attention/utils.py | 24 +++- 6 files changed, 126 insertions(+), 48 deletions(-) diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 418e824c10..e2d771cfd0 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -30,13 +30,13 @@ do # Build Flash Attention if [ "${fa_version}" \< "3.0.0" ] then - pip3 install flash-attn==${fa_version} + pip3 install flash-attn==${fa_version} --no-build-isolation else git clone https://github.com/Dao-AILab/flash-attention.git - cd flash-attention/ && git checkout 27f501d && cd hopper/ && python setup.py install + cd flash-attention/hopper && python setup.py install python_path=`python -c "import site; print(site.getsitepackages()[0])"` mkdir -p $python_path/flash_attn_3 - wget -P $python_path/flash_attn_3 https://raw.githubusercontent.com/Dao-AILab/flash-attention/27f501dbe011f4371bff938fe7e09311ab3002fa/hopper/flash_attn_interface.py + cp flash_attn_interface.py $python_path/flash_attn_3/ cd ../../ fi diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 4f4cad97db..4aedcff1b8 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -117,7 +117,14 @@ def reset_global_fp8_state(): @pytest.mark.parametrize("swa", [False]) @pytest.mark.parametrize("pad_between_seqs", [False]) def test_dot_product_attention( - dtype, model_configs, model, ckpt_attn, workspace_opt, qkv_layout, swa, pad_between_seqs + dtype, + model_configs, + model, + ckpt_attn, + workspace_opt, + qkv_layout, + swa, + pad_between_seqs, ): """Test DotProductAttention module""" @@ -308,6 +315,31 @@ def test_dpa_max_logit(dtype, model_configs, model, qkv_layout): test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) +model_configs_num_splits = { + # test: ModelConfig(b, sq, hq, dqk) + "num_splits_1_0": ModelConfig(2, 2048, 24, 128, num_splits=2), + "num_splits_1_1": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096, num_splits=4), +} + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_num_splits]) +@pytest.mark.parametrize("model", model_configs_num_splits.keys()) +def test_dpa_num_splits(dtype, model_configs, model): + """Test DotProductAttention with FlashAttention-3 num_splits enabled""" + test_dot_product_attention( + dtype, + model_configs, + model, + False, + True, + None, + False, + False, + ) + + model_configs_softmax = { # test: ModelConfig(b, sq, hq, dqk) "softmax_1_0": ModelConfig(2, 2048, 64, 64, num_gqa_groups=8), @@ -1152,6 +1184,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: core_attention_bias=bias, alibi_slopes=alibi_slopes, fast_zero_fill=True, + # Only pass num_splits when exercising the FlashAttention path + num_splits=config.num_splits if backend == "FlashAttention" else 1, ) max_logit = None if config.return_max_logit: @@ -1786,9 +1820,10 @@ def test_mha_fp8_vs_f16( fp8_meta=fp8_meta, is_training=is_training, ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if flash_attn_supported + fused_attn_supported < 1: + flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends + if flash_attn_supported + fused_attn_supported_fp8 < 1: pytest.skip("No FP8 attention backend available.") + fused_attn_supported_f16 = False if not fp8_dpa_bwd: available_backends, _, fused_attn_backends = get_available_attention_backends( config, @@ -1796,8 +1831,8 @@ def test_mha_fp8_vs_f16( qkv_layout=qkv_format.replace("hd", "h3d"), is_training=is_training, ) - _, fused_attn_supported, _ = available_backends - if not fused_attn_supported: + _, fused_attn_supported_f16, _ = available_backends + if not fused_attn_supported_f16: pytest.skip("No attention backend available.") if flash_attn_supported: @@ -1809,23 +1844,28 @@ def test_mha_fp8_vs_f16( dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe ) - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") - fused_attn_fwd_fp8, param_names, fused_attn_bwd_fp8 = _run_mha_fp8_vs_f16( - dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe - ) + if fused_attn_supported_fp8: + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") + fused_attn_fwd_fp8, param_names, fused_attn_bwd_fp8 = _run_mha_fp8_vs_f16( + dtype, config, True, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe + ) - logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = False") - fused_attn_fwd_f16, param_names, fused_attn_bwd_f16 = _run_mha_fp8_vs_f16( - dtype, config, False, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe - ) + if fused_attn_supported_f16: + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = False") + fused_attn_fwd_f16, param_names, fused_attn_bwd_f16 = _run_mha_fp8_vs_f16( + dtype, config, False, qkv_format, input_layernorm, RoPE, is_training, fp8_recipe + ) atol = 5e-1 rtol = 5e-1 rmse_tol = 0.15 - if flash_attn_supported: + if flash_attn_supported and fused_attn_supported_f16: logging.debug("========== {:^25s} ==========".format("flash fp8 vs fused f16:")) logging.debug("========== {:^25s} ==========".format("forward output")) compare_and_assert( @@ -1838,32 +1878,33 @@ def test_mha_fp8_vs_f16( rmse_tol, True, ) - logging.debug("========== {:^25s} ==========".format("fused fp8 vs fused f16:")) - logging.debug("========== {:^25s} ==========".format("forward output")) - compare_and_assert( - fused_attn_fwd_fp8, - fused_attn_fwd_f16, - "fused_attn_fwd_fp8", - "fused_attn_fwd_f16", - atol, - rtol, - rmse_tol, - True, - ) + if fused_attn_supported_fp8 and fused_attn_supported_f16: + logging.debug("========== {:^25s} ==========".format("fused fp8 vs fused f16:")) + logging.debug("========== {:^25s} ==========".format("forward output")) + compare_and_assert( + fused_attn_fwd_fp8, + fused_attn_fwd_f16, + "fused_attn_fwd_fp8", + "fused_attn_fwd_f16", + atol, + rtol, + rmse_tol, + True, + ) - if is_training: - for i in range(len(param_names[:1])): - logging.debug("========== {:^25s} ==========".format(param_names[i])) - compare_and_assert( - fused_attn_bwd_fp8[i], - fused_attn_bwd_f16[i], - f"fused_attn_bwd_fp8[{i}]", - f"fused_attn_bwd_f16[{i}]", - atol, - rtol, - rmse_tol, - True, - ) + if is_training: + for i in range(len(param_names[:1])): + logging.debug("========== {:^25s} ==========".format(param_names[i])) + compare_and_assert( + fused_attn_bwd_fp8[i], + fused_attn_bwd_f16[i], + f"fused_attn_bwd_fp8[{i}]", + f"fused_attn_bwd_f16[{i}]", + atol, + rtol, + rmse_tol, + True, + ) def _run_mha_fp8_vs_f16( diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 485c739c03..bdf469c59a 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -8,6 +8,7 @@ import os from contextlib import contextmanager from typing import Optional, Tuple, Dict, Any, List +from packaging.version import Version as PkgVersion import torch @@ -210,6 +211,7 @@ def __init__( max_ctx_len: int = None, num_layers: int = 1, eps: float = 1e-5, + num_splits=1, ): self.batch_size = batch_size self.max_seqlen_q = max_seqlen_q @@ -239,6 +241,7 @@ def __init__( self.max_ctx_len = max_ctx_len self.num_layers = num_layers self.eps = eps + self.num_splits = num_splits @contextmanager @@ -321,6 +324,9 @@ def test(): inference_params=inference_params, softmax_type=config.softmax_type, return_max_logit=config.return_max_logit, + # allow all backends to pass so they can be used for testing; + # check for FA3 availability later + num_splits=1, ) ( use_flash_attention, @@ -330,6 +336,10 @@ def test(): use_unfused_attention, available_backends, ) = get_attention_backend(attention_params) + # Check if FA3 is an available backend when num_splits != 1 + if available_backends[0]: + if config.num_splits != 1 and not flash_attention_backend > PkgVersion("3.0.0b"): + available_backends[0] = False # Set attention.py _attention_backends var using return value # from get_attention_backend() _attention_backends["use_flash_attention"] = use_flash_attention diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 1480f900fd..c1ff46c75a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -681,6 +681,7 @@ def forward( inference_params: Optional[InferenceParams] = None, flash_attention_backend: Optional[PkgVersion] = PkgVersion("0"), fp8_output: bool = False, + num_splits: Optional[int] = 1, ) -> torch.Tensor: """flash-attn fprop""" @@ -957,6 +958,7 @@ def forward( else: fa_3_optional_forward_kwargs = {} fa_3_optional_forward_kwargs["window_size"] = window_size + fa_3_optional_forward_kwargs["num_splits"] = num_splits if inference_params is None: fa_3_optional_forward_kwargs["deterministic"] = self.deterministic else: 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 83528330eb..1f60ae020e 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 @@ -799,6 +799,7 @@ def forward( inference_params: Optional[InferenceParams] = None, pad_between_seqs: Optional[bool] = None, fp8_output: Optional[bool] = False, + num_splits: Optional[int] = 1, ) -> torch.Tensor: """ Dot Product Attention Layer. @@ -973,6 +974,10 @@ def forward( If true, there are padding tokens between individual sequences in a packed batch. fp8_output: Optional[bool], default = `False` Whether to enforce output to be in FP8 or not. + num_splits: Optional[int], default = 1 + Optional split control for FlashAttention-3 only. When set, this value is forwarded + to the FA3 backend to control internal kernel splitting behavior for non-context-parallel + cases. It is ignored for other backends and when context parallelism is enabled. """ with self.prepare_forward( @@ -1315,6 +1320,7 @@ def forward( softmax_type=self.softmax_type, return_max_logit=self.return_max_logit, cuda_graph=is_graph_capturing(), + num_splits=num_splits, ) global _attention_backends if is_in_onnx_export_mode(): @@ -1413,6 +1419,7 @@ def forward( inference_params=inference_params, flash_attention_backend=flash_attention_backend, fp8_output=fp8_output, + num_splits=num_splits, ) 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 a08ba14196..7a61c60094 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -135,7 +135,7 @@ class FlashAttentionUtils: # Please follow these instructions to install FA3 v3_installation_steps = """\ (1) git clone https://github.com/Dao-AILab/flash-attention.git -(2) cd flash-attention/ && git checkout 3ba6f82 && git submodule update --init && cd hopper/ && python setup.py install +(2) cd flash-attention/hopper && python setup.py install (3) python_path=`python -c "import site; print(site.getsitepackages()[0])"` (4) mkdir -p $python_path/flash_attn_3 (5) cp flash_attn_interface.py $python_path/flash_attn_3/flash_attn_interface.py""" @@ -233,6 +233,8 @@ class AttentionParams: Whether to output max_logit. cuda_graph: bool, default = `False` Whether support for cuda graph capture is needed or not. + num_splits: int, default = 1 + The number of kernels to split attention to. """ qkv_type: Union[torch.Tensor, Float8Tensor] = torch.Tensor @@ -263,6 +265,7 @@ class AttentionParams: softmax_type: str = "vanilla" return_max_logit: bool = False cuda_graph: bool = False + num_splits: int = 1 def __eq__(self, other): """ @@ -338,6 +341,7 @@ def get_attention_backend( softmax_type = attention_params.softmax_type return_max_logit = attention_params.return_max_logit cuda_graph = attention_params.cuda_graph + num_splits = attention_params.num_splits # Run config logger = logging.getLogger("DotProductAttention") @@ -511,6 +515,18 @@ def get_attention_backend( use_flash_attention = False use_fused_attention = False + # Filter: num_splits + if num_splits != 1: + if use_flash_attention_2 and FlashAttentionUtils.is_installed: + logger.debug("Disabling FlashAttention 2 for num_splits") + use_flash_attention_2 = False + if use_fused_attention: + logger.debug("Disabling FusedAttention for num_splits") + use_fused_attention = False + if use_unfused_attention: + logger.debug("Disabling UnfusedDotProductAttention for num_splits") + use_unfused_attention = False + # Filter: Return max_logit if return_max_logit: if use_flash_attention: @@ -1566,8 +1582,9 @@ def _pack_tensor( """ Packs the given tensor using the `indices`. """ + dtype = tensor.dtype if not isinstance(tensor, Float8Tensor) else torch.uint8 padding_indice = torch.zeros( - 1, tensor.shape[1], tensor.shape[2], dtype=tensor.dtype, device=tensor.device + 1, tensor.shape[1], tensor.shape[2], dtype=dtype, device=tensor.device ) indices = indices.repeat(1, tensor.shape[1], tensor.shape[2]) if isinstance(tensor, Float8Tensor): @@ -1622,8 +1639,9 @@ def _unpack_tensor( Inverse of `_pack_tensor`. """ indices = indices.repeat(1, tensor.shape[1], tensor.shape[2]) + dtype = tensor.dtype if not isinstance(tensor, Float8Tensor) else torch.uint8 unpacked = torch.zeros( - dim0 + 1, tensor.shape[1], tensor.shape[2], dtype=tensor.dtype, device=tensor.device + dim0 + 1, tensor.shape[1], tensor.shape[2], dtype=dtype, device=tensor.device ) if isinstance(tensor, Float8Tensor): unpacked.scatter_(0, indices, tensor._data) From 15cefbc5555bf0cadb5f2aefe558143845e43e19 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 18 Nov 2025 15:04:28 +0100 Subject: [PATCH 074/521] [JAX] Add support for sink attention in JAX (#2225) * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * removed packed versions Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * jax Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix: Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * sofmtax_fusion -> softmax_fusion_type Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix 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> --- tests/jax/test_distributed_fused_attn.py | 56 +++- tests/jax/test_distributed_softmax.py | 45 ++- tests/jax/test_fused_attn.py | 95 +++++- tests/jax/test_layer.py | 18 ++ tests/jax/test_softmax.py | 51 +-- tests/jax/utils.py | 47 +++ transformer_engine/jax/attention.py | 73 ++++- .../jax/cpp_extensions/attention.py | 292 ++++++++++++++---- .../jax/cpp_extensions/softmax.py | 74 ++++- transformer_engine/jax/csrc/extensions.h | 24 +- .../jax/csrc/extensions/attention.cpp | 177 ++++++----- .../jax/csrc/extensions/pybind.cpp | 5 + transformer_engine/jax/flax/module.py | 56 +++- transformer_engine/jax/flax/transformer.py | 109 ++++++- transformer_engine/jax/softmax.py | 24 +- .../dot_product_attention/softmax.py | 6 +- 16 files changed, 909 insertions(+), 243 deletions(-) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index ef8e370b6e..5372018ae8 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -18,6 +18,7 @@ is_fused_attn_kernel_available, AttnBiasType, AttnMaskType, + AttnSoftmaxType, QKVLayout, QKVFormat, reorder_causal_load_balancing, @@ -66,6 +67,7 @@ def impl_test_self_attn( bias_shape, attn_mask_type, dtype, + softmax_type, use_shardy, ): jax.config.update("jax_use_shardy_partitioner", use_shardy) @@ -80,6 +82,7 @@ def impl_test_self_attn( QKVLayout.BS3HD, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, num_head, num_head, @@ -109,6 +112,7 @@ def impl_test_self_attn( hidden, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, is_training, @@ -142,6 +146,14 @@ def impl_test_self_attn( ], ) @pytest.mark.parametrize("dtype", DTYPES) + @pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), + pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), + ], + ) def test_self_attn( self, device_count, @@ -153,6 +165,7 @@ def test_self_attn( bias_shape, attn_mask_type, dtype, + softmax_type, ): self.impl_test_self_attn( device_count, @@ -164,6 +177,7 @@ def test_self_attn( bias_shape, attn_mask_type, dtype, + softmax_type, use_shardy=False, ) @@ -175,8 +189,23 @@ def test_self_attn( pytest.param(AttnBiasType.PRE_SCALE_BIAS, BiasShape._1HSS, id="PRE_SCALE_BIAS-1HSS"), ], ) + @pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), + pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), + ], + ) def test_self_attn_shardy( - self, device_count, mesh_shape, mesh_axes, mesh_resource, attn_bias_type, bias_shape + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + attn_bias_type, + bias_shape, + softmax_type, ): data_shape = (32, 512, 12, 64) self.impl_test_self_attn( @@ -189,6 +218,7 @@ def test_self_attn_shardy( bias_shape, AttnMaskType.PADDING_MASK, jnp.bfloat16, + softmax_type, use_shardy=True, ) @@ -213,8 +243,24 @@ def generate_collectives_count_ref(self): "attn_mask_type", [AttnMaskType.PADDING_MASK, AttnMaskType.CAUSAL_MASK] ) @pytest.mark.parametrize("dtype", DTYPES) + @pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), + pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), + ], + ) def test_cross_attn( - self, device_count, mesh_shape, mesh_axes, mesh_resource, data_shape, attn_mask_type, dtype + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + attn_mask_type, + dtype, + softmax_type, ): attn_bias_type = AttnBiasType.NO_BIAS bias_shape = None @@ -230,6 +276,7 @@ def test_cross_attn( QKVLayout.BSHD_BS2HD, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, num_head, num_head, @@ -252,6 +299,7 @@ def test_cross_attn( hidden, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, is_training, @@ -322,6 +370,8 @@ def impl_test_context_parallel_attn( bias_shape = None dropout_prob = 0.0 is_training = True + # Context parallel does not support softmax_offset + softmax_type = AttnSoftmaxType.VANILLA_SOFTMAX dp_size, cp_size, tp_size = mesh_shape batch, seqlen, num_head, hidden = data_shape @@ -343,6 +393,7 @@ def impl_test_context_parallel_attn( hidden, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, is_training, @@ -366,6 +417,7 @@ def check_has_backend_for_mask(mask_type): qkv_layout, attn_bias_type, mask_type, + softmax_type, dropout_prob, num_head, num_kv_heads, diff --git a/tests/jax/test_distributed_softmax.py b/tests/jax/test_distributed_softmax.py index f1ae6c9e49..8cdd4c3f59 100644 --- a/tests/jax/test_distributed_softmax.py +++ b/tests/jax/test_distributed_softmax.py @@ -16,7 +16,7 @@ from distributed_test_base import compare_ops from utils import make_causal_mask, make_self_mask from transformer_engine.jax import autocast -from transformer_engine.jax.softmax import SoftmaxType, softmax +from transformer_engine.jax.softmax import SoftmaxFusionType, softmax DTYPES = [jnp.float16, jnp.bfloat16] @@ -29,12 +29,12 @@ def generate_collectives_count_ref(self): return generate_collectives_count(allreduce=all_reduce_loss_bytes, allgather=0, other=0) def generate_inputs( - self, shape, mesh_resource, softmax_type, dtype, bad_sharding, broadcast_batch_mask + self, shape, mesh_resource, softmax_fusion_type, dtype, bad_sharding, broadcast_batch_mask ): batch, _, sqelen, _ = shape x = random.normal(random.PRNGKey(1124), shape, dtype=dtype) - if softmax_type == SoftmaxType.SCALED_UPPER_TRIANG_MASKED: + if softmax_fusion_type == SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED: mask = make_causal_mask(batch, sqelen) else: mask = make_self_mask(1 if broadcast_batch_mask else batch, sqelen) @@ -56,8 +56,10 @@ def generate_inputs( return (x, mask), (x_pspec, mask_pspec) @staticmethod - def target_func(x, mask, scale_factor=1.0, softmax_type=SoftmaxType.SCALED): - return jnp.mean(softmax(x, mask, scale_factor=scale_factor, softmax_type=softmax_type)) + def target_func(x, mask, scale_factor=1.0, softmax_fusion_type=SoftmaxFusionType.SCALED): + return jnp.mean( + softmax(x, mask, scale_factor=scale_factor, softmax_fusion_type=softmax_fusion_type) + ) @staticmethod def ref_func(x, mask, scale_factor=1.0, dtype=jnp.float16): @@ -80,24 +82,29 @@ def impl_test_softmax( mesh_axes, mesh_resource, data_shape, - softmax_type, + softmax_fusion_type, scale_factor, dtype, bad_sharding, broadcast_batch_mask, use_shardy, ): - if broadcast_batch_mask and softmax_type != SoftmaxType.SCALED_MASKED: + if broadcast_batch_mask and softmax_fusion_type != SoftmaxFusionType.SCALED_MASKED: pytest.skip("Softmax type has no mask.") jax.config.update("jax_use_shardy_partitioner", use_shardy) target_func = partial( - self.target_func, scale_factor=scale_factor, softmax_type=softmax_type + self.target_func, scale_factor=scale_factor, softmax_fusion_type=softmax_fusion_type ) ref_func = partial(self.ref_func, scale_factor=scale_factor, dtype=dtype) (x, mask), (x_pspec, mask_pspec) = self.generate_inputs( - data_shape, mesh_resource, softmax_type, dtype, bad_sharding, broadcast_batch_mask + data_shape, + mesh_resource, + softmax_fusion_type, + dtype, + bad_sharding, + broadcast_batch_mask, ) collective_count_ref = self.generate_collectives_count_ref() devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) @@ -139,8 +146,12 @@ def impl_test_softmax( @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) @pytest.mark.parametrize("data_shape", [[32, 12, 128, 128], [8, 8, 1024, 1024]]) @pytest.mark.parametrize( - "softmax_type", - [SoftmaxType.SCALED, SoftmaxType.SCALED_MASKED, SoftmaxType.SCALED_UPPER_TRIANG_MASKED], + "softmax_fusion_type", + [ + SoftmaxFusionType.SCALED, + SoftmaxFusionType.SCALED_MASKED, + SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED, + ], ) @pytest.mark.parametrize("scale_factor", [1.0, 3.0]) @pytest.mark.parametrize("dtype", DTYPES) @@ -153,7 +164,7 @@ def test_softmax( mesh_axes, mesh_resource, data_shape, - softmax_type, + softmax_fusion_type, scale_factor, dtype, bad_sharding, @@ -165,7 +176,7 @@ def test_softmax( mesh_axes, mesh_resource, data_shape, - softmax_type, + softmax_fusion_type, scale_factor, dtype, bad_sharding, @@ -174,7 +185,9 @@ def test_softmax( ) @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) - @pytest.mark.parametrize("softmax_type", [SoftmaxType.SCALED, SoftmaxType.SCALED_MASKED]) + @pytest.mark.parametrize( + "softmax_fusion_type", [SoftmaxFusionType.SCALED, SoftmaxFusionType.SCALED_MASKED] + ) @pytest.mark.parametrize("bad_sharding", [False, True]) @pytest.mark.parametrize("broadcast_batch_mask", [False, True]) def test_softmax_gspmd( @@ -183,7 +196,7 @@ def test_softmax_gspmd( mesh_shape, mesh_axes, mesh_resource, - softmax_type, + softmax_fusion_type, bad_sharding, broadcast_batch_mask, ): @@ -193,7 +206,7 @@ def test_softmax_gspmd( mesh_axes, mesh_resource, data_shape=[32, 12, 128, 128], - softmax_type=softmax_type, + softmax_fusion_type=softmax_fusion_type, scale_factor=1.0, dtype=DTYPES[0], bad_sharding=bad_sharding, diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index a5d73d9605..f4caaef165 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -27,6 +27,7 @@ from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, + AttnSoftmaxType, QKVLayout, QKVFormat, reorder_causal_load_balancing, @@ -59,14 +60,16 @@ def init(): yield -@partial(jax.jit, static_argnums=(5, 6, 7, 9)) +@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11)) def general_dot_product_attention( query: ArrayLike, key: ArrayLike, value: ArrayLike, + softmax_offset: Optional[ArrayLike], bias: ArrayLike, mask: ArrayLike, deterministic: bool, + softmax_type: AttnSoftmaxType, scale_factor: float, dropout_rate: float, dropout_rng: ArrayLike, @@ -99,7 +102,25 @@ def general_dot_product_attention( mask = jnp.expand_dims(mask, axis=-3) logits = jnp.where(mask, jnp.finfo(dtype).min, logits) - softmax_out = jax.nn.softmax(logits).astype(dtype) + match softmax_type: + case AttnSoftmaxType.VANILLA_SOFTMAX: + softmax_out = jax.nn.softmax(logits).astype(dtype) + case AttnSoftmaxType.OFF_BY_ONE_SOFTMAX: + # Softmax with +1 in denominator: exp(x_i) / (sum(exp(x_j)) + 1) + # Append a zero logit, apply standard softmax, then remove last column + zero_logit = jnp.zeros(logits.shape[:-1] + (1,), dtype=logits.dtype) + logits_with_extra = jnp.concatenate([logits, zero_logit], axis=-1) + softmax_with_extra = jax.nn.softmax(logits_with_extra, axis=-1) + softmax_out = softmax_with_extra[..., :-1].astype(dtype) + case AttnSoftmaxType.LEARNABLE_SOFTMAX: + # Append learnable offset logit, apply standard softmax, then remove last column + learnable_logit = softmax_offset.reshape(1, h_kv, num_groups, 1, 1) + learnable_logit = jnp.broadcast_to(learnable_logit, logits.shape[:-1] + (1,)) + logits_with_extra = jnp.concatenate([logits, learnable_logit], axis=-1) + softmax_with_extra = jax.nn.softmax(logits_with_extra, axis=-1) + softmax_out = softmax_with_extra[..., :-1].astype(dtype) + case _: + raise NotImplementedError(f"Unknown {softmax_type=}") if not deterministic and dropout_rate > 0.0: keep_prob = 1.0 - dropout_rate @@ -238,7 +259,7 @@ def _split_valid_and_invalid(primitive, reference, pad): return primitive_valid, primitive_invalid, reference_valid, reference_invalid -def jax_dpa(query, key, value, bias, mask, dropout_rng, **kwargs): +def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs): """ JAX native dot product attention implementation """ @@ -246,11 +267,13 @@ def jax_dpa(query, key, value, bias, mask, dropout_rng, **kwargs): query, key, value, + softmax_offset, bias, mask, deterministic=not kwargs["is_training"], scale_factor=kwargs["scaling_factor"], dropout_rate=kwargs["dropout_probability"], + softmax_type=kwargs["softmax_type"], dropout_rng=dropout_rng, dtype=jnp.float32, ) @@ -262,6 +285,7 @@ def customcall_fused_dpa( key, value, bias, + softmax_offset, sequence_descriptor, dropout_rng, **kwargs, @@ -283,9 +307,9 @@ def customcall_fused_dpa( qkv_args = (query, key, value) case _: raise ValueError(f"Unsupported {qkv_layout=}") - return fused_attn(qkv_args, bias, sequence_descriptor, dropout_rng, **kwargs).astype( - query.dtype - ) + return fused_attn( + qkv_args, bias, sequence_descriptor, dropout_rng, softmax_offset=softmax_offset, **kwargs + ).astype(query.dtype) class BiasShape(Enum): @@ -320,6 +344,7 @@ class FusedAttnRunner: head_dim_v: int attn_bias_type: AttnBiasType attn_mask_type: AttnMaskType + softmax_type: AttnSoftmaxType dropout_prob: float dtype: DTypeLike is_training: bool @@ -402,6 +427,7 @@ def _check_configs(self): self.qkv_layout, self.attn_bias_type, self.attn_mask_type, + self.softmax_type, self.dropout_prob, self.num_heads_q, self.num_heads_kv, @@ -439,7 +465,7 @@ def _setup_inputs(self): self.tp_size = self.mesh.shape.get(self.mesh_resource.tpsp_resource, 1) key = jax.random.PRNGKey(0) - q_key, k_key, v_key, bias_key, dropout_key = jax.random.split(key, 5) + q_key, k_key, v_key, bias_key, dropout_key, softmax_key = jax.random.split(key, 6) q_shape = (self.batch_size, self.max_seqlen_q, self.num_heads_q, self.head_dim_qk) k_shape = (self.batch_size, self.max_seqlen_kv, self.num_heads_kv, self.head_dim_qk) @@ -490,6 +516,13 @@ def _setup_inputs(self): else: pad_ratio = 0.0 + if self.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX: + self.softmax_offset = jax.random.uniform( + softmax_key, (1, self.num_heads_q, 1, 1), jnp.float32, -1.0 + ) + else: + self.softmax_offset = None + def gen_valid(bs, max_seqlen, pad_ratio): pad_len = int(max_seqlen * pad_ratio) valid_len = max_seqlen - pad_len @@ -713,6 +746,16 @@ def to_dp_shardings(x): self.bias_pspec = PartitionSpec() self.bias_sharding = NamedSharding(self.mesh, self.bias_pspec) + # Softmax offset sharding (1, num_heads, 1, 1) + # Use the same logic as HEAD_AXES: tpsp_resource if enabled, else tp_resource + head_resource = ( + self.mesh_resource.tpsp_resource + if self.mesh_resource.tpsp_resource is not None + else self.mesh_resource.tp_resource + ) + self.softmax_offset_pspec = PartitionSpec(None, head_resource, None, None) + self.softmax_offset_sharding = NamedSharding(self.mesh, self.softmax_offset_pspec) + self.dropout_rng_pspec = PartitionSpec( None, ) @@ -732,7 +775,7 @@ def test_forward(self): """ self._setup_inputs() - args = [self.q, self.k, self.v, self.bias, self.mask, self.dropout_rng] + args = [self.q, self.k, self.v, self.bias, self.softmax_offset, self.mask, self.dropout_rng] customcall_args = [ # Put test data onto each GPU for distributed. @@ -742,12 +785,14 @@ def test_forward(self): jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding), jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding), jax.device_put(self.bias, self.bias_sharding), + jax.device_put(self.softmax_offset, self.softmax_offset_sharding), jax.device_put(self.sequence_desciptor, self.seq_desc_sharding), jax.device_put(self.dropout_rng, self.dropout_rng_sharding), ] kwargs = { "attn_bias_type": self.attn_bias_type, "attn_mask_type": self.attn_mask_type, + "softmax_type": self.softmax_type, "scaling_factor": self.scaling_factor, "dropout_probability": self.dropout_prob, "is_training": self.is_training, @@ -766,6 +811,7 @@ def test_forward(self): self.qkvo_sharding, self.qkvo_sharding, self.bias_sharding, + self.softmax_offset_sharding, self.seq_desc_sharding, self.dropout_rng_sharding, ], @@ -826,7 +872,7 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): jnp.mean(ret_valid.astype(jnp.float32), dtype=jnp.float32) * gradient_multiplier ).astype(self.dtype) - args = [self.q, self.k, self.v, self.bias, self.mask, self.dropout_rng] + args = [self.q, self.k, self.v, self.bias, self.softmax_offset, self.mask, self.dropout_rng] customcall_args = [ # TODO(mgoldfarb-nvidia): We will need to add reordering for bias, mas and # THD params once we support those features on CP. @@ -834,12 +880,14 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding), jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding), jax.device_put(self.bias, self.bias_sharding), + jax.device_put(self.softmax_offset, self.softmax_offset_sharding), jax.device_put(self.sequence_desciptor, self.seq_desc_sharding), jax.device_put(self.dropout_rng, self.dropout_rng_sharding), ] kwargs = { "attn_bias_type": self.attn_bias_type, "attn_mask_type": self.attn_mask_type, + "softmax_type": self.softmax_type, "scaling_factor": self.scaling_factor, "dropout_probability": self.dropout_prob, "is_training": self.is_training, @@ -866,8 +914,16 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): # Use FP16/BF16 to sum the results may cause overflow, use FP32 for the summation jitted_primitive = jit( value_and_grad( - lambda q, k, v, bias, *args: grad_func( - customcall_fused_dpa, q, k, v, bias, *args, cp_reverse_out=True, **kwargs + lambda q, k, v, bias, softmax_offset, *args: grad_func( + customcall_fused_dpa, + q, + k, + v, + bias, + softmax_offset, + *args, + cp_reverse_out=True, + **kwargs, ), arg_nums, ), @@ -876,6 +932,7 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): self.qkvo_sharding, self.qkvo_sharding, self.bias_sharding, + self.softmax_offset_sharding, self.seq_desc_sharding, self.dropout_rng_sharding, ), @@ -883,7 +940,9 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): ) jitted_reference = jit( value_and_grad( - lambda q, k, v, bias, *args: grad_func(jax_dpa, q, k, v, bias, *args, **kwargs), + lambda q, k, v, bias, softmax_offset, *args: grad_func( + jax_dpa, q, k, v, bias, softmax_offset, *args, **kwargs + ), arg_nums, ) ) @@ -976,6 +1035,14 @@ def check_dqkv(primitive, reference, pad, idx): ), ], ) +@pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), + pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), + ], +) @pytest.mark.parametrize( "qkv_layout", [ @@ -1084,6 +1151,7 @@ def _test_forward( d_v, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, is_training, @@ -1110,6 +1178,7 @@ def _test_forward( d_v, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, is_training, @@ -1138,6 +1207,7 @@ def test_backward( d_v, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, qkv_layout, @@ -1161,6 +1231,7 @@ def test_backward( d_v, attn_bias_type, attn_mask_type, + softmax_type, dropout_prob, dtype, True, diff --git a/tests/jax/test_layer.py b/tests/jax/test_layer.py index b51d6b2136..2fc9f688ab 100644 --- a/tests/jax/test_layer.py +++ b/tests/jax/test_layer.py @@ -83,6 +83,7 @@ def enable_fused_attn(): _KEY_OF_USE_BIAS = "use_bias" _KEY_OF_RELATIVE_EMBEDDING = "enable_relative_embedding" _KEY_OF_WINDOW_SIZE = "window_size" +_KEY_OF_SOFTMAX_TYPE = "softmax_type" BASE_ATTRS = { _KEY_OF_TRANSPOSE_BS: True, @@ -276,6 +277,14 @@ def enable_fused_attn(): _KEY_OF_RELATIVE_EMBEDDING: True, _KEY_OF_SELF_ATTN_BIAS_TYPE: "post_scale_bias", }, + # attrs31 + { + _KEY_OF_SOFTMAX_TYPE: "off_by_one", + }, + # attrs31 + { + _KEY_OF_SOFTMAX_TYPE: "learnable", + }, ] ATTRS = [{**BASE_ATTRS, **attr} for attr in ATTRS] @@ -418,6 +427,9 @@ class EncoderRunner(BaseRunner): "attention/qkv/ln_bias": "pre_attention_layer_norm/ln_bias", "attention/query/scale": "pre_attention_layer_norm/scale", "attention/query/ln_bias": "pre_attention_layer_norm/ln_bias", + "attention/DotProductAttention_0/_UnfusedDotProductAttention_0/softmax_offset": ( + "attention/DotProductAttention_0/softmax_offset" + ), "mlp/wi_kernel": "mlp/wi/kernel", "mlp/wi_bias": "mlp/wi/bias", "mlp/wo_kernel": "mlp/wo/kernel", @@ -463,10 +475,16 @@ class DecoderRunner(BaseRunner): "encoder_decoder_attention/qkv/ln_bias": "pre_cross_attention_layer_norm/ln_bias", "encoder_decoder_attention/query/scale": "pre_cross_attention_layer_norm/scale", "encoder_decoder_attention/query/ln_bias": "pre_cross_attention_layer_norm/ln_bias", + "encoder_decoder_attention/DotProductAttention_0/_UnfusedDotProductAttention_0/softmax_offset": ( + "encoder_decoder_attention/DotProductAttention_0/softmax_offset" + ), "self_attention/qkv/scale": "pre_self_attention_layer_norm/scale", "self_attention/qkv/ln_bias": "pre_self_attention_layer_norm/ln_bias", "self_attention/query/scale": "pre_self_attention_layer_norm/scale", "self_attention/query/ln_bias": "pre_self_attention_layer_norm/ln_bias", + "self_attention/DotProductAttention_0/_UnfusedDotProductAttention_0/softmax_offset": ( + "self_attention/DotProductAttention_0/softmax_offset" + ), "mlp/wi_kernel": "mlp/wi/kernel", "mlp/wi_bias": "mlp/wi/bias", "mlp/wo_kernel": "mlp/wo/kernel", diff --git a/tests/jax/test_softmax.py b/tests/jax/test_softmax.py index 09386c92ed..9dd03ea0fd 100644 --- a/tests/jax/test_softmax.py +++ b/tests/jax/test_softmax.py @@ -17,7 +17,8 @@ from utils import assert_allclose from transformer_engine.jax.cpp_extensions import is_softmax_kernel_available -from transformer_engine.jax.softmax import SoftmaxType, softmax +from transformer_engine.jax.cpp_extensions.attention import AttnSoftmaxType +from transformer_engine.jax.softmax import SoftmaxFusionType, softmax from transformer_engine.jax.flax.module import Softmax @@ -50,8 +51,9 @@ class SoftmaxRunner: max_seqlen_kv: int num_heads: int scale_factor: float - softmax_type: SoftmaxType + softmax_fusion_type: SoftmaxFusionType dtype: DTypeLike + softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX @staticmethod def reference_softmax(logits, mask, scale_factor, **_): @@ -68,6 +70,7 @@ def reference_softmax(logits, mask, scale_factor, **_): def _is_support(self): return is_softmax_kernel_available( + self.softmax_fusion_type, self.softmax_type, self.batch_size, self.num_heads, @@ -85,22 +88,22 @@ def _setup_inputs(self): self.logits = jax.random.uniform(logits_key, logits_shape, self.dtype, -1.0) - match self.softmax_type: - case SoftmaxType.SCALED: + match self.softmax_fusion_type: + case SoftmaxFusionType.SCALED: self.mask = None - case SoftmaxType.SCALED_MASKED: + case SoftmaxFusionType.SCALED_MASKED: self.mask = jax.random.bernoulli(mask_key, shape=mask_shape).astype(jnp.uint8) - case SoftmaxType.SCALED_UPPER_TRIANG_MASKED: + case SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED: self.mask = (1.0 - jnp.tril(jnp.ones_like(self.logits))).astype(jnp.uint8) case _: - raise ValueError(f"Unknown {self.softmax_type=}") + raise ValueError(f"Unknown {self.softmax_fusion_type=}") def test_forward(self): """ Test transformer_engine.jax.softmax.softmax fwd rule """ self._setup_inputs() - primitive_out = softmax(self.logits, self.mask, self.scale_factor, self.softmax_type) + primitive_out = softmax(self.logits, self.mask, self.scale_factor, self.softmax_fusion_type) reference_out = __class__.reference_softmax(self.logits, self.mask, self.scale_factor) assert_allclose(primitive_out, reference_out, dtype=self.dtype) @@ -117,7 +120,7 @@ def grad_func(func, *args, **kwargs): args = [self.logits, self.mask] kwargs = { "scale_factor": self.scale_factor, - "softmax_type": self.softmax_type, + "softmax_fusion_type": self.softmax_fusion_type, } # Use FP16/BF16 to sum the results may cause overflow, use FP32 for the summation @@ -175,7 +178,7 @@ def test_forward(self): rng = jax.random.PRNGKey(0) softmax_module = Softmax( scale_factor=runner.scale_factor, - softmax_type=runner.softmax_type, + softmax_fusion_type=runner.softmax_fusion_type, ) softmax_vars = softmax_module.init(rng, runner.logits, runner.mask) module_out = softmax_module.apply(softmax_vars, runner.logits, runner.mask) @@ -194,11 +197,11 @@ def test_forward(self): ) @pytest.mark.parametrize("scale_factor", [0.125]) @pytest.mark.parametrize( - "softmax_type", + "softmax_fusion_type", [ - pytest.param(SoftmaxType.SCALED, id="SCALED"), - pytest.param(SoftmaxType.SCALED_MASKED, id="SCALED_MASKED"), - pytest.param(SoftmaxType.SCALED_UPPER_TRIANG_MASKED, id="SCALED_UPPER_TRIANG_MASKED"), + pytest.param(SoftmaxFusionType.SCALED, id="SCALED"), + pytest.param(SoftmaxFusionType.SCALED_MASKED, id="SCALED_MASKED"), + pytest.param(SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED, id="SCALED_UPPER_TRIANG_MASKED"), ], ) @pytest.mark.parametrize( @@ -214,19 +217,19 @@ class TestSoftmaxPrimitives: """ @staticmethod - def test_forward(b, s_q, s_kv, h, scale_factor, softmax_type, dtype): + def test_forward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): """ Test forward with parameterized configs """ - runner = SoftmaxPrimitivesRunner(b, s_q, s_kv, h, scale_factor, softmax_type, dtype) + runner = SoftmaxPrimitivesRunner(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype) runner.test_forward() @staticmethod - def test_backward(b, s_q, s_kv, h, scale_factor, softmax_type, dtype): + def test_backward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): """ Test forward with parameterized configs """ - runner = SoftmaxPrimitivesRunner(b, s_q, s_kv, h, scale_factor, softmax_type, dtype) + runner = SoftmaxPrimitivesRunner(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype) runner.test_backward() @@ -243,11 +246,11 @@ def test_backward(b, s_q, s_kv, h, scale_factor, softmax_type, dtype): ) @pytest.mark.parametrize("scale_factor", [0.125]) @pytest.mark.parametrize( - "softmax_type", + "softmax_fusion_type", [ - pytest.param(SoftmaxType.SCALED, id="SCALED"), - pytest.param(SoftmaxType.SCALED_MASKED, id="SCALED_MASKED"), - pytest.param(SoftmaxType.SCALED_UPPER_TRIANG_MASKED, id="SCALED_UPPER_TRIANG_MASKED"), + pytest.param(SoftmaxFusionType.SCALED, id="SCALED"), + pytest.param(SoftmaxFusionType.SCALED_MASKED, id="SCALED_MASKED"), + pytest.param(SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED, id="SCALED_UPPER_TRIANG_MASKED"), ], ) @pytest.mark.parametrize( @@ -263,11 +266,11 @@ class TestSoftmaxModule: """ @staticmethod - def test_forward(b, s_q, s_kv, h, scale_factor, softmax_type, dtype): + def test_forward(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype): """ Test forward with parameterized configs """ - module_runner = SoftmaxRunner(b, s_q, s_kv, h, scale_factor, softmax_type, dtype) + module_runner = SoftmaxRunner(b, s_q, s_kv, h, scale_factor, softmax_fusion_type, dtype) bias = None runner = SoftmaxModuleRunner(module_runner, bias) runner.test_forward() diff --git a/tests/jax/utils.py b/tests/jax/utils.py index bbe8e65829..7194e387c7 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -21,6 +21,7 @@ import pytest from transformer_engine.jax.attention import ( + AttnSoftmaxType, canonicalize_attn_mask_type, make_swa_mask, ) @@ -162,6 +163,7 @@ class DotProductAttention(nn.Module): dropout_rate: float = 0.0 dtype: DType = jnp.float32 float32_logits: bool = False + softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX """Computes dot-product attention given query, key, and value. This is the core function for applying attention based on @@ -211,6 +213,24 @@ def __call__( assert key.shape[-2] == value.shape[-2], "k, v num_heads must match." assert query.shape[-1] == key.shape[-1], "q, k head_dim must match." + # Infer number of attention heads from query shape + # query shape: [..., h, d] where h is num_attention_heads + num_attention_heads = query.shape[-2] + + # Initialize softmax_offset for off-by-one or learnable softmax + softmax_offset = None + if self.softmax_type == AttnSoftmaxType.OFF_BY_ONE_SOFTMAX: + # For off-by-one softmax, use zeros with shape (1, h, 1, 1) + softmax_offset = jnp.zeros((1, num_attention_heads, 1, 1), dtype=input_dtype) + elif self.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX: + # For learnable softmax, create a learnable parameter with shape (1, h, 1, 1) + softmax_offset = self.param( + "softmax_offset", + nn.initializers.zeros, + (1, num_attention_heads, 1, 1), + jnp.float32, + ) + if self.scale_attn_logits: head_dim = query.shape[-1] depth_scaling = jnp.sqrt(head_dim).astype(input_dtype) @@ -241,9 +261,23 @@ def __call__( if bias is not None: attn_weights = attn_weights + bias.astype(attn_weights.dtype) + # Add attention sink to the last column if not vanilla softmax + if self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + # Add extra column with softmax_offset + # softmax_offset shape: (1, h, 1, 1), attn_weights shape: [b, h, q, k] + extra_col = jnp.broadcast_to( + softmax_offset, + (attn_weights.shape[0], attn_weights.shape[1], attn_weights.shape[2], 1), + ) + attn_weights = jnp.concatenate([attn_weights, extra_col], axis=-1) + # Normalize the attention weights across `kv_length` dimension. attn_weights = jax_nn.softmax(attn_weights).astype(input_dtype) + # Remove the extra column after softmax if not vanilla softmax + if self.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + attn_weights = attn_weights[..., :-1] + # Apply attention dropout. if not deterministic and self.dropout_rate > 0.0: keep_prob = 1.0 - self.dropout_rate @@ -535,6 +569,7 @@ class MultiHeadAttention(nn.Module): rotary_pos_emb_group_method: str = "consecutive" fuse_qkv: bool = True use_bias: bool = False + softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX def __post_init__(self): if self.kernel_init is None: @@ -801,6 +836,7 @@ def qkv_init(key, shape, dtype): dropout_rate=self.dropout_rate, dtype=self.dtype, float32_logits=self.float32_logits, + softmax_type=self.softmax_type, )(query, key, value, bias=attention_bias, deterministic=deterministic) x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3])) @@ -1058,6 +1094,7 @@ class EncoderLayer(nn.Module): self_attn_bias_type: Any = None self_attn_mask_type: str = "no_mask" window_size: Tuple[int, int] = (-1, -1) + softmax_type: str = "vanilla" def __post_init__(self): if self.num_gqa_groups is None: @@ -1111,6 +1148,9 @@ def __call__(self, inputs, encoder_mask=None, deterministic=False): else: x = inputs + # Convert softmax_type string to AttnSoftmaxType enum + attn_softmax_type = AttnSoftmaxType.from_str(self.softmax_type) + # [batch, length, emb_dim] -> [batch, length, emb_dim] x = MultiHeadAttention( num_heads=self.num_attention_heads, @@ -1126,6 +1166,7 @@ def __call__(self, inputs, encoder_mask=None, deterministic=False): enable_rotary_pos_emb=self.enable_rotary_pos_emb, rotary_pos_emb_group_method=self.rotary_pos_emb_group_method, use_bias=self.use_bias, + softmax_type=attn_softmax_type, name="attention", )(x, x, encoder_mask, encoder_bias, deterministic=deterministic) x = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)( @@ -1222,6 +1263,7 @@ class DecoderLayer(nn.Module): self_attn_bias_type: Any = None self_attn_mask_type: str = "no_mask" window_size: Tuple[int, int] = (-1, -1) + softmax_type: str = "vanilla" def __post_init__(self): if self.num_gqa_groups is None: @@ -1290,6 +1332,9 @@ def __call__( else: x = inputs + # Convert softmax_type string to AttnSoftmaxType enum + attn_softmax_type = AttnSoftmaxType.from_str(self.softmax_type) + # Self-attention block x = MultiHeadAttention( num_heads=self.num_attention_heads, @@ -1305,6 +1350,7 @@ def __call__( rotary_pos_emb_group_method=self.rotary_pos_emb_group_method, fuse_qkv=self.fuse_qkv_params, use_bias=self.use_bias, + softmax_type=attn_softmax_type, name="self_attention", )(x, x, decoder_mask, decoder_bias, deterministic=deterministic, decode=decode) x = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)( @@ -1343,6 +1389,7 @@ def __call__( rotary_pos_emb_group_method=self.rotary_pos_emb_group_method, fuse_qkv=self.fuse_qkv_params, use_bias=self.use_bias, + softmax_type=attn_softmax_type, name="encoder_decoder_attention", )(y, encoded, encoder_decoder_mask, deterministic=deterministic) y = nn.Dropout(rate=self.hidden_dropout, broadcast_dims=self.hidden_dropout_dims)( diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 1ce44a2b93..57b118d635 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -18,6 +18,7 @@ from transformer_engine_jax import NVTE_QKV_Layout from transformer_engine_jax import NVTE_QKV_Format from transformer_engine_jax import nvte_get_qkv_format +from transformer_engine_jax import NVTE_Softmax_Type from . import cpp_extensions as tex @@ -74,6 +75,35 @@ def is_bottom_right(self): ] +class AttnSoftmaxType(Enum): + """ + VANILLA_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), + OFF_BY_ONE_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), + LEARNABLE_SOFTMAX: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), + where alpha is a learnable parameter in shape [H]. + """ + + VANILLA_SOFTMAX = NVTE_Softmax_Type.NVTE_VANILLA_SOFTMAX + OFF_BY_ONE_SOFTMAX = NVTE_Softmax_Type.NVTE_OFF_BY_ONE_SOFTMAX + LEARNABLE_SOFTMAX = NVTE_Softmax_Type.NVTE_LEARNABLE_SOFTMAX + + @classmethod + def from_str(cls, softmax_type: str) -> "AttnSoftmaxType": + """Convert string to AttnSoftmaxType: 'vanilla', 'off_by_one', or 'learnable'.""" + softmax_type_map = { + "vanilla": cls.VANILLA_SOFTMAX, + "off_by_one": cls.OFF_BY_ONE_SOFTMAX, + "learnable": cls.LEARNABLE_SOFTMAX, + } + result = softmax_type_map.get(softmax_type) + if result is None: + raise ValueError( + f"Unknown softmax_type: {softmax_type}. " + "Valid options: 'vanilla', 'off_by_one', 'learnable'" + ) + return result + + class QKVFormat(Enum): """ SBHD: q,k,v memory layout with [s, b, ..., h, d] @@ -301,6 +331,7 @@ def is_fused_attn_kernel_available( qkv_layout, attn_bias_type, attn_mask_type, + softmax_type, dropout_probability, q_num_heads, kv_num_heads, @@ -313,6 +344,7 @@ def is_fused_attn_kernel_available( """ To check whether the fused attention kernel is supported """ + window_size_tuple = (-1, -1) if window_size is None else window_size def make_helper(attn_mask_type): return tex.FusedAttnHelper( @@ -322,6 +354,7 @@ def make_helper(attn_mask_type): qkv_layout, attn_bias_type, attn_mask_type, + softmax_type, dropout_probability, q_num_heads, kv_num_heads, @@ -329,7 +362,7 @@ def make_helper(attn_mask_type): kv_max_seqlen, head_dim_qk, head_dim_v, - (-1, -1) if window_size is None else window_size, + window_size_tuple, ) return make_helper(attn_mask_type).is_fused_attn_kernel_available() @@ -786,6 +819,7 @@ def _legacy_fused_attn( attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType, qkv_layout: QKVLayout, + softmax_type: AttnSoftmaxType, scaling_factor: float, dropout_probability: float, is_training: bool, @@ -793,6 +827,7 @@ def _legacy_fused_attn( context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT, context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", + softmax_offset: Optional[jnp.ndarray] = None, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -815,6 +850,7 @@ def _legacy_fused_attn( seed (Optional[jnp.ndarray]): Optional random seed for dropout. attn_bias_type (AttnBiasType): Type of attention bias. attn_mask_type (AttnMaskType): Type of attention mask. + softmax_type (AttnSoftmaxType): Type of attention softmax. qkv_layout (QKVLayout): Layout of the QKV tensors. scaling_factor (float): Scaling factor for the attention scores. dropout_probability (float): Dropout probability to apply during attention. @@ -863,10 +899,12 @@ def _legacy_fused_attn( output = _fused_attn( qkv, bias, + softmax_offset, SequenceDescriptor.from_seqlens((q_seq_lens, kv_seq_lens)), seed, attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, + softmax_type=softmax_type, qkv_layout=qkv_layout, scaling_factor=scaling_factor, dropout_probability=dropout_probability, @@ -900,6 +938,7 @@ def fused_attn_thd( context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT, context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", + softmax_offset: Optional[jnp.ndarray] = None, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -937,6 +976,7 @@ def fused_attn_thd( output = _fused_attn( qkv, bias, + softmax_offset, SequenceDescriptor.from_seqlens_and_offsets( (q_seq_lens, kv_seq_lens), (q_seq_offsets, kv_seq_offsets) ), @@ -945,6 +985,7 @@ def fused_attn_thd( attn_mask_type=attn_mask_type, qkv_layout=qkv_layout, scaling_factor=scaling_factor, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, dropout_probability=dropout_probability, is_training=is_training, max_segments_per_seq=max_segments_per_seq, @@ -957,15 +998,17 @@ def fused_attn_thd( return output -@partial(jax.custom_vjp, nondiff_argnums=(4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15)) +@partial(jax.custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17)) def _fused_attn( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], + softmax_offset: Optional[jnp.ndarray], sequence_descriptor: SequenceDescriptor, seed: Optional[jnp.ndarray], attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType, qkv_layout: QKVLayout, + softmax_type: AttnSoftmaxType, scaling_factor: float, dropout_probability: float, is_training: bool, @@ -979,11 +1022,13 @@ def _fused_attn( output, _ = _fused_attn_fwd_rule( qkv, bias, + softmax_offset, sequence_descriptor, seed, attn_bias_type, attn_mask_type, qkv_layout, + softmax_type, scaling_factor, dropout_probability, is_training, @@ -1000,11 +1045,13 @@ def _fused_attn( def _fused_attn_fwd_rule( qkv, bias, + softmax_offset, sequence_descriptor, seed, attn_bias_type, attn_mask_type, qkv_layout, + softmax_type, scaling_factor, dropout_probability, is_training, @@ -1018,10 +1065,12 @@ def _fused_attn_fwd_rule( output, softmax_aux, rng_state = tex.fused_attn_fwd( qkv, bias, + softmax_offset, sequence_descriptor, seed, attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, + softmax_type=softmax_type, qkv_layout=qkv_layout, scaling_factor=scaling_factor, dropout_probability=dropout_probability, @@ -1041,6 +1090,7 @@ def _fused_attn_fwd_rule( sequence_descriptor, softmax_aux, rng_state, + softmax_offset, output, ) @@ -1049,6 +1099,7 @@ def _fused_attn_bwd_rule( attn_bias_type, attn_mask_type, qkv_layout, + softmax_type, scaling_factor, dropout_probability, is_training, @@ -1068,11 +1119,13 @@ def _fused_attn_bwd_rule( sequence_descriptor, softmax_aux, rng_state, + softmax_offset, output, ) = ctx - grad_qkv, grad_bias = tex.fused_attn_bwd( + grad_qkv, grad_bias, grad_softmax_offset = tex.fused_attn_bwd( qkv, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -1080,6 +1133,7 @@ def _fused_attn_bwd_rule( sequence_descriptor, attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, + softmax_type=softmax_type, qkv_layout=qkv_layout, scaling_factor=scaling_factor, dropout_probability=dropout_probability, @@ -1092,9 +1146,12 @@ def _fused_attn_bwd_rule( ) if attn_bias_type == AttnBiasType.NO_BIAS: grad_bias = None + if softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX: + grad_softmax_offset = None return ( grad_qkv, grad_bias, + grad_softmax_offset, None, None, ) @@ -1111,6 +1168,7 @@ def fused_attn( attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType, qkv_layout: QKVLayout, + softmax_type: AttnSoftmaxType, scaling_factor: float, dropout_probability: float, is_training: bool, @@ -1120,6 +1178,7 @@ def fused_attn( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", context_checkpoint_name: str = "context", + softmax_offset: Optional[jnp.ndarray] = None, ): """ Perform cuDNN fused attention. @@ -1139,6 +1198,7 @@ def fused_attn( seed (Optional[jnp.ndarray]): Optional random seed for dropout. attn_bias_type (AttnBiasType): Type of attention bias. attn_mask_type (AttnMaskType): Type of attention mask. + softmax_type (AttnSoftmaxType): Type of attention softmax. qkv_layout (QKVLayout): Layout of the QKV tensors. scaling_factor (float): Scaling factor for the attention scores. dropout_probability (float): Dropout probability to apply during attention. @@ -1153,6 +1213,9 @@ def fused_attn( Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. context_checkpoint_name (str): The name of the context checkpoint for the custom VJP forward pass. + softmax_offset (Optional[jnp.ndarray]): An optional learnable softmax offset tensor with shape + [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX. + If provided, this parameter will receive gradients during backpropagation. Returns: (jnp.ndarray): The output tensor from the fused attention. @@ -1200,6 +1263,7 @@ def fused_attn( seed, attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, + softmax_type=softmax_type, qkv_layout=qkv_layout, scaling_factor=scaling_factor, dropout_probability=dropout_probability, @@ -1208,15 +1272,18 @@ def fused_attn( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + softmax_offset=softmax_offset, ) output = _fused_attn( qkv, bias, + softmax_offset, sequence_descriptor, seed, attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, qkv_layout=qkv_layout, + softmax_type=softmax_type, scaling_factor=scaling_factor, dropout_probability=dropout_probability, is_training=is_training, diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 6a21480d8d..f0778bfd29 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -20,11 +20,13 @@ from transformer_engine.jax.attention import ( AttnBiasType, AttnMaskType, + AttnSoftmaxType, QKVLayout, QKVFormat, CPStrategy, SequenceDescriptor, ) +from ..sharding import with_sharding_constraint_by_logical_axes, HEAD_AXES from .base import BasePrimitive, register_primitive from .misc import ( @@ -61,6 +63,7 @@ meta_fields=[ "attn_bias_type", "attn_mask_type", + "softmax_type", "qkv_layout", "scaling_factor", "dropout_probability", @@ -80,6 +83,7 @@ class _FusedAttnConfig: attn_bias_type: AttnBiasType attn_mask_type: AttnMaskType + softmax_type: AttnSoftmaxType qkv_layout: QKVLayout scaling_factor: float dropout_probability: float @@ -103,6 +107,7 @@ class FusedAttnHelper: qkv_layout: QKVLayout attn_bias_type: AttnBiasType attn_mask_type: AttnMaskType + softmax_type: AttnSoftmaxType dropout_probability: float q_num_heads: int kv_num_heads: int @@ -125,6 +130,7 @@ def get_fused_attn_backend(self): self.qkv_layout.value, self.attn_bias_type.value, self.attn_mask_type.value, + self.softmax_type.value, self.dropout_probability, self.q_num_heads, self.kv_num_heads, @@ -254,7 +260,7 @@ class FusedAttnFwdPrimitive(BasePrimitive): name = "te_fused_attn_forward_ffi" multiple_results = True - impl_static_args = (13,) + impl_static_args = (14,) inner_primitive = None outer_primitive = None @@ -264,6 +270,7 @@ def abstract( k_aval, v_aval, bias_aval, + softmax_offset_aval, seed_aval, q_seqlen_or_cu_seqlen_aval, kv_seqlen_or_cu_seqlen_aval, @@ -312,6 +319,7 @@ def abstract( config.qkv_layout, config.attn_bias_type, config.attn_mask_type, + config.softmax_type, config.dropout_probability, attn_heads, num_gqa_groups, @@ -375,6 +383,7 @@ def abstract( config.dropout_probability, config.attn_bias_type.value, config.attn_mask_type.value, + config.softmax_type.value, config.qkv_layout.value, jax_dtype_to_te_dtype(q_aval.dtype), config.is_training, @@ -386,6 +395,12 @@ def abstract( shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) ) + assert softmax_offset_aval.dtype == jnp.float32 + if config.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + assert softmax_offset_aval.shape == (1, attn_heads, 1, 1) + else: + assert softmax_offset_aval.shape == (0,) + return out_aval, softmax_aux_aval, rng_state_aval, wkspace_aval @staticmethod @@ -405,6 +420,7 @@ def lowering( k, v, bias, + softmax_offset, seed, q_cu_seqlen, kv_cu_seqlen, @@ -453,6 +469,7 @@ def lowering( k, v, bias, + softmax_offset, seed, q_cu_seqlen, kv_cu_seqlen, @@ -481,6 +498,7 @@ def lowering( deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), window_size_left=window_size_left, window_size_right=window_size_right, + softmax_type=int(config.softmax_type.value), ) @staticmethod @@ -489,6 +507,7 @@ def impl( k, v, bias, + softmax_offset, seed, q_seqlen, kv_seqlen, @@ -579,6 +598,7 @@ def convert_to_2d(offsets, batch, max_seqlen): k, v, bias, + softmax_offset, seed, q_cu_seqlen, kv_cu_seqlen, @@ -596,7 +616,7 @@ def convert_to_2d(offsets, batch, max_seqlen): def batcher(batched_args, batch_dims, *, config): check_valid_batch_dims(batch_dims) assert FusedAttnFwdPrimitive.outer_primitive is not None - q_bdim, _, _, _, seed_bdim, *_ = batch_dims + q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims out_bdims = q_bdim, q_bdim, seed_bdim return ( @@ -662,7 +682,7 @@ def partition(config, mesh, arg_infos, result_infos): mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] - arg_shardings[4] = seed_sharding + arg_shardings[5] = seed_sharding arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) @@ -710,7 +730,7 @@ class FusedAttnBwdPrimitive(BasePrimitive): name = "te_fused_attn_backward_ffi" multiple_results = True - impl_static_args = (16,) + impl_static_args = (17,) inner_primitive = None outer_primitive = None @@ -720,6 +740,7 @@ def abstract( k_aval, v_aval, bias_aval, + softmax_offset_aval, softmax_aux_aval, rng_state_aval, output_aval, @@ -781,6 +802,7 @@ def abstract( config.dropout_probability, config.attn_bias_type.value, config.attn_mask_type.value, + config.softmax_type.value, config.qkv_layout.value, jax_dtype_to_te_dtype(q_aval.dtype), config.is_training, @@ -798,15 +820,39 @@ def abstract( shape=wkspace_shape, dtype=te_dtype_to_jax_dtype(wkspace_dtype) ) - return dq_aval, dk_aval, dv_aval, dbias_aval, wkspace_aval + # Validate incoming softmax_offset shape and dtype + assert ( + softmax_offset_aval.dtype == jnp.float32 + ), f"Incorrect softmax_offset dtype: {softmax_offset_aval.dtype}, expected: {jnp.float32}" + if config.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + assert softmax_offset_aval.shape == (1, attn_heads, 1, 1), ( + f"Incorrect softmax_offset shape for {config.softmax_type}:" + f" {softmax_offset_aval.shape}, expected: (1, {attn_heads}, 1, 1)" + ) + else: + assert softmax_offset_aval.shape == (0,), ( + f"Incorrect softmax_offset shape for {config.softmax_type}:" + f" {softmax_offset_aval.shape}, expected: (0,)" + ) + + if config.softmax_type == AttnSoftmaxType.VANILLA_SOFTMAX: + dsoftmax_offset_aval = q_aval.update( + shape=softmax_offset_aval.shape, dtype=softmax_offset_aval.dtype + ) + else: + dsoftmax_offset_aval = q_aval.update(shape=(1, attn_heads, 1, 1), dtype=jnp.float32) + + return dq_aval, dk_aval, dv_aval, dbias_aval, dsoftmax_offset_aval, wkspace_aval @staticmethod def outer_abstract(*args, **kwargs): """ Fused attention fwd outer primitive abstract """ - dq_aval, dk_aval, dv_aval, dbias_aval, _ = FusedAttnBwdPrimitive.abstract(*args, **kwargs) - return dq_aval, dk_aval, dv_aval, dbias_aval + dq_aval, dk_aval, dv_aval, dbias_aval, dsoftmax_offset_aval, _ = ( + FusedAttnBwdPrimitive.abstract(*args, **kwargs) + ) + return dq_aval, dk_aval, dv_aval, dbias_aval, dsoftmax_offset_aval @staticmethod def lowering( @@ -815,6 +861,7 @@ def lowering( k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -866,6 +913,7 @@ def lowering( k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -897,6 +945,7 @@ def lowering( deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), window_size_left=window_size_left, window_size_right=window_size_right, + softmax_type=int(config.softmax_type.value), ) @staticmethod @@ -905,6 +954,7 @@ def impl( k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -993,11 +1043,12 @@ def convert_to_2d(offsets, batch, max_seqlen): q_cu_seqlen = generate_cu_seqlen(q_seqlen.flatten()) kv_cu_seqlen = generate_cu_seqlen(kv_seqlen.flatten()) - dq, dk, dv, dbias, _ = FusedAttnBwdPrimitive.inner_primitive.bind( + dq, dk, dv, dbias, dsoftmax_offset, _ = FusedAttnBwdPrimitive.inner_primitive.bind( q, k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -1012,15 +1063,15 @@ def convert_to_2d(offsets, batch, max_seqlen): _kv_segment_pos, config=config, ) - return dq, dk, dv, dbias + return dq, dk, dv, dbias, dsoftmax_offset @staticmethod def batcher(batched_args, batch_dims, *, config): check_valid_batch_dims(batch_dims) assert FusedAttnBwdPrimitive.outer_primitive is not None - q_bdim, k_bdim, v_bdim, *_ = batch_dims + q_bdim, k_bdim, v_bdim, bias_bdim, softmax_offset_bdim, *_ = batch_dims - out_bdims = q_bdim, k_bdim, v_bdim, q_bdim + out_bdims = q_bdim, k_bdim, v_bdim, bias_bdim, softmax_offset_bdim return ( FusedAttnBwdPrimitive.outer_primitive.bind(*batched_args, config=config), out_bdims, @@ -1033,11 +1084,13 @@ def infer_sharding_from_operands(config, mesh, arg_infos, result_infos): k_spec = get_padded_spec(arg_infos[1]) v_spec = get_padded_spec(arg_infos[2]) bias_spec = get_padded_spec(arg_infos[3]) + softmax_offset_spec = get_padded_spec(arg_infos[4]) dq_sharding = NamedSharding(mesh, PartitionSpec(*q_spec)) dk_sharding = NamedSharding(mesh, PartitionSpec(*k_spec)) dv_sharding = NamedSharding(mesh, PartitionSpec(*v_spec)) dbias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) - return (dq_sharding, dk_sharding, dv_sharding, dbias_sharding) + dsoftmax_offset_sharding = NamedSharding(mesh, PartitionSpec(*softmax_offset_spec)) + return (dq_sharding, dk_sharding, dv_sharding, dbias_sharding, dsoftmax_offset_sharding) @staticmethod def partition(config, mesh, arg_infos, result_infos): @@ -1046,21 +1099,30 @@ def partition(config, mesh, arg_infos, result_infos): k_spec = get_padded_spec(arg_infos[1]) v_spec = get_padded_spec(arg_infos[2]) bias_spec = get_padded_spec(arg_infos[3]) + softmax_offset_spec = get_padded_spec(arg_infos[4]) dq_sharding = NamedSharding(mesh, PartitionSpec(*q_spec)) dk_sharding = NamedSharding(mesh, PartitionSpec(*k_spec)) dv_sharding = NamedSharding(mesh, PartitionSpec(*v_spec)) dbias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) + dsoftmax_offset_sharding = NamedSharding(mesh, PartitionSpec(*softmax_offset_spec)) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (dq_sharding, dk_sharding, dv_sharding, dbias_sharding) + out_shardings = ( + dq_sharding, + dk_sharding, + dv_sharding, + dbias_sharding, + dsoftmax_offset_sharding, + ) def sharded_impl( q, k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -1074,36 +1136,43 @@ def sharded_impl( _q_segment_pos, _kv_segment_pos, ): - local_dq, local_dk, local_dv, local_dbias = FusedAttnBwdPrimitive.impl( - q, - k, - v, - bias, - softmax_aux, - rng_state, - output, - doutput, - q_cu_seqlen, - kv_cu_seqlen, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=config, + local_dq, local_dk, local_dv, local_dbias, local_dsoftmax_offset = ( + FusedAttnBwdPrimitive.impl( + q, + k, + v, + bias, + softmax_offset, + softmax_aux, + rng_state, + output, + doutput, + q_cu_seqlen, + kv_cu_seqlen, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=config, + ) ) global_dbias = local_dbias if config.attn_bias_type is not AttnBiasType.NO_BIAS: global_dbias = all_reduce_sum_along_dp_fsdp(local_dbias, mesh) - return local_dq, local_dk, local_dv, global_dbias + + global_dsoftmax_offset = local_dsoftmax_offset + if config.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX: + global_dsoftmax_offset = all_reduce_sum_along_dp_fsdp(local_dsoftmax_offset, mesh) + + return local_dq, local_dk, local_dv, global_dbias, global_dsoftmax_offset return mesh, sharded_impl, out_shardings, arg_shardings @staticmethod def shardy_sharding_rule(config, mesh, value_types, result_types): del config, mesh - # We only care about the four first arguments. # Keep in sync with `infer_sharding_from_operands`. input_spec = tuple((f"…{x}",) for x in range(len(value_types))) output_spec = tuple((f"…{x}",) for x in range(len(result_types))) @@ -1229,6 +1298,11 @@ def check_supported(self): if self.config.dropout_probability != 0.0: raise ValueError(f"{header} does not support dropout") + if self.config.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + raise ValueError( + f"{header} only supports VANILLA_SOFTMAX, got: {self.config.softmax_type}" + ) + def get_adjusted_mask(self): """Converts the mask for context parallelism.""" if self.config.attn_mask_type == AttnMaskType.CAUSAL_MASK: @@ -1240,6 +1314,7 @@ def get_step_config(self) -> _FusedAttnConfig: return _FusedAttnConfig( attn_bias_type=self.config.attn_bias_type, attn_mask_type=self.get_adjusted_mask(), + softmax_type=self.config.softmax_type, qkv_layout=self.config.qkv_layout, scaling_factor=self.config.scaling_factor, dropout_probability=self.config.dropout_probability, @@ -1376,7 +1451,7 @@ def partition(config, mesh, arg_infos, result_infos): mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] - arg_shardings[4] = seed_sharding + arg_shardings[5] = seed_sharding arg_shardings = tuple(arg_shardings) out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) @@ -1385,6 +1460,7 @@ def impl( k, v, bias, + softmax_offset, seed, q_seqlen, kv_seqlen, @@ -1404,7 +1480,7 @@ def impl( # meeting the expectation of the SPMD model. # TODO(mgoldfarb-nvidia): When cuDNN supports we should be able to make use of a padding # mask/sequence length tensor to avoid this unrolled loop. - def _cross_attn(idx, q, k, v, bias, q_seqlen, kv_seqlen, seed): + def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): kv_max_seqlen = k.shape[1] kv_seqlen_per_subrank = kv_max_seqlen // (cp_size * 2) assert kv_max_seqlen % cp_size == 0, "sequence length must evenly divide cp size" @@ -1431,6 +1507,7 @@ def _cross_attn(idx, q, k, v, bias, q_seqlen, kv_seqlen, seed): k_unmasked, v_unmasked, bias, + softmax_offset, seed, q_seqlen_for_step, kv_seqlen_for_step, @@ -1453,7 +1530,9 @@ def _cross_attn(idx, q, k, v, bias, q_seqlen, kv_seqlen, seed): k_ag, v_ag = helper.all_gather_kv(k, v) functions = [ - partial(_cross_attn, idx, q, k_ag, v_ag, bias, q_seqlen, kv_seqlen, seed) + partial( + _cross_attn, idx, q, k_ag, v_ag, bias, softmax_offset, q_seqlen, kv_seqlen, seed + ) for idx in range(cp_size) ] @@ -1492,18 +1571,27 @@ def partition(config, mesh, arg_infos, result_infos): k_spec = get_padded_spec(arg_infos[1]) v_spec = get_padded_spec(arg_infos[2]) bias_spec = get_padded_spec(arg_infos[3]) + softmax_offset_spec = get_padded_spec(arg_infos[4]) dq_sharding = NamedSharding(mesh, PartitionSpec(*q_spec)) dk_sharding = NamedSharding(mesh, PartitionSpec(*k_spec)) dv_sharding = NamedSharding(mesh, PartitionSpec(*v_spec)) dbias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) + dsoftmax_offset_sharding = NamedSharding(mesh, PartitionSpec(*softmax_offset_spec)) arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) - out_shardings = (dq_sharding, dk_sharding, dv_sharding, dbias_sharding) + out_shardings = ( + dq_sharding, + dk_sharding, + dv_sharding, + dbias_sharding, + dsoftmax_offset_sharding, + ) def impl( q, k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -1527,6 +1615,7 @@ def _cross_attn_bwd( k, v, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -1562,11 +1651,12 @@ def _cross_attn_bwd( num_kv_chunks = kv_max_seqlen // kv_seqlens_for_rank[sub_idx] kv_seqlen_for_step = (kv_seqlen // (cp_size * 2)) * num_kv_chunks - dq_local, dk_local, dv_local, dbias_local = FusedAttnBwdPrimitive.impl( + dq_local, dk_local, dv_local, dbias_local, _ = FusedAttnBwdPrimitive.impl( q_split[sub_idx], k_unmasked, v_unmasked, bias, + softmax_offset, softmax_aux_split[sub_idx], rng_state, output_split[sub_idx], @@ -1604,6 +1694,7 @@ def _cross_attn_bwd( k_ag, v_ag, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -1621,7 +1712,9 @@ def _cross_attn_bwd( dq, dk_local, dv_local, dbias = lax.switch(cp_rank, functions) dk, dv = helper.reduce_scatter_dkv(dk_local, dv_local) - return dq, dk, dv, dbias + # Return dummy dsoftmax_offset for arity matching (all-gather CP doesn't use it) + dummy_dsoftmax_offset = jnp.empty_like(softmax_offset) + return dq, dk, dv, dbias, dummy_dsoftmax_offset return mesh, impl, out_shardings, arg_shardings @@ -1679,6 +1772,11 @@ def check_supported(self): if self.config.dropout_probability != 0.0: raise ValueError(f"{header} does not support dropout") + if self.config.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + raise ValueError( + f"{header} only supports VANILLA_SOFTMAX, got: {self.config.softmax_type}" + ) + # We want to encourage use of scan loop to minimize unrolling and ensure more # predictable scheduling from XLA. The unrolled flavor will be supported but # not the prefered implementation. @@ -1703,6 +1801,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: return _FusedAttnConfig( attn_bias_type=self.config.attn_bias_type, attn_mask_type=attn_mask_type, + softmax_type=self.config.softmax_type, qkv_layout=QKVLayout.BSHD_BS2HD, scaling_factor=self.config.scaling_factor, dropout_probability=self.config.dropout_probability, @@ -1783,7 +1882,7 @@ def partition(config, mesh, arg_infos, result_infos): mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] - arg_shardings[4] = seed_sharding + arg_shardings[5] = seed_sharding # Ensure segment_pos gets same sharding as ID. arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] @@ -1795,6 +1894,7 @@ def ring_attn_fwd_impl( k, v, bias, + _softmax_offset, seed, q_seqlen, kv_seqlen, @@ -1840,6 +1940,7 @@ def mask_compute(attn_mask_type): kv, _not_used, bias, + _softmax_offset, seed, q_seqlen_per_step, kv_seqlen_per_step, @@ -1865,6 +1966,7 @@ def half_kv_no_mask_compute(): kv_part, _not_used, bias, + _softmax_offset, seed, q_seqlen_per_step, kv_seqlen_per_step, @@ -1887,6 +1989,7 @@ def half_q_no_mask_compute(): kv, _not_used, bias, + _softmax_offset, seed, q_seqlen_per_step, kv_seqlen_per_step, @@ -1990,18 +2093,24 @@ def partition(config, mesh, arg_infos, result_infos): k_spec = get_padded_spec(arg_infos[1]) v_spec = get_padded_spec(arg_infos[2]) bias_spec = get_padded_spec(arg_infos[3]) + softmax_offset_spec = get_padded_spec(arg_infos[4]) dq_sharding = NamedSharding(mesh, PartitionSpec(*q_spec)) dk_sharding = NamedSharding(mesh, PartitionSpec(*k_spec)) dv_sharding = NamedSharding(mesh, PartitionSpec(*v_spec)) dbias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) - + # Ring attention doesn't use dsoftmax_offset, but we need to return it for arity matching + dsoftmax_offset_sharding = NamedSharding(mesh, PartitionSpec(*softmax_offset_spec)) arg_shardings = [arg_i.sharding for arg_i in arg_infos] - # Ensure segment_pos gets same sharding as ID. arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - - out_shardings = (dq_sharding, dk_sharding, dv_sharding, dbias_sharding) + out_shardings = ( + dq_sharding, + dk_sharding, + dv_sharding, + dbias_sharding, + dsoftmax_offset_sharding, + ) helper = _FusedAttnCPWithP2PHelper(mesh, config) helper.check_supported() @@ -2011,6 +2120,7 @@ def ring_attn_bwd_impl( k, v, bias, + _softmax_offset, softmax_aux, rng_state, output, @@ -2054,11 +2164,12 @@ def scan_kv_block(idx, carry): def mask_compute(attn_mask_type): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) - dq_per_step, dk_dv_per_step, _, dbias_per_step = FusedAttnBwdPrimitive.impl( + dq_per_step, dk_dv_per_step, _, dbias_per_step, _ = FusedAttnBwdPrimitive.impl( q, kv, _not_used, bias, + _softmax_offset, softmax_aux, rng_state, output, @@ -2082,11 +2193,12 @@ def half_kv_no_mask_compute(): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) // 2 kv_part = lax.slice_in_dim(kv, 0, kv_max_seqlen // 2, axis=1) - dq_per_step, dk_dv_per_step, _, dbias_per_step = FusedAttnBwdPrimitive.impl( + dq_per_step, dk_dv_per_step, _, dbias_per_step, _ = FusedAttnBwdPrimitive.impl( q, kv_part, _not_used, bias, + _softmax_offset, softmax_aux, rng_state, output, @@ -2120,11 +2232,12 @@ def half_q_no_mask_compute(): softmax_aux, q_max_seqlen // 2, q_max_seqlen, axis=2 ) - dq_per_step, dk_dv_per_step, _, dbias_per_step = FusedAttnBwdPrimitive.impl( + dq_per_step, dk_dv_per_step, _, dbias_per_step, _ = FusedAttnBwdPrimitive.impl( q_part, kv, _not_used, bias, + _softmax_offset, softmax_aux_part, rng_state, output_part, @@ -2184,7 +2297,9 @@ def jax_cond_wrap(): global_dbias = all_reduce_sum_along_dp_fsdp(dbias, mesh) dk, dv = helper.unstack_kv(dk_dv) - return dq, dk, dv, global_dbias + # Return dummy dsoftmax_offset for arity matching (ring attention doesn't use it) + dummy_dsoftmax_offset = jnp.empty_like(_softmax_offset) + return dq, dk, dv, global_dbias, dummy_dsoftmax_offset return mesh, ring_attn_bwd_impl, out_shardings, arg_shardings @@ -2273,7 +2388,7 @@ def partition(config, mesh, arg_infos, result_infos): mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] - arg_shardings[4] = seed_sharding + arg_shardings[5] = seed_sharding # Ensure segment_pos gets same sharding as ID. arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] @@ -2285,6 +2400,7 @@ def fwd_impl( k, v, bias, + _softmax_offset, seed, q_seqlen, kv_seqlen, @@ -2336,6 +2452,7 @@ def compute(config): kv, _not_used, bias, + _softmax_offset, seed, q_seqlen, kv_seqlen, @@ -2345,7 +2462,7 @@ def compute(config): kv_segment_ids, q_segment_pos, kv_segment_pos, - config, + config=config, ) if config.window_size != (-1, -1): @@ -2420,8 +2537,8 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - # dq, dk, dv, dbias sharding = q, k, v, bias sharding - out_shardings = tuple(arg.sharding for arg in arg_infos[:4]) + # dq, dk, dv, dbias, dsoftmax_offset sharding = q, k, v, bias, softmax_offset sharding + out_shardings = tuple(arg.sharding for arg in arg_infos[:5]) helper = _FusedAttnCPWithP2PHelper(mesh, config) helper.check_supported() @@ -2431,6 +2548,7 @@ def bwd_impl( k, v, bias, + _softmax_offset, softmax_aux, rng_state, output, @@ -2478,11 +2596,12 @@ def scan_kv_block(idx, carry): kv_segment_pos_next = helper.permute_kv(kv_segment_pos, cp_perm) def compute(config): - dq_per_step, dkv_per_step, _, dbias_per_step = FusedAttnBwdPrimitive.impl( + dq_per_step, dkv_per_step, _, dbias_per_step, _ = FusedAttnBwdPrimitive.impl( q, kv, _not_used, bias, + _softmax_offset, softmax_aux, rng_state, output, @@ -2536,7 +2655,9 @@ def compute(config): global_dbias = all_reduce_sum_along_dp_fsdp(dbias, mesh) dk, dv = helper.unstack_kv(dkv) - return dq, dk, dv, global_dbias + # Return dummy dsoftmax_offset for arity matching (ring attention doesn't use it) + dummy_dsoftmax_offset = jnp.empty_like(_softmax_offset) + return dq, dk, dv, global_dbias, dummy_dsoftmax_offset return mesh, bwd_impl, out_shardings, arg_shardings @@ -2557,10 +2678,12 @@ def _maybe_context_parallel_axis(cp_axis: str): def fused_attn_fwd( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], + softmax_offset: Optional[jnp.ndarray], sequence_descriptor: SequenceDescriptor, seed: Optional[jnp.ndarray], attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType, + softmax_type: AttnSoftmaxType, qkv_layout: QKVLayout, scaling_factor: float, dropout_probability: float, @@ -2585,6 +2708,7 @@ def fused_attn_fwd( query has a different shape (e.g., cross-attention). - `(query, key, value)`: For separate query, key, and value tensors. bias (Optional[jnp.ndarray]): An optional bias tensor to be added to the attention scores. + softmax_offset (Optional[jnp.ndarray]): An optional softmax offset tensor. q_seqlen (jnp.ndarray): Sequence lengths for the query, with shape [batch,]. kv_seqlen (jnp.ndarray): Sequence lengths for the key and value, with shape [batch,]. q_seq_offsets (jnp.ndarray): @@ -2594,6 +2718,7 @@ def fused_attn_fwd( seed (Optional[jnp.ndarray]): Optional random seed for dropout. attn_bias_type (AttnBiasType): Type of attention bias. attn_mask_type (AttnMaskType): Type of attention mask. + softmax_type (AttnSoftmaxType): Type of softmax. qkv_layout (QKVLayout): Layout of the QKV tensors. scaling_factor (float): Scaling factor for the attention scores. dropout_probability (float): Dropout probability to apply during attention. @@ -2633,10 +2758,36 @@ def fused_attn_fwd( assert bias is None bias = jnp.zeros(0, dtype=qkv[0].dtype) + if softmax_offset is None: + assert ( + softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX + ), f"Softmax type {softmax_type} is not supported when softmax_offset is None" + if softmax_type == AttnSoftmaxType.OFF_BY_ONE_SOFTMAX: + num_heads = qkv[0].shape[-2] + # Create tensor [1, h, 1, 1] filled with zeros (logit value = 0) + # This adds exp(0 - x_max) = exp(-x_max) to the denominator, + # which contributes exactly 1 after normalization, giving: exp(x_i) / (sum(exp(x_j)) + 1) + softmax_offset = jnp.zeros((1, num_heads, 1, 1), dtype=jnp.float32) + # Shard by heads dimension + softmax_offset = with_sharding_constraint_by_logical_axes( + softmax_offset, (None, HEAD_AXES, None, None) + ) + else: + assert softmax_type == AttnSoftmaxType.VANILLA_SOFTMAX + softmax_offset = jnp.zeros(0, dtype=jnp.float32) + else: + assert softmax_offset.dtype == jnp.float32 + # Shard by heads dimension if not VANILLA_SOFTMAX + if softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + softmax_offset = with_sharding_constraint_by_logical_axes( + softmax_offset, (None, HEAD_AXES, None, None) + ) + fused_config = _FusedAttnConfig( attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, qkv_layout=qkv_layout, + softmax_type=softmax_type, scaling_factor=scaling_factor, dropout_probability=dropout_probability, is_training=is_training, @@ -2662,6 +2813,7 @@ def fused_attn_fwd( output, softmax_aux, rng_state = primitive.bind( *qkv_for_primitive, bias, + softmax_offset, seed, *seq_desc_flatten, config=fused_config, @@ -2673,6 +2825,7 @@ def fused_attn_fwd( def fused_attn_bwd( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], + softmax_offset: Optional[jnp.ndarray], softmax_aux: jnp.ndarray, rng_state: jnp.ndarray, output: jnp.ndarray, @@ -2681,6 +2834,7 @@ def fused_attn_bwd( attn_bias_type: AttnBiasType, attn_mask_type: AttnMaskType, qkv_layout: QKVLayout, + softmax_type: AttnSoftmaxType, scaling_factor: float, dropout_probability: float, is_training: bool, @@ -2702,6 +2856,7 @@ def fused_attn_bwd( query has a different shape (e.g., cross-attention). - `(query, key, value)`: For separate query, key, and value tensors. bias (Optional[jnp.ndarray]): An optional bias tensor to be added to the attention scores. + softmax_offset (Optional[jnp.ndarray]): An optional softmax offset tensor. softmax_aux (jnp.ndarray): Auxiliary tensors from the softmax step used in the forward pass. rng_state (jnp.ndarray): Auxiliary tensors to save the random state in the forward pass. output (jnp.ndarray): The output tensor from the forward pass. @@ -2714,6 +2869,7 @@ def fused_attn_bwd( The offsets in the sequence dim for the query, with shape [batch + 1,]. attn_bias_type (AttnBiasType): Type of attention bias. attn_mask_type (AttnMaskType): Type of attention mask. + softmax_type (AttnSoftmaxType): Type of softmax. qkv_layout (QKVLayout): Layout of the QKV tensors. scaling_factor (float): Scaling factor for the attention scores. dropout_probability (float): Dropout probability to apply during attention. @@ -2755,6 +2911,28 @@ def fused_attn_bwd( assert bias is None bias = jnp.zeros(0, dtype=qkv[0].dtype) + if softmax_offset is None: + assert softmax_type != AttnSoftmaxType.LEARNABLE_SOFTMAX, f"Unknown {softmax_type=}" + if softmax_type == AttnSoftmaxType.OFF_BY_ONE_SOFTMAX: + num_heads = qkv[0].shape[-2] + # Create tensor [1, h, 1, 1] filled with zeros + softmax_offset = jnp.zeros((1, num_heads, 1, 1), dtype=jnp.float32) + # Shard by heads dimension + softmax_offset = with_sharding_constraint_by_logical_axes( + softmax_offset, (None, HEAD_AXES, None, None) + ) + elif softmax_type == AttnSoftmaxType.VANILLA_SOFTMAX: + softmax_offset = jnp.zeros(0, dtype=jnp.float32) + else: + raise NotImplementedError(f"Unknown {softmax_type=}") + else: + softmax_offset = softmax_offset.astype(jnp.float32) + # Shard by heads dimension if not VANILLA_SOFTMAX + if softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + softmax_offset = with_sharding_constraint_by_logical_axes( + softmax_offset, (None, HEAD_AXES, None, None) + ) + # TODO(KshitijLakhani): Add a check for cuDNN version when determinism does get supported on # sm100+ compute_capabilities = get_all_device_compute_capability() @@ -2767,6 +2945,7 @@ def fused_attn_bwd( attn_bias_type=attn_bias_type, attn_mask_type=attn_mask_type, qkv_layout=qkv_layout, + softmax_type=softmax_type, scaling_factor=scaling_factor, dropout_probability=dropout_probability, is_training=is_training, @@ -2788,9 +2967,10 @@ def fused_attn_bwd( primitive = FusedRingAttnBwdPrimitive.outer_primitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - *qkv_grads, bias_grad = primitive.bind( + *qkv_grads, bias_grad, softmax_offset_grad = primitive.bind( *qkv_for_primitive, bias, + softmax_offset, softmax_aux, rng_state, output, @@ -2798,4 +2978,4 @@ def fused_attn_bwd( *seq_desc_flatten, config=fused_config, ) - return tuple(qkv_grads[: len(qkv)]), bias_grad + return tuple(qkv_grads[: len(qkv)]), bias_grad, softmax_offset_grad diff --git a/transformer_engine/jax/cpp_extensions/softmax.py b/transformer_engine/jax/cpp_extensions/softmax.py index 575a2dd3ab..6d8b24b07d 100644 --- a/transformer_engine/jax/cpp_extensions/softmax.py +++ b/transformer_engine/jax/cpp_extensions/softmax.py @@ -11,10 +11,11 @@ import jax.numpy as jnp from jax import dtypes, ffi from jax.sharding import PartitionSpec, NamedSharding +from .attention import AttnSoftmaxType from .base import BasePrimitive, register_primitive from .misc import get_padded_spec, check_valid_batch_dims -from ..softmax import SoftmaxType +from ..softmax import SoftmaxFusionType __all__ = [ @@ -32,7 +33,8 @@ def is_softmax_kernel_available( - softmax_type: SoftmaxType, + softmax_fusion_type: SoftmaxFusionType, + softmax_type: AttnSoftmaxType, batch: int, heads: int, q_seqlen: int, @@ -40,15 +42,18 @@ def is_softmax_kernel_available( dtype: jnp.dtype, ): """check softmax available""" - if softmax_type is SoftmaxType.SCALED: + if softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: + return False + + if softmax_fusion_type is SoftmaxFusionType.SCALED: return ScaledSoftmaxFwdPrimitive.is_kernel_available( batch, heads, q_seqlen, k_seqlen, dtype ) - if softmax_type is SoftmaxType.SCALED_MASKED: + if softmax_fusion_type is SoftmaxFusionType.SCALED_MASKED: return ScaledMaskedSoftmaxFwdPrimitive.is_kernel_available( batch, heads, q_seqlen, k_seqlen, dtype ) - if softmax_type is SoftmaxType.SCALED_UPPER_TRIANG_MASKED: + if softmax_fusion_type is SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED: return ScaledUpperTriangMaskedSoftmaxFwdPrimitive.is_kernel_available( batch, heads, q_seqlen, k_seqlen, dtype ) @@ -792,26 +797,77 @@ def shardy_sharding_rule(*args): register_primitive(ScaledUpperTriangMaskedSoftmaxBwdPrimitive) -def jax_scaled_softmax(logits: jnp.ndarray, scale_factor: float): +def jax_scaled_softmax( + logits: jnp.ndarray, scale_factor: float, softmax_offset: jnp.ndarray | float | None = None +): """ JAX based implementation of scaled softmax """ + if softmax_offset is not None: + return jax_general_softmax(scale_factor * logits, offset=softmax_offset) return jax.nn.softmax(scale_factor * logits) -def jax_scaled_masked_softmax(logits: jnp.ndarray, mask: jnp.ndarray, scale_factor: float): +def jax_scaled_masked_softmax( + logits: jnp.ndarray, + mask: jnp.ndarray, + scale_factor: float, + softmax_offset: jnp.ndarray | float | None = None, +): """ JAX based implementation of scaled and masked softmax """ + if softmax_offset is not None: + return jax_general_softmax(logits * scale_factor, offset=softmax_offset, where=mask != 1) return jax.nn.softmax(logits * scale_factor, where=mask != 1) -def jax_scaled_upper_triang_masked_softmax(logits: jnp.ndarray, scale_factor: float): +def jax_scaled_upper_triang_masked_softmax( + logits: jnp.ndarray, scale_factor: float, softmax_offset: jnp.ndarray | float | None = None +): """ JAX based implementation of scaled and upper triangle masked softmax """ mask = 1 - jnp.tril(jnp.ones_like(logits)) - return jax_scaled_masked_softmax(logits, mask, scale_factor) + return jax_scaled_masked_softmax(logits, mask, scale_factor, softmax_offset) + + +def jax_general_softmax( + x: jnp.ndarray, + axis: int = -1, + where: jnp.ndarray | None = None, + initial: jnp.ndarray = -jnp.inf, + offset: jnp.ndarray | float | None = None, +) -> jnp.ndarray: + """ + JAX based implementation of general softmax with optional masking and offset. + """ + # Compute max of x + x_max = jnp.max(x, axis, where=where, initial=initial, keepdims=True) + + if offset is not None: + # Cast offset to x.dtype to prevent type promotion + if isinstance(offset, (int, float)): + offset = jnp.array(offset, dtype=x.dtype) + else: + offset = offset.astype(x.dtype) + + # Include offset in max: x_max = max(x_max, offset) + # This is equivalent to computing max over [x..., offset] + x_max = jnp.maximum(x_max, offset) + + x_safe = x if where is None else jnp.where(where, x, initial) + unnormalized = jnp.exp(x_safe - x_max) + denominator = jnp.sum(unnormalized, axis, where=where, keepdims=True) + + if offset is not None: + # Add exp(offset - x_max) to denominator + denominator = denominator + jnp.exp(offset - x_max) + + result = unnormalized / denominator + if where is not None: + result = jnp.where(where, result, 0) + return result def scaled_softmax_fwd(logits: jnp.ndarray, scale_factor: float) -> jnp.ndarray: diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index c1c7e0d665..75d22fbf53 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -108,28 +108,28 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedAttnBackwardHandler); -NVTE_Fused_Attn_Backend GetFusedAttnBackend(bool is_training, DType q_dtype, DType kv_dtype, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, float dropout_probability, - size_t q_num_heads, size_t kv_num_heads, - size_t q_max_seqlen, size_t kv_max_seqlen, - size_t qk_head_dim, size_t v_head_dim, - int64_t window_size_left, int64_t window_size_right); +NVTE_Fused_Attn_Backend GetFusedAttnBackend( + bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, + size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, + int64_t window_size_right); pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_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); + 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); pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_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); + 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); // GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmHandler); diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index ac7eba5c87..a834273035 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -11,14 +11,12 @@ namespace transformer_engine { namespace jax { -NVTE_Fused_Attn_Backend GetFusedAttnBackend(bool is_training, DType q_dtype, DType kv_dtype, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, float dropout_probability, - size_t q_attn_heads, size_t kv_attn_heads, - size_t q_max_seqlen, size_t kv_max_seqlen, - size_t qk_head_dim, size_t v_head_dim, - int64_t window_size_left, int64_t window_size_right) { - NVTE_Softmax_Type softmax_type = NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX; +NVTE_Fused_Attn_Backend GetFusedAttnBackend( + bool is_training, DType q_dtype, DType kv_dtype, NVTE_QKV_Layout qkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, + size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, + int64_t window_size_right) { auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, @@ -39,7 +37,8 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t const size_t kv_max_seqlen, DType dtype, NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend, void *softmax_buf, void *rng_state_buf = nullptr, - void *bias_buf = nullptr) { + 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; @@ -67,10 +66,12 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t softmax_aux_data.shape.data[3] = 1; // {B,H,Qs,Ks} -> {B,H,Qs,1} softmax_aux_data.dtype = static_cast(DType::kFloat32); + int size = 2; // Start at 2 (we have softmax and rng_state at indices 0, 1) + // include bias if enabled if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) { - tensor_pack->size = 3; - NVTETensor &bias_aux = tensor_pack->tensors[2]; + NVTETensor &bias_aux = tensor_pack->tensors[size]; + size++; NVTEBasicTensor bias_aux_data; bias_aux_data.data_ptr = bias_buf; bias_aux_data.shape.ndim = 4; @@ -81,6 +82,24 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t bias_aux_data.dtype = static_cast(dtype); nvte_set_tensor_param(&bias_aux, kNVTERowwiseData, &bias_aux_data); } + + // include softmax_offset if provided + if (softmax_offset_buf != nullptr) { + NVTETensor &softmax_offset_aux = tensor_pack->tensors[size]; + size++; + NVTEBasicTensor softmax_offset_aux_data; + softmax_offset_aux_data.data_ptr = softmax_offset_buf; + softmax_offset_aux_data.shape.ndim = 4; + softmax_offset_aux_data.shape.data[0] = 1; + softmax_offset_aux_data.shape.data[1] = attn_heads; + softmax_offset_aux_data.shape.data[2] = 1; + softmax_offset_aux_data.shape.data[3] = 1; + softmax_offset_aux_data.dtype = static_cast(DType::kFloat32); + nvte_set_tensor_param(&softmax_offset_aux, kNVTERowwiseData, &softmax_offset_aux_data); + } + + // Set final size + tensor_pack->size = size; } nvte_set_tensor_param(&softmax_aux, kNVTERowwiseData, &softmax_aux_data); } @@ -98,14 +117,16 @@ void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_ const size_t bias_heads, const size_t q_max_seqlen, const size_t kv_max_seqlen, DType dtype, NVTE_Fused_Attn_Backend backend, void *softmax_buf, - void *rng_state_buf, void *bias_buf) { + void *rng_state_buf, void *bias_buf, + void *softmax_offset_buf = nullptr) { // Backward calls put everything into the tensor pack for every backend // so we set dummy bias_type and backend choices here to follow the correct code path auto dummy_bias_type = NVTE_Bias_Type::NVTE_POST_SCALE_BIAS; auto dummy_backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; PrepareFusedAttnForwardAuxTensors(tensor_pack, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type, - dummy_backend, softmax_buf, rng_state_buf, bias_buf); + 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) { @@ -121,8 +142,9 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_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) { + 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) { auto q_shape = 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}; @@ -141,7 +163,6 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); auto dummy_softmax_offset_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kFloat32); - NVTE_Softmax_Type softmax_type = NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX; NVTETensorPack aux_output_tensors; nvte_tensor_pack_create(&aux_output_tensors); @@ -208,18 +229,21 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( auto layout_group = nvte_get_qkv_layout_group(qkv_layout); static void FusedAttnForwardImpl( - cudaStream_t stream, void *q, void *k, void *v, void *bias, void *seed, void *q_cu_seqlens, - void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output, void *softmax_aux, - void *rng_state, void *workspace, size_t input_batch, size_t bias_batch, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, - size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, size_t wkspace_size, - float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, - bool is_training, bool deterministic, int64_t window_size_left, int64_t window_size_right) { + cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, void *seed, + void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output, + void *softmax_aux, void *rng_state, void *workspace, size_t input_batch, size_t bias_batch, + size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, + size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, + size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, + DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, + int64_t window_size_left, int64_t window_size_right) { FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ auto bias_tensor = TensorWrapper(bias, bias_shape, dtype); + auto softmax_offset_tensor = + TensorWrapper(softmax_offset, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); if (is_ragged) { auto output_size = input_batch * q_max_seqlen * attn_heads * v_head_dim; @@ -238,10 +262,6 @@ static void FusedAttnForwardImpl( /* Prepare RNG state */ auto rng_state_tensor = TensorWrapper(rng_state, std::vector{2}, DType::kInt64); - auto dummy_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1}, DType::kFloat32); - NVTE_Softmax_Type softmax_type = NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX; - auto backend = nvte_get_fused_attn_backend( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, @@ -254,7 +274,7 @@ static void FusedAttnForwardImpl( nvte_tensor_pack_create(&aux_output_tensors); PrepareFusedAttnForwardAuxTensors(&aux_output_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, bias_type, - backend, softmax_aux); + backend, softmax_aux, softmax_offset); /* Call the underlying NVTE API */ auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); @@ -303,7 +323,7 @@ static void FusedAttnForwardImpl( nvte_fused_attn_fwd( q_tensor.data(), k_tensor.data(), v_tensor.data(), bias_tensor.data(), - dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, + softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, 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, @@ -332,6 +352,8 @@ static void FusedAttnForwardImpl( static_cast(get_attr_value(attrs, "bias_type")); \ NVTE_Mask_Type mask_type = \ static_cast(get_attr_value(attrs, "mask_type")); \ + NVTE_Softmax_Type softmax_type = \ + static_cast(get_attr_value(attrs, "softmax_type")); \ NVTE_QKV_Layout qkv_layout = \ static_cast(get_attr_value(attrs, "qkv_layout")); \ bool is_training = get_attr_value(attrs, "is_training"); \ @@ -342,7 +364,8 @@ static void FusedAttnForwardImpl( DType wkspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type()); Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, - Buffer_Type v_buf, Buffer_Type bias_buf, Buffer_Type seed_buf, + Buffer_Type v_buf, Buffer_Type bias_buf, + Buffer_Type softmax_offset_buf, Buffer_Type seed_buf, Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf, Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type _unused_args, Result_Type output_buf, @@ -352,15 +375,15 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty FusedAttnForwardImpl( stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(), - bias_buf.untyped_data(), seed_buf.untyped_data(), q_cu_seqlens_buf.untyped_data(), - kv_cu_seqlens_buf.untyped_data(), is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, + bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), seed_buf.untyped_data(), + q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), + is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, output_buf->untyped_data(), softmax_aux_buf->untyped_data(), rng_state_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, max_segments_per_seq, wkspace_size, scaling_factor, - dropout_probability, bias_type, mask_type, qkv_layout, dtype, wkspace_dtype, is_training, - deterministic, window_size_left, window_size_right); - + dropout_probability, bias_type, mask_type, softmax_type, qkv_layout, dtype, wkspace_dtype, + is_training, deterministic, window_size_left, window_size_right); return ffi_with_cuda_error_check(); } @@ -371,6 +394,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI, .Arg() // k .Arg() // v .Arg() // bias + .Arg() // softmax_offset .Arg() // seed_buf .Arg() // q_cu_seqlens .Arg() // kv_cu_seqlens @@ -388,9 +412,9 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_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) { + 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) { auto q_shape = 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); @@ -425,9 +449,14 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( // For cuDNN < 9.3.0, it requires to run all possible seqlens to address act_seqlen = 0 min_num_segments = input_batch * max_segments_per_seq; } - auto dummy_d_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1}, DType::kFloat32); - NVTE_Softmax_Type softmax_type = NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX; + + TensorWrapper dummy_d_softmax_offset_tensor; + if (softmax_type == NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX || + softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX) { + dummy_d_softmax_offset_tensor = + TensorWrapper(nullptr, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); + } + for (auto num_segments = min_num_segments; num_segments <= max_num_segments; ++num_segments) { // the last one is the largest which will be the returned workspace size auto q_cu_seqlens_tensor = @@ -457,15 +486,16 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( } static void FusedAttnBackwardImpl( - cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_aux, void *rng_state, - void *output, void *doutput, void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, - void *k_seq_offsets, void *dq, void *dk, void *dv, void *dbias, void *workspace, - size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, - size_t attn_heads, size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, - size_t v_head_dim, size_t max_segments_per_seq, size_t wkspace_size, float scaling_factor, - float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, bool is_training, - bool deterministic, int64_t window_size_left, int64_t window_size_right) { + cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, + void *softmax_aux, void *rng_state, void *output, void *doutput, void *q_cu_seqlens, + void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *dq, void *dk, void *dv, + void *dbias, void *dsoftmax_offset, void *workspace, size_t input_batch, size_t bias_batch, + size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, + size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, + size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, + DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, + int64_t window_size_left, int64_t window_size_right) { FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ @@ -476,9 +506,13 @@ static void FusedAttnBackwardImpl( /* Output tensors */ auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); // not used in F16 auto dbias_tensor = TensorWrapper(dbias, bias_shape, dtype); - auto dummy_d_softmax_offset_tensor = - TensorWrapper(nullptr, std::vector{1}, DType::kFloat32); - NVTE_Softmax_Type softmax_type = NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX; + + TensorWrapper dsoftmax_offset_tensor; + if (softmax_type == NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX || + softmax_type == NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX) { + dsoftmax_offset_tensor = + TensorWrapper(dsoftmax_offset, std::vector{1, attn_heads, 1, 1}, DType::kFloat32); + } /* Auxiliary tensors (propagated from the forward pass) */ NVTETensorPack aux_input_tensors; @@ -490,7 +524,7 @@ static void FusedAttnBackwardImpl( false, false); PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, - softmax_aux, rng_state, bias); + softmax_aux, rng_state, bias, softmax_offset); /* Call the underly NVTE API */ // Prepare Q, K, V pointers and shapes based on layout @@ -564,7 +598,7 @@ static void FusedAttnBackwardImpl( 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(), + 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, deterministic, false, workspace_tensor.data(), stream); @@ -574,26 +608,29 @@ static void FusedAttnBackwardImpl( Error_Type FusedAttnBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, Buffer_Type v_buf, Buffer_Type bias_buf, - Buffer_Type softmax_aux_buf, Buffer_Type rng_state_buf, - Buffer_Type output_buf, Buffer_Type doutput_buf, - Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf, - Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, - Variadic_Buffer_Type _unused_args, Result_Type dq_buf, - Result_Type dk_buf, Result_Type dv_buf, Result_Type dbias_buf, + Buffer_Type softmax_offset_buf, Buffer_Type softmax_aux_buf, + Buffer_Type rng_state_buf, Buffer_Type output_buf, + Buffer_Type doutput_buf, Buffer_Type q_cu_seqlens_buf, + Buffer_Type kv_cu_seqlens_buf, Buffer_Type q_seq_offsets_buf, + Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type _unused_args, + Result_Type dq_buf, Result_Type dk_buf, Result_Type dv_buf, + Result_Type dbias_buf, Result_Type dsoftmax_offset_buf, Result_Type workspace_buf, Dictionary attrs) { FUSED_ATTN_FFI_GET_ATTRS; FusedAttnBackwardImpl( stream, q_buf.untyped_data(), k_buf.untyped_data(), v_buf.untyped_data(), - bias_buf.untyped_data(), softmax_aux_buf.untyped_data(), rng_state_buf.untyped_data(), - output_buf.untyped_data(), doutput_buf.untyped_data(), q_cu_seqlens_buf.untyped_data(), - kv_cu_seqlens_buf.untyped_data(), is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, + bias_buf.untyped_data(), softmax_offset_buf.untyped_data(), softmax_aux_buf.untyped_data(), + rng_state_buf.untyped_data(), output_buf.untyped_data(), doutput_buf.untyped_data(), + q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), + is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, dq_buf->untyped_data(), dk_buf->untyped_data(), dv_buf->untyped_data(), dbias_buf->untyped_data(), - workspace_buf->untyped_data(), input_batch, bias_batch, q_max_seqlen, kv_max_seqlen, - attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, max_segments_per_seq, - wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, qkv_layout, dtype, - wkspace_dtype, is_training, deterministic, window_size_left, window_size_right); + dsoftmax_offset_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, + q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, + max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, + softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, deterministic, window_size_left, + window_size_right); return ffi_with_cuda_error_check(); } @@ -605,6 +642,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnBackwardHandler, FusedAttnBackwardFFI, .Arg() // k .Arg() // v .Arg() // bias + .Arg() // softmax_offset .Arg() // softmax_aux .Arg() // rng_state .Arg() // output @@ -618,6 +656,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnBackwardHandler, FusedAttnBackwardFFI, .Ret() // dk .Ret() // dv .Ret() // dbias + .Ret() // dsoftmax_offset .Ret() // workspace .Attrs(), FFI_CudaGraph_Traits); diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index e57d07872e..9784565cc9 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -142,6 +142,11 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("NVTE_BSHD", NVTE_QKV_Format::NVTE_BSHD) .value("NVTE_THD", NVTE_QKV_Format::NVTE_THD); + pybind11::enum_(m, "NVTE_Softmax_Type", pybind11::module_local()) + .value("NVTE_VANILLA_SOFTMAX", NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) + .value("NVTE_OFF_BY_ONE_SOFTMAX", NVTE_Softmax_Type::NVTE_OFF_BY_ONE_SOFTMAX) + .value("NVTE_LEARNABLE_SOFTMAX", NVTE_Softmax_Type::NVTE_LEARNABLE_SOFTMAX); + pybind11::enum_(m, "NVTE_Activation_Type", pybind11::module_local()) .value("GELU", NVTE_Activation_Type::GELU) .value("GEGLU", NVTE_Activation_Type::GEGLU) diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index b5f1590229..19e4c57ce2 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -7,6 +7,7 @@ from functools import reduce import operator from typing import Any, Callable, Iterable, List, Sequence, Tuple, Union, NewType, Optional +import warnings import numpy as np import jax.numpy as jnp @@ -23,8 +24,9 @@ from ..layernorm_dense import layernorm_dense from ..layernorm_mlp import layernorm_mlp from ..activation import activation -from ..softmax import softmax, SoftmaxType +from ..softmax import softmax, SoftmaxFusionType from ..sharding import with_sharding_constraint_by_logical_axes +from ..attention import AttnSoftmaxType from ..cpp_extensions import ( is_softmax_kernel_available, jax_scaled_softmax, @@ -171,15 +173,20 @@ class Softmax(nn.Module): # pylint: disable=too-few-public-methods ---------- scale_factor : float, default = 1.0 Scalar for the input to softmax. - softmax_type : SoftmaxType, default = SoftmaxType.SCALED + softmax_fusion_type : SoftmaxFusionType, default = SoftmaxFusionType.SCALED + Indicate the type of softmax. + softmax_type : AttnSoftmaxType, default = AttnSoftmaxType.VANILLA_SOFTMAX Indicate the type of softmax. """ scale_factor: float = 1.0 - softmax_type: SoftmaxType = SoftmaxType.SCALED + softmax_fusion_type: SoftmaxFusionType = SoftmaxFusionType.SCALED + softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX @nn.compact - def __call__(self, inputs: Array, mask: Array = None, bias: Array = None) -> jnp.ndarray: + def __call__( + self, inputs: Array, mask: Array = None, bias: Array = None, softmax_offset: Array = None + ) -> jnp.ndarray: batch = inputs.shape[0] heads = inputs.shape[1] q_seqlen = inputs.shape[2] @@ -187,33 +194,52 @@ def __call__(self, inputs: Array, mask: Array = None, bias: Array = None) -> jnp input_dtype = inputs.dtype logits = inputs + if softmax_offset is not None: + assert self.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX + if self.softmax_type == AttnSoftmaxType.OFF_BY_ONE_SOFTMAX: + softmax_offset = 0.0 + # use primitives if is_softmax_kernel_available( - self.softmax_type, batch, heads, q_seqlen, k_seqlen, input_dtype + self.softmax_fusion_type, + self.softmax_type, + batch, + heads, + q_seqlen, + k_seqlen, + input_dtype, ): if bias is not None: logits = logits + bias.astype(input_dtype) mask_ = mask - if self.softmax_type is not SoftmaxType.SCALED_MASKED: + if self.softmax_fusion_type is not SoftmaxFusionType.SCALED_MASKED: mask_ = None - outputs = softmax(logits, mask_, self.scale_factor, self.softmax_type) + outputs = softmax(logits, mask_, self.scale_factor, self.softmax_fusion_type) # use default jax based implementation else: + warnings.warn( + "Using unfused JAX softmax implementation instead of TE fused primitives. ", + UserWarning, + stacklevel=2, + ) + if bias is not None: logits = logits + bias.astype(input_dtype) - if self.softmax_type is SoftmaxType.SCALED: - outputs = jax_scaled_softmax(logits, self.scale_factor) - elif self.softmax_type is SoftmaxType.SCALED_MASKED: - outputs = jax_scaled_masked_softmax(logits, mask, self.scale_factor) - elif self.softmax_type is SoftmaxType.SCALED_UPPER_TRIANG_MASKED: - outputs = jax_scaled_upper_triang_masked_softmax(logits, self.scale_factor) + if self.softmax_fusion_type is SoftmaxFusionType.SCALED: + outputs = jax_scaled_softmax(logits, self.scale_factor, softmax_offset) + elif self.softmax_fusion_type is SoftmaxFusionType.SCALED_MASKED: + outputs = jax_scaled_masked_softmax(logits, mask, self.scale_factor, softmax_offset) + elif self.softmax_fusion_type is SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED: + outputs = jax_scaled_upper_triang_masked_softmax( + logits, self.scale_factor, softmax_offset + ) else: raise ValueError( - f"Unsupported softmax type: {self.softmax_type}. softmax_type must be [SCALED," - " SCALED_MASKED, SCALED_UPPER_TRIANG_MASKED]" + f"Unsupported softmax fusion: {self.softmax_fusion_type}. softmax_fusion_type" + " must be [SCALED, SCALED_MASKED, SCALED_UPPER_TRIANG_MASKED]" ) assert input_dtype == outputs.dtype return outputs diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index d096e7997c..edf5f37227 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -23,11 +23,17 @@ from .module import DenseGeneral, LayerNormDenseGeneral, LayerNormMLP from .module import LayerNorm, Softmax -from ..attention import AttnBiasType, AttnMaskType, QKVLayout, SequenceDescriptor +from ..attention import ( + AttnBiasType, + AttnMaskType, + AttnSoftmaxType, + QKVLayout, + SequenceDescriptor, +) from ..attention import is_fused_attn_kernel_available, make_swa_mask, canonicalize_attn_mask_type from ..attention import fused_attn from ..attention import CPStrategy -from ..softmax import SoftmaxType +from ..softmax import SoftmaxFusionType from ..sharding import num_of_devices from ..sharding import get_sharding_map_logic_axis_to_mesh_axis from ..sharding import with_sharding_constraint_by_logical_axes @@ -120,6 +126,7 @@ class _UnfusedDotProductAttention(nn.Module): # pylint: disable=too-few-public- scale_factor: Optional[float] = None transpose_batch_sequence: bool = True window_size: Optional[Tuple[int, int]] = None + softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX @nn.compact def __call__( @@ -145,6 +152,22 @@ def __call__( input_dtype = query.dtype + # Infer number of attention heads from query shape + # query shape: [..., h, d] where h is num_attention_heads + num_attention_heads = query.shape[-2] + + # Initialize softmax_offset for learnable softmax + # Note: OFF_BY_ONE_SOFTMAX is handled internally by the Softmax module + softmax_offset = None + if self.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX: + # For learnable softmax, create a learnable parameter with proper sharding and shape (1, h, 1, 1) + softmax_offset = self.param( + "softmax_offset", + nn.with_logical_partitioning(nn.initializers.zeros, (None, HEAD_AXES, None, None)), + (1, num_attention_heads, 1, 1), + jnp.float32, + ) + if self.scale_factor is None: scale_factor = 1.0 / sqrt(query.shape[-1]) else: @@ -213,8 +236,8 @@ def apply_swa_mask(original_mask: Array) -> Array: new_mask = jnp.where(original_mask == 0, swa_mask, original_mask) return new_mask - def convert_to_softmax_type(attn_mask_type, mask): - """Convert the attn_mask_type to SoftmaxType""" + def convert_to_softmax_fusion_type(attn_mask_type, mask): + """Convert the attn_mask_type to SoftmaxFusionType""" # mask is ignored for no_mask and causal_mask without sliding window if attn_mask_type == AttnMaskType.NO_MASK: mask = None @@ -224,21 +247,23 @@ def convert_to_softmax_type(attn_mask_type, mask): mask = apply_swa_mask(mask) # Currently cuDNN backend only supports SWA for causal/padding_causal, follow this if mask is not None: - return SoftmaxType.SCALED_MASKED, mask + return SoftmaxFusionType.SCALED_MASKED, mask if attn_mask_type is AttnMaskType.CAUSAL_MASK: - return SoftmaxType.SCALED_UPPER_TRIANG_MASKED, mask + return SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED, mask if attn_mask_type is AttnMaskType.NO_MASK: - return SoftmaxType.SCALED, mask + return SoftmaxFusionType.SCALED, mask raise ValueError( f"Unsupported {attn_mask_type=}, supported attn_mask_type=" "{'no_mask', 'padding', 'causal', 'padding_causal', 'causal_padding'}" ) - softmax_type, mask = convert_to_softmax_type(self.attn_mask_type, mask) + softmax_fusion_type, mask = convert_to_softmax_fusion_type(self.attn_mask_type, mask) - attn_weights = Softmax(softmax_type=softmax_type, scale_factor=fused_scale_factor)( - attn_weights, mask, bias - ).astype(input_dtype) + attn_weights = Softmax( + softmax_fusion_type=softmax_fusion_type, + softmax_type=self.softmax_type, + scale_factor=fused_scale_factor, + )(attn_weights, mask, bias, softmax_offset=softmax_offset).astype(input_dtype) if is_gqa: attn_weights = attn_weights.reshape(attn_weights_with_groups_shape) @@ -279,6 +304,7 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me context_parallel_axis: str = "" context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT context_checkpoint_name: str = "context" + softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX @nn.compact def __call__( @@ -303,6 +329,17 @@ def __call__( scale_factor = self.scale_factor del self.scale_factor + num_attention_heads = query.shape[-2] + softmax_offset = None + if self.softmax_type == AttnSoftmaxType.LEARNABLE_SOFTMAX: + # For learnable softmax, create a learnable parameter with proper sharding and shape (1, h, 1, 1) + softmax_offset = self.param( + "softmax_offset", + nn.with_logical_partitioning(nn.initializers.zeros, (None, HEAD_AXES, None, None)), + (1, num_attention_heads, 1, 1), + jnp.float32, + ) + if self.qkv_layout.is_qkvpacked(): """qkvpacked format, treat query: qkvpacked tensor, shape = [..., 3, h, d] @@ -320,6 +357,7 @@ def __call__( attn_mask_type=self.attn_mask_type, attn_bias_type=self.attn_bias_type, qkv_layout=self.qkv_layout, + softmax_type=self.softmax_type, scaling_factor=scale_factor, dropout_probability=self.attention_dropout, is_training=not deterministic, @@ -329,6 +367,7 @@ def __call__( context_parallel_axis=self.context_parallel_axis, context_parallel_strategy=self.context_parallel_strategy, context_checkpoint_name=self.context_checkpoint_name, + softmax_offset=softmax_offset, ) elif self.qkv_layout.is_kvpacked(): """kvpacked format, treat @@ -348,6 +387,7 @@ def __call__( attn_mask_type=self.attn_mask_type, attn_bias_type=self.attn_bias_type, qkv_layout=self.qkv_layout, + softmax_type=self.softmax_type, scaling_factor=scale_factor, dropout_probability=self.attention_dropout, is_training=not deterministic, @@ -357,6 +397,7 @@ def __call__( context_parallel_axis=self.context_parallel_axis, context_parallel_strategy=self.context_parallel_strategy, context_checkpoint_name=self.context_checkpoint_name, + softmax_offset=softmax_offset, ) elif self.qkv_layout.is_separate(): if self.transpose_batch_sequence: @@ -371,6 +412,7 @@ def __call__( attn_mask_type=self.attn_mask_type, attn_bias_type=self.attn_bias_type, qkv_layout=self.qkv_layout, + softmax_type=self.softmax_type, scaling_factor=scale_factor, dropout_probability=self.attention_dropout, is_training=not deterministic, @@ -380,6 +422,7 @@ def __call__( context_parallel_axis=self.context_parallel_axis, context_parallel_strategy=self.context_parallel_strategy, context_checkpoint_name=self.context_checkpoint_name, + softmax_offset=softmax_offset, ) else: raise ValueError(f"Unsupported {self.qkv_layout=}.") @@ -514,6 +557,17 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods context_parallel_axis (str): The name of the context parallel axis. context_parallel_strategy (CPStrategy): The strategy of context parallel. 0: DEFAULT, 1: ALL_GATHER, 2: RING. context_checkpoint_name (str): The name of the context checkpoint in the forward pass of fused attention. + softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' + softmax type as described in this paper: + `Efficient Streaming Language Models with Attention Sinks + `_. + For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], + 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), + 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and + 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), + where alpha is a learnable parameter in shape [h]. + 'off-by-one' and 'learnable' softmax types are also called sink attention + ('zero sink' and 'learnable sink'). Optimization parameters ----------------------- @@ -539,6 +593,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods context_parallel_axis: str = "" context_parallel_strategy: str = "DEFAULT" context_checkpoint_name: str = "context" + softmax_type: str = "vanilla" @nn.compact def __call__( @@ -595,6 +650,7 @@ def __call__( attn_bias_type = AttnBiasType[self.attn_bias_type.upper()] attn_mask_type = canonicalize_attn_mask_type(self.attn_mask_type) qkv_layout = QKVLayout[self.qkv_layout.upper()] + softmax_type = AttnSoftmaxType.from_str(self.softmax_type) del self.attn_bias_type, self.attn_mask_type, self.qkv_layout if attn_bias_type == AttnBiasType.NO_BIAS: @@ -626,6 +682,7 @@ def __call__( qkv_layout, attn_bias_type, attn_mask_type, + softmax_type, self.attention_dropout, self.num_attention_heads, self.num_gqa_groups, @@ -702,6 +759,7 @@ def __call__( scale_factor=scale_factor, transpose_batch_sequence=self.transpose_batch_sequence, window_size=self.window_size, + softmax_type=softmax_type, )( query, key, @@ -726,6 +784,7 @@ def __call__( context_parallel_axis=self.context_parallel_axis, context_parallel_strategy=context_parallel_strategy, context_checkpoint_name=self.context_checkpoint_name, + softmax_type=softmax_type, )( query, key, @@ -1005,6 +1064,17 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods Deprecated. Please refer `fuse_qkv_params` window_size: Optional[Tuple[int, int]], default = None Sliding window size. Default value is no sliding window. + softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' + softmax type as described in this paper: + `Efficient Streaming Language Models with Attention Sinks + `_. + For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], + 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), + 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and + 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), + where alpha is a learnable parameter in shape [h]. + 'off-by-one' and 'learnable' softmax types are also called sink attention + ('zero sink' and 'learnable sink'). """ head_dim: int @@ -1036,6 +1106,7 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods scaled_query_init: bool = True float32_logits: bool = False window_size: Optional[Tuple[int, int]] = None + softmax_type: str = "vanilla" # Deprecated parameters num_heads: Optional[int] = None @@ -1440,6 +1511,7 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): scale_factor=scale_factor, transpose_batch_sequence=self.transpose_batch_sequence, window_size=self.window_size, + softmax_type=self.softmax_type, )(*dpa_args, mask, bias, deterministic=deterministic) x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3])) @@ -1721,6 +1793,18 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods Whether to enable sequence parallelism to operations except dot. window_size: Optional[Tuple[int, int]], default = None Sliding window size. Default value is no sliding window. + softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' + Softmax type as described in this paper: + `Efficient Streaming Language Models with Attention Sinks + `_. + For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], + 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), + 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and + 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), + where alpha is a learnable parameter in shape [h]. + 'off-by-one' and 'learnable' softmax types are also called sink attention + ('zero sink' and 'learnable sink'). + Only supported for fused attention backend. Optimization parameters ----------------------- @@ -1786,6 +1870,7 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods scale_attn_logits: bool = False scaled_query_init: bool = True window_size: Optional[Tuple[int, int]] = None + softmax_type: str = "vanilla" def __post_init__(self): if self.mha_kernel_init is None: @@ -1946,6 +2031,7 @@ def generate_batch_seqlen_logical_axes(is_shared_seq=None): bias_init=self.bias_init, name=mha_name, window_size=self.window_size, + softmax_type=self.softmax_type, )(inputs, inputs, attention_mask, attn_bias, deterministic=deterministic, decode=decode) def hidden_dropout(x, deterministic): @@ -2024,6 +2110,7 @@ def hidden_dropout(x, deterministic): bias_init=self.bias_init, name="encoder_decoder_attention", window_size=self.window_size, + softmax_type=self.softmax_type, )(x, encoded, encoder_decoder_mask, deterministic=deterministic) y = with_sharding_constraint_by_logical_axes( diff --git a/transformer_engine/jax/softmax.py b/transformer_engine/jax/softmax.py index 9b32002388..24fca6bc71 100644 --- a/transformer_engine/jax/softmax.py +++ b/transformer_engine/jax/softmax.py @@ -12,8 +12,8 @@ from . import cpp_extensions as tex -class SoftmaxType(Enum): - """SoftmaxType.""" +class SoftmaxFusionType(Enum): + """SoftmaxFusionType.""" SCALED = "scaled" SCALED_MASKED = "scaled_masked" @@ -24,27 +24,27 @@ def softmax( logits: jnp.ndarray, mask: Optional[jnp.ndarray] = None, scale_factor: Optional[float] = 1.0, - softmax_type: Optional[SoftmaxType] = SoftmaxType.SCALED, + softmax_fusion_type: Optional[SoftmaxFusionType] = SoftmaxFusionType.SCALED, ): """ Softmax wrapper """ - output = _softmax(logits, mask, scale_factor, softmax_type) + output = _softmax(logits, mask, scale_factor, softmax_fusion_type) return output @partial(jax.custom_vjp, nondiff_argnums=(2, 3)) -def _softmax(logits, mask, scale_factor, softmax_type): +def _softmax(logits, mask, scale_factor, softmax_fusion_type): - output, _ = _softmax_fwd_rule(logits, mask, scale_factor, softmax_type) + output, _ = _softmax_fwd_rule(logits, mask, scale_factor, softmax_fusion_type) return output -def _softmax_fwd_rule(logits, mask, scale_factor, softmax_type): - if softmax_type is SoftmaxType.SCALED_MASKED: +def _softmax_fwd_rule(logits, mask, scale_factor, softmax_fusion_type): + if softmax_fusion_type is SoftmaxFusionType.SCALED_MASKED: assert mask is not None output = tex.scaled_masked_softmax_fwd(logits, mask, scale_factor) - elif softmax_type is SoftmaxType.SCALED_UPPER_TRIANG_MASKED: + elif softmax_fusion_type is SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED: output = tex.scaled_upper_triang_masked_softmax_fwd(logits, scale_factor) else: output = tex.scaled_softmax_fwd(logits, scale_factor) @@ -52,12 +52,12 @@ def _softmax_fwd_rule(logits, mask, scale_factor, softmax_type): return output, (output, logits, mask) -def _softmax_bwd_rule(scale_factor, softmax_type, ctx, dz): +def _softmax_bwd_rule(scale_factor, softmax_fusion_type, ctx, dz): (softmax_output, logits, mask) = ctx - if softmax_type is SoftmaxType.SCALED_MASKED: + if softmax_fusion_type is SoftmaxFusionType.SCALED_MASKED: dgrad = tex.scaled_masked_softmax_bwd(dz, softmax_output, logits, mask, scale_factor) - elif softmax_type is SoftmaxType.SCALED_UPPER_TRIANG_MASKED: + elif softmax_fusion_type is SoftmaxFusionType.SCALED_UPPER_TRIANG_MASKED: dgrad = tex.scaled_upper_triang_masked_softmax_bwd(dz, softmax_output, logits, scale_factor) else: dgrad = tex.scaled_softmax_bwd(dz, softmax_output, logits, scale_factor) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index df10fc7905..fd799957b4 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -156,7 +156,9 @@ def __init__( softmax_in_fp32: bool = True, ) -> None: super().__init__() - self.scaled_masked_softmax_fusion = bool(int(os.getenv("NVTE_MASKED_SOFTMAX_FUSION", "1"))) + self.scaled_masked_softmax_fusion_type = bool( + int(os.getenv("NVTE_MASKED_SOFTMAX_FUSION", "1")) + ) self.mask_func = mask_func self.softmax_in_fp32 = softmax_in_fp32 @@ -189,7 +191,7 @@ def is_kernel_available(self, mask: torch.Tensor, b: int, np: int, sq: int, sk: """Check FusedScaleMaskSoftmax kernel availability based on size""" attn_batches = b * np - if not self.scaled_masked_softmax_fusion: + if not self.scaled_masked_softmax_fusion_type: return False # user doesn't want to fuse if not self.input_in_float16: return False # input must be fp16 From d677a26903bc931fd06ce911f94b51792accaa11 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 18 Nov 2025 12:23:12 -0800 Subject: [PATCH 075/521] Show quickstart_jax.ipynb along with quickstart.ipynb on html documentation (#2394) Signed-off-by: tdophung --- docs/getting_started.rst | 16 ++++++++++++++++ docs/index.rst | 2 +- 2 files changed, 17 insertions(+), 1 deletion(-) create mode 100644 docs/getting_started.rst diff --git a/docs/getting_started.rst b/docs/getting_started.rst new file mode 100644 index 0000000000..2e8047763a --- /dev/null +++ b/docs/getting_started.rst @@ -0,0 +1,16 @@ +.. + Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Getting Started +=============== + +Choose your framework to get started with Transformer Engine: + +.. toctree:: + :maxdepth: 1 + + PyTorch + JAX + diff --git a/docs/index.rst b/docs/index.rst index 2c04810f4d..277259edf0 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -29,7 +29,7 @@ Transformer Engine documentation :caption: Getting Started installation - examples/quickstart.ipynb + getting_started faq .. toctree:: From e1221735c9104ee6a718c2d9a2002e587b2b97dc Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 18 Nov 2025 16:57:55 -0500 Subject: [PATCH 076/521] [PyTorch] Cache RHT device tensors properly (#2395) * Cache device tensors properly Signed-off-by: Kirthi Shankar Sivamani * Fix annotation and add test Signed-off-by: Kirthi Shankar Sivamani * skip nvfp4 test if not supported Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/distributed/test_sanity.py | 34 ++++++++++++++++++- .../pytorch/tensor/nvfp4_tensor.py | 32 +++++++++-------- 2 files changed, 50 insertions(+), 16 deletions(-) diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index 81e49f3c54..46f8a5b29a 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -7,7 +7,16 @@ import pytest import torch import transformer_engine -from transformer_engine.pytorch import DotProductAttention, TransformerLayer, Linear, GroupedLinear +from transformer_engine.pytorch import ( + DotProductAttention, + TransformerLayer, + Linear, + GroupedLinear, + NVFP4Quantizer, + autocast, + is_nvfp4_available, +) +from transformer_engine.common import recipe _current_file = pathlib.Path(__file__).resolve() sys.path.append(str(_current_file.parent.parent)) @@ -17,6 +26,8 @@ "small": ModelConfig(2, 10, 2, 16), } +nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True) + @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize( @@ -138,3 +149,24 @@ def test_current_device(model, module): assert ( tensor_device_grad == tensor_device ), "The gradient tensor should be the same as the input tensors!" + + +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +def test_nvfp4_rht_cache(): + """Ensure correct RHT cache for NVFP4.""" + + num_devices = torch.cuda.device_count() + assert num_devices > 1, "This test requires more than one GPU!" + + # Populate cache on last device. + with torch.cuda.device(num_devices - 1): + _ = NVFP4Quantizer() + + hidden_size = 128 + dtype = torch.bfloat16 + + model = Linear(hidden_size, hidden_size, params_dtype=dtype) + inp = torch.randn(hidden_size, hidden_size, device=torch.cuda.current_device(), dtype=dtype) + fp4_recipe = recipe.NVFP4BlockScaling() + with autocast(recipe=fp4_recipe): + _ = model(inp) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 652163295c..5ee9441529 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -28,9 +28,9 @@ aten = torch.ops.aten -def get_no_random_sign_vector() -> torch.Tensor: +def get_no_random_sign_vector(device: int) -> torch.Tensor: """Non-random sign vector for Hadamard transform.""" - return torch.tensor([1], dtype=torch.float32, device="cuda") + return torch.tensor([1], dtype=torch.float32, device=device) def get_sign_from_vector(vector: torch.Tensor) -> int: @@ -45,7 +45,7 @@ def get_sign_from_vector(vector: torch.Tensor) -> int: return mask.item() -def get_wgrad_sign_vector() -> torch.Tensor: +def get_wgrad_sign_vector(device: int) -> torch.Tensor: """Hard-coded random signs for Hadamard transform. https://xkcd.com/221/ @@ -54,11 +54,11 @@ def get_wgrad_sign_vector() -> torch.Tensor: return torch.tensor( [1, 1, 1, -1, 1, -1, -1, -1, -1, -1, -1, 1, -1, 1, -1, -1], dtype=torch.float32, - device="cuda", + device=device, ) -def get_hadamard_matrix(hadamard_dimension: int) -> torch.Tensor: +def get_hadamard_matrix(hadamard_dimension: int, device: int) -> torch.Tensor: """Construct a 16x16 Hadamard matrix.""" assert hadamard_dimension == 16, "Only hadamard dimension 16 is supported." hadamard_scale = 1 / math.sqrt(hadamard_dimension) @@ -83,30 +83,30 @@ def get_hadamard_matrix(hadamard_dimension: int) -> torch.Tensor: [1, -1, -1, 1, -1, 1, 1, -1, -1, 1, 1, -1, 1, -1, -1, 1], ], dtype=torch.float32, - device="cuda", + device=device, ) * hadamard_scale ) @functools.lru_cache(maxsize=None) -def get_rht_matrix(with_random_sign_mask: bool) -> torch.Tensor: +def get_rht_matrix(with_random_sign_mask: bool, device: int) -> torch.Tensor: """Construct matrix used in random Hadamard transform.""" hadamard_dimension = 16 if with_random_sign_mask: - signs = get_wgrad_sign_vector() + signs = get_wgrad_sign_vector(device=device) else: - signs = get_no_random_sign_vector() - sign_matrix = signs * torch.eye(hadamard_dimension, dtype=torch.float32, device="cuda") - rht_matrix = sign_matrix @ get_hadamard_matrix(hadamard_dimension) + signs = get_no_random_sign_vector(device=device) + sign_matrix = signs * torch.eye(hadamard_dimension, dtype=torch.float32, device=device) + rht_matrix = sign_matrix @ get_hadamard_matrix(hadamard_dimension, device=device) return rht_matrix.to(dtype=torch.bfloat16) @functools.lru_cache(maxsize=None) -def get_random_sign_mask_for_rht(with_random_sign_mask: bool) -> int: +def get_random_sign_mask_for_rht(with_random_sign_mask: bool, device: int) -> int: """Sign mask for random Hadamard transform.""" if with_random_sign_mask: - return get_sign_from_vector(get_wgrad_sign_vector()) + return get_sign_from_vector(get_wgrad_sign_vector(device=device)) return 0 @@ -152,8 +152,10 @@ def __init__( self.amax_reduction_group = amax_reduction_group self.with_2d_quantization = with_2d_quantization self.stochastic_rounding = stochastic_rounding - self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht(with_random_sign_mask) - self.rht_matrix = get_rht_matrix(with_random_sign_mask) + self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( + with_random_sign_mask, torch.cuda.current_device() + ) + self.rht_matrix = get_rht_matrix(with_random_sign_mask, torch.cuda.current_device()) def update_quantized( self, From 30c0120b34c9d2fbb78e454167beb122ad44ab35 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 18 Nov 2025 23:00:41 +0100 Subject: [PATCH 077/521] [PyTorch] Fix small errors (#2396) * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski --- tests/pytorch/distributed/run_gemm_with_overlap.py | 13 ++----------- transformer_engine/pytorch/cpu_offload.py | 2 ++ transformer_engine/pytorch/distributed.py | 6 ++++++ transformer_engine/pytorch/tensor/mxfp8_tensor.py | 4 +++- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/tests/pytorch/distributed/run_gemm_with_overlap.py b/tests/pytorch/distributed/run_gemm_with_overlap.py index df0e4a216e..073fa08117 100644 --- a/tests/pytorch/distributed/run_gemm_with_overlap.py +++ b/tests/pytorch/distributed/run_gemm_with_overlap.py @@ -24,10 +24,8 @@ MXFP8Quantizer, ) import transformer_engine.pytorch.cpp_extensions as tex -from transformer_engine.pytorch.module.base import ( - fill_userbuffers_buffer_for_all_gather, - get_cublas_workspace_size_bytes, -) +from transformer_engine.pytorch.cpp_extensions.gemm import get_cublas_workspace_size_bytes +from transformer_engine.pytorch.module.base import fill_userbuffers_buffer_for_all_gather warnings.filterwarnings("ignore", category=DeprecationWarning) warnings.filterwarnings("ignore", category=FutureWarning) @@ -417,10 +415,6 @@ def dist_print(msg, src=None, info=False, error=False, section=False, group=None std=opts.std, ) - # Allocate cuBLAS workspace - workspace_size = 3 * get_cublas_workspace_size_bytes() - workspace = torch.empty(workspace_size, dtype=torch.uint8, device="cuda") - # Gather global tensors and calculate reference result (need these first for Fp8 scales) if opts.bulk_overlap: ker_g = torch.transpose(kernel_t, 0, 1) @@ -617,7 +611,6 @@ def _fp8_gemm(): return tex.general_gemm( kernel_t_fp8, gemm_inp, - workspace, out_dtype=torch.float8_e4m3fn if opts.fp8_output else torch.bfloat16, quantization_params=out_quantizer, use_split_accumulator=te.module.base._2X_ACC_FPROP, @@ -635,7 +628,6 @@ def _fp8_gemm2(gemm1_out): return tex.general_gemm( kernel2_t_fp8, gemm2_inp, - workspace, out_dtype=torch.float8_e4m3fn if opts.fp8_output else torch.bfloat16, quantization_params=out2_quantizer, use_split_accumulator=te.module.base._2X_ACC_FPROP, @@ -648,7 +640,6 @@ def _gemm(): return tex.general_gemm( kernel_t, gemm_inp, - workspace, out_dtype=torch.bfloat16, use_split_accumulator=te.module.base._2X_ACC_FPROP, ub=ub_obj, diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index bfdee34752..241cd0e9a8 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -471,6 +471,8 @@ def fwd_step(self) -> int: """ if self.num_of_fwds in [None, self.num_layers - 1]: # reset the offload synchronizer + for layer_id in self.layer_states: + self.layer_states[layer_id].release_all_memory() self.num_of_fwds = 0 else: self.num_of_fwds += 1 diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 620ea83013..8ce54d7f64 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -948,7 +948,13 @@ def _all_gather_fp8( if isinstance(inp, Float8Tensor): dtype = inp.dtype device = inp.device + # Temporarily ensure rowwise usage for output tensor creation + # since we're gathering rowwise data, not the transpose + init_rowwise_usage = quantizer.rowwise_usage + init_columnwise_usage = quantizer.columnwise_usage + quantizer.set_usage(rowwise=True, columnwise=init_columnwise_usage) out = quantizer.make_empty(out_shape, dtype=dtype, device=device) + quantizer.set_usage(rowwise=init_rowwise_usage, columnwise=init_columnwise_usage) elif isinstance(inp, Float8Tensor): out = inp.make_like(inp, shape=out_shape) out._data = torch.empty( diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index b42b328091..cf65814656 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -134,7 +134,9 @@ def make_empty( columnwise_data = None columnwise_scale_inv = None if self.columnwise_usage: - columnwise_data = torch.empty_like(data, pin_memory=pin_memory) + 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), From 05bfa3f8a69ecbed5db0d74a852a1df0f3241768 Mon Sep 17 00:00:00 2001 From: Jaime <102792198+jaimec00@users.noreply.github.com> Date: Tue, 18 Nov 2025 18:41:41 -0500 Subject: [PATCH 078/521] [PyTorch] Implement Selective Activation Checkpointing for LayerNormMLP with checkpoint flag (#2311) * custom tests for selective activation checkpointing for layernorm mlp Signed-off-by: Jaime Cardenas * add selective layernorm mlp to te.pytorch Signed-off-by: Jaime Cardenas * update test and fix SLNMLP bug Signed-off-by: Jaime Cardenas * implement slnmlp Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Jaime Cardenas * fix tests pointed out by greptile app bot, still pass Signed-off-by: Jaime Cardenas * minor formatting change in tests/pytorch/selective_layernorm_mlp/distributed/run_numerics.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Jaime <102792198+jaimec00@users.noreply.github.com> Signed-off-by: Jaime Cardenas * remove duplicate import in test/pytorch/selective_layernorm_mlp/test_recipe.py Signed-off-by: Jaime Cardenas * clean up tests, remove unused imports Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Jaime Cardenas * remove unused paths in test_deffered_init Signed-off-by: Jaime Cardenas * fix issue with zero_centered_gamma in test_numerics reference implementation Signed-off-by: Jaime Cardenas * clean up tests Signed-off-by: Jaime Cardenas * make comparison.py more extensive, cleaner output Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Jaime Cardenas * fix small typo in tests/pytorch/selective_layernorm_mlp/compare.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Jaime <102792198+jaimec00@users.noreply.github.com> Signed-off-by: Jaime Cardenas * fix typo by grepbot in compare.py Signed-off-by: Jaime Cardenas * make selectiuve activation checkpointing optional in slnmlp via checkpoint flag Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Jaime Cardenas * add comments to clarify logic Signed-off-by: Jaime Cardenas * add checkpoint param to pytests, change compare.py to compare checkppoint=False vs checkpoint=True, skip cuda graph tests for checkpoint=True Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Jaime Cardenas * refactor tests to call modified LayerNormMLP Signed-off-by: Jaime Cardenas * refactor to implement selective activation checkpointing directly into LayerNormMLP, also fix bug to reach cleanup logic in fwd Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix skip explanation for cuda_graphs.py Signed-off-by: Jaime Cardenas * make _recompute deal with lists instead of tuples Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix MOST cuda graph failures by initializing identical quantizers during fwd. Float8CurrentScaling with bf16 and fp16 still fail with checkpointing Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix cuda graphs issue, all tests pass now Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix small logic bugs, clean up Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * integrate tests into main testing scripts Signed-off-by: Jaime Cardenas * incorporate rng state tracking in checkpointing Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * clean up tests Signed-off-by: Jaime Cardenas * fix return type mismatches Signed-off-by: Jaime Cardenas * remove checkpoint test from test_recipe, add sperate test in test_numerics Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor typo fix Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Jaime <102792198+jaimec00@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * clear up assertions in tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add license and copyright info Signed-off-by: Jaime Cardenas * fix lint issues in layernorm_mlp Signed-off-by: Jaime Cardenas * fix cpu_offload_v1 error Signed-off-by: Jaime Cardenas * possibly fix recomputation in cuda graph bug Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * skip cuda graphs test for SLNMLP with SM>=10.0 and using delayed scaling Signed-off-by: Jaime Cardenas * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix typo for setting IS_FIRST_FP8_MODULE Signed-off-by: Jaime Cardenas --------- Signed-off-by: Jaime Cardenas Signed-off-by: Jaime <102792198+jaimec00@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: Kirthi Shankar Sivamani --- tests/pytorch/distributed/run_numerics.py | 1 + tests/pytorch/distributed/test_numerics.py | 2 +- .../test_selective_activation_checkpoint.py | 175 +++++++ tests/pytorch/test_cuda_graphs.py | 30 +- tests/pytorch/test_numerics.py | 62 ++- tests/pytorch/test_recipe.py | 1 - tests/pytorch/test_sanity.py | 3 + .../pytorch/module/layernorm_mlp.py | 464 ++++++++++++++---- 8 files changed, 624 insertions(+), 114 deletions(-) create mode 100644 tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index 63ecb548bd..c109f463d8 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -1030,6 +1030,7 @@ def test_layernorm_mlp(): {"return_bias": True}, {"return_layernorm_output": True}, {"delay_wgrad_compute": True}, + {"checkpoint": True}, ] for kwargs in kwargs_list: diff --git a/tests/pytorch/distributed/test_numerics.py b/tests/pytorch/distributed/test_numerics.py index 97a69e779e..05b3e54280 100644 --- a/tests/pytorch/distributed/test_numerics.py +++ b/tests/pytorch/distributed/test_numerics.py @@ -13,7 +13,7 @@ """ Distributed numerics tests - These tests test the numerical corectness of the TransformerEngine layers. + These tests test the numerical correctness of the TransformerEngine layers. Tests are parametrized by the layer and fp8 precision. One test consists of running multiple configurations from file run_numerics.py Such design is due to the fact the initialization of one test is long diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py new file mode 100644 index 0000000000..8ec8a29d80 --- /dev/null +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -0,0 +1,175 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch +from transformer_engine.pytorch import LayerNormMLP +import pytest + +torch.manual_seed(1234) +device = torch.device("cuda") + + +class _Sequential(torch.nn.Sequential): + """Sequential model that forwards keyword arguments to modules""" + + def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor: + x = input_ + for module in self: + x = module(x, **kwargs) + return x + + +class ModelConfig: + def __init__( + self, + hidden_size: int = 128, + ffn_hidden_size: int = 512, + layers: int = 1, + ): + self._hidden_size = hidden_size + self._ffn_hidden_size = ffn_hidden_size + self._layers = layers + + def build(self): + + ln_list, sln_list = [], [] + for _ in range(self._layers): + ln = LayerNormMLP(self._hidden_size, self._ffn_hidden_size, checkpoint=False).to(device) + sln = LayerNormMLP(self._hidden_size, self._ffn_hidden_size, checkpoint=True).to(device) + with torch.no_grad(): + sln.layer_norm_weight = torch.nn.Parameter(ln.layer_norm_weight.clone()) + sln.layer_norm_bias = torch.nn.Parameter(ln.layer_norm_bias.clone()) + sln.fc1_weight = torch.nn.Parameter(ln.fc1_weight.clone()) + sln.fc2_weight = torch.nn.Parameter(ln.fc2_weight.clone()) + sln.fc1_bias = torch.nn.Parameter(ln.fc1_bias.clone()) + sln.fc2_bias = torch.nn.Parameter(ln.fc2_bias.clone()) + ln_list.append(ln) + sln_list.append(sln) + + ln_model = _Sequential(*ln_list) + sln_model = _Sequential(*sln_list) + + return ln_model, sln_model + + +config = { + "small": ModelConfig(128, 512, 12), + "medium": ModelConfig(512, 2048, 12), + "large": ModelConfig(1024, 4096, 12), + "huge": ModelConfig(2048, 8192, 12), +} + +seq_sizes = [2**7, 2**10, 2**14, 2**16] + + +def _warmup(model, tensor): + for _ in range(3): + model(tensor).sum().backward() + + +def _run_fwd(model, tensor): + + torch.cuda.reset_peak_memory_stats(device) + start_time, end_time = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + + torch.cuda.synchronize() + start_mem = torch.cuda.memory_allocated(device) + start_time.record() + out = model(tensor) + end_time.record() + end_time.synchronize() + elapsed = start_time.elapsed_time(end_time) + peak_mem = torch.cuda.max_memory_allocated(device) + mem = float(peak_mem - start_mem) + + return out, elapsed, mem + + +def _run_bwd(model, out): + + model.zero_grad(set_to_none=False) + loss = out.sum() + + torch.cuda.reset_peak_memory_stats(device) + start_time, end_time = torch.cuda.Event(enable_timing=True), torch.cuda.Event( + enable_timing=True + ) + + torch.cuda.synchronize() + start_mem = torch.cuda.memory_allocated(device) + start_time.record() + loss.backward() + end_time.record() + end_time.synchronize() + elapsed = start_time.elapsed_time(end_time) + peak_mem = torch.cuda.max_memory_allocated(device) + mem = float(peak_mem - start_mem) + + param_grads = _collect_param_grads(model) + return param_grads, elapsed, mem + + +def _max_diff(ref, other): + """Return max absolute difference between two tensors or collections.""" + if ref is None or other is None: + return 0.0 + if isinstance(ref, (list, tuple)): + diffs = [_max_diff(r, o) for r, o in zip(ref, other)] + return max(diffs) if diffs else 0.0 + return torch.max(torch.abs(ref.detach() - other.detach())).item() + + +def _collect_param_grads(model): + grads = {} + for name, param in model.named_parameters(): + if param.grad is None: + continue + key = _param_key(name) + if key is not None: + grads[key] = param.grad.detach().clone() + return grads + + +def _param_key(name): + return name.split(".")[-1] + + +@pytest.mark.parametrize("size", config.keys()) +@pytest.mark.parametrize("seq_size", seq_sizes) +def test_selective_activation_checkpoint(size, seq_size): + + ln_model, sln_model = config[size].build() + data = torch.randn((seq_size, config[size]._hidden_size), device=device) + + _warmup(ln_model, data) + ln_fwd_out, ln_fwd_time, ln_fwd_mem = _run_fwd(ln_model, data) + ln_grads, ln_bwd_time, ln_bwd_mem = _run_bwd(ln_model, ln_fwd_out) + + _warmup(sln_model, data) + sln_fwd_out, sln_fwd_time, sln_fwd_mem = _run_fwd(sln_model, data) + sln_grads, sln_bwd_time, sln_bwd_mem = _run_bwd(sln_model, sln_fwd_out) + + assert ln_fwd_mem > 6 * sln_fwd_mem, ( + "selective activation checkpointing does not reduce forward memory by 6X, only by" + f" {ln_fwd_mem/sln_fwd_mem}!" + ) + assert ln_bwd_time < sln_bwd_time, ( + "selective activation activation checkpointing backward pass is NOT slower than native!" + f" got Native LayerNormMLP Backward Time: {ln_bwd_time} ms and Selective Activation" + f" Checkpointed LayerNormMLP Backward Time: {sln_bwd_time} ms" + ) + diff = _max_diff(ln_fwd_out, sln_fwd_out) + assert diff == 0.0, f"outputs are not equal! maximum difference {diff}" + for key in [ + "layer_norm_weight", + "layer_norm_bias", + "fc1_weight", + "fc1_bias", + "fc2_weight", + "fc2_bias", + ]: + diff = _max_diff(ln_grads[key], sln_grads[key]) + assert diff == 0.0, f"gradients for {key} are not equal! maximum difference: {diff}" diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index eacbf5168e..3ddf33b16a 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -190,7 +190,8 @@ def forward(self, input_: torch.Tensor, **kwargs) -> torch.Tensor: # creating TMA descriptor for MXFP8 quantization. "linear", "transformer", - "layernorm_mlp", + "layernorm_mlp_nocheckpoint", + "layernorm_mlp_checkpoint", "layernorm_linear", "mha", "linear_op", @@ -232,12 +233,23 @@ def _test_cuda_graphs( ) for _ in range(num_layers) ] - elif module == "layernorm_mlp": + elif module == "layernorm_mlp_nocheckpoint": modules = [ LayerNormMLP( model_config.hidden_size, model_config.hidden_size, params_dtype=dtype, + checkpoint=False, + ) + for _ in range(num_layers) + ] + elif module == "layernorm_mlp_checkpoint": + modules = [ + LayerNormMLP( + model_config.hidden_size, + model_config.hidden_size, + params_dtype=dtype, + checkpoint=True, ) for _ in range(num_layers) ] @@ -376,6 +388,17 @@ def test_make_graphed_callables( ) if fp8_params: pytest.skip("NVFP4 params not supported") + if ( + fp8 + and fp8_recipe.delayed() + and torch.cuda.get_device_capability() >= (10, 0) + and module == "layernorm_mlp_checkpoint" + ): + pytest.skip( + "CUDA graphs not supported for LayerNormMLP " + "with checkpoint=True, SM>=10, " + "and DelayedScaling recipe" + ) # Run model with different CUDA graph settings. model_config = model_configs[model_config] @@ -402,7 +425,8 @@ def test_make_graphed_callables( _test_make_graphed_callables_with_fp8_weight_caching_modules = [ "transformer", - "layernorm_mlp", + "layernorm_mlp_nocheckpoint", + "layernorm_mlp_checkpoint", "layernorm_linear", "linear", "mha", diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 1925f2e2e9..66f08d5409 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -185,7 +185,7 @@ def dtype_tols(dtype: torch.dtype) -> Dict[str, float]: return dict(rtol=1e-3, atol=1e-5) if dtype == torch.bfloat16: return dict(rtol=1.6e-2, atol=1e-5) - raise ValueError(f"Unsuppored dtype ({dtype})") + raise ValueError(f"Unsupported dtype ({dtype})") def assert_allclose( @@ -1363,7 +1363,7 @@ def test_linear_accuracy_save_original_input(dtype, model, recipe): te_outputs = _test_granular_accuracy(te_linear, bs, dtype, config, recipe=recipe) te_outputs_ref = _test_granular_accuracy(te_linear_ref, bs, dtype, config, recipe=recipe) - # Shoule be bit-wise match + # Should be bit-wise match for i, (o, o_ref) in enumerate(zip(te_outputs, te_outputs_ref)): torch.testing.assert_close(o, o_ref, rtol=0, atol=0) @@ -1696,7 +1696,11 @@ def test_layernorm_mlp_accuracy(dtype, bs, model, activation, normalization, ret @pytest.mark.parametrize("bias", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) def test_layernorm_mlp_accuracy_delay_wgrad_compute( - dtype, bs, model, bias, fuse_wgrad_accumulation + dtype, + bs, + model, + bias, + fuse_wgrad_accumulation, ): config = model_configs[model] @@ -1747,6 +1751,58 @@ def test_layernorm_mlp_accuracy_delay_wgrad_compute( torch.testing.assert_close(o, o_ref, rtol=0, atol=0) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("bs", [2]) +@pytest.mark.parametrize("model", ["small"]) +@pytest.mark.parametrize("bias", all_boolean) +def test_layernorm_mlp_accuracy_checkpoint( + dtype, + bs, + model, + bias, +): + config = model_configs[model] + + ln_mlp = LayerNormMLP( + hidden_size=config.hidden_size, + ffn_hidden_size=4 * config.hidden_size, + eps=config.eps, + bias=bias, + params_dtype=dtype, + device="cuda", + checkpoint=True, + ).eval() + + ln_mlp_ref = LayerNormMLP( + hidden_size=config.hidden_size, + ffn_hidden_size=4 * config.hidden_size, + eps=config.eps, + bias=bias, + params_dtype=dtype, + device="cuda", + checkpoint=False, + ).eval() + + # Share params + with torch.no_grad(): + ln_mlp_ref.layer_norm_weight = Parameter(ln_mlp.layer_norm_weight.clone()) + ln_mlp_ref.layer_norm_bias = Parameter(ln_mlp.layer_norm_bias.clone()) + ln_mlp_ref.fc1_weight = Parameter(ln_mlp.fc1_weight.clone()) + ln_mlp_ref.fc2_weight = Parameter(ln_mlp.fc2_weight.clone()) + if bias: + ln_mlp_ref.fc1_bias = Parameter(ln_mlp.fc1_bias.clone()) + ln_mlp_ref.fc2_bias = Parameter(ln_mlp.fc2_bias.clone()) + + te_outputs = _test_granular_accuracy(ln_mlp, bs, dtype, config, delay_wgrad_compute=False) + te_outputs_ref = _test_granular_accuracy( + ln_mlp_ref, bs, dtype, config, delay_wgrad_compute=False + ) + + # Shoule be bit-wise match + for i, (o, o_ref) in enumerate(zip(te_outputs, te_outputs_ref)): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + def _test_grouped_linear_accuracy( block, num_gemms, diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 71032d23fb..ea26f0b108 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -29,7 +29,6 @@ ) import transformer_engine.pytorch.ops as te_ops from transformer_engine.common.recipe import DelayedScaling, Float8BlockScaling, MXFP8BlockScaling -import transformer_engine_torch as tex # Check if FP8 is supported fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 5c116496ef..fc4a5d6515 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -525,6 +525,7 @@ def test_sanity_grouped_linear( @pytest.mark.parametrize("activation", all_activations) @pytest.mark.parametrize("normalization", all_normalizations) @pytest.mark.parametrize("microbatching", all_boolean) +@pytest.mark.parametrize("checkpoint", all_boolean) def test_sanity_layernorm_mlp( dtype, fp8_recipe, @@ -535,6 +536,7 @@ def test_sanity_layernorm_mlp( activation, normalization, microbatching, + checkpoint, ): config = model_configs[model] @@ -559,6 +561,7 @@ def test_sanity_layernorm_mlp( normalization=normalization, params_dtype=dtype, device="cuda", + checkpoint=checkpoint, ) _test_sanity_common(block, dtype, config, fp8_recipe, skip_wgrad, skip_dgrad, microbatching) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 35dcb10f34..c1e3d8d2c9 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -56,6 +56,8 @@ use_reentrant_activation_recompute, in_fp8_activation_recompute_phase, _fsdp_scatter_tensors, + _get_cuda_rng_state, + _set_cuda_rng_state, ) from ..constants import dist_group_type from ..jit import no_torch_dynamo @@ -165,7 +167,7 @@ class _LayerNormMLP(torch.autograd.Function): """ @staticmethod - def forward( + def _forward( ctx, inp: torch.Tensor, ln_weight: torch.Tensor, @@ -226,9 +228,103 @@ def forward( module, skip_fp8_weight_update, symmetric_ar_type, + checkpoint, debug, + recompute_for_bwd, ) = non_tensor_args + # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take + if is_grad_enabled and not recompute_for_bwd: + ctx.checkpoint = checkpoint + if checkpoint: + # save the state of autocast and quantizers for recomputation + ctx.autocast_state = ( + FP8GlobalStateManager.get_autocast_state() + ) # to restore autocast state during recomputation + if ( + fp8 + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ + == "DelayedScaling" + ): # only applicable for delayed scaling + FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute( + module.fp8_meta + ) # to restore quantizers during recomputation + # save the rng states + ctx.cpu_rng_state = torch.get_rng_state() + ctx.cuda_rng_state = _get_cuda_rng_state() + + # whether to save activations regularly, or save inputs for recomputation in bwd + save_for_checkpoint = checkpoint and is_grad_enabled and not recompute_for_bwd + + # whether we are in the forward stage, or recomputing in the bwd stage (false if not checkpointing) + is_recomputation = checkpoint and is_grad_enabled and recompute_for_bwd + + # save the initial state for recomputation by bwd + if save_for_checkpoint: + + # save tensors + tensors_to_save, tensor_objects = prepare_for_saving( + inp, + ln_weight, + ln_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias, + ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + + ctx.other_args = { + "eps": eps, + "is_first_microbatch": is_first_microbatch, + "fp8": fp8, + "fp8_calibration": fp8_calibration, + "wgrad_store": wgrad_store, + "fuse_wgrad_accumulation": fuse_wgrad_accumulation, + "fc1_input_quantizer": fc1_input_quantizer, + "fc1_weight_quantizer": fc1_weight_quantizer, + "fc1_output_quantizer": fc1_output_quantizer, + "fc1_grad_input_quantizer": fc1_grad_input_quantizer, + "fc1_grad_weight_quantizer": fc1_grad_weight_quantizer, + "fc1_grad_output_quantizer": fc1_grad_output_quantizer, + "fc2_input_quantizer": fc2_input_quantizer, + "fc2_weight_quantizer": fc2_weight_quantizer, + "fc2_output_quantizer": fc2_output_quantizer, + "fc2_grad_input_quantizer": fc2_grad_input_quantizer, + "fc2_grad_weight_quantizer": fc2_grad_weight_quantizer, + "fc2_grad_output_quantizer": fc2_grad_output_quantizer, + "cpu_offloading": cpu_offloading, + "tp_group": tp_group, + "tp_size": tp_size, + "sequence_parallel": sequence_parallel, + "tensor_parallel": tensor_parallel, + "activation_dtype": activation_dtype, + "return_layernorm_output": return_layernorm_output, + "return_layernorm_output_gathered": return_layernorm_output_gathered, + "bias_gelu_fusion": bias_gelu_fusion, + "set_parallel_mode": set_parallel_mode, + "is_grad_enabled": is_grad_enabled, + "fwd_ln_sm_margin": fwd_ln_sm_margin, + "bwd_ln_sm_margin": bwd_ln_sm_margin, + "zero_centered_gamma": zero_centered_gamma, + "activation": activation, + "activation_params": activation_params, + "normalization": normalization, + "ub_overlap_ag": ub_overlap_ag, + "ub_overlap_rs": ub_overlap_rs, + "ub_overlap_rs_dgrad": ub_overlap_rs_dgrad, + "ub_bulk_wgrad": ub_bulk_wgrad, + "ub_bulk_dgrad": ub_bulk_dgrad, + "gemm_gelu_fusion": gemm_gelu_fusion, + "fsdp_group": fsdp_group, + "module": module, + "skip_fp8_weight_update": skip_fp8_weight_update, + "symmetric_ar_type": symmetric_ar_type, + "checkpoint": checkpoint, + "debug": debug, + "recompute_for_bwd": True, # set this to true for recomputation phase + } # Make sure input dimensions are compatible in_features, inp_shape = ln_weight.numel(), inp.shape assert inp_shape[-1] == in_features, "GEMM not possible" @@ -250,7 +346,14 @@ def forward( start_offload(inputmat) tp_world_size = get_distributed_world_size(tp_group) - backwards_needs_fc1_input = is_grad_enabled and fc1_weight.requires_grad + + # bwd needs fc1 input when grad is enabled, fc1 needs grad, and either + # 1) no checkpointing + # or 2) doing the recomputation with checkpointing + backwards_needs_fc1_input = fc1_weight.requires_grad and ( + (is_grad_enabled and not checkpoint) or is_recomputation + ) + device = inp.device # Configure Userbuffers communication (comm+GEMM overlap) @@ -308,7 +411,9 @@ def forward( zero_centered_gamma, ) ln_out_return = None - if return_layernorm_output or return_layernorm_output_gathered: + + # do not return layernorm output unless 1) no checkpointing or 2) checkpointing but not recomputing + if (return_layernorm_output or return_layernorm_output_gathered) and not is_recomputation: ln_out_return = ln_out # Prepare GEMM input @@ -316,7 +421,9 @@ def forward( ln_out_total = None ub_obj_lnout = None if sequence_parallel: - if return_layernorm_output_gathered: + + # do not return ln output if checkpointing and in recomputation, not necessary + if return_layernorm_output_gathered and not is_recomputation: # Perform all-gather in high precision if gathered # norm output will be returned ln_out_total, _ = gather_along_first_dim(ln_out, tp_group) @@ -459,7 +566,12 @@ def forward( # ------------------------------------------------------ # Deallocate FC1 GEMM input tensor if no longer needed - if not is_grad_enabled and (ln_out_total is not ln_out_return): + # first part of if statement means that we only clear ln_out_total if + # 1) checkpointing and not recomputing (in the forward stage, not bwd recompute stage) + # 2) not checkpointing and grad disabled + if ((checkpoint and not is_recomputation) or not is_grad_enabled) and ( + ln_out_total is not ln_out_return + ): clear_tensor_data(ln_out_total) # ACTIVATION - sometimes activation is fused with the GEMM above. @@ -497,89 +609,88 @@ def forward( else: act_out = activation_func(fc1_out, fc2_input_quantizer, **act_params) - if not is_grad_enabled: - clear_tensor_data(fc1_out) - if not fp8 and fp8_calibration: if fc2_input_quantizer is not None: fc2_input_quantizer.calibrate(act_out) - if fc2_weight_quantizer is not None: - fc2_weight_quantizer.calibrate(fc2_weight) - - # Configure Userbuffers reduce-scatter if needed - ub_obj_fc2out = None - reduce_scatter_out = None - if ub_overlap_rs: - ub_obj_fc2out = get_ub("fc2_fprop", fp8) - dim_size = list(act_out.size()) - dim_size[0] //= tp_world_size - dim_size[-1] = fc2_weight.size(0) - reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device) - # ------------------------------------------------------ - # FC2 GEMM - # ------------------------------------------------------ - gemm_out, *_, reduce_scatter_out = general_gemm( - fc2_weight_final, - act_out, - out_dtype=activation_dtype, - bias=fc2_bias, - quantization_params=fc2_output_quantizer, - use_split_accumulator=use_split_accumulator, - ub=ub_obj_fc2out, - ub_type=tex.CommOverlapType.RS if ub_overlap_rs else None, - extra_output=reduce_scatter_out, - ) - # ------------------------------------------------------ - # Finished FC2 GEMM... - # ------------------------------------------------------ + # we want to skip fc2 computation if we are checkpointing and recomputing, + # otherwise we compute fc2 + if not (is_recomputation and checkpoint): - # Deallocate tensors if no longer needed - if not is_grad_enabled: - clear_tensor_data(act_out, fc1_out_without_bias, fc1_out) - - # Prepare output tensor - # Note: Perform tensor-parallel communication if needed - fc2_out = None - if ub_overlap_rs: - fc2_out = reduce_scatter_out - elif set_parallel_mode and sequence_parallel: - fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) - elif set_parallel_mode and tensor_parallel: - if symmetric_ar_type is not None: - fc2_out, _ = symmetric_all_reduce( - gemm_out, tp_group, all_reduce_type=symmetric_ar_type - ) - else: - fc2_out, _ = allreduce(gemm_out, tp_group) - else: - fc2_out = gemm_out - fc2_out = fc2_out.view(-1, *inp_shape[1:-1], fc2_out.shape[-1]) - - # Cache state for backward pass - if is_grad_enabled: - if cpu_offloading: - mark_activation_offload( - inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out - ) + # if we get to this point, we know this is not bwd recomputation + # so we must be in the fwd + # now is_grad_enabled can be true or false + # if false, can safely delete + # if true, we can only delete if checkpoint is true, since we will recompute anyways, + # otherwise, checkpoint is false, so cant delete + if ( + checkpoint or not is_grad_enabled + ): # we can safely get rid of these if this is the case + clear_tensor_data(fc1_out) - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already - # shards/unshards the base weights so we don't do it ourselves - ctx.fsdp_group = fsdp_group - ctx.fsdp_shapes = _fsdp_scatter_tensors( - fsdp_group, - mu, - rsigma, - ln_out, - fc1_out_without_bias if bias_gelu_fusion else fc1_out, + if not fp8 and fp8_calibration: + + if fc2_weight_quantizer is not None: + fc2_weight_quantizer.calibrate(fc2_weight) + + # Configure Userbuffers reduce-scatter if needed + ub_obj_fc2out = None + reduce_scatter_out = None + if ub_overlap_rs: + ub_obj_fc2out = get_ub("fc2_fprop", fp8) + dim_size = list(act_out.size()) + dim_size[0] //= tp_world_size + dim_size[-1] = fc2_weight.size(0) + reduce_scatter_out = torch.empty(dim_size, dtype=activation_dtype, device=device) + + # ------------------------------------------------------ + # FC2 GEMM + # ------------------------------------------------------ + gemm_out, *_, reduce_scatter_out = general_gemm( + fc2_weight_final, act_out, - fc1_weight_final if fp8 and not isinstance(fc1_weight, Float8Tensor) else None, - fc2_weight_final if fp8 and not isinstance(fc2_weight, Float8Tensor) else None, + out_dtype=activation_dtype, + bias=fc2_bias, + quantization_params=fc2_output_quantizer, + use_split_accumulator=use_split_accumulator, + ub=ub_obj_fc2out, + ub_type=tex.CommOverlapType.RS if ub_overlap_rs else None, + extra_output=reduce_scatter_out, ) + # ------------------------------------------------------ + # Finished FC2 GEMM... + # ------------------------------------------------------ + + # Deallocate tensors if no longer needed, again, can safely deallocate + if checkpoint or not is_grad_enabled: # same logic as last clear_tensor_data block + clear_tensor_data(act_out, fc1_out_without_bias, fc1_out) + + # Prepare output tensor + # Note: Perform tensor-parallel communication if needed + fc2_out = None + if ub_overlap_rs: + fc2_out = reduce_scatter_out + elif set_parallel_mode and sequence_parallel: + fc2_out, _ = reduce_scatter_along_first_dim(gemm_out, tp_group) + elif set_parallel_mode and tensor_parallel: + if symmetric_ar_type is not None: + fc2_out, _ = symmetric_all_reduce( + gemm_out, tp_group, all_reduce_type=symmetric_ar_type + ) + else: + fc2_out, _ = allreduce(gemm_out, tp_group) + else: + fc2_out = gemm_out + fc2_out = fc2_out.view(-1, *inp_shape[1:-1], fc2_out.shape[-1]) + + # now saving stuff for bwd: + # if we are using checkpointing, this information will be saved in the bwd recomputation stage, so can skip it in fwd + # if we are not checkpointing, then we must save this if grad is enabled + if is_grad_enabled and not save_for_checkpoint: ctx.fc1_weight_quantizer = fc1_weight_quantizer ctx.fc2_weight_quantizer = fc2_weight_quantizer + if not fc1_weight.requires_grad: if not return_layernorm_output: clear_tensor_data(ln_out) @@ -588,34 +699,69 @@ def forward( clear_tensor_data(act_out) act_out = None - if cpu_offloading: - mark_not_offload( + if not checkpoint: # regular path, no selective activation checkpointing + + if cpu_offloading: + mark_activation_offload( + inputmat, mu, rsigma, ln_out, fc1_out, fc1_out_without_bias, act_out + ) + + # Scatter intermediate/activation tensors saved for the backward pass + # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already + # shards/unshards the base weights so we don't do it ourselves + ctx.fsdp_group = fsdp_group + + ctx.fsdp_shapes = ( + _fsdp_scatter_tensors( # again, ony relevant if we have activations to save + fsdp_group, + mu, + rsigma, + ln_out, + fc1_out_without_bias if bias_gelu_fusion else fc1_out, + act_out, + ( + fc1_weight_final + if fp8 and not isinstance(fc1_weight, Float8Tensor) + else None + ), + ( + fc2_weight_final + if fp8 and not isinstance(fc2_weight, Float8Tensor) + else None + ), + ) + ) + + if cpu_offloading: + mark_not_offload( + ln_weight, + ln_bias, + fc1_weight_final, + fc1_weight, + fc1_bias, + fc2_weight_final, + fc2_weight, + fc2_bias, + ) + tensors_to_save, tensor_objects = prepare_for_saving( + inputmat, ln_weight, - ln_bias, + 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, ) - tensors_to_save, tensor_objects = prepare_for_saving( - inputmat, - 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, - ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects if fuse_wgrad_accumulation: # This check is needed to ensure that main_grad is not created @@ -633,9 +779,6 @@ def forward( ctx.fc1_main_grad_func = lambda: fc1_weight.main_grad ctx.fc2_main_grad_func = lambda: fc2_weight.main_grad - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None ctx.fc1_grad_input_quantizer = fc1_grad_input_quantizer ctx.fc1_grad_weight_quantizer = fc1_grad_weight_quantizer @@ -690,11 +833,30 @@ def forward( ): _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): + if in_fp8_activation_recompute_phase() or is_recomputation: FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module ctx.wgrad_store = wgrad_store + if is_recomputation: # return the recomputed tensors + return ( + ctx, + inputmat, + 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, + ) + # we only get to this point if we are not recomputing for bwd, since that would have returned in the block above if return_layernorm_output: if return_layernorm_output_gathered: shape = list(inp_shape) @@ -703,14 +865,101 @@ def forward( return fc2_out, ln_out_return.view(inp_shape) return fc2_out + @staticmethod + def forward( + ctx, + inp: torch.Tensor, + ln_weight: torch.Tensor, + ln_bias: torch.Tensor, + fc1_weight: torch.Tensor, + fc1_bias: torch.Tensor, + fc2_weight: torch.Tensor, + fc2_bias: torch.Tensor, + non_tensor_args: Tuple, + ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: + # pylint: disable=missing-function-docstring + + # add recompute_for_bwd + non_tensor_args += (False,) + + return _LayerNormMLP._forward( + ctx, + inp, + ln_weight, + ln_bias, + fc1_weight, + fc1_bias, + fc2_weight, + fc2_bias, + non_tensor_args, + ) + + @staticmethod + def _recompute(ctx): + # pylint: disable=missing-function-docstring + + saved_tensors = ctx.saved_tensors + tensors = restore_from_saved(ctx.tensor_objects, saved_tensors) + # Delete the references to tensor objects once they've been consumed + # by the `restore_from_saved` method to construct back the actual tensors. + ctx.tensor_objects = None + + if ctx.checkpoint: # do recomputation from the original args + + # backward is not in autocast context, so we set the state here + # we also have to set the quantizer states to what they were before the forward pass (only relevant for DelayedScaling recipe) + final_autocast_state = ( + FP8GlobalStateManager.get_autocast_state() + ) # get current autocast state + FP8GlobalStateManager.set_autocast_state(ctx.autocast_state) # set old autocast state + if ( + ctx.other_args["fp8"] + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): # only applicable for delayed scaling + FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute( + ctx.other_args["module"].fp8_meta + ) # set old quantizer state + + # get current rng state + final_cpu_rng_state = torch.get_rng_state() + final_cuda_rng_state = _get_cuda_rng_state() + + # set rng state for fwd + torch.set_rng_state(ctx.cpu_rng_state) + _set_cuda_rng_state(ctx.cuda_rng_state) + + out = _LayerNormMLP._forward( # recompute + ctx, + *tensors, + tuple(ctx.other_args.values()), + ) + + FP8GlobalStateManager.set_autocast_state(final_autocast_state) # restore autocast state + if ( + ctx.other_args["fp8"] + and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" + ): + FP8GlobalStateManager.restore_fp8_meta_tensors( + ctx.other_args["module"].fp8_meta + ) # restore quantizers + + # set rng state for fwd + torch.set_rng_state(final_cpu_rng_state) + _set_cuda_rng_state(final_cuda_rng_state) + + return out + + # load from saved (return ctx is just because the other branch does too) + return tuple([ctx] + tensors) + @staticmethod def backward( ctx, *grad_outputs: Tuple[torch.Tensor, ...] ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_LayerNormMLP_backward"): - saved_tensors = ctx.saved_tensors ( # pylint: disable=unbalanced-tuple-unpacking + ctx, inputmat, ln_weight, ln_out, @@ -725,11 +974,7 @@ def backward( fc2_bias, mu, rsigma, - ) = restore_from_saved(ctx.tensor_objects, saved_tensors) - - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None + ) = _LayerNormMLP._recompute(ctx) # Since main_grad can be modified inplace, it should not be a part of saved_tensors fc1_weight_main_grad = ( @@ -1512,6 +1757,10 @@ class LayerNormMLP(TransformerEngineBaseModule): This can help in latency bound communication situations. Requires PyTorch version 2.7.0 or higher. When set to None, standard all-reduce is used. + checkpoint: bool, default = False + whether to use selective activation checkpointing, where activations are not saved for bwd, + and instead are recomputed (skipping fc2, as it is not needed for backward). Trades compute + for memory. default is false, in which activations are saved in fwd. not supported for onnx forward """ def __init__( @@ -1547,6 +1796,7 @@ def __init__( ub_bulk_wgrad: bool = False, delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, + checkpoint: bool = False, ) -> None: super().__init__() @@ -1567,6 +1817,7 @@ def __init__( self.set_parallel_mode = set_parallel_mode self.zero_centered_gamma = zero_centered_gamma self.symmetric_ar_type = symmetric_ar_type + self.checkpoint = checkpoint # GEMM-GELU fusion is currently only supported with split GEMM-AG overlap self.gemm_gelu_fusion = ( @@ -1896,6 +2147,7 @@ def forward( self, skip_fp8_weight_update, self.symmetric_ar_type, + self.checkpoint, debug, ) out = fwd_fn( From 41fb9bcf34dddac6288e556be2ed74925808e673 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 18 Nov 2025 17:03:55 -0800 Subject: [PATCH 079/521] [PyTorch] fix `test_current_device` test (#2398) * fix test_current_device 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> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/distributed/test_sanity.py | 60 ++++++++++++++++++++---- 1 file changed, 50 insertions(+), 10 deletions(-) diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index 46f8a5b29a..f7c0e1fe88 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -55,7 +55,29 @@ def test_current_device(model, module): self_attn_mask_type="padding", device=f"cuda:{tensor_device}", ) - num_tokens = torch.randint(0, config.max_seqlen_q, (1,)).item() + seqlens_q = torch.randint( + 1, + config.max_seqlen_q, + [config.batch_size], + dtype=torch.int32, + device=f"cuda:{tensor_device}", + ) + cu_seqlens_q = torch.zeros( + config.batch_size + 1, dtype=torch.int32, device=f"cuda:{tensor_device}" + ) + cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) + seqlens_kv = torch.randint( + 1, + config.max_seqlen_kv, + [config.batch_size], + dtype=torch.int32, + device=f"cuda:{tensor_device}", + ) + cu_seqlens_kv = torch.zeros( + config.batch_size + 1, dtype=torch.int32, device=f"cuda:{tensor_device}" + ) + cu_seqlens_kv[1:] = torch.cumsum(seqlens_kv, dim=0) + num_tokens = cu_seqlens_q[-1] args = [ torch.randn( (num_tokens, config.hidden_size), @@ -64,9 +86,6 @@ def test_current_device(model, module): requires_grad=True, ) ] - cu_seqlens_q, cu_seqlens_kv = [ - torch.Tensor([0, 2, 3]).to(dtype=torch.int32, device=tensor_device) for _ in range(2) - ] kwargs["cu_seqlens_q"] = cu_seqlens_q kwargs["cu_seqlens_kv"] = cu_seqlens_kv kwargs["max_seqlen_q"] = config.max_seqlen_q @@ -75,26 +94,47 @@ def test_current_device(model, module): model = DotProductAttention( config.num_heads, config.head_dim_qk, qkv_format="thd", attn_mask_type="padding" ) - num_tokens = torch.randint(0, config.max_seqlen_q, (1,)).item() + seqlens_q = torch.randint( + 1, + config.max_seqlen_q, + [config.batch_size], + dtype=torch.int32, + device=f"cuda:{tensor_device}", + ) + cu_seqlens_q = torch.zeros( + config.batch_size + 1, dtype=torch.int32, device=f"cuda:{tensor_device}" + ) + cu_seqlens_q[1:] = torch.cumsum(seqlens_q, dim=0) + seqlens_kv = torch.randint( + 1, + config.max_seqlen_kv, + [config.batch_size], + dtype=torch.int32, + device=f"cuda:{tensor_device}", + ) + cu_seqlens_kv = torch.zeros( + config.batch_size + 1, dtype=torch.int32, device=f"cuda:{tensor_device}" + ) + cu_seqlens_kv[1:] = torch.cumsum(seqlens_kv, dim=0) + num_tokens = cu_seqlens_q[-1] args = [ torch.randn( num_tokens, config.num_heads, config.head_dim_qk, dtype=dtype, - device=tensor_device, + device=f"cuda:{tensor_device}", requires_grad=True, ) for _ in range(3) ] - cu_seqlens_q, cu_seqlens_kv = [ - torch.Tensor([0, 2, 3]).to(dtype=torch.int32, device=tensor_device) for _ in range(2) - ] kwargs["cu_seqlens_q"] = cu_seqlens_q kwargs["cu_seqlens_kv"] = cu_seqlens_kv kwargs["max_seqlen_q"] = config.max_seqlen_q kwargs["max_seqlen_kv"] = config.max_seqlen_kv - bwd_args = [torch.randn(num_tokens, config.hidden_size, dtype=dtype, device=tensor_device)] + bwd_args = [ + torch.randn(num_tokens, config.hidden_size, dtype=dtype, device=f"cuda:{tensor_device}") + ] elif module == "Linear": model = Linear( config.hidden_size, From 877b7966d15b98e598d4a1f7eec3e151c1decb44 Mon Sep 17 00:00:00 2001 From: Jianbing Date: Wed, 19 Nov 2025 20:38:40 +0800 Subject: [PATCH 080/521] Feature fast cast-only mxfp8 (#2062) * refactor mxfp8_cast_only kernel Signed-off-by: Jianbing Dong * fix ptx.cuh after format Signed-off-by: Jianbing Dong --------- Signed-off-by: Jianbing Dong Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 68 + .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 1618 +++++++++++++++++ .../cast/mxfp8/specialized/state_counter.cuh | 61 + .../common/cast/mxfp8/specialized/swizzle.cuh | 90 + transformer_engine/common/util/ptx.cuh | 679 +++++++ 5 files changed, 2516 insertions(+) create mode 100644 transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh create mode 100644 transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh create mode 100644 transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 5505de6050..cbb46f3f28 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -21,6 +21,7 @@ #include "../../util/ptx.cuh" #include "../../utils.cuh" #include "../core/common.cuh" +#include "specialized/quantize_mxfp8.cuh" namespace transformer_engine { namespace dispatch { @@ -619,6 +620,73 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( output->dtype(), OType, + if (specialized::hasSpec()) { + switch (scaling_type) { + case ScalingType::ROWWISE: { + using traits = specialized::CastTraits; + auto kernel = specialized::quantize_mxfp8_kernel_cast_only; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + traits::smem); + + dim3 block(traits::threadLayout::num, traits::warpLayout::N, traits::warpLayout::M); + dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN, + (rows + traits::blockDimM - 1) / traits::blockDimM); + kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output->data.dptr), + scales_rowwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); + + 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); + // 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; + create_2D_tensor_map(tensor_map_input, input.data, rows, cols, + traits::blockIterDim::M, traits::blockIterDim::N, + /*stride_elems=*/cols, + /*offset_elems=*/0, input_type_bit_size, + traits::input_swizzle_pattern); + + alignas(64) CUtensorMap tensor_map_rowwise_output{}; + alignas(64) CUtensorMap tensor_map_colwise_output{}; + constexpr size_t output_type_bit_size = TypeInfo::size; + create_2D_tensor_map(tensor_map_rowwise_output, output->data, rows, cols, + traits::blockIterDim::M, traits::blockIterDim::N, + /*stride_elems=*/cols, + /*offset_elems=*/0, output_type_bit_size, + traits::output_swizzle_pattern); + create_2D_tensor_map(tensor_map_colwise_output, output->columnwise_data, rows, cols, + traits::blockIterDim::M, traits::blockIterDim::N, cols, 0, + output_type_bit_size, traits::output_swizzle_pattern); + + dim3 block(traits::rowThreadLayout::num, traits::numWarps); + dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N, + (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); + kernel<<>>( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, + scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + + break; + } + default: { + NVTE_ERROR("Invalid scaling type."); + } + } + return; + } + alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_act_input{}; alignas(64) CUtensorMap tensor_map_output_rowwise{}; diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh new file mode 100644 index 0000000000..4a39e54a35 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -0,0 +1,1618 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_mxfp8_spec.cuh + * \brief CUDA kernels to cast MXFP8. + */ + +#ifndef TRANSFORMER_ENGINE_SPECIALIZED_QUANTIZE_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_SPECIALIZED_QUANTIZE_MXFP8_CUH_ + +#include + +#include "../../../util/ptx.cuh" +#include "state_counter.cuh" +#include "swizzle.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace quantize_kernel { +namespace specialized { + +namespace ptx = transformer_engine::ptx; +namespace { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + +#if defined(_ENABLE_MXFMA) +template +struct _Quantized_Limits; + +template <> +struct _Quantized_Limits { + static constexpr uint16_t max_norm_rcp{0}; +}; + +template <> +struct _Quantized_Limits { + static constexpr uint16_t max_norm_rcp{0}; +}; + +template <> +struct _Quantized_Limits { + static constexpr uint16_t max_norm_rcp{0x125}; +}; + +template <> +struct _Quantized_Limits { + static constexpr uint16_t max_norm_rcp{0x3792}; +}; + +template <> +struct _Quantized_Limits { + static constexpr uint16_t max_norm_rcp{0x1892}; +}; + +template <> +struct _Quantized_Limits { + static constexpr uint16_t max_norm_rcp{0x3b12}; +}; +#endif // #if defined(_ENABLE_MXFMA) + +template +__device__ __forceinline__ e8m0_t to_e8m0(IType amax) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) && (defined _ENABLE_MXFMA) + constexpr uint16_t max_norm_rcp = _Quantized_Limits::max_norm_rcp; + + float amax_fp32; + if constexpr (std::is_same_v) { + ptx::fma_f32_f16(amax_fp32, reinterpret_cast(amax), max_norm_rcp); + } else if constexpr (std::is_same_v) { + ptx::fma_f32_bf16(amax_fp32, reinterpret_cast(amax), max_norm_rcp); + } else { + amax_fp32 = 0.0f; + __trap(); + } + return ptx::float_to_e8m0(amax_fp32); +#else + if constexpr (std::is_same_v) { + return ptx::float_to_e8m0(__fmaf_ieee_rn(amax, Quantized_Limits::max_norm_rcp, 0.0f)); + } else { + float amax_fp32 = static_cast(amax); + return ptx::float_to_e8m0( + __fmaf_ieee_rn(amax_fp32, Quantized_Limits::max_norm_rcp, 0.0f)); + } +#endif +} + +#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; +} + +// IType could be [fp16, bf16] +// OType could be [fp8e5m2, fp8e4m3] +template <> +inline bool hasSpec() { + return is_cast_only_enabled(); +} +template <> +inline bool hasSpec() { + return is_cast_only_enabled(); +} +template <> +inline bool hasSpec() { + return is_cast_only_enabled(); +} +template <> +inline bool hasSpec() { + return is_cast_only_enabled(); +} + +template +struct Layout { + static constexpr int32_t M = _M; // row + static constexpr int32_t N = _N; // col + static constexpr int32_t num = M * N; +}; + +template +struct CastTraits; + +// 1x32 +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> { + static constexpr bool isRowwise = true; + static constexpr bool isColwise = false; + using IType = _IType; + using OType = _OType; + + static constexpr int32_t chunkElems = 32; + using threadLayout = Layout<1, 32>; + static constexpr int32_t numThreadsPerChunk = 1; + static constexpr int32_t warpDimM = threadLayout::M; + static constexpr int32_t warpDimN = threadLayout::N * chunkElems; + using inputUnitType = uint4; + static constexpr int32_t numUnitsPerChunk = chunkElems * sizeof(IType) / sizeof(inputUnitType); + using outputUnitType = uint4; + static constexpr int32_t numOutUnitsPerChunk = + chunkElems * sizeof(OType) / sizeof(outputUnitType); + + using warpLayout = Layout<4, 1>; + static constexpr int32_t blockIterDimM = warpLayout::M * warpDimM; + static constexpr int32_t blockIterDimN = warpLayout::N * warpDimN; + + using iterLayout = Layout<1, 1>; + static constexpr int32_t blockDimM = iterLayout::M * blockIterDimM; + static constexpr int32_t blockDimN = iterLayout::N * blockIterDimN; + + static constexpr int32_t numStages = 1; + static constexpr int32_t numPrefetch = numStages - 1; + + static constexpr bool _use_cvt_4x = true; + static constexpr bool _cache_rowwise_scale_in_smem = true; + + static constexpr int32_t numThreads = warpLayout::num * 32; + + static constexpr size_t smem_rowwise_scale = + _cache_rowwise_scale_in_smem ? (blockDimM * (blockDimN / chunkElems) * sizeof(e8m0_t)) : 0ul; + static constexpr size_t smem = smem_rowwise_scale; +}; + +// 1x32 +template = 0> +__global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__restrict__ input, + typename CastTraits::OType *__restrict__ output, + e8m0_t *__restrict__ scales_rowwise, int32_t rows, + int32_t cols, int32_t scale_stride_rowwise, + int32_t scale_stride_colwise) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using IType = typename CastTraits::IType; + using OType = typename CastTraits::OType; + using inputUnitType = typename CastTraits::inputUnitType; + using outputUnitType = typename CastTraits::outputUnitType; + + using IType2 = typename ptx::FPx2; + constexpr int32_t numItersIType2 = sizeof(inputUnitType) / sizeof(IType2); + using OType2 = typename ptx::FPx2; + + e8m0_t *sRowwiseScale = nullptr; + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + extern __shared__ char smem[]; + sRowwiseScale = reinterpret_cast(smem); + } + + int2 block_coords; + block_coords.y = blockIdx.y * CastTraits::blockDimM + threadIdx.z * CastTraits::warpDimM + + (threadIdx.x / CastTraits::threadLayout::N); + block_coords.x = blockIdx.x * CastTraits::blockDimN + threadIdx.y * CastTraits::warpDimN + + (threadIdx.x % CastTraits::threadLayout::N) * CastTraits::chunkElems; + + int32_t rowwise_scale_smem_base_offset; + constexpr int32_t rowwise_scale_stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + rowwise_scale_smem_base_offset = + threadIdx.z * CastTraits::warpDimM * rowwise_scale_stride_in_smem + + threadIdx.y * (CastTraits::warpDimN / CastTraits::chunkElems) + + (threadIdx.x / CastTraits::threadLayout::N) * rowwise_scale_stride_in_smem + + (threadIdx.x % CastTraits::threadLayout::N); + } + + inputUnitType rInput[CastTraits::numStages][CastTraits::numUnitsPerChunk]; +// prologue +#pragma unroll + for (int32_t iter = 0; iter < CastTraits::numPrefetch; iter++) { + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDimM; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDimN; + + if (coords.y < rows && coords.x < cols) { + size_t offset = coords.y * static_cast(cols) + coords.x; + inputUnitType *input_units = reinterpret_cast(input + offset); + +#pragma unroll + for (int32_t i = 0; i < CastTraits::numUnitsPerChunk; i++) { + rInput[iter][i] = input_units[i]; + } + } + } +// mainloop +#pragma unroll + for (int32_t iter = CastTraits::numPrefetch; iter < CastTraits::iterLayout::num; iter++) { + { + // load data + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDimM; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDimN; + + if (coords.y < rows && coords.x < cols) { + size_t offset = coords.y * static_cast(cols) + coords.x; + inputUnitType *input_units = reinterpret_cast(input + offset); + +#pragma unroll + for (int32_t i = 0; i < CastTraits::numUnitsPerChunk; i++) { + rInput[iter % CastTraits::numStages][i] = input_units[i]; + } + } + } + int32_t process_iter = iter - CastTraits::numPrefetch; + int32_t iter_m = process_iter / CastTraits::iterLayout::N; + int32_t iter_n = process_iter % CastTraits::iterLayout::N; + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDimM; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDimN; + if (coords.y >= rows || coords.x >= cols) { + return; + } + + if constexpr (std::is_same_v) { + float thread_amax = 0.f; + IType2 *rInput2 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); +#pragma unroll + for (int32_t j = 0; j < numItersIType2 * CastTraits::numUnitsPerChunk; j++) { + ptx::abs_max_2x(thread_amax, thread_amax, rInput2[j].x, rInput2[j].y); + } + e8m0_t biased_exponent = to_e8m0(thread_amax); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + int32_t rowwise_scale_offset = + rowwise_scale_smem_base_offset + + iter_m * CastTraits::blockIterDimM * rowwise_scale_stride_in_smem + + iter_n * (CastTraits::blockIterDimN / CastTraits::chunkElems); + sRowwiseScale[rowwise_scale_offset] = biased_exponent; + } else { + scales_rowwise[coords.y * static_cast(scale_stride_rowwise) + + coords.x / CastTraits::chunkElems] = biased_exponent; + } + + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; + if constexpr (CastTraits::_use_cvt_4x) { + using OType4 = ptx::FPx4; + using IType4 = ptx::FPx4; + IType4 *rInput4 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); + OType4 *rOutput4 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t j = 0; j < CastTraits::chunkElems / 4; j++) { + IType4 in = rInput4[j]; + OType4 out; + ptx::mul_cvt_4x(out, in, block_scale_inverse_2x); + rOutput4[j] = out; + } + } else { + OType2 *rOutput2 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t j = 0; j < CastTraits::chunkElems / 2; j++) { + IType2 in = rInput2[j]; + OType2 out; + ptx::mul_cvt_2x(out, in, block_scale_inverse_2x); + rOutput2[j] = out; + } + } + outputUnitType *output_units = + reinterpret_cast(output + coords.y * cols + coords.x); +#pragma unroll + for (int32_t j = 0; j < CastTraits::numOutUnitsPerChunk; j++) { + output_units[j] = rOutput[j]; + } + } else { + IType2 thread_amax2{0.f, 0.f}; + IType2 *rInput2 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); +#pragma unroll + for (int32_t j = 0; j < numItersIType2 * CastTraits::numUnitsPerChunk; j++) { + ptx::abs_max_2x(thread_amax2, thread_amax2, rInput2[j]); + } + IType thread_amax = ptx::get_amax(thread_amax2.x, thread_amax2.y); + e8m0_t biased_exponent = to_e8m0(thread_amax); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + int32_t rowwise_scale_offset = + rowwise_scale_smem_base_offset + + iter_m * CastTraits::blockIterDimM * rowwise_scale_stride_in_smem + + iter_n * (CastTraits::blockIterDimN / CastTraits::chunkElems); + sRowwiseScale[rowwise_scale_offset] = biased_exponent; + } else { + scales_rowwise[coords.y * static_cast(scale_stride_rowwise) + + coords.x / CastTraits::chunkElems] = biased_exponent; + } + + // scaling input + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; + if constexpr (CastTraits::_use_cvt_4x) { + using OType4 = ptx::FPx4; + using IType4 = ptx::FPx4; + IType4 *rInput4 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); + OType4 *rOutput4 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t i = 0; i < CastTraits::chunkElems / 4; i++) { + IType4 in = rInput4[i]; + OType4 out; + ptx::mul_cvt_4x(out, in, block_scale_inverse_2x); + rOutput4[i] = out; + } + } else { + OType2 *rOutput2 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t i = 0; i < CastTraits::chunkElems / 2; i++) { + IType2 in = rInput2[i]; + OType2 out; + ptx::mul_cvt_2x(out, in, block_scale_inverse_2x); + rOutput2[i] = out; + } + } + outputUnitType *output_units = + reinterpret_cast(output + coords.y * cols + coords.x); +#pragma unroll + for (int32_t j = 0; j < CastTraits::numOutUnitsPerChunk; j++) { + output_units[j] = rOutput[j]; + } + } + } + +// epilogue +#pragma unroll + for (int32_t iter = CastTraits::iterLayout::num; + iter < CastTraits::iterLayout::num + CastTraits::numPrefetch; iter++) { + int32_t process_iter = iter - CastTraits::numPrefetch; + int32_t iter_m = process_iter / CastTraits::iterLayout::N; + int32_t iter_n = process_iter % CastTraits::iterLayout::N; + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDimM; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDimN; + if (coords.y >= rows || coords.x >= cols) { + return; + } + + if constexpr (std::is_same_v) { + float thread_amax = 0.f; + IType2 *rInput2 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); +#pragma unroll + for (int32_t j = 0; j < numItersIType2 * CastTraits::numUnitsPerChunk; j++) { + ptx::abs_max_2x(thread_amax, thread_amax, rInput2[j].x, rInput2[j].y); + } + e8m0_t biased_exponent = to_e8m0(thread_amax); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + int32_t rowwise_scale_offset = + rowwise_scale_smem_base_offset + + iter_m * CastTraits::blockIterDimM * rowwise_scale_stride_in_smem + + iter_n * (CastTraits::blockIterDimN / CastTraits::chunkElems); + sRowwiseScale[rowwise_scale_offset] = biased_exponent; + } else { + scales_rowwise[coords.y * static_cast(scale_stride_rowwise) + + coords.x / CastTraits::chunkElems] = biased_exponent; + } + + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; + if constexpr (CastTraits::_use_cvt_4x) { + using OType4 = ptx::FPx4; + using IType4 = ptx::FPx4; + IType4 *rInput4 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); + OType4 *rOutput4 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t j = 0; j < CastTraits::chunkElems / 4; j++) { + IType4 in = rInput4[j]; + OType4 out; + ptx::mul_cvt_4x(out, in, block_scale_inverse_2x); + rOutput4[j] = out; + } + } else { + OType2 *rOutput2 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t j = 0; j < CastTraits::chunkElems / 2; j++) { + IType2 in = rInput2[j]; + OType2 out; + ptx::mul_cvt_2x(out, in, block_scale_inverse_2x); + rOutput2[j] = out; + } + } + outputUnitType *output_units = + reinterpret_cast(output + coords.y * cols + coords.x); +#pragma unroll + for (int32_t j = 0; j < CastTraits::numOutUnitsPerChunk; j++) { + output_units[j] = rOutput[j]; + } + } else { + IType2 thread_amax2{0.f, 0.f}; + IType2 *rInput2 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); +#pragma unroll + for (int32_t j = 0; j < numItersIType2 * CastTraits::numUnitsPerChunk; j++) { + ptx::abs_max_2x(thread_amax2, thread_amax2, rInput2[j]); + } + IType thread_amax = ptx::get_amax(thread_amax2.x, thread_amax2.y); + e8m0_t biased_exponent = to_e8m0(thread_amax); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + int32_t rowwise_scale_offset = + rowwise_scale_smem_base_offset + + iter_m * CastTraits::blockIterDimM * rowwise_scale_stride_in_smem + + iter_n * (CastTraits::blockIterDimN / CastTraits::chunkElems); + sRowwiseScale[rowwise_scale_offset] = biased_exponent; + } else { + scales_rowwise[coords.y * static_cast(scale_stride_rowwise) + + coords.x / CastTraits::chunkElems] = biased_exponent; + } + + // scaling input + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; + if constexpr (CastTraits::_use_cvt_4x) { + using OType4 = ptx::FPx4; + using IType4 = ptx::FPx4; + IType4 *rInput4 = reinterpret_cast(&rInput[process_iter % CastTraits::numStages]); + OType4 *rOutput4 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t i = 0; i < CastTraits::chunkElems / 4; i++) { + IType4 in = rInput4[i]; + OType4 out; + ptx::mul_cvt_4x(out, in, block_scale_inverse_2x); + rOutput4[i] = out; + } + } else { + OType2 *rOutput2 = reinterpret_cast(&rOutput); +#pragma unroll + for (int32_t i = 0; i < CastTraits::chunkElems / 2; i++) { + IType2 in = rInput2[i]; + OType2 out; + ptx::mul_cvt_2x(out, in, block_scale_inverse_2x); + rOutput2[i] = out; + } + } + outputUnitType *output_units = + reinterpret_cast(output + coords.y * cols + coords.x); +#pragma unroll + for (int32_t j = 0; j < CastTraits::numOutUnitsPerChunk; j++) { + output_units[j] = rOutput[j]; + } + } + } + + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + __syncthreads(); + + int32_t warpId = threadIdx.z * CastTraits::warpLayout::N + threadIdx.y; + + block_coords.y = blockIdx.y * CastTraits::blockDimM; + block_coords.x = blockIdx.x * CastTraits::blockDimN; + + constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems; + using PreferredDataType = std::conditional_t< + stride_in_smem % 16 == 0, uint4, + std::conditional_t< + stride_in_smem % 8 == 0, uint2, + std::conditional_t>>>; + + int2 end_coords; + end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows); + end_coords.x = std::min((block_coords.x + CastTraits::blockDimN) / CastTraits::chunkElems, + scale_stride_rowwise); + int2 valid_coords; + valid_coords.y = end_coords.y - block_coords.y; + valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems); + + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::chunkElems); + + for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * 32) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::chunkElems); + + for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * 32) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +enum class ColwiseReduceMax : int32_t { + Atom = 0, + Red = 1, // it's actually the same to Atom + RedAsync = 2, + Redux = 3, + Num = 4 +}; + +// 32x32 +template +struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> { + static constexpr bool isRowwise = true; + static constexpr bool isColwise = true; + using IType = _IType; + using OType = _OType; + + static constexpr int32_t rowChunkElems = 32; + static constexpr int32_t colChunkElems = 32; + + using rowThreadLayout = Layout<32, 1>; // 32x1 + using colThreadLayout = Layout; // 1x32 + static_assert(rowThreadLayout::num == colThreadLayout::num, + "rowThreadLayout::num must be equal to colThreadLayout::num"); + static_assert(rowThreadLayout::num == 32, "rowThreadLayout::num must be 32"); + + using rowWarpDim = Layout; + using colWarpDim = Layout; + using warpDim = + Layout; + + static constexpr bool _tma_swizzle = true; + using warpLayout = Layout<1, 2>; + static_assert(_tma_swizzle ? (warpLayout::N == 2) : true); + static constexpr CUtensorMapSwizzle input_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + static constexpr CUtensorMapSwizzle output_swizzle_pattern = + _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B + : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE; + + using blockIterDim = Layout; + + using iterLayout = Layout<1, 4>; + using blockDIM = Layout; + + static constexpr int32_t numStages = 2; + + using inputUnitType = uint4; + static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType); + static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit; + // TODO: set condition for float + using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>; + using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>; + + using colIndexSwz = swz::Swizzle<5, 0, 5>; + + using rowOutputUnitType = uint4; + static constexpr int32_t rowNumOutUnitsPerChunk = + rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType); + static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType); + + using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>; + using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>; + + static constexpr bool _use_cvt_4x = true; + static constexpr bool _use_warp_specialization = false; + static constexpr bool _need_wait_group = iterLayout::num > numStages; + static constexpr bool _reuse_input_out_smem = false; + static_assert(_reuse_input_out_smem == false, "Just don't use it"); + static constexpr bool _cache_rowwise_scale_in_smem = true; + + static constexpr bool _colwise_source_coming_from_rowwise = true; + static constexpr ColwiseReduceMax _colwise_reduce_max = ColwiseReduceMax::Redux; + static_assert(_colwise_reduce_max != ColwiseReduceMax::RedAsync, + "It requires aligned smem pointer"); + + static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization; + static constexpr int32_t numThreads = numWarps * 32; + static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024"); + + static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType); + static constexpr size_t smemInputPerBlock = smemInputPerWarp * warpLayout::num; + + static constexpr size_t smemRowwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemRowwiseOutputPerBlock = smemRowwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemColwiseOutputPerWarp = warpDim::num * sizeof(OType); + static constexpr size_t smemColwiseOutputPerBlock = smemColwiseOutputPerWarp * warpLayout::num; + + static constexpr size_t smemInput = smemInputPerBlock * numStages; + static constexpr size_t smemRowwiseOutput = smemRowwiseOutputPerBlock * numStages; + static constexpr size_t smemColwiseOutput = smemColwiseOutputPerBlock * numStages; + + static constexpr size_t smem_rowwise_scale = + _cache_rowwise_scale_in_smem ? (blockDIM::M * (blockDIM::N / rowChunkElems) * sizeof(e8m0_t)) + : 0ul; + + using ColwiseReduceDataType = float; + static constexpr bool _need_smem_for_colwise_reduce = + _colwise_source_coming_from_rowwise; // && _colwise_reduce_max != ColwiseReduceMax::Redux; + static constexpr size_t smem_colwise_reduce = + _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul; + + static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul; + static constexpr size_t smem = _reuse_input_out_smem + ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + + smem_alignment + smem_rowwise_scale + smem_colwise_reduce) + : (smemInput + smemRowwiseOutput + smemColwiseOutput + + smem_alignment + smem_rowwise_scale + smem_colwise_reduce); +}; + +__device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) { + return (x + align - 1) & ~((align)-1); +} + +// 32x32 +template = 0, + std::enable_if_t = 0> +// __launch_bounds__(CastTraits::numThreads) +__global__ void quantize_mxfp8_kernel_cast_only( + const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_rowwise_output, + const __grid_constant__ CUtensorMap tensor_map_colwise_output, + e8m0_t *__restrict__ scales_rowwise, e8m0_t *__restrict__ scales_colwise, int32_t rows, + int32_t cols, int32_t scale_stride_rowwise, int32_t scale_stride_colwise) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using IType = typename CastTraits::IType; + using OType = typename CastTraits::OType; + using inputUnitType = typename CastTraits::inputUnitType; + using rowOutputUnitType = typename CastTraits::rowOutputUnitType; + using ColwiseReduceDataType = typename CastTraits::ColwiseReduceDataType; + + using IType2 = typename ptx::FPx2; + using OType2 = typename ptx::FPx2; + constexpr int32_t numItersIType2 = sizeof(inputUnitType) / sizeof(IType2); + + int32_t warpId = threadIdx.y; + int32_t leader = ptx::elect_one_sync(); + int2 block_coords; + block_coords.y = blockIdx.y * CastTraits::blockDIM::M; + block_coords.x = blockIdx.x * CastTraits::blockDIM::N; + + extern __shared__ char smem[]; + char *smemAligned = reinterpret_cast( + align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + + IType *sInput = reinterpret_cast(smemAligned); + inputUnitType *sInputUnit = reinterpret_cast(sInput); + + OType *sRowOutput = + reinterpret_cast(sInput + CastTraits::blockIterDim::num * CastTraits::numStages); + rowOutputUnitType *sRowOutputUnit = reinterpret_cast(sRowOutput); + + OType *sColOutput = + reinterpret_cast(sRowOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + rowOutputUnitType *sColOutputUnit = reinterpret_cast(sColOutput); + + e8m0_t *sRowwiseScale = nullptr; + ColwiseReduceDataType *sColwiseReduce = nullptr; + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + sRowwiseScale = reinterpret_cast(sColOutput + CastTraits::blockIterDim::num * + CastTraits::numStages); + if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce = reinterpret_cast( + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + sColwiseReduce += warpId * 32; + } + } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce = reinterpret_cast( + sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + sColwiseReduce += warpId * 32; + } + + // TODO: maybe we can assign a different barrier for each warp + __shared__ uint64_t ldg_producer[CastTraits::numStages], ldg_consumer[CastTraits::numStages]; + __shared__ uint64_t stg_producer[CastTraits::numStages], stg_consumer[CastTraits::numStages]; + + if (warpId == 0 && leader) { +#pragma unroll + for (int32_t i = 0; i < CastTraits::numStages; i++) { + ptx::mbarrier_init(&ldg_producer[i], 1); + ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * 32); + ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * 32); + ptx::mbarrier_init(&stg_consumer[i], 1); + } + ptx::fence_mbarrier_init_release_cluster(); + } + __syncthreads(); + + if (warpId == CastTraits::warpLayout::num) { + if (leader) { + PipeState write_state; +#pragma unroll 1 + for (int32_t iter = 0; iter < CastTraits::iterLayout::num; iter++) { + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDim::M; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N; + + if (coords.x >= cols || coords.y >= rows) { + break; + } + + ptx::mbarrier_wait_parity(&ldg_consumer[write_state.index()], write_state.phase()); + + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(sInput + + write_state.index() * CastTraits::blockIterDim::num), + reinterpret_cast(&tensor_map_input), static_cast(coords.x), + static_cast(coords.y), &ldg_producer[write_state.index()]); + ptx::mbarrier_arrive_expect_tx(&ldg_producer[write_state.index()], + CastTraits::blockIterDim::num * sizeof(IType)); + write_state++; + } + } + } else if (warpId == CastTraits::warpLayout::num + 1) { + if (leader) { + PipeState read_state; + +#pragma unroll 1 + for (int32_t iter = 0; iter < CastTraits::numStages - 1; iter++) { + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDim::M; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N; + + size_t gmem_offset = + static_cast(read_state.index()) * CastTraits::blockIterDim::num; + + if (coords.x >= cols || coords.y >= rows) { + break; + } + + ptx::mbarrier_wait_parity(&stg_producer[read_state.index()], read_state.phase()); + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_rowwise_output), + static_cast(coords.x), static_cast(coords.y), + reinterpret_cast(sRowOutput + gmem_offset)); + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_colwise_output), + static_cast(coords.x), static_cast(coords.y), + reinterpret_cast(sColOutput + gmem_offset)); + ptx::cp_async_bulk_commit_group(); + read_state++; + } + +#pragma unroll 1 + for (int32_t iter = CastTraits::numStages - 1; iter < CastTraits::iterLayout::num; iter++) { + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDim::M; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N; + + size_t gmem_offset = + static_cast(read_state.index()) * CastTraits::blockIterDim::num; + + if (coords.x >= cols || coords.y >= rows) { + break; + } + + ptx::mbarrier_wait_parity(&stg_producer[read_state.index()], read_state.phase()); + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_rowwise_output), + static_cast(coords.x), static_cast(coords.y), + reinterpret_cast(sRowOutput + gmem_offset)); + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_colwise_output), + static_cast(coords.x), static_cast(coords.y), + reinterpret_cast(sColOutput + gmem_offset)); + ptx::cp_async_bulk_commit_group(); + read_state++; + + ptx::cp_async_bulk_wait_group_read(); + ptx::mbarrier_arrive_expect_tx(&stg_consumer[read_state.index()], 0u); + } + } + ptx::cp_async_bulk_wait_group_read<0>(); + } else { + PipeState read_state; + + int2 warp_coords; + warp_coords.y = (warpId / CastTraits::warpLayout::N) * CastTraits::warpDim::M; + warp_coords.x = (warpId % CastTraits::warpLayout::N) * CastTraits::warpDim::N; + + int32_t warp_base_offset = warp_coords.y * CastTraits::blockIterDim::N + warp_coords.x; + + int32_t thread_base_offset = + (threadIdx.x / CastTraits::rowThreadLayout::N) * + (CastTraits::blockIterDim::N / CastTraits::rowNumElemsPerUnit) + + (threadIdx.x % CastTraits::rowThreadLayout::N) * + (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit); + + size_t rowwise_scale_base_offset = + (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) * + static_cast(scale_stride_rowwise) + + (block_coords.x + warp_coords.x + + (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) / + CastTraits::rowChunkElems; + size_t colwise_scale_base_offset = + ((block_coords.y + warp_coords.y + + (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / + CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N)); + + constexpr int32_t rowwise_scale_stride_in_smem = + CastTraits::blockDIM::N / CastTraits::rowChunkElems; + int32_t rowwise_scale_smem_base_offset = + (warpId / CastTraits::warpLayout::N) * CastTraits::warpDim::M * + rowwise_scale_stride_in_smem + + (warpId % CastTraits::warpLayout::N) * + (CastTraits::warpDim::N / CastTraits::rowChunkElems) + + (threadIdx.x / CastTraits::rowThreadLayout::N) * rowwise_scale_stride_in_smem + + (threadIdx.x % CastTraits::rowThreadLayout::N); + +#pragma unroll 1 + for (int32_t iter = 0; iter < CastTraits::iterLayout::num; iter++) { + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + if (block_coords.x + iter_n * CastTraits::blockIterDim::N >= cols || + block_coords.y + iter_m * CastTraits::blockIterDim::M >= rows) { + break; + } + + ptx::mbarrier_wait_parity(&ldg_producer[read_state.index()], read_state.phase()); + + { + int32_t warp_offset = warp_base_offset + read_state.index() * CastTraits::blockIterDim::num; + static_assert(CastTraits::_colwise_source_coming_from_rowwise); + if constexpr (CastTraits::_colwise_source_coming_from_rowwise) { + if constexpr (CastTraits::_need_smem_for_colwise_reduce && + CastTraits::_colwise_reduce_max != ColwiseReduceMax::Redux) { + sColwiseReduce[threadIdx.x] = 0; + } + + IType rInput[CastTraits::rowChunkElems]; + { + inputUnitType *rInputUnit = reinterpret_cast(rInput); + int32_t base = thread_base_offset + warp_offset / CastTraits::rowNumElemsPerUnit; +#pragma unroll + for (int32_t i = 0; i < CastTraits::rowNumUnitsPerChunk; i++) { + rInputUnit[i] = sInputUnit[CastTraits::inputUnitSwz::swz(base + i)]; + } + ptx::mbarrier_arrive_expect_tx(&ldg_consumer[read_state.index()], 0u); + } + + if constexpr (std::is_same_v) { + } else { + static_assert(CastTraits::_colwise_reduce_max == ColwiseReduceMax::Redux, + "Only Redux is implemented"); + + float row_scale_inverse; + + IType2 *rInput2 = reinterpret_cast(&rInput); + float2 *sColwiseReduce_2x = reinterpret_cast(sColwiseReduce); + + IType2 row_amax2{0.0f, 0.0f}; +#pragma unroll + for (int32_t i = 0; i < CastTraits::rowChunkElems / 2; i++) { + ptx::abs_max_2x(row_amax2, row_amax2, rInput2[i]); + + float2 values = ptx::up_cast(rInput2[i]); + + float2 amaxs; + ptx::reduce_sync_max_abs_f32(amaxs.x, values.x); + ptx::reduce_sync_max_abs_f32(amaxs.y, values.y); + if (leader) { + sColwiseReduce_2x[i] = amaxs; + } + } + { + IType row_amax = ptx::get_amax(row_amax2.x, row_amax2.y); + e8m0_t row_biased_exponent = to_e8m0(row_amax); + row_scale_inverse = ptx::exp2f_rcp(row_biased_exponent); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + int32_t rowwise_scale_offset = + rowwise_scale_smem_base_offset + + iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; + } else { + size_t rowwise_scale_offset = + rowwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M) * + static_cast(scale_stride_rowwise) + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + scales_rowwise[rowwise_scale_offset] = row_biased_exponent; + } + } + { + __syncwarp(); + float col_amax = sColwiseReduce[threadIdx.x]; + e8m0_t col_biased_exponent = to_e8m0(col_amax); + float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); + sColwiseReduce[threadIdx.x] = col_scale_inverse; + size_t colwise_scale_offset = + colwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + iter_n * CastTraits::blockIterDim::N; + scales_colwise[colwise_scale_offset] = col_biased_exponent; + __syncwarp(); + } + // rowwise & colwise scaling + { + rowOutputUnitType rRowOutputUnit[CastTraits::rowNumOutUnitsPerChunk]; + rowOutputUnitType rColOutputUnit[CastTraits::rowNumOutUnitsPerChunk]; + + ptx::floatx2 row_scale_inverse_2{row_scale_inverse, row_scale_inverse}; + if constexpr (CastTraits::_use_cvt_4x) { + using OType4 = ptx::FPx4; + using IType4 = ptx::FPx4; + + ptx::floatx4 col_scale_inverse_4[2]; + ptx::floatx4 *sColwiseScale4x = reinterpret_cast(sColwiseReduce); + col_scale_inverse_4[0] = sColwiseScale4x[0]; + + IType4 *rInput4 = reinterpret_cast(&rInput); + OType4 *rRowOutput4 = reinterpret_cast(&rRowOutputUnit); + OType4 *rColOutput4 = reinterpret_cast(&rColOutputUnit); +#pragma unroll + for (int32_t i = 1; i < CastTraits::rowChunkElems / 4; i++) { + { + col_scale_inverse_4[i % 2] = sColwiseScale4x[i]; + } + + IType4 in = rInput4[i - 1]; + ptx::floatx4 in_fp4 = ptx::up_cast(in); + + OType4 row_out; + ptx::mul_cvt_4x(row_out, in_fp4, row_scale_inverse_2); + rRowOutput4[i - 1] = row_out; + + OType4 col_out; + ptx::mul_cvt_4x(col_out, in_fp4, col_scale_inverse_4[(i - 1) % 2]); + rColOutput4[i - 1] = col_out; + } + { + constexpr int32_t i = (CastTraits::rowChunkElems / 4) - 1; + IType4 in = rInput4[i]; + ptx::floatx4 in_fp4 = ptx::up_cast(in); + + OType4 row_out; + ptx::mul_cvt_4x(row_out, in_fp4, row_scale_inverse_2); + rRowOutput4[i] = row_out; + + OType4 col_out; + ptx::mul_cvt_4x(col_out, in_fp4, col_scale_inverse_4[i % 2]); + rColOutput4[i] = col_out; + } + } else { + ptx::floatx2 col_scale_inverse_2[2]; + ptx::floatx2 *sColwiseScale2x = reinterpret_cast(sColwiseReduce); + col_scale_inverse_2[0] = sColwiseScale2x[0]; + + IType2 *rInput2 = reinterpret_cast(&rInput); + OType2 *rRowOutput2 = reinterpret_cast(&rRowOutputUnit); + OType2 *rColOutput2 = reinterpret_cast(&rColOutputUnit); +#pragma unroll + for (int32_t i = 1; i < CastTraits::rowChunkElems / 2; i++) { + { + col_scale_inverse_2[i % 2] = sColwiseScale2x[i]; + } + + IType2 in = rInput2[i - 1]; + ptx::floatx2 in_fp2 = ptx::up_cast(in); + + OType2 row_out; + mul_cvt_2x(row_out, in_fp2, row_scale_inverse_2); + rRowOutput2[i - 1] = row_out; + + OType2 col_out; + mul_cvt_2x(col_out, in_fp2, col_scale_inverse_2[(i - 1) % 2]); + rColOutput2[i - 1] = col_out; + } + { + constexpr int32_t i = (CastTraits::rowChunkElems / 2) - 1; + IType2 in = rInput2[i]; + ptx::floatx2 in_fp2 = ptx::up_cast(in); + + OType2 row_out; + mul_cvt_2x(row_out, in_fp2, row_scale_inverse_2); + rRowOutput2[i] = row_out; + + OType2 col_out; + mul_cvt_2x(col_out, in_fp2, col_scale_inverse_2[i % 2]); + rColOutput2[i] = col_out; + } + } + { + ptx::mbarrier_wait_parity(&stg_consumer[read_state.index()], + read_state.phase() ^ 1); + + int32_t base = thread_base_offset / (CastTraits::rowOutNumElemsPerUnit / + CastTraits::rowNumElemsPerUnit) + + warp_offset / CastTraits::rowOutNumElemsPerUnit; +#pragma unroll + for (int32_t i = 0; i < CastTraits::rowNumOutUnitsPerChunk; i++) { + int32_t offset = CastTraits::rowOutputChunkSwz::swz(base + i); + sRowOutputUnit[offset] = rRowOutputUnit[i]; + sColOutputUnit[offset] = rColOutputUnit[i]; + } + } + } + } + } + } + ptx::fence_proxy_async_shared_cta(); + + ptx::mbarrier_arrive_expect_tx(&stg_producer[read_state.index()], 0u); + read_state++; + } + + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u); + + constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; + using PreferredDataType = std::conditional_t< + stride_in_smem % 16 == 0, uint4, + std::conditional_t< + stride_in_smem % 8 == 0, uint2, + std::conditional_t>>>; + + int2 end_coords; + end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); + end_coords.x = + std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); + int2 valid_coords; + valid_coords.y = end_coords.y - block_coords.y; + valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); + + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * 32) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * 32) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = + sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + } +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +template = 0, + std::enable_if_t = 0> +__global__ void quantize_mxfp8_kernel_cast_only( + const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_rowwise_output, + const __grid_constant__ CUtensorMap tensor_map_colwise_output, e8m0_t *scales_rowwise, + e8m0_t *scales_colwise, int32_t rows, int32_t cols, int32_t scale_stride_rowwise, + int32_t scale_stride_colwise) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using IType = typename CastTraits::IType; + using OType = typename CastTraits::OType; + using inputUnitType = typename CastTraits::inputUnitType; + using rowOutputUnitType = typename CastTraits::rowOutputUnitType; + using ColwiseReduceDataType = typename CastTraits::ColwiseReduceDataType; + + using IType2 = typename ptx::FPx2; + using OType2 = typename ptx::FPx2; + + int32_t warpId = threadIdx.y; + int32_t leader = ptx::elect_one_sync(); + int2 block_coords; + block_coords.y = blockIdx.y * CastTraits::blockDIM::M; + block_coords.x = blockIdx.x * CastTraits::blockDIM::N; + + extern __shared__ char smem[]; + char *smemAligned = reinterpret_cast( + align_to(reinterpret_cast(smem), CastTraits::smem_alignment)); + IType *sInput = reinterpret_cast(smemAligned); + inputUnitType *sInputUnit = reinterpret_cast(sInput); + + OType *sRowOutput = + reinterpret_cast(sInput + CastTraits::blockIterDim::num * CastTraits::numStages); + rowOutputUnitType *sRowOutputUnit = reinterpret_cast(sRowOutput); + + // colwise output will reuse input buffer + OType *sColOutput; + e8m0_t *sRowwiseScale = nullptr; + ColwiseReduceDataType *sColwiseReduce = nullptr; + if constexpr (CastTraits::_reuse_input_out_smem) { + sColOutput = reinterpret_cast(sInput); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + sRowwiseScale = reinterpret_cast(sRowOutput + CastTraits::blockIterDim::num * + CastTraits::numStages); + if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce = reinterpret_cast( + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + } + } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce = reinterpret_cast( + sRowOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + } + } else { + sColOutput = reinterpret_cast(sRowOutput + + CastTraits::blockIterDim::num * CastTraits::numStages); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + sRowwiseScale = reinterpret_cast(sColOutput + CastTraits::blockIterDim::num * + CastTraits::numStages); + if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce = reinterpret_cast( + sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t)); + } + } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce = reinterpret_cast( + sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages); + } + } + rowOutputUnitType *sColOutputUnit = reinterpret_cast(sColOutput); + + if constexpr (CastTraits::_need_smem_for_colwise_reduce) { + sColwiseReduce += warpId * 32; + } + + __shared__ uint64_t producer[CastTraits::numStages]; + uint64_t *colwise_reduce_barrier = nullptr; + if constexpr (CastTraits::_colwise_source_coming_from_rowwise && + CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { + __shared__ uint64_t colwise_reduce_bar[CastTraits::warpLayout::num]; + colwise_reduce_barrier = &colwise_reduce_bar[warpId]; + } + + if (leader) { + if (warpId == 0) { +#pragma unroll + for (int32_t i = 0; i < CastTraits::numStages; i++) { + ptx::mbarrier_init(&producer[i], 1); + } + } + if constexpr (CastTraits::_colwise_source_coming_from_rowwise && + CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { + ptx::mbarrier_init(colwise_reduce_barrier, 32); + } + + ptx::fence_mbarrier_init_release_cluster(); + } + __syncthreads(); + + PipeState states; + + int2 warp_coords; + warp_coords.y = (warpId / CastTraits::warpLayout::N) * CastTraits::warpDim::M; + warp_coords.x = (warpId % CastTraits::warpLayout::N) * CastTraits::warpDim::N; + + int32_t warp_base_offset = warp_coords.y * CastTraits::blockIterDim::N + warp_coords.x; + + int32_t thread_base_offset = (threadIdx.x / CastTraits::rowThreadLayout::N) * + (CastTraits::blockIterDim::N / CastTraits::rowNumElemsPerUnit) + + (threadIdx.x % CastTraits::rowThreadLayout::N) * + (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit); + + size_t rowwise_scale_base_offset = + (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) * + static_cast(scale_stride_rowwise) + + (block_coords.x + warp_coords.x + + (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) / + CastTraits::rowChunkElems; + size_t colwise_scale_base_offset = + ((block_coords.y + warp_coords.y + + (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) / + CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N)); + + constexpr int32_t rowwise_scale_stride_in_smem = + CastTraits::blockDIM::N / CastTraits::rowChunkElems; + int32_t rowwise_scale_smem_base_offset = + (warpId / CastTraits::warpLayout::N) * CastTraits::warpDim::M * rowwise_scale_stride_in_smem + + (warpId % CastTraits::warpLayout::N) * (CastTraits::warpDim::N / CastTraits::rowChunkElems) + + (threadIdx.x / CastTraits::rowThreadLayout::N) * rowwise_scale_stride_in_smem + + (threadIdx.x % CastTraits::rowThreadLayout::N); + + if (warpId == 0 && leader) { +#pragma unroll 1 + for (int32_t iter = 0; iter < CastTraits::numStages - 1; iter++) { + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDim::M; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N; + if (coords.x >= cols || coords.y >= rows) { + break; + } + + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(sInput + iter * CastTraits::blockIterDim::num), + reinterpret_cast(&tensor_map_input), static_cast(coords.x), + static_cast(coords.y), &producer[iter]); + ptx::mbarrier_arrive_expect_tx(&producer[iter], + CastTraits::blockIterDim::num * sizeof(IType)); + } + } +#pragma unroll 1 + for (int32_t iter = 0; iter < CastTraits::iterLayout::num; iter++) { + { + int32_t next = iter + (CastTraits::numStages - 1); + int32_t next_stage = next % CastTraits::numStages; + int32_t iter_m = next / CastTraits::iterLayout::N; + int32_t iter_n = next % CastTraits::iterLayout::N; + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDim::M; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N; + if (coords.x < cols && coords.y < rows) { + if (warpId == 0 && leader) { + if constexpr (CastTraits::_need_wait_group) { + ptx::cp_async_bulk_wait_group_read(); + } + + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(sInput + next_stage * CastTraits::blockIterDim::num), + reinterpret_cast(&tensor_map_input), + static_cast(coords.x), static_cast(coords.y), + &producer[next_stage]); + ptx::mbarrier_arrive_expect_tx(&producer[next_stage], + CastTraits::blockIterDim::num * sizeof(IType)); + } + } + } + + int32_t iter_m = iter / CastTraits::iterLayout::N; + int32_t iter_n = iter % CastTraits::iterLayout::N; + + int2 coords; + coords.y = block_coords.y + iter_m * CastTraits::blockIterDim::M; + coords.x = block_coords.x + iter_n * CastTraits::blockIterDim::N; + + if (coords.x >= cols || coords.y >= rows) { + break; + } + + ptx::mbarrier_wait_parity(&producer[states.index()], states.phase()); + + int32_t warp_offset = warp_base_offset + states.index() * CastTraits::blockIterDim::num; + static_assert(CastTraits::_colwise_source_coming_from_rowwise); + if constexpr (CastTraits::_colwise_source_coming_from_rowwise) { + if constexpr (CastTraits::_need_smem_for_colwise_reduce && + CastTraits::_colwise_reduce_max != ColwiseReduceMax::Redux) { + sColwiseReduce[threadIdx.x] = 0.0f; + } + + IType rInput[CastTraits::rowChunkElems]; + { + inputUnitType *rInputUnit = reinterpret_cast(rInput); + int32_t base = thread_base_offset + warp_offset / CastTraits::rowNumElemsPerUnit; +#pragma unroll + for (int32_t i = 0; i < CastTraits::rowNumUnitsPerChunk; i++) { + rInputUnit[i] = sInputUnit[CastTraits::inputUnitSwz::swz(base + i)]; + } + } + + if constexpr (std::is_same_v) { + if constexpr (CastTraits::_colwise_reduce_max == ColwiseReduceMax::Atom || + CastTraits::_colwise_reduce_max == ColwiseReduceMax::Red) { + } else if constexpr (CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) { + } else if constexpr (CastTraits::_colwise_reduce_max == ColwiseReduceMax::Redux) { + } + } else { + float row_scale_inverse; + static_assert(CastTraits::_colwise_reduce_max == ColwiseReduceMax::Redux); + if constexpr (CastTraits::_colwise_reduce_max == ColwiseReduceMax::Redux) { + IType2 *rInput2 = reinterpret_cast(&rInput); + float2 *sColwiseReduce_2x = reinterpret_cast(sColwiseReduce); + + IType2 row_amax2{0.0f, 0.0f}; +#pragma unroll + for (int32_t i = 0; i < CastTraits::rowChunkElems / 2; i++) { + ptx::abs_max_2x(row_amax2, row_amax2, rInput2[i]); + + ptx::floatx2 values = ptx::up_cast(rInput2[i]); + + float2 amaxs; + ptx::reduce_sync_max_abs_f32(amaxs.x, values.x); + ptx::reduce_sync_max_abs_f32(amaxs.y, values.y); + + if (leader) { + sColwiseReduce_2x[i] = amaxs; + } + } + + { + IType row_amax = ptx::get_amax(row_amax2.x, row_amax2.y); + e8m0_t row_biased_exponent = to_e8m0(row_amax); + row_scale_inverse = ptx::exp2f_rcp(row_biased_exponent); + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + int32_t rowwise_scale_offset = + rowwise_scale_smem_base_offset + + iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + sRowwiseScale[rowwise_scale_offset] = row_biased_exponent; + } else { + size_t rowwise_scale_offset = + rowwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M) * + static_cast(scale_stride_rowwise) + + iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems); + scales_rowwise[rowwise_scale_offset] = row_biased_exponent; + } + } + { + __syncwarp(); + float col_amax = sColwiseReduce[threadIdx.x]; + e8m0_t col_biased_exponent = to_e8m0(col_amax); + float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); + sColwiseReduce[threadIdx.x] = col_scale_inverse; + size_t colwise_scale_offset = + colwise_scale_base_offset + + iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) * + static_cast(scale_stride_colwise) + + iter_n * CastTraits::blockIterDim::N; + scales_colwise[colwise_scale_offset] = col_biased_exponent; + __syncwarp(); + } + } + // row & colwise + { + rowOutputUnitType rRowOutputUnit[CastTraits::rowNumOutUnitsPerChunk]; + rowOutputUnitType rColOutputUnit[CastTraits::rowNumOutUnitsPerChunk]; + + ptx::floatx2 row_scale_inverse_2{row_scale_inverse, row_scale_inverse}; + if constexpr (CastTraits::_use_cvt_4x) { + using OType4 = ptx::FPx4; + using IType4 = ptx::FPx4; + + ptx::floatx4 col_scale_inverse_4[2]; + ptx::floatx4 *sColwiseScale4x = reinterpret_cast(sColwiseReduce); + col_scale_inverse_4[0] = sColwiseScale4x[0]; + + IType4 *rInput4 = reinterpret_cast(&rInput); + OType4 *rRowOutput4 = reinterpret_cast(&rRowOutputUnit); + OType4 *rColOutput4 = reinterpret_cast(&rColOutputUnit); +#pragma unroll + for (int32_t i = 1; i < CastTraits::rowChunkElems / 4; i++) { + { + col_scale_inverse_4[i % 2] = sColwiseScale4x[i]; + } + + IType4 in = rInput4[i - 1]; + ptx::floatx4 in_fp4 = ptx::up_cast(in); + + OType4 row_out; + ptx::mul_cvt_4x(row_out, in_fp4, row_scale_inverse_2); + rRowOutput4[i - 1] = row_out; + + OType4 col_out; + ptx::mul_cvt_4x(col_out, in_fp4, col_scale_inverse_4[(i - 1) % 2]); + rColOutput4[i - 1] = col_out; + } + { + constexpr int32_t i = (CastTraits::rowChunkElems / 4) - 1; + IType4 in = rInput4[i]; + ptx::floatx4 in_fp4 = ptx::up_cast(in); + + OType4 row_out; + ptx::mul_cvt_4x(row_out, in_fp4, row_scale_inverse_2); + rRowOutput4[i] = row_out; + + OType4 col_out; + ptx::mul_cvt_4x(col_out, in_fp4, col_scale_inverse_4[i % 2]); + rColOutput4[i] = col_out; + } + } else { + ptx::floatx2 col_scale_inverse_2[2]; + ptx::floatx2 *sColwiseScale2x = reinterpret_cast(sColwiseReduce); + col_scale_inverse_2[0] = sColwiseScale2x[0]; + + IType2 *rInput2 = reinterpret_cast(&rInput); + OType2 *rRowOutput2 = reinterpret_cast(&rRowOutputUnit); + OType2 *rColOutput2 = reinterpret_cast(&rColOutputUnit); +#pragma unroll + for (int32_t i = 1; i < CastTraits::rowChunkElems / 2; i++) { + { + col_scale_inverse_2[i % 2] = sColwiseScale2x[i]; + } + + IType2 in = rInput2[i - 1]; + ptx::floatx2 in_fp2 = ptx::up_cast(in); + + OType2 row_out; + mul_cvt_2x(row_out, in_fp2, row_scale_inverse_2); + rRowOutput2[i - 1] = row_out; + + OType2 col_out; + mul_cvt_2x(col_out, in_fp2, col_scale_inverse_2[(i - 1) % 2]); + rColOutput2[i - 1] = col_out; + } + { + constexpr int32_t i = (CastTraits::rowChunkElems / 2) - 1; + IType2 in = rInput2[i]; + ptx::floatx2 in_fp2 = ptx::up_cast(in); + + OType2 row_out; + mul_cvt_2x(row_out, in_fp2, row_scale_inverse_2); + rRowOutput2[i] = row_out; + + OType2 col_out; + mul_cvt_2x(col_out, in_fp2, col_scale_inverse_2[i % 2]); + rColOutput2[i] = col_out; + } + } + + { + int32_t base = thread_base_offset / (CastTraits::rowOutNumElemsPerUnit / + CastTraits::rowNumElemsPerUnit) + + warp_offset / CastTraits::rowOutNumElemsPerUnit; +#pragma unroll + for (int32_t i = 0; i < CastTraits::rowNumOutUnitsPerChunk; i++) { + int32_t offset = CastTraits::rowOutputChunkSwz::swz(base + i); + sRowOutputUnit[offset] = rRowOutputUnit[i]; + sColOutputUnit[offset] = rColOutputUnit[i]; + } + } + } + } + } + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); + + if (warpId == 0 && leader) { + size_t gmem_offset = static_cast(states.index()) * CastTraits::blockIterDim::num; + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_rowwise_output), + static_cast(coords.x), static_cast(coords.y), + reinterpret_cast(sRowOutput + gmem_offset)); + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_colwise_output), + static_cast(coords.x), static_cast(coords.y), + reinterpret_cast(sColOutput + gmem_offset)); + ptx::cp_async_bulk_commit_group(); + } + states++; + } + + if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { + constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems; + using PreferredDataType = std::conditional_t< + stride_in_smem % 16 == 0, uint4, + std::conditional_t< + stride_in_smem % 8 == 0, uint2, + std::conditional_t>>>; + + int2 end_coords; + end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows); + end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems, + scale_stride_rowwise); + int2 valid_coords; + valid_coords.y = end_coords.y - block_coords.y; + valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems); + + if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) { + using DataType = int32_t; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * 32) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + } + } else { + using DataType = PreferredDataType; + constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t); + constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group; + + int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group); + int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group; + + DataType *sScales = reinterpret_cast(sRowwiseScale); + DataType *gScales = + reinterpret_cast(scales_rowwise + block_coords.y * scale_stride_rowwise + + block_coords.x / CastTraits::rowChunkElems); + + for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row); + i += CastTraits::warpLayout::num * 32) { + int32_t row = i / num_threads_per_row; + int32_t col = i % num_threads_per_row; + gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col]; + } + } + } + + ptx::cp_async_bulk_wait_group_read<0>(); + +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +} // namespace specialized +} // namespace quantize_kernel +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // #ifndef TRANSFORMER_ENGINE_SPECIALIZED_QUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh b/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh new file mode 100644 index 0000000000..5073de5b11 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh @@ -0,0 +1,61 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file state_counter.cuh + * \brief CUDA kernels to count state. + */ + +#ifndef TRANSFORMER_ENGINE_SPECIALIZED_STATE_COUNTER_CUH_ +#define TRANSFORMER_ENGINE_SPECIALIZED_STATE_COUNTER_CUH_ + +#include + +namespace transformer_engine { + +template +struct PipeState { + int2 _storage; // x: index, y: phase + + __device__ __forceinline__ PipeState() : _storage{0, 0} { + if constexpr (Flip) { + _storage.y ^= 1; + } + } + + __device__ __forceinline__ int32_t index() const { return _storage.x; } + + __device__ __forceinline__ int32_t phase() const { return _storage.y; } + + __device__ __forceinline__ void operator++(int32_t) { + if constexpr (numStages > 0) { + _storage.x++; + if (_storage.x == numStages) { + _storage.x = 0; + _storage.y ^= 1; + } + } + } +}; + +template +struct PipeStateCounter { + int32_t _counter; + + __device__ __forceinline__ PipeStateCounter() : _counter(0) {} + + __device__ __forceinline__ int32_t index() const { return _counter; } + + __device__ __forceinline__ void operator++(int32_t) { + if constexpr (numStages > 0) { + _counter++; + _counter = _counter == numStages ? 0 : _counter; + } + } +}; + +} // namespace transformer_engine + +#endif // #ifndef TRANSFORMER_ENGINE_SPECIALIZED_STATE_COUNTER_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh b/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh new file mode 100644 index 0000000000..63fc1d6bd9 --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh @@ -0,0 +1,90 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file swizzle.cuh + * \brief CUDA kernels to swizzle. + */ + +#ifndef TRANSFORMER_ENGINE_SPECIALIZED_SWIZZLE_CUH_ +#define TRANSFORMER_ENGINE_SPECIALIZED_SWIZZLE_CUH_ + +#include +#include + +namespace transformer_engine { +namespace swz { + +template +struct C { + using type = C; + static constexpr auto value = v; + using value_type = decltype(v); + + __device__ __host__ __forceinline__ constexpr operator value_type() const noexcept { + return value; + } +}; + +template +using constant = C; + +template +__host__ __device__ __forceinline__ constexpr T shiftr(T x) { + if constexpr (std::is_same_v) { + return x >> s; + } else if constexpr (std::is_same_v) { + if constexpr (s >= 0) { + return x >> s; + } else { + return x << -s; + } + } +} + +template +struct Swizzle { + static constexpr int32_t num_bits = BBits; // number of rows + static constexpr int32_t num_base = MBase; // number of elements within a chunk + static constexpr int32_t num_shft = SShift; // number of columns, at the granularity of a chunk + + static_assert(num_base >= 0, "MBase must be non-negative"); + static_assert(num_bits >= 0, "BBits must be non-negative"); + static_assert(abs(num_shft) >= num_bits, "abs(SShift) must be greater than or equal to num_bits"); + + using bit_mask = constant; + using yyy_mask = + constant; + using zzz_mask = + constant; + using msk_shft = constant; + static constexpr int32_t swz_code = int32_t(yyy_mask{} | zzz_mask{}); + + template + __host__ __device__ __forceinline__ constexpr static int32_t apply(Offset const &offset) { + return offset ^ + shiftr(offset & yyy_mask{}); + } + + __host__ __device__ __forceinline__ constexpr static int32_t swz(int32_t const &offset) { + return apply(offset); + } +}; + +struct Linear { + template + __host__ __device__ __forceinline__ constexpr static int32_t apply(Offset const &offset) { + return offset; + } + + __host__ __device__ __forceinline__ constexpr static int32_t swz(int32_t const &offset) { + return offset; + } +}; + +} // namespace swz +} // namespace transformer_engine + +#endif // #ifndef TRANSFORMER_ENGINE_SPECIALIZED_SWIZZLE_CUH_ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 6605d9cad1..754cbd900a 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -826,6 +826,685 @@ __device__ __forceinline__ void abs_max_2x(fp16x2 &dst, const fp16x2 &p1, const #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 890) } +__device__ __forceinline__ int32_t elect_one_sync(uint32_t mask = 0xFFFFFFFFu) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + int32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %px; \n" + "elect.sync _|%px, %1; \n" + "selp.b32 %0, 1, 0, %px; \n" + "\n\t}" + : "=r"(pred) + : "r"(mask)); + return pred; +#else + NVTE_DEVICE_ERROR("elect_one_sync is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void numbered_barrier_sync(uint32_t num_threads, + uint32_t barrier_id = 1u) { + asm volatile("bar.sync %0, %1;\n" ::"r"(barrier_id), "r"(num_threads)); +} + +__device__ __forceinline__ void fma_f32_f16(float &out, uint16_t const &a, uint16_t const &b, + float const &c = 0.0f) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile("fma.rn.f32.f16 %0, %1, %2, %3;" : "=f"(out) : "h"(a), "h"(b), "f"(c) : "memory"); +#else + NVTE_DEVICE_ERROR("fma_f32_f16 is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void fma_f32_bf16(float &out, uint16_t const &a, uint16_t const &b, + float const &c = 0.0f) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile("fma.rn.f32.bf16 %0, %1, %2, %3;" : "=f"(out) : "h"(a), "h"(b), "f"(c) : "memory"); +#else + NVTE_DEVICE_ERROR("fma_f32_bf16 is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void reduce_sync_max_abs_f32(float &out, float const &in) { +#if ((__CUDA_ARCH_HAS_FEATURE__(SM100_ALL)) || (__CUDA_ARCH_HAS_FEATURE__(SM101_ALL)) || \ + (__CUDA_ARCH_HAS_FEATURE__(SM120_ALL))) + asm volatile("redux.sync.max.abs.f32 %0, %1, 0xFFFFFFFF;" : "=f"(out) : "f"(in)); +#else + asm volatile( + "{\n\t" + ".reg.b32 val;\n" + "abs.f32 val, %1;\n" + "redux.sync.max.u32 %0, val, 0xFFFFFFFF;\n" + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "f"(in)); +#endif +} + +__device__ __forceinline__ bf16 get_amax(bf16 a, bf16 b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + bf16 r; + asm volatile("max.xorsign.abs.bf16 %0, %1, %2;" + : "=h"(*reinterpret_cast(&r)) + : "h"(*reinterpret_cast(&a)), "h"(*reinterpret_cast(&b))); + return r; +#else + NVTE_DEVICE_ERROR("get_amax is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ fp16 get_amax(fp16 a, fp16 b) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + fp16 r; + asm volatile("max.xorsign.abs.f16 %0, %1, %2;" + : "=h"(*reinterpret_cast(&r)) + : "h"(*reinterpret_cast(&a)), "h"(*reinterpret_cast(&b))); + return r; +#else + NVTE_DEVICE_ERROR("get_amax is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const bf16x4 &in, + const ptx::floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::bf16x2 const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "prmt.b32 val2, 0x0, %1, 0x7632;\n\t" + "prmt.b32 val1, 0x0, %1, 0x5410;\n\t" + "prmt.b32 val4, 0x0, %2, 0x7632;\n\t" + "prmt.b32 val3, 0x0, %2, 0x5410;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %3, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e4m3x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale)), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const bf16x4 &in, const floatx4 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::bf16x2 const *in2 = reinterpret_cast(&in); + ptx::floatx2 const *scale2 = reinterpret_cast(&scale); + asm volatile( + "{\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "prmt.b32 val2, 0x0, %1, 0x7632;\n\t" + "prmt.b32 val1, 0x0, %1, 0x5410;\n\t" + "prmt.b32 val4, 0x0, %2, 0x7632;\n\t" + "prmt.b32 val3, 0x0, %2, 0x5410;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %4, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e4m3x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale2[0])), + "l"(reinterpret_cast(scale2[1])), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const bf16x4 &in, + const ptx::floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::bf16x2 const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "prmt.b32 val2, 0x0, %1, 0x7632;\n\t" + "prmt.b32 val1, 0x0, %1, 0x5410;\n\t" + "prmt.b32 val4, 0x0, %2, 0x7632;\n\t" + "prmt.b32 val3, 0x0, %2, 0x5410;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %3, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e5m2x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale)), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const bf16x4 &in, const floatx4 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::bf16x2 const *in2 = reinterpret_cast(&in); + ptx::floatx2 const *scale2 = reinterpret_cast(&scale); + asm volatile( + "{\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "prmt.b32 val2, 0x0, %1, 0x7632;\n\t" + "prmt.b32 val1, 0x0, %1, 0x5410;\n\t" + "prmt.b32 val4, 0x0, %2, 0x7632;\n\t" + "prmt.b32 val3, 0x0, %2, 0x5410;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %4, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e5m2x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale2[0])), + "l"(reinterpret_cast(scale2[1])), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const fp16x4 &in, + const ptx::floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::fp16x2 const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + ".reg.b16 val1_f16;\n\t" + ".reg.b16 val2_f16;\n\t" + ".reg.b16 val3_f16;\n\t" + ".reg.b16 val4_f16;\n\t" + "mov.b32 {val1_f16, val2_f16}, %1;\n\t" + "mov.b32 {val3_f16, val4_f16}, %2;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "cvt.f32.f16 val1, val1_f16;\n\t" + "cvt.f32.f16 val2, val2_f16;\n\t" + "cvt.f32.f16 val3, val3_f16;\n\t" + "cvt.f32.f16 val4, val4_f16;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %3, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e4m3x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale)), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const fp16x4 &in, const floatx4 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::fp16x2 const *in2 = reinterpret_cast(&in); + ptx::floatx2 const *scale2 = reinterpret_cast(&scale); + asm volatile( + "{\n\t" + ".reg.b16 val1_f16;\n\t" + ".reg.b16 val2_f16;\n\t" + ".reg.b16 val3_f16;\n\t" + ".reg.b16 val4_f16;\n\t" + "mov.b32 {val1_f16, val2_f16}, %1;\n\t" + "mov.b32 {val3_f16, val4_f16}, %2;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "cvt.f32.f16 val1, val1_f16;\n\t" + "cvt.f32.f16 val2, val2_f16;\n\t" + "cvt.f32.f16 val3, val3_f16;\n\t" + "cvt.f32.f16 val4, val4_f16;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %4, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e4m3x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale2[0])), + "l"(reinterpret_cast(scale2[1])), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const fp16x4 &in, + const ptx::floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::fp16x2 const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + ".reg.b16 val1_f16;\n\t" + ".reg.b16 val2_f16;\n\t" + ".reg.b16 val3_f16;\n\t" + ".reg.b16 val4_f16;\n\t" + "mov.b32 {val1_f16, val2_f16}, %1;\n\t" + "mov.b32 {val3_f16, val4_f16}, %2;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "cvt.f32.f16 val1, val1_f16;\n\t" + "cvt.f32.f16 val2, val2_f16;\n\t" + "cvt.f32.f16 val3, val3_f16;\n\t" + "cvt.f32.f16 val4, val4_f16;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %3, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e5m2x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale)), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const fp16x4 &in, const floatx4 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::fp16x2 const *in2 = reinterpret_cast(&in); + ptx::floatx2 const *scale2 = reinterpret_cast(&scale); + asm volatile( + "{\n\t" + ".reg.b16 val1_f16;\n\t" + ".reg.b16 val2_f16;\n\t" + ".reg.b16 val3_f16;\n\t" + ".reg.b16 val4_f16;\n\t" + "mov.b32 {val1_f16, val2_f16}, %1;\n\t" + "mov.b32 {val3_f16, val4_f16}, %2;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "cvt.f32.f16 val1, val1_f16;\n\t" + "cvt.f32.f16 val2, val2_f16;\n\t" + "cvt.f32.f16 val3, val3_f16;\n\t" + "cvt.f32.f16 val4, val4_f16;\n\t" + ".reg.b64 val_1_2;\n\t" + ".reg.b64 val_3_4;\n\t" + "mov.b64 val_1_2, {val1, val2};\n\t" + "mov.b64 val_3_4, {val3, val4};\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + "fma.rn.f32x2 val_1_2, val_1_2, %3, zeros;\n\t" + "fma.rn.f32x2 val_3_4, val_3_4, %4, zeros;\n\t" + "mov.b64 {val1, val2}, val_1_2;\n\t" + "mov.b64 {val3, val4}, val_3_4;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e5m2x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "r"(reinterpret_cast(in2[0])), + "r"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale2[0])), + "l"(reinterpret_cast(scale2[1])), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, floatx4 const &in, + const ptx::floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::floatx2 const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + ".reg.b64 re1;\n\t" + ".reg.b64 re2;\n\t" + "fma.rn.f32x2 re1, %1, %3, zeros;\n\t" + "fma.rn.f32x2 re2, %2, %3, zeros;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "mov.b64 {val1, val2}, re1;\n\t" + "mov.b64 {val3, val4}, re2;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e5m2x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in2[0])), + "l"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale)), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, floatx4 const &in, + const floatx4 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::floatx2 const *in2 = reinterpret_cast(&in); + ptx::floatx2 const *scale2 = reinterpret_cast(&scale); + asm volatile( + "{\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + ".reg.b64 re1;\n\t" + ".reg.b64 re2;\n\t" + "fma.rn.f32x2 re1, %1, %3, zeros;\n\t" + "fma.rn.f32x2 re2, %2, %4, zeros;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "mov.b64 {val1, val2}, re1;\n\t" + "mov.b64 {val3, val4}, re2;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e5m2x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e5m2x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in2[0])), + "l"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale2[0])), + "l"(reinterpret_cast(scale2[1])), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, floatx4 const &in, + const ptx::floatx2 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::floatx2 const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + ".reg.b64 re1;\n\t" + ".reg.b64 re2;\n\t" + "fma.rn.f32x2 re1, %1, %3, zeros;\n\t" + "fma.rn.f32x2 re2, %2, %3, zeros;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "mov.b64 {val1, val2}, re1;\n\t" + "mov.b64 {val3, val4}, re2;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e4m3x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in2[0])), + "l"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale)), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, floatx4 const &in, + const floatx4 &scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + ptx::floatx2 const *in2 = reinterpret_cast(&in); + ptx::floatx2 const *scale2 = reinterpret_cast(&scale); + asm volatile( + "{\n\t" + ".reg.b64 zeros;\n\t" + "mov.b64 zeros, {0x0, 0x0};\n\t" + ".reg.b64 re1;\n\t" + ".reg.b64 re2;\n\t" + "fma.rn.f32x2 re1, %1, %3, zeros;\n\t" + "fma.rn.f32x2 re2, %2, %4, zeros;\n\t" + ".reg.b32 val1;\n\t" + ".reg.b32 val2;\n\t" + ".reg.b32 val3;\n\t" + ".reg.b32 val4;\n\t" + "mov.b64 {val1, val2}, re1;\n\t" + "mov.b64 {val3, val4}, re2;\n\t" +#if (defined _LOOSE_PRECISION) + "cvt.rs.satfinite.e4m3x4.f32 %0, {val4, val3, val2, val1}, %4;\n\t" +#else + ".reg.b16 r1;\n\t" + ".reg.b16 r2;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r1, val2, val1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 r2, val4, val3;\n\t" + "mov.b32 %0, {r1, r2};\n\t" +#endif + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in2[0])), + "l"(reinterpret_cast(in2[1])), + "l"(reinterpret_cast(scale2[0])), + "l"(reinterpret_cast(scale2[1])), "r"(0x80008000)); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void abs_max_2x(float &dst, const float &p1, const float &p2, + const float &p3) { +#if (defined CUDA_VERSION) && (CUDA_VERSION >= 12090) + asm volatile("max.abs.f32 %0, %1, %2, %3;" : "=f"(dst) : "f"(p1), "f"(p2), "f"(p3)); +#else + asm volatile( + "max.xorsign.abs.f32 %0, %2, %3;" + "max.xorsign.abs.f32 %0, %0, %1;" + : "+f"(dst) + : "f"(p1), "f"(p2), "f"(p3)); +#endif +} + +__device__ __forceinline__ ptx::floatx2 up_cast(const ptx::fp16x2 &in) { + ptx::floatx2 out; + asm volatile( + "{\n\t" + ".reg.b16 f16_1;\n\t" + ".reg.b16 f16_2;\n\t" + "mov.b32 {f16_1, f16_2}, %2;\n\t" + "cvt.f32.f16 %0, f16_1;\n\t" + "cvt.f32.f16 %1, f16_2;\n\t" + "}\n\t" + : "=f"(out.x), "=f"(out.y) + : "r"(reinterpret_cast(in))); + return out; +} + +__device__ __forceinline__ floatx4 up_cast(const fp16x4 &in) { + floatx4 out; + asm volatile( + "{\n\t" + ".reg.b16 f16_1;\n\t" + ".reg.b16 f16_2;\n\t" + ".reg.b16 f16_3;\n\t" + ".reg.b16 f16_4;\n\t" + "mov.b64 {f16_1, f16_2, f16_3, f16_4}, %4;\n\t" + "cvt.f32.f16 %0, f16_1;\n\t" + "cvt.f32.f16 %1, f16_2;\n\t" + "cvt.f32.f16 %2, f16_3;\n\t" + "cvt.f32.f16 %3, f16_4;\n\t" + "}\n\t" + : "=f"(out.x1), "=f"(out.x2), "=f"(out.x3), "=f"(out.x4) + : "l"(reinterpret_cast(in))); + return out; +} + +__device__ __forceinline__ ptx::floatx2 up_cast(const ptx::bf16x2 &in) { + ptx::floatx2 out; + asm volatile( + "{\n\t" + "prmt.b32 %1, 0x0, %2, 0x7632;\n\t" + "prmt.b32 %0, 0x0, %2, 0x5410;\n\t" + "}\n\t" + : "=r"(reinterpret_cast(out.x)), "=r"(reinterpret_cast(out.y)) + : "r"(reinterpret_cast(in))); + return out; +} + +__device__ __forceinline__ floatx4 up_cast(const bf16x4 &in) { + floatx4 out; + int32_t const *in2 = reinterpret_cast(&in); + asm volatile( + "{\n\t" + "prmt.b32 %1, 0x0, %4, 0x7632;\n\t" + "prmt.b32 %0, 0x0, %4, 0x5410;\n\t" + "prmt.b32 %3, 0x0, %5, 0x7632;\n\t" + "prmt.b32 %2, 0x0, %5, 0x5410;\n\t" + "}\n\t" + : "=r"(reinterpret_cast(out.x1)), "=r"(reinterpret_cast(out.x2)), + "=r"(reinterpret_cast(out.x3)), "=r"(reinterpret_cast(out.x4)) + : "r"(in2[0]), "r"(in2[1])); + return out; +} + } // namespace ptx namespace { From e6da012a80e9745d215a7d07703761049862478e Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 19 Nov 2025 07:53:38 -0800 Subject: [PATCH 081/521] [PyTorch] Disable Flash Attention backend in Userbuffers tests (#2399) Disable Flash attention in Userbuffers tests Signed-off-by: Tim Moon --- tests/pytorch/distributed/test_comm_gemm_overlap.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index ddb31c30f9..7134e36a6a 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -120,12 +120,14 @@ def _run_layer_with_overlap( os.environ["PYTORCH_JIT"] = "0" os.environ["NVTE_TORCH_COMPILE"] = "0" os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + os.environ["NVTE_FLASH_ATTN"] = "0" result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) os.unsetenv("PYTORCH_JIT") os.unsetenv("NVTE_TORCH_COMPILE") os.unsetenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO") + os.unsetenv("NVTE_FLASH_ATTN") if ( result.returncode != 0 From 49f7c1db03605d15999cbeae1cc7404b76b855c5 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 19 Nov 2025 10:52:16 -0800 Subject: [PATCH 082/521] Avoid autogenerating docs for Python files with leading underscore (#2397) * Avoid autogenerating docs for Python files with leading underscore Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Do not exclude __init__.py files from doc generation 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> --- docs/conf.py | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/conf.py b/docs/conf.py index 4083bfd242..7f5966d717 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -101,3 +101,4 @@ autoapi_generate_api_docs = False autoapi_dirs = [root_path / "transformer_engine"] +autoapi_ignore = ["*/_[!_]*"] From 8ef8285c40542c8c3724f9b3eadbb006793958f0 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 19 Nov 2025 14:23:36 -0500 Subject: [PATCH 083/521] Minor improvements to CPU overhead (#2400) * Minor CPU overhead changes Signed-off-by: Kirthi Shankar Sivamani * Cache per device Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- .../common/transformer_engine.cpp | 19 ++++--- .../pytorch/cpp_extensions/gemm.py | 51 +++++++++---------- 2 files changed, 35 insertions(+), 35 deletions(-) diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 35e8b683ad..314ba3b40f 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -717,11 +717,16 @@ void nvte_destroy_quantization_config(NVTEQuantizationConfig config) { } int nvte_is_non_tn_fp8_gemm_supported() { - int deviceComputeCapability = - transformer_engine::cuda::sm_arch(transformer_engine::cuda::current_device()); - - // Note: this is temporary restriction and should be lifted in the future. - // (remove the note once it's done.) - return (deviceComputeCapability >= 100 && deviceComputeCapability < 120) || - deviceComputeCapability >= 130; + int num_devices = transformer_engine::cuda::num_devices(); + static std::vector cache(num_devices, -1); + static std::vector flags(num_devices); + int device_id = transformer_engine::cuda::current_device(); + std::call_once(flags[device_id], [&]() { + int deviceComputeCapability = transformer_engine::cuda::sm_arch(device_id); + // Note: this is temporary restriction and should be lifted in the future. + // (remove the note once it's done.) + cache[device_id] = (deviceComputeCapability >= 100 && deviceComputeCapability < 120) || + deviceComputeCapability >= 130; + }); + return cache[device_id]; } diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 76a0e449c0..1a2d619b0f 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -12,10 +12,7 @@ from ..constants import TE_DType from ..utils import get_sm_count, _empty_tensor -from ..quantized_tensor import Quantizer, QuantizedTensor, QuantizedTensorStorage -from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage -from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage -from ..tensor.storage.float8_tensor_storage import Float8TensorStorage +from ..quantized_tensor import Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm @@ -46,8 +43,10 @@ def get_cublas_workspace(device: int, ub: bool, grouped_gemm: bool) -> torch.Ten if ub: return torch.empty( - get_cublas_workspace_size_bytes(), dtype=torch.uint8, device=device - ).repeat(_NUM_MAX_UB_STREAMS) + get_cublas_workspace_size_bytes() * _NUM_MAX_UB_STREAMS, + dtype=torch.uint8, + device=device, + ) if grouped_gemm: _multi_stream_cublas_workspace = [] for _ in range(tex.get_num_cublas_streams()): @@ -69,29 +68,25 @@ def validate_gemm_scale(scale: Optional[float], required: bool) -> float: def get_tensor_device(tensor: torch.Tensor) -> int: - """Returns tensor device as an integer""" - if not isinstance(tensor, QuantizedTensorStorage): - return tensor.device.index - if isinstance(tensor, QuantizedTensor): + """ + Returns tensor device as an integer. + + This method is used because checking instances of + QuantizedTensor or Storage incurs more CPU overhead. + The order of attributes checked is important to also + minimize overhead. + """ + if hasattr(tensor, "device"): return tensor.device.index - if isinstance(tensor, (Float8BlockwiseQTensorStorage, MXFP8TensorStorage, NVFP4TensorStorage)): - return ( - tensor._rowwise_data.device.index - if tensor._rowwise_data is not None - else tensor._columnwise_data.device.index - ) - if isinstance(tensor, Float8TensorStorage): - return ( - tensor._data.device.index - if tensor._data is not None - else tensor._transpose.device.index - ) - try: - return ( - tensor._data.device.index if tensor._data is not None else tensor._data_t.device.index - ) - except AttributeError: - return torch.cuda.current_device() + if hasattr(tensor, "_rowwise_data") and tensor._rowwise_data is not None: + return tensor._rowwise_data.device.index + if hasattr(tensor, "_columnwise_data") and tensor._columnwise_data is not None: + return tensor._columnwise_data.device.index + if hasattr(tensor, "_data") and tensor._data is not None: + return tensor._data.device.index + if hasattr(tensor, "_transpose") and tensor._transpose is not None: + return tensor._transpose.device.index + return torch.cuda.current_device() def general_gemm( From 4142547656460f9e2cec9c1dfd817bc496cc20aa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 21 Nov 2025 12:09:06 +0100 Subject: [PATCH 084/521] [PyTorch] Fix ONNX export errors (#2406) * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix 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> --- transformer_engine/pytorch/onnx_extensions.py | 4 +++- transformer_engine/pytorch/ops/basic/rmsnorm.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/onnx_extensions.py b/transformer_engine/pytorch/onnx_extensions.py index 38df5fc54a..79f9a9fb47 100644 --- a/transformer_engine/pytorch/onnx_extensions.py +++ b/transformer_engine/pytorch/onnx_extensions.py @@ -356,7 +356,9 @@ def onnx_layernorm( ) if normalization == "RMSNorm": - ln_out = torch.nn.functional.rms_norm(inp, inp.shape[-1:], ln_weight, eps) + variance = inp.pow(2).mean(-1, keepdim=True) + ln_out = inp * torch.rsqrt(variance + eps) + ln_out = ln_out * ln_weight else: ln_out = torch.nn.functional.layer_norm( inp, inp.shape[-1:], ln_weight, layer_norm_bias, eps diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 8c3f029747..d91091eb02 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -249,4 +249,6 @@ def op_onnx_forward( ) -> torch.Tensor: """Every operand in this function has a defined ONNX translation.""" weight = self.weight + 1 if self.zero_centered_gamma else self.weight - return torch.nn.functional.rms_norm(input_, input_.shape[-1:], weight, self.eps) + variance = input_.pow(2).mean(-1, keepdim=True) + normalized = input_ * torch.rsqrt(variance + self.eps) + return normalized * weight From 15dead118feab1918262e4698746b77d4ea47bae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 21 Nov 2025 13:35:27 +0100 Subject: [PATCH 085/521] [PyTorch] Fix for CPU offloading (#2403) * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/cpu_offload.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 241cd0e9a8..9e6d577235 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -748,6 +748,11 @@ def get_cpu_offload_context( double_buffering=double_buffering, ) + if not enabled: + if manual_synchronization: + return contextlib.nullcontext(), lambda x: x, None + return contextlib.nullcontext(), lambda x: x + if not offload_weights and not offload_activations: raise ValueError( "CPU Offloading is enabled while it is not " @@ -763,6 +768,8 @@ def get_cpu_offload_context( # Weights offloading is deprecated but we maintain backward compatibility by doing nothing. if not offload_activations: + if manual_synchronization: + return contextlib.nullcontext(), lambda x: x, None return contextlib.nullcontext(), lambda x: x if TEDebugState.debug_enabled: @@ -848,15 +855,13 @@ def hook(_): cpu_offload_context = _CpuOffloadContext() - if enabled: - if manual_synchronization: - return ( - cpu_offload_context, - cpu_offload_context.synchronization_function, - offload_synchronizer, - ) + if manual_synchronization: return ( cpu_offload_context, cpu_offload_context.synchronization_function, + offload_synchronizer, ) - return contextlib.nullcontext(), lambda x: x + return ( + cpu_offload_context, + cpu_offload_context.synchronization_function, + ) From 6f4bc3348ac7aea4d02308a4d4efee8e52f64ec6 Mon Sep 17 00:00:00 2001 From: Jack Date: Fri, 21 Nov 2025 22:35:13 +0800 Subject: [PATCH 086/521] Make grad_output contiguous in cross_entropy.py (#2402) Signed-off-by: Jack --- transformer_engine/pytorch/triton/cross_entropy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index d7e2256e24..1a5756105a 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -121,7 +121,7 @@ def cross_entropy_backward( element_mul_kernel[(n_rows,)]( _input, _input.stride(-2), - grad_output, + grad_output.contiguous(), 1 if grad_output.numel() > 1 else 0, V, BLOCK_SIZE=BLOCK_SIZE, From 632c4c3ebeb0b994667b3d91ed413096eb4b5a70 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Fri, 21 Nov 2025 17:36:01 +0100 Subject: [PATCH 087/521] ci: Build and attach bdist wheels to release page (#2138) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * ci: Build and attach bdist wheels to release page Signed-off-by: oliver könig * free up space Signed-off-by: oliver könig * cleanup Signed-off-by: oliver könig * test Signed-off-by: oliver könig * test Signed-off-by: oliver könig * test Signed-off-by: oliver könig * fix Signed-off-by: oliver könig * test Signed-off-by: oliver könig * fix Signed-off-by: oliver könig * fix Signed-off-by: oliver könig * fix Signed-off-by: oliver könig * fix Signed-off-by: oliver könig * c28619d8999a147d5e09c1199f84ff6af6ad5794 Signed-off-by: oliver könig * c28619d8999a147d5e09c1199f84ff6af6ad5794 Signed-off-by: oliver könig * Reduce months to check from 7 to 5 Signed-off-by: oliver könig * Update .github/scripts/check_for_ngc_images.sh Signed-off-by: Kirthi Shankar Sivamani * Update .github/actions/build-pytorch-wheel/build.sh Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: oliver könig Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Kirthi Shankar Sivamani --- .../actions/build-pytorch-wheel/Dockerfile | 49 +++++ .../actions/build-pytorch-wheel/action.yml | 118 +++++++++++ .github/actions/build-pytorch-wheel/build.sh | 26 +++ .github/scripts/check_for_ngc_images.sh | 69 ++++++ .../workflows/attach-wheels-to-release.yml | 198 ++++++++++++++++++ .github/workflows/build.yml | 89 ++++++-- .gitignore | 1 + transformer_engine/pytorch/setup.py | 22 +- 8 files changed, 548 insertions(+), 24 deletions(-) create mode 100644 .github/actions/build-pytorch-wheel/Dockerfile create mode 100644 .github/actions/build-pytorch-wheel/action.yml create mode 100644 .github/actions/build-pytorch-wheel/build.sh create mode 100644 .github/scripts/check_for_ngc_images.sh create mode 100644 .github/workflows/attach-wheels-to-release.yml diff --git a/.github/actions/build-pytorch-wheel/Dockerfile b/.github/actions/build-pytorch-wheel/Dockerfile new file mode 100644 index 0000000000..5bf0960fa7 --- /dev/null +++ b/.github/actions/build-pytorch-wheel/Dockerfile @@ -0,0 +1,49 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +FROM ubuntu:22.04 + +ENV DEBIAN_FRONTEND=noninteractive + +ENV CUDA_HOME=/usr/local/cuda +ENV PATH=$PATH:$CUDA_HOME/bin +ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH +ENV TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6;9.0" + +ARG PYTHON_VERSION=3.12 +ARG TORCH_VERSION=2.9.1 +ARG CUDA_VERSION=12.9.1 +ARG CUDNN_MAJOR_VERSION=9 +ENV PATH=/opt/venv/bin:$PATH +ENV PYTHONUNBUFFERED=1 +ARG AARCH=x86_64 + +# Install Python +RUN apt-get update && \ + apt-get install -y software-properties-common wget && \ + add-apt-repository ppa:deadsnakes/ppa -y && \ + apt-get install -y python$PYTHON_VERSION-dev python$PYTHON_VERSION-venv python3-pip && \ + python$PYTHON_VERSION -m venv /opt/venv + + +# Install cuda-toolkit +RUN CUDA_MAJOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1'}) && \ + CUDA_MINOR_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $2'}) && \ + rm /etc/apt/sources.list.d/cuda*.list || true && \ + rm /etc/apt/sources.list.d/nvidia-cuda.list || true && \ + wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/${AARCH}/cuda-keyring_1.1-1_all.deb && \ + dpkg -i cuda-keyring_1.1-1_all.deb && \ + rm cuda-keyring_1.1-1_all.deb && \ + apt-get update && \ + apt-get install -y cuda-toolkit-${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION} cudnn-cuda-$CUDA_MAJOR_VERSION libcudnn$CUDNN_MAJOR_VERSION-cuda-$CUDA_MAJOR_VERSION libnccl2 libnccl-dev cmake + +# Install PyTorch +RUN export MATRIX_CUDA_VERSION=$(echo $CUDA_VERSION | awk -F \. {'print $1 $2'}) && \ + export MATRIX_TORCH_VERSION=$(echo $TORCH_VERSION | awk -F \. {'print $1 "." $2'}) && \ + export TORCH_CUDA_VERSION=$(python -c "from os import environ as env; \ + minv = {'2.5': 118, '2.6': 118, '2.7': 118, '2.8': 126, '2.9': 126}[env['MATRIX_TORCH_VERSION']]; \ + maxv = {'2.5': 124, '2.6': 126, '2.7': 128, '2.8': 129, '2.9': 130}[env['MATRIX_TORCH_VERSION']]; \ + print(minv if int(env['MATRIX_CUDA_VERSION']) < 120 else maxv)" \ + ) && \ + pip install --no-cache-dir torch==${TORCH_VERSION} --index-url https://download.pytorch.org/whl/cu${TORCH_CUDA_VERSION} \ No newline at end of file diff --git a/.github/actions/build-pytorch-wheel/action.yml b/.github/actions/build-pytorch-wheel/action.yml new file mode 100644 index 0000000000..a49b12227d --- /dev/null +++ b/.github/actions/build-pytorch-wheel/action.yml @@ -0,0 +1,118 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +name: Build PyTorch Wheel +description: Builds a PyTorch wheel for TransformerEngine + +inputs: + release-version: + description: 'The release version to use for the build' + required: true + python-version: + description: 'The Python version to use for the build' + required: true + cuda-version: + description: 'The CUDA version to use for the build' + required: true + cudnn-version: + description: 'The cuDNN version to use for the build' + required: true + torch-version: + description: 'The PyTorch version to use for the build' + required: true + cxx11_abi: + description: 'Enable torch flag C++11 ABI (TRUE/FALSE)' + required: true + base-image: + description: 'The base image to use for the build' + required: false + aarch: + description: 'The architecture to use for the build' + required: true +outputs: + wheel_name: + description: 'The name of the built wheel' + value: ${{ steps.build_wheel.outputs.wheel_name }} + +runs: + using: 'composite' + steps: + - name: Move /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" + + - name: Maximize build space + uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 + with: + root-reserve-mb: 5120 + temp-reserve-mb: 32 + swap-size-mb: 10240 + remove-dotnet: 'true' + remove-android: 'true' + remove-haskell: 'true' + remove-codeql: 'true' + build-mount-path: '/var/lib/docker/' + + - name: Restore /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" + + - name: Checkout + uses: actions/checkout@v4 + with: + ref: ${{ inputs.release-version }} + submodules: recursive + + - name: Checkout build tools + uses: actions/checkout@v4 + with: + path: build-tools + submodules: recursive + + - name: Build image + shell: bash -euxo pipefail {0} + env: + BASE_IMAGE: ${{ inputs.base-image }} + run: | + if [[ "${BASE_IMAGE}" == "" ]]; then + docker build \ + -t transformer-engine-build \ + -f build-tools/.github/actions/build-pytorch-wheel/Dockerfile \ + --build-arg PYTHON_VERSION=${{ inputs.python-version }} \ + --build-arg TORCH_VERSION=${{ inputs.torch-version }} \ + --build-arg CUDA_VERSION=${{ inputs.cuda-version }} \ + --build-arg CUDNN_MAJOR_VERSION=${{ inputs.cudnn-version }} \ + --build-arg AARCH=${{ inputs.aarch }} \ + . + else + docker pull ${BASE_IMAGE} + docker tag ${BASE_IMAGE} transformer-engine-build + fi + - name: Build wheel + shell: bash -euxo pipefail {0} + id: build_wheel + env: + CXX11_ABI: ${{ inputs.cxx11_abi }} + run: | + echo ::group::Build wheel + + EXIT_CODE=$(docker run \ + --rm \ + --shm-size=64g \ + --workdir /workspace/transformer_engine/pytorch \ + --volume $(pwd):/workspace \ + --volume $GITHUB_OUTPUT:$GITHUB_OUTPUT \ + -e PIP_CONSTRAINT= \ + -e CXX11_ABI=$CXX11_ABI \ + -e GITHUB_OUTPUT=$GITHUB_OUTPUT \ + transformer-engine-build bash /workspace/build-tools/.github/actions/build-pytorch-wheel/build.sh | tail -n 1) + + # Do not fail the job if timeout killed the build + exit $EXIT_CODE + echo ::endgroup:: + + - name: Log Built Wheels + shell: bash -euxo pipefail {0} + run: | + ls transformer_engine/pytorch/dist diff --git a/.github/actions/build-pytorch-wheel/build.sh b/.github/actions/build-pytorch-wheel/build.sh new file mode 100644 index 0000000000..8a219a5959 --- /dev/null +++ b/.github/actions/build-pytorch-wheel/build.sh @@ -0,0 +1,26 @@ +#!/bin/bash + +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +set -eoxu pipefail + +export NVTE_PYTORCH_FORCE_BUILD=TRUE +export NVTE_NO_LOCAL_VERSION=1 +export NVTE_PYTORCH_FORCE_CXX11_ABI=$CXX11_ABI +export PIP_CONSTRAINT= + +pip install wheel packaging nvidia-mathdx ninja pybind11 + +# 5h timeout since GH allows max 6h and we want some buffer +EXIT_CODE=0 +timeout 5h python setup.py bdist_wheel --dist-dir=dist || EXIT_CODE=$? + +if [ $EXIT_CODE -eq 0 ]; then + wheel_name=$(python -c "import setup; print(setup.get_wheel_url()[1])" | tail -n 1) + ls dist/*whl |xargs -I {} mv {} dist/${wheel_name} + echo "wheel_name=${wheel_name}" | tee -a "$GITHUB_OUTPUT" +fi + +echo $EXIT_CODE diff --git a/.github/scripts/check_for_ngc_images.sh b/.github/scripts/check_for_ngc_images.sh new file mode 100644 index 0000000000..f065541838 --- /dev/null +++ b/.github/scripts/check_for_ngc_images.sh @@ -0,0 +1,69 @@ +#!/bin/bash + +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Configuration +BASE_IMAGE="nvcr.io/nvidia/pytorch" +TAG_SUFFIX="-py3" +MONTHS_TO_CHECK=5 # Check current month and previous 4 months (total 5) + +# Initialize an array to store existing tags +EXISTING_TAGS=() + +echo "Checking for existence of the last ${MONTHS_TO_CHECK} NGC PyTorch images: ${BASE_IMAGE}:YY.MM${TAG_SUFFIX}" +echo "---------------------------------------------------------------------" + +# Loop through the last N months +for i in $(seq 0 $((MONTHS_TO_CHECK - 1))); do + # Calculate Year and Month for the tag + CURRENT_YEAR=$(date +%Y) + CURRENT_MONTH=$(date +%m) + + # Calculate target month and year + TARGET_DATE=$(date -d "$CURRENT_YEAR-$CURRENT_MONTH-01 -$i months" +%y.%m) + + # Construct the full image tag and the tag-only string + IMAGE_TAG="${TARGET_DATE}${TAG_SUFFIX}" + FULL_IMAGE="${BASE_IMAGE}:${IMAGE_TAG}" + + echo "Checking: ${FULL_IMAGE}" + + # Use 'docker manifest inspect' to check for image existence without pulling. + if docker manifest inspect "${FULL_IMAGE}" > /dev/null 2>&1; then + echo "✅ EXISTS: Found." + # Add the tag-only string to the array + EXISTING_TAGS+=("nvcr.io/nvidia/pytorch:${IMAGE_TAG}") + else + echo "❌ MISSING: Not found." + fi +done + +echo "---------------------------------------------------------------------" + +## JSON Output Generation +# This uses the collected array to build a JSON string. + +# 1. Convert the shell array to a newline-separated string. +TAGS_NL_SEP=$(printf "%s\n" "${EXISTING_TAGS[@]}") + +# 2. Use jq to read the newline-separated list and format it into a JSON array. +# . | split("\n") | .[:-1] reads the input, splits it by newline, and removes the trailing empty element. +if command -v jq &> /dev/null; then + JSON_STRING=$(echo -e "${TAGS_NL_SEP}" | jq -R -s 'split("\n") | .[:-1]') + + echo "Generated JSON String of Existing Tags:" + echo "${JSON_STRING}" + + # Optional: Save the JSON string to a variable for further use + # echo "JSON_STRING is now available in the shell if you source this script." +else + echo "WARNING: 'jq' is not installed. Cannot format output as JSON." + echo "Found Tags: ${EXISTING_TAGS[*]}" +fi + +echo "---" +echo "Check complete." + +echo "${JSON_STRING}" > ngc_images.json diff --git a/.github/workflows/attach-wheels-to-release.yml b/.github/workflows/attach-wheels-to-release.yml new file mode 100644 index 0000000000..c7d31a7c7d --- /dev/null +++ b/.github/workflows/attach-wheels-to-release.yml @@ -0,0 +1,198 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# This workflow will: +# - Create a new Github release +# - Build wheels for supported architectures +# - Deploy the wheels to the Github release +# - Release the static code to PyPi +# For more information see: https://help.github.com/en/actions/language-and-framework-guides/using-python-with-github-actions#publishing-to-package-registries + +name: Attach wheels to release + +on: + release: + types: [published] + workflow_dispatch: + inputs: + runs-on: + description: 'The runner to use for the build' + required: true + type: string + default: ubuntu-22.04 + release-version: + description: 'Release version' + required: true + default: '0.1.0' + python-version: + description: 'Python version' + required: true + default: '3.12' + torch-version: + description: 'Torch version' + required: true + default: '2.8.0' + cuda-version: + description: 'CUDA version' + required: true + default: '12.9.1' + cudnn-version: + description: 'CUDNN version' + required: true + default: '9' + cxx11_abi: + description: 'C++11 ABI' + required: true + type: choice + default: 'TRUE' + options: + - 'TRUE' + - 'FALSE' + ngc-image: + description: 'NGC PyTorch image (will take precedence over the source build)' + required: false + type: string + default: '' +jobs: + pre-flight: + runs-on: ubuntu-latest + outputs: + build-wheel-matrix: ${{ steps.matrix.outputs.matrix }} + release-assets-url: ${{ steps.release-assets-url.outputs.upload_url }} + ngc-images: ${{ steps.check_for_ngc_images.outputs.IMAGES }} + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Build release matrix + id: matrix + env: + EVENT: ${{ github.event_name }} + run: | + if [[ "$EVENT" == "release" ]]; then + MATRIX=$(echo '{ + "os": ["ubuntu-22.04", "ubuntu-22.04-arm"], + "release-version": ["${{ github.event.release.tag_name }}"], + "python-version": ["3.12"], + "torch-version": ["2.8.0"], + "cuda-version": ["12.9.1"], + "cudnn-version": ["9"], + "cxx11_abi": ["TRUE"] + }' | jq -rc) + else + MATRIX=$(echo '{ + "os": ["${{ inputs.runs-on }}"], + "release-version": ["${{ inputs.release-version }}"], + "python-version": ["${{ inputs.python-version }}"], + "torch-version": ["${{ inputs.torch-version }}"], + "cuda-version": ["${{ inputs.cuda-version }}"], + "cudnn-version": ["${{ inputs.cudnn-version }}"], + "cxx11_abi": ["${{ inputs.cxx11_abi }}"] + }' | jq -rc) + fi + + echo "matrix=$MATRIX" | tee -a "$GITHUB_OUTPUT" + + - name: Get Release with tag + id: get_current_release + uses: joutvhu/get-release@v1 + if: ${{ github.event_name == 'workflow_dispatch' }} + with: + tag_name: ${{ inputs.release-version }} + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + + - name: Get release assets url + env: + EVENT: ${{ github.event_name }} + if: ${{ (success() || !failure()) && !cancelled()}} + id: release-assets-url + run: | + if [[ "$EVENT" == "release" ]]; then + echo "upload_url=${{ github.event.release.upload_url }}" | tee -a "$GITHUB_OUTPUT" + else + echo "upload_url=${{ steps.get_current_release.outputs.upload_url }}" | tee -a "$GITHUB_OUTPUT" + fi + + - name: Check for NGC PyTorch images + id: check_for_ngc_images + if: ${{ (success() || !failure()) && !cancelled()}} + env: + EVENT: ${{ github.event_name }} + run: | + if [[ "$EVENT" == "release" ]]; then + bash ./.github/scripts/check_for_ngc_images.sh + echo "IMAGES=$(cat ngc_images.json | jq -cr)" | tee -a $GITHUB_OUTPUT + else + echo 'IMAGES=["${{ inputs.ngc-image }}"]' | tee -a "$GITHUB_OUTPUT" + fi + + build_wheels: + name: Build Wheel + runs-on: ${{ matrix.os }} + needs: pre-flight + if: ${{ github.event_name == 'release' || inputs.ngc-image == '' }} + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.pre-flight.outputs.build-wheel-matrix) }} + steps: + - name: 'Checkout' + uses: actions/checkout@v3 + + - name: 'Build PyTorch Wheel' + uses: ./.github/actions/build-pytorch-wheel + id: build-pytorch-wheel + with: + release-version: ${{ matrix.release-version }} + python-version: ${{ matrix.python-version }} + cuda-version: ${{ matrix.cuda-version }} + cudnn-version: ${{ matrix.cudnn-version }} + torch-version: ${{ matrix.torch-version }} + cxx11_abi: ${{ matrix.cxx11_abi }} + aarch: ${{ matrix.os == 'ubuntu-22.04' && 'x86_64' || 'sbsa' }} + env: + NVTE_FRAMEWORK: pytorch + MAX_JOBS: 1 + + - name: Upload Release Asset + id: upload_release_asset + uses: actions/upload-release-asset@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + upload_url: ${{ needs.pre-flight.outputs.release-assets-url }} + asset_path: ./transformer_engine/pytorch/dist/${{ steps.build-pytorch-wheel.outputs.wheel_name }} + asset_name: ${{ steps.build-pytorch-wheel.outputs.wheel_name }} + asset_content_type: application/* + + build_wheels_for_ngc: + name: Build Wheels for NGC PyTorch images + runs-on: ${{ matrix.os }} + needs: pre-flight + if: ${{ github.event_name == 'release' || inputs.ngc-image != '' }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-22.04] + container-image: ${{ fromJson(needs.pre-flight.outputs.ngc-images) }} + steps: + - name: 'Checkout' + uses: actions/checkout@v3 + + - name: 'Build PyTorch Wheel' + uses: ./.github/actions/build-pytorch-wheel + id: build-pytorch-wheel + with: + base-image: ${{ matrix.container-image }} + + - name: Upload Release Asset + id: upload_release_asset + uses: actions/upload-release-asset@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + upload_url: ${{ needs.pre-flight.outputs.release-assets-url }} + asset_path: ./transformer_engine/pytorch/dist/${{ steps.build-pytorch-wheel.outputs.wheel_name }} + asset_name: ${{ steps.build-pytorch-wheel.outputs.wheel_name }} + asset_content_type: application/* diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 42c5f0342e..51036e40bd 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -35,26 +35,52 @@ jobs: pytorch: name: 'PyTorch' runs-on: ubuntu-latest - container: - image: nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 - options: --user root steps: - - name: 'Dependencies' - run: | - apt-get update - apt-get install -y git python3.9 pip cudnn9-cuda-12 - pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript + - name: Move /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" + + - name: Maximize build space + uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 + with: + root-reserve-mb: 5120 + temp-reserve-mb: 32 + swap-size-mb: 10240 + remove-dotnet: 'true' + remove-android: 'true' + remove-haskell: 'true' + remove-codeql: 'true' + build-mount-path: '/var/lib/docker/' + + - name: Restore /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" + - name: 'Checkout' uses: actions/checkout@v3 with: submodules: recursive + + - name: Start named container + run: | + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d nvcr.io/nvidia/cuda:12.8.0-devel-ubuntu22.04 sleep infinity + + - name: 'Dependencies' + run: | + docker exec builder bash -c '\ + apt-get update && \ + apt-get install -y git python3.9 pip cudnn9-cuda-12 && \ + pip install cmake torch ninja pydantic importlib-metadata>=1.0 packaging pybind11 numpy einops onnxscript && \ + apt-get clean \ + ' + - name: 'Build' - run: pip install --no-build-isolation . -v --no-deps + run: docker exec builder bash -c 'pip install --no-build-isolation . -v --no-deps' env: NVTE_FRAMEWORK: pytorch MAX_JOBS: 1 - name: 'Sanity check' - run: python3 tests/pytorch/test_sanity_import.py + run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py' jax: name: 'JAX' runs-on: ubuntu-latest @@ -78,22 +104,47 @@ jobs: all: name: 'All' runs-on: ubuntu-latest - container: - image: ghcr.io/nvidia/jax:jax - options: --user root steps: - - name: 'Dependencies' - run: | - pip install pybind11[global] einops onnxscript - pip install torch --index-url https://download.pytorch.org/whl/cu130 + - name: Move /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo mv /var/lib/docker/ "${GITHUB_WORKSPACE}/docker" + + - name: Maximize build space + uses: easimon/maximize-build-space@c28619d8999a147d5e09c1199f84ff6af6ad5794 + with: + root-reserve-mb: 5120 + temp-reserve-mb: 32 + swap-size-mb: 10240 + remove-dotnet: 'true' + remove-android: 'true' + remove-haskell: 'true' + remove-codeql: 'true' + build-mount-path: '/var/lib/docker/' + + - name: Restore /var/lib/docker/ + shell: bash -euxo pipefail {0} + run: sudo sh -c "mv ${GITHUB_WORKSPACE}/docker/* /var/lib/docker" + - name: 'Checkout' uses: actions/checkout@v3 with: submodules: recursive + + - name: Start named container + run: | + docker run -v $(pwd):$(pwd) -w $(pwd) --name builder -d ghcr.io/nvidia/jax:jax sleep infinity + + - name: 'Dependencies' + run: | + docker exec builder bash -c '\ + pip install pybind11[global] einops onnxscript && \ + pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 + ' + - name: 'Build' - run: pip install --no-build-isolation . -v --no-deps + run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' env: NVTE_FRAMEWORK: all MAX_JOBS: 1 - name: 'Sanity check' - run: python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py + run: docker exec builder bash -c 'python3 tests/pytorch/test_sanity_import.py && python3 tests/jax/test_sanity_import.py' diff --git a/.gitignore b/.gitignore index 5da08d3638..74acd6ad7f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ +.venv *.o *.swp *.ii diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py index 7a81550047..9719ccb35c 100644 --- a/transformer_engine/pytorch/setup.py +++ b/transformer_engine/pytorch/setup.py @@ -75,21 +75,29 @@ def get_platform(): def get_wheel_url(): """Construct the wheel URL for the current platform.""" - torch_version_raw = parse(torch.__version__) python_version = f"cp{sys.version_info.major}{sys.version_info.minor}" platform_name = get_platform() nvte_version = te_version() - torch_version = f"{torch_version_raw.major}.{torch_version_raw.minor}" cxx11_abi = str(torch._C._GLIBCXX_USE_CXX11_ABI).upper() # Determine the version numbers that will be used to determine the correct wheel # We're using the CUDA version used to build torch, not the one currently installed # _, cuda_version_raw = get_cuda_bare_metal_version(CUDA_HOME) torch_cuda_version = parse(torch.version.cuda) - # For CUDA 11, we only compile for CUDA 11.8, and for CUDA 12 we only compile for CUDA 12.3 + # For CUDA 12 we only compile for CUDA 12.3 # to save CI time. Minor versions should be compatible. - torch_cuda_version = parse("11.8") if torch_cuda_version.major == 11 else parse("12.3") - # cuda_version = f"{cuda_version_raw.major}{cuda_version_raw.minor}" + if torch_cuda_version.major == 12: + torch_cuda_version = parse("12.3") + elif torch_cuda_version.major == 13: + torch_cuda_version = parse("13.0") + else: + raise ValueError(f"CUDA version {torch_cuda_version} not supported") + + if os.environ.get("NVIDIA_PRODUCT_NAME", "") == "PyTorch": + torch_version = str(os.environ.get("NVIDIA_PYTORCH_VERSION")) + else: + torch_version = f"{torch.__version__}" + cuda_version = f"{torch_cuda_version.major}" # Determine wheel URL based on CUDA version, torch version, python version and OS @@ -109,8 +117,10 @@ class CachedWheelsCommand(_bdist_wheel): """ def run(self): + """Acts a proxy before _bdist_wheel.run() and downloads a prebuilt wheel if available.""" if FORCE_BUILD: super().run() + return wheel_url, wheel_filename = get_wheel_url() print("Guessing wheel URL: ", wheel_url) @@ -129,10 +139,12 @@ def run(self): wheel_path = os.path.join(self.dist_dir, archive_basename + ".whl") print("Raw wheel path", wheel_path) os.rename(wheel_filename, wheel_path) + return except (urllib.error.HTTPError, urllib.error.URLError): print("Precompiled wheel not found. Building from source...") # If the wheel could not be downloaded, build from source super().run() + return if __name__ == "__main__": From b14f417a7c3b840abdea4ae9caea5bff4d7325ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jan=20Bernl=C3=B6hr?= Date: Fri, 21 Nov 2025 17:43:06 +0100 Subject: [PATCH 088/521] [PyTorch] Fix assertion error message formatting in DotProductAttention (#2103) Signed-off-by: janbernloehr Co-authored-by: Kirthi Shankar Sivamani --- .../dot_product_attention.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) 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 1f60ae020e..550c2a4a5a 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 @@ -1033,14 +1033,14 @@ def forward( query_layer.shape[-1] == key_layer.shape[-1] ), "Queries and keys must have the same head dimension!" head_dim_qk, head_dim_v = query_layer.shape[-1], value_layer.shape[-1] - assert ( - head_dim_qk == self.hidden_size_per_attention_head_k - ), f"Keys have head_dim = {head_dim_qk}, " - "but expected head_dim = {self.hidden_size_per_attention_head_k}!" - assert ( - head_dim_v == self.hidden_size_per_attention_head_v - ), f"Values have head_dim = {head_dim_v}, " - "but expected head_dim = {self.hidden_size_per_attention_head_v}!" + assert head_dim_qk == self.hidden_size_per_attention_head_k, ( + f"Keys have head_dim = {head_dim_qk}, but expected head_dim =" + f" {self.hidden_size_per_attention_head_k}!" + ) + assert head_dim_v == self.hidden_size_per_attention_head_v, ( + f"Values have head_dim = {head_dim_v}, but expected head_dim =" + f" {self.hidden_size_per_attention_head_v}!" + ) assert num_gqa_groups == self.num_gqa_groups_per_partition, ( "Keys and values must have num_gqa_group =" f" {self.num_gqa_groups_per_partition} heads! Found {num_gqa_groups}." From beed55b9b2893d39be317b0092d7838d04d39bf1 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Fri, 21 Nov 2025 11:07:08 -0800 Subject: [PATCH 089/521] [JAX] Set BSHD as default in Unfused DPA, DPA and MHA API calls (#2392) * Make BSHD default for Unfused DPA, DPA and MHA in TE JAX Signed-off-by: Kshitij Janardan Lakhani * Remove explicit transpose_batch set for BSHD for DPA in JAX quickstart Signed-off-by: Kshitij Janardan Lakhani * Add warnings in DPA and MHA to warn users of change defaults to BSHD instead of SBHD Signed-off-by: Kshitij Janardan Lakhani * Minimize the scope of when to trigger warnings for changed defaults for transpose_batch_sequence Signed-off-by: Kshitij Janardan Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Janardan Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/examples/quickstart_jax.ipynb | 2 -- transformer_engine/jax/flax/transformer.py | 34 ++++++++++++++++++---- 2 files changed, 28 insertions(+), 8 deletions(-) diff --git a/docs/examples/quickstart_jax.ipynb b/docs/examples/quickstart_jax.ipynb index 0bf928d6ee..7146a95f4a 100644 --- a/docs/examples/quickstart_jax.ipynb +++ b/docs/examples/quickstart_jax.ipynb @@ -368,7 +368,6 @@ " num_gqa_groups=self.num_attention_heads, # No GQA\n", " attention_dropout=self.attention_dropout,\n", " attn_mask_type='causal',\n", - " transpose_batch_sequence=False, # Input format is [batch, seq_len, ...]\n", " )\n", " x = attention(q, k, v, sequence_descriptor=None, deterministic=deterministic)\n", " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", @@ -628,7 +627,6 @@ " num_gqa_groups=self.num_attention_heads, \n", " attention_dropout=self.attention_dropout,\n", " attn_mask_type='causal',\n", - " transpose_batch_sequence=False, # Input format is [batch, seq_len, ...]\n", " )\n", " x = attention(q, k, v, sequence_descriptor=None, deterministic=deterministic)\n", " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index edf5f37227..e51cc3691e 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -124,7 +124,7 @@ class _UnfusedDotProductAttention(nn.Module): # pylint: disable=too-few-public- dtype: DType = jnp.float32 float32_logits: bool = False scale_factor: Optional[float] = None - transpose_batch_sequence: bool = True + transpose_batch_sequence: bool = False window_size: Optional[Tuple[int, int]] = None softmax_type: AttnSoftmaxType = AttnSoftmaxType.VANILLA_SOFTMAX @@ -544,9 +544,10 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods Scale factor to apply on query. When :attr:`None` is present, the scale factor is equal to :math:`\frac{1}{\sqrt{head\_dim}}`. This is useful for model like T5X, which doesn't need to apply scale on query, which is to set :attr:`scale_factor=1.`. - transpose_batch_sequence: bool, default = True + TODO(KshitijLakhani): Reset this to bool only with default False arg in TransformerEngine v2.12 + transpose_batch_sequence: bool | None, default = None (however, default is forced to False in post_init) Indicate whether the input tensors were switched axis of batch - and sequence length dimension. if set to True, the input tensors + and sequence length dimension. If set to True, the input tensors should be in (seqlen, batch, ...), otherwise (batch, seqlen, ...). window_size: Optional[Tuple[int, int]], default = None Sliding window size. The default value is no sliding window. @@ -586,7 +587,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods float32_logits: bool = False qkv_layout: str = "bshd_bshd_bshd" scale_factor: Optional[float] = None - transpose_batch_sequence: bool = True + transpose_batch_sequence: bool | None = None window_size: Optional[Tuple[int, int]] = None max_segments_per_seq: Optional[int] = 1 context_parallel_causal_load_balanced: bool = False @@ -595,6 +596,17 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods context_checkpoint_name: str = "context" softmax_type: str = "vanilla" + def __post_init__(self): + # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 + # None implies that the user is relying on defaults, hence warn the user and set the new defaults + if self.transpose_batch_sequence is None: + warnings.warn( + "transpose_batch_sequence defaults to False in DotProductAttention starting" + " TransformerEngine v2.10" + ) + self.transpose_batch_sequence = False + super().__post_init__() + @nn.compact def __call__( self, @@ -1047,7 +1059,8 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods If set to True, this module exposes a single fused parameter for query-key-value for self-attention and key-value for cross-attention. - transpose_batch_sequence: bool, default = True + TODO(KshitijLakhani): Reset this to bool only with default False arg in TransformerEngine v2.12 + transpose_batch_sequence: bool | None, default = None (however, default is forced to False in post_init) Indicate whether the input tensors were switched axis of batch and sequence length dimension. if set to True, the input tensors should be in (seqlen, batch, hidden), otherwise (batch, seqlen, hidden). @@ -1100,7 +1113,7 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods low_rank_adaptation_alpha: float = None dtype: DType = jnp.float32 fuse_qkv_params: bool = True - transpose_batch_sequence: bool = True + transpose_batch_sequence: bool | None = None enable_sequence_parallel: bool = False scale_attn_logits: bool = False scaled_query_init: bool = True @@ -1116,6 +1129,15 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods fuse_qkv: Optional[bool] = None def __post_init__(self): + # Deal with changed defaults in API + # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 + # None implies that the user is relying on defaults, hence warn the user and set the new defaults + if self.transpose_batch_sequence is None: + warnings.warn( + "transpose_batch_sequence defaults to False in MultiHeadAttention starting" + " TransformerEngine v2.10" + ) + self.transpose_batch_sequence = False # Deal with the deprecated parameters if self.num_heads is not None: self.num_attention_heads = self.num_heads From 4654b70a86a450ee805bba317d62a902e20cd0cb Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Fri, 21 Nov 2025 11:08:08 -0800 Subject: [PATCH 090/521] [JAX] Remove unnecessary SWA calculation in _segment_ids_pos_to_seqlens_offsets() (#2201) * Remove unnecessary SWA calculation from _segment_ids_pos_to_seqlens_offsets 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> --- transformer_engine/jax/attention.py | 20 +++++--------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 57b118d635..0a32be9679 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -530,6 +530,11 @@ def _segment_ids_pos_to_seqlens_offsets( # # 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 + # 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() @@ -591,21 +596,6 @@ def _segment_ids_pos_to_seqlens_offsets( ) attn_mask = jnp.logical_and(segment_mask, causal_mask) - # TODO(KshitijLakhani): Evaluate if swa_mask is needed to procure seqlen and offsets - swa_mask = ( - make_swa_mask( - segment_pos_q, - segment_pos_kv, - window_size, - dtype=jnp.bool, - segment_ids_q=segment_ids_q, - segment_ids_kv=segment_ids_kv, - ) - if attn_mask_type.is_bottom_right() - else make_swa_mask(segment_pos_q, segment_pos_kv, window_size, dtype=jnp.bool) - ) - attn_mask = jnp.logical_and(attn_mask, swa_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 From a75da0ca15000f786d908a4d285f9e67edc3555c Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 21 Nov 2025 15:13:20 -0800 Subject: [PATCH 091/521] Enable SWA with CP for THD input format (#2220) * Add support for THD+CP+SWA through A2A comms Signed-off-by: Sudhakar Singh * unblock the `padding`+`THD`+`CP(A2A)` with SWA case in A2A forward Signed-off-by: Sudhakar Singh * add proper support for thd Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * bug fix Signed-off-by: Sudhakar Singh * enable thd+cp tests as essential Signed-off-by: Sudhakar Singh * add cp+thd+a2a test to essential Signed-off-by: Sudhakar Singh * fix comments from greptile Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add proper skip for flash attention Signed-off-by: Sudhakar Singh * fix the test to create separate tensors for flash and fused attention backend scenarios Signed-off-by: Sudhakar Singh * remove redundant compare Signed-off-by: Sudhakar Singh * simplify code Signed-off-by: Sudhakar Singh * add note for cu_seqlens_kv and cu_seqlens_kv_padded Signed-off-by: Sudhakar Singh * Update tests/pytorch/attention/test_attention_with_cp.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Sudhakar Singh * Update transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Sudhakar Singh * fixo Signed-off-by: Sudhakar Singh * fix docs Signed-off-by: Sudhakar Singh * fix the argument name 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> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../attention/run_attention_with_cp.py | 45 +-- .../attention/test_attention_with_cp.py | 38 ++- .../dot_product_attention/context_parallel.py | 288 +++++++++++++++--- 3 files changed, 303 insertions(+), 68 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 5ed67c3d5e..e58b2da3a8 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -89,40 +89,47 @@ def generate_input_shapes( cu_seqlens_q_padded = None cu_seqlens_kv_padded = None elif qkv_format == "thd": + seqlens_q = torch.randint(0, config.max_seqlen_q + 1, [config.batch_size]).to(torch.int32) + seqlens_q_padded = (seqlens_q + 2 * world_size - 1) // (world_size * 2) * (world_size * 2) + cu_seqlens_q_padded = torch.cat( + [ + torch.zeros([1], dtype=torch.int32), + seqlens_q_padded.cumsum(0, dtype=torch.int32), + ] + ).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": + 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` + # will not be the same as `cu_seqlens_q` and `cu_seqlens_q_padded` respectively. + cu_seqlens_kv = cu_seqlens_q + cu_seqlens_kv_padded = cu_seqlens_q_padded + + total_tokens = cu_seqlens_q_padded[-1] + q_input_shape = ( - config.batch_size * config.max_seqlen_q, + total_tokens, config.num_heads, config.head_dim_qk, ) k_input_shape = ( - config.batch_size * config.max_seqlen_q, + total_tokens, config.num_gqa_groups, config.head_dim_qk, ) v_input_shape = ( - config.batch_size * config.max_seqlen_q, + total_tokens, config.num_gqa_groups, config.head_dim_v, ) attn_output_shape = ( - config.batch_size * config.max_seqlen_q, + total_tokens, config.num_heads * config.head_dim_v, ) - seqlens_q = torch.randint(0, config.max_seqlen_q + 1, [config.batch_size]).to(torch.int32) - seqlens_q_padded = (seqlens_q + 2 * world_size - 1) // (world_size * 2) * (world_size * 2) - cu_seqlens_q_padded = torch.cat( - [ - torch.zeros([1], dtype=torch.int32), - seqlens_q_padded.cumsum(0, dtype=torch.int32), - torch.tensor([q_input_shape[0]], dtype=torch.int32), - ] - ).cuda() - cu_seqlens_q = torch.clone(cu_seqlens_q_padded) - if kernel_backend == "FusedAttention": - cu_seqlens_q[1:-1] = seqlens_q.cumsum(0, dtype=torch.int32).cuda() - cu_seqlens_q[-1] = cu_seqlens_q[-2] - cu_seqlens_kv = cu_seqlens_q - cu_seqlens_kv_padded = cu_seqlens_q_padded else: assert False, f"{qkv_format=} is not supported!" diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index e5c856acd8..2d4fe69e32 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -7,7 +7,7 @@ import sys import pathlib import logging - +import copy import pytest import torch from transformer_engine.pytorch import ( @@ -73,7 +73,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_2_1", "cp_3_2", "cp_3_3"] + configs = ["cp_1_0", "cp_1_2", "cp_2_1", "cp_3_2", "cp_3_3"] model_configs_flash_attn = {k: model_configs_flash_attn[k] for k in configs} dtypes = ["bf16"] qkv_formats = ["sbhd", "thd"] @@ -96,12 +96,16 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, 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 qkv_format == "thd": - pytest.skip("CP implementation with KV all-gather does not support THD format 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 qkv_format == "thd": - pytest.skip("CP implementation with QKVO A2A does not support THD format 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): @@ -183,7 +187,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): qkv_formats = ["bshd", "sbhd", "thd"] cp_comm_types = ["p2p", "all_gather", "a2a", "a2a+p2p"] if test_essential: - configs = ["cp_1_0", "cp_1_1", "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"] + configs = ["cp_1_0", "cp_1_1", "cp_1_4", "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"] model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs} dtypes = ["bf16", "fp8"] qkv_formats = ["sbhd", "thd"] @@ -224,10 +228,14 @@ def test_cp_with_fused_attention( 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" and cp_comm_type == "all_gather": - pytest.skip("CP implementation with KV all-gather does not support THD format yet!") - if qkv_format == "thd" and "a2a" in cp_comm_type: - pytest.skip("CP implementation with QKVO A2A does not support THD format 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!" @@ -281,6 +289,14 @@ def test_cp_with_fused_attention( ) dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16} + + if qkv_format == "thd": + config = copy.deepcopy(config) + if "causal" in config.attn_mask_type: + config.attn_mask_type = "padding_causal" + else: + config.attn_mask_type = "padding" + fp8_meta = {} fp8_meta["recipe"] = None fp8_meta["local_recipes"] = [] 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 00d609ab9e..1bcff966b7 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4,6 +4,7 @@ """Context Parallelism.""" import os +import itertools from typing import List, Union, Tuple import torch import transformer_engine_torch as tex @@ -260,6 +261,146 @@ def reorder_seq_chunks_for_a2a_after_attn(x, chunk_ids_for_a2a, seq_dim, cp_size return x +def reorder_seq_chunks_before_a2a_after_attn_thd(x, cu_seqlens, cp_size, seq_dim=0): + """ + Reorder sequence chunks for A2A communication that happens after attention + compute. + + Args: + x: The input tensor to be reordered. + cu_seqlens: The cumulative sequence lengths of the input tensor. + cp_size: The number of ranks participating in context parallelism. + seq_dim: The dimension in which to reorder. + + Returns: + The reordered tensor. + + Example: + x: [ 0., 1., 2., 3., 4., 5., 6., 7., 0., 1., 2., 3., 4., 5., + 6., 7., 0., 1., 2., 3., 4., 5., 6., 7., 0., 1., 2., 3., + 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15.] + cu_seqlens: [ 0, 8, 16, 24, 40] + cp_size: 4 + + Returns: [ 0., 7., 0., 7., 0., 7., 0., 1., 14., 15., 1., 6., 1., 6., + 1., 6., 2., 3., 12., 13., 2., 5., 2., 5., 2., 5., 4., 5., + 10., 11., 3., 4., 3., 4., 3., 4., 6., 7., 8., 9.] + + + This logic is similar to how the DualChunking is done to split the sequence + for each rank. Here, the indices of sequence chunks for all those ranks + are concatenated together. So the returned tensor ends up looking like as if + the chunks from all the ranks are concatenated together. + + e.g. [ + 0., 7., 0., 7., 0., 7., 0., 1., 14., 15., # chunk on rank 0 + 1., 6., 1., 6., 1., 6., 2., 3., 12., 13., # chunk on rank 1 + 2., 5., 2., 5., 2., 5., 4., 5., 10., 11., # chunk on rank 2 + 3., 4., 3., 4., 3., 4., 6., 7., 8., 9. # chunk on rank 3 + ] + """ + total_slices_of_any_sequence = 2 * cp_size + slice_sizes = (cu_seqlens[1:] - cu_seqlens[:-1]) // total_slices_of_any_sequence + + indices = [ + ( + # 1st segment + torch.arange( + seq_start + (cp_rank * slice_size), + seq_start + ((cp_rank + 1) * slice_size), + device=cu_seqlens.device, + ), + # 2nd segment + torch.arange( + seq_start + ((total_slices_of_any_sequence - cp_rank - 1) * slice_size), + seq_start + ((total_slices_of_any_sequence - cp_rank) * slice_size), + device=cu_seqlens.device, + ), + ) + for cp_rank in range(cp_size) + for slice_size, seq_start in zip(slice_sizes, cu_seqlens[:-1]) + ] + + # flatten the list of tuples to a list + indices = list(itertools.chain(*indices)) + indices = torch.cat(indices) + return x.index_select(seq_dim, indices) + + +def reorder_seq_chunks_after_a2a_before_attn_thd(x, cu_seqlens, seq_chunk_ids, cp_size, seq_dim=0): + """ + Reorder sequence chunks for A2A communication that happens before attention + compute. + + Args: + x: The input tensor to be reordered. + cu_seqlens: The cumulative sequence lengths of the input tensor. + seq_chunk_ids: The sequence chunk ids of the input `x` which is to be reordered. + cp_size: The number of ranks participating in context parallelism. + seq_dim: The dimension in which to reorder. + + Returns: + The reordered tensor. + + Example: + x: [ 0., 7., 0., 7., 0., 7., 0., 1., 14., 15., 1., 6., 1., 6., + 1., 6., 2., 3., 12., 13., 2., 5., 2., 5., 2., 5., 4., 5., + 10., 11., 3., 4., 3., 4., 3., 4., 6., 7., 8., 9.] + cu_seqlens: [ 0, 8, 16, 24, 40] + seq_chunk_ids: [ 0, 2, 4, 6, 7, 5, 3, 1] + cp_size: 4 + + Returns: [ 0., 1., 2., 3., 4., 5., 6., 7., 0., 1., 2., 3., 4., 5., + 6., 7., 0., 1., 2., 3., 4., 5., 6., 7., 0., 1., 2., 3., + 4., 5., 6., 7., 8., 9., 10., 11., 12., 13., 14., 15.] + + Note that the input sequences (x) are arranged after A2A communication as if DualChunked + chunks on all the ranks are concatenated together in the `seq_dim`. + + e.g. [ + 0., 7., 0., 7., 0., 7., 0., 1., 14., 15., # chunk on rank 0 + 1., 6., 1., 6., 1., 6., 2., 3., 12., 13., # chunk on rank 1 + 2., 5., 2., 5., 2., 5., 4., 5., 10., 11., # chunk on rank 2 + 3., 4., 3., 4., 3., 4., 6., 7., 8., 9. # chunk on rank 3 + ] + + Then the logic to serialize the sequences is: + 1. For every sequence segment on any rank (denoted by `start` and `end`): + 1a. For every chunk (in `chunk_id` and the total of those are twice as many as the number of CP ranks) : + 1aa. The first `cp_size` number of chunks form the first half of the whole sequence. Get those indices. + 1ab. The second `cp_size` number of chunks form the second half of the whole sequence. Get those indices. + 1b. Concatenate the indices of the first half and the second half. + 2. Reorder the entire input tensor by those indices. + """ + + max_cum_seqlen_per_cp_rank = cu_seqlens[-1] // cp_size + cu_seqlens_on_any_cp_rank = cu_seqlens // cp_size + + # Go through all the sequence segments (the sizes should be the same from all the ranks) + indices = [ + torch.arange( + # Calculate 'left' boundary + ( + start + max_cum_seqlen_per_cp_rank * (chunk_id // 2) + if loc < cp_size + else (start + end) // 2 + max_cum_seqlen_per_cp_rank * (chunk_id // 2) + ), + # Calculate 'right' boundary + ( + (start + end) // 2 + max_cum_seqlen_per_cp_rank * (chunk_id // 2) + if loc < cp_size + else end + max_cum_seqlen_per_cp_rank * (chunk_id // 2) + ), + device=cu_seqlens.device, + ) + for start, end in zip(cu_seqlens_on_any_cp_rank[:-1], cu_seqlens_on_any_cp_rank[1:]) + for loc, chunk_id in enumerate(seq_chunk_ids) + ] + + indices = torch.cat(indices) + return x.index_select(seq_dim, indices) + + def flash_attn_a2a_communicate( a2a_inputs: Union[torch.Tensor, List[torch.Tensor]], chunk_ids_for_a2a: torch.Tensor, @@ -268,8 +409,14 @@ def flash_attn_a2a_communicate( cp_group: dist_group_type, cp_stream: torch.cuda.Stream, before_attn: bool, + qkv_format: str = "bshd", + cu_seqlens_padded: torch.Tensor = 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!" 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) if before_attn: @@ -283,20 +430,33 @@ def flash_attn_a2a_communicate( with torch.cuda.stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] - # 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, h//cp, d] -> [b, cp*s, h//cp, d] - # or [cp*2, s//2, b, h//cp, d] -> [cp*s, b, h//cp, d] - a2a_outputs[i - 2] = x.view(*x.shape[:seq_dim], -1, *x.shape[(seq_dim + 2) :]) + if qkv_format in ["bshd", "sbhd"]: + # 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] + 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] + x = x.view(-1, *x.shape[2:]) + # reorder the sequence chunks + a2a_outputs[i - 2] = reorder_seq_chunks_after_a2a_before_attn_thd( + x, cu_seqlens_padded, chunk_ids_for_a2a, cp_size + ) + if i < len(a2a_inputs): x = a2a_inputs[i] - # [b, s, h, d] -> [b, s, cp, h//cp, d] - # or [s, b, h, d] -> [s, b, cp, h//cp, d] + # [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, h//cp, d] -> [cp, b, s, h//cp, d] - # or [s, b, cp, h//cp, d] -> [cp, s, b, h//cp, d] + # [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() else: for i in range(len(a2a_inputs) + 2): @@ -307,22 +467,30 @@ def flash_attn_a2a_communicate( ) if i < len(a2a_inputs): x = a2a_inputs[i] - # [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] - 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 - ) + 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] + 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" + # 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] + 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, h//cp, d] -> [b, 2, s//2, cp, h//cp, d] - # or [cp, 2, s//2, b, h//cp, d] -> [2, s//2, b, cp, h//cp, d] + # [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, h//cp, d] -> [b*s, h, d] - # or [2, s//2, b, cp, h//cp, d] -> [s*b, h, d] + # [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] 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 @@ -3145,7 +3313,9 @@ def forward( causal = "causal" in attn_mask_type padding = "padding" in attn_mask_type - assert not padding, f"{attn_mask_type} mask type is not supported!" + 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!" assert ( @@ -3196,11 +3366,14 @@ def forward( 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!" - assert qkv_format != "thd", f"{qkv_format} format is not supported!" qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format - batch_dim = qkv_format.index("b") - seq_dim = qkv_format.index("s") + 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!" @@ -3246,7 +3419,15 @@ def forward( 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, cp_size, cp_group, cp_stream, True + [q, k, v], + chunk_ids_for_a2a, + seq_dim, + cp_size, + cp_group, + cp_stream, + before_attn=True, + qkv_format=qkv_format, + cu_seqlens_padded=cu_seqlens_q_padded, ) if softmax_type != "vanilla": softmax_offset = flash_attn_a2a_communicate_softmax_offset( @@ -3337,7 +3518,15 @@ def forward( 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, cp_size, cp_group, cp_stream, False + out_, + chunk_ids_for_a2a, + seq_dim, + cp_size, + cp_group, + cp_stream, + before_attn=False, + qkv_format=qkv_format, + cu_seqlens_padded=cu_seqlens_q_padded, ) if return_max_logit: max_logit = flash_attn_a2a_communicate_softmax_offset( @@ -3454,9 +3643,15 @@ def backward(ctx, dout, *_args): cu_seqlens_kv_padded, *aux_ctx_tensors, ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) - qkv_layout = ctx.qkv_format + "_" + ctx.qkv_format + "_" + ctx.qkv_format + + qkv_format = ctx.qkv_format + qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format causal = "causal" in ctx.attn_mask_type - seq_dim = ctx.qkv_format.index("s") + + if qkv_format in ["bshd", "sbhd"]: + seq_dim = qkv_format.index("s") + else: # qkv_format == "thd" + seq_dim = qkv_format.index("t") bwd_nominal_dtype = ctx.fwd_nominal_dtype dqkv_te_dtype = None @@ -3486,14 +3681,23 @@ def backward(ctx, dout, *_args): fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"] if not ctx.use_fused_attention: - out = out.view(ctx.batch_size, -1, *out.shape[-2:]) - dout = dout.view(ctx.batch_size, -1, *dout.shape[-2:]) + 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) 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, cp_size, ctx.cp_group, ctx.cp_stream, True + dout, + chunk_ids_for_a2a, + seq_dim, + cp_size, + ctx.cp_group, + ctx.cp_stream, + before_attn=True, + qkv_format=qkv_format, + cu_seqlens_padded=cu_seqlens_q_padded, ) flash_attn_bwd = None @@ -3510,7 +3714,7 @@ def backward(ctx, dout, *_args): fa_backward_kwargs["window_size"] = ctx.window_size fa_backward_kwargs["deterministic"] = ctx.deterministic else: - if ctx.qkv_format == "thd": + if qkv_format == "thd": from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_varlen_bwd, ) @@ -3579,7 +3783,7 @@ def backward(ctx, dout, *_args): fa_backward_args_thd = get_fa_args( False, ctx.use_flash_attn_3, - ctx.qkv_format, + qkv_format, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=ctx.max_seqlen_q, @@ -3604,12 +3808,20 @@ def backward(ctx, dout, *_args): 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, cp_size, ctx.cp_group, ctx.cp_stream, False + [dq, dk, dv], + chunk_ids_for_a2a, + seq_dim, + cp_size, + ctx.cp_group, + ctx.cp_stream, + before_attn=False, + qkv_format=qkv_format, + cu_seqlens_padded=cu_seqlens_q_padded, ) - if ctx.qkv_format == "bshd": + if 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": + elif qkv_format == "sbhd": dq, dk, dv = [x.view(-1, ctx.batch_size, *x.shape[-2:]) for x in [dq, dk, dv]] d_bias = None From f8cb598c1c9b6228d39bf77cec3ea27c7ab6a1d7 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Fri, 21 Nov 2025 15:46:29 -0800 Subject: [PATCH 092/521] [PyTorch] Only disable Flash Attention in Userbuffers test on SM 8.0 (#2401) Only disable Flash Attention in Userbuffers test on A100 Signed-off-by: Tim Moon --- tests/pytorch/distributed/test_comm_gemm_overlap.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 7134e36a6a..3f4848e105 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -120,7 +120,11 @@ def _run_layer_with_overlap( os.environ["PYTORCH_JIT"] = "0" os.environ["NVTE_TORCH_COMPILE"] = "0" os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" - os.environ["NVTE_FLASH_ATTN"] = "0" + if te.get_device_compute_capability() <= (8, 0): + # We've experienced numerical discrepancies in Flash Attention + # backward when running with Userbuffers on A100s. This does + # not show up in more recent GPUs. + os.environ["NVTE_FLASH_ATTN"] = "0" result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False) From 0056b9813ef882f537a439340c1fb55159a5c955 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 25 Nov 2025 06:59:51 -0800 Subject: [PATCH 093/521] [PyTorch] Change arguments order in triton kernels to make jax-triton work (#2416) * Change order of arguments to make jax works Signed-off-by: tdophung * make num_experts a tl.constepxr again Signed-off-by: tdophung --------- Signed-off-by: tdophung --- .../common/triton/permutation.py | 43 ++++++++++--------- .../pytorch/triton/permutation.py | 28 ++++++------ 2 files changed, 36 insertions(+), 35 deletions(-) diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index 3a3a32014f..e8c43f52d2 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -81,10 +81,8 @@ def _argsort(x, indices, n_dims: tl.constexpr): @triton.jit def _row_id_map_pass_1_kernel( - # pointers + # input pointers routing_map_ptr, - row_id_map_ptr, - workspace_ptr, # sizes num_tokens, # strides @@ -92,6 +90,9 @@ def _row_id_map_pass_1_kernel( stride_routing_map_expert, stride_row_id_map_token, stride_row_id_map_expert, + # output pointers + row_id_map_ptr, + workspace_ptr, # metas BLOCK_SIZE: tl.constexpr, ): @@ -155,12 +156,11 @@ def _row_id_map_pass_2_kernel( def _row_id_map_pass_3_kernel( # pointers row_id_map_ptr, - # sizes - num_experts: tl.constexpr, # strides stride_row_id_map_token, stride_row_id_map_expert, # metas + num_experts: tl.constexpr, LOAD_SIZE: tl.constexpr, ): pid = tl.program_id(0) @@ -194,17 +194,13 @@ def _row_id_map_pass_3_kernel( @triton.jit def _permute_kernel( - # pointers + # input pointers input_ptr, - output_ptr, row_id_map_ptr, probs_ptr, scale_ptr, - permuted_probs_ptr, permuted_scale_ptr, # sizes - num_experts: tl.constexpr, - hidden_size: tl.constexpr, scale_hidden_dim, # strides stride_row_id_map_token, @@ -220,7 +216,12 @@ def _permute_kernel( stride_permuted_probs_token, stride_permuted_scale_token, stride_permuted_scale_hidden, + # output pointers + output_ptr, + permuted_probs_ptr, # metas + num_experts: tl.constexpr, + hidden_size: tl.constexpr, PERMUTE_PROBS: tl.constexpr, PERMUTE_SCALE: tl.constexpr, BLOCK_SIZE: tl.constexpr, @@ -291,16 +292,11 @@ def _permute_kernel( @triton.jit def _unpermute_kernel( - # pointers + # input pointers input_ptr, - output_ptr, row_id_map_ptr, merging_probs_ptr, permuted_probs_ptr, - unpermuted_probs_ptr, - # sizes - num_experts: tl.constexpr, - hidden_size: tl.constexpr, # strides stride_row_id_map_token, stride_row_id_map_expert, @@ -313,7 +309,12 @@ def _unpermute_kernel( stride_permuted_probs_token, stride_unpermuted_probs_token, stride_unpermuted_probs_expert, + # output pointers + output_ptr, + unpermuted_probs_ptr, # metas + num_experts: tl.constexpr, + hidden_size: tl.constexpr, PROBS_LOAD_WIDTH: tl.constexpr, WITH_MERGING_PROBS: tl.constexpr, PERMUTE_PROBS: tl.constexpr, @@ -546,14 +547,10 @@ def _make_chunk_sort_map_kernel( @triton.jit def _sort_chunks_by_map_kernel( - # pointers + # input pointers input_ptr, - output_ptr, row_id_map_ptr, probs_ptr, - permuted_probs_ptr, - # sizes - hidden_size: tl.constexpr, # strides stride_input_token, stride_input_hidden, @@ -561,7 +558,11 @@ def _sort_chunks_by_map_kernel( stride_output_hidden, stride_probs_token, stride_permuted_probs_token, + # output pointers + output_ptr, + permuted_probs_ptr, # metas + hidden_size: tl.constexpr, PERMUTE_PROBS: tl.constexpr, BLOCK_SIZE: tl.constexpr, FORWARD: tl.constexpr, diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index da22299fe5..741dd60c06 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -72,13 +72,13 @@ def make_row_id_map( # [0, 0, 0, r, r, r, r]] _row_id_map_pass_1_kernel[grid]( routing_map, - row_id_map, - workspace_tensor, num_tokens, routing_map.stride(0), routing_map.stride(1), row_id_map.stride(0), row_id_map.stride(1), + row_id_map, + workspace_tensor, block_size, ) @@ -110,9 +110,9 @@ def make_row_id_map( grid = (num_tokens,) _row_id_map_pass_3_kernel[grid]( row_id_map, - num_experts, row_id_map.stride(0), row_id_map.stride(1), + num_experts, triton.next_power_of_2(num_experts), ) return row_id_map @@ -169,14 +169,10 @@ def permute_with_mask_map( grid = lambda META: (num_tokens, triton.cdiv(hidden_size, META["BLOCK_SIZE"])) _permute_kernel[grid]( inp, - output, row_id_map, probs, scale, - permuted_probs, permuted_scale, - num_experts, - hidden_size, scale_hidden_dim, row_id_map.stride(0), row_id_map.stride(1), @@ -191,6 +187,10 @@ def permute_with_mask_map( permuted_probs.stride(0) if permuted_probs is not None else None, permuted_scale.stride(0) if permuted_scale is not None else None, permuted_scale.stride(1) if permuted_scale is not None else None, + output, + permuted_probs, + num_experts, + hidden_size, PERMUTE_PROBS=probs is not None, PERMUTE_SCALE=scale is not None, ) @@ -238,13 +238,9 @@ def unpermute_with_mask_map( grid = lambda META: (num_tokens, triton.cdiv(hidden_size, META["BLOCK_SIZE"])) _unpermute_kernel[grid]( inp, - output, row_id_map, merging_probs, permuted_probs, - unpermuted_probs, - num_experts, - hidden_size, row_id_map.stride(0), row_id_map.stride(1), inp.stride(0), @@ -256,6 +252,10 @@ def unpermute_with_mask_map( permuted_probs.stride(0) if permuted_probs is not None else None, unpermuted_probs.stride(0) if unpermuted_probs is not None else None, unpermuted_probs.stride(1) if unpermuted_probs is not None else None, + output, + unpermuted_probs, + num_experts, + hidden_size, PROBS_LOAD_WIDTH=triton.next_power_of_2(num_experts), WITH_MERGING_PROBS=merging_probs is not None, PERMUTE_PROBS=permuted_probs is not None, @@ -395,17 +395,17 @@ def sort_chunks_by_map( grid = lambda META: (num_tokens, triton.cdiv(hidden_size, META["BLOCK_SIZE"])) _sort_chunks_by_map_kernel[grid]( inp, - output, row_id_map, probs, - permuted_probs, - hidden_size, inp.stride(0), inp.stride(1), output.stride(0), output.stride(1), probs.stride(0) if probs is not None else None, permuted_probs.stride(0) if permuted_probs is not None else None, + output, + permuted_probs, + hidden_size, PERMUTE_PROBS=probs is not None, FORWARD=is_forward, ) From f612b749fab2c82b41789c903d84b50dfb20d28b Mon Sep 17 00:00:00 2001 From: satias10 Date: Tue, 25 Nov 2025 17:27:14 +0200 Subject: [PATCH 094/521] docs: Document NVTE_CUDA_ARCHS environment variable in README (#2414) Add:: NVTE_CUDA_ARCHS to README Signed-off-by: Shoval Atias Co-authored-by: Shoval Atias --- README.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/README.rst b/README.rst index 50c1dcd807..d2bb2f4056 100644 --- a/README.rst +++ b/README.rst @@ -259,6 +259,7 @@ These environment variables can be set before installation to customize the buil * **NVTE_FRAMEWORK**: Comma-separated list of frameworks to build for (e.g., ``pytorch,jax``) * **MAX_JOBS**: Limit number of parallel build jobs (default varies by system) * **NVTE_BUILD_THREADS_PER_JOB**: Control threads per build job +* **NVTE_CUDA_ARCHS**: Semicolon-separated list of CUDA compute architectures to compile for (e.g., ``80;90`` for A100 and H100). If not set, automatically determined based on CUDA version. Setting this can significantly reduce build time and binary size. Compiling with FlashAttention ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ From 66ae3030460f86a702e804d57e6b3140d0fba813 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 25 Nov 2025 11:56:39 -0500 Subject: [PATCH 095/521] [JAX] Allow DP + FSDP and fixed sr_rng_state partitioning (#2418) * allow dp + fsdp and fixed sr_rng_state partitioning Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup for lint test Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../jax/cpp_extensions/quantization.py | 34 ++++++++++--------- transformer_engine/jax/sharding.py | 12 ------- 2 files changed, 18 insertions(+), 28 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index d16dab6d6c..b55fa20790 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -497,6 +497,7 @@ def partition( x_spec = get_padded_spec(arg_infos[0]) amax_spec = get_padded_spec(arg_infos[2]) + sr_rng_state_spec = get_padded_spec(arg_infos[3]) out_sharding = NamedSharding( mesh, PartitionSpec(*x_spec), @@ -551,11 +552,14 @@ def partition( ) arg_shardings = list(arg_i.sharding for arg_i in arg_infos) - arg_shardings[3] = NamedSharding( - mesh, - PartitionSpec(tuple(x for x in x_spec if x is not None), None), - desc="BaseDBiasQuantizePrimitive.sr_rng_state", - ) + if len(sr_rng_state_spec) > 1: + # sr_rng_state shape [n_devices, state_per_device] + sr_rng_state_spec = (*tuple(x for x in x_spec if x is not None), None) + arg_shardings[3] = NamedSharding( + mesh, + PartitionSpec(*sr_rng_state_spec), + desc="BaseDBiasQuantizePrimitive.sr_rng_state", + ) arg_shardings = tuple(arg_shardings) out_shardings = ( out_sharding, @@ -654,10 +658,12 @@ def shardy_sharding_rule( dbias = input_spec[flatten_axis:] if is_dbias else (prefix + "_dbias",) amax = (BATCHING + prefix + "_amax",) scale = (BATCHING + prefix + "_scale",) - sr_rng_state = ( - BATCHING + prefix + "_sr_rng_state_partition_axis", - BATCHING + prefix + "sr_rng_state_data_axis", - ) + sr_rng_state = (BATCHING + prefix + "_sr_rng_state",) + if value_types[3].shape != [0]: + sr_rng_state = ( + BATCHING + prefix + "_sr_rng_state_devices", + prefix + "sr_rng_state_data", + ) post_rht_amax = (BATCHING + prefix + "_post_rht_amax",) rht_matrix = (BATCHING + prefix + "_rht_matrix_1", BATCHING + prefix + "_rht_matrix_2") @@ -849,7 +855,7 @@ def _quantize_dbias_impl( if force_1x_quantization: q_layout = QuantizeLayout.ROWWISE - sr_rng_state = None + sr_rng_state = jnp.empty((0,), jnp.uint32) if quantizer.scaling_mode.is_nvfp4_scaling: # Only NVFP4 scaling modes support stochastic rounding if quantizer.stochastic_rounding_rng_state is not None: @@ -866,11 +872,7 @@ def _quantize_dbias_impl( x.data, scale, amax, - ( - sr_rng_state - if sr_rng_state is not None - else jnp.empty((get_num_devices_in_mesh(), 1), jnp.uint32) - ), + sr_rng_state, post_rht_amax if post_rht_amax is not None else jnp.zeros((1,), jnp.float32), rht_matrix, out_dtype=quantizer.q_dtype, @@ -880,7 +882,7 @@ def _quantize_dbias_impl( scale_dtype=quantizer.get_scale_dtype(), is_dbias=is_dbias if not quantizer.scaling_mode.is_nvfp4_scaling else False, is_outer=True, - stochastic_rounding=sr_rng_state is not None, + stochastic_rounding=sr_rng_state.size != 0, use_rht=use_rht, ) # For DelayedScaling2x, the scale buffer is shared between rowwise and colwise diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 7f204e768b..6cb0dd257c 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -44,9 +44,6 @@ def _get_mesh_info(resource: str, mesh: jax.sharding.Mesh): def _validate_mesh_resource_configuration(mesh_resource): """Validate that the mesh resource configuration is consistent and conflict-free.""" - is_dp_enabled = ( - mesh_resource.dp_resource is not None and get_mesh_axis_size(mesh_resource.dp_resource) > 1 - ) is_tp_enabled = ( mesh_resource.tp_resource is not None and get_mesh_axis_size(mesh_resource.tp_resource) > 1 ) @@ -54,16 +51,7 @@ def _validate_mesh_resource_configuration(mesh_resource): mesh_resource.tpsp_resource is not None and get_mesh_axis_size(mesh_resource.tpsp_resource) > 1 ) - is_fsdp_enabled = ( - mesh_resource.fsdp_resource is not None - and get_mesh_axis_size(mesh_resource.fsdp_resource) > 1 - ) - assert not (is_dp_enabled and is_fsdp_enabled), ( - "Data parallelism and full-sharded data parallelism cannot be enabled at the same time." - f" Got dp_resource={mesh_resource.dp_resource} and" - f" fsdp_resource={mesh_resource.fsdp_resource}" - ) assert not (is_tp_enabled and is_tpsp_enabled), ( "Tensor parallelism and tensor sequence parallelism cannot be enabled at the same time." f" Got tp_resource={mesh_resource.tp_resource} and" From 3b8d9a8ae9213b3699dac7299f8abcc080d28dd3 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 25 Nov 2025 22:33:38 +0530 Subject: [PATCH 096/521] [Pytorch] remove redundant error check in Linear module (#2420) remove linear redundant check Signed-off-by: Varun Thumbe --- transformer_engine/pytorch/module/linear.py | 14 -------------- 1 file changed, 14 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b3f8165a77..3c804ffaa8 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1536,25 +1536,11 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage def _get_weight_and_bias_tensors(self) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: # Get concatenated weight and bias tensors unfused_weights = self._get_weight_tensors() - if any(isinstance(w, QuantizedTensor) for w in unfused_weights): - if self.fp8: - if len(unfused_weights) != 1: - raise RuntimeError( - "Splitting QuantizedTensor into multiple params is not supported" - ) - else: - warnings.warn( - "You are using quantized weights without quantized compute. " - "Please make sure this is intentional." - ) - unfused_weights = [w.dequantize() for w in unfused_weights] - weight_tensor = noop_cat(unfused_weights) if self.use_bias: bias_tensor = noop_cat([getattr(self, name) for name in self.bias_names]) else: bias_tensor = None - return weight_tensor, bias_tensor def onnx_forward( From 89cc2a7e59df130475c595c3aa72c57630e87f62 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Tue, 25 Nov 2025 11:32:57 -0800 Subject: [PATCH 097/521] [PyTorch][NVFP4][MOE] NVFP4 Grouped Hadamard Amax Kernel (#2351) * minor fix of torch view dtype Signed-off-by: Zhongbo Zhu * multi-tensor RHT amax, compiles Signed-off-by: Zhongbo Zhu * setup multi_tensor_quantize_nvfp4_impl Signed-off-by: Zhongbo Zhu * wire things up and run without crash Signed-off-by: Zhongbo Zhu * numerical test Signed-off-by: Zhongbo Zhu * unit test passing Signed-off-by: Zhongbo Zhu * finish unit test of split quantize api Signed-off-by: Zhongbo Zhu * bump up padding to 64 for nvfp4 grouped quantize Signed-off-by: Zhongbo Zhu * fix stochastic rounding Signed-off-by: Zhongbo Zhu * lint Signed-off-by: Zhongbo Zhu * change error message Signed-off-by: Zhongbo Zhu * clean up Signed-off-by: Zhongbo Zhu * enable multi-amax without RHT Signed-off-by: Zhongbo Zhu * fix col-only quantize mode Signed-off-by: Zhongbo Zhu * improve benchmark script Signed-off-by: Zhongbo Zhu * add NCU example script Signed-off-by: Zhongbo Zhu * add larger test case Signed-off-by: Zhongbo Zhu * add contiguous_data_and_scale check to bulk allocator Signed-off-by: Zhongbo Zhu * unified naming and differentiate between group_ and multi_ Signed-off-by: Zhongbo Zhu * move regular amax into multi_tensor.h Signed-off-by: Zhongbo Zhu * Disentangle logic for split-quantize and general multi-tensor quantize Signed-off-by: Tim Moon * Use size_t for split sections Signed-off-by: Tim Moon * Suggestions from @greptile-apps 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: Zhongbo Zhu 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 | 76 ++- .../nvfp4/test_nvfp4_group_quantize.py | 308 +++++++++ tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py | 195 ++++++ tests/pytorch/test_numerics.py | 13 +- transformer_engine/common/CMakeLists.txt | 1 + .../group_hadamard_transform.cu | 604 ++++++++++++++++++ .../hadamard_transform/hadamard_transform.cu | 175 +---- .../hadamard_transform_utils.cuh | 198 ++++++ .../transformer_engine/hadamard_transform.h | 25 + .../include/transformer_engine/multi_tensor.h | 16 + transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/cast.cpp | 375 +++++++++-- .../pytorch/module/fp8_padding.py | 12 +- .../pytorch/module/fp8_unpadding.py | 12 +- transformer_engine/pytorch/quantization.py | 10 + 15 files changed, 1750 insertions(+), 272 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_group_quantize.py create mode 100644 transformer_engine/common/hadamard_transform/group_hadamard_transform.cu create mode 100644 transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh diff --git a/benchmarks/linear/benchmark_grouped_linear.py b/benchmarks/linear/benchmark_grouped_linear.py index d4bbad75cd..02e2bcf4b9 100644 --- a/benchmarks/linear/benchmark_grouped_linear.py +++ b/benchmarks/linear/benchmark_grouped_linear.py @@ -45,6 +45,16 @@ --trace=cuda,nvtx,cudnn,cublas \ python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe nvfp4 +# Example for jagged input benchmark to simulate unbalanced token splits +python benchmarks/linear/benchmark_grouped_linear.py --recipe nvfp4 --jagged-input "15296,8960,14656,14784,11712,7936,14080,10880" + +# Example to look at a single kernel target with NCU, like the fused hadamard amax kernel for NVFP4 recipe +ncu -f -o ./benchmarks/linear/ncu_b200_numgemm_8_nvfp4_rht_amax \ + --set=full \ + --kernel-name "GroupHadamardAmaxTmaKernel" \ + -s 5 -c 5 \ + python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe nvfp4 --profile + """ RECIPES = { @@ -163,7 +173,7 @@ def benchmark_linear( return timing_ms -def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): +def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None): data = [] assert not use_bias, "Bias is not supported for GroupedLinear benchmark" @@ -173,12 +183,13 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): 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)] assert m % num_gemms == 0 - m_splits = [m // num_gemms] * num_gemms + m_splits = [m // num_gemms] * num_gemms if m_splits is None else m_splits # Bias is not supported for GroupedLinear benchmark bias = None # Run the benchmark print(f"fwd_m={m}, fwd_k={k}, fwd_n={n}") + print(f"m_splits: {m_splits}") grouped_fwd_bwd_timing_ms = benchmark_linear( x, @@ -235,8 +246,35 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): default="bf16", help="Recipe to use, options are fp8_sub_channel, mxfp8, bf16, or all", ) + # add an argument for the jagged input + # example: [15296, 8960, 14656, 14784, 11712, 7936, 14080, 10880] => sums up to 98304 + parser.add_argument( + "--jagged-input", + type=str, + default=None, + help="Jagged input to use, example: [15296, 8960, 14656, 14784, 11712, 7936, 14080, 10880]", + ) + parser.add_argument( + "--hidden-dim", + type=int, + default=7168, + help="Hidden dimension to use, default is 7168", + ) + parser.add_argument( + "--output-dim", + type=int, + default=2048, + help="Output dimension to use, default is 2048", + ) args = parser.parse_args() + jagged_input_splits = None + if args.jagged_input is not None: + jagged_input_splits = [int(x) for x in args.jagged_input.split(",")] + print(f"Jagged input splits: {jagged_input_splits}") + print(f"Jagged input splits sum: {sum(jagged_input_splits)}") + print(f"Jagged input splits num_gemms: {len(jagged_input_splits)}") + use_bias = False # Set the MKN values to benchmark # Deepseek V3 EP64, SEQ_LEN=8192, topK8 @@ -256,11 +294,28 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): # 4 or 8local experts per rank num_gemms_list = [4, 8] + if jagged_input_splits is not None: + num_gemms_list = [len(jagged_input_splits)] + + token_dim_list = [65536] + hidden_dim_list = [7168] + output_dim_list = [2048] + + # override the default targets to benchmark if specified + if jagged_input_splits is not None: + token_dim_list = [sum(jagged_input_splits)] + + if args.hidden_dim is not None: + hidden_dim_list = [args.hidden_dim] + + if args.output_dim is not None: + output_dim_list = [args.output_dim] + # MKN for group linear mkns = [] - for m in [65536]: - for k in [7168]: - for n in [2048]: + for m in token_dim_list: + for k in hidden_dim_list: + for n in output_dim_list: mkns.append((m, k, n)) # default recipes to run if not specified @@ -272,14 +327,20 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): recipe_list = [args.recipe] if args.profile: - mkns = [(8192 * 8, 7168, 2048)] + num_gemms_list = [8] + hidden_dim_to_profile = 7168 if args.hidden_dim is None else args.hidden_dim + output_dim_to_profile = 2048 if args.output_dim is None else args.output_dim + token_dim_to_profile = 8192 * 8 + if jagged_input_splits is not None: + num_gemms_list = [len(jagged_input_splits)] + token_dim_to_profile = sum(jagged_input_splits) + mkns = [(token_dim_to_profile, hidden_dim_to_profile, output_dim_to_profile)] # in profile mode, only run one recipe specified in args.recipe assert args.recipe != "all", ( "In profile mode, only one recipe can be specified, please specify the recipe as" " fp8_sub_channel, mxfp8, nvfp4, or bf16" ) recipe_list = [args.recipe] - num_gemms_list = [8] torch.autograd.profiler.emit_nvtx(record_shapes=True).__enter__() # Initialize a dataframe to store the results @@ -310,6 +371,7 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4): recipe_name, use_bias, num_gemms=num_gemms, + m_splits=jagged_input_splits, ) df_linears = pd.concat([df_linears, df]) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py new file mode 100644 index 0000000000..10aa3eb505 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -0,0 +1,308 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# NOTE: This file is dependent on the success of test_nvfp4_quantize_exact.py +# and also the test_nvfp4_rht_quantize_exact.py. +# Separate to make sure all the functionalities are working as expected. +# Otherwise reference implementation will get messy. + +# Due to the structure of NVFP4Quantizer, we need to test the RHT functionality +# together with the quantization functionality. + +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 import utils +from transformer_engine.pytorch.constants import TE_DType +from transformer_engine.common.recipe import NVFP4BlockScaling + +import pytest +import torch +import random +import math + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def generate_random_multiples_sum(total=8192, n=4, multiple=64): + if total % multiple != 0: + raise ValueError(f"Total ({total}) must be a multiple of {multiple}") + if (total // multiple) < n: + raise ValueError("Total too small for given n and multiple.") + + # Work in units of multiples + total_units = total // multiple + + # choose n−1 random cut points in [1, total_units−1) + cuts = sorted(random.sample(range(1, total_units), n - 1)) + + # convert to segment lengths + parts = ( + [cuts[0]] + [cuts[i] - cuts[i - 1] for i in range(1, len(cuts))] + [total_units - cuts[-1]] + ) + + # convert back to multiples + return [p * multiple for p in parts] + + +def generate_split_sections(M: int, N: int, edge_cases: str) -> list[int]: + least_multiple = 64 + num_chunks = 4 + split_sections = None + + avg_split = M // num_chunks + + if M == 0 or N == 0: + # all zeros + return [0] * num_chunks + if edge_cases == "regular": + split_sections = [avg_split] * num_chunks + elif edge_cases == "zero_tokens_front": + split_sections = [0] + [avg_split] * (num_chunks - 2) + [avg_split * 2] + elif edge_cases == "zero_tokens_end": + split_sections = [avg_split * 2] + [avg_split] * (num_chunks - 2) + [0] + elif edge_cases == "zero_tokens_middle": + split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2] + elif edge_cases == "random_uneven_split": + split_sections = generate_random_multiples_sum(M, num_chunks, least_multiple) + else: + raise ValueError(f"Invalid edge case: {edge_cases}") + + # adds up the split_sections to make it M + assert sum(split_sections) == M, "The split_sections do not add up to M" + + # make sure every split_section is a multiple of least_multiple + for split_section in split_sections: + assert ( + split_section % least_multiple == 0 + ), "The split_sections are not multiples of least_multiple" + + return split_sections + + +# Calculate the shape of the scaling tensor for NVFP4 1D blockwise quantization without padding +def get_nvfp4_scale_shape_no_padding(shape, columnwise): + M, K = 1, 1 + M = math.prod(shape[:-1]) + K = shape[-1] + + if columnwise: + outer = K + inner = math.ceil(M / 16) + return (outer, inner) + # rowwise + outer = M + inner = math.ceil(K / 16) + return (outer, inner) + + +def reference_group_quantize( + x: torch.Tensor, + quantizers: list[NVFP4Quantizer], + split_sections: list[int], + return_identity: bool, + return_transpose: bool, +) -> torch.Tensor: + x_view = x.reshape(-1, x.size(-1)) + x_chunks = torch.split(x, split_sections) + + # rowwise quantization + x_qx = [] + x_sx = [] + x_amax_rowwise = [] + # columnwise quantization + x_qx_t = [] + x_sx_t = [] + x_amax_colwise = [] + + for i in range(len(x_chunks)): + x_chunk = x_chunks[i] + x_nvfp4_res = quantizers[i](x_chunk) + if return_identity: + x_qx.append(x_nvfp4_res._rowwise_data.view(dtype=torch.uint8)) + x_sx.append(x_nvfp4_res._rowwise_scale_inv) + x_amax_rowwise.append(x_nvfp4_res._amax_rowwise) + else: + x_qx.append(None) + x_sx.append(None) + x_amax_rowwise.append(None) + if return_transpose: + x_qx_t.append(x_nvfp4_res._columnwise_data.view(dtype=torch.uint8)) + x_sx_t.append(x_nvfp4_res._columnwise_scale_inv) + x_amax_colwise.append(x_nvfp4_res._amax_columnwise) + else: + x_qx_t.append(None) + x_sx_t.append(None) + x_amax_colwise.append(None) + + return x_qx, x_sx, x_amax_rowwise, x_qx_t, x_sx_t, x_amax_colwise + + +def assert_same_shape_and_dtype(x: torch.Tensor, y: torch.Tensor) -> None: + assert x.shape == y.shape + assert x.dtype == y.dtype + + +def check_group_quantization_nvfp4_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + return_identity: bool, + return_transpose: bool, + split_sections: list[int], + with_rht: bool = True, + with_post_rht_amax: bool = True, + with_random_sign_mask: bool = True, +) -> None: + + te_dtype = tex.DType.kFloat4E2M1 + + # Setup device and random seed + device = "cuda" + seed = 0 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + # Input + x = torch.randn((M, N), dtype=x_dtype, device=device) + num_chunks = len(split_sections) + + x_splits = torch.split(x, split_sections) + + # Quantize + quantizers = [ + NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=return_identity, + columnwise=return_transpose, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=with_rht, + with_post_rht_amax=with_post_rht_amax, + with_random_sign_mask=with_random_sign_mask, + ) + for _ in range(len(split_sections)) + ] + x_qx_ref, x_sx_ref, x_amax_rowwise_ref, x_qx_t_ref, x_sx_t_ref, x_amax_colwise_ref = ( + reference_group_quantize(x, quantizers, split_sections, return_identity, return_transpose) + ) + + split_quantize_outputs = tex.split_quantize(x, split_sections, quantizers) + + if return_identity: + x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] + x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] + x_amax_rowwise = [output._amax_rowwise for output in split_quantize_outputs] + + for i in range(len(x_qx)): + if split_sections[i] == 0: + # then just assert the same same and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_amax_rowwise[i], x_amax_rowwise_ref[i]) + assert_same_shape_and_dtype(x_qx[i], x_qx_ref[i]) + assert_same_shape_and_dtype(x_sx[i], x_sx_ref[i]) + else: + torch.testing.assert_close( + x_amax_rowwise[i], x_amax_rowwise_ref[i], atol=0.0, rtol=0.0 + ) + torch.testing.assert_close(x_qx[i], x_qx_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, False) + x_sx_valid = x_sx[i][: valid_scale_shape[0], : valid_scale_shape[1]] + x_sx_ref_valid = x_sx_ref[i][: valid_scale_shape[0], : valid_scale_shape[1]] + torch.testing.assert_close(x_sx_valid, x_sx_ref_valid, atol=0.0, rtol=0.0) + + if return_transpose: + x_qx_t = [ + output._columnwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs + ] + x_sx_t = [output._columnwise_scale_inv for output in split_quantize_outputs] + x_amax_colwise = [output._amax_columnwise for output in split_quantize_outputs] + # assert with zero tolerance + for i in range(len(x_qx_t)): + if split_sections[i] == 0: + # then just assert the same same and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_amax_colwise[i], x_amax_colwise_ref[i]) + assert_same_shape_and_dtype(x_qx_t[i], x_qx_t_ref[i]) + assert_same_shape_and_dtype(x_sx_t[i], x_sx_t_ref[i]) + else: + torch.testing.assert_close( + x_amax_colwise[i], x_amax_colwise_ref[i], atol=0.0, rtol=0.0 + ) + torch.testing.assert_close(x_qx_t[i], x_qx_t_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, True) + x_sx_t_valid = x_sx_t[i][: valid_scale_shape[0], : valid_scale_shape[1]] + x_sx_t_ref_valid = x_sx_t_ref[i][: valid_scale_shape[0], : valid_scale_shape[1]] + torch.testing.assert_close(x_sx_t_valid, x_sx_t_ref_valid, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # edge case, zero tokens for all + (0, 512), + # full tile cases + (256, 1024), + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 16384), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + ], +) +@pytest.mark.parametrize( + "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] +) +@pytest.mark.parametrize( + "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] +) +@pytest.mark.parametrize("with_rht", [True, False], ids=["with_rht", "no_rht"]) +def test_rht_with_quantization_block_tiling_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, + quantize_mode: str, + with_random_sign_mask: bool, + with_rht: bool, +) -> None: + + split_sections = generate_split_sections(M, N, edge_cases) + + # currently disable pre-RHT amax + with_post_rht_amax = with_rht + + if quantize_mode == "quantize": + return_identity = True + return_transpose = False + elif quantize_mode == "quantize_transpose": + return_identity = True + return_transpose = True + elif quantize_mode == "quantize_colwise_only": + return_identity = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + check_group_quantization_nvfp4_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + return_identity=return_identity, + return_transpose=return_transpose, + split_sections=split_sections, + with_rht=with_rht, + with_post_rht_amax=with_post_rht_amax, + with_random_sign_mask=with_random_sign_mask, + ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py index 0842de9ea4..1407c1d8bf 100755 --- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py @@ -2,9 +2,14 @@ # # See LICENSE for license information. +from typing import List, Tuple + import pytest import torch import transformer_engine.pytorch as te + +import transformer_engine_torch as tex + from transformer_engine.pytorch import NVFP4Quantizer recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -151,6 +156,74 @@ def quantize_fp4( return qx, sx, qx_t, sx_t +def group_quantize_fp4( + x: torch.Tensor, + use_stochastic_rounding: bool, + use_2D: bool, + use_RHT: bool, + split_sections: list[int], + use_tex_split_quantize: bool = True, +) -> Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]: + """ + Group quantize function with toggle between tex.split_quantize and manual split/call methods. + + Args: + x (torch.Tensor): Input tensor. + use_stochastic_rounding (bool): Use stochastic rounding. + use_2D (bool): Use 2D quantization. + use_RHT (bool): Use RHT. + split_sections (list[int]): Split sizes for inputs. + use_tex_split_quantize (bool): Toggle method. If True, use tex.split_quantize, else use manual split and per-quantizer invocation. + + Returns: + tuple: Lists of quantized tensors and scale tensors for all sections. + """ + num_tensors = len(split_sections) + nvfp4_quantizers = [ + NVFP4Quantizer( + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=use_RHT, + with_post_rht_amax=True, + stochastic_rounding=use_stochastic_rounding, + with_2d_quantization=use_2D, + ) + for _ in range(num_tensors) + ] + + if use_tex_split_quantize: + outputs = tex.split_quantize(x, split_sections, nvfp4_quantizers) + qx_list = [output._rowwise_data.view(dtype=torch.uint8) for output in outputs] + sx_list = [output._rowwise_scale_inv for output in outputs] + qx_t_list = [output._columnwise_data.view(dtype=torch.uint8) for output in outputs] + sx_t_list = [output._columnwise_scale_inv for output in outputs] + else: + x_chunks = torch.split(x, split_sections) + qx_list = [] + sx_list = [] + qx_t_list = [] + sx_t_list = [] + for i in range(num_tensors): + x_chunk = x_chunks[i] + x_nvfp4_sut = nvfp4_quantizers[i](x_chunk) + assert x_nvfp4_sut._rowwise_data is not None + qx = x_nvfp4_sut._rowwise_data.view(dtype=torch.uint8) + assert x_nvfp4_sut._rowwise_scale_inv is not None + sx = x_nvfp4_sut._rowwise_scale_inv + assert x_nvfp4_sut._columnwise_data is not None + qx_t = x_nvfp4_sut._columnwise_data.view(dtype=torch.uint8) + assert x_nvfp4_sut._columnwise_scale_inv is not None + sx_t = x_nvfp4_sut._columnwise_scale_inv + qx_list.append(qx) + sx_list.append(sx) + qx_t_list.append(qx_t) + sx_t_list.append(sx_t) + + return qx_list, sx_list, qx_t_list, sx_t_list + + def check_quantization_nvfp4_versus_reference( x_dtype: torch.dtype, M: int, N: int, use_2D: bool, use_RHT: bool ) -> None: @@ -209,6 +282,92 @@ def check_quantization_nvfp4_versus_reference( assert me_t_sr < me_t_rn, "Stochastic rounding failed - error larger than the round to nearest." +def check_group_quantization_nvfp4_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + use_2D: bool, + use_RHT: bool, + num_splits: int, + use_tex_split_quantize: bool = True, +) -> None: + device = "cuda" + torch.manual_seed(seed) + n_iters = 50 + + split_sections = [M // num_splits] * num_splits + x_total = torch.randn((M, N), dtype=x_dtype, device=device) * 2 - 1 + x_splits = torch.split(x_total, split_sections) + + q_rn_list, s_rn_list, q_t_rn_list, s_t_rn_list = group_quantize_fp4( + x_total, + use_stochastic_rounding=False, + use_2D=use_2D, + use_RHT=use_RHT, + split_sections=split_sections, + use_tex_split_quantize=use_tex_split_quantize, + ) + sr_n_iter_results = [] + for i in range(n_iters): + q_sr_list, s_sr_list, q_t_sr_list, s_t_sr_list = group_quantize_fp4( + x_total, + use_stochastic_rounding=True, + use_2D=use_2D, + use_RHT=use_RHT, + split_sections=split_sections, + use_tex_split_quantize=use_tex_split_quantize, + ) + sr_n_iter_results.append((q_sr_list, s_sr_list, q_t_sr_list, s_t_sr_list)) + + for i, x in enumerate(x_splits): + y = x.t().contiguous() + if use_RHT: + y = RHT(y) + amax = torch.max(torch.abs(x)).float() + + # fetch q_rn, s_rn, q_t_rn, s_t_rn + q_rn = q_rn_list[i] + s_rn = s_rn_list[i] + q_t_rn = q_t_rn_list[i] + s_t_rn = s_t_rn_list[i] + + dq_rn = dequantize_fp4(q_rn, s_rn, amax) + dq_t_rn = dequantize_fp4(q_t_rn, s_t_rn, amax) + error_rn = (dq_rn - x).float() + me_rn = torch.sqrt((error_rn * error_rn).mean()) + error_t_rn = (dq_t_rn - y).float() + me_t_rn = torch.sqrt((error_t_rn * error_t_rn).mean()) + sr_result = torch.zeros_like(x).float() + sr_t_result = torch.zeros_like(x).float().t().contiguous() + for iter_idx in range(n_iters): + result_sr = sr_n_iter_results[iter_idx] + q_sr = result_sr[0][i] + s_sr = result_sr[1][i] + q_t_sr = result_sr[2][i] + s_t_sr = result_sr[3][i] + + dq_sr = dequantize_fp4(q_sr, s_sr, amax) + dq_t_sr = dequantize_fp4(q_t_sr, s_t_sr, amax) + sr_result += dq_sr.float() + sr_t_result += dq_t_sr.float() + + # Get the mean result of the stochastic rounding + # It should be more accurate than the RN result + sr_result /= n_iters + error_sr = (sr_result - x).float() + me_sr = torch.sqrt((error_sr * error_sr).mean()) + sr_t_result /= n_iters + error_t_sr = (sr_t_result - y).float() + me_t_sr = torch.sqrt((error_t_sr * error_t_sr).mean()) + + print(f"RMSE SR: {me_sr:.3e} | RMSE RN: {me_rn:.3e}") + print(f"RMSE SR_t: {me_t_sr:.3e} | RMSE RN_t: {me_t_rn:.3e}") + assert me_sr < me_rn, "Stochastic rounding failed - error larger than the round to nearest." + assert ( + me_t_sr < me_t_rn + ), "Stochastic rounding failed - error larger than the round to nearest." + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, N", @@ -236,3 +395,39 @@ def test_quantization_block_tiling_versus_reference( M=M, N=N, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + (8192, 8192), + (4096, 7168), + (16384, 2048), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize("use_2D", [False], ids=str) +@pytest.mark.parametrize("use_RHT", [True], ids=str) +@pytest.mark.parametrize("num_splits", [4, 8], ids=str) +@pytest.mark.parametrize("use_tex_split_quantize", [True, False], ids=str) +def test_group_stochastic_rounding_quantization_versus_reference( + x_dtype: torch.dtype, + use_2D: bool, + use_RHT: bool, + num_splits: int, + use_tex_split_quantize: bool, + M: int, + N: int, +) -> None: + if x_dtype == torch.float32 and use_RHT: + pytest.skip("RHT is only supported with bfloat16") + check_group_quantization_nvfp4_versus_reference( + x_dtype=x_dtype, + use_2D=use_2D, + use_RHT=use_RHT, + M=M, + N=N, + num_splits=num_splits, + use_tex_split_quantize=use_tex_split_quantize, + ) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 66f08d5409..f809618346 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -12,7 +12,10 @@ import torch.nn as nn from torch.nn import Parameter -from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.quantization import ( + FP8GlobalStateManager, + get_align_size_for_quantization, +) from transformer_engine.pytorch.utils import ( init_method_normal, scaled_init_method_normal, @@ -1829,9 +1832,7 @@ def _test_grouped_linear_accuracy( if num_gemms > 1: split_size = 1 if fp8: - split_size = 16 - if recipe.mxfp8() or recipe.nvfp4(): - split_size = 32 + 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 @@ -2137,9 +2138,7 @@ def test_grouped_linear_accuracy_single_gemm(recipe): 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 = 16 - if recipe.mxfp8() or recipe.nvfp4(): - align_size = 32 + 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 diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 62b769c77e..d3532b8c45 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -175,6 +175,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources transpose/quantize_transpose_square_blockwise.cu transpose/quantize_transpose_vector_blockwise_fp4.cu hadamard_transform/hadamard_transform.cu + hadamard_transform/group_hadamard_transform.cu hadamard_transform/hadamard_transform_cast_fusion.cu) # Compiling the files with the worst compilation time first to hopefully overlap diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu new file mode 100644 index 0000000000..84eb6bb5c3 --- /dev/null +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -0,0 +1,604 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "common/common.h" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" +#include "hadamard_transform_utils.cuh" + +namespace transformer_engine { +namespace { + +constexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed +struct MultiAmaxArgs { + // (output) Amax buffer for pre-RHT amax buffer + void* output_pre_rht_amax_list[kMaxTensorsPerKernel]; + // (output) Amax buffer for RHT identity amax buffer + void* output_identity_amax_list[kMaxTensorsPerKernel]; + // (output) Amax buffer for RHT transpose amax buffer + void* output_transpose_amax_list[kMaxTensorsPerKernel]; + // Prefix sum (with leading zero) of split_sections of each tensor of input + int split_sections_range[kMaxTensorsPerKernel + 1]; + // Number of tensors (splits) being processed by kernel + int num_tensors; +}; + +constexpr int kThreadsPerWarp = 32; + +template +__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4], + IType* in_sh_ptr, uint32_t& local_pre_rht_amax_reg, + uint32_t& local_amax_reg, + uint32_t& local_amax_t_reg) { + uint32_t a_frag[4]; // A matrix fragment + uint32_t c_frag[4]; // Result fragment + + int warp_id = threadIdx.x / kThreadsPerWarp; + int local_rank = (threadIdx.x % kThreadsPerWarp); + + int ld_row_idx = local_rank % kHadamardDimension; + int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; + int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); + + uint32_t temp_amax_reg; + uint32_t temp_amax_t_reg; + + if (kReturnIdentityAmax) { + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + + mma_m16_n16_k16_b16_b16_b16_noacc( + a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2], + b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(local_amax_reg) + : "r"(local_amax_reg), "r"(temp_amax_reg)); + } + + if (kReturnTransposedAmax) { + // TODO(Frank): This is not efficient, since we could directly load the + // matrix in transposed layout. + if (!kReturnIdentityAmax) { + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } + + matrix_transpose_m8_n8_b16_inplace(a_frag[0]); + matrix_transpose_m8_n8_b16_inplace(a_frag[1]); + matrix_transpose_m8_n8_b16_inplace(a_frag[2]); + matrix_transpose_m8_n8_b16_inplace(a_frag[3]); + + mma_m16_n16_k16_b16_b16_b16_noacc( + a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], + b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(local_amax_t_reg) + : "r"(local_amax_t_reg), "r"(temp_amax_t_reg)); + } + + if (kReturnPreRhtAmax) { + if (!kReturnIdentityAmax && !kReturnTransposedAmax) { + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } + + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(a_frag[0]) + : "r"(a_frag[0]), "r"(a_frag[1])); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(a_frag[2]) + : "r"(a_frag[2]), "r"(a_frag[3])); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(a_frag[0]) + : "r"(a_frag[0]), "r"(a_frag[2])); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(local_pre_rht_amax_reg) + : "r"(a_frag[0]), "r"(local_pre_rht_amax_reg)); + } +} + +template +__device__ __host__ constexpr int NextPowerOf2() { + static_assert(kN > 0, "kN must be > 0"); + // Round up to the next power of 2 by counting leading zeros. + return 1 << (32 - __builtin_clz(kN - 1)); +} + +template +__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax, + const float transpose_amax, float* staging_for_pre_rht, + float* staging_for_identity, float* staging_for_transpose, + float* output_pre_rht_amax_ptr, + float* output_identity_amax_ptr, + float* output_transpose_amax_ptr, const int warpid) { + // intra-warp reduction + constexpr int kWarpSize = 32; + int local_rank = threadIdx.x % 32; + float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f; + float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f; + float warp_transpose_amax = + kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f; + + // inter-warp reduction + if (threadIdx.x % 32 == 0) { + if (kReturnPreRhtAmax) { + staging_for_pre_rht[warpid] = warp_pre_rht_amax; + } + if (kReturnIdentityAmax) { + staging_for_identity[warpid] = warp_identity_amax; + } + if (kReturnTransposedAmax) { + staging_for_transpose[warpid] = warp_transpose_amax; + } + } + __syncthreads(); + constexpr int kNumWarpsPow2 = NextPowerOf2(); + if (warpid == 0) { + if (kReturnIdentityAmax) { + float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f; + identity_accum = warp_reduce_max(identity_accum); + if (local_rank == 0) { + atomicMaxFloat(output_identity_amax_ptr, identity_accum); + } + } + } + if (warpid == 1) { + if (kReturnTransposedAmax) { + float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f; + transpose_accum = warp_reduce_max(transpose_accum); + if (local_rank == 0) { + atomicMaxFloat(output_transpose_amax_ptr, transpose_accum); + } + } + } + if (warpid == 2) { + if (kReturnPreRhtAmax) { + float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f; + pre_rht_accum = warp_reduce_max(pre_rht_accum); + if (local_rank == 0) { + atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum); + } + } + } +} + +// args: the mult-tensor amax arguments +__global__ void MultiZeroAmaxKernel(MultiAmaxArgs args) { + int num_tensors = args.num_tensors; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (; tid < num_tensors; tid += stride) { + float* output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tid]); + float* output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tid]); + float* output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tid]); + if (output_pre_rht_amax_ptr != nullptr) { + *output_pre_rht_amax_ptr = 0; + } + if (output_identity_amax_ptr != nullptr) { + *output_identity_amax_ptr = 0; + } + if (output_transpose_amax_ptr != nullptr) { + *output_transpose_amax_ptr = 0; + } + } +} + +// args: the mult-tensor amax arguments +__global__ void MultiAmaxMemcpyD2DKernelPreRHT(MultiAmaxArgs args) { + int num_tensors = args.num_tensors; + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + for (; tid < num_tensors; tid += stride) { + float* output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tid]); + float* output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tid]); + float* output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tid]); + if (output_pre_rht_amax_ptr != nullptr) { + float pre_rht_amax = *output_pre_rht_amax_ptr; + if (output_identity_amax_ptr != nullptr) { + *output_identity_amax_ptr = pre_rht_amax; + } + if (output_transpose_amax_ptr != nullptr) { + *output_transpose_amax_ptr = pre_rht_amax; + } + } + } +} + +template +__global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor_map_input, + const MultiAmaxArgs args, uint16_t random_sign_mask, + uint16_t random_sign_mask_t, uint64_t num_rows, + uint64_t row_length) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + + float* output_pre_rht_amax_ptr; + float* output_identity_amax_ptr; + float* output_transpose_amax_ptr; + + // calculate the global offset in Y direction to access the correct amax buffer + int global_offset_y = blockIdx.y * CHUNK_DIM_Y; + int tensor_id = 0; + while (args.split_sections_range[tensor_id + 1] <= global_offset_y) { + ++tensor_id; + } + output_pre_rht_amax_ptr = static_cast(args.output_pre_rht_amax_list[tensor_id]); + output_identity_amax_ptr = static_cast(args.output_identity_amax_list[tensor_id]); + output_transpose_amax_ptr = static_cast(args.output_transpose_amax_list[tensor_id]); + + static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0); + static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0); + + constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y; + constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X; + + constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp; + + const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X; + + extern __shared__ __align__(128) 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! + uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL); + + // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned + constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType); + IType* in_sh_0 = reinterpret_cast(dshmem); + dshmem += in_buff_size; + IType* in_sh_1 = reinterpret_cast(dshmem); + dshmem += in_buff_size; + + IType* in_shs[2] = {in_sh_0, in_sh_1}; + + constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType); + + const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0); + + // Initialize shared memory barrier with the number of threads participating in the barrier. +#pragma nv_diag_suppress static_var_with_dynamic_init + uint64_t* mbar = reinterpret_cast(dshmem); + dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y); + + float* max_staging_identity = reinterpret_cast(dshmem); + dshmem += sizeof(float) * kNumWarps; + float* max_staging_transpose = reinterpret_cast(dshmem); + dshmem += sizeof(float) * kNumWarps; + float* max_staging_pre_rht = reinterpret_cast(dshmem); + dshmem += sizeof(float) * kNumWarps; + + initialize_barriers(mbar, + is_master_thread); + + copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input), + input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0], + is_master_thread); + + uint32_t had_frag_i[4]; + uint32_t had_frag_t[4]; + get_hadamard_matrix_fragment( + had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t); + + float local_pre_rht_amax = 0.0; + float local_amax = 0.0; + float local_amax_t = 0.0; + uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax); + uint32_t local_amax_reg = *reinterpret_cast(&local_amax); + uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t); + + for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) { + for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) { + int stage = STAGES_X * stage_y + stage_x; + + const int next_stage = stage + 1; + const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1; + const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y; + + if (next_stage < STAGES_X * STAGES_Y) { + const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y; + const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X; + + copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong + reinterpret_cast(&tensor_map_input), input_global_offset_X, + input_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); + + const size_t compute_stage_x_num = + BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)); + const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y); + + const size_t in_row_stride = BUFF_DIM_X; + + IType* in_sh_ptr = in_shs[stage % 2]; + +#pragma unroll + for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) { + const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y + + threadIdx.y * kHadamardDimension); + const int in_row_offset = row_idx_offset * in_row_stride; + +#pragma unroll + for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) { + ComputeKernel( + had_frag_i, had_frag_t, + in_sh_ptr + in_row_offset + + (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)), + local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); + } + + // Ensure all threads have finished their computation before new data over-writes the shared + // memory. + __syncthreads(); + } + } + } + + const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp; + + if constexpr (kReturnPreRhtAmax) { + unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax); + } + if constexpr (kReturnIdentityAmax) { + unpack_max_of_packed_bf16(local_amax_reg, local_amax); + } + if constexpr (kReturnTransposedAmax) { + unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t); + } + + ReduceMax( + local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity, + max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr, + output_transpose_amax_ptr, warpid); + + destroy_barriers(mbar, is_master_thread); +#else + NVTE_DEVICE_ERROR("Kernel is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +} // namespace + +// broadcast_pre_rht_amax: when it's true, hadamard transform will be disabled +// if at this time, the amax buffers for output expects both amax_rowwise and amax_colwise +// then call MultiAmaxMemcpyD2DKernelPreRHT to D2D copy the amax values +void group_hadamard_transform_amax(const Tensor& input_, std::vector& output_list, + const size_t* split_sections, size_t num_tensors, + uint16_t random_sign_mask, uint16_t random_sign_mask_t, + bool broadcast_pre_rht_amax, cudaStream_t stream) { + NVTE_API_CALL(group_hadamard_transform_amax); +#if CUDA_VERSION >= 12080 + + // Check input tensor + NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING, + "Input tensor must be BF16 tensor, but scaling mode is ", + to_string(input_.scaling_mode), "."); + NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16, + "Input tensor must be BF16 tensor, but dtype is ", to_string(input_.dtype()), "."); + NVTE_CHECK(input_.dim() >= 2, "Input must be a 2D tensor."); + const SimpleTensor& input = input_.data; + + // TODO: validate num_tensors and split_sections + // assert if num_tensors is greater than kMaxTensorsPerKernel + // will expand 64 to higher value if needed + // if input size is going to exceed 4KB kernel launch limit, will then support multi-launch + NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel, + "Number of tensors should be less than or equal to ", kMaxTensorsPerKernel); + + // check split_sections + // TODO: support m_splits_tensor for device initiated API + NVTE_CHECK(split_sections != nullptr, "split_sections should not be nullptr"); + + MultiAmaxArgs kernel_args; + kernel_args.num_tensors = 0; + kernel_args.split_sections_range[0] = 0; + bool all_return_pre_rht_amax = true; + bool all_return_identity_amax = true; + bool all_return_transposed_amax = true; + for (size_t i = 0; i < num_tensors; ++i) { + void* output_pre_rht_amax_ptr = output_list[i]->amax.dptr; + // disable RHT(x) for now, only RHT_T(x) should be used + void* output_identity_amax_ptr = nullptr; + void* output_transpose_amax_ptr = output_list[i]->columnwise_amax.dptr; + all_return_pre_rht_amax &= (output_pre_rht_amax_ptr != nullptr); + all_return_identity_amax &= (output_identity_amax_ptr != nullptr); + all_return_transposed_amax &= (output_transpose_amax_ptr != nullptr); + // sanity check split_sections component to see if it's 64 multiple for each element + NVTE_CHECK(split_sections[i] % 64 == 0, "component ", i, + " of split_sections should be 64 multiple"); + // also skip adding this tensor to the kernel args there are zero elements in this split + if (split_sections[i] == 0) { + continue; + } + // fill in kernel arguments + kernel_args.output_pre_rht_amax_list[kernel_args.num_tensors] = output_pre_rht_amax_ptr; + kernel_args.output_identity_amax_list[kernel_args.num_tensors] = output_identity_amax_ptr; + kernel_args.output_transpose_amax_list[kernel_args.num_tensors] = output_transpose_amax_ptr; + kernel_args.split_sections_range[kernel_args.num_tensors + 1] = + kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i]; + kernel_args.num_tensors++; + } + + NVTE_CHECK(all_return_pre_rht_amax || all_return_identity_amax || all_return_transposed_amax, + "At least one of return_pre_rht_amax, return_identity_amax, or return_transposed_amax " + "must be true"); + // currently we haven't supported all_return_identity_amax, assert error if it's mistakenly enabled + NVTE_CHECK(!all_return_identity_amax, + "Currently RHT transform should only be applied to transposed input"); + + if (broadcast_pre_rht_amax) { + NVTE_CHECK(all_return_pre_rht_amax, + "broadcast_pre_rht_amax is only supported when we compute pre-RHT amax"); + // if all_return_identity_amax and all_return_transposed_amax both are false, there is no need to broadcast anything + broadcast_pre_rht_amax &= (all_return_identity_amax || all_return_transposed_amax); + } + + // Multi zero out multiple amaxes if needed + // Curretly don't support multi-launch when num_tensors is larger than kMaxTensorsPerKernel + // let the number of threads equal to number of tensors, use 1 block, kMaxTensorsPerKernel threads per block + dim3 block_setup_amax(kMaxTensorsPerKernel); + dim3 grid_setup_amax(1); + MultiZeroAmaxKernel<<>>(kernel_args); + NVTE_CHECK_CUDA(cudaGetLastError()); + + checkCuDriverContext(stream); + + using IType = bf16; + + const size_t ndim = input.shape.size(); + const size_t row_length = input.shape[ndim - 1]; + size_t num_rows = 1; + for (size_t i = 0; i < ndim - 1; ++i) { + num_rows *= input.shape[i]; + } + + constexpr int kHadamardDimension = 16; + NVTE_CHECK(row_length % kHadamardDimension == 0, + "row_length must be divisible by hadamard_dimension."); + NVTE_CHECK(num_rows % kHadamardDimension == 0, + "num_rows must be divisible by hadamard_dimension"); + + // four (1x4) 64x64 sub-tiles for ping-pong overlap + constexpr uint64_t kChunkBlockXSmall = 256; + constexpr uint64_t kChunkBlockYSmall = 64; + constexpr uint64_t kBuffDimX = 64; + constexpr uint64_t kBuffDimY = 64; + + alignas(64) CUtensorMap tensor_map_input{}; + + create_2D_tensor_map( + /*tensorMap=*/tensor_map_input, + /*tensor=*/input, + /*globalY=*/num_rows, + /*globalX=*/row_length, + /*shmemY=*/kBuffDimY, + /*shmemX=*/kBuffDimX, + /*stride_elems=*/row_length, + /*offset_elems=*/0, + /*type_num_bits=*/sizeof(IType) * 8, + /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B); + + constexpr uint64_t kThreadBlockX = 4; + constexpr uint64_t kThreadBlockY = 1; + constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY; + + dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY); + + dim3 grid(DIVUP(row_length, kChunkBlockXSmall), DIVUP(num_rows, kChunkBlockYSmall)); + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + (all_return_transposed_amax && !broadcast_pre_rht_amax), kReturnTransposedAmax, + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + (all_return_identity_amax && !broadcast_pre_rht_amax), kReturnIdentityAmax, + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + all_return_pre_rht_amax, kReturnPreRhtAmax, + + // *2 for ping-pong + size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType); + size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) * + (kChunkBlockYSmall / kBuffDimY); + size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3; + // Add padding in case shmem ptr is not aligned to 128 bytes. + shmem_bytes = (shmem_bytes + 128); + + auto kernel = GroupHadamardAmaxTmaKernel< + IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY, + kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax, + kReturnIdentityAmax, kReturnTransposedAmax>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + shmem_bytes); + + kernel<<>>(tensor_map_input, kernel_args, + random_sign_mask, random_sign_mask_t, + num_rows, row_length); + if (broadcast_pre_rht_amax) { + MultiAmaxMemcpyD2DKernelPreRHT<<>>( + kernel_args); + }))); + + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is ", + CUDA_VERSION); +#endif // CUDA_VERSION >= 12080 +} + +} // namespace transformer_engine + +// Naming convention: "Group" kernels here means contiguous input concatenated +// While "Multi" kernels are processing a list of pointers, like the zero amax kernel + +// Group hadamard transform API is unlike other multi-input & multi-output APIs +// Group hadamard transform will take in a single input tensor, and directly calculate amax +// with optional RHT transform. That's because we can assume the input tensor list to be +// contiguous in memory, so the tensors are only splitted in dimension 0. +// RHT transform is 16x16, so as long as each split of the input has 16 multiple shape +// in dimension 0, we can treat the entire input as a single tensor. +// Although mathmatically 16 multple is enough for this function to be correct, +// for this kernel, we required 64 multiple of 16 in dimension 0 for better performance. +void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outputs, + const size_t* split_sections, size_t num_tensors, + int random_sign_mask, int random_sign_mask_t, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_hadamard_transform_amax); + using namespace transformer_engine; + if (num_tensors == 0) { + return; + } + + Tensor* input_tensor = convertNVTETensorCheck(input); + std::vector output_list(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + output_list[i] = convertNVTETensorCheck(outputs[i]); + } + // Call the group tensor Hadamard transform amax implementation. + group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors, + static_cast(random_sign_mask), + static_cast(random_sign_mask_t), false, stream); +} + +// Grouped-tensor amax without doing hadamard transform +void nvte_group_amax(const NVTETensor input, NVTETensor* outputs, const size_t* split_sections, + size_t num_tensors, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_amax); + using namespace transformer_engine; + if (num_tensors == 0) { + return; + } + + Tensor* input_tensor = convertNVTETensorCheck(input); + std::vector output_list(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + output_list[i] = convertNVTETensorCheck(outputs[i]); + } + + group_hadamard_transform_amax(*input_tensor, output_list, split_sections, num_tensors, 0, 0, true, + stream); +} diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform.cu b/transformer_engine/common/hadamard_transform/hadamard_transform.cu index 9d4bec41d5..c01ce7b78f 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform.cu @@ -16,185 +16,12 @@ #include "common/common.h" #include "common/util/ptx.cuh" #include "common/utils.cuh" +#include "hadamard_transform_utils.cuh" namespace transformer_engine { namespace { constexpr int kThreadsPerWarp = 32; -constexpr float k16x16HadamardScale = 0.25f; - -template -__device__ __forceinline__ void ldmatrix_x4_m8n8_shared_b16(uint32_t& a0, uint32_t& a1, - uint32_t& a2, uint32_t& a3, - void* addr) { - auto smem_addr = static_cast(__cvta_generic_to_shared(addr)); - if constexpr (kTranspose) { - asm volatile("ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) - : "r"(smem_addr)); - } else { - asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0,%1,%2,%3}, [%4];\n" - : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) - : "r"(smem_addr)); - } -} - -template -__device__ __forceinline__ void load_matrix_16x16_from_shared(uint32_t& a0, uint32_t& a1, - uint32_t& a2, uint32_t& a3, - void* addr, uint32_t stride) { - if constexpr (kTranspose) { - asm volatile( - "wmma.load.a.sync.aligned.col.m16n16k16.shared::cta.bf16 " - "{%0,%1,%2,%3}, [%4], %5;\n" - : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) - : "l"(addr), "r"(stride)); - } else { - asm volatile( - "wmma.load.a.sync.aligned.row.m16n16k16.shared::cta.bf16 " - "{%0,%1,%2,%3}, [%4], %5;\n" - : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) - : "l"(addr), "r"(stride)); - } -} - -template -__device__ __forceinline__ void store_matrix_16x16_to_global(uint32_t& a0, uint32_t& a1, - uint32_t& a2, uint32_t& a3, void* addr, - uint32_t stride) { - if constexpr (kTranspose) { - asm volatile("wmma.store.d.sync.aligned.col.m16n16k16.global.f16 [%0], {%1, %2, %3, %4}, %5;\n" - : - : "l"(addr), "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(stride)); - } else { - asm volatile("wmma.store.d.sync.aligned.row.m16n16k16.global.f16 [%0], {%1, %2, %3, %4}, %5;\n" - : - : "l"(addr), "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(stride)); - } -} - -__device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(uint32_t& a0) { - asm volatile( - "movmatrix.sync.aligned.m8n8.trans.b16 " - "%0, %1;\n\t" - : "=r"(a0) - : "r"(a0)); -} - -__device__ __forceinline__ void unpack_max_of_packed_bf16(uint32_t& packed_bf16, float& float_dst) { - __nv_bfloat162 bf16x2 = *reinterpret_cast<__nv_bfloat162*>(&packed_bf16); - float f_a = __bfloat162float(bf16x2.x); - float f_b = __bfloat162float(bf16x2.y); - asm volatile("max.xorsign.abs.f32 %0, %1, %2;\n\t" : "=f"(float_dst) : "f"(f_a), "f"(f_b)); - float_dst = fabsf(float_dst); -} - -template -__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc( - uint32_t& a0, uint32_t& a1, uint32_t& a2, uint32_t& a3, uint32_t& b0, uint32_t& b1, - uint32_t& b2, uint32_t& b3, uint32_t& c0, uint32_t& c1, uint32_t& c2, uint32_t& c3, - uint32_t& amax_result) { - uint32_t zero = 0; - uint32_t temp0, temp1, temp2, temp3, temp4, temp5, temp6, temp7; - asm volatile( - "wmma.mma.sync.aligned.row.row.m16n16k16.f32.bf16.bf16.f32 \n" - "{%0, %1, %2, %3, %4, %5, %6, %7}, \n" - "{%8, %9, %10, %11}, \n" - "{%12, %13, %14, %15}, \n" - "{%16, %17, %18, %19, %20, %21, %22, %23};\n\t" - : "=r"(temp0), "=r"(temp1), "=r"(temp2), "=r"(temp3), "=r"(temp4), "=r"(temp5), "=r"(temp6), - "=r"(temp7) - : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(b2), "r"(b3), "r"(zero), - "r"(zero), "r"(zero), "r"(zero), "r"(zero), "r"(zero), "r"(zero), "r"(zero)); - asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c0) : "r"(temp1), "r"(temp0)); - asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c1) : "r"(temp3), "r"(temp2)); - asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c2) : "r"(temp5), "r"(temp4)); - asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c3) : "r"(temp7), "r"(temp6)); - if constexpr (kCalculateAmax) { - uint32_t max_even; - uint32_t max_odd; - // Reduction tree to amax(abs(result)) into bf16x2 reg outparam. - asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" : "=r"(max_even) : "r"(c0), "r"(c2)); - asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" : "=r"(max_odd) : "r"(c1), "r"(c3)); - // N.B. mma is only called up to once per thread for identity and transpose respectively, so - // we don't have to accumulate into amax_result and can directly store into it. - asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" - : "=r"(amax_result) - : "r"(max_even), "r"(max_odd)); - } -} - -template -__device__ __forceinline__ void get_hadamard_matrix_fragment(uint32_t* had_frag_i, - uint16_t random_sign_mask, - uint32_t* had_frag_t, - uint16_t random_sign_mask_t) { - int32_t tid = threadIdx.x % 32; // Local tid - float temp_i[2]; - float temp_t[2]; -#pragma unroll - for (int i = 0; i < 2; i++) { - // i is the vertical fragment index. - // For a 16x16 matrix matrix fragment, 4 threads fill a fragment of 8 BF16 vals. - uint32_t r = i * 8 + tid / 4; - -#pragma unroll - for (int j = 0; j < 2; j++) { -#pragma unroll - for (int k = 0; k < 2; k++) { - // k is column position [0, 1] within a quad of 2 BF16s stored together in 32 bits. - // j is the column fragment idx selecting between even and odd fragments. - // j increments 8 columns by switching fragments. - uint32_t c = j * 8 + k + tid % 4 * 2; - // 1 -> -1.0f, 0 -> 1.0f - int32_t base_sign = __popc(r & c); - if constexpr (kReturnIdentity) { - int32_t sign_i; - // Because tensor cores want the dot product dimension, - // contiguous, the regular, non-inverse hadamard swaps - // signs of columns and rows for inverse. In a simple reference, - // x.reshape(-1, 16) @ sign @ H16, this would be opposite but - // (sign @ H16) is transposed in this fragment. - if constexpr (kInverseHadamardIdentity) { - sign_i = ((random_sign_mask >> r) ^ base_sign); - } else { - sign_i = ((random_sign_mask >> c) ^ base_sign); - } - temp_i[k] = copysignf(k16x16HadamardScale, __int_as_float(sign_i << 31)); - } - if constexpr (kReturnTransposed) { - int32_t sign_t; - if constexpr (kInverseHadamardTransposed) { - sign_t = ((random_sign_mask_t >> r) ^ base_sign); - } else { - sign_t = ((random_sign_mask_t >> c) ^ base_sign); - } - temp_t[k] = copysignf(k16x16HadamardScale, __int_as_float(sign_t << 31)); - } - } - - if constexpr (kReturnIdentity) { - asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" - : "=r"(had_frag_i[i * 2 + j]) - : "f"(temp_i[1]), "f"(temp_i[0])); - } - if constexpr (kReturnTransposed) { - asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" - : "=r"(had_frag_t[i * 2 + j]) - : "f"(temp_t[1]), "f"(temp_t[0])); - } - } - } -} - -__device__ __forceinline__ uint32_t swizzle_128B_atom_32B(uint32_t gmem_row_idx, - uint32_t gmem_col_idx) { - uint32_t smem_row_idx = gmem_row_idx; - uint32_t xor_factor = (smem_row_idx * 2) % 8; - uint32_t smem_col_idx = gmem_col_idx ^ xor_factor; - return smem_row_idx * 8 + smem_col_idx; -} template diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh b/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh new file mode 100644 index 0000000000..ad3bbf5cd7 --- /dev/null +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh @@ -0,0 +1,198 @@ +/************************************************************************* +* Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +* +* See LICENSE for license information. +************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_HADAMARD_TRANSFORM_UTILS_CUH_ +#define TRANSFORMER_ENGINE_HADAMARD_TRANSFORM_UTILS_CUH_ + +#include +#include +#include +#include + +#include "common/common.h" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" + +namespace transformer_engine { + +constexpr float k16x16HadamardScale = 0.25f; + +template +__device__ __forceinline__ void ldmatrix_x4_m8n8_shared_b16(uint32_t& a0, uint32_t& a1, + uint32_t& a2, uint32_t& a3, + void* addr) { + auto smem_addr = static_cast(__cvta_generic_to_shared(addr)); + if constexpr (kTranspose) { + asm volatile("ldmatrix.sync.aligned.x4.trans.m8n8.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) + : "r"(smem_addr)); + } else { + asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0,%1,%2,%3}, [%4];\n" + : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) + : "r"(smem_addr)); + } +} + +template +__device__ __forceinline__ void load_matrix_16x16_from_shared(uint32_t& a0, uint32_t& a1, + uint32_t& a2, uint32_t& a3, + void* addr, uint32_t stride) { + if constexpr (kTranspose) { + asm volatile( + "wmma.load.a.sync.aligned.col.m16n16k16.shared::cta.bf16 " + "{%0,%1,%2,%3}, [%4], %5;\n" + : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) + : "l"(addr), "r"(stride)); + } else { + asm volatile( + "wmma.load.a.sync.aligned.row.m16n16k16.shared::cta.bf16 " + "{%0,%1,%2,%3}, [%4], %5;\n" + : "=r"(a0), "=r"(a1), "=r"(a2), "=r"(a3) + : "l"(addr), "r"(stride)); + } +} + +template +__device__ __forceinline__ void store_matrix_16x16_to_global(uint32_t& a0, uint32_t& a1, + uint32_t& a2, uint32_t& a3, void* addr, + uint32_t stride) { + if constexpr (kTranspose) { + asm volatile("wmma.store.d.sync.aligned.col.m16n16k16.global.f16 [%0], {%1, %2, %3, %4}, %5;\n" + : + : "l"(addr), "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(stride)); + } else { + asm volatile("wmma.store.d.sync.aligned.row.m16n16k16.global.f16 [%0], {%1, %2, %3, %4}, %5;\n" + : + : "l"(addr), "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(stride)); + } +} + +__device__ __forceinline__ void matrix_transpose_m8_n8_b16_inplace(uint32_t& a0) { + asm volatile( + "movmatrix.sync.aligned.m8n8.trans.b16 " + "%0, %1;\n\t" + : "=r"(a0) + : "r"(a0)); +} + +__device__ __forceinline__ void unpack_max_of_packed_bf16(uint32_t& packed_bf16, float& float_dst) { + __nv_bfloat162 bf16x2 = *reinterpret_cast<__nv_bfloat162*>(&packed_bf16); + float f_a = __bfloat162float(bf16x2.x); + float f_b = __bfloat162float(bf16x2.y); + asm volatile("max.xorsign.abs.f32 %0, %1, %2;\n\t" : "=f"(float_dst) : "f"(f_a), "f"(f_b)); + float_dst = fabsf(float_dst); +} + +template +__device__ __forceinline__ void mma_m16_n16_k16_b16_b16_b16_noacc( + uint32_t& a0, uint32_t& a1, uint32_t& a2, uint32_t& a3, uint32_t& b0, uint32_t& b1, + uint32_t& b2, uint32_t& b3, uint32_t& c0, uint32_t& c1, uint32_t& c2, uint32_t& c3, + uint32_t& amax_result) { + uint32_t zero = 0; + uint32_t temp0, temp1, temp2, temp3, temp4, temp5, temp6, temp7; + asm volatile( + "wmma.mma.sync.aligned.row.row.m16n16k16.f32.bf16.bf16.f32 \n" + "{%0, %1, %2, %3, %4, %5, %6, %7}, \n" + "{%8, %9, %10, %11}, \n" + "{%12, %13, %14, %15}, \n" + "{%16, %17, %18, %19, %20, %21, %22, %23};\n\t" + : "=r"(temp0), "=r"(temp1), "=r"(temp2), "=r"(temp3), "=r"(temp4), "=r"(temp5), "=r"(temp6), + "=r"(temp7) + : "r"(a0), "r"(a1), "r"(a2), "r"(a3), "r"(b0), "r"(b1), "r"(b2), "r"(b3), "r"(zero), + "r"(zero), "r"(zero), "r"(zero), "r"(zero), "r"(zero), "r"(zero), "r"(zero)); + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c0) : "r"(temp1), "r"(temp0)); + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c1) : "r"(temp3), "r"(temp2)); + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c2) : "r"(temp5), "r"(temp4)); + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" : "=r"(c3) : "r"(temp7), "r"(temp6)); + if constexpr (kCalculateAmax) { + uint32_t max_even; + uint32_t max_odd; + // Reduction tree to amax(abs(result)) into bf16x2 reg outparam. + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" : "=r"(max_even) : "r"(c0), "r"(c2)); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" : "=r"(max_odd) : "r"(c1), "r"(c3)); + // N.B. mma is only called up to once per thread for identity and transpose respectively, so + // we don't have to accumulate into amax_result and can directly store into it. + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(amax_result) + : "r"(max_even), "r"(max_odd)); + } +} + +template +__device__ __forceinline__ void get_hadamard_matrix_fragment(uint32_t* had_frag_i, + uint16_t random_sign_mask, + uint32_t* had_frag_t, + uint16_t random_sign_mask_t) { + int32_t tid = threadIdx.x % 32; // Local tid + float temp_i[2]; + float temp_t[2]; +#pragma unroll + for (int i = 0; i < 2; i++) { + // i is the vertical fragment index. + // For a 16x16 matrix matrix fragment, 4 threads fill a fragment of 8 BF16 vals. + uint32_t r = i * 8 + tid / 4; + +#pragma unroll + for (int j = 0; j < 2; j++) { +#pragma unroll + for (int k = 0; k < 2; k++) { + // k is column position [0, 1] within a quad of 2 BF16s stored together in 32 bits. + // j is the column fragment idx selecting between even and odd fragments. + // j increments 8 columns by switching fragments. + uint32_t c = j * 8 + k + tid % 4 * 2; + // 1 -> -1.0f, 0 -> 1.0f + int32_t base_sign = __popc(r & c); + if constexpr (kReturnIdentity) { + int32_t sign_i; + // Because tensor cores want the dot product dimension, + // contiguous, the regular, non-inverse hadamard swaps + // signs of columns and rows for inverse. In a simple reference, + // x.reshape(-1, 16) @ sign @ H16, this would be opposite but + // (sign @ H16) is transposed in this fragment. + if constexpr (kInverseHadamardIdentity) { + sign_i = ((random_sign_mask >> r) ^ base_sign); + } else { + sign_i = ((random_sign_mask >> c) ^ base_sign); + } + temp_i[k] = copysignf(k16x16HadamardScale, __int_as_float(sign_i << 31)); + } + if constexpr (kReturnTransposed) { + int32_t sign_t; + if constexpr (kInverseHadamardTransposed) { + sign_t = ((random_sign_mask_t >> r) ^ base_sign); + } else { + sign_t = ((random_sign_mask_t >> c) ^ base_sign); + } + temp_t[k] = copysignf(k16x16HadamardScale, __int_as_float(sign_t << 31)); + } + } + + if constexpr (kReturnIdentity) { + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" + : "=r"(had_frag_i[i * 2 + j]) + : "f"(temp_i[1]), "f"(temp_i[0])); + } + if constexpr (kReturnTransposed) { + asm volatile("cvt.rn.bf16x2.f32 %0, %1, %2;\n\t" + : "=r"(had_frag_t[i * 2 + j]) + : "f"(temp_t[1]), "f"(temp_t[0])); + } + } + } +} + +__device__ __forceinline__ uint32_t swizzle_128B_atom_32B(uint32_t gmem_row_idx, + uint32_t gmem_col_idx) { + uint32_t smem_row_idx = gmem_row_idx; + uint32_t xor_factor = (smem_row_idx * 2) % 8; + uint32_t smem_col_idx = gmem_col_idx ^ xor_factor; + return smem_row_idx * 8 + smem_col_idx; +} + +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_HADAMARD_TRANSFORM_UTILS_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index a0dd325da0..05541fe30c 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -61,6 +61,31 @@ void nvte_hadamard_transform_cast_fusion_columnwise(const NVTETensor input, NVTE const NVTEQuantizationConfig quant_config, cudaStream_t stream); +/*! \brief Split a tensor along dimension 0 and compute RHT amaxes for each split. + * + * This function is experimental and the API is not stable. + * + * This is intended for quantizing to NVFP4 with random Hadamard + * transforms (RHT). For each tensor split, compute the maximum + * absolute value (amax) and populate the row-wise amax of the + * corresponding output tensor. Also, compute the amax after a + * transposed RHT and populate the column-wise amax of the + * corresponding output tensor. + * + * \param[in] input Input tensor. + * \param[in,out] outputs Array of NVFP4 output tensors. Only the row-wise and + * column-wise amaxes are updated. + * \param[in] split_sections Size of each tensor split along dimension 0. + * \param[in] num_tensors Number of tensor splits. + * \param[in] random_sign_mask 16-bit sign mask for RHT. + * \param[in] random_sign_mask_t 16-bit sign mask for transposed RHT. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outputs, + const size_t* split_sections, size_t num_tensors, + int random_sign_mask, int random_sign_mask_t, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index a01b2e5da0..af3f51d46f 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -265,6 +265,22 @@ void nvte_multi_tensor_compute_scale_and_scale_inv_cuda(int chunk_size, NVTETens float max_fp8, int force_pow_2_scales, float epsilon, cudaStream_t stream); +/*! \brief Split a tensor along dimension 0 and compute the amax for each split. + * + * This function is experimental and the API is not stable. + * + * For each tensor split, compute the maximum absolute value (amax) + * and populate the amax of the corresponding output tensor. + * + * \param[in] input Input tensor. + * \param[in,out] amaxes Array of output tensors. Only the amax is updated. + * \param[in] split_sections Size of each tensor split along dimension 0. + * \param[in] num_tensors Number of tensor splits. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, + size_t num_tensors, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 77fb348589..44c49b20bc 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -253,7 +253,7 @@ std::vector multi_tensor_quantize(const std::vector &ten std::vector quantizer_list); std::vector split_quantize(const at::Tensor &tensor, - const std::vector &split_sections, + const std::vector &split_sections, std::vector quantizer_list); /*************************************************************************************************** diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 7d15e436ea..b12da7542b 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -6,6 +6,7 @@ #include "transformer_engine/cast.h" +#include #include #include #include @@ -494,13 +495,15 @@ std::tuple, std::vector> bulk_allocate_mx // allocate fp4 data, fp8 scalings, and amax values // layout: [fp4_data0, ..., fp4_dataN, fp8_scaling0, ..., fp8_scalingN, amax0, ..., amaxN] // amax buffer will be zeroed out by later amax kernels, so we can use empty to allocate -std::tuple, std::vector> bulk_allocate_nvfp4_tensors( +std::tuple, std::vector, bool> bulk_allocate_nvfp4_tensors( std::vector> &shape_list, std::vector &quantizer_py_list, std::vector &quantizer_cpp_list) { init_extension(); - std::tuple, std::vector> retval; + std::tuple, std::vector, bool> retval; auto &tensor_py_list = std::get<0>(retval); auto &tensor_cpp_list = std::get<1>(retval); + auto &contiguous_data_and_scale = std::get<2>(retval); + contiguous_data_and_scale = true; // Number of tensors const size_t num_tensors = shape_list.size(); @@ -555,22 +558,29 @@ std::tuple, std::vector> bulk_allocate_nv size_t buffer_size = 0; std::vector data_offsets, scale_offsets, amax_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); - // Store ceil(product/2) bytes for fp4 (since each element is 4 bits = 0.5 bytes). - // Integer arithmetic: ceil(product / 2) == (product + 1) / 2. - buffer_size += (product(rowwise_data_shapes[i]) + 1) / 2; + // FP4 data is aligned to 256B + const auto offset = roundup(buffer_size, 256); + if (offset != buffer_size) { + 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) { - 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; + // Scales are aligned to 16B + const auto offset = roundup(buffer_size, 16); + if (offset != buffer_size) { + contiguous_data_and_scale = false; + } + scale_offsets.push_back(offset); + buffer_size = offset + product(rowwise_scale_shapes[i]) * scale_elem_size; } for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 16); // align to 16B - amax_offsets.push_back(buffer_size); - // amax is scalar in fp32, 4 bytes each - buffer_size += 4; + // 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; } // Allocate full buffer @@ -584,7 +594,7 @@ std::tuple, std::vector> bulk_allocate_nv rowwise_scale_list.emplace_back( make_torch_view(buffer, rowwise_scale_shapes[i], scale_offsets[i], torch::kUInt8)); amax_rowwise_list.emplace_back( - make_torch_view(buffer, std::vector{1}, amax_offsets[i], torch::kUInt8)); + make_torch_view(buffer, std::vector{1}, amax_offsets[i], torch::kFloat32)); } } @@ -610,22 +620,29 @@ std::tuple, std::vector> bulk_allocate_nv size_t buffer_size = 0; std::vector data_offsets, scale_offsets, amax_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); - // Store ceil(product/2) bytes for fp4 (since each element is 4 bits = 0.5 bytes). - // Integer arithmetic: ceil(product / 2) == (product + 1) / 2. - buffer_size += (product(columnwise_data_shapes[i]) + 1) / 2; + // FP4 data is aligned to 256B + const auto offset = roundup(buffer_size, 256); + if (offset != buffer_size) { + 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) { - 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; + // Scales are aligned to 16B + const auto offset = roundup(buffer_size, 16); + if (offset != buffer_size) { + contiguous_data_and_scale = false; + } + scale_offsets.push_back(offset); + buffer_size = offset + product(columnwise_scale_shapes[i]) * scale_elem_size; } for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 16); // align to 16B - amax_offsets.push_back(buffer_size); - // amax is scalar in fp32, 4 bytes each - buffer_size += 4; + // 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; } // Allocate full buffer @@ -639,7 +656,7 @@ std::tuple, std::vector> bulk_allocate_nv 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::kUInt8)); + make_torch_view(buffer, std::vector{1}, amax_offsets[i], torch::kFloat32)); } } @@ -692,10 +709,209 @@ std::tuple, std::vector> bulk_allocate_nv return retval; } +void split_quantize_nvfp4_impl(const TensorWrapper &input, + const std::vector &input_list, + std::vector &output_list, + const std::vector &split_sections, + const std::vector &quantizers) { + // Check tensor lists + const size_t num_tensors = split_sections.size(); + NVTE_CHECK(input_list.size() == num_tensors, "Expected ", num_tensors, " input tensors, but got ", + input_list.size(), "."); + NVTE_CHECK(output_list.size() == num_tensors, "Expected ", num_tensors, + " output tensors, but got ", output_list.size(), "."); + NVTE_CHECK(quantizers.size() == num_tensors, "Expected ", num_tensors, + " NVFP4 quantizers, but got ", quantizers.size(), "."); + + // Trivial cases + if (num_tensors == 0) { + return; + } + if (input.numel() == 0) { + for (const auto &tensor : input_list) { + NVTE_CHECK(tensor.numel() == 0, + "Input tensor has zero elements but got split with non-zero elements"); + } + return; + } + + // Assume all quantizers have identical config + const auto &quantizer = *quantizers.front(); + NVTE_CHECK(!quantizer.with_2d_quantization, + "NVFP4 split-quantize does not support 2D quantization"); + NVTE_CHECK(!quantizer.with_amax_reduction, + "NVFP4 split-quantize does not support amax reduction"); + + // Check input tensor shape + const size_t input_last_dim = input.ndim() > 0 ? input.size(input.ndim() - 1) : 1; + NVTE_CHECK(input_last_dim % 128 == 0, + "NVFP4 multi-quantize requires inner dim to be multiple of 128."); + + // CUDA stream + auto stream = at::cuda::getCurrentCUDAStream(); + + // Objects for TE C API + std::vector nvte_tensor_input_list; + std::vector nvte_tensor_output_list; + std::vector quant_config_list; + for (size_t i = 0; i < num_tensors; ++i) { + nvte_tensor_input_list.push_back(input_list[i].data()); + nvte_tensor_output_list.push_back(output_list[i].data()); + quant_config_list.emplace_back(QuantizationConfigWrapper()); + } + + // Stochastic rounding + std::vector te_rng_state_list; + at::Tensor rng_states_tensor; + if (quantizer.stochastic_rounding) { + // TODO(zhongbo): remove the for loop of generating rng states with a single call + // with rng_elts_per_thread = 1024 * num_tensors + // Change to the bulk generate rng states api when grouped quantize is available + const size_t rng_elts_per_thread = 1024; // Wild guess, probably can be tightened + auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); + rng_states_tensor = torch::empty({static_cast(2 * num_tensors)}, opts); + for (size_t i = 0; i < num_tensors; ++i) { + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); + int64_t *rng_state_ptr = static_cast(rng_states_tensor.data_ptr()) + i * 2; + philox_unpack(philox_args, rng_state_ptr); + te_rng_state_list.push_back(makeTransformerEngineTensor( + static_cast(rng_state_ptr), std::vector{2}, DType::kInt64)); + quant_config_list[i].set_rng_state(te_rng_state_list[i].data()); + quant_config_list[i].set_stochastic_rounding(true); + } + } + + // Perform multi-tensor quantization + if (quantizer.with_rht) { // Quantize row-wise data, RHT+quantize column-wise data + // Check that config is supported + NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); + + // Compute amaxes + if (quantizer.with_post_rht_amax) { + // We need: + // 1. Rowwise amax = amax for input + // 2. Columnwise amax = amax for RHT(input.t) + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_amax( + input.data(), reinterpret_cast(nvte_tensor_output_list.data()), + split_sections.data(), num_tensors, 0, quantizer.rht_matrix_random_sign_mask_t, stream); + }); + } else { + // RHT is enabled, but amax is pre-RHT amax + NVTE_ERROR("NVFP4 split-quantize does not yet support pre-RHT amax"); + } + + // Check that RHT matrix is available + NVTE_CHECK(quantizer.rht_matrix.defined() && quantizer.rht_matrix.numel() > 0, + "RHT matrix is not available."); + auto rht_matrix_nvte = makeTransformerEngineTensor(quantizer.rht_matrix); + + // Quantize tensors individually + NVTE_SCOPED_GIL_RELEASE({ + for (size_t i = 0; i < num_tensors; i++) { + if (input_list[i].numel() == 0) { + continue; // Skip tensors with no elements + } + + // Direct NVFP4 quantization for row-wise data + if (quantizer.rowwise_usage) { + auto out_rowwise_data = output_list[i].get_rowwise_data(); + auto out_rowwise_scale_inv = output_list[i].get_rowwise_scale_inv(); + auto out_rowwise_amax = output_list[i].get_amax(); + TensorWrapper out_rowwise(output_list[i].scaling_mode()); + out_rowwise.set_rowwise_data(out_rowwise_data.data_ptr, + static_cast(out_rowwise_data.dtype), + out_rowwise_data.shape); + out_rowwise.set_rowwise_scale_inv(out_rowwise_scale_inv.data_ptr, + static_cast(out_rowwise_scale_inv.dtype), + out_rowwise_scale_inv.shape); + out_rowwise.set_amax(out_rowwise_amax.data_ptr, + static_cast(out_rowwise_amax.dtype), out_rowwise_amax.shape); + nvte_quantize_v2(input_list[i].data(), out_rowwise.data(), quant_config_list[i], stream); + } + + // RHT + NVFP4 quantize for column-wise data + if (quantizer.columnwise_usage) { + // Get the output column-wise data, scale_inv, and amax + auto out_columnwise_data = output_list[i].get_columnwise_data(); + auto out_columnwise_scale_inv = output_list[i].get_columnwise_scale_inv(); + auto out_columnwise_amax = output_list[i].get_columnwise_amax(); + + // Flatten column-wise data to 2D + 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 i = 1; i < colwise_data_shape.ndim; ++i) { + last_dim *= colwise_data_shape.data[i]; + } + colwise_data_shape_2d.push_back(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(output_list[i].scaling_mode()); + out_transpose.set_rowwise_data(out_columnwise_data.data_ptr, + static_cast(out_columnwise_data.dtype), + colwise_data_shape_2d); + out_transpose.set_rowwise_scale_inv(out_columnwise_scale_inv.data_ptr, + static_cast(out_columnwise_scale_inv.dtype), + out_columnwise_scale_inv.shape); + out_transpose.set_amax(out_columnwise_amax.data_ptr, + static_cast(out_columnwise_amax.dtype), + out_columnwise_amax.shape); + + // RHT + NVFP4 quantize kernel + nvte_hadamard_transform_cast_fusion_columnwise(input_list[i].data(), out_transpose.data(), + rht_matrix_nvte.data(), + quant_config_list[i], stream); + } + } + }); + + } else { // NVFP4 quantize + // We need: + // 1. Rowwise amax = amax for input + // 2. Columnwise amax = amax for input too + // Columnwise amax will be filled with a fused D2D copy from rowwise amax + // Note that the multi compute amax API expects rowwise amax pointer to be not null + // So we need to set the pointer accordingly to make colwise-only quantization work + std::vector orig_amax_ptr_list; + for (size_t i = 0; i < num_tensors; i++) { + auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr; + orig_amax_ptr_list.push_back(rowwise_amax_ptr); + auto columnwise_amax_ptr = output_list[i].get_columnwise_amax().data_ptr; + void *amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr; + NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer"); + output_list[i].set_amax(amax_ptr, DType::kFloat32, std::vector{1}); + } + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_amax(input.data(), reinterpret_cast(nvte_tensor_output_list.data()), + split_sections.data(), num_tensors, stream); + }); + for (size_t i = 0; i < num_tensors; i++) { + output_list[i].set_amax(orig_amax_ptr_list[i], DType::kFloat32, std::vector{1}); + } + + // Quantize tensors individually + NVTE_SCOPED_GIL_RELEASE({ + for (size_t i = 0; i < num_tensors; i++) { + // skip this round if input is empty + if (input_list[i].numel() == 0) { + continue; + } + nvte_quantize_v2(input_list[i].data(), output_list[i].data(), quant_config_list[i], stream); + } + }); + } +} + } // namespace std::vector split_quantize(const at::Tensor &tensor, - const std::vector &split_sections, + const std::vector &split_sections, std::vector quantizer_list) { init_extension(); @@ -726,8 +942,6 @@ std::vector split_quantize(const at::Tensor &tensor, const size_t dim0_stride = input_shape[0] == 0 ? 0 : input_py.element_size() * input_size / input_shape[0]; for (size_t i = 0; i < num_splits; ++i) { - NVTE_CHECK(split_sections[i] >= 0, "Attempted to split tensor with shape=", input_shape, - " along dim 0 with split_sections=", split_sections); NVTE_CHECK(dim0_offset + split_sections[i] <= input_shape[0], "Attempted to split tensor with shape=", input_shape, " along dim 0 with split_sections=", split_sections); @@ -745,65 +959,96 @@ std::vector split_quantize(const at::Tensor &tensor, quantizer_cpp_list.push_back(convert_quantizer(quantizer_list[i])); } - // For FP8 block-scaling, we construct output tensors with bulk allocations - // For MXFP8, we also use bulk allocations - bool use_fused_bulk_alloc = true; - for (size_t i = 0; i < quantizer_list.size(); i++) { - if (!detail::IsFloat8BlockwiseQuantizers(quantizer_list[i].ptr()) && - !detail::IsMXFP8Quantizers(quantizer_list[i].ptr()) && - !detail::IsNVFP4Quantizers(quantizer_list[i].ptr())) { - use_fused_bulk_alloc = false; - break; - } + // Choose implementation for allocating and populating tensors + enum class AllocationMethod { UNFUSED, BULK_FP8_BLOCKWISE, BULK_MXFP8, BULK_NVFP4 }; + enum class QuantizationMethod { UNFUSED, FUSED_NVFP4 }; + AllocationMethod allocation_method = AllocationMethod::UNFUSED; + QuantizationMethod quantization_method = QuantizationMethod::UNFUSED; + if (std::all_of(quantizer_list.begin(), quantizer_list.end(), + [](const py::handle &quantizer) -> bool { + return detail::IsFloat8BlockwiseQuantizers(quantizer.ptr()); + })) { + allocation_method = AllocationMethod::BULK_FP8_BLOCKWISE; + } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(), + [](const py::handle &quantizer) -> bool { + return detail::IsMXFP8Quantizers(quantizer.ptr()); + })) { + allocation_method = AllocationMethod::BULK_MXFP8; + } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(), + [](const py::handle &quantizer) -> bool { + return detail::IsNVFP4Quantizers(quantizer.ptr()); + })) { + allocation_method = AllocationMethod::BULK_NVFP4; + quantization_method = QuantizationMethod::FUSED_NVFP4; } // Allocate output tensors std::vector output_cpp_list; std::vector output_py_list; - if (!use_fused_bulk_alloc) { - // Allocate output tensors individually - for (size_t i = 0; i < num_splits; ++i) { - auto [output_cpp, output_py] = - quantizer_cpp_list[i]->create_tensor(split_shapes[i], input_dtype); - output_cpp_list.emplace_back(std::move(output_cpp)); - output_py_list.emplace_back(std::move(output_py)); - } - } else { - // TODO(zhongbo): make a better api to make this part less hacky - bool is_fp8_blockwise = detail::IsFloat8BlockwiseQuantizers(quantizer_list[0].ptr()); - bool is_mxfp8 = detail::IsMXFP8Quantizers(quantizer_list[0].ptr()); - bool is_nvfp4 = detail::IsNVFP4Quantizers(quantizer_list[0].ptr()); - if (is_fp8_blockwise) { - // FP8 block-scaling: construct output tensors with bulk allocations + switch (allocation_method) { + case AllocationMethod::BULK_FP8_BLOCKWISE: { + // Bulk allocation for FP8 block-scaling tensors std::vector blockwise_quantizers; for (auto &quantizer : quantizer_cpp_list) { blockwise_quantizers.push_back(static_cast(quantizer.get())); } std::tie(output_py_list, output_cpp_list) = bulk_allocate_fp8_blockwise_tensors(split_shapes, quantizer_list, blockwise_quantizers); - } else if (is_mxfp8) { - // MXFP8: construct output tensors with bulk allocations + break; + } + case AllocationMethod::BULK_MXFP8: { + // Bulk allocation for MXFP8 tensors std::vector mxfp8_quantizers; for (auto &quantizer : quantizer_cpp_list) { mxfp8_quantizers.push_back(static_cast(quantizer.get())); } std::tie(output_py_list, output_cpp_list) = bulk_allocate_mxfp8_tensors(split_shapes, quantizer_list, mxfp8_quantizers); - } else if (is_nvfp4) { - // NVFP4: construct output tensors with bulk allocations + break; + } + case AllocationMethod::BULK_NVFP4: { + // Bulk allocation for NVFP4 tensors std::vector nvfp4_quantizers; for (auto &quantizer : quantizer_cpp_list) { nvfp4_quantizers.push_back(static_cast(quantizer.get())); } - std::tie(output_py_list, output_cpp_list) = + bool contiguous_data_and_scale; + std::tie(output_py_list, output_cpp_list, contiguous_data_and_scale) = bulk_allocate_nvfp4_tensors(split_shapes, quantizer_list, nvfp4_quantizers); - } else { - NVTE_CHECK(false, "Expected either FP8 block-scaling or MXFP8 quantizer"); + if (!contiguous_data_and_scale) { + // Avoid fused quantize kernel if data is not contiguous + quantization_method = QuantizationMethod::UNFUSED; + } + break; + } + default: { + // Allocate output tensors individually + for (size_t i = 0; i < num_splits; ++i) { + auto [output_cpp, output_py] = + quantizer_cpp_list[i]->create_tensor(split_shapes[i], input_dtype); + output_cpp_list.emplace_back(std::move(output_cpp)); + output_py_list.emplace_back(std::move(output_py)); + } } } - // Perform multi-tensor quantization - multi_tensor_quantize_impl(input_list, quantizer_list, quantizer_cpp_list, output_cpp_list); + // Quantize into output tensors + switch (quantization_method) { + case QuantizationMethod::FUSED_NVFP4: { + // Fused NVFP4 quantize kernel + auto input_nvte = makeTransformerEngineTensor(input_dptr, input_shape, input_dtype); + std::vector nvfp4_quantizers; + for (auto &quantizer : quantizer_cpp_list) { + nvfp4_quantizers.push_back(static_cast(quantizer.get())); + } + split_quantize_nvfp4_impl(input_nvte, input_list, output_cpp_list, split_sections, + nvfp4_quantizers); + break; + } + default: + // General multi-tensor quantization + multi_tensor_quantize_impl(input_list, quantizer_list, quantizer_cpp_list, output_cpp_list); + } return output_py_list; } diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index fd9b9b4377..d59a26ca33 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex -from ..quantization import FP8GlobalStateManager +from ..quantization import FP8GlobalStateManager, get_align_size_for_quantization from ..jit import no_torch_dynamo @@ -114,14 +114,8 @@ def forward( assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." if self.align_size is None: - self.align_size = ( - 32 - if ( - FP8GlobalStateManager.get_fp8_recipe().mxfp8() - or FP8GlobalStateManager.get_fp8_recipe().nvfp4() - ) - else 16 - ) + recipe = FP8GlobalStateManager.get_fp8_recipe() + self.align_size = get_align_size_for_quantization(recipe) # FP8 padding calculate padded_m_splits = [ diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index 58187c20ea..6f9702a96f 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -10,7 +10,7 @@ import transformer_engine_torch as tex -from ..quantization import FP8GlobalStateManager +from ..quantization import FP8GlobalStateManager, get_align_size_for_quantization from ..jit import no_torch_dynamo @@ -112,14 +112,8 @@ def forward( assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." if self.align_size is None: - self.align_size = ( - 32 - if ( - FP8GlobalStateManager.get_fp8_recipe().mxfp8() - or FP8GlobalStateManager.get_fp8_recipe().nvfp4() - ) - else 16 - ) + recipe = FP8GlobalStateManager.get_fp8_recipe() + self.align_size = get_align_size_for_quantization(recipe) # FP8 padding calculate padded_m_splits = [ diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 030370b9db..acc2e55320 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -40,6 +40,7 @@ "is_fp8_block_scaling_available", "is_nvfp4_available", "get_default_recipe", + "get_align_size_for_quantization", ] @@ -114,6 +115,15 @@ def get_default_recipe() -> Recipe: return get_default_fp8_recipe() +def get_align_size_for_quantization(recipe: Recipe) -> int: + """Get the alignment size for quantization.""" + if recipe.mxfp8(): + return 32 + if recipe.nvfp4(): + return 64 + return 16 + + def get_fp8_torch_dtype(fp8_recipe: Recipe, fprop_tensor: bool = True) -> torch.dtype: """Get fp8 data type according to recipe and tensor""" if fp8_recipe.fp8_format == Format.E4M3 or ( From d52ed471b1f20f541ba907fecfcf15c14dbc7571 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Wed, 26 Nov 2025 01:07:56 +0530 Subject: [PATCH 098/521] FSDP2 Allgather Perf improvement and support for FusedAdam with FSDP2 (#2370) * fix ci issue Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert back testing changes Signed-off-by: Varun Thumbe * remove quantizer copy + fused adam working 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 mxfp8 bug, god knows who created it Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/pytorch/optimizers/fused_adam.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update comment Signed-off-by: Tim Moon --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Signed-off-by: Tim Moon 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 --- .../pytorch/optimizers/fused_adam.py | 7 +++- .../pytorch/tensor/float8_tensor.py | 28 +++++++++---- .../pytorch/tensor/mxfp8_tensor.py | 41 +++++++++++-------- 3 files changed, 48 insertions(+), 28 deletions(-) diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 73b312ec28..b5c87b4815 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -11,6 +11,7 @@ import warnings import torch +from torch.distributed._tensor import DTensor import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer from .multi_tensor_apply import multi_tensor_applier @@ -567,8 +568,10 @@ def step(self, closure=None, grad_scaler=None): unscaled_lists[name].append(unscaled) scaled_lists[name].append(state[name]) state_scales[name].append(self._scales[p][name]) - - if isinstance(p, Float8Tensor): + if isinstance(p, Float8Tensor) or ( + isinstance(p, DTensor) and isinstance(p._local_tensor, Float8Tensor) + ): + 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) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 80e7ed4674..f44ed33b9e 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -713,9 +713,8 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [transpose, t_shape] + list(args[2:]), kwargs, ) - # deep copy the scale inverse tensor and quantizer as well. scale_inv = tensor._scale_inv.detach().clone() - quantizer = tensor._quantizer.copy() + quantizer = tensor._quantizer # Deep-copied in constructor out_tensor = Float8Tensor( data=func_out, shape=func_out.shape, @@ -820,7 +819,7 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m # sure that updated Quantized weight tensor have same scale inverse across all shards. self._quantizer.amax_reduction_group = mesh.get_group() self._quantizer.with_amax_reduction = True - quantizer = self._quantizer.copy() # quantizer to be used for allgathered weights + fsdp_state = _get_module_fsdp_state(module) reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward # If weights are resharded after forward pass, then its enough to set the quantizer usages @@ -833,9 +832,13 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m is_backward_pass = training_state == TrainingState.PRE_BACKWARD # In case of hopper/L40, only one of data/transpose is needed # based on forward or backward pass. So setting the quantizer usages appropriately. - quantizer.set_usage(rowwise=not is_backward_pass, columnwise=is_backward_pass) + rowwise_usage = not is_backward_pass + columnwise_usage = is_backward_pass + else: + rowwise_usage = True + columnwise_usage = self._quantizer.columnwise_usage sharded_tensors = (self._data,) - metadata = (self._scale_inv, self._fp8_dtype, quantizer) + metadata = (self._scale_inv, rowwise_usage, columnwise_usage, self._fp8_dtype) return sharded_tensors, metadata def fsdp_post_all_gather( @@ -861,7 +864,7 @@ def fsdp_post_all_gather( """ (data,) = all_gather_outputs - (fp8_scale_inv, fp8_dtype, quantizer) = metadata + (fp8_scale_inv, rowwise_usage, columnwise_usage, fp8_dtype) = metadata orig_shape = data.size() # Quantizer has only columnwise usage set for backward pass # In Blackwell+ architectures, transpose is not needed at all, @@ -870,20 +873,27 @@ def fsdp_post_all_gather( if out is not None: out._data = data else: + # We ll be here when post all gather is called the first time. + # Float8Tensor constructor makes a copy of the quantizer to + # save as its own quantizer. For the consequent iterations, + # the same quantizer is used. Copy is needed in the first iteration, + # since we need different quantizers for sharded and allgathered tensors. + # and self._quantizer belongs to the sharded parameter. fp8_args = { "shape": orig_shape, "dtype": param_dtype, "fp8_scale_inv": fp8_scale_inv, "fp8_dtype": fp8_dtype, - "quantizer": quantizer, + "quantizer": self._quantizer, "requires_grad": False, "data": data, } out = Float8Tensor(**fp8_args) + out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) out.update_usage( - rowwise_usage=quantizer.rowwise_usage, - columnwise_usage=quantizer.columnwise_usage, + rowwise_usage=rowwise_usage, + columnwise_usage=columnwise_usage, ) return out, all_gather_outputs diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index cf65814656..3b9bdd2fe7 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -552,7 +552,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): rowwise_scale_inv=rowwise_scale_inv, columnwise_data=columnwise_data, columnwise_scale_inv=columnwise_scale_inv, - quantizer=tensor._quantizer.copy(), + quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, ) @@ -583,7 +583,6 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m fsdp_state = _get_module_fsdp_state(module) reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward - quantizer = self._quantizer.copy() # Remove padding from scale inverses before allgather # Rowwise scale_inv should be divisible by [128,4], columnwise by [4, 128] rowwise_scale_inv = self._rowwise_scale_inv @@ -601,9 +600,8 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m if columnwise_scale_inv.size(0) != flattened_in_shape0: columnwise_scale_inv = columnwise_scale_inv[:flattened_in_shape0] - sharded_tensors = (self._rowwise_data, rowwise_scale_inv) - # If weights are resharded after forward pass, then its enough to set the quantizer usages - # based on whether its forward or backward pass for the allgathered weights. + # If weights are resharded after forward pass, then its enough to send one row/col + # usage based on whether its forward or backward pass for the allgathered weights. # If not resharded after forward pass, the same weights allgathered in forward # are used again in backward. And hence if we need the columnwise data/scale_inv, # we need to send them as well for allgather in forward pass itself. @@ -611,18 +609,24 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m training_state = fsdp_state._fsdp_param_group._training_state is_backward_pass = training_state == TrainingState.PRE_BACKWARD # Allgather only the necessary tensors based on forward/backward pass - quantizer.set_usage(rowwise=not is_backward_pass, columnwise=is_backward_pass) + rowwise_usage = not is_backward_pass + columnwise_usage = is_backward_pass sharded_tensors = ( (self._columnwise_data, columnwise_scale_inv) if is_backward_pass - else sharded_tensors + else (self._rowwise_data, rowwise_scale_inv) ) else: - if quantizer.columnwise_usage: + # rowwise usage is always needed for forward pass. + rowwise_usage = True + sharded_tensors = (self._rowwise_data, rowwise_scale_inv) + columnwise_usage = self._quantizer.columnwise_usage + if columnwise_usage: # If weights are not resharded after forward, then both # rowwise and columnwise data/scale_inv need to be allgathered. sharded_tensors += (self._columnwise_data, columnwise_scale_inv) - metadata = (self._fp8_dtype, quantizer) + + metadata = (self._fp8_dtype, rowwise_usage, columnwise_usage) return sharded_tensors, metadata def fsdp_post_all_gather( @@ -645,12 +649,10 @@ def fsdp_post_all_gather( Tuple[MXFP8Tensor, Tuple[torch.Tensor, ...]]: Allgathered MXFP8Tensor and tuple of internal tensors used by the MXFP8Tensor that was being computed after allgather. """ - fp8_dtype, quantizer = metadata - rowwise_data, rowwise_scale_inv = ( - all_gather_outputs[:2] if quantizer.rowwise_usage else (None, None) - ) + fp8_dtype, rowwise_usage, columnwise_usage = metadata + rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] if rowwise_usage else (None, None) columnwise_data, columnwise_scale_inv = ( - all_gather_outputs[-2:] if quantizer.columnwise_usage else (None, None) + all_gather_outputs[-2:] if columnwise_usage else (None, None) ) # Add padding to scale_inv tensors to be multiples of [128, 4]for rowwise and [4, 128] for columnwise @@ -675,8 +677,13 @@ def fsdp_post_all_gather( out._rowwise_scale_inv = rowwise_scale_inv out._columnwise_data = columnwise_data out._columnwise_scale_inv = columnwise_scale_inv - out._quantizer = quantizer else: + # We ll be here when post all gather is called the first time. + # MXFP8Tensor constructor makes a copy of the quantizer to + # save as its own quantizer. For the consequent iterations, + # the same quantizer is used. Copy is needed in the first iteration, + # since we need different quantizers for sharded and allgathered tensors. + # and self._quantizer belongs to the sharded parameter. out = MXFP8Tensor( rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, @@ -685,9 +692,9 @@ def fsdp_post_all_gather( fp8_dtype=fp8_dtype, dtype=param_dtype, shape=rowwise_data.shape if rowwise_data is not None else columnwise_data.shape, - quantizer=quantizer, + quantizer=self._quantizer, ) - + out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) return out, all_gather_outputs @classmethod From b3c25057405fc35d10be8109b635696e341ccf86 Mon Sep 17 00:00:00 2001 From: Pingtian Li <158665726+Wohox@users.noreply.github.com> Date: Wed, 26 Nov 2025 05:45:31 +0800 Subject: [PATCH 099/521] [Pytorch] Fix backward_dw cuda graph order (#2376) * fix backward_dw cuda graph order Signed-off-by: Pingtian Li * add validation for num_layers_per_chunk Signed-off-by: Pingtian Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pingtian Li Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/graph.py | 114 ++++++++++++++++++++++------ 1 file changed, 90 insertions(+), 24 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index f55f1dd128..3e515eecd5 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -7,6 +7,7 @@ import contextlib import gc import warnings +from math import ceil from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union import torch @@ -127,6 +128,8 @@ def _make_graphed_callables( ) # Check sizes of args + _order_without_wgrad = None + delay_wgrad_compute = False if _order is None: assert len(sample_args) == len(callables) assert len(sample_kwargs) == len(callables) @@ -145,17 +148,34 @@ def _make_graphed_callables( # values indicate backward passes. Each # entry in sample_args corresponds to one of the forward # passes. - num_model_chunks = max(_order) - num_microbatches = len(_order) // num_model_chunks // 2 - assert num_model_chunks * num_microbatches * 2 == len(_order) + _order_without_wgrad = [] + for c_id in _order: + if ceil(c_id) != c_id: + delay_wgrad_compute = True + continue + _order_without_wgrad.append(c_id) + num_model_chunks = max(_order_without_wgrad) + num_microbatches = len(_order_without_wgrad) // num_model_chunks // 2 + assert num_model_chunks * num_microbatches * 2 == len(_order_without_wgrad) + + # When delay_wgrad_compute is enabled, each layer is treated as a model chunk, which + # allows for fine-grained graph capture order. + if delay_wgrad_compute: + assert ( + _num_layers_per_chunk is not None + ), "'_num_layers_per_chunk' must be provided when delay_wgrad_compute is True." + for num_layers in _num_layers_per_chunk: + assert ( + num_layers == 1 + ), "Each model chunk must have only one layer when delay_wgrad_compute is True." # Determine number of layers in each model chunk. if _num_layers_per_chunk is None: - assert len(sample_args) * 2 >= len(_order) and ( - len(sample_args) * 2 % len(_order) == 0 + assert len(sample_args) * 2 >= len(_order_without_wgrad) and ( + len(sample_args) * 2 % len(_order_without_wgrad) == 0 ), ( - f"{len(sample_args)} * 2 >= {len(_order)} and {len(sample_args)} * 2 %" - f" {len(_order)} == 0" + f"{len(sample_args)} * 2 >= {len(_order_without_wgrad)} and {len(sample_args)} * 2" + f" % {len(_order_without_wgrad)} == 0" ) num_layers = len(sample_args) // num_model_chunks // num_microbatches _num_layers_per_chunk = [num_layers] * num_model_chunks @@ -175,7 +195,7 @@ def _make_graphed_callables( + f"entries when order input is provided but got {len(callables)}." ) assert len(sample_args) == total_num_layers * num_microbatches, ( - f"Expected {total_num_layers * num_microbatches}" + f"Expected {total_num_layers * num_microbatches} " + f"args tuple, but got {len(sample_args)}." ) @@ -214,7 +234,7 @@ def _make_graphed_callables( consumed_sample_q = {} fwd_idx = [0] * num_model_chunks for c_id in _order: - m_chunk = abs(c_id) - 1 + m_chunk = abs(ceil(c_id)) - 1 if c_id > 0: sample_start_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( @@ -241,6 +261,8 @@ def _make_graphed_callables( sample_args[per_callable_fwd_idx] = sample_args[reuse_fwd_idx] sample_kwargs[per_callable_fwd_idx] = sample_kwargs[reuse_fwd_idx] fwd_idx[m_chunk] += 1 + elif ceil(c_id) != c_id: + continue else: num_consumed_samples = min( len(fwd_sample_qs[m_chunk]), _num_layers_per_chunk[m_chunk] @@ -477,9 +499,11 @@ def hook_fn( fwd_idx = [0] * num_model_chunks bwd_idx = [0] * num_model_chunks static_grad_outputs_dict = {} + wgrad_validation_list = [None] * len(_order) previous_chunk_last_callable_bwd_idx = None - for c_id in _order: + for i, c_id in enumerate(_order): if c_id > 0: + assert isinstance(c_id, int), "Forward order value must be an integer." # Capture forward graph for model chunk c_id, microbatch fwd_idx[c_id-1] m_chunk = c_id - 1 for l_no in range(_num_layers_per_chunk[m_chunk]): @@ -499,12 +523,65 @@ def hook_fn( fwd_idx[m_chunk] += 1 else: # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1] - m_chunk = -c_id - 1 + m_chunk = -ceil(c_id) - 1 previous_per_callable_bwd_idx = None for l_no in list(reversed(range(_num_layers_per_chunk[m_chunk]))): per_callable_bwd_idx = (_prefix_num_layers[m_chunk] * num_microbatches) + ( bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no ) + if ceil(c_id) == c_id and need_bwd_dw_graph[per_callable_bwd_idx]: + # Check if bwd graph has corresponding wgrad graph: + # Number of dgrad backward graphs should be equal to number of + # wgrad backward graphs. + # Note: For MCore, the validation rule is more strict (the next backward + # of dgrad graph must be corresponding wgrad graph). + if wgrad_validation_list[i] is None: + same_bwd_c_id_list = [i] + num_wgrad_c_id = 0 + for idx in range(i + 1, len(_order)): + if _order[idx] > 0: + continue + if _order[idx] == c_id: + same_bwd_c_id_list.append(idx) + if _order[idx] + 0.5 == c_id: + num_wgrad_c_id += 1 + if len(same_bwd_c_id_list) == num_wgrad_c_id: + for same_c_id_idx in same_bwd_c_id_list: + wgrad_validation_list[same_c_id_idx] = True + break + if len(same_bwd_c_id_list) < num_wgrad_c_id: + # It's impossible to have more wgrad than dgrad. + wgrad_validation_list[i] = False + break + if wgrad_validation_list[i] is None: + wgrad_validation_list[i] = False + assert wgrad_validation_list[i], ( + f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number " + f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}." + ) + elif ceil(c_id) != c_id: + per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk] + assert is_training, "Only training mode supports backward_dw." + # If no one module needs the backward_dw, the bwd_dw_graph will be empty. + # So skip capturing it. For backward_dw, the order value is c_id - 0.5 to indicate + # the specific order of backward_dw. + assert ceil(c_id) - c_id == 0.5, ( + "The order diff of wgrad and dgrad must be 0.5, " + f"get {ceil(c_id) - c_id}." + ) + assert need_bwd_dw_graph[ + per_callable_bwd_idx + ], "No module needs wgrad computation but get float in order" + bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx] + with _graph_context_wrapper(bwd_dw_graph, pool=mempool): + for module in visited_te_modules[per_callable_bwd_idx]: + if ( + hasattr(module, "need_backward_dw") + and module.need_backward_dw() + ): + module.backward_dw() + continue + static_input_surface = per_callable_static_input_surfaces[per_callable_bwd_idx] static_outputs = per_callable_static_outputs[per_callable_bwd_idx] bwd_graph = bwd_graphs[per_callable_bwd_idx] @@ -537,17 +614,6 @@ def hook_fn( allow_unused=allow_unused_input, retain_graph=retain_graph_in_backward, ) - # If no one module needs the backward_dw, the bwd_dw_graph will be empty. - # So skip capturing it. - if need_bwd_dw_graph[per_callable_bwd_idx]: - bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx] - with _graph_context_wrapper(bwd_dw_graph, pool=mempool): - for module in visited_te_modules[per_callable_bwd_idx]: - if ( - hasattr(module, "need_backward_dw") - and module.need_backward_dw() - ): - module.backward_dw() # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs # that don't require grad. I couldn't think of a one-liner for this pattern. @@ -596,8 +662,8 @@ def hook_fn( per_callable_static_grad_inputs[idx] ) previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx - - bwd_idx[m_chunk] += 1 + if ceil(c_id) == c_id: + bwd_idx[m_chunk] += 1 else: # Capture forward graphs per_callable_static_outputs = [] From 9f61f8a594bf61dc26e8706303fae2ec38107749 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, 26 Nov 2025 00:58:31 +0100 Subject: [PATCH 100/521] [PyTorch Debug] Debug support for GroupedLinear (#1953) * main Signed-off-by: Pawel Gadzinski * docs Signed-off-by: Pawel Gadzinski * add Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * test fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix 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> --- docs/debug/2_config_file_structure.rst | 2 + tests/pytorch/debug/test_log.py | 22 +++ tests/pytorch/test_numerics.py | 15 ++ .../debug/pytorch/debug_quantization.py | 23 ++- .../pytorch/cpp_extensions/gemm.py | 34 +++- transformer_engine/pytorch/module/base.py | 19 +-- .../pytorch/module/grouped_linear.py | 149 ++++++++++++++---- 7 files changed, 213 insertions(+), 51 deletions(-) diff --git a/docs/debug/2_config_file_structure.rst b/docs/debug/2_config_file_structure.rst index f1069b0c80..fe82df638f 100644 --- a/docs/debug/2_config_file_structure.rst +++ b/docs/debug/2_config_file_structure.rst @@ -107,6 +107,8 @@ The ``TransformerLayer`` in Transformer Engine is a composition of multiple sub- depending on the configuration. Some layers, like ``LayerNormLinear``, are fusions of two layers: ``LayerNorm`` and ``Linear``. When referring to such layers in precision debug tools, only the ``Linear`` part is affected. +For `GroupedLinear` layer, the names of underlying GEMMS are of the form `layer_name.gemm_n`, where `n` is the index of the GEMM. + Below is an example ``TransformerLayer`` with four linear layers that can be influenced by the precision debug tools. .. figure:: ./img/names.svg diff --git a/tests/pytorch/debug/test_log.py b/tests/pytorch/debug/test_log.py index 0f833d41fb..5456ab820b 100644 --- a/tests/pytorch/debug/test_log.py +++ b/tests/pytorch/debug/test_log.py @@ -363,6 +363,28 @@ def test_log_every_3_or_5_layers(layer, configs_dir, feature_dirs): TEDebugState._reset() +def test_log_grouped_gemm(feature_dirs): + if not fp8_available: + pytest.skip(reason_for_no_fp8) + + log_all_stats_config = LOG_QUANTIZED_CONFIG_BASE.format(stats=", ".join(all_stats)) + with debug_session(log_all_stats_config, feature_dirs) as log_dir: + model = te.GroupedLinear(3, 128, 128, name="linear1", params_dtype=torch.bfloat16) + inp = torch.randn((1, 128, 128), dtype=torch.bfloat16).cuda() + m_splits = [64, 32, 32] + with te.fp8_autocast(fp8_recipe=recipe.DelayedScaling()): + output = model(inp, m_splits=m_splits) + loss = output.sum() + loss.backward() + debug_api.step() + + output = read_log(log_dir) + + assert "gemm_0" in output, "gemm0 not found in output" + assert "gemm_1" in output, "gemm1 not found in output" + assert "gemm_2" in output, "gemm2 not found in output" + + def test_compute_max_blockwise_dynamic_range_direct(): """Direct unit test for compute_max_blockwise_dynamic_range function. diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index f809618346..c30c7bcc76 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1270,6 +1270,9 @@ def test_linear_accuracy(dtype, bs, model, return_bias, bias): @pytest.mark.parametrize("bias", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) def test_linear_accuracy_delay_wgrad_compute(dtype, bs, model, bias, fuse_wgrad_accumulation): + if NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("Delayed wgrad compute is not supported in debug mode.") + config = model_configs[model] te_linear_ref = Linear( @@ -1566,6 +1569,9 @@ def test_layernorm_linear_accuracy( def test_layernorm_linear_accuracy_delay_wgrad_compute( dtype, bs, model, normalization, zero_centered_gamma, bias, fuse_wgrad_accumulation ): + if NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("Delayed wgrad compute is not supported in debug mode.") + config = model_configs[model] ln_linear_ref = LayerNormLinear( @@ -1705,6 +1711,9 @@ def test_layernorm_mlp_accuracy_delay_wgrad_compute( bias, fuse_wgrad_accumulation, ): + if NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("Delayed wgrad compute is not supported in debug mode.") + config = model_configs[model] ln_mlp = LayerNormMLP( @@ -1899,6 +1908,8 @@ def test_grouped_linear_accuracy( 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.") config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -2041,6 +2052,8 @@ def test_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") + if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: + pytest.skip("Delayed wgrad compute is not supported in debug mode.") config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -2757,6 +2770,7 @@ def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass): A, B, out, + [None] * z, dtype, m_splits=m_splits, grad=grad, @@ -2904,6 +2918,7 @@ def test_fp8_grouped_gemm(shape, accumulate): A_fp8, B_fp8, out, + [None] * z, dtype, m_splits=m_splits, accumulate=accumulate, diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 7f45a24e20..caa5f5a7ee 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -9,7 +9,7 @@ """ from __future__ import annotations -from typing import Optional, Tuple, Iterable, Union +from typing import Optional, Tuple, Iterable, Union, List import torch import transformer_engine_torch as tex @@ -556,6 +556,23 @@ def set_usage(self, rowwise: bool = None, columnwise: bool = None): if not self.output_tensor: self._update_parent_quantizer_usage() + @classmethod + def multi_tensor_quantize( + cls, + tensor: torch.Tensor, + quantizers: List[Quantizer], + m_splits: List[int], + activation_dtype: torch.dtype, + ) -> List[DebugQuantizedTensor]: + """ + Splits a tensor into a list of tensors and quantizes each tensor using a list of quantizers. + """ + tensors = torch.split(tensor, m_splits) + output = [] + for tensor, quantizer in zip(tensors, quantizers): + output.append(quantizer.quantize(tensor, dtype=activation_dtype)) + return output + class DebugQuantizedTensor(QuantizedTensorStorage): """ @@ -623,9 +640,9 @@ def get_tensor(self, transpose: bool): """Is used in the python gemm() to get tensor or transpose of the tensor.""" return self.rowwise_gemm_tensor if not transpose else self.columnwise_gemm_tensor - def size(self): + def size(self, *args): """Size of the tensor.""" - return self.rowwise_gemm_tensor.size() + return self.rowwise_gemm_tensor.size(*args) def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None): """Update usage of the tensor.""" diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 1a2d619b0f..d4ff0b96d9 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -114,7 +114,6 @@ def general_gemm( assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." transa = layout[0] == "T" transb = layout[1] == "T" - # assert quantization_params is None, "FP8 output not supported yet" alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) @@ -215,6 +214,7 @@ def general_grouped_gemm( A: List[torch.Tensor], B: List[torch.Tensor], out: List[torch.Tensor], + quantization_params: List[Optional[Quantizer]], out_dtype: torch.dtype, layout: str = "TN", m_splits: Optional[List[int]] = None, @@ -247,7 +247,7 @@ def general_grouped_gemm( if grad and use_bias: grad_bias = [ - torch.empty(B[i].shape[1], dtype=out[0].dtype, device="cuda") for i in range(num_gemms) + torch.empty(B[i].size(1), dtype=out[0].dtype, device="cuda") for i in range(num_gemms) ] else: grad_bias = empty_tensors @@ -257,6 +257,36 @@ def general_grouped_gemm( else: bias_dtype = TE_DType[torch.bfloat16] + if isinstance(quantization_params[0], DebugQuantizer): + assert not gelu, "GELU not supported in debug mode" + if single_output: + out_init = out[0] + start_idx = 0 + out = [None] * num_gemms + for i in range(num_gemms): + size = m_splits[i] + out[i] = out_init[start_idx : start_idx + size] + start_idx += size + for i in range(num_gemms): + _, bias_or_grad, _, _ = general_gemm( + A[i], + B[i], + quantization_params=quantization_params[i], + out_dtype=out[0].dtype, + layout=layout, + accumulate=accumulate, + out=out[i], + bias=bias[i] if use_bias else None, + use_split_accumulator=use_split_accumulator, + grad=grad, + ) + if grad and use_bias: + grad_bias[i] = bias_or_grad + if single_output: + out = out_init + + return out, grad_bias if grad else bias, None + if gelu: gelu_input = [ torch.empty_like(o, dtype=bias_dtype, memory_format=torch.contiguous_format) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 6d1d8c3540..40589a82d5 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1165,18 +1165,7 @@ def grad_output_preprocess( # bgrad only if wgrad is in FP8, otherwise it is fused with wgrad and we return None if ctx.debug: grad_output_ = quantizer(grad_output) - if ( - isinstance( - grad_output_.get_tensor(True), - ( - QuantizedTensor, - Float8TensorStorage, - MXFP8TensorStorage, - Float8BlockwiseQTensorStorage, - ), - ) - and ctx.use_bias - ): + if ctx.use_bias: grad_bias = grad_output.view(-1, grad_output.shape[-1]).sum(dim=0) else: grad_bias = None @@ -1540,6 +1529,12 @@ def is_debug_iter(self) -> bool: # we use the debug value from the first invocation in the iteration. debug = self.debug_enabled_in_this_iteration + self.debug_last_iteration = TEDebugState.get_iteration() + + if self.wgrad_store is not None: + if debug and self.wgrad_store.delay_wgrad_compute(): + raise RuntimeError("Delayed wgrad compute is not supported in debug mode.") + return debug def no_debug_features_active(self, quantizers): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index b3a96df399..da3ead631b 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -4,6 +4,7 @@ """GroupedLinear API""" from typing import Union, Optional, Callable, Tuple, List +from itertools import chain import warnings import functools @@ -49,6 +50,8 @@ prepare_for_saving, restore_from_saved, ) +from ...debug.pytorch.debug_quantization import DebugQuantizer +from ...debug.pytorch.debug_state import TEDebugState __all__ = ["GroupedLinear"] @@ -58,6 +61,7 @@ class _GroupedLinear(torch.autograd.Function): Calls custom cuda extensions. """ + # pylint: disable=keyword-arg-before-vararg @staticmethod def forward( ctx, @@ -79,6 +83,8 @@ def forward( input_quantizers, weight_quantizers, output_quantizers, + grad_input_quantizers, + grad_weight_quantizers, grad_output_quantizers, fuse_wgrad_accumulation, cpu_offloading, @@ -88,6 +94,7 @@ def forward( module, skip_fp8_weight_update, save_original_input, + debug, ) = non_tensor_args num_gemms = len(m_splits) @@ -135,8 +142,12 @@ def forward( ) inp_view = inp.reshape(-1, in_features) inputmats: list - if fp8: + if fp8 and not debug: inputmats = tex.split_quantize(inp_view, m_splits, input_quantizers) + elif debug: + inputmats = DebugQuantizer.multi_tensor_quantize( + inp_view, input_quantizers, m_splits, activation_dtype + ) else: inputmats = torch.split(cast_if_needed(inp_view, activation_dtype), m_splits) @@ -145,7 +156,7 @@ def forward( # Initialize weights weights_fp8: list - if fp8: + if fp8 or debug: # FP8 cast to workspace buffer weights_fp8 = [] update_workspace = is_first_microbatch is None or is_first_microbatch @@ -156,6 +167,7 @@ def forward( cache_name=(None if is_first_microbatch is None else f"weight{i}"), update_workspace=update_workspace, skip_update_flag=skip_fp8_weight_update, + workspace_dtype=activation_dtype, ) weights_fp8.append(weight_fp8) @@ -167,7 +179,6 @@ def forward( if fp8 and activation_dtype == torch.float32: bias_dtype = torch.bfloat16 # FP8 GEMM only supports BF16/FP16 bias biases = [cast_if_needed(bias, bias_dtype) for bias in biases] if use_bias else biases - # Initialize output tensor out = torch.empty( [sum(m_splits), weights_fp8[0].size(0)], @@ -183,10 +194,11 @@ def forward( use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator # Perform GEMM - _ = general_grouped_gemm( + general_grouped_gemm( weights_fp8, inputmats, [out], + output_quantizers, activation_dtype, single_output=True, m_splits=m_splits, @@ -244,6 +256,10 @@ def forward( ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects + 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 = weights[0].requires_grad if fuse_wgrad_accumulation and ctx.weights_requires_grad: # This check is needed to ensure that main_grad is not created @@ -259,7 +275,7 @@ def forward( else: ctx.main_grad_funcs = [lambda: None for i in range(num_gemms)] ctx.device = device - ctx.grad_output_quantizers = grad_output_quantizers + ctx.output_quantizers = output_quantizers ctx.m_splits = m_splits ctx.num_gemms = num_gemms ctx.activation_dtype = activation_dtype @@ -279,6 +295,7 @@ def forward( or FP8GlobalStateManager.is_first_fp8_module() ) ctx.wgrad_store = wgrad_store + ctx.debug = debug ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers @@ -311,7 +328,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) grad_output = [None] * ctx.num_gemms grad_biases = [None] * ctx.num_gemms - if ctx.fp8: + if ctx.fp8 and not ctx.debug: if ctx.use_bias: grad_output_mats = torch.split(grad_output_view, ctx.m_splits) recipe = ctx.fp8_recipe @@ -338,6 +355,13 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ctx.m_splits, ctx.grad_output_quantizers, ) + elif ctx.debug: + grad_output_mats = torch.split(grad_output_view, ctx.m_splits) + for i in range(ctx.num_gemms): + grad_biases[i] = grad_output_mats[i].sum(dim=0) + grad_output = DebugQuantizer.multi_tensor_quantize( + grad_output_view, ctx.grad_output_quantizers, ctx.m_splits, ctx.activation_dtype + ) else: # Only split grad output. Grad bias is fused with # wgrad GEMM. @@ -355,7 +379,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if ctx.requires_dgrad: dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: + if ctx.fp8 or ctx.debug: recipe = ctx.fp8_recipe if hasattr(recipe, "fp8_gemm_dgrad"): dgrad_gemm_use_split_accumulator = ( @@ -375,6 +399,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], weights, grad_output, [dgrad], + ctx.grad_input_quantizers, ctx.activation_dtype, single_output=True, layout="NN", @@ -412,15 +437,19 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], else: input_quantizer.set_usage(rowwise=False, columnwise=True) inputmats: list - if ctx.fp8: + if ctx.fp8 and not ctx.debug: inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) + elif ctx.debug: + inputmats = DebugQuantizer.multi_tensor_quantize( + inp_view, ctx.input_quantizers, ctx.m_splits, ctx.activation_dtype + ) else: inputmats = torch.split( cast_if_needed(inp_view, ctx.activation_dtype), ctx.m_splits ) - grouped_gemm_wgrad = functools.partial( general_grouped_gemm, + quantization_params=ctx.grad_weight_quantizers, out_dtype=ctx.activation_dtype, layout="NT", grad=True, @@ -576,6 +605,7 @@ def __init__( ub_name: Optional[str] = None, delay_wgrad_compute: bool = False, save_original_input: bool = False, + name: Optional[str] = None, ) -> None: super().__init__() @@ -596,6 +626,7 @@ def __init__( ), "GroupedLinear doesn't support Userbuffer overlap." self.get_rng_state_tracker = get_rng_state_tracker self.rng_tracker_name = rng_tracker_name + self.name = name self.wgrad_store = WeightGradStore(delay_wgrad_compute) @@ -745,6 +776,8 @@ def forward( first microbatch (since it is the first gradient being produced) """ + debug = self.is_debug_iter() + assert not isinstance( inp, QuantizedTensorStorage ), "GroupedLinear doesn't support input tensor in FP8." @@ -756,31 +789,24 @@ def forward( weight_tensors = self._get_weight_tensors() bias_tensors = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] - weight_quantizers = self._get_weight_quantizers() - input_quantizers, output_quantizers = ( - [None] * self.num_gemms, - [None] * self.num_gemms, - ) - grad_output_quantizers, _ = [None] * self.num_gemms, [None] * self.num_gemms - if self.fp8: - input_quantizers = [ - self.quantizers["scaling_fwd"][ - self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"] - ] - for i in range(self.num_gemms) - ] - # TODO: use internal after #1638 is merged. # pylint: disable=fixme - for i in range(self.num_gemms): - input_quantizers[i].internal = False - if is_grad_enabled: - grad_output_quantizers = [ - self.quantizers["scaling_bwd"][ - self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"] - ] - for i in range(self.num_gemms) - ] - for i in range(self.num_gemms): - grad_output_quantizers[i].internal = True + quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() + + if debug: + if self.no_debug_features_active(list(chain(*quantizers))): + debug = False + quantizers = self._get_quantizers() + + if isinstance(weight_tensors, QuantizedTensorStorage): + raise RuntimeError("FP8 weights are not supported in debug mode.") + + ( + input_quantizers, + weight_quantizers, + output_quantizers, + grad_input_quantizers, + grad_weight_quantizers, + grad_output_quantizers, + ) = quantizers if is_grad_enabled: linear_fn = _GroupedLinear.apply @@ -799,6 +825,8 @@ def forward( input_quantizers, weight_quantizers, output_quantizers, + grad_input_quantizers, + grad_weight_quantizers, grad_output_quantizers, self.fuse_wgrad_accumulation, is_cpu_offload_enabled(), @@ -808,6 +836,7 @@ def forward( self, None, # skip_fp8_weight_update self.save_original_input, + debug, ) out = linear_fn(*autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors) @@ -898,3 +927,55 @@ def _get_weight_quantizers(self) -> List[Quantizer]: for i in range(self.num_gemms): weight_quantizers[i].internal = True return weight_quantizers + + def _get_quantizers(self): + weight_quantizers = self._get_weight_quantizers() + input_quantizers, output_quantizers = ( + [None] * self.num_gemms, + [None] * self.num_gemms, + ) + grad_input_quantizers, grad_weight_quantizers, grad_output_quantizers = ( + [None] * self.num_gemms, + [None] * self.num_gemms, + [None] * self.num_gemms, + ) + if self.fp8: + input_quantizers = [ + self.quantizers["scaling_fwd"][ + self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["fwd"] + ] + for i in range(self.num_gemms) + ] + # TODO: use internal after #1638 is merged. # pylint: disable=fixme + for i in range(self.num_gemms): + input_quantizers[i].internal = False + if torch.is_grad_enabled(): + grad_output_quantizers = [ + self.quantizers["scaling_bwd"][ + self._offsets["input"] + i * self._num_fp8_tensors_per_gemm["bwd"] + ] + for i in range(self.num_gemms) + ] + for i in range(self.num_gemms): + grad_output_quantizers[i].internal = True + return ( + input_quantizers, + weight_quantizers, + output_quantizers, + grad_input_quantizers, + grad_weight_quantizers, + grad_output_quantizers, + ) + + def _get_debug_quantizers(self): + original_quantizers = self._get_quantizers() + assert TEDebugState.debug_enabled + + names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"] + return tuple( + [ + DebugQuantizer(self.name + f".gemm_{q_id}", name, q, self.tp_group) + for q_id, q in enumerate(qs) + ] + for name, qs in zip(names, original_quantizers) + ) From 9ca89e97e9d9b7d84ccd91b7ab74dc20a9824518 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Tue, 25 Nov 2025 17:13:12 -0800 Subject: [PATCH 101/521] [PyTorch] Avoid initializing recipe state in fusible op base class constructor (#2421) Do not initialize recipe state in base op class Op attrs may not be set. Move recipe state initialization to linear op constructor. Signed-off-by: Tim Moon --- tests/pytorch/test_fusible_ops.py | 10 +++++----- transformer_engine/pytorch/ops/basic/basic_linear.py | 4 +++- transformer_engine/pytorch/ops/op.py | 3 --- 3 files changed, 8 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index d2770347aa..735cc9b953 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -901,15 +901,15 @@ def _test_basic_linear( dtype=dtype, accumulate_into_main_grad=accumulate_into_main_grad, ) + forward = te_ops.Sequential( + te_ops.Quantize(forward=quantized_input, backward=quantized_grad_input), + op, + te_ops.Quantize(forward=quantized_output, backward=quantized_grad_output), + ) with torch.no_grad(): op.weight.copy_(w_test) del w_test op.weight.main_grad = torch.full_like(op.weight, 0.5, dtype=torch.float32) - forward = te_ops.Sequential( - te_ops.Quantize(forward=quantized_input, backward=quantized_grad_input), - op, - te_ops.Quantize(forward=quantized_output, backward=quantized_grad_output), - ) with te.autocast(enabled=quantized_compute, recipe=recipe): y_test = forward(x_test) y_test.backward(dy_test) diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 749ab7a650..c629d0158d 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -137,8 +137,10 @@ def __init__( out_features=out_features, ) - # Whether weight tensor is natively quantized + # Initialize recipe state if needed for natively quantized weight self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters() + if self._with_quantized_weight: + self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe()) # Initialize parameters if needed weight = torch.empty( diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 639817ada7..6ae49dcd4e 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -188,9 +188,6 @@ def __init__(self) -> None: # Objects for quantization self._fp8_metas: Optional[dict[str, dict[str, Any]]] = None self._quantizers: Optional[dict[str, list[Quantizer]]] = None - with_fp8_parameters = FP8GlobalStateManager.with_fp8_parameters() - recipe = FP8GlobalStateManager.get_fp8_recipe() if with_fp8_parameters else None - self.reset_recipe_state(recipe=recipe) @property def is_fused_op(self) -> bool: From ca468ebe64c8af3fb49ef11cb47b09d07ff4fd2c Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Wed, 26 Nov 2025 16:59:04 +0100 Subject: [PATCH 102/521] Extend docs with quantizers/quantized_tensors/custom_recipe (#2428) * Extend docs with quantizers/quantized_tensors/custom_recipe Signed-off-by: Evgeny * Bring structure, reduce redundant members Signed-off-by: Evgeny --------- Signed-off-by: Evgeny --- docs/api/common.rst | 2 ++ docs/api/pytorch.rst | 48 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 50 insertions(+) diff --git a/docs/api/common.rst b/docs/api/common.rst index 3edd7cae21..728dcd6ed0 100644 --- a/docs/api/common.rst +++ b/docs/api/common.rst @@ -17,3 +17,5 @@ Common API .. autoapiclass:: transformer_engine.common.recipe.Float8CurrentScaling(fp8_format=Format.HYBRID) .. autoapiclass:: transformer_engine.common.recipe.Float8BlockScaling(fp8_format=Format.E4M3) + +.. autoapiclass:: transformer_engine.common.recipe.CustomRecipe(qfactory, fp8_dpa=False, fp8_mha=False) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index c456f1a6ad..391e52de95 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -85,3 +85,51 @@ pyTorch .. autoapiclass:: transformer_engine.pytorch.UserBufferQuantizationMode :members: FP8, NONE + +Quantized tensors +----------------- + +.. autoapiclass:: transformer_engine.pytorch.QuantizedTensorStorage + :members: update_usage, prepare_for_saving, restore_from_saved + +.. autoapiclass:: transformer_engine.pytorch.QuantizedTensor(shape, dtype, *, requires_grad=False, device=None) + :members: dequantize, quantize_ + +.. autoapiclass:: transformer_engine.pytorch.Float8TensorStorage(data, fp8_scale_inv, fp8_dtype, data_transpose=None, quantizer=None) + +.. autoapiclass:: transformer_engine.pytorch.MXFP8TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer) + +.. autoapiclass:: transformer_engine.pytorch.Float8BlockwiseQTensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer, is_2D_scaled, data_format) + +.. autoapiclass:: transformer_engine.pytorch.NVFP4TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) + +.. autoapiclass:: transformer_engine.pytorch.Float8Tensor(shape, dtype, data, fp8_scale_inv, fp8_dtype, requires_grad=False, data_transpose=None, quantizer=None) + +.. autoapiclass:: transformer_engine.pytorch.MXFP8Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer) + +.. autoapiclass:: transformer_engine.pytorch.Float8BlockwiseQTensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer, is_2D_scaled, data_format) + +.. autoapiclass:: transformer_engine.pytorch.NVFP4Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer) + +Quantizers +---------- + +.. autoapiclass:: transformer_engine.pytorch.Quantizer(rowwise, columnwise) + :members: update_quantized, quantize + +.. autoapiclass:: transformer_engine.pytorch.Float8Quantizer(scale, amax, fp8_dtype, *, rowwise=True, columnwise=True) + +.. autoapiclass:: transformer_engine.pytorch.Float8CurrentScalingQuantizer(fp8_dtype, device, *, rowwise=True, columnwise=True, **kwargs) + +.. autoapiclass:: transformer_engine.pytorch.MXFP8Quantizer(fp8_dtype, *, rowwise=True, columnwise=True) + +.. autoapiclass:: transformer_engine.pytorch.Float8BlockQuantizer(fp8_dtype, *, rowwise, columnwise, **kwargs) + +.. autoapiclass:: transformer_engine.pytorch.NVFP4Quantizer(fp4_dtype, *, rowwise=True, columnwise=True, **kwargs) + +Tensor saving and restoring functions +------------------------------------- + +.. autoapifunction:: transformer_engine.pytorch.prepare_for_saving + +.. autoapifunction:: transformer_engine.pytorch.restore_from_saved From df39a7c2eba40eacd94da75302ea0c1ffb72d33b 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, 26 Nov 2025 18:01:55 +0100 Subject: [PATCH 103/521] Docs fix (#2301) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * init Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * lines lenght Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * subtitle --- fix in many files: Signed-off-by: Pawel Gadzinski * cross entropy _input -> input rename Signed-off-by: Pawel Gadzinski * cross entropy _input -> input rename Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * a lot of small fixes Signed-off-by: Pawel Gadzinski * torch_version() change Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add missing module and fix warnings Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * removed training whitespace: Signed-off-by: Pawel Gadzinski * Update docs/api/pytorch.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * Fix import Signed-off-by: Kirthi Shankar Sivamani * Fix more imports Signed-off-by: Kirthi Shankar Sivamani * Fix NumPy docstring parameter spacing and indentation - Standardize parameter documentation to use 'param : type' format (space before and after colon) per NumPy style guide - Fix inconsistent indentation in cpu_offload.py docstring - Modified 51 Python files across transformer_engine/pytorch Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani 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: Kirthi Shankar Sivamani --- .github/workflows/docs.yml | 4 +- docs/api/jax.rst | 6 +- docs/api/pytorch.rst | 34 +- docs/conf.py | 28 +- docs/debug.rst | 3 +- docs/debug/1_getting_started.rst | 15 +- docs/debug/2_config_file_structure.rst | 15 +- docs/debug/3_api_debug_setup.rst | 7 +- docs/debug/3_api_features.rst | 2 +- docs/debug/4_distributed.rst | 13 +- docs/debug/api.rst | 3 +- docs/examples/advanced_optimizations.ipynb | 4 +- docs/examples/attention/attention.ipynb | 12 +- docs/examples/quickstart_jax.ipynb | 2 +- .../tutorial_generation_gemma_with_te.ipynb | 2 +- docs/index.rst | 2 +- docs/installation.rst | 4 +- .../common/fused_attn/kv_cache.cu | 2 +- .../include/transformer_engine/fused_attn.h | 2 +- transformer_engine/common/recipe/__init__.py | 38 +- .../jax/cpp_extensions/activation.py | 2 +- transformer_engine/jax/cpp_extensions/gemm.py | 2 +- transformer_engine/jax/cpp_extensions/misc.py | 2 +- .../jax/cpp_extensions/normalization.py | 2 +- .../jax/cpp_extensions/quantization.py | 2 +- transformer_engine/jax/dense.py | 2 +- transformer_engine/jax/flax/module.py | 94 ++--- transformer_engine/jax/flax/transformer.py | 291 ++++++++------ transformer_engine/pytorch/__init__.py | 10 +- .../dot_product_attention.py | 328 ++++++++-------- .../attention/dot_product_attention/utils.py | 164 ++++---- .../pytorch/attention/inference.py | 34 +- .../pytorch/attention/multi_head_attention.py | 307 ++++++++------- transformer_engine/pytorch/attention/rope.py | 46 +-- .../pytorch/cpp_extensions/fused_attn.py | 132 +++---- transformer_engine/pytorch/cpu_offload.py | 58 +-- transformer_engine/pytorch/cpu_offload_v1.py | 12 +- transformer_engine/pytorch/cross_entropy.py | 84 +++- transformer_engine/pytorch/distributed.py | 40 +- transformer_engine/pytorch/export.py | 2 +- transformer_engine/pytorch/graph.py | 22 +- transformer_engine/pytorch/jit.py | 2 +- transformer_engine/pytorch/module/base.py | 65 ++-- .../pytorch/module/grouped_linear.py | 44 ++- .../pytorch/module/layernorm.py | 21 +- .../pytorch/module/layernorm_linear.py | 60 +-- .../pytorch/module/layernorm_mlp.py | 90 ++--- transformer_engine/pytorch/module/linear.py | 60 +-- transformer_engine/pytorch/module/rmsnorm.py | 23 +- transformer_engine/pytorch/ops/_common.py | 2 +- .../pytorch/ops/basic/activation.py | 8 +- .../pytorch/ops/basic/all_gather.py | 2 +- .../pytorch/ops/basic/all_reduce.py | 2 +- .../pytorch/ops/basic/basic_linear.py | 18 +- transformer_engine/pytorch/ops/basic/bias.py | 10 +- .../pytorch/ops/basic/l2normalization.py | 6 +- .../pytorch/ops/basic/layer_norm.py | 8 +- .../pytorch/ops/basic/quantize.py | 4 +- .../pytorch/ops/basic/reduce_scatter.py | 2 +- .../pytorch/ops/basic/reshape.py | 2 +- .../pytorch/ops/basic/rmsnorm.py | 8 +- .../ops/fused/backward_activation_bias.py | 6 +- .../pytorch/ops/fused/backward_add_rmsnorm.py | 4 +- .../pytorch/ops/fused/backward_linear_add.py | 4 +- .../ops/fused/backward_linear_scale.py | 4 +- .../fused/forward_linear_bias_activation.py | 4 +- .../ops/fused/forward_linear_bias_add.py | 4 +- .../ops/fused/forward_linear_scale_add.py | 4 +- .../ops/fused/userbuffers_backward_linear.py | 4 +- .../ops/fused/userbuffers_forward_linear.py | 4 +- transformer_engine/pytorch/ops/fuser.py | 2 +- transformer_engine/pytorch/ops/linear.py | 20 +- transformer_engine/pytorch/ops/op.py | 2 +- transformer_engine/pytorch/permutation.py | 44 +-- transformer_engine/pytorch/quantization.py | 22 +- .../pytorch/quantized_tensor.py | 8 +- transformer_engine/pytorch/router.py | 44 +-- .../pytorch/tensor/float8_blockwise_tensor.py | 12 +- .../pytorch/tensor/float8_tensor.py | 16 +- .../pytorch/tensor/mxfp8_tensor.py | 8 +- .../pytorch/tensor/nvfp4_tensor.py | 18 +- transformer_engine/pytorch/tensor/utils.py | 2 +- transformer_engine/pytorch/torch_version.py | 15 + transformer_engine/pytorch/transformer.py | 366 +++++++++--------- .../pytorch/triton/permutation.py | 76 ++-- transformer_engine/pytorch/utils.py | 8 +- 86 files changed, 1606 insertions(+), 1366 deletions(-) create mode 100644 transformer_engine/pytorch/torch_version.py diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 3c4229a888..5beeeb8879 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -22,10 +22,10 @@ jobs: sudo apt-get install -y pandoc graphviz doxygen export GIT_SHA=$(git show-ref --hash HEAD) - name: 'Build docs' - run: | + run: | # SPHINXOPTS="-W" errors out on warnings doxygen docs/Doxyfile cd docs - make html + make html SPHINXOPTS="-W" - name: 'Upload docs' uses: actions/upload-artifact@v4 with: diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 789b27e59c..99782f99c7 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -4,7 +4,7 @@ See LICENSE for license information. Jax -======= +=== Pre-defined Variable of Logical Axes ------------------------------------ @@ -20,11 +20,11 @@ Variables are available in `transformer_engine.jax.sharding`. Checkpointing ------------------------------------- +------------- When using checkpointing with Transformer Engine JAX, please be aware of the checkpointing policy being applied to your model. Any JAX checkpointing policy using `dot`, such as `jax.checkpoint_policies.dots_with_no_batch_dims`, may not work with GEMMs provided by Transformer Engine as they do not always use the `jax.lax.dot_general` primitive. Instead, you can use `transformer_engine.jax.checkpoint_policies.dots_and_te_gemms_with_no_batch_dims` or similar policies that are designed to work with Transformer Engine's GEMMs and `jax.lax.dot_general` GEMMs. You may also use any JAX policies that do not filter by primitive, such as `jax.checkpoint_policies.save_only_these_names` or `jax.checkpoint_policies.everything_saveable`. Modules ------------------------------------- +------- .. autoapiclass:: transformer_engine.jax.flax.TransformerLayerType .. autoapiclass:: transformer_engine.jax.MeshResource() diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 391e52de95..18abe0f2c2 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -3,7 +3,7 @@ See LICENSE for license information. -pyTorch +PyTorch ======= .. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs) @@ -37,9 +37,6 @@ pyTorch .. autoapiclass:: transformer_engine.pytorch.CudaRNGStatesTracker() :members: reset, get_states, set_states, add, fork -.. autoapifunction:: transformer_engine.pytorch.fp8_autocast - -.. autoapifunction:: transformer_engine.pytorch.fp8_model_init .. autoapifunction:: transformer_engine.pytorch.autocast @@ -47,6 +44,16 @@ pyTorch .. autoapifunction:: transformer_engine.pytorch.checkpoint + +.. autoapifunction:: transformer_engine.pytorch.make_graphed_callables + +.. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context + +.. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy + +Recipe availability +------------------- + .. autoapifunction:: transformer_engine.pytorch.is_fp8_available .. autoapifunction:: transformer_engine.pytorch.is_mxfp8_available @@ -63,9 +70,8 @@ pyTorch .. autoapifunction:: transformer_engine.pytorch.get_default_recipe -.. autoapifunction:: transformer_engine.pytorch.make_graphed_callables - -.. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context +Mixture of Experts (MoE) functions +---------------------------------- .. autoapifunction:: transformer_engine.pytorch.moe_permute @@ -75,10 +81,12 @@ pyTorch .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index -.. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy - .. autoapifunction:: transformer_engine.pytorch.moe_sort_chunks_by_index_with_probs + +Communication-computation overlap +--------------------------------- + .. autoapifunction:: transformer_engine.pytorch.initialize_ub .. autoapifunction:: transformer_engine.pytorch.destroy_ub @@ -86,6 +94,7 @@ pyTorch .. autoapiclass:: transformer_engine.pytorch.UserBufferQuantizationMode :members: FP8, NONE + Quantized tensors ----------------- @@ -133,3 +142,10 @@ Tensor saving and restoring functions .. autoapifunction:: transformer_engine.pytorch.prepare_for_saving .. autoapifunction:: transformer_engine.pytorch.restore_from_saved + +Deprecated functions +-------------------- + +.. autoapifunction:: transformer_engine.pytorch.fp8_autocast + +.. autoapifunction:: transformer_engine.pytorch.fp8_model_init diff --git a/docs/conf.py b/docs/conf.py index 7f5966d717..479c1f8948 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -61,7 +61,11 @@ ] templates_path = ["_templates"] -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] +exclude_patterns = [ + "_build", + "Thumbs.db", + "sphinx_rtd_theme", +] source_suffix = ".rst" @@ -94,6 +98,7 @@ ("Values", "params_style"), ("Graphing parameters", "params_style"), ("FP8-related parameters", "params_style"), + ("Quantization parameters", "params_style"), ] breathe_projects = {"TransformerEngine": root_path / "docs" / "doxygen" / "xml"} @@ -101,4 +106,23 @@ autoapi_generate_api_docs = False autoapi_dirs = [root_path / "transformer_engine"] -autoapi_ignore = ["*/_[!_]*"] +autoapi_ignore = ["*test*"] + + +# There are 2 warnings about the same namespace (transformer_engine) in two different c++ api +# docs pages. This seems to be the only way to suppress these warnings. +def setup(app): + """Custom Sphinx setup to filter warnings.""" + import logging + + # Filter out duplicate C++ declaration warnings + class DuplicateDeclarationFilter(logging.Filter): + def filter(self, record): + message = record.getMessage() + if "Duplicate C++ declaration" in message and "transformer_engine" in message: + return False + return True + + # Apply filter to Sphinx logger + logger = logging.getLogger("sphinx") + logger.addFilter(DuplicateDeclarationFilter()) diff --git a/docs/debug.rst b/docs/debug.rst index d33568ea3b..527f30ed02 100644 --- a/docs/debug.rst +++ b/docs/debug.rst @@ -2,8 +2,9 @@ Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. + Precision debug tools -============================================== +===================== .. toctree:: :caption: Precision debug tools diff --git a/docs/debug/1_getting_started.rst b/docs/debug/1_getting_started.rst index 906c625567..a5cdc1a6b1 100644 --- a/docs/debug/1_getting_started.rst +++ b/docs/debug/1_getting_started.rst @@ -4,7 +4,7 @@ See LICENSE for license information. Getting started -============== +=============== .. note:: @@ -38,7 +38,7 @@ To start debugging, one needs to create a configuration YAML file. This file lis one - ``UserProvidedPrecision`` - is a custom feature implemented by the user. Nvidia-DL-Framework-Inspect inserts features into the layers according to the config. Example training script ----------------------- +----------------------- Let's look at a simple example of training a Transformer layer using Transformer Engine with FP8 precision. This example demonstrates how to set up the layer, define an optimizer, and perform a few training iterations using synthetic data. @@ -81,7 +81,7 @@ We will demonstrate two debug features on the code above: 2. Logging statistics for other GEMM operations, such as gradient statistics for data gradient GEMM within the LayerNormLinear sub-layer of the TransformerLayer. Config file ----------- +----------- We need to prepare the configuration YAML file, as below @@ -114,7 +114,8 @@ We need to prepare the configuration YAML file, as below Further explanation on how to create config files is in the :doc:`next part of the documentation <2_config_file_structure>`. Adjusting Python file --------------------- +--------------------- + .. code-block:: python @@ -145,7 +146,8 @@ In the modified code above, the following changes were made: 3. Added ``debug_api.step()`` after each of the forward-backward pass. Inspecting the logs ------------------- +------------------- + Let's look at the files with the logs. Two files will be created: @@ -213,7 +215,8 @@ The second log file (``nvdlfw_inspect_statistics_logs/nvdlfw_inspect_globalrank- INFO - transformer_layer.self_attention.layernorm_qkv_activation_l1_norm iteration=000004 value=130776.7969 Logging using TensorBoard ------------------------- +------------------------- + Precision debug tools support logging using `TensorBoard `_. To enable it, one needs to pass the argument ``tb_writer`` to the ``debug_api.initialize()``. Let's modify ``train.py`` file. diff --git a/docs/debug/2_config_file_structure.rst b/docs/debug/2_config_file_structure.rst index fe82df638f..2132cd32c7 100644 --- a/docs/debug/2_config_file_structure.rst +++ b/docs/debug/2_config_file_structure.rst @@ -4,13 +4,14 @@ See LICENSE for license information. Config File Structure -==================== +===================== To enable debug features, create a configuration YAML file to specify the desired behavior, such as determining which GEMMs (General Matrix Multiply operations) should run in higher precision rather than FP8 and defining which statistics to log. Below, we outline how to structure the configuration YAML file. General Format -------------- +-------------- + A config file can have one or more sections, each containing settings for specific layers and features: @@ -55,7 +56,8 @@ Sections may have any name and must contain: 3. Additional fields describing features for those layers. Layer Specification ------------------- +------------------- + Debug layers can be identified by a ``name`` parameter: @@ -89,7 +91,8 @@ Examples: (...) Names in Transformer Layers --------------------------- +--------------------------- + There are three ways to assign a name to a layer in the Transformer Engine: @@ -156,7 +159,7 @@ Below is an example ``TransformerLayer`` with four linear layers that can be inf Structured Configuration for GEMMs and Tensors ---------------------------------------------- +---------------------------------------------- Sometimes a feature is parameterized by a list of tensors or by a list of GEMMs. There are multiple ways of describing this parameterization. @@ -218,7 +221,7 @@ We can use both structs for tensors and GEMMs. The tensors_struct should be nest gemm_feature_param1: value Enabling or Disabling Sections and Features ------------------------------------------- +------------------------------------------- Debug features can be enabled or disabled with the ``enabled`` keyword: diff --git a/docs/debug/3_api_debug_setup.rst b/docs/debug/3_api_debug_setup.rst index bda8f096d6..176bc13d32 100644 --- a/docs/debug/3_api_debug_setup.rst +++ b/docs/debug/3_api_debug_setup.rst @@ -11,7 +11,8 @@ Please refer to the Nvidia-DL-Framework-Inspect `documentation `_ for more details. @@ -61,7 +62,7 @@ If the tensor reduction group is not specified, then statistics are reduced acro # activation/gradient tensor statistics are reduced along pipeline_parallel_group set_weight_tensor_tp_group_reduce() ---------------------------------- +----------------------------------- By default, weight tensor statistics are reduced within the tensor parallel group. This function allows you to disable that behavior; for more details, see `reduction group section <./4_distributed.rst#reduction-groups>`_. diff --git a/docs/debug/3_api_features.rst b/docs/debug/3_api_features.rst index b31c437b2d..8cdbde8edd 100644 --- a/docs/debug/3_api_features.rst +++ b/docs/debug/3_api_features.rst @@ -4,7 +4,7 @@ See LICENSE for license information. Debug features -========== +============== .. autoapiclass:: transformer_engine.debug.features.log_tensor_stats.LogTensorStats .. autoapiclass:: transformer_engine.debug.features.log_fp8_tensor_stats.LogFp8TensorStats diff --git a/docs/debug/4_distributed.rst b/docs/debug/4_distributed.rst index 6f69f2712c..764fee6541 100644 --- a/docs/debug/4_distributed.rst +++ b/docs/debug/4_distributed.rst @@ -4,7 +4,7 @@ See LICENSE for license information. Distributed training -=================== +==================== Nvidia-Pytorch-Inspect with Transformer Engine supports multi-GPU training. This guide describes how to run it and how the supported features work in the distributed setting. @@ -14,7 +14,8 @@ To use precision debug tools in multi-GPU training, one needs to: 2. If one wants to log stats, one may want to invoke ``debug_api.set_tensor_reduction_group`` with a proper reduction group. Behavior of the features ------------------------ +------------------------ + In a distributed setting, **DisableFP8GEMM** and **DisableFP8Layer** function similarly to the single-GPU case, with no notable differences. @@ -28,7 +29,8 @@ In a distributed setting, **DisableFP8GEMM** and **DisableFP8Layer** function si Logging-related features are more complex and will be discussed further in the next sections. Reduction groups --------------- +---------------- + In setups with tensor, data, or pipeline parallelism, some tensors are distributed across multiple GPUs, requiring a reduction operation to compute statistics for these tensors. @@ -65,7 +67,8 @@ Below, we illustrate configurations for a 4-node setup with tensor parallelism s Microbatching ------------ +------------- + Let's dive into how statistics collection works with microbatching. By microbatching, we mean invoking multiple ``forward()`` calls for each ``debug_api.step()``. The behavior is as follows: @@ -73,7 +76,7 @@ Let's dive into how statistics collection works with microbatching. By microbatc - For other tensors, the stats are accumulated. Logging to files and TensorBoard ------------------------------- +-------------------------------- In a single-node setup with ``default_logging_enabled=True``, all logs are saved by default to ``log_dir/nvdlfw_inspect_statistics_logs/nvdlfw_inspect_globalrank-0.log``. In multi-GPU training, each node writes its reduced statistics to its unique file, named ``log_dir/nvdlfw_inspect_statistics_logs/nvdlfw_inspect_globalrank-i.log`` for rank i. Because these logs contain reduced statistics, the logged values are identical for all nodes within a reduction group. diff --git a/docs/debug/api.rst b/docs/debug/api.rst index ac593d353a..6ccb32cc8b 100644 --- a/docs/debug/api.rst +++ b/docs/debug/api.rst @@ -2,8 +2,9 @@ Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. + API -============ +=== .. toctree:: :caption: Precision debug tools API diff --git a/docs/examples/advanced_optimizations.ipynb b/docs/examples/advanced_optimizations.ipynb index 5dc9cb92f9..7c08bb6586 100644 --- a/docs/examples/advanced_optimizations.ipynb +++ b/docs/examples/advanced_optimizations.ipynb @@ -100,7 +100,7 @@ "\n", "\n", "\n", - "A variety of parallelism strategies can be used to enable multi-GPU training of Transformer models, often based on different approaches to distribute their $\\text{sequence_length} \\times \\text{batch_size} \\times \\text{hidden_size}$ activation tensors. The most common approach is data parallelism, which distributes along the $\\text{batch_size}$ dimension. By storing duplicate copies of the model on each GPU, the forward and backward passes of the training step can be done independently, followed by a gradient synchronization. A more advanced strategy is tensor parallelism, a type of model parallelism that distributes along the $\\text{hidden_size}$ dimension. This allows us to scale past the limits of data parallelism (typically $\\text{hidden_size} > \\text{batch_size}$) and to reduce the per-GPU memory usage (since model parameters are also distributed), but it also incurs the overhead of communicating activation tensors between GPUs at every step. For a more detailed explanation, please see the [Megatron-LM paper](https://arxiv.org/pdf/1909.08053.pdf). Finally, sequence parallelism distributes along the $\\text{sequence_length}$ dimension. This can be used when tensor parallelism is enabled in order to parallelize operations that run outside the tensor-parallel region (e.g. layer norm). For more details, please see [this paper](https://arxiv.org/pdf/2205.05198.pdf).\n", + "A variety of parallelism strategies can be used to enable multi-GPU training of Transformer models, often based on different approaches to distribute their $\\text{sequence_length} \\cdot \\text{batch_size} \\cdot \\text{hidden_size}$ activation tensors. The most common approach is data parallelism, which distributes along the $\\text{batch_size}$ dimension. By storing duplicate copies of the model on each GPU, the forward and backward passes of the training step can be done independently, followed by a gradient synchronization. A more advanced strategy is tensor parallelism, a type of model parallelism that distributes along the $\\text{hidden_size}$ dimension. This allows us to scale past the limits of data parallelism (typically $\\text{hidden_size} > \\text{batch_size}$) and to reduce the per-GPU memory usage (since model parameters are also distributed), but it also incurs the overhead of communicating activation tensors between GPUs at every step. For a more detailed explanation, please see the [Megatron-LM paper](https://arxiv.org/pdf/1909.08053.pdf). Finally, sequence parallelism distributes along the $\\text{sequence_length}$ dimension. This can be used when tensor parallelism is enabled in order to parallelize operations that run outside the tensor-parallel region (e.g. layer norm). For more details, please see [this paper](https://arxiv.org/pdf/2205.05198.pdf).\n", "\n", "To show this in action, let's first initialize NCCL with a trivial process group:" ] @@ -131,7 +131,7 @@ "id": "1f2b80d0", "metadata": {}, "source": [ - "We only initialize with one GPU to keep this example simple. Please consult the documentation [torch.distributed](https://pytorch.org/docs/stable/distributed.html) for guidance on running with multiple GPUs. Note that we require that each distributed process corresponds to exactly one GPU, so we treat them interchangeably. In practice, there are multiple factors that can affect the optimal parallel layout: the system hardware, the network topology, usage of other parallelism schemes like pipeline parallelism. A rough rule-of-thumb is to interpret the GPUs as a 2D grid with dimensions of $\\text{num_nodes} \\times \\text{gpus_per_node}$. The rows are tensor-parallel groups and the columns are data-parallel groups.\n", + "We only initialize with one GPU to keep this example simple. Please consult the documentation [torch.distributed](https://pytorch.org/docs/stable/distributed.html) for guidance on running with multiple GPUs. Note that we require that each distributed process corresponds to exactly one GPU, so we treat them interchangeably. In practice, there are multiple factors that can affect the optimal parallel layout: the system hardware, the network topology, usage of other parallelism schemes like pipeline parallelism. A rough rule-of-thumb is to interpret the GPUs as a 2D grid with dimensions of $\\text{num_nodes} \\cdot \\text{gpus_per_node}$. The rows are tensor-parallel groups and the columns are data-parallel groups.\n", "\n", "Enabling data parallelism with Transformer Engine is similar to enabling data parallelism with standard PyTorch models: simply wrap the modules with [torch.nn.parallel.DistributedDataParallel](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html). Transformer Engine modules also have native support for tensor and sequence parallelism. If the user provides a process group for tensor parallelism, the modules will distribute the data and perform communication internally. If sequence parallelism is enabled, it will be applied for operations that are not amenable to tensor parallelism and it will use the tensor-parallel process group.\n", "\n", diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 61a6ad949f..4b2ed80497 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -174,7 +174,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "50852cb5", "metadata": {}, "outputs": [ @@ -266,7 +266,7 @@ }, { "cell_type": "code", - "execution_count": 24, + "execution_count": null, "id": "906b8cf1", "metadata": {}, "outputs": [ @@ -299,7 +299,7 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "id": "d3637094", "metadata": {}, "outputs": [ @@ -509,10 +509,10 @@ "\n", "* PyTorch: When both options are provided by the user, `cu_seqlens` is preferred as there is no extra conversion needed.\n", " - `cu_seqlens`: Users can provide cumulative sequence length tensors `cu_seqlens_q` and `cu_seqlens_kv` for `q` and `k`/`v` to the flash-attention or cuDNN attention backend. An example of `cu_seqlens` is `[0, 2, 6, 7]` for a batch of 3 `[aa000, bbbb0, c0000]`.\n", - " - `attention_mask`: Users can also provide `attention_mask` as an alternative, which will then be converted to `cu_seqlens`. For self-attention, `attention_mask` should be one single tensor in shape `[batch_size, 1, 1, seqlen_q]`, and for cross-attention, `attention_mask` should be a list of two tensors in shapes `[batch_size, 1, 1, seqlen_q]` and `[batch_size, 1, 1, seqlen_kv]`, respectively.\n", + " - `attention_mask`: Users can also provide `attention_mask` as an alternative, which will then be converted to `cu_seqlens`. For self-attention, `attention_mask` should be one single tensor of shape `[batch_size, 1, 1, seqlen_q]`, and for cross-attention, `attention_mask` should be a list of two tensors of shapes `[batch_size, 1, 1, seqlen_q]` and `[batch_size, 1, 1, seqlen_kv]`, respectively.\n", "\n", "\n", - "* JAX: Users should provide the `attention_mask` tensor in shape `[batch_size, 1, seqlen_q, seqlen_kv]`.\n", + "* JAX: Users should provide the `attention_mask` tensor of shape `[batch_size, 1, seqlen_q, seqlen_kv]`.\n", "\n", "**qkv_format=thd:** Transformer Engine extracts the max sequence length information from `q`, `k`, `v` if `max_seqlen_q` and `max_seqlen_kv` are not provided. This requires GPU-CPU copy and synchronization operations. For performance reasons, please set `max_seqlen_q` and `max_seqlen_kv` to their appropriate values for `thd` QKV format.\n", "\n", @@ -521,7 +521,7 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": null, "id": "a1f25a9b", "metadata": {}, "outputs": [ diff --git a/docs/examples/quickstart_jax.ipynb b/docs/examples/quickstart_jax.ipynb index 7146a95f4a..dc500afd1b 100644 --- a/docs/examples/quickstart_jax.ipynb +++ b/docs/examples/quickstart_jax.ipynb @@ -502,7 +502,7 @@ "\n", "\n", "\n", - "Enabling FP8 support is very simple in Transformer Engine. We just need to wrap the modules within an [autocast](.../api/jax.rst#transformer_engine.jax.fp8_autocast) context manager. See the [FP8 tutorial](fp8_primer.ipynb) for a detailed explanation of FP8 recipes and the supported options.\n", + "Enabling FP8 support is very simple in Transformer Engine. We just need to wrap the modules within an [autocast](../api/jax.rst#transformer_engine.jax.fp8_autocast) context manager. See the [FP8 tutorial](fp8_primer.ipynb) for a detailed explanation of FP8 recipes and the supported options.\n", "\n", "
\n", "\n", diff --git a/docs/examples/te_gemma/tutorial_generation_gemma_with_te.ipynb b/docs/examples/te_gemma/tutorial_generation_gemma_with_te.ipynb index c31e272b25..1ce60840b6 100755 --- a/docs/examples/te_gemma/tutorial_generation_gemma_with_te.ipynb +++ b/docs/examples/te_gemma/tutorial_generation_gemma_with_te.ipynb @@ -38,7 +38,7 @@ "\n", "For those seeking a deeper understanding of text generation mechanisms in Transformers, it is recommended to check out the [HuggingFace generation tutorial](https://huggingface.co/docs/transformers/llm_tutorial).\n", "\n", - "In a previous tutorial on [Llama](../te_llama/tutorial_accelerate_hf_llama_finetuning_with_te.ipynb), it was demonstrated how finetuning of an open-source Llama model can be accelerated using Transformer Engine's `TransformerLayer`. Building on that foundation, this tutorial showcases how to accelerate the token generation from the open-source Hugging Face Gemma 7B model.\n", + "In a previous tutorial on [Llama](../te_llama/tutorial_accelerate_hf_llama_with_te.ipynb), it was demonstrated how finetuning of an open-source Llama model can be accelerated using Transformer Engine's `TransformerLayer`. Building on that foundation, this tutorial showcases how to accelerate the token generation from the open-source Hugging Face Gemma 7B model.\n", "\n", "This tutorial introduces several features of the Transformer Engine library that contribute towards this goal. A brief explanation is as follows:\n", "\n", diff --git a/docs/index.rst b/docs/index.rst index 277259edf0..4fd55d241c 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -4,7 +4,7 @@ See LICENSE for license information. Transformer Engine documentation -============================================== +================================= .. ifconfig:: "dev" in release diff --git a/docs/installation.rst b/docs/installation.rst index a8bb74fd1a..24563c456e 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -28,7 +28,7 @@ on `NVIDIA GPU Cloud `_. pip - from PyPI ------------------------ +--------------- Transformer Engine can be directly installed from `our PyPI `_, e.g. @@ -47,7 +47,7 @@ The core package from Transformer Engine (without any framework extensions) can By default, this will install the core library compiled for CUDA 12. The cuda major version can be specified by modified the extra dependency to `core_cu12` or `core_cu13`. pip - from GitHub ------------------------ +----------------- Additional Prerequisites ^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/transformer_engine/common/fused_attn/kv_cache.cu b/transformer_engine/common/fused_attn/kv_cache.cu index 67119c323b..3b78cab239 100644 --- a/transformer_engine/common/fused_attn/kv_cache.cu +++ b/transformer_engine/common/fused_attn/kv_cache.cu @@ -278,7 +278,7 @@ void convert_bshd_to_thd(Tensor tensor, Tensor cu_seqlens, Tensor new_tensor, in /*************************************************************************************************** * KV Cache: Copy new KV tokens to the KV cache * 1. new_k and new_v are in qkv_format; k_cache and v_cache are in 'bshd' format - * 2. cu_new_lens and cu_cached_lens are in shape [b + 1]; cu_cached_lens include the added lens + * 2. cu_new_lens and cu_cached_lens are of shape [b + 1]; cu_cached_lens include the added lens * in current step * 3. Non-paged KV cache is a special case of paged KV cache, with page_table = [b, 1] and * max_pages_per_seq = 1. We use the same underlying kernel for both non-paged and paged. diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 298dc63900..6622019280 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -131,7 +131,7 @@ enum NVTE_Mask_Type { * NVTE_VANILLA_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), * NVTE_OFF_BY_ONE_SOFTMAX: S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and * NVTE_LEARNABLE_SOFTMAX: S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - * where alpha is a learnable parameter in shape [H]. + * where alpha is a learnable parameter of shape [H]. */ enum NVTE_Softmax_Type { /*! Vanilla softmax */ diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 7bc39f0745..98e2a29df8 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -50,7 +50,7 @@ class MMParams: Parameters ---------- - use_split_accumulator : bool, default = `True` + use_split_accumulator : bool, default = True Use FP8 fast accumulation on Hopper or Ada. For more details, see CUBLASLT_MATMUL_DESC_FAST_ACCUM option for cublasLtMatmul. """ @@ -159,7 +159,7 @@ def scaling_factor_compute(amax: Tensor, recipe: DelayedScaling) -> Tensor where `Tensor` is a framework tensor type. - reduce_amax: bool, default = `True` + reduce_amax: bool, default = True By default, if `torch.distributed` is initialized, the `amax` value for FP8 tensors is reduced across the `amax_reduction_group` (specified in the `autocast` call). This keeps the amaxes and scaling factors synced across the given @@ -167,13 +167,13 @@ def scaling_factor_compute(amax: Tensor, GPU maintains local amaxes and scaling factors. To ensure results are numerically identical across checkpointing boundaries in this case, all ranks must checkpoint in order to store the local tensors. - fp8_dpa: bool, default = `False` + fp8_dpa: bool, default = False Whether to enable FP8 dot product attention (DPA). When the model is placed in an `autocast(enabled=True)` region and `fp8_dpa` is set to `True`, DPA casts the inputs from higher precision to FP8, performs attention in FP8, and casts tensors back to higher precision as outputs. FP8 DPA currently is only supported in the `FusedAttention` backend. - fp8_mha: bool, default = `False` + fp8_mha: bool, default = False Whether to enable FP8 multi-head attention (MHA). When `True`, it removes the casting operations mentioned above at the DPA boundaries. Currently only standard MHA modules i.e. `LayerNormLinear/Linear + DPA + Linear`, are supported for this feature. When @@ -422,11 +422,11 @@ class NVFP4BlockScaling(Recipe): ---------- fp4_format : {Format.E2M1}, default = Format.E2M1 FP4 data type. - disable_rht : bool, default = `False` + disable_rht : bool, default = False If set to `True`, random Hadamard transforms are not applied to any tensor. - disable_stochastic_rounding : bool, default = `False` + disable_stochastic_rounding : bool, default = False If set to `True`, stochastic rounding is disabled during quantization for all tensors. - disable_2d_quantization : bool, default = `False` + disable_2d_quantization : bool, default = False If set to `True`, 1D block scaling with block size 16 is used for all tensors. """ @@ -492,17 +492,19 @@ 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: - qfactory( - role: str, - ) - - Where `role` is one of the following strings for e.g. te.Linear - (stable public contract): - - forward: "linear_input", "linear_weight", "linear_output" - - backward: "linear_grad_output", "linear_grad_input" + Factory callable that returns a quantizer instance for a + given semantic tensor role. + The callable is typically invoked as:: + + qfactory( + role: str, + ) + + Where `role` is one of the following strings for e.g. te.Linear + (stable public contract): + + - forward: "linear_input", "linear_weight", "linear_output" + - backward: "linear_grad_output", "linear_grad_input" """ qfactory: Callable[..., Any] diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index e8249de170..c5fb85041e 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -32,9 +32,9 @@ from ..quantize import ScaledTensor, ScaledTensorFactory, NoScaleTensor from ..quantize import ( Quantizer, - QuantizeLayout, DelayedScaleQuantizer, ScalingMode, + QuantizeLayout, ) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index c00b816f2e..76a8b225ba 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -39,12 +39,12 @@ Quantizer, GroupedQuantizer, QuantizerSet, - QuantizeLayout, noop_quantizer_set, is_fp8_gemm_with_all_layouts_supported, apply_padding_to_scale_inv, get_quantize_config_with_recipe, get_global_quantize_recipe, + QuantizeLayout, ) from .misc import get_padded_spec, is_all_reduce_in_float32 from ..sharding import ( diff --git a/transformer_engine/jax/cpp_extensions/misc.py b/transformer_engine/jax/cpp_extensions/misc.py index f15fe72bad..225d577cd3 100644 --- a/transformer_engine/jax/cpp_extensions/misc.py +++ b/transformer_engine/jax/cpp_extensions/misc.py @@ -116,7 +116,7 @@ def multidim_transpose(shape, static_axis_boundary=-1, transpose_axis=-1): transpose. Note, transpose_axis should be greater than static_axis_boundary examples: - X in shape (dim0, dim1, dim2, dim3, dim4) + X of shape (dim0, dim1, dim2, dim3, dim4) static_axis_boundary == -1, transpose_axis == 2 Xt = (dim2, dim3, dim4, dim0, dim1) diff --git a/transformer_engine/jax/cpp_extensions/normalization.py b/transformer_engine/jax/cpp_extensions/normalization.py index 92efb91a76..862780620e 100644 --- a/transformer_engine/jax/cpp_extensions/normalization.py +++ b/transformer_engine/jax/cpp_extensions/normalization.py @@ -35,9 +35,9 @@ from ..quantize import ScaledTensor, ScaledTensorFactory, NoScaleTensor from ..quantize import ( Quantizer, - QuantizeLayout, DelayedScaleQuantizer, ScalingMode, + QuantizeLayout, ) diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index b55fa20790..b3f24e9337 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -40,11 +40,11 @@ GroupedScaledTensor1x, Quantizer, GroupedQuantizer, - QuantizeLayout, ScalingMode, compute_scale_from_amax, NoScaleTensor, get_rht_matrix, + QuantizeLayout, ) diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 613455b6c3..c499b0651e 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -21,12 +21,12 @@ ScaledTensorFactory, ScaledTensor, ScalingMode, - QuantizeLayout, QuantizerSet, noop_quantizer_set, with_sharding_constraint_by_logical_axes, is_fp8_gemm_with_all_layouts_supported, TensorUsage, + QuantizeLayout, ) diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 19e4c57ce2..58df85fa52 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -279,26 +279,26 @@ class LayerNorm(nn.Module): # pylint: disable=too-few-public-methods layernorm_type : {'layernorm', 'rmsnorm'}, default = 'layernorm' Indicate the type of layer normalization. zero_centered_gamma : bool, default = False - If set to `True`, the LayerNorm formula changes to + If set to ``True``, the LayerNorm formula changes to .. math:: - y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * + y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} \cdot (1 + \gamma) + \beta - This parameter is only applicable for 'layernorm'. - The default of `scale_init` will also be changed. See `scale_init`. + This parameter is only applicable for ``'layernorm'``. + The default of ``scale_init`` will also be changed. See ``scale_init``. scale_init : Initializer, default = None Used for initializing scale factors :math:`\gamma`. - If `None` is provided, scale_init is set according to the value of zero_centered_gamma. - If zero_centered_gamma is set to `True`, then scale_init is `flax.linen.initializers.zeros`. - Otherwise, scale_init is `flax.linen.initializers.ones`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + If ``None`` is provided, scale_init is set according to the value of zero_centered_gamma. + If zero_centered_gamma is set to ``True``, then scale_init is ``flax.linen.initializers.zeros``. + Otherwise, scale_init is ``flax.linen.initializers.ones``. + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. scale_axes : Tuple[str, ...], default = ('embed', ) The name of axes used to shard the scale factors :math:`\gamma` with a corresponding mesh. bias_init : Initializer, default = flax.linen.initializers.zeros Used for initializing shift factors :math:`\beta`, only used when :attr:`layernorm_type='layernorm'`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. bias_axes : Tuple[str, ...], default = ('embed', ) The name of axes used to shard the shift factors :math:`\beta` with a corresponding mesh. only used when :attr:`layernorm_type='layernorm'`. @@ -424,15 +424,15 @@ class DenseGeneral(TransformerEngineBase): kernel_init : Initializer, default = flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal') Used for initializing weights. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. kernel_axes : Tuple[str, ...], default = () The name of axes used to shard the weights with a corresponding mesh. use_bias: bool, default = False Indicate whether to enable bias shifting. - If set to False, the layer will not learn an additive bias. + If set to ``False``, the layer will not learn an additive bias. bias_init: Initializer, default = flax.linen.initializers.zeros Used for initializing bias, only used when :attr:`use_bias=True`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. bias_axes: Tuple[str, ...], default = () The name of axes used to shard bias with a corresponding mesh, only used when :attr:`use_bias=True`. @@ -443,12 +443,12 @@ class DenseGeneral(TransformerEngineBase): :attr:`enable_low_rank_adaptation=True` low_rank_adaptation_alpha: float, default = None The alpha for computing the scaling factor of LoRA output. - :math:`\frac{alpha}{rank} * lora_output`. None means no scaling. + :math:`\frac{alpha}{rank} \cdot lora\_output`. ``None`` means no scaling. axis: Union[Iterable[int], int], default = -1 An integer tuple with axes to apply the transformation on. input_axes: Tuple[str, ...], default = None Indicate the logical axes of sharding constraint to the input, like - (BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES). Default is None, which means not to insert + ``(BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES)``. Default is ``None``, which means not to insert sharding constraint. Optimization parameters @@ -597,48 +597,48 @@ class LayerNormDenseGeneral(TransformerEngineBase): epsilon : float, default = 1e-6 A value added to the denominator of layer normalization for numerical stability. zero_centered_gamma : bool, default = False - If set to `True`, the LayerNorm formula changes to + If set to ``True``, the LayerNorm formula changes to .. math:: - y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * + y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} \cdot (1 + \gamma) + \beta - This parameter is only applicable for 'layernorm'. - The default of `scale_init` will also be changed. See `scale_init` + This parameter is only applicable for ``'layernorm'``. + The default of ``scale_init`` will also be changed. See ``scale_init`` scale_init : Initializer, default = None Used for initializing scale factors :math:`\gamma`. - If `None` is provided, scale_init is set according to the value of zero_centered_gamma. - If zero_centered_gamma is set to `True`, then scale_init is `flax.linen.initializers.zeros`. - Otherwise, scale_init is `flax.linen.initializers.ones`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + If ``None`` is provided, scale_init is set according to the value of zero_centered_gamma. + If zero_centered_gamma is set to ``True``, then scale_init is ``flax.linen.initializers.zeros``. + Otherwise, scale_init is ``flax.linen.initializers.ones``. + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. scale_axes : Tuple[str, ...], default = ('embed', ) The name of axes used to shard the scale factors :math:`\gamma` with a corresponding mesh, only used when :attr:`enable_layernorm=True`. ln_bias_init: Initializer, default = flax.linen.initializers.zeros Used for initializing shift factors :math:`\beta`, only used when :attr:`enable_layernorm=True` and :attr:`layernorm_type='layernorm'`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. ln_bias_axes: Tuple[str, ...], default = ('embed', ) The name of axes used to shard the shift factors :math:`\beta` with a corresponding mesh. It is only used when :attr:`enable_layernorm=True` and :attr:`layernorm_type='layernorm'`. kernel_init : Initializer, default = flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal') Used for initializing weights. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. kernel_axes : Tuple[str, ...], default = () The name of axes used to shard the weights with a corresponding mesh. use_bias: bool, default = False Indicate whether to enable bias shifting. - If set to False, the layer will not learn an additive bias. + If set to ``False``, the layer will not learn an additive bias. bias_init: Initializer, default = flax.linen.initializers.zeros Used for initializing bias, only used when :attr:`use_bias=True`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. bias_axes: Tuple[str, ...], default = () The name of axes used to shard bias with a corresponding mesh, only used when :attr:`use_bias=True`. return_layernorm_output: bool, default = False Indicate whether to return the output of layer normalization. - If set False, return None as the second tensor in outputs. + If set ``False``, return ``None`` as the second tensor in outputs. enable_low_rank_adaptation: bool, default = False Indicate whether to enable low rank adaptation for each dense layer. low_rank_adaptation_dim: int, default = 32 @@ -646,16 +646,16 @@ class LayerNormDenseGeneral(TransformerEngineBase): :attr:`enable_low_rank_adaptation=True` low_rank_adaptation_alpha: float, default = None The alpha for computing the scaling factor of LoRA output. - :math:`\frac{alpha}{rank} * lora_output`. None means no scaling. + :math:`\frac{alpha}{rank} \cdot lora\_output`. ``None`` means no scaling. axis: Union[Iterable[int], int], default = -1 An integer tuple with axes to apply the transformation on. layernorm_input_axes: Tuple[str, ...], default = None Indicate the logical axes of sharding constraint to the input of layernorm, like - (BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES). Default is None, which means not to insert + ``(BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES)``. Default is ``None``, which means not to insert sharding constraint. dot_input_axes: Tuple[str, ...], default = None Indicate the logical axes of sharding constraint to the input of dot, like - (BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES). Default is None, which means not to insert + ``(BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES)``. Default is ``None``, which means not to insert sharding constraint. Optimization parameters @@ -887,34 +887,34 @@ class LayerNormMLP(TransformerEngineBase): epsilon : float, default = 1e-6 A value added to the denominator of layer normalization for numerical stability. zero_centered_gamma : bool, default = False - If set to `True`, the LayerNorm formula changes to + If set to ``True``, the LayerNorm formula changes to .. math:: - y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * + y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} \cdot (1 + \gamma) + \beta - This parameter is only applicable for 'layernorm'. - The default of `scale_init` will also be changed. See `scale_init`. + This parameter is only applicable for ``'layernorm'``. + The default of ``scale_init`` will also be changed. See ``scale_init``. scale_init : Initializer, default = None Used for initializing scale factors :math:`\gamma`. - If `None` is provided, scale_init is set according to the value of zero_centered_gamma. - If zero_centered_gamma is set to `True`, then scale_init is `flax.linen.initializers.zeros`. - Otherwise, scale_init is `flax.linen.initializers.ones`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + If ``None`` is provided, scale_init is set according to the value of zero_centered_gamma. + If zero_centered_gamma is set to ``True``, then scale_init is ``flax.linen.initializers.zeros``. + Otherwise, scale_init is ``flax.linen.initializers.ones``. + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. scale_axes : Tuple[str, ...], default = ('embed', ) The name of axes used to shard the scale factors :math:`\gamma` with a corresponding mesh, only used when :attr:`enable_layernorm=True`. ln_bias_init: Initializer, default = flax.linen.initializers.zeros Used for initializing shift factors :math:`\beta`, only used when :attr:`enable_layernorm=True` and :attr:`layernorm_type='layernorm'`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. ln_bias_axes: Tuple[str, ...], default = ('embed', ) The name of axes used to shard the shift factors :math:`\beta` with a corresponding mesh. Only used when :attr:`enable_layernorm=True` and :attr:`layernorm_type='layernorm'`. kernel_init : Initializer, default = flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal') Used for initializing the weights of both dense layer transformations. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. kernel_axes_1 : Tuple[str, ...], default = ('embed', 'act', 'mlp') The name of axes used to shard the weights with a corresponding mesh for the weight of the first dense layer transformation. @@ -923,10 +923,10 @@ class LayerNormMLP(TransformerEngineBase): the weight of the second dense layer transformation. use_bias: bool, default = False Indicate whether to enable bias shifting. - If set to False, the layer will not learn an additive bias. + If set to ``False``, the layer will not learn an additive bias. bias_init: Initializer, default = flax.linen.initializers.zeros Used for initializing bias, only used when :attr:`use_bias=True`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. bias_axes_1: Tuple[str, ...], default = ('mlp',) The name of axes used to shard bias with a corresponding mesh for the weight of the first dense layer transformation. @@ -937,7 +937,7 @@ class LayerNormMLP(TransformerEngineBase): Only used when :attr:`use_bias=True`. return_layernorm_output: bool, default = False Indicate whether to return the output of layer normalization. - If set False, return None as the second tensor in outputs. + If set ``False``, return ``None`` as the second tensor in outputs. activations: Sequence[Union[str, Callable]], default = ('gelu',) The sequence of activation functions to apply after the first dense layer transformation. Each activation has its own transformation layer. @@ -958,20 +958,20 @@ class LayerNormMLP(TransformerEngineBase): :attr:`enable_low_rank_adaptation=True`. low_rank_adaptation_alpha: float, default = None The alpha for computing the scaling factor of LoRA output. - :math:`\frac{alpha}{rank} * lora_output`. None means no scaling. + :math:`\frac{alpha}{rank} \cdot lora\_output`. ``None`` means no scaling. axis: Union[Iterable[int], int], default = -1 An integer tuple with axes to apply the transformation on. layernorm_input_axes: Tuple[str, ...], default = None Indicate the logical axes of sharding constraint to the input of layernorm, like - (BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES). Default is None, which means not to insert + ``(BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES)``. Default is ``None``, which means not to insert sharding constraint. dot_1_input_axes: Tuple[str, ...], default = None Indicate the logical axes of sharding constraint to the input of 1st dot, like - (BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES). Default is None, which means not to insert + ``(BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES)``. Default is ``None``, which means not to insert sharding constraint. dot_2_input_axes: Tuple[str, ...], default = None Indicate the logical axes of sharding constraint to the input of 2nd dot, like - (BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES). Default is None, which means not to insert + ``(BATCH_AXES, SEQLEN_AXES, HIDDEN_AXES)``. Default is ``None``, which means not to insert sharding constraint. ffn1_ckpt_name: str = "ffn1" Checkpoint name for the output of the first fully-connected layer in the MLP block. diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index e51cc3691e..d0190f54c5 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -469,7 +469,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods The hidden dimension of each attention head. num_attention_heads: int The number of attention heads. - num_gqa_groups: int, default = `None` + num_gqa_groups: int, default = None Number of GQA groups. When `None` is present, it is equal to num_attention_heads. Grouped Query Attention is described in `this paper `_. @@ -482,32 +482,45 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods attn_mask_type: str, default = 'causal' This parameter specifies the type of attention mask to be applied during the softmax operation. - Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'} + Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'}. Each described below: - * no_mask: No attention mask is applied. This means the attention will consider the + * ``no_mask``: No attention mask is applied. This means the attention will consider the full sequence without any restrictions. - * padding: Indicates the presence of padding at the end of each sequence. - Users must provide a mask with the shape [batch, 1, max_seqlen_q, max_seqlen_kv] in the + * ``padding``: Indicates the presence of padding at the end of each sequence. + Users must provide a mask with the shape ``[batch, 1, max_seqlen_q, max_seqlen_kv]`` in the :attr:`__call__` method to specify the padding positions. - * causal: An upper triangular mask is applied to the softmax inputs, + * ``causal``: An upper triangular mask is applied to the softmax inputs, ensuring that the prediction for a certain position is only dependent on known outputs from positions before it. - * causal_padding / padding_causal: A combination of both causal and padding masks. - Both 'causal_padding' and 'padding_causal' are acceptable and have the same effect. + * ``causal_padding`` / ``padding_causal``: A combination of both causal and padding masks. + Both ``'causal_padding'`` and ``'padding_causal'`` are acceptable and have the same effect. - .. note:: :attr:`mask` in :attr:`__call__` is ignored for 'no_mask' and 'causal'. + | - .. note:: THD format only supports 'padding' or 'causal_padding' mask type. + .. note:: :attr:`mask` in :attr:`__call__` is ignored for ``'no_mask'`` and ``'causal'``. - attn_mask_type mask/sequence_descriptor SWA softmax type - -------------------------------------------------------------------------------------------- - no_mask None None SCALED - causal None None SCALED_UPPER_TRIANG_MASKED - causal None Yes SCALED_MASKED - padding Required Yes/No SCALED_MASKED - padding_causal Required Yes/No SCALED_MASKED + | + + .. note:: THD format only supports ``'padding'`` or ``'causal_padding'`` mask type. + + | + + .. table:: + :widths: auto + + ================== ============ ========== ============================== + attn_mask_type mask/sd SWA softmax type + ================== ============ ========== ============================== + no_mask None None SCALED + causal None None SCALED_UPPER_TRIANG_MASKED + causal None Yes SCALED_MASKED + padding Required Yes/No SCALED_MASKED + padding_causal Required Yes/No SCALED_MASKED + ================== ============ ========== ============================== + + where sd stands for sequence_descriptor. attn_bias_type: Optional[str], default = None Type of the attention bias passed in the attention. @@ -553,22 +566,40 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods Sliding window size. The default value is no sliding window. max_segments_per_seq: Optional[int], default = 1 The maximum number of segments per sequence, also used for THD format (sequence packing). - context_parallel_causal_load_balanced (bool): - Indicates the sequences are ordered for causal mask load balancing when running context parallelism. - context_parallel_axis (str): The name of the context parallel axis. - context_parallel_strategy (CPStrategy): The strategy of context parallel. 0: DEFAULT, 1: ALL_GATHER, 2: RING. - context_checkpoint_name (str): The name of the context checkpoint in the forward pass of fused attention. + context_parallel_causal_load_balanced: bool + Indicates the sequences are ordered for causal mask load balancing when running context parallelism. + context_parallel_axis: str + The name of the context parallel axis. + context_parallel_strategy: CPStrategy + The strategy of context parallel. 0: DEFAULT, 1: ALL_GATHER, 2: RING. + context_checkpoint_name: str + The name of the context checkpoint in the forward pass of fused attention. softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' - softmax type as described in this paper: + Softmax type as described in the paper `Efficient Streaming Language Models with Attention Sinks `_. - For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], - 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - where alpha is a learnable parameter in shape [h]. - 'off-by-one' and 'learnable' softmax types are also called sink attention - ('zero sink' and 'learnable sink'). + + For a given attention score :math:`S = Q \cdot K^T`, of shape ``[b, h, s_q, s_kv]``: + + * ``'vanilla'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{\sum_j \exp(S_{:,:,:,j})} + + * ``'off-by-one'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{1 + \sum_j \exp(S_{:,:,:,j})} + + * ``'learnable'``: + + .. math:: + Softmax(S)_{:,h,:,i} = \frac{\exp(S_{:,h,:,i})}{\exp(\alpha_h) + \sum_j \exp(S_{:,h,:,j})} + + where :math:`\alpha` is a learnable parameter of shape ``[h]``. + + ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention + (``'zero sink'`` and ``'learnable sink'``). Optimization parameters ----------------------- @@ -631,7 +662,7 @@ def __call__( mask: jax.numpy.ndarray, default = None Boolean tensor used to mask out the attention softmax input. :attr:`True` means to mask out the corresponding values. - Ignored when :attr:`self.attn_mask_type` is either 'no_mask' or 'causal'. + Ignored when :attr:`self.attn_mask_type` is either ``'no_mask'`` or ``'causal'``. bias: jax.numpy.ndarray, default = None A tensor used to shift attention softmax input. *: @@ -818,7 +849,7 @@ def rotary_pos_emb( ): """ Rotary Positional Embedding - x should be in shape of + x should be of shape [Batch, Seqlen, ..., Heads, Hidden] if transpose_batch_sequence is False, or [Seqlen, Batch, ..., Heads, Hidden] if transpose_batch_sequence is True. """ @@ -956,7 +987,7 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods The hidden dimension of each attention head. num_attention_heads: int The number of attention heads. - num_gqa_groups: int, default = `None` + num_gqa_groups: int, default = None Number of GQA groups. When `None` is present, it is equal to num_attention_heads. Grouped Query Attention is described in `this paper `_. @@ -969,28 +1000,28 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods attn_mask_type: str, default = 'causal' This parameter specifies the type of attention mask to be applied during the softmax operation. - Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'} + Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'}. Each described below: - * no_mask: No attention mask is applied. This means the attention will consider the + * ``no_mask``: No attention mask is applied. This means the attention will consider the full sequence without any restrictions. - * padding: Indicates the presence of padding at the end of each sequence. - Users must provide a mask with the shape [batch, 1, max_seqlen_q, max_seqlen_kv] in the + * ``padding``: Indicates the presence of padding at the end of each sequence. + Users must provide a mask with the shape ``[batch, 1, max_seqlen_q, max_seqlen_kv]`` in the :attr:`__call__` method to specify the padding positions. - * causal: An upper triangular mask is applied to the softmax inputs, + * ``causal``: An upper triangular mask is applied to the softmax inputs, ensuring that the prediction for a certain position is only dependent on known outputs from positions before it. - * causal_padding / padding_causal: A combination of both causal and padding masks. - Both 'causal_padding' and 'padding_causal' are acceptable and have the same effect. + * ``causal_padding`` / ``padding_causal``: A combination of both causal and padding masks. + Both ``'causal_padding'`` and ``'padding_causal'`` are acceptable and have the same effect. - .. note:: :attr:`mask` in :attr:`__call__` is ignored for 'no_mask' and 'causal'. + .. note:: :attr:`mask` in :attr:`__call__` is ignored for ``'no_mask'`` and ``'causal'``. attn_bias_type: Optional[str], default = None Type of the attention bias passed in the attention. - Available options: {'no_bias', 'pre_scale_bias', 'post_scale_bias'}. + Available options: ``{'no_bias', 'pre_scale_bias', 'post_scale_bias'}``. When default is present, the type is automatically decided by the MHA's bias parameter. - Where it is `post_scale_bias` if there is bias. Otherwise `no_bias` is used. + Where it is ``'post_scale_bias'`` if there is bias. Otherwise ``'no_bias'`` is used. dropout_rng_name: str, default = 'dropout' The key in given RNGs via flax.linen.Module.apply that is used to generate Dropout masks in the core attention. @@ -999,27 +1030,27 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods layernorm_epsilon: float, default = 1e-6 A value added to the denominator of layer normalization for numerical stability. zero_centered_gamma: bool, default = False - If set to `True`, the LayerNorm formula changes to + If set to ``True``, the LayerNorm formula changes to .. math:: - y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * + y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} \cdot (1 + \gamma) + \beta - This parameter is only applicable for 'layernorm'. + This parameter is only applicable for ``'layernorm'``. kernel_init: Initializer, default = - flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'normal') + ``flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'normal')`` Used for initializing the QKV and output projection weights. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. use_bias: bool, default = False Indicate whether or not to enable bias shifting for QKV and output projections. - If set to False, the layer will not learn additive biases. - bias_init: Initializer, default = flax.linen.initializers.zeros + If set to ``False``, the layer will not learn additive biases. + bias_init: Initializer, default = ``flax.linen.initializers.zeros`` Used for initializing bias of QKVO projections, only used when :attr:`use_bias=True`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. input_layernorm: bool, default = True - If set to False, layer normalization to the input is not applied. + If set to ``False``, layer normalization to the input is not applied. return_layernorm_output: bool, default = False - If set to True, output of layernorm is returned from the forward together with the output + If set to ``True``, output of layernorm is returned from the forward together with the output of the linear transformation. Example use case: residual connection for transformer module is taken post layernorm. enable_rotary_pos_emb: bool, default = False @@ -1029,17 +1060,17 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods only used when :attr:`enable_rotary_pos_emb=True` rotary_pos_emb_group_method: str, default = 'consecutive' Indicate the method to coupled the coordinates. It should be one of - ['consecutive', 'alternate']. 'alternate' is to pair index :math:`i` with :math:`i + d/2` - , d is the hidden dimension. 'consecutive' pairs index :math:`i` with :math:`i + 1`. + ``['consecutive', 'alternate']``. ``'alternate'`` is to pair index :math:`i` with :math:`i + d/2` + , d is the hidden dimension. ``'consecutive'`` pairs index :math:`i` with :math:`i + 1`. low_rank_adaptation_scope: str, default = 'none' Indicate the scope to apply low rank adaptation. It should be one of - ['none', 'all', 'qkv_proj', 'output_proj', 'exclude_qkv_proj', 'exclude_output_proj'] + ``['none', 'all', 'qkv_proj', 'output_proj', 'exclude_qkv_proj', 'exclude_output_proj']`` low_rank_adaptation_dim: int, default = 32 The dimension for low rank adaptation, only used when :attr:`enable_low_rank_adaptation=True` low_rank_adaptation_alpha: float, default = None The alpha for computing the scaling factor of LoRA output. - :math:`\frac{alpha}{rank} * lora_output`. None means no scaling. + :math:`\frac{alpha}{rank} \cdot lora\_output`. ``None`` means no scaling. enable_sequence_parallel: bool, default = False Whether to enable sequence parallelism to operations except dot. num_heads: int, default = None @@ -1066,8 +1097,8 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods should be in (seqlen, batch, hidden), otherwise (batch, seqlen, hidden). scale_attn_logits: bool, default = False Indicate whether to scale attention logits. - If set to True, :math:`\frac{Q}{\sqrt{head\_dim}*K}`, - else :math:`Q*K` + If set to True, :math:`\frac{Q \cdot K^T}{\sqrt{head\_dim}}`, + else :math:`Q \cdot K^T` scaled_query_init: bool, default = True Whether to scale WQ on initialization by :math:`\frac{1}{\sqrt{head\_dim}}` float32_logits: bool, default = False @@ -1078,16 +1109,31 @@ class MultiHeadAttention(nn.Module): # pylint: disable=too-few-public-methods window_size: Optional[Tuple[int, int]], default = None Sliding window size. Default value is no sliding window. softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' - softmax type as described in this paper: + Softmax type as described in the paper `Efficient Streaming Language Models with Attention Sinks `_. - For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], - 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - where alpha is a learnable parameter in shape [h]. - 'off-by-one' and 'learnable' softmax types are also called sink attention - ('zero sink' and 'learnable sink'). + + For a given attention score :math:`S = Q \cdot K^T`, of shape ``[b, h, s_q, s_kv]``: + + * ``'vanilla'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{\sum_j \exp(S_{:,:,:,j})} + + * ``'off-by-one'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{1 + \sum_j \exp(S_{:,:,:,j})} + + * ``'learnable'``: + + .. math:: + Softmax(S)_{:,h,:,i} = \frac{\exp(S_{:,h,:,i})}{\exp(\alpha_h) + \sum_j \exp(S_{:,h,:,j})} + + where :math:`\alpha` is a learnable parameter of shape ``[h]``. + + ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention + (``'zero sink'`` and ``'learnable sink'``). """ head_dim: int @@ -1202,7 +1248,7 @@ def __call__( mask: jax.numpy.ndarray, default = None Boolean tensor used to mask out the attention softmax input. :attr:`True` means mask out the corresponding values. - Ignored when :attr:`self.attn_mask_type` is either 'no_mask' or 'causal'. + Ignored when :attr:`self.attn_mask_type` is either ``'no_mask'`` or ``'causal'``. bias: jax.numpy.ndarray, default = None A tensor used to shift the attention softmax input. * @@ -1688,7 +1734,7 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods Intermediate size to which input samples are projected. num_attention_heads: int, default = 8 Number of attention heads in the transformer layer. - num_gqa_groups: int, default = `None` + num_gqa_groups: int, default = None Number of GQA groups. When `None` is present, it is equal to num_attention_heads. Grouped Query Attention is described in `this paper `_. @@ -1722,31 +1768,31 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods The key in given RNGs via flax.linen.Module.apply that for generating Dropout masks in the Multi-Head Attention. mha_kernel_init: Initializer, default = - flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'normal') + ``flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'normal')`` Used for initializing weights of QKV and Output projection weights. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. mlp_kernel_init: Initializer, default = - flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal') + ``flax.linen.initializers.variance_scaling(1.0, 'fan_in', 'truncated_normal')`` Used for initializing weights of FC1 and FC2 layers. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. mlp_activations: Sequence[str], default = ('gelu', ) The sequence of activation functions to apply after the first linear transformation. Each activation has its own transformation layer. mlp_activation_params: dict = None - This is only used when ('clamped_silu', 'clamped_linear') is in :attr:`mlp_activations`. At the moment - ClampedSwiglu is the only activation that requires parameters. + This is only used when ``('clamped_silu', 'clamped_linear')`` is in :attr:`mlp_activations`. At the moment + ``ClampedSwiglu`` is the only activation that requires parameters. use_bias: bool, default = False Indicate whether to enable bias shifting for QKVO projections, FC1 and FC2. - If set to False, the layer will not learn additive biases. - bias_init: Initializer, default = flax.linen.initializers.zeros + If set to ``False``, the layer will not learn additive biases. + bias_init: Initializer, default = ``flax.linen.initializers.zeros`` Used for initializing bias of QKVO projections, FC1 and FC2. It is only used when :attr:`use_bias=True`. - It should be a callable object with three arguments (jax.random.PRNGKey, shape, dtype). + It should be a callable object with three arguments ``(jax.random.PRNGKey, shape, dtype)``. apply_residual_connection_post_layernorm: bool, default = False - If set to True, residual connections are taken from the output + If set to ``True``, residual connections are taken from the output of layer norm (default is taken from input of layer norm) output_layernorm: bool, default = False - If set to True, layer normalization is applied on the output side, + If set to ``True``, layer normalization is applied on the output side, after the final dropout-add. default behavior is to apply layer normalization on the input side, before the QKV transformation. float32_attention_logits: bool, default = False @@ -1754,43 +1800,43 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods For fused attention backend, the accumulation is always float32 without the perf overhead. layer_type: TransformerLayerType, default = TransformerLayerType.ENCODER If set to TransformerLayerType.DECODER, an additional cross-attention block - is added after self-attention.this can be used for structures like `T5` + is added after self-attention.this can be used for structures like T5 Transformer in conjunction with the TransformerLayerType.ENCODER option. self_attn_mask_type: str, default = 'causal' This parameter specifies the type of attention mask to be applied during the softmax operation in the self attention. - Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'} + Available options are {'no_mask', 'padding', 'causal', 'causal_padding', 'padding_causal'}. Each described below: - * no_mask: No attention mask is applied. This means the self attention will consider the + * ``no_mask``: No attention mask is applied. This means the self attention will consider the full sequence without any restrictions. - * padding: Indicates the presence of padding at the end of each sequence. - Users must provide a mask with the shape [batch, 1, max_seqlen_q, max_seqlen_kv] in the + * ``padding``: Indicates the presence of padding at the end of each sequence. + Users must provide a mask with the shape ``[batch, 1, max_seqlen_q, max_seqlen_kv]`` in the :attr:`__call__` method to specify the padding positions. - * causal: An upper triangular mask is applied to the softmax inputs, + * ``causal``: An upper triangular mask is applied to the softmax inputs, ensuring that the prediction for a certain position is only dependent on known outputs from positions before it. - * causal_padding / padding_causal: A combination of both causal and padding masks. - Both 'causal_padding' and 'padding_causal' are acceptable and have the same effect. + * ``causal_padding`` / ``padding_causal``: A combination of both causal and padding masks. + Both ``'causal_padding'`` and ``'padding_causal'`` are acceptable and have the same effect. - .. note:: :attr:`attention_mask` in :attr:`__call__` is ignored for 'no_mask' and 'causal'. + .. note:: :attr:`attention_mask` in :attr:`__call__` is ignored for ``'no_mask'`` and ``'causal'``. self_attn_bias_type: Optional[str], default = None Type of the attention bias passed into the self attention. - Available options: {'no_bias', 'pre_scale_bias', 'post_scale_bias'}. + Available options: ``{'no_bias', 'pre_scale_bias', 'post_scale_bias'}``. When default is present, the type is automatically decided by the MHA's bias parameter. - Where it is `post_scale_bias` if there is bias. Otherwise `no_bias` is used. + Where it is ``'post_scale_bias'`` if there is bias. Otherwise ``'no_bias'`` is used. enable_relative_embedding: bool, default = True Whether to enable relative embedding as shifting of attention logits. relative_embedding: flax.linen.Module, default = None The module for relative embedding execution, only used when - :attr:`enable_relative_embedding=True`. Default is None, which will create + :attr:`enable_relative_embedding=True`. Default is ``None``, which will create an instance of RelativePositionBiases if :attr:`enable_relative_embedding=True`. - Default: RelativePositionBiases( num_buckets=32, max_distance=128, + Default: ``RelativePositionBiases( num_buckets=32, max_distance=128, num_attention_heads=self.num_attention_heads, dtype=self.dtype, embedding_init=flax.linen.initializers.variance_scaling(1.0, 'fan_avg', 'uniform'), - name='relpos_bias') + name='relpos_bias')`` enable_rotary_pos_emb: bool, default = False Whether to enable rotary position embedding to projected query and key in MHA. rotary_pos_emb_windows: Tuple[int, int], default = (1, 10000) @@ -1798,34 +1844,49 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods only used when :attr:`enable_rotary_pos_emb=True` rotary_pos_emb_group_method: str, default = 'consecutive' Indicate the method to couple the coordinates. It should be one of - ['consecutive', 'alternate']. 'alternate' is to pair index :math:`i` with :math:`i + d/2`, - where :math:`d` is the hidden dimension. 'consecutive' pairs index :math:`i` with + ``['consecutive', 'alternate']``. ``'alternate'`` is to pair index :math:`i` with :math:`i + d/2`, + where :math:`d` is the hidden dimension. ``'consecutive'`` pairs index :math:`i` with :math:`i + 1`. low_rank_adaptation_scope: str, default = 'none' Indicate the scope to apply low rank adaptation. It should be one of - ['none', 'all', 'qkv_proj', 'output_proj', 'mlp', 'exclude_qkv_proj', - 'exclude_output_proj', 'exclude_mlp'] + ``['none', 'all', 'qkv_proj', 'output_proj', 'mlp', 'exclude_qkv_proj', + 'exclude_output_proj', 'exclude_mlp']`` low_rank_adaptation_dim: int, default = 32 The dimension for low rank adaptation, only used when :attr:`enable_low_rank_adaptation=True` low_rank_adaptation_alpha: float, default = None The alpha for computing the scaling factor of LoRA output. - :math:`\frac{alpha}{rank} * lora\_output`. None means no scaling. + :math:`\frac{alpha}{rank} \cdot lora\_output`. ``None`` means no scaling. enable_sequence_parallel: bool, default = False Whether to enable sequence parallelism to operations except dot. window_size: Optional[Tuple[int, int]], default = None Sliding window size. Default value is no sliding window. softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' - Softmax type as described in this paper: + Softmax type as described in the paper `Efficient Streaming Language Models with Attention Sinks `_. - For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], - 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - where alpha is a learnable parameter in shape [h]. - 'off-by-one' and 'learnable' softmax types are also called sink attention - ('zero sink' and 'learnable sink'). + + For a given attention score :math:`S = Q \cdot K^T`, of shape ``[b, h, s_q, s_kv]``: + + * ``'vanilla'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{\sum_j \exp(S_{:,:,:,j})} + + * ``'off-by-one'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{1 + \sum_j \exp(S_{:,:,:,j})} + + * ``'learnable'``: + + .. math:: + Softmax(S)_{:,h,:,i} = \frac{\exp(S_{:,h,:,i})}{\exp(\alpha_h) + \sum_j \exp(S_{:,h,:,j})} + + where :math:`\alpha` is a learnable parameter of shape ``[h]``. + + ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention + (``'zero sink'`` and ``'learnable sink'``). Only supported for fused attention backend. Optimization parameters @@ -1836,19 +1897,19 @@ class TransformerLayer(nn.Module): # pylint: disable=too-few-public-methods When > 0.0, applies stochastic depth per sample in the main path of the residual block. fuse_qkv_params: bool, default = True - If set to True, `TransformerLayer` module exposes a single fused + If set to ``True``, ``TransformerLayer`` module exposes a single fused parameter for query-key-value for self-attention and key-value for cross-attention. transpose_batch_sequence: bool, default = False Indicate whether the input tensors were switched axis of batch - and sequence length dimension. if set to True, the input tensors - should be in (seqlen, batch, hidden), otherwise (batch, seqlen, hidden). + and sequence length dimension. if set to ``True``, the input tensors + should be in ``(seqlen, batch, hidden)``, otherwise ``(batch, seqlen, hidden)``. scale_attn_logits: bool, default = False Indicate whether to scale attention logits. - if set to True, :math:`\frac{Q}{\sqrt{head_dim}*K}`, - else :math:`Q*K` - scaled_query_init: bool, default = `True` - Whether to scale WQ on initialization by :math:`\sqrt{head_dim}` + if set to ``True``, :math:`\frac{Q \cdot K^T}{\sqrt{head\_dim}}`, + else :math:`Q \cdot K^T` + scaled_query_init: bool, default = True + Whether to scale WQ on initialization by :math:`\sqrt{head\_dim}` """ hidden_size: int = 512 @@ -1931,7 +1992,7 @@ def __call__( attention_mask : jax.numpy.ndarray, default = None Boolean tensor used to mask out self-attention softmax input. :attr:`True` means mask out the corresponding values. - Ignored when :attr:`self.self_attn_mask_type` is either 'no_mask' or 'causal'. + Ignored when :attr:`self.self_attn_mask_type` is either ``'no_mask'`` or ``'causal'``. encoder_decoder_mask: jax.numpy.ndarray, default = None Boolean tensor used to mask out cross-attention softmax input when :attr:`layer_type=TransformerLayerType.DECODER`. diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 9d894a389b..5341af3d74 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -7,22 +7,14 @@ # pylint: disable=wrong-import-position import functools -from packaging.version import Version as PkgVersion import torch from transformer_engine.common import load_framework_extension - - -@functools.lru_cache(maxsize=None) -def torch_version() -> tuple[int, ...]: - """Get PyTorch version""" - return PkgVersion(str(torch.__version__)).release - +from transformer_engine.pytorch.torch_version import torch_version assert torch_version() >= (2, 1), f"Minimum torch version 2.1 required. Found {torch_version()}." - load_framework_extension("torch") from transformer_engine.pytorch.module import LayerNormLinear from transformer_engine.pytorch.module import Linear 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 550c2a4a5a..f506035c1e 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 @@ -152,25 +152,25 @@ class DotProductAttention(TransformerEngineBaseModule): - """Allows the model to jointly attend to information from different + r"""Allows the model to jointly attend to information from different representation subspaces as described in the paper: `Attention Is All You Need `_. .. note:: - Argument :attr:`attention_mask` in the `forward` call is only used when - :attr:`attn_mask_type` includes '"padding"' or `"arbitrary"`. + Argument :attr:`attention_mask` in the ``forward`` call is only used when + :attr:`attn_mask_type` includes '"padding"' or ``"arbitrary"``. .. warning:: FlashAttention uses a non-deterministic algorithm for optimal performance. To observe - deterministic behavior at the cost of performance, use FlashAttention version >= `2.4.1` + deterministic behavior at the cost of performance, use FlashAttention version >= ``2.4.1`` and set the environment variable :attr:`NVTE_ALLOW_NONDETERMINISTIC_ALGO=0`. In order - to disable`flash-attn` entirely, set :attr:`NVTE_FLASH_ATTN=0`. + to disable ``flash-attn`` entirely, set :attr:`NVTE_FLASH_ATTN=0`. .. note:: - Transformer Engine stores the FP8 metadata under a `._extra_state` key when checkpointing. + Transformer Engine stores the FP8 metadata under a ``._extra_state`` key when checkpointing. As the FP8 attention support expands from one backend to multiple backends, the location of that key has also shifted (see `FP8 checkpoint compatibility `_). @@ -182,118 +182,137 @@ class DotProductAttention(TransformerEngineBaseModule): kv_channels : Union[int, Tuple[int, int]] the head size in key and value tensors. If the same, :attr:`kv_channels` can be an integer; if not, :attr:`kv_channels` should be a tuple of two integers. - num_gqa_groups : Optional[int] = None + num_gqa_groups : Optional[int], default = None number of GQA groups in the transformer layer. Grouped Query Attention is described in `this paper `_. This only affects the keys and values, not the queries. GQA-1 is equivalent to Multi-Query Attention (`MQA `_), while GQA-H - is equivalent to MHA, i.e. `num_gqa_groups = num_attention_heads`. - attention_dropout: float, default = 0.0 + is equivalent to MHA, i.e. ``num_gqa_groups = num_attention_heads``. + attention_dropout : float, default = 0.0 dropout probability for the dropout op during multi-head attention. - attn_mask_type: str, default = `causal` - type of attention mask passed into softmax operation, options are "`no_mask`", - "`padding`", "`causal`", "`padding,causal`", "`causal,padding`", - "`padding_causal`", "`causal_bottom_right`", "`padding_causal_bottom_right`", and - "`arbitrary`", where "`padding,causal`", "`causal,padding`" and "`padding_causal`" + attn_mask_type : str, default = "causal" + type of attention mask passed into softmax operation, options are ``"no_mask"``, + ``"padding"``, ``"causal"``, ``"padding,causal"``, ``"causal,padding"``, + ``"padding_causal"``, ``"causal_bottom_right"``, ``"padding_causal_bottom_right"``, and + ``"arbitrary"``, where ``"padding,causal"``, ``"causal,padding"`` and ``"padding_causal"`` are equivalent. This arg can be overridden by :attr:`attn_mask_type` in the - `forward` method. It is useful for cases involving compilation/tracing, e.g. + :meth:`forward` method. It is useful for cases involving compilation/tracing, e.g. ONNX export, and the forward arg is useful for dynamically changing mask types, e.g. a different mask for training and inference. - 1. For "`no_mask`", no attention mask is applied. - 2. For "`causal`", "`causal_bottom_right`", or the causal mask in - "`padding_causal`" and "`padding_causal_bottom_right`", Transformer Engine - calculates and applies an upper triangular mask to the softmax input. - No user input is needed. Causal masks without the "`bottom_right`" appendix align - the diagonal line to the top left corner of the softmax matrix. With - "`bottom_right`", the causal mask is aligned to the bottom right corner, which is - often used in inference/KV caching. - 3. For "`padding`", or the padding mask in "`padding_causal`" and - "`padding_causal_bottom_right`", users need to provide the locations of padded - tokens, either via :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv` (both in shape - [batch_size + 1]), or via :attr:`attention_mask` (one tensor for self-attention - in shape [batch_size, 1, 1, max_seqlen_q], or two tensors in a tuple for - cross-attention in shapes [batch_size, 1, 1, max_seqlen_q] and - [batch_size, 1, 1, max_seqlen_kv]). - 4. For "`arbitrary`", users need to provide a mask that is broadcastable to - the shape of softmax input [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]. - window_size: Optional[Tuple[int, int]], default = `None` + + 1. For ``"no_mask"``, no attention mask is applied. + 2. For ``"causal"``, ``"causal_bottom_right"``, or the causal mask in + ``"padding_causal"`` and ``"padding_causal_bottom_right"``, Transformer Engine + calculates and applies an upper triangular mask to the softmax input. + No user input is needed. Causal masks without the ``"bottom_right"`` appendix align + the diagonal line to the top left corner of the softmax matrix. With + ``"bottom_right"``, the causal mask is aligned to the bottom right corner, which is + often used in inference/KV caching. + 3. For ``"padding"``, or the padding mask in ``"padding_causal"`` and + ``"padding_causal_bottom_right"``, users need to provide the locations of padded + tokens, either via :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv` (both of shape + ``[batch_size + 1]``), or via :attr:`attention_mask` (one tensor for self-attention + of shape ``[batch_size, 1, 1, max_seqlen_q]``, or two tensors in a tuple for + cross-attention of shapes ``[batch_size, 1, 1, max_seqlen_q]`` and + ``[batch_size, 1, 1, max_seqlen_kv]``). + 4. For ``"arbitrary"``, users need to provide a mask that is broadcastable to + the shape of softmax input ``[batch_size, num_heads, max_seqlen_q, max_seqlen_kv]``. + + window_size : Optional[Tuple[int, int]], default = None sliding window size for local attention, where query at position i attends to keys - in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q - + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding - window and causal mask specifically. Both `causal` and `causal_bottom_right` masks - map to `window_size = (-1, 0)` and Transformer Engine distinguishes them based on - `attn_mask_type`. Similar to :attr:`attn_mask_type`, `window_size` can - be overridden by :attr:`window_size` in `forward` as well. - attention_type: str, default = `self` - type of attention, either "`self`" and "`cross`". - layer_number: int, default = `None` - layer number of the current `DotProductAttention` when multiple such modules + in ``[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + + window_size[1]] inclusive. Special cases ``(-1, -1)`` and ``(-1, 0)`` mean no sliding + window and causal mask specifically. Both ``causal`` and ``causal_bottom_right`` masks + map to ``window_size = (-1, 0)`` and Transformer Engine distinguishes them based on + ``attn_mask_type``. Similar to :attr:`attn_mask_type`, ``window_size`` can + be overridden by :attr:`window_size` in ``forward`` as well. + attention_type : str, default = "self" + type of attention, either ``"self"`` and ``"cross"``. + layer_number : int, default = None + layer number of the current ``DotProductAttention`` when multiple such modules are concatenated, for instance in consecutive transformer blocks. - qkv_format: str, default = `sbhd` - dimension format for `query_layer`, `key_layer` and `value_layer`, - {`sbhd`, `bshd`, `thd`}. `s` stands for the sequence length, `b` batch size, - `h` the number of heads, `d` head size, and `t` the total number of tokens - in a batch, with `t = sum(s_i), for i = 0...b-1`. `sbhd` and `bshd` formats + qkv_format : str, default = "sbhd" + dimension format for ``query_layer``, ``key_layer`` and ``value_layer``, + {``"sbhd"``, ``"bshd"``, ``"thd"``}. ``s`` stands for the sequence length, ``b`` batch size, + ``h`` the number of heads, ``d`` head size, and ``t`` the total number of tokens + in a batch, with ``t = sum(s_i), for i = 0...b-1``. ``"sbhd"`` and ``"bshd"`` formats are used for when sequences in a batch are of equal length or padded to - equal length, and the `thd` format is used for when sequences in a batch + equal length, and the ``"thd"`` format is used for when sequences in a batch have different lengths. Please note that these formats do not reflect how - tensors `query_layer`, `key_layer`, `value_layer` are laid out in memory. - For that, please use `get_qkv_layout` to gain the layout information. - softmax_scale: Optional[float], default = `None` - softmax scale for the attention scores. If `None`, defaults to - `1.0/math.sqrt(kv_channels if isinstance(kv_channels, int) else kv_channels[0])`. - softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' - softmax type as described in this paper: + tensors ``query_layer``, ``key_layer``, ``value_layer`` are laid out in memory. + For that, please use ``get_qkv_layout`` to gain the layout information. + softmax_scale : Optional[float], default = None + softmax scale for the attention scores. If ``None``, defaults to + ``1.0/math.sqrt(kv_channels if isinstance(kv_channels, int) else kv_channels[0])``. + softmax_type : str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' + Softmax type as described in the paper `Efficient Streaming Language Models with Attention Sinks `_. - For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], - 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - where alpha is a learnable parameter in shape [h]. - 'off-by-one' and 'learnable' softmax types are also called sink attention - ('zero sink' and 'learnable sink'). - return_max_logit: Optional[bool], default = `False` + + For a given attention score :math:`S = Q \cdot K^T`, of shape ``[b, h, s_q, s_kv]``: + + * ``'vanilla'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{\sum_j \exp(S_{:,:,:,j})} + + * ``'off-by-one'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{1 + \sum_j \exp(S_{:,:,:,j})} + + * ``'learnable'``: + + .. math:: + Softmax(S)_{:,h,:,i} = \frac{\exp(S_{:,h,:,i})}{\exp(\alpha_h) + \sum_j \exp(S_{:,h,:,j})} + + where :math:`\alpha` is a learnable parameter of shape ``[h]``. + + ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention + (``'zero sink'`` and ``'learnable sink'``). + + return_max_logit : Optional[bool], default = False If true, returns the maximum attention score that can be used in a Muon optimizer to rescale the Q and K projection weights (see `Muon is Scalable for LLM Training `_). - max_logit = max(S), where S = mask(Q*K^T*softmax_scale + bias) in shape [b, h, s_q, s_kv], - and max_logit is in shape [h]. + :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]``. Parallelism parameters ---------------------- - sequence_parallel : bool, default = `False` - if set to `True`, uses sequence parallelism. + sequence_parallel : bool, default = False + if set to ``True``, uses sequence parallelism. tp_size : int, default = 1 tensor parallel world size. - tp_group : ProcessGroup, default = `None` + tp_group : ProcessGroup, default = None tensor parallel process group. - cp_group : Union[ProcessGroup, List[ProcessGroup]], default = `None` + cp_group : Union[ProcessGroup, List[ProcessGroup]], default = None context parallel process group. - ProcessGroup is for cp_comm_type of "p2p", "all_gather", and "a2a". - List[ProcessGroup] is for cp_comm_type of "a2a+p2p", where cp_group[0] - and cp_group[1] are for a2a and p2p communications respectively. - cp_global_ranks : list of global rank IDs, default = `None` - global rank IDs of GPUs that are in cp_group. - cp_stream : CUDA stream, default = `None` + ``ProcessGroup`` is for :attr:`cp_comm_type` of ``"p2p"``, ``"all_gather"``, and ``"a2a"``. + ``List[ProcessGroup]`` is for :attr:`cp_comm_type` of ``"a2a+p2p"``, where :attr:`cp_group[0]` + and :attr:`cp_group[1]` are for ``"a2a"`` and ``"p2p"`` communications respectively. + cp_global_ranks : list of global rank IDs, default = None + global rank IDs of GPUs that are in ``cp_group``. + cp_stream : CUDA stream, default = None context parallelism splits flash attention into multiple steps for compute and communication overlapping. To address the wave quantization issue of each split step, we add an additional CUDA stream so that we can overlap two flash attention kernels. - cp_comm_type : str, default = `p2p` + cp_comm_type : str, default = "p2p" inter-gpu communication type for context parallelism. - Can be "p2p" or "all_gather" or "a2a" or "a2a+p2p". - "p2p": Exchange KV chunks with P2P communications in ring topology. - P2P is async and can be overlapped with attention compute. - "all_gather": All-gather to get full sequence of KV before attention. - The all-gather is not async, and cannot be overlapped. - "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP - group, and gather to get full sequence of QKV. - "a2a+p2p": hierarchical CP implementation. First applying a2a to QKV - across each CP sub-group (e.g., via NVLink), then exchanging KV with - p2p between sub-groups (e.g., via IBLink). + Can be ``"p2p"`` or ``"all_gather"`` or ``"a2a"`` or ``"a2a+p2p"``. + + - ``"p2p"``: Exchange KV chunks with P2P communications in ring topology. + P2P is async and can be overlapped with attention compute. + - ``"all_gather"``: All-gather to get full sequence of KV before attention. + The all-gather is not async, and cannot be overlapped. + - ``"a2a"``: Like DeepSpeed Ulysses, scatter attention heads across the CP + group, and gather to get full sequence of QKV. + - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV + across each CP sub-group (e.g., via NVLink), then exchanging KV with + p2p between sub-groups (e.g., via IBLink). """ def __init__( @@ -468,8 +487,8 @@ def _load_from_state_dict( ): """ This function helps to load Transformer Engine 1.6 and 1.7 checkpoints, where FP8 attention - metadata is stored under the `core_attention.fused_attention._extra_state` key and not the - `core_attention._extra_state` key. Please see `FP8 checkpoint compatibility + metadata is stored under the ``core_attention.fused_attention._extra_state`` key and not the + ``core_attention._extra_state`` key. Please see `FP8 checkpoint compatibility `_ for more details. """ fused_attn_key = False @@ -522,25 +541,26 @@ def set_context_parallel_group( ---------- cp_group : Union[ProcessGroup, List[ProcessGroup]] context parallel process group. - ProcessGroup is for cp_comm_type of "p2p", "all_gather", and "a2a". - List[ProcessGroup] is for cp_comm_type of "a2a+p2p", where cp_group[0] - and cp_group[1] are for a2a and p2p communications respectively. + ``ProcessGroup`` is for :attr:`cp_comm_type` of ``"p2p"``, ``"all_gather"``, and ``"a2a"``. + ``List[ProcessGroup]`` is for :attr:`cp_comm_type` of ``"a2a+p2p"``, where :attr:`cp_group[0]` + and :attr:`cp_group[1]` are for ``"a2a"`` and ``"p2p"`` communications respectively. cp_global_ranks : List[int] list of global ranks in the context group. cp_stream : torch.cuda.Stream cuda stream for context parallel execution. - cp_comm_type : str, default = `p2p` + cp_comm_type : str, default = "p2p" inter-gpu communication type for context parallelism. - Can be "p2p" or "all_gather" or "a2a" or "a2a+p2p". - "p2p": Exchange KV chunks with P2P communications in ring topology. - P2P is async and can be overlapped with attention compute. - "all_gather": All-gather to get full sequence of KV before attention. - The all-gather is not async, and cannot be overlapped. - "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP - group, and gather to get full sequence of QKV. - "a2a+p2p": hierarchical CP implementation. First applying a2a to QKV - across each CP sub-group (e.g., via NVLink), then exchanging KV with - p2p between sub-groups (e.g., via IBLink). + Can be ``"p2p"`` or ``"all_gather"`` or ``"a2a"`` or ``"a2a+p2p"``. + + - ``"p2p"``: Exchange KV chunks with P2P communications in ring topology. + P2P is async and can be overlapped with attention compute. + - ``"all_gather"``: All-gather to get full sequence of KV before attention. + The all-gather is not async, and cannot be overlapped. + - ``"a2a"``: Like DeepSpeed Ulysses, scatter attention heads across the CP + group, and gather to get full sequence of QKV. + - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV + across each CP sub-group (e.g., via NVLink), then exchanging KV with + p2p between sub-groups (e.g., via IBLink). """ self.cp_group = cp_group self.cp_global_ranks = cp_global_ranks @@ -801,13 +821,13 @@ def forward( fp8_output: Optional[bool] = False, num_splits: Optional[int] = 1, ) -> torch.Tensor: - """ + r""" Dot Product Attention Layer. .. note:: Argument :attr:`attention_mask` is only used when :attr:`attn_mask_type` - includes '"padding"' or `"arbitrary"`. + includes ``"padding"`` or ``"arbitrary"``. .. note:: @@ -846,24 +866,24 @@ def forward( Pass in :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`, or :attr:`attention_mask` (which will be converted to :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`), to provide the real sequence length information. For example, a batch of 3 sequences - [a a a b b c c c c] can be padded to [a a a PAD b b PAD PAD c c c c], and the cumulative + ``[a a a b b c c c c]`` can be padded to ``[a a a PAD b b PAD PAD c c c c]``, and the cumulative sequence length tensors would be - :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = [0, 3, 5, 9] for self-attention. + :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = ``[0, 3, 5, 9]`` for self-attention. 2. Do not perform padding on training data. Use :attr:`qkv_format` = "thd" and :attr:`attn_mask_type` = {"padding", "padding_causal", "padding_causal_bottom_right"}. Pass in :attr:`cu_seqlens_q` and :attr:`cu_seqlens_kv`, or :attr:`attention_mask`, - as in option 1. For example, a batch of 3 sequences [a a a b b c c c c] can be processed + as in option 1. For example, a batch of 3 sequences ``[a a a b b c c c c]`` can be processed without any padding, and the sequence length tensors would be - :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = [0, 3, 5, 9] for self-attention. + :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = ``[0, 3, 5, 9]`` for self-attention. In certain use cases, a varying number of identifier tokens are inserted between sequences. These tokens do not participate in the attention calculation. :attr:`cu_seqlens_q_padded` and :attr:`cu_seqlens_kv_padded` must be specified in such cases to correctly identify the start and end of each sequence in a batch. - For example, a batch of 3 sequences [a a a 1 b b 2 2 c c c c 3] would have - :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = [0, 3, 5, 9], and - :attr:`cu_seqlens_q_padded` = :attr:`cu_seqlens_kv_padded` = [0, 4, 8, 13] + For example, a batch of 3 sequences ``[a a a 1 b b 2 2 c c c c 3]`` would have + :attr:`cu_seqlens_q` = :attr:`cu_seqlens_kv` = ``[0, 3, 5, 9]``, and + :attr:`cu_seqlens_q_padded` = :attr:`cu_seqlens_kv_padded` = ``[0, 4, 8, 13]`` for self-attention. .. note:: @@ -898,81 +918,81 @@ def forward( value_layer : torch.Tensor Value tensor. attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]], - default = `None`. Boolean tensor(s) used to mask out attention softmax input. - It should be `None` for causal masks and "`no_mask`". For padding masks, it should be - a single tensor of [batch_size, 1, 1, seqlen_q] for self-attention, and a tuple of - two tensors in shapes [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv] - for cross-attention. For "`arbitrary`" mask, it should be in a shape broadcastable - to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]. A `True` value means - the corresponding position is masked out and a `False` means that position + default = None. Boolean tensor(s) used to mask out attention softmax input. + It should be ``None`` for causal masks and ``"no_mask"``. For padding masks, it should be + a single tensor of ``[batch_size, 1, 1, seqlen_q]`` for self-attention, and a tuple of + two tensors of shapes ``[batch_size, 1, 1, seqlen_q]`` and ``[batch_size, 1, 1, seqlen_kv]`` + for cross-attention. For ``"arbitrary"`` mask, it should be of a shape broadcastable + to ``[batch_size, num_heads, max_seqlen_q, max_seqlen_kv]``. A ``True`` value means + the corresponding position is masked out and a ``False`` means that position is allowed to participate in attention. - qkv_format: str, default = `None` + qkv_format: str, default = None If provided, overrides :attr:`qkv_format` from initialization. - cu_seqlens_q: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (without offset) in a batch for `query_layer`, + cu_seqlens_q: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (without offset) in a batch for ``query_layer``, with shape [batch_size + 1] and dtype torch.int32. See :ref:`note` for more details. - cu_seqlens_kv: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (without offset) in a batch for `key_layer` - and `value_layer`, with shape [batch_size + 1] and dtype torch.int32. + cu_seqlens_kv: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (without offset) in a batch for ``key_layer`` + and ``value_layer``, with shape [batch_size + 1] and dtype torch.int32. See :ref:`note` for more details. - cu_seqlens_q_padded: Optional[torch.Tensor], default = `None` + cu_seqlens_q_padded: Optional[torch.Tensor], default = None Cumulative sum of sequence lengths (with offset) in a batch for - `query_layer`, with shape [batch_size + 1] and dtype torch.int32. + ``query_layer``, with shape ``[batch_size + 1]`` and dtype torch.int32. When there is no padding between sequences in a batch, - `cu_seqlens_q_padded = cu_seqlens_q`. + :attr:`cu_seqlens_q_padded` = :attr:`cu_seqlens_q`. See :ref:`note` for more details. - cu_seqlens_kv_padded: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (with offset) in a batch for `key_layer` - and `value_layer`, with shape [batch_size + 1] and dtype torch.int32. + cu_seqlens_kv_padded: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (with offset) in a batch for ``key_layer`` + and ``value_layer``, with shape ``[batch_size + 1]`` and dtype torch.int32. When there is no padding between sequences in a batch, - `cu_seqlens_kv_padded = cu_seqlens_kv`. + :attr:`cu_seqlens_kv_padded` = :attr:`cu_seqlens_kv`. See :ref:`note` for more details. - max_seqlen_q: Optional[int], default = `None` - Maximum sequence length in `query_layer`. + max_seqlen_q: Optional[int], default = None + Maximum sequence length in ``query_layer``. See :ref:`note` for more details. - max_seqlen_kv: Optional[int], default = `None` - Maximum sequence length in `key_layer` and `value_layer`. + max_seqlen_kv: Optional[int], default = None + Maximum sequence length in ``key_layer`` and ``value_layer``. See :ref:`note` for more details. attn_mask_type: {'no_mask', 'padding', 'causal', 'padding,causal', 'causal,padding', 'padding_causal', 'causal_bottom_right', 'padding_causal_bottom_right', - 'arbitrary'}, default = `None`. Type of attention mask passed into + 'arbitrary'}, default = None. Type of attention mask passed into softmax operation. 'padding,causal', 'causal,padding' and 'padding_causal' are equivalent. By default, causal masks are aligned to the top left corner - of the softmax matrix. When "`bottom_right`" is specified in the mask type, + of the softmax matrix. When ``"bottom_right"`` is specified in the mask type, causal masks are aligned to the bottom right corner. - window_size: Optional[Tuple[int, int]], default = `None` + window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention. - checkpoint_core_attention : bool, default = `False` + checkpoint_core_attention : bool, default = False If true, forward activations for attention are recomputed during the backward pass in order to save memory that would otherwise be occupied to store the forward activations until backprop. - core_attention_bias_type: str, default = `no_bias` - Bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`} - core_attention_bias: Optional[torch.Tensor], default = `None` - Bias tensor for Q * K.T, shape [1, num_head, max_seqlen_q, max_seqlen_kv]. - It should be 'None' for 'no_bias' and 'alibi' bias types. - alibi_slopes: Optional[torch.Tensor], default = `None` - ALiBi slopes in FP32 and shape [nheads] or [batch_size, nheads]. + core_attention_bias_type: str, default = "no_bias" + Bias type, {``"no_bias"``, ``"pre_scale_bias"``, ``"post_scale_bias"``, ``"alibi"``} + core_attention_bias: Optional[torch.Tensor], default = None + Bias tensor for :math:`Q \cdot K^T`, shape ``[1, num_head, max_seqlen_q, max_seqlen_kv]``. + It should be ``None`` for ``"no_bias"`` and ``"alibi"`` bias types. + alibi_slopes: Optional[torch.Tensor], default = None + ALiBi slopes in FP32 and shape ``[nheads]`` or ``[batch_size, nheads]``. It adds a bias of (-alibi_slope * (i + seqlen_k - seqlen_q - j)) to the attention score of query i and key j. - fast_zero_fill: bool, default = `True` + fast_zero_fill: bool, default = True Whether to use the fast path to set output tensors to 0 or not. - inference_params: Optional[InferenceParams], default = `None` + inference_params: Optional[InferenceParams], default = None Optimizes execution performance during inference by caching Keys and Values of the current decoding iteration. These cached values are appended to the K and V values computed in previous iterations, eliminating the need to recalculate them for the entire sequence. - Initialization of `inference_params` is required prior to use to ensure sufficient + Initialization of ``inference_params`` is required prior to use to ensure sufficient memory allocation. Adjustments of the sequence_len_offset should be done after a complete forward pass. If rotary positional embeddings (RoPE) are utilized, they must be prepared beforehand. Supports "sbhd" and "bshd" layouts, with the "sbhd" layout being more efficient. - pad_between_seqs: Optional[bool], default = `None` - If None, inferred from qkv_format, cu_seqlens and cu_seqlens_padded. - If true, there are padding tokens between individual sequences in a packed batch. - fp8_output: Optional[bool], default = `False` + pad_between_seqs: Optional[bool], default = None + If ``None``, inferred from qkv_format, cu_seqlens and cu_seqlens_padded. + If ``True``, there are padding tokens between individual sequences in a packed batch. + fp8_output: Optional[bool], default = False Whether to enforce output to be in FP8 or not. num_splits: Optional[int], default = 1 Optional split control for FlashAttention-3 only. When set, this value is forwarded diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7a61c60094..8c6b6afc90 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -175,65 +175,65 @@ class AttentionParams: Parameters ---------- - qkv_type: Union[torch.Tensor, Float8Tensor], default = `torch.Tensor` + qkv_type : Union[torch.Tensor, Float8Tensor], default = torch.Tensor Type of query/key/value tensors, {`torch.Tensor`, `Float8Tensor`}. - qkv_dtype: torch.dtype, default = `torch.bfloat16` + qkv_dtype : torch.dtype, default = torch.bfloat16 Data type of query/key/value tensors. - qkv_layout: str, default = "sbh3d" + qkv_layout : str, default = "sbh3d" Query/key/value tensor memory layout. - batch_size: int, default = 1 + batch_size : int, default = 1 Batch size. - num_heads: int, default = 16 + num_heads : int, default = 16 Number of attention heads in the query tensor. - num_gqa_groups: int, default = 16 + num_gqa_groups : int, default = 16 Number of attention heads in key and value tensors. - max_seqlen_q: int, default = 128 + max_seqlen_q : int, default = 128 Maximum sequence length of the query tensor. - max_seqlen_kv: int, default = 128 + max_seqlen_kv : int, default = 128 Maximum sequence length of the key and value tensors. - head_dim_qk: int, default = 64 + head_dim_qk : int, default = 64 The size of each attention head in query and key tensors. - head_dim_v: int, default = 64 + head_dim_v : int, default = 64 The size of each attention head in the value tensor. - attn_mask_type: str, default = `no_mask` + attn_mask_type : str, default = no_mask Attention mask type, {`no_mask`, `padding`, `causal`, `padding_causal`, `causal_bottom_right`, `padding_causal_bottom_right`, `arbitrary`} - window_size: Tuple[int, int], default = None + window_size : Tuple[int, int], default = None Sliding window attention size. - alibi_slopes_shape: Optional[Union[torch.Size, List]], default = `None` + alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. - core_attention_bias_type: str, default = `no_bias` + core_attention_bias_type : str, default = no_bias Attention bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`}. - core_attention_bias_shape: str, default = `1hss` + core_attention_bias_shape : str, default = 1hss Attention bias shape, {`1hss`, `b1ss`, `bhss`}. - core_attention_bias_requires_grad: bool, default = `True` + core_attention_bias_requires_grad : bool, default = True Whether attention bias requires gradient. - pad_between_seqs: bool, default = `False` + pad_between_seqs : bool, default = False Whether there is padding between sequences in a batch. This only applies to `qkv_format=thd`. - attention_dropout: float, default = 0.0 + attention_dropout : float, default = 0.0 Attention dropout. - context_parallel: bool, default = `False` + context_parallel : bool, default = False Whether context parallelism is used or not. - cp_comm_type: str, default = "p2p" + cp_comm_type : str, default = "p2p" The communication type of context parallelism. - deterministic: bool, default = `False` + deterministic : bool, default = False Whether to run `DotProductAttention` with determinism or not. - is_training: bool, default = `True` + is_training : bool, default = True Whether in training mode (`True`) or inference mode (`False`) - fp8: bool, default = `False` + fp8 : bool, default = False Whether `DotProductAttention` is in an `autocast` region. - fp8_meta: Optional[Dict[str Any]], default = `None` + fp8_meta : Optional[Dict[str Any]], default = None The FP8 metadata tensor of `DotProductAttention`. - inference_params: Optional[InferenceParams], default = `None` + inference_params : Optional[InferenceParams], default = None Inference-related parameters. See InferenceParams for details. - softmax_type: str, default = "vanilla" + softmax_type : str, default = "vanilla" The type of softmax operation. See DotProductAttention for details. - return_max_logit: bool, default = `False` + return_max_logit : bool, default = False Whether to output max_logit. - cuda_graph: bool, default = `False` + cuda_graph : bool, default = `False` Whether support for cuda graph capture is needed or not. - num_splits: int, default = 1 + num_splits : int, default = 1 The number of kernels to split attention to. """ @@ -298,15 +298,15 @@ def get_attention_backend( Returns ---------- - use_flash_attention: bool + use_flash_attention : bool Whether the `FlashAttention` backend has been selected. - use_fused_attention: bool + use_fused_attention : bool Whether the `FusedAttention` backend has been selected. - fused_attention_backend: tex.NVTE_Fused_Attn_Backend + fused_attention_backend : tex.NVTE_Fused_Attn_Backend If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`. - use_unfused_attention: bool + use_unfused_attention : bool Whether the `UnfusedDotProductAttention` backend has been selected. - available_backends: List[bool] + available_backends : List[bool] All available backends that could support the provided input. A list of Booleans in the form of [use_flash_attention, use_fused_attention, use_unfused_attention]. """ @@ -835,8 +835,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # ---------------------------------------------------------------------------------------- # no_mask | None | All # padding | | All - # self-attention | One tensor in shape [b, 1, 1, sq] | - # cross-attention | Tuple of two tensors in shapes | + # self-attention | One tensor of shape [b, 1, 1, sq] | + # cross-attention | Tuple of two tensors of shapes | # | [b, 1, 1, sq] and [b, 1, 1, skv] | # causal | None | # self-attention | | All @@ -846,7 +846,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # cross-attention | | FusedAttention, UnfusedDotProductAttention # causal_bottom_right | None | All # padding_causal_bottom_right | Same as "padding" | All - # arbitrary | One tensor in shape broadcastable to | UnfusedDotProductAttention + # arbitrary | One tensor of shape broadcastable to | UnfusedDotProductAttention # | [b, h, sq, skv] | if attn_mask_type == "arbitrary": if (use_flash_attention_2 and FlashAttentionUtils.is_installed) or ( @@ -1271,42 +1271,42 @@ def get_full_mask( Parameters ---------- - max_seqlen_q: int + max_seqlen_q : int Maximum sequence length for queries. - max_seqlen_kv: int + max_seqlen_kv : int Maximum sequence length for keys and values. - attn_mask_type: str, default = `no_mask` - Attention mask type, {"`no_mask`", "`padding`", "`causal`", "`padding_causal`", - "`causal_bottom_right`", "`padding_causal_bottom_right`", "`arbitrary`"} - attention_mask: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], - default = `None` + attn_mask_type : str, default = no_mask + Attention mask type, {``"no_mask"``, ``"padding"``, ``"causal"``, ``"padding_causal"``, + ``"causal_bottom_right"``, ``"padding_causal_bottom_right"``, ``"arbitrary"``} + attention_mask : Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], + default = None Boolean tensor(s) used to mask out attention softmax input. Please see DotProductAttention for the requirements of `attention_mask` for different `attn_mask_type`s. - window_size: Tuple[int, int], default = `None` + window_size : Tuple[int, int], default = None Sliding window size for local attention, where query at position i attends to keys in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding window and causal mask specifically. Both `causal` and `causal_bottom_right` masks map to `window_size = (-1, 0)` and Transformer Engine distinguishes them based on `attn_mask_type`. - attention_type: str, default = "self" + attention_type : str, default = "self" Attention type, {"self", "cross"} - bottom_right_alignment: bool, default = `True` + bottom_right_alignment : bool, default = True Whether to align the diagonal of the sliding window attention to the bottom right (`True`) or top left (`False`) corner of the softmax matrix. Ignored if `attn_mask_type` explicitly specifies "causal" or "causal_bottom_right". Returns ---------- - attn_mask_type: str + attn_mask_type : str For sliding window attention (>=0, >0), "arbitrary"; otherwise, the same as input `attn_mask_type` - attention_mask: torch.Tensor + attention_mask : torch.Tensor The full attention mask based on `attn_mask_type`, `attention_mask` and `window_size` - actual_seqlens_q: torch.Tensor - For padding masks, the actual sequence lengths for queries, in shape [batch_size]. + actual_seqlens_q : torch.Tensor + For padding masks, the actual sequence lengths for queries, of shape [batch_size]. For other masks, `None`. - actual_seqlens_kv: Optional[torch.Tensor], default = `None` - For padding masks, the actual sequence lengths for keys and values, in shape [batch_size]. + actual_seqlens_kv : Optional[torch.Tensor], default = None + For padding masks, the actual sequence lengths for keys and values, of shape [batch_size]. For other masks, `None`. """ # perform basic checks @@ -1392,29 +1392,29 @@ def get_alibi( """ Parameters ---------- - num_heads: int + num_heads : int Number of heads. - max_seqlen_q: int + max_seqlen_q : int Maximum sequence length for queries. - max_seqlen_kv: int + max_seqlen_kv : int Maximum sequence length for keys and values. - actual_seqlens_q: Optional[torch.Tensor], default = `None` - Actual sequence lengths for queries, in shape [batch_size]. - actual_seqlens_kv: Optional[torch.Tensor], default = `None` - Actual sequence lengths for keys and values, in shape [batch_size]. - alibi_slopes: Optional[torch.Tensor], default = `None` - Custom ALiBi slopes, FP32, CUDA tensor, in shape [num_heads] or [batch_size, num_heads]. - bias_dtype: Optional[torch.dtype], default = `None` + actual_seqlens_q : Optional[torch.Tensor], default = None + Actual sequence lengths for queries, of shape [batch_size]. + actual_seqlens_kv : Optional[torch.Tensor], default = None + Actual sequence lengths for keys and values, of shape [batch_size]. + alibi_slopes : Optional[torch.Tensor], default = None + Custom ALiBi slopes, FP32, CUDA tensor, of shape [num_heads] or [batch_size, num_heads]. + bias_dtype : Optional[torch.dtype], default = None Dtype of the generated ALiBi bias. If None, use torch.float32. - bottom_right_alignment: bool, default = `True` + bottom_right_alignment : bool, default = True Whether to align the diagonal of the ALiBi bias to the bottom right corner of the matrix (`True`) or top left (`False`). Returns ---------- - alibi_slopes: torch.Tensor + alibi_slopes : torch.Tensor ALiBi slopes in FP32 and shape [num_heads] or [batch_size, num_heads]. - alibi_bias: torch.Tensor + alibi_bias : torch.Tensor ALiBi bias in FP32 or `bias_dtype`. Its shape is (1) [1, num_heads, max_seqlen_q, max_seqlen_kv] if `alibi_slopes` is in [num_heads] shape, and `actual_seqlens_q` and `actual_seqlens_kv` are `None`; or @@ -1818,18 +1818,18 @@ def get_qkv_format( Parameters ---------- - qkv_layout: str + qkv_layout : str Memory layout of `q`, `k` and `v`. See get_qkv_layout() for more details. - inference_params: InferenceParams, default = `None` + inference_params : InferenceParams, default = None InferenceParams related to KV caching. Returns ---------- - qkv_format: str, default = `sbhd` + qkv_format : str, default = sbhd Dimension format for `q`, `k` and `v`, {`sbhd`, `bshd`, `thd`}. - q_format: str + q_format : str Format of the `q` tensor, {`bshd`, `sbhd`, `thd`}. - kv_format: str + kv_format : str Format of the `k` and `v` tensors, {`bshd`, `sbhd`, `thd`}. """ splited = qkv_layout.replace("paged_kv_", "").split("_") @@ -1855,23 +1855,23 @@ def get_qkv_layout( Parameters ---------- - q: torch.Tensor + q : torch.Tensor Query tensor. - k: torch.Tensor + k : torch.Tensor Key tensor. - v: torch.Tensor + v : torch.Tensor Value tensor. - qkv_format: str, default = `sbhd` + qkv_format : str, default = sbhd Dimension format for `q`, `k` and `v`, {`sbhd`, `bshd`, `thd`}. `s` stands for the sequence length dimension, `b` batch size, `h` the number of attention heads, `d` head size, and `t` the total number of tokens in a batch, i.e. `t = sum(s_i) for i = 0...b-1`. - inference_params: InferenceParams, default = `None` + inference_params : InferenceParams, default = None InferenceParams related to KV caching. Returns ---------- - qkv_layout: str + qkv_layout : str Memory layout of `q`, `k` and `v`. Each `qkv_layout` maps to a pair of `q_format` and `kv_format` in {`bshd`, `sbhd`, `thd`}. The `paged_kv_` prefix is used to indicate that paged KV caching is in play. A few examples of the layouts are as follows. @@ -1893,18 +1893,18 @@ def get_qkv_layout( `thd_2bshd`: {`thd_bshd_bshd`, `paged_kv_thd_bshd_bshd`} `thd_2sbhd`: {`thd_sbhd_sbhd`, `paged_kv_thd_sbhd_sbhd`} - q: torch.Tensor + q : torch.Tensor Query tensor. It may be different from input `q` as we try to fit tensors to a supported layout. - k: torch.Tensor + k : torch.Tensor Key tensor. It may be different from input `k` as we try to fit tensors to a supported layout. - v: torch.Tensor + v : torch.Tensor Value tensor. It may be different from input `v` as we try to fit tensors to a supported layout. - q_format: str + q_format : str Format of the query tensor, {`bshd`, `sbhd`, `thd`}. - kv_format: str + kv_format : str Format of the key and value tensors, {`bshd`, `sbhd`, `thd`}. """ diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index f0ef8d0bd5..4ae1bd09a1 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -98,29 +98,29 @@ class DotProductAttention: Parameters ---------- - max_batch_size: int + max_batch_size : int Maximum batch size in inference - max_sequence_length: int + max_sequence_length : int Maximum sequence length in inference - num_heads_kv: int + num_heads_kv : int Number of attention heads in keys and values - head_dim_k: int + head_dim_k : int Head size for keys - dtype: torch.dtype + dtype : torch.dtype Data type of the KV cache - head_dim_v: int, default = None + head_dim_v : int, default = None Head size for values. If None, initialized as head_dim_k. - is_paged: bool, default = False + is_paged : bool, default = False Whether the KV cache is paged (True) or non-paged (False) - total_num_pages: int, default = None + total_num_pages : int, default = None Total number of pages in the KV cache. Required for is_paged = True. - page_size: int, default = None + page_size : int, default = None Page size of the KV cache. Required for is_paged = True. - max_ctx_len: int, default = None + max_ctx_len : int, default = None Maximum context length in inference. 1 <= max_ctx_len <= max_sequence_length. - qkv_format: str, default = "bshd" + qkv_format : str, default = "bshd" Format of the incoming query/key/value tensors in current iteration - custom_cache_manager: KVCacheManager, default = None + custom_cache_manager : KVCacheManager, default = None Custom cache manager, with KVCacheManager as the base class. """ @@ -525,9 +525,9 @@ def step( new_v: torch.Tensor New value tokens for layer_number in current inference iteration cu_new_seqlens: torch.Tensor - Cumulative sequence lengths for new_k and new_v, in shape [batch_size + 1] + Cumulative sequence lengths for new_k and new_v, of shape [batch_size + 1] cu_cached_seqlens: torch.Tensor - Cumulative sequence lengths for k_cache and v_cache (after new tokens are copied in), in shape [batch_size + 1] + Cumulative sequence lengths for k_cache and v_cache (after new tokens are copied in), of shape [batch_size + 1] qkv_format: str Format of new_k and new_v tensors, {'bshd', 'sbhd', 'thd'} @@ -701,7 +701,7 @@ def get_page_list(self, seq: int): return [x.page_id for x in self.allocated_pages[seq]] def get_page_table(self, sequences: List[int]): - """Get the page table, in shape [batch_size, max_pages_per_seq]""" + """Get the page table, of shape [batch_size, max_pages_per_seq]""" page_table = torch.Tensor( [ self.get_page_list(seq) + [0] * (self.max_pages_per_seq - self.get_page_count(seq)) @@ -783,9 +783,9 @@ def step( new_v: torch.Tensor New value tokens for layer_number in current inference iteration cu_new_seqlens: torch.Tensor - Cumulative sequence lengths for new_k and new_v, in shape [batch_size + 1] + Cumulative sequence lengths for new_k and new_v, of shape [batch_size + 1] cu_cached_seqlens: torch.Tensor - Cumulative sequence lengths for k_cache and v_cache (after new tokens are copied in), in shape [batch_size + 1] + Cumulative sequence lengths for k_cache and v_cache (after new tokens are copied in), of shape [batch_size + 1] qkv_format: str Format of new_k and new_v tensors, {'bshd', 'sbhd', 'thd'} diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 2440693df4..beb13b7f1e 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -50,8 +50,8 @@ class MultiheadAttention(torch.nn.Module): .. note:: - Argument :attr:`attention_mask` in the `forward` call is only used when - :attr:`attn_mask_type` includes '"padding"' or `"arbitrary"`. + Argument :attr:`attention_mask` in the :meth:`forward() ` method is only used when + :attr:`attn_mask_type` includes ``"padding"`` or ``"arbitrary"``. Parameters ---------- @@ -59,57 +59,56 @@ class MultiheadAttention(torch.nn.Module): size of each input sample. num_attention_heads : int number of attention heads in the transformer layer. - kv_channels: int, default = `None` + kv_channels : int, default = None number of key-value channels. defaults to - :attr:`hidden_size` / :attr:`num_attention_heads` if `None`. - attention_dropout: float, default = 0.1 + :attr:`hidden_size` / :attr:`num_attention_heads` if ``None``. + attention_dropout : float, default = 0.1 dropout probability for the dropout op during multi-head attention. layernorm_epsilon : float, default = 1e-5 a value added to the denominator of layer normalization for numerical stability. - init_method : Callable, default = `None` + init_method : Callable, default = None used for initializing weights of QKV and 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)`. - output_layer_init_method : Callable, default = `None` + ``init_method(weight)``. When set to ``None``, defaults to + ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + output_layer_init_method : Callable, default = None used for initializing weights of PROJ and FC2 in the following way: - `output_layer_init_method(weight)`. When set to `None`, defaults to - `torch.nn.init.normal_(mean=0.0, std=0.023)`. - layer_number: int, default = `None` - layer number of the current `TransformerLayer` when multiple such modules are + ``output_layer_init_method(weight)``. When set to ``None``, defaults to + ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + layer_number : int, default = None + layer number of the current ``TransformerLayer`` when multiple such modules are concatenated to form a transformer block. - attn_mask_type: {'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right', + attn_mask_type : {'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right', 'padding_causal_bottom_right','arbitrary'}, - default = `causal` + default = "causal" type of attention mask passed into softmax operation. Overridden by - :attr:`attn_mask_type` in the `forward` method. The forward + :attr:`attn_mask_type` in the :meth:`forward` method. The :meth:`forward` arg is useful for dynamically changing mask types, e.g. a different - mask for training and inference. The init arg is useful for cases + mask for training and inference. The :meth:`__init__` arg is useful for cases involving compilation/tracing, e.g. ONNX export. - window_size: Optional[Tuple[int, int]], default = `None` + window_size : Optional[Tuple[int, int]], default = None sliding window size for local attention, where query at position i attends to keys - in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q - + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding - window and causal mask specifically. Both `causal` and `causal_bottom_right` masks - map to `window_size = (-1, 0)` and Transformer Engine distinguishes them based on - `attn_mask_type`. Similar to :attr:`attn_mask_type`, `window_size` can - be overridden by :attr:`window_size` in `forward` as well. - num_gqa_groups : int, default = `None` + in ``[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]]`` inclusive. Special cases ``(-1, -1)`` and ``(-1, 0)`` mean no sliding + window and causal mask specifically. Both ``"causal"`` and ``"causal_bottom_right"`` masks + map to ``window_size = (-1, 0)`` and Transformer Engine distinguishes them based on + ``attn_mask_type``. Similar to :attr:`attn_mask_type`, ``window_size`` can + be overridden by :attr:`window_size` in :meth:`forward` as well. + num_gqa_groups : int, default = None number of GQA groups in the transformer layer. Grouped Query Attention is described in `this paper `_. This only affects the keys and values, not the querys. GQA-1 is equivalent to Multi-Query Attention (`MQA `_), while GQA-H - is equivalent to MHA, i.e. `num_gqa_groups = num_attention_heads`. - return_layernorm_output : bool, default = `False` - if set to `True`, output of layernorm is returned from the forward + is equivalent to MHA, i.e. ``num_gqa_groups = num_attention_heads``. + return_layernorm_output : bool, default = False + if set to ``True``, output of layernorm is returned from the :meth:`forward` method together with the output of the linear transformation. Example use case: residual connection for transformer module is taken post layernorm. - input_layernorm: bool, default = `False` - if set to `True`, layer normalization to the input is applied. - attention_type: { 'self', 'cross' }, default = 'self' + input_layernorm : bool, default = False + if set to ``True``, layer normalization to the input is applied. + attention_type : { 'self', 'cross' }, default = 'self' type of attention applied. zero_centered_gamma : bool, default = 'False' if set to 'True', gamma parameter in LayerNorm is initialized to 0 and @@ -120,103 +119,118 @@ class MultiheadAttention(torch.nn.Module): (1 + \gamma) + \beta normalization : { 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm' type of normalization applied. - qkv_weight_interleaved : bool, default = `True` - if set to `False`, the QKV weight is interpreted as a concatenation of - query, key, and value weights along the `0th` dimension. The default - interpretation is that the individual `q`, `k`, and `v` weights for each - attention head are interleaved. This parameter is set to `False` when + qkv_weight_interleaved : bool, default = True + if set to ``False``, the QKV weight is interpreted as a concatenation of + query, key, and value weights along the ``0th`` dimension. The default + interpretation is that the individual ``q``, ``k``, and ``v`` weights for each + attention head are interleaved. This parameter is set to ``False`` when using :attr:`fuse_qkv_params=False`. - rotary_pos_interleaved : bool, default = `False` + rotary_pos_interleaved : bool, default = False whether to use interleaved rotary position embeddings. - bias : bool, default = `True` - if set to `False`, the transformer layer will not learn any additive biases. + bias : bool, default = True + if set to ``False``, the transformer layer will not learn any additive biases. 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 forward pass. - qkv_format: str, default = `sbhd` - dimension format for `query_layer`, `key_layer` and `value_layer`, - {`sbhd`, `bshd`}. `s` stands for the sequence length, `b` batch size, - `h` the number of heads and `d` head size. `sbhd` and `bshd` formats + qkv_format : str, default = "sbhd" + dimension format for ``query_layer``, ``key_layer`` and ``value_layer``, + {``"sbhd"``, ``"bshd"``}. ``s`` stands for the sequence length, ``b`` batch size, + ``h`` the number of heads and ``d`` head size. ``"sbhd"`` and ``"bshd"`` formats are used for when sequences in a batch are of equal length or padded to equal length. Please note that these formats do not reflect how - tensors `query_layer`, `key_layer`, `value_layer` are laid out in memory. - For that, please use `get_qkv_layout` to gain the layout information. - name: str, default = `None` + tensors ``query_layer``, ``key_layer``, ``value_layer`` are laid out in memory. + For that, please use ``get_qkv_layout`` to gain the layout information. + name : str, default = None name of the module, currently used for debugging purposes. - softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' - softmax type as described in this paper: + softmax_type : str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' + Softmax type as described in the paper `Efficient Streaming Language Models with Attention Sinks `_. - For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], - 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - where alpha is a learnable parameter in shape [h]. - 'off-by-one' and 'learnable' softmax types are also called sink attention - ('zero sink' and 'learnable sink'). + + For a given attention score :math:`S = Q \cdot K^T`, of shape ``[b, h, s_q, s_kv]``: + + * ``'vanilla'``: + + .. math:: + S_{:,:,:,i} = = \frac{\exp(S_{:,:,:,i})}{\sum_j \exp(S_{:,:,:,j})} + + * ``'off-by-one'``: + + .. math:: + S_{:,:,:,i} = = \frac{\exp(S_{:,:,:,i})}{1 + \sum_j \exp(S_{:,:,:,j})} + + * ``'learnable'``: + + .. math:: + S_{:,:,:,i} = = \frac{\exp(S_{:,h,:,i})}{\exp(\alpha_h) + \sum_j \exp(S_{:,h,:,j})} + + where :math:`\alpha` is a learnable parameter of shape ``[h]``. + + ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention + (``'zero sink'`` and ``'learnable sink'``). Parallelism parameters ---------------------- - set_parallel_mode : bool, default = `False` - if set to `True`, QKV and FC1 layers are used as Column Parallel + set_parallel_mode : bool, default = False + if set to ``True``, QKV and FC1 layers are used as Column Parallel whereas PROJ and FC2 is used as Row Parallel as described `here `_. - sequence_parallel : bool, default = `False` - if set to `True`, uses sequence parallelism. - tp_group : ProcessGroup, default = `None` + sequence_parallel : bool, default = False + if set to ``True``, uses sequence parallelism. + tp_group : ProcessGroup, default = None tensor parallel process group. tp_size : int, default = 1 used as TP (tensor parallel) world size when TP groups are not formed during initialization. In this case, users must call the - `set_tensor_parallel_group(tp_group)` method on the initialized module before the + ``set_tensor_parallel_group(tp_group)`` method on the initialized module before the forward pass to supply the tensor parallel group needed for tensor and sequence parallel collectives. Optimization parameters ----------------------- fuse_wgrad_accumulation : bool, default = 'False' - if set to `True`, enables fusing of creation and accumulation of + if set to ``True``, enables fusing of creation and accumulation of the weight gradient. When enabled, it is assumed that the weights - have an additional `main_grad` attribute (used instead of the - regular `grad`) which is a pre-allocated buffer of the correct + have an additional ``main_grad`` attribute (used instead of the + regular ``grad``) which is a pre-allocated buffer of the correct size to accumulate gradients in. - params_dtype : torch.dtype, default = `torch.get_default_dtype()` + params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. - return_bias : bool, default = `False` - when set to `True`, this module will not apply the additive bias itself, but - instead return the bias value during the forward pass together with the + return_bias : bool, default = False + when set to ``True``, this module will not apply the additive bias itself, but + instead return the bias value during the :meth:`forward` method together with the output of the linear transformation :math:`y = xA^T`. This is useful when the bias addition can be fused to subsequent operations. - fuse_qkv_params: bool, default = 'False' - if set to `True`, `TransformerLayer` module exposes a single fused + fuse_qkv_params : bool, default = 'False' + if set to ``True``, ``TransformerLayer`` module exposes a single fused parameter for query-key-value. This enables optimizations such as QKV fusion without concatentations/splits and also enables the argument - `fuse_wgrad_accumulation`. - qk_norm_type: Optional[str], default = None + ``fuse_wgrad_accumulation``. + qk_norm_type : Optional[str], default = None type of normalization to apply to query and key tensors. - Options: None, 'L2Normalization', 'RMSNorm', 'LayerNorm'. When None, no normalization is applied. - When 'L2Normalization', L2 normalization is applied to query and key tensors. - When 'RMSNorm', RMS normalization is applied to query and key tensors. - When 'LayerNorm', layer normalization is applied to query and key tensors. + Options: ``None``, ``'L2Normalization'``, ``'RMSNorm'``, ``'LayerNorm'``. When ``None``, no normalization is applied. + When ``'L2Normalization'``, L2 normalization is applied to query and key tensors. + When ``'RMSNorm'``, RMS normalization is applied to query and key tensors. + When ``'LayerNorm'``, layer normalization is applied to query and key tensors. Normalization is applied after RoPE (if applicable) but before attention computation - when `qk_norm_before_rope` is False. This follows the e.g. Llama4 approach + when ``qk_norm_before_rope`` is ``False``. This follows the e.g. Llama4 approach for QK normalization to improve training stability and model performance. - qk_norm_eps: float, default = 1e-6 + qk_norm_eps : float, default = 1e-6 epsilon value for normalization of query and key tensors. - Only used when `qk_norm_type` is not None. - qk_norm_before_rope: bool, default = `False` - if set to `True`, query and key normalization is applied before rotary position - embedding. When `False` (default), normalization is applied after RoPE. + Only used when ``qk_norm_type`` is not ``None``. + qk_norm_before_rope : bool, default = False + if set to ``True``, query and key normalization is applied before rotary position + embedding. When ``False`` (default), normalization is applied after RoPE. This parameter allows supporting different architectural variants that apply QK normalization at different points. - seq_length: Optional[int], default = `None` + seq_length : Optional[int], default = None sequence length of input samples. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propagation and activation recompute phase. - micro_batch_size: Optional[int], default = `None` + micro_batch_size : Optional[int], default = None batch size per training step. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propagation and activation recompute phase. @@ -535,7 +549,7 @@ def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> N Parameters ---------- - tp_group : ProcessGroup, default = `None` + tp_group : ProcessGroup, default = None tensor parallel process group. """ self.tp_group = tp_group @@ -555,25 +569,26 @@ def set_context_parallel_group( ---------- cp_group : Union[ProcessGroup, List[ProcessGroup]] context parallel process group. - ProcessGroup is for cp_comm_type of "p2p", "all_gather", and "a2a". - List[ProcessGroup] is for cp_comm_type of "a2a+p2p", where cp_group[0] - and cp_group[1] are for a2a and p2p communications respectively. + ``ProcessGroup`` is for :attr:`cp_comm_type` of ``"p2p"``, ``"all_gather"``, and ``"a2a"``. + ``List[ProcessGroup]`` is for :attr:`cp_comm_type` of ``"a2a+p2p"``, where :attr:`cp_group[0]` + and :attr:`cp_group[1]` are for ``"a2a"`` and ``"p2p"`` communications respectively. cp_global_ranks : List[int] list of global ranks in the context group. cp_stream : torch.cuda.Stream cuda stream for context parallel execution. - cp_comm_type : str, default = `p2p` + cp_comm_type : str, default = "p2p" inter-gpu communication type for context parallelism. - Can be "p2p" or "all_gather" or "a2a", "a2a+p2p". - "p2p": Exchange KV chunks with P2P communications in ring topology. - P2P is async and can be overlapped with attention compute. - "all_gather": All-gather to get full sequence of KV before attention. - The all-gather is not async, and cannot be overlapped. - "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP - group, and gather to get full sequence of QKV. - "a2a+p2p": hierarchical CP implementation. First applying a2a to QKV - across each CP sub-group (e.g., via NVLink), then exchanging KV with - p2p between sub-groups (e.g., via IBLink). + Can be ``"p2p"`` or ``"all_gather"`` or ``"a2a"`` or ``"a2a+p2p"``. + + - ``"p2p"``: Exchange KV chunks with P2P communications in ring topology. + P2P is async and can be overlapped with attention compute. + - ``"all_gather"``: All-gather to get full sequence of KV before attention. + The all-gather is not async, and cannot be overlapped. + - ``"a2a"``: Like DeepSpeed Ulysses, scatter attention heads across the CP + group, and gather to get full sequence of QKV. + - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV + across each CP sub-group (e.g., via NVLink), then exchanging KV with + p2p between sub-groups (e.g., via IBLink). """ if isinstance(cp_group, dist_group_type): self.cp_size = get_distributed_world_size(cp_group) @@ -622,39 +637,39 @@ def forward( fast_zero_fill: bool = True, pad_between_seqs: Optional[bool] = None, ) -> Tuple[Union[torch.Tensor, None], ...]: - """ + r""" Forward propagation for MultiheadAttention layer. .. note:: Argument :attr:`attention_mask` is only used when :attr:`attn_mask_type` - includes `"padding"` or `"arbitrary"`. + includes ``"padding"`` or ``"arbitrary"``. Parameters ---------- hidden_states : torch.Tensor Input tensor. attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]], - default = `None`. Boolean tensor(s) used to mask out attention softmax input. - It should be `None` for causal masks and "`no_mask`". For padding masks, it should be - a single tensor of [batch_size, 1, 1, seqlen_q] for self-attention, and a tuple of - two tensors in shapes [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv] - for cross-attention. For "`arbitrary`" mask, it should be in a shape broadcastable to - [batch_size, num_heads, max_seqlen_q, max_seqlen_kv]. A `True` value means - the corresponding position is masked out and a `False` means that position + default = None. Boolean tensor(s) used to mask out attention softmax input. + It should be ``None`` for causal masks and ``"no_mask"``. For padding masks, it should be + a single tensor of ``[batch_size, 1, 1, seqlen_q]`` for self-attention, and a tuple of + two tensors of shapes ``[batch_size, 1, 1, seqlen_q]`` and ``[batch_size, 1, 1, seqlen_kv]`` + for cross-attention. For ``"arbitrary"`` mask, it should be of a shape broadcastable to + ``[batch_size, num_heads, max_seqlen_q, max_seqlen_kv]``. A ``True`` value means + the corresponding position is masked out and a ``False`` means that position is allowed to participate in attention. attn_mask_type: {'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right', 'padding_causal_bottom_right','arbitrary'}, - default = `None` + default = None type of attention mask passed into softmax operation. By default, causal masks are aligned to the top left corner of the softmax matrix. - When "`bottom_right`" is specified in the mask type, causal masks are + When ``"bottom_right"`` is specified in the mask type, causal masks are aligned to the bottom right corner. - window_size: Optional[Tuple[int, int]], default = `None` + window_size: Optional[Tuple[int, int]], default = None sliding window size for local attention. - encoder_output : Optional[torch.Tensor], default = `None` + encoder_output : Optional[torch.Tensor], default = None Output of the encoder block to be fed into the decoder block if using - `layer_type="decoder"`. + ``layer_type="decoder"``. is_first_microbatch : {True, False, None}, default = None During training using either gradient accumulation or pipeline parallelism a minibatch of data is further split @@ -668,46 +683,46 @@ def forward( * it also allows skipping gradient accumulation during the first microbatch (since it is the first gradient being produced) - checkpoint_core_attention: bool, default = `False` - If true, forward activations for core attention are recomputed + checkpoint_core_attention: bool, default = False + If ``True``, forward activations for core attention are recomputed during the backward pass in order to save memory that would otherwise be occupied to store the forward activations until backprop. - rotary_pos_emb: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = `None` + rotary_pos_emb: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = None Embeddings for query and key tensors for applying rotary position embedding. By default no input embedding is applied. - core_attention_bias_type: str, default = `no_bias` - Bias type, {`no_bias`, `pre_scale_bias`, 'post_scale_bias`, `alibi`} - core_attention_bias: Optional[torch.Tensor], default = `None` - Bias tensor for Q * K.T, shape [1, num_head, max_seqlen_q, max_seqlen_kv]. - It should be 'None' for 'no_bias' and 'alibi' bias types. - alibi_slopes: Optional[torch.Tensor], default = `None` - ALiBi slopes in FP32 and shape [nheads] or [batch_size, nheads]. - It adds a bias of (-alibi_slope * (i + seqlen_k - seqlen_q - j)) + core_attention_bias_type: str, default = "no_bias" + Bias type, {``"no_bias"``, ``"pre_scale_bias"``, ``"post_scale_bias"``, ``"alibi"``} + core_attention_bias: Optional[torch.Tensor], default = None + Bias tensor for :math:`Q \cdot K^T`, shape ``[1, num_head, max_seqlen_q, max_seqlen_kv]``. + It should be ``None`` for ``"no_bias"`` and ``"alibi"`` bias types. + alibi_slopes: Optional[torch.Tensor], default = None + ALiBi slopes in FP32 and shape ``[nheads]`` or ``[batch_size, nheads]``. + It adds a bias of ``(-alibi_slope * (i + seqlen_k - seqlen_q - j))`` to the attention score of query i and key j. - cu_seqlens_q: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (without offset) in a batch for `query_layer`, - with shape [batch_size + 1] and dtype torch.int32. - cu_seqlens_kv: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (without offset) in a batch for `key_layer` - and `value_layer`, with shape [batch_size + 1] and dtype torch.int32. - cu_seqlens_q_padded: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (with offset) in a batch for `query_layer`, - with shape [batch_size + 1] and dtype torch.int32. - cu_seqlens_kv_padded: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (with offset) in a batch for `key_layer` - and `value_layer`, with shape [batch_size + 1] and dtype torch.int32. - max_seqlen_q: Optional[int], default = `None` - Maximum sequence length in `query_layer`. - Calculated from `cu_seqlens_q` if not provided. - max_seqlen_kv: Optional[int], default = `None` - Maximum sequence length in `key_layer` and `value_layer`. - Calculated from `cu_seqlens_kv` if not provided. - fast_zero_fill: bool, default = `True` + cu_seqlens_q: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (without offset) in a batch for ``query_layer``, + with shape ``[batch_size + 1]`` and dtype torch.int32. + cu_seqlens_kv: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (without offset) in a batch for ``key_layer`` + and ``value_layer``, with shape ``[batch_size + 1]`` and dtype torch.int32. + cu_seqlens_q_padded: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (with offset) in a batch for ``query_layer``, + with shape ``[batch_size + 1]`` and dtype torch.int32. + cu_seqlens_kv_padded: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (with offset) in a batch for ``key_layer`` + and ``value_layer``, with shape ``[batch_size + 1]`` and dtype torch.int32. + max_seqlen_q: Optional[int], default = None + Maximum sequence length in ``query_layer``. + Calculated from ``cu_seqlens_q`` if not provided. + max_seqlen_kv: Optional[int], default = None + Maximum sequence length in ``key_layer`` and ``value_layer``. + Calculated from ``cu_seqlens_kv`` if not provided. + fast_zero_fill: bool, default = True Whether to set output tensors to 0 or not before use. - pad_between_seqs: Optional[bool], default = `None` - If None, inferred from qkv_format, cu_seqlens and cu_seqlens_padded. - If true, there are padding tokens between individual sequences in a packed batch. + pad_between_seqs: Optional[bool], default = None + If ``None``, inferred from qkv_format, cu_seqlens and cu_seqlens_padded. + If ``True``, there are padding tokens between individual sequences in a packed batch. """ # hidden_states: [sq, b, h] diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 0e1222c22f..a32b2d3edb 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -287,16 +287,16 @@ def _apply_rotary_pos_emb_base( Parameters ---------- - t: torch.Tensor + t : torch.Tensor Input tensor of shape `[s, b, h, d]` or `[b, s, h, d]`, on which rotary positional embedding will be applied. - freqs: torch.Tensor + freqs : torch.Tensor Rotary positional embedding tensor of shape `[s2, 1, 1, d2]` or `[s2, b, 1, d2]` and dtype 'float', with `s2 >= s` and `d2 <= d`. - tensor_format: {'sbhd', 'bshd'}, default = 'sbhd' + tensor_format : {'sbhd', 'bshd'}, default = 'sbhd' Should be `bshd` if `t` is of shape `[bs, seq, ...]`, or `sbhd` if `t` is of shape `[seq, bs, ...]`. - interleaved: bool, default = False + interleaved : bool, default = False Whether to use interleaved rotary position embedding. """ # [seq, 1, 1, dim] -> [1, seq, 1, dim] or @@ -324,7 +324,7 @@ def _get_freqs_on_this_cp_rank( """Get the position embedding on the current context parallel rank. Args: - freqs: torch.Tensor. Positional embedding tensor in shape `[s2, 1, 1, d2]`. + freqs: torch.Tensor. Positional embedding tensor of shape `[s2, 1, 1, d2]`. seqlen: int. Length of the current sequence. cp_size: int. Context parallel world size. cp_rank: int. Context parallel rank. @@ -372,29 +372,29 @@ def apply_rotary_pos_emb( Parameters ---------- - t: torch.Tensor + t : torch.Tensor Input tensor of shape `[s, b, h, d]`, `[b, s, h, d]` or `[t, h, d]`, on which rotary positional embedding will be applied. - freqs: torch.Tensor + freqs : torch.Tensor Rotary positional embedding tensor of shape `[s2, 1, 1, d2]` and dtype 'float', with `s2 >= s` and `d2 <= d`. - start_positions: torch.Tensor, default = None. + start_positions : torch.Tensor, default = None. Tokens in a sequence `i` should be applied with position encoding offset by `start_positions[i]`. If `start_positions=None`, there's no offset. - tensor_format: {'sbhd', 'bshd', 'thd'}, default = 'sbhd' + tensor_format : {'sbhd', 'bshd', 'thd'}, default = 'sbhd' is `bshd` if `t` is of shape `[bs, seq, ...]`, or `sbhd` if `t` is of shape `[seq, bs, ...]`. 'thd' is only supported when `fused` is True. - interleaved: bool, default = False + interleaved : bool, default = False Whether to use interleaved rotary position embedding. - fused: bool, default = False + fused : bool, default = False Whether to use a fused applying RoPE implementation. - cu_seqlens: torch.Tensor, default = None. + cu_seqlens : torch.Tensor, default = None. Cumulative sum of sequence lengths in a batch for `t`, with shape [b + 1] and dtype torch.int32. Only valid when `tensor_format` is 'thd'. Should be `cu_seqlens_padded` when cp_size > 1. - cp_size: int, default = 1. + cp_size : int, default = 1. Context parallel world size. Only valid when `tensor_format` is 'thd' and `fused` is True. - cp_rank: int, default = 0. + cp_rank : int, default = 0. Context parallel rank. Only valid when `tensor_format` is 'thd' and `fused` is True. """ assert ( @@ -492,32 +492,32 @@ def apply_fused_qkv_rotary_pos_emb( Parameters ---------- - qkv: torch.Tensor + qkv : torch.Tensor Input tensor of shape `[s, b, h, d]` or `[b, s, h, d]`, on which rotary positional embedding will be applied. This tensor has q, k, v concatenated along the last dimension. - q_freqs: torch.Tensor + q_freqs : torch.Tensor Rotary positional embedding Q tensor of shape `[s2, 1, 1, d2]` and dtype 'float', with `s2 >= s` and `d2 <= d`. - k_freqs: torch.Tensor + k_freqs : torch.Tensor Rotary positional embedding K tensor of shape `[s2, 1, 1, d2]` and dtype 'float', with `s2 >= s` and `d2 <= d`. - qkv_split_arg_list: List[int] + qkv_split_arg_list : List[int] List of integers that specify the split of the qkv tensor. The list should have 3 elements, the first element is the number of elements in the q tensor, the second element is the number of elements in the k tensor, and the third element is the number of elements in the v tensor. The sum of the elements in the list should be equal to the last dimension of the qkv tensor. - start_positions: torch.Tensor, default = None. + start_positions : torch.Tensor, default = None. Tokens in a sequence `i` should be applied with position encoding offset by `start_positions[i]`. If `start_positions=None`, there's no offset. - tensor_format: {'sbhd', 'bshd'}, default = 'sbhd' + tensor_format : {'sbhd', 'bshd'}, default = 'sbhd' is `bshd` if `qkv` is of shape `[bs, seq, ...]`, or `sbhd` if `qkv` is of shape `[seq, bs, ...]`. - interleaved: bool, default = False + interleaved : bool, default = False Whether to use interleaved rotary position embedding. - cp_size: int, default = 1. + cp_size : int, default = 1. Context parallel world size. - cp_rank: int, default = 0. + cp_rank : int, default = 0. Context parallel rank. """ diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index e55ea2a54a..88c223eb46 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -146,89 +146,89 @@ def fused_attn_fwd( Parameters ---------- - is_training: bool + is_training : bool if True, runs training and produces auxiliary tensors aux_ctx_tensors for the backward; if False, runs inference and doesn't produce aux_ctx_tensors - max_seqlen_q: int + max_seqlen_q : int max sequence length for Q, used for padding; may be larger than max(seqlens_q), seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - max_seqlen_kv: int + max_seqlen_kv : int max sequence length for K and V, used for padding; may be larger than max(seqlens_kv), seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - cu_seqlens_q: torch.Tensor + cu_seqlens_q : torch.Tensor cumulative sequence lengths for Q; shape [batch_size + 1] - cu_seqlens_kv: torch.Tensor + cu_seqlens_kv : torch.Tensor cumulative sequence lengths for K and V; shape [batch_size + 1] - q: torch.Tensor + q : torch.Tensor input tensor Q; shape sbhd, bshd or thd (see `qkv_layout` for details) - k: torch.Tensor + k : torch.Tensor input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) - v: torch.Tensor + v : torch.Tensor input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - fake_dtype: tex.DType + 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 - fused_attention_backend: tex.NVTE_Fused_Attn_Backend + fused_attention_backend : tex.NVTE_Fused_Attn_Backend please see FusedAttention module for details on supported backends. - attn_bias: torch.Tensor, default = None + attn_bias : torch.Tensor, default = None input tensor Bias when attn_bias_type is "pre_scale_bias" or "post_scale_bias"; shape [1, num_heads, max_seqlen_q, max_seqlen_kv], same data type as q, k and v - cu_seqlens_q_padded: torch.Tensor, default = None + cu_seqlens_q_padded : torch.Tensor, default = None cumulative sequence offsets for Q; shape [batch_size + 1] - cu_seqlens_kv_padded: torch.Tensor, default = None + cu_seqlens_kv_padded : torch.Tensor, default = None cumulative sequence offsets for KV; shape [batch_size + 1] - page_table_k: torch.Tensor, default = None + page_table_k : torch.Tensor, default = None page table for K cache; shape [batch_size, max_pages_per_seq_k] - page_table_v: torch.Tensor, default = None + page_table_v : torch.Tensor, default = None page table for V cache; shape [batch_size, max_pages_per_seq_v] - s_quantizer: Quantizer, default = None + s_quantizer : Quantizer, default = None Quantizer object for the intermediate value S. - o_quantizer: Quantizer, default = None + o_quantizer : Quantizer, default = None Quantizer object for the output of the attention. - attn_scale: float, default = None + 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 : 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 - fast_zero_fill: bool, default = True + fast_zero_fill : bool, default = True if True, initializes the output tensor O to zero using the fast filling method; if False, uses PyTorch's .fill_() method - qkv_layout: str, default = "sbh3d" + qkv_layout : str, default = "sbh3d" layout of Q, K and V; {"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"} - attn_bias_type: str, default = "no_bias" + 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" + attn_mask_type : str, default = "padding" type of the attention mask; {"padding", "causal", "padding_causal", "no_mask"} - softmax_type: str, default = "vanilla" + softmax_type : str, default = "vanilla" type of the attention softmax; {"vanilla", "off-by-one", "learnable"} - window_size: Tuple[int, int], default = (-1, -1) + window_size : Tuple[int, int], default = (-1, -1) sliding window size for local attention, where query at position i attends to keys in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding window and causal mask specifically. - rng_gen: torch.Generator, default = None + rng_gen : torch.Generator, default = None random number generator; if None, uses the default CUDA generator from PyTorch; otherwise, uses rng_gen - softmax_offset: torch.Tensor, default = None - softmax offset tensor in shape [1, h_q, 1, 1]. + softmax_offset : torch.Tensor, default = None + softmax offset tensor of shape [1, h_q, 1, 1]. See softmax_type in DotProductAttention for details. - return_max_logit: bool, default = False + return_max_logit : bool, default = False whether to return the maximum attention score - cuda_graph: bool, default = False + cuda_graph : bool, default = False whether or not cuda graph capture is enabled. Returns ---------- - o: torch.Tensor + o : torch.Tensor output tensor O, of the attention calculation; same data type as Q, K and V; same shape as Q - aux_ctx_tensors: List[torch.Tensor] + aux_ctx_tensors : List[torch.Tensor] auxiliary output tensors used for the backward; if is_training is True, aux_ctx_tensors = [softmax-related tensors, rng_state] if is_training is False, aux_ctx_tensors = None @@ -252,7 +252,7 @@ def fused_attn_fwd( rng_state: torch.Tensor, optional, if backend is not F16_max512_seqlen 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 + max_logit : if return_max_logit = True, shape [h] and same data type as O; otherwise None """ if attn_scale is None: @@ -377,89 +377,89 @@ def fused_attn_bwd( Parameters ---------- - max_seqlen_q: int + max_seqlen_q : int max sequence length for Q, used for padding; may be larger than max(seqlens_q), seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] - max_seqlen_kv: int + max_seqlen_kv : int max sequence length for K and V, used for padding; may be larger than max(seqlens_kv), seqlens_kv = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] - cu_seqlens_q: torch.Tensor + cu_seqlens_q : torch.Tensor cumulative sequence lengths for Q; shape [batch_size + 1] - cu_seqlens_kv: torch.Tensor + cu_seqlens_kv : torch.Tensor cumulative sequence lengths for K and V; shape [batch_size + 1] - q: torch.Tensor + q : torch.Tensor input tensor Q; shape sbhd, bshd or thd (see `qkv_layout` for details) - k: torch.Tensor + k : torch.Tensor input tensor K; shape sbhd, bshd or thd (see `qkv_layout` for details) - v: torch.Tensor + v : torch.Tensor input tensor V; shape sbhd, bshd or thd (see `qkv_layout` for details) - o: torch.Tensor + o : torch.Tensor input tensor O (output of forward); same data type as Q, K and V; same shape as Q - d_o: torch.Tensor + d_o : torch.Tensor input tensor dO (gradient of O); same data type as Q, K and V; same shape as Q - fake_dtype: tex.DType + 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 + dqkv_dtype : tex.DType data type of dQ, dK and dV; in tex.DType, not torch.dtype - aux_ctx_tensors: List[torch.Tensor] + 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] - fused_attention_backend: tex.NVTE_Fused_Attn_Backend + 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 + cu_seqlens_q_padded : torch.Tensor, default = None cumulative sequence offsets for Q; shape [batch_size + 1] - cu_seqlens_kv_padded: torch.Tensor, default = None + cu_seqlens_kv_padded : torch.Tensor, default = None cumulative sequence offsets for KV; shape [batch_size + 1] - s_quantizer: Quantizer, default = None + s_quantizer : Quantizer, default = None Quantizer object for the intermediate value S. - dp_quantizer: Quantizer, default = None + dp_quantizer : Quantizer, default = None Quantizer object for the intermediate value dP. - dqkv_quantizer: Quantizer, default = None + dqkv_quantizer : Quantizer, default = None Quantizer object for the output values of the fused_attn_bwd. - dropout: float, default = 0.0 + 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 - fast_zero_fill: bool, default = True + fast_zero_fill : bool, default = True if True, initializes the output tensor O to zero using the fast filling method; if False, uses PyTorch's .fill_() method - qkv_layout: str, default = "sbh3d" + qkv_layout : str, default = "sbh3d" layout of Q, K and V; {"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"} - attn_bias_type: str, default = "no_bias" + 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" + attn_mask_type : str, default = "padding" type of the attention mask; {"padding", "causal", "padding_causal", "no_mask"} - softmax_type: str, default = "vanilla" + softmax_type : str, default = "vanilla" type of the attention softmax; {"vanilla", "off-by-one", "learnable"} - window_size: Tuple[int, int], default = (-1, -1) + window_size : Tuple[int, int], default = (-1, -1) sliding window size for local attention, where query at position i attends to keys in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding window and causal mask specifically. - deterministic: bool, default = False + deterministic : bool, default = False whether to execute the backward pass with deterministic behaviours. - cuda_graph: bool, default = False + cuda_graph : bool, default = False whether or not cuda graph capture is enabled. Returns ---------- - d_q: torch.Tensor + d_q : torch.Tensor gradient tensor of Q; same data type and shape as Q - d_k: torch.Tensor + d_k : torch.Tensor gradient tensor of K; same data type and shape as K - d_v: torch.Tensor + d_v : torch.Tensor gradient tensor of V; same data type and shape as V - d_bias: torch.Tensor, optional + d_bias : torch.Tensor, optional gradient tensor of Bias when attn_bias_type is "pre_scale_bias" or "post_scale_bias"; same data type and shape as Bias - d_softmax_offset: torch.Tensor, optional - gradient tensor of softmax offset in shape [1, h_q, 1, 1]. + d_softmax_offset : torch.Tensor, optional + gradient tensor of softmax offset of shape [1, h_q, 1, 1]. See softmax_type in DotProductAttention for details. """ if attn_scale is None: diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 9e6d577235..58ed063066 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -657,60 +657,64 @@ def get_cpu_offload_context( Parameters ---------- - enabled: bool, default = `False` + enabled : bool, default = False When set to True, CPU Offloading functionality is enabled. - num_layers: int, default = 1 + num_layers : int, default = 1 Determines the number of layers you want to offload activations/weights for. - model_layers: int, default = 1 + model_layers : int, default = 1 Number of layers in the model that will be used under this context. - offload_activations: bool, default = `True` + offload_activations : bool, default = True Deprecated. - offload_weights: bool, default = `True` + offload_weights : bool, default = True Deprecated. - double_buffering: bool, default = `False` + double_buffering : bool, default = False Deprecated. - retain_pinned_cpu_buffers: bool, default = `False` + retain_pinned_cpu_buffers : bool, default = False If True, the pinned CPU buffers are retained after offloading and reused for the next iteration. It is useful for cuda graphs capture. - manual_synchronization: bool, default = `False` + manual_synchronization : bool, default = False If True, the synchronization is done manually by the user. Additional argument manual_controller is returned. See more in manual control section. - offload_stream: torch.cuda.Stream, default = `None` + offload_stream : torch.cuda.Stream, default = None If provided, the offload stream is used for offloading and reloading. Otherwise, a new stream is allocated internally. It can be other than None only if manual_synchronization is True. - Manual synchronization - ---------- + Notes + ----- + **Manual synchronization:** + By default, layers are offloaded/reloaded asynchronously with respect to the current forward/backward stream with predefined synchronization, to ensure that activation memory usage is equal to - `(num_layers - num_offloaded_layers) * T`, where `T` is the memory footprint of a layer. + ``(num_layers - num_offloaded_layers) * T``, where ``T`` is the memory footprint of a layer. - For more control over the offloading and reloading process, you can set `manual_synchronization=True`. - In this case, an additional argument, `manual_controller`, is returned. + For more control over the offloading and reloading process, you can set ``manual_synchronization=True``. + In this case, an additional argument, ``manual_controller``, is returned. - The `manual_controller` provides the following methods: - - `start_offload_layer(layer_id: int)` - - `release_activation_forward_gpu_memory(layer_id: int)` - - `start_reload_layer(layer_id: int)` + The ``manual_controller`` provides the following methods: + - ``start_offload_layer(layer_id: int)`` + - ``release_activation_forward_gpu_memory(layer_id: int)`` + - ``start_reload_layer(layer_id: int)`` If none of these methods are invoked for a given layer, that layer will not be offloaded or reloaded. - If `start_offload_layer()` is called for a layer, offload copies for that layer begin asynchronously on the offload stream. + If ``start_offload_layer()`` is called for a layer, offload copies for that layer begin asynchronously on the offload stream. Since GPU activations must be kept in memory until the copy is finished, pointers to all activations are stored. - To release this memory, you need to call `release_activation_forward_gpu_memory(layer_id)`. + To release this memory, you need to call ``release_activation_forward_gpu_memory(layer_id)``. This method makes the current stream wait for an event recorded on the offload stream after all tensors from the layer have been offloaded. - The `start_reload_layer()` method is used to start reloading a layer. - Each tensor reload is awaited to finish before `tensor_pop()` for that tensor is called on the current stream. + The ``start_reload_layer()`` method is used to start reloading a layer. + Each tensor reload is awaited to finish before ``tensor_pop()`` for that tensor is called on the current stream. - You can provide an `offload_stream` to be used for offload and reload operations. + You can provide an ``offload_stream`` to be used for offload and reload operations. This allows for more detailed synchronization, such as delaying the start of offloading. - Example: + **Example:** + .. code-block:: python + offload_stream = torch.cuda.Stream() cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context( enabled=True, model_layers=num_layers, manual_synchronization=True, offload_stream=offload_stream) @@ -732,10 +736,10 @@ def get_cpu_offload_context( for i in range(num_layers): out[i].sum().backward() - V1 code path - ---------- + **V1 code path:** + If you want to use the v1 code path for offloading, - please set the environment variable NVTE_CPU_OFFLOAD_V1 to 1. + please set the environment variable ``NVTE_CPU_OFFLOAD_V1`` to 1. """ if NVTE_CPU_OFFLOAD_V1: diff --git a/transformer_engine/pytorch/cpu_offload_v1.py b/transformer_engine/pytorch/cpu_offload_v1.py index 9f904864ab..e79e37b019 100644 --- a/transformer_engine/pytorch/cpu_offload_v1.py +++ b/transformer_engine/pytorch/cpu_offload_v1.py @@ -685,18 +685,18 @@ def get_cpu_offload_context( Parameters ---------- - enabled: bool, default = `False` + enabled : bool, default = `False` When set to True, CPU Offloading functionality is enabled. - num_layers: int, default = 1 + num_layers : int, default = 1 Determines the number of transformer layers you want to offload activations/weights for. - model_layers: int, default = 1 + model_layers : int, default = 1 Number of layers in the model that will be used under this context. - offload_activations: bool, default = `True` + offload_activations : bool, default = `True` When set to `True`, offloads the activations for the TE layer. - offload_weights: bool, default = `True` + offload_weights : bool, default = `True` When set to `True`, offloads the weights for the TE layer. - double_buffering: bool, default = `False` + double_buffering : bool, default = `False` When set to `True`, uses double buffering for offloading. """ diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 076dbec0dc..30002cdbfd 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -4,6 +4,9 @@ """Cross Entropy Loss API""" +from typing import Optional +import warnings + import torch import transformer_engine.pytorch.triton.cross_entropy as triton_cross_entropy @@ -23,7 +26,7 @@ class CrossEntropyFunction(torch.autograd.Function): @staticmethod def forward( ctx, - _input, + inp, target, label_smoothing=0.0, reduce_loss=False, @@ -37,7 +40,7 @@ def forward( Parameters: ctx : The context object. - _input (tensor): The input tensor of shape (B, SQ, V) or (SQ, B, V) where B is batch size, SQ is sequence length, V is vocab size. + inp (tensor): The input tensor of shape (B, SQ, V) or (SQ, B, V) where B is batch size, SQ is sequence length, V is vocab size. target (tensor): The target tensor of shape (B,SQ) or (SQ, B) where each value is in [0, V-1]. label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing. reduce_loss (bool): If true, returns the averaged loss across the B*SQ dimension. @@ -47,8 +50,8 @@ def forward( Returns: tensor: The computed loss. """ - loss, _input = triton_cross_entropy.cross_entropy_forward( - _input, + loss, inp = triton_cross_entropy.cross_entropy_forward( + inp, target, label_smoothing, reduce_loss, @@ -56,7 +59,7 @@ def forward( ignore_idx, ) - ctx.save_for_backward(_input.detach()) + ctx.save_for_backward(inp.detach()) ctx.is_cg_capturable = is_cg_capturable return loss @@ -72,12 +75,10 @@ def backward(ctx, grad_output): Returns: tuple: A tuple with the gradients with respect to the inputs. The elements are tensors or None. """ - (_input,) = ctx.saved_tensors - _input = triton_cross_entropy.cross_entropy_backward( - _input, grad_output, ctx.is_cg_capturable - ) + (inp,) = ctx.saved_tensors + inp = triton_cross_entropy.cross_entropy_backward(inp, grad_output, ctx.is_cg_capturable) return ( - _input, + inp, None, None, None, @@ -87,4 +88,65 @@ def backward(ctx, grad_output): ) -parallel_cross_entropy = CrossEntropyFunction.apply +def parallel_cross_entropy( + inp: torch.Tensor, + target: torch.Tensor, + label_smoothing: float = 0.0, + reduce_loss: bool = False, + dist_process_group: Optional[torch.distributed.ProcessGroup] = None, + ignore_idx: int = -100, + is_cg_capturable: bool = False, + *, + _input: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """ + Cross Entropy loss with optional distributed reduction. + + The input tensor can be in BF16/FP32, the loss and gradient calculation happens in + FP32 only. The returned loss is always in FP32, the input gradients are upcasted + to the datatype of the input. + + If ``dist_process_group`` is passed for distributed loss calculation, the input to each + distributed rank should be ``(*, V/world_size)``. Note that each of the ranks should + get equal shards along the V dimension. + + Parameters + ---------- + inp : torch.Tensor + The input tensor of shape ``(B, SQ, V)`` or ``(SQ, B, V)`` where B is batch size, + SQ is sequence length, V is vocab size. + target : torch.Tensor + The target tensor of shape ``(B, SQ)`` or ``(SQ, B)`` where each value is in ``[0, V-1]``. + label_smoothing : float, default = 0.0 + The amount of smoothing when computing the loss, where 0.0 means no smoothing. + reduce_loss : bool, default = False + If True, returns the averaged loss across the B*SQ dimension. + dist_process_group : torch.distributed.ProcessGroup, default = None + The distributed process group the loss computation is split across, None if on 1 device. + ignore_idx : int, default = -100 + The index for which loss and gradients are made to zero. + is_cg_capturable : bool, default = False + Whether the operation is CUDA graph capturable. + + Returns + ------- + torch.Tensor + The computed loss. + """ + # Handle backward compatibility with _input parameter + if _input is not None: + warnings.warn( + "The '_input' parameter is deprecated. Please use 'inp' instead.", + FutureWarning, + ) + inp = _input + + return CrossEntropyFunction.apply( + inp, + target, + label_smoothing, + reduce_loss, + dist_process_group, + ignore_idx, + is_cg_capturable, + ) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 8ce54d7f64..5284b297e2 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -30,7 +30,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.triton.pad import pad_columnwise_scale_inv -from . import torch_version +from .torch_version import torch_version from .utils import ( is_non_tn_fp8_gemm_supported, safely_set_viewless_tensor_data, @@ -642,18 +642,18 @@ def checkpoint( Parameters ---------- - function: Callable + function : Callable pytorch module used to run the forward and backward passes using the specified :attr:`args` and :attr:`kwargs`. - distribute_saved_activations: bool, default = False - if set to `True` and `use_reentrant=True`, first tensor argument is distributed - across the specified tensor parallel group (`tp_group`) before saving it for the - backward pass. This has no effect when `use_reentrant=False`. - get_rng_state_tracker: `Callable`, default = None - python callable which returns an instance of :func:`CudaRNGStatesTracker`. + distribute_saved_activations : bool, default = False + if set to ``True`` and ``use_reentrant=True``, first tensor argument is distributed + across the specified tensor parallel group (``tp_group``) before saving it for the + backward pass. This has no effect when ``use_reentrant=False``. + get_rng_state_tracker : Callable, default = None + python callable which returns an instance of :class:`CudaRNGStatesTracker`. tp_group : ProcessGroup, default = None - tensor parallel process group. Used only when `distribute_saved_activations=True` - and `use_reentrant=True`. If `None`, it falls back to the default group. + tensor parallel process group. Used only when ``distribute_saved_activations=True`` + and ``use_reentrant=True``. If ``None``, it falls back to the default group. use_reentrant : bool, default = True perform checkpointing in reentrant mode. args : tuple @@ -778,8 +778,8 @@ class CudaRNGStatesTracker: For model parallelism, multiple RNG states need to simultaneously exist in order to execute operations in or out of the model parallel region. This class keeps track of the various RNG states and provides utility methods to maintain them and - execute parts of the model under a given RNG setting. Using the `add` method, a - cuda rng state is initialized based on the input `seed` and is assigned to `name`. + execute parts of the model under a given RNG setting. Using the :meth:`add` method, a + cuda rng state is initialized based on the input ``seed`` and is assigned to ``name``. Later, by forking the rng state, we can perform operations and return to our starting cuda state. """ @@ -812,7 +812,9 @@ def set_states(self, states: Dict[str, torch.Tensor]) -> None: Set the rng states. For efficiency purposes, we do not check the size of seed for compatibility. - states: Dict[str, torch.Tensor] + Parameters + ---------- + states : Dict[str, torch.Tensor] A mapping from string names to RNG states. """ self.states_ = states @@ -821,9 +823,11 @@ def add(self, name: str, seed: int) -> None: """ Adds a new RNG state. - name: str + Parameters + ---------- + name : str string identifier for the RNG state. - seed: int + seed : int PyTorch seed for the RNG state. """ # Check seed is not already used. @@ -857,7 +861,9 @@ def fork(self, name: str = "model-parallel-rng"): Fork the cuda rng state, perform operations, and exit with the original state. - name: str + Parameters + ---------- + name : str string identifier for the RNG state. """ # Check if we have added the state @@ -2003,7 +2009,7 @@ def prepare_te_modules_for_fsdp(fsdp_root: torch.nn.Module) -> None: Parameters ---------- - fsdp_root: torch.nn.Module + fsdp_root : torch.nn.Module FSDP-wrapped root module that may contain FSDP-wrapped TE modules. """ assert isinstance(fsdp_root, FSDP), "Root module must be FSDP-wrapped." diff --git a/transformer_engine/pytorch/export.py b/transformer_engine/pytorch/export.py index f75271e2cc..a86f8ee58c 100644 --- a/transformer_engine/pytorch/export.py +++ b/transformer_engine/pytorch/export.py @@ -28,7 +28,7 @@ def onnx_export(enabled: bool = False) -> Generator[None, None, None]: Parameters ---------- - enabled: bool, default = `False` + enabled : bool, default = False whether or not to enable export """ diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 3e515eecd5..1baa67414d 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1016,38 +1016,38 @@ def make_graphed_callables( Positional arguments to callable(s). num_warmup_iters: int, default = 3 Number of warmup iterations. - allow_unused_input: bool, default = `False` + allow_unused_input: bool, default = False Whether to handle case where callable inputs and outputs are disconnected in compute graph. sample_kwargs: (tuple of) dict, optional Keyword arguments to callable(s) - pool: (tuple of) int, default = `None`, optional + pool: (tuple of) int, default = None, optional An instance returned from function `torch.cuda.graph_pool_handle` that hints this graph may share memory with the indicated pool. - retain_graph_in_backward: bool, default = `False` + retain_graph_in_backward: bool, default = False Whether to set retain_graph=True in backward graph capture. - _reuse_graph_input_output_buffers: bool, default = `False` + _reuse_graph_input_output_buffers: bool, default = False Reduce memory usage by reusing input/output data buffers between graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. when `_order` is provided. All callables in `modules` are assumed to have inputs and outputs with the same dtype and shape. - Quantization related parameters - ---------------------- - enabled: (tuple of) bool, default = `False` + Quantization parameters + ----------------------- + enabled: (tuple of) bool, default = False whether or not to enable low precision quantization (FP8/FP4). If tuple, the length must match the number of modules. - calibrating: bool, default = `False` + calibrating: bool, default = False calibration mode allows collecting statistics such as amax and scale data of quantized tensors even when executing without quantization enabled. This is useful for saving an inference ready checkpoint while training using a higher precision. - recipe: recipe.Recipe, default = `None` + recipe: recipe.Recipe, default = None recipe used for low precision quantization. - amax_reduction_group: torch._C._distributed_c10d.ProcessGroup, default = `None` + amax_reduction_group: torch._C._distributed_c10d.ProcessGroup, default = None distributed group over which amaxes for the quantized tensors are reduced at the end of each training step. - cache_quantized_params: bool, default = `False` + cache_quantized_params: bool, default = False Whether or not to cache quantized weights across microbatches. if set to `True`, the `is_first_microbatch` boolean argument must be passed into the forward method for TransformerEngine modules. When storing primary weights in low precision diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index f0f77621e5..e9a65a72ff 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -8,7 +8,7 @@ from typing import Callable, Optional, Tuple import torch -from . import torch_version +from .torch_version import torch_version from .export import is_in_onnx_export_mode from .utils import gpu_autocast_ctx diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 40589a82d5..acf9233281 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -20,7 +20,6 @@ from torch.distributed.tensor import DTensor import transformer_engine_torch as tex -from transformer_engine.common.recipe import Recipe from ._common import _ParameterInitMeta, noop_cat from ..quantization import ( @@ -104,55 +103,55 @@ def initialize_ub( ) -> None: r""" Initialize the Userbuffers communicator for overlapping tensor-parallel communications with - GEMM compute in te.Linear, te.LayerNormLinear and te.LayerNormMLP modules. + GEMM compute in ``te.Linear``, ``te.LayerNormLinear`` and ``te.LayerNormMLP`` modules. Parameters ---------- shape : list shape of the communication buffer, typically set to be the same as the global shape of - the input tensor to a te.TransformerLayer forward pass, with the sequence and batch - dimensions collapsed together -- i.e.: `(sequence_length * batch_size, hidden_size)` + the input tensor to a ``te.TransformerLayer`` forward pass, with the sequence and batch + dimensions collapsed together -- i.e.: ``(sequence_length * batch_size, hidden_size)`` tp_size : int number of GPUs in the tensor-parallel process group use_fp8 : bool = False allocate the communication buffer for FP8 GEMM inputs/outputs. - DEPRECATED: Please use `quantization_modes` instead. + DEPRECATED: Please use ``quantization_modes`` instead. quantization_modes : List[UserBufferQuantizationMode] = None if a list of UserBufferQuantizationMode is provided, a UB communicator is created for each quantization setting in the list. - falls back to the legacy `use_fp8` parameter if `None` is provided. + falls back to the legacy ``use_fp8`` parameter if ``None`` is provided. dtype : torch.dtype = torch.bfloat16 - non-FP8 data type of the communication buffer when `use_fp8 = False` - ub_cfgs: dict = None - Configuration dictionary with the structure - ``` - { - : { - "method": <"ring_exchange" or "pipeline">, - "is_reduce_scatter": bool, - "num_sm": int, - "cga_size": int, - "set_sm_margin": bool, - "num_splits": int, - "aggregate": bool, - "atomic_gemm": bool, - "use_ce": bool, - "fp8_buf": bool, - } - } - ``` - for `te.TransformerLayer` GEMM layers in `["qkv_fprop", "qkv_dgrad", "qkv_wgrad", + non-FP8 data type of the communication buffer when ``use_fp8 = False`` + ub_cfgs : dict = None + Configuration dictionary with the structure:: + + { + : { + "method": <"ring_exchange" or "pipeline">, + "is_reduce_scatter": bool, + "num_sm": int, + "cga_size": int, + "set_sm_margin": bool, + "num_splits": int, + "aggregate": bool, + "atomic_gemm": bool, + "use_ce": bool, + "fp8_buf": bool, + } + } + + for ``te.TransformerLayer`` GEMM layers in ``["qkv_fprop", "qkv_dgrad", "qkv_wgrad", "proj_fprop", "proj_dgrad", "proj_wgrad", "fc1_fprop", "fc1_dgrad", "fc2_dgrad", - "fc2_fprop", "fc2_wgrad"]`. - a list may be provided to specify different overlap configurations for different the quantization settings in `quantization_modes` + "fc2_fprop", "fc2_wgrad"]``. + a list may be provided to specify different overlap configurations for different the quantization settings in ``quantization_modes`` bootstrap_backend : str = None - `torch.distributed` communication backend for the all-gather, broadcast and + ``torch.distributed`` communication backend for the all-gather, broadcast and barrier collectives during Userbuffers initialization. Not all backends are valid for every cluster configuration and distributed launch method even if they are available in PyTorch. When left unset, the initialization prefers to use the MPI backend, falling back first on Gloo and then NCCL if MPI is - not available. Setting `NVTE_UB_WITH_MPI=1` when building TE overrides this + not available. Setting ``NVTE_UB_WITH_MPI=1`` when building TE overrides this option and always initializes Userbuffers with direct MPI calls in C++, - which also requires `MPI_HOME=/path/to/mpi/root` to be set at compile time. + which also requires ``MPI_HOME=/path/to/mpi/root`` to be set at compile time. """ if not tex.device_supports_multicast(): assert bool(int(os.getenv("UB_SKIPMC", "0"))), ( @@ -951,7 +950,7 @@ def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> N Parameters ---------- - tp_group : ProcessGroup, default = `None` + tp_group : ProcessGroup, default = None tensor parallel process group. """ self.tp_group = tp_group @@ -1345,7 +1344,7 @@ def get_weight_workspace( workspace is being constructed or updated. cache_name: str, optional Key for caching. - update_workspace: bool, default = `True` + update_workspace: bool, default = True Update workspace with values from `tensor`. skip_update_flag: torch.Tensor, optional GPU flag to skip updating the workspace. Take precedence diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index da3ead631b..004c95c372 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -537,14 +537,14 @@ class GroupedLinear(TransformerEngineBaseModule): size of each input sample. out_features : int size of each output sample. - bias : bool, default = `True` - if set to `False`, the layer will not learn an additive bias. - init_method : Callable, default = `None` - used for initializing weights in the following way: `init_method(weight)`. - When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`. - get_rng_state_tracker : Callable, default = `None` + bias : bool, default = True + if set to ``False``, the layer will not learn an additive bias. + init_method : Callable, default = None + used for initializing weights in the following way: ``init_method(weight)``. + When set to ``None``, defaults to ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + get_rng_state_tracker : Callable, default = None used to get the random number generator state tracker for initializing weights. - rng_tracker_name : str, default = `None` + rng_tracker_name : str, default = None the param passed to get_rng_state_tracker to get the specific rng tracker. device : Union[torch.device, str], default = "cuda" The device on which the parameters of the model will be allocated. It is the user's @@ -553,34 +553,36 @@ class GroupedLinear(TransformerEngineBaseModule): Optimization parameters ----------------------- - fuse_wgrad_accumulation : bool, default = 'False' - if set to `True`, enables fusing of creation and accumulation of + fuse_wgrad_accumulation : bool, default = False + if set to ``True``, enables fusing of creation and accumulation of the weight gradient. When enabled, it is assumed that the weights - have an additional `main_grad` attribute (used instead of the - regular `grad`) which is a pre-allocated buffer of the correct + have an additional ``main_grad`` attribute (used instead of the + regular ``grad``) which is a pre-allocated buffer of the correct size to accumulate gradients in. This argument along with weight tensor having attribute 'overwrite_main_grad' set to True - will overwrite `main_grad` instead of accumulating. - return_bias : bool, default = `False` - when set to `True`, this module will not apply the additive bias itself, but + will overwrite ``main_grad`` instead of accumulating. + return_bias : bool, default = False + when set to ``True``, this module will not apply the additive bias itself, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when the bias addition can be fused to subsequent operations. - params_dtype : torch.dtype, default = `torch.get_default_dtype()` + params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. - delay_wgrad_compute : bool, default = `False` + delay_wgrad_compute : bool, default = False Whether to delay weight gradient computation - save_original_input : bool, default = `False` - If set to `True`, always saves the original input tensor rather than the + save_original_input : bool, default = False + If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. Cannot work with FP8 DelayedScaling recipe. - Note: GroupedLinear doesn't really handle the TP communications inside. The `tp_size` and - `parallel_mode` are used to determine the shapes of weights and biases. - The TP communication should be handled in the dispatch and combine stages of MoE models. + Notes + ----- + GroupedLinear doesn't really handle the TP communications inside. The ``tp_size`` and + ``parallel_mode`` are used to determine the shapes of weights and biases. + The TP communication should be handled in the dispatch and combine stages of MoE models. """ def __init__( diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py index 6d13544e4f..52802c618c 100644 --- a/transformer_engine/pytorch/module/layernorm.py +++ b/transformer_engine/pytorch/module/layernorm.py @@ -28,33 +28,30 @@ class LayerNorm(_LayerNormOp): Parameters ---------- - normalized_shape: int or iterable of int + normalized_shape : int or iterable of int Inner dimensions of input tensor eps : float, default = 1e-5 A value added to the denominator of layer normalization for numerical stability - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype zero_centered_gamma : bool, default = 'False' - If `True`, the :math:`\gamma` parameter is initialized to zero + If ``True``, the :math:`\gamma` parameter is initialized to zero and the calculation changes to .. math:: y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta - sm_margin: int or dict, default = 0 + sm_margin : int or dict, default = 0 Number of SMs to exclude when launching CUDA kernels. This helps overlap with other kernels, e.g. communication kernels. For more fine-grained control, provide a dict with the SM - margin at each compute stage ("forward", "backward", - "inference"). - - Legacy - ------ - sequence_parallel: bool - Set a bool attr named `sequence_parallel` in the parameters. + margin at each compute stage (``"forward"``, ``"backward"``, + ``"inference"``). + sequence_parallel : bool + **Legacy parameter.** Set a bool attr named ``sequence_parallel`` in the parameters. This is custom logic for Megatron-LM integration. """ diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 3adbbc22e9..667c199c49 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -15,7 +15,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch import torch_version +from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, @@ -1045,20 +1045,20 @@ class LayerNormLinear(TransformerEngineBaseModule): size of each output sample. eps : float, default = 1e-5 a value added to the denominator of layer normalization for numerical stability. - bias : bool, default = `True` - if set to `False`, the layer will not learn an additive bias. + bias : bool, default = True + if set to ``False``, the layer will not learn an additive bias. normalization : { 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm' type of normalization applied. - init_method : Callable, default = `None` - used for initializing weights in the following way: `init_method(weight)`. - When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`. - return_layernorm_output : bool, default = `False` - if set to `True`, output of layernorm is returned from the forward + init_method : Callable, default = None + used for initializing weights in the following way: ``init_method(weight)``. + When set to ``None``, defaults to ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + return_layernorm_output : bool, default = False + if set to ``True``, output of layernorm is returned from the forward together with the output of the linear transformation. Example use case: residual connection for transformer module is taken post layernorm. - return_layernorm_output_gathered : bool, default = `False` - if set to `True`, output of layernorm is returned after the all + return_layernorm_output_gathered : bool, default = False + if set to ``True``, output of layernorm is returned after the all gather operation. Ignored if return_layernorm_output is False. Example use case: with sequence parallel, input to residual connection for transformer module (e.g. LoRA) will need to be gathered. @@ -1069,10 +1069,10 @@ class LayerNormLinear(TransformerEngineBaseModule): they are used to make the names of equally-sized parameters. If a dict (preferably an OrderedDict) is provided, the keys are used as names and values as split sizes along dim 0. The resulting parameters will have - names that end in `_weight` or `_bias`, so trailing underscores are + names that end in ``_weight`` or ``_bias``, so trailing underscores are stripped from any provided names. zero_centered_gamma : bool, default = 'False' - if set to 'True', gamma parameter in LayerNorm is initialized to 0 and + if set to ``'True'``, gamma parameter in LayerNorm is initialized to 0 and the LayerNorm formula changes to .. math:: @@ -1082,53 +1082,53 @@ class LayerNormLinear(TransformerEngineBaseModule): 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 forward pass. - name: str, default = `None` + name : str, default = None name of the module, currently used for debugging purposes. Parallelism parameters ---------------------- - sequence_parallel : bool, default = `False` - if set to `True`, uses sequence parallelism. - tp_group : ProcessGroup, default = `None` + sequence_parallel : bool, default = False + if set to ``True``, uses sequence parallelism. + tp_group : ProcessGroup, default = None tensor parallel process group. tp_size : int, default = 1 used as TP (tensor parallel) world size when TP groups are not formed during initialization. In this case, users must call the - `set_tensor_parallel_group(tp_group)` method on the initialized module before the + ``set_tensor_parallel_group(tp_group)`` method on the initialized module before the forward pass to supply the tensor parallel group needed for tensor and sequence parallel collectives. - parallel_mode : {None, 'column', 'row'}, default = `None` + parallel_mode : {None, 'column', 'row'}, default = None used to decide whether this Linear layer is Column Parallel Linear or Row Parallel Linear as described `here `_. - When set to `None`, no communication is performed. + When set to ``None``, no communication is performed. Optimization parameters ----------------------- fuse_wgrad_accumulation : bool, default = 'False' - if set to `True`, enables fusing of creation and accumulation of + if set to ``True``, enables fusing of creation and accumulation of the weight gradient. When enabled, it is assumed that the weights - have an additional `main_grad` attribute (used instead of the - regular `grad`) which is a pre-allocated buffer of the correct + have an additional ``main_grad`` attribute (used instead of the + regular ``grad``) which is a pre-allocated buffer of the correct size to accumulate gradients in. This argument along with weight tensor having attribute 'overwrite_main_grad' set to True - will overwrite `main_grad` instead of accumulating. - return_bias : bool, default = `False` - when set to `True`, this module will not apply the additive bias itself, but + will overwrite ``main_grad`` instead of accumulating. + return_bias : bool, default = False + when set to ``True``, this module will not apply the additive bias itself, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when the bias addition can be fused to subsequent operations. - params_dtype : torch.dtype, default = `torch.get_default_dtype()` + params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. - delay_wgrad_compute : bool, default = `False` - Whether or not to delay weight gradient computation. If set to `True`, - it's the user's responsibility to call `module.backward_dw` to compute + delay_wgrad_compute : bool, default = False + Whether or not to delay weight gradient computation. If set to ``True``, + it's the user's responsibility to call ``module.backward_dw`` to compute weight gradients. symmetric_ar_type : {None, 'multimem_all_reduce', 'two_shot', 'one_shot'}, default = None Type of symmetric memory all-reduce to use during the forward pass. This can help in latency bound communication situations. - Requires PyTorch version 2.7.0 or higher. When set to None, standard all-reduce + Requires PyTorch version 2.7.0 or higher. When set to ``None``, standard all-reduce is used. """ diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index c1e3d8d2c9..56e050fe88 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -16,7 +16,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch import torch_version +from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.tensor.utils import is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, @@ -1661,38 +1661,38 @@ class LayerNormMLP(TransformerEngineBaseModule): intermediate size to which input samples are projected. eps : float, default = 1e-5 a value added to the denominator of layer normalization for numerical stability. - bias : bool, default = `True` - if set to `False`, the FC1 and FC2 layers will not learn an additive bias. + bias : bool, default = True + if set to ``False``, the FC1 and FC2 layers will not learn an additive bias. normalization : { 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm' type of normalization applied. activation : str, default = 'gelu' activation function used. - Options: 'gelu', 'geglu', 'qgelu', 'qgeglu', 'relu', 'reglu', 'srelu', 'sreglu', - 'silu', 'swiglu', and 'clamped_swiglu'. - activation_params : dict, default = `None` + Options: ``'gelu'``, ``'geglu'``, ``'qgelu'``, ``'qgeglu'``, ``'relu'``, ``'reglu'``, ``'srelu'``, ``'sreglu'``, + ``'silu'``, ``'swiglu'``, and ``'clamped_swiglu'``. + 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. - 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)`. - output_layer_init_method : Callable, default = `None` + At the moment, only used for ``'clamped_swiglu'`` activation which + supports ``'limit'`` and ``'alpha'`` 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)``. + output_layer_init_method : Callable, default = None used for initializing FC2 weights in the following way: - `output_layer_init_method(weight)`. When set to `None`, defaults to - `torch.nn.init.normal_(mean=0.0, std=0.023)`. - return_layernorm_output : bool, default = `False` - if set to `True`, output of layernorm is returned from the forward + ``output_layer_init_method(weight)``. When set to ``None``, defaults to + ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + return_layernorm_output : bool, default = False + if set to ``True``, output of layernorm is returned from the :meth:`forward` method together with the output of the linear transformation. Example use case: residual connection for transformer module is taken post layernorm. - return_layernorm_output_gathered : bool, default = `False` - if set to `True`, output of layernorm is returned after the all - gather operation. Ignored if return_layernorm_output is False. + return_layernorm_output_gathered : bool, default = False + if set to ``True``, output of layernorm is returned after the all + gather operation. Ignored if ``return_layernorm_output`` is False. Example use case: with sequence parallel, input to residual connection for transformer module (e.g. LoRA) will need to be gathered. Returning layernorm output gathered will prevent a redundant gather. - zero_centered_gamma : bool, default = 'False' - if set to 'True', gamma parameter in LayerNorm is initialized to 0 and + zero_centered_gamma : bool, default = False + if set to ``True``, gamma parameter in LayerNorm is initialized to 0 and the LayerNorm formula changes to .. math:: @@ -1702,62 +1702,62 @@ class LayerNormMLP(TransformerEngineBaseModule): 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 forward pass. - name: str, default = `None` + name : str, default = None name of the module, currently used for debugging purposes. Parallelism parameters ---------------------- - set_parallel_mode : bool, default = `False` - if set to `True`, FC1 is used as Column Parallel and FC2 is used as Row + set_parallel_mode : bool, default = False + if set to ``True``, FC1 is used as Column Parallel and FC2 is used as Row Parallel as described `here `_. - sequence_parallel : bool, default = `False` - if set to `True`, uses sequence parallelism. - tp_group : ProcessGroup, default = `None` + sequence_parallel : bool, default = False + if set to ``True``, uses sequence parallelism. + tp_group : ProcessGroup, default = None tensor parallel process group. tp_size : int, default = 1 used as TP (tensor parallel) world size when TP groups are not formed during initialization. In this case, users must call the - `set_tensor_parallel_group(tp_group)` method on the initialized module before the + ``set_tensor_parallel_group(tp_group)`` method on the initialized module before the forward pass to supply the tensor parallel group needed for tensor and sequence parallel collectives. Optimization parameters ----------------------- - fuse_wgrad_accumulation : bool, default = 'False' - if set to `True`, enables fusing of creation and accumulation of + fuse_wgrad_accumulation : bool, default = False + if set to ``True``, enables fusing of creation and accumulation of the weight gradient. When enabled, it is assumed that the weights - have an additional `main_grad` attribute (used instead of the - regular `grad`) which is a pre-allocated buffer of the correct + have an additional ``main_grad`` attribute (used instead of the + regular ``grad``) which is a pre-allocated buffer of the correct size to accumulate gradients in. This argument along with - weight tensor having attribute 'overwrite_main_grad' set to True - will overwrite `main_grad` instead of accumulating. - return_bias : bool, default = `False` - when set to `True`, this module will not apply the additive bias for FC2, but + weight tensor having attribute ``'overwrite_main_grad'`` set to True + will overwrite ``main_grad`` instead of accumulating. + return_bias : bool, default = False + when set to ``True``, this module will not apply the additive bias for FC2, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when the bias addition can be fused to subsequent operations. - params_dtype : torch.dtype, default = `torch.get_default_dtype()` + params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. - seq_length: int + seq_length : int sequence length of input samples. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propogation and activation recompute phase. - micro_batch_size: int + micro_batch_size : int batch size per training step. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propogation and activation recompute phase. - delay_wgrad_compute : bool, default = `False` - Whether or not to delay weight gradient computation. If set to `True`, - it's the user's responsibility to call `module.backward_dw` to compute + delay_wgrad_compute : bool, default = False + Whether or not to delay weight gradient computation. If set to ``True``, + it's the user's responsibility to call :meth:`backward_dw` to compute weight gradients. symmetric_ar_type : {None, 'multimem_all_reduce', 'two_shot', 'one_shot'}, default = None Type of symmetric memory all-reduce to use during the forward pass. This can help in latency bound communication situations. - Requires PyTorch version 2.7.0 or higher. When set to None, standard all-reduce + Requires PyTorch version 2.7.0 or higher. When set to ``None``, standard all-reduce is used. - checkpoint: bool, default = False + checkpoint : bool, default = False whether to use selective activation checkpointing, where activations are not saved for bwd, and instead are recomputed (skipping fc2, as it is not needed for backward). Trades compute for memory. default is false, in which activations are saved in fwd. not supported for onnx forward @@ -2235,7 +2235,7 @@ def onnx_forward( self, inp: torch.Tensor, is_grad_enabled: bool ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ - ONNX-compatible version of the forward function that provides numerical equivalence + ONNX-compatible version of the :meth:`forward` method that provides numerical equivalence while only using operations that have defined ONNX symbolic translations. This simplified implementation is designed specifically for inference scenarios. """ diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 3c804ffaa8..b65f7005eb 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -13,7 +13,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch import torch_version +from transformer_engine.pytorch.torch_version import torch_version from .base import ( fill_userbuffers_buffer_for_all_gather, @@ -985,7 +985,7 @@ def wgrad_gemm( class Linear(TransformerEngineBaseModule): """Applies a linear transformation to the incoming data :math:`y = xA^T + b` - On NVIDIA GPUs it is a drop-in replacement for `torch.nn.Linear`. + On NVIDIA GPUs it is a drop-in replacement for ``torch.nn.Linear``. Parameters ---------- @@ -993,14 +993,14 @@ class Linear(TransformerEngineBaseModule): size of each input sample. out_features : int size of each output sample. - bias : bool, default = `True` - if set to `False`, the layer will not learn an additive bias. - init_method : Callable, default = `None` - used for initializing weights in the following way: `init_method(weight)`. - When set to `None`, defaults to `torch.nn.init.normal_(mean=0.0, std=0.023)`. - get_rng_state_tracker : Callable, default = `None` + bias : bool, default = True + if set to ``False``, the layer will not learn an additive bias. + init_method : Callable, default = None + used for initializing weights in the following way: ``init_method(weight)``. + When set to ``None``, defaults to ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + get_rng_state_tracker : Callable, default = None used to get the random number generator state tracker for initializing weights. - rng_tracker_name : str, default = `None` + rng_tracker_name : str, default = None the param passed to get_rng_state_tracker to get the specific rng tracker. parameters_split : Optional[Union[Tuple[str, ...], Dict[str, int]]], default = None Configuration for splitting the weight and bias tensors along dim 0 into @@ -1008,62 +1008,62 @@ class Linear(TransformerEngineBaseModule): they are used to make the names of equally-sized parameters. If a dict (preferably an OrderedDict) is provided, the keys are used as names and values as split sizes along dim 0. The resulting parameters will have - names that end in `_weight` or `_bias`, so trailing underscores are + names that end in ``_weight`` or ``_bias``, so trailing underscores are stripped from any provided names. 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 forward pass. - name: str, default = `None` + name : str, default = None name of the module, currently used for debugging purposes. Parallelism parameters ---------------------- - sequence_parallel : bool, default = `False` - if set to `True`, uses sequence parallelism. - tp_group : ProcessGroup, default = `None` + sequence_parallel : bool, default = False + if set to ``True``, uses sequence parallelism. + tp_group : ProcessGroup, default = None tensor parallel process group. tp_size : int, default = 1 used as TP (tensor parallel) world size when TP groups are not formed during initialization. In this case, users must call the - `set_tensor_parallel_group(tp_group)` method on the initialized module before the + ``set_tensor_parallel_group(tp_group)`` method on the initialized module before the forward pass to supply the tensor parallel group needed for tensor and sequence parallel collectives. - parallel_mode : {None, 'column', 'row'}, default = `None` + parallel_mode : {None, 'column', 'row'}, default = None used to decide whether this Linear layer is Column Parallel Linear or Row Parallel Linear as described `here `_. - When set to `None`, no communication is performed. + When set to ``None``, no communication is performed. Optimization parameters ----------------------- fuse_wgrad_accumulation : bool, default = 'False' - if set to `True`, enables fusing of creation and accumulation of + if set to ``True``, enables fusing of creation and accumulation of the weight gradient. When enabled, it is assumed that the weights - have an additional `main_grad` attribute (used instead of the - regular `grad`) which is a pre-allocated buffer of the correct + have an additional ``main_grad`` attribute (used instead of the + regular ``grad``) which is a pre-allocated buffer of the correct size to accumulate gradients in. This argument along with weight tensor having attribute 'overwrite_main_grad' set to True - will overwrite `main_grad` instead of accumulating. - return_bias : bool, default = `False` - when set to `True`, this module will not apply the additive bias itself, but + will overwrite ``main_grad`` instead of accumulating. + return_bias : bool, default = False + when set to ``True``, this module will not apply the additive bias itself, but instead return the bias value during the forward pass together with the output of the linear transformation :math:`y = xA^T`. This is useful when the bias addition can be fused to subsequent operations. - params_dtype : torch.dtype, default = `torch.get_default_dtype()` + params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. - delay_wgrad_compute : bool, default = `False` - Whether or not to delay weight gradient computation. If set to `True`, - it's the user's responsibility to call `module.backward_dw` to compute + delay_wgrad_compute : bool, default = False + Whether or not to delay weight gradient computation. If set to ``True``, + it's the user's responsibility to call ``module.backward_dw`` to compute weight gradients. symmetric_ar_type : {None, 'multimem_all_reduce', 'two_shot', 'one_shot'}, default = None Type of symmetric memory all-reduce to use during the forward pass. This can help in latency bound communication situations. - Requires PyTorch version 2.7.0 or higher. When set to None, standard all-reduce + Requires PyTorch version 2.7.0 or higher. When set to ``None``, standard all-reduce is used. - save_original_input : bool, default = `False` - If set to `True`, always saves the original input tensor rather than the + save_original_input : bool, default = False + If set to ``True``, always saves the original input tensor rather than the cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. Cannot work with FP8 DelayedScaling recipe. diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py index fb267d8a9b..cac3e18220 100644 --- a/transformer_engine/pytorch/module/rmsnorm.py +++ b/transformer_engine/pytorch/module/rmsnorm.py @@ -33,32 +33,29 @@ class RMSNorm(_RMSNormOp): Parameters ---------- - normalized_shape: int or iterable of int + normalized_shape : int or iterable of int Inner dimensions of input tensor eps : float, default = 1e-5 A value added to the denominator for numerical stability - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype - zero_centered_gamma : bool, default = 'False' - If `True`, the :math:`\gamma` parameter is initialized to zero + zero_centered_gamma : bool, default = False + If ``True``, the :math:`\gamma` parameter is initialized to zero and the calculation changes to .. math:: y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) - sm_margin: int, default = 0 + sm_margin : int, default = 0 Number of SMs to exclude when launching CUDA kernels. This helps overlap with other kernels, e.g. communication kernels. For more fine-grained control, provide a dict with the SM - margin at each compute stage ("forward", "backward", - "inference"). - - Legacy - ------ - sequence_parallel: bool - Set a bool attr named `sequence_parallel` in the parameters. + margin at each compute stage (``"forward"``, ``"backward"``, + ``"inference"``). + sequence_parallel : bool + **Legacy parameter.** Set a bool attr named ``sequence_parallel`` in the parameters. This is custom logic for Megatron-LM integration. """ diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index a07ffea43f..103e537dd0 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -10,7 +10,7 @@ import torch from transformer_engine_torch import FP8TensorMeta -from .. import torch_version +from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager from ..tensor.float8_tensor import Float8Tensor from ..quantized_tensor import QuantizedTensorStorage diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 8a754c6382..a444facd0a 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -53,7 +53,7 @@ class _ActivationOperation(BasicOperation, metaclass=abc.ABCMeta): Parameters ---------- - cache_quantized_input: bool, default = False + cache_quantized_input : bool, default = False Quantize input tensor when caching for use in the backward pass. This will typically reduce memory usage but require extra compute and increase numerical error. This feature is @@ -408,11 +408,11 @@ class ClampedSwiGLU(_ActivationOperation): Parameters ---------- - limit: float + limit : float The clamp limit. - alpha: float + alpha : float The scaling factor for the sigmoid function used in the activation. - cache_quantized_input: bool, default = False + cache_quantized_input : bool, default = False Quantize input tensor when caching for use in the backward pass. """ diff --git a/transformer_engine/pytorch/ops/basic/all_gather.py b/transformer_engine/pytorch/ops/basic/all_gather.py index bcd3c1417e..fc768ad83b 100644 --- a/transformer_engine/pytorch/ops/basic/all_gather.py +++ b/transformer_engine/pytorch/ops/basic/all_gather.py @@ -23,7 +23,7 @@ class AllGather(BasicOperation): Parameters ---------- - process_group: torch.distributed.ProcessGroup, default = world group + process_group : torch.distributed.ProcessGroup, default = world group Process group for communication """ diff --git a/transformer_engine/pytorch/ops/basic/all_reduce.py b/transformer_engine/pytorch/ops/basic/all_reduce.py index d8c1eb0069..d9a253924c 100644 --- a/transformer_engine/pytorch/ops/basic/all_reduce.py +++ b/transformer_engine/pytorch/ops/basic/all_reduce.py @@ -24,7 +24,7 @@ class AllReduce(BasicOperation): Parameters ---------- - process_group: torch.distributed.ProcessGroup, default = world group + process_group : torch.distributed.ProcessGroup, default = world group Process group for communication """ diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index c629d0158d..9f09e6634b 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -53,27 +53,27 @@ class BasicLinear(BasicOperation): Parameters ---------- - in_features: int + in_features : int Inner dimension of input tensor - out_features: int + out_features : int Inner dimension of output tensor - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype - tensor_parallel_mode: {`None`, "column", "row"}, default = `None` + tensor_parallel_mode : {`None`, "column", "row"}, default = `None` Mode for tensor parallelism - tensor_parallel_group: torch.distributed.ProcessGroup, default = world group + tensor_parallel_group : torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel: bool, default = `False` + sequence_parallel : bool, default = `False` Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) - rng_state_tracker_function: callable + rng_state_tracker_function : callable Function that returns `CudaRNGStatesTracker`, which is used for model-parallel weight initialization - accumulate_into_main_grad: bool, default = `False` + accumulate_into_main_grad : bool, default = `False` Whether to directly accumulate weight gradients into the weight's `main_grad` attribute instead of relying on PyTorch autograd. The weight's `main_grad` must be set externally and diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 5ec0d2ce5e..6910163825 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -22,16 +22,16 @@ class Bias(BasicOperation): Parameters ---------- - size: int + size : int Inner dimension of input tensor - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype - tensor_parallel: bool, default = `False` + tensor_parallel : bool, default = `False` Whether to distribute input tensor and bias tensors along inner dimension - tensor_parallel_group: torch.distributed.ProcessGroup, default = world group + tensor_parallel_group : torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism """ diff --git a/transformer_engine/pytorch/ops/basic/l2normalization.py b/transformer_engine/pytorch/ops/basic/l2normalization.py index 440fee34d1..ff4f923819 100644 --- a/transformer_engine/pytorch/ops/basic/l2normalization.py +++ b/transformer_engine/pytorch/ops/basic/l2normalization.py @@ -10,7 +10,7 @@ import torch -from ... import torch_version +from ...torch_version import torch_version from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload from ...jit import ( l2normalization_fused, @@ -40,11 +40,11 @@ class L2Normalization(BasicOperation): ---------- eps : float, default = 1e-6 A value added to the denominator for numerical stability - seq_length: int, default = None + seq_length : int, default = None sequence length of input samples. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propagation and activation recompute phase. - micro_batch_size: int, default = None + micro_batch_size : int, default = None batch size per training step. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propagation and activation recompute phase. diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 91e6de07d7..3922f85cad 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -42,14 +42,14 @@ class LayerNorm(BasicOperation): Parameters ---------- - normalized_shape: int or iterable of int + normalized_shape : int or iterable of int Inner dimensions of input tensor eps : float, default = 1e-5 A value added to the denominator of layer normalization for numerical stability - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype zero_centered_gamma : bool, default = 'False' If `True`, the :math:`\gamma` parameter is initialized to zero @@ -58,7 +58,7 @@ class LayerNorm(BasicOperation): .. math:: y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta - sm_margin: int or dict, default = 0 + sm_margin : int or dict, default = 0 Number of SMs to exclude when launching CUDA kernels. This helps overlap with other kernels, e.g. communication kernels. For more fine-grained control, provide a dict with the SM diff --git a/transformer_engine/pytorch/ops/basic/quantize.py b/transformer_engine/pytorch/ops/basic/quantize.py index 87c65d4b29..1278701a9b 100644 --- a/transformer_engine/pytorch/ops/basic/quantize.py +++ b/transformer_engine/pytorch/ops/basic/quantize.py @@ -23,9 +23,9 @@ class Quantize(BasicOperation): Parameters ---------- - forward: bool, default = `True` + forward : bool, default = `True` Perform quantization in forward pass - backward: bool, default = `False` + backward : bool, default = `False` Perform quantization in backward pass """ diff --git a/transformer_engine/pytorch/ops/basic/reduce_scatter.py b/transformer_engine/pytorch/ops/basic/reduce_scatter.py index e0017853f6..eabbb461bc 100644 --- a/transformer_engine/pytorch/ops/basic/reduce_scatter.py +++ b/transformer_engine/pytorch/ops/basic/reduce_scatter.py @@ -23,7 +23,7 @@ class ReduceScatter(BasicOperation): Parameters ---------- - process_group: torch.distributed.ProcessGroup, default = world group + process_group : torch.distributed.ProcessGroup, default = world group Process group for communication """ diff --git a/transformer_engine/pytorch/ops/basic/reshape.py b/transformer_engine/pytorch/ops/basic/reshape.py index 50af9fcfff..fcdb3b0bbe 100644 --- a/transformer_engine/pytorch/ops/basic/reshape.py +++ b/transformer_engine/pytorch/ops/basic/reshape.py @@ -24,7 +24,7 @@ class Reshape(BasicOperation): Parameters ---------- - shape: iterable of int + shape : iterable of int Output tensor dimensions. If one dimension is -1, it is inferred based on input tensor dimensions. diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index d91091eb02..316c292c53 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -42,13 +42,13 @@ class RMSNorm(BasicOperation): Parameters ---------- - normalized_shape: int or iterable of int + normalized_shape : int or iterable of int Inner dimensions of input tensor eps : float, default = 1e-5 A value added to the denominator for numerical stability - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype zero_centered_gamma : bool, default = 'False' If `True`, the :math:`\gamma` parameter is initialized to zero @@ -57,7 +57,7 @@ class RMSNorm(BasicOperation): .. math:: y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) - sm_margin: int, default = 0 + sm_margin : int, default = 0 Number of SMs to exclude when launching CUDA kernels. This helps overlap with other kernels, e.g. communication kernels. For more fine-grained control, provide a dict with the SM diff --git a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py index 7897ef164e..a33ef4acf8 100644 --- a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py +++ b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py @@ -90,15 +90,15 @@ def fuse_backward_activation_bias( Parameters ---------- - ops: list of tuples + ops : list of tuples Backward pass operations and the indices of the corresponding basic operations. - recipe: Recipe, optional + recipe : Recipe, optional Used quantization recipe Returns ------- - ops: list of tuples + ops : list of tuples Updated backward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index 54a23395af..1df55b83a0 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -87,13 +87,13 @@ def fuse_backward_add_rmsnorm( Parameters ---------- - ops: list of tuples + ops : list of tuples Backward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated backward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index a86745a686..0c12e3ab3e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -119,13 +119,13 @@ def fuse_backward_linear_add( Parameters ---------- - ops: list of tuples + ops : list of tuples Backward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated backward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py index 832e51de83..39ee4ab2fa 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py @@ -119,13 +119,13 @@ def fuse_backward_linear_scale( Parameters ---------- - ops: list of tuples + ops : list of tuples Backward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated backward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 74bd3d1b32..ca3d57ac98 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -142,13 +142,13 @@ def fuse_forward_linear_bias_activation( Parameters ---------- - ops: list of tuples + ops : list of tuples Forward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated forward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 6d5d553391..8a0f77dd56 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -139,13 +139,13 @@ def fuse_forward_linear_bias_add( Parameters ---------- - ops: list of tuples + ops : list of tuples Forward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated forward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index 24788bcdfb..fe93410707 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -118,13 +118,13 @@ def fuse_forward_linear_scale_add( Parameters ---------- - ops: list of tuples + ops : list of tuples Forward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated forward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 32e4ee3657..5149aa1ffb 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -589,13 +589,13 @@ def fuse_userbuffers_backward_linear( Parameters ---------- - ops: list of tuples + ops : list of tuples Backward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated backward pass operations """ diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index d50d031ba7..517632d651 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -377,13 +377,13 @@ def fuse_userbuffers_forward_linear( Parameters ---------- - ops: list of tuples + ops : list of tuples Forward pass operations and the indices of the corresponding basic operations. Returns ------- - ops: list of tuples + ops : list of tuples Updated forward pass operations """ diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 6026a40b65..fecf28f0a9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -310,7 +310,7 @@ class OperationFuser: Parameters ---------- - ops: list of FusibleOperation + ops : list of FusibleOperation Pipeline of operations """ diff --git a/transformer_engine/pytorch/ops/linear.py b/transformer_engine/pytorch/ops/linear.py index 325126a3d4..d1e6382291 100644 --- a/transformer_engine/pytorch/ops/linear.py +++ b/transformer_engine/pytorch/ops/linear.py @@ -27,29 +27,29 @@ class Linear(FusedOperation): Parameters ---------- - in_features: int + in_features : int Inner dimension of input tensor - out_features: int + out_features : int Inner dimension of output tensor - bias: bool, default = `True` + bias : bool, default = `True` Apply additive bias - device: torch.device, default = default CUDA device + device : torch.device, default = default CUDA device Tensor device - dtype: torch.dtype, default = default dtype + dtype : torch.dtype, default = default dtype Tensor datatype - tensor_parallel_mode: {`None`, "column", "row"}, default = `None` + tensor_parallel_mode : {`None`, "column", "row"}, default = `None` Mode for tensor parallelism - tensor_parallel_group: torch.distributed.ProcessGroup, default = world group + tensor_parallel_group : torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel: bool, default = `False` + sequence_parallel : bool, default = `False` Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) - rng_state_tracker_function: callable + rng_state_tracker_function : callable Function that returns CudaRNGStatesTracker, which is used for model-parallel weight initialization - accumulate_into_main_grad: bool, default = `False` + accumulate_into_main_grad : bool, default = `False` Whether to directly accumulate weight gradients into the weight's `main_grad` attribute instead of relying on PyTorch autograd. The weight's `main_grad` must be set externally and diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 6ae49dcd4e..421c92b823 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -684,7 +684,7 @@ class FusedOperation(FusibleOperation): Parameters ---------- - basic_ops: iterable of FusibleOperation + basic_ops : iterable of FusibleOperation Basic ops that are interchangeable with this op """ diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index f73bc9a966..9fce9cefcf 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -514,22 +514,22 @@ def moe_permute( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. - routing_map: torch.Tensor + routing_map : torch.Tensor The token to expert mapping tensor. If map_type is 'mask', routing_map is of shape [num_tokens, num_experts] and dtype 'int32'. The values in it: 1 means the token is routed to this expert and 0 means not. If map_type is 'index', routing_map is of shape [num_tokens, topK] and dtype 'int32'. The values in it are the routed expert indices. - num_out_tokens: int, default = -1 + num_out_tokens : int, default = -1 The effective output token count, representing the number of tokens not dropped. By default, set to '-1', meaning no tokens are dropped. - max_token_num: int, default = -1 + max_token_num : int, default = -1 The maximum number of tokens, used for workspace allocation. By default, set to '-1', meaning the calculation of the size of workspace is automatically taken over by the operator. - map_type: str, default = 'mask' + map_type : str, default = 'mask' Type of the routing map tensor. Options are: 'mask', 'index'. Refer to `routing_map` for more details. @@ -556,16 +556,16 @@ def moe_permute_with_probs( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. - probs: torch.Tensor + probs : torch.Tensor The tensor of probabilities corresponding to the permuted tokens and is of shape [num_tokens, num_experts]. It will be permuted with the tokens according to the routing_map. - routing_map: torch.Tensor + routing_map : torch.Tensor The token to expert mapping tensor of shape [num_tokens, num_experts] and dtype 'int32'. The values in it: 1 means the token is routed to this expert and 0 means not. - num_out_tokens: int, default = -1 + num_out_tokens : int, default = -1 The effective output token count, representing the number of tokens not dropped. By default, set to '-1', meaning no tokens are dropped. """ @@ -589,21 +589,21 @@ def moe_unpermute( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor with permuted tokens of shape `[num_tokens, hidden_size]` to be unpermuted. - row_id_map: torch.Tensor + row_id_map : torch.Tensor The tensor of a mapping table for sorted indices used to unpermute the tokens, which is the second output tensor of `Permute`. - merging_probs: torch.Tensor, default = None + merging_probs : torch.Tensor, default = None The tensor of probabilities corresponding to the permuted tokens. If provided, the unpermuted tokens will be merged with their respective probabilities. By default, set to an empty tensor, which means that the tokens are directly merged by accumulation. - restore_shape: torch.Size, default = None + restore_shape : torch.Size, default = None The output shape after the unpermute operation. - map_type: str, default = 'mask' + map_type : str, default = 'mask' Type of the routing map tensor. Should be the same as the value passed to moe_permute. Options are: 'mask', 'index'. - probs: torch.Tensor, default = None + probs : torch.Tensor, default = None Renamed to merging_probs. Keep for backward compatibility. """ if probs is not None: @@ -733,11 +733,11 @@ def moe_sort_chunks_by_index( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. - split_sizes: torch.Tensor + split_sizes : torch.Tensor Chunk sizes of the inp tensor along the 0-th dimension. - sorted_indices: torch.Tensor + sorted_indices : torch.Tensor Chunk indices used to permute the chunks. """ output, _ = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, None) @@ -757,15 +757,15 @@ def moe_sort_chunks_by_index_with_probs( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. - probs: torch.Tensor + probs : torch.Tensor The tensor of probabilities corresponding to the permuted tokens and is of shape [num_tokens]. It will be permuted with the tokens according to the split_sizes and sorted_indices. - split_sizes: torch.Tensor + split_sizes : torch.Tensor Chunk sizes of the inp tensor along the 0-th dimension. - sorted_indices: torch.Tensor + sorted_indices : torch.Tensor Chunk indices used to permute the chunks. """ output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index acc2e55320..fbe2ee6d1c 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -26,8 +26,8 @@ NVFP4BlockScaling, CustomRecipe, ) - from .constants import dist_group_type + from .utils import get_device_compute_capability from .jit import jit_fuser @@ -678,7 +678,7 @@ def fp8_model_init( .. warning:: fp8_model_init is deprecated and will be removed in a future release. Use - quantized_model_init(enabled=..., recipe=..., preserve_high_precision_init_val=...) instead. + ``quantized_model_init(enabled=..., recipe=..., preserve_high_precision_init_val=...)`` instead. """ @@ -723,7 +723,7 @@ def quantized_model_init( Parameters ---------- - enabled: bool, default = `True` + enabled : bool, default = True when enabled, Transformer Engine modules created inside this `quantized_model_init` region will hold only quantized copies of its parameters, as opposed to the default behavior where both higher precision and quantized copies are present. Setting this @@ -734,9 +734,9 @@ def quantized_model_init( precision copies of weights are already present in the optimizer. * inference, where only the quantized copies of the parameters are used. * LoRA-like fine-tuning, where the main parameters of the model do not change. - recipe: transformer_engine.common.recipe.Recipe, default = `None` + recipe : transformer_engine.common.recipe.Recipe, default = None Recipe used to create the parameters. If left to None, it uses the default recipe. - preserve_high_precision_init_val: bool, default = `False` + preserve_high_precision_init_val : bool, default = False when enabled, store the high precision tensor used to initialize quantized parameters in CPU memory, and add two function attributes named `get_high_precision_init_val()` and `clear_high_precision_init_val()` to quantized parameters to get/clear this high @@ -773,8 +773,8 @@ def fp8_autocast( """ .. warning:: - fp8_autocast is deprecated and will be removed in a future release. - Use autocast(enabled=..., calibrating=..., recipe=..., group=..., _graph=...) instead. + ``fp8_autocast`` is deprecated and will be removed in a future release. + Use ``autocast(enabled=..., calibrating=..., recipe=..., group=..., _graph=...)`` instead. """ @@ -828,16 +828,16 @@ def autocast( Parameters ---------- - enabled: bool, default = `True` + enabled : bool, default = True whether or not to enable low precision quantization (FP8/FP4). - calibrating: bool, default = `False` + calibrating : bool, default = False calibration mode allows collecting statistics such as amax and scale data of quantized tensors even when executing without quantization enabled. This is useful for saving an inference ready checkpoint while training using a higher precision. - recipe: recipe.Recipe, default = `None` + recipe : recipe.Recipe, default = None recipe used for low precision quantization. - amax_reduction_group: torch._C._distributed_c10d.ProcessGroup, default = `None` + amax_reduction_group : torch._C._distributed_c10d.ProcessGroup, default = None distributed group over which amaxes for the quantized tensors are reduced at the end of each training step. """ diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 3e3f460b41..c9a4467a82 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -27,7 +27,7 @@ class QuantizedTensorStorage: - r"""Base class for all *TensorStorage classes. + r"""Base class for all TensorStorage classes. This class (and its subclasses) are optimization for when the full QuantizedTensor is not needed (when it is fully @@ -54,11 +54,11 @@ def update_usage( Parameters ---------- - rowwise_usage : Optional[bool[, default = `None` + rowwise_usage : Optional[bool[, default = None Whether to create or keep the data needed for using the tensor in rowwise fashion (e.g. as B argument in TN GEMM). Leaving it as `None` preserves the original value in the tensor. - columnwise_usage : Optional[bool], default = `None` + columnwise_usage : Optional[bool], default = None Whether to create or keep the data needed for using the tensor in columnwise fashion (e.g. as A argument in TN GEMM). Leaving it as `None` preserves the original value in the tensor. @@ -128,7 +128,7 @@ def prepare_for_saving( ]: """Prepare tensors for saving. Needed because save_for_backward accepts only torch.Tensor/torch.nn.Parameter types, while we want to be able to save - the internal *TensorStorage types too.""" + the internal TensorStorage types too.""" tensor_list, tensor_objects_list = [], [] for tensor in tensors: diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index db5114ae04..a6030dd9df 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -92,24 +92,24 @@ def fused_topk_with_score_function( Fused topk with score function router. Parameters ---------- - logits: torch.Tensor - topk: int - use_pre_softmax: bool + logits : torch.Tensor + topk : int + use_pre_softmax : bool if enabled, the computation order: softmax -> topk - num_groups: int + num_groups : int used in the group topk - group_topk: int + group_topk : int used in the group topk - scaling_factor: float - score_function: str + scaling_factor : float + score_function : str currently only support softmax and sigmoid - expert_bias: torch.Tensor + expert_bias : torch.Tensor could be used in the sigmoid Returns ------- - probs: torch.Tensor - routing_map: torch.Tensor + probs : torch.Tensor + routing_map : torch.Tensor """ if logits.dtype == torch.float64: raise ValueError("Current TE does not support float64 router type") @@ -186,15 +186,15 @@ def fused_compute_score_for_moe_aux_loss( Fused compute scores for MoE aux loss, subset of the fused_topk_with_score_function. Parameters ---------- - logits: torch.Tensor - topk: int - score_function: str + logits : torch.Tensor + topk : int + score_function : str currently only support softmax and sigmoid Returns ------- - routing_map: torch.Tensor - scores: torch.Tensor + routing_map : torch.Tensor + scores : torch.Tensor """ return FusedComputeScoresForMoEAuxLoss.apply(logits, topk, score_function) @@ -258,18 +258,18 @@ def fused_moe_aux_loss( Fused MoE aux loss. Parameters ---------- - probs: torch.Tensor - tokens_per_expert: torch.Tensor + probs : torch.Tensor + tokens_per_expert : torch.Tensor the number of tokens per expert - total_num_tokens: int + total_num_tokens : int the total number of tokens, involved in the aux loss calculation - num_experts: int - topk: int - coeff: float + num_experts : int + topk : int + coeff : float the coefficient of the aux loss Returns ------- - aux_loss: torch.scalar + aux_loss : torch.scalar """ return FusedAuxLoss.apply(probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff) diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 069565f388..01e03e5355 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -307,18 +307,18 @@ class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor): Parameters ---------- - rowwise_data: torch.Tensor + rowwise_data : torch.Tensor FP8 data in a uint8 tensor matching shape of dequantized tensor. - rowwise_scale_inv: torch.Tensor + rowwise_scale_inv : torch.Tensor FP32 dequantization scales in GEMM format for dequantizing rowwise_data. - columnwise_data: Optional[torch.Tensor] + columnwise_data : Optional[torch.Tensor] FP8 data in a uint8 tensor matching shape of dequantized tensor transpose. - columnwise_scale_inv: Optional[torch.Tensor] + columnwise_scale_inv : Optional[torch.Tensor] FP32 dequantization scales in GEMM format for dequantizing columnwise_data. - fp8_dtype: transformer_engine_torch.DType, default = kFloat8E4M3 + fp8_dtype : transformer_engine_torch.DType, default = kFloat8E4M3 FP8 format. - quantizer: Quantizer - the Float8BlockQuantizer that quantized this tensor and + quantizer : Quantizer - the Float8BlockQuantizer that quantized this tensor and holds configuration about quantization and dequantization modes. """ diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index f44ed33b9e..e3ca110dfa 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -453,23 +453,23 @@ class Float8Tensor(Float8TensorStorage, QuantizedTensor): Parameters ---------- - shape: int or iterable of int + shape : int or iterable of int Tensor dimensions. - dtype: torch.dtype + dtype : torch.dtype Nominal tensor datatype. - requires_grad: bool, optional = False + requires_grad : bool, optional = False Whether to compute gradients for this tensor. - data: torch.Tensor + data : torch.Tensor Raw FP8 data in a uint8 tensor - fp8_scale_inv: torch.Tensor + fp8_scale_inv : torch.Tensor Reciprocal of the scaling factor applied when casting to FP8, i.e. the scaling factor that must be applied when casting from FP8 to higher precision. - fp8_dtype: transformer_engine_torch.DType + fp8_dtype : transformer_engine_torch.DType FP8 format. - data_transpose: torch.Tensor, optional + data_transpose : torch.Tensor, optional FP8 transpose data in a uint8 tensor - quantizer: Float8Quantizer, Float8CurrentScalingQuantizer, optional + quantizer : Float8Quantizer, Float8CurrentScalingQuantizer, optional Builder class for FP8 tensors """ diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 3b9bdd2fe7..a41079080c 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -204,16 +204,16 @@ class MXFP8Tensor(MXFP8TensorStorage, QuantizedTensor): Parameters ---------- - data: torch.Tensor + data : torch.Tensor Raw FP8 data in a uint8 tensor - fp8_dtype: transformer_engine_torch.DType, default = kFloat8E4M3 + fp8_dtype : transformer_engine_torch.DType, default = kFloat8E4M3 FP8 format. - fp8_scale_inv: torch.Tensor + fp8_scale_inv : torch.Tensor Reciprocal of the scaling factor applied when casting to FP8, i.e. the scaling factor that must be applied when casting from FP8 to higher precision. - dtype: torch.dtype, default = torch.float32 + dtype : torch.dtype, default = torch.float32 Nominal tensor datatype. """ diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 5ee9441529..0c244628d6 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -382,26 +382,26 @@ class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor): Parameters ---------- - rowwise_data: torch.Tensor + rowwise_data : torch.Tensor Raw FP4 data in a uint8 tensor (rowwise layout). - rowwise_scale_inv: torch.Tensor + rowwise_scale_inv : torch.Tensor Reciprocal of the scaling factor applied when casting to FP4, i.e. the scaling factor that must be applied when casting from FP4 to higher precision (rowwise). - columnwise_data: torch.Tensor, optional + columnwise_data : torch.Tensor, optional Raw FP4 data in a uint8 tensor (columnwise layout). - columnwise_scale_inv: torch.Tensor, optional + columnwise_scale_inv : torch.Tensor, optional Reciprocal of the scaling factor for columnwise FP4 data. - amax_rowwise: torch.Tensor, optional + amax_rowwise : torch.Tensor, optional Rowwise amax tracking tensor. - amax_columnwise: torch.Tensor, optional + amax_columnwise : torch.Tensor, optional Columnwise amax tracking tensor. - fp4_dtype: TE_DType + fp4_dtype : TE_DType The FP4 data type used for quantization. - quantizer: Quantizer + quantizer : Quantizer The quantizer instance used for this tensor. - dtype: torch.dtype, default = torch.float32 + dtype : torch.dtype, default = torch.float32 Nominal tensor datatype, used in dequantize. """ diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 20aba6c2bf..9773e17e64 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -74,7 +74,7 @@ def cast_master_weights_to_fp8( fsdp_shard_model_weights : list of FSDP shard model weights. If None, it means that the model weights are not sharded. Otherwise, it means that the model weights are sharded and we get target model weights data storage using the FSDP shard model weights. - manual_post_all_gather_processing: bool, default = `False`. + manual_post_all_gather_processing : bool, default = `False`. If False, post processing will be automatically triggered during next forward. If True, the timing of calling post_all_gather_processing is left to the user. Note that users must call `post_all_gather_processing` if it's set to True, diff --git a/transformer_engine/pytorch/torch_version.py b/transformer_engine/pytorch/torch_version.py new file mode 100644 index 0000000000..ff1a0abb89 --- /dev/null +++ b/transformer_engine/pytorch/torch_version.py @@ -0,0 +1,15 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch version utilities""" +from __future__ import annotations +import functools +import torch +from packaging.version import Version as PkgVersion + + +@functools.lru_cache(maxsize=None) +def torch_version() -> tuple[int, ...]: + """Get PyTorch version""" + return PkgVersion(str(torch.__version__)).release diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 4c7599ad80..b3ad8ccc55 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -10,7 +10,7 @@ import torch -from transformer_engine.pytorch import torch_version +from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm from transformer_engine.debug.pytorch.debug_state import TEDebugState from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention @@ -75,8 +75,8 @@ class TransformerLayer(torch.nn.Module): .. note:: - Argument :attr:`attention_mask` in the `forward` call is only used when - :attr:`self_attn_mask_type` includes `"padding"` or `"arbitrary"`. + Argument :attr:`attention_mask` in the :meth:`forward` call is only used when + :attr:`self_attn_mask_type` includes ``"padding"`` or ``"arbitrary"``. Parameters ---------- @@ -86,76 +86,76 @@ class TransformerLayer(torch.nn.Module): intermediate size to which input samples are projected. num_attention_heads : int number of attention heads in the transformer layer. - num_gqa_groups : int, default = `None` + num_gqa_groups : int, default = None number of GQA groups in the transformer layer. Grouped Query Attention is described in `this paper `_. This only affects the keys and values, not the querys. GQA-1 is equivalent to Multi-Query Attention (`MQA `_), while GQA-H - is equivalent to MHA, i.e. `num_gqa_groups = num_attention_heads`. + is equivalent to MHA, i.e. ``num_gqa_groups = num_attention_heads``. layernorm_epsilon : float, default = 1e-5 a value added to the denominator of layer normalization for numerical stability. - hidden_dropout: float, default = 0.1 + hidden_dropout : float, default = 0.1 dropout probability for the dropout op after FC2 layer. - attention_dropout: float, default = 0.1 + attention_dropout : float, default = 0.1 dropout probability for the dropout op during multi-head attention. - init_method : Callable, default = `None` + init_method : Callable, default = None used for initializing weights of QKV and 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)`. - output_layer_init_method : Callable, default = `None` + ``init_method(weight)``. When set to ``None``, defaults to + ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + output_layer_init_method : Callable, default = None used for initializing weights of PROJ and FC2 in the following way: - `output_layer_init_method(weight)`. When set to `None`, defaults to - `torch.nn.init.normal_(mean=0.0, std=0.023)`. - apply_residual_connection_post_layernorm : bool, default = `False` - if set to `True`, residual connections are taken + ``output_layer_init_method(weight)``. When set to ``None``, defaults to + ``torch.nn.init.normal_(mean=0.0, std=0.023)``. + apply_residual_connection_post_layernorm : bool, default = False + if set to ``True``, residual connections are taken from the output of layer norm (default is taken from input of layer norm) - layer_number: int, default = `None` - layer number of the current `TransformerLayer` when multiple such modules are + layer_number : int, default = None + layer number of the current :class:`TransformerLayer` when multiple such modules are concatenated to form a transformer block. - output_layernorm: bool, default = `False` - if set to `True`, layer normalization is applied on the output side, + output_layernorm : bool, default = False + if set to ``True``, layer normalization is applied on the output side, after the final dropout-add. default behavior is to apply layer normalization on the input side, before the QKV transformation. - parallel_attention_mlp: bool, default = `False` - if set to `True`, self-attention and feedforward network are computed + parallel_attention_mlp : bool, default = False + if set to ``True``, self-attention and feedforward network are computed based on the same input (in parallel) instead of sequentially. Both blocks have an independent normalization. This architecture is used in `Falcon` models. - layer_type: {'encoder', 'decoder'}, default = `encoder` - if set to `decoder`, an additional cross-attn block is added after self-attn. + layer_type : {'encoder', 'decoder'}, default = "encoder" + if set to ``"decoder"``, an additional cross-attn block is added after self-attn. This can be used for structures like `T5` Transformer in conjunction with the - `encoder` option. - kv_channels: int, default = `None` + ``"encoder"`` option. + kv_channels : int, default = None number of query-key-value channels per attention head. defaults to - :attr:`hidden_size` / :attr:`num_attention_heads` if `None`. - self_attn_mask_type: {'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right', + :attr:`hidden_size` / :attr:`num_attention_heads` if ``None``. + self_attn_mask_type : {'no_mask', 'padding', 'causal', 'padding_causal', 'causal_bottom_right', 'padding_causal_bottom_right', 'arbitrary'}, - default = `causal` + default = "causal" type of attention mask passed into softmax operation for encoder. - Overridden by :attr:`self_attn_mask_type` in the `forward` method. - The forward arg is useful for dynamically changing mask types, e.g. - a different mask for training and inference. The init arg is useful + Overridden by :attr:`self_attn_mask_type` in the :meth:`forward` method. + The :meth:`forward` arg is useful for dynamically changing mask types, e.g. + a different mask for training and inference. The :meth:`__init__` arg is useful for cases involving compilation/tracing, e.g. ONNX export. - window_size: Optional[Tuple[int, int]], default = `None` + window_size : Optional[Tuple[int, int]], default = None sliding window size for local attention in encoder, where query at position i - attends to keys in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - - seqlen_q + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean - no sliding window and causal mask specifically. Both `causal` and - `causal_bottom_right` masks map to `window_size = (-1, 0)` and Transformer Engine - distinguishes them based on `self_attn_mask_type` or `enc_dec_attn_mask_type`. - Similar to :attr:`self_attn_mask_type`, `window_size` can be overridden by - :attr:`window_size` in `forward` as well. - enc_dec_attn_mask_type: {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'}, - default = `no_mask` + attends to keys in ``[i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k + - seqlen_q + window_size[1]]`` inclusive. Special cases ``(-1, -1)`` and ``(-1, 0)`` mean + no sliding window and causal mask specifically. Both ``"causal"`` and + ``"causal_bottom_right"`` masks map to :attr:`window_size` = ``(-1, 0)`` and Transformer Engine + distinguishes them based on :attr:`self_attn_mask_type` or :attr:`enc_dec_attn_mask_type`. + Similar to :attr:`self_attn_mask_type`, :attr:`window_size` can be overridden by + :attr:`window_size` in :meth:`forward` as well. + enc_dec_attn_mask_type : {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'}, + default = "no_mask" type of attention mask passed into softmax operation for decoder. - enc_dec_window_size: Optional[Tuple[int, int]], default = `None` + enc_dec_window_size : Optional[Tuple[int, int]], default = None sliding window size for local attention in decoder. - zero_centered_gamma : bool, default = 'False' - if set to 'True', gamma parameter in LayerNorm is initialized to 0 and + zero_centered_gamma : bool, default = False + if set to ``True``, gamma parameter in LayerNorm is initialized to 0 and the LayerNorm formula changes to .. math:: @@ -163,111 +163,126 @@ class TransformerLayer(torch.nn.Module): (1 + \gamma) + \beta normalization : { 'LayerNorm', 'RMSNorm' }, default = 'LayerNorm' type of normalization applied. - qkv_weight_interleaved : bool, default = `True` - if set to `False`, the QKV weight is interpreted as a concatenation of - query, key, and value weights along the `0th` dimension. The default - interpretation is that the individual `q`, `k`, and `v` weights for each - attention head are interleaved. This parameter is set to `False` when + qkv_weight_interleaved : bool, default = True + if set to ``False``, the QKV weight is interpreted as a concatenation of + query, key, and value weights along the ``0th`` dimension. The default + interpretation is that the individual ``q``, ``k``, and ``v`` weights for each + attention head are interleaved. This parameter is set to ``False`` when using :attr:`fuse_qkv_params=False`. - rotary_pos_interleaved : bool, default = `False` + rotary_pos_interleaved : bool, default = False whether to use interleaved rotary position embeddings. - bias : bool, default = `True` - if set to `False`, the transformer layer will not learn any additive biases. + bias : bool, default = True + if set to ``False``, the transformer layer will not learn any additive biases. activation : str, default = 'gelu' Type of activation used in MLP block. - Options are: 'gelu', 'geglu', 'qgelu', 'qgeglu', 'relu', 'reglu', 'srelu', 'sreglu', - 'silu', 'swiglu', and 'clamped_swiglu'. - activation_params : Optional[dict], default = `None` + Options are: ``'gelu'``, ``'geglu'``, ``'qgelu'``, ``'qgeglu'``, ``'relu'``, ``'reglu'``, ``'srelu'``, ``'sreglu'``, + ``'silu'``, ``'swiglu'``, and ``'clamped_swiglu'``. + 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}`. + 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}``. 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 forward pass. - attn_input_format: {'sbhd', 'bshd', 'thd'}, default = 'sbhd' - This controls whether the dimensions of the - intermediate hidden states is 'sequence first' ('sbhd'), 'batch first' ('bshd'), - or 'token first' ('thd'). `s` stands for the sequence length, `b` batch size, - `t` the total number of tokens, `h` the number of heads, `d` head size. - Note that these formats are very closely - related to the `qkv_format` in the `MultiHeadAttention` - and `DotProductAttention` modules. - name: str, default = `None` + attn_input_format : {'sbhd', 'bshd', 'thd'}, default = 'sbhd' + This controls whether the dimensions of the + intermediate hidden states is 'sequence first' (``'sbhd'``), 'batch first' (``'bshd'``), + or 'token first' (``'thd'``). ``s`` stands for the sequence length, ``b`` batch size, + ``t`` the total number of tokens, ``h`` the number of heads, ``d`` head size. + Note that these formats are very closely + related to the :attr:`qkv_format` parameter in the :class:`MultiHeadAttention` + and :class:`DotProductAttention` modules. + name : str, default = None name of the module, currently used for debugging purposes. - softmax_type: str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' - softmax type as described in this paper: + softmax_type : str = {'vanilla', 'off-by-one', 'learnable'}, default = 'vanilla' + Softmax type as described in the paper `Efficient Streaming Language Models with Attention Sinks `_. - For a given attention score S = Q*K^T, of shape [b, h, s_q, s_kv], - 'vanilla': S[:,:,:,i] = exp(S[:,:,:,i])/sum(exp(S[:,:,:,:]), dim=-1), - 'off-by-one': S[:,:,:,i] = exp(S[:,:,:,i])/(1 + sum(exp(S[:,:,:,:]), dim=-1)), and - 'learnable': S[:,j,:,i] = exp(S[:,j,:,i])/(exp(alpha[j]) + sum(exp(S[:,j,:,:]), dim=-1)), - where alpha is a learnable parameter in shape [h]. - 'off-by-one' and 'learnable' softmax types are also called sink attention - ('zero sink' and 'learnable sink'). + + For a given attention score :math:`S = Q \cdot K^T`, of shape ``[b, h, s_q, s_kv]``: + + * ``'vanilla'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{\sum_j \exp(S_{:,:,:,j})} + + * ``'off-by-one'``: + + .. math:: + Softmax(S)_{:,:,:,i} = \frac{\exp(S_{:,:,:,i})}{1 + \sum_j \exp(S_{:,:,:,j})} + + * ``'learnable'``: + + .. math:: + Softmax(S)_{:,h,:,i} = \frac{\exp(S_{:,h,:,i})}{\exp(\alpha_h) + \sum_j \exp(S_{:,h,:,j})} + + where :math:`\\alpha` is a learnable parameter of shape ``[h]``. + + ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention + (``'zero sink'`` and ``'learnable sink'``). Parallelism parameters ---------------------- - set_parallel_mode : bool, default = `False` - if set to `True`, QKV and FC1 layers are used as Column Parallel + set_parallel_mode : bool, default = False + if set to ``True``, QKV and FC1 layers are used as Column Parallel whereas PROJ and FC2 is used as Row Parallel as described `here `_. - sequence_parallel : bool, default = `False` - if set to `True`, uses sequence parallelism. - tp_group : ProcessGroup, default = `None` + sequence_parallel : bool, default = False + if set to ``True``, uses sequence parallelism. + tp_group : ProcessGroup, default = None tensor parallel process group. tp_size : int, default = 1 used as TP (tensor parallel) world size when TP groups are not formed during initialization. In this case, users must call the - `set_tensor_parallel_group(tp_group)` method on the initialized module before the + :meth:`set_tensor_parallel_group` method on the initialized module before the forward pass to supply the tensor parallel group needed for tensor and sequence parallel collectives. Optimization parameters ----------------------- - fuse_wgrad_accumulation : bool, default = 'False' - if set to `True`, enables fusing of creation and accumulation of + fuse_wgrad_accumulation : bool, default = False + if set to ``True``, enables fusing of creation and accumulation of the weight gradient. When enabled, it is assumed that the weights - have an additional `main_grad` attribute (used instead of the - regular `grad`) which is a pre-allocated buffer of the correct + have an additional :attr:`main_grad` attribute (used instead of the + regular :attr:`grad`) which is a pre-allocated buffer of the correct size to accumulate gradients in. - params_dtype : torch.dtype, default = `torch.get_default_dtype()` + params_dtype : torch.dtype, default = torch.get_default_dtype() it controls the type used to allocate the initial parameters. Useful when the model is trained with lower precision and the original FP32 parameters would not fit in GPU memory. - seq_length: int + seq_length : int sequence length of input samples. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propogation and activation recompute phase. - micro_batch_size: int + micro_batch_size : int batch size per training step. Needed for JIT Warmup, a technique where jit fused functions are warmed up before training to ensure same kernels are used for forward propogation and activation recompute phase. - drop_path_rate: float, default = 0.0 + drop_path_rate : float, default = 0.0 when > 0.0, applies stochastic depth per sample in the main path of the residual block. - fuse_qkv_params: bool, default = 'False' - if set to `True`, `TransformerLayer` module exposes a single fused + fuse_qkv_params : bool, default = False + if set to ``True``, :class:`TransformerLayer` module exposes a single fused parameter for query-key-value. This enables optimizations such as QKV fusion without concatentations/splits and also enables the argument - `fuse_wgrad_accumulation`. - qk_norm_type: Optional[str], default = None + :attr:`fuse_wgrad_accumulation`. + qk_norm_type : Optional[str], default = None type of normalization to apply to query and key tensors. - Options: None, 'L2Normalization', 'RMSNorm', 'LayerNorm'. When None, no normalization is applied. - When 'L2Normalization', L2 normalization is applied to query and key tensors. - When 'RMSNorm', RMS normalization is applied to query and key tensors. - When 'LayerNorm', layer normalization is applied to query and key tensors. + Options: ``None``, ``'L2Normalization'``, ``'RMSNorm'``, ``'LayerNorm'``. When ``None``, no normalization is applied. + When ``'L2Normalization'``, L2 normalization is applied to query and key tensors. + When ``'RMSNorm'``, RMS normalization is applied to query and key tensors. + When ``'LayerNorm'``, layer normalization is applied to query and key tensors. Normalization is applied after RoPE (if applicable) but before attention computation - when `qk_norm_before_rope` is False. This follows the e.g. Llama4 approach for + when ``qk_norm_before_rope`` is ``False``. This follows the e.g. Llama4 approach for QK normalization to improve training stability and model performance. - qk_norm_eps: float, default = 1e-6 + qk_norm_eps : float, default = 1e-6 epsilon value for normalization of query and key tensors. - Only used when `qk_norm_type` is not None. - qk_norm_before_rope: bool, default = `False` - if set to `True`, query and key normalization is applied before rotary position - embedding. When `False` (default), normalization is applied after RoPE. + Only used when ``qk_norm_type`` is not ``None``. + qk_norm_before_rope : bool, default = False + if set to ``True``, query and key normalization is applied before rotary position + embedding. When ``False`` (default), normalization is applied after RoPE. This parameter allows supporting different architectural variants that apply QK normalization at different points. """ @@ -523,7 +538,7 @@ def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> N Parameters ---------- - tp_group : ProcessGroup, default = `None` + tp_group : ProcessGroup, default = None tensor parallel process group. """ # Deep iterate but skip self to avoid infinite recursion. @@ -549,7 +564,7 @@ def set_context_parallel_group( cp_stream: torch.cuda.Stream, cp_comm_type: str = "p2p", ) -> None: - """ + r""" Set the context parallel attributes for the given module before executing the forward pass. @@ -557,25 +572,26 @@ def set_context_parallel_group( ---------- cp_group : Union[ProcessGroup, List[ProcessGroup]] context parallel process group. - ProcessGroup is for cp_comm_type of "p2p", "all_gather", and "a2a". - List[ProcessGroup] is for cp_comm_type of "a2a+p2p", where cp_group[0] - and cp_group[1] are for a2a and p2p communications respectively. + ProcessGroup is for cp_comm_type of ``"p2p"``, ``"all_gather"``, and ``"a2a"``. + List[ProcessGroup] is for cp_comm_type of ``"a2a+p2p"``, where ``cp_group[0]`` + and ``cp_group[1]`` are for a2a and p2p communications respectively. cp_global_ranks : List[int] list of global ranks in the context group. cp_stream : torch.cuda.Stream cuda stream for context parallel execution. - cp_comm_type : str, default = `p2p` + cp_comm_type : str, default = "p2p" inter-gpu communication type for context parallelism. - Can be "p2p" or "all_gather" or "a2a", or "a2a+p2p". - "p2p": Exchange KV chunks with P2P communications in ring topology. - P2P is async and can be overlapped with attention compute. - "all_gather": All-gather to get full sequence of KV before attention. - The all-gather is not async, and cannot be overlapped. - "a2a": Like DeepSpeed Ulysses, scatter attention heads across the CP - group, and gather to get full sequence of QKV. - "a2a+p2p": hierarchical CP implementation. First applying a2a to QKV - across each CP sub-group (e.g., via NVLink), then exchanging KV with - p2p between sub-groups (e.g., via IBLink). + Can be ``"p2p"`` or ``"all_gather"`` or ``"a2a"`` or ``"a2a+p2p"``. + + - ``"p2p"``: Exchange KV chunks with P2P communications in ring topology. + P2P is async and can be overlapped with attention compute. + - ``"all_gather"``: All-gather to get full sequence of KV before attention. + The all-gather is not async, and cannot be overlapped. + - ``"a2a"``: Like DeepSpeed Ulysses, scatter attention heads across the CP + group, and gather to get full sequence of QKV. + - ``"a2a+p2p"``: hierarchical CP implementation. First applying a2a to QKV + across each CP sub-group (e.g., via NVLink), then exchanging KV with + p2p between sub-groups (e.g., via IBLink). """ # Deep iterate but skip self to avoid infinite recursion. for index, child in enumerate(self.modules()): @@ -610,49 +626,49 @@ def forward( fast_zero_fill: bool = True, pad_between_seqs: Optional[bool] = None, ) -> torch.Tensor: - """ + r""" Transformer Layer: attention block and a feedforward network (MLP) .. note:: Argument :attr:`attention_mask` is only used when :attr:`self_attn_mask_type` - includes `"padding"` or `"arbitrary"`. + includes ``"padding"`` or ``"arbitrary"``. Parameters ---------- hidden_states : torch.Tensor Input tensor. - attention_mask : Optional[torch.Tensor], default = `None` + attention_mask : Optional[torch.Tensor], default = None Boolean tensor used to mask out self-attention softmax input. It should be - in [batch_size, 1, 1, seqlen_q] for padding masks, and broadcastable - to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] for "`arbitrary`" - mask. It should be `None` for causal masks and "`no_mask`" type. - A `True` value means the corresponding position is masked out and - a `False` means that position is allowed to participate in attention. + in ``[batch_size, 1, 1, seqlen_q]`` for padding masks, and broadcastable + to ``[batch_size, num_heads, max_seqlen_q, max_seqlen_kv]`` for ``"arbitrary"`` + mask. It should be ``None`` for causal masks and ``"no_mask"`` type. + A ``True`` value means the corresponding position is masked out and + a ``False`` means that position is allowed to participate in attention. self_attn_mask_type: {'no_mask', 'causal', 'padding', 'padding_causal', 'causal_bottom_right', 'padding_causal_bottom_right','arbitrary'}, - default = `causal` + default = "causal" Type of attention mask passed into softmax operation for encoder. By default, causal masks are aligned to the top left corner of - the softmax matrix. When "`bottom_right`" is specified in the mask type, + the softmax matrix. When ``"bottom_right"`` is specified in the mask type, causal masks are aligned to the bottom right corner. - window_size: Optional[Tuple[int, int]], default = `None` + window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention in encoder. - encoder_output : Optional[torch.Tensor], default = `None` + encoder_output : Optional[torch.Tensor], default = None Output of the encoder block to be fed into the decoder block if using - `layer_type="decoder"`. + :attr:`layer_type` = ``"decoder"``. enc_dec_attn_mask : Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]], - default = `None`. Boolean tensors used to mask out inter-attention softmax input if - using `layer_type="decoder"`. It should be a tuple of two masks in - [batch_size, 1, 1, seqlen_q] and [batch_size, 1, 1, seqlen_kv] for padding masks. - It should be broadcastable to [batch_size, num_heads, max_seqlen_q, max_seqlen_kv] - for "`arbitrary`" mask. It should be `None` for causal masks and "`no_mask`". - A `True` value means the corresponding position is masked out and a `False` + default = None. Boolean tensors used to mask out inter-attention softmax input if + using :attr:`layer_type` = ``"decoder"``. It should be a tuple of two masks in + ``[batch_size, 1, 1, seqlen_q]`` and ``[batch_size, 1, 1, seqlen_kv]`` for padding masks. + It should be broadcastable to ``[batch_size, num_heads, max_seqlen_q, max_seqlen_kv]`` + for ``"arbitrary"`` mask. It should be ``None`` for causal masks and ``"no_mask"``. + A ``True`` value means the corresponding position is masked out and a ``False`` means that position is allowed to participate in attention. enc_dec_attn_mask_type: {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'}, - default = `None` + default = None Type of attention mask passed into softmax operation for decoder. - enc_dec_window_size: Optional[Tuple[int, int]], default = `None` + enc_dec_window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention in decoder. is_first_microbatch : {True, False, None}, default = None During training using either gradient accumulation or @@ -667,53 +683,53 @@ def forward( * it also allows skipping gradient accumulation during the first microbatch (since it is the first gradient being produced) - checkpoint_core_attention: bool, default = `False` - If true, forward activations for core attention are recomputed + checkpoint_core_attention: bool, default = False + If ``True``, forward activations for core attention are recomputed during the backward pass in order to save memory that would otherwise be occupied to store the forward activations until backprop. - rotary_pos_emb: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = `None` + rotary_pos_emb: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]], default = None Embeddings for query and key tensors for applying rotary position embedding. By default no input embedding is applied. - core_attention_bias_type: str, default = `no_bias` - Bias type, {`no_bias`, `pre_scale_bias`, `post_scale_bias`, `alibi`} - core_attention_bias: Optional[torch.Tensor], default = `None` - Bias tensor for Q * K.T - alibi_slopes: Optional[torch.Tensor], default = `None` - ALiBi slopes in FP32 and shape [nheads] or [batch_size, nheads]. - It adds a bias of (-alibi_slope * (i + seqlen_k - seqlen_q - j)) + core_attention_bias_type: str, default = "no_bias" + Bias type, {``"no_bias"``, ``"pre_scale_bias"``, ``"post_scale_bias"``, ``"alibi"``} + core_attention_bias: Optional[torch.Tensor], default = None + Bias tensor for :math:`Q \cdot K^T` + alibi_slopes: Optional[torch.Tensor], default = None + ALiBi slopes in FP32 and shape ``[nheads]`` or ``[batch_size, nheads]``. + It adds a bias of :math:`(-\text{alibi_slope} \cdot (i + \text{seqlen_k} - \text{seqlen_q} - j))` to the attention score of query i and key j. - cu_seqlens_q: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (without offset) in a batch for `query_layer`, - with shape [batch_size + 1] and dtype torch.int32. + cu_seqlens_q: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (without offset) in a batch for query layer, + with shape ``[batch_size + 1]`` and dtype torch.int32. Used by encoders, or decoders' self-attention. - cu_seqlens_kv: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (without offset) in a batch for `key_layer` - and `value_layer`, with shape [batch_size + 1] and dtype torch.int32. + cu_seqlens_kv: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (without offset) in a batch for key layer + and value layer, with shape ``[batch_size + 1]`` and dtype torch.int32. Used by decoders' cross-attention. - cu_seqlens_q_padded: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (with offset) in a batch for `query_layer`, - with shape [batch_size + 1] and dtype torch.int32. Set to `cu_seqlens_q` if None. + cu_seqlens_q_padded: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (with offset) in a batch for query layer, + with shape ``[batch_size + 1]`` and dtype torch.int32. Set to :attr:`cu_seqlens_q` if ``None``. Used by encoders, or decoders' self-attention. - cu_seqlens_kv_padded: Optional[torch.Tensor], default = `None` - Cumulative sum of sequence lengths (with offset) in a batch for `key_layer` - and `value_layer`, with shape [batch_size + 1] and dtype torch.int32. - Set to `cu_seqlens_kv` if None. Used by decoders' cross-attention. - max_seqlen_q: Optional[int], default = `None` - Maximum sequence length in `query_layer`. - Calculated from `cu_seqlens_q_padded` if not provided. - max_seqlen_kv: Optional[int], default = `None` - Maximum sequence length in `key_layer` and `value_layer`. - Calculated from `cu_seqlens_kv_padded` if not provided. - fast_zero_fill: bool, default = `True` + cu_seqlens_kv_padded: Optional[torch.Tensor], default = None + Cumulative sum of sequence lengths (with offset) in a batch for key layer + and value layer, with shape ``[batch_size + 1]`` and dtype torch.int32. + Set to :attr:`cu_seqlens_kv` if ``None``. Used by decoders' cross-attention. + max_seqlen_q: Optional[int], default = None + Maximum sequence length in query layer. + Calculated from :attr:`cu_seqlens_q_padded` if not provided. + max_seqlen_kv: Optional[int], default = None + Maximum sequence length in key layer and value layer. + Calculated from :attr:`cu_seqlens_kv_padded` if not provided. + fast_zero_fill: bool, default = True Whether to set output tensors to 0 or not before use. inference_params: InferenceParams, default = None Inference parameters that are passed to the main model in order to efficiently calculate and store the context during inference. - pad_between_seqs: Optional[bool], default = `None` - If None, inferred from qkv_format, cu_seqlens and cu_seqlens_padded. - If true, there are padding tokens between individual sequences in a packed batch, - i.e. qkv_format = 'thd'. + pad_between_seqs: Optional[bool], default = None + If ``None``, inferred from :attr:`qkv_format`, cu_seqlens and cu_seqlens_padded. + If ``True``, there are padding tokens between individual sequences in a packed batch, + i.e. :attr:`qkv_format` = ``'thd'``. """ if self_attn_mask_type is None: diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 741dd60c06..39d6fdaa6a 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -31,18 +31,18 @@ def make_row_id_map( Parameters ---------- - routing_map: torch.Tensor + routing_map : torch.Tensor Input tensor of shape `[num_tokens, num_experts]`. It is a mask tensor that indicates which experts are routed to which tokens. The values in it: 1 means the token is routed to this expert and 0 means not. - num_tokens: int + num_tokens : int Number of tokens in the input tensor. - num_experts: int + num_experts : int Number of experts in the input tensor. Returns ------- - row_id_map: torch.Tensor + row_id_map : torch.Tensor The row_id_map for the permutation of shape `[num_tokens, num_experts * 2 + 1]`. For each token, the last item is the number of experts that are routed (n_routed). The first n_routed items are the destination row indices in the permuted tokens. @@ -134,23 +134,23 @@ def permute_with_mask_map( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. - row_id_map: torch.Tensor + row_id_map : torch.Tensor The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. - probs: torch.Tensor + probs : torch.Tensor The probabilities of the input tensor. If it is not None, it will be permuted. - scale: torch.Tensor + scale : torch.Tensor The scale of the input tensor. If it is not None, it will be permuted. - num_tokens: int + num_tokens : int Number of tokens in the input tensor. - num_experts: int + num_experts : int Number of experts in the input tensor. - num_out_tokens: int + num_out_tokens : int Number of tokens in the permuted tensor. - hidden_size: int + hidden_size : int Hidden size of the input tensor. - scale_hidden_dim: int + scale_hidden_dim : int Hidden size of the scale tensor. """ output = torch.empty((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda") @@ -211,20 +211,20 @@ def unpermute_with_mask_map( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_out_tokens, hidden_size]`. - row_id_map: torch.Tensor + row_id_map : torch.Tensor The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. - merging_probs: torch.Tensor + merging_probs : torch.Tensor The merging probabilities of the input tensor. If it is not None, it will be used as weights to reduce the unpermuted tokens. - permuted_probs: torch.Tensor + permuted_probs : torch.Tensor The permuted probabilities of the input tensor. If it is not None, it will be unpermuted. - num_tokens: int + num_tokens : int Number of tokens in the permuted tensor. - num_experts: int + num_experts : int Number of experts in the permuted tensor. - hidden_size: int + hidden_size : int Hidden size of the permuted tensor. """ output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda") @@ -278,21 +278,21 @@ def unpermute_with_mask_map_bwd_with_merging_probs( Parameters ---------- - fwd_output_grad: torch.Tensor + fwd_output_grad : torch.Tensor The gradient of the output tensor of shape `[num_tokens, hidden_size]`. - row_id_map: torch.Tensor + row_id_map : torch.Tensor The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. - fwd_input: torch.Tensor + fwd_input : torch.Tensor The input tensor of the forward pass of shape `[num_out_tokens, hidden_size]`. - merging_probs: torch.Tensor + merging_probs : torch.Tensor The merging probabilities of the input tensor of shape `[num_tokens, num_experts]`. - num_tokens: int + num_tokens : int Number of tokens in the permuted tensor. - num_experts: int + num_experts : int Number of experts in the permuted tensor. - num_out_tokens: int + num_out_tokens : int Number of tokens in the output tensor. - hidden_size: int + hidden_size : int Hidden size of the output tensor. """ act_grad = torch.empty( @@ -339,13 +339,13 @@ def make_chunk_sort_map( Parameters ---------- - split_sizes: torch.Tensor + split_sizes : torch.Tensor The sizes of the chunks of shape `[num_splits,]`. - sorted_indices: torch.Tensor + sorted_indices : torch.Tensor The indices of the sorted chunks of shape `[num_splits,]`. - num_tokens: int + num_tokens : int Number of tokens in the input tensor. - num_splits: int + num_splits : int Number of splits of split_sizes and sorted_indices. """ row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device="cuda") @@ -373,17 +373,17 @@ def sort_chunks_by_map( Parameters ---------- - inp: torch.Tensor + inp : torch.Tensor Input tensor of shape `[num_tokens, hidden_size]`. - row_id_map: torch.Tensor + row_id_map : torch.Tensor The token to expert mapping tensor of shape `[num_tokens,]`. - probs: torch.Tensor + probs : torch.Tensor The probabilities of the input tensor. If it is not None, it will be permuted. - num_tokens: int + num_tokens : int Number of tokens in the input tensor. - hidden_size: int + hidden_size : int Hidden size of the input tensor. - is_forward: bool + is_forward : bool Whether the sort is for forward or backward. """ output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device="cuda") diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 083117b7b4..16e126493f 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -12,8 +12,8 @@ import numpy as np import torch -from . import torch_version from .quantized_tensor import Quantizer +from .torch_version import torch_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor @@ -601,7 +601,7 @@ def get_nvtx_range_context(msg: str): Parameters ---------- - msg: str + msg : str Message to associate with profiling context. """ @@ -619,7 +619,7 @@ def nvtx_range_push(msg: str) -> None: Parameters ---------- - msg: str + msg : str Message to associate with range """ @@ -637,7 +637,7 @@ def nvtx_range_pop(msg: Optional[str] = None) -> None: Parameters ---------- - msg: str, optional + msg : str, optional Message associated with range """ From 3ff0b8d4934f1459d2ca84a5cb312b9d678b1a5a Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 26 Nov 2025 17:21:47 -0800 Subject: [PATCH 104/521] Change Flax MHA to DPA to remove the duplicated QKV projection step (#2429) Signed-off-by: tdophung --- docs/examples/quickstart_jax.ipynb | 129 +++++++++++++++++++++-------- 1 file changed, 95 insertions(+), 34 deletions(-) diff --git a/docs/examples/quickstart_jax.ipynb b/docs/examples/quickstart_jax.ipynb index dc500afd1b..369cc371f6 100644 --- a/docs/examples/quickstart_jax.ipynb +++ b/docs/examples/quickstart_jax.ipynb @@ -53,7 +53,7 @@ }, { "cell_type": "code", - "execution_count": 1, + "execution_count": 26, "id": "d5284a38", "metadata": {}, "outputs": [], @@ -67,7 +67,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": 27, "id": "a4d1cfdc", "metadata": {}, "outputs": [], @@ -116,19 +116,33 @@ " 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", - " # Reshape to [batch, seq_len, num_heads * head_dim] for Flax MultiHeadDotProductAttention\n", - " q_reshaped = q.reshape(q.shape[0], q.shape[1], self.hidden_size)\n", - " k_reshaped = k.reshape(k.shape[0], k.shape[1], self.hidden_size)\n", - " v_reshaped = v.reshape(v.shape[0], v.shape[1], self.hidden_size)\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", - " # Attention using Flax's MultiHeadDotProductAttention\n", - " attention = nn.MultiHeadDotProductAttention(\n", - " num_heads=self.num_attention_heads,\n", - " qkv_features=self.kv_channels,\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", + " 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", - " x = attention(q_reshaped, k_reshaped, v_reshaped, mask=attention_mask, deterministic=deterministic)\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", " x = res + x\n", " \n", " # Second residual connection\n", @@ -157,7 +171,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 28, "id": "8b44649d", "metadata": {}, "outputs": [], @@ -178,7 +192,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 29, "id": "e44ed26d", "metadata": {}, "outputs": [ @@ -187,7 +201,7 @@ "output_type": "stream", "text": [ "Pure Flax FlaxTransformerLayer initialized successfully!\n", - "Parameter shapes: {'params': {'Dense_0': {'bias': (12288,), 'kernel': (4096, 12288)}, '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,)}, 'MultiHeadDotProductAttention_0': {'key': {'bias': (32, 4), 'kernel': (4096, 32, 4)}, 'out': {'bias': (4096,), 'kernel': (32, 4, 4096)}, 'query': {'bias': (32, 4), 'kernel': (4096, 32, 4)}, 'value': {'bias': (32, 4), 'kernel': (4096, 32, 4)}}}}\n" + "Parameter shapes: {'params': {'Dense_0': {'bias': (12288,), 'kernel': (4096, 12288)}, '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" ] } ], @@ -208,7 +222,7 @@ }, { "cell_type": "code", - "execution_count": 5, + "execution_count": 30, "id": "de91af7a", "metadata": {}, "outputs": [ @@ -234,7 +248,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": 31, "id": "037bc8d9", "metadata": {}, "outputs": [ @@ -242,7 +256,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 17.708301544189453 ms\n" + "Mean time: 18.546080589294434 ms\n" ] } ], @@ -290,7 +304,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": 32, "id": "bed20d6b", "metadata": {}, "outputs": [], @@ -309,7 +323,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": 33, "id": "56105579", "metadata": {}, "outputs": [], @@ -414,7 +428,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": 34, "id": "4b67511f", "metadata": {}, "outputs": [ @@ -422,7 +436,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 16.505107879638672 ms\n" + "Mean time: 16.375374794006348 ms\n" ] } ], @@ -456,15 +470,39 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": 35, "id": "5146cd99", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/code/github/TransformerEngine/transformer_engine/jax/flax/transformer.py:634: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n", + "/code/github/TransformerEngine/transformer_engine/jax/flax/transformer.py:742: UserWarning: Fused attention is not enabled because there is no available kernel.\n", + "Fall back to the unfused attention.\n", + "Please try to update the cuDNN and TE to the latest version.\n", + "self.dtype=\n", + "qkv_layout=>\n", + "attn_bias_type=>\n", + "attn_mask_type=>\n", + "self.attention_dropout=0.1\n", + "self.num_attention_heads=32\n", + "self.num_gqa_groups=32\n", + "seqlen_q=2048\n", + "seqlen_kv=2048\n", + "head_dim_qk=128\n", + "head_dim_v=128\n", + "\n", + " warnings.warn(\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 12.80329704284668 ms\n" + "Mean time: 12.403340339660645 ms\n" ] } ], @@ -515,7 +553,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": 36, "id": "c2eee376", "metadata": {}, "outputs": [], @@ -527,7 +565,7 @@ }, { "cell_type": "code", - "execution_count": 12, + "execution_count": 37, "id": "de96827c", "metadata": {}, "outputs": [ @@ -535,7 +573,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 9.615030288696289 ms\n" + "Mean time: 9.396424293518066 ms\n" ] } ], @@ -588,7 +626,7 @@ }, { "cell_type": "code", - "execution_count": 13, + "execution_count": 38, "id": "11203785", "metadata": {}, "outputs": [], @@ -659,7 +697,7 @@ }, { "cell_type": "code", - "execution_count": 15, + "execution_count": 39, "id": "6b0c705e", "metadata": {}, "outputs": [ @@ -667,7 +705,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 9.331779479980469 ms\n" + "Mean time: 9.145426750183105 ms\n" ] } ], @@ -704,10 +742,33 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": 40, "id": "b2aaa8ef", "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/code/github/TransformerEngine/transformer_engine/jax/flax/transformer.py:742: UserWarning: Fused attention is not enabled because there is no available kernel.\n", + "Fall back to the unfused attention.\n", + "Please try to update the cuDNN and TE to the latest version.\n", + "self.dtype=\n", + "qkv_layout=>\n", + "attn_bias_type=>\n", + "attn_mask_type=>\n", + "self.attention_dropout=0.1\n", + "self.num_attention_heads=32\n", + "self.num_gqa_groups=32\n", + "seqlen_q=2048\n", + "seqlen_kv=2048\n", + "head_dim_qk=128\n", + "head_dim_v=128\n", + "\n", + " warnings.warn(\n" + ] + } + ], "source": [ "\n", "te_transformer = te_flax.TransformerLayer(\n", @@ -731,7 +792,7 @@ }, { "cell_type": "code", - "execution_count": 17, + "execution_count": 41, "id": "b9cdbf22", "metadata": {}, "outputs": [ @@ -739,7 +800,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 9.23741340637207 ms\n" + "Mean time: 9.020795822143555 ms\n" ] } ], From f1512b21c73951a7606375cf4a4ecaa9fc6a2969 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 2 Dec 2025 12:17:01 -0500 Subject: [PATCH 105/521] [JAX] Triton binding (#2437) * init triton binding with test case/example * added Triton as TE-JAX test dependency * grid with blocksize from autotune Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- build_tools/jax.py | 2 +- tests/jax/test_triton_custom_calls.py | 115 ++++++ .../jax/triton_extensions/__init__.py | 25 ++ .../jax/triton_extensions/utils.py | 330 ++++++++++++++++++ 4 files changed, 471 insertions(+), 1 deletion(-) create mode 100644 tests/jax/test_triton_custom_calls.py create mode 100644 transformer_engine/jax/triton_extensions/__init__.py create mode 100644 transformer_engine/jax/triton_extensions/utils.py diff --git a/build_tools/jax.py b/build_tools/jax.py index 1f9552eb69..df78bf3e2f 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -20,7 +20,7 @@ def install_requirements() -> List[str]: def test_requirements() -> List[str]: """Test dependencies for TE/JAX extensions.""" - return ["numpy"] + return ["numpy", "triton"] def xla_path() -> str: diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py new file mode 100644 index 0000000000..071b8a73a2 --- /dev/null +++ b/tests/jax/test_triton_custom_calls.py @@ -0,0 +1,115 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Tests for Triton-based custom calls in TE JAX.""" + +import jax +import jax.numpy as jnp +import pytest + +from utils import assert_allclose, pytest_parametrize_wrapper + +import triton +import triton.language as tl + +from transformer_engine.jax.cpp_extensions.base import BasePrimitive, register_primitive +from transformer_engine.jax.triton_extensions import triton_call_lowering + + +@pytest.fixture(autouse=True, scope="module") +def init(): + """WAR for CUDA uninitialize error""" + _ = jnp.zeros(0) + yield + + +class TestTritonBinding: + """Test Triton binding primitive.""" + + # Define autotuned Triton kernel + @staticmethod + @triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 256}), # Uses defaults: num_warps=4, num_stages=3 + triton.Config({"BLOCK_SIZE": 512}, num_warps=8), # Custom num_warps + ], + key=["n_elements"], # Autotune based on input size + ) + @triton.jit + def amax_kernel( + x_ptr, + amax_ptr, + n_elements: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + ): + """Compute amax using Triton with autotuning.""" + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask, other=0.0) + abs_x = tl.abs(x) + block_max = tl.max(abs_x) + + tl.atomic_max(amax_ptr, block_max) + + # Define test primitive + class AmaxTritonPrimitive(BasePrimitive): + """Test primitive using Triton kernel.""" + + name = "te_amax_triton_test" + multiple_results = False + impl_static_args = () + + @staticmethod + def abstract(x_aval): + return jax.core.ShapedArray((1,), jnp.float32) + + @staticmethod + def impl(x): + assert TestTritonBinding.AmaxTritonPrimitive.inner_primitive is not None + return TestTritonBinding.AmaxTritonPrimitive.inner_primitive.bind(x) + + @staticmethod + def lowering(ctx, x): + """MLIR lowering using Triton kernel.""" + n_elements = 1 + for dim in ctx.avals_in[0].shape: + n_elements *= dim + + # For autotuned kernels, use the minimum BLOCK_SIZE from configs + # to ensure all elements are processed by all configs + block_size = min( + config.kwargs.get("BLOCK_SIZE") for config in TestTritonBinding.amax_kernel.configs + ) + grid = (triton.cdiv(n_elements, block_size),) + + return triton_call_lowering( + ctx, + TestTritonBinding.amax_kernel, # Autotuned kernel + x, + grid=grid, + constexprs={"n_elements": n_elements}, + # BLOCK_SIZE comes from autotuner config, not passed here + ) + + register_primitive(AmaxTritonPrimitive) + + @staticmethod + def _triton_amax(x: jnp.ndarray) -> jnp.ndarray: + """Compute amax using Triton kernel.""" + return TestTritonBinding.AmaxTritonPrimitive.outer_primitive.bind(x) + + @pytest_parametrize_wrapper("shape", [(1024, 1024)]) + @pytest_parametrize_wrapper("dtype", [jnp.bfloat16]) + def test_triton_amax(self, shape, dtype): + """Test Triton amax with JIT.""" + key = jax.random.PRNGKey(0) + x = jax.random.uniform(key, shape, dtype) + + expected = jnp.max(jnp.abs(x), keepdims=False).astype(jnp.float32) + jitted_amax = jax.jit(self._triton_amax) + result = jitted_amax(x) + + assert_allclose(result, expected, dtype=jnp.float32) diff --git a/transformer_engine/jax/triton_extensions/__init__.py b/transformer_engine/jax/triton_extensions/__init__.py new file mode 100644 index 0000000000..7ce6c476c2 --- /dev/null +++ b/transformer_engine/jax/triton_extensions/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +""" +Triton extensions for Transformer Engine JAX. + +This module provides Triton kernel integration for TE primitives. + +IMPORTANT: This module requires Triton to be installed. If you don't have Triton, +use transformer_engine.jax.cpp_extensions instead (CUDA/FFI based primitives). + +Install Triton: pip install triton + + +Usage: + # Import utilities + from transformer_engine.jax.triton_extensions import triton_call_lowering + + # Use in your primitive's lowering + @staticmethod + def lowering(ctx, x, **kwargs): + return triton_call_lowering(ctx, my_kernel, x, ...) +""" + +from .utils import * diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py new file mode 100644 index 0000000000..accb316fec --- /dev/null +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -0,0 +1,330 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +""" +Triton utilities for JAX primitives. + +This module provides utility functions for integrating Triton kernels into +JAX primitives. Triton is only imported when this module is used. +""" + +import hashlib +from typing import Any, Callable, Mapping +import zlib + +from jax import core +import jax +import jax.numpy as jnp + + +try: + from jax._src.lib import gpu_triton + from triton.compiler import compiler as tc + from triton.backends.nvidia import compiler as cb + from triton.runtime import autotuner +except ImportError as e: + raise ImportError( + "Triton is required for transformer_engine.jax.triton_extensions. " + "Install with: pip install triton\n" + "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead." + ) from e + + +__all__ = ["triton_call_lowering"] + +# Triton kernel cache (module-level, shared across all kernels) +_TRITON_KERNEL_CACHE = {} + + +def get_triton_dtype(aval): + """Convert JAX dtype to Triton type string. + + Args: + aval: JAX ShapedArray + + Returns: + Triton type string (e.g., "*fp32" for pointer, "i32" for scalar) + """ + dtype_map = { + jnp.dtype("bfloat16"): "bf16", + jnp.dtype("float64"): "fp64", + jnp.dtype("float32"): "fp32", + jnp.dtype("float16"): "fp16", + jnp.dtype("float8_e4m3fn"): "fp8e4nv", + jnp.dtype("float8_e5m2"): "fp8e5", + jnp.dtype("int64"): "i64", + jnp.dtype("int32"): "i32", + jnp.dtype("int16"): "i16", + jnp.dtype("int8"): "i8", + jnp.dtype("uint64"): "u64", + jnp.dtype("uint32"): "u32", + jnp.dtype("uint16"): "u16", + jnp.dtype("uint8"): "u8", + jnp.dtype("bool"): "i1", + } + + assert isinstance(aval, core.ShapedArray), "aval must be a JAX ShapedArray" + return f"*{dtype_map[aval.dtype]}" + + +def compile_triton( + kernel_fn: Callable, + signature: Mapping[str, str], + constants: Mapping[str, Any], + num_warps: int, + num_stages: int, + num_ctas: int, + compute_capability: int, + enable_fp_fusion: bool = False, +): + """Compile a Triton kernel to PTX. + + Kernels are cached to avoid recompilation. + + Args: + kernel_fn: Triton kernel function (decorated with @triton.jit) + signature: Dict mapping arg names to types (e.g., {"x_ptr": "*fp32", "n": "i32"}) + constants: Dict of compile-time constants + num_warps: Number of warps per block + num_stages: Number of pipeline stages + num_ctas: Number of CTAs (cooperative thread arrays) + compute_capability: CUDA compute capability + enable_fp_fusion: Enable FP fusion optimizations (default False for accuracy) + + Returns: + TritonKernel object for JAX + """ + # Create cache key + cache_key = hashlib.md5( + str( + ( + kernel_fn.__name__, + tuple(sorted(signature.items())), + tuple(sorted(constants.items())), + num_warps, + num_stages, + num_ctas, + enable_fp_fusion, + compute_capability, + ) + ).encode() + ).hexdigest() + + if cache_key in _TRITON_KERNEL_CACHE: + return _TRITON_KERNEL_CACHE[cache_key] + + # Compile kernel + options = cb.CUDAOptions( + num_warps=num_warps, + num_stages=num_stages, + num_ctas=num_ctas, + cluster_dims=(1, 1, 1), + debug=False, + enable_fp_fusion=enable_fp_fusion, + ) + + # Mark constants as constexpr in signature + signature_with_constexpr = dict(signature) + for const_name in constants.keys(): + if const_name in signature_with_constexpr: + signature_with_constexpr[const_name] = "constexpr" + + src = tc.ASTSource( + fn=kernel_fn, + constexprs=constants, + signature=signature_with_constexpr, + ) + + compiled = tc.compile( + src, + target=tc.GPUTarget("cuda", compute_capability, 32), + options=options.__dict__, + ) + + # Create kernel object for JAX + kernel = gpu_triton.TritonKernel( + compiled.name, + num_warps, + compiled.metadata.shared, + compiled.asm["ptx"], + "", # ttir + compute_capability, + 1, + 1, + 1, # cluster_dims + ) + + _TRITON_KERNEL_CACHE[cache_key] = kernel + return kernel + + +def triton_call_lowering( + ctx, + kernel_fn: Callable, + *array_args, + grid, + input_output_aliases: Mapping[int, int] = None, + constexprs: Mapping[str, Any] = None, +): + """Helper for MLIR lowering that calls a Triton kernel. + + Use this in your primitive's lowering method to call Triton kernels. + + Args: + ctx: MLIR lowering context + kernel_fn: Triton kernel function + *array_args: Input arrays (from ctx) + grid: Grid dimensions (int or tuple) + input_output_aliases: Mapping of input to output aliases + constexprs: Compile-time constants for the kernel + + Returns: + MLIR lowering result + + Example: + @staticmethod + def lowering(ctx, x, *, block_size): + from ..triton_extensions import triton_call_lowering + n = ctx.avals_in[0].size + return triton_call_lowering( + ctx, my_kernel, x, + grid=(triton.cdiv(n, block_size),), + n_elements=n, + BLOCK_SIZE=block_size + ) + """ + # Get compute capability using gpu_triton + compute_capability = gpu_triton.get_compute_capability(0) # device 0 + + # Build signature dict: map arg names to types + # Get arg names from kernel function + if isinstance(kernel_fn, autotuner.Autotuner): + arg_names = kernel_fn.fn.arg_names + else: + arg_names = kernel_fn.arg_names + + # Build signature for inputs + outputs + all_avals = list(ctx.avals_in) + list(ctx.avals_out) + signature = {arg_names[i]: get_triton_dtype(aval) for i, aval in enumerate(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] + + # Default values for the kernel + 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_ctas = 1 + kernel_constexprs = constexprs if constexprs is not None else {} + + # Handle autotuned kernels - compile all configs + if isinstance(kernel_fn, autotuner.Autotuner): + # Compile all configs for runtime selection + kernel_calls = [] + actual_kernel_fn = kernel_fn.fn + + for config in kernel_fn.configs: + # Extract parameters from config + config_num_warps = config.num_warps if config.num_warps is not None else num_warps + config_num_stages = config.num_stages if config.num_stages is not None else num_stages + config_num_ctas = config.num_ctas if config.num_ctas is not None else num_ctas + + # Merge config kwargs with user constexprs + config_constexprs = {**config.kwargs, **(constexprs if constexprs else {})} + + # Compile this config + config_kernel = compile_triton( + actual_kernel_fn, + signature, + config_constexprs, + config_num_warps, + config_num_stages, + config_num_ctas, + compute_capability, + enable_fp_fusion=False, + ) + + # Create kernel call for this config + config_params = [] + for _ in list(ctx.avals_in) + list(ctx.avals_out): + config_params.append(gpu_triton.create_array_parameter(0, 16)) + + config_call = gpu_triton.TritonKernelCall( + config_kernel, + grid_tuple[0], + grid_tuple[1], + grid_tuple[2], + config_params, + ) + + kernel_calls.append((config_call, str(config))) + + # Create autotuned kernel call + # Convert input_output_aliases to format with sizes + if input_output_aliases is None: + input_output_aliases = {} + + input_output_aliases_with_sizes = tuple( + ( + input_idx, + output_idx, + ctx.avals_in[input_idx].size * ctx.avals_in[input_idx].dtype.itemsize, + ) + for input_idx, output_idx in input_output_aliases.items() + ) + + kernel_call = gpu_triton.TritonAutotunedKernelCall( + f"{actual_kernel_fn.__name__}_autotuned", + kernel_calls, + input_output_aliases_with_sizes, + ) + + else: + # Regular kernel: compile single config + kernel = compile_triton( + actual_kernel_fn, + signature, + kernel_constexprs, + num_warps, + num_stages, + num_ctas, + compute_capability, + enable_fp_fusion=False, + ) + + kernel_params = [] + for _ in list(ctx.avals_in) + list(ctx.avals_out): + kernel_params.append(gpu_triton.create_array_parameter(0, 16)) + + kernel_call = gpu_triton.TritonKernelCall( + kernel, + grid_tuple[0], + grid_tuple[1], + grid_tuple[2], + kernel_params, + ) + + serialized_metadata = b"" + call_proto = kernel_call.to_proto(actual_kernel_fn.__name__, serialized_metadata) + + if input_output_aliases is None: + input_output_aliases = {} + + # 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=input_output_aliases, + ) + + return rule(ctx, *array_args) From 14b53313e576c2bce5e420ec4dd881e484c655bb Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 2 Dec 2025 13:00:13 -0500 Subject: [PATCH 106/521] [Common] NVTEGroupedTensor class and helpers (#2388) * add grouped_tensor classes and helpers Signed-off-by: Phuong Nguyen * rm non-contiguous option and dptrs Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address comments + rework CheckIn/OutputGroupedTensor Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix for compilation Signed-off-by: Phuong Nguyen * make first_dims/last_dims optional + data.shape 2d Signed-off-by: Phuong Nguyen * added assertion Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rs conflicts Signed-off-by: Phuong Nguyen * add data.shape info Signed-off-by: Phuong Nguyen * added logical shape field Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * compilation fix Signed-off-by: Phuong Nguyen * fixed issues raised by greptile Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * return default dtype when grouped_tensor is empty Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * use has_data() for dim queries Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update comments Signed-off-by: Phuong Nguyen * fix index bound Signed-off-by: Phuong Nguyen * Update transformer_engine/common/transformer_engine.cpp Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Phuong Nguyen * Update transformer_engine/common/transformer_engine.cpp Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Phuong Nguyen * restore Tensor.has_data() + add experimental marks Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * restore Tensor::has_columnwise_data Signed-off-by: Phuong Nguyen * cleanup Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen 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> --- transformer_engine/common/common.h | 136 +++++++ .../transformer_engine/transformer_engine.h | 108 ++++++ .../common/transformer_engine.cpp | 334 ++++++++++++++++++ 3 files changed, 578 insertions(+) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 97b130952d..661f8b00e1 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -101,6 +101,7 @@ struct SimpleTensor { } return acc; } + bool has_data() const noexcept { return dptr != nullptr && numel() > 0; } void clear() { dptr = nullptr; @@ -154,9 +155,11 @@ struct Tensor { return acc; } + // TODO(Tim): Change this to use data.has_data() bool has_data() const noexcept { return data.dptr != nullptr; } // Check for size (not just pointer) for 0-dim or no token cases. + // TODO(Tim): Change this to use columnwise_data.has_data() bool has_columnwise_data() const noexcept { return columnwise_data.dptr != nullptr || columnwise_data.shape.size() != 0; } @@ -281,6 +284,129 @@ struct Tensor { } }; +struct GroupedTensor { + public: + /* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ + /* + Grouped tensor is a collection of tensors with different shapes but the same dtype and scaling mode + + Shape Representation: + - logical_shape: 2D shape representing the conceptual layouy, i.e. the shape when member tensors are flattened to 2D and stacked together (REQUIRED) + + When all_same_shape(): [num_tensors * M, N] where each tensor is (M, N) + + When varying_first_dim(): [~sum_of_first_dims, N] where N is common + + When varying_last_dim(): [M, ~sum_of_last_dims] where M is common + + When varying_both_dims(): [1, total_elements] (fully flattened) + + - first_dims and last_dims are OPTIONAL (empty if dimension is uniform) + + Empty first_dims: all tensors have the same first dimension + + Empty last_dims: all tensors have the same last dimension + + Both empty: all tensors have identical shapes + + Both set: each tensor has unique shape (first_dims[i], last_dims[i]) + + Data Layout: + - ALL data fields are stored as 1D flattened arrays (data, columnwise_data, scale_inv, etc.) + - logical_shape provides the conceptual 2D interpretation + - All data is stored on device in contiguous layout + */ + + SimpleTensor data; + SimpleTensor columnwise_data; + SimpleTensor scale_inv; + SimpleTensor columnwise_scale_inv; + SimpleTensor amax; + SimpleTensor columnwise_amax; + SimpleTensor scale; // for FP8-DS only + + // Shape information (OPTIONAL - empty if dimension is uniform across all tensors) + // first_dims[i] = first dimension of tensor i (empty if all tensors have same first dim) + // last_dims[i] = last dimension of tensor i (empty if all tensors have same last dim) + SimpleTensor first_dims; // Device pointer to int64_t array of length num_tensors (or empty) + SimpleTensor last_dims; // Device pointer to int64_t array of length num_tensors (or empty) + + // Offsets for indexing into contiguous 1D layout (OPTIONAL - not needed if all_same_shape()) + // tensor_offsets[i] = element offset to start of tensor i (cumulative sum of numel for tensors 0..i-1) + // Usage: tensor_i_ptr = (char*)data.dptr + tensor_offsets[i] * element_size + // If empty and all_same_shape(): offset[i] = i * M * N (where M, N are common dimensions) + SimpleTensor tensor_offsets; // Device pointer to int64_t array of length num_tensors (or empty) + + // Logical shape: conceptual 2D shape of the grouped data (REQUIRED) + // Represents how the 1D flattened data should be interpreted as 2D + // Always 2D with positive dimensions + NVTEShape logical_shape; + + NVTEScalingMode scaling_mode; + size_t num_tensors; + NVTEGroupedTensor nvte_tensor; + + GroupedTensor(NVTEScalingMode scaling_mode, size_t num_tensors) + : data(), + columnwise_data(), + scale_inv(), + columnwise_scale_inv(), + amax(), + columnwise_amax(), + scale(), + num_tensors(num_tensors), + first_dims(nullptr, {}, DType::kInt64), + last_dims(nullptr, {}, DType::kInt64), + tensor_offsets(nullptr, {}, DType::kInt64), + logical_shape(nvte_make_shape(nullptr, 0)), + scaling_mode(scaling_mode), + nvte_tensor(0) {} + + explicit operator NVTEGroupedTensor() const noexcept { return nvte_tensor; } + + bool has_data() const noexcept { return data.has_data(); } + bool has_columnwise_data() const noexcept { return columnwise_data.has_data(); } + + bool all_same_first_dim() const noexcept { return !first_dims.has_data(); } + bool all_same_last_dim() const noexcept { return !last_dims.has_data(); } + bool all_same_shape() const noexcept { return !first_dims.has_data() && !last_dims.has_data(); } + bool varying_both_dims() const noexcept { return first_dims.has_data() && last_dims.has_data(); } + + size_t get_common_first_dim() const { + NVTE_CHECK(all_same_first_dim(), "First dim varies across tensors"); + NVTE_CHECK(logical_shape.ndim == 2, "Logical shape must be 2D"); + if (all_same_shape()) { + // When both dims are uniform: logical_shape = [num_tensors * M, N] + return logical_shape.data[0] / num_tensors; + } else { + // When varying last dims but not first dim: logical_shape = [M, sum_of_last_dims] + return logical_shape.data[0]; + } + } + size_t get_common_last_dim() const { + NVTE_CHECK(all_same_last_dim(), "Last dim varies across tensors"); + NVTE_CHECK(logical_shape.ndim == 2, "Logical shape must be 2D"); + // For both uniform and varying first dim cases: logical_shape[1] is the common last dim + return logical_shape.data[1]; + } + + DType dtype() const { + if (has_data()) return data.dtype; + if (has_columnwise_data()) return columnwise_data.dtype; + // Fallback, used e.g. in workspace or when allow_empty=true + return data.dtype; + } + + void clear() { + data.clear(); + columnwise_data.clear(); + scale_inv.clear(); + columnwise_scale_inv.clear(); + amax.clear(); + columnwise_amax.clear(); + scale.clear(); + first_dims.clear(); + last_dims.clear(); + tensor_offsets.clear(); + logical_shape = nvte_make_shape(nullptr, 0); + num_tensors = 0; + scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + nvte_tensor = 0; + } +}; + struct QuantizationConfig { bool force_pow_2_scales = false; float amax_epsilon = 0.0f; @@ -779,6 +905,16 @@ std::vector> convert_tensor_array(NVTETensor **nvte_tensor Tensor *convertNVTETensor(const NVTETensor tensor); Tensor *convertNVTETensorCheck(const NVTETensor tensor); + +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, + bool allow_empty = false); + } // namespace transformer_engine #endif // TRANSFORMER_ENGINE_COMMON_COMMON_H_ diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 1a901ab82d..76cc636a35 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -393,6 +393,114 @@ int nvte_is_non_tn_fp8_gemm_supported(); */ void nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream); +/*! \brief TE Grouped Tensor type + * + * NVTEGroupedTensor is a collection of tensors with potentially different shapes + * but the same dtype and scaling mode. It does not own the memory it points to. + */ +typedef void *NVTEGroupedTensor; + +/*! \enum NVTEGroupedTensorParam + * \brief Indicates the kind of the grouped tensor parameter to set/get. + */ +enum NVTEGroupedTensorParam { + kNVTEGroupedRowwiseData = 0, /*!< Data usable in rowwise manner */ + kNVTEGroupedColumnwiseData = 1, /*!< Data usable in columnwise manner */ + kNVTEGroupedScale = 2, /*!< Scale tensor */ + kNVTEGroupedAmax = 3, /*!< Amax tensor */ + kNVTEGroupedRowwiseScaleInv = 4, /*!< Scale inverse tensor for decoding Rowwise Data */ + kNVTEGroupedColumnwiseScaleInv = 5, /*!< Scale inverse tensor for decoding Columnwise Data */ + kNVTEGroupedColumnwiseAmax = 6, /*!< Columnwise Amax tensor */ + kNVTEGroupedFirstDims = 7, /*!< First dimension sizes (device pointer to int64_t array) */ + kNVTEGroupedLastDims = 8, /*!< Last dimension sizes (device pointer to int64_t array) */ + kNVTEGroupedTensorOffsets = + 9, /*!< Tensor offsets for contiguous layout (device pointer to int64_t array) */ + kNVTENumGroupedTensorParams +}; + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Create a new TE grouped tensor. + * + * Create a new TE grouped tensor. Before use its parameters need to be set. + * TE grouped tensors are just wrappers on top of raw data and do not + * own memory. + * + * \param[in] scaling_mode Scaling mode of the grouped tensor. + * \param[in] num_tensors Number of tensors in the group (must be > 0). + * \param[in] logical_shape Logical 2D shape of the grouped data. + * + * \return A new TE grouped tensor. + */ +NVTEGroupedTensor nvte_create_grouped_tensor(NVTEScalingMode scaling_mode, size_t num_tensors, + NVTEShape logical_shape); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Destroy a TE grouped tensor. + * + * Since the TE grouped tensor does not own memory, the underlying + * data is not freed during this operation. + * + * \param[in] tensor Grouped tensor to be destroyed. + */ +void nvte_destroy_grouped_tensor(NVTEGroupedTensor tensor); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Set a parameter of the grouped tensor. + * + * \param[in/out] tensor Grouped tensor. + * \param[in] param_name The parameter to be set. + * \param[in] param The value to be set (NVTEBasicTensor). + */ +void nvte_set_grouped_tensor_param(NVTEGroupedTensor *tensor, NVTEGroupedTensorParam param_name, + const NVTEBasicTensor *param); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Get a value of the parameter of the grouped tensor. + * + * \param[in] tensor Grouped tensor. + * \param[in] param_name The parameter to be queried. + * + * \return NVTEBasicTensor containing the parameter data. + */ +NVTEBasicTensor nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, + NVTEGroupedTensorParam param_name); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Get the number of tensors in a grouped tensor. + * + * \param[in] tensor Grouped tensor. + * + * \return Number of tensors in the group. + */ +size_t nvte_grouped_tensor_num_tensors(const NVTEGroupedTensor tensor); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Get a grouped tensor's data type. + * + * \param[in] tensor Grouped tensor. + * + * \return A data type of the grouped tensor. + */ +NVTEDType nvte_grouped_tensor_type(const NVTEGroupedTensor tensor); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Get a scaling mode of the grouped tensor. + * + * \param[in] tensor Grouped tensor. + * + * \return Scaling mode of the grouped tensor. + */ +NVTEScalingMode nvte_grouped_tensor_scaling_mode(const NVTEGroupedTensor tensor); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Get the logical shape of a grouped tensor. + * + * \param[in] tensor Grouped tensor. + * + * \return Logical 2D shape. + */ +NVTEShape nvte_get_grouped_tensor_logical_shape(const NVTEGroupedTensor tensor); + #ifdef __cplusplus } // extern "C" diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 314ba3b40f..e9e1d5bfb1 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -273,6 +273,128 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt CheckScaleTensorShape(t, name); } +void CheckGroupedTensorShapeArrays(const GroupedTensor &t, const std::string &name) { + NVTE_CHECK(t.num_tensors > 0, "Grouped tensor ", name, " has no tensors!"); + + // Helper lambda to validate shape arrays + // All three arrays are OPTIONAL: + // - first_dims: empty if all tensors have same first dimension + // - last_dims: empty if all tensors have same last dimension + // - tensor_offsets: empty if all tensors have same shape (offsets are predictable) + auto check_shape_array = [&](const SimpleTensor &arr, const char *arr_name) { + if (arr.has_data()) { + NVTE_CHECK(arr.shape.size() == 1, "Grouped tensor ", name, " ", arr_name, " must be 1D"); + NVTE_CHECK(arr.dtype == DType::kInt64, "Grouped tensor ", name, " ", arr_name, + " must have dtype Int64"); + NVTE_CHECK(arr.shape[0] == t.num_tensors, "Grouped tensor ", name, " ", arr_name, " size (", + arr.shape[0], ") must equal num_tensors (", t.num_tensors, ")"); + } + }; + + // Validate shape arrays (all optional) + check_shape_array(t.first_dims, "first_dims"); + check_shape_array(t.last_dims, "last_dims"); + check_shape_array(t.tensor_offsets, "tensor_offsets"); + + // tensor_offsets is required if any dimension varies + // (i.e., required unless all_same_shape()) + if (!t.all_same_shape()) { + NVTE_CHECK( + t.tensor_offsets.dptr != nullptr, "Grouped tensor ", name, + " must have tensor_offsets when any dimension varies (first_dims or last_dims is set)"); + } + + // Validate logical_shape + NVTE_CHECK(t.logical_shape.ndim == 2, "Grouped tensor ", name, " logical_shape must be 2D"); + NVTE_CHECK(t.logical_shape.data[0] > 0 && t.logical_shape.data[1] > 0, "Grouped tensor ", name, + " logical_shape must have positive dimensions"); + + // Validate all data fields are 1D (flattened) + if (t.has_data()) { + NVTE_CHECK(t.data.shape.size() == 1, "Grouped tensor ", name, " data must be 1D"); + } + if (t.has_columnwise_data()) { + NVTE_CHECK(t.columnwise_data.shape.size() == 1, "Grouped tensor ", name, + " columnwise_data must be 1D"); + } + + // Validate data size matches logical_shape + size_t expected_numel = t.logical_shape.data[0] * t.logical_shape.data[1]; + if (t.has_data()) { + NVTE_CHECK(t.data.numel() == expected_numel, "Grouped tensor ", name, " data size (", + t.data.numel(), ") must match logical_shape size (", expected_numel, ")"); + } + if (t.has_columnwise_data()) { + NVTE_CHECK(t.columnwise_data.numel() == expected_numel, "Grouped tensor ", name, + " columnwise_data size (", t.columnwise_data.numel(), + ") must match logical_shape size (", expected_numel, ")"); + } +} + +// Helper function to check scale_inv for both input and output +static void CheckGroupedScaleInv(const GroupedTensor &t, const std::string &name, bool is_output) { + const char *tensor_type = is_output ? "output" : "input"; + + // Helper to check scale_inv for both rowwise and columnwise layouts + auto check_scales = [&](DType expected_dtype) { + if (t.has_data()) { + NVTE_CHECK(t.scale_inv.has_data(), tensor_type, " ", name, + " rowwise scale_inv must be allocated"); + NVTE_CHECK(t.scale_inv.dtype == expected_dtype, tensor_type, " ", name, + " rowwise scale_inv has invalid dtype (expected ", to_string(expected_dtype), + ", got ", to_string(t.scale_inv.dtype), ")"); + } + if (t.has_columnwise_data()) { + NVTE_CHECK(t.columnwise_scale_inv.has_data(), tensor_type, " ", name, + " columnwise scale_inv must be allocated"); + NVTE_CHECK(t.columnwise_scale_inv.dtype == expected_dtype, tensor_type, " ", name, + " columnwise scale_inv has invalid dtype (expected ", to_string(expected_dtype), + ", got ", to_string(t.columnwise_scale_inv.dtype), ")"); + } + }; + + // 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_mxfp8_scaling(t.scaling_mode)) { + check_scales(DType::kFloat8E8M0); + } else if (is_nvfp4_scaling(t.scaling_mode)) { + check_scales(DType::kFloat8E4M3); + } else { + // Non-quantized types should not have scale/scale_inv + NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv not supported for non-quantized ", tensor_type, + " ", name); + NVTE_CHECK(!t.columnwise_scale_inv.has_data(), "Scale_inv not supported for non-quantized ", + tensor_type, " ", name); + } +} + +void CheckInputGroupedTensor(const GroupedTensor &t, const std::string &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) { + if (!allow_empty) { + NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Output grouped tensor ", name, + " not allocated"); + } + + // Only perform dtype-specific validation if data is allocated + if (t.has_data() || t.has_columnwise_data()) { + // Amax validation for delayed scaling + if (is_fp8_dtype(t.dtype()) && t.scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { + NVTE_CHECK(t.amax.has_data(), "Output ", name, " amax must be allocated"); + NVTE_CHECK(t.amax.dtype == DType::kFloat32, "Output ", name, " amax must be Float32"); + } + CheckGroupedScaleInv(t, name, true); + } + + CheckGroupedTensorShapeArrays(t, name); +} + class TensorAllocator { public: static TensorAllocator &instance() { @@ -387,6 +509,89 @@ Tensor *convertNVTETensorCheck(const NVTETensor t) { return ptr; } +// GroupedTensor allocator - similar pattern to TensorAllocator +class GroupedTensorAllocator { + public: + static GroupedTensorAllocator &instance() { + static GroupedTensorAllocator allocator; + return allocator; + } + + ~GroupedTensorAllocator() {} + + NVTEGroupedTensor Allocate(NVTEScalingMode mode, size_t num_tensors, NVTEShape logical_shape) { + std::lock_guard lock(mutex); + if (!free_list.empty()) { + uintptr_t index = free_list.back(); + NVTEGroupedTensor ret = reinterpret_cast(index); + free_list.pop_back(); + // 1-based indexing - fully reinitialize the tensor to avoid stale data + memory[index - 1].scaling_mode = mode; + memory[index - 1].num_tensors = num_tensors; + memory[index - 1].logical_shape = logical_shape; + memory[index - 1].nvte_tensor = ret; + return ret; + } + if (memory.size() < memory.capacity()) { + memory.emplace_back(mode, num_tensors); + GroupedTensor &t = memory.back(); + size = memory.size(); + // 1-based indexing + uintptr_t index = memory.size(); + t.logical_shape = logical_shape; + t.nvte_tensor = reinterpret_cast(index); + return reinterpret_cast(index); + } + NVTE_ERROR( + "Cannot allocate a new NVTEGroupedTensor. Maximum number of grouped tensors reached: ", + MAX_GROUPED_TENSOR_NUM, ". There is probably a memory leak in your application."); + } + + void Free(NVTEGroupedTensor t) { + std::lock_guard lock(mutex); + uintptr_t index = reinterpret_cast(t); + if (index == 0) return; + NVTE_CHECK(index <= memory.size(), "Invalid grouped tensor."); + free_list.push_back(index); + // Clean up + memory[index - 1].clear(); + } + + GroupedTensor *convertNVTEGroupedTensor(NVTEGroupedTensor t) { + uintptr_t index = reinterpret_cast(t); + // 1-based indexing to enable 0-initialization of NVTEGroupedTensor + // to be invalid tensor + static_assert(nullptr == 0); + if (index != 0 && index <= size) { + return &(memory[index - 1]); + } + return nullptr; + } + + private: + GroupedTensorAllocator() { + std::lock_guard lock(mutex); + memory.reserve(MAX_GROUPED_TENSOR_NUM); + } + + std::mutex mutex; + std::atomic size; + // Allocate at most 20 MB for grouped tensors + const size_t MAX_GROUPED_TENSOR_NUM = 20 * 1024 * 1024 / sizeof(GroupedTensor); + std::vector free_list; + std::vector memory; +}; + +GroupedTensor *convertNVTEGroupedTensor(const NVTEGroupedTensor t) { + return GroupedTensorAllocator::instance().convertNVTEGroupedTensor(t); +} + +GroupedTensor *convertNVTEGroupedTensorCheck(const NVTEGroupedTensor t) { + GroupedTensor *ptr = GroupedTensorAllocator::instance().convertNVTEGroupedTensor(t); + NVTE_CHECK(ptr != nullptr, "Invalid grouped tensor."); + return ptr; +} + } // namespace transformer_engine NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode) { @@ -730,3 +935,132 @@ int nvte_is_non_tn_fp8_gemm_supported() { }); return cache[device_id]; } + +// Grouped Tensor C API implementations +NVTEGroupedTensor nvte_create_grouped_tensor(NVTEScalingMode scaling_mode, size_t num_tensors, + NVTEShape logical_shape) { + NVTE_CHECK(num_tensors > 0, "Number of tensors must be greater than 0"); + NVTE_CHECK(logical_shape.ndim == 2, "Logical shape must be 2D"); + NVTE_CHECK(logical_shape.data[0] > 0 && logical_shape.data[1] > 0, + "Logical shape must have positive dimensions"); + NVTEGroupedTensor ret = transformer_engine::GroupedTensorAllocator::instance().Allocate( + scaling_mode, num_tensors, logical_shape); + return ret; +} + +void nvte_destroy_grouped_tensor(NVTEGroupedTensor tensor) { + transformer_engine::GroupedTensorAllocator::instance().Free(tensor); +} + +void nvte_set_grouped_tensor_param(NVTEGroupedTensor *tensor, NVTEGroupedTensorParam param_name, + const NVTEBasicTensor *param) { + NVTE_CHECK(tensor != nullptr, "Grouped tensor pointer can't be NULL."); + auto *t = transformer_engine::convertNVTEGroupedTensor(*tensor); + NVTE_CHECK(t != nullptr, "Grouped tensor is not allocated."); + NVTE_CHECK(param != nullptr, "Grouped tensor param can't be NULL."); + + switch (param_name) { + case kNVTEGroupedRowwiseData: + t->data = *param; + break; + case kNVTEGroupedColumnwiseData: + t->columnwise_data = *param; + break; + case kNVTEGroupedScale: + t->scale = *param; + break; + case kNVTEGroupedAmax: + t->amax = *param; + break; + case kNVTEGroupedRowwiseScaleInv: + t->scale_inv = *param; + break; + case kNVTEGroupedColumnwiseScaleInv: + t->columnwise_scale_inv = *param; + break; + case kNVTEGroupedColumnwiseAmax: + t->columnwise_amax = *param; + break; + case kNVTEGroupedFirstDims: + t->first_dims = *param; + // Validate it's Int64 + NVTE_CHECK(t->first_dims.dtype == transformer_engine::DType::kInt64, + "first_dims must have dtype Int64"); + break; + case kNVTEGroupedLastDims: + t->last_dims = *param; + // Validate it's Int64 + NVTE_CHECK(t->last_dims.dtype == transformer_engine::DType::kInt64, + "last_dims must have dtype Int64"); + break; + case kNVTEGroupedTensorOffsets: + t->tensor_offsets = *param; + // Validate it's Int64 + NVTE_CHECK(t->tensor_offsets.dtype == transformer_engine::DType::kInt64, + "tensor_offsets must have dtype Int64"); + break; + default: + NVTE_ERROR("Unknown grouped tensor parameter!"); + } +} + +NVTEBasicTensor nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, + NVTEGroupedTensorParam param_name) { + if (tensor == nullptr) { + return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 0)}; + } + const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); + + switch (param_name) { + case kNVTEGroupedRowwiseData: + return t.data; + case kNVTEGroupedColumnwiseData: + return t.columnwise_data; + case kNVTEGroupedScale: + return t.scale; + case kNVTEGroupedAmax: + return t.amax; + case kNVTEGroupedRowwiseScaleInv: + return t.scale_inv; + case kNVTEGroupedColumnwiseScaleInv: + return t.columnwise_scale_inv; + case kNVTEGroupedColumnwiseAmax: + return t.columnwise_amax; + case kNVTEGroupedFirstDims: + return t.first_dims; + case kNVTEGroupedLastDims: + return t.last_dims; + case kNVTEGroupedTensorOffsets: + return t.tensor_offsets; + default: + NVTE_ERROR("Unknown grouped tensor parameter!"); + } +} + +size_t nvte_grouped_tensor_num_tensors(const NVTEGroupedTensor tensor) { + auto *t = transformer_engine::convertNVTEGroupedTensor(tensor); + if (t == nullptr) return 0; + return t->num_tensors; +} + +NVTEDType nvte_grouped_tensor_type(const NVTEGroupedTensor tensor) { + auto *t = transformer_engine::convertNVTEGroupedTensor(tensor); + if (t == nullptr) return kNVTEFloat32; + return static_cast(t->dtype()); +} + +NVTEScalingMode nvte_grouped_tensor_scaling_mode(const NVTEGroupedTensor tensor) { + if (tensor == nullptr) { + return NVTE_DELAYED_TENSOR_SCALING; + } + const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); + return t.scaling_mode; +} + +NVTEShape nvte_get_grouped_tensor_logical_shape(const NVTEGroupedTensor tensor) { + if (tensor == nullptr) { + return nvte_make_shape(nullptr, 0); + } + const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); + return t.logical_shape; +} From cc42a5771431ad29abb21c6a8215928a964717e7 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 2 Dec 2025 10:07:41 -0800 Subject: [PATCH 107/521] [JAX] Make test_layer.py tolerances stricter (#2306) * wip Signed-off-by: Jeremy Berchtold * wip Signed-off-by: Jeremy Berchtold * revert change to utils.py Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- tests/jax/test_layer.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/jax/test_layer.py b/tests/jax/test_layer.py index 2fc9f688ab..ca07e15742 100644 --- a/tests/jax/test_layer.py +++ b/tests/jax/test_layer.py @@ -552,7 +552,7 @@ def test_forward_with_fp8(self, data_shape, dtype, attrs, fp8_recipe): """Test forward with fp8 enabled""" # Empty MeshResource is used as we are running on a single device with autocast(enabled=True, recipe=fp8_recipe, mesh_resource=MeshResource()): - self.runner(attrs).test_forward(data_shape, dtype, rtol=1e-4, atol=1e-3) + self.runner(attrs).test_forward(data_shape, dtype) @pytest.mark.skipif(not is_fp8_supported, reason=reason) @pytest.mark.parametrize("fp8_recipe", QUANTIZE_RECIPES) @@ -560,7 +560,7 @@ def test_backward_with_fp8(self, data_shape, dtype, attrs, fp8_recipe): """Test backward with fp8 enabled""" # Empty MeshResource is used as we are running on a single device with autocast(enabled=True, recipe=fp8_recipe, mesh_resource=MeshResource()): - self.runner(attrs).test_backward(data_shape, dtype, rtol=1e-4, atol=1e-3) + self.runner(attrs).test_backward(data_shape, dtype) class TestEncoderLayer(BaseTester): From d126cdd6c0a8d6ce0419dcf1ffe027ef7aadf7e8 Mon Sep 17 00:00:00 2001 From: Kunlun Li <94586211+kunlunl@users.noreply.github.com> Date: Wed, 3 Dec 2025 03:24:31 +0800 Subject: [PATCH 108/521] Add primary weighs fp8 support for mxfp8 (#2055) * Add primary weighs fp8 support for mxfp8 Signed-off-by: kunlunl * Fix unit test and add better error log to unit test Signed-off-by: kunlunl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move post all-gather processing out of for loop Signed-off-by: kunlunl * Add descriptions and ASCII diagrams for partial cast and partial amax functions Signed-off-by: kunlunl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Minor fix based on greptile bot Signed-off-by: kunlunl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix compilation errors due to arch-specific PTX instructions Signed-off-by: Tim Moon * Remove unused noop flag from C API Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Expose test_partial_cast Signed-off-by: kunlunl * Skip mxfp8 partial cast test if mxfp8 is not available Signed-off-by: kunlunl * Fix pytest error Signed-off-by: kunlunl * pylint ignore unused manual_post_all_gather_processing Signed-off-by: kunlunl * Fix error when using is_mxfp8_available Signed-off-by: kunlunl --------- Signed-off-by: kunlunl Signed-off-by: Tim Moon 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 --- qa/L0_pytorch_unittest/test.sh | 1 + .../test_cast_master_weights_to_fp8.py | 168 ++++++++---- tests/pytorch/test_multi_tensor.py | 34 +++ tests/pytorch/test_partial_cast.py | 137 ++++++++++ transformer_engine/common/CMakeLists.txt | 15 +- .../include/transformer_engine/multi_tensor.h | 15 ++ .../include/transformer_engine/recipe.h | 183 +++++++++++++ .../common/multi_tensor/compute_scale.cu | 48 ++++ .../common/recipe/mxfp8_scaling.cu | 253 ++++++++++++++++++ transformer_engine/pytorch/csrc/extensions.h | 12 + ..._partial_cast.cpp => fp8_partial_cast.cpp} | 38 +++ .../extensions/multi_tensor/compute_scale.cpp | 10 + .../pytorch/csrc/extensions/pybind.cpp | 13 + transformer_engine/pytorch/tensor/utils.py | 141 +++++++++- 14 files changed, 1005 insertions(+), 63 deletions(-) create mode 100644 tests/pytorch/test_partial_cast.py create mode 100644 transformer_engine/common/recipe/mxfp8_scaling.cu rename transformer_engine/pytorch/csrc/extensions/{fp8_block_scaling_partial_cast.cpp => fp8_partial_cast.cpp} (53%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index e1ce680094..f1a48e421b 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -49,6 +49,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint 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" if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 0ff98e6cb7..51b920eab5 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -18,6 +18,7 @@ DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, Format, Recipe, ) @@ -25,9 +26,11 @@ from transformer_engine.pytorch import ( is_fp8_available, is_fp8_block_scaling_available, + is_mxfp8_available, QuantizedTensor, Float8Tensor, Float8BlockwiseQTensor, + MXFP8Tensor, ) from transformer_engine.pytorch.tensor import cast_master_weights_to_fp8 from transformer_engine.pytorch.tensor.utils import post_all_gather_processing, replace_raw_data @@ -42,17 +45,21 @@ def _get_quantization_recipe(quantization) -> Recipe: return Float8CurrentScaling(fp8_format=fp8_format) elif quantization == "fp8_block": return Float8BlockScaling(fp8_format=fp8_format) + elif quantization == "mxfp8": + return MXFP8BlockScaling() else: raise ValueError(f"Unsupported quantization: {quantization}") -def _get_raw_data(quantized_tensor): +def _get_raw_data(quantized_tensor, colwise=False): """Get the underlying data of a quantized tensor, used in zero-1 optimizer""" if isinstance(quantized_tensor, Float8Tensor): + assert not colwise, "Float8Tensor does not support get colwise data" assert hasattr(quantized_tensor, "_data"), "Float8Tensor does not have _data attribute" assert quantized_tensor._data.dtype == torch.uint8, "Float8Tensor _data must be uint8" return quantized_tensor._data elif isinstance(quantized_tensor, Float8BlockwiseQTensor): + assert not colwise, "Float8BlockwiseQTensor does not support get colwise data" assert hasattr( quantized_tensor, "_rowwise_data" ), "Float8BlockwiseQTensor does not have _rowwise_data attribute" @@ -60,6 +67,23 @@ def _get_raw_data(quantized_tensor): quantized_tensor._rowwise_data.dtype == torch.uint8 ), "Float8BlockwiseQTensor _rowwise_data must be uint8" return quantized_tensor._rowwise_data + elif isinstance(quantized_tensor, MXFP8Tensor): + if colwise: + assert hasattr( + quantized_tensor, "_columnwise_data" + ), "MXFP8Tensor does not have columnwise_data attribute" + assert ( + quantized_tensor._columnwise_data.dtype == torch.uint8 + ), "MXFP8Tensor columnwise_data must be uint8" + return quantized_tensor._columnwise_data + else: + assert hasattr( + quantized_tensor, "_rowwise_data" + ), "MXFP8Tensor does not have rowwise_data attribute" + assert ( + quantized_tensor._rowwise_data.dtype == torch.uint8 + ), "MXFP8Tensor rowwise_data must be uint8" + return quantized_tensor._rowwise_data else: raise ValueError(f"Unsupported quantized tensor type: {type(quantized_tensor)}") @@ -229,38 +253,43 @@ def step(self): end = start_offset + master_weight.numel() weight.data.view(-1)[start:end].copy_(master_weight) - # ----------------------------------------------------------------------------------------- - # Step 5: Copy the updated weights (not all weights) to the weight buffer - # ----------------------------------------------------------------------------------------- - for i in range(len(self.weights)): - master_weight = self.master_weights[i] - if master_weight is None: - continue - start_offset = self.start_offsets[i] - if isinstance(self.weights[i], QuantizedTensor): - weight = _get_raw_data(self.weights[i]) - else: - weight = self.weights[i] - weight_slice = weight.view(-1)[start_offset : start_offset + master_weight.numel()] - overlapping_start, overlapping_end = self.overlapping_areas[i] - self.weight_buffer[overlapping_start:overlapping_end].copy_(weight_slice) + colwise_list = [False] + if isinstance(self.weights[0], MXFP8Tensor): + colwise_list.append(True) - # ----------------------------------------------------------------------------------------- - # Step 6: Weight all-gather (FP8 or BF16) - # ----------------------------------------------------------------------------------------- - dist.all_gather_into_tensor( - self.weight_buffer, self.weight_buffer_slice, group=self.dp_group - ) + for colwise in colwise_list: + # ------------------------------------------------------------------------------------- + # Step 5: Copy the updated weights (not all weights) to the weight buffer + # ------------------------------------------------------------------------------------- + for i in range(len(self.weights)): + master_weight = self.master_weights[i] + if master_weight is None: + continue + start_offset = self.start_offsets[i] + if isinstance(self.weights[i], QuantizedTensor): + weight = _get_raw_data(self.weights[i], colwise) + else: + weight = self.weights[i] + weight_slice = weight.view(-1)[start_offset : start_offset + master_weight.numel()] + overlapping_start, overlapping_end = self.overlapping_areas[i] + self.weight_buffer[overlapping_start:overlapping_end].copy_(weight_slice) + + # ------------------------------------------------------------------------------------- + # Step 6: Weight all-gather (FP8 or BF16) + # ------------------------------------------------------------------------------------- + dist.all_gather_into_tensor( + self.weight_buffer, self.weight_buffer_slice, group=self.dp_group + ) - # ----------------------------------------------------------------------------------------- - # Step 7: Copy the gathered weights from weight buffer to the actual weights - # ----------------------------------------------------------------------------------------- - for weight, offset in zip(self.weights, self.offsets[:-1]): - start = offset - end = offset + weight.numel() - if isinstance(weight, QuantizedTensor): - weight = _get_raw_data(weight) - weight.view(-1).data.copy_(self.weight_buffer[start:end]) + # ------------------------------------------------------------------------------------- + # Step 7: Copy the gathered weights from weight buffer to the actual weights + # ------------------------------------------------------------------------------------- + for weight, offset in zip(self.weights, self.offsets[:-1]): + start = offset + end = offset + weight.numel() + if isinstance(weight, QuantizedTensor): + weight = _get_raw_data(weight, colwise) + weight.view(-1).data.copy_(self.weight_buffer[start:end]) if self.manual_post_all_gather_processing: quantized_weights = [ @@ -285,9 +314,15 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals else: raw_data_list = [w.view(-1) for w in weights] self.flatten_weight, original_length = self._flatten_tensors_with_pad(raw_data_list) + if isinstance(weights[0], MXFP8Tensor): + self.flatten_columnwise = self.flatten_weight.clone() + else: + self.flatten_columnwise = None # Split flattened weights into shards self.local_weight_shard = torch.chunk(self.flatten_weight, world_size)[rank] + if self.flatten_columnwise is not None: + self.local_columnwise_shard = torch.chunk(self.flatten_columnwise, world_size)[rank] self.local_main_grad_shard = torch.zeros_like( self.local_weight_shard, dtype=torch.float32, device="cuda" ) @@ -319,14 +354,25 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals self.shard_indices.append((None, None)) if isinstance(weights[idx], QuantizedTensor): - replace_raw_data( - weights[idx], self.flatten_weight[start:end].view(weights[idx].shape) - ) + if self.flatten_columnwise is not None: + new_rowwise_data = self.flatten_weight[start:end].view(weights[idx].shape) + new_rowwise_data.copy_(weights[idx]._rowwise_data) + weights[idx]._rowwise_data = new_rowwise_data + new_columnwise_data = self.flatten_columnwise[start:end].view( + weights[idx].shape + ) + new_columnwise_data.copy_(weights[idx]._columnwise_data) + weights[idx]._columnwise_data = new_columnwise_data + else: + replace_raw_data( + weights[idx], self.flatten_weight[start:end].view(weights[idx].shape) + ) else: weights[idx].data = self.flatten_weight[start:end].view(weights[idx].shape) # Initialize local model weights and high-precision master weights self.local_weights = [] + self.local_columnwise = [] self.master_weights = [] for i, weight in enumerate(self.weights): weight_start, weight_end = self.weight_indices[i] @@ -334,6 +380,11 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals if shard_start is not None and shard_end is not None: local_weight_shard = self.local_weight_shard[shard_start:shard_end] self.local_weights.append(local_weight_shard) + if self.flatten_columnwise is not None: + local_columnwise_shard = self.local_columnwise_shard[shard_start:shard_end] + else: + local_columnwise_shard = None + self.local_columnwise.append(local_columnwise_shard) if isinstance(weight, QuantizedTensor): high_precision_init_val = weight.get_high_precision_init_val().view(-1) @@ -345,6 +396,7 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals self.master_weights.append(master_weight_shard) else: self.local_weights.append(None) + self.local_columnwise.append(None) self.master_weights.append(None) setattr( weight, "main_grad", torch.zeros_like(weight, dtype=torch.float32, device="cuda") @@ -415,12 +467,12 @@ def step(self): # Step 3: Cast master weights to FP8 or BF16 precision if isinstance(self.weights[0], QuantizedTensor): local_weights = [] - for local_weight in self.local_weights: - if local_weight is None: - local_weights.append(None) - continue - - local_weights.append(local_weight) + for i, local_weight in enumerate(self.local_weights): + if self.flatten_columnwise is not None: + local_columnwise = self.local_columnwise[i] + local_weights.append((local_weight, local_columnwise)) + else: + local_weights.append(local_weight) cast_master_weights_to_fp8( self.weights, @@ -442,6 +494,10 @@ def step(self): dist.all_gather_into_tensor( self.flatten_weight, self.local_weight_shard, group=self.dp_group ) + if self.flatten_columnwise is not None: + dist.all_gather_into_tensor( + self.flatten_columnwise, self.local_columnwise_shard, group=self.dp_group + ) if self.manual_post_all_gather_processing: quantized_weights = [ @@ -513,15 +569,15 @@ def _test_cast_master_weights_to_fp8(quantization, dp_group, manual_post_all_gat preserve_high_precision_init_val=True, ): model_fp8 = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(128, 256 + 32, **linear_kwargs), + te.Linear(256 + 32, 256 * 3, **linear_kwargs), te.Linear(256 * 3, 128, **linear_kwargs), ) # Create model with BF16 weights model = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(128, 256 + 32, **linear_kwargs), + te.Linear(256 + 32, 256 * 3, **linear_kwargs), te.Linear(256 * 3, 128, **linear_kwargs), ) @@ -546,7 +602,7 @@ def _test_cast_master_weights_to_fp8(quantization, dp_group, manual_post_all_gat w.main_grad.zero_() inputs = [ - torch.randn(16, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) + torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) ] # Choose based on rank to make sure the inputs of different ranks are different. x = inputs[rank] @@ -577,7 +633,9 @@ def _test_cast_master_weights_to_fp8(quantization, dp_group, manual_post_all_gat optimizer_fp8.step() optimizer.step() - torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) + assert torch.allclose( + loss_fp8, loss, atol=0, rtol=0 + ), f"Loss mismatch at rank {rank}, step {i} for {quantization}" def _test_fsdp_cast_master_weights_to_fp8( @@ -609,15 +667,15 @@ def _test_fsdp_cast_master_weights_to_fp8( preserve_high_precision_init_val=True, ): model_fp8 = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(128, 256 + 32, **linear_kwargs), + te.Linear(256 + 32, 256 * 3, **linear_kwargs), te.Linear(256 * 3, 128, **linear_kwargs), ) # Create model with BF16 weights model = nn.Sequential( - te.Linear(128, 256 + 16, **linear_kwargs), - te.Linear(256 + 16, 256 * 3, **linear_kwargs), + te.Linear(128, 256 + 32, **linear_kwargs), + te.Linear(256 + 32, 256 * 3, **linear_kwargs), te.Linear(256 * 3, 128, **linear_kwargs), ) @@ -631,12 +689,12 @@ def _test_fsdp_cast_master_weights_to_fp8( ) optimizer = MiniFSDP([w for w in model.parameters()], 10.0, dp_group) - for _ in range(100): + for i in range(100): optimizer_fp8.zero_grad() optimizer.zero_grad() inputs = [ - torch.randn(16, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) + torch.randn(32, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) ] # Choose based on rank to make sure the inputs of different ranks are different. x = inputs[rank] @@ -667,7 +725,9 @@ def _test_fsdp_cast_master_weights_to_fp8( optimizer_fp8.step() optimizer.step() - torch.testing.assert_close(loss_fp8, loss, atol=0, rtol=0) + assert torch.allclose( + loss_fp8, loss, atol=0, rtol=0 + ), f"Loss mismatch at rank {rank}, step {i} for {quantization} (FSDP)" def run_parallel_tests() -> None: @@ -698,6 +758,8 @@ def run_parallel_tests() -> None: quantizations.extend(["fp8", "fp8_cs"]) if is_fp8_block_scaling_available(): quantizations.append("fp8_block") + if is_mxfp8_available(): + quantizations.append("mxfp8") manual_post_all_gather_processings = [False, True] diff --git a/tests/pytorch/test_multi_tensor.py b/tests/pytorch/test_multi_tensor.py index 46ba821879..94012354db 100644 --- a/tests/pytorch/test_multi_tensor.py +++ b/tests/pytorch/test_multi_tensor.py @@ -7,6 +7,7 @@ import transformer_engine.pytorch import transformer_engine_torch as tex +from transformer_engine.pytorch import is_mxfp8_available from transformer_engine.pytorch.optimizers import MultiTensorApply from references.quantize_scale_calc import scale_from_amax_tensor @@ -23,6 +24,7 @@ (555, 33333), ] appliers = [MultiTensorApply(2048 * 32), MultiTensorApply(333), MultiTensorApply(33333)] +mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) @pytest.mark.parametrize("input_size_pair", input_size_pairs) @@ -259,3 +261,35 @@ def test_multi_tensor_compute_scale_and_scale_inv( ) torch.testing.assert_close(scale, scale_ref, rtol=0, atol=0) torch.testing.assert_close(scale_inv, scale_inv_ref, rtol=0, atol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("input_size_pair", input_size_pairs + [(1, 1)]) +@pytest.mark.parametrize("applier", appliers) +@pytest.mark.parametrize("repeat", [1, 55]) +def test_multi_tensor_compute_scale_inv_e8m0(input_size_pair, applier, repeat): + sizea, sizeb = input_size_pair + device = torch.device("cuda") + a = torch.randn([sizea], dtype=torch.bfloat16, device=device).abs() + b = torch.randn([sizeb], dtype=torch.bfloat16, device=device).abs() + + amax_list = [] + for _ in range(repeat): + amax_list += [a.clone(), b.clone()] + scale_inv_list = [torch.empty_like(x).to(torch.uint8) for x in amax_list] + + applier( + tex.multi_tensor_compute_scale_inv_e8m0, + None, # overflow_buf + [amax_list, scale_inv_list], + ) + + max_fp8 = torch.finfo(torch.float8_e4m3fn).max + for amax, scale_inv in zip(amax_list, scale_inv_list): + scale_inv_u32 = (amax.float() / max_fp8).view(torch.int) + exponent = scale_inv_u32 // 2**23 + mantissa = scale_inv_u32 & 0x7FFFFF + exponent += ( + ((mantissa > 0) & (exponent != 0xFE)) & ~((exponent == 0) & (mantissa <= 0x400000)) + ).to(torch.int) + torch.testing.assert_close(exponent.to(torch.uint8), scale_inv) diff --git a/tests/pytorch/test_partial_cast.py b/tests/pytorch/test_partial_cast.py new file mode 100644 index 0000000000..cb0c4d75bd --- /dev/null +++ b/tests/pytorch/test_partial_cast.py @@ -0,0 +1,137 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine_torch import multi_tensor_compute_scale_inv_e8m0 +from transformer_engine.pytorch import is_mxfp8_available +from transformer_engine.pytorch.optimizers.multi_tensor_apply import multi_tensor_applier + + +mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) + + +def compute_partial_amax_reference(inp, amax_rowwise, amax_colwise, h, w, start_offset): + n = inp.view(-1).size(0) + if n == h * w: + full = inp.view(-1) + else: + full = torch.zeros(h * w, dtype=inp.dtype, device=inp.device) + full[start_offset : start_offset + n].copy_(inp) + full = torch.abs(full) + _amax_rowwise, _ = torch.max(full.view(h, w // 32, 32), dim=2) + amax_rowwise[:h, : (w // 32)].copy_(_amax_rowwise) + _amax_colwise, _ = torch.max(full.view(h // 32, 32, w), dim=1) + amax_colwise[: (h // 32), :w].copy_(_amax_colwise) + + +def partial_cast_reference( + inp, rowwise_out, colwise_out, rowwise_inv_scale, colwise_inv_scale, h, w, start_offset +): + rowwise_scale = ((254 - rowwise_inv_scale.int()) * 2**23).view(torch.float32) + colwise_scale = ((254 - colwise_inv_scale.int()) * 2**23).view(torch.float32) + n = inp.view(-1).size(0) + if n == h * w: + full = inp + else: + full = torch.empty(h * w, dtype=inp.dtype, device=inp.device) + full[start_offset : start_offset + n].copy_(inp) + full = full.float() + rowwise_scale = rowwise_scale[:h, : (w // 32)].contiguous().float() + colwise_scale = colwise_scale[: (h // 32), :w].contiguous().float() + scaled = (full.view(-1, 32) * rowwise_scale.view(-1, 1)).view(-1) + rowwise_out.copy_( + scaled[start_offset : start_offset + n].to(torch.float8_e4m3fn).view(rowwise_out.dtype) + ) + scaled = (full.view(h // 32, 32, w) * colwise_scale.view(h // 32, 1, w)).view(-1) + colwise_out.copy_( + scaled[start_offset : start_offset + n].to(torch.float8_e4m3fn).view(colwise_out.dtype) + ) + + +def run_one_case(n, h, w, start_offset): + inp = torch.randn(n, dtype=torch.bfloat16, device="cuda") + + rowwise_padding = [128, 4] + colwise_padding = [4, 128] + + def _pad(x, padding): + return (x + padding - 1) // padding * padding + + rowwise_shape = [_pad(h, rowwise_padding[0]), _pad(w // 32, rowwise_padding[1])] + colwise_shape = [_pad(h // 32, colwise_padding[0]), _pad(w, colwise_padding[1])] + + # Partial amax cuda kernel + amax_rowwise = torch.zeros(*rowwise_shape, dtype=inp.dtype, device=inp.device) + amax_colwise = torch.zeros(*colwise_shape, dtype=inp.dtype, device=inp.device) + tex.mxfp8_scaling_compute_partial_amax(inp, amax_rowwise, amax_colwise, h, w, start_offset) + + # Partial amax pytorch reference + amax_rowwise_ref = torch.zeros(*rowwise_shape, dtype=inp.dtype, device=inp.device) + amax_colwise_ref = torch.zeros(*colwise_shape, dtype=inp.dtype, device=inp.device) + compute_partial_amax_reference(inp, amax_rowwise_ref, amax_colwise_ref, h, w, start_offset) + + # Check partial amax + torch.testing.assert_close(amax_rowwise, amax_rowwise_ref, atol=0, rtol=0) + torch.testing.assert_close(amax_colwise, amax_colwise_ref, atol=0, rtol=0) + + # Calculate scales and scale_invs + scale_inv_rowwise = torch.empty_like(amax_rowwise).to(torch.uint8) + scale_inv_colwise = torch.empty_like(amax_colwise).to(torch.uint8) + multi_tensor_applier( + multi_tensor_compute_scale_inv_e8m0, + None, + [ + [amax_rowwise, amax_colwise], + [scale_inv_rowwise, scale_inv_colwise], + ], + ) + + # Partial cast cuda kernel + output_rowwise = torch.empty_like(inp).to(torch.uint8) + output_colwise = torch.empty_like(inp).to(torch.uint8) + tex.mxfp8_scaling_partial_cast( + inp, + output_rowwise, + output_colwise, + scale_inv_rowwise, + scale_inv_colwise, + h, + w, + start_offset, + ) + + # Partial cast pytorch reference + output_rowwise_ref = torch.empty_like(inp).to(torch.uint8) + output_colwise_ref = torch.empty_like(inp).to(torch.uint8) + partial_cast_reference( + inp, + output_rowwise_ref, + output_colwise_ref, + scale_inv_rowwise, + scale_inv_colwise, + h, + w, + start_offset, + ) + + # Check partial cast results + torch.testing.assert_close(output_rowwise, output_rowwise_ref, atol=0, rtol=0) + torch.testing.assert_close(output_colwise, output_colwise_ref, atol=0, rtol=0) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +def test_mxfp8_scaling_partial_cast(): + torch.cuda.manual_seed(1234) + + run_one_case(3, 32, 64, 31) + run_one_case(64 * 64 - 2, 64, 64, 1) + run_one_case(16384 * 6144, 16384, 6144, 0) + run_one_case(32768, 256, 128, 0) + run_one_case(131072, 768, 256, 0) + run_one_case(65536, 768, 256, 131072) + run_one_case(98304, 128, 768, 0) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index d3532b8c45..264f7f9a78 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -125,7 +125,6 @@ list(APPEND transformer_engine_cpp_sources list(APPEND transformer_engine_cuda_sources common.cu multi_tensor/adam.cu - multi_tensor/compute_scale.cu multi_tensor/l2norm.cu multi_tensor/scale.cu multi_tensor/sgd.cu @@ -167,16 +166,18 @@ list(APPEND transformer_engine_cuda_sources comm_gemm_overlap/userbuffers/userbuffers.cu) list(APPEND transformer_engine_cuda_arch_specific_sources - gemm/cutlass_grouped_gemm.cu - cast/cast.cu activation/gelu.cu activation/relu.cu activation/swiglu.cu - transpose/quantize_transpose_square_blockwise.cu - transpose/quantize_transpose_vector_blockwise_fp4.cu - hadamard_transform/hadamard_transform.cu + cast/cast.cu + gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu - hadamard_transform/hadamard_transform_cast_fusion.cu) + hadamard_transform/hadamard_transform.cu + hadamard_transform/hadamard_transform_cast_fusion.cu + multi_tensor/compute_scale.cu + recipe/mxfp8_scaling.cu + transpose/quantize_transpose_square_blockwise.cu + transpose/quantize_transpose_vector_blockwise_fp4.cu) # Compiling the files with the worst compilation time first to hopefully overlap # better with the faster-compiling cpp files diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index af3f51d46f..03d35dc2ed 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -265,6 +265,21 @@ void nvte_multi_tensor_compute_scale_and_scale_inv_cuda(int chunk_size, NVTETens float max_fp8, int force_pow_2_scales, float epsilon, cudaStream_t stream); +/*! \brief Compute E8M0 scale_inv for a list of tensors. + * + * \warning This API is **experimental** and subject to change. + * + * \param[in] chunk_size Number of tensor elements processed by a CUDA block. + * \param[in,out] tensor_lists 2D array of input tensors. + * \param[in] num_tensor_lists Size (dim0) of tensor_lists. + * \param[in] num_tensors_per_list Size (dim1) of tensor_lists. + * \param[in] stream CUDA stream used for this operation. + */ +void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor **tensor_lists, + const size_t num_tensor_lists, + const size_t num_tensors_per_list, + cudaStream_t stream); + /*! \brief Split a tensor along dimension 0 and compute the amax for each split. * * This function is experimental and the API is not stable. diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 6e1e9dd7ac..b1773a8db3 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -111,17 +111,200 @@ void nvte_compute_amax_with_config(const NVTETensor input, NVTETensor output, void nvte_compute_scale_from_amax(NVTETensor output, const NVTEQuantizationConfig config, cudaStream_t stream); +/*! \brief Compute partial amax for FP8 blockwise scaling. + * + * This function computes the maximum absolute values for each block of the original tensor. + * `inp` contains a continuous segment from the flattened original tensor. For each block, + * if it overlaps with the range [start_offset, start_offset+inp.length), the amax is + * computed from inp; otherwise, the amax is set to 0. + * + * Example: Original tensor (logically 512x512) divided into 16 blocks of size 128x128. + * `inp` contains continuous elements starting from position start_offset + * in the flattened original tensor. + * + * Logical view - Original Tensor (e.g., 512x512) divided into 16 blocks of size 128x128: + * ┌─────────┬─────────┬─────────┬─────────┐ + * │ Block0 │ Block1 │ Block2 │ Block3 │ Each block: 128x128 + * │ 128x128 │ 128x128 │ 128x128 │ 128x128 │ + * ├─────────┼─────────┼─────────┼─────────┤ + * │ Block4 │ Block5 │ Block6 │ Block7 │ + * ├─────────┼─────────┼─────────┼─────────┤ + * │ Block8 │ Block9 │ Block10 │ Block11 │ + * ├─────────┼─────────┼─────────┼─────────┤ + * │ Block12 │ Block13 │ Block14 │ Block15 │ + * └─────────┴─────────┴─────────┴─────────┘ + * + * Physical view - Flattened in row-major order: + * ┌────────────────────────────────────────────────────────────────┐ + * │[0...128][128...256][256...384][384...512]...[261632...262143] │ + * └────────────────────────────────────────────────────────────────┘ + * ^ ^ + * start_offset start_offset + inp.length + * + * For each 128x128 block, compute amax: + * - If the block overlaps with [start_offset, start_offset+inp.length), compute amax + * - If the block is completely outside this range, set amax = 0 + * + * amax output (one value per 128x128 block), block 1 and block 2 are non-zero because they + * overlap with the [start_offset, start_offset+inp.length) range: + * ┌───────┬───────┬───────┬───────┐ + * │ 0 │ amax │ amax │ 0 │ Block0-3 + * ├───────┼───────┼───────┼───────┤ + * │ 0 │ 0 │ 0 │ 0 │ Block4-7 + * ├───────┼───────┼───────┼───────┤ + * │ 0 │ 0 │ 0 │ 0 │ Block8-11 + * ├───────┼───────┼───────┼───────┤ + * │ 0 │ 0 │ 0 │ 0 │ Block12-15 + * └───────┴───────┴───────┴───────┘ + * + * \param[in] inp Input tensor (continuous slice of flattened original tensor). + * \param[in,out] amax Output tensor for maximum absolute values per block. + * \param[in] h Height dimension of the logical tensor. + * \param[in] w Width dimension of the logical tensor. + * \param[in] amax_stride_h Stride in height dimension for amax tensor. + * \param[in] amax_stride_w Stride in width dimension for amax tensor. + * \param[in] start_offset Starting offset in the flattened tensor. + * \param[in] block_len Length of a quantization block to process. + * \param[in] stream CUDA stream used for the operation. + */ void nvte_fp8_block_scaling_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h, size_t w, size_t amax_stride_h, size_t amax_stride_w, size_t start_offset, size_t block_len, cudaStream_t stream); +/*! \brief Perform partial FP8 casting with blockwise scaling. + * + * This function casts the input tensor to FP8 format using blockwise scaling factors. + * `inp` contains a continuous segment from the flattened original tensor. + * + * \param[in] inp Input tensor. + * \param[out] out Output tensor in FP8 format. + * \param[in] scale Scaling factors per block. + * \param[in] h Height dimension of the tensor. + * \param[in] w Width dimension of the tensor. + * \param[in] scale_stride_h Stride in height dimension for scale tensor. + * \param[in] scale_stride_w Stride in width dimension for scale tensor. + * \param[in] start_offset Starting offset for partial computation. + * \param[in] block_len Length of the block to process. + * \param[in] out_dtype Output FP8 datatype. + * \param[in] stream CUDA stream used for the operation. + */ void nvte_fp8_block_scaling_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, size_t block_len, const NVTEDType out_dtype, cudaStream_t stream); +/*! \brief Compute partial amax for MXFP8 scaling. + * + * This function computes the maximum absolute values along both row and column dimensions. + * input contains a continuous segment from the flattened original tensor. For each row/column + * block, if it overlaps with the range starting from start_offset, the amax is computed from + * `input`; otherwise, the amax is set to 0. + * + * Example: Original tensor (64 rows x 64 cols). + * Rowwise amax granularity: 1x32 (each row divided into 2 blocks) + * Columnwise amax granularity: 32x1 (each column divided into 2 blocks) + * input contains a continuous segment starting from start_offset. + * + * Logical view - Original Tensor (64x64) with 1x32 and 32x1 blocks: + * + * Rowwise blocks (1x32): Each row has 2 blocks + * ┌──────────────┬──────────────┐ + * row0 │ Block_r0_0 │ Block_r0_1 │ (cols 0-31, 32-63) + * ├──────────────┼──────────────┤ + * row1 │ Block_r1_0 │ Block_r1_1 │ + * ├──────────────┼──────────────┤ + * ... │ ... │ ... │ + * ├──────────────┼──────────────┤ + * row63│ Block_r63_0 │ Block_r63_1 │ + * └──────────────┴──────────────┘ + * + * Columnwise blocks (32x1): Each column has 2 blocks + * ┌───┬───┬─────┬───┬───┐ + * │c0 │c1 │ ... │c62│c63│ + * ┌────┼───┼───┼─────┼───┼───┤ + * │Blk0│ │ │ │ │ │ rows 0-31 + * ├────┼───┼───┼─────┼───┼───┤ + * │Blk1│ │ │ │ │ │ rows 32-63 + * └────┴───┴───┴─────┴───┴───┘ + * + * Physical view - Flattened in row-major order: + * Total elements: 64*64 = 4096 + * ┌──────────────────────────────────────────────────────┐ + * │[0...63][64...127][128...191]...[4032...4095] │ + * └──────────────────────────────────────────────────────┘ + * ^ ^ + * start_offset=60 start_offset + input.length=130 + * + * Row-wise amax output (one value per 1x32 block): + * ┌────────┬────────┐ + * │ amax │ amax │ row0 (block0 and block1 partially covered) + * ├────────┼────────┤ + * │ 0 │ 0 │ row1 (not covered) + * ├────────┼────────┤ + * │ ... │ ... │ + * ├────────┼────────┤ + * │ 0 │ 0 │ row63 (not covered) + * └────────┴────────┘ + * + * Column-wise amax output (one value per 32x1 block): + * ┌────────┬────────┬────────┬────────┬────────┬────────┬────────┐ + * │ amax │ amax │ amax │ amax │ amax │ amax │ amax │ ... row 0-31 + * ├────────┼────────┼────────┼────────┼────────┼────────┼────────┤ + * │ amax=0 │ amax=0 │ amax=0 │ amax=0 │ amax=0 │ amax=0 │ amax=0 │ ... row 32-62 + * └────────┴────────┴────────┴────────┴────────┴────────┴────────┘ + * col0 col1 col2 col3 col4 col5 col6 + * + * For each 1x32 or 32x1 block, if it overlaps with [start_offset, start_offset+input.length), + * compute amax; otherwise set to 0. + * + * \param[in] input Input tensor (continuous segment of flattened original tensor). + * \param[in,out] amax_rowwise Output tensor for row-wise maximum absolute values. + * \param[in,out] amax_colwise Output tensor for column-wise maximum absolute values. + * \param[in] rows Number of rows in the logical tensor. + * \param[in] cols Number of columns in the logical tensor. + * \param[in] start_offset Starting offset in the flattened tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor amax_rowwise, + NVTETensor amax_colwise, int rows, int cols, + size_t start_offset, cudaStream_t stream); + +/*! \brief Perform partial MXFP8 casting. + * + * This function casts the input tensor to MXFP8 format, producing both row-wise and + * column-wise scaled outputs. input contains a continuous segment from the flattened + * original tensor. + * + * \param[in] input Input (continuous segment of flattened original tensor). + * \param[out] output_rowwise Output tensor with row-wise scaling (MXFP8 format). + * \param[out] output_colwise Output tensor with column-wise scaling (MXFP8 format). + * \param[in] scale_inv_rowwise Inverse scaling factors for row-wise scaling. + * \param[in] scale_inv_colwise Inverse scaling factors for column-wise scaling. + * \param[in] rows Number of rows in the logical tensor. + * \param[in] cols Number of columns in the logical tensor. + * \param[in] start_offset Starting offset in the flattened tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_mxfp8_scaling_partial_cast(const NVTETensor input, NVTETensor output_rowwise, + NVTETensor output_colwise, const NVTETensor scale_inv_rowwise, + const NVTETensor scale_inv_colwise, int rows, int cols, + size_t start_offset, cudaStream_t stream); + +/*! \brief Compute per-tensor scaling factor for NVFP4 format. + * + * This function computes the scaling factor (alpha) for NVFP4 quantization based + * on the input tensors A and B, with options for using row-wise amax values. + * + * \param[in] inpA Input tensor A. + * \param[in] use_rowwise_amax_A Whether to use row-wise amax for tensor A. + * \param[in] inpB Input tensor B. + * \param[in] use_rowwise_amax_B Whether to use row-wise amax for tensor B. + * \param[in] alpha_in Input scaling factor. + * \param[out] alpha_out Output scaling factor. + * \param[in] stream CUDA stream used for the operation. + */ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_rowwise_amax_A, const NVTETensor inpB, const bool use_rowwise_amax_B, float alpha_in, NVTETensor alpha_out, cudaStream_t stream); diff --git a/transformer_engine/common/multi_tensor/compute_scale.cu b/transformer_engine/common/multi_tensor/compute_scale.cu index dc4eb87145..0ac9ab7371 100644 --- a/transformer_engine/common/multi_tensor/compute_scale.cu +++ b/transformer_engine/common/multi_tensor/compute_scale.cu @@ -14,6 +14,7 @@ #include #include "../recipe/recipe_common.cuh" +#include "../util/ptx.cuh" #include "../utils.cuh" #include "multi_tensor_apply.cuh" @@ -55,6 +56,28 @@ struct ComputeScaleAndScaleInvFunctor { } }; +struct ComputeScaleInvE8M0Functor { + __device__ __forceinline__ void operator()(int chunk_size, volatile int *unused, + TensorListMetadata<2> &tl) { + int tensor_loc = tl.block_to_tensor[blockIdx.x]; + int chunk_idx = tl.block_to_chunk[blockIdx.x]; + int n = tl.sizes[tensor_loc]; + + bf16 *amax = reinterpret_cast(tl.addresses[0][tensor_loc]); + amax += chunk_idx * chunk_size; + + e8m0_t *scale_inv = reinterpret_cast(tl.addresses[1][tensor_loc]); + scale_inv += chunk_idx * chunk_size; + + n -= chunk_idx * chunk_size; + + for (int i_start = threadIdx.x; i_start < n && i_start < chunk_size; i_start += blockDim.x) { + scale_inv[i_start] = ptx::float_to_e8m0(static_cast(amax[i_start]) * + Quantized_Limits::max_norm_rcp); + } + } +}; + void multi_tensor_compute_scale_and_scale_inv_cuda(int chunk_size, Tensor noop_flag, std::vector> tensor_lists, float max_fp8, bool force_pow_2_scales, @@ -65,6 +88,19 @@ void multi_tensor_compute_scale_and_scale_inv_cuda(int chunk_size, Tensor noop_f NVTE_CHECK_CUDA(cudaGetLastError()); } +void multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, + std::vector> tensor_lists, + cudaStream_t stream) { + NVTE_CHECK(tensor_lists[0][0]->data.dtype == DType::kBFloat16, "amax should be bf16"); + auto scale_inv_dtype = tensor_lists[1][0]->data.dtype; + NVTE_CHECK(scale_inv_dtype == DType::kByte || scale_inv_dtype == DType::kFloat8E8M0, + "scale_inv should be e8m0/uint8"); + Tensor dummy; + multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, dummy, tensor_lists, ComputeScaleInvE8M0Functor(), + stream); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + } // namespace multi_tensor_compute_scale } // namespace transformer_engine @@ -82,3 +118,15 @@ void nvte_multi_tensor_compute_scale_and_scale_inv_cuda(int chunk_size, NVTETens convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list), max_fp8, force_pow_2_scales, epsilon, stream); } + +void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor **tensor_lists, + const size_t num_tensor_lists, + const size_t num_tensors_per_list, + cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_compute_scale_inv_e8m0_cuda); + using namespace transformer_engine; + + multi_tensor_compute_scale::multi_tensor_compute_scale_inv_e8m0_cuda( + chunk_size, convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list), + stream); +} diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu new file mode 100644 index 0000000000..8a7ecc6b01 --- /dev/null +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -0,0 +1,253 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../common.h" +#include "../util/ptx.cuh" +#include "../utils.cuh" + +namespace transformer_engine { +namespace mxfp8_scaling_recipe { + +constexpr int rowwise_row_padding = 128; // Row padding of rowwise_scale and rowwise_amax +constexpr int rowwise_col_padding = 4; // Column padding of rowwise_scale and rowwise_amax +constexpr int colwise_row_padding = 4; // Row padding of colwise_scale and colwise_amax +constexpr int colwise_col_padding = 128; // Column padding of colwise_scale and colwise_amax + +constexpr int kRowsPerTile = 32; // Rows each block processes +constexpr int kColsPerTile = 128; // Columns each block processes + +constexpr int kThreadsPerBlock = 128; + +template +__global__ void __launch_bounds__(kThreadsPerBlock) + mxfp8_scaling_compute_partial_amax_kernel(const IType *input, IType *amax_rowwise, + IType *amax_colwise, int amax_rowwise_stride, + int amax_colwise_stride, int rows, int cols, + size_t start_offset, size_t len) { + __shared__ float smem_amax_rowwise[kRowsPerTile][kColsPerTile / 32]; + + size_t end_offset = start_offset + len; + const IType *input_minus_offset = input - start_offset; + int warp_idx = threadIdx.x / 32; + int lane_idx = threadIdx.x % 32; + int c = blockIdx.x * kColsPerTile + threadIdx.x; + int r = blockIdx.y * kRowsPerTile; + + float col_amax = 0.0f; +#pragma unroll + for (int i = 0; i < kRowsPerTile; i++) { + size_t idx = r * cols + c; + float row_amax = 0.0f; + + if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { + float abs_input = fabs(static_cast(input_minus_offset[idx])); + row_amax = fmaxf(row_amax, abs_input); + col_amax = fmaxf(col_amax, abs_input); + } + +#pragma unroll + for (int delta = 16; delta > 0; delta /= 2) { + float other_row_amax = __shfl_down_sync(0xFFFFFFFF, row_amax, delta); + row_amax = fmaxf(row_amax, other_row_amax); + } + + if (lane_idx == 0) { + smem_amax_rowwise[i][warp_idx] = row_amax; + } + + r++; + } + + amax_colwise[blockIdx.y * amax_colwise_stride + c] = static_cast(col_amax); + + __syncthreads(); + + int r_ = threadIdx.x / (kColsPerTile / 32); // rows in shared memory + int c_ = threadIdx.x % (kColsPerTile / 32); // cols in shared memory + r = blockIdx.y * kRowsPerTile + r_; + c = blockIdx.x * kColsPerTile / 32 + c_; + amax_rowwise[r * amax_rowwise_stride + c] = static_cast(smem_amax_rowwise[r_][c_]); +} + +template +__global__ void __launch_bounds__(kThreadsPerBlock) + mxfp8_scaling_partial_cast_kernel(const IType *input, OType *output_rowwise, + OType *output_colwise, const e8m0_t *scale_inv_rowwise, + const e8m0_t *scale_inv_colwise, int scale_inv_rowwise_stride, + int scale_inv_colwise_stride, int rows, int cols, + size_t start_offset, size_t len) { + __shared__ float smem_scales_rowwise[kRowsPerTile][kColsPerTile / 32]; + __shared__ float smem_scales_colwise[kColsPerTile]; + + // Load scales_rowwise + { + int r_ = threadIdx.x / (kColsPerTile / 32); // rows in shared memory + int c_ = threadIdx.x % (kColsPerTile / 32); // cols in shared memory + int r = blockIdx.y * kRowsPerTile + r_; + int c = blockIdx.x * kColsPerTile / 32 + c_; + size_t idx = r * scale_inv_rowwise_stride + c; + smem_scales_rowwise[r_][c_] = ptx::exp2f_rcp(scale_inv_rowwise[idx]); + } + + // Load scales_colwise + { + int c_ = threadIdx.x; + int r = blockIdx.y * kRowsPerTile / 32; + int c = blockIdx.x * kColsPerTile + c_; + size_t idx = r * scale_inv_colwise_stride + c; + smem_scales_colwise[c_] = ptx::exp2f_rcp(scale_inv_colwise[idx]); + } + + __syncthreads(); + + size_t end_offset = start_offset + len; + const IType *input_minus_offset = input - start_offset; + OType *output_rowwise_minus_offset = output_rowwise - start_offset; + OType *output_colwise_minus_offset = output_colwise - start_offset; + int warp_idx = threadIdx.x / 32; + int lane_idx = threadIdx.x % 32; + int c = blockIdx.x * kColsPerTile + threadIdx.x; + int r = blockIdx.y * kRowsPerTile; + +#pragma unroll + for (int i = 0; i < kRowsPerTile; i++) { + size_t idx = r * cols + c; + + if (r < rows && c < cols && idx >= start_offset && idx < end_offset) { + float inp = static_cast(input_minus_offset[idx]); + OType out_rowwise = static_cast(inp * smem_scales_rowwise[i][warp_idx]); + OType out_colwise = static_cast(inp * smem_scales_colwise[threadIdx.x]); + output_rowwise_minus_offset[idx] = out_rowwise; + output_colwise_minus_offset[idx] = out_colwise; + } + + r++; + } +} + +void mxfp8_scaling_compute_partial_amax(const Tensor input, Tensor amax_rowwise, + Tensor amax_colwise, int rows, int cols, + size_t start_offset, cudaStream_t stream) { + NVTE_CHECK(rows % 32 == 0, "rows must be divisible by 32"); + NVTE_CHECK(cols % 32 == 0, "cols must be divisible by 32"); + + NVTE_CHECK(input.data.shape.size() == 1, "input must be a 1D tensor"); + NVTE_CHECK(start_offset + input.data.shape[0] <= static_cast(rows) * cols, + "Invalid start_offset"); + + NVTE_CHECK(amax_rowwise.data.shape.size() == 2, "amax_rowwise must be a 2D tensor"); + NVTE_CHECK(amax_rowwise.data.shape[0] % rowwise_row_padding == 0, + "Wrong padding of amax_rowwise's rows"); + NVTE_CHECK(amax_rowwise.data.shape[0] >= rows, "Invalid rows"); + NVTE_CHECK(amax_rowwise.data.shape[1] % rowwise_col_padding == 0, + "Wrong padding of amax_rowwise's cols"); + NVTE_CHECK(amax_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); + NVTE_CHECK(amax_rowwise.dtype() == input.dtype(), "Wrong dtype of amax_rowwise"); + + NVTE_CHECK(amax_colwise.data.shape.size() == 2, "amax_colwise must be a 2D tensor"); + NVTE_CHECK(amax_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of amax_colwise's rows"); + NVTE_CHECK(amax_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(amax_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of amax_colwise's cols"); + NVTE_CHECK(amax_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(amax_colwise.dtype() == input.dtype(), "Wrong dtype of amax_colwise"); + + int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; + int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; + dim3 grid(blocks_x, blocks_y); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + mxfp8_scaling_compute_partial_amax_kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(amax_rowwise.data.dptr), + reinterpret_cast(amax_colwise.data.dptr), amax_rowwise.data.shape[1], + amax_colwise.data.shape[1], rows, cols, start_offset, input.data.shape[0]);) +} + +void mxfp8_scaling_partial_cast(const Tensor input, Tensor output_rowwise, Tensor output_colwise, + const Tensor scale_inv_rowwise, const Tensor scale_inv_colwise, + int rows, int cols, size_t start_offset, cudaStream_t stream) { + NVTE_CHECK(rows % 32 == 0, "rows must be divisible by 32"); + NVTE_CHECK(cols % 32 == 0, "cols must be divisible by 32"); + + NVTE_CHECK(input.data.shape.size() == 1, "input must be a 1D tensor"); + NVTE_CHECK(start_offset + input.data.shape[0] <= static_cast(rows) * cols, + "Invalid start_offset"); + + NVTE_CHECK(output_rowwise.data.shape.size() == 1, "output_rowwise must be a 1D tensor"); + NVTE_CHECK(output_colwise.data.shape.size() == 1, "output_colwise must be a 1D tensor"); + NVTE_CHECK(output_rowwise.data.shape[0] == input.data.shape[0], + "Size of input and output_rowwise mismatch"); + NVTE_CHECK(output_colwise.data.shape[0] == input.data.shape[0], + "Size of input and output_colwise mismatch"); + + NVTE_CHECK(output_rowwise.dtype() == DType::kFloat8E4M3 || output_rowwise.dtype() == DType::kByte, + "output_rowwise should be e4m3 or uint8"); + NVTE_CHECK(output_colwise.dtype() == DType::kFloat8E4M3 || output_colwise.dtype() == DType::kByte, + "output_colwise should be e4m3 or uint8"); + + NVTE_CHECK(scale_inv_rowwise.data.shape.size() == 2, "scale_inv_rowwise must be a 2D tensor"); + NVTE_CHECK(scale_inv_rowwise.data.shape[0] % rowwise_row_padding == 0, + "Wrong padding of scale_inv_rowwise's rows"); + NVTE_CHECK(scale_inv_rowwise.data.shape[0] >= rows, "Invalid rows"); + NVTE_CHECK(scale_inv_rowwise.data.shape[1] % rowwise_col_padding == 0, + "Wrong padding of scale_inv_rowwise's cols"); + NVTE_CHECK(scale_inv_rowwise.data.shape[1] >= cols / 32, "Invalid cols"); + NVTE_CHECK(scale_inv_rowwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_rowwise"); + + NVTE_CHECK(scale_inv_colwise.data.shape.size() == 2, "scale_inv_colwise must be a 2D tensor"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] % colwise_row_padding == 0, + "Wrong padding of scale_inv_colwise's rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[0] >= rows / 32, "Invalid rows"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] % colwise_col_padding == 0, + "Wrong padding of scale_inv_colwise's cols"); + NVTE_CHECK(scale_inv_colwise.data.shape[1] >= cols, "Invalid cols"); + NVTE_CHECK(scale_inv_colwise.dtype() == DType::kByte, "Wrong dtype of scale_inv_colwise"); + + int blocks_x = (cols + kColsPerTile - 1) / kColsPerTile; + int blocks_y = (rows + kRowsPerTile - 1) / kRowsPerTile; + dim3 grid(blocks_x, blocks_y); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input.dtype(), IType, + mxfp8_scaling_partial_cast_kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output_rowwise.data.dptr), + reinterpret_cast(output_colwise.data.dptr), + reinterpret_cast(scale_inv_rowwise.data.dptr), + reinterpret_cast(scale_inv_colwise.data.dptr), + scale_inv_rowwise.data.shape[1], scale_inv_colwise.data.shape[1], rows, cols, + start_offset, input.data.shape[0]);) +} + +} // namespace mxfp8_scaling_recipe +} // namespace transformer_engine + +void nvte_mxfp8_scaling_compute_partial_amax(const NVTETensor input, NVTETensor amax_rowwise, + NVTETensor amax_colwise, int rows, int cols, + size_t start_offset, cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp8_scaling_compute_partial_amax); + using namespace transformer_engine; + mxfp8_scaling_recipe::mxfp8_scaling_compute_partial_amax( + *convertNVTETensorCheck(input), *convertNVTETensorCheck(amax_rowwise), + *convertNVTETensorCheck(amax_colwise), rows, cols, start_offset, stream); +} + +void nvte_mxfp8_scaling_partial_cast(const NVTETensor input, NVTETensor output_rowwise, + NVTETensor output_colwise, const NVTETensor scale_inv_rowwise, + const NVTETensor scale_inv_colwise, int rows, int cols, + size_t start_offset, cudaStream_t stream) { + NVTE_API_CALL(nvte_mxfp8_scaling_partial_cast); + using namespace transformer_engine; + mxfp8_scaling_recipe::mxfp8_scaling_partial_cast( + *convertNVTETensorCheck(input), *convertNVTETensorCheck(output_rowwise), + *convertNVTETensorCheck(output_colwise), *convertNVTETensorCheck(scale_inv_rowwise), + *convertNVTETensorCheck(scale_inv_colwise), rows, cols, start_offset, stream); +} diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 44c49b20bc..80479dccf4 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -335,6 +335,15 @@ void fp8_block_scaling_partial_cast(const at::Tensor &inp, at::Tensor out, const size_t h, size_t w, size_t start_offset, size_t block_len, const DType out_dtype); +void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax_rowwise, + at::Tensor amax_colwise, int rows, int cols, + size_t start_offset); + +void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwise, + at::Tensor output_colwise, const at::Tensor &scale_inv_rowwise, + const at::Tensor &scale_inv_colwise, int rows, int cols, + size_t start_offset); + /*************************************************************************************************** * Rotary positional embedding **************************************************************************************************/ @@ -451,6 +460,9 @@ void multi_tensor_compute_scale_and_scale_inv_cuda( int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, float max_fp8, bool force_pow_2_scales, float epsilon); +void multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, const py::object &dummy, + std::vector> tensor_lists); + /*************************************************************************************************** * padding **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/fp8_block_scaling_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp similarity index 53% rename from transformer_engine/pytorch/csrc/extensions/fp8_block_scaling_partial_cast.cpp rename to transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp index bea6f8c907..3be2ca9396 100644 --- a/transformer_engine/pytorch/csrc/extensions/fp8_block_scaling_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp @@ -48,4 +48,42 @@ void fp8_block_scaling_partial_cast(const at::Tensor &inp, at::Tensor out, const start_offset, block_len, static_cast(out_dtype), at::cuda::getCurrentCUDAStream()); } +void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax_rowwise, + at::Tensor amax_colwise, int rows, int cols, + size_t start_offset) { + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(amax_rowwise.is_contiguous(), "amax_rowwise must be contiguous"); + TORCH_CHECK(amax_colwise.is_contiguous(), "amax_colwise must be contiguous"); + + const TensorWrapper input_cu = makeTransformerEngineTensor(input); + TensorWrapper amax_rowwise_cu = makeTransformerEngineTensor(amax_rowwise); + TensorWrapper amax_colwise_cu = makeTransformerEngineTensor(amax_colwise); + + nvte_mxfp8_scaling_compute_partial_amax(input_cu.data(), amax_rowwise_cu.data(), + amax_colwise_cu.data(), rows, cols, start_offset, + at::cuda::getCurrentCUDAStream()); +} + +void mxfp8_scaling_partial_cast(const at::Tensor &input, at::Tensor output_rowwise, + at::Tensor output_colwise, const at::Tensor &scale_inv_rowwise, + const at::Tensor &scale_inv_colwise, int rows, int cols, + size_t start_offset) { + TORCH_CHECK(input.is_contiguous(), "input must be contiguous"); + TORCH_CHECK(output_rowwise.is_contiguous(), "output_rowwise must be contiguous"); + TORCH_CHECK(output_colwise.is_contiguous(), "output_colwise must be contiguous"); + TORCH_CHECK(scale_inv_rowwise.is_contiguous(), "scale_inv_rowwise must be contiguous"); + TORCH_CHECK(scale_inv_colwise.is_contiguous(), "scale_inv_colwise must be contiguous"); + + const TensorWrapper input_cu = makeTransformerEngineTensor(input); + TensorWrapper output_rowwise_cu = makeTransformerEngineTensor(output_rowwise); + TensorWrapper output_colwise_cu = makeTransformerEngineTensor(output_colwise); + const TensorWrapper scale_inv_rowwise_cu = makeTransformerEngineTensor(scale_inv_rowwise); + const TensorWrapper scale_inv_colwise_cu = makeTransformerEngineTensor(scale_inv_colwise); + + nvte_mxfp8_scaling_partial_cast(input_cu.data(), output_rowwise_cu.data(), + output_colwise_cu.data(), scale_inv_rowwise_cu.data(), + scale_inv_colwise_cu.data(), rows, cols, start_offset, + at::cuda::getCurrentCUDAStream()); +} + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp index 8a1a34698b..e60b001f6f 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp @@ -20,4 +20,14 @@ void multi_tensor_compute_scale_and_scale_inv_cuda( force_pow_2_scales, epsilon, at::cuda::getCurrentCUDAStream()); } +void multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, const py::object &dummy, + std::vector> tensor_lists) { + NVTE_CHECK(dummy.is_none(), "No-op flag is not supported."); + auto [_, __, tensor_lists_ptr, num_lists, num_tensors] = + makeTransformerEngineTensorList(tensor_lists); + + nvte_multi_tensor_compute_scale_inv_e8m0_cuda(chunk_size, tensor_lists_ptr.data(), num_lists, + num_tensors, at::cuda::getCurrentCUDAStream()); +} + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 3b81393dbd..d0f450bc71 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -276,6 +276,16 @@ 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()); + m.def("mxfp8_scaling_compute_partial_amax", + &transformer_engine::pytorch::mxfp8_scaling_compute_partial_amax, + "Compute partial amax from master weights for fp8 mxfp8 scaling", py::arg("input"), + py::arg("amax_rowwise"), py::arg("amax_colwise"), py::arg("rows"), py::arg("cols"), + py::arg("start_offset"), py::call_guard()); + m.def("mxfp8_scaling_partial_cast", &transformer_engine::pytorch::mxfp8_scaling_partial_cast, + "Partial cast from master weights for fp8 mxfp8 scaling", py::arg("input"), + py::arg("output_rowwise"), py::arg("output_colwise"), py::arg("scale_inv_rowwise"), + py::arg("scale_inv_colwise"), py::arg("rows"), py::arg("cols"), py::arg("start_offset"), + py::call_guard()); m.def("fused_multi_row_padding", &transformer_engine::pytorch::fused_multi_row_padding, "Fused Multi-tensor padding", py::call_guard()); m.def("fused_multi_row_unpadding", &transformer_engine::pytorch::fused_multi_row_unpadding, @@ -427,6 +437,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("multi_tensor_compute_scale_and_scale_inv", &transformer_engine::pytorch::multi_tensor_compute_scale_and_scale_inv_cuda, "Fused compute scale and scale_inv from amax", py::call_guard()); + m.def("multi_tensor_compute_scale_inv_e8m0", + &transformer_engine::pytorch::multi_tensor_compute_scale_inv_e8m0_cuda, + "Fused compute E8M0 scale_inv from amax", py::call_guard()); // Comm+GEMM Overlap m.def("bulk_overlap_ag_with_external_gemm", diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 9773e17e64..94f761f2b0 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -8,7 +8,11 @@ import torch import transformer_engine_torch as tex -from transformer_engine_torch import multi_tensor_scale, multi_tensor_compute_scale_and_scale_inv +from transformer_engine_torch import ( + multi_tensor_scale, + multi_tensor_compute_scale_and_scale_inv, + multi_tensor_compute_scale_inv_e8m0, +) from ..quantized_tensor import QuantizedTensor, Quantizer, QuantizedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer @@ -85,6 +89,7 @@ def cast_master_weights_to_fp8( delayed_scaling_params = [] current_scaling_params = [] blockwise_scaling_params = [] + mxfp8_scaling_params = [] if fsdp_shard_model_weights is None: use_fsdp_shard_model_weights = False @@ -131,8 +136,8 @@ def cast_master_weights_to_fp8( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) elif isinstance(quantizer, MXFP8Quantizer): - raise NotImplementedError( - "cast_master_weights_to_fp8 for MXFP8BlockScaling is not supported yet" + mxfp8_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) else: raise ValueError( @@ -146,6 +151,8 @@ def cast_master_weights_to_fp8( _cast_master_weights_to_fp8_current_scaling(current_scaling_params, *extra_args) if len(blockwise_scaling_params) > 0: _cast_master_weights_to_fp8_blockwise_scaling(blockwise_scaling_params, *extra_args) + if len(mxfp8_scaling_params) > 0: + _cast_master_weights_to_fp8_mxfp8_scaling(mxfp8_scaling_params, *extra_args) def _cast_master_weights_to_fp8_delayed_scaling( @@ -467,6 +474,131 @@ def _cast_master_weights_to_fp8_blockwise_scaling( ) +def _cast_master_weights_to_fp8_mxfp8_scaling( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): # pylint: disable=unused-argument + r"""Helper function to cast master weights to FP8 primary weights for mxfp8 scaling. + + Parameters + ---------- + params : List of tuple, each tuple contains a model weight, a master weight, and an offset + indicating the starting index of the master weight in the model weight. + group : The distributed group to do amax reduction. Typically it's the data parallel + group. + use_fsdp_shard_model_weights : bool, if True, it means that the model weights are sharded. + """ + + # Parameter attributes + device = params[0][0].device + for _, master_weight, _, _ in params: + if master_weight is not None: + master_weight_dtype = master_weight.dtype + break + + # Get the total number of amax elements in all the model weights. + cu_rowwise_amax_sizes = [0] + cu_colwise_amax_sizes = [0] + for model_weight, _, _, _ in params: + rowwise_shape = model_weight._rowwise_scale_inv.shape + assert len(rowwise_shape) == 2 + colwise_shape = model_weight._columnwise_scale_inv.shape + assert len(colwise_shape) == 2 + cu_rowwise_amax_sizes.append( + cu_rowwise_amax_sizes[-1] + rowwise_shape[0] * rowwise_shape[1] + ) + cu_colwise_amax_sizes.append( + cu_colwise_amax_sizes[-1] + colwise_shape[0] * colwise_shape[1] + ) + + # Create a contiguous buffer to store amaxes temporarily, so we can perform all all-reduce + # NCCL kernels at once. + packed_amaxes = torch.zeros( + cu_rowwise_amax_sizes[-1] + cu_colwise_amax_sizes[-1], + dtype=master_weight_dtype, + device=device, + ) + + # --------------------------------------------------------------------------------------------- + # Step 1: Iterate through all the none empty master weights and compute amax of them. Store the + # amaxes in a contiguous buffer. If a block of a master weight is empty, the + # corresponding amax will be set to 0. + # --------------------------------------------------------------------------------------------- + amaxes_rowwise, scale_invs_rowwise = [], [] + amaxes_colwise, scale_invs_colwise = [], [] + for i, (model_weight, master_weight, start_offset, _) in enumerate(params): + rowwise_shape = model_weight._rowwise_scale_inv.shape + colwise_shape = model_weight._columnwise_scale_inv.shape + rowwise_start = cu_rowwise_amax_sizes[i] + rowwise_end = cu_rowwise_amax_sizes[i + 1] + colwise_start = cu_rowwise_amax_sizes[-1] + cu_colwise_amax_sizes[i] + colwise_end = cu_rowwise_amax_sizes[-1] + cu_colwise_amax_sizes[i + 1] + amax_rowwise = packed_amaxes[rowwise_start:rowwise_end].reshape(rowwise_shape) + amax_colwise = packed_amaxes[colwise_start:colwise_end].reshape(colwise_shape) + amaxes_rowwise.append(amax_rowwise) + amaxes_colwise.append(amax_colwise) + scale_invs_rowwise.append(model_weight._rowwise_scale_inv) + scale_invs_colwise.append(model_weight._columnwise_scale_inv) + + # Compute amax of the master weight and store it in packed_amaxes. + if master_weight is not None: + assert len(model_weight.shape) == 2 + h, w = model_weight.shape + tex.mxfp8_scaling_compute_partial_amax( + master_weight, amax_rowwise, amax_colwise, h, w, start_offset + ) + + # --------------------------------------------------------------------------------------------- + # Step 2: Perform all-reduce on packed_amaxes to get the global amax. + # --------------------------------------------------------------------------------------------- + torch.distributed.all_reduce(packed_amaxes, op=torch.distributed.ReduceOp.MAX, group=group) + + # --------------------------------------------------------------------------------------------- + # Step 3: Update scales and scale_invs. + # --------------------------------------------------------------------------------------------- + multi_tensor_applier( + multi_tensor_compute_scale_inv_e8m0, + None, # dummy_overflow_buf + [ + amaxes_rowwise + amaxes_colwise, + scale_invs_rowwise + scale_invs_colwise, + ], + ) + + # --------------------------------------------------------------------------------------------- + # Step 4: Cast master weights to FP8. + # --------------------------------------------------------------------------------------------- + for ( + (model_weight, master_weight, start_offset, model_weight_fragment), + scale_inv_rowwise, + scale_inv_colwise, + ) in zip(params, scale_invs_rowwise, scale_invs_colwise): + # If master weight is None, it means that the master weight of the current model weight + # is in other DP ranks. + if master_weight is None: + continue + + # Cast master weight to FP8 + end_offset = start_offset + master_weight.numel() + if use_fsdp_shard_model_weights: + rowwise_fragment = model_weight_fragment[0] + colwise_fragment = model_weight_fragment[1] + else: + rowwise_fragment = model_weight._rowwise_data.reshape(-1)[start_offset:end_offset] + colwise_fragment = model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] + assert len(model_weight.shape) == 2 + h, w = model_weight.shape + tex.mxfp8_scaling_partial_cast( + master_weight, + rowwise_fragment, + colwise_fragment, + scale_inv_rowwise, + scale_inv_colwise, + h, + w, + start_offset, + ) + + def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Tensor]]): """ Post-processing after all-gather for weights in distributed optimizer. @@ -485,6 +617,9 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten elif isinstance(model_weight, Float8BlockwiseQTensor): # Blockwise scaling: create column-wise storage. model_weight._create_columnwise() + elif isinstance(model_weight, MXFP8Tensor): + # MXFP8 scaling: no need to do anything. + pass elif isinstance(model_weight, QuantizedTensor): raise ValueError(f"post_processing for {type(model_weight)} is not supported") From 6182206102a553ca73aaa13b499d0fb3f40fdd73 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 4 Dec 2025 14:34:11 -0800 Subject: [PATCH 109/521] [Core] Fix inconsistent logic in C++ tensor class (#2330) * Initialize empty tensors with shape=[0] instead of shape=[]. Signed-off-by: Tim Moon * Fix runtime crash in LayerNorm Still seeing correctness issues. Signed-off-by: Tim Moon * Make sure norm workspace sizes are not zero Signed-off-by: Tim Moon * Remove assumption in swizzle kernel that data is available. Signed-off-by: Tim Moon * Remove assumption in multi-swizzle kernel that data is available. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove unnecessary explicit call to default constructor Signed-off-by: Tim Moon * Avoid accessing tensor data pointer if tensor has no entries Signed-off-by: Tim Moon * [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: 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 * Update transformer_engine/common/swizzle/swizzle.cu 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 * Review suggestions from @ptrendx and @greptile-apps Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Prefer using row-wise/col-wise shape based on which has data Signed-off-by: Tim Moon * Fix merge conflict, expand docs, fix inconsistency in dim function Signed-off-by: Tim Moon * Change Tensor::has_data to check whether tensor is initialized, not whether pointer is valid. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Review suggestion from @greptile-apps Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Debug incorrect tensor initialization in tests Signed-off-by: Tim Moon * Clarify comments that has_data does not guarantee safe pointer accesses Signed-off-by: Tim Moon * Debug test failure when computing amaxes Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon 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: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/cpp/test_common.cu | 107 +++++---- transformer_engine/common/common.h | 172 ++++++-------- .../transformer_engine/transformer_engine.h | 23 +- .../common/normalization/common.cpp | 16 +- .../common/normalization/layernorm/ln_api.cpp | 8 +- .../normalization/rmsnorm/rmsnorm_api.cpp | 12 +- transformer_engine/common/swizzle/swizzle.cu | 224 +++++++++++------- .../common/transformer_engine.cpp | 103 ++++---- .../common/transpose/cast_transpose_fusion.cu | 2 - .../common/transpose/transpose_fusion.cu | 2 - .../pytorch/csrc/extensions/cast.cpp | 24 +- .../pytorch/csrc/extensions/gemm.cpp | 8 +- .../pytorch/csrc/extensions/recipe.cpp | 4 +- transformer_engine/pytorch/csrc/util.cpp | 13 +- 14 files changed, 412 insertions(+), 306 deletions(-) diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index cdbfb05b3c..d70eb13536 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -278,52 +278,33 @@ std::pair get_scales(const NVTEShape& shape, Tensor::Tensor(const std::string& name, const NVTEShape &shape, const DType type, const bool rowwise, const bool columnwise, - const NVTEScalingMode &scaling_mode) { - name_ = name; + const NVTEScalingMode &scaling_mode) + : tensor_(scaling_mode), rowwise_{rowwise}, columnwise_{columnwise}, name_{name} { + // Initialize RNG const size_t seed = create_seed_from_tensor_name(name); gen_.seed(seed); - rowwise_ = rowwise; - columnwise_ = columnwise; - size_t total_size = bytes(shape, type); - void *dptr_rowwise = nullptr; - void *dptr_columnwise = nullptr; - cpu_data_rowwise_ = nullptr; - cpu_data_columnwise_ = nullptr; - amax_cpu_data_ = nullptr; - scale_cpu_data_ = nullptr; - rowwise_scale_inv_cpu_data_ = nullptr; - columnwise_scale_inv_cpu_data_ = nullptr; - float *amax = nullptr, *scale = nullptr; - float *rowwise_scale_inv = nullptr, *columnwise_scale_inv = nullptr; + + // Make sure shape is valid if (columnwise) { NVTE_CHECK(shape.ndim >= 2); } - std::vector normalized_shape_v = {product(shape, 0, shape.ndim - 1), - shape.data[shape.ndim - 1]}; - NVTEShape normalized_shape = convertShape(normalized_shape_v); - NVTEShape columnwise_shape = {}; - - std::vector columnwise_shape_vec; - if (scaling_mode == NVTE_DELAYED_TENSOR_SCALING - || scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D) { - // Transpose when tensor scaling - columnwise_shape_vec.emplace_back(shape.data[shape.ndim - 1]); - for (size_t i = 0; i < shape.ndim - 1; ++i) { - columnwise_shape_vec.emplace_back(shape.data[i]); - } - } else { - // Same shape for MX and NVFP4 - for (size_t i = 0; i < shape.ndim; ++i) { - columnwise_shape_vec.emplace_back(shape.data[i]); - } - } - if (columnwise) { - columnwise_shape = nvte_make_shape(columnwise_shape_vec.data(), columnwise_shape_vec.size()); + // Shape after flattening to 2D + NVTEShape flattened_shape; + { + std::vector flattened_shape_vec; + if (shape.ndim > 0) { + flattened_shape_vec.push_back(product(shape, 0, shape.ndim - 1)); + flattened_shape_vec.push_back(shape.data[shape.ndim - 1]); + } else { + flattened_shape_vec.resize(2, 1); + } + flattened_shape = convertShape(flattened_shape_vec); } - tensor_ = TensorWrapper(scaling_mode); - + // 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(*) @@ -339,11 +320,51 @@ Tensor::Tensor(const std::string& name, } } - const DType rowwise_type = (scaling_mode == NVTE_NVFP4_1D_SCALING) ? DType::kFloat4E2M1 : type; - const DType colwise_type = (scaling_mode == NVTE_NVFP4_1D_SCALING) ? DType::kFloat4E2M1 : type; - tensor_.set_rowwise_data(dptr_rowwise, rowwise_type, shape); - tensor_.set_columnwise_data(dptr_columnwise, colwise_type, columnwise_shape); + // Set tensor 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); + } + // Set tensor column-wise data + if (columnwise) { + // Determine shape of column-wise data + std::vector columnwise_shape_vec; + switch (scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: + case NVTE_BLOCK_SCALING_1D: + case NVTE_BLOCK_SCALING_2D: { + // Column-wise data shape is transposed + if (shape.ndim > 0) { + columnwise_shape_vec.emplace_back(shape.data[shape.ndim - 1]); + for (size_t i = 0; i < shape.ndim - 1; ++i) { + columnwise_shape_vec.emplace_back(shape.data[i]); + } + } + break; + } + case NVTE_MXFP8_1D_SCALING: + case NVTE_NVFP4_1D_SCALING: { + // Column-wise data matches shape + for (size_t i = 0; i < shape.ndim; ++i) { + columnwise_shape_vec.emplace_back(shape.data[i]); + } + break; + } + default: + NVTE_ERROR("Unrecognized scaling mode (", (size_t)scaling_mode, ")."); + } + 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); + } + + // Configure scales, amaxes, and other tensor buffers + float *amax = nullptr, *scale = nullptr; + float *rowwise_scale_inv = nullptr, *columnwise_scale_inv = nullptr; if (isFp8Type(type) || isFp4Type(type)) { if (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { cudaMalloc((void**)&amax, sizeof(float)); // NOLINT(*) @@ -375,7 +396,7 @@ Tensor::Tensor(const std::string& name, scale_cpu_data_ = std::make_shared(0); tensor_.set_scale(scale, DType::kFloat32, std::vector{1}); } - auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(normalized_shape, tensor_.scaling_mode()); + 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; diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 661f8b00e1..38b437b994 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -74,38 +74,49 @@ inline size_t product(const std::vector &shape) { return ret; } +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); + struct SimpleTensor { void *dptr; std::vector shape; DType dtype; - SimpleTensor(void *dptr, const std::vector &shape, DType dtype) - : dptr(dptr), shape(shape), dtype(dtype) {} + SimpleTensor(void *dptr, std::vector shape, DType dtype) + : dptr{dptr}, shape{std::move(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, {}, DType::kFloat32) {} + SimpleTensor() : SimpleTensor(nullptr, std::vector{0}, DType::kFloat32) {} operator NVTEBasicTensor() const { return {dptr, static_cast(dtype), nvte_make_shape(this->shape.data(), this->shape.size())}; } - size_t numel() const { - size_t acc = 1; - for (const auto &dim : shape) { - acc *= dim; - } - return acc; - } - bool has_data() const noexcept { return dptr != nullptr && numel() > 0; } + /*! Number of tensor elements. */ + size_t numel() const { return product(shape); } + + /*! Whether the tensor is initialized. + * + * Tensors with non-trivial shapes are considered initialized. This + * means that there is no guarantee that the data pointer can be + * safely accessed. + */ + bool has_data() const { return !(dptr == nullptr && shape.size() == 1 && shape[0] == 0); } + + /*! Buffer size in bytes. */ + size_t buffer_size_bytes() const { return get_buffer_size_bytes(numel(), dtype); } + /*! Reset to uninitialized tensor. */ void clear() { dptr = nullptr; - shape.resize(0); + shape.resize(1); + shape[0] = 0; dtype = DType::kFloat32; } }; @@ -123,17 +134,9 @@ struct Tensor { NVTEScalingMode scaling_mode; NVTETensor nvte_tensor; - Tensor() - : data(), - columnwise_data(), - amax(nullptr, {1}, DType::kFloat32), - columnwise_amax(nullptr, {1}, DType::kFloat32), - scale(nullptr, {1}, DType::kFloat32), - scale_inv(nullptr, {1}, DType::kFloat32), - columnwise_scale_inv(nullptr, {1}, DType::kFloat32), - scaling_mode(NVTE_DELAYED_TENSOR_SCALING), - nvte_tensor(0) {} + Tensor() : scaling_mode{NVTE_DELAYED_TENSOR_SCALING}, nvte_tensor{0} {} + /*! Reset tensor data. */ void clear() { data.clear(); columnwise_data.clear(); @@ -147,65 +150,62 @@ struct Tensor { explicit operator NVTETensor() const noexcept { return nvte_tensor; } + /*! Number of tensor elements. */ size_t numel() const { - size_t acc = 1; - for (const auto dim : shape()) { - acc *= dim; + if (!has_data() && has_columnwise_data()) { + return product(columnwise_data.shape); } - return acc; + return product(data.shape); } - // TODO(Tim): Change this to use data.has_data() - bool has_data() const noexcept { return data.dptr != nullptr; } + /*! Whether the tensor data buffer is not uninitialized. + * + * Buffers with non-trivial shapes are considered initialized. This + * means that there is no guarantee that the data pointer can be + * safely accessed. + */ + bool has_data() const { return data.has_data(); } - // Check for size (not just pointer) for 0-dim or no token cases. - // TODO(Tim): Change this to use columnwise_data.has_data() - bool has_columnwise_data() const noexcept { - return columnwise_data.dptr != nullptr || columnwise_data.shape.size() != 0; - } + /*! Whether the tensor column-wise data buffer is not uninitialized. + * + * Buffers with non-trivial shapes are considered initialized. This + * means that there is no guarantee that the data pointer can be + * safely accessed. + */ + bool has_columnwise_data() const { return columnwise_data.has_data(); } + /*! Datatype of tensor elements. */ DType dtype() const { - if (has_data()) return data.dtype; - if (has_columnwise_data()) return columnwise_data.dtype; - // Fallback, used e.g. in workspace + if (!has_data() && has_columnwise_data()) { + return columnwise_data.dtype; + } return data.dtype; } + /*! Number of tensor dimensions. */ size_t dim() const { if (!has_data() && has_columnwise_data()) { return columnwise_data.shape.size(); - } else { - return data.shape.size(); } + return data.shape.size(); } + /*! Tensor dimensions. + * + * This is the logical tensor shape. The underlying data may have a + * different shape, e.g. the column-wise data for some tensor + * formats are transposed. + */ std::vector shape() const { - /* Note: We sometimes experience spurious compiler errors - * (-Wstringop-overflow) from this function. It appears that GCC - * has some bugs with std::vector (see - * https://gcc.gnu.org/bugzilla/show_bug.cgi?id=109569). - */ + // Each tensor format interprets its data differently switch (scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: + case NVTE_BLOCK_SCALING_1D: + case NVTE_BLOCK_SCALING_2D: case NVTE_NVFP4_1D_SCALING: { - // Choose data buffer based on whether it is initialized - // Note: Uninitialized buffers currently have shape=[]. - // However, this is logically incorrect. 0-D tensors have 1 - // entry, and uninitialized tensors should have shape=[0]. - bool use_columnwise_shape = false; - if (data.dptr != nullptr) { - use_columnwise_shape = false; - } else if (columnwise_data.dptr != nullptr) { - use_columnwise_shape = true; - } else if (data.shape.size() != 0) { - use_columnwise_shape = false; - } else if (columnwise_data.shape.size() != 0) { - use_columnwise_shape = true; - } - - // Infer shape based on data - if (use_columnwise_shape) { - // Column-wise data is transposed + // 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; if (!columnwise_data.shape.empty()) { ret.reserve(columnwise_data.shape.size()); @@ -218,38 +218,16 @@ struct Tensor { } return data.shape; } - case NVTE_MXFP8_1D_SCALING: + case NVTE_MXFP8_1D_SCALING: { + // Row-wise and column-wise data shapes both match tensor + // logical shape if (!has_data() && has_columnwise_data()) { return columnwise_data.shape; - } else { - return data.shape; - } - break; - case NVTE_BLOCK_SCALING_1D: - case NVTE_BLOCK_SCALING_2D: { - if (!has_data() && has_columnwise_data()) { - std::vector shape; - size_t ndim = columnwise_data.shape.size(); - shape.reserve(ndim); - for (size_t i = 0; i + 1 < ndim; ++i) { - shape.push_back(columnwise_data.shape[i + 1]); - } - if (ndim > 0) { - shape.push_back(columnwise_data.shape[0]); - } - return shape; - } else { - // NOTE: We may have removed the data pointer from - // data by setting usage. In that case, we return - // the non-null shape. It is our best guess at the most - // recent shape. - return data.shape; } - break; + return data.shape; } default: NVTE_ERROR("Cannot parse tensor shape with scaling mode \"", to_string(scaling_mode), "\""); - return {}; } } @@ -347,10 +325,10 @@ struct GroupedTensor { columnwise_amax(), scale(), num_tensors(num_tensors), - first_dims(nullptr, {}, DType::kInt64), - last_dims(nullptr, {}, DType::kInt64), - tensor_offsets(nullptr, {}, DType::kInt64), - logical_shape(nvte_make_shape(nullptr, 0)), + first_dims(nullptr, std::vector{0}, DType::kInt64), + last_dims(nullptr, std::vector{0}, DType::kInt64), + tensor_offsets(nullptr, std::vector{0}, DType::kInt64), + logical_shape(nvte_make_shape(nullptr, 1)), scaling_mode(scaling_mode), nvte_tensor(0) {} @@ -383,9 +361,9 @@ struct GroupedTensor { } DType dtype() const { - if (has_data()) return data.dtype; - if (has_columnwise_data()) return columnwise_data.dtype; - // Fallback, used e.g. in workspace or when allow_empty=true + if (!has_data() && has_columnwise_data()) { + return columnwise_data.dtype; + } return data.dtype; } @@ -400,7 +378,7 @@ struct GroupedTensor { first_dims.clear(); last_dims.clear(); tensor_offsets.clear(); - logical_shape = nvte_make_shape(nullptr, 0); + logical_shape = nvte_make_shape(nullptr, 1); num_tensors = 0; scaling_mode = NVTE_DELAYED_TENSOR_SCALING; nvte_tensor = 0; @@ -869,10 +847,6 @@ inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) { size_t typeToSize(const DType type); size_t typeToNumBits(const DType type); -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); - void CheckNoopTensor(const Tensor &t, const std::string &name); void CheckInputTensor(const Tensor &t, const std::string &name); void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty = false); diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 76cc636a35..b2e04ba69f 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -142,7 +142,8 @@ void *nvte_tensor_columnwise_data(const NVTETensor tensor); /*! \brief Construct a shape from an array of dimension sizes. * - * \param[data] Pointer to start of shape array. + * \param[data] Pointer to start of shape array. If NULL, the shape + * will be filled with zeros. * \param[data] Number of dimensions (must be <= 14) * * \return A shape. The shape will own its own copy of the data. @@ -575,15 +576,22 @@ class TensorWrapper { */ TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr, - const NVTEShape scale_inv_shape = defaultShape, + NVTEShape scale_inv_shape = defaultShape, const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING) { tensor_ = nvte_create_tensor(scaling_mode); NVTEBasicTensor data = {dptr, static_cast(dtype), shape}; nvte_set_tensor_param(&tensor_, kNVTERowwiseData, &data); - NVTEBasicTensor amax = {amax_dptr, kNVTEFloat32, defaultShape}; + NVTEBasicTensor amax = {amax_dptr, kNVTEFloat32, + amax_dptr != nullptr ? defaultShape : emptyShape}; nvte_set_tensor_param(&tensor_, kNVTEAmax, &amax); - NVTEBasicTensor scale = {scale_dptr, kNVTEFloat32, defaultShape}; + NVTEBasicTensor scale = {scale_dptr, kNVTEFloat32, + scale_dptr != nullptr ? defaultShape : emptyShape}; nvte_set_tensor_param(&tensor_, kNVTEScale, &scale); + if (scale_inv_dptr == nullptr && scale_inv_shape.ndim == defaultShape.ndim && + scale_inv_shape.ndim == 1 && scale_inv_shape.data[0] == defaultShape.data[0]) { + // Scale-inv pointer has not been provided and shape matches default + scale_inv_shape = emptyShape; + } NVTEBasicTensor scale_inv = {scale_inv_dptr, kNVTEFloat32, scale_inv_shape}; nvte_set_tensor_param(&tensor_, kNVTERowwiseScaleInv, &scale_inv); } @@ -734,7 +742,7 @@ class TensorWrapper { */ const NVTEShape shape() const noexcept { if (tensor_ == nullptr) { - return nvte_make_shape(nullptr, 0); + return emptyShape; } return nvte_tensor_shape(tensor_); } @@ -745,7 +753,7 @@ class TensorWrapper { */ const NVTEShape columnwise_shape() const noexcept { if (tensor_ == nullptr) { - return nvte_make_shape(nullptr, 0); + return emptyShape; } return nvte_tensor_columnwise_shape(tensor_); } @@ -869,7 +877,7 @@ class TensorWrapper { */ const NVTEShape scale_inv_shape() const noexcept { if (tensor_ == nullptr) { - return nvte_make_shape(nullptr, 0); + return emptyShape; } return nvte_tensor_scale_inv_shape(tensor_); } @@ -888,6 +896,7 @@ class TensorWrapper { static constexpr size_t defaultData = 1; static constexpr NVTEShape defaultShape = { {defaultData, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 1}; + static constexpr NVTEShape emptyShape = {{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 1}; private: NVTEShape convertShape(const NVTEShape &s) { return s; } diff --git a/transformer_engine/common/normalization/common.cpp b/transformer_engine/common/normalization/common.cpp index 337b165080..70e814e806 100644 --- a/transformer_engine/common/normalization/common.cpp +++ b/transformer_engine/common/normalization/common.cpp @@ -127,7 +127,13 @@ void TeNormalizationPlan::_build() { template std::vector TeNormalizationPlan::getWorkspaceShape() const { - return {_launch_params.getTotalWorkspaceBytes(_is_layernorm)}; + size_t workspace_size = _launch_params.getTotalWorkspaceBytes(_is_layernorm); + if (workspace_size == 0) { + // Workspace size must not be zero since that corresponds to a + // workspace size query + workspace_size = 1; + } + return {workspace_size}; } template @@ -405,7 +411,13 @@ void CudnnNormalizationPlan::_build() { } std::vector CudnnNormalizationPlan::getWorkspaceShape() const { - return {static_cast(_graph.get_workspace_size())}; + size_t workspace_size = _graph.get_workspace_size(); + if (workspace_size == 0) { + // Workspace size must not be zero since that corresponds to a + // workspace size query + workspace_size = 1; + } + return {workspace_size}; } void CudnnNormalizationPlan::execute(Tensor* z, void* x_dptr, void* gamma_dptr, void* beta_dptr, diff --git a/transformer_engine/common/normalization/layernorm/ln_api.cpp b/transformer_engine/common/normalization/layernorm/ln_api.cpp index 5785fd2233..b83ae25f25 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -51,7 +51,7 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size "RSigma must be 1D tensor with shape (x.shape[0],)."); NVTE_CHECK(rsigma->data.dtype == DType::kFloat32, "RSigma must be a float32 tensor."); - if (!workspace->data.shape.empty()) { + if (workspace->data.numel() != 0) { CheckInputTensor(x, "x"); CheckInputTensor(gamma, "gamma"); CheckInputTensor(beta, "beta"); @@ -94,7 +94,7 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size multiprocessorCount, zero_centered_gamma, is_aligned, z->scaling_mode, training, gamma_in_weight_dtype); - if (workspace->data.shape.empty()) { + if (workspace->data.numel() == 0) { workspace->data.shape = plan->getWorkspaceShape(); workspace->data.dtype = DType::kByte; return; @@ -146,7 +146,7 @@ void layernorm_bwd(const Tensor& dz, const Tensor& x, const Tensor& mu, const Te NVTE_CHECK(dbeta->data.shape == gamma.data.shape); NVTE_CHECK(dbeta->data.dtype == gamma.data.dtype); - if (!workspace->data.shape.empty()) { + if (workspace->data.numel() != 0) { CheckInputTensor(dz, "dz"); CheckInputTensor(x, "x"); CheckInputTensor(mu, "mu"); @@ -179,7 +179,7 @@ void layernorm_bwd(const Tensor& dz, const Tensor& x, const Tensor& mu, const Te multiprocessorCount, zero_centered_gamma, is_aligned, NVTE_DELAYED_TENSOR_SCALING, true, gamma_in_weight_dtype); - if (workspace->data.shape.empty()) { + if (workspace->data.numel() == 0) { workspace->data.shape = plan->getWorkspaceShape(); workspace->data.dtype = DType::kByte; return; diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index a3b05f7a29..ea6c972bf5 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -39,7 +39,7 @@ void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tens "RSigma must be 1D tensor with shape (x.shape[0],)."); NVTE_CHECK(rsigma->data.dtype == DType::kFloat32, "RSigma must be a float32 tensor."); - if (!workspace->data.shape.empty()) { + if (workspace->data.numel() != 0) { CheckInputTensor(x, "x"); CheckInputTensor(gamma, "gamma"); @@ -79,7 +79,7 @@ void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tens multiprocessorCount, zero_centered_gamma, is_aligned, z->scaling_mode, training, gamma_in_weight_dtype); - if (workspace->data.shape.empty()) { + if (workspace->data.numel() == 0) { workspace->data.shape = plan->getWorkspaceShape(); workspace->data.dtype = DType::kByte; return; @@ -125,7 +125,7 @@ void rmsnorm_bwd(const Tensor &dz, const Tensor &x, const Tensor &rsigma, const NVTE_CHECK(dgamma->data.shape == gamma.data.shape); NVTE_CHECK(dgamma->data.dtype == gamma.data.dtype); - if (!workspace->data.shape.empty()) { + if (workspace->data.numel() != 0) { CheckInputTensor(dz, "dz"); CheckInputTensor(x, "x"); CheckInputTensor(rsigma, "rsigma"); @@ -156,7 +156,7 @@ void rmsnorm_bwd(const Tensor &dz, const Tensor &x, const Tensor &rsigma, const multiprocessorCount, zero_centered_gamma, is_aligned, NVTE_DELAYED_TENSOR_SCALING, true, gamma_in_weight_dtype); - if (workspace->data.shape.empty()) { + if (workspace->data.numel() == 0) { workspace->data.shape = plan->getWorkspaceShape(); workspace->data.dtype = DType::kByte; return; @@ -191,7 +191,7 @@ void rmsnorm_bwd_add(const Tensor &dz, const Tensor &x, const Tensor &add, const NVTE_CHECK(dgamma->data.shape == gamma.data.shape); NVTE_CHECK(dgamma->data.dtype == gamma.data.dtype); - if (!workspace->data.shape.empty()) { + if (workspace->data.numel() != 0) { CheckInputTensor(dz, "dz"); CheckInputTensor(x, "x"); CheckInputTensor(add, "add"); @@ -222,7 +222,7 @@ void rmsnorm_bwd_add(const Tensor &dz, const Tensor &x, const Tensor &add, const multiprocessorCount, zero_centered_gamma, is_aligned, NVTE_DELAYED_TENSOR_SCALING, true, gamma_in_weight_dtype); - if (workspace->data.shape.empty()) { + if (workspace->data.numel() == 0) { workspace->data.shape = plan->getWorkspaceShape(); workspace->data.dtype = DType::kByte; return; diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 06735e3104..2cb43e8f27 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -332,68 +332,118 @@ __global__ void multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_ } // namespace void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t stream) { - NVTE_CHECK( - input->scaling_mode == NVTE_MXFP8_1D_SCALING || input->scaling_mode == NVTE_NVFP4_1D_SCALING, - "Input tensor has invalid scaling mode (", to_string(input->scaling_mode), ")."); - NVTE_CHECK(is_fp8_dtype(input->dtype()) || is_fp4_dtype(input->dtype()), - "Input tensor has invalid dtype (", to_string(input->dtype()), ")."); - - // Do nothing if tensor is empty - if (input->data.numel() == 0) { - return; - } + // Check scaling mode + const auto& scaling_mode = input->scaling_mode; + NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING, + "Input tensor has invalid scaling mode (", to_string(input->scaling_mode), ")."); + // Check tensors CheckInputTensor(*input, "scaling_factor_input"); CheckInputTensor(*output, "scaling_factor_output"); + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP8, got ", + to_string(input->dtype()), ")."); + break; + case NVTE_NVFP4_1D_SCALING: + NVTE_CHECK(is_fp4_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP4, got ", + to_string(input->dtype()), ")."); + break; + default: + NVTE_ERROR("Invalid scaling mode"); + } - auto& scaling_mode = input->scaling_mode; - NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING, - "Unsupported scaling mode for swizzling."); - - bool nvfp4 = scaling_mode == NVTE_NVFP4_1D_SCALING; + // Check if scaling factors are non-trivial + const bool has_rowwise_scale_inv = input->scale_inv.has_data(); + const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); + NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, + "Input tensor has both row-wise and column-wise scaling factors"); + if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) { + return; + } - // 1D block scaling, row-wise or colum-wise - int m, k; - if (input->has_data()) { - m = input->scale_inv.shape[0]; - k = input->scale_inv.shape[1]; - } else { - if (nvfp4) { - m = input->columnwise_scale_inv.shape[0]; - k = input->columnwise_scale_inv.shape[1]; - } else { - m = input->columnwise_scale_inv.shape[1]; - k = input->columnwise_scale_inv.shape[0]; + // Deduce tensor dims + int m{0}, k{0}; + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + if (has_rowwise_scale_inv) { + NVTE_CHECK(input->scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", input->scale_inv.shape, "."); + m = input->scale_inv.shape[0]; + k = input->scale_inv.shape[1]; + } else if (has_columnwise_scale_inv) { + NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", input->columnwise_scale_inv.shape, + "."); + m = input->columnwise_scale_inv.shape[1]; + k = input->columnwise_scale_inv.shape[0]; + } + break; + } + case NVTE_NVFP4_1D_SCALING: { + if (has_rowwise_scale_inv) { + NVTE_CHECK(input->scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", input->scale_inv.shape, "."); + m = input->scale_inv.shape[0]; + k = input->scale_inv.shape[1]; + } else if (has_columnwise_scale_inv) { + NVTE_CHECK(input->columnwise_scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", input->columnwise_scale_inv.shape, + "."); + m = input->columnwise_scale_inv.shape[0]; + k = input->columnwise_scale_inv.shape[1]; + } + break; } + default: + NVTE_ERROR("Invalid scaling mode"); } + // Check dims constexpr int SF_TILE_DIM_M = 128; constexpr int SF_TILE_DIM_K = 4; - NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Input should be padded in M/N dimension!"); NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Input should be padded in K dimension!"); - NVTE_CHECK(k > 0, "Input scale inverse should be 2D!"); - if (output->has_data()) { - NVTE_CHECK(m * k == std::accumulate(output->scale_inv.shape.begin(), - output->scale_inv.shape.end(), 1, std::multiplies()), - "Input.scale_inv size is not equal to Output.scale_inv size!"); + + // Check that output tensor matches input tensor + if (has_rowwise_scale_inv) { + NVTE_CHECK(output->scale_inv.has_data(), + "Output tensor does not have row-wise scaling factors."); + NVTE_CHECK(m * k == output->scale_inv.numel(), "Expected output tensor to have ", m * k, + " row-wise scaling factors, but got shape=", output->scale_inv.shape, "."); } - if (output->has_columnwise_data()) { - NVTE_CHECK(m * k == std::accumulate(output->columnwise_scale_inv.shape.begin(), - output->columnwise_scale_inv.shape.end(), 1, - std::multiplies()), - "Input.columnwise_scale_inv size is not equal to " - "Output.columnwise_scale_inv size!"); + if (has_columnwise_scale_inv) { + NVTE_CHECK(output->columnwise_scale_inv.has_data(), + "Output tensor does not have column-wise scaling factors."); + NVTE_CHECK( + m * k == output->columnwise_scale_inv.numel(), "Expected output tensor to have ", m * k, + " column-wise scaling factors, but got shape=", output->columnwise_scale_inv.shape, "."); } - int num_tiles_m = m / SF_TILE_DIM_M; - int num_tiles_k = k / SF_TILE_DIM_K; + // Choose swizzle implementation + bool rowwise_swizzle{false}, columnwise_swizzle{false}; + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + rowwise_swizzle = has_rowwise_scale_inv; + columnwise_swizzle = has_columnwise_scale_inv; + break; + } + case NVTE_NVFP4_1D_SCALING: { + // NVFP4 column-wise data is transposed, so row-wise and + // column-wise scales have same swizzling format + rowwise_swizzle = true; + columnwise_swizzle = false; + break; + } + default: + NVTE_ERROR("Invalid scaling mode"); + } - // For NVFP4, the scale inverse for tranposed data needs rowwise swizzle. - const bool rowwise_swizzle = input->has_data() || nvfp4; - const bool columnwise_swizzle = input->has_columnwise_data() && !nvfp4; + const dim3 block_size(TB_DIM, TB_DIM); + const int num_tiles_m = m / SF_TILE_DIM_M; + const int num_tiles_k = k / SF_TILE_DIM_K; - dim3 block_size(TB_DIM, TB_DIM); + // 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 */ @@ -402,20 +452,32 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s 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, original_K; - void *input_scale_inv_ptr, *output_scale_inv_ptr; - - if (!nvfp4 || input->has_data()) { - int block_scale_size = nvfp4 ? NVFP4_BLOCK_SIZE : MXFP8_BLOCK_SIZE; - original_M = input->flat_first_dim(); - original_K = input->flat_last_dim() / block_scale_size; - input_scale_inv_ptr = input->scale_inv.dptr; - output_scale_inv_ptr = output->scale_inv.dptr; - } else { - original_M = input->flat_last_dim(); - original_K = input->flat_first_dim() / NVFP4_BLOCK_SIZE; - input_scale_inv_ptr = input->columnwise_scale_inv.dptr; - output_scale_inv_ptr = output->columnwise_scale_inv.dptr; + int original_M{0}, original_K{0}; + void *input_scale_inv_ptr{nullptr}, *output_scale_inv_ptr{nullptr}; + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + original_M = input->flat_first_dim(); + original_K = input->flat_last_dim() / MXFP8_BLOCK_SIZE; + input_scale_inv_ptr = input->scale_inv.dptr; + output_scale_inv_ptr = output->scale_inv.dptr; + break; + } + case NVTE_NVFP4_1D_SCALING: { + if (has_rowwise_scale_inv) { + original_M = input->flat_first_dim(); + original_K = input->flat_last_dim() / NVFP4_BLOCK_SIZE; + input_scale_inv_ptr = input->scale_inv.dptr; + output_scale_inv_ptr = output->scale_inv.dptr; + } else if (has_columnwise_scale_inv) { + original_M = input->flat_last_dim(); + original_K = input->flat_first_dim() / NVFP4_BLOCK_SIZE; + input_scale_inv_ptr = input->columnwise_scale_inv.dptr; + output_scale_inv_ptr = output->columnwise_scale_inv.dptr; + } + break; + } + default: + NVTE_ERROR("Invalid scaling mode"); } switch (vec_load_size) { @@ -447,7 +509,10 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s 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 */ @@ -456,8 +521,6 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s 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; - // NVFP4 shouldn't end up here because it only needs rowwise swizzle - NVTE_CHECK(!nvfp4, "NVFP4 shouldn't end up here because it only needs rowwise swizzle"); switch (vec_load_size) { case 4: @@ -491,9 +554,8 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_ERROR("Not valid vec_load_size."); break; } + NVTE_CHECK_CUDA(cudaGetLastError()); } - - NVTE_CHECK_CUDA(cudaGetLastError()); } template @@ -595,17 +657,18 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)), "Not implemented scaling mode " + to_string(scaling_mode) + "."); // We don't allow empty tensors. They should be filtered out before calling this function. - if (input[i]->data.numel() == 0) { - NVTE_ERROR("Tensor input[" + std::to_string(i) + "] is empty."); - } + 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) + "]"); - all_has_data &= input[i]->has_data(); - all_has_columnwise_data &= input[i]->has_columnwise_data(); - all_nvfp4 &= is_nvfp4_scaling(scaling_mode); + 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()); + all_nvfp4 = all_nvfp4 && is_nvfp4_scaling(scaling_mode); } NVTE_CHECK(all_has_data || all_has_columnwise_data, "All tensors should have data or columnwise data."); + NVTE_CHECK(!all_has_data || !all_has_columnwise_data, + "All tensors have both data and columnwise data."); const bool rowwise_swizzle = all_has_data || all_nvfp4; const bool columnwise_swizzle = all_has_columnwise_data && !all_nvfp4; @@ -644,18 +707,19 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Input should be padded in K dimension!"); NVTE_CHECK(k > 0, "Input scale inverse should be 2D!"); - if (output[i]->has_data()) { - NVTE_CHECK( - m * k == std::accumulate(output[i]->scale_inv.shape.begin(), - output[i]->scale_inv.shape.end(), 1, std::multiplies()), - "Input.scale_inv size is not equal to Output.scale_inv size!"); + if (all_has_data) { + NVTE_CHECK(output[i]->scale_inv.has_data(), "Output tensor ", i, + " does not have row-wise scaling factors."); + NVTE_CHECK(m * k == output[i]->scale_inv.numel(), "Expected output tensor ", i, " to have ", + m * k, " row-wise scaling factors, but got shape=", output[i]->scale_inv.shape, + "."); } - if (output[i]->has_columnwise_data()) { - NVTE_CHECK(m * k == std::accumulate(output[i]->columnwise_scale_inv.shape.begin(), - output[i]->columnwise_scale_inv.shape.end(), 1, - std::multiplies()), - "Input.columnwise_scale_inv size is not equal to " - "Output.columnwise_scale_inv size!"); + if (all_has_columnwise_data) { + NVTE_CHECK(output[i]->columnwise_scale_inv.has_data(), "Output tensor ", i, + " does not have column-wise scaling factors."); + NVTE_CHECK(m * k == output[i]->columnwise_scale_inv.numel(), "Expected output tensor ", i, + " to have ", m * k, " column-wise scaling factors, but got shape=", + output[i]->columnwise_scale_inv.shape, "."); } int num_tiles_k = k / SF_TILE_DIM_K; diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index e9e1d5bfb1..8d9563b789 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -77,7 +77,7 @@ std::string to_string(const NVTEScalingMode &mode) { } void CheckNoopTensor(const Tensor &t, const std::string &name) { - if (t.data.dptr != nullptr) { + if (t.data.has_data()) { NVTE_CHECK(t.numel() == 1, "Expected 1 element for ", name, " noop, but found ", t.numel(), "."); NVTE_CHECK(t.data.dtype == DType::kFloat32, "Found wrong dtype for ", name, @@ -88,15 +88,30 @@ void CheckNoopTensor(const Tensor &t, const std::string &name) { void CheckScaleTensorShape(const Tensor &t, const std::string &name) { NVTE_CHECK(t.scaling_mode != NVTE_INVALID_SCALING, "Invalid scaling mode!"); if (is_tensor_scaling(t.scaling_mode)) { - // per-tensor scaling - if (t.has_data()) { - NVTE_CHECK(t.scale_inv.numel() == 1, "Tensor \"", name, - "\" has invalid scale_inv shape (expected (1), got ", t.scale_inv.shape, ")"); - } - if (t.has_columnwise_data()) { - NVTE_CHECK(t.columnwise_scale_inv.numel() == 1, "Tensor \"", name, - "\" has invalid columnwise_scale_inv shape (expected (1), got ", - t.columnwise_scale_inv.shape, ")"); + if (is_fp8_dtype(t.dtype())) { + // FP8 tensor with tensor scaling + if (t.has_data()) { + NVTE_CHECK(t.scale_inv.numel() == 1, "Tensor \"", name, + "\" has invalid scale_inv shape (expected 1 entry, got ", t.scale_inv.shape, + ")"); + } + if (t.has_columnwise_data()) { + NVTE_CHECK(t.columnwise_scale_inv.numel() == 1, "Tensor \"", name, + "\" has invalid columnwise_scale_inv shape (expected 1 entry, got ", + t.columnwise_scale_inv.shape, ")"); + } + } else { + // High-precision tensor + if (t.has_data()) { + NVTE_CHECK(t.scale_inv.numel() == 0, "Tensor \"", name, + "\" has invalid scale_inv shape (expected 0 entries, got ", t.scale_inv.shape, + ")"); + } + if (t.has_columnwise_data()) { + NVTE_CHECK(t.columnwise_scale_inv.numel() == 0, "Tensor \"", name, + "\" has invalid columnwise_scale_inv shape (expected 0 entries, got ", + t.columnwise_scale_inv.shape, ")"); + } } } else { if (t.scaling_mode == NVTE_MXFP8_1D_SCALING) { @@ -159,7 +174,7 @@ void CheckInputTensor(const Tensor &t, const std::string &name) { if (is_fp8_dtype(type)) { // FP8 input needs to have scale_inv if (t.has_data()) { - NVTE_CHECK(t.scale_inv.dptr != nullptr, "FP8 scaling factor input ", name, + NVTE_CHECK(t.scale_inv.has_data(), "FP8 scaling factor input ", name, "_scale_inverse must be allocated"); NVTE_CHECK(t.scale_inv.dtype == DType::kFloat32 || t.scale_inv.dtype == DType::kFloat8E8M0, "FP8 scaling factor input ", name, @@ -168,7 +183,7 @@ void CheckInputTensor(const Tensor &t, const std::string &name) { to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { - NVTE_CHECK(t.columnwise_scale_inv.dptr != nullptr, "FP8 scaling factor input ", name, + NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP8 scaling factor input ", name, "_columnwise_scale_inverse must be allocated"); NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat32 || t.columnwise_scale_inv.dtype == DType::kFloat8E8M0, @@ -181,7 +196,7 @@ void CheckInputTensor(const Tensor &t, const std::string &name) { // TODO(ksivaman): Fix this to check for amaxes and other details. // For now only needed for swizzle. if (t.has_data()) { - NVTE_CHECK(t.scale_inv.dptr != nullptr, "FP4 scaling factor input ", name, + NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor input ", name, "_scale_inverse must be allocated"); NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor input ", name, "_scale_inverse has invalid dtype " @@ -189,7 +204,7 @@ void CheckInputTensor(const Tensor &t, const std::string &name) { to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { - NVTE_CHECK(t.columnwise_scale_inv.dptr != nullptr, "FP4 scaling factor input ", name, + NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor input ", name, "_columnwise_scale_inverse must be allocated"); NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3, "FP8 scaling factor input ", name, @@ -198,11 +213,10 @@ void CheckInputTensor(const Tensor &t, const std::string &name) { to_string(t.columnwise_scale_inv.dtype), ")"); } } else { - NVTE_CHECK(t.scale.dptr == nullptr, "Scale is not supported for non-FP8 input ", name); - NVTE_CHECK(t.amax.dptr == nullptr, "Amax is not supported for non-FP8 input ", name); - NVTE_CHECK(t.scale_inv.dptr == nullptr, "Scale_inv is not supported for non-FP8 input ", name); - NVTE_CHECK(t.columnwise_scale_inv.dptr == nullptr, - "Scale_inv is not supported for non-FP8 input ", name); + NVTE_CHECK(!t.scale.has_data(), "Scale is not supported for non-FP8 input ", name); + NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv is not supported for non-FP8 input ", name); + NVTE_CHECK(!t.columnwise_scale_inv.has_data(), "Scale_inv is not supported for non-FP8 input ", + name); } NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Input ", name, " is not allocated!"); @@ -213,14 +227,14 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt const DType type = t.dtype(); if (is_fp8_dtype(type)) { // FP8 output needs to have scale, scale_inv and (if delayed scaling) amax - if (t.scaling_mode == NVTE_DELAYED_TENSOR_SCALING && t.amax.dptr != nullptr) { + if (t.scaling_mode == NVTE_DELAYED_TENSOR_SCALING && t.amax.has_data()) { NVTE_CHECK(t.amax.dtype == DType::kFloat32, "Invalid amax dtype (expected ", to_string(DType::kFloat32), ", got ", to_string(t.amax.dtype), ")"); - NVTE_CHECK(product(t.amax.shape) == 1, "Invalid shape of amax in output ", name, + NVTE_CHECK(t.amax.numel() == 1, "Invalid shape of amax in output ", name, " (expected 1 entry, got shape=", t.amax.shape, ")"); } if (t.has_data()) { - NVTE_CHECK(t.scale_inv.dptr != nullptr, "FP8 scaling factor output ", name, + NVTE_CHECK(t.scale_inv.has_data(), "FP8 scaling factor output ", name, "_scale_inverse must be allocated"); NVTE_CHECK(t.scale_inv.dtype == DType::kFloat32 || t.scale_inv.dtype == DType::kFloat8E8M0, "FP8 scaling factor output ", name, @@ -229,7 +243,7 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { - NVTE_CHECK(t.columnwise_scale_inv.dptr != nullptr, "FP8 scaling factor output ", name, + NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP8 scaling factor output ", name, "_columnwise_scale_inverse must be allocated"); NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat32 || t.columnwise_scale_inv.dtype == DType::kFloat8E8M0, @@ -241,7 +255,7 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt } else if (is_fp4_dtype(type)) { // FP4 output needs to have the scale_inv if (t.has_data()) { - NVTE_CHECK(t.scale_inv.dptr != nullptr, "FP4 scaling factor output ", name, + NVTE_CHECK(t.scale_inv.has_data(), "FP4 scaling factor output ", name, "_scale_inverse must be allocated"); NVTE_CHECK(t.scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor output ", name, "_scale_inverse has invalid dtype " @@ -249,7 +263,7 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt to_string(t.scale_inv.dtype), ")"); } if (t.has_columnwise_data()) { - NVTE_CHECK(t.columnwise_scale_inv.dptr != nullptr, "FP4 scaling factor output ", name, + NVTE_CHECK(t.columnwise_scale_inv.has_data(), "FP4 scaling factor output ", name, "_columnwise_scale_inverse must be allocated"); NVTE_CHECK(t.columnwise_scale_inv.dtype == DType::kFloat8E4M3, "FP4 scaling factor output ", name, @@ -258,12 +272,10 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt to_string(t.columnwise_scale_inv.dtype), ")"); } } else { - NVTE_CHECK(t.scale.dptr == nullptr, "Scale is not supported for non-FP8 output ", name); - // Unfused quant with level 2 nvfp4 scaling will produce high precision tensors with amax. - // NVTE_CHECK(t.amax.dptr == nullptr, "Amax is not supported for non-FP8 output ", name); - NVTE_CHECK(t.scale_inv.dptr == nullptr, "Scale_inv is not supported for non-FP8 output ", name); - NVTE_CHECK(t.columnwise_scale_inv.dptr == nullptr, - "Scale_inv is not supported for non-FP8 input ", name); + NVTE_CHECK(!t.scale.has_data(), "Scale is not supported for non-FP8 output ", name); + NVTE_CHECK(!t.scale_inv.has_data(), "Scale_inv is not supported for non-FP8 output ", name); + NVTE_CHECK(!t.columnwise_scale_inv.has_data(), "Scale_inv is not supported for non-FP8 input ", + name); } if (!allow_empty) { @@ -622,7 +634,11 @@ NVTEShape nvte_make_shape(const size_t *data, size_t ndim) { NVTE_CHECK(ndim <= sizeof(ret.data) / sizeof(ret.data[0]), "Too many dims for NVTEShape (requested: ", ndim, ", max: ", sizeof(ret.data) / sizeof(ret.data[0]), ")"); - std::copy(data, data + ndim, ret.data); + if (data == nullptr) { + std::fill(ret.data, ret.data + ndim, 0); + } else { + std::copy(data, data + ndim, ret.data); + } ret.ndim = ndim; return ret; } @@ -729,7 +745,7 @@ void *nvte_tensor_columnwise_scale_inv(const NVTETensor tensor) { NVTEShape nvte_tensor_scale_inv_shape(const NVTETensor tensor) { auto *t = transformer_engine::convertNVTETensor(tensor); if (t == nullptr) { - return nvte_make_shape(nullptr, 0); + return nvte_make_shape(nullptr, 1); } return nvte_make_shape(t->scale_inv.shape.data(), t->scale_inv.shape.size()); } @@ -768,7 +784,7 @@ void nvte_set_tensor_param(NVTETensor *tensor, NVTETensorParam param_name, NVTEBasicTensor nvte_get_tensor_param(const NVTETensor tensor, NVTETensorParam param_name) { if (tensor == nullptr) { - return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 0)}; + return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 1)}; } const auto &t = *transformer_engine::convertNVTETensorCheck(tensor); switch (param_name) { @@ -813,14 +829,21 @@ void nvte_tensor_pack_destroy(NVTETensorPack *pack) { void nvte_zero_tensor(const NVTETensor tensor, cudaStream_t stream) { if (tensor == nullptr) return; const auto &t = *transformer_engine::convertNVTETensorCheck(tensor); + // Zero out tensor data if allocated if (t.data.dptr != nullptr) { - const size_t size_in_bytes = nvte_tensor_size_bytes(tensor); - NVTE_CHECK_CUDA(cudaMemsetAsync(t.data.dptr, 0, size_in_bytes, stream)); + const auto size = t.data.buffer_size_bytes(); + if (size > 0) { + NVTE_CHECK_CUDA(cudaMemsetAsync(t.data.dptr, 0, size, stream)); + } } - // Set amax to 0 if allocated + + // Zero out amax if allocated if (t.amax.dptr != nullptr) { - NVTE_CHECK_CUDA(cudaMemsetAsync(t.amax.dptr, 0, sizeof(float), stream)); + const auto size = t.amax.buffer_size_bytes(); + if (size > 0) { + NVTE_CHECK_CUDA(cudaMemsetAsync(t.amax.dptr, 0, size, stream)); + } } } @@ -1007,7 +1030,7 @@ void nvte_set_grouped_tensor_param(NVTEGroupedTensor *tensor, NVTEGroupedTensorP NVTEBasicTensor nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, NVTEGroupedTensorParam param_name) { if (tensor == nullptr) { - return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 0)}; + return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 1)}; } const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); @@ -1059,7 +1082,7 @@ NVTEScalingMode nvte_grouped_tensor_scaling_mode(const NVTEGroupedTensor tensor) NVTEShape nvte_get_grouped_tensor_logical_shape(const NVTEGroupedTensor tensor) { if (tensor == nullptr) { - return nvte_make_shape(nullptr, 0); + return nvte_make_shape(nullptr, 1); } const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); return t.logical_shape; diff --git a/transformer_engine/common/transpose/cast_transpose_fusion.cu b/transformer_engine/common/transpose/cast_transpose_fusion.cu index 6329e79ae7..a04ab902ca 100644 --- a/transformer_engine/common/transpose/cast_transpose_fusion.cu +++ b/transformer_engine/common/transpose/cast_transpose_fusion.cu @@ -198,8 +198,6 @@ void populate_cast_transpose_dbias_workspace_config(const Tensor &cast_output, / workspace->data.dtype); const size_t required_size = get_buffer_size_bytes(num_rows_partial_dbias, row_length, DType::kFloat32); - NVTE_CHECK(!workspace->data.shape.empty(), "Invalid workspace dims (expected (", - num_rows_partial_dbias, ",", row_length, "), found ())"); NVTE_CHECK(workspace_size >= required_size, "Invalid workspace (expected dims=(", num_rows_partial_dbias, ",", row_length, "), dtype=", to_string(DType::kFloat32), "; found dims=", workspace->data.shape, diff --git a/transformer_engine/common/transpose/transpose_fusion.cu b/transformer_engine/common/transpose/transpose_fusion.cu index 3c51ce3dab..75fa07a5b3 100644 --- a/transformer_engine/common/transpose/transpose_fusion.cu +++ b/transformer_engine/common/transpose/transpose_fusion.cu @@ -388,8 +388,6 @@ void populate_transpose_dbias_workspace_config(const Tensor &input, /*cast*/ workspace->data.dtype); const size_t required_size = get_buffer_size_bytes(num_rows_partial_dbias, row_length, DType::kFloat32); - NVTE_CHECK(!workspace->data.shape.empty(), "Invalid workspace dims (expected (", - num_rows_partial_dbias, ",", row_length, "), found ())"); NVTE_CHECK(workspace_size >= required_size, "Invalid workspace (expected dims=(", num_rows_partial_dbias, ",", row_length, "), dtype=", to_string(DType::kFloat32), "; found dims=", workspace->data.shape, diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index b12da7542b..fb5f0b55d4 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -334,12 +334,12 @@ std::tuple, std::vector> bulk_allocate_fp tensor_cpp_list.emplace_back(makeTransformerEngineTensor( rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr, columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_data_shapes[i] : std::vector{}, - columnwise_usage ? columnwise_data_shapes[i] : std::vector{}, fp8_dtype, nullptr, + rowwise_usage ? rowwise_data_shapes[i] : std::vector{0}, + columnwise_usage ? columnwise_data_shapes[i] : std::vector{0}, fp8_dtype, nullptr, nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr, columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_scale_shapes[i] : std::vector{}, - columnwise_usage ? columnwise_scale_shapes[i] : std::vector{}, scaling_mode)); + rowwise_usage ? rowwise_scale_shapes[i] : std::vector{0}, + columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode)); } return retval; @@ -481,12 +481,12 @@ std::tuple, std::vector> bulk_allocate_mx tensor_cpp_list.emplace_back(makeTransformerEngineTensor( rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr, columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_data_shapes[i] : std::vector{}, - columnwise_usage ? columnwise_data_shapes[i] : std::vector{}, fp8_dtype, nullptr, + rowwise_usage ? rowwise_data_shapes[i] : std::vector{0}, + columnwise_usage ? columnwise_data_shapes[i] : std::vector{0}, fp8_dtype, nullptr, nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr, columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_scale_shapes[i] : std::vector{}, - columnwise_usage ? columnwise_scale_shapes[i] : std::vector{}, scaling_mode)); + rowwise_usage ? rowwise_scale_shapes[i] : std::vector{0}, + columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode)); } return retval; @@ -685,13 +685,13 @@ std::tuple, std::vector, bool> bulk_alloc auto tensor_wrapper = makeTransformerEngineTensor( rowwise_usage ? rowwise_data_list[i].data_ptr() : nullptr, columnwise_usage ? columnwise_data_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_data_shapes[i] : std::vector{}, - columnwise_usage ? columnwise_data_shapes[i] : std::vector{}, fp4_dtype, + rowwise_usage ? rowwise_data_shapes[i] : std::vector{0}, + columnwise_usage ? columnwise_data_shapes[i] : std::vector{0}, fp4_dtype, /*amax_ptr=*/nullptr, /*scale_ptr=*/nullptr, rowwise_usage ? rowwise_scale_list[i].data_ptr() : nullptr, columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, - rowwise_usage ? rowwise_scale_shapes[i] : std::vector{}, - columnwise_usage ? columnwise_scale_shapes[i] : std::vector{}, scaling_mode); + rowwise_usage ? rowwise_scale_shapes[i] : std::vector{0}, + columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode); // Set the amax rowwise and amax columnwise if available if (rowwise_usage) { diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 13e8bfb6e5..335052296f 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -43,10 +43,10 @@ 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 = product(A_shape, 0, A_shape.ndim - 1); - const size_t A1 = A_shape.data[A_shape.ndim - 1]; - const size_t B0 = product(B_shape, 0, B_shape.ndim - 1); - const size_t B1 = B_shape.data[B_shape.ndim - 1]; + 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; // 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/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp index 8d1d865604..63c26ee303 100644 --- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp +++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp @@ -22,8 +22,8 @@ void compute_amax(const at::Tensor& tensor, at::Tensor& amax) { TORCH_CHECK(amax.numel() == 1, "amax must have exactly one element"); auto* amax_ptr = amax.data_ptr(); TensorWrapper fake_te_output( - nullptr, te_input.shape(), - DType::kFloat8E4M3, // It doesn't matter because we only compute amax. + /*dptr=*/nullptr, te_input.shape(), + DType::kFloat32, // It doesn't matter because we only compute amax. amax_ptr); nvte_compute_amax(te_input.data(), fake_te_output.data(), at::cuda::getCurrentCUDAStream()); diff --git a/transformer_engine/pytorch/csrc/util.cpp b/transformer_engine/pytorch/csrc/util.cpp index 134185ac82..ce547d302e 100644 --- a/transformer_engine/pytorch/csrc/util.cpp +++ b/transformer_engine/pytorch/csrc/util.cpp @@ -142,13 +142,20 @@ std::optional multi_tensor_swizzle_scaling_factors( auto& tensor = tensors[i]; void* scale_inv_dptr = scale_inv_dptrs[i]; void* swizzled_scale_inv_dptr = getDataPtr(buffer, scale_inv_offsets[i]); - // auto input_shape = nvte_shape_to_vector(tensor.shape()); + + // Empty tensors don't require scale swizzling + if (tensor.numel() == 0) { + continue; + } + + // Tensor shape NVTEShape nvte_input_shape; if (rowwise) { nvte_input_shape = tensor.shape(); } else { nvte_input_shape = tensor.get_columnwise_data().shape; } + auto input_shape = nvte_shape_to_vector(nvte_input_shape); // Reconstruct input only to avoid swizzling both directions if not needed. // Use any 8 bit type, it's irrelevant. @@ -202,14 +209,14 @@ at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapp size_t data_flat_last_dim = 1; if (rowwise) { data = input.get_rowwise_data(); - for (int i = 0; i < data.shape.ndim - 1; ++i) { + 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 (int i = 1; i < data.shape.ndim; ++i) { + for (size_t i = 1; i < data.shape.ndim; ++i) { data_flat_last_dim *= data.shape.data[i]; } } From 50be029948410ef21184b11d91aacb4ace2f8636 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 4 Dec 2025 16:40:04 -0800 Subject: [PATCH 110/521] [JAX] Enable TE/JAX test timings in CI (#2475) enable jax test timings in qa scripts Signed-off-by: Jeremy Berchtold --- qa/L0_jax_distributed_unittest/test.sh | 2 ++ qa/L0_jax_unittest/test.sh | 2 ++ qa/L1_jax_distributed_unittest/test.sh | 2 ++ qa/L2_jax_distributed_unittest/test.sh | 2 ++ qa/L2_jax_unittest/test.sh | 2 ++ 5 files changed, 10 insertions(+) diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index ae45f398e8..5268f9ba0e 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -16,6 +16,8 @@ function test_fail() { RET=0 FAILED_CASES="" +export NVTE_JAX_TEST_TIMING=1 + : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index cb097d492a..92f7dd2525 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -18,6 +18,8 @@ function test_fail() { RET=0 FAILED_CASES="" +export NVTE_JAX_TEST_TIMING=1 + pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 886f27747e..5751d28200 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -11,6 +11,8 @@ function test_fail() { RET=0 FAILED_CASES="" +export NVTE_JAX_TEST_TIMING=1 + : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index de5624a596..b81331dbc2 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -4,6 +4,8 @@ set -xe +export NVTE_JAX_TEST_TIMING=1 + : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index f933a0732e..1b35596fd5 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -18,6 +18,8 @@ function test_fail() { RET=0 FAILED_CASES="" +export NVTE_JAX_TEST_TIMING=1 + pip3 install "nltk>=3.8.2" || error_exit "Failed to install nltk" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" From f0572aa5c3acd9043a7c687751cb46c0218b923c Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 4 Dec 2025 18:08:56 -0800 Subject: [PATCH 111/521] Fix bugs from refactoring C++ tensor class (#2481) Remve assumption in quantize/activation kernels that data buffer is initialized Signed-off-by: Tim Moon --- transformer_engine/common/cast/dispatch/dequantize.cuh | 6 +++--- transformer_engine/common/cast/dispatch/gated.cuh | 10 +++++----- .../common/cast/mxfp8/dequantize_mxfp8.cuh | 5 ++--- 3 files changed, 10 insertions(+), 11 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index b8547915c7..9138fda040 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -27,9 +27,9 @@ inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t switch (input.scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: { - NVTE_CHECK(is_fp8_dtype(input.data.dtype), "Input must have FP8 type."); - NVTE_CHECK(!is_fp8_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."); + NVTE_CHECK(is_fp8_dtype(input.dtype()), "Input must have FP8 type."); + NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision."); + NVTE_CHECK(output->shape() == input.shape(), "Input and output shapes need to match."); fp8::dequantize(input, output, stream); break; } diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index 4373090b72..f08b09317e 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -98,8 +98,8 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte const size_t rows = gated_input.flat_first_dim(); const size_t cols = gated_input.flat_last_dim() / 2; - NVTE_CHECK(!is_fp8_dtype(grad.data.dtype), "Grad input must be in higher precision."); - NVTE_CHECK(grad.data.dtype == gated_input.data.dtype, "Types of both inputs must match."); + NVTE_CHECK(!is_fp8_dtype(grad.dtype()), "Grad input must be in higher precision."); + NVTE_CHECK(grad.dtype() == gated_input.dtype(), "Types of both inputs must match."); NVTE_CHECK(grad.flat_first_dim() == rows, "Wrong Grad shape. Expected first dimension (after flattening) [", rows, ", *], got [", @@ -116,9 +116,9 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte NVTE_CHECK(output->flat_last_dim() == cols * 2, "Wrong output shape. Expected (after flattening) [*, ", cols * 2, "], got [", output->flat_first_dim(), ", ", output->flat_last_dim(), "]."); - NVTE_CHECK(gated_input.data.shape == output->data.shape, - "Gated input and output shapes must match. Input shape: ", gated_input.data.shape, - ", output shape: ", output->data.shape, "."); + NVTE_CHECK(gated_input.shape() == output->shape(), + "Gated input and output shapes must match. Input shape: ", gated_input.shape(), + ", output shape: ", output->shape(), "."); switch (output->scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: { diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index fb43fce96b..c56ebe172c 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -227,8 +227,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) bool use_colwise_scaling = input.has_columnwise_data(); checkCuDriverContext(stream); - const auto &input_shape = input.data.shape; - NVTE_CHECK(input_shape.size() >= 2, "Input must have at least 2 dimensions."); + NVTE_CHECK(input.dim() >= 2, "Input must have at least 2 dimensions."); if (use_rowwise_scaling) { NVTE_CHECK(input.has_data(), "Cannot dequantize tensor without rowwise data."); @@ -241,7 +240,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) } NVTE_CHECK(!is_fp8_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."); + NVTE_CHECK(output->shape() == input.shape(), "Input and output shapes need to match."); // TODO: Make more general const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1; From fd0cd12e912a101173275b15e10c3906f20d4551 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Fri, 5 Dec 2025 17:10:34 -0800 Subject: [PATCH 112/521] [JAX] Add CP + THD + AG + Striped>1 + SWA support (#2379) * Add generic stripe_height support for load balancing Signed-off-by: Kshitij Lakhani * Fix imports in test for deprecated jax.experimental.pjit Signed-off-by: Kshitij Lakhani * Add test case for stripe_height greater than 1. Add stripe_height arg to reordering methods Signed-off-by: Kshitij Janardan Lakhani * Add Striped 1 and 4 test cases. Refactor the Load Balancing test case. Fix the incorrect shape in striping inverser reordering Signed-off-by: Kshitij Janardan Lakhani * Modify test code for CP + AG + THD + stripe height greater than 1 Signed-off-by: Kshitij Janardan Lakhani * Add stripe_height arg to fused attn and fused attn fwd API. Add appropriate mask checks for AG+THD+CP and pick BRCM to be executed per rank. Add Fused Attn Primitive for CP + THD +AG + Striping. Add a method to reorder and all gather segment ids and offsets for kv Signed-off-by: Kshitij Janardan Lakhani * TMP: Throwaway testing commit Signed-off-by: Kshitij Janardan Lakhani * Add comments in primitive registration process Signed-off-by: Kshitij Janardan Lakhani * TMP: Throwaway test commit Signed-off-by: Kshitij Lakhani * Undoing incorrect rebase/merge leftovers Signed-off-by: Kshitij Janardan Lakhani * TMP: Throwaway test commits Signed-off-by: Kshitij Janardan Lakhani * Add support for calculating q and kv seqlens and offsets per rank for CP+THD+AG+SW+Striped>1 primitive Signed-off-by: Kshitij Janardan Lakhani * Augment jax primitive register code comments Signed-off-by: Kshitij Janardan Lakhani * Fix the array sizes and padding values returned for seqlens and offsets to fit what the fused attn primitive non cp computation Signed-off-by: Kshitij Lakhani * Add support in new primitive for softmax_offset related changes. Put in missing primitive registering line in again. Increase the seqoffsets arrays lengths by 1 Signed-off-by: Kshitij Janardan Lakhani * Add new set of helper functions for seqlens and seqoffsets fo AG+THD+CP+Stripe>1 which accounts for batching and seq offsets size b+1 Signed-off-by: Kshitij Lakhani * Add backward primitive for CP+THD+AG+Striped>1 Signed-off-by: Kshitij Lakhani * Modify tests for backward primitive for CP+THD+AG+Striped>1 Signed-off-by: Kshitij Lakhani * Move stripe_height along with other static args in fused_attn_bwd rule. Fix typo in CP+AG+TH+Striped>1 primitive Signed-off-by: Kshitij Janardan Lakhani * Code clean up: remove older version for calculating seqlens and offsets for CP+AG+THD+striped>1 primitive Signed-off-by: Kshitij Lakhani * Add test for CP+THD+AG+Striped>1 Signed-off-by: Kshitij Lakhani * Fix missing var Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add SWA tests for AG+Striped>1+CP+THD+SWA Signed-off-by: Kshitij Janardan Lakhani * Restoring test code Signed-off-by: Kshitij Janardan Lakhani * Remove assert preventing SWA code path in CP+AG+Striped primitive Signed-off-by: Kshitij Janardan Lakhani * Parametrize num_segments_per_seq in tests Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clean up test code Signed-off-by: Kshitij Lakhani Clean up test code in TE common Signed-off-by: Kshitij Lakhani Clean up debug statements Signed-off-by: Kshitij Lakhani * Rename stripe_height to stripe_size Signed-off-by: Kshitij Lakhani * Code clean up and add additional comments Signed-off-by: Kshitij Lakhani nit: Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Fix type on fused attn tests Signed-off-by: Kshitij Janardan Lakhani * Fix seqoffsets length to be passed onto FusedAttn primitive as it is b and not b+1 needed by cuDNN Signed-off-by: Kshitij Janardan Lakhani * Remove commented code Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Fix linting issues Signed-off-by: Kshitij Lakhani Fix incorrect greptile change Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip THD test cases for CP + AG + Dual chunk. Skip BSHD cases for CP + AG + Striped>1. Correct the layout and shapr parameters passed to the tests Signed-off-by: Kshitij Janardan Lakhani * Pass stripe_size explicitly for ring attn tests for THD cases Signed-off-by: Kshitij Janardan Lakhani * Remove TODO Signed-off-by: Kshitij Janardan Lakhani * Explicitly fail if THD + AG is being used with a non padding causal mask Signed-off-by: Kshitij Janardan Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Correct the ID for the test dist fused attn tests to account for cp*2 which is done under the hood Signed-off-by: Kshitij Lakhani * Set num_segments_per_seq defaults to None instead of 0 Signed-off-by: Kshitij Lakhani * Augment comments. Add ValueError for stripe_size=0 Signed-off-by: Kshitij Lakhani * Test only 1 num_segments_per_seq combination for CP+AG+THD+Striped>1+SWA instead of 2. Modify the num segments and window size to easily to debug values Signed-off-by: Kshitij Lakhani * Default stripe_size to None instead of 0. Modify stripe_size check for <=0 instead of ==0 Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove incorrectly added file Signed-off-by: Kshitij Lakhani * Explicitly pass zero sized arrays for seg ids and pos in the CP + AG + Striped primitive rather than using the seqlens or the offsets as placeholders Signed-off-by: Kshitij Lakhani * Fix linting errors Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add a deep dive doc for CP+THD+AG+Stripe>1+SWA regarding design considerations and decisions Signed-off-by: Kshitij Janardan Lakhani * Put docs and pngs into it's separate dir Signed-off-by: Kshitij Janardan Lakhani * Replace png screenshots with markdown coe blocks for the attention patterns. Remove unecessary pngs Signed-off-by: Kshitij Janardan Lakhani * Add doc file to index.rst. Fix grammatical errors Signed-off-by: Kshitij Janardan Lakhani --------- Signed-off-by: Kshitij Lakhani Signed-off-by: Kshitij Janardan Lakhani Signed-off-by: Kshitij Janardan Lakhani Signed-off-by: Kshitij Janardan Lakhani Co-authored-by: Kshitij Janardan Lakhani Co-authored-by: Kshitij Janardan Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../cp_ag_thd_dpa_jax_deep_dive.ipynb | 255 +++++++ docs/index.rst | 1 + tests/jax/test_distributed_fused_attn.py | 127 +++- tests/jax/test_fused_attn.py | 18 +- transformer_engine/jax/attention.py | 57 +- .../jax/cpp_extensions/attention.py | 675 +++++++++++++++++- transformer_engine/jax/cpp_extensions/base.py | 6 + 7 files changed, 1089 insertions(+), 50 deletions(-) create mode 100644 docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb diff --git a/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb b/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb new file mode 100644 index 0000000000..56bc3b13cf --- /dev/null +++ b/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb @@ -0,0 +1,255 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "14efeb1e", + "metadata": {}, + "source": [ + "## Deep Dive into CP + THD + AG + Striped>1 + SWA support for Transformer Engine JAX\n", + "This feature was merged as part of [PR 2379](https://github.com/NVIDIA/TransformerEngine/pull/2379/) and was made available in Transformer Engine v2.11. This document addresses 3 fundamental questions about the design considerations and the implementation logic for this feature." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "16f738c7", + "metadata": { + "vscode": { + "languageId": "plaintext" + } + }, + "outputs": [], + "source": [] + }, + { + "cell_type": "markdown", + "id": "2f31119f", + "metadata": {}, + "source": [ + "### Question 1: Why choose Striped>1 ?\n", + "\n", + "Prior to the addition of this feature, Transformer Engine JAX attention already supported load balancing via a striping pattern, i.e., `stripe_size=1` for `CP + THD + P2P(Ring) + Striped + SWA`. However, this reordering technique does not lend itself well to an all-gathered (post-AG) pattern. The following example illustrates this distinction. For this example, `cp_size=4`, `num_segments=4`, `window_size=(8,0)`, and the pattern is for a single rank after striped reordering has been performed: \n", + "\n", + "#### I. Striped (`stripe_size=1`)\n", + "- Such a staggered pattern is not supported by cuDNN\n", + "- One possible way to express this with cuDNN support is by treating each `q` token as a segment, thereby producing 16 segments with varying `kv` token counts. However, this is very inefficient and does not scale well as max_seqlens increases\n", + "\n", + "```\n", + "1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 1 1 1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - 1 1 1 1 1 1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 4 4 4 - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 4 4 4 4 4 4 4 - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 4 4 4 4 4 4 4 4 4 - - -\n", + "```\n", + "
\n", + "
Figure 1: Post load balancing using stripe_size=1 and post AG attention pattern for a single cp rank
\n", + "
\n", + "\n", + "\n", + "#### II. Striped > 1 (`stripe_size > 1`)\n", + "- This pattern is supported by cuDNN, with a suggested `stripe_size=128`\n", + "- The mask type supported by `CP + THD + AG + Striped>1 + SWA` is `PADDING_CAUSAL_MASK`; however, to express the pattern below, each rank executes THD + SWA using `PADDING_BOTTOM_RIGHT_CAUSAL_MASK`\n", + "- `max_num_segments_for_rank` needs to be estimated. The estimation formula used is: `max_seqlens // (stripe_size * cp_size) + max_num_segments`\n", + "\n", + "```\n", + "1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 4 - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - 4 4 - - - - - - - - - - - -\n", + "```\n", + "
\n", + "
Figure 2: Post load balancing using stripe_size=4 and post AG attention pattern for a single cp rank
\n", + "
" + ] + }, + { + "cell_type": "markdown", + "id": "6eddfa7a", + "metadata": {}, + "source": [ + "### Question 2: Why is there a need for separate helper functions for calculating seqlens and offsets ?\n", + "\n", + "The seqlens and offsets are calculated by the fused attn JAX primitives (both, CP and non-CP) so that they can be passed down to `fused_attn_arbitrary_seqlen_fwd_impl()` / `fused_attn_arbitrary_seqlen_bwd_impl()`, where it is translated before passing down to the cuDNN FE layer. The current (Transformer Engine v2.10) calculation of seqlens and offsets entails the CP primitive passing the sharded segment_ids, segment_pos, seq_lens, seq_offsets stuffed in a SequenceDescriptor object (a convenience class provided for packing these 4 tensors) to the `FusedAttnPrimitive`, which in turn calls `get_seqlens_and_offsets()` on the SequenceDescriptor object. \n", + "\n", + "If `get_seqlens_and_offsets()` receives a SequenceDescriptor object with seq_lens and seq_offsets populated and, segment_ids, segment_pos with size=0, it returns the seq_lens and seq_ofsets as it is (for e.g. `CP + BSHD + AG`). However, if `get_seqlens_and_offsets()` receives a SequenceDescriptor object with segment_ids and segment_pos populated and, seq_lens, seq_offsets with size=0, it first constructs a mask using the segment_ids and segment_pos and then extracts the seq_lens and seq_offsets from it and then returns it (for e.g. `CP + THD + P2P`).\n", + "\n", + "The problem with the current approach of calculating a mask followed by extracting the seq_lens and seq_offsets is that it is unable to express the patterns seen in `CP + THD + AG`. Below is one such example: \n", + "\n", + "```\n", + "1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 3 - - - - - - - - - - - -\n", + "```\n", + "
\n", + "
Figure 3: Example 1 for problem using mask path in get_seqlens_and_offsets() for attention pattern (post striping and AG) .
\n", + "
\n", + "\n", + "Here, ideally, the two sections of the segment 3 should be split into two different segments (segment 3_1 formed using rows 9-12 and segment 3_2 formed using rows 13-16) as cuDNN does not support segment 3's entire staggered shape (as discussed earlier) , however, the mask route is unable to make this distinction, and it ends up treating it as one large segment thereby performing unnecessary computations of the padded regions in segment 3(rows 9-12 )\n", + "\n", + "In the below example, the mask route takes the `kv_seqlens` for segment 1 to be 6 and masks it using Bottom Right Causal Mask rather than taking `kv_seqlens` of 4 and masks it using Bottom Right Causal Mask, resulting in incorrect results\n", + "\n", + "```\n", + "1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "```\n", + "
\n", + "
Figure 4: Example 2 for problem using mask path in get_seqlens_and_offsets() for attention pattern (post striping and AG)
\n", + "
\n", + "\n", + "The second case can be resolved in the mask path, but that would require adding CP specific details to the non-CP FusedAttn primitive which would contaminate it. Besides, resolving the first case would be even trickier with this approach. Due to it being incompatible with the design of FusedAttn primitive and inadequate to express the pattern needed for `CP + THD + AG` fully, separate helper functions were created which calculate the seqlens and seqoffsets, without creating a mask, hence also being O(N) space." + ] + }, + { + "cell_type": "markdown", + "id": "3cc4a12c", + "metadata": {}, + "source": [ + "### Question 3: What is the implementation logic for the separate helper functions ?\n", + "\n", + "This section discusses the implementation logic for two of these four helper functions which serve as a reference, as the other two are using similar principles. Consider the test example in the code block, for which, `cp_size=4`, `stripe_size=4`, `max_seqlens=64`, `num_segments=2` and no SWA for simplicity. seg_1 has 8 valid tokens + 13 padded tokens and seg_2 has 31 valid tokens + 1 padded token. The 0 is used to explicitly show the padded region of seg_1 which is reordered, but for computation purposes it is equivalent to any of the `-` marked elements.\n", + "\n", + "```\n", + "segment_ids_q_0_reordered = segment_ids_kv_0_reordered = jnp.array([[1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2]])\n", + "\n", + "segment_pos_q_0_reordered = segment_pos_kv_0_reordered = jnp.array([[0, 1, 2, 3, 16, 17, 18, 19, 11, 12, 13, 14, 27, 28, 29, 30]])\n", + "\n", + "segment_ids_kv_0_seed12_ag_inv_reordered = jnp.array([[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])\n", + "\n", + "segment_pos_kv_0_seed12_ag_inv_reordered= jnp.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])\n", + "```\n", + "\n", + "```\n", + "1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "1 1 1 1 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - - -\n", + "- - - - - - - - - - - - - - - - - - - - - 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 - - - - - - - - - - - -\n", + "```\n", + "
\n", + "
Figure 5: An example of post striped reordering and AG attention pattern on a single rank.
\n", + "
\n", + "\n", + "#### I. Implementation logic for q_seqlens_for_striped_for_rank()\n", + "**What is the objective/logic ?**\n", + "- Create a new set of segment ids for this rank such that:\n", + " - It gets rid of padding information as it does not contribute to the seqlens calculation\n", + " - It has the ability to identify ”new segments” being created from the same original segment\n", + "- Use this new set of segment ids to calculate the seqlens\n", + "\n", + "**Example walkthrough**\n", + "1. Calculate the non-zero indices (where seg ids !=0)\n", + "2. Calculate the valid seg ids and valid seg pos (i.e. index into seg ids and seg pos using the non-zero indices)\n", + " - `valid_segment_ids=[[1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0]]`\n", + " - `valid_segment_pos=[[0, 1, 2, 3, 11, 12, 13, 14, 27, 28, 29, 30, 0, 0, 0, 0]]`\n", + " - Ignore the 0s at the end of the two arrays as they are just for padding to a static length\n", + "3. Find locations where a q segment change/break happens. A segment change happens when: \n", + " - there is a change in valid_segment_ids OR \n", + " - `valid_segment_pos[i+1] != valid_segment_pos[i]`\n", + " - `segment_changes=[[True, False, False, False, True, False, False, False, True, False, False, False, True, True, True, True]]`\n", + "4. Perform a cumulative sum on the segment changes: \n", + " - `new_segment_ids=[[1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 5, 6, 7]]`\n", + "5. Filter out the valid indices only and pad at the end with 0s upto static length (these are our “new” segment indices without padding)\n", + " - `new_segment_ids_filtered=[[1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0, 0, 0]]`\n", + " - Notice here that the large chunk of 8 q token rows (rows 9-16 in Fig 5) gets broken down into 2 \"new\" segments of 4 q token rows each,\n", + " which is a pattern that cuDNN supports and it ensures that wasted computation for padded regions of rows 9-12 is not performed, which was the\n", + " case in Fig 3\n", + "6. Perform a bin count and pad with -1s upto `max_num_segments_per_seq_for_rank`\n", + " - `seqlens_with_neg1_padding[[ 4, 4, 4, -1, -1, -1, -1]]`\n", + "\n", + "\n", + "#### II. Implementation logic for kv_seqoffsets_for_striped_for_rank()\n", + "**What is the objective/logic ?**\n", + "- Get the original segment ids for those locations where segment changes happen (arr1)\n", + " - Each segment has a known kv offset, hence if we know which original segment id a \"new\" segment is associated with we can find it's kv offset\n", + " - So, for e.g., in Fig 5, all valid tokens of seg_3 have the same kv offset, so even if this gets split into a 2 \"new\" segments, we can procure the offset for both using a mapping of original seg-ids to kv offset \n", + "- Get the segment ids for those locations where segment changes happen in the AG tensor (arr2)\n", + " - This is used to create a kind of mapping between original seg-ids to kv offset\n", + "- Pick values from arr2 mapping for the \"new\" segment ids collected in arr1\n", + "\n", + "**Example walkthrough**\n", + "1. Find locations where a kv segment pos change/break happens and mask out zero seg ids. A segment change happens when: \n", + " - `kv_segment_pos[i+1] != kv_segment_pos[i]`\n", + " - `segment_changes_masked=[[ True, False, False, False, False, False, False, False, True, False, False, False, True, False, False, False]]`\n", + "2. Get the indices where the segment changes happen and the segment ids associated with them:\n", + " - `segment_changes_indices=[[0, 8, 12, -1, -1, -1, -1, -1, -1]]`\n", + " - `[[1, 2, 2, -1, -1, -1, -1, -1, -1]]`\n", + "3. Find the segment pos changes/break for the AG seg pos and mask out zero seg ids\n", + " - `segment_changes_masked_ag=[[True, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, True, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False, False]]`\n", + "4. Get indices where the segment changes happen for the AG seg pos (this works as a mapping between segment ids and kv offsets)\n", + " - `segment_changes_ag_indices=[[0, 21, -1, -1, -1, -1, -1, -1, -1]]`\n", + "5. Get the seq offsets by indexing into segment_changes_ag_indices using segment_changes_indices :\n", + " - `kv_seq_offsets[[0, 21, 21, -1, -1, -1, -1, -1, -1]]`\n", + "\n", + "The implementation details for `q_seqoffsets_for_striped_for_rank()` and `kv_seqlens_for_striped_for_rank()` can be found in [PR 2379](https://github.com/NVIDIA/TransformerEngine/pull/2379/)" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/index.rst b/docs/index.rst index 4fd55d241c..37d21c2a5d 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -56,3 +56,4 @@ Transformer Engine documentation api/c/index debug examples/attention/attention.ipynb + examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 5372018ae8..6b7e04f124 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -327,9 +327,9 @@ def test_cross_attn( ] DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES = [ - # Sequence lengths will be scaled by CP so that we don't run with tiny sizes. - pytest.param([2, 128, 8, 128], id="2-128xCP-8-128"), - pytest.param([4, 256, 16, 64], id="4-256xCP-16-64"), + # Sequence lengths will be scaled by CP*2 so that we don't run with tiny sizes. + pytest.param([2, 128, 8, 128], id="2-128xCPx2-8-128"), + pytest.param([4, 256, 16, 64], id="4-256xCPx2-16-64"), ] @@ -351,12 +351,14 @@ def impl_test_context_parallel_attn( use_shardy, use_scan_ring=False, window_size=None, + stripe_size=None, + num_segments_per_seq=None, ): if qkv_layout.is_thd(): - if cp_strategy == CPStrategy.ALL_GATHER: - pytest.skip("THD doesn't support all gather context parallelism.") - if not load_balanced and cp_strategy == CPStrategy.RING: - pytest.skip("THD + ring doesn't support unbalanced context parallelism.") + if not load_balanced and ( + cp_strategy == CPStrategy.RING or cp_strategy == CPStrategy.ALL_GATHER + ): + pytest.skip(f"THD + {cp_strategy=} doesn't support unbalanced context parallelism.") assert not use_scan_ring or cp_strategy == CPStrategy.RING @@ -382,7 +384,6 @@ def impl_test_context_parallel_attn( data_shape = batch, seqlen, num_head, hidden num_kv_heads = num_head // kv_groups - runner = FusedAttnRunner( batch, seqlen, @@ -401,6 +402,8 @@ def impl_test_context_parallel_attn( bias_shape, window_size, SeqDescFormat.SegmentIDs, + stripe_size=stripe_size, + num_segments_per_seq=num_segments_per_seq, number_of_devices=device_count, mesh_shape=mesh_shape, mesh_axes=mesh_axes, @@ -453,7 +456,7 @@ def check_has_backend_for_mask(mask_type): "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), ) - @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) + @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES) @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) @pytest.mark.parametrize( "qkv_layout, attn_mask_type", @@ -470,6 +473,8 @@ def test_context_parallel_allgather_attn_shardy( dtype, qkv_layout, ): + if qkv_layout.is_thd(): + pytest.skip("Only BSHD layout is supported for CP + AG + Dual chunk attention") kv_groups = 8 self.impl_test_context_parallel_attn( device_count, @@ -486,6 +491,72 @@ def test_context_parallel_allgather_attn_shardy( use_shardy=True, ) + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) + @pytest.mark.parametrize("kv_groups", [1, 8]) + @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type", + DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS, + ) + @pytest.mark.parametrize( + "load_balanced", + [pytest.param(True, id="BALANCED")], + ) + @pytest.mark.parametrize( + "stripe_size", + [pytest.param(64, id="STRIPE-64"), pytest.param(128, id="STRIPE-128")], + ) + @pytest.mark.parametrize( + "window_size", + [ + pytest.param((-1, -1), id="window_size(-1, -1)"), + pytest.param((5, 0), id="window_size(8, 0)"), + ], + ) + @pytest.mark.parametrize( + "num_segments_per_seq", + [pytest.param(5, id="SEG-5")], + ) + def test_context_parallel_allgather_striped_attn( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + kv_groups, + attn_mask_type, + dtype, + qkv_layout, + load_balanced, + window_size, + stripe_size, + num_segments_per_seq, + ): + if not qkv_layout.is_thd(): + pytest.skip("Only THD layout is supported for CP + AG + Striped attention") + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + kv_groups, + attn_mask_type, + dtype, + qkv_layout, + load_balanced, + CPStrategy.ALL_GATHER, + use_shardy=False, + window_size=window_size, + stripe_size=stripe_size, + num_segments_per_seq=num_segments_per_seq, + ) + @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), @@ -514,6 +585,8 @@ def test_context_parallel_allgather_attn( qkv_layout, load_balanced, ): + if qkv_layout.is_thd(): + pytest.skip("Only BSHD layout is supported for CP + AG + Dual chunk attention") self.impl_test_context_parallel_attn( device_count, mesh_shape, @@ -577,6 +650,8 @@ def test_context_parallel_ring_attn( "When context parallelism and sliding window attention are used, " "scanloop is not supported" ) + # Set the stripe size to 1 (ring attention only support stripe_size=1) + stripe_size = 1 if qkv_layout.is_thd() else None self.impl_test_context_parallel_attn( device_count, mesh_shape, @@ -592,6 +667,7 @@ def test_context_parallel_ring_attn( use_shardy=False, use_scan_ring=use_scan, window_size=window_size, + stripe_size=stripe_size, ) @pytest_parametrize_wrapper( @@ -616,6 +692,8 @@ def test_context_parallel_ring_attn_shardy( qkv_layout, ): kv_groups = 8 + # Set the stripe size to 1 (ring attention only support stripe_size=1) + stripe_size = 1 if qkv_layout.is_thd() else None self.impl_test_context_parallel_attn( device_count, mesh_shape, @@ -630,6 +708,7 @@ def test_context_parallel_ring_attn_shardy( cp_strategy=CPStrategy.RING, use_shardy=False, use_scan_ring=True, + stripe_size=stripe_size, ) @@ -639,31 +718,39 @@ def test_context_parallel_ring_attn_shardy( "L2": [[4, 32, 12, 32], [1, 16, 1, 1]], } +REORDER_STRATEGY = [ + pytest.param(ReorderStrategy.DualChunkSwap, None, id="DualChunkSwap"), + pytest.param(ReorderStrategy.Striped, 1, id="Striped-1"), + pytest.param(ReorderStrategy.Striped, 4, id="Striped-4"), +] + class TestReorderCausalLoadBalancing: @pytest.mark.parametrize("cp_size", [2, 4, 8]) @pytest_parametrize_wrapper("shape", REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES) - @pytest.mark.parametrize("qkv_format", [QKVFormat.BSHD, QKVFormat.SBHD]) + @pytest.mark.parametrize("qkv_format", [QKVFormat.BSHD, QKVFormat.SBHD, QKVFormat.THD]) @pytest.mark.parametrize( - "reorder_strategy", - [ - pytest.param(ReorderStrategy.DualChunkSwap, id="DualChunkSwap"), - pytest.param(ReorderStrategy.Striped, id="Striped"), - ], + "reorder_strategy, stripe_size", + REORDER_STRATEGY, ) - def test(self, cp_size, shape, qkv_format, reorder_strategy): + def test(self, cp_size, shape, qkv_format, reorder_strategy, stripe_size): tensor = random.normal(random.PRNGKey(1124), shape, dtype=jnp.bfloat16) seq_dim = 1 if qkv_format == QKVFormat.SBHD: tensor = tensor.swapaxes(0, 1) seq_dim = 0 + if reorder_strategy == ReorderStrategy.Striped: + seq_lens = shape[seq_dim] + if seq_lens < (cp_size * stripe_size): + pytest.skip(f"{seq_lens=} must be larger than {cp_size*stripe_size=}") + ref = tensor.copy() - reorder = jax.jit(reorder_causal_load_balancing, static_argnums=[1, 2, 3]) - inverse = jax.jit(inverse_reorder_causal_load_balancing, static_argnums=[1, 2, 3]) + reorder = jax.jit(reorder_causal_load_balancing, static_argnums=[1, 2, 3, 4]) + inverse = jax.jit(inverse_reorder_causal_load_balancing, static_argnums=[1, 2, 3, 4]) - reordered = reorder(tensor, reorder_strategy, cp_size, seq_dim) - inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim) + reordered = reorder(tensor, reorder_strategy, cp_size, seq_dim, stripe_size) + inversed = inverse(reordered, reorder_strategy, cp_size, seq_dim, stripe_size) assert jnp.array_equal(inversed, ref) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index f4caaef165..49372fda1d 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -352,6 +352,8 @@ class FusedAttnRunner: bias_shape: BiasShape window_size: Tuple[int, int] seq_desc_format: SeqDescFormat + stripe_size: int | None = None + num_segments_per_seq: int | None = None # Specifies sharding resources for distributed tests number_of_devices: int = 1 @@ -366,6 +368,14 @@ class FusedAttnRunner: # dictionary of expected collective comm bytes coll_count_ref: Optional[Dict[str, int]] = None + def __post_init__(self): + # Reset defaults for num_segments_per_seq if not explicitly passed + if self.num_segments_per_seq is None: + if self.qkv_layout.is_thd(): + self.num_segments_per_seq = 2 + else: + self.num_segments_per_seq = 1 + # See https://docs.nvidia.com/deeplearning/cudnn/latest/release-notes.html#cudnn-9-4-0 for known issue # generating zero-length ragged tensors. This setting adjusts the test to avoid the zero-length cases. def _get_max_segments_per_sequence(self): @@ -577,7 +587,6 @@ def generate_random_segment_ids( return segment_ids, segment_pos, segment_pad if self.qkv_layout.is_thd(): - self.num_segments_per_seq = 2 self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_random_segment_ids( self.batch_size, self.max_seqlen_q, self.num_segments_per_seq, seed=42 ) @@ -603,7 +612,6 @@ def generate_random_segment_ids( ) self.seqlens_kv, self.offsets_kv = get_seqlens_and_offsets(self.segment_ids_kv) else: - self.num_segments_per_seq = 1 self.segment_ids_q, self.pad_q = gen_valid( self.batch_size, self.max_seqlen_q, pad_ratio ) @@ -635,12 +643,14 @@ def generate_random_segment_ids( strategy=reorder_strategy, cp_size=self.cp_size, seq_dim=seq_dim, + stripe_size=self.stripe_size, ) self.cp_inverse_reorder_fn = partial( inverse_reorder_causal_load_balancing, strategy=reorder_strategy, cp_size=self.cp_size, seq_dim=seq_dim, + stripe_size=self.stripe_size, ) else: # no-ops for non cp or non load balanced @@ -771,7 +781,7 @@ def to_dp_shardings(x): def test_forward(self): """ - Test forward without JIT + Test forward with JITted primitive and unJITted reference """ self._setup_inputs() @@ -801,6 +811,7 @@ def test_forward(self): "window_size": self.window_size, "context_parallel_strategy": self.cp_strategy, "context_parallel_causal_load_balanced": self.cp_load_balanced, + "stripe_size": self.stripe_size, } customcall_fused_dpa_jit = jit( @@ -896,6 +907,7 @@ def grad_func(func, *args, cp_reverse_out=False, **kwargs): "window_size": self.window_size, "context_parallel_strategy": self.cp_strategy, "context_parallel_causal_load_balanced": self.cp_load_balanced, + "stripe_size": self.stripe_size, } # We can compute dBias only for the [1, h, s, s] layout diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 0a32be9679..21680dc805 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -386,23 +386,57 @@ def _obtain_batch_and_max_seqlen(qkv, qkv_layout): return batch, q_max_seqlen, kv_max_seqlen -def reorder_causal_load_balancing(tensor, strategy: ReorderStrategy, cp_size: int, seq_dim: int): +def reorder_causal_load_balancing( + tensor, strategy: ReorderStrategy, cp_size: int, seq_dim: int, stripe_size: int | None = None +): """Reorders a tensor for load balancing the compute of causal attention.""" if strategy == ReorderStrategy.DualChunkSwap: + if stripe_size is not None: + raise ValueError( + f"Incorrect value for CP dual chunk reordering {stripe_size=}. stripe_size must be" + " None" + ) return tex.attention.reorder_causal_dual_chunk_swap(tensor, cp_size, seq_dim, False) if strategy == ReorderStrategy.Striped: - return tex.attention.reorder_causal_striped(tensor, cp_size, seq_dim, False) + # stripe_size > 1 is only supported for CP+THD+AG+Striped>1+SWA + # stripe_size = 128 is recommended for CP+THD+AG+Striped>1+SWA + if stripe_size is not None and stripe_size <= 0: + raise ValueError( + f"Incorrect value for CP striped reordering {stripe_size=}. stripe_size must be a" + " positive integer" + ) + # Supporting old API defaults of stripe_size=1 + effective_stripe_size = 1 if stripe_size is None else stripe_size + return tex.attention.reorder_causal_striped( + tensor, cp_size, seq_dim, False, effective_stripe_size + ) raise ValueError(f"Unsupported {strategy=}") def inverse_reorder_causal_load_balancing( - tensor, strategy: ReorderStrategy, cp_size: int, seq_dim: int + tensor, strategy: ReorderStrategy, cp_size: int, seq_dim: int, stripe_size: int | None = None ): """Inverse operation of `reorder_causal_load_balancing`.""" if strategy == ReorderStrategy.DualChunkSwap: + if stripe_size is not None: + raise ValueError( + f"Incorrect value for CP dual chunk reordering {stripe_size=}. stripe_size must be" + " None" + ) return tex.attention.reorder_causal_dual_chunk_swap(tensor, cp_size, seq_dim, True) if strategy == ReorderStrategy.Striped: - return tex.attention.reorder_causal_striped(tensor, cp_size, seq_dim, True) + # stripe_size > 1 is only supported for CP+THD+AG+Striped>1+SWA + # stripe_size = 128 is recommended for CP+THD+AG+Striped>1+SWA + if stripe_size is not None and stripe_size <= 0: + raise ValueError( + f"Incorrect value for CP reordering {stripe_size=}. stripe_size must be a positive" + " integer" + ) + # Supporting old API defaults of stripe_size=1 + effective_stripe_size = 1 if stripe_size is None else stripe_size + return tex.attention.reorder_causal_striped( + tensor, cp_size, seq_dim, True, effective_stripe_size + ) raise ValueError(f"Unsupported {strategy=}") @@ -988,7 +1022,7 @@ def fused_attn_thd( return output -@partial(jax.custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17)) +@partial(jax.custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18)) def _fused_attn( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], @@ -1008,6 +1042,7 @@ def _fused_attn( context_parallel_causal_load_balanced: bool, context_parallel_axis: str, context_checkpoint_name: str = "context", + stripe_size: int | None = None, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1028,6 +1063,7 @@ def _fused_attn( context_parallel_causal_load_balanced, context_parallel_axis, context_checkpoint_name=context_checkpoint_name, + stripe_size=stripe_size, ) return output @@ -1051,6 +1087,7 @@ def _fused_attn_fwd_rule( context_parallel_causal_load_balanced, context_parallel_axis, context_checkpoint_name, + stripe_size, ): output, softmax_aux, rng_state = tex.fused_attn_fwd( qkv, @@ -1070,6 +1107,7 @@ def _fused_attn_fwd_rule( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + stripe_size=stripe_size, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) @@ -1099,6 +1137,7 @@ def _fused_attn_bwd_rule( context_parallel_causal_load_balanced, context_parallel_axis, context_checkpoint_name, + stripe_size, ctx, dz, ): @@ -1133,6 +1172,7 @@ def _fused_attn_bwd_rule( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + stripe_size=stripe_size, ) if attn_bias_type == AttnBiasType.NO_BIAS: grad_bias = None @@ -1169,6 +1209,7 @@ def fused_attn( context_parallel_axis: str = "", context_checkpoint_name: str = "context", softmax_offset: Optional[jnp.ndarray] = None, + stripe_size: int | None = None, ): """ Perform cuDNN fused attention. @@ -1206,6 +1247,11 @@ def fused_attn( softmax_offset (Optional[jnp.ndarray]): An optional learnable softmax offset tensor with shape [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX. If provided, this parameter will receive gradients during backpropagation. + stripe_size (int | None): + Indicates the striping size to be used when using ReorderStrategy.Striped. + Currently, a stripe_size > 1 is only supported for CP + THD + Striped + AG, whereas a stripe_size=1 + is supported for both, CP + THD + Striped + AG and CP + THD + Striped + P2P(Ring) + None indicates no striping strategy Returns: (jnp.ndarray): The output tensor from the fused attention. @@ -1283,5 +1329,6 @@ def fused_attn( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, context_checkpoint_name=context_checkpoint_name, + stripe_size=stripe_size, ) return output diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index f0778bfd29..7e0070dd43 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -73,6 +73,7 @@ "context_parallel_load_balanced", "cp_axis", "cp_striped_window_size", + "stripe_size", ], ) @dataclass(frozen=True) @@ -92,7 +93,10 @@ class _FusedAttnConfig: window_size: Tuple[int, int] context_parallel_load_balanced: bool cp_axis: str - cp_striped_window_size: Tuple[int, int] # Only for CP + Ring + THD + SWA + cp_striped_window_size: Tuple[int, int] # Only for CP + Ring P2P + THD + SWA + stripe_size: ( + int | None + ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. @dataclass(frozen=True) @@ -527,7 +531,6 @@ def impl( segment_ids=(_q_segment_ids, _kv_segment_ids), segment_pos=(_q_segment_pos, _kv_segment_pos), ) - (q_seqlen, kv_seqlen), (q_seq_offsets, k_seq_offsets) = ( sequence_descriptor.get_seqlens_and_offsets( config.attn_mask_type, @@ -536,7 +539,6 @@ def impl( config.max_segments_per_seq, ) ) - if config.qkv_layout.is_thd(): def _fix_len_take(x, condition, fill_value=-1): @@ -1234,31 +1236,38 @@ def reorder_causal_dual_chunk_swap(tensor, cp_size: int, seq_dim: int, to_contig return combined.reshape(ori_tensor_shape) -def reorder_causal_striped(tensor, cp_size: int, seq_dim: int, is_inverse: bool): +def reorder_causal_striped( + tensor, cp_size: int, seq_dim: int, is_inverse: bool, stripe_size: int = 1 +): """Reorders a tensor for load balancing with striped pattern""" origin_shape = tensor.shape - if origin_shape[seq_dim] % cp_size != 0: + if stripe_size <= 0: + raise ValueError( + f"Incorrect value for CP reordering {stripe_size=}. stripe_size must be a positive" + " integer" + ) + if origin_shape[seq_dim] % (cp_size * stripe_size) != 0: raise ValueError( - "Expected origin_shape[seq_dim] is multiple of cp_size but got" - f" {origin_shape[seq_dim]=} and {cp_size=}" + "Expected origin_shape[seq_dim] is multiple of cp_size*stripe_size but got" + f" {origin_shape[seq_dim]=}, {cp_size=}, {stripe_size=}, {cp_size*stripe_size=}" ) if not is_inverse: new_shape = [ *origin_shape[:seq_dim], - *[origin_shape[seq_dim] // cp_size, cp_size], + *[origin_shape[seq_dim] // (cp_size * stripe_size), cp_size, stripe_size], *origin_shape[seq_dim + 1 :], ] else: new_shape = [ *origin_shape[:seq_dim], - *[cp_size, origin_shape[seq_dim] // cp_size], + *[cp_size, origin_shape[seq_dim] // (cp_size * stripe_size), stripe_size], *origin_shape[seq_dim + 1 :], ] - chunked_tensor = tensor.reshape(new_shape) - reordered_chunked_tensor = jnp.swapaxes(chunked_tensor, seq_dim, seq_dim + 1) - return reordered_chunked_tensor.reshape(origin_shape) + striped_tensor = tensor.reshape(new_shape) + reordered_striped_tensor = jnp.swapaxes(striped_tensor, seq_dim, seq_dim + 1) + return reordered_striped_tensor.reshape(origin_shape) @dataclass(frozen=True) @@ -1272,26 +1281,47 @@ def check_supported(self): """Checks if the context parallel implementation is supported by the given arguments.""" header = "Context parallel fused attention" - allowed_layouts = [QKVLayout.BSHD_BS2HD, QKVLayout.BSHD_BSHD_BSHD] + allowed_layouts = [ + QKVLayout.BSHD_BS2HD, + QKVLayout.BSHD_BSHD_BSHD, + QKVLayout.THD_T2HD, + QKVLayout.THD_THD_THD, + ] if self.config.qkv_layout not in allowed_layouts: raise ValueError( f"{header} only supports layouts:" f" {','.join(map(str, allowed_layouts))} got: {self.config.qkv_layout}" ) + if (not self.config.qkv_layout.is_thd() and self.config.stripe_size is not None) or ( + self.config.qkv_layout.is_thd() and self.config.stripe_size is None + ): + raise ValueError( + f"{header} only supports Dual Chunk load balancing with BSHD layouts and Striped" + " load balancing with THD layouts" + ) + if self.config.attn_bias_type != AttnBiasType.NO_BIAS: raise ValueError(f"{header} does not support bias got: {self.config.attn_bias_type}") allowed_masks = [AttnMaskType.NO_MASK, AttnMaskType.CAUSAL_MASK] + if self.config.qkv_layout.is_thd(): + allowed_masks.append(AttnMaskType.PADDING_CAUSAL_MASK) if self.config.attn_mask_type not in allowed_masks: raise ValueError( f"{header} only supports masking types: " f" {','.join(map(str, allowed_masks))} got: {self.config.attn_mask_type}" ) + # Do not allow CP + AG + THD + Striped with NO_MASK + if ( + self.config.attn_mask_type is not AttnMaskType.PADDING_CAUSAL_MASK + and self.config.qkv_layout.is_thd() + ): + raise ValueError(f"{header} only supports PADDING_CAUSAL_MASK for THD types") - if self.config.max_segments_per_seq != 1: + if self.config.max_segments_per_seq != 1 and (not self.config.qkv_layout.is_thd): raise ValueError( - f"{header} only supports max_segments_per_seq == 1 got:" + f"{header} only supports max_segments_per_seq == 1 for BSHD layouts, got:" f" {self.config.max_segments_per_seq}" ) @@ -1305,10 +1335,25 @@ def check_supported(self): def get_adjusted_mask(self): """Converts the mask for context parallelism.""" - if self.config.attn_mask_type == AttnMaskType.CAUSAL_MASK: + if ( + self.config.attn_mask_type == AttnMaskType.CAUSAL_MASK + and not self.config.qkv_layout.is_thd() + ): # BSHD AG case only return AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK + if ( + self.config.attn_mask_type == AttnMaskType.PADDING_CAUSAL_MASK + and self.config.qkv_layout.is_thd() + ): # THD AG case only + return AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK return self.config.attn_mask_type + def get_adjusted_max_segments_per_seq(self, max_seqlen, cp_size): + """Converts the max segments per seq for context parallelism AG + THD.""" + # Estimating adjusted max segments per seq + return ( + max_seqlen // (self.config.stripe_size * cp_size) + ) + self.config.max_segments_per_seq + def get_step_config(self) -> _FusedAttnConfig: """Returns a _FusedAttnConfig for single CP step call to fused attention.""" return _FusedAttnConfig( @@ -1324,10 +1369,29 @@ def get_step_config(self) -> _FusedAttnConfig: context_parallel_load_balanced=self.config.context_parallel_load_balanced, cp_axis=self.config.cp_axis, cp_striped_window_size=None, + stripe_size=self.config.stripe_size, + ) + + def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: + """Returns a _FusedAttnConfig for single CP step call (made via a striped AG primitive) to fused attention.""" + return _FusedAttnConfig( + attn_bias_type=self.config.attn_bias_type, + attn_mask_type=self.get_adjusted_mask(), + softmax_type=self.config.softmax_type, + qkv_layout=self.config.qkv_layout, + scaling_factor=self.config.scaling_factor, + dropout_probability=self.config.dropout_probability, + is_training=self.config.is_training, + max_segments_per_seq=self.get_adjusted_max_segments_per_seq(max_seqlen, cp_size), + window_size=self.config.window_size, + context_parallel_load_balanced=self.config.context_parallel_load_balanced, + cp_axis=self.config.cp_axis, + cp_striped_window_size=None, + stripe_size=self.config.stripe_size, ) def all_gather_kv(self, k, v): - """Performs a all-gather of k and v over context parallel ranks.""" + """Performs an all-gather of k and v over context parallel ranks.""" def ag(x): x = lax_paral_op( @@ -1335,7 +1399,10 @@ def ag(x): ) if self.config.context_parallel_load_balanced: cp_size = get_mesh_axis_size(self.config.cp_axis, self.mesh) - x = reorder_causal_dual_chunk_swap(x, cp_size, 1, to_contiguous=True) + if self.config.qkv_layout.is_thd(): + x = reorder_causal_striped(x, cp_size, 1, True, self.config.stripe_size) + else: + x = reorder_causal_dual_chunk_swap(x, cp_size, 1, to_contiguous=True) return x if self.config.qkv_layout.is_kvpacked(): @@ -1345,13 +1412,36 @@ def ag(x): return k, v # fall through + def all_gather_segment_ids_and_pos(self, kv_segment_ids, kv_segment_pos): + """Performs an all-gather of kv segment ids and kv segment pos over context parallel ranks.""" + kv_segment_ids = lax_paral_op( + kv_segment_ids, lax.all_gather, self.config.cp_axis, mesh=self.mesh, axis=1, tiled=True + ) + kv_segment_pos = lax_paral_op( + kv_segment_pos, lax.all_gather, self.config.cp_axis, mesh=self.mesh, axis=1, tiled=True + ) + if self.config.context_parallel_load_balanced: + cp_size = get_mesh_axis_size(self.config.cp_axis, self.mesh) + if self.config.qkv_layout.is_thd(): + kv_segment_ids_ag = reorder_causal_striped( + kv_segment_ids, cp_size, 1, True, self.config.stripe_size + ) + kv_segment_pos_ag = reorder_causal_striped( + kv_segment_pos, cp_size, 1, True, self.config.stripe_size + ) + return kv_segment_ids_ag, kv_segment_pos_ag + return kv_segment_ids, kv_segment_pos # fall through + def reduce_scatter_dkv(self, dk, dv): """Performs a reduce-scatter of dk and dv over context parallel ranks.""" def rs(x): if self.config.context_parallel_load_balanced: cp_size = get_mesh_axis_size(self.config.cp_axis, self.mesh) - x = reorder_causal_dual_chunk_swap(x, cp_size, 1, to_contiguous=False) + if self.config.qkv_layout.is_thd(): + x = reorder_causal_striped(x, cp_size, 1, False, self.config.stripe_size) + else: + x = reorder_causal_dual_chunk_swap(x, cp_size, 1, to_contiguous=False) return lax_paral_op( x, @@ -1424,6 +1514,227 @@ def pad(x, npad): return dk, dv # fall through + # Below are the sharded post AG q seg ids and pos for a given rank: + # q_segment_ids = [[1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2]] + # q_segment_pos = [[0, 1, 2, 3, 16, 17, 18, 19, 11, 12, 13, 14, 27, 28, 29, 30]] + # max_segments_per_seq = 7 + # Below are some intermediate representations: + # non_zero_indices = [[ 0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, -1, -1]] + # segment_changes = [[ True, False, False, False, True, False, False, False, True, False, False, False, True, True, True, True]] + # seqlens_pre = [[1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 0, 0, 0, 0]] + # seqlens_all_pad_neg = [[ 4, 4, 4, -1, -1, -1, -1]] + def q_seqlens_for_striped_for_rank(self, q_segment_ids, q_segment_pos, max_segments_per_seq): + """Extract the q seqlens for striped primitive (post AG) from the sharded q seg ids and seg pos""" + # Create mask for non-zero seg ids and get the non-zero indices associated with the same + non_zero_mask = q_segment_ids != 0 + max_size = q_segment_ids.shape[-1] + non_zero_indices = jax.vmap( + lambda mask_row: jnp.where(mask_row, size=max_size, fill_value=-1)[0] + )(non_zero_mask) + + # Pick non-zero seg ids and seg pos using take_along_axis to index within the seg ids and pos + # Clip -1 to 0 for safe indexing + clipped_indices = jnp.clip(non_zero_indices, 0, None) + valid_segment_ids = jnp.where( + non_zero_indices >= 0, jnp.take_along_axis(q_segment_ids, clipped_indices, axis=-1), 0 + ) + valid_segment_pos = jnp.where( + non_zero_indices >= 0, jnp.take_along_axis(q_segment_pos, clipped_indices, axis=-1), 0 + ) + # Create a mask for actual valid entries (not padding) + actual_valid = valid_segment_ids != 0 + # First element is True only if it's actually valid + first_is_segment = actual_valid[..., 0:1] + + # Detect segment breaks in the valid tokens only (not full seq) + # Padding will always be true as the segment change condition is being applied + # on the valid segments (which have padding at the end so they'll always trigger True) + segment_changes = jnp.concatenate( + [ + first_is_segment, # First valid element starts a segment + (valid_segment_ids[..., 1:] != valid_segment_ids[..., :-1]) + | (valid_segment_pos[..., 1:] != valid_segment_pos[..., :-1] + 1), + ], + axis=-1, + ) + new_segment_ids = jnp.cumsum(segment_changes, axis=-1) + seqlens_pre = jax.vmap( + lambda av_row, nsi_row: jnp.where(av_row, nsi_row, 0).astype(jnp.int32) + )(actual_valid, new_segment_ids) + seqlens_all = jax.vmap( + lambda sp_row: jnp.bincount(sp_row, length=max_segments_per_seq + 1)[1:] + )(seqlens_pre) + seqlens_all_pad_neg = jnp.where(seqlens_all == 0, -1, seqlens_all) + return seqlens_all_pad_neg + + # Below are the sharded post AG q seg ids and pos for a given rank: + # q_segment_ids = [[1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2]] + # q_segment_pos = [[0, 1, 2, 3, 16, 17, 18, 19, 11, 12, 13, 14, 27, 28, 29, 30]] + # max_segments_per_seq = 7 + # Below are some intermediate representations: + # segment_changes = [[ True, False, False, False, True, False, False, False, True, False, False, False, True, False, False, False]] + # segment_changes_masked = [[ True, False, False, False, False, False, False, False, True, False, False, False, True, False, False, False]] + # seq_offsets = [[ 0, 8, 12, -1, -1, -1, -1, -1]] + def q_seqoffsets_for_striped_for_rank(self, q_segment_ids, q_segment_pos, max_segments_per_seq): + """Extract the q seqoffets for striped primitive (post AG) from the sharded q seg ids and seg pos""" + segment_changes = jnp.concatenate( + [ + jnp.full( + (q_segment_pos.shape[0], 1), True, dtype=bool + ), # First valid element starts a segment + (q_segment_pos[..., 1:] != q_segment_pos[..., :-1] + 1), # Segment pos changed + ], + axis=-1, + ) + # Remove any padded region segment changes + segment_changes_masked = jnp.where(q_segment_ids != 0, segment_changes, False) + # Get the indices for segment changes (these are the offsets) + seq_offsets = jax.vmap( + lambda scm_row: jnp.where(scm_row, size=max_segments_per_seq, fill_value=-1)[0] + )(segment_changes_masked) + return seq_offsets + + # Below are the sharded post AG q seg ids and pos for a given rank: + # kv_segment_ids = [[1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2]] + # kv_segment_pos = [[0, 1, 2, 3, 16, 17, 18, 19, 11, 12, 13, 14, 27, 28, 29, 30]] + # max_segments_per_seq = 7 + # Below are some intermediate representations: + # non_zero_mask = [[ True, True, True, True, False, False, False, False, True, True, True, True, True, True, True, True]] + # non_zero_indices = [[ 0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15, -1, -1, -1, -1]] + # segment_changes = [[False, False, False, True, False, False, False, True, False, False, False, True, True, True, True, False]] + # selected_values = [[ 4, 15, 31, -1, -1, -1, -1, -1]] + def kv_seqlens_for_striped_for_rank(self, kv_segment_ids, kv_segment_pos, max_segments_per_seq): + """Extract the kv seqlens for striped primitive (post AG) from the sharded kv seg ids and seg pos""" + # Create mask for non-zero seg ids and get the non-zero indices associated with the same + non_zero_mask = kv_segment_ids != 0 + max_size = kv_segment_ids.shape[-1] + non_zero_indices = jax.vmap( + lambda mask_row: jnp.where(mask_row, size=max_size, fill_value=-1)[0] + )(non_zero_mask) + + # Pick non zero seg ids and seg pos using take_along_axis + # Clip -1 to 0 for safe indexing + clipped_indices = jnp.clip(non_zero_indices, 0, None) + valid_segment_ids = jnp.where( + non_zero_indices >= 0, jnp.take_along_axis(kv_segment_ids, clipped_indices, axis=-1), 0 + ) + valid_segment_pos = jnp.where( + non_zero_indices >= 0, jnp.take_along_axis(kv_segment_pos, clipped_indices, axis=-1), 0 + ) + actual_valid = valid_segment_ids != 0 + + # Detect segment breaks (only for non-zero segments) + segment_changes = jnp.concatenate( + [ + ( + (valid_segment_ids[..., 1:] != valid_segment_ids[..., :-1]) + & actual_valid[..., 1:] + ) + | (valid_segment_pos[..., 1:] != valid_segment_pos[..., :-1] + 1), + actual_valid[..., -1:], + ], + axis=-1, + ) + # Get the indices for segment changes + segment_changes_valid = jax.vmap( + lambda sc_row, av_row: jnp.where( + sc_row & av_row, size=max_segments_per_seq, fill_value=-1 + )[0] + )(segment_changes, actual_valid) + safe_indices = jnp.maximum(segment_changes_valid, 0) + # Select values using take_along_axis per row + selected_values = jnp.where( + segment_changes_valid >= 0, + jnp.take_along_axis(valid_segment_pos, safe_indices, axis=-1) + 1, + -1, + ) + return selected_values + + # Below are the sharded post AG q seg ids and pos for a given rank: + # kv_segment_ids = [[1, 1, 1, 1, 0, 0, 0, 0, 2, 2, 2, 2, 2, 2, 2, 2]] + # kv_segment_pos = [[0, 1, 2, 3, 16, 17, 18, 19, 11, 12, 13, 14, 27, 28, 29, 30]] + # kv_segment_ids_ag = [[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + # 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + # 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] + # kv_segment_pos_ag = [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, + # 18, 19, 20, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, + # 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, + # 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]] + # max_segments_per_seq = 7 + # Below are some intermediate representations: + # segment_changes_first_true_masked = [[ True, False, False, False, False, False, False, False, True, + # False, False, False, True, False, False, False]] + # segment_changes_indices = [[ 0, 8, 12, -1, -1, -1, -1, -1, -1]] + # segment_ids = [[ 1, 2, 2, -1, -1, -1, -1, -1, -1]] + # segment_changes_ag_first_true_masked = [[ True, False, False, False, False, False, False, False, False, + # False, False, False, False, False, False, False, False, False, + # False, False, False, True, False, False, False, False, False, + # False, False, False, False, False, False, False, False, False, + # False, False, False, False, False, False, False, False, False, + # False, False, False, False, False, False, False, False, False, + # False, False, False, False, False, False, False, False, False, + # False] + # segment_changes_ag_indices = [[ 0, 21, -1, -1, -1, -1, -1, -1, -1]] + # seq_offsets = [[ 0, 21, 21, -1, -1, -1, -1, -1, -1]] + def kv_seqoffsets_for_striped_for_rank( + self, + kv_segment_pos, + kv_segment_ids, + kv_segment_pos_ag, + kv_segment_ids_ag, + max_segments_per_seq, + ): + """Extract the kv seqoffsets for striped primitive (post AG) from the sharded kv seg ids and seg pos, + AG kv seg ids and seg pos.""" + # Calculate the segment pos change mask + segment_changes_first_true = jnp.concatenate( + [ + jnp.full( + (kv_segment_pos.shape[0], 1), True, dtype=bool + ), # Assume valid element starts a segment and mask afterwards + (kv_segment_pos[..., 1:] != kv_segment_pos[..., :-1] + 1), # Segment pos changed + ], + axis=-1, + ) + segment_changes_first_true_masked = jnp.where( + kv_segment_ids != 0, segment_changes_first_true, False + ) + + # Get segment change indices for rank + segment_changes_indices = jax.vmap( + lambda sc_row: jnp.where(sc_row, size=max_segments_per_seq, fill_value=-1)[0] + )(segment_changes_first_true_masked) + # Get segment ids associated with the segment_changes_indices for rank + segment_ids = jax.vmap( + lambda sci_row, ksi_row: jnp.where(sci_row >= 0, ksi_row[sci_row], -1) + )(segment_changes_indices, kv_segment_ids) + + # Get segment change indices for AG + segment_changes_ag_first_true = jnp.concatenate( + [ + jnp.full( + (kv_segment_pos.shape[0], 1), True, dtype=bool + ), # Assume valid element starts a segment and mask afterwards + ( + kv_segment_pos_ag[..., 1:] != kv_segment_pos_ag[..., :-1] + 1 + ), # Segment pos changed + ], + axis=-1, + ) + segment_changes_ag_first_true_masked = jnp.where( + kv_segment_ids_ag != 0, segment_changes_ag_first_true, False + ) + # Get segment change indices for AG + segment_changes_ag_indices = jax.vmap( + lambda scag_row: jnp.where(scag_row, size=max_segments_per_seq, fill_value=-1)[0] + )(segment_changes_ag_first_true_masked) + + # Use the segment ids picked per rank to get the offsets from the AG indices + seq_offsets = jax.vmap( + lambda si_row, sca_row: jnp.where(si_row > 0, sca_row[si_row - 1], -1) + )(segment_ids, segment_changes_ag_indices) + return seq_offsets + class FusedAttnCPWithAllGatherFwdPrimitive(FusedAttnFwdPrimitive): """ @@ -1501,7 +1812,6 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): q_seqlen_for_step = q_seqlen / (cp_size * 2) num_kv_chunks = kv_max_seqlen // kv_seqlens_for_rank[sub_idx] kv_seqlen_for_step = (kv_seqlen / (cp_size * 2)) * num_kv_chunks - output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl( q_split[sub_idx], k_unmasked, @@ -1722,6 +2032,314 @@ def _cross_attn_bwd( register_primitive(FusedAttnCPWithAllGatherBwdPrimitive) +class FusedAttnCPStripedWithAllGatherFwdPrimitive(FusedAttnFwdPrimitive): + """ + Fused Attention Forward with Context Parallelism and Striped Load Balancing Primitive + + This context parallel implementation uses all-gather to collect KV inputs from context parallel ranks. + """ + + @staticmethod + def partition(config, mesh, arg_infos, result_infos): + # Call base implementation for non-context parallel mesh to avoid unecessary work. + is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1 + if not is_context_parallel: + return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos) + + helper = _FusedAttnCPWithAllGatherHelper(mesh, config) + helper.check_supported() + + out_sharding = result_infos[0].sharding + softmax_aux_sharding = result_infos[1].sharding + rng_state_sharding = seed_sharding = NamedSharding( + mesh, PartitionSpec(get_all_mesh_axes(), None) + ) + arg_shardings = [arg_i.sharding for arg_i in arg_infos] + arg_shardings[5] = seed_sharding + arg_shardings = tuple(arg_shardings) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + + def impl( + q, + k, + v, + bias, + softmax_offset, + seed, + q_seqlen, + kv_seqlen, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + ): # pylint: disable=unused-argument + cp_size = get_mesh_axis_size(config.cp_axis, mesh) + cp_rank = get_mesh_axis_rank(config.cp_axis, mesh) + + # cuDNN does not support right-aligned masking with dynamic sequence length padding. + # Therefore we must explicitly instantiate each CP rank slicing and use a runtime switch + # to select the appropriate computation. Each case generates a [..., SEQ/CP, ..] tensor + # meeting the expectation of the SPMD model. + # TODO(mgoldfarb-nvidia): When cuDNN supports we should be able to make use of a padding + # mask/sequence length tensor to avoid this unrolled loop. + + # Each rank receives the ag k and v along with the ag kv seg ids and kv seg offsets + # Each rank sees the sharded view for 5 tensors -> q, _q_segment_ids, _q_segment_pos, + # _kv_segment_ids, _kv_segment_pos -> Note these have also been reordered before passing in. + def _cross_attn( + q, k, v, bias, softmax_offset, kv_segment_ids_ag, kv_segment_pos_ag, seed + ): + # Helper generates the seqlens and offsets for q and kv and then pass them down to the FusedAttnFwdPrimitive + # Unset the segment_ids and segment_pos by passing placeholders so that the seqlens_from_segment_ids_pos() + # does not go down that route but instead just picks the pre-computed seqlens and offsets passed onto it + + kv_max_seqlen = k.shape[1] + # Estimate an adjusted max_segments_per_seq per rank based on the global max_segments_per_seq + adjusted_max_segments_per_seq = helper.get_adjusted_max_segments_per_seq( + max_seqlen=kv_max_seqlen, cp_size=cp_size + ) + q_seqlens_for_rank = helper.q_seqlens_for_striped_for_rank( + _q_segment_ids, _q_segment_pos, adjusted_max_segments_per_seq + ) + q_seq_offsets_for_rank = helper.q_seqoffsets_for_striped_for_rank( + q_segment_ids=_q_segment_ids, + q_segment_pos=_q_segment_pos, + max_segments_per_seq=adjusted_max_segments_per_seq, + ) + kv_seqlens_for_rank = helper.kv_seqlens_for_striped_for_rank( + kv_segment_ids=_kv_segment_ids, + kv_segment_pos=_kv_segment_pos, + max_segments_per_seq=adjusted_max_segments_per_seq, + ) + kv_seq_offsets_for_rank = helper.kv_seqoffsets_for_striped_for_rank( + kv_segment_pos=_kv_segment_pos, + kv_segment_ids=_kv_segment_ids, + kv_segment_pos_ag=kv_segment_pos_ag, + kv_segment_ids_ag=kv_segment_ids_ag, + max_segments_per_seq=adjusted_max_segments_per_seq, + ) + + output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl( + q, # sharded for rank + k, # ag + v, # ag + bias, + softmax_offset, + seed, + q_seqlens_for_rank, + kv_seqlens_for_rank, + q_seq_offsets_for_rank, + kv_seq_offsets_for_rank, + jnp.zeros(0), + jnp.zeros(0), + jnp.zeros(0), + jnp.zeros(0), + config=helper.get_step_config_for_striped( + max_seqlen=kv_max_seqlen, cp_size=cp_size + ), + ) + return output, softmax_aux, rng_state + + # AG the k, v, kv_segment_ids and kv_segment_pos + k_ag, v_ag = helper.all_gather_kv(k, v) + _kv_segment_ids_ag, _kv_segment_pos_ag = helper.all_gather_segment_ids_and_pos( + _kv_segment_ids, _kv_segment_pos + ) + functions = [ + partial( + _cross_attn, + q, + k_ag, + v_ag, + bias, + softmax_offset, + _kv_segment_ids_ag, + _kv_segment_pos_ag, + seed, + ) + for _ in range(cp_size) + ] + return lax.switch(cp_rank, functions) + + return mesh, impl, out_shardings, arg_shardings + + +register_primitive(FusedAttnCPStripedWithAllGatherFwdPrimitive) + + +class FusedAttnCPStripedWithAllGatherBwdPrimitive(FusedAttnBwdPrimitive): + """ + Fused Attention Backward with Context Parallelism and Striped Load Balancing Primitive. + + This context parallel implementation uses all-gather to collect KV and dKV inputs from context parallel ranks. + The gradients are subsequently reduce-scattered back to each context parallel rank. + """ + + @staticmethod + def partition(config, mesh, arg_infos, result_infos): + # Call base implementation for non-context parallel mesh to avoid unecessary work. + is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1 + if not is_context_parallel: + return FusedAttnBwdPrimitive.partition(config, mesh, arg_infos, result_infos) + + # Ensure we can support this configuration with context parallelism. + helper = _FusedAttnCPWithAllGatherHelper(mesh, config) + helper.check_supported() + + del result_infos + q_spec = get_padded_spec(arg_infos[0]) + k_spec = get_padded_spec(arg_infos[1]) + v_spec = get_padded_spec(arg_infos[2]) + bias_spec = get_padded_spec(arg_infos[3]) + softmax_offset_spec = get_padded_spec(arg_infos[4]) + dq_sharding = NamedSharding(mesh, PartitionSpec(*q_spec)) + dk_sharding = NamedSharding(mesh, PartitionSpec(*k_spec)) + dv_sharding = NamedSharding(mesh, PartitionSpec(*v_spec)) + dbias_sharding = NamedSharding(mesh, PartitionSpec(*bias_spec)) + dsoftmax_offset_sharding = NamedSharding(mesh, PartitionSpec(*softmax_offset_spec)) + arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + out_shardings = ( + dq_sharding, + dk_sharding, + dv_sharding, + dbias_sharding, + dsoftmax_offset_sharding, + ) + + def impl( + q, + k, + v, + bias, + softmax_offset, + softmax_aux, + rng_state, + output, + doutput, + q_seqlen, + kv_seqlen, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + ): # pylint: disable=unused-argument + cp_size = get_mesh_axis_size(config.cp_axis, mesh) + cp_rank = get_mesh_axis_rank(config.cp_axis, mesh) + + # See comment in FusedAttnCPFwdPrimitive.partition for why we define this function. + def _cross_attn_bwd( + q, + k, + v, + bias, + softmax_offset, + softmax_aux, + rng_state, + output, + doutput, + _q_segment_ids, + kv_segment_ids_ag, + _q_segment_pos, + kv_segment_pos_ag, + ): + # Helper generates the seqlens and offsets for q and kv and then pass them down to the FusedAttnFwdPrimitive + # Unset the segment_ids and segment_pos by passing placeholders so that the seqlens_from_segment_ids_pos() + # does not go down that route but instead just picks the pre-computed seqlens and offsets passed onto it + + kv_max_seqlen = k.shape[1] + # Estimate an adjusted max_segments_per_seq per rank based on the global max_segments_per_seq + adjusted_max_segments_per_seq = helper.get_adjusted_max_segments_per_seq( + max_seqlen=kv_max_seqlen, cp_size=cp_size + ) + q_seqlens_for_rank = helper.q_seqlens_for_striped_for_rank( + _q_segment_ids, _q_segment_pos, adjusted_max_segments_per_seq + ) + q_seq_offsets_for_rank = helper.q_seqoffsets_for_striped_for_rank( + q_segment_ids=_q_segment_ids, + q_segment_pos=_q_segment_pos, + max_segments_per_seq=adjusted_max_segments_per_seq, + ) + kv_seqlens_for_rank = helper.kv_seqlens_for_striped_for_rank( + kv_segment_ids=_kv_segment_ids, + kv_segment_pos=_kv_segment_pos, + max_segments_per_seq=adjusted_max_segments_per_seq, + ) + kv_seq_offsets_for_rank = helper.kv_seqoffsets_for_striped_for_rank( + kv_segment_pos=_kv_segment_pos, + kv_segment_ids=_kv_segment_ids, + kv_segment_pos_ag=kv_segment_pos_ag, + kv_segment_ids_ag=kv_segment_ids_ag, + max_segments_per_seq=adjusted_max_segments_per_seq, + ) + + dq_local, dk_local, dv_local, dbias_local, _ = FusedAttnBwdPrimitive.impl( + q, # sharded for rank + k, # ag + v, # ag + bias, + softmax_offset, + softmax_aux, + rng_state, + output, + doutput, + q_seqlens_for_rank, + kv_seqlens_for_rank, + q_seq_offsets_for_rank, + kv_seq_offsets_for_rank, + jnp.zeros(0), + jnp.zeros(0), + jnp.zeros(0), + jnp.zeros(0), + config=helper.get_step_config_for_striped( + max_seqlen=kv_max_seqlen, cp_size=cp_size + ), + ) + return dq_local, dk_local, dv_local, dbias_local + + # AG the k, v, kv_segment_ids and kv_segment_pos + k_ag, v_ag = helper.all_gather_kv(k, v) + _kv_segment_ids_ag, _kv_segment_pos_ag = helper.all_gather_segment_ids_and_pos( + _kv_segment_ids, _kv_segment_pos + ) + + functions = [ + partial( + _cross_attn_bwd, + q, + k_ag, + v_ag, + bias, + softmax_offset, + softmax_aux, + rng_state, + output, + doutput, + _q_segment_ids, + _kv_segment_ids_ag, + _q_segment_pos, + _kv_segment_pos_ag, + ) + for _ in range(cp_size) + ] + + dq, dk_local, dv_local, dbias = lax.switch(cp_rank, functions) + # RS the dk and dv + dk, dv = helper.reduce_scatter_dkv(dk_local, dv_local) + + # Return dummy dsoftmax_offset for arity matching (all-gather CP doesn't use it) + dummy_dsoftmax_offset = jnp.empty_like(softmax_offset) + return dq, dk, dv, dbias, dummy_dsoftmax_offset + + return mesh, impl, out_shardings, arg_shardings + + +register_primitive(FusedAttnCPStripedWithAllGatherBwdPrimitive) + + @dataclass(frozen=True) class _FusedAttnCPWithP2PHelper: """Helper class to assist with running the P2P ring strategy for CP attention.""" @@ -1811,6 +2429,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: context_parallel_load_balanced=self.config.context_parallel_load_balanced, cp_axis=self.config.cp_axis, cp_striped_window_size=None, + stripe_size=self.config.stripe_size, ) def stack_kv(self, k, v): @@ -2693,6 +3312,7 @@ def fused_attn_fwd( context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT, context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", + stripe_size: int | None = None, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -2731,6 +3351,7 @@ def fused_attn_fwd( context_parallel_causal_load_balanced (bool): Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. + stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing Returns: (jnp.ndarray): The output tensor from the fused attention. """ @@ -2796,12 +3417,16 @@ def fused_attn_fwd( context_parallel_load_balanced=context_parallel_causal_load_balanced, cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, + stripe_size=stripe_size, ) primitive = None match context_parallel_strategy: case CPStrategy.DEFAULT | CPStrategy.ALL_GATHER: - primitive = FusedAttnCPWithAllGatherFwdPrimitive.outer_primitive + if qkv_layout.is_thd(): + primitive = FusedAttnCPStripedWithAllGatherFwdPrimitive.outer_primitive + else: + primitive = FusedAttnCPWithAllGatherFwdPrimitive.outer_primitive case CPStrategy.RING: # We must use stripe attention for THD-RING if qkv_layout.is_thd(): @@ -2843,6 +3468,7 @@ def fused_attn_bwd( context_parallel_strategy: CPStrategy = CPStrategy.DEFAULT, context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", + stripe_size: int | None = None, ): """ Perform the backward pass of the cuDNN fused attention implementations. @@ -2882,6 +3508,7 @@ def fused_attn_bwd( context_parallel_causal_load_balanced (bool): Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. + stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing Returns: Tuple[jnp.ndarray, ...], jnp.ndarray: - The first tuple contains the gradients with respect to the input `qkv` tensors in the @@ -2954,12 +3581,16 @@ def fused_attn_bwd( context_parallel_load_balanced=context_parallel_causal_load_balanced, cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, + stripe_size=stripe_size, ) primitive = None match context_parallel_strategy: case CPStrategy.DEFAULT | CPStrategy.ALL_GATHER: - primitive = FusedAttnCPWithAllGatherBwdPrimitive.outer_primitive + if qkv_layout.is_thd(): + primitive = FusedAttnCPStripedWithAllGatherBwdPrimitive.outer_primitive + else: + primitive = FusedAttnCPWithAllGatherBwdPrimitive.outer_primitive case CPStrategy.RING: if qkv_layout.is_thd(): primitive = FusedRingAttnStripedBwdPrimitive.outer_primitive diff --git a/transformer_engine/jax/cpp_extensions/base.py b/transformer_engine/jax/cpp_extensions/base.py index 556b587191..61deab5b80 100644 --- a/transformer_engine/jax/cpp_extensions/base.py +++ b/transformer_engine/jax/cpp_extensions/base.py @@ -176,6 +176,9 @@ def shardy_sharding_rule(*args): def register_primitive(cls, outer_only=False): """ Register a JAX primitive and add it to the internal registry. + Inner primitive - single device, no sharding awareness, eager mode fallback + Outer primitive - multi device, sharding aware, partition() distributes work, + used when there's a dev mesh context """ _primitive_registry[cls.__name__] = cls @@ -190,14 +193,17 @@ def name_of_wrapper_p(): inner_p = core.Primitive(cls.name) dispatch.prim_requires_devices_during_lowering.add(inner_p) inner_p.multiple_results = cls.multiple_results + # Define eager execution implementation (by invoking it's MLIR lowering) inner_p.def_impl(partial(xla.apply_primitive, inner_p)) inner_p.def_abstract_eval(cls.abstract) mlir.register_lowering(inner_p, cls.lowering, platform="cuda") cls.inner_primitive = inner_p + # Create the outer primitive for distributed execution outer_p = core.Primitive(name_of_wrapper_p()) dispatch.prim_requires_devices_during_lowering.add(outer_p) outer_p.multiple_results = cls.multiple_results + # Define the eager execution implementation outer_p.def_impl(cls.outer_impl) outer_p.def_abstract_eval(cls.outer_abstract) batching.primitive_batchers[outer_p] = cls.batcher From fd91bae314710a30e62e7c9863d221b779c355dd Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 8 Dec 2025 10:13:41 -0800 Subject: [PATCH 113/521] Changed VERSION to 2.12.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 5b70b33bd8..d5e1cb2914 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.11.0.dev0 +2.12.0.dev0 From c09411d82042faf0a1a5a9d4fd1953877de8f577 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 9 Dec 2025 05:16:45 +0530 Subject: [PATCH 114/521] [Pytorch][Bug]MXFP8 Split tensor Bug fix (#2427) * bug fixed, test added Signed-off-by: Varun Thumbe * fix contigous Signed-off-by: Varun Thumbe * revert unecessary change Signed-off-by: Varun Thumbe * revert another change 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 * Update transformer_engine/pytorch/tensor/mxfp8_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * address review comments Signed-off-by: Varun Thumbe * missed adding renamed file Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix minor issue Signed-off-by: Varun Thumbe * fix ci issue Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the test for bfloat16 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: 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> --- qa/L0_pytorch_unittest/test.sh | 2 +- ...oat8tensor.py => test_quantized_tensor.py} | 99 ++++++++++++++++++- .../pytorch/tensor/float8_tensor.py | 26 +++-- .../pytorch/tensor/mxfp8_tensor.py | 13 ++- 4 files changed, 123 insertions(+), 17 deletions(-) rename tests/pytorch/{test_float8tensor.py => test_quantized_tensor.py} (80%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f1a48e421b..512c01db42 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -32,7 +32,7 @@ PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED 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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_nvfp4.xml $TE_PATH/tests/pytorch/nvfp4 || test_fail "test_nvfp4" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8tensor.xml $TE_PATH/tests/pytorch/test_float8tensor.py || test_fail "test_float8tensor.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_quantized_tensor.xml $TE_PATH/tests/pytorch/test_quantized_tensor.py || test_fail "test_quantized_tensor.py" 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" diff --git a/tests/pytorch/test_float8tensor.py b/tests/pytorch/test_quantized_tensor.py similarity index 80% rename from tests/pytorch/test_float8tensor.py rename to tests/pytorch/test_quantized_tensor.py index b7ddf0e8a6..bdf355c7bc 100644 --- a/tests/pytorch/test_float8tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -13,9 +13,15 @@ import transformer_engine.pytorch as te from transformer_engine.pytorch import ( Float8Quantizer, - Float8Tensor, Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, + Float8Tensor, + MXFP8Tensor, + NVFP4Tensor, ) + from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported import transformer_engine_torch as tex @@ -47,6 +53,12 @@ def _to_list(x: Union[Iterable, Any]) -> List: # Check if FP8 is supported fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + # delayed scaling def to_float8( @@ -452,3 +464,88 @@ def test_quantize_dequantize( # Make sure we are not trivially passing the test with pytest.raises(AssertionError): torch.testing.assert_close(x_fp8_dequantized, -x_hp, **_tols[fp8_dtype]) + + +class TestAllQuantizedTensors: + @staticmethod + def setup_class(cls) -> None: + # Configure RNG + seed = 1234 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + @pytest.mark.parametrize("quantization", ["fp8", "mxfp8", "nvfp4", "fp8_blockwise"]) + @pytest.mark.parametrize("dim", [0, 1]) + def test_chunk( + self, + quantization: str, + dim: int, + shape: Iterable[int] = (128, 128), + chunks: int = 2, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + # Skip invalid configs + if quantization == "fp8" and not fp8_available: + pytest.skip(reason_for_no_fp8) + if quantization == "fp8_blockwise" and not fp8_block_scaling_available: + pytest.skip(reason_for_no_fp8_block_scaling) + if quantization == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if quantization == "nvfp4" and not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + # Create 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 == "mxfp8": + quantizer = MXFP8Quantizer(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 == "nvfp4": + quantizer = NVFP4Quantizer( + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=False, + ) + else: + raise ValueError(f"Unknown quantizer ({quantizer})") + # Create reference and quantized tensor + ref_tensor = torch.randn(shape, device=device, dtype=dtype) + quantized_tensor = quantizer(ref_tensor) + ref_tensor.copy_(quantized_tensor) + + # Chunk tensors + ref_splits = torch.chunk(ref_tensor, chunks, dim=dim) + quantized_splits = torch.chunk(quantized_tensor, chunks, dim=dim) + # Check splits + for ref_split, quantized_split in zip(ref_splits, quantized_splits): + # Check split shapes + assert ref_split.size() == quantized_split.size() + + # Check that splits are quantized when expected + if quantization == "fp8": + assert isinstance(quantized_split, Float8Tensor) + expected_value = quantized_split.dequantize() + elif quantization == "mxfp8" and dim == 0: + assert isinstance(quantized_split, MXFP8Tensor) + expected_value = quantized_split.dequantize() + else: + # Otherwise torch dispatch would default to base implementation + # dequantize and computing output and hence output from torch chunk + # is already dequantized. + expected_value = quantized_split + # Check values + torch.testing.assert_close(expected_value, ref_split) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e3ca110dfa..1077fe818f 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -493,10 +493,10 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: # Convert PyTorch dtype to TE dtype if dtype is None: dtype = self.dtype - + tensor = self.contiguous() if torch.is_grad_enabled(): - return _FromFloat8Func.apply(self, dtype) - return _FromFloat8Func.forward(None, self, dtype) + return _FromFloat8Func.apply(tensor, dtype) + return _FromFloat8Func.forward(None, tensor, dtype) def quantize_( self, @@ -554,13 +554,19 @@ def contiguous( Returns `self` if data is already in correct memory format. """ - if self._data is not None and self._data.is_contiguous(memory_format=memory_format): - return self - if self._transpose is not None and self._transpose.is_contiguous( - memory_format=memory_format - ): - return self - return Float8Tensor.make_like(tensor=self, data=self._data.contiguous()) + # requires_grad remains unaltered when calling contiguous on + # torch tensor and so should be the case for our custom float8 tensor + # as well. + return Float8Tensor.make_like( + tensor=self, + data=self._data.contiguous(memory_format=memory_format), + data_transpose=( + self._transpose.contiguous(memory_format=memory_format) + if self._transpose is not None + else None + ), + requires_grad=self.requires_grad, + ) # raise ValueError("Float8Tensor does not support different memory formats!") diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index a41079080c..6dcf9ae79a 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -434,13 +434,16 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if scale_inv is not None else None ) + scale_inv_out = list(scale_inv_out) if scale_inv_out is not None else None # Pad scale_inv_out to be a multiple of pad_multiple if scale_inv_out is not None: - current_shape = scale_inv_out.shape - pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple - if pad_dim0 > 0: - scale_inv_out = torch.nn.functional.pad(scale_inv_out, (0, 0, 0, pad_dim0)) - + for idx, split_scale_inv_out in enumerate(scale_inv_out): + current_shape = split_scale_inv_out.shape + pad_dim0 = (pad_multiple - current_shape[0] % pad_multiple) % pad_multiple + if pad_dim0 > 0: + scale_inv_out[idx] = torch.nn.functional.pad( + split_scale_inv_out, (0, 0, 0, pad_dim0) + ) out_data.append(scale_inv_out) return [ MXFP8Tensor( From 8ef3a33de472ec4a9bdde75f05341d3d0671686f Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 9 Dec 2025 12:06:16 +0530 Subject: [PATCH 115/521] Fix runtime lib loading logic (#2297) Fixes to runtime loading logic and add missing deps Signed-off-by: Kirthi Shankar Sivamani --- build_tools/pytorch.py | 2 +- build_tools/utils.py | 10 +- transformer_engine/common/__init__.py | 206 ++++++++++++-------------- 3 files changed, 101 insertions(+), 117 deletions(-) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 3d44d8740c..302816c6fd 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -14,7 +14,7 @@ def install_requirements() -> List[str]: """Install dependencies for TE/PyTorch extensions.""" - return ["torch>=2.1", "einops", "onnxscript", "onnx"] + return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic"] def test_requirements() -> List[str]: diff --git a/build_tools/utils.py b/build_tools/utils.py index 395b41261b..50ba007594 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -241,13 +241,9 @@ def get_cuda_include_dirs() -> Tuple[str, str]: cuda_root = Path(nvidia.__file__).parent return [ - cuda_root / "cuda_nvcc" / "include", - cuda_root / "cublas" / "include", - cuda_root / "cuda_runtime" / "include", - cuda_root / "cudnn" / "include", - cuda_root / "cuda_cccl" / "include", - cuda_root / "nvtx" / "include", - cuda_root / "cuda_nvrtc" / "include", + subdir / "include" + for subdir in cuda_root.iterdir() + if subdir.is_dir() and (subdir / "include").is_dir() ] diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 3ffe1c7b1d..2d7932d5aa 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -235,31 +235,6 @@ def _get_sys_extension() -> str: raise RuntimeError(f"Unsupported operating system ({system})") -@functools.lru_cache(maxsize=None) -def _load_nvidia_cuda_library(lib_name: str): - """ - Attempts to load shared object file installed via pip. - - `lib_name`: Name of package as found in the `nvidia` dir in python environment. - """ - - so_paths = glob.glob( - os.path.join( - sysconfig.get_path("purelib"), - f"nvidia/{lib_name}/lib/lib*{_get_sys_extension()}.*[0-9]", - ) - ) - - path_found = len(so_paths) > 0 - ctypes_handles = [] - - if path_found: - for so_path in so_paths: - ctypes_handles.append(ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)) - - return path_found, ctypes_handles - - @functools.lru_cache(maxsize=None) def _nvidia_cudart_include_dir() -> str: """Returns the include directory for cuda_runtime.h if exists in python environment.""" @@ -279,101 +254,102 @@ def _nvidia_cudart_include_dir() -> str: @functools.lru_cache(maxsize=None) -def _load_cudnn(): - """Load CUDNN shared library.""" +def _load_cuda_library_from_python(lib_name: str, strict: bool = False): + """ + Attempts to load shared object file installed via python packages. - # Attempt to locate cuDNN in CUDNN_HOME or CUDNN_PATH, if either is set - cudnn_home = os.environ.get("CUDNN_HOME") or os.environ.get("CUDNN_PATH") - if cudnn_home: - libs = glob.glob(f"{cudnn_home}/**/libcudnn{_get_sys_extension()}*", recursive=True) - libs.sort(reverse=True, key=os.path.basename) - if libs: - return ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) + `lib_name` : Name of package as found in the `nvidia` dir in python environment. + `strict` : If set to `True`, throw an error if lib is not found. + """ - # Attempt to locate cuDNN in CUDA_HOME, CUDA_PATH or /usr/local/cuda - cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") or "/usr/local/cuda" - libs = glob.glob(f"{cuda_home}/**/libcudnn{_get_sys_extension()}*", recursive=True) - libs.sort(reverse=True, key=os.path.basename) - if libs: - return ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) + ext = _get_sys_extension() + nvidia_dir = os.path.join(sysconfig.get_path("purelib"), "nvidia") - # Attempt to locate cuDNN in Python dist-packages - found, handle = _load_nvidia_cuda_library("cudnn") - if found: - return handle + # PyPI packages provided by nvidia libs exist + # in 4 possible locations inside `nvidia`. + # Check by order of priority. + path_found = False + if os.path.isdir(os.path.join(nvidia_dir, "cu13", lib_name)): + so_paths = glob.glob(os.path.join(nvidia_dir, "cu13", lib_name, f"lib/lib*{ext}.*[0-9]")) + path_found = len(so_paths) > 0 + + if not path_found and os.path.isdir(os.path.join(nvidia_dir, "cu13")): + so_paths = glob.glob(os.path.join(nvidia_dir, "cu13", f"lib/lib{lib_name}*{ext}.*[0-9]")) + path_found = len(so_paths) > 0 + + if not path_found and os.path.isdir(os.path.join(nvidia_dir, lib_name)): + so_paths = glob.glob(os.path.join(nvidia_dir, lib_name, f"lib/lib*{ext}.*[0-9]")) + path_found = len(so_paths) > 0 - # Attempt to locate libcudnn via ldconfig - libs = subprocess.check_output(["ldconfig", "-p"]) - libs = libs.decode("utf-8").split("\n") - sos = [] - for lib in libs: - if "libcudnn" in lib and "=>" in lib: - sos.append(lib.split(">")[1].strip()) - if sos: - return ctypes.CDLL(sos[0], mode=ctypes.RTLD_GLOBAL) + if not path_found: + so_paths = glob.glob(os.path.join(nvidia_dir, f"cuda_{lib_name}", f"lib/lib*{ext}.*[0-9]")) + path_found = len(so_paths) > 0 - # If all else fails, assume that it is in LD_LIBRARY_PATH and error out otherwise - return ctypes.CDLL(f"libcudnn{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) + ctypes_handles = [] + + if path_found: + for so_path in so_paths: + ctypes_handles.append(ctypes.CDLL(so_path, mode=ctypes.RTLD_GLOBAL)) + + if strict and not path_found: + raise RuntimeError(f"{lib_name} shared object not found.") + + return path_found, ctypes_handles @functools.lru_cache(maxsize=None) -def _load_nvrtc(): - """Load NVRTC shared library.""" - # Attempt to locate NVRTC in CUDA_HOME, CUDA_PATH or /usr/local/cuda - cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") or "/usr/local/cuda" - libs = glob.glob(f"{cuda_home}/**/libnvrtc{_get_sys_extension()}*", recursive=True) - libs = list(filter(lambda x: not ("stub" in x or "libnvrtc-builtins" in x), libs)) - libs.sort(reverse=True, key=os.path.basename) - if libs: - return ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) - - # Attempt to locate NVRTC in Python dist-packages - found, handle = _load_nvidia_cuda_library("cuda_nvrtc") - if found: - return handle +def _load_cuda_library_from_system(lib_name: str): + """ + Attempts to load shared object file installed via system/cuda-toolkit. + + `lib_name`: Name of library to load without extension or `lib` prefix. + """ - # Attempt to locate NVRTC via ldconfig - libs = subprocess.check_output(["ldconfig", "-p"]) - libs = libs.decode("utf-8").split("\n") - sos = [] - for lib in libs: - if "libnvrtc" in lib and "=>" in lib: - sos.append(lib.split(">")[1].strip()) - if sos: - return ctypes.CDLL(sos[0], mode=ctypes.RTLD_GLOBAL) + # Where to look for the shared lib in decreasing order of preference. + paths = ( + os.environ.get(f"{lib_name.upper()}_HOME"), + os.environ.get(f"{lib_name.upper()}_PATH"), + os.environ.get("CUDA_HOME"), + os.environ.get("CUDA_PATH"), + "/usr/local/cuda", + ) + + for path in paths: + if path is None: + continue + libs = glob.glob(f"{path}/**/lib{lib_name}{_get_sys_extension()}*", recursive=True) + libs = [lib for lib in libs if "stub" not in lib] + libs.sort(reverse=True, key=os.path.basename) + if libs: + return True, ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) - # If all else fails, assume that it is in LD_LIBRARY_PATH and error out otherwise - return ctypes.CDLL(f"libnvrtc{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) + # Search in LD_LIBRARY_PATH. + try: + _lib_handle = ctypes.CDLL(f"lib{lib_name}{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) + return True, _lib_handle + except OSError: + return False, None @functools.lru_cache(maxsize=None) -def _load_curand(): - """Load cuRAND shared library.""" - # Attempt to locate cuRAND in CUDA_HOME, CUDA_PATH or /usr/local/cuda - cuda_home = os.environ.get("CUDA_HOME") or os.environ.get("CUDA_PATH") or "/usr/local/cuda" - libs = glob.glob(f"{cuda_home}/**/libcurand{_get_sys_extension()}*", recursive=True) - libs = list(filter(lambda x: not ("stub" in x), libs)) - libs.sort(reverse=True, key=os.path.basename) - if libs: - return ctypes.CDLL(libs[0], mode=ctypes.RTLD_GLOBAL) - - # Attempt to locate cuRAND in Python dist-packages - found, handle = _load_nvidia_cuda_library("curand") +def _load_cuda_library(lib_name: str): + """ + Load given shared library. + Prioritize loading from system/toolkit + before checking python packages. + """ + + # Attempt to locate library in system. + found, handle = _load_cuda_library_from_system(lib_name) if found: - return handle + return True, handle - # Attempt to locate cuRAND via ldconfig - libs = subprocess.check_output(["ldconfig", "-p"]) - libs = libs.decode("utf-8").split("\n") - sos = [] - for lib in libs: - if "libcurand" in lib and "=>" in lib: - sos.append(lib.split(">")[1].strip()) - if sos: - return ctypes.CDLL(sos[0], mode=ctypes.RTLD_GLOBAL) + # Attempt to locate library in Python dist-packages. + found, handle = _load_cuda_library_from_python(lib_name) + if found: + return False, handle - # If all else fails, assume that it is in LD_LIBRARY_PATH and error out otherwise - return ctypes.CDLL(f"libcurand{_get_sys_extension()}", mode=ctypes.RTLD_GLOBAL) + raise RuntimeError(f"{lib_name} shared object not found.") @functools.lru_cache(maxsize=None) @@ -384,11 +360,23 @@ def _load_core_library(): if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))): sanity_checks_for_pypi_installation() - _CUDNN_LIB_CTYPES = _load_cudnn() - _NVRTC_LIB_CTYPES = _load_nvrtc() - _CURAND_LIB_CTYPES = _load_curand() - _CUBLAS_LIB_CTYPES = _load_nvidia_cuda_library("cublas") - _CUDART_LIB_CTYPES = _load_nvidia_cuda_library("cuda_runtime") + + # `_load_cuda_library` is used for packages that must be loaded + # during runtime. Both system and pypi packages are searched + # and an error is thrown if not found. + _, _CUDNN_LIB_CTYPES = _load_cuda_library("cudnn") + system_nvrtc, _NVRTC_LIB_CTYPES = _load_cuda_library("nvrtc") + system_curand, _CURAND_LIB_CTYPES = _load_cuda_library("curand") + + # This additional step is necessary to be able to install TE wheels + # and import TE (without any guards) in an environment where the cuda + # toolkit might be absent without being guarded + load_libs_for_no_ctk = not system_nvrtc and not system_curand + if load_libs_for_no_ctk: + _CUBLAS_LIB_CTYPES = _load_cuda_library_from_python("cublas", strict=True) + _CUDART_LIB_CTYPES = _load_cuda_library_from_python("cudart", strict=True) + _CUDNN_ALL_LIB_CTYPES = _load_cuda_library_from_python("cudnn", strict=True) + _TE_LIB_CTYPES = _load_core_library() # Needed to find the correct headers for NVRTC kernels. From e05f87e193b9aa9cc385c958a42e9498ea0c68cd Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 9 Dec 2025 10:14:10 -0800 Subject: [PATCH 116/521] [PyTorch] Change order of args in another permutation triton kernel (#2488) change order Signed-off-by: tdophung --- transformer_engine/common/triton/permutation.py | 12 ++++++------ transformer_engine/pytorch/triton/permutation.py | 8 ++++---- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index e8c43f52d2..87a9c24533 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -402,16 +402,11 @@ def _unpermute_kernel( @triton.jit def _unpermute_bwd_with_merging_probs_kernel( - # pointers + # input pointers fwd_output_grad_ptr, - fwd_input_grad_ptr, fwd_input_ptr, merging_probs_ptr, - merging_probs_grad_ptr, row_id_map_ptr, - # sizes - num_experts: tl.constexpr, - hidden_size: tl.constexpr, # strides stride_row_id_map_token, stride_row_id_map_expert, @@ -425,7 +420,12 @@ def _unpermute_bwd_with_merging_probs_kernel( stride_merging_probs_expert, stride_merging_probs_grad_token, stride_merging_probs_grad_expert, + # output pointers + fwd_input_grad_ptr, + merging_probs_grad_ptr, # metas + num_experts: tl.constexpr, + hidden_size: tl.constexpr, PROBS_LOAD_WIDTH: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 39d6fdaa6a..8f953e9c31 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -304,13 +304,9 @@ def unpermute_with_mask_map_bwd_with_merging_probs( grid = (num_tokens,) _unpermute_bwd_with_merging_probs_kernel[grid]( fwd_output_grad, - act_grad, fwd_input, merging_probs, - merging_probs_grad, row_id_map, - num_experts, - hidden_size, row_id_map.stride(0), row_id_map.stride(1), fwd_output_grad.stride(0), @@ -323,6 +319,10 @@ def unpermute_with_mask_map_bwd_with_merging_probs( merging_probs.stride(1), merging_probs_grad.stride(0), merging_probs_grad.stride(1), + act_grad, + merging_probs_grad, + num_experts, + hidden_size, PROBS_LOAD_WIDTH=triton.next_power_of_2(num_experts), ) return act_grad, merging_probs_grad From dbaa02d08ef4b180f9cd61796049e82a67f636a1 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Tue, 9 Dec 2025 10:40:00 -0800 Subject: [PATCH 117/521] Fix the sm120 compilation with CUDA 12 (#2482) Signed-off-by: Przemek Tredak --- transformer_engine/common/util/ptx.cuh | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 754cbd900a..7f296c9e38 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -867,19 +867,19 @@ __device__ __forceinline__ void fma_f32_bf16(float &out, uint16_t const &a, uint } __device__ __forceinline__ void reduce_sync_max_abs_f32(float &out, float const &in) { -#if ((__CUDA_ARCH_HAS_FEATURE__(SM100_ALL)) || (__CUDA_ARCH_HAS_FEATURE__(SM101_ALL)) || \ - (__CUDA_ARCH_HAS_FEATURE__(SM120_ALL))) - asm volatile("redux.sync.max.abs.f32 %0, %1, 0xFFFFFFFF;" : "=f"(out) : "f"(in)); -#else - asm volatile( - "{\n\t" - ".reg.b32 val;\n" - "abs.f32 val, %1;\n" - "redux.sync.max.u32 %0, val, 0xFFFFFFFF;\n" - "}\n\t" - : "=r"(reinterpret_cast(out)) - : "f"(in)); -#endif + constexpr bool is_sm_100f = NVTE_CUDA_ARCH_MATCHES(ptx::FamilySpecific<100>); + if constexpr (is_sm_100f) { + asm volatile("redux.sync.max.abs.f32 %0, %1, 0xFFFFFFFF;" : "=f"(out) : "f"(in)); + } else { + asm volatile( + "{\n\t" + ".reg.b32 val;\n" + "abs.f32 val, %1;\n" + "redux.sync.max.u32 %0, val, 0xFFFFFFFF;\n" + "}\n\t" + : "=r"(reinterpret_cast(out)) + : "f"(in)); + } } __device__ __forceinline__ bf16 get_amax(bf16 a, bf16 b) { From 46c6ef31bd1216b88cad7d27394ca04c439e8e1a Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 9 Dec 2025 14:37:13 -0800 Subject: [PATCH 118/521] Jax primitives for permutation on single GPU (#2473) * branch off of initial permutation jax-triton PR Signed-off-by: tdophung * Set 0 as the size of dummy tensors to reduce memory usage. Signed-off-by: tdophung * Correct setting of permuted_probs_stride_token, unpermuted_probs_stride_token and unpermuted_probs_stride_expert in unpermutation Signed-off-by: tdophung * Implement primitives, wrapper, test for wrapper, edit trit on binding to accomodate scalars Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change implemementation of VJP functions to match correct pattern. Deduce some static scalar args from shapes of inputs. Accept B, S instead of num_tokens. Change test to use value_and_grad to test vjp funcs properly Signed-off-by: tdophung * formatting Signed-off-by: tdophung * fix pylint Signed-off-by: tdophung * fix test to compare to the correct reference impl. relax 1 tol for grad compare, fix lint the rightway Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test_permutation to use value_and_grad for reference impl, tighten tols, and add unpermute with probs for token combine bwd rule Signed-off-by: tdophung * added forgotten file in prev commit Signed-off-by: tdophung * format Signed-off-by: tdophung * merge with_probs to without_probs Signed-off-by: tdophung * add aserts and fix lint Signed-off-by: tdophung --------- Signed-off-by: tdophung Co-authored-by: Ming Huang Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_permutation.py | 694 ++++++++++ transformer_engine/jax/cpp_extensions/amax.py | 4 +- transformer_engine/jax/permutation.py | 401 ++++++ .../jax/triton_extensions/__init__.py | 4 + .../jax/triton_extensions/permutation.py | 1136 +++++++++++++++++ .../jax/triton_extensions/utils.py | 18 +- 6 files changed, 2250 insertions(+), 7 deletions(-) create mode 100644 tests/jax/test_permutation.py create mode 100644 transformer_engine/jax/permutation.py create mode 100644 transformer_engine/jax/triton_extensions/permutation.py diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py new file mode 100644 index 0000000000..23d9f50609 --- /dev/null +++ b/tests/jax/test_permutation.py @@ -0,0 +1,694 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for permutation Triton kernels and high-level APIs""" + +import jax +import jax.numpy as jnp +import pytest + +# High-level API with VJP support +from transformer_engine.jax.permutation import ( + token_dispatch, + token_combine, + sort_chunks_by_index, +) +from utils import assert_allclose + + +def reference_make_row_id_map( + routing_map: jnp.ndarray, + num_tokens: int, + num_experts: int, +) -> jnp.ndarray: + """ + Reference implementation of make_row_id_map using JAX primitives. + + Parameters + ---------- + routing_map : jnp.ndarray + Input tensor of shape [num_tokens, num_experts]. Mask indicating which experts + are routed to which tokens (1 = routed, 0 = not routed). + num_tokens : int + Number of tokens in the input tensor. + num_experts : int + Number of experts in the input tensor. + + Returns + ------- + row_id_map : jnp.ndarray + The row_id_map for the permutation of shape [num_tokens, num_experts * 2 + 1]. + """ + row_id_map = jnp.full((num_tokens, num_experts * 2 + 1), -1, dtype=jnp.int32) + + # For each expert, compute cumulative sum to get destination indices + cumsum_per_expert = jnp.cumsum(routing_map, axis=0) + + # Compute total tokens per expert + tokens_per_expert = jnp.sum(routing_map, axis=0) + expert_offsets = jnp.concatenate([jnp.array([0]), jnp.cumsum(tokens_per_expert)[:-1]]) + + # Build the row_id_map + for token_idx in range(num_tokens): + routed_experts = jnp.where(routing_map[token_idx] == 1)[0] + n_routed = len(routed_experts) + + # Store number of routed experts in the last position + row_id_map = row_id_map.at[token_idx, -1].set(n_routed) + + # For each routed expert, compute destination row and store it + dest_rows = [] + expert_indices = [] + for expert_idx in routed_experts: + # Destination row = expert offset + (cumsum - 1) + dest_row = expert_offsets[expert_idx] + cumsum_per_expert[token_idx, expert_idx] - 1 + dest_rows.append(dest_row) + expert_indices.append(expert_idx) + + # Sort by destination row + if n_routed > 0: + sort_indices = jnp.argsort(-jnp.array(dest_rows)) # Negative for descending sort + sorted_dest_rows = jnp.array(dest_rows)[sort_indices] + sorted_expert_indices = jnp.array(expert_indices)[sort_indices] + + # Store sorted destination rows and expert indices + for i in range(n_routed): + row_id_map = row_id_map.at[token_idx, i].set(sorted_dest_rows[i]) + row_id_map = row_id_map.at[token_idx, num_experts + i].set(sorted_expert_indices[i]) + + return row_id_map + + +def _reference_permute_impl( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + probs: jnp.ndarray, + num_tokens: int, + num_experts: int, + num_out_tokens: int, + hidden_size: int, +) -> tuple: + """ + Internal helper for reference permutation implementation. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [num_tokens, hidden_size]. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape [num_tokens, num_experts * 2 + 1]. + probs : jnp.ndarray + The probabilities of the input tensor. + num_tokens : int + Number of tokens in the input tensor. + num_experts : int + Number of experts. + num_out_tokens : int + Number of tokens in the permuted tensor. + hidden_size : int + Hidden size of the input tensor. + + Returns + ------- + output : jnp.ndarray + Permuted output tensor of shape [num_out_tokens, hidden_size]. + permuted_probs : jnp.ndarray + Permuted probabilities if probs was provided, None otherwise. + """ + output = jnp.zeros((num_out_tokens, hidden_size), dtype=inp.dtype) + permuted_probs = None if probs is None else jnp.zeros((num_out_tokens,), dtype=probs.dtype) + + for token_idx in range(num_tokens): + n_routed = int(row_id_map[token_idx, -1]) # int() needed for Python range() + for i in range(n_routed): + # Don't use int() here - JAX can index with traced values, + # and int() breaks autodiff gradient tracking + dest_row = row_id_map[token_idx, i] + expert_idx = row_id_map[token_idx, num_experts + i] + + # Get probability for this expert + if probs is not None: + if probs.ndim == 1: + prob = probs[token_idx] + else: + prob = probs[token_idx, expert_idx] + + # Match kernel behavior: if prob == 0.0, zero out the output (padding indicator) + if prob == 0.0: + output = output.at[dest_row].set(0.0) + else: + output = output.at[dest_row].set(inp[token_idx]) + + permuted_probs = permuted_probs.at[dest_row].set(prob) + else: + output = output.at[dest_row].set(inp[token_idx]) + + return output, permuted_probs + + +def _reference_unpermute_impl( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: jnp.ndarray, + permuted_probs: jnp.ndarray, + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> tuple: + """ + Internal helper for reference unpermutation implementation. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [num_out_tokens, hidden_size]. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape [num_tokens, num_experts * 2 + 1]. + merging_probs : jnp.ndarray + The merging probabilities for weighted reduction. + permuted_probs : jnp.ndarray + The permuted probabilities. + num_tokens : int + Number of tokens. + num_experts : int + Number of experts. + hidden_size : int + Hidden size. + + Returns + ------- + output : jnp.ndarray + Unpermuted output tensor of shape [num_tokens, hidden_size]. + unpermuted_probs : jnp.ndarray + Unpermuted probabilities if permuted_probs was provided, None otherwise. + """ + output = jnp.zeros((num_tokens, hidden_size), dtype=inp.dtype) + unpermuted_probs = ( + None + if permuted_probs is None + else jnp.zeros((num_tokens, num_experts), dtype=permuted_probs.dtype) + ) + + for token_idx in range(num_tokens): + n_routed = int(row_id_map[token_idx, -1]) # int() needed for Python range() + for i in range(n_routed): + # Don't use int() here - JAX can index with traced values, + # and int() breaks autodiff gradient tracking + src_row = row_id_map[token_idx, i] + expert_idx = row_id_map[token_idx, num_experts + i] + + if merging_probs is not None: + weight = merging_probs[token_idx, expert_idx] + output = output.at[token_idx].add(inp[src_row] * weight) + else: + output = output.at[token_idx].add(inp[src_row]) + + if permuted_probs is not None: + unpermuted_probs = unpermuted_probs.at[token_idx, expert_idx].set( + permuted_probs[src_row] + ) + + return output, unpermuted_probs + + +def reference_token_dispatch( + inp: jnp.ndarray, + routing_map: jnp.ndarray, + num_out_tokens: int, + probs: jnp.ndarray = None, +) -> tuple: + """ + Reference implementation of token_dispatch using JAX primitives. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [num_tokens, hidden_size]. + routing_map : jnp.ndarray + Routing mask of shape [num_tokens, num_experts]. + num_out_tokens : int + Number of tokens in the permuted tensor. + probs : jnp.ndarray, optional + The probabilities of shape [num_tokens, num_experts]. + + Returns + ------- + output : jnp.ndarray + Permuted output tensor of shape [num_out_tokens, hidden_size]. + permuted_probs : jnp.ndarray or None + Permuted probabilities of shape [num_out_tokens], or None if probs not provided. + row_id_map : jnp.ndarray + The row_id_map for the permutation. + """ + num_tokens, num_experts = routing_map.shape + hidden_size = inp.shape[1] + + row_id_map = reference_make_row_id_map(routing_map, num_tokens, num_experts) + output, permuted_probs = _reference_permute_impl( + inp, row_id_map, probs, num_tokens, num_experts, num_out_tokens, hidden_size + ) + + return output, permuted_probs, row_id_map + + +def reference_token_combine( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: jnp.ndarray, +) -> jnp.ndarray: + """ + Reference implementation of token_combine using JAX primitives. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [num_out_tokens, hidden_size]. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape [num_tokens, num_experts * 2 + 1]. + merging_probs : jnp.ndarray + The merging probabilities for weighted reduction. + + Returns + ------- + output : jnp.ndarray + Unpermuted output tensor of shape [num_tokens, hidden_size]. + """ + num_tokens = row_id_map.shape[0] + num_experts = (row_id_map.shape[1] - 1) // 2 + hidden_size = inp.shape[1] + + output, _ = _reference_unpermute_impl( + inp, row_id_map, merging_probs, None, num_tokens, num_experts, hidden_size + ) + + return output + + +def reference_make_chunk_sort_map( + split_sizes: jnp.ndarray, + sorted_indices: jnp.ndarray, + num_tokens: int, + num_splits: int, +) -> jnp.ndarray: + """ + Reference implementation of make_chunk_sort_map using JAX primitives. + + Parameters + ---------- + split_sizes : jnp.ndarray + The sizes of the chunks of shape [num_splits,]. + sorted_indices : jnp.ndarray + The indices of the sorted chunks of shape [num_splits,]. + num_tokens : int + Number of tokens. + num_splits : int + Number of splits. + + Returns + ------- + row_id_map : jnp.ndarray + Row ID map for chunk sorting of shape [num_tokens,]. + """ + row_id_map = jnp.zeros((num_tokens,), dtype=jnp.int32) + + # Compute cumulative positions + cumsum_sizes = jnp.concatenate([jnp.array([0]), jnp.cumsum(split_sizes)]) + + # For each chunk, compute the destination indices + dest_offset = 0 + for sorted_idx in sorted_indices: + chunk_start = cumsum_sizes[sorted_idx] + chunk_end = cumsum_sizes[sorted_idx + 1] + chunk_size = chunk_end - chunk_start + + # Map source positions to destination positions + for i in range(chunk_size): + row_id_map = row_id_map.at[chunk_start + i].set(dest_offset + i) + + dest_offset += chunk_size + + return row_id_map + + +def reference_sort_chunks_by_map( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + probs: jnp.ndarray, + num_tokens: int, + hidden_size: int, + is_forward: bool, +) -> tuple: + """ + Reference implementation of sort_chunks_by_map using JAX primitives. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [num_tokens, hidden_size]. + row_id_map : jnp.ndarray + The token to destination mapping of shape [num_tokens,]. + probs : jnp.ndarray + The probabilities. + num_tokens : int + Number of tokens. + hidden_size : int + Hidden size. + is_forward : bool + Whether this is forward or backward. + + Returns + ------- + output : jnp.ndarray + Sorted output tensor of shape [num_tokens, hidden_size]. + permuted_probs : jnp.ndarray + Sorted probabilities if probs was provided, None otherwise. + """ + output = jnp.zeros((num_tokens, hidden_size), dtype=inp.dtype) + permuted_probs = None if probs is None else jnp.zeros((num_tokens,), dtype=probs.dtype) + + if is_forward: + # Forward: src -> dest + for src_idx in range(num_tokens): + # Don't use int() - JAX can index with traced values + dest_idx = row_id_map[src_idx] + output = output.at[dest_idx].set(inp[src_idx]) + if probs is not None: + permuted_probs = permuted_probs.at[dest_idx].set(probs[src_idx]) + else: + # Backward: dest -> src + for dest_idx in range(num_tokens): + # Don't use int() - JAX can index with traced values + src_idx = row_id_map[dest_idx] + output = output.at[dest_idx].set(inp[src_idx]) + if probs is not None: + permuted_probs = permuted_probs.at[dest_idx].set(probs[src_idx]) + + return output, permuted_probs + + +class TestHighLevelPermutationAPI: + """Test high-level permutation APIs (token_dispatch, token_combine, etc.) + + These tests compare the high-level APIs against reference implementations + to verify correctness of both forward and backward passes. + """ + + @staticmethod + def generate_routing_map( + num_tokens: int, + num_experts: int, + tokens_per_expert: int = 2, + key: jax.Array = None, + ): + """Generate random routing map for testing""" + if key is None: + key = jax.random.PRNGKey(0) + + routing_map = jnp.zeros((num_tokens, num_experts), dtype=jnp.int32) + for token_idx in range(num_tokens): + key, subkey = jax.random.split(key) + expert_indices = jax.random.choice( + subkey, num_experts, shape=(tokens_per_expert,), replace=False + ) + routing_map = routing_map.at[token_idx, expert_indices].set(1) + + return routing_map + + # ========================================================================= + # token_dispatch tests + # ========================================================================= + + @pytest.mark.parametrize( + "num_tokens,num_experts,hidden_size,tokens_per_expert", + [ + (32, 8, 256, 2), + (64, 16, 512, 3), + ], + ) + @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) + def test_token_dispatch(self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype): + """Test token_dispatch forward and backward pass against reference""" + key = jax.random.PRNGKey(42) + + # Generate routing map + routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) + num_out_tokens = int(jnp.sum(routing_map)) + + # Generate input data + key, inp_key = jax.random.split(key) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + + # Define loss functions + def loss_fn(x): + output, _, _ = token_dispatch(x, routing_map, num_out_tokens) + return jnp.sum(output**2) + + def ref_loss_fn(x): + output, _, _ = reference_token_dispatch(x, routing_map, num_out_tokens) + return jnp.sum(output**2) + + loss_val, computed_grad = jax.value_and_grad(loss_fn)(inp) + ref_loss_val, ref_grad = jax.value_and_grad(ref_loss_fn)(inp) + + # Compare forward outputs + output, _, _ = token_dispatch(inp, routing_map, num_out_tokens) + ref_output, _, _ = reference_token_dispatch(inp, routing_map, num_out_tokens) + assert_allclose(output, ref_output) + + # Compare loss and gradient + assert_allclose(loss_val, ref_loss_val) + assert_allclose(computed_grad, ref_grad) + + # ========================================================================= + # token_dispatch with probs tests + # ========================================================================= + + @pytest.mark.parametrize( + "num_tokens,num_experts,hidden_size,tokens_per_expert", + [ + (32, 8, 256, 2), + (64, 16, 512, 3), + ], + ) + @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) + def test_token_dispatch_with_probs( + self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype + ): + """Test token_dispatch with probs forward and backward pass against reference""" + key = jax.random.PRNGKey(42) + + # Generate routing map + routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) + num_out_tokens = int(jnp.sum(routing_map)) + + # Generate input data and probs + key, inp_key, prob_key = jax.random.split(key, 3) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + probs = jax.random.uniform( + prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.0, maxval=1.0 + ) + + # Define loss function that uses token_dispatch with probs + # We compute gradients w.r.t. both inp and probs + def loss_fn(x, p): + output, permuted_probs, _ = token_dispatch(x, routing_map, num_out_tokens, probs=p) + return jnp.sum(output**2) + jnp.sum(permuted_probs**2) + + def ref_loss_fn(x, p): + output, permuted_probs, _ = reference_token_dispatch( + x, routing_map, num_out_tokens, probs=p + ) + return jnp.sum(output**2) + jnp.sum(permuted_probs**2) + + loss_val, (inp_grad, probs_grad) = jax.value_and_grad(loss_fn, argnums=(0, 1))(inp, probs) + ref_loss_val, (ref_inp_grad, ref_probs_grad) = jax.value_and_grad( + ref_loss_fn, argnums=(0, 1) + )(inp, probs) + + output, permuted_probs, _ = token_dispatch(inp, routing_map, num_out_tokens, probs=probs) + + ref_output, ref_permuted_probs, _ = reference_token_dispatch( + inp, routing_map, num_out_tokens, probs=probs + ) + + # Compare forward outputs + assert_allclose(output, ref_output) + assert_allclose(permuted_probs, ref_permuted_probs) + + # Compare loss and gradients + assert_allclose(loss_val, ref_loss_val) + assert_allclose(inp_grad, ref_inp_grad) + assert_allclose(probs_grad, ref_probs_grad) + + # ========================================================================= + # token_combine tests + # ========================================================================= + + @pytest.mark.parametrize( + "num_tokens,num_experts,hidden_size,tokens_per_expert", + [ + (32, 8, 256, 2), + (64, 16, 512, 3), + ], + ) + @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) + @pytest.mark.parametrize("with_merging_probs", [True, False]) + def test_token_combine( + self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype, with_merging_probs + ): + """Test token_combine forward and backward pass against reference""" + key = jax.random.PRNGKey(42) + + # Generate routing map + routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) + num_out_tokens = int(jnp.sum(routing_map)) + + # Get row_id_map from reference_token_dispatch + key, dummy_key = jax.random.split(key) + dummy_inp = jax.random.uniform( + dummy_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + _, _, row_id_map = reference_token_dispatch(dummy_inp, routing_map, num_out_tokens) + + # Generate input data (from expert outputs) + key, inp_key, merge_key = jax.random.split(key, 3) + inp = jax.random.uniform( + inp_key, (num_out_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + + if with_merging_probs: + merging_probs = jax.random.uniform( + merge_key, (num_tokens, num_experts), dtype=dtype, minval=0.0, maxval=1.0 + ) + # Normalize per token + merging_probs = merging_probs / (jnp.sum(merging_probs, axis=1, keepdims=True) + 1e-8) + else: + merging_probs = None + + # Define loss functions + def loss_fn(x): + output = token_combine(x, row_id_map, merging_probs) + return jnp.sum(output**2) + + def ref_loss_fn(x): + output = reference_token_combine(x, row_id_map, merging_probs) + return jnp.sum(output**2) + + loss_val, computed_grad = jax.value_and_grad(loss_fn)(inp) + ref_loss_val, ref_grad = jax.value_and_grad(ref_loss_fn)(inp) + + # Compare forward outputs + output = token_combine(inp, row_id_map, merging_probs) + ref_output = reference_token_combine(inp, row_id_map, merging_probs) + assert_allclose(output, ref_output) + + # Compare loss and gradient + assert_allclose(loss_val, ref_loss_val) + assert_allclose(computed_grad, ref_grad) + + # ========================================================================= + # sort_chunks_by_index tests + # ========================================================================= + + @pytest.mark.parametrize( + "num_splits,total_tokens,hidden_size", + [ + (4, 128, 256), + (8, 256, 512), + ], + ) + @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) + def test_sort_chunks_by_index(self, num_splits, total_tokens, hidden_size, dtype): + """Test sort_chunks_by_index forward and backward pass against reference""" + key = jax.random.PRNGKey(42) + + # Generate random split sizes + key, size_key = jax.random.split(key) + split_sizes = jax.random.randint(size_key, (num_splits,), 10, total_tokens // num_splits) + split_sizes = split_sizes.at[-1].set(total_tokens - jnp.sum(split_sizes[:-1])) + + # Generate sorted indices + key, sort_key = jax.random.split(key) + sorted_indices = jax.random.permutation(sort_key, num_splits) + + # Generate input data + key, inp_key = jax.random.split(key) + inp = jax.random.uniform( + inp_key, (total_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + + row_id_map = reference_make_chunk_sort_map( + split_sizes, sorted_indices, total_tokens, num_splits + ) + + # Define loss functions + def loss_fn(x): + output, _ = sort_chunks_by_index(x, split_sizes, sorted_indices) + return jnp.sum(output**2) + + def ref_loss_fn(x): + output, _ = reference_sort_chunks_by_map( + x, row_id_map, None, total_tokens, hidden_size, is_forward=True + ) + return jnp.sum(output**2) + + loss_val, computed_grad = jax.value_and_grad(loss_fn)(inp) + ref_loss_val, ref_grad = jax.value_and_grad(ref_loss_fn)(inp) + + # Compare forward outputs + output, _ = sort_chunks_by_index(inp, split_sizes, sorted_indices) + ref_output, _ = reference_sort_chunks_by_map( + inp, row_id_map, None, total_tokens, hidden_size, is_forward=True + ) + assert_allclose(output, ref_output) + + # Compare loss and gradient + assert_allclose(loss_val, ref_loss_val) + assert_allclose(computed_grad, ref_grad) + + # ========================================================================= + # Round-trip tests (token_dispatch -> expert processing -> token_combine) + # ========================================================================= + + @pytest.mark.parametrize( + "num_tokens,num_experts,hidden_size,tokens_per_expert", + [ + (32, 8, 256, 2), + (64, 16, 512, 3), + ], + ) + @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) + def test_dispatch_combine_roundtrip( + self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype + ): + """Test that token_dispatch followed by token_combine recovers original input""" + key = jax.random.PRNGKey(42) + + # Generate routing map + routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) + num_out_tokens = int(jnp.sum(routing_map)) + + # Generate input data + key, inp_key = jax.random.split(key) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + + # Create uniform merging probs (equal weight for all routed experts) + merging_probs = routing_map.astype(dtype) / jnp.maximum( + jnp.sum(routing_map, axis=1, keepdims=True), 1.0 + ) + + # Dispatch tokens to experts (returns output, permuted_probs, row_id_map) + dispatched, _, row_id_map = token_dispatch(inp, routing_map, num_out_tokens) + + # Combine tokens back (with uniform merging) (new signature) + combined = token_combine(dispatched, row_id_map, merging_probs) + + # Compare with original input + assert_allclose(combined, inp) diff --git a/transformer_engine/jax/cpp_extensions/amax.py b/transformer_engine/jax/cpp_extensions/amax.py index 2f3bc402ec..afc248a0ad 100644 --- a/transformer_engine/jax/cpp_extensions/amax.py +++ b/transformer_engine/jax/cpp_extensions/amax.py @@ -73,7 +73,7 @@ def abstract( transpose_batch_sequence, ): """ - amax calcuation abstract + amax calculation abstract """ del amax_scope, transpose_batch_sequence @@ -251,7 +251,7 @@ def impl( flatten_axis, ): """ - amax calcuation implementation + amax calculation implementation """ assert RHTAmaxCalculationPrimitive.inner_primitive is not None ( diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py new file mode 100644 index 0000000000..55a59a1650 --- /dev/null +++ b/transformer_engine/jax/permutation.py @@ -0,0 +1,401 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""MoE Permutation API for JAX. + +This module provides high-level token dispatch and combine operations for +Mixture of Experts (MoE) models with proper automatic differentiation support. + +Token Dispatch (Permute): + - Forward: Permute tokens according to routing map (scatter to experts) + - Backward: Unpermute gradients (gather from experts) + +Token Combine (Unpermute): + - Forward: Unpermute tokens and merge with weights (gather from experts) + - Backward: Permute gradients (scatter to experts) +""" + +from functools import partial +from typing import 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, + unpermute_with_mask_map, + unpermute_bwd_with_merging_probs, + make_chunk_sort_map, + sort_chunks_by_map, +) + +__all__ = [ + "token_dispatch", + "token_combine", + "sort_chunks_by_index", +] + + +def token_dispatch( + inp: jnp.ndarray, + routing_map: jnp.ndarray, + num_out_tokens: int, + probs: Optional[jnp.ndarray] = None, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: + """ + Dispatch tokens to experts based on routing map. + + This is the forward pass of the MoE permutation. Tokens are scattered + to their designated experts according to the routing map. The row_id_map + is computed internally from the routing_map. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [batch, sequence, hidden_size] or [num_tokens, hidden_size]. + routing_map : jnp.ndarray + Routing mask of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. + Values: 1 = routed, 0 = not routed. + num_out_tokens : int + The number of output tokens after permutation. This should equal the sum of + routing_map and must be provided explicitly for JIT compatibility. + probs : Optional[jnp.ndarray] + Optional routing probabilities of shape [batch, sequence, num_experts] or + [num_tokens, num_experts]. If provided, permuted_probs will be returned. + + Returns + ------- + output : jnp.ndarray + Permuted output tensor of shape [num_out_tokens, hidden_size]. + permuted_probs : Optional[jnp.ndarray] + Permuted probabilities of shape [num_out_tokens], or None if probs was not provided. + row_id_map : jnp.ndarray + Row ID map for use in token_combine (shape [num_tokens, num_experts * 2 + 1]). + """ + return _token_dispatch(inp, routing_map, probs, num_out_tokens) + + +@partial(jax.custom_vjp, nondiff_argnums=(1, 3)) +def _token_dispatch( + inp: jnp.ndarray, + routing_map: jnp.ndarray, + probs: Optional[jnp.ndarray], + num_out_tokens: int, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: + """Internal token_dispatch with custom VJP.""" + (output, permuted_probs, row_id_map), _ = _token_dispatch_fwd_rule( + inp, routing_map, probs, num_out_tokens + ) + return output, permuted_probs, row_id_map + + +def _token_dispatch_fwd_rule( + inp: jnp.ndarray, + routing_map: jnp.ndarray, + probs: Optional[jnp.ndarray], + num_out_tokens: int, +) -> Tuple[ + Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray], + Tuple[jnp.ndarray, int, int, int, bool], +]: + """Forward pass rule for token_dispatch.""" + # Validate input dimensions + assert inp.ndim in [2, 3], f"inp must be 2D or 3D, got {inp.ndim}D" + assert routing_map.ndim in [2, 3], f"routing_map must be 2D or 3D, got {routing_map.ndim}D" + + # Infer dimensions from input shapes + num_tokens = inp.shape[0] * inp.shape[1] if inp.ndim == 3 else inp.shape[0] + hidden_size = inp.shape[-1] + num_experts = routing_map.shape[-1] + + # Verify consistency between inp and routing_map + routing_num_tokens = ( + routing_map.shape[0] * routing_map.shape[1] + if routing_map.ndim == 3 + else routing_map.shape[0] + ) + assert num_tokens == routing_num_tokens, ( + f"Token count mismatch: inp has {num_tokens} tokens, " + f"routing_map has {routing_num_tokens} tokens" + ) + + # Always compute row_id_map internally from routing_map + row_id_map = make_row_id_map(routing_map, num_tokens, num_experts) + + with_probs = probs is not None + + output, permuted_probs = permute_with_mask_map( + inp, + row_id_map, + probs, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) + + # Return (primals, residuals) + # Include with_probs flag to know how to handle backward pass + residuals = (row_id_map, num_tokens, num_experts, hidden_size, with_probs) + return (output, permuted_probs, row_id_map), residuals + + +def _token_dispatch_bwd_rule( + _routing_map: jnp.ndarray, + _num_out_tokens: int, + residuals: Tuple[jnp.ndarray, int, int, int, bool], + g: Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray], +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """Backward pass rule for token_dispatch.""" + row_id_map, num_tokens, num_experts, hidden_size, with_probs = residuals + output_grad, permuted_probs_grad, _ = g # Ignore row_id_map gradient + + # Backward: unpermute gradients (gather from experts back to tokens) + inp_grad, probs_grad = unpermute_with_mask_map( + output_grad, + row_id_map, + None, # No merging probs + permuted_probs_grad if with_probs else None, + num_tokens, + num_experts, + hidden_size, + ) + + return inp_grad, probs_grad if with_probs else None + + +_token_dispatch.defvjp(_token_dispatch_fwd_rule, _token_dispatch_bwd_rule) + + +# ============================================================================= +# Token Combine (Unpermute) with VJP +# ============================================================================= + + +def token_combine( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: Optional[jnp.ndarray] = None, +) -> jnp.ndarray: + """ + Combine tokens from experts back to original token positions. + + This is the forward pass of MoE unpermutation. Tokens are gathered from + experts and merged (optionally weighted by merging_probs). + + Parameters + ---------- + inp : jnp.ndarray + Input tensor from experts of shape [num_out_tokens, hidden_size]. + row_id_map : jnp.ndarray + Row ID map from token_dispatch of shape [num_tokens, num_experts * 2 + 1]. + merging_probs : Optional[jnp.ndarray] + Merging weights of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. + If provided, tokens from different experts are weighted-summed. + If None, tokens are summed directly. + + Returns + ------- + output : jnp.ndarray + Combined output tensor of shape [num_tokens, hidden_size]. + """ + return _token_combine(inp, row_id_map, merging_probs) + + +@partial(jax.custom_vjp, nondiff_argnums=(1,)) +def _token_combine( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: Optional[jnp.ndarray], +) -> jnp.ndarray: + """Internal token_combine with custom VJP.""" + output, _ = _token_combine_fwd_rule(inp, row_id_map, merging_probs) + return output + + +def _token_combine_fwd_rule( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: Optional[jnp.ndarray], +) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray, Optional[jnp.ndarray], int, int, int, int]]: + """Forward pass rule for token_combine.""" + # Infer dimensions from row_id_map shape: [num_tokens, num_experts * 2 + 1] + num_tokens = row_id_map.shape[0] + num_experts = (row_id_map.shape[1] - 1) // 2 + hidden_size = inp.shape[-1] + num_out_tokens = inp.shape[0] + + # Call triton extension + output, _ = unpermute_with_mask_map( + inp, + row_id_map, + merging_probs, + None, # No permuted probs to unpermute + num_tokens, + num_experts, + hidden_size, + ) + + # Return (primal, residuals) + # Include inp in residuals for backward with merging_probs + residuals = ( + row_id_map, + inp, + merging_probs, + num_tokens, + num_experts, + hidden_size, + num_out_tokens, + ) + return output, residuals + + +def _token_combine_bwd_rule( + row_id_map: jnp.ndarray, + residuals: Tuple[jnp.ndarray, jnp.ndarray, Optional[jnp.ndarray], int, int, int, int], + g: jnp.ndarray, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """Backward pass rule for token_combine.""" + ( + row_id_map, + fwd_input, + merging_probs, + num_tokens, + num_experts, + hidden_size, + num_out_tokens, + ) = residuals + output_grad = g + + with_merging_probs = merging_probs is not None + + if with_merging_probs: + # Use specialized backward kernel that properly scales by merging_probs + inp_grad, merging_probs_grad = unpermute_bwd_with_merging_probs( + output_grad, + row_id_map, + fwd_input, + merging_probs, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) + else: + # Simple case: just permute gradients back + inp_grad, _ = permute_with_mask_map( + output_grad, + row_id_map, + None, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) + merging_probs_grad = None + + return inp_grad, merging_probs_grad + + +_token_combine.defvjp(_token_combine_fwd_rule, _token_combine_bwd_rule) + + +# ============================================================================= +# Chunk Sort with VJP +# ============================================================================= + + +def sort_chunks_by_index( + inp: jnp.ndarray, + split_sizes: jnp.ndarray, + sorted_indices: jnp.ndarray, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Sort chunks of tokens according to sorted indices. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape [batch, sequence, hidden_size] or [num_tokens, hidden_size]. + split_sizes : jnp.ndarray + Sizes of each chunk of shape [num_splits]. + sorted_indices : jnp.ndarray + Permutation indices for chunks of shape [num_splits]. + + Returns + ------- + output : jnp.ndarray + Sorted output tensor of shape [num_tokens, hidden_size]. + row_id_map : jnp.ndarray + Row ID map for reversing the sort. + """ + return _sort_chunks_by_index(inp, split_sizes, sorted_indices) + + +@partial(jax.custom_vjp, nondiff_argnums=(1, 2)) +def _sort_chunks_by_index( + inp: jnp.ndarray, + split_sizes: jnp.ndarray, + sorted_indices: jnp.ndarray, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Internal sort_chunks_by_index with custom VJP.""" + (output, row_id_map), _ = _sort_chunks_by_index_fwd_rule(inp, split_sizes, sorted_indices) + return output, row_id_map + + +def _sort_chunks_by_index_fwd_rule( + inp: jnp.ndarray, + split_sizes: jnp.ndarray, + sorted_indices: jnp.ndarray, +) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, int, int]]: + """Forward pass rule for sort_chunks_by_index.""" + # Validate input dimensions + assert inp.ndim in [2, 3], f"inp must be 2D or 3D, got {inp.ndim}D" + + # Infer dimensions from input shape + num_tokens = inp.shape[0] * inp.shape[1] if inp.ndim == 3 else inp.shape[0] + hidden_size = inp.shape[-1] + num_splits = split_sizes.shape[0] + + row_id_map = make_chunk_sort_map(split_sizes, sorted_indices, num_tokens, num_splits) + + output, _ = sort_chunks_by_map( + inp, + row_id_map, + None, # No probs + num_tokens, + hidden_size, + is_forward=True, + ) + + # Return (primals, residuals) + residuals = (row_id_map, num_tokens, hidden_size) + return (output, row_id_map), residuals + + +def _sort_chunks_by_index_bwd_rule( + _split_sizes: jnp.ndarray, + _sorted_indices: jnp.ndarray, + residuals: Tuple[jnp.ndarray, int, int], + g: Tuple[jnp.ndarray, jnp.ndarray], +) -> Tuple[jnp.ndarray]: + """Backward pass rule for sort_chunks_by_index.""" + row_id_map, num_tokens, hidden_size = residuals + output_grad, _ = g + + # Backward: reverse the sort + inp_grad, _ = sort_chunks_by_map( + output_grad, + row_id_map, + None, + num_tokens, + hidden_size, + is_forward=False, + ) + + return (inp_grad,) + + +_sort_chunks_by_index.defvjp(_sort_chunks_by_index_fwd_rule, _sort_chunks_by_index_bwd_rule) diff --git a/transformer_engine/jax/triton_extensions/__init__.py b/transformer_engine/jax/triton_extensions/__init__.py index 7ce6c476c2..13a36421bf 100644 --- a/transformer_engine/jax/triton_extensions/__init__.py +++ b/transformer_engine/jax/triton_extensions/__init__.py @@ -20,6 +20,10 @@ @staticmethod def lowering(ctx, x, **kwargs): return triton_call_lowering(ctx, my_kernel, x, ...) + + # Use permutation functions + from transformer_engine.jax.triton_extensions import make_row_id_map, permute_with_mask_map """ from .utils import * +from .permutation import * diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py new file mode 100644 index 0000000000..4f59f65a87 --- /dev/null +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -0,0 +1,1136 @@ +# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX/TE custom ops for permutation in MOE using Triton kernels.""" + +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +import triton + +from transformer_engine.jax.cpp_extensions.base import BasePrimitive, register_primitive +from transformer_engine.common.triton.permutation import ( + _row_id_map_pass_1_kernel, + _row_id_map_pass_2_kernel, + _row_id_map_pass_3_kernel, + _permute_kernel, + _unpermute_kernel, + _unpermute_bwd_with_merging_probs_kernel, + _make_chunk_sort_map_kernel, + _sort_chunks_by_map_kernel, +) +from .utils import triton_call_lowering + + +__all__ = [ + "make_row_id_map", + "permute_with_mask_map", + "unpermute_with_mask_map", + "unpermute_bwd_with_merging_probs", + "make_chunk_sort_map", + "sort_chunks_by_map", +] + +DEFAULT_BLOCK_SIZE = 1024 + + +def _get_min_block_size(kernel, default=128): + if hasattr(kernel, "configs"): + return min(config.kwargs.get("BLOCK_SIZE", default) for config in kernel.configs) + return default + + +class RowIdMapPass1Primitive(BasePrimitive): + """ + Pass 1 of row_id_map generation: block cumsum. + + For each expert, compute the cumsum of every block_size tokens. + """ + + name = "te_row_id_map_pass1_triton" + multiple_results = True + impl_static_args = (1, 2, 3) # num_tokens, num_experts, block_size + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract(routing_map_aval, *, num_tokens, num_experts, block_size): + """Shape/dtype inference for pass 1.""" + del block_size # Only affects grid, not output shape + + assert routing_map_aval.shape == ( + num_tokens, + num_experts, + ), f"routing_map shape mismatch: expected ({num_tokens}, {num_experts})" + + row_id_map_shape = (num_tokens, num_experts * 2 + 1) + workspace_shape = ( + num_experts, + triton.cdiv(num_tokens, DEFAULT_BLOCK_SIZE), + ) + + return ( + jax.core.ShapedArray(row_id_map_shape, jnp.int32), + jax.core.ShapedArray(workspace_shape, jnp.int32), + ) + + @staticmethod + def impl(routing_map, num_tokens, num_experts, block_size): + """Forward to inner primitive.""" + assert RowIdMapPass1Primitive.inner_primitive is not None + return RowIdMapPass1Primitive.inner_primitive.bind( + routing_map, + num_tokens=num_tokens, + num_experts=num_experts, + block_size=block_size, + ) + + @staticmethod + def lowering(ctx, routing_map, *, num_tokens, num_experts, block_size): + """MLIR lowering using triton_call_lowering.""" + # Compute strides + routing_stride_token = num_experts + routing_stride_expert = 1 + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + + grid = (num_experts, triton.cdiv(num_tokens, block_size)) + + # All scalar arguments must be passed as constexprs + return triton_call_lowering( + ctx, + _row_id_map_pass_1_kernel, + routing_map, # Only tensor arguments here + grid=grid, + constexprs={ + "num_tokens": num_tokens, + "stride_routing_map_token": routing_stride_token, + "stride_routing_map_expert": routing_stride_expert, + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(RowIdMapPass1Primitive) + + +class RowIdMapPass2Primitive(BasePrimitive): + """ + Pass 2 of row_id_map generation: cumsum all and process the mask. + """ + + name = "te_row_id_map_pass2_triton" + multiple_results = True + impl_static_args = (2, 3, 4) # num_tokens, num_experts, block_size + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract(row_id_map_aval, workspace_aval, *, num_tokens, num_experts, block_size): + """Shape/dtype inference for pass 2 (in-place operation).""" + del row_id_map_aval, workspace_aval + del block_size + + row_id_map_shape = (num_tokens, num_experts * 2 + 1) + workspace_shape = (num_experts, triton.cdiv(num_tokens, DEFAULT_BLOCK_SIZE)) + + return ( + jax.core.ShapedArray(row_id_map_shape, jnp.int32), + jax.core.ShapedArray(workspace_shape, jnp.int32), + ) + + @staticmethod + def impl(row_id_map, workspace, num_tokens, num_experts, block_size): + """Forward to inner primitive.""" + assert RowIdMapPass2Primitive.inner_primitive is not None + return RowIdMapPass2Primitive.inner_primitive.bind( + row_id_map, + workspace, + num_tokens=num_tokens, + num_experts=num_experts, + block_size=block_size, + ) + + @staticmethod + def lowering(ctx, row_id_map, workspace, *, num_tokens, num_experts, block_size): + """MLIR lowering using triton_call_lowering.""" + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + + grid = (num_experts, triton.cdiv(num_tokens, block_size)) + workspace_load_width = triton.next_power_of_2( + num_experts * triton.cdiv(num_tokens, block_size) + ) + + return triton_call_lowering( + ctx, + _row_id_map_pass_2_kernel, + row_id_map, + workspace, + grid=grid, + input_output_aliases={0: 0, 1: 1}, + constexprs={ + "num_tokens": num_tokens, + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "WORKSPACE_LOAD_WIDTH": workspace_load_width, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(RowIdMapPass2Primitive) + + +class RowIdMapPass3Primitive(BasePrimitive): + """ + Pass 3 of row_id_map generation: make the row_id_map from sparse to dense structure. + """ + + name = "te_row_id_map_pass3_triton" + multiple_results = False + impl_static_args = (1, 2) # num_tokens, num_experts + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract(row_id_map_aval, *, num_tokens, num_experts): + """Shape/dtype inference for pass 3 (in-place operation).""" + del row_id_map_aval + row_id_map_shape = (num_tokens, num_experts * 2 + 1) + return jax.core.ShapedArray(row_id_map_shape, jnp.int32) + + @staticmethod + def impl(row_id_map, num_tokens, num_experts): + """Forward to inner primitive.""" + assert RowIdMapPass3Primitive.inner_primitive is not None + return RowIdMapPass3Primitive.inner_primitive.bind( + row_id_map, + num_tokens=num_tokens, + num_experts=num_experts, + ) + + @staticmethod + def lowering(ctx, row_id_map, *, num_tokens, num_experts): + """MLIR lowering using triton_call_lowering.""" + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + + grid = (num_tokens,) + load_size = triton.next_power_of_2(num_experts) + + return triton_call_lowering( + ctx, + _row_id_map_pass_3_kernel, + row_id_map, + grid=grid, + input_output_aliases={0: 0}, + constexprs={ + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "num_experts": num_experts, + "LOAD_SIZE": load_size, + }, + ) + + +register_primitive(RowIdMapPass3Primitive) + + +class PermuteWithMaskMapPrimitive(BasePrimitive): + """ + Permute the input tensor based on the row_id_map. + """ + + name = "te_permute_with_mask_map_triton" + multiple_results = True + # scale and permuted_scale are dummy inputs (not used when PERMUTE_SCALE=False) + # but they need to be in the signature for the kernel call + impl_static_args = ( + 5, + 6, + 7, + 8, + 9, + ) # num_tokens, num_experts, num_out_tokens, hidden_size, with_probs + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + inp_aval, + row_id_map_aval, + probs_aval, + scale_aval, # dummy, same shape as inp + permuted_scale_aval, # dummy, same shape as inp + *, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + with_probs, + ): + """Shape/dtype inference for permute.""" + del row_id_map_aval, scale_aval, permuted_scale_aval + del num_tokens, num_experts + + output_shape = (num_out_tokens, hidden_size) + output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) + + if with_probs: + permuted_probs_aval = jax.core.ShapedArray((num_out_tokens,), probs_aval.dtype) + else: + permuted_probs_aval = jax.core.ShapedArray((0,), inp_aval.dtype) + + return output_aval, permuted_probs_aval + + @staticmethod + def impl( + inp, + row_id_map, + probs, + scale, + permuted_scale, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + with_probs, + ): + """Forward to inner primitive.""" + assert PermuteWithMaskMapPrimitive.inner_primitive is not None + return PermuteWithMaskMapPrimitive.inner_primitive.bind( + inp, + row_id_map, + probs, + scale, + permuted_scale, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + with_probs=with_probs, + ) + + @staticmethod + def lowering( + ctx, + inp, + row_id_map, + probs, + scale, + permuted_scale, + *, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + with_probs, + ): + """MLIR lowering using triton_call_lowering.""" + del num_out_tokens + inp_stride_token = hidden_size + inp_stride_hidden = 1 + output_stride_token = hidden_size + output_stride_hidden = 1 + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + permuted_probs_stride_token = 1 + + if with_probs: + # Check if probs is 2D [num_tokens, num_experts] or 1D [num_tokens] + probs_aval = ctx.avals_in[2] + if len(probs_aval.shape) > 1: + probs_stride_token = num_experts + probs_stride_expert = 1 + else: + probs_stride_token = 1 + probs_stride_expert = 1 + else: + 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 + block_size = _get_min_block_size(_permute_kernel) + grid = (num_tokens, triton.cdiv(hidden_size, block_size)) + + return triton_call_lowering( + ctx, + _permute_kernel, + inp, + row_id_map, + probs, + scale, + permuted_scale, + grid=grid, + constexprs={ + "scale_hidden_dim": 0, + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "stride_input_token": inp_stride_token, + "stride_input_hidden": inp_stride_hidden, + "stride_output_token": output_stride_token, + "stride_output_hidden": output_stride_hidden, + "stride_probs_token": probs_stride_token, + "stride_probs_expert": probs_stride_expert, + "stride_scale_token": hidden_size, + "stride_scale_hidden": 1, + "stride_permuted_probs_token": permuted_probs_stride_token, + "stride_permuted_scale_token": hidden_size, + "stride_permuted_scale_hidden": 1, + "num_experts": num_experts, + "hidden_size": hidden_size, + "PERMUTE_PROBS": with_probs, + "PERMUTE_SCALE": False, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(PermuteWithMaskMapPrimitive) + + +class UnpermuteWithMaskMapPrimitive(BasePrimitive): + """ + Unpermute the input tensor based on the row_id_map. + """ + + name = "te_unpermute_with_mask_map_triton" + multiple_results = True + impl_static_args = ( + 4, + 5, + 6, + 7, + 8, + ) # num_tokens, num_experts, hidden_size, with_merging_probs, with_probs + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + inp_aval, + row_id_map_aval, + merging_probs_aval, + permuted_probs_aval, + *, + num_tokens, + num_experts, + hidden_size, + with_merging_probs, + with_probs, + ): + """Shape/dtype inference for unpermute.""" + del row_id_map_aval, merging_probs_aval, with_merging_probs + + output_shape = (num_tokens, hidden_size) + output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) + + if with_probs: + unpermuted_probs_shape = (num_tokens, num_experts) + unpermuted_probs_aval = jax.core.ShapedArray( + unpermuted_probs_shape, permuted_probs_aval.dtype + ) + else: + unpermuted_probs_aval = jax.core.ShapedArray((0,), inp_aval.dtype) + + return output_aval, unpermuted_probs_aval + + @staticmethod + def impl( + inp, + row_id_map, + merging_probs, + permuted_probs, + num_tokens, + num_experts, + hidden_size, + with_merging_probs, + with_probs, + ): + """Forward to inner primitive.""" + assert UnpermuteWithMaskMapPrimitive.inner_primitive is not None + return UnpermuteWithMaskMapPrimitive.inner_primitive.bind( + inp, + row_id_map, + merging_probs, + permuted_probs, + num_tokens=num_tokens, + num_experts=num_experts, + hidden_size=hidden_size, + with_merging_probs=with_merging_probs, + with_probs=with_probs, + ) + + @staticmethod + def lowering( + ctx, + inp, + row_id_map, + merging_probs, + permuted_probs, + *, + num_tokens, + num_experts, + hidden_size, + with_merging_probs, + with_probs, + ): + """MLIR lowering using triton_call_lowering.""" + # Compute strides + inp_stride_token = hidden_size + inp_stride_hidden = 1 + output_stride_token = hidden_size + output_stride_hidden = 1 + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + + if with_merging_probs: + merging_probs_stride_token = num_experts + merging_probs_stride_expert = 1 + else: + merging_probs_stride_token = 0 + merging_probs_stride_expert = 0 + + permuted_probs_stride_token = 1 + unpermuted_probs_stride_token = num_experts + unpermuted_probs_stride_expert = 1 + + # Grid - use minimum BLOCK_SIZE from autotune configs + block_size = _get_min_block_size(_unpermute_kernel) + grid = (num_tokens, triton.cdiv(hidden_size, block_size)) + + return triton_call_lowering( + ctx, + _unpermute_kernel, + inp, + row_id_map, + merging_probs, + permuted_probs, + grid=grid, + constexprs={ + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "stride_input_token": inp_stride_token, + "stride_input_hidden": inp_stride_hidden, + "stride_output_token": output_stride_token, + "stride_output_hidden": output_stride_hidden, + "stride_merging_probs_token": merging_probs_stride_token, + "stride_merging_probs_expert": merging_probs_stride_expert, + "stride_permuted_probs_token": permuted_probs_stride_token, + "stride_unpermuted_probs_token": unpermuted_probs_stride_token, + "stride_unpermuted_probs_expert": unpermuted_probs_stride_expert, + "num_experts": num_experts, + "hidden_size": hidden_size, + "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), + "WITH_MERGING_PROBS": with_merging_probs, + "PERMUTE_PROBS": with_probs, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(UnpermuteWithMaskMapPrimitive) + + +class UnpermuteBwdWithMergingProbsPrimitive(BasePrimitive): + """ + Backward pass for unpermute with merging probabilities. + + This kernel computes gradients for both the input and merging_probs. + """ + + name = "te_unpermute_bwd_with_merging_probs_triton" + multiple_results = True + impl_static_args = (4, 5, 6, 7) # num_tokens, num_experts, num_out_tokens, hidden_size + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + fwd_output_grad_aval, + fwd_input_aval, + merging_probs_aval, + row_id_map_aval, + *, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ): + """Shape/dtype inference for unpermute backward with merging probs.""" + del fwd_input_aval, row_id_map_aval + + # fwd_input_grad has same shape as fwd_input + fwd_input_grad_shape = (num_out_tokens, hidden_size) + fwd_input_grad_aval = jax.core.ShapedArray(fwd_input_grad_shape, fwd_output_grad_aval.dtype) + + # merging_probs_grad has same shape as merging_probs + merging_probs_grad_shape = (num_tokens, num_experts) + merging_probs_grad_aval = jax.core.ShapedArray( + merging_probs_grad_shape, merging_probs_aval.dtype + ) + + return fwd_input_grad_aval, merging_probs_grad_aval + + @staticmethod + def impl( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ): + """Forward to inner primitive.""" + assert UnpermuteBwdWithMergingProbsPrimitive.inner_primitive is not None + return UnpermuteBwdWithMergingProbsPrimitive.inner_primitive.bind( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + ) + + @staticmethod + def lowering( + ctx, + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + *, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ): + """MLIR lowering using triton_call_lowering.""" + del num_out_tokens + + # Compute strides + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + fwd_output_grad_stride_token = hidden_size + fwd_output_grad_stride_hidden = 1 + fwd_input_grad_stride_token = hidden_size + fwd_input_grad_stride_hidden = 1 + fwd_input_stride_token = hidden_size + fwd_input_stride_hidden = 1 + merging_probs_stride_token = num_experts + merging_probs_stride_expert = 1 + merging_probs_grad_stride_token = num_experts + merging_probs_grad_stride_expert = 1 + + # Grid - one program per token + grid = (num_tokens,) + + # Get min block size from autotune configs for consistency + block_size = _get_min_block_size(_unpermute_bwd_with_merging_probs_kernel) + + # Pass inputs in kernel argument order: fwd_output_grad, fwd_input, merging_probs, row_id_map + return triton_call_lowering( + ctx, + _unpermute_bwd_with_merging_probs_kernel, + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + grid=grid, + constexprs={ + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "stride_fwd_output_grad_token": fwd_output_grad_stride_token, + "stride_fwd_output_grad_hidden": fwd_output_grad_stride_hidden, + "stride_fwd_input_grad_token": fwd_input_grad_stride_token, + "stride_fwd_input_grad_hidden": fwd_input_grad_stride_hidden, + "stride_fwd_input_token": fwd_input_stride_token, + "stride_fwd_input_hidden": fwd_input_stride_hidden, + "stride_merging_probs_token": merging_probs_stride_token, + "stride_merging_probs_expert": merging_probs_stride_expert, + "stride_merging_probs_grad_token": merging_probs_grad_stride_token, + "stride_merging_probs_grad_expert": merging_probs_grad_stride_expert, + "num_experts": num_experts, + "hidden_size": hidden_size, + "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(UnpermuteBwdWithMergingProbsPrimitive) + + +def unpermute_bwd_with_merging_probs( + fwd_output_grad: jnp.ndarray, + row_id_map: jnp.ndarray, + fwd_input: jnp.ndarray, + merging_probs: jnp.ndarray, + num_tokens: int, + num_experts: int, + num_out_tokens: int, + hidden_size: int, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Backward pass for unpermute with merging probabilities. + + This computes gradients for both the input tensor and merging_probs. + + Parameters + ---------- + fwd_output_grad : jnp.ndarray + Gradient of the forward output of shape `[num_tokens, hidden_size]`. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. + fwd_input : jnp.ndarray + The input tensor from the forward pass of shape `[num_out_tokens, hidden_size]`. + merging_probs : jnp.ndarray + The merging probabilities of shape `[num_tokens, num_experts]`. + num_tokens : int + Number of tokens in the unpermuted tensor. + num_experts : int + Number of experts. + num_out_tokens : int + Number of tokens in the permuted tensor. + hidden_size : int + Hidden size. + + Returns + ------- + fwd_input_grad : jnp.ndarray + Gradient w.r.t. the input tensor of shape `[num_out_tokens, hidden_size]`. + merging_probs_grad : jnp.ndarray + Gradient w.r.t. merging_probs of shape `[num_tokens, num_experts]`. + """ + # Pass arguments in kernel order: fwd_output_grad, fwd_input, merging_probs, row_id_map + return UnpermuteBwdWithMergingProbsPrimitive.outer_primitive.bind( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + ) + + +class MakeChunkSortMapPrimitive(BasePrimitive): + """ + Make a row_id_map for chunk sort. + """ + + name = "te_make_chunk_sort_map_triton" + multiple_results = False + impl_static_args = (2, 3) # num_tokens, num_splits + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract(split_sizes_aval, sorted_indices_aval, *, num_tokens, num_splits): + """Shape/dtype inference.""" + del sorted_indices_aval + assert split_sizes_aval.shape == (num_splits,) + return jax.core.ShapedArray((num_tokens,), jnp.int32) + + @staticmethod + def impl(split_sizes, sorted_indices, num_tokens, num_splits): + """Forward to inner primitive.""" + assert MakeChunkSortMapPrimitive.inner_primitive is not None + return MakeChunkSortMapPrimitive.inner_primitive.bind( + split_sizes, + sorted_indices, + num_tokens=num_tokens, + num_splits=num_splits, + ) + + @staticmethod + def lowering(ctx, split_sizes, sorted_indices, *, num_tokens, num_splits): + """MLIR lowering using triton_call_lowering.""" + grid = (num_tokens,) + + return triton_call_lowering( + ctx, + _make_chunk_sort_map_kernel, + split_sizes, + sorted_indices, + grid=grid, + constexprs={ + "num_splits": num_splits, + "IDX_LOAD_WIDTH": triton.next_power_of_2(num_splits), + }, + ) + + +register_primitive(MakeChunkSortMapPrimitive) + + +class SortChunksByMapPrimitive(BasePrimitive): + """ + Sort chunks with row_id_map. + """ + + name = "te_sort_chunks_by_map_triton" + multiple_results = True + impl_static_args = (3, 4, 5, 6) # num_tokens, hidden_size, is_forward, with_probs + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + inp_aval, row_id_map_aval, probs_aval, *, num_tokens, hidden_size, is_forward, with_probs + ): + """Shape/dtype inference.""" + del row_id_map_aval, is_forward + + output_aval = jax.core.ShapedArray((num_tokens, hidden_size), inp_aval.dtype) + + if with_probs: + permuted_probs_aval = jax.core.ShapedArray((num_tokens,), probs_aval.dtype) + else: + permuted_probs_aval = jax.core.ShapedArray((0,), inp_aval.dtype) + + return output_aval, permuted_probs_aval + + @staticmethod + def impl(inp, row_id_map, probs, num_tokens, hidden_size, is_forward, with_probs): + """Forward to inner primitive.""" + assert SortChunksByMapPrimitive.inner_primitive is not None + return SortChunksByMapPrimitive.inner_primitive.bind( + inp, + row_id_map, + probs, + num_tokens=num_tokens, + hidden_size=hidden_size, + is_forward=is_forward, + with_probs=with_probs, + ) + + @staticmethod + def lowering(ctx, inp, row_id_map, probs, *, num_tokens, hidden_size, is_forward, with_probs): + """MLIR lowering using triton_call_lowering.""" + # Compute strides + inp_stride_token = hidden_size + inp_stride_hidden = 1 + output_stride_token = hidden_size + output_stride_hidden = 1 + probs_stride_token = 1 + permuted_probs_stride_token = 1 + + # Grid - use minimum BLOCK_SIZE from autotune configs + block_size = _get_min_block_size(_sort_chunks_by_map_kernel) + grid = (num_tokens, triton.cdiv(hidden_size, block_size)) + + return triton_call_lowering( + ctx, + _sort_chunks_by_map_kernel, + inp, + row_id_map, + probs, + grid=grid, + constexprs={ + "stride_input_token": inp_stride_token, + "stride_input_hidden": inp_stride_hidden, + "stride_output_token": output_stride_token, + "stride_output_hidden": output_stride_hidden, + "stride_probs_token": probs_stride_token, + "stride_permuted_probs_token": permuted_probs_stride_token, + "hidden_size": hidden_size, + "PERMUTE_PROBS": with_probs, + "FORWARD": is_forward, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(SortChunksByMapPrimitive) + + +def make_row_id_map( + routing_map: jnp.ndarray, + num_tokens: int, + num_experts: int, +) -> jnp.ndarray: + """ + Prepare the row_id_map for the permutation. + + This function chains 3 Triton kernel passes together. + + Parameters + ---------- + routing_map : jnp.ndarray + Input tensor of shape `[num_tokens, num_experts]`. It is a mask tensor that indicates + which experts are routed to which tokens. The values in it: 1 means the token is routed to + this expert and 0 means not. + num_tokens : int + Number of tokens in the input tensor. + num_experts : int + Number of experts in the input tensor. + + Returns + ------- + row_id_map : jnp.ndarray + The row_id_map for the permutation of shape `[num_tokens, num_experts * 2 + 1]`. + For each token, the last item is the number of experts that are routed (n_routed). + The first n_routed items are the destination row indices in the permuted tokens. + The [num_experts, num_experts + n_routed) items are the indices of the experts corresponding + to the first n_routed row indices above. + """ + block_size = DEFAULT_BLOCK_SIZE + + # Pass 1: Block cumsum + row_id_map_pass1, workspace_tensor = RowIdMapPass1Primitive.outer_primitive.bind( + routing_map, + num_tokens=num_tokens, + num_experts=num_experts, + block_size=block_size, + ) + + # Pass 2: Cumsum all and process the mask + row_id_map_pass2, _ = RowIdMapPass2Primitive.outer_primitive.bind( + row_id_map_pass1, + workspace_tensor, + num_tokens=num_tokens, + num_experts=num_experts, + block_size=block_size, + ) + + # Initialize columns [num_experts:] to -1 since Pass 1/2 only wrote to [0:num_experts] + # Reference implementation expects -1 for invalid entries + row_id_map = row_id_map_pass2.at[:, num_experts:].set(-1) + + # Pass 3: Make the row_id_map from sparse to dense structure + row_id_map = RowIdMapPass3Primitive.outer_primitive.bind( + row_id_map, + num_tokens=num_tokens, + num_experts=num_experts, + ) + + return row_id_map + + +def permute_with_mask_map( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + probs: Optional[jnp.ndarray], + num_tokens: int, + num_experts: int, + num_out_tokens: int, + hidden_size: int, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """ + Permute the input tensor based on the row_id_map. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. + probs : Optional[jnp.ndarray] + The probabilities of the input tensor. If it is not None, it will be permuted. + num_tokens : int + Number of tokens in the input tensor. + num_experts : int + Number of experts in the input tensor. + num_out_tokens : int + Number of tokens in the permuted tensor. + hidden_size : int + Hidden size of the input tensor. + + Returns + ------- + output : jnp.ndarray + Permuted output tensor of shape `[num_out_tokens, hidden_size]`. + permuted_probs : Optional[jnp.ndarray] + Permuted probabilities if probs was provided, None otherwise. + """ + with_probs = probs is not None + + # Handle None probs by creating dummy tensor + if not with_probs: + probs = jnp.zeros((0,), dtype=inp.dtype) + + # Create dummy scale tensors (not used when PERMUTE_SCALE=False, but required by kernel signature) + dummy_scale = inp + dummy_permuted_scale = inp + + output, permuted_probs = PermuteWithMaskMapPrimitive.outer_primitive.bind( + inp, + row_id_map, + probs, + dummy_scale, + dummy_permuted_scale, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + with_probs=with_probs, + ) + + if not with_probs: + permuted_probs = None + + return output, permuted_probs + + +def unpermute_with_mask_map( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: Optional[jnp.ndarray], + permuted_probs: Optional[jnp.ndarray], + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """ + Unpermute the input tensor based on the row_id_map. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape `[num_out_tokens, hidden_size]`. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. + merging_probs : Optional[jnp.ndarray] + The merging probabilities of the input tensor. If it is not None, it will be used as weights + to reduce the unpermuted tokens. + permuted_probs : Optional[jnp.ndarray] + The permuted probabilities of the input tensor. If it is not None, it will be unpermuted. + num_tokens : int + Number of tokens in the permuted tensor. + num_experts : int + Number of experts in the permuted tensor. + hidden_size : int + Hidden size of the permuted tensor. + + Returns + ------- + output : jnp.ndarray + Unpermuted output tensor of shape `[num_tokens, hidden_size]`. + unpermuted_probs : Optional[jnp.ndarray] + Unpermuted probabilities if permuted_probs was provided, None otherwise. + """ + with_merging_probs = merging_probs is not None + with_probs = permuted_probs is not None + + # Handle None inputs by creating dummy tensors + if not with_merging_probs: + merging_probs = jnp.zeros((0,), dtype=inp.dtype) + if not with_probs: + permuted_probs = jnp.zeros((0,), dtype=inp.dtype) + + output, unpermuted_probs = UnpermuteWithMaskMapPrimitive.outer_primitive.bind( + inp, + row_id_map, + merging_probs, + permuted_probs, + num_tokens=num_tokens, + num_experts=num_experts, + hidden_size=hidden_size, + with_merging_probs=with_merging_probs, + with_probs=with_probs, + ) + + if not with_probs: + unpermuted_probs = None + + return output, unpermuted_probs + + +def make_chunk_sort_map( + split_sizes: jnp.ndarray, + sorted_indices: jnp.ndarray, + num_tokens: int, + num_splits: int, +) -> jnp.ndarray: + """ + Make a row_id_map for chunk sort. + + Parameters + ---------- + split_sizes : jnp.ndarray + The sizes of the chunks of shape `[num_splits,]`. + sorted_indices : jnp.ndarray + The indices of the sorted chunks of shape `[num_splits,]`. + num_tokens : int + Number of tokens in the input tensor. + num_splits : int + Number of splits of split_sizes and sorted_indices. + + Returns + ------- + row_id_map : jnp.ndarray + Row ID map for chunk sorting of shape `[num_tokens,]`. + """ + return MakeChunkSortMapPrimitive.outer_primitive.bind( + split_sizes, + sorted_indices, + num_tokens=num_tokens, + num_splits=num_splits, + ) + + +def sort_chunks_by_map( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + probs: Optional[jnp.ndarray], + num_tokens: int, + hidden_size: int, + is_forward: bool, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """ + Sort chunks with row_id_map. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape `[num_tokens, hidden_size]`. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens,]`. + probs : Optional[jnp.ndarray] + The probabilities of the input tensor. If it is not None, it will be permuted. + num_tokens : int + Number of tokens in the input tensor. + hidden_size : int + Hidden size of the input tensor. + is_forward : bool + Whether the sort is for forward or backward. + + Returns + ------- + output : jnp.ndarray + Sorted output tensor of shape `[num_tokens, hidden_size]`. + permuted_probs : Optional[jnp.ndarray] + Sorted probabilities if probs was provided, None otherwise. + """ + with_probs = probs is not None + + # Handle None probs by creating dummy tensor + if not with_probs: + probs = jnp.zeros((0,), dtype=inp.dtype) + + output, permuted_probs = SortChunksByMapPrimitive.outer_primitive.bind( + inp, + row_id_map, + probs, + num_tokens=num_tokens, + hidden_size=hidden_size, + is_forward=is_forward, + with_probs=with_probs, + ) + + if not with_probs: + permuted_probs = None + + return output, permuted_probs diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index accb316fec..12d6a9e3de 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -176,7 +176,9 @@ def triton_call_lowering( *array_args: Input arrays (from ctx) grid: Grid dimensions (int or tuple) input_output_aliases: Mapping of input to output aliases - constexprs: Compile-time constants for the kernel + constexprs: Compile-time constants for the kernel. This includes both + tl.constexpr arguments AND scalar runtime arguments (like + num_tokens, strides) that are known at JAX trace time. Returns: MLIR lowering result @@ -189,8 +191,10 @@ def lowering(ctx, x, *, block_size): return triton_call_lowering( ctx, my_kernel, x, grid=(triton.cdiv(n, block_size),), - n_elements=n, - BLOCK_SIZE=block_size + constexprs={ + "n_elements": n, # scalar arg (not tl.constexpr in kernel) + "BLOCK_SIZE": block_size, # tl.constexpr arg + }, ) """ # Get compute capability using gpu_triton @@ -203,9 +207,13 @@ def lowering(ctx, x, *, block_size): else: arg_names = kernel_fn.arg_names - # Build signature for inputs + outputs + # Build signature for tensor arguments only (inputs + outputs) + # Scalar arguments should be passed via constexprs and will be + # specialized into the kernel at compile time all_avals = list(ctx.avals_in) + list(ctx.avals_out) - signature = {arg_names[i]: get_triton_dtype(aval) for i, aval in enumerate(all_avals)} + constexpr_names = set(constexprs.keys()) if constexprs else set() + 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): From 5afbb0e14f58e068a6370eb9e5dbcbb96bc7ac04 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 9 Dec 2025 19:25:44 -0800 Subject: [PATCH 119/521] [JAX] Make softmax_type in FFI optional (#2491) * Make softmax_type in FFI optional Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add warn message Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../jax/csrc/extensions/attention.cpp | 57 ++++++++++--------- transformer_engine/jax/csrc/extensions/ffi.h | 15 +++++ 2 files changed, 44 insertions(+), 28 deletions(-) diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index a834273035..79436fb8b7 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -333,34 +333,35 @@ static void FusedAttnForwardImpl( nvte_tensor_pack_destroy(&aux_output_tensors); } -#define FUSED_ATTN_FFI_GET_ATTRS \ - size_t input_batch = get_attr_value(attrs, "input_batch"); \ - size_t bias_batch = get_attr_value(attrs, "bias_batch"); \ - size_t q_max_seqlen = get_attr_value(attrs, "q_max_seqlen"); \ - size_t kv_max_seqlen = get_attr_value(attrs, "kv_max_seqlen"); \ - size_t attn_heads = get_attr_value(attrs, "attn_heads"); \ - size_t num_gqa_groups = get_attr_value(attrs, "num_gqa_groups"); \ - size_t bias_heads = get_attr_value(attrs, "bias_heads"); \ - size_t qk_head_dim = get_attr_value(attrs, "qk_head_dim"); \ - size_t v_head_dim = get_attr_value(attrs, "v_head_dim"); \ - size_t max_segments_per_seq = get_attr_value(attrs, "max_segments_per_seq"); \ - auto window_size_left = get_attr_value(attrs, "window_size_left"); \ - auto window_size_right = get_attr_value(attrs, "window_size_right"); \ - float scaling_factor = get_attr_value(attrs, "scaling_factor"); \ - float dropout_probability = get_attr_value(attrs, "dropout_probability"); \ - NVTE_Bias_Type bias_type = \ - static_cast(get_attr_value(attrs, "bias_type")); \ - NVTE_Mask_Type mask_type = \ - static_cast(get_attr_value(attrs, "mask_type")); \ - NVTE_Softmax_Type softmax_type = \ - static_cast(get_attr_value(attrs, "softmax_type")); \ - NVTE_QKV_Layout qkv_layout = \ - static_cast(get_attr_value(attrs, "qkv_layout")); \ - bool is_training = get_attr_value(attrs, "is_training"); \ - bool deterministic = get_attr_value(attrs, "deterministic"); \ - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ - size_t wkspace_size = product(workspace_buf->dimensions()); \ - DType dtype = convert_ffi_datatype_to_te_dtype(q_buf.element_type()); \ +#define FUSED_ATTN_FFI_GET_ATTRS \ + size_t input_batch = get_attr_value(attrs, "input_batch"); \ + size_t bias_batch = get_attr_value(attrs, "bias_batch"); \ + size_t q_max_seqlen = get_attr_value(attrs, "q_max_seqlen"); \ + size_t kv_max_seqlen = get_attr_value(attrs, "kv_max_seqlen"); \ + size_t attn_heads = get_attr_value(attrs, "attn_heads"); \ + size_t num_gqa_groups = get_attr_value(attrs, "num_gqa_groups"); \ + size_t bias_heads = get_attr_value(attrs, "bias_heads"); \ + size_t qk_head_dim = get_attr_value(attrs, "qk_head_dim"); \ + size_t v_head_dim = get_attr_value(attrs, "v_head_dim"); \ + size_t max_segments_per_seq = get_attr_value(attrs, "max_segments_per_seq"); \ + auto window_size_left = get_attr_value(attrs, "window_size_left"); \ + auto window_size_right = get_attr_value(attrs, "window_size_right"); \ + float scaling_factor = get_attr_value(attrs, "scaling_factor"); \ + float dropout_probability = get_attr_value(attrs, "dropout_probability"); \ + NVTE_Bias_Type bias_type = \ + static_cast(get_attr_value(attrs, "bias_type")); \ + NVTE_Mask_Type mask_type = \ + static_cast(get_attr_value(attrs, "mask_type")); \ + NVTE_Softmax_Type softmax_type = \ + static_cast(get_attr_value_or_default( \ + attrs, "softmax_type", static_cast(NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX))); \ + NVTE_QKV_Layout qkv_layout = \ + static_cast(get_attr_value(attrs, "qkv_layout")); \ + bool is_training = get_attr_value(attrs, "is_training"); \ + bool deterministic = get_attr_value(attrs, "deterministic"); \ + auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ + size_t wkspace_size = product(workspace_buf->dimensions()); \ + DType dtype = convert_ffi_datatype_to_te_dtype(q_buf.element_type()); \ DType wkspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type()); Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Type k_buf, diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index 0fc2e83898..d4f76a011a 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -75,6 +75,21 @@ T get_attr_value(Dictionary& attrs, std::string attr_name, return attr.value(); } +template +T get_attr_value_or_default(Dictionary& attrs, std::string attr_name, T default_value, + const source_location& loc = source_location::current()) { + auto attr = attrs.get(attr_name); + if (attr.has_error()) { + NVTE_WARN("Failure in getting attribute value of '", attr_name, "'\n", + "Called from: ", loc.file_name(), ":", loc.line(), "\n", + "In function: ", loc.function_name(), "\n", + "Please ensure the attribute name and datatype match between C++ and Python APIs. " + "Currently falling back to a default value."); + return default_value; + } + return attr.value(); +} + inline size_t product(const xla::ffi::Span& data, size_t start_idx = 0, size_t end_idx = 0) { end_idx = (end_idx == 0) ? data.size() : end_idx; From e411547b6124cc44d74733f6a79841f9b85ba076 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, 10 Dec 2025 12:16:01 +0100 Subject: [PATCH 120/521] [PyTorch Debug] Add nvdlfw-inspect to dependencies (#2173) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- build_tools/pytorch.py | 13 ++++++++++--- docs/debug/1_getting_started.rst | 2 +- qa/L0_pytorch_debug_unittest/test.sh | 5 ----- qa/L1_pytorch_distributed_unittest/test.sh | 6 ------ qa/L1_pytorch_onnx_unittest/test.sh | 4 ---- 5 files changed, 11 insertions(+), 19 deletions(-) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 302816c6fd..93fef60f02 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -14,12 +14,19 @@ def install_requirements() -> List[str]: """Install dependencies for TE/PyTorch extensions.""" - return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic"] + return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic", "nvdlfw-inspect"] def test_requirements() -> List[str]: - """Test dependencies for TE/JAX extensions.""" - return ["numpy", "torchvision", "transformers", "torchao==0.13"] + """Test dependencies for TE/PyTorch extensions.""" + return [ + "numpy", + "torchvision", + "transformers", + "torchao==0.13", + "onnxruntime", + "onnxruntime_extensions", + ] def setup_pytorch_extension( diff --git a/docs/debug/1_getting_started.rst b/docs/debug/1_getting_started.rst index a5cdc1a6b1..58adf73cef 100644 --- a/docs/debug/1_getting_started.rst +++ b/docs/debug/1_getting_started.rst @@ -21,7 +21,7 @@ Transformer Engine provides a set of precision debug tools which allow you to ea There are 4 things one needs to do to use Transformer Engine debug features: 1. Create a configuration YAML file to configure the desired features. -2. Import, initialize, and install the `Nvidia-DL-Framework-Inspect `_ tool. +2. Import and initialize the `Nvidia-DL-Framework-Inspect `_ tool, which is installed as a dependency of Transformer Engine. 3. One can pass ``name="..."`` when creating TE layers to easier identify layer names. If this is not provided, names will be inferred automatically. 4. Invoke ``debug_api.step()`` at the end of one forward-backward pass. diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index b6c42109b5..a176d21b15 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -26,11 +26,6 @@ mkdir -p "$XML_LOG_DIR" # Nvinspect will be disabled if no feature is active. : ${NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE:=$TE_PATH/tests/pytorch/debug/test_configs/dummy_feature.yaml} -# It is not installed as a requirement, -# because it is not available on PyPI. -pip uninstall -y nvdlfw-inspect -pip install git+https://github.com/NVIDIA/nvidia-dlfw-inspect.git - pip install pytest==8.2.1 || error_exit "Failed to install pytest" pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity.xml $TE_PATH/tests/pytorch/debug/test_sanity.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || test_fail "test_sanity.py" diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index e698e997a6..e0abdd281b 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -20,12 +20,6 @@ FAILED_CASES="" : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" - -# It is not installed as a requirement, -# because it is not available on PyPI. -pip uninstall -y nvdlfw-inspect -pip install git+https://github.com/NVIDIA/nvidia-dlfw-inspect.git - pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" 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" diff --git a/qa/L1_pytorch_onnx_unittest/test.sh b/qa/L1_pytorch_onnx_unittest/test.sh index 7fce13a3dc..303c5c281a 100644 --- a/qa/L1_pytorch_onnx_unittest/test.sh +++ b/qa/L1_pytorch_onnx_unittest/test.sh @@ -2,10 +2,6 @@ # # See LICENSE for license information. - -pip3 install onnxruntime -pip3 install onnxruntime_extensions - : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" From 93c5c65b5f4b34326c059233e90f74df676bbf65 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Wed, 10 Dec 2025 15:56:24 -0800 Subject: [PATCH 121/521] [PyTorch] Add THD support for max_logit/MuonClip (#2480) * update FE; initial pass at thd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * produce Stats+Max instead of Max+Sum_Exp Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "produce Stats+Max instead of Max+Sum_Exp" This reverts commit c7d2b77b2da9ff3f68344097284187ac427eeb6a. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- 3rdparty/cudnn-frontend | 2 +- .../common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu | 6 +++--- .../pytorch/attention/dot_product_attention/utils.py | 6 ------ 3 files changed, 4 insertions(+), 10 deletions(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index be6c079be8..0258951d4d 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit be6c079be8aaffa0fc079fcf039887e637c289c7 +Subproject commit 0258951d4d512f4714eb1574496f4d57669b1b93 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 14468b543a..efa4c78439 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 @@ -1101,7 +1101,7 @@ void fused_attn_arbitrary_seqlen_fwd( Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_Max->data.shape = {max_tokens_q, num_attn_heads, 1}; + output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; } @@ -1109,7 +1109,7 @@ void fused_attn_arbitrary_seqlen_fwd( Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Sum_Exp->data.dptr = nullptr; if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_Sum_Exp->data.shape = {max_tokens_q, num_attn_heads, 1}; + output_Sum_Exp->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; } @@ -1118,7 +1118,7 @@ void fused_attn_arbitrary_seqlen_fwd( Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { - output_S->data.shape = {max_tokens_q, num_attn_heads, 1}; + output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; } diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 8c6b6afc90..10a06ed965 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -532,9 +532,6 @@ def get_attention_backend( if use_flash_attention: use_flash_attention = False logger.debug("Disabling FlashAttention for max_logit") - if use_fused_attention and qkv_format == "thd": - use_fused_attention = False - logger.debug("Disabling FusedAttention for max_logit with qkv_format = thd") if fp8 and fp8_meta["recipe"].fp8_dpa: use_flash_attention = False use_fused_attention = False @@ -677,9 +674,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # Filter: QKV layout if qkv_format == "thd": - if use_unfused_attention: - logger.debug("Disabling UnfusedDotProductAttention for qkv_format = thd") - use_unfused_attention = False if pad_between_seqs: if (use_flash_attention_2 and FlashAttentionUtils.is_installed) or ( use_flash_attention_3 and FlashAttentionUtils.v3_is_installed From a5694f261048e58a63baf66816ed93a29694e056 Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Thu, 11 Dec 2025 11:49:20 +0100 Subject: [PATCH 122/521] Add separate RNG states for column-wise quantization with Stochastic Rounding (#2487) * Add separate RNG states for columnwise quantization with Stochastic Rounding Signed-off-by: Evgeny * Fix single tensor path Signed-off-by: Evgeny --------- Signed-off-by: Evgeny --- .../pytorch/csrc/extensions/cast.cpp | 34 ++++++++++++++++++- transformer_engine/pytorch/csrc/quantizer.cpp | 34 ++++++++++++++++--- 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index fb5f0b55d4..ac541435c7 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -761,8 +761,16 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, } // Stochastic rounding + // When both rowwise and columnwise quantization are used, + // we need separate RNG states for each to ensure they use different random numbers. std::vector te_rng_state_list; + std::vector te_rng_state_columnwise_list; + std::vector quant_config_columnwise_list; at::Tensor rng_states_tensor; + at::Tensor rng_states_columnwise_tensor; + const bool need_separate_columnwise_rng = + quantizer.stochastic_rounding && quantizer.with_rht && quantizer.columnwise_usage; + if (quantizer.stochastic_rounding) { // TODO(zhongbo): remove the for loop of generating rng states with a single call // with rng_elts_per_thread = 1024 * num_tensors @@ -770,9 +778,18 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, const size_t rng_elts_per_thread = 1024; // Wild guess, probably can be tightened auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); rng_states_tensor = torch::empty({static_cast(2 * num_tensors)}, opts); + + // Allocate columnwise RNG resources when separate RNG is needed + if (need_separate_columnwise_rng) { + rng_states_columnwise_tensor = torch::empty({static_cast(2 * num_tensors)}, opts); + for (size_t i = 0; i < num_tensors; ++i) { + quant_config_columnwise_list.emplace_back(QuantizationConfigWrapper()); + } + } for (size_t i = 0; i < num_tensors; ++i) { auto gen = at::get_generator_or_default( std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + // Generate RNG state for rowwise quantization at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); int64_t *rng_state_ptr = static_cast(rng_states_tensor.data_ptr()) + i * 2; philox_unpack(philox_args, rng_state_ptr); @@ -780,6 +797,18 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, static_cast(rng_state_ptr), std::vector{2}, DType::kInt64)); quant_config_list[i].set_rng_state(te_rng_state_list[i].data()); quant_config_list[i].set_stochastic_rounding(true); + + // Generate separate RNG state for columnwise quantization + if (need_separate_columnwise_rng) { + at::PhiloxCudaState philox_args_columnwise = init_philox_state(gen, rng_elts_per_thread); + int64_t *rng_state_columnwise_ptr = + static_cast(rng_states_columnwise_tensor.data_ptr()) + i * 2; + philox_unpack(philox_args_columnwise, rng_state_columnwise_ptr); + te_rng_state_columnwise_list.push_back(makeTransformerEngineTensor( + static_cast(rng_state_columnwise_ptr), std::vector{2}, DType::kInt64)); + quant_config_columnwise_list[i].set_rng_state(te_rng_state_columnwise_list[i].data()); + quant_config_columnwise_list[i].set_stochastic_rounding(true); + } } } @@ -864,9 +893,12 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, out_columnwise_amax.shape); // RHT + NVFP4 quantize kernel + // Use separate RNG state for columnwise to ensure different random numbers than rowwise + auto &columnwise_quant_config = + need_separate_columnwise_rng ? quant_config_columnwise_list[i] : quant_config_list[i]; nvte_hadamard_transform_cast_fusion_columnwise(input_list[i].data(), out_transpose.data(), rht_matrix_nvte.data(), - quant_config_list[i], stream); + columnwise_quant_config, stream); } } }); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index d7e8912ac7..c73c09b317 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1468,17 +1468,37 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou } size_t cols = input.size(input.ndim() - 1); + // Stochastic rounding + // When both rowwise and columnwise quantization are used with RHT, + // we need separate RNG states for each to ensure they use different random numbers. TensorWrapper te_rng_state; + TensorWrapper te_rng_state_columnwise; + QuantizationConfigWrapper quant_config_columnwise; + const bool need_separate_columnwise_rng = + this->stochastic_rounding && this->with_rht && this->columnwise_usage; + if (this->stochastic_rounding) { const size_t rng_elts_per_thread = 1024; // Wild guess, probably can be tightened auto gen = at::get_generator_or_default( std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); - at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); + + // Generate RNG state for rowwise quantization + at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); auto rng_state = torch::empty({2}, opts); philox_unpack(philox_args, static_cast(rng_state.data_ptr())); te_rng_state = makeTransformerEngineTensor(rng_state); quant_config.set_rng_state(te_rng_state.data()); + + // Generate separate RNG state for columnwise quantization + if (need_separate_columnwise_rng) { + at::PhiloxCudaState philox_args_columnwise = init_philox_state(gen, rng_elts_per_thread); + auto rng_state_columnwise = torch::empty({2}, opts); + philox_unpack(philox_args_columnwise, static_cast(rng_state_columnwise.data_ptr())); + te_rng_state_columnwise = makeTransformerEngineTensor(rng_state_columnwise); + quant_config_columnwise.set_stochastic_rounding(true); + quant_config_columnwise.set_rng_state(te_rng_state_columnwise.data()); + } } // Restriction for the RHT cast fusion kernel. @@ -1605,6 +1625,10 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou static_cast(out_columnwise_amax.dtype), out_columnwise_amax.shape); + // Use separate RNG state for columnwise to ensure different random numbers than rowwise + auto& columnwise_quant_config = + need_separate_columnwise_rng ? quant_config_columnwise : quant_config; + if (!eligible_for_rht_cast_fusion) { // Invoking fallback RHT kernel. @@ -1629,7 +1653,8 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // Quantize kernel will treat everything as rowwise input/output, which is // intended. NVTE_SCOPED_GIL_RELEASE({ - nvte_quantize_v2(rht_output_t_cpp.data(), out_transpose.data(), quant_config, stream); + nvte_quantize_v2(rht_output_t_cpp.data(), out_transpose.data(), columnwise_quant_config, + stream); }); } else { // RHT cast fusion kernel. @@ -1637,8 +1662,9 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou "RHT matrix is not set"); auto rht_matrix_nvte = makeTransformerEngineTensor(this->rht_matrix); NVTE_SCOPED_GIL_RELEASE({ - nvte_hadamard_transform_cast_fusion_columnwise( - input.data(), out_transpose.data(), rht_matrix_nvte.data(), quant_config, stream); + nvte_hadamard_transform_cast_fusion_columnwise(input.data(), out_transpose.data(), + rht_matrix_nvte.data(), + columnwise_quant_config, stream); }); } } From 811e090859a94d79148c5d902f62cf8d9539af20 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Fri, 12 Dec 2025 01:57:49 +0800 Subject: [PATCH 123/521] [PyTorch] Update RNG global states in tracker set_states (#2501) set_all_rng_states in set_states Signed-off-by: Robin Zhang Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/distributed.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 5284b297e2..deb9b3ff91 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -818,6 +818,8 @@ def set_states(self, states: Dict[str, torch.Tensor]) -> None: A mapping from string names to RNG states. """ self.states_ = states + # Update global states. + set_all_rng_states(self.states_) def add(self, name: str, seed: int) -> None: """ From 50352325811043208f5b2a5edbb941ca560df2d0 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Fri, 12 Dec 2025 01:58:02 +0800 Subject: [PATCH 124/521] [PyTorch] Convert sample tuple to list in cudagraph input reuse (#2426) Convert sample tuple to list in reuse Signed-off-by: Robin Zhang Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/graph.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 1baa67414d..92826735f9 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -218,9 +218,10 @@ def _make_graphed_callables( assert ( is_training ), "`_reuse_graph_input_output_buffers` is only available in training mode." - assert isinstance( - sample_args, list - ), "sample_args must be a list for _reuse_graph_input_output_buffers." + if isinstance(sample_args, tuple): + sample_args = list(sample_args) + if isinstance(sample_kwargs, tuple): + sample_kwargs = list(sample_kwargs) # Reorganize args and kwargs for input tensor reuse. # fwd_sample_qs is keyed by model chunk index. The value is a queue of tuples. From 887a4fca9e2ab33ed0a48b81fa768e9a24d6234f Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 11 Dec 2025 10:17:26 -0800 Subject: [PATCH 125/521] [JAX] Unset NVTE_FUSED_RING_ATTENTION_USE_SCAN by default (#2503) * Unset NVTE_FUSED_RING_ATTENTION_USE_SCAN by default Signed-off-by: Kshitij Lakhani * Add TODO Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change the warning check in P2P helper to warn against using scan loop 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> --- .../jax/cpp_extensions/attention.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 7e0070dd43..c272ab6671 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -2350,7 +2350,8 @@ class _FusedAttnCPWithP2PHelper: @staticmethod def use_scanloop(): """Returns true if the implementation will use a scan loop for iteration.""" - use_scan = bool(int(os.getenv("NVTE_FUSED_RING_ATTENTION_USE_SCAN", "1"))) + # TODO(KshitijLakhani): Reset default to 1, once the extra kv permute op issue is resolved + use_scan = bool(int(os.getenv("NVTE_FUSED_RING_ATTENTION_USE_SCAN", "0"))) return use_scan def check_supported(self): @@ -2395,13 +2396,15 @@ def check_supported(self): f"{header} only supports VANILLA_SOFTMAX, got: {self.config.softmax_type}" ) - # We want to encourage use of scan loop to minimize unrolling and ensure more - # predictable scheduling from XLA. The unrolled flavor will be supported but - # not the prefered implementation. - if not self.use_scanloop(): + # TODO(KshitijLakhani): Flip the condition to check for disabled scan loop and warn + # against using unrolled loops once the scan issue is resolved. + # We want to discourage the use of scan loop as additional kv permute op observed. + # The scan loop flavor will be supported but not the prefered implementation until + # a resolution for the additional kv permute op, which degrades perf, is found. + if self.use_scanloop(): warnings.warn( - "Scan loop is disabled for fused ring attention. To enable set" - " NVTE_FUSED_RING_ATTENTION_USE_SCAN=1 in your environment" + "Scan loop is enabled for fused ring attention. To disable set" + " NVTE_FUSED_RING_ATTENTION_USE_SCAN=0 in your environment" ) # If using scanloop, idx in scan_kv_block() will be a traced device value, but From 8c9f7c25e7fe80c31559ed96e7529fe2480c3550 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 12 Dec 2025 19:14:07 +0530 Subject: [PATCH 126/521] [PyTorch] Add triton requirement (#2490) * Add triton dep Signed-off-by: Kirthi Shankar Sivamani * Fix Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Teddy Do --- build_tools/pytorch.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 93fef60f02..b03ef04fa4 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -14,7 +14,16 @@ def install_requirements() -> List[str]: """Install dependencies for TE/PyTorch extensions.""" - return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic", "nvdlfw-inspect"] + return [ + "torch>=2.1", + "einops", + "onnxscript", + "onnx", + "packaging", + "pydantic", + "nvdlfw-inspect", + "triton", + ] def test_requirements() -> List[str]: From 36f2dfd28a5a1965e59fc6fecd17d1167f25d4e1 Mon Sep 17 00:00:00 2001 From: Yashaswi Karnati <144376261+yashaswikarnati@users.noreply.github.com> Date: Mon, 15 Dec 2025 10:21:31 -0800 Subject: [PATCH 127/521] fix ce loss calculation when some tokens are ignored (#2476) * fix ce loss with ignore idx Signed-off-by: ykarnati * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: ykarnati * remove fix comments Signed-off-by: ykarnati * fallback divisor to 1 Signed-off-by: ykarnati * have arg for n_rows and n_non_ignore Signed-off-by: ykarnati * fuse n_non_ignore to softmax kernel Signed-off-by: ykarnati * fix incorrect arg Signed-off-by: ykarnati --------- Signed-off-by: ykarnati Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_parallel_cross_entropy.py | 15 ++++++++++++++- transformer_engine/common/triton/cross_entropy.py | 14 ++++++++++++-- .../pytorch/triton/cross_entropy.py | 13 +++++++++++-- 3 files changed, 37 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/test_parallel_cross_entropy.py b/tests/pytorch/test_parallel_cross_entropy.py index e325146b7e..7028dc993a 100644 --- a/tests/pytorch/test_parallel_cross_entropy.py +++ b/tests/pytorch/test_parallel_cross_entropy.py @@ -89,7 +89,7 @@ def one_iteration_test( # Check that loss and grad input match tols = dtype_tols(dtype) test_loss = test_loss.to(dtype=torch.float64, device="cpu") - ref_loss = test_loss.to(dtype=torch.float64, device="cpu") + ref_loss = ref_loss.to(dtype=torch.float64, device="cpu") ref_loss = ref_loss.reshape(test_loss.size()) test_grad_input = self.input_test.grad.to(dtype=torch.float64, device="cpu") ref_grad_input = self.input_ref.grad.to(dtype=torch.float64, device="cpu") @@ -154,3 +154,16 @@ def test_ignore_idx(self): reduce_loss=False, ignore_idx=True, ) + + def test_ignore_idx_reduced_loss(self): + """Test ignore_idx with reduce_loss=True""" + self.generate_iters(5) + self.generate_infra(True, 0) # reduce_loss=True + for i in range(self.iters): + self.one_iteration_test( + dtype=torch.float32, + swap_dim=random.choice([True, False]), + label_smoothing=0, + reduce_loss=True, + ignore_idx=True, + ) diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index fc49ac20b7..282b23cda8 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -18,6 +18,8 @@ def online_softmax_kernel( m_d_X_y_stride, rank, n_cols, + ignore_idx, + n_non_ignore, BLOCK_SIZE: tl.constexpr, ): """ @@ -32,6 +34,8 @@ def online_softmax_kernel( m_d_X_y_stride (int): The stride of the m/d/X_y tensor. rank (int): The rank of this device in the TP group. n_cols (int): The number of columns in the input tensor. + ignore_idx (int): The index to ignore for loss calculation. + n_non_ignore: The number of non-ignored elements in the batch. BLOCK_SIZE (int): The block size for Triton operations. """ @@ -44,6 +48,9 @@ def online_softmax_kernel( Y_ptr += program_id * Y_stride y = tl.load(Y_ptr) + if y != ignore_idx: + tl.atomic_add(n_non_ignore, 1) + vocab_start_idx = rank * n_cols vocab_end_idx = (rank + 1) * n_cols if y >= vocab_start_idx: @@ -89,6 +96,7 @@ def cross_entropy_kernel( world_size, ignore_idx, n_cols, + n_rows, n_non_ignore, reduce_loss: tl.constexpr, label_smoothing: tl.constexpr, @@ -110,12 +118,14 @@ def cross_entropy_kernel( world_size (int): The size of world involved in this distributed loss calculation. ignore_idx (int): Tokens to be ignored for loss and gradient calculation. n_cols (int): The number of columns in the input tensor. - n_non_ignore (int): The number of non-ignored elements in the batch. + n_rows (int): The number of rows in the batch (B * SQ), used for buffer indexing. + n_non_ignore: The number of non-ignored elements in the batch. label_smoothing (float): The amount of smoothing when computing the loss, where 0.0 means no smoothing. BLOCK_SIZE (int): The block size for Triton operations. """ program_id = tl.program_id(0).to(tl.int64) + n_non_ignore = tl.load(n_non_ignore) # locate the start index X_ptr += program_id * X_stride @@ -140,7 +150,7 @@ def cross_entropy_kernel( ori_X_y = tl.load(m_d_X_y_ptr + (2 * m_d_X_y_stride)) for i in range(1, world_size): - offset = i * 3 * n_non_ignore * m_d_X_y_stride + offset = i * 3 * n_rows * m_d_X_y_stride access_ptr = m_d_X_y_ptr + offset m_new = tl.load(access_ptr) d_new = tl.load(access_ptr + m_d_X_y_stride) diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index 1a5756105a..e6d0397673 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -46,6 +46,8 @@ def cross_entropy_forward( # tensor to hold this rank's m/d/X_y values m_d_X_y = torch.zeros(n_rows * 3, dtype=torch.float32, device=_input.device) + n_non_ignore = torch.zeros(1, dtype=torch.int64, device=_input.device) + # ensure _input and target are contiguous in the last dimension if _input.stride(-1) != 1: _input = _input.contiguous() @@ -63,10 +65,14 @@ def cross_entropy_forward( m_d_X_y_stride=m_d_X_y.stride(-1), rank=rank, n_cols=V, + ignore_idx=ignore_idx, + n_non_ignore=n_non_ignore, BLOCK_SIZE=BLOCK_SIZE, num_warps=32, ) + n_non_ignore = torch.clamp(n_non_ignore, min=1) + world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group) if world_size > 1: @@ -90,14 +96,17 @@ def cross_entropy_forward( world_size=world_size, ignore_idx=ignore_idx, n_cols=V, - n_non_ignore=n_rows, + n_rows=n_rows, + n_non_ignore=n_non_ignore, reduce_loss=reduce_loss, label_smoothing=label_smoothing, BLOCK_SIZE=BLOCK_SIZE, num_warps=32, ) - loss = torch.reshape(loss_1d, (B, SQ)) if not reduce_loss else (torch.sum(loss_1d) / n_rows) + loss = ( + torch.reshape(loss_1d, (B, SQ)) if not reduce_loss else (torch.sum(loss_1d) / n_non_ignore) + ) return loss, _input From b215116a506c5f80aeb882e87aeda1b72c236879 Mon Sep 17 00:00:00 2001 From: kwyss-nvidia Date: Mon, 15 Dec 2025 15:03:38 -0800 Subject: [PATCH 128/521] Check calling convention for amax switch. (#2506) * Check calling convention for amax switch. Wgrad gemms with colwise x colwise require rowwise data via general_gemm. Since dy has both for dgrad and wgrad, the brittleness has likely not affected results. Signed-off-by: Keith Wyss * Clear rowwise data when applicable. Signed-off-by: Keith Wyss * Update test with columnwise cases. Signed-off-by: Keith Wyss * Check enum value rather than implicit cast. Signed-off-by: Keith Wyss --------- Signed-off-by: Keith Wyss --- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 23 ++++++++++++++----- .../common/gemm/cublaslt_gemm.cu | 10 +++++--- 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 6009643ffa..9f860551d0 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -122,8 +122,15 @@ def check_nvfp4_gemm_versus_reference( ) # Create reference quantized tensors needed by reference GEMM - x_nvfp4_ref = ref_quantizer.quantize(x) - w_nvfp4_ref = ref_quantizer.quantize(w) + # Reference GEMM is only rowwise. + if x_columnwise: + x_nvfp4_ref = ref_quantizer.quantize(x.t().contiguous()) + else: + x_nvfp4_ref = ref_quantizer.quantize(x) + if w_columnwise: + w_nvfp4_ref = ref_quantizer.quantize(w.t().contiguous()) + else: + w_nvfp4_ref = ref_quantizer.quantize(w) # Reference GEMM using quantizer's qgemm method y_ref = ref_quantizer.qgemm( @@ -155,6 +162,10 @@ def check_nvfp4_gemm_versus_reference( use_grad = False use_split_accumulator = False + if x_columnwise: + 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. @@ -212,11 +223,11 @@ def check_nvfp4_gemm_versus_reference( @pytest.mark.parametrize( "is_x_columnwise, is_w_columnwise", [ - (False, False), # Only rowwise x rowwise is supported by reference GEMM - # Note: Reference GEMM expects inputs as (M,K) x (N,K) with rowwise quantization - # Columnwise layouts are not supported by the reference implementation + (False, False), # TN + (True, False), # NN + (True, True), # NT ], - ids=["rowxrow"], + ids=["rowxrow", "colxrow", "colxcol"], ) def test_nvfp4_gemm_versus_reference( M: int, diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 97e8ec9a3e..118bf19335 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -363,7 +363,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, // TODO: Check whether scales are on CPU/GPU or add API to control. // Currently scales are assumed to be on CPU when amax is provided // and on GPU when not provided, but this is brittle. - if (use_fp4 && (inputA->amax.dptr != nullptr || inputB->amax.dptr != nullptr)) { + if (use_fp4 && + ((transa == CUBLAS_OP_T ? inputA->amax.dptr : inputA->columnwise_amax.dptr) != nullptr || + (transb == CUBLAS_OP_T ? inputB->columnwise_amax.dptr : inputB->amax.dptr) != nullptr)) { // Reserve some workspace for alpha scale NVTE_CHECK(workspaceSize >= 4, "NVFP4 GEMM requires at least 4 byte workspace for alpha scale, but only has ", @@ -378,8 +380,10 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, // tensor scales in matmul output, instead of in matmul inputs. float old_alpha = *reinterpret_cast(alpha); // Assumed to be on CPU TensorWrapper new_alpha_tensor(new_alpha_ptr, std::vector{1}, DType::kFloat32); - nvte_nvfp4_compute_per_tensor_scale(inputA->nvte_tensor, transa, inputB->nvte_tensor, !transb, - old_alpha, new_alpha_tensor.data(), stream); + bool a_rowwise_amax = transa == CUBLAS_OP_T; + bool b_rowwise_amax = transb != CUBLAS_OP_T; + nvte_nvfp4_compute_per_tensor_scale(inputA->nvte_tensor, a_rowwise_amax, inputB->nvte_tensor, + b_rowwise_amax, old_alpha, new_alpha_tensor.data(), stream); alpha = new_alpha_ptr; // Make sure beta scale is on device From 2886cbce11aeb189b7e90374d8c3297ff4bbeaea Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 16 Dec 2025 00:33:37 +0100 Subject: [PATCH 129/521] [PyTorch debug] Fix test for debug tools (#2507) * Skip delayed wgrad tests in distributed numerics when debug mode is enabled Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- qa/L1_pytorch_distributed_unittest/test.sh | 2 +- tests/pytorch/distributed/run_numerics.py | 9 ++++++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index e0abdd281b..b1e3a3e15c 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -44,7 +44,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_ pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_distributed.xml $TE_PATH/tests/pytorch/debug/test_distributed.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS || test_fail "debug test_distributed.py" # standard numerics tests with initialized debug -NVTE_TEST_NVINSPECT_ENABLED=True NVTE_TEST_NVINSPECT_CONFIG_FILE=$NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE NVTE_TEST_NVINSPECT_FEATURE_DIRS=$NVTE_TEST_NVINSPECT_FEATURE_DIRS pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_2.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "debug test_numerics.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 pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_2.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "debug test_numerics.py" if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index c109f463d8..6cad80fde7 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -38,8 +38,9 @@ NCCL_WORLD = None LOSS_FN = nn.MSELoss() QUANTIZATION = None +NVTE_TEST_NVINSPECT_ENABLED = int(os.environ.get("NVTE_TEST_NVINSPECT_ENABLED") or "0") -if os.environ.get("NVTE_TEST_NVINSPECT_ENABLED", False): +if NVTE_TEST_NVINSPECT_ENABLED: # The numerics of all the layers should work the same, # when debug=True. I fed them with dummy feature # to prevent switching off debug, which can happen if @@ -745,6 +746,8 @@ def test_linear(): for kwargs in kwargs_list: if kwargs.get("save_original_input", False) and QUANTIZATION == "fp8": continue + if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: + continue for parallel_mode in ["column", "row"]: for sequence_parallel in [False, True]: _test_linear(parallel_mode, sequence_parallel, **kwargs) @@ -924,6 +927,8 @@ def test_layernorm_linear(): ] for kwargs in kwargs_list: + if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: + continue for parallel_mode in ["column"]: for sequence_parallel in [False, True]: _test_layernorm_linear(parallel_mode, sequence_parallel, **kwargs) @@ -1034,6 +1039,8 @@ def test_layernorm_mlp(): ] for kwargs in kwargs_list: + if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED: + continue for set_parallel_mode in [True]: for sequence_parallel in [False, True]: _test_layernorm_mlp(set_parallel_mode, sequence_parallel, **kwargs) From eac8af6a42f265e81f8ebdbcc8932c7d4ea6b757 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Tue, 16 Dec 2025 13:24:24 -0800 Subject: [PATCH 130/521] Remove test skip logic for GEMM-AR tests (#2516) * Use GEMM-AR fallback on newer cuBLASMp Signed-off-by: Vladimir Cherepanov * Remove test skip logic completely Signed-off-by: Vladimir Cherepanov --------- Signed-off-by: Vladimir Cherepanov --- tests/cpp_distributed/test_comm_gemm.cu | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index 884faa4748..0a20aa1cea 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -63,12 +63,6 @@ int main(int argc, char* argv[]) { return ret; } -bool IsMulticastSupported(int device_id) { - int supported = 0; - CHECK_CU(cuDeviceGetAttribute(&supported, CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, device_id)); - return supported; -} - int GetDeviceComputeCapability(int device_id) { int major{}; int minor{}; @@ -369,11 +363,6 @@ struct GemmAr : public CommGemmFixure { nvte_gemm_all_reduce(ctx_, m, n, k, a, b, d, bias, pre_act_out, transa, transb, grad, accumulate, comm_sm_count, stream, kNVTECommGemmAlgoDefault); } - - void SetUp() override { - if (!IsMulticastSupported(rank_)) - GTEST_SKIP() << "Multicast is not supported on device " << rank_; - } }; TEST_P(AgGemm, Gemm) { From dbd0197e7c1e4a8a82fcad9d87426c1a62147508 Mon Sep 17 00:00:00 2001 From: Jinhang Choi Date: Tue, 16 Dec 2025 20:23:18 -0800 Subject: [PATCH 131/521] Reset cache logic of weight workspace for NVFP4TensorStorage (#2524) reset weight ws cache for NVFP4TensorStorage Signed-off-by: Jinhang Choi --- transformer_engine/pytorch/module/base.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index acf9233281..ab7cd9ab47 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -45,6 +45,7 @@ from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage +from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..utils import ( is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, @@ -1388,6 +1389,11 @@ def get_weight_workspace( reset_cache = True elif quantizer.columnwise_usage and out._columnwise_data is None: reset_cache = True + elif isinstance(out, NVFP4TensorStorage): + if quantizer.rowwise_usage and out._rowwise_data is None: + reset_cache = True + elif quantizer.columnwise_usage and out._columnwise_data is None: + reset_cache = True if isinstance(out, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer): reset_cache = True if reset_cache: From 5c2f2ff56089d44df49fe675ce4b3c6929b7f607 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 17 Dec 2025 08:14:52 -0800 Subject: [PATCH 132/521] Add ccache support to TE and use it in GitHub actions (#2444) * Add ccache support to TE and use it in GitHub actions Signed-off-by: Przemek Tredak * Move to allowed action with sccache Signed-off-by: Przemek Tredak * Properly handle sccache Signed-off-by: Przemek Tredak * Fix typo Signed-off-by: Przemek Tredak * Removing ccache from the custom docker workflows where we can't run the action in the container Signed-off-by: Przemek Tredak * JAX already uses same cmake options to build the extension so there is no need to set CXX too Signed-off-by: Przemek Tredak * Removed the unnecessary env variables Signed-off-by: Przemek Tredak --------- Signed-off-by: Przemek Tredak --- .github/workflows/build.yml | 12 +++++++++--- build_tools/build_ext.py | 6 ++++++ 2 files changed, 15 insertions(+), 3 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 51036e40bd..2ed8766de0 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -24,11 +24,14 @@ jobs: uses: actions/checkout@v3 with: submodules: recursive + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad - name: 'Build' - run: pip install --no-build-isolation . -v + run: NVTE_USE_CCACHE=1 NVTE_CCACHE_BIN=sccache pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: none MAX_JOBS: 1 + SCCACHE_GHA_ENABLED: "true" - name: 'Sanity check' run: python3 -c "import transformer_engine" working-directory: / @@ -94,11 +97,15 @@ jobs: uses: actions/checkout@v3 with: submodules: recursive + - name: ccache + uses: mozilla-actions/sccache-action@7d986dd989559c6ecdb630a3fd2557667be217ad - name: 'Build' - run: pip install --no-build-isolation . -v + run: | + NVTE_CCACHE_BIN=sccache NVTE_USE_CCACHE=1 pip install --no-build-isolation . -v env: NVTE_FRAMEWORK: jax MAX_JOBS: 1 + SCCACHE_GHA_ENABLED: "true" - name: 'Sanity check' run: python3 tests/jax/test_sanity_import.py all: @@ -140,7 +147,6 @@ jobs: pip install pybind11[global] einops onnxscript && \ pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 ' - - name: 'Build' run: docker exec builder bash -c 'pip install --no-cache-dir --no-build-isolation . -v --no-deps' env: diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index 349858ac49..c269a29874 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -61,6 +61,12 @@ def _build_cmake(self, build_dir: Path, install_dir: Path) -> None: f"-DCMAKE_BUILD_TYPE={build_type}", f"-DCMAKE_INSTALL_PREFIX={install_dir}", ] + if bool(int(os.getenv("NVTE_USE_CCACHE", "0"))): + ccache_bin = os.getenv("NVTE_CCACHE_BIN", "ccache") + configure_command += [ + f"-DCMAKE_CXX_COMPILER_LAUNCHER={ccache_bin}", + f"-DCMAKE_CUDA_COMPILER_LAUNCHER={ccache_bin}", + ] configure_command += self.cmake_flags import pybind11 From 442513c59686672549419f24aaa2c5b40106ee92 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 17 Dec 2025 11:59:55 -0800 Subject: [PATCH 133/521] [JAX] Add tutorial for integrating TE/JAX quantization into an existing framework (#2423) * Tutorial for integration te/jax quantization into an existing framework Signed-off-by: Jeremy Berchtold * add todos Signed-off-by: Jeremy Berchtold * support nvfp4 sr rng key, move wrapper module into TE itself, fix bfloat16 cast Signed-off-by: Jeremy Berchtold * update docstrings Signed-off-by: Jeremy Berchtold * Fix QKV proj and out proj in Flax example transformer layer Signed-off-by: Jeremy Berchtold * Use fused attention in quickstart_jax example Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remat policy Signed-off-by: Jeremy Berchtold * add tutorial to docs Signed-off-by: Jeremy Berchtold * update title Signed-off-by: Jeremy Berchtold * remove unused dtype from TE DPA module Signed-off-by: Jeremy Berchtold * Fix notebook title Signed-off-by: Jeremy Berchtold * Fix lint Signed-off-by: Jeremy Berchtold * Add explanation of flax module wrapper Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/examples/quickstart_jax.ipynb | 161 +++--- docs/examples/quickstart_jax_utils.py | 35 +- docs/examples/te_jax_integration.ipynb | 462 ++++++++++++++++++ docs/index.rst | 1 + .../jax/cpp_extensions/attention.py | 4 +- transformer_engine/jax/flax/__init__.py | 3 + transformer_engine/jax/flax/module.py | 84 ++++ transformer_engine/jax/flax/transformer.py | 39 +- transformer_engine/jax/sharding.py | 8 + 9 files changed, 688 insertions(+), 109 deletions(-) create mode 100644 docs/examples/te_jax_integration.ipynb diff --git a/docs/examples/quickstart_jax.ipynb b/docs/examples/quickstart_jax.ipynb index 369cc371f6..5b4e439c00 100644 --- a/docs/examples/quickstart_jax.ipynb +++ b/docs/examples/quickstart_jax.ipynb @@ -53,7 +53,7 @@ }, { "cell_type": "code", - "execution_count": 26, + "execution_count": null, "id": "d5284a38", "metadata": {}, "outputs": [], @@ -67,7 +67,7 @@ }, { "cell_type": "code", - "execution_count": 27, + "execution_count": 2, "id": "a4d1cfdc", "metadata": {}, "outputs": [], @@ -142,6 +142,9 @@ " \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)(x)\n", " \n", " x = res + x\n", " \n", @@ -171,7 +174,7 @@ }, { "cell_type": "code", - "execution_count": 28, + "execution_count": 3, "id": "8b44649d", "metadata": {}, "outputs": [], @@ -192,7 +195,7 @@ }, { "cell_type": "code", - "execution_count": 29, + "execution_count": 4, "id": "e44ed26d", "metadata": {}, "outputs": [ @@ -201,7 +204,7 @@ "output_type": "stream", "text": [ "Pure Flax FlaxTransformerLayer initialized successfully!\n", - "Parameter shapes: {'params': {'Dense_0': {'bias': (12288,), 'kernel': (4096, 12288)}, '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" + "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" ] } ], @@ -222,7 +225,7 @@ }, { "cell_type": "code", - "execution_count": 30, + "execution_count": 5, "id": "de91af7a", "metadata": {}, "outputs": [ @@ -248,7 +251,7 @@ }, { "cell_type": "code", - "execution_count": 31, + "execution_count": 6, "id": "037bc8d9", "metadata": {}, "outputs": [ @@ -256,7 +259,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 18.546080589294434 ms\n" + "Mean time: 19.258604049682617 ms\n" ] } ], @@ -270,8 +273,8 @@ " variables=params,\n", " input=x,\n", " output_grad=dy,\n", - " dropout_key=dropout_key,\n", " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + " rngs={\"dropout\": dropout_key},\n", ")" ] }, @@ -304,7 +307,7 @@ }, { "cell_type": "code", - "execution_count": 32, + "execution_count": 7, "id": "bed20d6b", "metadata": {}, "outputs": [], @@ -323,14 +326,11 @@ }, { "cell_type": "code", - "execution_count": 33, + "execution_count": 8, "id": "56105579", "metadata": {}, "outputs": [], "source": [ - "from transformer_engine.jax.flax.transformer import DotProductAttention as TEDotProductAttention\n", - "\n", - "\n", "class TEUnfusedMLP(nn.Module):\n", " hidden_size : int\n", " ffn_hidden_size: int\n", @@ -360,11 +360,7 @@ " 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", + " ) -> jnp.ndarray: \n", " res = x\n", " x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x)\n", "\n", @@ -376,14 +372,19 @@ " # Attention - either TE or Flax implementation\n", " if self.use_te_attention:\n", " # Use TE's DotProductAttention\n", - " attention = TEDotProductAttention(\n", + " attention = te_flax.DotProductAttention(\n", " head_dim=self.kv_channels,\n", " num_attention_heads=self.num_attention_heads,\n", " num_gqa_groups=self.num_attention_heads, # No GQA\n", " attention_dropout=self.attention_dropout,\n", " attn_mask_type='causal',\n", " )\n", - " x = attention(q, k, v, sequence_descriptor=None, deterministic=deterministic)\n", + " x = attention(\n", + " q, k, v,\n", + " # Causal mask does not need an explicit instatiated mask as specialized kernels exist to handle it\n", + " sequence_descriptor=None, \n", + " deterministic=deterministic\n", + " )\n", " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", " x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))\n", " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x)\n", @@ -393,6 +394,10 @@ " q_reshaped = q.reshape(q.shape[0], q.shape[1], self.hidden_size)\n", " k_reshaped = k.reshape(k.shape[0], k.shape[1], self.hidden_size)\n", " v_reshaped = v.reshape(v.shape[0], v.shape[1], self.hidden_size)\n", + "\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", " attention = nn.MultiHeadDotProductAttention(\n", " num_heads=self.num_attention_heads,\n", @@ -428,7 +433,7 @@ }, { "cell_type": "code", - "execution_count": 34, + "execution_count": 9, "id": "4b67511f", "metadata": {}, "outputs": [ @@ -436,7 +441,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 16.375374794006348 ms\n" + "Mean time: 16.003193855285645 ms\n" ] } ], @@ -455,8 +460,8 @@ " variables=te_params, # Ensure the correct `params` is passed\n", " input=x,\n", " output_grad=dy,\n", - " dropout_key=dropout_key,\n", " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + " rngs={\"dropout\": dropout_key},\n", ")\n" ] }, @@ -470,7 +475,7 @@ }, { "cell_type": "code", - "execution_count": 35, + "execution_count": 10, "id": "5146cd99", "metadata": {}, "outputs": [ @@ -478,23 +483,9 @@ "name": "stderr", "output_type": "stream", "text": [ - "/code/github/TransformerEngine/transformer_engine/jax/flax/transformer.py:634: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", " warnings.warn(\n", - "/code/github/TransformerEngine/transformer_engine/jax/flax/transformer.py:742: UserWarning: Fused attention is not enabled because there is no available kernel.\n", - "Fall back to the unfused attention.\n", - "Please try to update the cuDNN and TE to the latest version.\n", - "self.dtype=\n", - "qkv_layout=>\n", - "attn_bias_type=>\n", - "attn_mask_type=>\n", - "self.attention_dropout=0.1\n", - "self.num_attention_heads=32\n", - "self.num_gqa_groups=32\n", - "seqlen_q=2048\n", - "seqlen_kv=2048\n", - "head_dim_qk=128\n", - "head_dim_v=128\n", - "\n", + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", " warnings.warn(\n" ] }, @@ -502,7 +493,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 12.403340339660645 ms\n" + "Mean time: 8.897695541381836 ms\n" ] } ], @@ -520,8 +511,8 @@ " variables=te_params, # Ensure the correct `params` is passed\n", " input=x,\n", " output_grad=dy,\n", - " dropout_key=dropout_key,\n", " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", + " rngs={\"dropout\": dropout_key},\n", ")" ] }, @@ -553,7 +544,7 @@ }, { "cell_type": "code", - "execution_count": 36, + "execution_count": 11, "id": "c2eee376", "metadata": {}, "outputs": [], @@ -565,15 +556,27 @@ }, { "cell_type": "code", - "execution_count": 37, + "execution_count": 12, "id": "de96827c", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n", + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n", + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 9.396424293518066 ms\n" + "Mean time: 5.651178359985352 ms\n" ] } ], @@ -589,9 +592,9 @@ " variables=te_unfused_params, # Ensure the correct `params` is passed\n", " input=x,\n", " output_grad=dy,\n", - " dropout_key=dropout_key,\n", " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe}\n", + " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe},\n", + " rngs={\"dropout\": dropout_key},\n", ")" ] }, @@ -626,7 +629,7 @@ }, { "cell_type": "code", - "execution_count": 38, + "execution_count": 13, "id": "11203785", "metadata": {}, "outputs": [], @@ -659,7 +662,7 @@ " q, k, v = jnp.split(qkv, 3, axis=3)\n", "\n", " # Attention using TE's DotProductAttention\n", - " attention = TEDotProductAttention(\n", + " attention = te_flax.DotProductAttention(\n", " head_dim=self.kv_channels,\n", " num_attention_heads=self.num_attention_heads,\n", " num_gqa_groups=self.num_attention_heads, \n", @@ -697,15 +700,27 @@ }, { "cell_type": "code", - "execution_count": 39, + "execution_count": 14, "id": "6b0c705e", "metadata": {}, "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n", + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n", + "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", + " warnings.warn(\n" + ] + }, { "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 9.145426750183105 ms\n" + "Mean time: 5.493879318237305 ms\n" ] } ], @@ -726,9 +741,9 @@ " variables=te_fused_params,\n", " input=x,\n", " output_grad=dy,\n", - " dropout_key=dropout_key,\n", " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe}\n", + " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe},\n", + " rngs={\"dropout\": dropout_key},\n", ")" ] }, @@ -742,35 +757,11 @@ }, { "cell_type": "code", - "execution_count": 40, + "execution_count": 15, "id": "b2aaa8ef", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/code/github/TransformerEngine/transformer_engine/jax/flax/transformer.py:742: UserWarning: Fused attention is not enabled because there is no available kernel.\n", - "Fall back to the unfused attention.\n", - "Please try to update the cuDNN and TE to the latest version.\n", - "self.dtype=\n", - "qkv_layout=>\n", - "attn_bias_type=>\n", - "attn_mask_type=>\n", - "self.attention_dropout=0.1\n", - "self.num_attention_heads=32\n", - "self.num_gqa_groups=32\n", - "seqlen_q=2048\n", - "seqlen_kv=2048\n", - "head_dim_qk=128\n", - "head_dim_v=128\n", - "\n", - " warnings.warn(\n" - ] - } - ], + "outputs": [], "source": [ - "\n", "te_transformer = te_flax.TransformerLayer(\n", " hidden_size=hidden_size,\n", " mlp_hidden_size=ffn_hidden_size, \n", @@ -782,7 +773,7 @@ " intermediate_dropout=0.0,\n", " enable_relative_embedding=False,\n", " self_attn_bias_type='no_bias',\n", - " hidden_dropout=0.0\n", + " hidden_dropout=0.0,\n", ")\n", "\n", "with te.autocast(enabled=True, recipe=fp8_recipe):\n", @@ -792,7 +783,7 @@ }, { "cell_type": "code", - "execution_count": 41, + "execution_count": 20, "id": "b9cdbf22", "metadata": {}, "outputs": [ @@ -800,7 +791,7 @@ "name": "stdout", "output_type": "stream", "text": [ - "Mean time: 9.020795822143555 ms\n" + "Mean time: 5.334172248840332 ms\n" ] } ], @@ -811,9 +802,9 @@ " variables=te_transformer_params,\n", " input=x,\n", " output_grad=dy,\n", - " dropout_key=dropout_key,\n", " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe }\n", + " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe },\n", + " rngs={\"dropout\": dropout_key},\n", ")" ] } diff --git a/docs/examples/quickstart_jax_utils.py b/docs/examples/quickstart_jax_utils.py index 138427338d..f1ff9b7d99 100644 --- a/docs/examples/quickstart_jax_utils.py +++ b/docs/examples/quickstart_jax_utils.py @@ -19,12 +19,12 @@ def speedometer( variables: Any, input: jnp.ndarray, output_grad: jnp.ndarray, - dropout_key: jax.random.PRNGKey, model_init_fn: Callable = None, forward_kwargs: dict = {}, autocast_kwargs: Optional[dict] = None, timing_iters: int = 50, warmup_iters: int = 50, + rngs: Dict[str, jax.random.PRNGKey] = None, ) -> None: """Measure average runtime for a JAX module Perform forward and backward passes . @@ -33,19 +33,21 @@ def speedometer( autocast_kwargs = {"enabled": False} model_init_fn = None + if rngs is None: + rngs = {} + train_step_fn = create_train_step_fn(model_apply_fn, autocast_kwargs, forward_kwargs) # Warm up runs - key = dropout_key for _ in range(warmup_iters): - key, step_key = jax.random.split(key) - loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_key) + rngs, step_rngs = _split_step_rngs(rngs) + loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs) # Timing runs start = time.time() for _ in range(timing_iters): - key, step_key = jax.random.split(key) - loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_key) + rngs, step_rngs = _split_step_rngs(rngs) + loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs) end = time.time() print(f"Mean time: {(end - start) * 1000 / timing_iters} ms") @@ -63,8 +65,12 @@ def create_train_step_fn( if forward_kwargs is None: forward_kwargs = {} - def loss_fn(variables: Any, inp: jnp.ndarray, grad_target: jnp.ndarray, dropout_key): - rngs = {"dropout": dropout_key} + def loss_fn( + variables: Any, + inp: jnp.ndarray, + grad_target: jnp.ndarray, + rngs: Dict[str, jax.random.PRNGKey], + ): with te.autocast(**autocast_kwargs): # Forward Pass: Apply the model using current parameters and variables call_kwargs = {**forward_kwargs, "rngs": rngs} @@ -84,3 +90,16 @@ def fwd_bwd_fn(*args, **kwargs): # JIT-compile the fwd_bwd_fn return jax.jit(fwd_bwd_fn) + + +def _split_step_rngs( + rngs: Dict[str, jax.random.PRNGKey], +) -> Tuple[Dict[str, jax.random.PRNGKey], Dict[str, jax.random.PRNGKey]]: + """Splits each RNG in the rngs dictionary for a new step.""" + step_rngs = {} + new_rngs = {} + for name, key in rngs.items(): + new_key, step_key = jax.random.split(key) + new_rngs[name] = new_key + step_rngs[name] = step_key + return new_rngs, step_rngs diff --git a/docs/examples/te_jax_integration.ipynb b/docs/examples/te_jax_integration.ipynb new file mode 100644 index 0000000000..70647e421a --- /dev/null +++ b/docs/examples/te_jax_integration.ipynb @@ -0,0 +1,462 @@ +{ + "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/index.rst b/docs/index.rst index 37d21c2a5d..7a3ab9f6fd 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -48,6 +48,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 .. toctree:: :hidden: diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index c272ab6671..ef921c2762 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -26,7 +26,7 @@ CPStrategy, SequenceDescriptor, ) -from ..sharding import with_sharding_constraint_by_logical_axes, HEAD_AXES +from ..sharding import with_sharding_constraint_by_logical_axes, HEAD_AXES, is_mesh_available from .base import BasePrimitive, register_primitive from .misc import ( @@ -3288,7 +3288,7 @@ def compute(config): def _maybe_context_parallel_axis(cp_axis: str): - if not cp_axis: + if not cp_axis and is_mesh_available(): gmr = global_mesh_resource() if gmr is not None: cp_axis = gmr.cp_resource diff --git a/transformer_engine/jax/flax/__init__.py b/transformer_engine/jax/flax/__init__.py index a40ccc500f..d1a9cb47f8 100644 --- a/transformer_engine/jax/flax/__init__.py +++ b/transformer_engine/jax/flax/__init__.py @@ -4,6 +4,7 @@ """Transformer Engine bindings for JAX""" from .module import DenseGeneral, LayerNorm from .module import LayerNormDenseGeneral, LayerNormMLP +from .module import wrap_function_in_te_state_module, make_dot_general_cls from .transformer import extend_logical_axis_rules from .transformer import DotProductAttention, MultiHeadAttention, RelativePositionBiases from .transformer import TransformerLayer, TransformerLayerType @@ -13,6 +14,8 @@ "LayerNorm", "LayerNormDenseGeneral", "LayerNormMLP", + "wrap_function_in_te_state_module", + "make_dot_general_cls", "extend_logical_axis_rules", "DotProductAttention", "MultiHeadAttention", diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 58df85fa52..dcfb812896 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -1354,3 +1354,87 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args): assert out.dtype == input_dtype return out, ln_output # Output, layer_norm_output + + +def wrap_function_in_te_state_module(f, quantization_recipe, name: Optional[str] = None): + """Wraps the given function `f` to support TransformerEngine quantization. + + This method does a couple things: + + 1. Wraps the given function in a Flax linen module. This module does not store any Flax parameters + but can store Flax variables for quantizers if required by the recipe. + + 2. When the wrapper is called, it provides an additional argument to the given function `f`, 'generate_quantizer_set' as the first argument. 'generate_quantizer_set' is a function that can be called to generate a TransformerEngine/JAX quantizer set object used in TransformerEngine/JAX APIs. 'generate_quantizer_set' will generate quantizers based on the recipe of this TransformerEngineQuantizer object. + + Args: + f: The function to wrap. The first argument must be 'generate_quantizer_set'. + name: The name of this wrapped operation. If unspecified, will use `f.__name__`. + + Returns: + A Flax linen module that wraps the given function. + """ + + import transformer_engine.jax as te + + class TEWrapper(te.flax.module.TransformerEngineBase): + """Wrapper Flax module for TransformerEngine quantization support.""" + + def generate_quantizer_set(self, postfix: str = ""): + OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient" + return super().generate_quantizer_set( + postfix=postfix, + variable_collection=OVERWRITE_WITH_GRADIENT, + fp8_recipe=quantization_recipe, + ) + + @nn.compact + def __call__(self, *args, **kwargs): + return f(self.generate_quantizer_set, *args, **kwargs) + + TEWrapper.__name__ = f"TEWrapper_{name if name else f.__name__}" + + return TEWrapper + + +def make_dot_general_cls(quantization_recipe): + """Creates a Flax module class that performs a dot_general operation with the arguments x and kernel using the given quantization recipe. + + This is intended for usage when you already have model parameters initialized and sharded for the kernel weights and you want to replace the GEMM implementation with TE's quantized GEMM using a given recipe. + + For example, + ``` + te_dot_general_cls = make_dot_general_cls(DelayedScaling()) + dense = nn.Dense(..., dot_general=te_dot_general_cls()) + ``` + + If you would like a drop-in replacement for nn.Dense that manages the model weights itself, please use TE's DenseGeneral module. + + Args: + quantization_recipe: The quantization recipe to use for the dot_general operation. + Returns: + A Flax module class that performs a dot_general operation with the given quantization recipe. + """ + import transformer_engine.jax as te + from transformer_engine.common.recipe import NVFP4BlockScaling + + def te_dot_general(generate_quantizer_set, x, kernel, dims, **kwargs): + """Performs a dot_general operation using TransformerEngine with quantization.""" + del kwargs # Unused + contracting_dims, batch_dims = dims + assert batch_dims == ((), ()), "Batch dimensions must be empty for TransformerEngine dot." + + quantizer_set = generate_quantizer_set() + + if isinstance(quantization_recipe, NVFP4BlockScaling): + # NVFP4 RHT requires inputs to be in bfloat16 + x = x.astype(jnp.bfloat16) + kernel = kernel.astype(jnp.bfloat16) + + return te.dense.dense( + x, + kernel, + contracting_dims=contracting_dims, + quantizer_set=quantizer_set, + ) + + return wrap_function_in_te_state_module(te_dot_general, quantization_recipe, "dot_general") diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index d0190f54c5..3c1d8ef9eb 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -121,7 +121,6 @@ class _UnfusedDotProductAttention(nn.Module): # pylint: disable=too-few-public- attention_dropout: float = 0.0 attn_mask_type: AttnMaskType = AttnMaskType.CAUSAL_MASK attn_bias_type: Optional[AttnBiasType] = None - dtype: DType = jnp.float32 float32_logits: bool = False scale_factor: Optional[float] = None transpose_batch_sequence: bool = False @@ -294,7 +293,6 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me attention_dropout: float = 0.0 attn_mask_type: AttnMaskType = AttnMaskType.CAUSAL_MASK attn_bias_type: Optional[AttnBiasType] = None - dtype: DType = jnp.float32 qkv_layout: QKVLayout = QKVLayout.BSHD_BSHD_BSHD scale_factor: Optional[float] = None transpose_batch_sequence: bool = False @@ -600,11 +598,6 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention (``'zero sink'`` and ``'learnable sink'``). - - Optimization parameters - ----------------------- - dtype: jax.numpy.dtype, default = jax.numpy.float32 - The data type used to allocate the initial parameters. """ head_dim: int @@ -613,7 +606,6 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods attention_dropout: float = 0.0 attn_mask_type: AttnMaskType = "causal" attn_bias_type: AttnBiasType = None - dtype: DType = jnp.float32 dropout_rng_name: str = "dropout" float32_logits: bool = False qkv_layout: str = "bshd_bshd_bshd" @@ -638,6 +630,24 @@ def __post_init__(self): self.transpose_batch_sequence = False super().__post_init__() + def _assert_dtypes(self, query: Array, key: Array, value: Array, qkv_layout: QKVLayout): + """Asserts that the dtypes of query, key, and value dtypes are consistent.""" + if qkv_layout.is_qkvpacked(): + pass # No need to check dtypes for key and value since it is packed + elif qkv_layout.is_kvpacked(): + assert ( + key.dtype == query.dtype + ), f"Expected kv dtype={key.dtype} to match query dtype={query.dtype}." + elif qkv_layout.is_separate(): + assert ( + key.dtype == query.dtype + ), f"Expected key dtype={key.dtype} to match query dtype={query.dtype}." + assert ( + value.dtype == query.dtype + ), f"Expected value dtype={value.dtype} to match query dtype={query.dtype}." + else: + raise ValueError(f"Unsupported {qkv_layout=}.") + @nn.compact def __call__( self, @@ -700,6 +710,9 @@ def __call__( assert bias is None else: assert bias is not None + bias = bias.astype(input_dtype) + + self._assert_dtypes(query, key, value, qkv_layout) # Use fused attn (if kernel check below passes) by default enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) @@ -720,8 +733,9 @@ def __call__( has_fused_attn_kernel = is_fused_attn_kernel_available( # This needs to be fixed: TE-Jax has historically correlated training mode with deterministic mode. not deterministic, - self.dtype, - self.dtype, + input_dtype, + # self._assert_dtypes enforces Q, K, V, bias to have the same dtype so using input_dtype as kv dtype is sufficient + input_dtype, qkv_layout, attn_bias_type, attn_mask_type, @@ -743,7 +757,7 @@ def __call__( "Fused attention is not enabled because there is no available kernel.\n" "Fall back to the unfused attention.\n" "Please try to update the cuDNN and TE to the latest version.\n" - f"{self.dtype=}\n{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\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.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\n" ) @@ -797,7 +811,6 @@ def __call__( attention_dropout=self.attention_dropout, attn_mask_type=attn_mask_type, attn_bias_type=attn_bias_type, - dtype=self.dtype, float32_logits=self.float32_logits, scale_factor=scale_factor, transpose_batch_sequence=self.transpose_batch_sequence, @@ -817,7 +830,6 @@ def __call__( attention_dropout=self.attention_dropout, attn_mask_type=attn_mask_type, attn_bias_type=attn_bias_type, - dtype=self.dtype, scale_factor=scale_factor, transpose_batch_sequence=self.transpose_batch_sequence, qkv_layout=qkv_layout, @@ -1572,7 +1584,6 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): attn_mask_type=self.attn_mask_type, attn_bias_type=self.attn_bias_type, attention_dropout=self.attention_dropout, - dtype=self.dtype, dropout_rng_name=self.dropout_rng_name, float32_logits=self.float32_logits, qkv_layout=qkv_layout.name, diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 6cb0dd257c..c8daebbabd 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -59,6 +59,14 @@ def _validate_mesh_resource_configuration(mesh_resource): ) +def is_mesh_available() -> bool: + """ + Check if a physical mesh is available. + """ + mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + return mesh is not None and not mesh.empty + + def get_sharding_map_logic_axis_to_mesh_axis(): """ Generate a dict to map logical axes to mesh axes. From 14ddb43035c2bf8398481c197ab9c2f25794f8b4 Mon Sep 17 00:00:00 2001 From: LucienXian Date: Thu, 18 Dec 2025 21:15:03 +0800 Subject: [PATCH 134/521] Fix meta device check failure when passing torch.device objects (#2519) * Fix meta device check failure when passing torch.device objects Signed-off-by: LucienXian * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: LucienXian Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/module/grouped_linear.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 004c95c372..c4d35a9c2c 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -702,7 +702,8 @@ def __init__( if self.primary_weights_in_fp8: self.init_fp8_metadata(num_gemms=self.num_gemms) - self.reset_parameters(defer_init=device == "meta") + is_meta = torch.device(device).type == "meta" + self.reset_parameters(defer_init=is_meta) if self.wgrad_store.delay_wgrad_compute(): for name, param in self.named_parameters(): From 3e6939707737dff0002500c597e05ea5a5b09da7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?oliver=20k=C3=B6nig?= Date: Thu, 18 Dec 2025 22:47:55 +0100 Subject: [PATCH 135/521] ci: Use whitelisted sha for `get-release` (#2531) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: oliver könig --- .github/workflows/attach-wheels-to-release.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/attach-wheels-to-release.yml b/.github/workflows/attach-wheels-to-release.yml index c7d31a7c7d..6b97d4bb9b 100644 --- a/.github/workflows/attach-wheels-to-release.yml +++ b/.github/workflows/attach-wheels-to-release.yml @@ -96,7 +96,7 @@ jobs: - name: Get Release with tag id: get_current_release - uses: joutvhu/get-release@v1 + uses: joutvhu/get-release@9a8271732adc3299a22f8ad09b0a67eb3aa836ac if: ${{ github.event_name == 'workflow_dispatch' }} with: tag_name: ${{ inputs.release-version }} From 6fd620985658db4634832d6b50d4bbe85c02d4ee Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Thu, 18 Dec 2025 18:08:03 -0800 Subject: [PATCH 136/521] [PyTorch] Make sure Float8Tensor.contiguous supports autograd (#2533) * add early return back (removed in 2427) Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make sure Float8Tensor.contiguous supports autograd Expand quantized tensor tests to check identity ops. 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: Sudhakar Singh Signed-off-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon --- tests/pytorch/test_quantized_tensor.py | 223 +++++++++++++----- .../pytorch/tensor/float8_tensor.py | 37 +-- 2 files changed, 187 insertions(+), 73 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index bdf355c7bc..0353944ed6 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -20,6 +20,7 @@ Float8Tensor, MXFP8Tensor, NVFP4Tensor, + QuantizedTensor, ) from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported @@ -50,14 +51,22 @@ def _to_list(x: Union[Iterable, Any]) -> List: # Types that can be interpreted as tensor dims DimsType = Union[Iterable[int], int] -# Check if FP8 is supported +# Supported quantization recipes fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) - fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( return_reason=True ) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +_quantization_list: List[str] = [] +if fp8_available: + _quantization_list.append("fp8") +if fp8_block_scaling_available: + _quantization_list.append("fp8_blockwise") +if mxfp8_available: + _quantization_list.append("mxfp8") +if nvfp4_available: + _quantization_list.append("nvfp4") # delayed scaling @@ -98,6 +107,79 @@ def to_float8_CS( return quantizer(tensor) +@torch.no_grad() +def make_reference_and_test_tensors( + shape: int | Iterable[int], + quantization: Optional[str] = None, + ref_dtype: torch.dtype = torch.float64, + ref_device: torch.device = "cpu", + test_dtype: torch.dtype = torch.float32, + test_device: torch.device = "cuda", + requires_grad: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Construct tensors with the same values + + The reference tensor is intended for use in plain PyTorch + operations in high precision. The test tensor is intended for use + in Transformer Engine operations. + + If a quantization scheme is provided, the tensor values are + quantized so that they are representable. + + """ + + # Random reference tensor + ref = torch.rand(shape, dtype=ref_dtype, device=ref_device) + + # Construct test tensor from reference tensor + test = ref.to(device=test_device, dtype=test_dtype) + if quantization is None: + if test.data_ptr() == ref.data_ptr(): + test = test.clone() + elif quantization in ("fp8", "fp8_delayed_scaling"): + quantizer = Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device=test_device).squeeze(), + amax=torch.zeros(1, dtype=torch.float32, device=test_device), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + test = quantizer(test) + elif quantization == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=test_device, + ) + test = quantizer(test) + 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, + ) + test = quantizer(test) + elif quantization == "mxfp8": + test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) + elif quantization == "nvfp4": + test = NVFP4Quantizer( + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=False, + )(test) + else: + raise ValueError(f"Unsupported quantization scheme ({quantization})") + + # Make sure reference and test tensors match each other + ref.copy_(test) + + ref.requires_grad_(requires_grad) + test.requires_grad_(requires_grad) + return ref, test + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) class TestFloat8Tensor: @@ -466,7 +548,7 @@ def test_quantize_dequantize( torch.testing.assert_close(x_fp8_dequantized, -x_hp, **_tols[fp8_dtype]) -class TestAllQuantizedTensors: +class TestQuantizedTensor: @staticmethod def setup_class(cls) -> None: # Configure RNG @@ -474,10 +556,69 @@ def setup_class(cls) -> None: torch.manual_seed(seed) torch.cuda.manual_seed(seed) - @pytest.mark.parametrize("quantization", ["fp8", "mxfp8", "nvfp4", "fp8_blockwise"]) + @pytest.mark.parametrize("op", ("clone", "view", "reshape", "contiguous")) + @pytest.mark.parametrize("quantization", _quantization_list) + def test_identity_op( + self, + *, + op: str, + quantization: str, + shape: Iterable[int] = (128, 128), + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + ) -> None: + """Test operations that do not affect tensor values. + + These operations are must produce outputs that are bit-wise + equivalent to the inputs. They must support autograd. + + """ + + # Create reference and quantized tensor + x_ref, x_test = make_reference_and_test_tensors( + shape=shape, + quantization=quantization, + test_dtype=dtype, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + shape=shape, + test_dtype=dtype, + requires_grad=False, + ) + + # Apply identity operation + if op == "clone": + y_ref = x_ref.clone() + y_test = x_test.clone() + elif op == "view": + y_ref = x_ref.view(shape) + y_test = x_test.view(shape) + elif op == "reshape": + y_ref = x_ref.reshape(shape) + y_test = x_test.reshape(shape) + elif op == "contiguous": + y_ref = x_ref.contiguous() + y_test = x_test.contiguous() + + # Check autograd + y_test.backward(dy_test) + assert x_test.grad is not None + + # Check values + tols = dict(rtol=0, atol=0) + if isinstance(y_test, QuantizedTensor): + y_test = y_test.dequantize() + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dx_ref = dy_ref + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(dx_test, dx_ref, **tols) + + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("dim", [0, 1]) def test_chunk( self, + *, quantization: str, dim: int, shape: Iterable[int] = (128, 128), @@ -485,67 +626,33 @@ def test_chunk( dtype: torch.dtype = torch.bfloat16, device: torch.device = "cuda", ) -> None: - # Skip invalid configs - if quantization == "fp8" and not fp8_available: - pytest.skip(reason_for_no_fp8) - if quantization == "fp8_blockwise" and not fp8_block_scaling_available: - pytest.skip(reason_for_no_fp8_block_scaling) - if quantization == "mxfp8" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) - if quantization == "nvfp4" and not nvfp4_available: - pytest.skip(reason_for_no_nvfp4) - # Create 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 == "mxfp8": - quantizer = MXFP8Quantizer(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 == "nvfp4": - quantizer = NVFP4Quantizer( - with_rht=False, - with_post_rht_amax=False, - with_2d_quantization=False, - stochastic_rounding=False, - with_random_sign_mask=False, - ) - else: - raise ValueError(f"Unknown quantizer ({quantizer})") + # Create reference and quantized tensor - ref_tensor = torch.randn(shape, device=device, dtype=dtype) - quantized_tensor = quantizer(ref_tensor) - ref_tensor.copy_(quantized_tensor) + x_ref, x_test = make_reference_and_test_tensors( + shape=shape, + quantization=quantization, + test_dtype=dtype, + ) # Chunk tensors - ref_splits = torch.chunk(ref_tensor, chunks, dim=dim) - quantized_splits = torch.chunk(quantized_tensor, chunks, dim=dim) + ys_ref = torch.chunk(x_ref, chunks, dim=dim) + ys_test = torch.chunk(x_test, chunks, dim=dim) + # Check splits - for ref_split, quantized_split in zip(ref_splits, quantized_splits): + for y_ref, y_test in zip(ys_ref, ys_test): + # Check split shapes - assert ref_split.size() == quantized_split.size() + assert y_ref.size() == y_test.size() # Check that splits are quantized when expected if quantization == "fp8": - assert isinstance(quantized_split, Float8Tensor) - expected_value = quantized_split.dequantize() + assert isinstance(y_test, Float8Tensor) + y_test = y_test.dequantize() elif quantization == "mxfp8" and dim == 0: - assert isinstance(quantized_split, MXFP8Tensor) - expected_value = quantized_split.dequantize() - else: - # Otherwise torch dispatch would default to base implementation - # dequantize and computing output and hence output from torch chunk - # is already dequantized. - expected_value = quantized_split + assert isinstance(y_test, MXFP8Tensor) + y_test = y_test.dequantize() + # Check values - torch.testing.assert_close(expected_value, ref_split) + tols = dict(rtol=0, atol=0) # Chunking is exact + y_test = y_test.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 1077fe818f..67df40c047 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -551,24 +551,31 @@ def contiguous( ) -> Float8Tensor: """Returns tensor with data in provided memory format - Returns `self` if data is already in correct memory format. + Returns ``self`` if data is already in correct memory format. """ - # requires_grad remains unaltered when calling contiguous on - # torch tensor and so should be the case for our custom float8 tensor - # as well. - return Float8Tensor.make_like( - tensor=self, - data=self._data.contiguous(memory_format=memory_format), - data_transpose=( - self._transpose.contiguous(memory_format=memory_format) - if self._transpose is not None - else None - ), - requires_grad=self.requires_grad, - ) - # raise ValueError("Float8Tensor does not support different memory formats!") + # Check if tensor already has correct memory format + if self._data is not None and not self._data.is_contiguous(memory_format=memory_format): + pass + elif self._transpose is not None and not self._transpose.is_contiguous( + memory_format=memory_format + ): + pass + else: + # Tensor has correct memory format, so return immediately + return self + + # Construct tensor with correct data format + data, data_transpose = None, None + if self._data is not None: + data = self._data.contiguous(memory_format=memory_format) + if self._transpose is not None and not self._transpose_invalid: + data_transpose = self._transpose.contiguous(memory_format=memory_format) + return _IdentityFunc.apply( + self, + {"data": data, "data_transpose": data_transpose}, + ) def _reset_caches(self) -> None: """ From d46d5db449c270b79c99ba0a0091ae5c44ff7c50 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 19 Dec 2025 11:35:13 -0800 Subject: [PATCH 137/521] [JAX] Handle meshs set with jax.set_mesh (#2532) * Handle meshs set with jax.set_mesh Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/jax/sharding.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index c8daebbabd..b4b8c42027 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -37,6 +37,15 @@ W_JOINED_AXES = "nvte_w_joined" +def _get_mesh(): + # Handle Mesh's set via `with mesh:` + mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + if mesh is not None and not mesh.empty: + return mesh + # Handle Mesh's set via `jax.set_mesh(mesh)` + return jax.sharding.get_abstract_mesh() + + def _get_mesh_info(resource: str, mesh: jax.sharding.Mesh): assert resource in mesh.axis_names, f"{resource} is not in the axis_names of Mesh {mesh}." return mesh.shape[resource], resource @@ -63,7 +72,7 @@ def is_mesh_available() -> bool: """ Check if a physical mesh is available. """ - mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + mesh = _get_mesh() return mesh is not None and not mesh.empty @@ -71,7 +80,7 @@ def get_sharding_map_logic_axis_to_mesh_axis(): """ Generate a dict to map logical axes to mesh axes. """ - mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + mesh = _get_mesh() if mesh is None or mesh.empty: # If no mesh is defined, return an empty dict and do not require a MeshResource context to be present return {} @@ -130,7 +139,7 @@ def with_sharding_constraint(x: jnp.array, pspec: PartitionSpec): if pspec is None: return x - mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + mesh = _get_mesh() if mesh.empty: return x @@ -211,7 +220,7 @@ def get_all_mesh_axes(): """ Get all name of mesh axes """ - mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + mesh = _get_mesh() return mesh.axis_names @@ -251,7 +260,7 @@ def get_num_devices_in_mesh(mesh=None): by the global mesh. """ if mesh is None: - mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + mesh = _get_mesh() if mesh.empty: return 1 return np.prod(list(mesh.shape.values())) @@ -264,7 +273,7 @@ def get_mesh_axis_size(axis, mesh=None): by the global mesh. """ if mesh is None: - mesh = _PXLA_THREAD_RESOURCES.env.physical_mesh + mesh = _get_mesh() if axis is None: return 1 From 47902e96c5414750b5e61c8e49968b0dc1b944e5 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 19 Dec 2025 18:45:35 -0800 Subject: [PATCH 138/521] [JAX] Remove unused TE DPA module dtype which fixes cuDNN backend detection to properly use input dtypes (#2485) * Remove unused TE DPA module dtype which fixes cuDNN backend detection to properly use input dtypes Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Warning fallback Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * adjust test tolerances slightly for encoder tests due to change in backend Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../encoder/test_model_parallel_encoder.py | 4 +-- transformer_engine/jax/flax/transformer.py | 28 +++++++++++++++++-- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index a3935da9ff..8618c2be87 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -535,7 +535,7 @@ def test_te_delayed_scaling_fp8_with_sp(self): self.args.use_fp8 = True self.args.fp8_recipe = "DelayedScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.36 and actual[1] > 0.84 + assert actual[0] < 0.361 and actual[1] > 0.84 @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) def test_te_mxfp8_with_sp(self): @@ -569,7 +569,7 @@ def test_te_delayed_scaling_fp8_shardy(self): self.args.use_fp8 = True self.args.fp8_recipe = "DelayedScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.36 and actual[1] > 0.84 + assert actual[0] < 0.361 and actual[1] > 0.84 @unittest.skipIf(not is_fp8_supported, fp8_reason) def test_te_delayed_scaling_fp8_with_sp_shardy(self): diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 3c1d8ef9eb..1395976b9f 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -598,6 +598,11 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods ``'off-by-one'`` and ``'learnable'`` softmax types are also called sink attention (``'zero sink'`` and ``'learnable sink'``). + + Optimization parameters + ----------------------- + dtype(deprecated): jax.numpy.dtype, default = None + This dtype is deprecated and will be removed in a future release. DPA will use the dtype of the inputs instead as this module does not have any parameters. """ head_dim: int @@ -606,6 +611,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods attention_dropout: float = 0.0 attn_mask_type: AttnMaskType = "causal" attn_bias_type: AttnBiasType = None + dtype: Optional[DType] = None # Deprecated dropout_rng_name: str = "dropout" float32_logits: bool = False qkv_layout: str = "bshd_bshd_bshd" @@ -637,14 +643,14 @@ def _assert_dtypes(self, query: Array, key: Array, value: Array, qkv_layout: QKV elif qkv_layout.is_kvpacked(): assert ( key.dtype == query.dtype - ), f"Expected kv dtype={key.dtype} to match query dtype={query.dtype}." + ), f"Expected kv {key.dtype=} to match query {query.dtype=}." elif qkv_layout.is_separate(): assert ( key.dtype == query.dtype - ), f"Expected key dtype={key.dtype} to match query dtype={query.dtype}." + ), f"Expected key {key.dtype=} to match query {query.dtype=}." assert ( value.dtype == query.dtype - ), f"Expected value dtype={value.dtype} to match query dtype={query.dtype}." + ), f"Expected value {value.dtype=} to match query {query.dtype=}." else: raise ValueError(f"Unsupported {qkv_layout=}.") @@ -713,6 +719,22 @@ def __call__( bias = bias.astype(input_dtype) self._assert_dtypes(query, key, value, qkv_layout) + if self.dtype is not None: + if self.dtype == input_dtype: + warnings.warn( + "The dtype argument is deprecated and will be removed in a future release." + " DotProductAttention will use the dtype of the inputs instead as this module" + f" does not have any parameters. Module dtype specified {self.dtype=} matches" + " dtype of inputs so behavior is unchanged. Please remove the dtype argument" + " within the next few releases." + ) + else: + raise ValueError( + "The DotProductAttention module dtype is deprecated and will be removed in a" + " future release. DotProductAttention will use the dtype of the inputs instead" + " as this module does not have any parameters. Module dtype specified" + f" {self.dtype=} does not match dtype of inputs {input_dtype=}." + ) # Use fused attn (if kernel check below passes) by default enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) From eb8e792b38228e3c221a2d6b69babe3a91acd1f4 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Sat, 20 Dec 2025 09:46:56 -0800 Subject: [PATCH 139/521] [PyTorch][NVFP4][MOE] NVFP4 Grouped Quantize with Hadamard Transform (#2411) * rowwise colwise RHT group quant v1 Signed-off-by: Zhongbo Zhu * remove local array RW Signed-off-by: Zhongbo Zhu * change wait_barrier Signed-off-by: Zhongbo Zhu * fast math options Signed-off-by: Zhongbo Zhu * use mult to replace div Signed-off-by: Zhongbo Zhu * format Signed-off-by: Zhongbo Zhu * bulk move random states Signed-off-by: Zhongbo Zhu * greptile Signed-off-by: Zhongbo Zhu * lint Signed-off-by: Zhongbo Zhu * revert to use divides Signed-off-by: Zhongbo Zhu * avoid fp32 bf16 round-trip in RHT cast fusion Signed-off-by: Zhongbo Zhu * trigger fastmath by toggle NVTE_RHT_CAST_FUSION_USE_FAST_MATH Signed-off-by: Zhongbo Zhu * integrate row col rht fusion, functional Signed-off-by: Zhongbo Zhu * numerics aligned Signed-off-by: Zhongbo Zhu * style Signed-off-by: Zhongbo Zhu * remove device sync Signed-off-by: Zhongbo Zhu * 128 padding Signed-off-by: Zhongbo Zhu * revert colwise rng state creation because of row-col fused kernel Signed-off-by: Zhongbo Zhu * fix CI, linter Signed-off-by: Zhongbo Zhu * refactor RS for generating two random values Signed-off-by: Zhongbo Zhu * Avoid invalid configs with templated kernel Signed-off-by: Tim Moon * fix acc pipeline init with 0 arrival count Signed-off-by: Zhongbo Zhu * restore rowwise-only mode Signed-off-by: Zhongbo Zhu * switch to dynamic atomic scheduler Signed-off-by: Zhongbo Zhu * Avoid instantiating group RHT+cast kernel without row-wise or col-wise output Signed-off-by: Tim Moon * Include fast math option in quantization config Signed-off-by: Tim Moon * Fix linter warnings and review nits Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use TE license Signed-off-by: Tim Moon * Fix bug where kernel is always launched on stream Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore BF16 intermediate downcast in fused RHT-cast kernels Signed-off-by: Tim Moon * fix numerical test of grouped kernel Signed-off-by: Zhongbo Zhu * Make sure row-wise and col-wise quantization use different RNG seeds Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Restore autoformatter 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: Zhongbo Zhu Signed-off-by: Tim Moon Signed-off-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> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- benchmarks/linear/benchmark_grouped_linear.py | 29 +- .../nvfp4/test_nvfp4_group_quantize.py | 5 +- transformer_engine/common/CMakeLists.txt | 2 + transformer_engine/common/cast/cast.cu | 15 + .../common/cast/dispatch/quantize.cuh | 65 + .../nvfp4/group_quantize_transpose_nvfp4.cuh | 904 ++++++++++ transformer_engine/common/common.h | 4 +- .../customized_pipeline.cuh | 222 +++ .../group_hadamard_transform.cu | 2 +- .../group_hadamard_transform_cast_fusion.cu | 1003 +++++++++++ ...cast_col_hadamard_transform_cast_fusion.cu | 1499 +++++++++++++++++ .../hadamard_transform_cast_fusion.cu | 91 +- .../common/include/transformer_engine/cast.h | 14 + .../transformer_engine/hadamard_transform.h | 37 + .../transformer_engine/transformer_engine.h | 12 + .../common/transformer_engine.cpp | 20 +- .../pytorch/csrc/extensions/cast.cpp | 509 ++++-- transformer_engine/pytorch/csrc/quantizer.cpp | 2 +- transformer_engine/pytorch/quantization.py | 2 +- 19 files changed, 4205 insertions(+), 232 deletions(-) create mode 100644 transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh create mode 100644 transformer_engine/common/hadamard_transform/customized_pipeline.cuh create mode 100644 transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu create mode 100644 transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu diff --git a/benchmarks/linear/benchmark_grouped_linear.py b/benchmarks/linear/benchmark_grouped_linear.py index 02e2bcf4b9..f559928f8c 100644 --- a/benchmarks/linear/benchmark_grouped_linear.py +++ b/benchmarks/linear/benchmark_grouped_linear.py @@ -53,7 +53,7 @@ --set=full \ --kernel-name "GroupHadamardAmaxTmaKernel" \ -s 5 -c 5 \ - python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe nvfp4 --profile + python benchmarks/linear/benchmark_grouped_linear.py --profile --recipe nvfp4 """ @@ -173,7 +173,9 @@ def benchmark_linear( return timing_ms -def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None): +def run_benchmark_linear( + mkns, recipe_name, use_bias, num_gemms=4, m_splits_provided=None, fwd_only=False +): data = [] assert not use_bias, "Bias is not supported for GroupedLinear benchmark" @@ -182,14 +184,14 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None device = "cuda" 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)] - assert m % num_gemms == 0 - m_splits = [m // num_gemms] * num_gemms if m_splits is None else m_splits + m_splits = [m // num_gemms] * num_gemms if m_splits_provided is None else m_splits_provided # Bias is not supported for GroupedLinear benchmark bias = None # Run the benchmark print(f"fwd_m={m}, fwd_k={k}, fwd_n={n}") print(f"m_splits: {m_splits}") + print(f"fwd_only: {fwd_only}") grouped_fwd_bwd_timing_ms = benchmark_linear( x, @@ -197,7 +199,7 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None m_splits, bias, recipe_name, - mode="fwd_bwd", + mode="fwd_only" if fwd_only else "fwd_bwd", num_gemms=num_gemms, ) @@ -213,6 +215,8 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None ] ) + timing_notation = "grouped_fwd_time_ms" if fwd_only else "grouped_fwd_bwd_time_ms" + df = pd.DataFrame( data=data, columns=[ @@ -221,7 +225,7 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None "n", "recipe", "num_gemms", - "grouped_fwd_bwd_time_ms", + timing_notation, ], ) @@ -234,7 +238,7 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None parser = argparse.ArgumentParser() parser.add_argument("--profile", action="store_true", help="Enable profiling mode") parser.add_argument( - "--output_dir", + "--output-dir", type=str, default="benchmark_output/", help="output path for report", @@ -266,6 +270,12 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None default=2048, help="Output dimension to use, default is 2048", ) + parser.add_argument( + "--fwd-only", + action="store_true", + default=False, + help="Run forward pass only, default is both forward and backward passes", + ) args = parser.parse_args() jagged_input_splits = None @@ -297,7 +307,7 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None if jagged_input_splits is not None: num_gemms_list = [len(jagged_input_splits)] - token_dim_list = [65536] + token_dim_list = [16384, 32768, 65536, 98304] hidden_dim_list = [7168] output_dim_list = [2048] @@ -371,7 +381,8 @@ def run_benchmark_linear(mkns, recipe_name, use_bias, num_gemms=4, m_splits=None recipe_name, use_bias, num_gemms=num_gemms, - m_splits=jagged_input_splits, + m_splits_provided=jagged_input_splits, + fwd_only=args.fwd_only, ) df_linears = pd.concat([df_linears, df]) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 10aa3eb505..a29dcb4279 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -198,7 +198,7 @@ def check_group_quantization_nvfp4_versus_reference( for i in range(len(x_qx)): if split_sections[i] == 0: - # then just assert the same same and dtype because the buffer won't be zero out + # then just assert the same shape and dtype because the buffer won't be zero out assert_same_shape_and_dtype(x_amax_rowwise[i], x_amax_rowwise_ref[i]) assert_same_shape_and_dtype(x_qx[i], x_qx_ref[i]) assert_same_shape_and_dtype(x_sx[i], x_sx_ref[i]) @@ -221,7 +221,7 @@ def check_group_quantization_nvfp4_versus_reference( # assert with zero tolerance for i in range(len(x_qx_t)): if split_sections[i] == 0: - # then just assert the same same and dtype because the buffer won't be zero out + # then just assert the same shape and dtype because the buffer won't be zero out assert_same_shape_and_dtype(x_amax_colwise[i], x_amax_colwise_ref[i]) assert_same_shape_and_dtype(x_qx_t[i], x_qx_t_ref[i]) assert_same_shape_and_dtype(x_sx_t[i], x_sx_t_ref[i]) @@ -247,6 +247,7 @@ def check_group_quantization_nvfp4_versus_reference( (1024, 256), # larger sizes (8192, 1024), + (16384, 8192), (16384, 16384), ], ) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 264f7f9a78..79948e28f7 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -174,6 +174,8 @@ list(APPEND transformer_engine_cuda_arch_specific_sources hadamard_transform/group_hadamard_transform.cu hadamard_transform/hadamard_transform.cu hadamard_transform/hadamard_transform_cast_fusion.cu + hadamard_transform/group_hadamard_transform_cast_fusion.cu + hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu multi_tensor/compute_scale.cu recipe/mxfp8_scaling.cu transpose/quantize_transpose_square_blockwise.cu diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 1ed46a3359..73467d7275 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -100,3 +100,18 @@ 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 +// TODO (zhongbo): find a better way to make it a more generalized API +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_helper(input, outputs, split_sections, + num_tensors, quant_config, stream); +} diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 9f7a4a9b01..6d4454402c 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -19,6 +19,7 @@ #include "../core/common.cuh" #include "../fp8/quantize_fp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" +#include "../nvfp4/group_quantize_transpose_nvfp4.cuh" #include "../nvfp4/quantize_nvfp4.cuh" #include "../nvfp4/quantize_transpose_nvfp4.cuh" @@ -320,6 +321,70 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens } } +template +void group_quantize_fwd_helper(const NVTETensor input, NVTETensor *outputs, + const size_t *split_sections, const size_t num_tensors, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + using namespace detail; + + const Tensor *input_tensor = convertNVTETensorCheck(input); + std::vector output_tensors; + for (size_t i = 0; i < num_tensors; ++i) { + output_tensors.push_back(convertNVTETensorCheck(outputs[i])); + } + + // Quantization config + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Noop flag + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + // Check for unsupported options + if (quant_config_cpp.stochastic_rounding) { + NVTE_CHECK(output_tensors[0]->scaling_mode == NVTE_NVFP4_1D_SCALING, + "Stochastic rounding is only supported for NVFP4 quantization."); + } + + // Take the scaling mode of the first output tensor + auto scaling_mode = output_tensors[0]->scaling_mode; + + // Dispatch to quantization kernel depending on data format + switch (scaling_mode) { + case NVTE_NVFP4_1D_SCALING: { + NVTE_CHECK(!IS_ACT, "IS_ACT is not supported by FWD NVTE_NVFP4_1D_SCALING"); + + // Check tensors + CheckNoopTensor(*noop_tensor, "cast_noop"); + CheckInputTensor(*input_tensor, "input"); + // Skip checking output tensor list + // 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(); + auto dtype = input_tensor->dtype(); + + NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, + "2D quantization is not supported for group quantize."); + + // Launch NVFP4 group quantize kernel + nvfp4::group_quantize_transpose( + *input_tensor, noop_tensor, output_tensors, split_sections, num_tensors, + &quant_config_cpp, stream); + break; + } + default: + NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); + } +} + } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh new file mode 100644 index 0000000000..28b47e32d2 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -0,0 +1,904 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_transpose_nvfp4.cuh + * \brief CUDA kernels to cast to NVFP4 and transpose. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_QUANTIZE_TRANSPOSE_NVFP4_CUH_ +#define TRANSFORMER_ENGINE_GROUP_QUANTIZE_TRANSPOSE_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 group_quantize_transpose_kernel { + +using namespace quantization_and_transposition_SF; +using namespace core; +using namespace ptx; + +#if FP4_TYPE_SUPPORTED + +constexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed +struct MultiAmaxCastTransposeFusionArgs { + // Amax buffer for rowwise scaling + void *rowwise_amax_list[kMaxTensorsPerKernel]; + // Rowwise scale pointers with 128x4 padding included for rowwise scaling + void *output_rowwise_scale_inv_list[kMaxTensorsPerKernel]; + // (Unused for rowwise only scaling) Amax buffer for colwise scaling + void *colwise_amax_list[kMaxTensorsPerKernel]; + // (Unused for rowwise only scaling) output data pointers for fp4 transposed output + void *output_colwise_data_list[kMaxTensorsPerKernel]; + // (Unused for rowwise only scaling) output scale inverse pointers for each tensor + void *output_colwise_scale_inv_list[kMaxTensorsPerKernel]; + // (Unused for rowwise only scaling) output scale stride for colwise scaling + int output_colwise_scale_stride[kMaxTensorsPerKernel]; + // Prefix sum (with leading zero) of split_sections of each tensor of input + int split_sections_range[kMaxTensorsPerKernel + 1]; + // Number of tensors (splits) being processed by kernel + int num_tensors; +}; + +__device__ __forceinline__ int GetTensorId(MultiAmaxCastTransposeFusionArgs *kernel_args_ptr, + int offset) { + // check the kernel args and get the corresponding id + int tensor_id = 0; + while (kernel_args_ptr->split_sections_range[tensor_id + 1] <= offset) { + ++tensor_id; + } + return tensor_id; +} + +// Helper to get tensor id at offset, and also whether [offset_start, offset_end) crosses a split boundary. +__device__ __forceinline__ int GetTensorIdAndBoundary( + MultiAmaxCastTransposeFusionArgs *kernel_args_ptr, int offset_start, int offset_end, + bool *cross_boundary) { + int tensor_id_start = 0; + while (kernel_args_ptr->split_sections_range[tensor_id_start + 1] <= offset_start) { + ++tensor_id_start; + } + int tensor_id_end = tensor_id_start; + if (offset_end != offset_start) { + if (kernel_args_ptr->split_sections_range[tensor_id_start + 1] < offset_end) { + tensor_id_end = tensor_id_start + 1; + } + } + if (cross_boundary) { + *cross_boundary = (tensor_id_start != tensor_id_end); + } + return tensor_id_start; +} + +__device__ __forceinline__ void UpdateEncodeDecodeScaleFP32(float *amax_ptr, float *s_enc_ptr, + float *s_dec_ptr) { + float s_env_value = + (amax_ptr == nullptr) ? 1.0f : compute_global_encode_scaling_factor_FP4(*amax_ptr); + float s_dec_value = 1.0 / s_env_value; + *s_enc_ptr = s_env_value; + *s_dec_ptr = s_dec_value; + return; +} + +constexpr size_t SCALE_DIM = 16; // NVFP4 block (x16 elts) + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_NUM = 128; + +constexpr size_t SCALES_PER_CHUNK_Y = CHUNK_DIM_Y / SCALE_DIM; +constexpr size_t SCALES_PER_CHUNK_X = CHUNK_DIM_X / SCALE_DIM; + +constexpr size_t SCALES_PER_THREAD = 2 * (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM / THREADS_NUM; + +// Each call generates 4x uint32_t random numbers +constexpr size_t RNG_GENS_PER_THREAD = SCALES_PER_THREAD / 4; + +constexpr size_t TILE_DIM_Y = 32; +constexpr size_t TILE_DIM_X = 128; + +// SHould this be SCALE_DIM or BLOCK_DIM? Both are 16, should work for both 1D and 2D +constexpr size_t SCALES_PER_TILE_Y = TILE_DIM_Y / SCALE_DIM; +constexpr size_t SCALES_PER_TILE_X = TILE_DIM_X / SCALE_DIM; // 128 / 16 = 8 + +constexpr size_t TILES_Y = CHUNK_DIM_Y / TILE_DIM_Y; +constexpr size_t TILES_X = CHUNK_DIM_X / TILE_DIM_X; +constexpr size_t STAGES = TILES_Y * TILES_X; + +constexpr size_t BUFFS_NUM = 2; +constexpr size_t BUFF_DIM_Y = TILE_DIM_Y; +constexpr size_t BUFF_DIM_X = TILE_DIM_X; +constexpr size_t BUFF_SIZE = BUFF_DIM_Y * BUFF_DIM_X; +constexpr size_t BUFF_SIZE_TOTAL = BUFF_SIZE * BUFFS_NUM; + +// Input buffer (BF16) +constexpr size_t BUFF_IN_DIM_Y = BUFF_DIM_Y; +constexpr size_t BUFF_IN_DIM_X = BUFF_DIM_X; +constexpr size_t BUFF_IN_SIZE = BUFF_IN_DIM_Y * BUFF_IN_DIM_X; + +// Output buffer (NVFP4) +constexpr size_t BUFF_OUT_DIM_Y = BUFF_DIM_Y; +constexpr size_t BUFF_OUT_DIM_X = (BUFF_DIM_X * 4) / 8; +constexpr size_t BUFF_OUT_SIZE = BUFF_OUT_DIM_Y * BUFF_OUT_DIM_X; + +// Output transpose buffer (NVFP4) +constexpr size_t BUFF_OUT_T_DIM_Y = BUFF_DIM_X; +constexpr size_t BUFF_OUT_T_DIM_X = (BUFF_DIM_Y * 4) / 8; +constexpr size_t BUFF_OUT_T_SIZE = BUFF_OUT_T_DIM_Y * BUFF_OUT_T_DIM_X; + +// Manual swizzling parameters to reduce SHMEM bank conflicts +constexpr size_t PACK_SIZE = 8; +constexpr size_t WAVES = SCALE_DIM / PACK_SIZE; + +constexpr size_t SCALING_FACTORS_PER_TILE_X = TILE_DIM_X / SCALE_DIM; +constexpr size_t THREADS_X_ROWWISE = SCALING_FACTORS_PER_TILE_X; // 128 / 16 = 8 +constexpr size_t THREADS_Y_ROWWISE = THREADS_NUM / THREADS_X_ROWWISE; // 128 / 8 = 16 + +constexpr size_t ITERATIONS_NORMAL = BUFF_DIM_Y / THREADS_Y_ROWWISE; // 32/ 16 = 2 +constexpr size_t ITERATIONS_TRANSPOSE = BUFF_IN_DIM_Y / SCALE_DIM; +constexpr size_t BUFF_OUT_IT_OFFSET = BUFF_OUT_T_DIM_X / ITERATIONS_TRANSPOSE; + +static_assert(BUFF_DIM_Y >= SCALE_DIM && + "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"); + +// 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; // 8 = 128 / 16 + +template +__global__ void __launch_bounds__(THREADS_NUM) + group_quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, + const __grid_constant__ CUtensorMap tensor_map_output, + nvfp4_scale_t *const scales_ptr, const float *noop, + const size_t rows, const size_t cols, + const size_t scale_stride, const size_t *rng_state, + MultiAmaxCastTransposeFusionArgs kernel_args) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + 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; + } + } + + const size_t rng_sequence = + threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; + // Index of the random number. It increments each time when used and resets to 0 if reaches 4x + int rnd_idx = 0; + + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS; + + const size_t block_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const size_t block_offset_X = blockIdx.x * CHUNK_DIM_X; + + // TODO(zhongbo): add back when transpose is supported + // const size_t block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; + // const size_t block_offset_X_t = blockIdx.y * CHUNK_DIM_Y; + + const size_t chunk_rows = rows - block_offset_Y; + + const size_t scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; + const size_t scales_block_offset_X_rowwise = blockIdx.x * SCALES_PER_CHUNK_X; + // TODO(zhongbo): add back when transpose is supported + // const size_t scales_block_offset_Y_t = blockIdx.x * CHUNK_DIM_X; + // const size_t scales_block_offset_X_t = blockIdx.y * SCALES_PER_CHUNK_Y; + + const size_t tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; + const size_t tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; + const size_t tid_X_colwise = threadIdx.x; + const size_t tid_Y_t = tid_X_colwise; + // const size_t tid_X_t = 0; + + const size_t thread_offset_Y_rowwise = tid_Y_rowwise; + const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM; + const size_t thread_offset_X_colwise = tid_X_colwise; + + const size_t row_base_rowwise = block_offset_Y + thread_offset_Y_rowwise; + const size_t row_base_colwise = block_offset_Y; + const size_t col_base_colwise = block_offset_X + thread_offset_X_colwise; + + const bool col_out_of_bounds_colwise = (col_base_colwise >= cols); + + const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; + const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; + // TODO(zhongbo): add back when transpose is supported + // const size_t scales_offset_Y_t = scales_block_offset_Y_t + tid_Y_t; + // const size_t scales_offset_X_t = scales_block_offset_X_t; + + const size_t SFs_per_row = cols / SCALE_DIM; + + const bool rowwise_scale_is_within_bounds_X = scales_offset_X_rowwise < SFs_per_row; + + // TODO(zhongbo): add back when transpose is supported + // const bool colwise_scale_is_within_bounds_Y = scales_offset_Y_t < 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 = + DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); + + constexpr size_t in_mem = buff_size_aligned_in; + + constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; + constexpr size_t out_mem_colwise_data = buff_size_aligned_out; + constexpr size_t out_mem_rowwise_scales = 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_data_sh = reinterpret_cast(dshmem + in_mem); + fp4e2m1x2 *out_t_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + + nvfp4_scale_t *out_rowwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer + + constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + const bool is_master_thread = (threadIdx.x == 0); + + // TODO (zhongbo): finish this + float *amax_rowwise_ptr = nullptr; + float *amax_colwise_ptr = nullptr; + nvfp4_scale_t *split_rowwise_scale_ptr = nullptr; + + // suppose the amax is fixed for the current 128x128 tile (need 128 padding) + bool need_update_tensor_id = true; + int tensor_id = GetTensorIdAndBoundary(&kernel_args, block_offset_Y, block_offset_Y + CHUNK_DIM_Y, + &need_update_tensor_id); + size_t split_start = kernel_args.split_sections_range[tensor_id]; + size_t split_end = kernel_args.split_sections_range[tensor_id + 1]; + amax_rowwise_ptr = reinterpret_cast(kernel_args.rowwise_amax_list[tensor_id]); + split_rowwise_scale_ptr = + reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); + + float S_enc_rowwise = 1.0f; + float S_dec_rowwise = 1.0f; + UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); + + // TODO (zhongbo): colwise scaling disabled for now because of transpose + float S_enc_colwise = 1.0f; + float S_dec_colwise = 1.0f; + if (amax_colwise_ptr != nullptr) { + UpdateEncodeDecodeScaleFP32(amax_colwise_ptr, &S_enc_colwise, &S_dec_colwise); + } else { + S_enc_colwise = S_enc_rowwise; + S_dec_colwise = S_dec_rowwise; + } + + 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 (size_t stage = 0; stage < STAGES; ++stage) { + const size_t buff = stage % BUFFS_NUM; + const size_t next_stage = stage + 1; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + + const size_t buff_offset_in = buff * BUFF_IN_SIZE; + const size_t buff_offset_out = buff * BUFF_OUT_SIZE; + const size_t buff_offset_out_t = buff * BUFF_OUT_T_SIZE; + + // for stages from 1 to STAGES - 1, we need to update the tensor id + // skip updating tensor id if it's the last CTA, and some stages will be out of bounds + if (need_update_tensor_id && stage > 0 && (block_offset_Y + stage_offset_Y < rows)) { + int new_tensor_id = GetTensorId(&kernel_args, block_offset_Y + stage_offset_Y); + if (new_tensor_id != tensor_id) { + tensor_id = new_tensor_id; + split_start = kernel_args.split_sections_range[tensor_id]; + split_end = kernel_args.split_sections_range[tensor_id + 1]; + amax_rowwise_ptr = reinterpret_cast(kernel_args.rowwise_amax_list[tensor_id]); + UpdateEncodeDecodeScaleFP32(amax_rowwise_ptr, &S_enc_rowwise, &S_dec_rowwise); + split_rowwise_scale_ptr = + reinterpret_cast(kernel_args.output_rowwise_scale_inv_list[tensor_id]); + // TODO (zhongbo): colwise scaling disabled for now because of transpose + // Skip fetching colwise amax pointer and scaling factor updates + } + } + + 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 size_t next_buff = next_stage % BUFFS_NUM; + const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_buff_offset = next_buff * BUFF_IN_SIZE; + + 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; + + // COLWISE scaling + if constexpr (RETURN_TRANSPOSE) { +#pragma unroll + for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { + const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; + const size_t in_thread_offset_X = thread_offset_X_colwise; + + const size_t out_t_thread_offset_Y = thread_offset_X_colwise; + const size_t out_t_thread_offset_X = 0 + it * BUFF_OUT_IT_OFFSET; + + const size_t shmem_offset_base_colwise_in = + buff_offset_in + in_thread_offset_Y * BUFF_IN_DIM_X + in_thread_offset_X; + const size_t shmem_offset_base_colwise_out_t = + buff_offset_out_t + out_t_thread_offset_Y * BUFF_OUT_T_DIM_X + out_t_thread_offset_X; + + block_amax = 0.0f; + float in_compute_colwise[SCALE_DIM]; + IType in_colwise_IType[SCALE_DIM]; + // 1. Read/Compute elements. Find NVFP4-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; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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; ++i) { + const int shmem_offset_colwise = shmem_offset_base_colwise_in + 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 E4M3 scaling factor + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax, S_enc_colwise); + + // Store scaling factors through SHMEM + const size_t scale_idx_sh = + tid_Y_t * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + it; + out_colwise_scales_sh[scale_idx_sh] = S_dec_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_colwise), float_max); // S_enc_b_fp8 + const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; + + // 3. Scale elements + fp4e2m1x4 regs[SCALE_DIM / 4]; + +#pragma unroll + for (int e = 0; e < SCALE_DIM / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); + regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); + const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); + regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } + } + + const int group = thread_lane / 16; + uint32_t val[2]; + uint32_t *regs_4x = reinterpret_cast(regs); + + // Helps reducing bank conflicts + switch (group) { + case 0: + val[0] = regs_4x[0]; + val[1] = regs_4x[1]; + break; + case 1: + val[0] = regs_4x[1]; + val[1] = regs_4x[0]; + + break; + } + uint32_t *out_t_data_sh_as_uint32_t = + reinterpret_cast(&out_t_data_sh[shmem_offset_base_colwise_out_t]); + out_t_data_sh_as_uint32_t[group] = val[0]; // idx1 = (group + 0) % 2; + out_t_data_sh_as_uint32_t[(group + 1) & 1] = val[1]; // idx2 = (group + 1) % 2; + } + } + + // ROWWISE scaling + { + const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; +#pragma unroll + for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { + const size_t it_thread_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; + + const size_t shmem_offset_base_rowwise_in = + buff_offset_in + it_thread_offset_Y_rowwise * BUFF_IN_DIM_X; + const size_t shmem_offset_base_rowwise_out = + buff_offset_out + it_thread_offset_Y_rowwise * BUFF_OUT_DIM_X; + + const size_t it_offset_Y = stage_offset_Y + it * THREADS_Y_ROWWISE; + + block_amax = 0.0f; + float in_compute_rowwise[SCALE_DIM]; + 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t 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 size_t 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 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 bool rowwise_scale_is_within_bounds_Y = + (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE) < chunk_rows; + + // TODO(zhongbo): depending on input padding multiple (whether 128 or 64), use either scale_ptr or split_rowwise_scale_ptr + // const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + // if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { + // scales_ptr[scale_idx_global] = S_dec_b_fp8; + // } + + // Map to local split coordinates + const size_t split_rows = split_end - split_start; + const size_t local_scale_row = scales_offset_Y - split_start; + + // Local bounds: 0 <= local_scale_row < split_rows + const bool local_rowwise_scale_is_within_bounds_Y = local_scale_row < split_rows; + + // Index inside this split’s scale buffer + const size_t scale_idx_local = local_scale_row * scale_stride + scales_offset_X; + + if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y && + local_rowwise_scale_is_within_bounds_Y) { + split_rowwise_scale_ptr[scale_idx_local] = S_dec_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 +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + Vec out; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + IType2 in01; + IType2 in23; + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_IType[w].data.elt[2 * e]); + out.data.elt[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else if constexpr (IS_CACHED_ACT_OP) { + const uint64_t elts = *reinterpret_cast(&in_cached[w].data.elt[4 * e]); + out.data.elt[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const int j = w * PACK_SIZE + 4 * e; + const float2 in01 = make_float2(in_compute_rowwise[j], in_compute_rowwise[j + 1]); + const float2 in23 = make_float2(in_compute_rowwise[j + 2], in_compute_rowwise[j + 3]); + out.data.elt[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } + } + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM; + const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; + out.store_to(&out_data_sh[shmem_offset_rowwise]); + } + } + } + + __builtin_assume(thread_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 size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + + // TODO(zhongbo): add back when transpose is supported + // const size_t global_offset_Y_t = block_offset_Y_t; + // const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, + reinterpret_cast(&out_data_sh[buff_offset_out])); + + // TODO(zhongbo): add back when transpose is supported + // if constexpr (RETURN_TRANSPOSE) { + // ptx::cp_async_bulk_tensor_2d_shared_to_global( + // reinterpret_cast(&tensor_map_output_t), global_offset_X_t, + // global_offset_Y_t, reinterpret_cast(&out_t_data_sh[buff_offset_out_t])); + // } + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + } + } // end of stages + + // TODO(zhongbo): add back when transpose is supported + // Vectorized store scaling factors through SHMEM + // if (RETURN_TRANSPOSE && colwise_scale_is_within_bounds_Y) { + // using ScalesVec = Vec; + // const size_t scale_idx_sh = tid_Y_t * SCALES_PER_CHUNK_Y; + // ScalesVec &scales_vec = *reinterpret_cast(&out_colwise_scales_sh[scale_idx_sh]); + // const size_t scale_idx_global = scales_offset_Y_t * scale_stride_t + scales_offset_X_t; + // const size_t count = // number of scales in Y dimension of this chunk + // (chunk_rows >= CHUNK_DIM_Y) ? SCALES_PER_CHUNK_Y : (chunk_rows / SCALE_DIM); + // nvfp4_scale_t *dst = &scales_t_ptr[scale_idx_global]; + // constexpr size_t vec_bytes = SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t); + // if (count == SCALES_PER_CHUNK_Y && (reinterpret_cast(dst) % vec_bytes == 0)) { + // // Fast path: vectorized store when destination is properly aligned + // scales_vec.store_to(dst); + // } else { + // // Safe path: element-wise store for tails or unaligned destinations + // scales_vec.store_to_elts(dst, 0, count); + // } + // } + + destroy_barriers(mbar, is_master_thread); +#else + NVTE_DEVICE_ERROR("sm_100 or higher is required."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +#endif // FP4_TYPE_SUPPORTED +} // namespace group_quantize_transpose_kernel + +template +void group_quantize_transpose(const Tensor &input, const Tensor *noop, + std::vector &output_list, const size_t *split_sections, + size_t num_tensors, const QuantizationConfig *quant_config, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace group_quantize_transpose_kernel; + using namespace ptx; + bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; + + NVTE_CHECK(num_tensors == output_list.size(), + "Number of output tensors should match number of tensors."); + NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel, + "Number of tensors should be less than or equal to ", kMaxTensorsPerKernel); + + Tensor *output = nullptr; + // loop over the list to find the first non-empty tensor + for (size_t i = 0; i < num_tensors; ++i) { + if (output_list[i]->has_data()) { + output = output_list[i]; + break; + } + } + NVTE_CHECK(output != nullptr, "No output tensor found."); + // also check that the output has not null data pointer + NVTE_CHECK(output->data.dptr != nullptr, "Output data pointer is null."); + + // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to + // return the transposed data. + bool return_transpose = output->has_columnwise_data(); + // forbid return transpose for now because group quantize transpose is not supported yet + NVTE_CHECK(!return_transpose, "Return transpose is not supported for group quantize transpose."); + + // output_List is contiguous in memory, so take the first tensor as the contiguous output + auto output_contiguous = output->data; + + constexpr bool COMPUTE_ACTIVATIONS = false; + using ParamOP = Empty; + constexpr float (*OP)(float, const ParamOP &) = nullptr; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + + 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(); + + NVTE_CHECK(rows % 32 == 0, + "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA + NVTE_CHECK(cols % 32 == 0, + "Number of tensor cols must be a multiple of 32"); // 16B alignment for TMA + + // process the output list and produce the multi-tensor args for grouped kernel + MultiAmaxCastTransposeFusionArgs kernel_args; + kernel_args.num_tensors = 0; + kernel_args.split_sections_range[0] = 0; + for (size_t i = 0; i < num_tensors; ++i) { + if (split_sections[i] == 0) { + continue; + } + kernel_args.rowwise_amax_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->amax.dptr); + kernel_args.output_rowwise_scale_inv_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->scale_inv.dptr); + // kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i]; + kernel_args.split_sections_range[kernel_args.num_tensors + 1] = + kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i]; + // check overflow + NVTE_CHECK(kernel_args.split_sections_range[kernel_args.num_tensors + 1] >= 0, + "split_sections_range overflow the int32_t"); + kernel_args.num_tensors++; + } + + 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_NUM; + + // Note (zhongbo): for group quantize of [x1, x2, ..., xn] + // for the rowwise sclaing, scaling factor stride is shared between all tensors + // for the colwise scaling, scaling factor stride is different for each tensor because of transpose + // since transpose puts token dimension splits in the last dimension of the tensor + const size_t scale_stride = output->scale_inv.shape[1]; + // const size_t scale_stride_transpose = + // return_transpose ? output->columnwise_scale_inv.shape[1] : 0; + + nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); + + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + const NVTETensor rng_state_tensor = (quant_config != nullptr) ? quant_config->rng_state : nullptr; + const size_t *rng_state = nullptr; + if (rng_state_tensor != nullptr) { + 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}, + "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); + } + + using IType = bf16; + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_output{}; + // alignas(64) CUtensorMap tensor_map_output_transpose{}; + + 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, output_contiguous, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, + cols, 0, 4); + // if (return_transpose) { + // create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, + // BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); + // } + 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 = + DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); + + constexpr size_t in_mem = buff_size_aligned_in; + + constexpr size_t out_data_mem = buff_size_aligned_out; + constexpr size_t out_data_transpose_mem = buff_size_aligned_out; + constexpr size_t out_scales_transpose_mem = buff_size_scales; + + constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; + + constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, + + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + auto kernel = + group_quantize_transpose_nvfp4_kernel; + + if constexpr (use_2d_quantization) { + NVTE_ERROR("2D quantization is not supported for group quantize transpose."); + } + + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + kernel<<>>(tensor_map_input, tensor_map_output, + scales_ptr, noop_ptr, rows, cols, + scale_stride, rng_state, kernel_args); + 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_GROUP_QUANTIZE_TRANSPOSE_NVFP4_CUH_ diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 38b437b994..0e264eaae3 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -394,6 +394,7 @@ struct QuantizationConfig { NVTETensor rng_state = nullptr; bool nvfp4_2d_quantization = false; bool stochastic_rounding = false; + bool use_fast_math = false; static constexpr size_t attr_sizes[] = { sizeof(bool), // force_pow_2_scales @@ -402,7 +403,8 @@ struct QuantizationConfig { sizeof(Float8BlockScaleTensorFormat), // float8_block_scale_tensor_format sizeof(NVTETensor), // rng_seed and offset sizeof(bool), // nvfp4_2d_quantization - sizeof(bool) // stochastic_rounding + sizeof(bool), // stochastic_rounding + sizeof(bool) // use_fast_math }; }; diff --git a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh new file mode 100644 index 0000000000..b6f6799a49 --- /dev/null +++ b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh @@ -0,0 +1,222 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_HADAMARD_TRANSFORM_CUSTOMIZED_PIPELINE_CUH_ +#define TRANSFORMER_ENGINE_COMMON_HADAMARD_TRANSFORM_CUSTOMIZED_PIPELINE_CUH_ + +#include "cutlass/pipeline/sm100_pipeline.hpp" + +namespace cutlass { + +using namespace cute; +namespace detail { +// Producer-consumer pipeline implementation +// for UMMA producer. In this case, UMMA barrier arrives are used +// by producer_commit. Use case, accumulator generation as +// the result of MMA instructions. +template , + class AtomThrShape_MNK_ = Shape<_1, _1, _1> > +class CustomizedPipelineTmaUmmaAsync { + public: + static constexpr uint32_t Stages = Stages_; + using AtomThrShape_MNK = AtomThrShape_MNK_; + + private: + using Impl = PipelineTmaAsync; + + public: + using FullBarrier = typename Impl::FullBarrier; + using EmptyBarrier = typename Impl::EmptyBarrier; + using ProducerBarrierType = typename Impl::ProducerBarrierType; + using ConsumerBarrierType = typename Impl::ConsumerBarrierType; + using PipelineState = typename Impl::PipelineState; + using SharedStorage = typename Impl::SharedStorage; + using ThreadCategory = typename Impl::ThreadCategory; + using Params = typename Impl::Params; + + using McastDirection = McastDirection; + + // Helper function to initialize barriers + static CUTLASS_DEVICE void init_barriers(SharedStorage& storage, Params params, + ClusterShape cluster_shape) { + int warp_idx = canonical_warp_idx_sync(); + if (warp_idx == params.initializing_warp) { + // Barrier FULL and EMPTY init + constexpr int producer_arv_cnt = 1; + auto atom_thr_shape = AtomThrShape_MNK{}; + + uint32_t multicast_consumer_arrival_count = params.num_consumers; // If cluster_size is 1 + if (cute::size(cluster_shape) > 1) { + multicast_consumer_arrival_count = + ((cute::size<0>(cluster_shape) / cute::size<0>(atom_thr_shape)) + + (cute::size<1>(cluster_shape) / cute::size<1>(atom_thr_shape)) - 1) * + params.num_consumers; + } + CUTLASS_ASSERT(multicast_consumer_arrival_count > 0 && + "Multicast consumer arrival count must be non-zero"); + CUTLASS_ASSERT(producer_arv_cnt > 0 && "Producer arrival count must be non-zero"); + cutlass::arch::detail::initialize_barrier_array_pair_aligned< + decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>( + storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, + multicast_consumer_arrival_count); + } + cutlass::arch::fence_barrier_init(); + } + + CUTLASS_DEVICE + void init_masks(ClusterShape cluster_shape, + dim3 block_id_in_cluster = cute::block_id_in_cluster()) { + // Calculate consumer mask + if (params_.role == ThreadCategory::Consumer) { + block_id_mask_ = detail::calculate_multicast_mask( + cluster_shape, AtomThrShape_MNK{}, block_id_in_cluster); + } + } + + CUTLASS_DEVICE + void init_masks(ClusterShape cluster_shape, McastDirection mcast_direction) { + // Calculate consumer mask + dim3 block_id_in_cluster = cute::block_id_in_cluster(); + if (mcast_direction == McastDirection::kRow) { + block_id_mask_ = detail::calculate_multicast_mask( + cluster_shape, AtomThrShape_MNK{}, block_id_in_cluster); + } else { + block_id_mask_ = detail::calculate_multicast_mask( + cluster_shape, AtomThrShape_MNK{}, block_id_in_cluster); + } + } + + // Constructor by default initializes barriers and calculates masks. + // These operations can be explicity deferred by specifying InitBarriers and InitMasks. + // If deferred, user code needs to guarantee init_masks and/or init_barriers is/are called. + template + CUTLASS_DEVICE CustomizedPipelineTmaUmmaAsync(SharedStorage& storage, Params params, + ClusterShape cluster_shape, InitBarriers = {}, + InitMasks = {}) + : impl_(storage, params, cluster_shape, cute::false_type{}, InitMasks{}), + params_(params), + empty_barrier_ptr_(&storage.empty_barrier_[0]), + full_barrier_ptr_(&storage.full_barrier_[0]) { + static_assert(cute::is_same_v || + cute::is_same_v); + if constexpr (cute::is_same_v) { + init_barriers(storage, params_, cluster_shape); + } + + static_assert(cute::is_same_v || + cute::is_same_v); + if constexpr (cute::is_same_v) { + init_masks(cluster_shape); + } + } + + //////////////////// + // Producer APIs + //////////////////// + // Four member functions are always used in pairs: + // + // * producer_try_acquire and producer_acquire, and + // * consumer_try_wait and consumer_wait. + // + // The two functions with "try" in their names are called "try" functions, + // and the other two are conceptually "finalize" functions. + // The "try" function in each pair starts the process of waiting on the barrier to flip. + // It opportunistically waits for an implementation-dependent timeout. + // Whether or not the barrier has flipped yet, the try function will return a token. + // If the token indicates that the barrier has not flipped, + // then the token must be passed into the corresponding "finalize" function. + // The finalize function will then block until the barrier has flipped. + // If the token indicates that the barrier _has_ flipped, + // then it is still correct to pass it into the finalize function. + // The finalize function will return immediately in that case. + CUTLASS_DEVICE + ProducerToken producer_try_acquire(PipelineState state, uint32_t skip_wait = false) { + return impl_.producer_try_acquire(state, skip_wait); + } + + CUTLASS_DEVICE + void producer_acquire(PipelineState state, + ProducerToken barrier_token = {BarrierStatus::WaitAgain}) { + impl_.producer_acquire(state, barrier_token); + } + + CUTLASS_DEVICE + void producer_expect_transaction(PipelineState state, uint32_t transaction_bytes) { + impl_.producer_expect_transaction(state, transaction_bytes); + } + + // NOP for TMA based mainloop + CUTLASS_DEVICE + void producer_commit(PipelineState state, uint32_t bytes) { impl_.producer_commit(state, bytes); } + + // Prevents early exit of producer blocks in Cluster. + // This should be called once before kernel exits. + CUTLASS_DEVICE + void producer_tail(PipelineState state) { impl_.producer_tail(state); } + + CUTLASS_DEVICE + ProducerBarrierType* producer_get_barrier(PipelineState state) { + return impl_.producer_get_barrier(state); + } + + //////////////////// + // Consumer APIs + //////////////////// + CUTLASS_DEVICE + ConsumerToken consumer_try_wait(PipelineState state, uint32_t skip_wait = false) { + return impl_.consumer_try_wait(state, skip_wait); + } + + CUTLASS_DEVICE + void consumer_wait(PipelineState state, + ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) { + impl_.consumer_wait(state, barrier_token); + } + + CUTLASS_DEVICE + void umma_consumer_release(PipelineState state) { umma_consumer_release(state.index(), false); } + CUTLASS_DEVICE + void consumer_release(PipelineState state) { impl_.consumer_release(state); } + + private: + Impl impl_; + Params params_; + EmptyBarrier* empty_barrier_ptr_; + FullBarrier* full_barrier_ptr_; + uint16_t block_id_mask_ = 0; + static constexpr bool is_2sm_mma = size(AtomThrShape_MNK{}) > 1; + + // Consumer signalling Producer of completion + // Ensures all blocks in the Same Row and Column get notified. + CUTLASS_DEVICE + void umma_consumer_release(uint32_t stage, uint32_t skip) { + detail::pipeline_check_is_consumer(params_.role); + uint64_t* smem_ptr = reinterpret_cast(&empty_barrier_ptr_[stage]); + // {$nv-release-never begin} + // TODO: Needs to be updated once Blackwell specialized pipeline is implemented. + // XMMA style bar_peek will be tested. We will need to revisit skip interface and + // what skip means when we have bar_peek functionality. + // A separate MR will implement MMA_2x1SM specialized pipeline. + // {$nv-release-never end} + if constexpr (is_2sm_mma) { // Mma cluster shape is 2x1 + if (!skip) { + cutlass::arch::umma_arrive_multicast_2x1SM(smem_ptr, block_id_mask_); + } + } else { + if (!skip) { + if constexpr (cute::is_static_v && size(ClusterShape{}) == 1) { + cutlass::arch::umma_arrive(smem_ptr); + } else { + cutlass::arch::umma_arrive_multicast(smem_ptr, block_id_mask_); + } + } + } + } +}; +} // namespace detail +} // namespace cutlass + +#endif // TRANSFORMER_ENGINE_COMMON_HADAMARD_TRANSFORM_CUSTOMIZED_PIPELINE_CUH_ diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index 84eb6bb5c3..ea5e22bbfb 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -459,7 +459,7 @@ void group_hadamard_transform_amax(const Tensor& input_, std::vector& o } // Multi zero out multiple amaxes if needed - // Curretly don't support multi-launch when num_tensors is larger than kMaxTensorsPerKernel + // Currently don't support multi-launch when num_tensors is larger than kMaxTensorsPerKernel // let the number of threads equal to number of tensors, use 1 block, kMaxTensorsPerKernel threads per block dim3 block_setup_amax(kMaxTensorsPerKernel); dim3 grid_setup_amax(1); 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 new file mode 100644 index 0000000000..6e071ec79f --- /dev/null +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu @@ -0,0 +1,1003 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "common/common.h" +#include "common/util/cuda_runtime.h" +#include "common/util/curanddx.hpp" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" +#include "cutlass/arch/barrier.h" +#include "cutlass/cutlass.h" +#include "cutlass/gemm/collective/builders/sm100_common.inl" +#include "cutlass/numeric_conversion.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/print_error.hpp" + +namespace transformer_engine { +namespace detail { +namespace { + +using namespace cute; +using cute:: + Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor + +using Stride2D = cute::Stride>; + +constexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB, expand 64 if needed +struct MultiAmaxHadamardCastFusionArgs { + // (output) Amax buffer for pre-RHT amax buffer + void *global_amax_list[kMaxTensorsPerKernel]; + // output C pointers for each tensor + void *output_colwise_list[kMaxTensorsPerKernel]; + // output scale inverse pointers for each tensor + void *output_colwise_scale_inv_list[kMaxTensorsPerKernel]; + // split sections of each tensor of input + int split_sections[kMaxTensorsPerKernel]; + // Prefix sum (with leading zero) of split_sections of each tensor of input + int split_sections_range[kMaxTensorsPerKernel + 1]; + // stride 2D struct for CUTE + Stride2D output_stride2d_list[kMaxTensorsPerKernel]; + // Number of tensors (splits) being processed by kernel + int num_tensors; +}; + +__device__ __forceinline__ float *GetGlobalAmaxPtrByTensorId( + MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, int tensor_id) { + // directly returns the global amax pointer by tensor id + if (tensor_id < 0 || tensor_id >= kernel_args_ptr->num_tensors) { + return nullptr; + } + return reinterpret_cast(kernel_args_ptr->global_amax_list[tensor_id]); +} + +__device__ __forceinline__ int GetTensorId(MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, + int offset) { + // Check the kernel args and get the corresponding id + const int num_tensors = kernel_args_ptr->num_tensors; + if (offset >= kernel_args_ptr->split_sections_range[num_tensors]) { + return num_tensors - 1; + } + int tensor_id = 0; + while (kernel_args_ptr->split_sections_range[tensor_id + 1] <= offset) { + ++tensor_id; + } + return tensor_id; +} + +// calculate the global encode scale factor for a given global amax. +__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) { + constexpr float kFP8E4M3Max = 448.0f; + constexpr float kFP4E2M1Max = 6.0f; + // If scale is infinity, return max value of float32 + float global_encode_scale = cutlass::minimum_with_nan_propagation{}( + kFP8E4M3Max * kFP4E2M1Max / global_amax, cutlass::platform::numeric_limits::max()); + // If global amax is 0 or infinity, return 1 + return (global_amax == 0.f || global_encode_scale == 0.f) ? 1.f : global_encode_scale; +} + +template +struct SharedStorage { + static constexpr int AccumulatorPipelineStageCount = 16; + using AtomThrShapeMNK = cute::Shape<_1, _1, _1>; + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage; + + static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = + cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>; + using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage; + + alignas(16) AccumulatorPipelineStorage accumulator; + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) cute::uint64_t tma_barrier[1]; + uint32_t tmem_base_ptr; + + struct TensorStorage : cute::aligned_struct<128, _1> { + // cute::array_aligned> smem_A; + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + } tensors; +}; + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverterBase( + cutlass::Array const &input, cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + auto output_ptr = reinterpret_cast(&output); + asm volatile( + "{\n" + "cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\n" + "cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\n" + "}" + : "=h"(output_ptr[0]), "=h"(output_ptr[1]) + : "f"(input[0]), "f"(input[1]), "f"(input[2]), "f"(input[3]), "f"(input[4]), "f"(input[5]), + "f"(input[6]), "f"(input[7]), "r"(rbits[0]), "r"(rbits[1])); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return output; +} + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverter( + cutlass::Array const &input, cutlass::Array const *rbits) { + using result_type = cutlass::Array; + result_type output; + cutlass::Array *result_ptr = + reinterpret_cast *>(&output); + cutlass::Array const *source_ptr = + reinterpret_cast const *>(&input); + cutlass::Array const *rbits_ptr = + reinterpret_cast const *>(rbits); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; i++) { + result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]); + } + return output; +} + +template +__global__ static void group_rht_gemm_device( + MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, TA const *A, AStride dA, + ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a, TB const *B, BStride dB, + BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, CSmemLayout, TiledMMA mma, + MultiAmaxHadamardCastFusionArgs kernel_args, const size_t *rng_state) { + using namespace cute; + using X = Underscore; + // static constexpr bool kApplyStochasticRounding = true; + using ElementAccumulator = float; + static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static constexpr uint32_t kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); + + static constexpr int kTmaRhtTensorTransactionBytes = + cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v); + static constexpr int AccumulatorPipelineStageCount = 16; + + static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = + cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static constexpr int VectorSize = 16; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; + // Preconditions + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + + // Represent the full tensors + Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, N)); + Tensor mB = tma_load_b.get_tma_tensor(make_shape(16, 16)); + + using TensorC = decltype(make_tensor(subbyte_iterator(recast_ptr(nullptr)), // engine + make_shape(int{}, int{}), // (M, N_i) + Stride2D{} // stride (dM, dN) + )); + + using TensorSFC = decltype(make_tensor( + make_gmem_ptr(recast_ptr(nullptr)), + make_layout(make_shape(int{}, // M + make_shape(make_shape(Int<16>{}, _4{}), // (16, 4) + int{}) // n_tiles = split / 64 + ), + make_stride(int{}, // dM = (split / 16) + make_stride(make_stride(_0{}, _1{}), // inner (16,4) layout + _4{}) // tiles stride + )))); + + auto cluster_shape = Shape<_1, _1, _1>{}; + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + const int K_TILE_MAX = min(N, K) / 64; + uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile); + uint32_t tiles_in_n = (N + 64 - 1) / 64; + uint32_t linear_tile_idx = blockIdx.x; + uint32_t tile_idx_m = linear_tile_idx % tiles_in_m; + uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + + auto mainloop_tiler = Shape<_128, _16, _64>{}; + auto epilogue_tiler = Shape<_128, _64, _64>{}; + Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gB_nk = + local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) + // Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + using TensorGC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){}, + make_coord(_, _, _), Step<_1, _1, X>{})); + + using TensorGSFC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){}, + make_coord(_, _, _), Step<_1, _1, X>{})); + + // Allocate SMEM + extern __shared__ char shared_memory[]; + using SharedStorage = SharedStorage; + SharedStorage &shared_storage = *reinterpret_cast(shared_memory); + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + // + // MMA: Define C accumulators and A/B partitioning + // + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + auto mma_epilogue = make_tiled_mma( + SM100_MMA_F16BF16_SS{}, + Layout>{}); + ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster); + + using TiledMmaEpilogue = decltype(mma_epilogue); + Tensor tCgA = thr_mma.partition_A(gA_mk); + // Allocate "fragments" -- these are actually umma smem descriptors + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0, 2>(ClusterTileShape{})); + auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0, 2>(epilogue_tiler)); + + auto bulk_tmem_mma = + TiledMMA::make_fragment_C(append(acc_shape_mma, Int{})); + + auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C( + append(acc_shape_epilogue, Int{})); + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + 32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = + tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); + + auto [tBgB, tBsB] = + tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7); + + // if (is_epilogue_warp && elect_one_sync()) { + // // prefetch to make the global amax in cache + // for (size_t i = 0; i < kernel_args.num_tensors; ++i) { + // cute::prefetch(raw_pointer_cast(kernel_args.global_amax_list[i])); + // } + // } + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, + cluster_shape, cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = + cutlass::make_producer_start_state(); + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = + cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128; + accumulator_pipeline_params.initializing_warp = 1; + AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, accumulator_pipeline_params, + cluster_shape, + cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + if (is_dma_warp) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], + kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), + tBsB(_, 0)); + } + + do { + bool is_first_wave = linear_tile_idx == blockIdx.x; + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_, tile_idx_m, _); + int k_tile = 0; + auto barrier_token = + mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait); + + CUTE_NO_UNROLL + while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) { + int k_tile_idx_n = tile_idx_n + k_tile; + ++k_tile; + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType *tma_barrier = + mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + barrier_token = + mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait); + if (cute::elect_one_sync()) { + copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), + tAsA(_, write_stage)); + } + } + linear_tile_idx += gridDim.x; + tile_idx_m = linear_tile_idx % tiles_in_m; + tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + auto barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + CUTE_NO_UNROLL + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n;) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_, _, _, read_stage); + auto tCrB_nk = tCrB(_, _, 0, 0); + CUTE_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block) { + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTE_UNROLL + for (int i = 0; i < 4; i++) { + auto accumulators = + bulk_tmem_mma(_, _, _, accumulator_pipe_producer_state.index() * 4 + i); + gemm(mma, tCrA_mk(_, _, k_block * 4 + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; + } + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state); + } + + linear_tile_idx += gridDim.x; + tile_idx_m = linear_tile_idx % tiles_in_m; + tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } else if (is_epilogue_warp) { + static constexpr int FragmentSize = 256 / sizeof_bits_v; + + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int thread_idx = threadIdx.x % 128; + + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); + auto tiled_r2g = + make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(thread_idx); + auto thr_r2g = tiled_r2g.get_slice(thread_idx); + + // NVFP4 non-E8 recipe constants and global scales + static constexpr float fp4_max = 6.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + + // get global amax pointer + int tensor_id = GetTensorId(&kernel_args, tile_idx_n * 64); + float *global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, tensor_id); + + TC *cur_output_colwise_ptr = reinterpret_cast(kernel_args.output_colwise_list[tensor_id]); + TSFC *cur_output_colwise_scale_inv_ptr = + reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); + int cur_output_colwise_n = kernel_args.split_sections[tensor_id]; + + TensorC cur_mC = + cute::make_tensor(cute::subbyte_iterator(cur_output_colwise_ptr), + cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i) + kernel_args.output_stride2d_list[tensor_id]); + + auto cur_sfc_shape = + make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64)); + + auto cur_sfc_stride = + make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{})); + + TensorSFC cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr), + make_layout(cur_sfc_shape, cur_sfc_stride)); + + TensorGC cur_gC_mn = + local_tile(cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N) + ); + + TensorGSFC cur_gSFC_mn = local_tile( + cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N-like) + ); + + Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); + + float global_amax_val = *global_amax_ptr; + float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); + + // Scaling factor for fast math path + float global_encode_scale_multiplier = 1.0f; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + + float global_decode_scale = 1.0f / global_encode_scale; + + auto sfd_converter = cutlass::NumericConverter{}; + + do { + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) { + // get the starting index of current k-tile in global tensor, to query the correct global amax + int cur_k_tile_global_elem_idx = (tile_idx_n + k_tile) * 64; + int new_tensor_id = GetTensorId(&kernel_args, cur_k_tile_global_elem_idx); + // float* new_global_amax_ptr = GetGlobalAmaxPtr(&kernel_args, cur_k_tile_global_elem_idx); + global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, new_tensor_id); + // update the scaling factors when it's no longer the same amax pointer + // TODO(zhongbo): the math operations are very expensive + // since the kernel is persistent, we can have a cache for all the possible scaling factors + if (tensor_id != new_tensor_id) { + global_amax_val = *global_amax_ptr; + global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + global_decode_scale = 1.0f / global_encode_scale; + tensor_id = new_tensor_id; + // went through the cute operations to update the local tensors + cur_output_colwise_ptr = + reinterpret_cast(kernel_args.output_colwise_list[tensor_id]); + cur_output_colwise_scale_inv_ptr = + reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); + cur_output_colwise_n = kernel_args.split_sections[tensor_id]; + + cur_mC = cute::make_tensor( + cute::subbyte_iterator(cur_output_colwise_ptr), + cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i) + kernel_args.output_stride2d_list[tensor_id]); + + cur_sfc_shape = + make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64)); + + cur_sfc_stride = + make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{})); + + cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr), + make_layout(cur_sfc_shape, cur_sfc_stride)); + + cur_gC_mn = local_tile( + cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N) + ); + + cur_gSFC_mn = local_tile(cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} + // (BLK_M, BLK_N-like) + ); + + tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); + } + // maybe udpated to the new tensor id + int tensor_start_elem = kernel_args.split_sections_range[tensor_id]; + int local_tile_idx_n = (cur_k_tile_global_elem_idx - tensor_start_elem) / 64; + + Tensor tCgC_mn = tCgC(_, _, _, tile_idx_m, local_tile_idx_n); + Tensor tCgSFC_mn = cur_gSFC_mn(_, _, tile_idx_m, local_tile_idx_n); + + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto tCtC = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); + Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + + Tensor tTR_rAcc = + make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDrC = make_tensor(shape(tDgC)); + Tensor tTR_rAcc_frag = + recast>(coalesce(tTR_rAcc)); + Tensor tDrC_frag = recast>(coalesce(tDrC)); + + Tensor src = thr_r2g.retile_S(tDrC); + Tensor dst = thr_r2g.retile_D(tDgC); + + Tensor tCgSFC = make_tensor( + tCgSFC_mn.data(), make_layout(make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{}))); + + Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC)); + Tensor tDrSFC = make_tensor(shape(tDgSFC)); + + static constexpr int NumVecs = size(tDgC) / VectorSize; + Tensor tC_rRowSFD_frg = recast>(tDrSFC); + + cutlass::maximum_absolute_value_reduction, + true> + amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // TMEM_LOAD + copy(tiled_t2r, tDtC, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with unfused + // kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + } + + auto compute_frgs = reinterpret_cast *>( + tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrC_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); + } + + if constexpr (kUseFastMath) { + // Fast math: multiply with precomputed reciprocal + pvscales = cutlass::multiplies>{}( + vec_maxs, global_encode_scale_multiplier); + } else { + // Accurate math: perform division + pvscales = + cutlass::divides>{}(vec_maxs, fp4_max); + pvscales = cutlass::multiplies>{}( + pvscales, global_encode_scale); + } + auto pvscales_cvted = + cutlass::NumericArrayConverter{}(pvscales); + + tC_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}( + tC_rRowSFD_frg(_0{})); + auto qpvscale_scaled = cutlass::multiplies>{}( + qpvscale_ups, global_decode_scale); + cutlass::Array acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = + cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = + cutlass::divides>{}(1.0, qpvscale_scaled); + } + + // Initialize RNG for tile + const size_t rng_sequence = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256; + + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], cutlass::platform::numeric_limits::max()); + // auto acc_scale = acc_scales[v]; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale), + reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale)); + } + } + + copy(tiled_r2g, src, dst); + + // copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrC, tDgC); + + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC); + } + linear_tile_idx += gridDim.x; + tile_idx_m = linear_tile_idx % tiles_in_m; + tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + } +} + +// this function computes RHT-GEMM for +// A: m x n: col-major +// B: 16 x 16: row-major +// C: m x n: row-major +// SFC: m x (n/16): row-major +template +void group_rht_gemm_ntt_w_sfc(int m, int n, TA const *A, TB const *B, + MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, + const size_t *rng_state, uint32_t sm_count, cudaStream_t stream, + int k_tile_size = 2048) { + using namespace cute; + + // Define shapes (dynamic) + auto M = static_cast(m); + auto N = static_cast(n); + + // Define strides (mixed) + auto dA = make_stride(Int<1>{}, m); // (dM,dK) + auto dB = make_stride(Int<1>{}, 16); // (dN,dK) + for (size_t i = 0; i < kernel_args_ptr->num_tensors; ++i) { + kernel_args_ptr->output_stride2d_list[i] = + make_stride(kernel_args_ptr->split_sections[i], Int<1>{}); + } + + auto cga_shape = Shape<_1, _1, _1>{}; + auto cga_tile_shape = Shape<_128, _16, _16>{}; + auto cluster_tile_mainloop = Shape<_128, _16, _64>{}; + + // Construct the MMA + auto mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}, + Layout>{}); + + // MMA in CGA Layout XXX: Need to generalize synchro? {$nv-release-never} + + // Assert that the TiledMMA uses all CTAs in the CGA. + CUTE_STATIC_ASSERT_V(size(cga_shape) == size(mma)); + CUTE_STATIC_ASSERT_V(evenly_divides(cga_tile_shape, tile_shape(mma))); + + // Determine the A and B shapes + auto mma_shape_B = + partition_shape_B(mma, make_shape(size<1>(cga_tile_shape), size<2>(cga_tile_shape))); + + using TiledMma = decltype(mma); + using AtomThrID = typename TiledMma::AtomThrID; + + using SmemShape_M = decltype(shape_div( + shape<0>(cga_tile_shape), + shape_div(shape<0>(cga_tile_shape), size<0>(cga_tile_shape) / size(AtomThrID{})))); + using SmemShape_N = decltype(shape_div( + shape<1>(cga_tile_shape), + shape_div(shape<1>(cga_tile_shape), size<1>(cga_tile_shape) / size(AtomThrID{})))); + using SmemShape_K = decltype(cute::get<2>(cga_tile_shape)); + + using SmemLayoutAtomB = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + + auto mma_shape_A = partition_shape_A( + mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop))); + using SmemShape_M_A = + decltype(shape_div(shape<0>(cluster_tile_mainloop), + shape_div(shape<0>(cluster_tile_mainloop), + size<0>(cluster_tile_mainloop) / size(AtomThrID{})))); + using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop)); + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>()); + + // Define the smem layouts (static) + // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory + constexpr int kBlackwellSmemSize = 232448; // 232KB in bytes + constexpr int kBytesPerStage = + cute::size(mma_shape_A) * sizeof(TA) + cute::size(mma_shape_B) * sizeof(TB); + constexpr int kReservedBytes = 256; // Reserve for barriers and other uses + constexpr int kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage; + auto sP = Int{}; // SMEM pipelines + auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, + append(mma_shape_A, sP)); // (MMA,MMA_M,MMA_K,PIPE) + auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{}, + append(mma_shape_B, sP)); // (MMA,MMA_N,MMA_K,PIPE) + auto sC = Layout<_1>{}; // XXX Dummy + + // Create GMEM tensors + Tensor tensorA = make_tensor(A, make_layout(make_shape(M, N), dA)); // (M,N) + Tensor tensorB = make_tensor(B, make_layout(make_shape(16, 16), dB)); // (16,16) + + // Create the TiledCopy + + auto tma_load_a = + make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma); + auto tma_load_b = + make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cga_tile_shape, mma); + + // Assert checks on tile sizes -- no predication + NVTE_CHECK(M % size<0>(cga_tile_shape) == 0, "Inner dimension must be divisible by ", + static_cast(size<0>(cga_tile_shape)), " but got ", M, "."); + NVTE_CHECK(N % (4 * size<1>(cga_tile_shape)) == 0, "Outer dimension must be divisible by ", + 4 * static_cast(size<1>(cga_tile_shape)), " but got ", N, "."); + + uint32_t tiles = size(ceil_div(M, get<0>(cga_tile_shape))) * size(ceil_div(N, k_tile_size)); + + tiles = (tiles < sm_count) ? tiles : sm_count; + + dim3 dimBlock(256); + dim3 dimCluster(size<0>(cga_shape), size<1>(cga_shape), size<2>(cga_shape)); + dim3 dimGrid(tiles, 1, 1); + + int smem_size = sizeof(SharedStorage); + auto *kernel_ptr = &group_rht_gemm_device< + decltype(M), decltype(N), decltype(k_tile_size), decltype(cga_tile_shape), TA, decltype(dA), + decltype(sA), decltype(tma_load_a), TB, decltype(dB), decltype(sB), decltype(tma_load_b), TC, + Stride2D, decltype(sC), TSFC, decltype(mma), kEnableStochasticRounding, kUseFastMath>; + + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + (*kernel_ptr)<<>>(M, N, k_tile_size, cga_tile_shape, A, dA, + sA, tma_load_a, B, dB, sB, tma_load_b, sC, + mma, *kernel_args_ptr, rng_state); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// this function is used to wrap the group_rht_gemm_ntt_w_sfc function +// to transpose the input tensor A +template +void group_rht_gemm_ttt_wrapper(int m, int n, TA const *A, TB const *B, + MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, + const size_t *rng_state, uint32_t sm_count, cudaStream_t stream, + int k_tile_size = 1024) { + // in addition to transpose the input tensor A + // we also need to reshape m, n to at best + // ultilize as many SMs as possible while keeping + // a relatively large contiguous dimension. + // for example, after swapping m, n for transpose purposes, + // the input / output tensor shapes for RHT-GEMM are: + // A: n x m: col-major + // B: 16 x 16: row-major + // C: n x m: row-major + // SFC: n x (m/16): row-major + group_rht_gemm_ntt_w_sfc( + n, m, A, B, kernel_args_ptr, rng_state, sm_count, stream, k_tile_size); +} + +} // namespace +} // namespace detail + +void group_hadamard_transform_cast_fusion_columnwise( + const Tensor &input_, std::vector &output_list, const size_t *split_sections, + size_t num_tensors, const Tensor &hadamard_matrix_, QuantizationConfig &quant_config, + cudaStream_t stream) { + NVTE_API_CALL(group_hadamard_transform_cast_fusion_columnwise); + + using transformer_engine::detail::kMaxTensorsPerKernel; + using transformer_engine::detail::MultiAmaxHadamardCastFusionArgs; + + NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16, + "Input tensor must be BF16 tensor, but dtype is ", to_string(input_.dtype()), "."); + NVTE_CHECK(input_.dim() >= 2, "Input must be a 2D tensor."); + const SimpleTensor &input = input_.data; + + NVTE_CHECK(output_list.size() == num_tensors, + "Number of output tensors should match number of tensors."); + + NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel, + "Number of tensors should be less than or equal to ", kMaxTensorsPerKernel); + + // construct the multi-tensor args + MultiAmaxHadamardCastFusionArgs kernel_args; + kernel_args.num_tensors = 0; + kernel_args.split_sections_range[0] = 0; + for (size_t i = 0; i < num_tensors; ++i) { + NVTE_CHECK(split_sections[i] % 64 == 0, "component ", i, + " of split_sections should be 64 multiple"); + if (split_sections[i] == 0) { + continue; + } + kernel_args.global_amax_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->amax.dptr); + // TODO(zhongbo): should we change API assumption to use columnwise_data instead of data? + kernel_args.output_colwise_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->data.dptr); + kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] = + reinterpret_cast(output_list[i]->scale_inv.dptr); + kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i]; + kernel_args.split_sections_range[kernel_args.num_tensors + 1] = + kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i]; + kernel_args.num_tensors++; + } + + // Stochastic rounding config + const bool use_stochastic_rounding = quant_config.stochastic_rounding; + const size_t *rng_state = nullptr; + if (quant_config.rng_state != nullptr) { + 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}, + "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); + rng_state = reinterpret_cast(rng_state_tensor.data.dptr); + } + + // Template arguments + using TA = cute::bfloat16_t; + using TB = cute::bfloat16_t; + using TC = cutlass::float_e2m1_t; + using TSFC = cutlass::float_ue4m3_t; + + checkCuDriverContext(stream); + + // Check Hadamard matrix + constexpr int kHadamardDimension = 16; + + NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16, + "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(), "."); + const size_t hadamard_dimension = hadamard_matrix.shape[0]; + + const size_t ndim = input.shape.size(); + const size_t n = input.shape[ndim - 1]; + size_t m = 1; + for (size_t i = 0; i < ndim - 1; ++i) { + m *= input.shape[i]; + } + + auto sm_count = transformer_engine::cuda::sm_count(); + + NVTE_CHECK(n % hadamard_dimension == 0, "row_length must be divisible by hadamard_dimension."); + + NVTE_CHECK(m % hadamard_dimension == 0, "num_rows must be divisible by hadamard_dimension"); + + int k_tile_size = 1024; + + if (m == 8192 && n == 5120) { + k_tile_size = 512; + } else if (m == 8192 && n == 10240) { + k_tile_size = 1024; + } else if (m == 8192 && n == 2560) { + k_tile_size = 1280; + } else if (m == 8192 && n == 11328) { + k_tile_size = 1024; + } else if (m == 8192 && n == 512) { + k_tile_size = 256; + } else if (m == 8192 && n == 3584) { + k_tile_size = 512; + } else if (m == 11328 && n == 8192) { + k_tile_size = 1024; + } else if (m == 5120 && n == 8192) { + k_tile_size = 512; + } else if (m == 10240 && n == 8192) { + k_tile_size = 1024; + } else if (m == 2560 && n == 8192) { + k_tile_size = 1280; + } else if (m == 512 && n == 8192) { + k_tile_size = 256; + } else if (m == 3584 && n == 8192) { + k_tile_size = 512; + } else if (m < 1024 || n < 1024) { + k_tile_size = 512; + } + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_stochastic_rounding, kUseStochasticRounding, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + detail::group_rht_gemm_ttt_wrapper( + /*m=*/m, /*n=*/n, /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*kernel_args_ptr=*/&kernel_args, /*rng_state=*/rng_state, /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size););); +} + +} // namespace transformer_engine + +void nvte_group_hadamard_transform_cast_fusion_columnwise( + const NVTETensor input, NVTETensor *outputs, const NVTETensor hadamard_matrix, + const size_t *split_sections, const size_t num_tensors, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_hadamard_transform_cast_fusion_columnwise); + using namespace transformer_engine; + NVTE_CHECK(num_tensors > 0, "Number of tensors should be greater than 0."); + + Tensor *input_tensor = convertNVTETensorCheck(input); + std::vector output_list(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + output_list[i] = convertNVTETensorCheck(outputs[i]); + } + + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Call the multi-tensor Hadamard transform amax implementation. + group_hadamard_transform_cast_fusion_columnwise( + *input_tensor, output_list, split_sections, num_tensors, + *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, stream); +} 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 new file mode 100644 index 0000000000..3932b328ae --- /dev/null +++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -0,0 +1,1499 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "common/common.h" +#include "common/util/cuda_runtime.h" +#include "common/util/curanddx.hpp" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" +#include "customized_pipeline.cuh" +#include "cutlass/arch/barrier.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/cluster_launch.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cutlass/fast_math.h" +#include "cutlass/float8.h" +#include "cutlass/float_subbyte.h" +#include "cutlass/gemm/collective/builders/sm100_common.inl" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/platform/platform.h" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/print_error.hpp" + +namespace transformer_engine { +namespace detail { +namespace { + +using namespace cute; + +// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor +using cute::Tensor; + +constexpr int kMaxTensorsPerKernel = 64; + +struct MultiAmaxHadamardCastFusionArgs { + // (output) Amax buffer for input A amax buffer + void *global_a_amax_list[kMaxTensorsPerKernel]; + // (output) Amax buffer for pre-RHT amax buffer + void *global_d_amax_list[kMaxTensorsPerKernel]; + // output D pointers for each tensor + void *output_colwise_list[kMaxTensorsPerKernel]; + // output SFD inverse pointers for each tensor + void *output_colwise_scale_inv_list[kMaxTensorsPerKernel]; + // split sections of each tensor of input + int split_sections[kMaxTensorsPerKernel]; + // Prefix sum (with leading zero) of split_sections of each tensor of input + int split_sections_range[kMaxTensorsPerKernel + 1]; + + // Number of tensors (splits) being processed by kernel + int num_tensors; +}; + +__device__ __forceinline__ int GetGroupIdx(MultiAmaxHadamardCastFusionArgs *kernel_args_ptr, + int offset) { + // Check the kernel args and get the corresponding id + const int num_tensors = kernel_args_ptr->num_tensors; + if (offset >= kernel_args_ptr->split_sections_range[num_tensors]) { + return num_tensors - 1; + } + int group_idx = 0; + while (kernel_args_ptr->split_sections_range[group_idx + 1] <= offset) { + ++group_idx; + } + return group_idx; +} + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverterBase( + cutlass::Array const &input, cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + auto output_ptr = reinterpret_cast(&output); + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + "cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\n" + "cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\n" + "}" + : "=h"(output_ptr[0]), "=h"(output_ptr[1]) + : "f"(input[0]), "f"(input[1]), "f"(input[2]), "f"(input[3]), "f"(input[4]), "f"(input[5]), + "f"(input[6]), "f"(input[7]), "r"(rbits[0]), "r"(rbits[1])); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return output; +} + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverter( + cutlass::Array const &input, cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + cutlass::Array *result_ptr = + reinterpret_cast *>(&output); + cutlass::Array const *source_ptr = + reinterpret_cast const *>(&input); + cutlass::Array const *rbits_ptr = + reinterpret_cast const *>(&rbits); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; i++) { + result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]); + } + return output; +} + +template +struct SharedStorage { + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_; + using AtomThrShapeMNK = cute::Shape<_1, _1, _1>; + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage; + + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = + cutlass::detail::CustomizedPipelineTmaUmmaAsync, + AtomThrShapeMNK>; + using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage; + using SchedPipeline = cutlass::PipelineCLCFetchAsync; + using SchedPipelineStorage = typename SchedPipeline::SharedStorage; + using SchedThrottlePipeline = cutlass::PipelineAsync; + using SchedThrottlePipelineStorage = typename SchedThrottlePipeline::SharedStorage; + + struct TensorStorage : cute::aligned_struct<128, _1> { + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + } tensors; + + alignas(16) AccumulatorPipelineStorage accumulator; + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) cute::uint64_t tma_barrier[1]; + alignas(16) SchedPipelineStorage sched; + alignas(16) SchedThrottlePipelineStorage sched_throttle; + alignas(16) int32_t atomic_tile_id[SchedulerPipelineStageCount_]; + alignas(16) float global_a_amax[kMaxTensorsPerKernel]; + alignas(16) float global_d_amax[kMaxTensorsPerKernel]; + uint32_t atomic_tile_counter[SchedulerPipelineStageCount_]; + uint32_t tmem_base_ptr; +}; + +// Main RHT GEMM kernel entry -- highly templated for flexible architecture/config support +template +__launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( + MShape M, NShape packed_N, KShape K, ClusterShape cluster_shape, ClusterTileShape cluster_tile, + TA const *A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a, + TB const *B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, + TQA *QA, QAStride dQA, TSFA *SFA, TSFALayout sfa_layout, MultiAmaxHadamardCastFusionArgs args, + uint32_t *tile_scheduler_workspace, TiledMMA mma, const size_t *rng_state) { + using namespace cute; + + // Abort immediately if compilation is not supported + constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; + if constexpr (!is_blackwell_arch) { + NVTE_DEVICE_ERROR( + "group_row_col_rht_gemm_device is only supported on Blackwell " + "with architecture-specific compilation. " + "Try recompiling with sm_100a or similar."); + return; + } + static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, + "group_row_col_rht_gemm_device must generate row-wise " + "and/or column-wise output."); +#if !defined(CUTLASS_ARCH_CLC_ENABLED) + CUTLASS_NOT_IMPLEMENTED(); + return; +#endif + + using X = Underscore; + // Accumulator data type for main computation + using ElementAccumulator = float; + static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); + static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; + static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; + static constexpr bool kEnableRowQuant = kEnableRowQuant_; + static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_; + static constexpr bool kUseFastMath = kUseFastMath_; + + // Constant for RHT tensor processing (tile size etc) + static int constexpr RhtTensorSize = 16; + + // Transaction bytes for TMA transfer on RHT tensor blocks + static int constexpr kTmaRhtTensorTransactionBytes = + cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; + + // Mainloop pipeline stage calculation, vectorization parameters for scaling factors + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + static int constexpr SFVecSize = 16; + // Swizzle output layout for scaling factor arrays + using SwizzledSFALayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + + // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling + using MainloopPipeline = + cutlass::detail::CustomizedPipelineTmaUmmaAsync; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + using SchedPipeline = cutlass::PipelineCLCFetchAsync; + using SchedPipelineState = typename SchedPipeline::PipelineState; + using SchedThrottlePipeline = cutlass::PipelineAsync; + using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState; + + static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, "ClusterShape must be Shape<_1,_1,_1>"); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static int constexpr VectorSize = RhtTensorSize; + + // Compile-time safety: static shapes required for shared memory layouts + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + // CUTE_STATIC_ASSERT(is_static::value); + + auto cluster_size = size<0>(cluster_shape); + auto mainloop_tiler = Shape<_128, _16, _128>{}; + auto epilogue_tiler = Shape<_128, _128, _128>{}; + + static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler); + + struct TileScheduler { + uint32_t tiles_in_m = 0; + uint32_t tiles_in_n = 0; + uint32_t linear_idx = 0; + uint32_t next_linear_idx = 0; + uint32_t start_idx = 0; + uint32_t tile_m_idx = 0; + uint32_t tile_n_idx = 0; + int k_tile_max = 0; + uint32_t *atomic_tile_index_; + uint32_t *smem_tile_counter; + uint32_t atomic_offset; + cutlass::FastDivmodU64 divmod_tiles_in_m; + + CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, + uint32_t *atomic_tile_index, uint32_t *smem_tile_counter) + : tiles_in_m(tiles_m), + tiles_in_n(tiles_n), + linear_idx(blockIdx.x), + next_linear_idx(blockIdx.x), + start_idx(blockIdx.x), + k_tile_max(kmax), + atomic_tile_index_(atomic_tile_index), + smem_tile_counter(smem_tile_counter), + atomic_offset(gridDim.x), + divmod_tiles_in_m(uint64_t(tiles_m)) { + update_tile_idx(); + } + CUTLASS_DEVICE void update_tile_idx() { + uint64_t q, r; + divmod_tiles_in_m(q, r, uint64_t(linear_idx)); + tile_m_idx = static_cast(r); + tile_n_idx = static_cast(q) * uint32_t(k_tile_max); + } + CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; } + CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; } + CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; } + + CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; } + + CUTLASS_DEVICE bool is_valid() const { + return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()), + cute::make_coord(tiles_in_m, tiles_in_n)); + } + + CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; } + + CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; } + + // Fetch a new tile_id using atomics. + CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) { + uint32_t tile_id_counter = 0; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p atom.global.add.u32 %0, [%1], 1; \n\t" + "}" + : "=r"(tile_id_counter) + : "l"(atomic_tile_index_), "r"(pred)); + + return tile_id_counter; + } + + CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_consumer_state) { + sched_pipeline.consumer_wait(sched_pipeline_consumer_state); + next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()]; + cutlass::arch::fence_view_async_shared(); + sched_pipeline.consumer_release(sched_pipeline_consumer_state); + return; + } + + CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_producer_state) { + uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state); + // Wait for clcID buffer to become empty with a flipped phase + sched_pipeline.producer_acquire(sched_pipeline_producer_state); + auto is_leading_thread = cute::elect_one_sync(); + uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset; + uint32_t smem_addr = + cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]); + if (is_leading_thread) { + cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0); + } + + ++sched_pipeline_producer_state; + return sched_pipeline_producer_state; + } + + CUTLASS_DEVICE auto update_work_tile_info() { + linear_idx = next_linear_idx; + update_tile_idx(); + return; + } + }; + + // Allocate and alias shared memory to the kernel's shared storage type + extern __shared__ char shared_memory[]; + using SharedStorage = + SharedStorage; + SharedStorage &shared_storage = *reinterpret_cast(shared_memory); + + // Compute the number of tiles in M and N after tiling and assign scheduler + uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); + uint32_t tiles_in_n = uint32_t( + size(ceil_div(args.split_sections_range[args.num_tensors], size<2>(epilogue_tiler)))); + + TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace, + shared_storage.atomic_tile_counter); + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + + // Shapes for accumulated tiles in mainloop and epilogue + auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{}); + auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{}); + + // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended + auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); + auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); + + // Number of threads assigned for various epilogue roles depending on quantization settings + static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; + static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; + static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0; + static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0; + static int constexpr NumSchedThreads = 32; + static int constexpr NumMainloopLoadThreads = 32; + static int constexpr NumEpilogueThreads = + NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + NumMmaThreadCount + NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + // warp assignment + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_sched_warp = (warp_idx == 2); + bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); + bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + + MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, + cluster_shape, cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = + cutlass::make_producer_start_state(); + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + using AccumulatorPipelineInitBarriers = cute::bool_constant; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = + cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_col_quant_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = + size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; + accumulator_pipeline_params.initializing_warp = 1; + AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, accumulator_pipeline_params, + cluster_shape, AccumulatorPipelineInitBarriers{}, + cute::true_type{}); // Delay mask calculation + typename SchedPipeline::Params sched_pipeline_params; + if (is_sched_warp) { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer; + } else { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer; + } + sched_pipeline_params.producer_blockid = 0; + sched_pipeline_params.producer_arv_count = 1; + sched_pipeline_params.consumer_arv_count = + NumSchedThreads + + cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); + sched_pipeline_params.transaction_bytes = sizeof(uint32_t); + sched_pipeline_params.initializing_warp = 3; + SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape); + SchedPipelineState sched_pipeline_consumer_state; + SchedPipelineState sched_pipeline_producer_state = + cutlass::make_producer_start_state(); + + typename SchedThrottlePipeline::Params sched_throttle_pipeline_params; + if (is_dma_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer; + } + if (is_sched_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer; + } + sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; + sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; + sched_throttle_pipeline_params.dst_blockid = 0; + sched_throttle_pipeline_params.initializing_warp = 4; + + SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle, + sched_throttle_pipeline_params); + SchedThrottlePipelineState sched_pipeline_throttle_consumer_state; + SchedThrottlePipelineState sched_pipeline_throttle_producer_state = + cutlass::make_producer_start_state(); + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + + // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer + if (is_dma_warp) { + // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access). + cutlass::arch::warpgroup_reg_dealloc<32>(); + // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory. + Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N)); + Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); + + // Partition tensors for tiling according to the mainloop and cluster tilers. + Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gB_nk = + local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) + + // Shared memory tensors for pipeline + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + // Determine warp/tile positioning + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Partition global to local fragments for A and B + Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) + Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = + tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = + tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); + + auto [tBgB, tBsB] = + tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + if constexpr (kEnableRHTColQuant) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], + kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), + tBsB(_, 0)); + } + } + + do { + // is_first_wave indicates whether this scheduler wave is the first among a group. + bool is_first_wave = scheduler.is_first_wave(); + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_, scheduler.tile_m(), _); + int k_tile = 0; + + sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state); + sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state); + ++sched_pipeline_throttle_producer_state; + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { + int k_tile_idx_n = scheduler.tile_n_base() + k_tile; + ++k_tile; + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType *tma_barrier = + mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + if (cute::elect_one_sync()) { + copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), + tAsA(_, write_stage)); + } + } + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + // scheduler.advance(); + } while (scheduler.is_valid()); + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform. + cutlass::arch::warpgroup_reg_dealloc<32>(); + if constexpr (kEnableRHTColQuant) { + // Setup shared memory fragments for A and B tiles. + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Allocate "fragments" -- these are actually umma smem descriptors + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + // Wait until the B (Hadamard) tensor copy is complete + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + + auto barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_, _, _, read_stage); + auto tCrB_nk = tCrB(_, _, 0, 0); + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) { + int accumulator_k_block = + accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; + int tCrA_k_block = k_block * EpilogueUnrollFactor; + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < EpilogueUnrollFactor; i++) { + auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i); + gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; + } + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); + barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + } else if (is_sched_warp) { + // Scheduler warp manages tile assignment and pipeline progress for warps + cutlass::arch::warpgroup_reg_dealloc<32>(); + do { + sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state); + sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state); + ++sched_pipeline_throttle_consumer_state; + sched_pipeline_producer_state = + scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } else if (is_epilogue_col_quant_warp) { + // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage, + // and writing result tensors/scales to global memory. + cutlass::arch::warpgroup_reg_alloc<192>(); + if constexpr (kEnableRHTColQuant) { + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + auto acc_epilogue_pipelined_shape = + append(acc_shape_epilogue, Int{}); + auto bulk_tmem_epilogue_layout = make_layout( + acc_epilogue_pipelined_shape, + make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler))); + auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); + + // Use 256-bit fragments for aligned bulk stores + static int constexpr FragmentSize = 256 / sizeof_bits_v; + + // Wait for TMEM allocation for this pipeline to finish + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; + // g2s load all global_d_amax + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueColQuantThreadCount) { + shared_storage.global_d_amax[g] = + __ldg(reinterpret_cast(args.global_d_amax_list[g])); + } + + size_t rng_seed = 0; + size_t rng_offset = 0; + // Setup RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); + + // Determine quantization scale factor layouts/output splits for this group + TSFDLayout sfd_layout; + int cur_N = args.split_sections[group_idx]; + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // Build output tensors for columns and their quant scales + Tensor mD = make_tensor( + cute::subbyte_iterator(reinterpret_cast(args.output_colwise_list[group_idx])), + make_shape(M, cur_N), DStride{}); // (M,packed_N) + Tensor gD_mn = + local_tile(mD, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + Tensor mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( + args.output_colwise_scale_inv_list[group_idx])), + sfd_layout); + Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + + // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); + auto tiled_r2g = + make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); + auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); + + cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float c_global_amax_val = shared_storage.global_d_amax[group_idx]; + float global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + float global_decode_scale = 1.0f / global_encode_scale; + + // Scaling factor for fast math path + float global_encode_scale_multiplier = 1.0f; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + + do { + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); + ++k_tile) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset); + + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + c_global_amax_val = shared_storage.global_d_amax[group_idx]; + // update amax + global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + cur_N = args.split_sections[group_idx]; + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = + tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = + make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // update tensor + mD = make_tensor(cute::subbyte_iterator( + reinterpret_cast(args.output_colwise_list[group_idx])), + make_shape(M, cur_N), DStride{}); + gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( + args.output_colwise_scale_inv_list[group_idx])), + sfd_layout); + gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + } + int group_start_offset = args.split_sections_range[group_idx]; + int local_tile_n_idx = + (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler); + Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx); + + Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx); + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); + Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + + Tensor tTR_rAcc = + make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDrD = make_tensor(shape(tDgD)); + Tensor tTR_rAcc_frag = + recast>(coalesce(tTR_rAcc)); + Tensor tDrD_frag = recast>(coalesce(tDrD)); + + Tensor src = thr_r2g.retile_S(tDrD); + Tensor dst = thr_r2g.retile_D(tDgD); + + Tensor tDgSFD_view = make_tensor( + tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); + Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); + Tensor tDrSFD = make_tensor(shape(tDgSFD)); + + static int constexpr NumVecs = size(tDgD) / VectorSize; + Tensor tD_rRowSFD_frg = recast>(tDrSFD); + + // Compute amax and quantization scales for this tile + cutlass::maximum_absolute_value_reduction, + true> + amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // Copy from TMEM to registers + copy(tiled_t2r, tDtAcc, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with + // unfused kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); + } + + auto compute_frgs = reinterpret_cast *>( + tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); + } + + if constexpr (kUseFastMath) { + // Fast math: multiply with precomputed reciprocal + pvscales = cutlass::multiplies>{}( + vec_maxs, global_encode_scale_multiplier); + } else { + // Accurate math: perform division + pvscales = + cutlass::divides>{}(vec_maxs, fp4_max); + pvscales = cutlass::multiplies>{}( + pvscales, global_encode_scale); + } + auto pvscales_cvted = + cutlass::NumericArrayConverter{}(pvscales); + + tD_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}( + tD_rRowSFD_frg(_0{})); + auto qpvscale_scaled = cutlass::multiplies>{}( + qpvscale_ups, global_decode_scale); + cutlass::Array acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = + cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides>{}( + 1.0, qpvscale_scaled); + } + + // Prepare stochastic rounding random state if enabled + uint4 random_uint4 = uint4{0, 0, 0, 0}; + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + // "Prefetch" a stochastic rounding state for the first tile + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + // Apply round/quantize to each fragment, with or without stochastic rounding + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], cutlass::platform::numeric_limits::max()); + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale)); + } + } + + // Write quantized FP4 tile and dequant scale to gmem + copy(tiled_r2g, src, dst); + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + } else if (is_epilogue_row_quant_warp) { + // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage. + cutlass::arch::warpgroup_reg_alloc<136>(); + if constexpr (kEnableRowQuant) { + using S2RVectorType = uint128_t; + + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % 256; + size_t rng_seed = 0; + size_t rng_offset = 0; + // g2s load all global_a_amax for all groups/tensors + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueRowQuantThreadCount) { + shared_storage.global_a_amax[g] = + __ldg(reinterpret_cast(args.global_a_amax_list[g])); + } + // RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + // Input/output tensors/partitions for row quant warp + Tensor mQA = + make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA)); + Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); + + Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _), + Step<_1, X, _1>{}); // (BLK_M,BLK_N) + // Swizzled shared memory A tile, with layout + Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>( + coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) + + // Set up layouts for partitioning – tile-by-warp, with vector granularity + using S2RWarpLayout = Layout>; + using WarpGroupLayout = Layout>; + using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); + using S2RValLayout = Layout, _1>>; + using S2RAtomA = Copy_Atom; + using R2GAtomQA = Copy_Atom; + using R2GAtomSFA = Copy_Atom; + auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{}); + + auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); + auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); + auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx); + Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) + + // Allocate temporary register tensors for copying quantization => output + Tensor tQArA = make_tensor_like( + make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) + Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn); + Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); + + Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); + Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); + + int row_quant_barrier_id = 10; + cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); + + int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); + float a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + + float global_decode_scale = 1.0f / global_encode_scale; + float global_encode_scale_multiplier = 1.0f; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + auto sfa_converter = cutlass::NumericConverter{}; + do { + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset); + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Update group quantization parameters/scaling + global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + } + + auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state); + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); + cutlass::arch::fence_view_async_shared(); + mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state); + ++mainloop_pipe_consumer_state; + ++k_tile; + + // static int constexpr NumVecs = size(tQArA) / VectorSize; + cutlass::maximum_absolute_value_reduction, + true> + amax_reduction; + auto compute_frgs = reinterpret_cast *>(tQArA.data()); + auto output_frgs = + reinterpret_cast *>(raw_pointer_cast(tQArQA.data())); + Tensor amax = + make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); + Tensor pvscales = make_tensor_like(amax); + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + + tiles_in_m * tiles_in_n * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) { + auto amax_view = group_modes<1, rank(amax)>(amax); + auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales); + auto compute_frgs_up = + cutlass::NumericArrayConverter{}( + compute_frgs[v]); + amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up); + if constexpr (kUseFastMath) { + // Fast math: multiply with precomputed reciprocal + pvscales_view(_0{}, v) = cutlass::multiplies{}( + amax_view(_0{}, v), global_encode_scale_multiplier); + } else { + // Accurate math: perform division + pvscales_view(_0{}, v) = + cutlass::divides{}(amax_view(_0{}, v), fp4_max); + pvscales_view(_0{}, v) = cutlass::multiplies{}( + pvscales_view(_0{}, v), global_encode_scale); + } + filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v)); + auto qpvscale_ups = + cutlass::NumericConverter{}(filter(tQArSFA)(v)); + auto qpvscale_scaled = + cutlass::multiplies{}(qpvscale_ups, global_decode_scale); + ElementAccumulator acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = + cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + } + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales, cutlass::platform::numeric_limits::max()); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = + cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale)); + } + } + copy(tiled_r2g_QA, tQArQA, tQAgQA_mn); + copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn)); + } + // scheduler.advance(); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + + } else { + cutlass::arch::warpgroup_reg_dealloc<32>(); + } +} // NOLINT(readability/fn_size) + +template +void group_row_col_rht_gemm_ntt_w_sfc(int packed_sequence_length, int hidden_size, TA const *A, + TB const *B, TQA *QA, TSFA *SFA, + MultiAmaxHadamardCastFusionArgs &args, + const size_t *rng_state, uint32_t sm_count, + cudaStream_t stream, int k_tile_size = 1024) { + using namespace cute; + static int constexpr SFVecSize = 16; + static int constexpr RhtTensorSize = 16; + + static_assert(RhtTensorSize == 16, "RhtTensorSize must be 16"); + using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0), + make_stride(make_stride(_0{}, _1{}), 0))); + using LinearSFDLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)), + make_stride(0, make_stride(_0{}, _1{})))); + + using SwizzledSFALayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFALayout = decltype(tile_to_shape( + SwizzledSFALayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{})); + using SwizzledSFDLayout = decltype(tile_to_shape( + SwizzledSFDLayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{})); + + using SFALayout = cute::conditional_t; + using SFDLayout = cute::conditional_t; + SFALayout sfa_layout; + SFDLayout sfd_layout; + + if constexpr (kEnableSwizzleSFOutput) { + sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{}, + make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{}); + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, + make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{}); + } else { + sfa_layout = make_layout( + make_shape(make_shape(Int{}, hidden_size / SFVecSize), packed_sequence_length), + make_stride(make_stride(_0{}, _1{}), hidden_size / SFVecSize)); + sfd_layout = make_layout( + make_shape(hidden_size, make_shape(Int{}, packed_sequence_length / SFVecSize)), + make_stride(packed_sequence_length / SFVecSize, make_stride(_0{}, _1{}))); + } + + // Define shapes (dynamic) + auto M = hidden_size; + auto N = packed_sequence_length; + Tensor tensorA = make_tensor(A, make_shape(hidden_size, packed_sequence_length), LayoutLeft{}); + Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{}); + Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, packed_sequence_length), LayoutLeft{}); + Tensor tensorSFA = make_tensor(SFA, sfa_layout); + + // Define strides (from tensors) + auto dA = stride(tensorA); // (dM,dK) + auto dB = stride(tensorB); // (dN,dK) + auto dD = LayoutRight{}; // (dM,dN) + auto dQA = stride(tensorQA); // (dM,dK) + using ClusterShape = Shape<_1, _1, _1>; + auto cluster_shape = ClusterShape{}; + auto cluster_tile_shape = Shape<_128, Int, Int>{}; + auto cluster_tile_mainloop = Shape<_128, Int, _128>{}; + + // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles + static int constexpr EpilogueUnrollFactor = + size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape); + // Construct the MMA + auto mma = make_tiled_mma( + SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape), + UMMA::Major::MN, UMMA::Major::MN>{}, + Layout>{}); + + // Assert that the TiledMMA uses all CTAs in the CGA. + CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma)); + CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma))); + + // Determine the A and B shapes + auto mma_shape_B = + partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape))); + + using TiledMma = decltype(mma); + using AtomThrID = typename TiledMma::AtomThrID; + + using SmemShape_M = decltype(shape_div( + shape<0>(cluster_tile_shape), + shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{})))); + using SmemShape_N = decltype(shape_div( + shape<1>(cluster_tile_shape), + shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{})))); + using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape)); + + using SmemLayoutAtomB = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + + auto mma_shape_A = partition_shape_A( + mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop))); + using SmemShape_M_A = + decltype(shape_div(shape<0>(cluster_tile_mainloop), + shape_div(shape<0>(cluster_tile_mainloop), + size<0>(cluster_tile_mainloop) / size(AtomThrID{})))); + using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop)); + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>()); + + static uint32_t constexpr TotalTmemRows = 128; + static uint32_t constexpr Sm100TmemCapacityColumns = 512; + static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns; + static uint32_t constexpr AccumulatorPipelineStageCount = + TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape)); + + // Define the smem layouts (static) + // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory + constexpr int SchedulerPipelineStageCount = 4; + static int constexpr MainloopPipelineBytes = sizeof( + typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<1, Shape<_1, _1, _1>, + Shape<_1, _1, _1>>::SharedStorage); + + static int constexpr SchedulerWorkspaceBytes = sizeof(int) * SchedulerPipelineStageCount; + static int constexpr SchedulerThrottlePipelineBytes = + sizeof(typename cutlass::PipelineAsync::SharedStorage); + static int constexpr SchedulerPipelineBytes = + sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage); + + static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier); + static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB); + static int constexpr AccPipelineBytes = sizeof( + typename cutlass::PipelineUmmaAsync>::SharedStorage); + static int constexpr TmemBasePtrsBytes = sizeof(uint32_t); + static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes + static int constexpr kBytesPerStage = + cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes; + static int constexpr kReservedBytes = SchedulerWorkspaceBytes + SchedulerThrottlePipelineBytes + + SchedulerPipelineBytes + TmemBasePtrsBytes + + TmemDeallocBytes + BTensorBytes + + AccPipelineBytes; // Reserve for barriers and other uses + static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage; + auto sP = Int{}; // SMEM pipelines + + auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP), + Step<_2, _1, _3>{}); // (MMA,MMA_M,MMA_K,PIPE) + auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{}, + append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1) + auto sD = Layout<_1>{}; // XXX Dummy + + auto tma_load_a = + make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma); + auto tma_load_b = + make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cluster_tile_shape, mma); + + // Assert checks on tile sizes -- no predication + assert(M % size<0>(cluster_tile_shape) == 0); + assert(N % size<1>(cluster_tile_shape) == 0); + + dim3 dimBlock(512); + dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape)); + dim3 dimGrid(sm_count, 1, 1); + + int smem_size = sizeof( + SharedStorage); + + auto *kernel_ptr = &group_row_col_rht_gemm_device< + decltype(M), decltype(N), decltype(k_tile_size), decltype(cluster_shape), + decltype(cluster_tile_shape), TA, decltype(dA), decltype(sA), decltype(tma_load_a), TB, + decltype(dB), decltype(sB), decltype(tma_load_b), TD, decltype(dD), decltype(sD), TSFD, + decltype(sfd_layout), TQA, decltype(dQA), TSFA, decltype(sfa_layout), decltype(mma), + AccumulatorPipelineStageCount, SchedulerPipelineStageCount, kEnableStochasticRounding, + kEnableRHTColQuant, kEnableRowQuant, kEnableSwizzleSFOutput, kUseFastMath>; + + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + // Allocate workspace and set to zero + void *tile_scheduler_workspace = nullptr; + NVTE_CHECK_CUDA(cudaMallocAsync(&tile_scheduler_workspace, sizeof(uint32_t), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(tile_scheduler_workspace, 0, sizeof(uint32_t), stream)); + + // Launch kernel + cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream}; + cutlass::Status status = cutlass::launch_kernel_on_cluster( + params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, A, dA, + sA, tma_load_a, B, dB, sB, tma_load_b, QA, dQA, SFA, sfa_layout, args, + tile_scheduler_workspace, mma, rng_state); + NVTE_CHECK_CUDA(cudaGetLastError()); + NVTE_CHECK(status == cutlass::Status::kSuccess, "Kernel launch failed."); + + NVTE_CHECK_CUDA(cudaFreeAsync(tile_scheduler_workspace, stream)); +} + +} // namespace +} // namespace detail + +void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector &output_list, + const size_t *split_sections, size_t num_tensors, + const Tensor &hadamard_matrix_, + QuantizationConfig &quant_config, cudaStream_t stream) { + NVTE_API_CALL(group_hadamard_transform_cast_fusion); + + using transformer_engine::detail::kMaxTensorsPerKernel; + using transformer_engine::detail::MultiAmaxHadamardCastFusionArgs; + + NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16, + "Input tensor must be BF16 tensor, but dtype is ", to_string(input_.dtype()), "."); + NVTE_CHECK(input_.dim() >= 2, "Input must be a 2D tensor."); + const SimpleTensor &input = input_.data; + + NVTE_CHECK(output_list.size() == num_tensors, + "Number of output tensors should match number of tensors."); + + NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel, + "Number of tensors should be less than or equal to ", kMaxTensorsPerKernel); + + // construct the multi-tensor args + MultiAmaxHadamardCastFusionArgs kernel_args; + kernel_args.num_tensors = 0; + kernel_args.split_sections_range[0] = 0; + bool all_has_row_quant = true; + bool all_has_col_quant = true; + void *rowwise_data_base_ptr = nullptr; + void *rowwise_scale_inv_base_ptr = nullptr; + for (size_t i = 0; i < num_tensors; ++i) { + NVTE_CHECK(split_sections[i] % 128 == 0, "component ", i, + " of split_sections should be 128 multiple"); + if (split_sections[i] == 0) { + continue; + } + bool has_row_quant = output_list[i]->data.dptr != nullptr; + bool has_col_quant = output_list[i]->columnwise_data.dptr != nullptr; + all_has_row_quant = all_has_row_quant && has_row_quant; + all_has_col_quant = all_has_col_quant && has_col_quant; + // sanity check, the two bool flags cannot be both false + NVTE_CHECK(has_row_quant || has_col_quant, + "At least one of the output tensors must have row or column quant."); + void *amax_rowwise_ptr = + has_row_quant ? reinterpret_cast(output_list[i]->amax.dptr) : nullptr; + void *amax_colwise_ptr = + has_col_quant ? reinterpret_cast(output_list[i]->columnwise_amax.dptr) : nullptr; + void *rowwise_data_ptr = + has_row_quant ? reinterpret_cast(output_list[i]->data.dptr) : nullptr; + void *rowwise_scale_inv_ptr = + has_row_quant ? reinterpret_cast(output_list[i]->scale_inv.dptr) : nullptr; + if (all_has_row_quant && + (rowwise_data_base_ptr == nullptr || rowwise_scale_inv_base_ptr == nullptr)) { + rowwise_data_base_ptr = rowwise_data_ptr; + rowwise_scale_inv_base_ptr = rowwise_scale_inv_ptr; + } + void *output_colwise_ptr = + has_col_quant ? reinterpret_cast(output_list[i]->columnwise_data.dptr) : nullptr; + void *output_colwise_scale_inv_ptr = + has_col_quant ? reinterpret_cast(output_list[i]->columnwise_scale_inv.dptr) + : nullptr; + kernel_args.global_a_amax_list[kernel_args.num_tensors] = amax_rowwise_ptr; + kernel_args.global_d_amax_list[kernel_args.num_tensors] = amax_colwise_ptr; + kernel_args.output_colwise_list[kernel_args.num_tensors] = output_colwise_ptr; + kernel_args.output_colwise_scale_inv_list[kernel_args.num_tensors] = + output_colwise_scale_inv_ptr; + kernel_args.split_sections[kernel_args.num_tensors] = split_sections[i]; + kernel_args.split_sections_range[kernel_args.num_tensors + 1] = + kernel_args.split_sections_range[kernel_args.num_tensors] + split_sections[i]; + kernel_args.num_tensors++; + } + + // Stochastic rounding config + const bool use_stochastic_rounding = quant_config.stochastic_rounding; + const size_t *rng_state = nullptr; + if (use_stochastic_rounding) { + NVTE_CHECK(quant_config.rng_state != nullptr, + "Enabled stochastic rounding without providing RNG state"); + 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}, + "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); + rng_state = reinterpret_cast(rng_state_tensor.data.dptr); + } + + // Template arguments + using TA = cute::bfloat16_t; + using TB = cute::bfloat16_t; + using TD = cutlass::float_e2m1_t; + using TSFD = cutlass::float_ue4m3_t; + using TQA = TD; + using TSFA = TSFD; + + checkCuDriverContext(stream); + + // Check Hadamard matrix + constexpr int kHadamardDimension = 16; + + NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16, + "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(), "."); + const size_t hadamard_dimension = hadamard_matrix.shape[0]; + + const size_t ndim = input.shape.size(); + const size_t n = input.shape[ndim - 1]; + size_t m = 1; + for (size_t i = 0; i < ndim - 1; ++i) { + m *= input.shape[i]; + } + + auto sm_count = transformer_engine::cuda::sm_count(); + + NVTE_CHECK(n % hadamard_dimension == 0, "row_length must be divisible by hadamard_dimension."); + + NVTE_CHECK(m % hadamard_dimension == 0, "num_rows must be divisible by hadamard_dimension"); + + int k_tile_size = 1024; + + const bool use_swizzle_sf_output = false; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_stochastic_rounding, kEnableStochasticRounding, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + all_has_col_quant, kEnableRhtColQuant, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + all_has_row_quant, kEnableRowQuant, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_swizzle_sf_output, kEnableSwizzleSFOutput, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + + if constexpr (kEnableRhtColQuant || kEnableRowQuant) { + detail::group_row_col_rht_gemm_ntt_w_sfc< + kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, + kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>( + /*packed_sequence_length=*/m, /*hidden_size=*/n, + /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*QA=*/reinterpret_cast(rowwise_data_base_ptr), + /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), + /*args=*/kernel_args, + /*rng_state=*/rng_state, /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size); + } else { + NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", + kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, ")."); + } + + ););););); +} + +} // namespace transformer_engine + +void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETensor *outputs, + const NVTETensor hadamard_matrix, + const size_t *split_sections, + const size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion); + using namespace transformer_engine; + NVTE_CHECK(num_tensors > 0, "Number of tensors should be greater than 0."); + + Tensor *input_tensor = convertNVTETensorCheck(input); + std::vector output_list(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + output_list[i] = convertNVTETensorCheck(outputs[i]); + } + + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Call the multi-tensor Hadamard transform amax implementation. + group_hadamard_transform_cast_fusion(*input_tensor, output_list, split_sections, num_tensors, + *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, + stream); +} 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 12f02dba6b..11325041ae 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -29,7 +29,6 @@ #include "cutlass/pipeline/pipeline.hpp" #include "cutlass/util/GPU_Clock.hpp" #include "cutlass/util/command_line.h" -#include "cutlass/util/helper_cuda.hpp" #include "cutlass/util/print_error.hpp" // clang-format off @@ -129,7 +128,8 @@ template + bool kEnableStochasticRounding = false, + bool kUseFastMath = false> __global__ static void rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, @@ -426,7 +426,13 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, const float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); const float global_decode_scale = 1.0f / global_encode_scale; - auto sfd_converter = cutlass::NumericConverter{}; + + // Scaling factor for fast math path + float global_encode_scale_multiplier = 1.0f; + if constexpr (kUseFastMath) { + static constexpr float fp4_max_inv = 1.0f / fp4_max; + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } do { for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) { @@ -469,10 +475,13 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, ++accumulator_pipe_consumer_state; - // Cast data from FP32 to BF16 to FP32. - auto convert_accum_to_bf16 = cutlass::NumericArrayConverter{}; - auto convert_bf16_to_accum = cutlass::NumericArrayConverter{}; - tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with unfused + // kernels + auto convert_accum_to_bf16 = cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + } auto compute_frgs = reinterpret_cast *>(tTR_rAcc_frag.data()); auto output_frgs = reinterpret_cast *>(tDrC_frag.data()); @@ -481,14 +490,27 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); } - pvscales = cutlass::divides>{}(vec_maxs, fp4_max); - pvscales = cutlass::multiplies>{}(pvscales, global_encode_scale); + if constexpr (kUseFastMath) { + // Fast math: multiply with precomputed reciprocal + pvscales = cutlass::multiplies>{}(vec_maxs, global_encode_scale_multiplier); + } else { + // Accurate math: perform division + pvscales = cutlass::divides>{}(vec_maxs, fp4_max); + pvscales = cutlass::multiplies>{}(pvscales, global_encode_scale); + } auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales); tC_rRowSFD_frg(_0{}) = pvscales_cvted; auto qpvscale_ups = cutlass::NumericArrayConverter{}(tC_rRowSFD_frg(_0{})); auto qpvscale_scaled = cutlass::multiplies>{}(qpvscale_ups, global_decode_scale); - auto acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled); + cutlass::Array acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled); + } // Initialize RNG for tile const size_t rng_sequence @@ -532,7 +554,7 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, // B: 16 x 16: row-major // C: m x n: row-major // SFC: m x (n/16): row-major -template +template void rht_gemm_ntt_w_sfc(int m, int n, TA const* A, @@ -644,16 +666,15 @@ rht_gemm_ntt_w_sfc(int m, int n, TC, decltype(dC), decltype(sC), TSFC, decltype(mma), - kEnableStochasticRounding>; + kEnableStochasticRounding, + kUseFastMath>; - bool status = cudaFuncSetAttribute(*kernel_ptr, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_size); + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(*kernel_ptr, + cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_size) + ); - if (status != cudaSuccess) { - std::cerr << "Error: Failed to set Shared Memory size." << std::endl; - return; - } (*kernel_ptr) <<< dimGrid, dimBlock, smem_size, stream >>> (M, N, k_tile_size, cga_tile_shape, @@ -663,11 +684,12 @@ rht_gemm_ntt_w_sfc(int m, int n, SFC, mma, global_amax, rng_state); + NVTE_CHECK_CUDA(cudaGetLastError()); } // this function is used to wrap the rht_gemm_ntt_w_sfc function //to transpose the input tensor A -template +template void rht_gemm_ttt_wrapper(int m, int n, TA const* A, @@ -690,7 +712,7 @@ rht_gemm_ttt_wrapper(int m, int n, // B: 16 x 16: row-major // C: n x m: row-major // SFC: n x (m/16): row-major - rht_gemm_ntt_w_sfc( + rht_gemm_ntt_w_sfc( n, m, A, B, C, SFC, global_amax, @@ -800,20 +822,23 @@ void hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &out } else if (m < 1024 || n < 1024) { k_tile_size = 512; } + TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, kUseStochasticRounding, - detail::rht_gemm_ttt_wrapper( - /*m=*/m, - /*n=*/n, - /*A=*/reinterpret_cast(input.dptr), - /*B=*/reinterpret_cast(hadamard_matrix.dptr), - /*C=*/reinterpret_cast(output_t.dptr), - /*SFC=*/reinterpret_cast(scale_inv_t.dptr), - /*global_amax=*/reinterpret_cast(global_amax.dptr), - /*rng_state=*/rng_state, - /*sm_count=*/sm_count, - /*stream=*/stream, - /*k_tile_size=*/k_tile_size);); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + detail::rht_gemm_ttt_wrapper( + /*m=*/m, + /*n=*/n, + /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*C=*/reinterpret_cast(output_t.dptr), + /*SFC=*/reinterpret_cast(scale_inv_t.dptr), + /*global_amax=*/reinterpret_cast(global_amax.dptr), + /*rng_state=*/rng_state, + /*sm_count=*/sm_count, + /*stream=*/stream, + /*k_tile_size=*/k_tile_size););); } } // namespace transformer_engine diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index a3235e84f1..19fbe431aa 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -270,6 +270,20 @@ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, const NVTEQuantizationConfig quant_config, const size_t num_tensors, cudaStream_t stream); +/*! \brief Casts grouped input tensor to quantized output tensors. + * + * \param[in] input Input tensor to be cast. + * \param[in,out] outputs Output quantized tensors. + * \param[in] split_sections Split sections of the input tensor. + * \param[in] num_tensors Number of output tensors. + * \param[in] quant_config (Optional) Quantization configurations. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *outputs, + const size_t *split_sections, size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index 05541fe30c..112cb9b54d 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -86,6 +86,43 @@ void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outp int random_sign_mask, int random_sign_mask_t, cudaStream_t stream); +/*! + * \brief Perform the grouped-tensor columnwise Hadamard transform cast fusion operation. + * + * This function is experimental and the API is not stable. Group_ prefix means contiguous input concatenated + * + * \param[in] input Input tensor to apply Hadamard transform. + * \param[in,out] outputs Array of output tensors. + * \param[in] hadamard_matrix Hadamard matrix to use for transformation. + * \param[in] split_sections Array specifying splits in dimension 0 for each output tensor. + * \param[in] num_tensors Number of output tensors, must be > 0. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_hadamard_transform_cast_fusion_columnwise( + const NVTETensor input, NVTETensor* outputs, const NVTETensor hadamard_matrix, + const size_t* split_sections, size_t num_tensors, const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + +/*! + * \brief Perform the grouped-tensor row quantize (without Hadamard) and columnwise Hadamard transform cast fusion operation. + * + * This function is experimental and the API is not stable. Group_ prefix means contiguous input concatenated + * + * \param[in] input Input tensor to apply Hadamard transform. + * \param[in,out] outputs Array of output tensors. + * \param[in] hadamard_matrix Hadamard matrix to use for transformation. + * \param[in] split_sections Array specifying splits in dimension 0 for each output tensor. + * \param[in] num_tensors Number of output tensors, must be > 0. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETensor* outputs, + const NVTETensor hadamard_matrix, + const size_t* split_sections, size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index b2e04ba69f..19cb646be2 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -337,6 +337,12 @@ enum NVTEQuantizationConfigAttribute { kNVTEQuantizationConfigNVFP42DQuantization = 5, /*! Whether to enable stochastic rounding */ kNVTEQuantizationConfigStochasticRounding = 6, + /*! Whether to enable fast math operations with reduced accuracy. + * + * Optimizations are kernel-specific and they may be applied + * inconsistently between kernels. + */ + kNVTEQuantizationConfigUseFastMath = 7, kNVTEQuantizationConfigNumAttributes }; @@ -997,6 +1003,12 @@ class QuantizationConfigWrapper { &stochastic_rounding, sizeof(bool)); } + /*! \brief Set whether to enable fast math operations */ + void set_use_fast_math(bool use_fast_math) { + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigUseFastMath, + &use_fast_math, sizeof(bool)); + } + private: /*! \brief Wrapped NVTEQuantizationConfig. */ NVTEQuantizationConfig config_ = nullptr; diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 8d9563b789..4a140b4376 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -857,9 +857,10 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, // Write attribute size NVTE_CHECK(attr < kNVTEQuantizationConfigNumAttributes, "Invalid NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); - NVTE_CHECK(size_written != nullptr, "Invalid size_written (got NULL)"); const auto &attr_size = transformer_engine::QuantizationConfig::attr_sizes[attr]; - *size_written = attr_size; + if (size_written != nullptr) { + *size_written = attr_size; + } // Return immediately if buffer is not provided if (buf == nullptr) { @@ -889,6 +890,18 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigFloat8BlockScaleTensorFormat: std::memcpy(buf, &config_.float8_block_scale_tensor_format, attr_size); break; + case kNVTEQuantizationConfigRNGState: + std::memcpy(buf, &config_.rng_state, attr_size); + break; + case kNVTEQuantizationConfigNVFP42DQuantization: + std::memcpy(buf, &config_.nvfp4_2d_quantization, attr_size); + break; + case kNVTEQuantizationConfigStochasticRounding: + std::memcpy(buf, &config_.stochastic_rounding, attr_size); + break; + case kNVTEQuantizationConfigUseFastMath: + std::memcpy(buf, &config_.use_fast_math, attr_size); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } @@ -933,6 +946,9 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigStochasticRounding: std::memcpy(&config_.stochastic_rounding, buf, attr_size); break; + case kNVTEQuantizationConfigUseFastMath: + std::memcpy(&config_.use_fast_math, buf, attr_size); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index ac541435c7..aa9d800c7b 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -16,6 +16,7 @@ #include "../extensions.h" #include "common.h" +#include "common/util/system.h" #include "pybind.h" #include "transformer_engine/transformer_engine.h" @@ -709,179 +710,225 @@ std::tuple, std::vector, bool> bulk_alloc return retval; } -void split_quantize_nvfp4_impl(const TensorWrapper &input, - const std::vector &input_list, - std::vector &output_list, - const std::vector &split_sections, - const std::vector &quantizers) { - // Check tensor lists - const size_t num_tensors = split_sections.size(); - NVTE_CHECK(input_list.size() == num_tensors, "Expected ", num_tensors, " input tensors, but got ", - input_list.size(), "."); - NVTE_CHECK(output_list.size() == num_tensors, "Expected ", num_tensors, - " output tensors, but got ", output_list.size(), "."); - NVTE_CHECK(quantizers.size() == num_tensors, "Expected ", num_tensors, - " NVFP4 quantizers, but got ", quantizers.size(), "."); +// Owns all allocations/wrappers backing quant_config_list[*].set_rng_state(...). +struct StochasticRngStateResources { + at::Tensor rng_states_tensor; // [2 * num_tensors], int64, CUDA + at::Tensor rng_states_tensor_colwise; // optional, same shape/dtype/device + std::vector te_rng_state_list; + std::vector te_rng_state_list_colwise; + + bool enabled{false}; + bool need_separate_rng_states{false}; + bool with_bulk_generate_rng_states{false}; +}; + +// Populates quant_config_list (+ optional colwise list) with rng_state pointers and stochastic flag. +static StochasticRngStateResources setup_stochastic_rounding_rng_states_helper( + size_t num_tensors, bool stochastic_rounding, bool with_bulk_generate_rng_states, + bool need_separate_rng_states, + std::vector &quant_config_list_rowwise, + std::vector &quant_config_list_colwise) { + // the return object will be used to keep rng states alive + StochasticRngStateResources res; + res.enabled = stochastic_rounding; + res.need_separate_rng_states = need_separate_rng_states; + res.with_bulk_generate_rng_states = with_bulk_generate_rng_states; + + if (!stochastic_rounding) return res; + + // Basic sanity: caller usually pre-sizes these to num_tensors. + TORCH_CHECK(quant_config_list_rowwise.size() == num_tensors, + "quant_config_list_rowwise must be sized to num_tensors"); + if (need_separate_rng_states) { + TORCH_CHECK(quant_config_list_colwise.size() == num_tensors, + "quant_config_list_colwise must be sized to num_tensors when " + "need_separate_rng_states=true"); + } - // Trivial cases - if (num_tensors == 0) { - return; + const size_t rng_elts_per_thread = + res.with_bulk_generate_rng_states ? (1024 * num_tensors) : 1024; + + auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); + res.rng_states_tensor = torch::empty({static_cast(2 * num_tensors)}, opts); + if (need_separate_rng_states) { + res.rng_states_tensor_colwise = torch::empty({static_cast(2 * num_tensors)}, opts); } - if (input.numel() == 0) { - for (const auto &tensor : input_list) { - NVTE_CHECK(tensor.numel() == 0, - "Input tensor has zero elements but got split with non-zero elements"); + + res.te_rng_state_list.reserve(num_tensors); + if (need_separate_rng_states) res.te_rng_state_list_colwise.reserve(num_tensors); + + for (size_t i = 0; i < num_tensors; ++i) { + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + + // Rowwise RNG state + at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); + int64_t *rng_state_ptr = static_cast(res.rng_states_tensor.data_ptr()) + i * 2; + philox_unpack(philox_args, rng_state_ptr); + + res.te_rng_state_list.push_back(makeTransformerEngineTensor( + static_cast(rng_state_ptr), std::vector{2}, DType::kInt64)); + quant_config_list_rowwise[i].set_rng_state(res.te_rng_state_list[i].data()); + quant_config_list_rowwise[i].set_stochastic_rounding(true); + + // Colwise RNG state (only if you truly need a different sequence) + if (need_separate_rng_states) { + // re-initialize philox_args for colwise RNG state + at::PhiloxCudaState philox_args_col = init_philox_state(gen, rng_elts_per_thread); + int64_t *rng_state_ptr_colwise = + static_cast(res.rng_states_tensor_colwise.data_ptr()) + i * 2; + + philox_unpack(philox_args_col, rng_state_ptr_colwise); + + res.te_rng_state_list_colwise.push_back(makeTransformerEngineTensor( + static_cast(rng_state_ptr_colwise), std::vector{2}, DType::kInt64)); + quant_config_list_colwise[i].set_rng_state(res.te_rng_state_list_colwise[i].data()); + quant_config_list_colwise[i].set_stochastic_rounding(true); } - return; - } - // Assume all quantizers have identical config - const auto &quantizer = *quantizers.front(); - NVTE_CHECK(!quantizer.with_2d_quantization, - "NVFP4 split-quantize does not support 2D quantization"); - NVTE_CHECK(!quantizer.with_amax_reduction, - "NVFP4 split-quantize does not support amax reduction"); + // break the loop if we are using bulk generate rng states + if (res.with_bulk_generate_rng_states) break; + } - // Check input tensor shape - const size_t input_last_dim = input.ndim() > 0 ? input.size(input.ndim() - 1) : 1; - NVTE_CHECK(input_last_dim % 128 == 0, - "NVFP4 multi-quantize requires inner dim to be multiple of 128."); + return res; +} - // CUDA stream - auto stream = at::cuda::getCurrentCUDAStream(); +// Implements split-quantize NVFP4 with Row/Column-wise Hadamard Transform (RHT) +void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, + const std::vector &input_list, + std::vector &output_list, + const std::vector &split_sections, + const std::vector &quantizers, + cudaStream_t stream) { + const size_t num_tensors = split_sections.size(); + const auto &quantizer = *quantizers.front(); - // Objects for TE C API std::vector nvte_tensor_input_list; std::vector nvte_tensor_output_list; - std::vector quant_config_list; for (size_t i = 0; i < num_tensors; ++i) { nvte_tensor_input_list.push_back(input_list[i].data()); nvte_tensor_output_list.push_back(output_list[i].data()); + } + + // trigger the row-col fusion when the split-sections shapes are all 128 aligned for max performance + bool all_aligned_token_dim = + std::all_of(split_sections.begin(), split_sections.end(), + [](size_t split_section) { return split_section % 128 == 0; }); + + // in the case when rowwise and colwise cannot be fused, we have to generate the RNG states twice + // so that rowwise and colwise will have different random numbers + bool need_separate_rng_states = + (!all_aligned_token_dim) && quantizer.rowwise_usage && quantizer.columnwise_usage; + + // Objects for TE C API + std::vector quant_config_list; + std::vector quant_config_list_colwise; + for (size_t i = 0; i < num_tensors; ++i) { quant_config_list.emplace_back(QuantizationConfigWrapper()); + quant_config_list_colwise.emplace_back(QuantizationConfigWrapper()); } + // this is true because we have already built grouped kernels for rowwise and colwise quantization with RHT + bool with_bulk_generate_rng_states = true; + // Stochastic rounding - // When both rowwise and columnwise quantization are used, - // we need separate RNG states for each to ensure they use different random numbers. - std::vector te_rng_state_list; - std::vector te_rng_state_columnwise_list; - std::vector quant_config_columnwise_list; - at::Tensor rng_states_tensor; - at::Tensor rng_states_columnwise_tensor; - const bool need_separate_columnwise_rng = - quantizer.stochastic_rounding && quantizer.with_rht && quantizer.columnwise_usage; - - if (quantizer.stochastic_rounding) { - // TODO(zhongbo): remove the for loop of generating rng states with a single call - // with rng_elts_per_thread = 1024 * num_tensors - // Change to the bulk generate rng states api when grouped quantize is available - const size_t rng_elts_per_thread = 1024; // Wild guess, probably can be tightened - auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); - rng_states_tensor = torch::empty({static_cast(2 * num_tensors)}, opts); - - // Allocate columnwise RNG resources when separate RNG is needed - if (need_separate_columnwise_rng) { - rng_states_columnwise_tensor = torch::empty({static_cast(2 * num_tensors)}, opts); - for (size_t i = 0; i < num_tensors; ++i) { - quant_config_columnwise_list.emplace_back(QuantizationConfigWrapper()); - } + bool need_stochastic_rounding = quantizer.stochastic_rounding; + auto stochastic_rng_state_resources = setup_stochastic_rounding_rng_states_helper( + num_tensors, need_stochastic_rounding, with_bulk_generate_rng_states, + need_separate_rng_states, quant_config_list, quant_config_list_colwise); + + // Enable NVFP4 kernels to use math operations that sacrifice + // accuracy for performance. These optimizations are experimental + // and inconsistently implemented. + const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); + if (use_fast_math) { + for (auto &config : quant_config_list) { + config.set_use_fast_math(true); } - for (size_t i = 0; i < num_tensors; ++i) { - auto gen = at::get_generator_or_default( - std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); - // Generate RNG state for rowwise quantization - at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); - int64_t *rng_state_ptr = static_cast(rng_states_tensor.data_ptr()) + i * 2; - philox_unpack(philox_args, rng_state_ptr); - te_rng_state_list.push_back(makeTransformerEngineTensor( - static_cast(rng_state_ptr), std::vector{2}, DType::kInt64)); - quant_config_list[i].set_rng_state(te_rng_state_list[i].data()); - quant_config_list[i].set_stochastic_rounding(true); - - // Generate separate RNG state for columnwise quantization - if (need_separate_columnwise_rng) { - at::PhiloxCudaState philox_args_columnwise = init_philox_state(gen, rng_elts_per_thread); - int64_t *rng_state_columnwise_ptr = - static_cast(rng_states_columnwise_tensor.data_ptr()) + i * 2; - philox_unpack(philox_args_columnwise, rng_state_columnwise_ptr); - te_rng_state_columnwise_list.push_back(makeTransformerEngineTensor( - static_cast(rng_state_columnwise_ptr), std::vector{2}, DType::kInt64)); - quant_config_columnwise_list[i].set_rng_state(te_rng_state_columnwise_list[i].data()); - quant_config_columnwise_list[i].set_stochastic_rounding(true); - } + for (auto &config : quant_config_list_colwise) { + config.set_use_fast_math(true); } } - // Perform multi-tensor quantization - if (quantizer.with_rht) { // Quantize row-wise data, RHT+quantize column-wise data - // Check that config is supported - NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); - - // Compute amaxes - if (quantizer.with_post_rht_amax) { - // We need: - // 1. Rowwise amax = amax for input - // 2. Columnwise amax = amax for RHT(input.t) - NVTE_SCOPED_GIL_RELEASE({ - nvte_group_hadamard_transform_amax( - input.data(), reinterpret_cast(nvte_tensor_output_list.data()), - split_sections.data(), num_tensors, 0, quantizer.rht_matrix_random_sign_mask_t, stream); - }); - } else { - // RHT is enabled, but amax is pre-RHT amax - NVTE_ERROR("NVFP4 split-quantize does not yet support pre-RHT amax"); - } + auto &quant_config_list_colwise_to_use = + need_separate_rng_states ? quant_config_list_colwise : quant_config_list; - // Check that RHT matrix is available - NVTE_CHECK(quantizer.rht_matrix.defined() && quantizer.rht_matrix.numel() > 0, - "RHT matrix is not available."); - auto rht_matrix_nvte = makeTransformerEngineTensor(quantizer.rht_matrix); + // Compute amaxes + if (quantizer.with_post_rht_amax) { + // We need: + // 1. Rowwise amax = amax for input + // 2. Columnwise amax = amax for RHT(input.t) + nvte_group_hadamard_transform_amax( + input.data(), reinterpret_cast(nvte_tensor_output_list.data()), + split_sections.data(), num_tensors, 0, quantizer.rht_matrix_random_sign_mask_t, stream); + } else { + // RHT is enabled, but amax is pre-RHT amax + NVTE_ERROR("NVFP4 split-quantize does not yet support pre-RHT amax"); + } - // Quantize tensors individually - NVTE_SCOPED_GIL_RELEASE({ - for (size_t i = 0; i < num_tensors; i++) { - if (input_list[i].numel() == 0) { - continue; // Skip tensors with no elements - } + // Check that RHT matrix is available + NVTE_CHECK(quantizer.rht_matrix.defined() && quantizer.rht_matrix.numel() > 0, + "RHT matrix is not available."); + auto rht_matrix_nvte = makeTransformerEngineTensor(quantizer.rht_matrix); - // Direct NVFP4 quantization for row-wise data - if (quantizer.rowwise_usage) { - auto out_rowwise_data = output_list[i].get_rowwise_data(); - auto out_rowwise_scale_inv = output_list[i].get_rowwise_scale_inv(); - auto out_rowwise_amax = output_list[i].get_amax(); - TensorWrapper out_rowwise(output_list[i].scaling_mode()); - out_rowwise.set_rowwise_data(out_rowwise_data.data_ptr, - static_cast(out_rowwise_data.dtype), - out_rowwise_data.shape); - out_rowwise.set_rowwise_scale_inv(out_rowwise_scale_inv.data_ptr, - static_cast(out_rowwise_scale_inv.dtype), - out_rowwise_scale_inv.shape); - out_rowwise.set_amax(out_rowwise_amax.data_ptr, - static_cast(out_rowwise_amax.dtype), out_rowwise_amax.shape); - nvte_quantize_v2(input_list[i].data(), out_rowwise.data(), quant_config_list[i], stream); + if (all_aligned_token_dim) { + // call the fully-fused grouped kernel for rowwise quantization & colwise RHT quantization transpose + nvte_group_hadamard_transform_cast_fusion( + input.data(), reinterpret_cast(nvte_tensor_output_list.data()), + rht_matrix_nvte.data(), split_sections.data(), num_tensors, quant_config_list[0], stream); + } else { + // Separate quantization for rowwise usage and columnwise usage + // Rowwise quantization fusion with grouped version + if (quantizer.rowwise_usage) { + std::vector out_identity_list; + std::vector nvte_tensor_out_identity_list; + for (size_t i = 0; i < num_tensors; i++) { + bool is_empty_split = input_list[i].numel() == 0; + TensorWrapper out_identity(output_list[i].scaling_mode()); + auto out_identity_data = output_list[i].get_rowwise_data(); + auto out_identity_scale_inv = output_list[i].get_rowwise_scale_inv(); + auto out_identity_amax = output_list[i].get_amax(); + if (!is_empty_split) { + out_identity.set_rowwise_data(out_identity_data.data_ptr, + static_cast(out_identity_data.dtype), + out_identity_data.shape); + out_identity.set_rowwise_scale_inv(out_identity_scale_inv.data_ptr, + static_cast(out_identity_scale_inv.dtype), + out_identity_scale_inv.shape); + out_identity.set_amax(out_identity_amax.data_ptr, + static_cast(out_identity_amax.dtype), + out_identity_amax.shape); } + out_identity_list.emplace_back(std::move(out_identity)); + nvte_tensor_out_identity_list.push_back(out_identity_list.back().data()); + } + nvte_group_nvfp4_quantize_with_amax(input.data(), nvte_tensor_out_identity_list.data(), + split_sections.data(), num_tensors, quant_config_list[0], + stream); + } - // RHT + NVFP4 quantize for column-wise data - if (quantizer.columnwise_usage) { - // Get the output column-wise data, scale_inv, and amax - auto out_columnwise_data = output_list[i].get_columnwise_data(); - auto out_columnwise_scale_inv = output_list[i].get_columnwise_scale_inv(); - auto out_columnwise_amax = output_list[i].get_columnwise_amax(); - - // Flatten column-wise data to 2D + // Columnwise RHT quantization fusion with grouped version + if (quantizer.columnwise_usage) { + std::vector out_transpose_list; + std::vector nvte_tensor_out_transpose_list; + for (size_t i = 0; i < num_tensors; i++) { + bool is_empty_split = input_list[i].numel() == 0; + auto out_columnwise_data = output_list[i].get_columnwise_data(); + auto out_columnwise_scale_inv = output_list[i].get_columnwise_scale_inv(); + auto out_columnwise_amax = output_list[i].get_columnwise_amax(); + + // 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 i = 1; i < colwise_data_shape.ndim; ++i) { - last_dim *= colwise_data_shape.data[i]; + 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); - // 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(output_list[i].scaling_mode()); out_transpose.set_rowwise_data(out_columnwise_data.data_ptr, static_cast(out_columnwise_data.dtype), colwise_data_shape_2d); @@ -891,53 +938,151 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, out_transpose.set_amax(out_columnwise_amax.data_ptr, static_cast(out_columnwise_amax.dtype), out_columnwise_amax.shape); - - // RHT + NVFP4 quantize kernel - // Use separate RNG state for columnwise to ensure different random numbers than rowwise - auto &columnwise_quant_config = - need_separate_columnwise_rng ? quant_config_columnwise_list[i] : quant_config_list[i]; - nvte_hadamard_transform_cast_fusion_columnwise(input_list[i].data(), out_transpose.data(), - rht_matrix_nvte.data(), - columnwise_quant_config, stream); } + out_transpose_list.emplace_back(std::move(out_transpose)); + nvte_tensor_out_transpose_list.push_back(out_transpose_list.back().data()); } - }); - - } else { // NVFP4 quantize - // We need: - // 1. Rowwise amax = amax for input - // 2. Columnwise amax = amax for input too - // Columnwise amax will be filled with a fused D2D copy from rowwise amax - // Note that the multi compute amax API expects rowwise amax pointer to be not null - // So we need to set the pointer accordingly to make colwise-only quantization work - std::vector orig_amax_ptr_list; - for (size_t i = 0; i < num_tensors; i++) { - auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr; - orig_amax_ptr_list.push_back(rowwise_amax_ptr); - auto columnwise_amax_ptr = output_list[i].get_columnwise_amax().data_ptr; - void *amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr; - NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer"); - output_list[i].set_amax(amax_ptr, DType::kFloat32, std::vector{1}); + nvte_group_hadamard_transform_cast_fusion_columnwise( + input.data(), reinterpret_cast(nvte_tensor_out_transpose_list.data()), + rht_matrix_nvte.data(), split_sections.data(), num_tensors, + quant_config_list_colwise_to_use[0], stream); } - NVTE_SCOPED_GIL_RELEASE({ - nvte_group_amax(input.data(), reinterpret_cast(nvte_tensor_output_list.data()), - split_sections.data(), num_tensors, stream); - }); - for (size_t i = 0; i < num_tensors; i++) { - output_list[i].set_amax(orig_amax_ptr_list[i], DType::kFloat32, std::vector{1}); + } +} + +void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, + const std::vector &input_list, + std::vector &output_list, + const std::vector &split_sections, + const std::vector &quantizers, + cudaStream_t stream) { + const size_t num_tensors = input_list.size(); + const auto &quantizer = *quantizers.front(); + + std::vector nvte_tensor_input_list; + std::vector nvte_tensor_output_list; + for (size_t i = 0; i < num_tensors; ++i) { + nvte_tensor_input_list.push_back(input_list[i].data()); + nvte_tensor_output_list.push_back(output_list[i].data()); + } + + // In this case without RHT, the rowwise and colwise quantization are fused + // we don't need separate rng states for rowwise and colwise + bool need_separate_rng_states = false; + + // Objects for TE C API + std::vector quant_config_list; + for (size_t i = 0; i < num_tensors; ++i) { + quant_config_list.emplace_back(QuantizationConfigWrapper()); + } + + // TODO: this is only true because the non-RHT path doesn't have grouped kernels yet, which we can be optimized + // so that we can generate all rng states at once + bool with_bulk_generate_rng_states = false; + + bool need_stochastic_rounding = quantizer.stochastic_rounding; + + // place holder for colwise rng states, which are not needed in this case + std::vector dummy_quant_config_list_colwise; + + auto stochastic_rng_state_resources = setup_stochastic_rounding_rng_states_helper( + num_tensors, need_stochastic_rounding, with_bulk_generate_rng_states, + need_separate_rng_states, quant_config_list, + dummy_quant_config_list_colwise); // colwise rng states are not needed in this case + + // We need: + // 1. Rowwise amax = amax for input + // 2. Columnwise amax = amax for input too + // Columnwise amax will be filled with a fused D2D copy from rowwise amax + // Note that the multi compute amax API expects rowwise amax pointer to be not null + // So we need to set the pointer accordingly to make colwise-only quantization work + std::vector orig_amax_ptr_list; + for (size_t i = 0; i < num_tensors; i++) { + auto rowwise_amax_ptr = output_list[i].get_amax().data_ptr; + orig_amax_ptr_list.push_back(rowwise_amax_ptr); + auto columnwise_amax_ptr = output_list[i].get_columnwise_amax().data_ptr; + void *amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr; + NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer"); + output_list[i].set_amax(amax_ptr, DType::kFloat32, std::vector{1}); + } + nvte_group_amax(input.data(), reinterpret_cast(nvte_tensor_output_list.data()), + split_sections.data(), num_tensors, stream); + for (size_t i = 0; i < num_tensors; i++) { + output_list[i].set_amax(orig_amax_ptr_list[i], DType::kFloat32, std::vector{1}); + } + + // Quantize tensors individually + for (size_t i = 0; i < num_tensors; i++) { + // skip this round if input is empty + if (input_list[i].numel() == 0) { + continue; } + nvte_quantize_v2(input_list[i].data(), output_list[i].data(), quant_config_list[i], stream); + } +} - // Quantize tensors individually - NVTE_SCOPED_GIL_RELEASE({ - for (size_t i = 0; i < num_tensors; i++) { - // skip this round if input is empty - if (input_list[i].numel() == 0) { - continue; - } - nvte_quantize_v2(input_list[i].data(), output_list[i].data(), quant_config_list[i], stream); - } - }); +void split_quantize_nvfp4_impl(const TensorWrapper &input, + const std::vector &input_list, + std::vector &output_list, + const std::vector &split_sections, + const std::vector &quantizers) { + // Check tensor lists + const size_t num_tensors = split_sections.size(); + NVTE_CHECK(input_list.size() == num_tensors, "Expected ", num_tensors, " input tensors, but got ", + input_list.size(), "."); + NVTE_CHECK(output_list.size() == num_tensors, "Expected ", num_tensors, + " output tensors, but got ", output_list.size(), "."); + NVTE_CHECK(quantizers.size() == num_tensors, "Expected ", num_tensors, + " NVFP4 quantizers, but got ", quantizers.size(), "."); + + // sanity check all the quantizers have the same scaling mode + bool all_same_scaling_mode = + std::all_of(quantizers.begin(), quantizers.end(), [&](const NVFP4Quantizer *quantizer) { + return quantizer->get_scaling_mode() == quantizers.front()->get_scaling_mode(); + }); + NVTE_CHECK(all_same_scaling_mode, "All quantizers must have the same scaling mode"); + + // Trivial cases + if (num_tensors == 0) { + return; + } + if (input.numel() == 0) { + for (const auto &tensor : input_list) { + NVTE_CHECK(tensor.numel() == 0, + "Input tensor has zero elements but got split with non-zero elements"); + } + return; } + + // Assume all quantizers have identical config + const auto &quantizer = *quantizers.front(); + NVTE_CHECK(!quantizer.with_2d_quantization, + "NVFP4 split-quantize does not support 2D quantization"); + NVTE_CHECK(!quantizer.with_amax_reduction, + "NVFP4 split-quantize does not support amax reduction"); + + // Check input tensor shape + const size_t input_last_dim = input.ndim() > 0 ? input.size(input.ndim() - 1) : 1; + NVTE_CHECK(input_last_dim % 128 == 0, + "NVFP4 multi-quantize requires inner dim to be multiple of 128."); + + // CUDA stream + auto stream = at::cuda::getCurrentCUDAStream(); + + // Perform multi-tensor quantization + NVTE_SCOPED_GIL_RELEASE({ + if (quantizer.with_rht) { // Quantize row-wise data, RHT+quantize column-wise data + // Check that config is supported + NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); + // Fuse the rowwise and colwise into one when the kernel is ready + split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections, + quantizers, stream); + } else { // NVFP4 quantize + // Fuse the rowwise and colwise into one when the kernel is ready + split_quantize_nvfp4_impl_helper(input, input_list, output_list, split_sections, quantizers, + stream); + } + }); } } // namespace diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index c73c09b317..fd748d1b21 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1501,7 +1501,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou } } - // Restriction for the RHT cast fusion kernel. + // 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; diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index fbe2ee6d1c..3f5995230c 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -120,7 +120,7 @@ def get_align_size_for_quantization(recipe: Recipe) -> int: if recipe.mxfp8(): return 32 if recipe.nvfp4(): - return 64 + return 128 return 16 From 97a09c29cffa36efd48cd57f765b0cdc02674eaf Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Mon, 22 Dec 2025 13:27:51 -0800 Subject: [PATCH 140/521] Fix ptxas compilation on sm103 for triton kernels (#2539) * add triton ptxas path for gb300 to find where it is to avoid compilation errors Signed-off-by: tdophung * add these flags in advance to preven future breaks when ops are extended to multi gpus Signed-off-by: tdophung * add this also to L1 Signed-off-by: tdophung --------- Signed-off-by: tdophung --- qa/L0_jax_distributed_unittest/test.sh | 1 + qa/L0_jax_unittest/test.sh | 1 + qa/L1_jax_distributed_unittest/test.sh | 1 + qa/L2_jax_distributed_unittest/test.sh | 1 + qa/L2_jax_unittest/test.sh | 1 + 5 files changed, 5 insertions(+) diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index 5268f9ba0e..58ce409add 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -1,6 +1,7 @@ # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas function error_exit() { echo "Error: $1" diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 92f7dd2525..c430e8d61b 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -1,6 +1,7 @@ # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas set -x diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index 5751d28200..b93224e64d 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -1,6 +1,7 @@ # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas function test_fail() { RET=1 diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index b81331dbc2..347ff35548 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -1,6 +1,7 @@ # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas set -xe diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index 1b35596fd5..d1eaaeb863 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -1,6 +1,7 @@ # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. +export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas set -x From 5ba01faadb6c05b2d9aed1e5e96b976afa65af4a Mon Sep 17 00:00:00 2001 From: xiaoxi-wangfj <690912414@qq.com> Date: Sat, 27 Dec 2025 09:57:20 +0800 Subject: [PATCH 141/521] [PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization (#1921) * [PyTorch] Fuse permute+pad and unpermute+unpad ops for FP8 optimization 1.Fused `moe_permute_with_probs` + `Fp8Padding` and fused `moe_unpermute` + `Fp8Unpadding`, that can remove the explicit padding/unpadding of moe expert, improved performance and reduced peak gpu memory usage. 2.Add tests of fused permute/pad and unpermute/unpad. Signed-off-by: xiaoxi-wangfj <690912414@qq.com> * [PyTorch/Common] Fuse permute+pad and unpermute+unpad support with_merging_probs Signed-off-by: xiaoxi-wangfj <690912414@qq.com> * [PyTorch]format code Signed-off-by: xiaoxi-wangfj <690912414@qq.com> * [Common]perf expert_idx loaded once Signed-off-by: xiaoxi-wangfj <690912414@qq.com> * fix: pad_offsets can be None Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: xiaoxi-wangfj <690912414@qq.com> * add padding + merging probs bwd support. Not tested Signed-off-by: tdophung * Fix garbage initialized act grad Signed-off-by: tdophung * all test passing for jax permutation + pad Signed-off-by: tdophung * change tokens_per_experts APIs to num_out_tokens with conservative allocation of worst case padding for output buffer Signed-off-by: tdophung * change test permutation to reduce test time Signed-off-by: tdophung * triggering PR refresh Signed-off-by: tdophung * format code Signed-off-by: tdophung * Remove some tests cases from pytorch side. Add a separate toekn_dispatch test for sanity in case combine accidentally undo an error on dispatch in the roundtrip test. Add distinction between L0 and L2 in test cases in jax Signed-off-by: tdophung * format code Signed-off-by: tdophung * remove chance for inefficiency in moving between CPU and GPU, remove redundant primitive using a new static bool for padding, add assert for align size Signed-off-by: tdophung * fix lint in jax Signed-off-by: tdophung * account for both jax newer and older than version 0.8.2. Adjusted gpu triton binding accordingly Signed-off-by: tdophung * format code Signed-off-by: tdophung * fix typo Signed-off-by: tdophung --------- Signed-off-by: xiaoxi-wangfj <690912414@qq.com> Signed-off-by: tdophung Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: tdophung --- tests/jax/test_permutation.py | 897 +++++++++++------- tests/pytorch/test_permutation.py | 658 ++++++++++++- .../common/triton/permutation.py | 38 +- transformer_engine/jax/permutation.py | 382 ++++++-- .../jax/triton_extensions/permutation.py | 547 ++++++++++- .../jax/triton_extensions/utils.py | 36 +- transformer_engine/pytorch/__init__.py | 1 + transformer_engine/pytorch/permutation.py | 104 +- .../pytorch/triton/permutation.py | 50 +- 9 files changed, 2233 insertions(+), 480 deletions(-) diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index 23d9f50609..9d1bcc820f 100644 --- a/tests/jax/test_permutation.py +++ b/tests/jax/test_permutation.py @@ -4,6 +4,8 @@ """Tests for permutation Triton kernels and high-level APIs""" +import functools + import jax import jax.numpy as jnp import pytest @@ -14,68 +16,117 @@ token_combine, sort_chunks_by_index, ) -from utils import assert_allclose +from utils import assert_allclose, pytest_parametrize_wrapper + + +# ============================================================================= +# Test parameter definitions with L0 (fast) and L2 (comprehensive) levels +# ============================================================================= + +# All dispatch/combine test cases +ALL_DISPATCH_COMBINE_CASES = [ + (128, 5, 128, 3), + (1024, 8, 128, 8), + (4096, 32, 1280, 2), + (4096, 256, 4096, 6), +] +DISPATCH_COMBINE_CASES = { + "L0": ALL_DISPATCH_COMBINE_CASES[0:2], + "L2": ALL_DISPATCH_COMBINE_CASES, +} + +# All sort chunks test cases +ALL_SORT_CHUNKS_CASES = [ + (8, 4096, 1280), + (64, 4096, 4096), + (256, 4096, 9216), +] +SORT_CHUNKS_CASES = { + "L0": ALL_SORT_CHUNKS_CASES[0:2], + "L2": ALL_SORT_CHUNKS_CASES, +} + +# All dispatch/combine with padding test cases +ALL_DISPATCH_COMBINE_PADDING_CASES = [ + (128, 5, 128, 3, 8), + (1024, 8, 128, 8, 16), + (4096, 32, 1280, 2, 128), + (4096, 256, 4096, 6, 16), +] +DISPATCH_COMBINE_PADDING_CASES = { + "L0": ALL_DISPATCH_COMBINE_PADDING_CASES[0:2], + "L2": ALL_DISPATCH_COMBINE_PADDING_CASES, +} + +# Dtypes for testing +ALL_DTYPES = [jnp.float32, jnp.bfloat16] +DTYPES = { + "L0": ALL_DTYPES, + "L2": ALL_DTYPES, +} + +# With probs options +ALL_WITH_PROBS = [True, False] +WITH_PROBS = { + "L0": [True], + "L2": ALL_WITH_PROBS, +} def reference_make_row_id_map( routing_map: jnp.ndarray, - num_tokens: int, - num_experts: int, ) -> jnp.ndarray: """ - Reference implementation of make_row_id_map using JAX primitives. + Vectorized reference implementation of make_row_id_map using JAX primitives. Parameters ---------- routing_map : jnp.ndarray Input tensor of shape [num_tokens, num_experts]. Mask indicating which experts are routed to which tokens (1 = routed, 0 = not routed). - num_tokens : int - Number of tokens in the input tensor. - num_experts : int - Number of experts in the input tensor. Returns ------- row_id_map : jnp.ndarray The row_id_map for the permutation of shape [num_tokens, num_experts * 2 + 1]. """ - row_id_map = jnp.full((num_tokens, num_experts * 2 + 1), -1, dtype=jnp.int32) + num_tokens, num_experts = routing_map.shape # For each expert, compute cumulative sum to get destination indices cumsum_per_expert = jnp.cumsum(routing_map, axis=0) - # Compute total tokens per expert + # Compute total tokens per expert and expert offsets tokens_per_expert = jnp.sum(routing_map, axis=0) expert_offsets = jnp.concatenate([jnp.array([0]), jnp.cumsum(tokens_per_expert)[:-1]]) - # Build the row_id_map - for token_idx in range(num_tokens): - routed_experts = jnp.where(routing_map[token_idx] == 1)[0] - n_routed = len(routed_experts) - - # Store number of routed experts in the last position - row_id_map = row_id_map.at[token_idx, -1].set(n_routed) - - # For each routed expert, compute destination row and store it - dest_rows = [] - expert_indices = [] - for expert_idx in routed_experts: - # Destination row = expert offset + (cumsum - 1) - dest_row = expert_offsets[expert_idx] + cumsum_per_expert[token_idx, expert_idx] - 1 - dest_rows.append(dest_row) - expert_indices.append(expert_idx) - - # Sort by destination row - if n_routed > 0: - sort_indices = jnp.argsort(-jnp.array(dest_rows)) # Negative for descending sort - sorted_dest_rows = jnp.array(dest_rows)[sort_indices] - sorted_expert_indices = jnp.array(expert_indices)[sort_indices] - - # Store sorted destination rows and expert indices - for i in range(n_routed): - row_id_map = row_id_map.at[token_idx, i].set(sorted_dest_rows[i]) - row_id_map = row_id_map.at[token_idx, num_experts + i].set(sorted_expert_indices[i]) + # Compute destination rows for all (token, expert) pairs + # dest_row[i, j] = expert_offsets[j] + cumsum_per_expert[i, j] - 1 if routed, else -1 + dest_rows_all = (expert_offsets[None, :] + cumsum_per_expert - 1) * routing_map + (-1) * ( + 1 - routing_map + ) + + # Count routed experts per token + n_routed_per_token = jnp.sum(routing_map, axis=1) + + # For each token, we need to sort by descending dest_row and pack into row_id_map + # Use a large negative value for non-routed experts so they sort to the end + sort_keys = jnp.where(routing_map == 1, -dest_rows_all, jnp.iinfo(jnp.int32).max) + sorted_expert_indices = jnp.argsort(sort_keys, axis=1) + + # Gather the sorted destination rows and expert indices using advanced indexing + # Create indices for gathering + token_idx = jnp.broadcast_to(jnp.arange(num_tokens)[:, None], (num_tokens, num_experts)) + sorted_dest_rows = dest_rows_all[token_idx, sorted_expert_indices] + + # Build row_id_map: [dest_row_0, ..., dest_row_{E-1}, expert_idx_0, ..., expert_idx_{E-1}, n_routed] + row_id_map = jnp.concatenate( + [ + sorted_dest_rows.astype(jnp.int32), + sorted_expert_indices.astype(jnp.int32), + n_routed_per_token.astype(jnp.int32)[:, None], + ], + axis=1, + ) return row_id_map @@ -84,13 +135,10 @@ def _reference_permute_impl( inp: jnp.ndarray, row_id_map: jnp.ndarray, probs: jnp.ndarray, - num_tokens: int, - num_experts: int, num_out_tokens: int, - hidden_size: int, ) -> tuple: """ - Internal helper for reference permutation implementation. + Vectorized internal helper for reference permutation implementation. Parameters ---------- @@ -100,14 +148,8 @@ def _reference_permute_impl( The token to expert mapping tensor of shape [num_tokens, num_experts * 2 + 1]. probs : jnp.ndarray The probabilities of the input tensor. - num_tokens : int - Number of tokens in the input tensor. - num_experts : int - Number of experts. num_out_tokens : int Number of tokens in the permuted tensor. - hidden_size : int - Hidden size of the input tensor. Returns ------- @@ -116,33 +158,63 @@ def _reference_permute_impl( permuted_probs : jnp.ndarray Permuted probabilities if probs was provided, None otherwise. """ + num_tokens, hidden_size = inp.shape + num_experts = (row_id_map.shape[1] - 1) // 2 + + # Extract destination rows, expert indices, and n_routed from row_id_map + dest_rows = row_id_map[:, :num_experts] # [num_tokens, num_experts] + expert_indices = row_id_map[:, num_experts : 2 * num_experts] # [num_tokens, num_experts] + n_routed = row_id_map[:, 2 * num_experts] # [num_tokens] + + # Create mask for valid entries: slot_idx < n_routed[token] + # The kernel's row_id_map only guarantees valid data in the first n_routed slots + # (slots beyond n_routed may contain garbage, not -1) + slot_indices = jnp.arange(num_experts)[None, :] # [1, num_experts] + valid_mask = slot_indices < n_routed[:, None] # [num_tokens, num_experts] + + # Flatten for scatter operations + flat_dest_rows = dest_rows.flatten() # [num_tokens * num_experts] + flat_valid_mask = valid_mask.flatten() + flat_token_indices = jnp.repeat(jnp.arange(num_tokens), num_experts) + flat_expert_indices = expert_indices.flatten() + + # Set invalid dest_rows to num_out_tokens (out of bounds, will be dropped) + # This avoids overwriting valid entries at index 0 with zeros + flat_dest_rows_clamped = jnp.where(flat_valid_mask, flat_dest_rows, num_out_tokens) + + # Gather input tokens and scatter to output output = jnp.zeros((num_out_tokens, hidden_size), dtype=inp.dtype) - permuted_probs = None if probs is None else jnp.zeros((num_out_tokens,), dtype=probs.dtype) - - for token_idx in range(num_tokens): - n_routed = int(row_id_map[token_idx, -1]) # int() needed for Python range() - for i in range(n_routed): - # Don't use int() here - JAX can index with traced values, - # and int() breaks autodiff gradient tracking - dest_row = row_id_map[token_idx, i] - expert_idx = row_id_map[token_idx, num_experts + i] - - # Get probability for this expert - if probs is not None: - if probs.ndim == 1: - prob = probs[token_idx] - else: - prob = probs[token_idx, expert_idx] - - # Match kernel behavior: if prob == 0.0, zero out the output (padding indicator) - if prob == 0.0: - output = output.at[dest_row].set(0.0) - else: - output = output.at[dest_row].set(inp[token_idx]) - - permuted_probs = permuted_probs.at[dest_row].set(prob) - else: - output = output.at[dest_row].set(inp[token_idx]) + gathered_inp = inp[flat_token_indices] # [num_tokens * num_experts, hidden_size] + + # Use segment_sum-like operation via scatter + # For each valid (token, expert) pair, write inp[token] to output[dest_row] + # Invalid entries target num_out_tokens and get dropped by mode="drop" + output = output.at[flat_dest_rows_clamped].set( + gathered_inp, + mode="drop", + ) + + permuted_probs = None + if probs is not None: + permuted_probs = jnp.zeros((num_out_tokens,), dtype=probs.dtype) + + # Vectorized approach: gather probs and scatter to permuted_probs + if probs.ndim == 1: + flat_probs = probs[flat_token_indices] + else: + # Clamp invalid expert indices to 0 to avoid wraparound indexing with -1 + # The result for invalid entries will be ignored anyway since they target num_out_tokens + # Cast to int32 explicitly for consistent indexing behavior + flat_expert_indices_clamped = jnp.where(flat_valid_mask, flat_expert_indices, 0).astype( + jnp.int32 + ) + flat_probs = probs[flat_token_indices.astype(jnp.int32), flat_expert_indices_clamped] + + # Invalid entries target num_out_tokens and get dropped by mode="drop" + permuted_probs = permuted_probs.at[flat_dest_rows_clamped.astype(jnp.int32)].set( + flat_probs, + mode="drop", + ) return output, permuted_probs @@ -152,12 +224,9 @@ def _reference_unpermute_impl( row_id_map: jnp.ndarray, merging_probs: jnp.ndarray, permuted_probs: jnp.ndarray, - num_tokens: int, - num_experts: int, - hidden_size: int, ) -> tuple: """ - Internal helper for reference unpermutation implementation. + Vectorized internal helper for reference unpermutation implementation. Parameters ---------- @@ -169,12 +238,6 @@ def _reference_unpermute_impl( The merging probabilities for weighted reduction. permuted_probs : jnp.ndarray The permuted probabilities. - num_tokens : int - Number of tokens. - num_experts : int - Number of experts. - hidden_size : int - Hidden size. Returns ------- @@ -183,31 +246,44 @@ def _reference_unpermute_impl( unpermuted_probs : jnp.ndarray Unpermuted probabilities if permuted_probs was provided, None otherwise. """ - output = jnp.zeros((num_tokens, hidden_size), dtype=inp.dtype) - unpermuted_probs = ( - None - if permuted_probs is None - else jnp.zeros((num_tokens, num_experts), dtype=permuted_probs.dtype) - ) + num_tokens = row_id_map.shape[0] + num_experts = (row_id_map.shape[1] - 1) // 2 - for token_idx in range(num_tokens): - n_routed = int(row_id_map[token_idx, -1]) # int() needed for Python range() - for i in range(n_routed): - # Don't use int() here - JAX can index with traced values, - # and int() breaks autodiff gradient tracking - src_row = row_id_map[token_idx, i] - expert_idx = row_id_map[token_idx, num_experts + i] - - if merging_probs is not None: - weight = merging_probs[token_idx, expert_idx] - output = output.at[token_idx].add(inp[src_row] * weight) - else: - output = output.at[token_idx].add(inp[src_row]) - - if permuted_probs is not None: - unpermuted_probs = unpermuted_probs.at[token_idx, expert_idx].set( - permuted_probs[src_row] - ) + # Extract source rows, expert indices, and n_routed from row_id_map + src_rows = row_id_map[:, :num_experts] # [num_tokens, num_experts] + expert_indices = row_id_map[:, num_experts : 2 * num_experts] # [num_tokens, num_experts] + n_routed = row_id_map[:, 2 * num_experts] # [num_tokens] + + # Create mask for valid entries: slot_idx < n_routed[token] + # The kernel's row_id_map only guarantees valid data in the first n_routed slots + slot_indices = jnp.arange(num_experts)[None, :] # [1, num_experts] + valid_mask = slot_indices < n_routed[:, None] # [num_tokens, num_experts] + + # Clamp invalid src_rows to 0 (they won't be used due to masking) + src_rows_clamped = jnp.where(valid_mask, src_rows, 0) + + # Gather input from permuted positions + gathered_inp = inp[src_rows_clamped] # [num_tokens, num_experts, hidden_size] + + # Apply merging probs if provided + if merging_probs is not None: + # Gather the merging weights for each (token, expert) pair using advanced indexing + token_idx = jnp.broadcast_to(jnp.arange(num_tokens)[:, None], (num_tokens, num_experts)) + weights = merging_probs[token_idx, expert_indices] # [num_tokens, num_experts] + gathered_inp = gathered_inp * weights[:, :, None] + + # Mask out invalid entries and sum across experts + gathered_inp = jnp.where(valid_mask[:, :, None], gathered_inp, 0.0) + output = jnp.sum(gathered_inp, axis=1) # [num_tokens, hidden_size] + + unpermuted_probs = None + if permuted_probs is not None: + gathered_probs = permuted_probs[src_rows_clamped] # [num_tokens, num_experts] + unpermuted_probs = jnp.zeros((num_tokens, num_experts), dtype=permuted_probs.dtype) + token_idx = jnp.broadcast_to(jnp.arange(num_tokens)[:, None], (num_tokens, num_experts)) + unpermuted_probs = unpermuted_probs.at[token_idx, expert_indices].set( + jnp.where(valid_mask, gathered_probs, 0.0) + ) return output, unpermuted_probs @@ -241,13 +317,8 @@ def reference_token_dispatch( row_id_map : jnp.ndarray The row_id_map for the permutation. """ - num_tokens, num_experts = routing_map.shape - hidden_size = inp.shape[1] - - row_id_map = reference_make_row_id_map(routing_map, num_tokens, num_experts) - output, permuted_probs = _reference_permute_impl( - inp, row_id_map, probs, num_tokens, num_experts, num_out_tokens, hidden_size - ) + row_id_map = reference_make_row_id_map(routing_map) + output, permuted_probs = _reference_permute_impl(inp, row_id_map, probs, num_out_tokens) return output, permuted_probs, row_id_map @@ -274,13 +345,7 @@ def reference_token_combine( output : jnp.ndarray Unpermuted output tensor of shape [num_tokens, hidden_size]. """ - num_tokens = row_id_map.shape[0] - num_experts = (row_id_map.shape[1] - 1) // 2 - hidden_size = inp.shape[1] - - output, _ = _reference_unpermute_impl( - inp, row_id_map, merging_probs, None, num_tokens, num_experts, hidden_size - ) + output, _ = _reference_unpermute_impl(inp, row_id_map, merging_probs, None) return output @@ -289,10 +354,9 @@ def reference_make_chunk_sort_map( split_sizes: jnp.ndarray, sorted_indices: jnp.ndarray, num_tokens: int, - num_splits: int, ) -> jnp.ndarray: """ - Reference implementation of make_chunk_sort_map using JAX primitives. + Vectorized reference implementation of make_chunk_sort_map using JAX primitives. Parameters ---------- @@ -302,45 +366,48 @@ def reference_make_chunk_sort_map( The indices of the sorted chunks of shape [num_splits,]. num_tokens : int Number of tokens. - num_splits : int - Number of splits. Returns ------- row_id_map : jnp.ndarray Row ID map for chunk sorting of shape [num_tokens,]. """ - row_id_map = jnp.zeros((num_tokens,), dtype=jnp.int32) + # Compute source chunk boundaries (cumulative sum of original split_sizes) + src_cumsum = jnp.concatenate([jnp.array([0]), jnp.cumsum(split_sizes)]) - # Compute cumulative positions - cumsum_sizes = jnp.concatenate([jnp.array([0]), jnp.cumsum(split_sizes)]) + # Compute destination chunk boundaries based on sorted order + sorted_sizes = split_sizes[sorted_indices] + dest_cumsum = jnp.concatenate([jnp.array([0]), jnp.cumsum(sorted_sizes)]) - # For each chunk, compute the destination indices - dest_offset = 0 - for sorted_idx in sorted_indices: - chunk_start = cumsum_sizes[sorted_idx] - chunk_end = cumsum_sizes[sorted_idx + 1] - chunk_size = chunk_end - chunk_start + # For each source chunk, compute its destination offset + # inverse_indices[i] = position of chunk i in sorted order + inverse_indices = jnp.argsort(sorted_indices) + dest_offsets = dest_cumsum[inverse_indices] - # Map source positions to destination positions - for i in range(chunk_size): - row_id_map = row_id_map.at[chunk_start + i].set(dest_offset + i) + # Create row_id_map: for each token position, compute its destination + # First, figure out which chunk each position belongs to + position_indices = jnp.arange(num_tokens) - dest_offset += chunk_size + # chunk_ids[i] = which chunk position i belongs to + chunk_ids = jnp.searchsorted(src_cumsum[1:], position_indices, side="right") - return row_id_map + # within_chunk_offset[i] = position i's offset within its chunk + within_chunk_offset = position_indices - src_cumsum[chunk_ids] + + # destination[i] = dest_offsets[chunk_ids[i]] + within_chunk_offset[i] + row_id_map = dest_offsets[chunk_ids] + within_chunk_offset + + return row_id_map.astype(jnp.int32) def reference_sort_chunks_by_map( inp: jnp.ndarray, row_id_map: jnp.ndarray, probs: jnp.ndarray, - num_tokens: int, - hidden_size: int, is_forward: bool, ) -> tuple: """ - Reference implementation of sort_chunks_by_map using JAX primitives. + Vectorized reference implementation of sort_chunks_by_map using JAX primitives. Parameters ---------- @@ -350,10 +417,6 @@ def reference_sort_chunks_by_map( The token to destination mapping of shape [num_tokens,]. probs : jnp.ndarray The probabilities. - num_tokens : int - Number of tokens. - hidden_size : int - Hidden size. is_forward : bool Whether this is forward or backward. @@ -364,25 +427,25 @@ def reference_sort_chunks_by_map( permuted_probs : jnp.ndarray Sorted probabilities if probs was provided, None otherwise. """ - output = jnp.zeros((num_tokens, hidden_size), dtype=inp.dtype) - permuted_probs = None if probs is None else jnp.zeros((num_tokens,), dtype=probs.dtype) + num_tokens = inp.shape[0] + hidden_size = inp.shape[1] if is_forward: - # Forward: src -> dest - for src_idx in range(num_tokens): - # Don't use int() - JAX can index with traced values - dest_idx = row_id_map[src_idx] - output = output.at[dest_idx].set(inp[src_idx]) - if probs is not None: - permuted_probs = permuted_probs.at[dest_idx].set(probs[src_idx]) + # Forward: scatter inp[src] to output[dest] where dest = row_id_map[src] + output = jnp.zeros((num_tokens, hidden_size), dtype=inp.dtype) + output = output.at[row_id_map].set(inp) + if probs is not None: + permuted_probs = jnp.zeros((num_tokens,), dtype=probs.dtype) + permuted_probs = permuted_probs.at[row_id_map].set(probs) + else: + permuted_probs = None else: - # Backward: dest -> src - for dest_idx in range(num_tokens): - # Don't use int() - JAX can index with traced values - src_idx = row_id_map[dest_idx] - output = output.at[dest_idx].set(inp[src_idx]) - if probs is not None: - permuted_probs = permuted_probs.at[dest_idx].set(probs[src_idx]) + # Backward: gather output[dest] = inp[src] where src = row_id_map[dest] + output = inp[row_id_map] + if probs is not None: + permuted_probs = probs[row_id_map] + else: + permuted_probs = None return output, permuted_probs @@ -415,20 +478,24 @@ def generate_routing_map( return routing_map - # ========================================================================= - # token_dispatch tests - # ========================================================================= - - @pytest.mark.parametrize( + @pytest_parametrize_wrapper( "num_tokens,num_experts,hidden_size,tokens_per_expert", - [ - (32, 8, 256, 2), - (64, 16, 512, 3), - ], + DISPATCH_COMBINE_CASES, ) - @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) - def test_token_dispatch(self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype): - """Test token_dispatch forward and backward pass against reference""" + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("with_probs", WITH_PROBS) + def test_token_dispatch( + self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype, with_probs + ): + """ + Individual test for token_dispatch forward and backward passes. + + This test validates dispatch in isolation to catch errors that might be + masked when combined with token_combine in the roundtrip test. + + Uses value_and_grad to validate both forward (via loss comparison) and + backward (via gradient comparison) passes against reference implementation. + """ key = jax.random.PRNGKey(42) # Generate routing map @@ -436,173 +503,231 @@ def test_token_dispatch(self, num_tokens, num_experts, hidden_size, tokens_per_e num_out_tokens = int(jnp.sum(routing_map)) # Generate input data - key, inp_key = jax.random.split(key) + key, inp_key, prob_key = jax.random.split(key, 3) inp = jax.random.uniform( inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 ) - # Define loss functions - def loss_fn(x): - output, _, _ = token_dispatch(x, routing_map, num_out_tokens) - return jnp.sum(output**2) + # Generate probs if needed (minval > 0 to avoid kernel's special prob==0 handling) + probs = None + if with_probs: + probs = jax.random.uniform( + prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 + ) - def ref_loss_fn(x): - output, _, _ = reference_token_dispatch(x, routing_map, num_out_tokens) - return jnp.sum(output**2) + # Generate reference row_id_map for comparison + ref_row_id_map = reference_make_row_id_map(routing_map) - loss_val, computed_grad = jax.value_and_grad(loss_fn)(inp) - ref_loss_val, ref_grad = jax.value_and_grad(ref_loss_fn)(inp) + # ===================================================================== + # Test forward and backward pass using value_and_grad + # (value validates forward, grad validates backward) + # ===================================================================== + if with_probs: - # Compare forward outputs - output, _, _ = token_dispatch(inp, routing_map, num_out_tokens) - ref_output, _, _ = reference_token_dispatch(inp, routing_map, num_out_tokens) - assert_allclose(output, ref_output) + @jax.jit + def dispatch_loss(x, p): + out, perm_probs, _, _, _ = token_dispatch(x, routing_map, num_out_tokens, probs=p) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) - # Compare loss and gradient - assert_allclose(loss_val, ref_loss_val) - assert_allclose(computed_grad, ref_grad) + @jax.jit + def ref_dispatch_loss(x, p): + out, perm_probs = _reference_permute_impl(x, ref_row_id_map, p, num_out_tokens) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + loss_val, (inp_grad, probs_grad) = jax.value_and_grad(dispatch_loss, argnums=(0, 1))( + inp, probs + ) + ref_loss_val, (ref_inp_grad, ref_probs_grad) = jax.value_and_grad( + ref_dispatch_loss, argnums=(0, 1) + )(inp, probs) + + # Validate forward loss matches + assert_allclose(loss_val, ref_loss_val, dtype=dtype) + # Validate gradients + assert_allclose(inp_grad, ref_inp_grad, dtype=dtype) + assert_allclose(probs_grad, ref_probs_grad, dtype=dtype) + else: + + @jax.jit + def dispatch_loss_no_probs(x): + out, _, _, _, _ = token_dispatch(x, routing_map, num_out_tokens) + return jnp.sum(out**2) + + @jax.jit + def ref_dispatch_loss_no_probs(x): + out, _ = _reference_permute_impl(x, ref_row_id_map, None, num_out_tokens) + return jnp.sum(out**2) + + loss_val, inp_grad = jax.value_and_grad(dispatch_loss_no_probs)(inp) + ref_loss_val, ref_inp_grad = jax.value_and_grad(ref_dispatch_loss_no_probs)(inp) + + # Validate forward loss matches + assert_allclose(loss_val, ref_loss_val, dtype=dtype) + # Validate gradients + assert_allclose(inp_grad, ref_inp_grad, dtype=dtype) # ========================================================================= - # token_dispatch with probs tests + # Consolidated dispatch + combine tests # ========================================================================= - @pytest.mark.parametrize( + @pytest_parametrize_wrapper( "num_tokens,num_experts,hidden_size,tokens_per_expert", - [ - (32, 8, 256, 2), - (64, 16, 512, 3), - ], + DISPATCH_COMBINE_CASES, ) - @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) - def test_token_dispatch_with_probs( - self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("with_probs", WITH_PROBS) + def test_dispatch_and_combine( + self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype, with_probs ): - """Test token_dispatch with probs forward and backward pass against reference""" + """ + Comprehensive test for token_dispatch and token_combine. + + Tests: + 1. Dispatch forward pass against reference (element-by-element) + 2. Dispatch backward pass against reference + 3. Combine forward pass against reference (element-by-element) + 4. Combine backward pass against reference + 5. Roundtrip: dispatch + combine recovers original input + 6. row_id_map n_routed column validation + 7. Probs permutation (when with_probs=True) + """ key = jax.random.PRNGKey(42) # Generate routing map routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) num_out_tokens = int(jnp.sum(routing_map)) - # Generate input data and probs - key, inp_key, prob_key = jax.random.split(key, 3) + # Generate input data + key, inp_key, prob_key, merge_key = jax.random.split(key, 4) inp = jax.random.uniform( inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 ) - probs = jax.random.uniform( - prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.0, maxval=1.0 - ) - # Define loss function that uses token_dispatch with probs - # We compute gradients w.r.t. both inp and probs - def loss_fn(x, p): - output, permuted_probs, _ = token_dispatch(x, routing_map, num_out_tokens, probs=p) - return jnp.sum(output**2) + jnp.sum(permuted_probs**2) - - def ref_loss_fn(x, p): - output, permuted_probs, _ = reference_token_dispatch( - x, routing_map, num_out_tokens, probs=p + # Generate probs if needed (minval > 0 to avoid kernel's special prob==0 handling) + probs = None + if with_probs: + probs = jax.random.uniform( + prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 ) - return jnp.sum(output**2) + jnp.sum(permuted_probs**2) - - loss_val, (inp_grad, probs_grad) = jax.value_and_grad(loss_fn, argnums=(0, 1))(inp, probs) - ref_loss_val, (ref_inp_grad, ref_probs_grad) = jax.value_and_grad( - ref_loss_fn, argnums=(0, 1) - )(inp, probs) - output, permuted_probs, _ = token_dispatch(inp, routing_map, num_out_tokens, probs=probs) - - ref_output, ref_permuted_probs, _ = reference_token_dispatch( - inp, routing_map, num_out_tokens, probs=probs + # Generate merging probs (normalized per token) + merging_probs = jax.random.uniform( + merge_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 ) - - # Compare forward outputs - assert_allclose(output, ref_output) - assert_allclose(permuted_probs, ref_permuted_probs) - - # Compare loss and gradients - assert_allclose(loss_val, ref_loss_val) - assert_allclose(inp_grad, ref_inp_grad) - assert_allclose(probs_grad, ref_probs_grad) - - # ========================================================================= - # token_combine tests - # ========================================================================= - - @pytest.mark.parametrize( - "num_tokens,num_experts,hidden_size,tokens_per_expert", - [ - (32, 8, 256, 2), - (64, 16, 512, 3), - ], - ) - @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) - @pytest.mark.parametrize("with_merging_probs", [True, False]) - def test_token_combine( - self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype, with_merging_probs - ): - """Test token_combine forward and backward pass against reference""" - key = jax.random.PRNGKey(42) - - # Generate routing map - routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) - num_out_tokens = int(jnp.sum(routing_map)) - - # Get row_id_map from reference_token_dispatch - key, dummy_key = jax.random.split(key) - dummy_inp = jax.random.uniform( - dummy_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + merging_probs = merging_probs * routing_map.astype(dtype) # Zero out non-routed + merging_probs = merging_probs / jnp.maximum( + jnp.sum(merging_probs, axis=1, keepdims=True), 1e-8 ) - _, _, row_id_map = reference_token_dispatch(dummy_inp, routing_map, num_out_tokens) - # Generate input data (from expert outputs) - key, inp_key, merge_key = jax.random.split(key, 3) - inp = jax.random.uniform( - inp_key, (num_out_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + # ===================================================================== + # Test 1: Dispatch forward pass + # ===================================================================== + output, permuted_probs, row_id_map, _, _ = token_dispatch( + inp, routing_map, num_out_tokens, probs=probs + ) + ref_output, ref_permuted_probs = _reference_permute_impl( + inp, row_id_map, probs, num_out_tokens ) - if with_merging_probs: - merging_probs = jax.random.uniform( - merge_key, (num_tokens, num_experts), dtype=dtype, minval=0.0, maxval=1.0 + # Validate row_id_map structure: n_routed column should match routing_map sum + n_routed_actual = row_id_map[:, -1] + n_routed_expected = jnp.sum(routing_map, axis=1) + assert jnp.array_equal( + n_routed_actual, n_routed_expected + ), "make_row_id_map n_routed column mismatch" + + # Compare dispatch output + assert_allclose(output, ref_output, dtype=dtype) + if with_probs: + assert_allclose(permuted_probs, ref_permuted_probs, dtype=dtype) + + # ===================================================================== + # Test 2: Dispatch backward pass + # ===================================================================== + if with_probs: + + @jax.jit + def dispatch_loss(x, p): + out, perm_probs, _, _, _ = token_dispatch(x, routing_map, num_out_tokens, probs=p) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + @jax.jit + def ref_dispatch_loss(x, p): + out, perm_probs = _reference_permute_impl(x, row_id_map, p, num_out_tokens) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + _, (inp_grad, probs_grad) = jax.value_and_grad(dispatch_loss, argnums=(0, 1))( + inp, probs ) - # Normalize per token - merging_probs = merging_probs / (jnp.sum(merging_probs, axis=1, keepdims=True) + 1e-8) + _, (ref_inp_grad, ref_probs_grad) = jax.value_and_grad( + ref_dispatch_loss, argnums=(0, 1) + )(inp, probs) + assert_allclose(inp_grad, ref_inp_grad, dtype=dtype) + assert_allclose(probs_grad, ref_probs_grad, dtype=dtype) else: - merging_probs = None - # Define loss functions - def loss_fn(x): - output = token_combine(x, row_id_map, merging_probs) - return jnp.sum(output**2) - - def ref_loss_fn(x): - output = reference_token_combine(x, row_id_map, merging_probs) - return jnp.sum(output**2) - - loss_val, computed_grad = jax.value_and_grad(loss_fn)(inp) - ref_loss_val, ref_grad = jax.value_and_grad(ref_loss_fn)(inp) + @jax.jit + def dispatch_loss_no_probs(x): + out, _, _, _, _ = token_dispatch(x, routing_map, num_out_tokens) + return jnp.sum(out**2) + + @jax.jit + def ref_dispatch_loss_no_probs(x): + out, _ = _reference_permute_impl(x, row_id_map, None, num_out_tokens) + return jnp.sum(out**2) + + _, inp_grad = jax.value_and_grad(dispatch_loss_no_probs)(inp) + _, ref_inp_grad = jax.value_and_grad(ref_dispatch_loss_no_probs)(inp) + assert_allclose(inp_grad, ref_inp_grad, dtype=dtype) + + # ===================================================================== + # Test 3: Combine forward pass + # ===================================================================== + combined = token_combine(output, row_id_map, merging_probs) + ref_combined = _reference_unpermute_impl(output, row_id_map, merging_probs, None)[0] + assert_allclose(combined, ref_combined, dtype=dtype) + + # ===================================================================== + # Test 4: Combine backward pass + # ===================================================================== + + @jax.jit + def combine_loss(x): + return jnp.sum(token_combine(x, row_id_map, merging_probs) ** 2) + + @jax.jit + def ref_combine_loss(x): + return jnp.sum(_reference_unpermute_impl(x, row_id_map, merging_probs, None)[0] ** 2) + + _, combine_grad = jax.value_and_grad(combine_loss)(output) + _, ref_combine_grad = jax.value_and_grad(ref_combine_loss)(output) + assert_allclose(combine_grad, ref_combine_grad, dtype=dtype) + + # ===================================================================== + # Test 5: Roundtrip (dispatch + combine = original) + # ===================================================================== + # Use uniform merging probs for perfect roundtrip + uniform_merging_probs = routing_map.astype(dtype) / jnp.maximum( + jnp.sum(routing_map, axis=1, keepdims=True), 1.0 + ) - # Compare forward outputs - output = token_combine(inp, row_id_map, merging_probs) - ref_output = reference_token_combine(inp, row_id_map, merging_probs) - assert_allclose(output, ref_output) + @jax.jit + def roundtrip(x): + dispatched, _, rid_map, _, _ = token_dispatch(x, routing_map, num_out_tokens) + return token_combine(dispatched, rid_map, uniform_merging_probs) - # Compare loss and gradient - assert_allclose(loss_val, ref_loss_val) - assert_allclose(computed_grad, ref_grad) + roundtrip_output = roundtrip(inp) + assert_allclose(roundtrip_output, inp, dtype=dtype) # ========================================================================= # sort_chunks_by_index tests # ========================================================================= - @pytest.mark.parametrize( + @pytest_parametrize_wrapper( "num_splits,total_tokens,hidden_size", - [ - (4, 128, 256), - (8, 256, 512), - ], + SORT_CHUNKS_CASES, ) - @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) + @pytest_parametrize_wrapper("dtype", DTYPES) def test_sort_chunks_by_index(self, num_splits, total_tokens, hidden_size, dtype): """Test sort_chunks_by_index forward and backward pass against reference""" key = jax.random.PRNGKey(42) @@ -622,73 +747,181 @@ def test_sort_chunks_by_index(self, num_splits, total_tokens, hidden_size, dtype inp_key, (total_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 ) - row_id_map = reference_make_chunk_sort_map( - split_sizes, sorted_indices, total_tokens, num_splits - ) + # Get reference row_id_map + row_id_map = reference_make_chunk_sort_map(split_sizes, sorted_indices, total_tokens) - # Define loss functions + # Define loss functions (JIT compiled for performance) + @jax.jit def loss_fn(x): output, _ = sort_chunks_by_index(x, split_sizes, sorted_indices) return jnp.sum(output**2) + @jax.jit def ref_loss_fn(x): - output, _ = reference_sort_chunks_by_map( - x, row_id_map, None, total_tokens, hidden_size, is_forward=True - ) + output, _ = reference_sort_chunks_by_map(x, row_id_map, None, is_forward=True) return jnp.sum(output**2) + # Test forward pass + output, _ = sort_chunks_by_index(inp, split_sizes, sorted_indices) + ref_output, _ = reference_sort_chunks_by_map(inp, row_id_map, None, is_forward=True) + + # Test backward pass with JIT loss_val, computed_grad = jax.value_and_grad(loss_fn)(inp) ref_loss_val, ref_grad = jax.value_and_grad(ref_loss_fn)(inp) - # Compare forward outputs - output, _ = sort_chunks_by_index(inp, split_sizes, sorted_indices) - ref_output, _ = reference_sort_chunks_by_map( - inp, row_id_map, None, total_tokens, hidden_size, is_forward=True - ) + # Compare forward and backward assert_allclose(output, ref_output) - - # Compare loss and gradient assert_allclose(loss_val, ref_loss_val) assert_allclose(computed_grad, ref_grad) # ========================================================================= - # Round-trip tests (token_dispatch -> expert processing -> token_combine) + # Consolidated dispatch + combine with padding tests # ========================================================================= - @pytest.mark.parametrize( - "num_tokens,num_experts,hidden_size,tokens_per_expert", - [ - (32, 8, 256, 2), - (64, 16, 512, 3), - ], + @pytest_parametrize_wrapper( + "num_tokens,num_experts,hidden_size,topk,align_size", + DISPATCH_COMBINE_PADDING_CASES, ) - @pytest.mark.parametrize("dtype", [jnp.float32, jnp.bfloat16]) - def test_dispatch_combine_roundtrip( - self, num_tokens, num_experts, hidden_size, tokens_per_expert, dtype + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("with_probs", WITH_PROBS) + def test_dispatch_and_combine_with_padding( + self, num_tokens, num_experts, hidden_size, topk, align_size, dtype, with_probs ): - """Test that token_dispatch followed by token_combine recovers original input""" + """ + Comprehensive test for token_dispatch and token_combine with padding/unpadding. + + Tests: + 1. Dispatch with padding: output shape and alignment + 2. Dispatch backward pass with padding + 3. Combine with unpad: output shape + 4. Combine backward pass with unpad + 5. Roundtrip with padding: dispatch + combine recovers original + 6. Probs permutation with padding (when with_probs=True) + """ key = jax.random.PRNGKey(42) # Generate routing map - routing_map = self.generate_routing_map(num_tokens, num_experts, tokens_per_expert, key) + routing_map = self.generate_routing_map(num_tokens, num_experts, topk, key) num_out_tokens = int(jnp.sum(routing_map)) + # Compute worst-case padded size + worst_case_size = ( + (num_out_tokens + num_experts * (align_size - 1)) // align_size + ) * align_size + # Generate input data - key, inp_key = jax.random.split(key) + key, inp_key, prob_key, merge_key = jax.random.split(key, 4) inp = jax.random.uniform( inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 ) - # Create uniform merging probs (equal weight for all routed experts) - merging_probs = routing_map.astype(dtype) / jnp.maximum( + # Generate probs if needed (minval > 0 to avoid kernel's special prob==0 handling) + probs = None + if with_probs: + probs = jax.random.uniform( + prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 + ) + + # Generate merging probs (normalized per token) + merging_probs = jax.random.uniform( + merge_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 + ) + merging_probs = merging_probs * routing_map.astype(dtype) # Zero out non-routed + merging_probs = merging_probs / jnp.maximum( + jnp.sum(merging_probs, axis=1, keepdims=True), 1e-8 + ) + + # ===================================================================== + # Test 1: Dispatch with padding - forward pass + # ===================================================================== + output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert = token_dispatch( + inp, routing_map, num_out_tokens, probs=probs, align_size=align_size + ) + + # Check output shape + assert output.shape == (worst_case_size, hidden_size) + if with_probs: + assert permuted_probs is not None + assert permuted_probs.shape == (worst_case_size,) + else: + assert permuted_probs is None + + # Check alignment: each expert's tokens should be aligned + for expert_idx in range(num_experts): + expert_tokens = int(target_tokens_per_expert[expert_idx]) + assert expert_tokens % align_size == 0 or expert_tokens == 0 + + # ===================================================================== + # Test 2: Dispatch with padding - backward pass + # ===================================================================== + if with_probs: + + @jax.jit + def dispatch_loss(x, p): + out, perm_probs, _, _, _ = token_dispatch( + x, routing_map, num_out_tokens, probs=p, align_size=align_size + ) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + inp_grad, probs_grad = jax.grad(dispatch_loss, argnums=(0, 1))(inp, probs) + assert inp_grad.shape == inp.shape + assert probs_grad.shape == probs.shape + assert not jnp.any(jnp.isnan(inp_grad)) + assert not jnp.any(jnp.isnan(probs_grad)) + else: + + @jax.jit + def dispatch_loss_no_probs(x): + out, _, _, _, _ = token_dispatch( + x, routing_map, num_out_tokens, align_size=align_size + ) + return jnp.sum(out**2) + + inp_grad = jax.grad(dispatch_loss_no_probs)(inp) + assert inp_grad.shape == inp.shape + assert not jnp.any(jnp.isnan(inp_grad)) + + # ===================================================================== + # Test 3: Combine with unpad - forward pass + # ===================================================================== + combined = token_combine(output, row_id_map, merging_probs, pad_offsets) + assert combined.shape == (num_tokens, hidden_size) + + # ===================================================================== + # Test 4: Combine with unpad - backward pass + # ===================================================================== + + @jax.jit + def combine_loss(x): + return jnp.sum(token_combine(x, row_id_map, merging_probs, pad_offsets) ** 2) + + combine_grad = jax.grad(combine_loss)(output) + assert combine_grad.shape == output.shape + assert not jnp.any(jnp.isnan(combine_grad)) + + # ===================================================================== + # Test 5: Roundtrip with padding (dispatch + combine = original) + # ===================================================================== + # Use uniform merging probs for perfect roundtrip + uniform_merging_probs = routing_map.astype(dtype) / jnp.maximum( jnp.sum(routing_map, axis=1, keepdims=True), 1.0 ) - # Dispatch tokens to experts (returns output, permuted_probs, row_id_map) - dispatched, _, row_id_map = token_dispatch(inp, routing_map, num_out_tokens) + @jax.jit + def roundtrip(x): + dispatched, _, rid_map, p_offsets, _ = token_dispatch( + x, routing_map, num_out_tokens, align_size=align_size + ) + return token_combine(dispatched, rid_map, uniform_merging_probs, p_offsets) + + roundtrip_output = roundtrip(inp) + assert_allclose(roundtrip_output, inp, dtype=dtype) - # Combine tokens back (with uniform merging) (new signature) - combined = token_combine(dispatched, row_id_map, merging_probs) + # Test roundtrip gradient + @jax.jit + def roundtrip_loss(x): + return jnp.sum(roundtrip(x) ** 2) - # Compare with original input - assert_allclose(combined, inp) + roundtrip_grad = jax.grad(roundtrip_loss)(inp) + assert roundtrip_grad.shape == inp.shape + assert not jnp.any(jnp.isnan(roundtrip_grad)) diff --git a/tests/pytorch/test_permutation.py b/tests/pytorch/test_permutation.py index e8a7bedc87..9a0cf6fb7c 100644 --- a/tests/pytorch/test_permutation.py +++ b/tests/pytorch/test_permutation.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. +import os import random import torch @@ -13,6 +14,7 @@ from transformer_engine.pytorch import ( moe_permute as te_permute, moe_permute_with_probs as te_permute_with_probs, + moe_permute_and_pad_with_probs as te_permute_and_pad_with_probs, moe_unpermute as te_unpermute, moe_sort_chunks_by_index as te_sort_chunks_by_index, moe_sort_chunks_by_index_with_probs as te_sort_chunks_by_index_with_probs, @@ -24,6 +26,7 @@ MXFP8Quantizer, ) import transformer_engine_torch as tex +from transformer_engine.pytorch import Fp8Padding, Fp8Unpadding import copy seed = 1234 @@ -653,6 +656,522 @@ def _test_permutation_mask_map( print(f"unpermute\tbwd: pytorch: {t1:.3f} ms, TE: {t2:.3f} ms") +def _test_permutation_and_padding_mask_map( + te_dtype, + num_tokens, + num_expert, + hidden_size, + topK, + num_out_tokens, + with_merging_probs=False, + align_size=16, + BENCHMARK=False, +): + if topK > num_expert: + pytest.skip("topK should be smaller than the number of experts.") + + if num_out_tokens is None: + num_out_tokens = num_tokens * topK + + print( + "permutation and padding:" + f" token:{num_tokens} hidden_size:{hidden_size} expert:{num_expert} topK:{topK}" + f" with_merging_probs:{with_merging_probs} align_size:{align_size} {te_dtype}" + ) + + # Convert TE dtypes to PyTorch dtypes + if te_dtype == tex.DType.kFloat32: + dtype = torch.float32 + elif te_dtype == tex.DType.kFloat16: + dtype = torch.float16 + elif te_dtype == tex.DType.kBFloat16: + dtype = torch.bfloat16 + else: + pytest.skip("Invalid dtype.") + + _tmp_tensor = torch.zeros((num_tokens * num_expert,)) + _tmp_tensor[: int(num_out_tokens)] = 1.0 + _tmp_idx = torch.randperm(num_tokens * num_expert) + routing_map = torch.reshape(_tmp_tensor[_tmp_idx], (num_tokens, num_expert)).bool().cuda() + + probs = torch.rand(num_tokens, num_expert).cuda() * routing_map + row_sums = probs.sum(dim=1, keepdim=True) + probs = probs / row_sums + probs = probs.to(dtype) + probs.requires_grad_(True) + + tokens_per_expert = routing_map.sum(dim=0).cpu() + target_tokens_per_expert = (torch.ceil(tokens_per_expert / align_size) * align_size).long() + num_permute_pad_out_tokens = target_tokens_per_expert.sum().item() + + permute_pad_fwd_input = torch.rand((num_tokens, hidden_size), dtype=dtype).cuda() + permute_pad_bwd_input = torch.rand( + (num_permute_pad_out_tokens, hidden_size), dtype=dtype + ).cuda() + unpermute_unpad_bwd_input = torch.rand((num_tokens, hidden_size), dtype=dtype).cuda() + permute_pad_fwd_input.requires_grad_(True) + + restore_shape = permute_pad_fwd_input.shape + ################################################################################################################################### + # + # moe_permute_with_probs and Fp8Padding, moe_unpermute and Fp8Unpadding + # + ################################################################################################################################### + # permute + padding + permuted_output, permuted_probs, row_id_map = te_permute_with_probs( + permute_pad_fwd_input, + probs, + routing_map, + num_out_tokens=num_out_tokens, + ) + tokens_per_expert_list = tokens_per_expert.tolist() + fp8_padding = Fp8Padding(num_expert, align_size) + permuted_paded_output, _ = fp8_padding(permuted_output, tokens_per_expert_list) + permuted_paded_probs, _ = fp8_padding(permuted_probs.unsqueeze(-1), tokens_per_expert_list) + + permuted_paded_output.backward(permute_pad_bwd_input, retain_graph=True) + + # unpadding + unpermute + + unpermute_unpad_fwd_input = permuted_paded_output.detach() + unpermute_unpad_fwd_input.requires_grad_(True) + + fp8_unpadding = Fp8Unpadding(num_expert, align_size) + unpaded_output = fp8_unpadding(unpermute_unpad_fwd_input, tokens_per_expert_list) + + probs_naive = probs + unpermuted_unpaded_output = te_unpermute( + unpaded_output, + row_id_map, + merging_probs=probs_naive if with_merging_probs else None, + restore_shape=restore_shape, + ) + + unpermuted_unpaded_output.backward(unpermute_unpad_bwd_input, retain_graph=True) + + ################################################################################################################################### + # + # fusion moe_permute_with_probs and Fp8Padding, fusion fusion moe_unpermute and Fp8Unpadding + # + ################################################################################################################################### + # fusion permute_and_pad + fusion_permute_and_pad_fwd_input = permute_pad_fwd_input.detach() + fusion_permute_and_pad_fwd_input.requires_grad_(True) + probs_fusion = probs_naive.detach().clone() + probs_fusion.requires_grad_(True) + + ( + fusion_permuted_padded_output, + fusion_permuted_padded_probs, + row_id_map, + pad_offsets, + target_tokens_per_expert, + ) = te_permute_and_pad_with_probs( + fusion_permute_and_pad_fwd_input, + probs_fusion, + routing_map, + tokens_per_expert, + align_size, + ) + fusion_permuted_padded_probs = fusion_permuted_padded_probs.unsqueeze(-1) + + fusion_permute_pad_bwd_input = permute_pad_bwd_input.detach() + fusion_permuted_padded_output.backward(fusion_permute_pad_bwd_input, retain_graph=True) + + # fusion unpad and unpermute + fusion_unpermute_unpad_fwd_input = fusion_permuted_padded_output.detach() + fusion_unpermute_unpad_fwd_input.requires_grad_(True) + + fusion_unpermuted_unpaded_output = te_unpermute( + fusion_unpermute_unpad_fwd_input, + row_id_map, + merging_probs=probs_fusion if with_merging_probs else None, + restore_shape=restore_shape, + pad_offsets=pad_offsets, + ) + + fusion_unpermute_bwd_input = unpermute_unpad_bwd_input.detach() + fusion_unpermuted_unpaded_output.backward(fusion_unpermute_bwd_input, retain_graph=True) + + ################################################################################################################################### + # + # Results Check + # + ################################################################################################################################### + tols = dtype_tols(te_dtype) + + permuted_paded_output_ = permuted_paded_output.float() + fusion_permuted_padded_output_ = fusion_permuted_padded_output.float() + permute_pad_fwd_input_grad = permute_pad_fwd_input.grad.float() + fusion_permute_and_pad_fwd_input_grad = fusion_permute_and_pad_fwd_input.grad.float() + + unpermuted_unpaded_output_ = unpermuted_unpaded_output.float() + fusion_unpermuted_unpaded_output_ = fusion_unpermuted_unpaded_output.float() + unpermute_unpad_fwd_input_grad = unpermute_unpad_fwd_input.grad.float() + fusion_unpermute_unpad_fwd_input_grad = fusion_unpermute_unpad_fwd_input.grad.float() + + if not BENCHMARK: + torch.testing.assert_close( + permuted_paded_output_, + fusion_permuted_padded_output_, + msg=f"Mismatch in te_permute_and_pad fwd", + **tols, + ) + torch.testing.assert_close( + permute_pad_fwd_input_grad, + fusion_permute_and_pad_fwd_input_grad, + msg=f"Mismatch in te_permute_and_pad bwd", + **tols, + ) + torch.testing.assert_close( + unpermuted_unpaded_output_, + fusion_unpermuted_unpaded_output_, + msg=f"Mismatch in te_unpermute fwd", + **tols, + ) + torch.testing.assert_close( + unpermute_unpad_fwd_input_grad, + fusion_unpermute_unpad_fwd_input_grad, + msg=f"Mismatch in te_unpermute bwd", + **tols, + ) + torch.testing.assert_close( + permuted_paded_probs.float(), + fusion_permuted_padded_probs.float(), + msg=f"Mismatch in te_permute_and_pad bwd", + **tols, + ) + if with_merging_probs: + torch.testing.assert_close( + probs_naive.grad.float(), + probs_fusion.grad.float(), + msg=f"Mismatch in te_unpermute bwd", + **tols, + ) + + ################################################################################################################################### + # + # Benchmark + # + ################################################################################################################################### + if BENCHMARK: + + def permute_and_pad(): + permuted_output, permuted_probs, row_id_map = te_permute_with_probs( + permute_pad_fwd_input, + probs, + routing_map, + num_out_tokens=num_out_tokens, + ) + fp8_padding(permuted_output, tokens_per_expert_list) + fp8_padding(permuted_probs.unsqueeze(-1), tokens_per_expert_list) + + def fusion_permute_and_pad(): + ( + fusion_permuted_padded_output, + fusion_permuted_padded_probs, + row_id_map, + pad_offsets, + target_tokens_per_expert, + ) = te_permute_and_pad_with_probs( + fusion_permute_and_pad_fwd_input, + probs, + routing_map, + tokens_per_expert, + align_size, + ) + fusion_permuted_padded_probs = fusion_permuted_padded_probs.unsqueeze(-1) + + t1 = perf_test_cuda_kernel(lambda: permute_and_pad()) + + t2 = perf_test_cuda_kernel(lambda: fusion_permute_and_pad()) + + print(f"permute_and_pad\t\tfwd: naive: {t1:.3f} ms, fusion: {t2:.3f} ms") + + t1 = perf_test_cuda_kernel( + lambda: backward_wrapper( + permuted_paded_output, + permute_pad_bwd_input, + forward_input=[permute_pad_fwd_input], + retain_graph=True, + accumulate_grad=False, + ) + ) + t2 = perf_test_cuda_kernel( + lambda: backward_wrapper( + fusion_permuted_padded_output, + fusion_permute_pad_bwd_input, + forward_input=[fusion_permute_and_pad_fwd_input], + retain_graph=True, + accumulate_grad=False, + ) + ) + print(f"permute_and_pad\t\tbwd: naive: {t1:.3f} ms, fusion: {t2:.3f} ms") + + def unpad_unpermute(): + unpaded_output = fp8_unpadding(unpermute_unpad_fwd_input, tokens_per_expert_list) + unpermuted_unpaded_output = te_unpermute( + unpaded_output, row_id_map, restore_shape=restore_shape + ) + + unpermuted_unpaded_output.backward(unpermute_unpad_bwd_input, retain_graph=True) + + t1 = perf_test_cuda_kernel(lambda: unpad_unpermute()) + t2 = perf_test_cuda_kernel( + lambda: te_unpermute( + fusion_unpermute_unpad_fwd_input, + row_id_map, + restore_shape=restore_shape, + pad_offsets=pad_offsets, + ) + ) + print(f"unpermute_and_unpad\tfwd: naive: {t1:.3f} ms, fusion: {t2:.3f} ms") + + t1 = perf_test_cuda_kernel( + lambda: backward_wrapper( + unpermuted_unpaded_output, + unpermute_unpad_bwd_input, + forward_input=([unpermute_unpad_fwd_input, probs]), + retain_graph=True, + accumulate_grad=False, + ) + ) + t2 = perf_test_cuda_kernel( + lambda: backward_wrapper( + fusion_unpermuted_unpaded_output, + fusion_unpermute_bwd_input, + forward_input=([fusion_unpermute_unpad_fwd_input, probs]), + retain_graph=True, + accumulate_grad=False, + ) + ) + print(f"unpermute_and_unpad\tbwd: naive: {t1:.3f} ms, fusion: {t2:.3f} ms") + + +def _test_permutation_and_padding_with_merging_probs( + te_dtype, + num_tokens, + num_expert, + hidden_size, + topK, + num_out_tokens, + align_size=16, + BENCHMARK=False, +): + """ + Test the combination of merging_probs AND pad_offsets together in moe_unpermute. + This specifically tests the backward pass fix where pad_offsets must be used + when computing gradients with merging_probs. + """ + if topK > num_expert: + pytest.skip("topK should be smaller than the number of experts.") + + if num_out_tokens == None: + num_out_tokens = num_tokens * topK + + print( + "permutation and padding with merging probs:" + f" token:{num_tokens} hidden_size:{hidden_size} expert:{num_expert} topK:{topK} align_size:{align_size} {te_dtype}" + ) + + # Convert TE dtypes to PyTorch dtypes + if te_dtype == tex.DType.kFloat32: + dtype = torch.float32 + elif te_dtype == tex.DType.kFloat16: + dtype = torch.float16 + elif te_dtype == tex.DType.kBFloat16: + dtype = torch.bfloat16 + else: + pytest.skip("Invalid dtype.") + + _tmp_tensor = torch.zeros((num_tokens * num_expert,)) + _tmp_tensor[: int(num_out_tokens)] = 1.0 + _tmp_idx = torch.randperm(num_tokens * num_expert) + routing_map = torch.reshape(_tmp_tensor[_tmp_idx], (num_tokens, num_expert)).bool().cuda() + + probs = torch.rand(num_tokens, num_expert).cuda() * routing_map + row_sums = probs.sum(dim=1, keepdim=True) + probs = probs / row_sums + probs = probs.to(dtype) + probs.requires_grad_(True) + + tokens_per_expert = routing_map.sum(dim=0).cpu() + target_tokens_per_expert = (torch.ceil(tokens_per_expert / align_size) * align_size).long() + num_permute_pad_out_tokens = target_tokens_per_expert.sum().item() + + permute_pad_fwd_input = torch.rand((num_tokens, hidden_size), dtype=dtype).cuda() + permute_pad_bwd_input = torch.rand( + (num_permute_pad_out_tokens, hidden_size), dtype=dtype + ).cuda() + unpermute_unpad_bwd_input = torch.rand((num_tokens, hidden_size), dtype=dtype).cuda() + permute_pad_fwd_input.requires_grad_(True) + + restore_shape = permute_pad_fwd_input.shape + ################################################################################################################################### + # + # Reference: moe_permute_with_probs + Fp8Padding, then Fp8Unpadding + moe_unpermute with merging_probs + # + ################################################################################################################################### + # permute + padding + permuted_output, permuted_probs, row_id_map = te_permute_with_probs( + permute_pad_fwd_input, + probs, + routing_map, + num_out_tokens=num_out_tokens, + ) + tokens_per_expert_list = tokens_per_expert.tolist() + fp8_padding = Fp8Padding(num_expert, align_size) + permuted_paded_output, _ = fp8_padding(permuted_output, tokens_per_expert_list) + + permuted_paded_output.backward(permute_pad_bwd_input, retain_graph=True) + + # Reference: unpadding + unpermute WITH merging_probs + ref_unpermute_fwd_input = permuted_paded_output.detach() + ref_unpermute_fwd_input.requires_grad_(True) + + ref_probs = probs.detach() + ref_probs.requires_grad_(True) + + fp8_unpadding = Fp8Unpadding(num_expert, align_size) + unpaded_output = fp8_unpadding(ref_unpermute_fwd_input, tokens_per_expert_list) + ref_unpermuted_output = te_unpermute( + unpaded_output, row_id_map, ref_probs, restore_shape=restore_shape + ) + + ref_unpermuted_output.backward(unpermute_unpad_bwd_input, retain_graph=True) + + ################################################################################################################################### + # + # Fused: moe_permute_and_pad_with_probs, then moe_unpermute with BOTH merging_probs AND pad_offsets + # + ################################################################################################################################### + # fusion permute_and_pad + fusion_permute_fwd_input = permute_pad_fwd_input.detach() + fusion_permute_fwd_input.requires_grad_(True) + fusion_probs = probs.detach() + fusion_probs.requires_grad_(True) + + ( + fusion_permuted_padded_output, + fusion_permuted_padded_probs, + fused_row_id_map, + pad_offsets, + _, + ) = te_permute_and_pad_with_probs( + fusion_permute_fwd_input, + fusion_probs, + routing_map, + tokens_per_expert, + align_size, + ) + + fusion_permute_pad_bwd_input = permute_pad_bwd_input.detach() + fusion_permuted_padded_output.backward(fusion_permute_pad_bwd_input, retain_graph=True) + + # Fused: unpermute with BOTH merging_probs AND pad_offsets + fusion_unpermute_fwd_input = fusion_permuted_padded_output.detach() + fusion_unpermute_fwd_input.requires_grad_(True) + + fusion_merging_probs = probs.detach() + fusion_merging_probs.requires_grad_(True) + + fusion_unpermuted_output = te_unpermute( + fusion_unpermute_fwd_input, + fused_row_id_map, + fusion_merging_probs, + restore_shape=restore_shape, + pad_offsets=pad_offsets, + ) + + fusion_unpermute_bwd_input = unpermute_unpad_bwd_input.detach() + fusion_unpermuted_output.backward(fusion_unpermute_bwd_input, retain_graph=True) + + ################################################################################################################################### + # + # Results Check + # + ################################################################################################################################### + tols = dtype_tols(te_dtype) + + # Check forward pass + ref_unpermuted_output_ = ref_unpermuted_output.float() + fusion_unpermuted_output_ = fusion_unpermuted_output.float() + + if not BENCHMARK: + torch.testing.assert_close( + ref_unpermuted_output_, + fusion_unpermuted_output_, + msg=f"Mismatch in te_unpermute with merging_probs and pad_offsets fwd", + **tols, + ) + + # Check backward pass - activation gradients + ref_unpermute_fwd_input_grad = ref_unpermute_fwd_input.grad.float() + fusion_unpermute_fwd_input_grad = fusion_unpermute_fwd_input.grad.float() + + torch.testing.assert_close( + ref_unpermute_fwd_input_grad, + fusion_unpermute_fwd_input_grad, + msg=f"Mismatch in te_unpermute with merging_probs and pad_offsets bwd (act_grad)", + **tols, + ) + + # Check backward pass - probs gradients + ref_probs_grad = ref_probs.grad.float() + fusion_probs_grad = fusion_merging_probs.grad.float() + + torch.testing.assert_close( + ref_probs_grad, + fusion_probs_grad, + msg=f"Mismatch in te_unpermute with merging_probs and pad_offsets bwd (probs_grad)", + **tols, + ) + + ################################################################################################################################### + # + # Benchmark + # + ################################################################################################################################### + if BENCHMARK: + + def ref_unpad_unpermute(): + unpaded = fp8_unpadding(ref_unpermute_fwd_input, tokens_per_expert_list) + return te_unpermute(unpaded, row_id_map, ref_probs, restore_shape=restore_shape) + + def fused_unpermute(): + return te_unpermute( + fusion_unpermute_fwd_input, + fused_row_id_map, + fusion_merging_probs, + restore_shape=restore_shape, + pad_offsets=pad_offsets, + ) + + t1 = perf_test_cuda_kernel(lambda: ref_unpad_unpermute()) + t2 = perf_test_cuda_kernel(lambda: fused_unpermute()) + print(f"unpermute_unpad_with_probs\tfwd: naive: {t1:.3f} ms, fusion: {t2:.3f} ms") + + t1 = perf_test_cuda_kernel( + lambda: backward_wrapper( + ref_unpermuted_output, + unpermute_unpad_bwd_input, + forward_input=[ref_unpermute_fwd_input, ref_probs], + retain_graph=True, + accumulate_grad=False, + ) + ) + t2 = perf_test_cuda_kernel( + lambda: backward_wrapper( + fusion_unpermuted_output, + fusion_unpermute_bwd_input, + forward_input=[fusion_unpermute_fwd_input, fusion_merging_probs], + retain_graph=True, + accumulate_grad=False, + ) + ) + print(f"unpermute_unpad_with_probs\tbwd: naive: {t1:.3f} ms, fusion: {t2:.3f} ms") + + def _test_permutation_mask_map_fp8( te_dtype, num_tokens, @@ -1126,7 +1645,7 @@ def perf_test_cuda_kernel(cuda_kernel_fn): @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) -@pytest.mark.parametrize("topK", [1, 2, 5]) +@pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) def test_permutation_index_map( te_dtype, @@ -1155,7 +1674,7 @@ def test_permutation_index_map( @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) -@pytest.mark.parametrize("topK", [1, 2, 5]) +@pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) def test_permutation_mask_map( te_dtype, @@ -1180,6 +1699,74 @@ def test_permutation_mask_map( ) +@pytest.mark.parametrize("te_dtype", _te_dtypes) +@pytest.mark.parametrize("num_out_tokens", [None]) +@pytest.mark.parametrize( + "num_tokens, num_expert, hidden_size, topK", + [ + (4096, 8, 1280, 2), + (4096, 64, 4096, 6), + (4096, 256, 7168, 6), + (4096, 512, 9216, 8), + ], +) +@pytest.mark.parametrize("with_merging_probs", [True, False]) +def test_permutation_and_padding_mask_map( + te_dtype, + num_tokens, + num_expert, + hidden_size, + topK, + num_out_tokens, + with_merging_probs, +): + BENCHMARK = False + + _test_permutation_and_padding_mask_map( + te_dtype=te_dtype, + num_tokens=num_tokens, + num_expert=num_expert, + hidden_size=hidden_size, + topK=topK, + num_out_tokens=num_out_tokens, + with_merging_probs=with_merging_probs, + BENCHMARK=BENCHMARK, + ) + + +@pytest.mark.parametrize("te_dtype", _te_dtypes) +@pytest.mark.parametrize("num_out_tokens", [None]) +@pytest.mark.parametrize( + "num_tokens, num_expert, hidden_size, topK", + [ + (4096, 8, 1280, 2), + (4096, 64, 4096, 6), + (4096, 256, 7168, 6), + (4096, 512, 9216, 8), + ], +) +def test_permutation_and_padding_with_merging_probs( + te_dtype, + num_tokens, + num_expert, + hidden_size, + topK, + num_out_tokens, +): + """Test moe_unpermute backward pass with BOTH merging_probs AND pad_offsets.""" + BENCHMARK = False + + _test_permutation_and_padding_with_merging_probs( + te_dtype=te_dtype, + num_tokens=num_tokens, + num_expert=num_expert, + hidden_size=hidden_size, + topK=topK, + num_out_tokens=num_out_tokens, + BENCHMARK=BENCHMARK, + ) + + @pytest.mark.parametrize("te_dtype", _te_dtypes) def test_permutation_mask_map_empty_input(te_dtype): with_probs = True @@ -1201,9 +1788,9 @@ def test_permutation_mask_map_empty_input(te_dtype): @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) -@pytest.mark.parametrize("topK", [1, 2, 5]) +@pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) -@pytest.mark.parametrize("tp_size", [1, 2, 8]) +@pytest.mark.parametrize("tp_size", [1, 2]) def test_permutation_mask_map_alongside_probs( te_dtype, num_tokens, @@ -1253,10 +1840,10 @@ def test_permutation_mask_map_alongside_probs_empty_input(te_dtype): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) @pytest.mark.parametrize("te_dtype", [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]) -@pytest.mark.parametrize("num_tokens", [2048]) +@pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) -@pytest.mark.parametrize("topK", [1, 2, 5]) +@pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) @pytest.mark.parametrize("recipe", fp8_recipes) def test_permutation_mask_map_fp8( @@ -1341,7 +1928,7 @@ def test_permutation_mask_map_topk1_no_probs( @pytest.mark.parametrize("te_dtype", _te_dtypes) @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) -@pytest.mark.parametrize("tp_size", [1, 2, 8]) +@pytest.mark.parametrize("tp_size", [2, 8]) @pytest.mark.parametrize("hidden_size", [4096]) def test_chunk_permutation( te_dtype, @@ -1376,6 +1963,10 @@ def test_chunk_permutation_empty_input(te_dtype): ) +@pytest.mark.skipif( + os.getenv("RUN_BENCHMARK_TESTS", "0") != "1", + reason="Benchmark test - run with: RUN_BENCHMARK_TESTS=1 pytest -k single_case", +) def test_permutation_single_case(): print("GPU:", torch.cuda.get_device_name(0)) @@ -1413,6 +2004,26 @@ def test_permutation_single_case(): BENCHMARK=Benchmark, ) + _test_permutation_and_padding_mask_map( + te_dtype=te_dtype, + num_tokens=num_tokens, + num_expert=num_expert, + hidden_size=hidden_size, + topK=topK, + num_out_tokens=num_out_tokens, + BENCHMARK=Benchmark, + ) + + _test_permutation_and_padding_with_merging_probs( + te_dtype=te_dtype, + num_tokens=num_tokens, + num_expert=num_expert, + hidden_size=hidden_size, + topK=topK, + num_out_tokens=num_out_tokens, + BENCHMARK=Benchmark, + ) + _test_moe_chunk_sort( te_dtype=te_dtype, num_tokens=num_tokens, @@ -1479,6 +2090,30 @@ def benchmark_single_case( ) torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_push("permutation_and_padding_mask_map") + _test_permutation_and_padding_mask_map( + te_dtype=te_dtype, + num_tokens=num_tokens, + num_expert=num_expert, + hidden_size=hidden_size, + topK=topK, + num_out_tokens=num_out_tokens, + BENCHMARK=True, + ) + torch.cuda.nvtx.range_pop() + + torch.cuda.nvtx.range_push("permutation_and_padding_with_merging_probs") + _test_permutation_and_padding_with_merging_probs( + te_dtype=te_dtype, + num_tokens=num_tokens, + num_expert=num_expert, + hidden_size=hidden_size, + topK=topK, + num_out_tokens=num_out_tokens, + BENCHMARK=True, + ) + torch.cuda.nvtx.range_pop() + torch.cuda.nvtx.range_push("permutation_mask_map_alongside_probs") _test_permutation_mask_map_alongside_probs( te_dtype=te_dtype, @@ -1495,7 +2130,12 @@ def benchmark_single_case( torch.cuda.nvtx.range_pop() -def benchmark_multiple_cases(): +@pytest.mark.skipif( + os.getenv("RUN_BENCHMARK_TESTS", "0") != "1", + reason="Benchmark test - run with: RUN_BENCHMARK_TESTS=1 pytest -k benchmark", +) +def test_benchmark_multiple_cases(): + """Benchmark test - skipped by default. Run with: RUN_BENCHMARK_TESTS=1 pytest -k benchmark""" print("GPU:", torch.cuda.get_device_name(0)) # te_dtype = tex.DType.kFloat32 @@ -1537,4 +2177,4 @@ def benchmark_multiple_cases(): if __name__ == "__main__": - benchmark_multiple_cases() + test_benchmark_multiple_cases() diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index 87a9c24533..de30c7c532 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -200,6 +200,7 @@ def _permute_kernel( probs_ptr, scale_ptr, permuted_scale_ptr, + pad_offsets_ptr, # sizes scale_hidden_dim, # strides @@ -224,8 +225,11 @@ def _permute_kernel( hidden_size: tl.constexpr, PERMUTE_PROBS: tl.constexpr, PERMUTE_SCALE: tl.constexpr, + FUSION_PAD: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): + expert_idx = 0 + pid_t = tl.program_id(0) pid_h = tl.program_id(1) cur_off = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) @@ -246,6 +250,15 @@ def _permute_kernel( dst_row = tl.load( row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert ).to(tl.int64) + if FUSION_PAD or PERMUTE_PROBS: + expert_idx = tl.load( + row_id_map_ptr + + pid_t * stride_row_id_map_token + + (num_experts + idx) * stride_row_id_map_expert + ) + if FUSION_PAD: + pad_off = tl.load(pad_offsets_ptr + expert_idx) + dst_row = dst_row + pad_off output_off = dst_row * stride_output_token + cur_off * stride_output_hidden if PERMUTE_SCALE: permuted_scale_off = ( @@ -253,11 +266,6 @@ def _permute_kernel( ) tl.store(permuted_scale_ptr + permuted_scale_off, scale, mask=mask_scale) if PERMUTE_PROBS: - expert_idx = tl.load( - row_id_map_ptr - + pid_t * stride_row_id_map_token - + (num_experts + idx) * stride_row_id_map_expert - ) prob_off = pid_t * stride_probs_token + expert_idx * stride_probs_expert prob = tl.load(probs_ptr + prob_off) if pid_h == 0: @@ -297,6 +305,7 @@ def _unpermute_kernel( row_id_map_ptr, merging_probs_ptr, permuted_probs_ptr, + pad_offsets_ptr, # strides stride_row_id_map_token, stride_row_id_map_expert, @@ -318,10 +327,12 @@ def _unpermute_kernel( PROBS_LOAD_WIDTH: tl.constexpr, WITH_MERGING_PROBS: tl.constexpr, PERMUTE_PROBS: tl.constexpr, + FUSION_UNPAD: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): data_type = input_ptr.dtype.element_ty compute_type = tl.float32 + expert_idx = 0 pid_t = tl.program_id(0) pid_h = tl.program_id(1) @@ -348,15 +359,19 @@ def _unpermute_kernel( src_row = tl.load( row_id_map_ptr + pid_t * stride_row_id_map_token + idx * stride_row_id_map_expert ).to(tl.int64) - input_off = src_row * stride_input_token + current_offset * stride_input_hidden - inp = tl.load(input_ptr + input_off, mask=mask) - inp = inp.to(compute_type) - if WITH_MERGING_PROBS: + if FUSION_UNPAD or WITH_MERGING_PROBS: expert_idx = tl.load( row_id_map_ptr + pid_t * stride_row_id_map_token + (num_experts + idx) * stride_row_id_map_expert ) + if FUSION_UNPAD: + pad_off = tl.load(pad_offsets_ptr + expert_idx) + src_row = src_row + pad_off + input_off = src_row * stride_input_token + current_offset * stride_input_hidden + inp = tl.load(input_ptr + input_off, mask=mask) + inp = inp.to(compute_type) + if WITH_MERGING_PROBS: merging_prob_off = ( pid_t * stride_merging_probs_token + expert_idx * stride_merging_probs_expert ) @@ -407,6 +422,7 @@ def _unpermute_bwd_with_merging_probs_kernel( fwd_input_ptr, merging_probs_ptr, row_id_map_ptr, + pad_offsets_ptr, # strides stride_row_id_map_token, stride_row_id_map_expert, @@ -427,6 +443,7 @@ def _unpermute_bwd_with_merging_probs_kernel( num_experts: tl.constexpr, hidden_size: tl.constexpr, PROBS_LOAD_WIDTH: tl.constexpr, + FUSION_UNPAD: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): data_type = fwd_output_grad_ptr.dtype.element_ty @@ -450,6 +467,9 @@ def _unpermute_bwd_with_merging_probs_kernel( + pid * stride_row_id_map_token + (num_experts + idx) * stride_row_id_map_expert ) + if FUSION_UNPAD: + pad_off = tl.load(pad_offsets_ptr + expert_idx) + dst_row = dst_row + pad_off prob_grad_accum = tl.zeros((BLOCK_SIZE,), dtype=compute_type) current_start = 0 while current_start < hidden_size: diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 55a59a1650..32de0b1a3c 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -25,8 +25,11 @@ 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, ) @@ -43,7 +46,14 @@ def token_dispatch( routing_map: jnp.ndarray, num_out_tokens: int, probs: Optional[jnp.ndarray] = None, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: + align_size: Optional[int] = None, +) -> Tuple[ + jnp.ndarray, + Optional[jnp.ndarray], + jnp.ndarray, + Optional[jnp.ndarray], + Optional[jnp.ndarray], +]: """ Dispatch tokens to experts based on routing map. @@ -51,6 +61,10 @@ def token_dispatch( to their designated experts according to the routing map. The row_id_map is computed internally from the routing_map. + Optionally supports fused padding for alignment when `align_size` is provided. + This is useful for efficient matrix multiplications that require aligned tensor + dimensions. The padding is computed internally from the routing_map. + Parameters ---------- inp : jnp.ndarray @@ -59,36 +73,99 @@ def token_dispatch( Routing mask of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. Values: 1 = routed, 0 = not routed. num_out_tokens : int - The number of output tokens after permutation. This should equal the sum of - routing_map and must be provided explicitly for JIT compatibility. + The number of output tokens after permutation (before padding). For the dropless + case, this should be equal to the sum of routing_map. Must be provided explicitly + for JIT compatibility since output shape must be known at compile time. probs : Optional[jnp.ndarray] Optional routing probabilities of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. If provided, permuted_probs will be returned. + align_size : Optional[int] + Optional alignment size for padding. If provided, outputs will be padded to + align each expert's tokens to a multiple of this size. The output buffer is + allocated with worst-case size, rounded down to align_size: + ((num_out_tokens + num_experts * (align_size - 1)) // align_size) * align_size + This enables full JIT compatibility. Returns ------- output : jnp.ndarray - Permuted output tensor of shape [num_out_tokens, hidden_size]. + Permuted output tensor of shape [num_out_tokens, hidden_size] without padding, + or [worst_case_padded_size, hidden_size] when using padding fusion. + With padding, the actual used portion may be smaller than the buffer; check + actual_num_out_tokens (sum of target_tokens_per_expert) for the actual size. permuted_probs : Optional[jnp.ndarray] - Permuted probabilities of shape [num_out_tokens], or None if probs was not provided. + Permuted probabilities of shape [num_out_tokens] or [worst_case_padded_size], + or None if probs was not provided. row_id_map : jnp.ndarray Row ID map for use in token_combine (shape [num_tokens, num_experts * 2 + 1]). + pad_offsets : Optional[jnp.ndarray] + Per-expert cumulative padding offsets of shape [num_experts] when using padding, + None otherwise. Pass this to token_combine when unpadding is needed. + target_tokens_per_expert : Optional[jnp.ndarray] + Aligned token counts per expert of shape [num_experts] when using padding, + None otherwise. + + Note + ---- + **JIT Compatibility:** + + This function is fully JIT-compatible. When using padding (align_size provided), + the output buffer is allocated with a fixed worst-case size that depends only on + compile-time constants (num_out_tokens, num_experts, align_size). The actual + padding offsets (pad_offsets) and aligned token counts (target_tokens_per_expert) + are computed internally from the routing_map and can be traced values. + + The worst-case output size is: + ((num_out_tokens + num_experts * (align_size - 1)) // align_size) * align_size + This accounts for the maximum possible padding when each expert needs (align_size - 1) + extra tokens to align, rounded down to align_size for buffer alignment. """ - return _token_dispatch(inp, routing_map, probs, num_out_tokens) + use_padding = align_size is not None + num_experts = routing_map.shape[-1] + if use_padding: + # Compute worst-case output size (compile-time constant) + # This is the maximum possible size when each expert needs max padding + worst_case_out_tokens = ( + (num_out_tokens + num_experts * (align_size - 1)) // align_size + ) * align_size + else: + worst_case_out_tokens = num_out_tokens + + return _token_dispatch( + inp, routing_map, probs, num_out_tokens, worst_case_out_tokens, align_size, use_padding + ) -@partial(jax.custom_vjp, nondiff_argnums=(1, 3)) + +@partial(jax.custom_vjp, nondiff_argnums=(1, 3, 4, 5, 6)) def _token_dispatch( inp: jnp.ndarray, routing_map: jnp.ndarray, probs: Optional[jnp.ndarray], num_out_tokens: int, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray]: + worst_case_out_tokens: int, + align_size: Optional[int], + use_padding: bool, +) -> Tuple[ + jnp.ndarray, + Optional[jnp.ndarray], + jnp.ndarray, + Optional[jnp.ndarray], + Optional[jnp.ndarray], +]: """Internal token_dispatch with custom VJP.""" - (output, permuted_probs, row_id_map), _ = _token_dispatch_fwd_rule( - inp, routing_map, probs, num_out_tokens + (output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert), _ = ( + _token_dispatch_fwd_rule( + inp, + routing_map, + probs, + num_out_tokens, + worst_case_out_tokens, + align_size, + use_padding, + ) ) - return output, permuted_probs, row_id_map + return output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert def _token_dispatch_fwd_rule( @@ -96,9 +173,18 @@ def _token_dispatch_fwd_rule( routing_map: jnp.ndarray, probs: Optional[jnp.ndarray], num_out_tokens: int, + worst_case_out_tokens: int, + align_size: Optional[int], + use_padding: bool, ) -> Tuple[ - Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray], - Tuple[jnp.ndarray, int, int, int, bool], + Tuple[ + jnp.ndarray, + Optional[jnp.ndarray], + jnp.ndarray, + Optional[jnp.ndarray], + Optional[jnp.ndarray], + ], + Tuple[jnp.ndarray, Optional[jnp.ndarray], int, int, int, bool], ]: """Forward pass rule for token_dispatch.""" # Validate input dimensions @@ -126,42 +212,102 @@ def _token_dispatch_fwd_rule( with_probs = probs is not None - output, permuted_probs = permute_with_mask_map( - inp, - row_id_map, - probs, - num_tokens, - num_experts, - num_out_tokens, - hidden_size, - ) + if use_padding: + # Compute tokens_per_expert internally from routing_map + # This can be a traced value since output shape uses worst_case_out_tokens + tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) + + # Calculate aligned token counts per expert + target_tokens_per_expert = (jnp.ceil(tokens_per_expert / align_size) * align_size).astype( + jnp.int32 + ) + + # Compute pad_offsets: cumulative padding for each expert + # pad_offsets[i] = sum of (target - actual) for experts 0..i-1 + 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]]) + + # Use worst_case_out_tokens as the output buffer size (compile-time constant) + # The actual used size is sum(target_tokens_per_expert), which may be smaller. + # Unused positions will be zero-initialized by the kernel. + output, permuted_probs = permute_with_mask_map_and_pad( + inp, + row_id_map, + probs, + pad_offsets, + num_tokens, + num_experts, + worst_case_out_tokens, + hidden_size, + ) + else: + # No padding + pad_offsets = None + target_tokens_per_expert = None + + output, permuted_probs = permute_with_mask_map( + inp, + row_id_map, + probs, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) # Return (primals, residuals) - # Include with_probs flag to know how to handle backward pass - residuals = (row_id_map, num_tokens, num_experts, hidden_size, with_probs) - return (output, permuted_probs, row_id_map), residuals + residuals = (row_id_map, pad_offsets, num_tokens, num_experts, hidden_size, with_probs) + return ( + output, + permuted_probs, + row_id_map, + pad_offsets, + target_tokens_per_expert, + ), residuals def _token_dispatch_bwd_rule( _routing_map: jnp.ndarray, _num_out_tokens: int, - residuals: Tuple[jnp.ndarray, int, int, int, bool], - g: Tuple[jnp.ndarray, Optional[jnp.ndarray], jnp.ndarray], + _worst_case_out_tokens: int, + _align_size: Optional[int], + _use_padding: bool, + residuals: Tuple[jnp.ndarray, Optional[jnp.ndarray], int, int, int, bool], + g: Tuple[ + jnp.ndarray, + Optional[jnp.ndarray], + jnp.ndarray, + Optional[jnp.ndarray], + Optional[jnp.ndarray], + ], ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: """Backward pass rule for token_dispatch.""" - row_id_map, num_tokens, num_experts, hidden_size, with_probs = residuals - output_grad, permuted_probs_grad, _ = g # Ignore row_id_map gradient + row_id_map, pad_offsets, num_tokens, num_experts, hidden_size, with_probs = residuals + output_grad, permuted_probs_grad, _, _, _ = g # Ignore row_id_map, pad_offsets, target grads # Backward: unpermute gradients (gather from experts back to tokens) - inp_grad, probs_grad = unpermute_with_mask_map( - output_grad, - row_id_map, - None, # No merging probs - permuted_probs_grad if with_probs else None, - num_tokens, - num_experts, - hidden_size, - ) + if pad_offsets is not None: + inp_grad, probs_grad = unpermute_with_mask_map_and_unpad( + output_grad, + row_id_map, + None, # No merging probs + permuted_probs_grad if with_probs else None, + pad_offsets, + num_tokens, + num_experts, + hidden_size, + ) + else: + inp_grad, probs_grad = unpermute_with_mask_map( + output_grad, + row_id_map, + None, # No merging probs + permuted_probs_grad if with_probs else None, + num_tokens, + num_experts, + hidden_size, + ) return inp_grad, probs_grad if with_probs else None @@ -178,6 +324,7 @@ def token_combine( inp: jnp.ndarray, row_id_map: jnp.ndarray, merging_probs: Optional[jnp.ndarray] = None, + pad_offsets: Optional[jnp.ndarray] = None, ) -> jnp.ndarray: """ Combine tokens from experts back to original token positions. @@ -185,33 +332,42 @@ def token_combine( This is the forward pass of MoE unpermutation. Tokens are gathered from experts and merged (optionally weighted by merging_probs). + Optionally supports fused unpadding when `pad_offsets` is provided (from + token_dispatch with padding enabled). + Parameters ---------- inp : jnp.ndarray - Input tensor from experts of shape [num_out_tokens, hidden_size]. + Input tensor from experts of shape [num_out_tokens, hidden_size] + (or [num_out_tokens_padded, hidden_size] when using unpadding). row_id_map : jnp.ndarray Row ID map from token_dispatch of shape [num_tokens, num_experts * 2 + 1]. merging_probs : Optional[jnp.ndarray] Merging weights of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. If provided, tokens from different experts are weighted-summed. If None, tokens are summed directly. + pad_offsets : Optional[jnp.ndarray] + Per-expert cumulative padding offsets of shape [num_experts] from token_dispatch. + If provided, fused unpadding will be performed. This should be the pad_offsets + returned by token_dispatch when using padding. Returns ------- output : jnp.ndarray Combined output tensor of shape [num_tokens, hidden_size]. """ - return _token_combine(inp, row_id_map, merging_probs) + return _token_combine(inp, row_id_map, merging_probs, pad_offsets) -@partial(jax.custom_vjp, nondiff_argnums=(1,)) +@jax.custom_vjp def _token_combine( inp: jnp.ndarray, row_id_map: jnp.ndarray, merging_probs: Optional[jnp.ndarray], + pad_offsets: Optional[jnp.ndarray], ) -> jnp.ndarray: """Internal token_combine with custom VJP.""" - output, _ = _token_combine_fwd_rule(inp, row_id_map, merging_probs) + output, _ = _token_combine_fwd_rule(inp, row_id_map, merging_probs, pad_offsets) return output @@ -219,7 +375,20 @@ def _token_combine_fwd_rule( inp: jnp.ndarray, row_id_map: jnp.ndarray, merging_probs: Optional[jnp.ndarray], -) -> Tuple[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray, Optional[jnp.ndarray], int, int, int, int]]: + pad_offsets: Optional[jnp.ndarray], +) -> Tuple[ + jnp.ndarray, + Tuple[ + jnp.ndarray, + Optional[jnp.ndarray], + jnp.ndarray, + Optional[jnp.ndarray], + int, + int, + int, + int, + ], +]: """Forward pass rule for token_combine.""" # Infer dimensions from row_id_map shape: [num_tokens, num_experts * 2 + 1] num_tokens = row_id_map.shape[0] @@ -227,21 +396,34 @@ def _token_combine_fwd_rule( hidden_size = inp.shape[-1] num_out_tokens = inp.shape[0] - # Call triton extension - output, _ = unpermute_with_mask_map( - inp, - row_id_map, - merging_probs, - None, # No permuted probs to unpermute - num_tokens, - num_experts, - hidden_size, - ) + # Call triton extension with or without unpadding + if pad_offsets is not None: + output, _ = unpermute_with_mask_map_and_unpad( + inp, + row_id_map, + merging_probs, + None, # No permuted probs to unpermute + pad_offsets, + num_tokens, + num_experts, + hidden_size, + ) + else: + output, _ = unpermute_with_mask_map( + inp, + row_id_map, + merging_probs, + None, # No permuted probs to unpermute + num_tokens, + num_experts, + hidden_size, + ) # Return (primal, residuals) # Include inp in residuals for backward with merging_probs residuals = ( row_id_map, + pad_offsets, inp, merging_probs, num_tokens, @@ -253,13 +435,26 @@ def _token_combine_fwd_rule( def _token_combine_bwd_rule( - row_id_map: jnp.ndarray, - residuals: Tuple[jnp.ndarray, jnp.ndarray, Optional[jnp.ndarray], int, int, int, int], + residuals: Tuple[ + jnp.ndarray, + Optional[jnp.ndarray], + jnp.ndarray, + Optional[jnp.ndarray], + int, + int, + int, + int, + ], g: jnp.ndarray, -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Backward pass rule for token_combine.""" +) -> Tuple[jnp.ndarray, None, Optional[jnp.ndarray], None]: + """Backward pass rule for token_combine. + + Returns gradients for: (inp, row_id_map, merging_probs, pad_offsets) + row_id_map and pad_offsets are integer arrays, so their gradients are None. + """ ( row_id_map, + pad_offsets, fwd_input, merging_probs, num_tokens, @@ -273,30 +468,63 @@ def _token_combine_bwd_rule( if with_merging_probs: # Use specialized backward kernel that properly scales by merging_probs - inp_grad, merging_probs_grad = unpermute_bwd_with_merging_probs( - output_grad, - row_id_map, - fwd_input, - merging_probs, - num_tokens, - num_experts, - num_out_tokens, - hidden_size, - ) + if pad_offsets is not None: + inp_grad, merging_probs_grad = unpermute_bwd_with_merging_probs_and_unpad( + output_grad, + row_id_map, + fwd_input, + merging_probs, + pad_offsets, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) + # The backward kernel only writes to positions that tokens map to. + # Padded positions may contain uninitialized (NaN) values - replace with zeros. + inp_grad = jnp.where(jnp.isnan(inp_grad), 0.0, inp_grad) + else: + inp_grad, merging_probs_grad = unpermute_bwd_with_merging_probs( + output_grad, + row_id_map, + fwd_input, + merging_probs, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) else: # Simple case: just permute gradients back - inp_grad, _ = permute_with_mask_map( - output_grad, - row_id_map, - None, - num_tokens, - num_experts, - num_out_tokens, - hidden_size, - ) + if pad_offsets is not None: + inp_grad, _ = permute_with_mask_map_and_pad( + output_grad, + row_id_map, + None, + pad_offsets, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) + # The permute kernel only writes to positions that tokens map to. + # Padded positions may contain uninitialized (NaN) values - replace with zeros. + inp_grad = jnp.where(jnp.isnan(inp_grad), 0.0, inp_grad) + else: + inp_grad, _ = permute_with_mask_map( + output_grad, + row_id_map, + None, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ) merging_probs_grad = None - return inp_grad, merging_probs_grad + # Return gradients for: inp, row_id_map, merging_probs, pad_offsets + # row_id_map and pad_offsets are integer arrays, so their gradients are None + return inp_grad, None, merging_probs_grad, None _token_combine.defvjp(_token_combine_fwd_rule, _token_combine_bwd_rule) diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index 4f59f65a87..01b15c5adc 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -27,8 +27,11 @@ __all__ = [ "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", ] @@ -243,20 +246,21 @@ def lowering(ctx, row_id_map, *, num_tokens, num_experts): class PermuteWithMaskMapPrimitive(BasePrimitive): """ - Permute the input tensor based on the row_id_map. + Permute the input tensor based on the row_id_map, optionally with fused padding. """ name = "te_permute_with_mask_map_triton" multiple_results = True - # scale and permuted_scale are dummy inputs (not used when PERMUTE_SCALE=False) - # but they need to be in the signature for the kernel call + # scale, permuted_scale are dummy inputs (not used when PERMUTE_SCALE=False) + # pad_offsets can be shape (0,) when not doing padding, or (num_experts,) when padding impl_static_args = ( - 5, 6, 7, 8, 9, - ) # num_tokens, num_experts, num_out_tokens, hidden_size, with_probs + 10, + 11, + ) # num_tokens, num_experts, num_out_tokens, hidden_size, with_probs, with_pad inner_primitive = None outer_primitive = None @@ -267,16 +271,18 @@ def abstract( probs_aval, scale_aval, # dummy, same shape as inp permuted_scale_aval, # dummy, same shape as inp + pad_offsets_aval, *, num_tokens, num_experts, num_out_tokens, hidden_size, with_probs, + with_pad, ): """Shape/dtype inference for permute.""" - del row_id_map_aval, scale_aval, permuted_scale_aval - del num_tokens, num_experts + del row_id_map_aval, scale_aval, permuted_scale_aval, pad_offsets_aval + del num_tokens, num_experts, with_pad output_shape = (num_out_tokens, hidden_size) output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) @@ -295,11 +301,13 @@ def impl( probs, scale, permuted_scale, + pad_offsets, num_tokens, num_experts, num_out_tokens, hidden_size, with_probs, + with_pad, ): """Forward to inner primitive.""" assert PermuteWithMaskMapPrimitive.inner_primitive is not None @@ -309,11 +317,13 @@ def impl( probs, scale, permuted_scale, + pad_offsets, num_tokens=num_tokens, num_experts=num_experts, num_out_tokens=num_out_tokens, hidden_size=hidden_size, with_probs=with_probs, + with_pad=with_pad, ) @staticmethod @@ -324,12 +334,14 @@ def lowering( probs, scale, permuted_scale, + pad_offsets, *, num_tokens, num_experts, num_out_tokens, hidden_size, with_probs, + with_pad, ): """MLIR lowering using triton_call_lowering.""" del num_out_tokens @@ -367,6 +379,7 @@ def lowering( probs, scale, permuted_scale, + pad_offsets, grid=grid, constexprs={ "scale_hidden_dim": 0, @@ -387,6 +400,7 @@ def lowering( "hidden_size": hidden_size, "PERMUTE_PROBS": with_probs, "PERMUTE_SCALE": False, + "FUSION_PAD": with_pad, "BLOCK_SIZE": block_size, }, ) @@ -403,11 +417,11 @@ class UnpermuteWithMaskMapPrimitive(BasePrimitive): name = "te_unpermute_with_mask_map_triton" multiple_results = True impl_static_args = ( - 4, 5, 6, 7, 8, + 9, ) # num_tokens, num_experts, hidden_size, with_merging_probs, with_probs inner_primitive = None outer_primitive = None @@ -418,6 +432,7 @@ def abstract( row_id_map_aval, merging_probs_aval, permuted_probs_aval, + pad_offsets_aval, # dummy, not used when FUSION_UNPAD=False *, num_tokens, num_experts, @@ -426,7 +441,7 @@ def abstract( with_probs, ): """Shape/dtype inference for unpermute.""" - del row_id_map_aval, merging_probs_aval, with_merging_probs + del row_id_map_aval, merging_probs_aval, with_merging_probs, pad_offsets_aval output_shape = (num_tokens, hidden_size) output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) @@ -447,6 +462,7 @@ def impl( row_id_map, merging_probs, permuted_probs, + pad_offsets, num_tokens, num_experts, hidden_size, @@ -460,6 +476,7 @@ def impl( row_id_map, merging_probs, permuted_probs, + pad_offsets, num_tokens=num_tokens, num_experts=num_experts, hidden_size=hidden_size, @@ -474,6 +491,7 @@ def lowering( row_id_map, merging_probs, permuted_probs, + pad_offsets, *, num_tokens, num_experts, @@ -505,6 +523,7 @@ def lowering( block_size = _get_min_block_size(_unpermute_kernel) grid = (num_tokens, triton.cdiv(hidden_size, block_size)) + # Pass all 5 inputs including pad_offsets (even though FUSION_UNPAD=False) return triton_call_lowering( ctx, _unpermute_kernel, @@ -512,6 +531,7 @@ def lowering( row_id_map, merging_probs, permuted_probs, + pad_offsets, grid=grid, constexprs={ "stride_row_id_map_token": row_id_stride_token, @@ -530,6 +550,7 @@ def lowering( "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), "WITH_MERGING_PROBS": with_merging_probs, "PERMUTE_PROBS": with_probs, + "FUSION_UNPAD": False, "BLOCK_SIZE": block_size, }, ) @@ -538,6 +559,155 @@ def lowering( register_primitive(UnpermuteWithMaskMapPrimitive) +class UnpermuteWithMaskMapAndUnpadPrimitive(BasePrimitive): + """ + Unpermute the input tensor based on the row_id_map with fused unpadding. + """ + + name = "te_unpermute_with_mask_map_and_unpad_triton" + multiple_results = True + impl_static_args = ( + 5, + 6, + 7, + 8, + 9, + ) # num_tokens, num_experts, hidden_size, with_merging_probs, with_probs + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + inp_aval, + row_id_map_aval, + merging_probs_aval, + permuted_probs_aval, + pad_offsets_aval, + *, + num_tokens, + num_experts, + hidden_size, + with_merging_probs, + with_probs, + ): + """Shape/dtype inference for unpermute with unpadding.""" + del row_id_map_aval, merging_probs_aval, with_merging_probs, pad_offsets_aval + + output_shape = (num_tokens, hidden_size) + output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) + + if with_probs: + unpermuted_probs_shape = (num_tokens, num_experts) + unpermuted_probs_aval = jax.core.ShapedArray( + unpermuted_probs_shape, permuted_probs_aval.dtype + ) + else: + unpermuted_probs_aval = jax.core.ShapedArray((0,), inp_aval.dtype) + + return output_aval, unpermuted_probs_aval + + @staticmethod + def impl( + inp, + row_id_map, + merging_probs, + permuted_probs, + pad_offsets, + num_tokens, + num_experts, + hidden_size, + with_merging_probs, + with_probs, + ): + """Forward to inner primitive.""" + assert UnpermuteWithMaskMapAndUnpadPrimitive.inner_primitive is not None + return UnpermuteWithMaskMapAndUnpadPrimitive.inner_primitive.bind( + inp, + row_id_map, + merging_probs, + permuted_probs, + pad_offsets, + num_tokens=num_tokens, + num_experts=num_experts, + hidden_size=hidden_size, + with_merging_probs=with_merging_probs, + with_probs=with_probs, + ) + + @staticmethod + def lowering( + ctx, + inp, + row_id_map, + merging_probs, + permuted_probs, + pad_offsets, + *, + num_tokens, + num_experts, + hidden_size, + with_merging_probs, + with_probs, + ): + """MLIR lowering using triton_call_lowering.""" + # Compute strides + inp_stride_token = hidden_size + inp_stride_hidden = 1 + output_stride_token = hidden_size + output_stride_hidden = 1 + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + + if with_merging_probs: + merging_probs_stride_token = num_experts + merging_probs_stride_expert = 1 + else: + merging_probs_stride_token = 0 + merging_probs_stride_expert = 0 + + permuted_probs_stride_token = 1 + unpermuted_probs_stride_token = num_experts + unpermuted_probs_stride_expert = 1 + + # Grid - use minimum BLOCK_SIZE from autotune configs + block_size = _get_min_block_size(_unpermute_kernel) + grid = (num_tokens, triton.cdiv(hidden_size, block_size)) + + return triton_call_lowering( + ctx, + _unpermute_kernel, + inp, + row_id_map, + merging_probs, + permuted_probs, + pad_offsets, + grid=grid, + constexprs={ + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "stride_input_token": inp_stride_token, + "stride_input_hidden": inp_stride_hidden, + "stride_output_token": output_stride_token, + "stride_output_hidden": output_stride_hidden, + "stride_merging_probs_token": merging_probs_stride_token, + "stride_merging_probs_expert": merging_probs_stride_expert, + "stride_permuted_probs_token": permuted_probs_stride_token, + "stride_unpermuted_probs_token": unpermuted_probs_stride_token, + "stride_unpermuted_probs_expert": unpermuted_probs_stride_expert, + "num_experts": num_experts, + "hidden_size": hidden_size, + "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), + "WITH_MERGING_PROBS": with_merging_probs, + "PERMUTE_PROBS": with_probs, + "FUSION_UNPAD": True, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(UnpermuteWithMaskMapAndUnpadPrimitive) + + class UnpermuteBwdWithMergingProbsPrimitive(BasePrimitive): """ Backward pass for unpermute with merging probabilities. @@ -547,7 +717,7 @@ class UnpermuteBwdWithMergingProbsPrimitive(BasePrimitive): name = "te_unpermute_bwd_with_merging_probs_triton" multiple_results = True - impl_static_args = (4, 5, 6, 7) # num_tokens, num_experts, num_out_tokens, hidden_size + impl_static_args = (5, 6, 7, 8) # num_tokens, num_experts, num_out_tokens, hidden_size inner_primitive = None outer_primitive = None @@ -557,6 +727,7 @@ def abstract( fwd_input_aval, merging_probs_aval, row_id_map_aval, + pad_offsets_aval, # dummy, not used when FUSION_UNPAD=False *, num_tokens, num_experts, @@ -564,7 +735,7 @@ def abstract( hidden_size, ): """Shape/dtype inference for unpermute backward with merging probs.""" - del fwd_input_aval, row_id_map_aval + del fwd_input_aval, row_id_map_aval, pad_offsets_aval # fwd_input_grad has same shape as fwd_input fwd_input_grad_shape = (num_out_tokens, hidden_size) @@ -584,6 +755,7 @@ def impl( fwd_input, merging_probs, row_id_map, + pad_offsets, num_tokens, num_experts, num_out_tokens, @@ -596,6 +768,7 @@ def impl( fwd_input, merging_probs, row_id_map, + pad_offsets, num_tokens=num_tokens, num_experts=num_experts, num_out_tokens=num_out_tokens, @@ -609,6 +782,7 @@ def lowering( fwd_input, merging_probs, row_id_map, + pad_offsets, *, num_tokens, num_experts, @@ -638,7 +812,7 @@ def lowering( # Get min block size from autotune configs for consistency block_size = _get_min_block_size(_unpermute_bwd_with_merging_probs_kernel) - # Pass inputs in kernel argument order: fwd_output_grad, fwd_input, merging_probs, row_id_map + # Pass all 5 inputs including pad_offsets (even though FUSION_UNPAD=False) return triton_call_lowering( ctx, _unpermute_bwd_with_merging_probs_kernel, @@ -646,6 +820,7 @@ def lowering( fwd_input, merging_probs, row_id_map, + pad_offsets, grid=grid, constexprs={ "stride_row_id_map_token": row_id_stride_token, @@ -663,6 +838,7 @@ def lowering( "num_experts": num_experts, "hidden_size": hidden_size, "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), + "FUSION_UNPAD": False, "BLOCK_SIZE": block_size, }, ) @@ -671,6 +847,145 @@ def lowering( register_primitive(UnpermuteBwdWithMergingProbsPrimitive) +class UnpermuteBwdWithMergingProbsAndUnpadPrimitive(BasePrimitive): + """ + Backward pass for unpermute with merging probabilities and fused unpadding. + + This kernel computes gradients for both the input and merging_probs, + while handling padded outputs. + """ + + name = "te_unpermute_bwd_with_merging_probs_and_unpad_triton" + multiple_results = True + impl_static_args = (5, 6, 7, 8) # num_tokens, num_experts, num_out_tokens, hidden_size + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + fwd_output_grad_aval, + fwd_input_aval, + merging_probs_aval, + row_id_map_aval, + pad_offsets_aval, + *, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ): + """Shape/dtype inference for unpermute backward with merging probs and unpadding.""" + del fwd_input_aval, row_id_map_aval, pad_offsets_aval + + # fwd_input_grad has same shape as fwd_input + fwd_input_grad_shape = (num_out_tokens, hidden_size) + fwd_input_grad_aval = jax.core.ShapedArray(fwd_input_grad_shape, fwd_output_grad_aval.dtype) + + # merging_probs_grad has same shape as merging_probs + merging_probs_grad_shape = (num_tokens, num_experts) + merging_probs_grad_aval = jax.core.ShapedArray( + merging_probs_grad_shape, merging_probs_aval.dtype + ) + + return fwd_input_grad_aval, merging_probs_grad_aval + + @staticmethod + def impl( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + pad_offsets, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ): + """Forward to inner primitive.""" + assert UnpermuteBwdWithMergingProbsAndUnpadPrimitive.inner_primitive is not None + return UnpermuteBwdWithMergingProbsAndUnpadPrimitive.inner_primitive.bind( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + pad_offsets, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + ) + + @staticmethod + def lowering( + ctx, + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + pad_offsets, + *, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + ): + """MLIR lowering using triton_call_lowering.""" + del num_out_tokens + + # Compute strides + row_id_stride_token = num_experts * 2 + 1 + row_id_stride_expert = 1 + fwd_output_grad_stride_token = hidden_size + fwd_output_grad_stride_hidden = 1 + fwd_input_grad_stride_token = hidden_size + fwd_input_grad_stride_hidden = 1 + fwd_input_stride_token = hidden_size + fwd_input_stride_hidden = 1 + merging_probs_stride_token = num_experts + merging_probs_stride_expert = 1 + merging_probs_grad_stride_token = num_experts + merging_probs_grad_stride_expert = 1 + + # Grid - one program per token + grid = (num_tokens,) + + # Get min block size from autotune configs for consistency + block_size = _get_min_block_size(_unpermute_bwd_with_merging_probs_kernel) + + return triton_call_lowering( + ctx, + _unpermute_bwd_with_merging_probs_kernel, + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + pad_offsets, + grid=grid, + constexprs={ + "stride_row_id_map_token": row_id_stride_token, + "stride_row_id_map_expert": row_id_stride_expert, + "stride_fwd_output_grad_token": fwd_output_grad_stride_token, + "stride_fwd_output_grad_hidden": fwd_output_grad_stride_hidden, + "stride_fwd_input_grad_token": fwd_input_grad_stride_token, + "stride_fwd_input_grad_hidden": fwd_input_grad_stride_hidden, + "stride_fwd_input_token": fwd_input_stride_token, + "stride_fwd_input_hidden": fwd_input_stride_hidden, + "stride_merging_probs_token": merging_probs_stride_token, + "stride_merging_probs_expert": merging_probs_stride_expert, + "stride_merging_probs_grad_token": merging_probs_grad_stride_token, + "stride_merging_probs_grad_expert": merging_probs_grad_stride_expert, + "num_experts": num_experts, + "hidden_size": hidden_size, + "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), + "FUSION_UNPAD": True, + "BLOCK_SIZE": block_size, + }, + ) + + +register_primitive(UnpermuteBwdWithMergingProbsAndUnpadPrimitive) + + def unpermute_bwd_with_merging_probs( fwd_output_grad: jnp.ndarray, row_id_map: jnp.ndarray, @@ -712,12 +1027,73 @@ def unpermute_bwd_with_merging_probs( merging_probs_grad : jnp.ndarray Gradient w.r.t. merging_probs of shape `[num_tokens, num_experts]`. """ - # Pass arguments in kernel order: fwd_output_grad, fwd_input, merging_probs, row_id_map + # Create dummy pad_offsets (not used when FUSION_UNPAD=False, but required by kernel signature) + dummy_pad_offsets = jnp.zeros((0,), dtype=jnp.int32) + # Pass arguments in kernel order: fwd_output_grad, fwd_input, merging_probs, row_id_map, pad_offsets return UnpermuteBwdWithMergingProbsPrimitive.outer_primitive.bind( fwd_output_grad, fwd_input, merging_probs, row_id_map, + dummy_pad_offsets, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + ) + + +def unpermute_bwd_with_merging_probs_and_unpad( + fwd_output_grad: jnp.ndarray, + row_id_map: jnp.ndarray, + fwd_input: jnp.ndarray, + merging_probs: jnp.ndarray, + pad_offsets: jnp.ndarray, + num_tokens: int, + num_experts: int, + num_out_tokens: int, + hidden_size: int, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Backward pass for unpermute with merging probabilities and fused unpadding. + + This computes gradients for both the input tensor and merging_probs, + while handling padded outputs. + + Parameters + ---------- + fwd_output_grad : jnp.ndarray + Gradient of the forward output of shape `[num_tokens, hidden_size]`. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. + fwd_input : jnp.ndarray + The input tensor from the forward pass of shape `[num_out_tokens, hidden_size]`. + merging_probs : jnp.ndarray + The merging probabilities of shape `[num_tokens, num_experts]`. + pad_offsets : jnp.ndarray + Per-expert cumulative padding offsets of shape `[num_experts]`. + num_tokens : int + Number of tokens in the unpermuted tensor. + num_experts : int + Number of experts. + num_out_tokens : int + Number of tokens in the permuted tensor (including padding). + hidden_size : int + Hidden size. + + Returns + ------- + fwd_input_grad : jnp.ndarray + Gradient w.r.t. the input tensor of shape `[num_out_tokens, hidden_size]`. + merging_probs_grad : jnp.ndarray + Gradient w.r.t. merging_probs of shape `[num_tokens, num_experts]`. + """ + return UnpermuteBwdWithMergingProbsAndUnpadPrimitive.outer_primitive.bind( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + pad_offsets, num_tokens=num_tokens, num_experts=num_experts, num_out_tokens=num_out_tokens, @@ -957,6 +1333,78 @@ def permute_with_mask_map( """ with_probs = probs is not None + # Handle None probs by creating dummy tensor + if not with_probs: + probs = jnp.zeros((0,), dtype=inp.dtype) + + # Create dummy scale tensors (not used when PERMUTE_SCALE=False, but required by kernel signature) + dummy_scale = inp + dummy_permuted_scale = inp + # Create dummy pad_offsets (not used when FUSION_PAD=False, but required by kernel signature) + dummy_pad_offsets = jnp.zeros((0,), dtype=jnp.int32) + + output, permuted_probs = PermuteWithMaskMapPrimitive.outer_primitive.bind( + inp, + row_id_map, + probs, + dummy_scale, + dummy_permuted_scale, + dummy_pad_offsets, + num_tokens=num_tokens, + num_experts=num_experts, + num_out_tokens=num_out_tokens, + hidden_size=hidden_size, + with_probs=with_probs, + with_pad=False, + ) + + if not with_probs: + permuted_probs = None + + return output, permuted_probs + + +def permute_with_mask_map_and_pad( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + probs: Optional[jnp.ndarray], + pad_offsets: jnp.ndarray, + num_tokens: int, + num_experts: int, + num_out_tokens: int, + hidden_size: int, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """ + Permute the input tensor based on the row_id_map with fused padding. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. + probs : Optional[jnp.ndarray] + The probabilities of the input tensor. If it is not None, it will be permuted. + pad_offsets : jnp.ndarray + Per-expert cumulative padding offsets of shape `[num_experts]`. + num_tokens : int + Number of tokens in the input tensor. + num_experts : int + Number of experts in the input tensor. + num_out_tokens : int + Number of tokens in the permuted tensor (including padding). + hidden_size : int + Hidden size of the input tensor. + + Returns + ------- + output : jnp.ndarray + Permuted and padded output tensor of shape `[num_out_tokens, hidden_size]`. + permuted_probs : Optional[jnp.ndarray] + Permuted probabilities if probs was provided, None otherwise. + """ + with_probs = probs is not None + # Handle None probs by creating dummy tensor if not with_probs: probs = jnp.zeros((0,), dtype=inp.dtype) @@ -971,11 +1419,13 @@ def permute_with_mask_map( probs, dummy_scale, dummy_permuted_scale, + pad_offsets, num_tokens=num_tokens, num_experts=num_experts, num_out_tokens=num_out_tokens, hidden_size=hidden_size, with_probs=with_probs, + with_pad=True, ) if not with_probs: @@ -1029,12 +1479,83 @@ def unpermute_with_mask_map( merging_probs = jnp.zeros((0,), dtype=inp.dtype) if not with_probs: permuted_probs = jnp.zeros((0,), dtype=inp.dtype) + # Create dummy pad_offsets (not used when FUSION_UNPAD=False, but required by kernel signature) + dummy_pad_offsets = jnp.zeros((0,), dtype=jnp.int32) output, unpermuted_probs = UnpermuteWithMaskMapPrimitive.outer_primitive.bind( inp, row_id_map, merging_probs, permuted_probs, + dummy_pad_offsets, + num_tokens=num_tokens, + num_experts=num_experts, + hidden_size=hidden_size, + with_merging_probs=with_merging_probs, + with_probs=with_probs, + ) + + if not with_probs: + unpermuted_probs = None + + return output, unpermuted_probs + + +def unpermute_with_mask_map_and_unpad( + inp: jnp.ndarray, + row_id_map: jnp.ndarray, + merging_probs: Optional[jnp.ndarray], + permuted_probs: Optional[jnp.ndarray], + pad_offsets: jnp.ndarray, + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """ + Unpermute the input tensor based on the row_id_map with fused unpadding. + + Parameters + ---------- + inp : jnp.ndarray + Input tensor of shape `[num_out_tokens, hidden_size]` (including padding). + row_id_map : jnp.ndarray + The token to expert mapping tensor of shape `[num_tokens, num_experts * 2 + 1]`. + merging_probs : Optional[jnp.ndarray] + The merging probabilities of the input tensor. If it is not None, it will be used as weights + to reduce the unpermuted tokens. + permuted_probs : Optional[jnp.ndarray] + The permuted probabilities of the input tensor. If it is not None, it will be unpermuted. + pad_offsets : jnp.ndarray + Per-expert cumulative padding offsets of shape `[num_experts]`. + num_tokens : int + Number of tokens in the unpermuted tensor. + num_experts : int + Number of experts. + hidden_size : int + Hidden size of the tensor. + + Returns + ------- + output : jnp.ndarray + Unpermuted output tensor of shape `[num_tokens, hidden_size]`. + unpermuted_probs : Optional[jnp.ndarray] + Unpermuted probabilities if permuted_probs was provided, None otherwise. + """ + with_merging_probs = merging_probs is not None + with_probs = permuted_probs is not None + + # Handle None inputs by creating dummy tensors + if not with_merging_probs: + merging_probs = jnp.zeros((0,), dtype=inp.dtype) + if not with_probs: + permuted_probs = jnp.zeros((0,), dtype=inp.dtype) + + output, unpermuted_probs = UnpermuteWithMaskMapAndUnpadPrimitive.outer_primitive.bind( + inp, + row_id_map, + merging_probs, + permuted_probs, + pad_offsets, num_tokens=num_tokens, num_experts=num_experts, hidden_size=hidden_size, diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 12d6a9e3de..41ce15303c 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -142,17 +142,31 @@ def compile_triton( ) # Create kernel object for JAX - kernel = gpu_triton.TritonKernel( - compiled.name, - num_warps, - compiled.metadata.shared, - compiled.asm["ptx"], - "", # ttir - compute_capability, - 1, - 1, - 1, # cluster_dims - ) + # From jax/jaxlib/gpu/triton_kernels.cc: + from packaging import version + + if version.parse(jax.__version__) >= version.parse("0.8.2"): + kernel = gpu_triton.TritonKernel( + compiled.name, # arg0: kernel_name (str) + num_warps, # arg1: num_warps (int) + num_ctas, # arg2: num_ctas (int) + compiled.metadata.shared, # arg3: shared_mem_bytes (int) + compiled.asm["ptx"], # arg4: ptx (str) + "", # arg5: ttir (str) - empty + compute_capability, # arg6: compute_capability (int) + ) + else: + kernel = gpu_triton.TritonKernel( + compiled.name, + num_warps, + compiled.metadata.shared, + compiled.asm["ptx"], + "", # ttir + compute_capability, + 1, + 1, + 1, + ) _TRITON_KERNEL_CACHE[cache_key] = kernel return kernel diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 5341af3d74..9f4a9678eb 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -34,6 +34,7 @@ from transformer_engine.pytorch.permutation import ( moe_permute, moe_permute_with_probs, + moe_permute_and_pad_with_probs, moe_unpermute, moe_sort_chunks_by_index, moe_sort_chunks_by_index_with_probs, diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index 9fce9cefcf..d15814585e 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""MoE Permutaion API""" +"""MoE Permutation API""" import warnings from typing import Optional, Tuple import torch @@ -191,6 +191,7 @@ def forward( routing_map: torch.Tensor, num_out_tokens: int, probs: torch.Tensor, + pad_offsets: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: # pylint: disable=missing-function-docstring if not inp.numel(): @@ -201,6 +202,8 @@ def forward( assert routing_map.is_cuda, "TransformerEngine needs CUDA." if probs is not None: assert probs.is_cuda, "TransformerEngine needs CUDA." + if pad_offsets is not None: + assert pad_offsets.is_cuda, "TransformerEngine needs CUDA." assert inp.size(0) == routing_map.size(0), "Permute not possible" num_tokens, hidden_size = inp.size() @@ -250,6 +253,7 @@ def forward( row_id_map, probs, fp8_scale, + pad_offsets, num_tokens, num_experts, num_out_tokens, @@ -292,7 +296,7 @@ def forward( requires_grad=output.requires_grad, ) - ctx.save_for_backward(row_id_map) + ctx.save_for_backward(row_id_map, pad_offsets) ctx.num_experts = num_experts ctx.num_tokens = num_tokens ctx.hidden_size = hidden_size @@ -307,12 +311,12 @@ def backward( ) -> Tuple[torch.Tensor, ...]: # pylint: disable=missing-function-docstring if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, ctx.probs + return permuted_act_grad, None, None, ctx.probs, None act_grad = None probs_grad = None if ctx.needs_input_grad[0]: - (row_id_map,) = ctx.saved_tensors + row_id_map, pad_offsets = ctx.saved_tensors assert not isinstance( permuted_act_grad, QuantizedTensor ), "The backward of moe_permute does not support FP8." @@ -321,13 +325,14 @@ def backward( row_id_map, None, permuted_probs_grad, + pad_offsets, ctx.num_tokens, ctx.num_experts, ctx.hidden_size, ) if not ctx.needs_input_grad[3]: probs_grad = None - return act_grad, None, None, probs_grad + return act_grad, None, None, probs_grad, None class _moe_unpermute_mask_map(torch.autograd.Function): @@ -340,6 +345,7 @@ def forward( row_id_map: torch.Tensor, merging_probs: Optional[torch.Tensor], restore_shape: Optional[torch.Size], + pad_offsets: Optional[torch.Tensor], ) -> torch.Tensor: # pylint: disable=missing-function-docstring if not inp.numel(): @@ -358,6 +364,8 @@ def forward( # Device check assert inp.is_cuda, "TransformerEngine needs CUDA." assert row_id_map.is_cuda, "TransformerEngine needs CUDA." + if pad_offsets is not None: + assert pad_offsets.is_cuda, "TransformerEngine needs CUDA." assert not isinstance( inp, QuantizedTensor @@ -367,15 +375,16 @@ def forward( row_id_map, merging_probs, None, + pad_offsets, num_tokens, num_experts, hidden_size, ) if with_probs: - ctx.save_for_backward(inp, row_id_map, merging_probs) + ctx.save_for_backward(inp, row_id_map, merging_probs, pad_offsets) else: - ctx.save_for_backward(row_id_map) + ctx.save_for_backward(row_id_map, pad_offsets) ctx.num_experts = num_experts ctx.num_tokens = num_tokens ctx.num_permuted_tokens = inp.size(0) @@ -387,15 +396,15 @@ def forward( def backward(ctx, unpermuted_act_grad): # pylint: disable=missing-function-docstring if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.merging_probs, None + return unpermuted_act_grad, None, ctx.merging_probs, None, None act_grad = None probs_grad = None if ctx.needs_input_grad[0]: if ctx.with_probs: - fwd_input, row_id_map, merging_probs = ctx.saved_tensors + fwd_input, row_id_map, merging_probs, pad_offsets = ctx.saved_tensors else: - (row_id_map,) = ctx.saved_tensors + row_id_map, pad_offsets = ctx.saved_tensors fp8 = isinstance(unpermuted_act_grad, QuantizedTensor) per_tensor_recipe = isinstance(unpermuted_act_grad, Float8Tensor) @@ -441,6 +450,7 @@ def backward(ctx, unpermuted_act_grad): row_id_map, fwd_input, merging_probs, + pad_offsets, ctx.num_tokens, ctx.num_experts, ctx.num_permuted_tokens, @@ -453,6 +463,7 @@ def backward(ctx, unpermuted_act_grad): row_id_map, None, fp8_scale, + pad_offsets, ctx.num_tokens, ctx.num_experts, ctx.num_permuted_tokens, @@ -497,7 +508,7 @@ def backward(ctx, unpermuted_act_grad): if not ctx.needs_input_grad[2]: probs_grad = None - return act_grad, None, probs_grad, None + return act_grad, None, probs_grad, None, None def moe_permute( @@ -537,7 +548,9 @@ def moe_permute( if map_type == "index": return _moe_permute_index_map.apply(inp, routing_map, num_out_tokens, max_token_num) if map_type == "mask": - output, row_id_map, _ = _moe_permute_mask_map.apply(inp, routing_map, num_out_tokens, None) + output, row_id_map, _ = _moe_permute_mask_map.apply( + inp, routing_map, num_out_tokens, None, None + ) return output, row_id_map raise ValueError("map_type should be one of 'mask' or 'index'") @@ -570,11 +583,67 @@ def moe_permute_with_probs( By default, set to '-1', meaning no tokens are dropped. """ output, row_id_map, permuted_probs = _moe_permute_mask_map.apply( - inp, routing_map, num_out_tokens, probs + inp, routing_map, num_out_tokens, probs, None ) return output, permuted_probs, row_id_map +def moe_permute_and_pad_with_probs( + inp: torch.Tensor, + probs: torch.Tensor, + routing_map: torch.Tensor, + tokens_per_expert: torch.Tensor, + align_size: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + """ + Permute the tokens and probs based on the routing_map. + Token with the same index will be grouped together. + Tokens with the same designated expert will be grouped together. + The routing_map indicates which experts were selected by each token. + + Parameters + ---------- + inp: torch.Tensor + Input tensor of shape `[num_tokens, hidden_size]`, on which permutation will be applied. + probs: torch.Tensor + The tensor of probabilities corresponding to the permuted tokens and is + of shape [num_tokens, num_experts]. It will be permuted with the tokens + according to the routing_map. + routing_map: torch.Tensor + The token to expert mapping tensor of shape [num_tokens, num_experts] and dtype 'int32'. + The values in it: 1 means the token is routed to this expert and 0 means not. + tokens_per_expert : torch.Tensor + Tensor of shape `[num_experts]` containing actual token counts per expert. + align_size : int + the alignment size for the input tensor. + """ + assert ( + tokens_per_expert is not None + ), "tokens_per_expert must be provided to the fused permute padding function." + assert align_size > 0, f"align_size must be positive, got {align_size}" + + # Ensure tokens_per_expert is on the same device as input to avoid device transfers + if tokens_per_expert.device != inp.device: + tokens_per_expert = tokens_per_expert.to(inp.device) + + # Calculate aligned token counts per expert + target_tokens_per_expert = (torch.ceil(tokens_per_expert / align_size) * align_size).long() + + if torch.equal(tokens_per_expert, target_tokens_per_expert): + pad_offsets = None + else: + pad_lengths = target_tokens_per_expert - tokens_per_expert + cum_pad = torch.cumsum(pad_lengths, dim=0) + pad_offsets = torch.cat( + [torch.zeros(1, dtype=cum_pad.dtype, device=inp.device), cum_pad[:-1]] + ) + + output, row_id_map, permuted_probs = _moe_permute_mask_map.apply( + inp, routing_map, target_tokens_per_expert.sum().item(), probs, pad_offsets + ) + return output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert + + def moe_unpermute( inp: torch.Tensor, row_id_map: torch.Tensor, @@ -582,6 +651,7 @@ def moe_unpermute( restore_shape: Optional[torch.Size] = None, map_type: str = "mask", probs: Optional[torch.Tensor] = None, + pad_offsets: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Unpermute a tensor with permuted tokens, and optionally merge the tokens with their @@ -605,6 +675,10 @@ def moe_unpermute( Options are: 'mask', 'index'. probs : torch.Tensor, default = None Renamed to merging_probs. Keep for backward compatibility. + pad_offsets : torch.Tensor, default = None + Tensor of per-expert cumulative padding offsets used to remove padding added + during permutation. This is the fourth output of `moe_permute_and_pad_with_probs` + and is required when unpermuting padded outputs. """ if probs is not None: if merging_probs is not None: @@ -616,7 +690,9 @@ def moe_unpermute( if map_type == "index": return _moe_unpermute_index_map.apply(inp, row_id_map, merging_probs) if map_type == "mask": - return _moe_unpermute_mask_map.apply(inp, row_id_map, merging_probs, restore_shape) + return _moe_unpermute_mask_map.apply( + inp, row_id_map, merging_probs, restore_shape, pad_offsets + ) raise ValueError("map_type should be one of 'mask' or 'index'") diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 8f953e9c31..27662e1b28 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -123,6 +123,7 @@ def permute_with_mask_map( row_id_map: torch.Tensor, probs: torch.Tensor, scale: torch.Tensor, + pad_offsets: torch.Tensor, num_tokens: int, num_experts: int, num_out_tokens: int, @@ -142,6 +143,9 @@ def permute_with_mask_map( The probabilities of the input tensor. If it is not None, it will be permuted. scale : torch.Tensor The scale of the input tensor. If it is not None, it will be permuted. + pad_offsets : torch.Tensor + Per-expert padding offsets of shape `[num_experts]` for FP8 fused padding. + If it is not None, it will be allocated output buffers with aligned sizes. num_tokens : int Number of tokens in the input tensor. num_experts : int @@ -153,18 +157,18 @@ def permute_with_mask_map( scale_hidden_dim : int Hidden size of the scale tensor. """ - output = torch.empty((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda") - if probs is not None: - permuted_probs = torch.empty((num_out_tokens,), dtype=probs.dtype, device="cuda") - else: - permuted_probs = None - - if scale is not None: - permuted_scale = torch.empty( - (num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda" - ) - else: - permuted_scale = None + # Use torch.zeros when pad_offsets is provided to ensure padding regions are zeroed, + # since the kernel doesn't write to padding positions. + alloc = torch.zeros if pad_offsets is not None else torch.empty + output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda") + permuted_probs = ( + alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None + ) + permuted_scale = ( + torch.empty((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda") + if scale is not None + else None + ) # pylint: disable=unnecessary-lambda-assignment grid = lambda META: (num_tokens, triton.cdiv(hidden_size, META["BLOCK_SIZE"])) _permute_kernel[grid]( @@ -173,6 +177,7 @@ def permute_with_mask_map( probs, scale, permuted_scale, + pad_offsets, scale_hidden_dim, row_id_map.stride(0), row_id_map.stride(1), @@ -193,6 +198,7 @@ def permute_with_mask_map( hidden_size, PERMUTE_PROBS=probs is not None, PERMUTE_SCALE=scale is not None, + FUSION_PAD=pad_offsets is not None, ) return output, permuted_scale, permuted_probs @@ -202,6 +208,7 @@ def unpermute_with_mask_map( row_id_map: torch.Tensor, merging_probs: Union[torch.Tensor, None], permuted_probs: Union[torch.Tensor, None], + pad_offsets: Union[torch.Tensor, None], num_tokens: int, num_experts: int, hidden_size: int, @@ -220,6 +227,9 @@ def unpermute_with_mask_map( to reduce the unpermuted tokens. permuted_probs : torch.Tensor The permuted probabilities of the input tensor. If it is not None, it will be unpermuted. + pad_offsets : torch.Tensor + Per-expert padding offsets of shape `[num_experts]` for FP8 fused unpadding. + If it is not None, it will remove the previously fused padding. num_tokens : int Number of tokens in the permuted tensor. num_experts : int @@ -241,6 +251,7 @@ def unpermute_with_mask_map( row_id_map, merging_probs, permuted_probs, + pad_offsets, row_id_map.stride(0), row_id_map.stride(1), inp.stride(0), @@ -259,6 +270,7 @@ def unpermute_with_mask_map( PROBS_LOAD_WIDTH=triton.next_power_of_2(num_experts), WITH_MERGING_PROBS=merging_probs is not None, PERMUTE_PROBS=permuted_probs is not None, + FUSION_UNPAD=pad_offsets is not None, ) return output, unpermuted_probs @@ -268,6 +280,7 @@ def unpermute_with_mask_map_bwd_with_merging_probs( row_id_map: torch.Tensor, fwd_input: torch.Tensor, merging_probs: torch.Tensor, + pad_offsets: Union[torch.Tensor, None], num_tokens: int, num_experts: int, num_out_tokens: int, @@ -286,6 +299,9 @@ def unpermute_with_mask_map_bwd_with_merging_probs( The input tensor of the forward pass of shape `[num_out_tokens, hidden_size]`. merging_probs : torch.Tensor The merging probabilities of the input tensor of shape `[num_tokens, num_experts]`. + pad_offsets : torch.Tensor + Per-expert padding offsets of shape `[num_experts]` for FP8 fused padding. + If it is not None, it will be allocated output buffers with aligned sizes. num_tokens : int Number of tokens in the permuted tensor. num_experts : int @@ -295,9 +311,11 @@ def unpermute_with_mask_map_bwd_with_merging_probs( hidden_size : int Hidden size of the output tensor. """ - act_grad = torch.empty( - (num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device="cuda" - ) + # Use zeros when pad_offsets is used because padding slots won't be written to + # by the kernel. This matches the behavior of Fp8Unpadding.backward which zeros + # out the padding slots. + alloc = torch.zeros if pad_offsets is not None else torch.empty + act_grad = alloc((num_out_tokens, hidden_size), dtype=fwd_output_grad.dtype, device="cuda") merging_probs_grad = torch.empty( (num_tokens, num_experts), dtype=merging_probs.dtype, device="cuda" ) @@ -307,6 +325,7 @@ def unpermute_with_mask_map_bwd_with_merging_probs( fwd_input, merging_probs, row_id_map, + pad_offsets, row_id_map.stride(0), row_id_map.stride(1), fwd_output_grad.stride(0), @@ -324,6 +343,7 @@ def unpermute_with_mask_map_bwd_with_merging_probs( num_experts, hidden_size, PROBS_LOAD_WIDTH=triton.next_power_of_2(num_experts), + FUSION_UNPAD=pad_offsets is not None, ) return act_grad, merging_probs_grad From 26c82db6cd04a5d33c650404b837987f81de79d2 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Wed, 31 Dec 2025 01:06:12 -0800 Subject: [PATCH 142/521] [JAX] Fix incorrect calculation of segment pos from segment ids in user-facing API (#2523) * Fix incorrect calculation of segment pos from segment ids for thd cases and load balanced cases in from_segment_ids_and_pos. Enforce passing of segment_pos for THD cases and lod balanced cases Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Correct the assert condition Signed-off-by: Kshitij Lakhani * Modify fused attn tests to pass new args to from_segment_ids_and_pos() Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Calculate seg ids before pos Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 1. Change the signature for from_segment_ids_and_pos() 2. Add support for THD in from_segment_ids_and_pos() 3. Assert if load balanced segment_ids is passed to generate a segment_pos Signed-off-by: Kshitij Janardan Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pass keyword-only args by name Signed-off-by: Kshitij Janardan Lakhani * nit: Fix typo to use seg_ids instead of segment_ids Signed-off-by: Kshitij Janardan Lakhani * nit: Fix comments Signed-off-by: Kshitij Janardan Lakhani * Modify the function call to differentiate between load balancing and actually reordered segment_ids and segment_pos 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 the is_segment_ids_reordered to be set only when CP and load balancing Signed-off-by: Kshitij Lakhani * Fix comments for from_segment_ids_and_pos() Signed-off-by: Kshitij Lakhani * Code clean up for more information, see https://pre-commit.ci Fix lint errors 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 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: Kirthi Shankar Sivamani --- tests/jax/test_fused_attn.py | 16 +++++- transformer_engine/jax/attention.py | 85 ++++++++++++++++++++++++++--- 2 files changed, 90 insertions(+), 11 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 49372fda1d..f7267af5b8 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -668,14 +668,24 @@ def generate_random_segment_ids( (self.offsets_q, self.offsets_kv), ) case SeqDescFormat.SegmentIDs: + # Exercise the path to generate the segment_pos in from_segment_ids_and_pos() + # if no CP and load balancing, else explicitly pass the segment_pos self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos( ( self.cp_reorder_fn(self.segment_ids_q), self.cp_reorder_fn(self.segment_ids_kv), ), ( - self.cp_reorder_fn(self.segment_pos_q), - self.cp_reorder_fn(self.segment_pos_kv), + ( + self.cp_reorder_fn(self.segment_pos_q), + self.cp_reorder_fn(self.segment_pos_kv), + ) + if self.cp_size > 1 and self.cp_load_balanced + else None + ), + is_thd=self.qkv_layout.is_thd(), + is_segment_ids_reordered=( + True if self.cp_size > 1 and self.cp_load_balanced else False ), ) case _: @@ -704,6 +714,8 @@ def generate_random_segment_ids( self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos( (self.segment_ids_q, self.segment_ids_kv), None, + is_thd=self.qkv_layout.is_thd(), + is_segment_ids_reordered=False, ) case _: raise ValueError(f"Unknown {self.seq_desc_format=}") diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 21680dc805..09a29f4cb8 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -658,7 +658,7 @@ class SequenceDescriptor: - SequenceDescriptor.from_seqlens_and_offsets For THD (packed) cases, where each batch may have not only 1 sequence. - SequenceDescriptor.from_segment_ids_and_pos - Experimental feature for THD (packed) cases with context parallelism. + Experimental feature for BSHD (with and without reordering) and THD (packed) cases without reordering """ seqlens: Optional[Tuple[jnp.ndarray, jnp.ndarray]] @@ -796,9 +796,14 @@ def from_segment_ids_and_pos( cls, segment_ids: Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]], segment_pos: Optional[Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]] = None, + *, + is_thd: bool, + is_segment_ids_reordered: bool, ) -> SequenceDescriptor: """ - Experimental factory method for inputs with segment IDs and optional positions. (THD) + Experimental factory method for inputs with segment IDs and optional positions. + segment_pos = None to be used only for: BSHD with or without load balancing and, + THD without load balancing Args: segment_ids(Tuple(jnp.ndarray, jnp.ndarray)) = (q_segment_ids, kv_segment_ids): - q_segment_ids (jnp.ndarray): @@ -812,22 +817,84 @@ def from_segment_ids_and_pos( The position inside each segment for query, with shape [batch, max_seqlen]. - kv_segment_pos (jnp.ndarray): The position inside each segment for key, value, with shape [batch, max_seqlen]. + is_thd(bool): If True, QKVLayout is of type THD, else it is BSHD + is_segment_ids_reordered(bool): If True, the segment ids have been reordered for load balancing. + Only THD with load balancing is expected to have this flag set to True Return: A SequenceDescriptor with segment_ids/segment_pos initialized. """ q_seg_ids, kv_seg_ids = cls._expand_to_pair(segment_ids) - if segment_pos is not None: - segment_pos = cls._expand_to_pair(segment_pos) - else: - - def generate_default_pos(segment_ids): - seqlen = segment_ids.shape[-1] - return jnp.broadcast_to(jnp.arange(seqlen), segment_ids.shape) + # Using defaults : segment pos has to be generated. + if segment_pos is None: + # THD + load balanced segment_ids are not supported in this function + # BSHD + load balanced segment_ids are incorrect as BSHD handles reordering within the primitive itself + if is_segment_ids_reordered: + assert not is_thd, ( + f"{segment_pos=} default arg is not supported for load balanced reordered" + " (Striped) THD inputs. Please pass the load balanced reordered segment_pos" + " and segment_ids explicitly to {from_segment_ids_and_pos.__qualname__}" + " using convenience function reorder_causal_load_balancing()" + ) + assert is_thd, ( + f"{segment_pos=} default arg is not supported for load balanced reordered (Dual" + " Chunk) BSHD inputs. BSHD segment_pos and segment_ids do not need to be load" + " balanced reordered. The reordering for these is performed within the" + " primitive" + ) + + # Generate the default pos for THD and BSHD non-reordered segment_ids + def generate_default_pos(seg_ids): + if is_thd: + batch_size, seq_size = seg_ids.shape + # Assume that the first token belongs to a segment and is not a padded token + first_is_segment = jnp.full((batch_size, 1), True, dtype=bool) + # Get segment start positions + segment_start = jnp.concatenate( + [ + first_is_segment, + (seg_ids[..., 1:] != seg_ids[..., :-1]) & (seg_ids[..., 1:] != 0), + ], + axis=-1, + ) + # Get offset for location where new segment starts + segment_start_idx = jax.vmap(lambda row: jnp.arange(row.size) * row)( + segment_start + ) + segment_start_offsets = jax.vmap(jnp.maximum.accumulate)(segment_start_idx) + + # Get the last non-zero index - after this everything is padding + # (B,) + last_nonzero_idx = jax.vmap( + lambda segids_row: jnp.max( + jnp.where(segids_row != 0, jnp.arange(seq_size), -1) + ) + )(seg_ids) + seg_pos_no_thd = jnp.arange(seq_size) + # Get a mask which can be used to zero out all the padding at the end (after the non-zero index) + mask = seg_pos_no_thd <= last_nonzero_idx[:, None] + + # Get the unmasked seg_pos for the THD sequence + seg_pos = ( + jnp.broadcast_to(jnp.arange(seq_size), seg_ids.shape) + - segment_start_offsets + ) + + # Use the mask to zero out the padding at the end (after the non-zero index) + segment_pos = jax.vmap( + lambda pos_row, mask_row: jnp.where(mask_row, pos_row, 0) + )(seg_pos, mask) + return segment_pos + + seqlen = seg_ids.shape[-1] + return jnp.broadcast_to(jnp.arange(seqlen), seg_ids.shape) q_seg_pos = generate_default_pos(q_seg_ids) kv_seg_pos = generate_default_pos(kv_seg_ids) segment_pos = (q_seg_pos, kv_seg_pos) + # Explicitly passed segment_pos + else: + segment_pos = cls._expand_to_pair(segment_pos) return cls( segment_ids=(q_seg_ids, kv_seg_ids), From 697b52cbde6a2d2b67f71879b63203d5082e1f8c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E4=BF=8A?= Date: Wed, 31 Dec 2025 19:51:24 +0800 Subject: [PATCH 143/521] Fix overflow of padding/unpadding kernel (#2548) Signed-off-by: fuyue.lj Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/common/util/padding.cu | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/transformer_engine/common/util/padding.cu b/transformer_engine/common/util/padding.cu index 0d92b243a7..4e569e1674 100644 --- a/transformer_engine/common/util/padding.cu +++ b/transformer_engine/common/util/padding.cu @@ -94,6 +94,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP #pragma unroll for (int i2 = 0; i2 < nvec; ++i2) { const int row = tile_row + i1 * nvec + i2; + size_t row_offset = static_cast(row) * row_length; const int col = tile_col + j1 * nvec; Vec local_input; Vec local_output; @@ -101,7 +102,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - local_input.data.elt[j2] = input[row * row_length + col + j2]; + local_input.data.elt[j2] = input[row_offset + col + j2]; } } } @@ -112,14 +113,14 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - output[row * row_length + col + j2] = local_output.data.elt[j2]; + output[row_offset + col + j2] = local_output.data.elt[j2]; } } } else if (row < padded_num_rows) { // padding for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - output[row * row_length + col + j2] = local_zero; + output[row_offset + col + j2] = local_zero; } } } @@ -178,6 +179,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult #pragma unroll for (int i2 = 0; i2 < nvec; ++i2) { const int row = tile_row + i1 * nvec + i2; + size_t row_offset = static_cast(row) * row_length; const int col = tile_col + j1 * nvec; Vec local_input; Vec local_output; @@ -185,7 +187,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - local_input.data.elt[j2] = input[row * row_length + col + j2]; + local_input.data.elt[j2] = input[row_offset + col + j2]; } } } @@ -196,7 +198,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - output[row * row_length + col + j2] = local_output.data.elt[j2]; + output[row_offset + col + j2] = local_output.data.elt[j2]; } } } From 324be3324278723bd8f66196ed1ccac29b94bd7f Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Thu, 1 Jan 2026 00:18:01 +0800 Subject: [PATCH 144/521] [PyTorch] Support cudagraph recomputation (#2518) * replace autograd.grad with autograd.backward Signed-off-by: Robin Zhang * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * get/set graphable rng state Signed-off-by: Robin Zhang * fix lint Signed-off-by: Robin Zhang --------- Signed-off-by: Robin Zhang Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/distributed.py | 41 +++++++++---- transformer_engine/pytorch/graph.py | 71 ++++++++++++++++------- 2 files changed, 81 insertions(+), 31 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index deb9b3ff91..9f589498a4 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -90,6 +90,11 @@ def graph_safe_rng_available() -> bool: ) +def is_graph_safe_rng_state(state: Union[torch.Tensor, torch.Generator]) -> bool: + """Returns whether the rng state is a graph safe version.""" + return graph_safe_rng_available() and isinstance(state, torch.Generator) + + def _get_cuda_rng_state( device: Union[int, str, torch.device] = "cuda", clone: bool = False, @@ -340,9 +345,16 @@ def forward( # Copy the rng states. ctx.fwd_cpu_rng_state = torch.get_rng_state() - ctx.fwd_cuda_rng_state = _get_cuda_rng_state(graph_safe=False) if get_rng_state_tracker is not None: ctx.fwd_cuda_rng_state_tracker = get_rng_state_tracker().get_states() + ctx.graph_safe_rng_state = ( + is_graph_safe_rng_state(next(iter(ctx.fwd_cuda_rng_state_tracker.values()))) + if ctx.fwd_cuda_rng_state_tracker + else False + ) + else: + ctx.graph_safe_rng_state = False + ctx.fwd_cuda_rng_state = _get_cuda_rng_state(graph_safe=ctx.graph_safe_rng_state) if context_fn is not None: forward_ctx, recompute_ctx = context_fn() @@ -406,13 +418,13 @@ def backward( # Store the current states. bwd_cpu_rng_state = torch.get_rng_state() - bwd_cuda_rng_state = _get_cuda_rng_state(graph_safe=False) + bwd_cuda_rng_state = _get_cuda_rng_state(graph_safe=ctx.graph_safe_rng_state) if get_rng_state_tracker is not None: bwd_cuda_rng_state_tracker = get_rng_state_tracker().get_states() # Set the states to what it used to be before the forward pass. torch.set_rng_state(ctx.fwd_cpu_rng_state) - _set_cuda_rng_state(ctx.fwd_cuda_rng_state, graph_safe=False) + _set_cuda_rng_state(ctx.fwd_cuda_rng_state, graph_safe=ctx.graph_safe_rng_state) if get_rng_state_tracker is not None: get_rng_state_tracker().set_states(ctx.fwd_cuda_rng_state_tracker) @@ -427,7 +439,7 @@ def backward( # Set the states back to what it was at the start of this function. torch.set_rng_state(bwd_cpu_rng_state) - _set_cuda_rng_state(bwd_cuda_rng_state, graph_safe=False) + _set_cuda_rng_state(bwd_cuda_rng_state, graph_safe=ctx.graph_safe_rng_state) if get_rng_state_tracker is not None: get_rng_state_tracker().set_states(bwd_cuda_rng_state_tracker) @@ -470,12 +482,21 @@ def __init__(self, recompute_fn: Callable, get_rng_state_tracker: Callable): def cache_rng_states(self, forward=True): """Cache fwd/bwd RNG states in the frame to restore later.""" - rng_states = ( - torch.get_rng_state(), - _get_cuda_rng_state(graph_safe=False), - ) + rng_states = (torch.get_rng_state(),) if self.get_rng_state_tracker is not None: - rng_states += (self.get_rng_state_tracker().get_states(),) + tracker_states = self.get_rng_state_tracker().get_states() + self.graph_safe_rng_state = ( + is_graph_safe_rng_state(next(iter(tracker_states.values()))) + if tracker_states + else False + ) + rng_states += ( + _get_cuda_rng_state(graph_safe=self.graph_safe_rng_state), + tracker_states, + ) + else: + self.graph_safe_rng_state = False + rng_states += (_get_cuda_rng_state(graph_safe=self.graph_safe_rng_state),) if forward: self.fwd_rng_states = rng_states @@ -490,7 +511,7 @@ def restore_rng_states(self, forward=True): rng_states = self.bwd_rng_states torch.set_rng_state(rng_states[0]) - _set_cuda_rng_state(rng_states[1], graph_safe=False) + _set_cuda_rng_state(rng_states[1], graph_safe=self.graph_safe_rng_state) if self.get_rng_state_tracker is not None: self.get_rng_state_tracker().set_states(rng_states[2]) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 92826735f9..1822c47d8b 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -62,6 +62,21 @@ def graph_pool_handle(): return _graph_pool_handle() +@contextlib.contextmanager +def _none_grad_context_wrapper(inputs): + """ + Wrapper to set the gradients of the inputs to None, + in case the backward pass makes grad accumulations. + """ + original_input_grads = [] + for input_tensor in inputs: + original_input_grads.append(input_tensor.grad) + input_tensor.grad = None + yield + for input_tensor, original_grad in zip(inputs, original_input_grads): + input_tensor.grad = original_grad + + @contextlib.contextmanager def _graph_context_wrapper(*args, **kwargs): """Wrapper around `torch.cuda.graph`. @@ -434,13 +449,15 @@ def hook_fn( for hook in hooks: hook.remove() if is_training: - grad_inputs = torch.autograd.grad( - outputs=tuple(o for o in outputs if o.requires_grad), - inputs=tuple(i for i in static_input_surface if i.requires_grad), - grad_outputs=tuple(torch.empty_like(o) for o in outputs if o.requires_grad), - only_inputs=True, - allow_unused=allow_unused_input, - ) + inputs = tuple(i for i in static_input_surface if i.requires_grad) + with _none_grad_context_wrapper(inputs): + torch.autograd.backward( + tuple(o for o in outputs if o.requires_grad), + grad_tensors=tuple( + torch.empty_like(o) for o in outputs if o.requires_grad + ), + ) + grad_inputs = tuple(input.grad for input in inputs) # Filter module params that get None grad from grad_inputs and remove them # from static_input_surface. This is to ensure that the backward hooks @@ -455,6 +472,14 @@ def hook_fn( module_params_with_grad = [] for grad_inputs_idx, inputs_idx in enumerate(required_grad_input_idx): if ( + grad_inputs[grad_inputs_idx] is None + and grad_inputs_idx < num_required_grad_sample_args + ): + assert allow_unused_input, ( + "The input tensor requires grad, but the grad is None after" + " backward pass." + ) + elif ( grad_inputs[grad_inputs_idx] is not None and grad_inputs_idx >= num_required_grad_sample_args ): @@ -606,15 +631,17 @@ def hook_fn( torch.empty_like(o) if o.requires_grad else None for o in static_outputs ) if is_training: - with _graph_context_wrapper(bwd_graph, pool=mempool): - grad_inputs = torch.autograd.grad( - outputs=tuple(o for o in static_outputs if o.requires_grad), - inputs=tuple(i for i in static_input_surface if i.requires_grad), - grad_outputs=tuple(o for o in static_grad_outputs if o is not None), - only_inputs=True, - allow_unused=allow_unused_input, + inputs = tuple(i for i in static_input_surface if i.requires_grad) + with _none_grad_context_wrapper(inputs), _graph_context_wrapper( + bwd_graph, pool=mempool + ): + torch.autograd.backward( + tuple(o for o in static_outputs if o.requires_grad), + grad_tensors=tuple(o for o in static_grad_outputs if o is not None), retain_graph=retain_graph_in_backward, ) + grad_inputs = tuple(input.grad for input in inputs) + # Constructs a tuple suitable for returning from Graphed.backward: # Pads out the actually-needed grads with Nones in gradient slots for inputs # that don't require grad. I couldn't think of a one-liner for this pattern. @@ -695,15 +722,17 @@ def hook_fn( torch.empty_like(o) if o.requires_grad else None for o in static_outputs ) if is_training: - with _graph_context_wrapper(bwd_graph, pool=mempool): - grad_inputs = torch.autograd.grad( - outputs=tuple(o for o in static_outputs if o.requires_grad), - inputs=tuple(i for i in static_input_surface if i.requires_grad), - grad_outputs=tuple(o for o in static_grad_outputs if o is not None), - only_inputs=True, - allow_unused=allow_unused_input, + inputs = tuple(i for i in static_input_surface if i.requires_grad) + with _none_grad_context_wrapper(inputs), _graph_context_wrapper( + bwd_graph, pool=mempool + ): + torch.autograd.backward( + tuple(o for o in static_outputs if o.requires_grad), + grad_tensors=tuple(o for o in static_grad_outputs if o is not None), retain_graph=retain_graph_in_backward, ) + grad_inputs = tuple(input.grad for input in inputs) + if need_bwd_dw_graph[bwd_idx]: with _graph_context_wrapper(bwd_dw_graph, pool=mempool): for module in visited_te_modules[bwd_idx]: From 830ef60fd89508ee7c372bcc1489c0beb86094cb Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 2 Jan 2026 11:03:26 +0530 Subject: [PATCH 145/521] Update copyright to include year 2026 (#2553) Update copyright to include 2026 Signed-off-by: Kirthi Shankar Sivamani --- .github/actions/build-pytorch-wheel/Dockerfile | 2 +- .github/actions/build-pytorch-wheel/action.yml | 2 +- .github/actions/build-pytorch-wheel/build.sh | 2 +- .github/scripts/check_for_ngc_images.sh | 2 +- .github/workflows/attach-wheels-to-release.yml | 2 +- .github/workflows/blossom-ci.yml | 2 +- .github/workflows/build.yml | 2 +- .github/workflows/deploy_nightly_docs.yml | 2 +- .github/workflows/docs.yml | 2 +- .github/workflows/license.yml | 2 +- .github/workflows/lint.yml | 2 +- .github/workflows/trigger-ci.yml | 2 +- .github/workflows/upload-ci-logs.yml | 2 +- CONTRIBUTING.rst | 2 +- CPPLINT.cfg | 2 +- README.rst | 2 +- benchmarks/attention/benchmark_attention.py | 2 +- benchmarks/benchmark_rht_cast.py | 2 +- benchmarks/linear/benchmark_grouped_linear.py | 2 +- build_tools/__init__.py | 2 +- build_tools/build_ext.py | 2 +- build_tools/jax.py | 2 +- build_tools/pytorch.py | 2 +- build_tools/te_version.py | 2 +- build_tools/utils.py | 2 +- build_tools/wheel_utils/Dockerfile.aarch | 2 +- build_tools/wheel_utils/Dockerfile.x86 | 2 +- build_tools/wheel_utils/build_wheels.sh | 2 +- build_tools/wheel_utils/launch_aarch.sh | 2 +- build_tools/wheel_utils/launch_x86.sh | 2 +- docs/api/c/activation.rst | 2 +- docs/api/c/cast.rst | 2 +- docs/api/c/cast_transpose_noop.rst | 2 +- docs/api/c/cudnn.rst | 2 +- docs/api/c/fused_attn.rst | 2 +- docs/api/c/fused_rope.rst | 2 +- docs/api/c/gemm.rst | 2 +- docs/api/c/index.rst | 2 +- docs/api/c/multi_tensor.rst | 2 +- docs/api/c/normalization.rst | 2 +- docs/api/c/padding.rst | 2 +- docs/api/c/permutation.rst | 2 +- docs/api/c/recipe.rst | 2 +- docs/api/c/softmax.rst | 2 +- docs/api/c/swizzle.rst | 2 +- docs/api/c/transformer_engine.rst | 2 +- docs/api/c/transpose.rst | 2 +- docs/api/common.rst | 2 +- docs/api/framework.rst | 2 +- docs/api/jax.rst | 2 +- docs/api/pytorch.rst | 2 +- docs/conf.py | 2 +- docs/debug.rst | 2 +- docs/debug/1_getting_started.rst | 2 +- docs/debug/2_config_file_structure.rst | 2 +- docs/debug/3_api_debug_setup.rst | 2 +- docs/debug/3_api_features.rst | 2 +- docs/debug/3_api_te_calls.rst | 2 +- docs/debug/4_distributed.rst | 2 +- docs/debug/api.rst | 2 +- docs/examples/attention/arbitrary_mask_to_post_scale_bias.py | 2 +- docs/examples/attention/example_attention.py | 2 +- docs/examples/onnx/utils.py | 2 +- docs/examples/quickstart_jax_utils.py | 2 +- docs/examples/quickstart_utils.py | 2 +- docs/examples/te_gemma/te_gemma.py | 2 +- docs/examples/te_gemma/te_gemma_loading_weights.py | 2 +- docs/examples/te_gemma/utils.py | 2 +- docs/examples/te_llama/te_llama.py | 2 +- docs/examples/te_llama/utils.py | 2 +- docs/faq.rst | 2 +- docs/getting_started.rst | 2 +- docs/index.rst | 2 +- docs/installation.rst | 2 +- examples/jax/collective_gemm/common.py | 2 +- examples/jax/collective_gemm/conftest.py | 2 +- examples/jax/collective_gemm/run_test_cgemm.sh | 2 +- examples/jax/collective_gemm/test_dense_grad.py | 2 +- examples/jax/collective_gemm/test_gemm.py | 2 +- examples/jax/collective_gemm/test_layernorm_mlp_grad.py | 2 +- examples/jax/encoder/common.py | 2 +- examples/jax/encoder/conftest.py | 2 +- examples/jax/encoder/run_test_multiprocessing_encoder.sh | 2 +- examples/jax/encoder/test_model_parallel_encoder.py | 2 +- examples/jax/encoder/test_multigpu_encoder.py | 2 +- examples/jax/encoder/test_multiprocessing_encoder.py | 2 +- examples/jax/encoder/test_single_gpu_encoder.py | 2 +- examples/jax/mnist/test_single_gpu_mnist.py | 2 +- examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py | 2 +- examples/pytorch/fsdp/README.md | 2 +- examples/pytorch/fsdp/fsdp.py | 2 +- examples/pytorch/mnist/main.py | 2 +- pyproject.toml | 2 +- qa/L0_cppunittest/test.sh | 2 +- qa/L0_jax_distributed_unittest/test.sh | 2 +- qa/L0_jax_lint/test.sh | 2 +- qa/L0_jax_unittest/test.sh | 2 +- qa/L0_jax_wheel/test.sh | 2 +- qa/L0_license/copyright_checker.py | 2 +- qa/L0_license/test.sh | 2 +- qa/L0_pytorch_debug_unittest/test.sh | 2 +- qa/L0_pytorch_lint/test.sh | 2 +- qa/L0_pytorch_unittest/test.sh | 2 +- qa/L0_pytorch_wheel/test.sh | 2 +- qa/L1_cpp_distributed/test.sh | 2 +- qa/L1_jax_distributed_unittest/test.sh | 2 +- qa/L1_pytorch_distributed_unittest/test.sh | 2 +- qa/L1_pytorch_mcore_integration/test.sh | 2 +- qa/L1_pytorch_onnx_unittest/test.sh | 2 +- qa/L1_pytorch_thunder_integration/test.sh | 2 +- qa/L2_jax_distributed_unittest/test.sh | 2 +- qa/L2_jax_unittest/test.sh | 2 +- qa/L3_pytorch_FA_versions_test/test.sh | 2 +- qa/format.sh | 2 +- setup.py | 2 +- tests/cpp/CMakeLists.txt | 2 +- tests/cpp/operator/CMakeLists.txt | 2 +- tests/cpp/operator/test_act.cu | 2 +- tests/cpp/operator/test_cast.cu | 2 +- tests/cpp/operator/test_cast_current_scaling.cu | 2 +- tests/cpp/operator/test_cast_dbias.cu | 2 +- tests/cpp/operator/test_cast_dbias_dgelu.cu | 2 +- tests/cpp/operator/test_cast_float8blockwise.cu | 2 +- tests/cpp/operator/test_cast_gated_swiglu.cu | 2 +- tests/cpp/operator/test_cast_mxfp8.cu | 2 +- tests/cpp/operator/test_cast_mxfp8_gated_swiglu.cu | 2 +- tests/cpp/operator/test_cast_nvfp4_transpose.cu | 2 +- tests/cpp/operator/test_cast_transpose.cu | 2 +- tests/cpp/operator/test_cast_transpose_current_scaling.cu | 2 +- tests/cpp/operator/test_cast_transpose_dbias.cu | 2 +- tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu | 2 +- tests/cpp/operator/test_cast_transpose_dgeglu.cu | 2 +- tests/cpp/operator/test_causal_softmax.cu | 2 +- tests/cpp/operator/test_dequantize_mxfp8.cu | 2 +- tests/cpp/operator/test_memset.cu | 2 +- tests/cpp/operator/test_multi_cast_transpose.cu | 2 +- tests/cpp/operator/test_multi_padding.cu | 2 +- tests/cpp/operator/test_multi_unpadding.cu | 2 +- tests/cpp/operator/test_normalization.cu | 2 +- tests/cpp/operator/test_normalization.h | 2 +- tests/cpp/operator/test_normalization_mxfp8.cu | 2 +- tests/cpp/operator/test_qdq.cu | 2 +- tests/cpp/operator/test_swap_first_dims.cu | 2 +- tests/cpp/operator/test_swizzle.cu | 2 +- tests/cpp/operator/test_transpose.cu | 2 +- tests/cpp/test_common.cu | 2 +- tests/cpp/test_common.h | 2 +- tests/cpp/util/CMakeLists.txt | 2 +- tests/cpp/util/test_nvrtc.cpp | 2 +- tests/cpp/util/test_string.cpp | 2 +- tests/cpp_distributed/CMakeLists.txt | 2 +- tests/cpp_distributed/test_comm_gemm.cu | 2 +- tests/jax/conftest.py | 2 +- tests/jax/distributed_test_base.py | 2 +- tests/jax/multi_process_launch.sh | 2 +- tests/jax/pytest.ini | 2 +- tests/jax/test_custom_call_compute.py | 2 +- tests/jax/test_distributed_dense.py | 2 +- tests/jax/test_distributed_fused_attn.py | 2 +- tests/jax/test_distributed_helper.py | 2 +- tests/jax/test_distributed_layernorm.py | 2 +- tests/jax/test_distributed_layernorm_mlp.py | 2 +- tests/jax/test_distributed_softmax.py | 2 +- tests/jax/test_functions.py | 2 +- tests/jax/test_fused_attn.py | 2 +- tests/jax/test_layer.py | 2 +- tests/jax/test_misc.py | 2 +- tests/jax/test_multi_process_distributed_grouped_gemm.py | 2 +- tests/jax/test_permutation.py | 2 +- tests/jax/test_recipe_characteristics.py | 2 +- tests/jax/test_sanity_import.py | 2 +- tests/jax/test_softmax.py | 2 +- tests/jax/test_triton_custom_calls.py | 2 +- tests/jax/utils.py | 2 +- tests/pytorch/attention/run_attention_with_cp.py | 2 +- tests/pytorch/attention/test_attention.py | 2 +- tests/pytorch/attention/test_attention_with_cp.py | 2 +- tests/pytorch/attention/test_cp_utils.py | 2 +- tests/pytorch/attention/test_kv_cache.py | 2 +- tests/pytorch/debug/conftest.py | 2 +- tests/pytorch/debug/run_distributed.py | 2 +- tests/pytorch/debug/test_api_features.py | 2 +- tests/pytorch/debug/test_config.py | 2 +- tests/pytorch/debug/test_distributed.py | 2 +- tests/pytorch/debug/test_log.py | 2 +- tests/pytorch/debug/test_numerics.py | 2 +- tests/pytorch/debug/test_perf.py | 2 +- tests/pytorch/debug/test_sanity.py | 2 +- tests/pytorch/debug/utils.py | 2 +- tests/pytorch/distributed/run_fsdp2_model.py | 2 +- tests/pytorch/distributed/run_gemm_with_overlap.py | 2 +- tests/pytorch/distributed/run_layer_with_overlap.py | 2 +- tests/pytorch/distributed/run_numerics.py | 2 +- tests/pytorch/distributed/run_numerics_exact.py | 2 +- tests/pytorch/distributed/test_cast_master_weights_to_fp8.py | 2 +- tests/pytorch/distributed/test_comm_gemm_overlap.py | 2 +- tests/pytorch/distributed/test_fusible_ops.py | 2 +- tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py | 2 +- tests/pytorch/distributed/test_numerics.py | 2 +- tests/pytorch/distributed/test_numerics_exact.py | 2 +- tests/pytorch/distributed/test_sanity.py | 2 +- tests/pytorch/distributed/test_torch_fsdp2.py | 2 +- .../layernorm_mlp/test_selective_activation_checkpoint.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_group_quantize.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_module_exact.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py | 2 +- tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py | 2 +- tests/pytorch/references/blockwise_fp8_gemm_reference.py | 2 +- tests/pytorch/references/blockwise_quantizer_reference.py | 2 +- tests/pytorch/references/quantize_scale_calc.py | 2 +- tests/pytorch/references/ref_per_tensor_cs.py | 2 +- tests/pytorch/test_checkpoint.py | 2 +- tests/pytorch/test_cpu_offloading.py | 2 +- tests/pytorch/test_cpu_offloading_v1.py | 2 +- tests/pytorch/test_cuda_graphs.py | 2 +- tests/pytorch/test_custom_recipe.py | 2 +- tests/pytorch/test_deferred_init.py | 2 +- tests/pytorch/test_float8_blockwise_gemm_exact.py | 2 +- tests/pytorch/test_float8_blockwise_scaling_exact.py | 2 +- tests/pytorch/test_float8_current_scaling_exact.py | 2 +- tests/pytorch/test_float8blockwisetensor.py | 2 +- tests/pytorch/test_fused_optimizer.py | 2 +- tests/pytorch/test_fused_rope.py | 2 +- tests/pytorch/test_fused_router.py | 2 +- tests/pytorch/test_fusible_ops.py | 2 +- tests/pytorch/test_gqa.py | 2 +- tests/pytorch/test_hf_integration.py | 2 +- tests/pytorch/test_jit.py | 2 +- tests/pytorch/test_multi_tensor.py | 2 +- tests/pytorch/test_numerics.py | 2 +- tests/pytorch/test_onnx_export.py | 2 +- tests/pytorch/test_parallel_cross_entropy.py | 2 +- tests/pytorch/test_partial_cast.py | 2 +- tests/pytorch/test_permutation.py | 2 +- tests/pytorch/test_qk_norm.py | 2 +- tests/pytorch/test_quantized_tensor.py | 2 +- tests/pytorch/test_recipe.py | 2 +- tests/pytorch/test_sanity.py | 2 +- tests/pytorch/test_sanity_import.py | 2 +- tests/pytorch/utils.py | 2 +- transformer_engine/__init__.py | 2 +- transformer_engine/common/CMakeLists.txt | 2 +- transformer_engine/common/__init__.py | 2 +- transformer_engine/common/activation/activation_template.h | 2 +- transformer_engine/common/activation/gelu.cu | 2 +- transformer_engine/common/activation/relu.cu | 2 +- transformer_engine/common/activation/swiglu.cu | 2 +- transformer_engine/common/cast/cast.cu | 2 +- transformer_engine/common/cast/core/common.cuh | 2 +- transformer_engine/common/cast/dispatch/dequantize.cuh | 2 +- transformer_engine/common/cast/dispatch/gated.cuh | 2 +- transformer_engine/common/cast/dispatch/quantize.cuh | 2 +- transformer_engine/common/cast/fp8/dequantize_fp8.cuh | 2 +- transformer_engine/common/cast/fp8/gated_fp8.cuh | 2 +- transformer_engine/common/cast/fp8/quantize_fp8.cuh | 2 +- transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh | 2 +- transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh | 2 +- transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh | 2 +- .../common/cast/mxfp8/specialized/quantize_mxfp8.cuh | 2 +- .../common/cast/mxfp8/specialized/state_counter.cuh | 2 +- transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh | 2 +- transformer_engine/common/cast/nvfp4/core_nvfp4.cuh | 2 +- transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh | 2 +- .../common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh | 2 +- transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh | 2 +- .../common/cast/nvfp4/quantize_transpose_nvfp4.cuh | 2 +- transformer_engine/common/comm_gemm/comm_gemm.cpp | 2 +- .../common/comm_gemm_overlap/comm_gemm_overlap.cpp | 2 +- .../common/comm_gemm_overlap/userbuffers/ipcsocket.cc | 2 +- .../common/comm_gemm_overlap/userbuffers/ipcsocket.h | 2 +- .../common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp | 2 +- .../common/comm_gemm_overlap/userbuffers/userbuffers.cu | 2 +- .../common/comm_gemm_overlap/userbuffers/userbuffers.h | 2 +- transformer_engine/common/common.cu | 2 +- transformer_engine/common/common.h | 2 +- transformer_engine/common/cudnn_utils.cpp | 2 +- transformer_engine/common/cudnn_utils.h | 2 +- transformer_engine/common/dropout/dropout.cu | 2 +- transformer_engine/common/fused_attn/context_parallel.cu | 2 +- transformer_engine/common/fused_attn/flash_attn.cu | 2 +- transformer_engine/common/fused_attn/fused_attn.cpp | 2 +- .../common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu | 2 +- .../common/fused_attn/fused_attn_f16_arbitrary_seqlen.h | 2 +- .../common/fused_attn/fused_attn_f16_max512_seqlen.cu | 2 +- .../common/fused_attn/fused_attn_f16_max512_seqlen.h | 2 +- transformer_engine/common/fused_attn/fused_attn_fp8.cu | 2 +- transformer_engine/common/fused_attn/fused_attn_fp8.h | 2 +- transformer_engine/common/fused_attn/kv_cache.cu | 2 +- transformer_engine/common/fused_attn/utils.cu | 2 +- transformer_engine/common/fused_attn/utils.h | 2 +- transformer_engine/common/fused_rope/fused_rope.cu | 2 +- transformer_engine/common/fused_router/fused_moe_aux_loss.cu | 2 +- .../common/fused_router/fused_score_for_moe_aux_loss.cu | 2 +- .../common/fused_router/fused_topk_with_score_function.cu | 2 +- transformer_engine/common/fused_router/utils.h | 2 +- .../fused_softmax/scaled_aligned_causal_masked_softmax.cu | 2 +- .../common/fused_softmax/scaled_masked_softmax.cu | 2 +- .../common/fused_softmax/scaled_upper_triang_masked_softmax.cu | 2 +- transformer_engine/common/gemm/config.cpp | 2 +- transformer_engine/common/gemm/config.h | 2 +- transformer_engine/common/gemm/cublaslt_gemm.cu | 2 +- transformer_engine/common/gemm/cutlass_grouped_gemm.cu | 2 +- transformer_engine/common/gemm/cutlass_grouped_gemm.cuh | 2 +- .../common/hadamard_transform/customized_pipeline.cuh | 2 +- .../common/hadamard_transform/group_hadamard_transform.cu | 2 +- .../hadamard_transform/group_hadamard_transform_cast_fusion.cu | 2 +- .../group_row_cast_col_hadamard_transform_cast_fusion.cu | 2 +- .../common/hadamard_transform/hadamard_transform.cu | 2 +- .../common/hadamard_transform/hadamard_transform_cast_fusion.cu | 2 +- .../common/hadamard_transform/hadamard_transform_utils.cuh | 2 +- .../common/include/transformer_engine/activation.h | 2 +- transformer_engine/common/include/transformer_engine/cast.h | 2 +- .../common/include/transformer_engine/cast_transpose_noop.h | 2 +- .../common/include/transformer_engine/comm_gemm.h | 2 +- .../common/include/transformer_engine/comm_gemm_overlap.h | 2 +- transformer_engine/common/include/transformer_engine/cudnn.h | 2 +- transformer_engine/common/include/transformer_engine/dropout.h | 2 +- .../common/include/transformer_engine/fused_attn.h | 2 +- .../common/include/transformer_engine/fused_rope.h | 2 +- .../common/include/transformer_engine/fused_router.h | 2 +- transformer_engine/common/include/transformer_engine/gemm.h | 2 +- .../common/include/transformer_engine/hadamard_transform.h | 2 +- .../common/include/transformer_engine/multi_stream.h | 2 +- .../common/include/transformer_engine/multi_tensor.h | 2 +- .../common/include/transformer_engine/normalization.h | 2 +- transformer_engine/common/include/transformer_engine/padding.h | 2 +- .../common/include/transformer_engine/permutation.h | 2 +- transformer_engine/common/include/transformer_engine/recipe.h | 2 +- transformer_engine/common/include/transformer_engine/softmax.h | 2 +- transformer_engine/common/include/transformer_engine/swizzle.h | 2 +- .../common/include/transformer_engine/transformer_engine.h | 2 +- .../common/include/transformer_engine/transpose.h | 2 +- transformer_engine/common/multi_tensor/adam.cu | 2 +- transformer_engine/common/multi_tensor/compute_scale.cu | 2 +- transformer_engine/common/multi_tensor/l2norm.cu | 2 +- transformer_engine/common/multi_tensor/multi_tensor_apply.cuh | 2 +- transformer_engine/common/multi_tensor/scale.cu | 2 +- transformer_engine/common/multi_tensor/sgd.cu | 2 +- transformer_engine/common/normalization/common.cpp | 2 +- transformer_engine/common/normalization/common.h | 2 +- transformer_engine/common/normalization/kernel_traits.h | 2 +- transformer_engine/common/normalization/layernorm/ln_api.cpp | 2 +- .../common/normalization/layernorm/ln_bwd_kernels.cuh | 2 +- .../common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu | 2 +- .../common/normalization/layernorm/ln_fwd_cuda_kernel.cu | 2 +- .../common/normalization/layernorm/ln_fwd_kernels.cuh | 2 +- transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp | 2 +- .../common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh | 2 +- .../normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu | 2 +- .../common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu | 2 +- .../common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh | 2 +- transformer_engine/common/nvshmem_api/CMakeLists.txt | 2 +- transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu | 2 +- transformer_engine/common/nvshmem_api/nvshmem_waitkernel.h | 2 +- transformer_engine/common/nvtx.h | 2 +- transformer_engine/common/permutation/permutation.cu | 2 +- transformer_engine/common/recipe/__init__.py | 2 +- transformer_engine/common/recipe/current_scaling.cu | 2 +- transformer_engine/common/recipe/delayed_scaling.cu | 2 +- transformer_engine/common/recipe/fp8_block_scaling.cu | 2 +- transformer_engine/common/recipe/mxfp8_scaling.cu | 2 +- transformer_engine/common/recipe/nvfp4.cu | 2 +- transformer_engine/common/recipe/recipe_common.cuh | 2 +- transformer_engine/common/swizzle/swizzle.cu | 2 +- transformer_engine/common/swizzle/swizzle_block_scaling.cu | 2 +- transformer_engine/common/transformer_engine.cpp | 2 +- transformer_engine/common/transpose/cast_transpose.cu | 2 +- transformer_engine/common/transpose/cast_transpose.h | 2 +- transformer_engine/common/transpose/cast_transpose_fusion.cu | 2 +- transformer_engine/common/transpose/multi_cast_transpose.cu | 2 +- .../common/transpose/quantize_transpose_square_blockwise.cu | 2 +- .../common/transpose/quantize_transpose_vector_blockwise.cu | 2 +- .../common/transpose/quantize_transpose_vector_blockwise_fp4.cu | 2 +- transformer_engine/common/transpose/rtc/cast_transpose.cu | 2 +- .../common/transpose/rtc/cast_transpose_fusion.cu | 2 +- transformer_engine/common/transpose/rtc/swap_first_dims.cu | 2 +- transformer_engine/common/transpose/rtc/transpose.cu | 2 +- transformer_engine/common/transpose/swap_first_dims.cu | 2 +- transformer_engine/common/transpose/transpose.cu | 2 +- transformer_engine/common/transpose/transpose_fusion.cu | 2 +- transformer_engine/common/triton/__init__.py | 2 +- transformer_engine/common/triton/cross_entropy.py | 2 +- transformer_engine/common/triton/pad.py | 2 +- transformer_engine/common/triton/permutation.py | 2 +- transformer_engine/common/util/cuda_driver.cpp | 2 +- transformer_engine/common/util/cuda_driver.h | 2 +- transformer_engine/common/util/cuda_nvml.cpp | 2 +- transformer_engine/common/util/cuda_nvml.h | 2 +- transformer_engine/common/util/cuda_runtime.cpp | 2 +- transformer_engine/common/util/cuda_runtime.h | 2 +- transformer_engine/common/util/curanddx.hpp | 2 +- transformer_engine/common/util/handle_manager.h | 2 +- transformer_engine/common/util/logging.h | 2 +- transformer_engine/common/util/math.h | 2 +- transformer_engine/common/util/multi_stream.cpp | 2 +- transformer_engine/common/util/multi_stream.h | 2 +- transformer_engine/common/util/padding.cu | 2 +- transformer_engine/common/util/ptx.cuh | 2 +- transformer_engine/common/util/pybind_helper.h | 2 +- transformer_engine/common/util/rtc.cpp | 2 +- transformer_engine/common/util/rtc.h | 2 +- transformer_engine/common/util/shared_lib_wrapper.h | 2 +- transformer_engine/common/util/string.h | 2 +- transformer_engine/common/util/string_header.h.in | 2 +- transformer_engine/common/util/system.h | 2 +- transformer_engine/common/util/vectorized_pointwise.h | 2 +- transformer_engine/common/utils.cuh | 2 +- transformer_engine/common/utils.py | 2 +- transformer_engine/debug/__init__.py | 2 +- transformer_engine/debug/features/__init__.py | 2 +- transformer_engine/debug/features/_test_dummy_feature.py | 2 +- transformer_engine/debug/features/api.py | 2 +- transformer_engine/debug/features/disable_fp8_gemm.py | 2 +- transformer_engine/debug/features/disable_fp8_layer.py | 2 +- transformer_engine/debug/features/fake_quant.py | 2 +- transformer_engine/debug/features/log_fp8_tensor_stats.py | 2 +- transformer_engine/debug/features/log_tensor_stats.py | 2 +- transformer_engine/debug/features/per_tensor_scaling.py | 2 +- transformer_engine/debug/features/utils/__init__.py | 2 +- transformer_engine/debug/features/utils/stats_buffer.py | 2 +- transformer_engine/debug/features/utils/stats_computation.py | 2 +- transformer_engine/debug/pytorch/__init__.py | 2 +- transformer_engine/debug/pytorch/debug_quantization.py | 2 +- transformer_engine/debug/pytorch/debug_state.py | 2 +- transformer_engine/debug/pytorch/utils.py | 2 +- transformer_engine/jax/__init__.py | 2 +- transformer_engine/jax/activation.py | 2 +- transformer_engine/jax/attention.py | 2 +- transformer_engine/jax/checkpoint_policies.py | 2 +- transformer_engine/jax/cpp_extensions/__init__.py | 2 +- transformer_engine/jax/cpp_extensions/activation.py | 2 +- transformer_engine/jax/cpp_extensions/amax.py | 2 +- transformer_engine/jax/cpp_extensions/attention.py | 2 +- transformer_engine/jax/cpp_extensions/base.py | 2 +- transformer_engine/jax/cpp_extensions/gemm.py | 2 +- transformer_engine/jax/cpp_extensions/misc.py | 2 +- transformer_engine/jax/cpp_extensions/normalization.py | 2 +- transformer_engine/jax/cpp_extensions/quantization.py | 2 +- transformer_engine/jax/cpp_extensions/softmax.py | 2 +- transformer_engine/jax/csrc/extensions.h | 2 +- transformer_engine/jax/csrc/extensions/activation.cpp | 2 +- transformer_engine/jax/csrc/extensions/amax.cpp | 2 +- transformer_engine/jax/csrc/extensions/attention.cpp | 2 +- transformer_engine/jax/csrc/extensions/cgemm_helper.cpp | 2 +- transformer_engine/jax/csrc/extensions/cgemm_helper.h | 2 +- transformer_engine/jax/csrc/extensions/cublas.cpp | 2 +- transformer_engine/jax/csrc/extensions/cudnn.cpp | 2 +- transformer_engine/jax/csrc/extensions/ffi.cpp | 2 +- transformer_engine/jax/csrc/extensions/ffi.h | 2 +- transformer_engine/jax/csrc/extensions/gemm.cpp | 2 +- transformer_engine/jax/csrc/extensions/misc.cpp | 2 +- transformer_engine/jax/csrc/extensions/misc.h | 2 +- transformer_engine/jax/csrc/extensions/normalization.cpp | 2 +- transformer_engine/jax/csrc/extensions/pybind.cpp | 2 +- transformer_engine/jax/csrc/extensions/quantization.cpp | 2 +- transformer_engine/jax/csrc/extensions/softmax.cpp | 2 +- transformer_engine/jax/csrc/extensions/utils.cpp | 2 +- transformer_engine/jax/csrc/extensions/utils.h | 2 +- transformer_engine/jax/dense.py | 2 +- transformer_engine/jax/flax/__init__.py | 2 +- transformer_engine/jax/flax/module.py | 2 +- transformer_engine/jax/flax/transformer.py | 2 +- transformer_engine/jax/layernorm.py | 2 +- transformer_engine/jax/layernorm_dense.py | 2 +- transformer_engine/jax/layernorm_mlp.py | 2 +- transformer_engine/jax/permutation.py | 2 +- transformer_engine/jax/pyproject.toml | 2 +- transformer_engine/jax/quantize/__init__.py | 2 +- transformer_engine/jax/quantize/dequantizer.py | 2 +- transformer_engine/jax/quantize/device_utils.py | 2 +- transformer_engine/jax/quantize/hadamard.py | 2 +- transformer_engine/jax/quantize/helper.py | 2 +- transformer_engine/jax/quantize/metadata.py | 2 +- transformer_engine/jax/quantize/misc.py | 2 +- transformer_engine/jax/quantize/quantizer.py | 2 +- transformer_engine/jax/quantize/scaling_modes.py | 2 +- transformer_engine/jax/quantize/tensor.py | 2 +- transformer_engine/jax/setup.py | 2 +- transformer_engine/jax/sharding.py | 2 +- transformer_engine/jax/softmax.py | 2 +- transformer_engine/jax/triton_extensions/__init__.py | 2 +- transformer_engine/jax/triton_extensions/permutation.py | 2 +- transformer_engine/jax/triton_extensions/utils.py | 2 +- transformer_engine/pytorch/__init__.py | 2 +- transformer_engine/pytorch/attention/__init__.py | 2 +- .../pytorch/attention/dot_product_attention/__init__.py | 2 +- .../pytorch/attention/dot_product_attention/backends.py | 2 +- .../pytorch/attention/dot_product_attention/context_parallel.py | 2 +- .../attention/dot_product_attention/dot_product_attention.py | 2 +- .../pytorch/attention/dot_product_attention/softmax.py | 2 +- .../pytorch/attention/dot_product_attention/utils.py | 2 +- transformer_engine/pytorch/attention/inference.py | 2 +- transformer_engine/pytorch/attention/multi_head_attention.py | 2 +- transformer_engine/pytorch/attention/rope.py | 2 +- transformer_engine/pytorch/constants.py | 2 +- transformer_engine/pytorch/cpp_extensions/__init__.py | 2 +- transformer_engine/pytorch/cpp_extensions/fused_attn.py | 2 +- transformer_engine/pytorch/cpp_extensions/gemm.py | 2 +- transformer_engine/pytorch/cpu_offload.py | 2 +- transformer_engine/pytorch/cpu_offload_v1.py | 2 +- transformer_engine/pytorch/cross_entropy.py | 2 +- transformer_engine/pytorch/csrc/common.cpp | 2 +- transformer_engine/pytorch/csrc/common.h | 2 +- transformer_engine/pytorch/csrc/extensions.h | 2 +- transformer_engine/pytorch/csrc/extensions/activation.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/apply_rope.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/attention.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/bias.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/cast.cpp | 2 +- .../pytorch/csrc/extensions/comm_gemm_overlap.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/dropout.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/gemm.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/misc.cpp | 2 +- .../pytorch/csrc/extensions/multi_tensor/adam.cpp | 2 +- .../pytorch/csrc/extensions/multi_tensor/compute_scale.cpp | 2 +- .../pytorch/csrc/extensions/multi_tensor/l2norm.cpp | 2 +- .../pytorch/csrc/extensions/multi_tensor/scale.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/multi_tensor/sgd.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/normalization.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/padding.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/permutation.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/pybind.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/recipe.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/router.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/softmax.cpp | 2 +- transformer_engine/pytorch/csrc/extensions/transpose.cpp | 2 +- transformer_engine/pytorch/csrc/pybind.h | 2 +- transformer_engine/pytorch/csrc/quantizer.cpp | 2 +- transformer_engine/pytorch/csrc/type_converters.cpp | 2 +- transformer_engine/pytorch/csrc/util.cpp | 2 +- transformer_engine/pytorch/csrc/util.h | 2 +- transformer_engine/pytorch/custom_recipes/__init__.py | 2 +- transformer_engine/pytorch/custom_recipes/gemm.py | 2 +- transformer_engine/pytorch/custom_recipes/quantization.py | 2 +- .../pytorch/custom_recipes/quantization_current_scaling.py | 2 +- transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py | 2 +- transformer_engine/pytorch/custom_recipes/utils.py | 2 +- transformer_engine/pytorch/distributed.py | 2 +- transformer_engine/pytorch/export.py | 2 +- transformer_engine/pytorch/float8_tensor.py | 2 +- transformer_engine/pytorch/fp8.py | 2 +- transformer_engine/pytorch/graph.py | 2 +- transformer_engine/pytorch/jit.py | 2 +- transformer_engine/pytorch/module/__init__.py | 2 +- transformer_engine/pytorch/module/_common.py | 2 +- transformer_engine/pytorch/module/base.py | 2 +- transformer_engine/pytorch/module/fp8_padding.py | 2 +- transformer_engine/pytorch/module/fp8_unpadding.py | 2 +- transformer_engine/pytorch/module/grouped_linear.py | 2 +- transformer_engine/pytorch/module/layernorm.py | 2 +- transformer_engine/pytorch/module/layernorm_linear.py | 2 +- transformer_engine/pytorch/module/layernorm_mlp.py | 2 +- transformer_engine/pytorch/module/linear.py | 2 +- transformer_engine/pytorch/module/rmsnorm.py | 2 +- transformer_engine/pytorch/numerics_debug.py | 2 +- transformer_engine/pytorch/onnx_extensions.py | 2 +- transformer_engine/pytorch/ops/__init__.py | 2 +- transformer_engine/pytorch/ops/_common.py | 2 +- transformer_engine/pytorch/ops/basic/__init__.py | 2 +- transformer_engine/pytorch/ops/basic/activation.py | 2 +- transformer_engine/pytorch/ops/basic/add_extra_input.py | 2 +- transformer_engine/pytorch/ops/basic/all_gather.py | 2 +- transformer_engine/pytorch/ops/basic/all_reduce.py | 2 +- transformer_engine/pytorch/ops/basic/basic_linear.py | 2 +- transformer_engine/pytorch/ops/basic/bias.py | 2 +- transformer_engine/pytorch/ops/basic/constant_scale.py | 2 +- transformer_engine/pytorch/ops/basic/dropout.py | 2 +- transformer_engine/pytorch/ops/basic/identity.py | 2 +- transformer_engine/pytorch/ops/basic/l2normalization.py | 2 +- transformer_engine/pytorch/ops/basic/layer_norm.py | 2 +- transformer_engine/pytorch/ops/basic/make_extra_output.py | 2 +- transformer_engine/pytorch/ops/basic/quantize.py | 2 +- transformer_engine/pytorch/ops/basic/reduce_scatter.py | 2 +- transformer_engine/pytorch/ops/basic/reshape.py | 2 +- transformer_engine/pytorch/ops/basic/rmsnorm.py | 2 +- transformer_engine/pytorch/ops/fused/__init__.py | 2 +- .../pytorch/ops/fused/backward_activation_bias.py | 2 +- transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py | 2 +- transformer_engine/pytorch/ops/fused/backward_linear_add.py | 2 +- transformer_engine/pytorch/ops/fused/backward_linear_scale.py | 2 +- .../pytorch/ops/fused/forward_linear_bias_activation.py | 2 +- transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py | 2 +- .../pytorch/ops/fused/forward_linear_scale_add.py | 2 +- .../pytorch/ops/fused/userbuffers_backward_linear.py | 2 +- .../pytorch/ops/fused/userbuffers_forward_linear.py | 2 +- transformer_engine/pytorch/ops/fuser.py | 2 +- transformer_engine/pytorch/ops/linear.py | 2 +- transformer_engine/pytorch/ops/op.py | 2 +- transformer_engine/pytorch/ops/sequential.py | 2 +- transformer_engine/pytorch/optimizers/__init__.py | 2 +- transformer_engine/pytorch/optimizers/fused_adam.py | 2 +- transformer_engine/pytorch/optimizers/fused_sgd.py | 2 +- transformer_engine/pytorch/optimizers/multi_tensor_apply.py | 2 +- transformer_engine/pytorch/permutation.py | 2 +- transformer_engine/pytorch/pyproject.toml | 2 +- transformer_engine/pytorch/quantization.py | 2 +- transformer_engine/pytorch/quantized_tensor.py | 2 +- transformer_engine/pytorch/router.py | 2 +- transformer_engine/pytorch/setup.py | 2 +- transformer_engine/pytorch/tensor/__init__.py | 2 +- transformer_engine/pytorch/tensor/_quantization_helpers.py | 2 +- transformer_engine/pytorch/tensor/float8_blockwise_tensor.py | 2 +- transformer_engine/pytorch/tensor/float8_tensor.py | 2 +- transformer_engine/pytorch/tensor/mxfp8_tensor.py | 2 +- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 2 +- transformer_engine/pytorch/tensor/storage/__init__.py | 2 +- .../pytorch/tensor/storage/float8_blockwise_tensor_storage.py | 2 +- .../pytorch/tensor/storage/float8_tensor_storage.py | 2 +- .../pytorch/tensor/storage/mxfp8_tensor_storage.py | 2 +- .../pytorch/tensor/storage/nvfp4_tensor_storage.py | 2 +- transformer_engine/pytorch/tensor/utils.py | 2 +- transformer_engine/pytorch/torch_version.py | 2 +- transformer_engine/pytorch/transformer.py | 2 +- transformer_engine/pytorch/triton/__init__.py | 2 +- transformer_engine/pytorch/triton/cross_entropy.py | 2 +- transformer_engine/pytorch/triton/pad.py | 2 +- transformer_engine/pytorch/triton/permutation.py | 2 +- transformer_engine/pytorch/utils.py | 2 +- 622 files changed, 622 insertions(+), 622 deletions(-) diff --git a/.github/actions/build-pytorch-wheel/Dockerfile b/.github/actions/build-pytorch-wheel/Dockerfile index 5bf0960fa7..a858307ab4 100644 --- a/.github/actions/build-pytorch-wheel/Dockerfile +++ b/.github/actions/build-pytorch-wheel/Dockerfile @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/actions/build-pytorch-wheel/action.yml b/.github/actions/build-pytorch-wheel/action.yml index a49b12227d..526f121d3b 100644 --- a/.github/actions/build-pytorch-wheel/action.yml +++ b/.github/actions/build-pytorch-wheel/action.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/actions/build-pytorch-wheel/build.sh b/.github/actions/build-pytorch-wheel/build.sh index 8a219a5959..9a0920e2be 100644 --- a/.github/actions/build-pytorch-wheel/build.sh +++ b/.github/actions/build-pytorch-wheel/build.sh @@ -1,6 +1,6 @@ #!/bin/bash -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/scripts/check_for_ngc_images.sh b/.github/scripts/check_for_ngc_images.sh index f065541838..9297c6a757 100644 --- a/.github/scripts/check_for_ngc_images.sh +++ b/.github/scripts/check_for_ngc_images.sh @@ -1,6 +1,6 @@ #!/bin/bash -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/attach-wheels-to-release.yml b/.github/workflows/attach-wheels-to-release.yml index 6b97d4bb9b..fb56622d77 100644 --- a/.github/workflows/attach-wheels-to-release.yml +++ b/.github/workflows/attach-wheels-to-release.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/blossom-ci.yml b/.github/workflows/blossom-ci.yml index 1402cc091a..88719231ef 100644 --- a/.github/workflows/blossom-ci.yml +++ b/.github/workflows/blossom-ci.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 2ed8766de0..427b2f27fa 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/deploy_nightly_docs.yml b/.github/workflows/deploy_nightly_docs.yml index 6470eee838..b4e015d2da 100644 --- a/.github/workflows/deploy_nightly_docs.yml +++ b/.github/workflows/deploy_nightly_docs.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 5beeeb8879..388c822eeb 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/license.yml b/.github/workflows/license.yml index d70c7def61..e12f50991f 100644 --- a/.github/workflows/license.yml +++ b/.github/workflows/license.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index ee6433d484..c5cb748c2c 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml index f12a95d79a..13ea45b070 100644 --- a/.github/workflows/trigger-ci.yml +++ b/.github/workflows/trigger-ci.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/.github/workflows/upload-ci-logs.yml b/.github/workflows/upload-ci-logs.yml index c9c7e4ef4d..a5fa93ddb7 100644 --- a/.github/workflows/upload-ci-logs.yml +++ b/.github/workflows/upload-ci-logs.yml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst index d92fd95675..14f1ee08d2 100644 --- a/CONTRIBUTING.rst +++ b/CONTRIBUTING.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/CPPLINT.cfg b/CPPLINT.cfg index ecfbbf3d0b..8062e18058 100644 --- a/CPPLINT.cfg +++ b/CPPLINT.cfg @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/README.rst b/README.rst index d2bb2f4056..9241e26cdd 100644 --- a/README.rst +++ b/README.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/benchmarks/attention/benchmark_attention.py b/benchmarks/attention/benchmark_attention.py index 1df16cc016..77b2da0b10 100644 --- a/benchmarks/attention/benchmark_attention.py +++ b/benchmarks/attention/benchmark_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/benchmarks/benchmark_rht_cast.py b/benchmarks/benchmark_rht_cast.py index 9c47856f71..badab1d199 100644 --- a/benchmarks/benchmark_rht_cast.py +++ b/benchmarks/benchmark_rht_cast.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/benchmarks/linear/benchmark_grouped_linear.py b/benchmarks/linear/benchmark_grouped_linear.py index f559928f8c..815e367f71 100644 --- a/benchmarks/linear/benchmark_grouped_linear.py +++ b/benchmarks/linear/benchmark_grouped_linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/__init__.py b/build_tools/__init__.py index 7669e4cfa6..bf3d8cd0c3 100644 --- a/build_tools/__init__.py +++ b/build_tools/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py index c269a29874..cbb8838b00 100644 --- a/build_tools/build_ext.py +++ b/build_tools/build_ext.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/jax.py b/build_tools/jax.py index df78bf3e2f..276c9943d6 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index b03ef04fa4..b4815a0942 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/te_version.py b/build_tools/te_version.py index 0aee63f647..f4a1a587ed 100644 --- a/build_tools/te_version.py +++ b/build_tools/te_version.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/utils.py b/build_tools/utils.py index 50ba007594..8a52440310 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/wheel_utils/Dockerfile.aarch b/build_tools/wheel_utils/Dockerfile.aarch index 404cb941cb..8c5b81d92b 100644 --- a/build_tools/wheel_utils/Dockerfile.aarch +++ b/build_tools/wheel_utils/Dockerfile.aarch @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/wheel_utils/Dockerfile.x86 b/build_tools/wheel_utils/Dockerfile.x86 index daa7f961cd..b77920250a 100644 --- a/build_tools/wheel_utils/Dockerfile.x86 +++ b/build_tools/wheel_utils/Dockerfile.x86 @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/wheel_utils/build_wheels.sh b/build_tools/wheel_utils/build_wheels.sh index d0055b791d..e9ec854dba 100644 --- a/build_tools/wheel_utils/build_wheels.sh +++ b/build_tools/wheel_utils/build_wheels.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/wheel_utils/launch_aarch.sh b/build_tools/wheel_utils/launch_aarch.sh index 85f754ca19..a6f30da62d 100644 --- a/build_tools/wheel_utils/launch_aarch.sh +++ b/build_tools/wheel_utils/launch_aarch.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/build_tools/wheel_utils/launch_x86.sh b/build_tools/wheel_utils/launch_x86.sh index 11fc522947..9fdc6871ed 100644 --- a/build_tools/wheel_utils/launch_x86.sh +++ b/build_tools/wheel_utils/launch_x86.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/api/c/activation.rst b/docs/api/c/activation.rst index 5b50aa513d..6bba2abf34 100644 --- a/docs/api/c/activation.rst +++ b/docs/api/c/activation.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/cast.rst b/docs/api/c/cast.rst index 2ae05a8456..2ffff8f701 100644 --- a/docs/api/c/cast.rst +++ b/docs/api/c/cast.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/cast_transpose_noop.rst b/docs/api/c/cast_transpose_noop.rst index ae80c5d2d4..0399a10f67 100644 --- a/docs/api/c/cast_transpose_noop.rst +++ b/docs/api/c/cast_transpose_noop.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/cudnn.rst b/docs/api/c/cudnn.rst index 5d93c4d6e4..279fb39126 100644 --- a/docs/api/c/cudnn.rst +++ b/docs/api/c/cudnn.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/fused_attn.rst b/docs/api/c/fused_attn.rst index 6db67f26fe..00333716b8 100644 --- a/docs/api/c/fused_attn.rst +++ b/docs/api/c/fused_attn.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/fused_rope.rst b/docs/api/c/fused_rope.rst index 289bb53d9b..d40b795319 100644 --- a/docs/api/c/fused_rope.rst +++ b/docs/api/c/fused_rope.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/gemm.rst b/docs/api/c/gemm.rst index 711733fc4c..30b48e5a94 100644 --- a/docs/api/c/gemm.rst +++ b/docs/api/c/gemm.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/index.rst b/docs/api/c/index.rst index 0499f52f05..82024e25b0 100644 --- a/docs/api/c/index.rst +++ b/docs/api/c/index.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/multi_tensor.rst b/docs/api/c/multi_tensor.rst index 8ba2d274c7..f8383adfe8 100644 --- a/docs/api/c/multi_tensor.rst +++ b/docs/api/c/multi_tensor.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/normalization.rst b/docs/api/c/normalization.rst index edbea00ac0..d2f418b3e8 100644 --- a/docs/api/c/normalization.rst +++ b/docs/api/c/normalization.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/padding.rst b/docs/api/c/padding.rst index 2141b874d2..c58d274eb3 100644 --- a/docs/api/c/padding.rst +++ b/docs/api/c/padding.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/permutation.rst b/docs/api/c/permutation.rst index bad6961621..32f56c26fd 100644 --- a/docs/api/c/permutation.rst +++ b/docs/api/c/permutation.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/recipe.rst b/docs/api/c/recipe.rst index 7c368f69b6..0b1243a14a 100644 --- a/docs/api/c/recipe.rst +++ b/docs/api/c/recipe.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/softmax.rst b/docs/api/c/softmax.rst index 55dc5d47de..2863334574 100644 --- a/docs/api/c/softmax.rst +++ b/docs/api/c/softmax.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/swizzle.rst b/docs/api/c/swizzle.rst index b2dd8f5977..dadf104fff 100644 --- a/docs/api/c/swizzle.rst +++ b/docs/api/c/swizzle.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/transformer_engine.rst b/docs/api/c/transformer_engine.rst index b5fd95e005..84c3c554f8 100644 --- a/docs/api/c/transformer_engine.rst +++ b/docs/api/c/transformer_engine.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/c/transpose.rst b/docs/api/c/transpose.rst index 9a3ba9e48b..327d97f680 100644 --- a/docs/api/c/transpose.rst +++ b/docs/api/c/transpose.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/common.rst b/docs/api/common.rst index 728dcd6ed0..3ddb66df98 100644 --- a/docs/api/common.rst +++ b/docs/api/common.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/framework.rst b/docs/api/framework.rst index 0ac1a0e34e..b950ef5d93 100644 --- a/docs/api/framework.rst +++ b/docs/api/framework.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/jax.rst b/docs/api/jax.rst index 99782f99c7..7a31c9d379 100644 --- a/docs/api/jax.rst +++ b/docs/api/jax.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 18abe0f2c2..d1d54c0dda 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/conf.py b/docs/conf.py index 479c1f8948..0734008137 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/debug.rst b/docs/debug.rst index 527f30ed02..7c3735ee12 100644 --- a/docs/debug.rst +++ b/docs/debug.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/1_getting_started.rst b/docs/debug/1_getting_started.rst index 58adf73cef..de72b2242d 100644 --- a/docs/debug/1_getting_started.rst +++ b/docs/debug/1_getting_started.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/2_config_file_structure.rst b/docs/debug/2_config_file_structure.rst index 2132cd32c7..3ade970b57 100644 --- a/docs/debug/2_config_file_structure.rst +++ b/docs/debug/2_config_file_structure.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/3_api_debug_setup.rst b/docs/debug/3_api_debug_setup.rst index 176bc13d32..9974df1afb 100644 --- a/docs/debug/3_api_debug_setup.rst +++ b/docs/debug/3_api_debug_setup.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/3_api_features.rst b/docs/debug/3_api_features.rst index 8cdbde8edd..fc48371a3f 100644 --- a/docs/debug/3_api_features.rst +++ b/docs/debug/3_api_features.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/3_api_te_calls.rst b/docs/debug/3_api_te_calls.rst index 1435d41d77..d9386e4390 100644 --- a/docs/debug/3_api_te_calls.rst +++ b/docs/debug/3_api_te_calls.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/4_distributed.rst b/docs/debug/4_distributed.rst index 764fee6541..51911fa71b 100644 --- a/docs/debug/4_distributed.rst +++ b/docs/debug/4_distributed.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/debug/api.rst b/docs/debug/api.rst index 6ccb32cc8b..a195734fcb 100644 --- a/docs/debug/api.rst +++ b/docs/debug/api.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py b/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py index 97f1bcd7ec..569af333dc 100644 --- a/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py +++ b/docs/examples/attention/arbitrary_mask_to_post_scale_bias.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/attention/example_attention.py b/docs/examples/attention/example_attention.py index cf650265bc..207d6ee974 100644 --- a/docs/examples/attention/example_attention.py +++ b/docs/examples/attention/example_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/onnx/utils.py b/docs/examples/onnx/utils.py index 7acf2ffc68..6dc4b32725 100644 --- a/docs/examples/onnx/utils.py +++ b/docs/examples/onnx/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/quickstart_jax_utils.py b/docs/examples/quickstart_jax_utils.py index f1ff9b7d99..3a2947f338 100644 --- a/docs/examples/quickstart_jax_utils.py +++ b/docs/examples/quickstart_jax_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/quickstart_utils.py b/docs/examples/quickstart_utils.py index 473fce7fe7..9b21807255 100644 --- a/docs/examples/quickstart_utils.py +++ b/docs/examples/quickstart_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/te_gemma/te_gemma.py b/docs/examples/te_gemma/te_gemma.py index d3de8a185d..aa9fa4b656 100755 --- a/docs/examples/te_gemma/te_gemma.py +++ b/docs/examples/te_gemma/te_gemma.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/te_gemma/te_gemma_loading_weights.py b/docs/examples/te_gemma/te_gemma_loading_weights.py index 36b0a5b739..f3ca34262c 100755 --- a/docs/examples/te_gemma/te_gemma_loading_weights.py +++ b/docs/examples/te_gemma/te_gemma_loading_weights.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/te_gemma/utils.py b/docs/examples/te_gemma/utils.py index 9b67f178fa..7297dbaafb 100755 --- a/docs/examples/te_gemma/utils.py +++ b/docs/examples/te_gemma/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/te_llama/te_llama.py b/docs/examples/te_llama/te_llama.py index 8297ac6d2e..b2d4d183ab 100644 --- a/docs/examples/te_llama/te_llama.py +++ b/docs/examples/te_llama/te_llama.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/examples/te_llama/utils.py b/docs/examples/te_llama/utils.py index 66f05701f5..4bc9f7e77a 100644 --- a/docs/examples/te_llama/utils.py +++ b/docs/examples/te_llama/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/docs/faq.rst b/docs/faq.rst index a9406ed459..0c55223fb1 100644 --- a/docs/faq.rst +++ b/docs/faq.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/getting_started.rst b/docs/getting_started.rst index 2e8047763a..f5ab2ae695 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/index.rst b/docs/index.rst index 7a3ab9f6fd..3f707d8904 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/docs/installation.rst b/docs/installation.rst index 24563c456e..cc48a0adac 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -1,5 +1,5 @@ .. - Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. See LICENSE for license information. diff --git a/examples/jax/collective_gemm/common.py b/examples/jax/collective_gemm/common.py index da79b21377..0d812da057 100644 --- a/examples/jax/collective_gemm/common.py +++ b/examples/jax/collective_gemm/common.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Shared functions for the comm_overlap tests""" diff --git a/examples/jax/collective_gemm/conftest.py b/examples/jax/collective_gemm/conftest.py index 83937971a4..5be5709ba7 100644 --- a/examples/jax/collective_gemm/conftest.py +++ b/examples/jax/collective_gemm/conftest.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/jax/collective_gemm/run_test_cgemm.sh b/examples/jax/collective_gemm/run_test_cgemm.sh index af263eb53d..388c878376 100644 --- a/examples/jax/collective_gemm/run_test_cgemm.sh +++ b/examples/jax/collective_gemm/run_test_cgemm.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/jax/collective_gemm/test_dense_grad.py b/examples/jax/collective_gemm/test_dense_grad.py index e14329d48f..94c7dc5b66 100644 --- a/examples/jax/collective_gemm/test_dense_grad.py +++ b/examples/jax/collective_gemm/test_dense_grad.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Collective Dense Gradient test on multi-GPU with tensor parallelism""" diff --git a/examples/jax/collective_gemm/test_gemm.py b/examples/jax/collective_gemm/test_gemm.py index ac86c551d7..d2994723bb 100644 --- a/examples/jax/collective_gemm/test_gemm.py +++ b/examples/jax/collective_gemm/test_gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Collective GEMM test on multi-GPU with tensor parallelism diff --git a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py index 407cec68a3..61c960a7aa 100644 --- a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py +++ b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Collective Dense Gradient test on multi-GPU with tensor parallelism""" diff --git a/examples/jax/encoder/common.py b/examples/jax/encoder/common.py index 819fdf443d..7906d44aec 100644 --- a/examples/jax/encoder/common.py +++ b/examples/jax/encoder/common.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Shared functions for the encoder tests""" diff --git a/examples/jax/encoder/conftest.py b/examples/jax/encoder/conftest.py index b1648892aa..083c1b4dce 100644 --- a/examples/jax/encoder/conftest.py +++ b/examples/jax/encoder/conftest.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/jax/encoder/run_test_multiprocessing_encoder.sh b/examples/jax/encoder/run_test_multiprocessing_encoder.sh index fa7102cb42..f2ef33da46 100644 --- a/examples/jax/encoder/run_test_multiprocessing_encoder.sh +++ b/examples/jax/encoder/run_test_multiprocessing_encoder.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index 8618c2be87..b534db8576 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Encoder training on multi-GPU with tesnor parallelism""" diff --git a/examples/jax/encoder/test_multigpu_encoder.py b/examples/jax/encoder/test_multigpu_encoder.py index 80a2b043cb..98184ccd75 100644 --- a/examples/jax/encoder/test_multigpu_encoder.py +++ b/examples/jax/encoder/test_multigpu_encoder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Encoder training on multi-GPU with data parallelism""" diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index f3092278e8..327540521c 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Encoder training with multi-GPU, multiprocessing, and tensor parallelism""" diff --git a/examples/jax/encoder/test_single_gpu_encoder.py b/examples/jax/encoder/test_single_gpu_encoder.py index 7835b08b23..82c7fed38e 100644 --- a/examples/jax/encoder/test_single_gpu_encoder.py +++ b/examples/jax/encoder/test_single_gpu_encoder.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Encoder training on single GPU""" diff --git a/examples/jax/mnist/test_single_gpu_mnist.py b/examples/jax/mnist/test_single_gpu_mnist.py index 62f7954e0d..0c76d51c37 100644 --- a/examples/jax/mnist/test_single_gpu_mnist.py +++ b/examples/jax/mnist/test_single_gpu_mnist.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """MNIST training on single GPU""" diff --git a/examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py b/examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py index 1fd40305c9..8b3fe542ad 100644 --- a/examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py +++ b/examples/pytorch/comm_gemm_overlap/te_layer_with_overlap.py @@ -1,6 +1,6 @@ #!/usr/bin/python3 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/pytorch/fsdp/README.md b/examples/pytorch/fsdp/README.md index f9a49af8d8..414e5e638e 100644 --- a/examples/pytorch/fsdp/README.md +++ b/examples/pytorch/fsdp/README.md @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/pytorch/fsdp/fsdp.py b/examples/pytorch/fsdp/fsdp.py index 789389757e..b469ef56b7 100644 --- a/examples/pytorch/fsdp/fsdp.py +++ b/examples/pytorch/fsdp/fsdp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/examples/pytorch/mnist/main.py b/examples/pytorch/mnist/main.py index f4a48bfc92..3754e3643b 100644 --- a/examples/pytorch/mnist/main.py +++ b/examples/pytorch/mnist/main.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/pyproject.toml b/pyproject.toml index 35a7c20727..4a8fded172 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_cppunittest/test.sh b/qa/L0_cppunittest/test.sh index cd46b0b63c..0b83747c0e 100755 --- a/qa/L0_cppunittest/test.sh +++ b/qa/L0_cppunittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index 58ce409add..3f25816600 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas diff --git a/qa/L0_jax_lint/test.sh b/qa/L0_jax_lint/test.sh index dbc1ed0a1d..3f804d3ef9 100644 --- a/qa/L0_jax_lint/test.sh +++ b/qa/L0_jax_lint/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index c430e8d61b..ee9ce130aa 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas diff --git a/qa/L0_jax_wheel/test.sh b/qa/L0_jax_wheel/test.sh index bf9e4a4619..fa50a6de68 100644 --- a/qa/L0_jax_wheel/test.sh +++ b/qa/L0_jax_wheel/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_license/copyright_checker.py b/qa/L0_license/copyright_checker.py index a0e137d1ef..86b22f824d 100644 --- a/qa/L0_license/copyright_checker.py +++ b/qa/L0_license/copyright_checker.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 # coding: utf-8 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_license/test.sh b/qa/L0_license/test.sh index 44b9469e55..b2826c59e2 100644 --- a/qa/L0_license/test.sh +++ b/qa/L0_license/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index a176d21b15..ce65bc4305 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_pytorch_lint/test.sh b/qa/L0_pytorch_lint/test.sh index e2c50c445e..f08dd8a03d 100644 --- a/qa/L0_pytorch_lint/test.sh +++ b/qa/L0_pytorch_lint/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 512c01db42..21eed28367 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L0_pytorch_wheel/test.sh b/qa/L0_pytorch_wheel/test.sh index 3056547ef2..fcf1a52b9c 100644 --- a/qa/L0_pytorch_wheel/test.sh +++ b/qa/L0_pytorch_wheel/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L1_cpp_distributed/test.sh b/qa/L1_cpp_distributed/test.sh index e074b46ae6..8d767a4efb 100755 --- a/qa/L1_cpp_distributed/test.sh +++ b/qa/L1_cpp_distributed/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L1_jax_distributed_unittest/test.sh b/qa/L1_jax_distributed_unittest/test.sh index b93224e64d..4f92d1c783 100644 --- a/qa/L1_jax_distributed_unittest/test.sh +++ b/qa/L1_jax_distributed_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index b1e3a3e15c..9d868d99cf 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L1_pytorch_mcore_integration/test.sh b/qa/L1_pytorch_mcore_integration/test.sh index a5130a52d3..06beba8864 100644 --- a/qa/L1_pytorch_mcore_integration/test.sh +++ b/qa/L1_pytorch_mcore_integration/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L1_pytorch_onnx_unittest/test.sh b/qa/L1_pytorch_onnx_unittest/test.sh index 303c5c281a..b3a520e129 100644 --- a/qa/L1_pytorch_onnx_unittest/test.sh +++ b/qa/L1_pytorch_onnx_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L1_pytorch_thunder_integration/test.sh b/qa/L1_pytorch_thunder_integration/test.sh index edf3f2eb84..8c3fdc8cdb 100644 --- a/qa/L1_pytorch_thunder_integration/test.sh +++ b/qa/L1_pytorch_thunder_integration/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/L2_jax_distributed_unittest/test.sh b/qa/L2_jax_distributed_unittest/test.sh index 347ff35548..04fbdf1643 100644 --- a/qa/L2_jax_distributed_unittest/test.sh +++ b/qa/L2_jax_distributed_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index d1eaaeb863..5822675663 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. export TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index e2d771cfd0..6e239bfb72 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/qa/format.sh b/qa/format.sh index 86fd8f1981..99608ed02a 100644 --- a/qa/format.sh +++ b/qa/format.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/setup.py b/setup.py index ce3805d2eb..18bb736f24 100644 --- a/setup.py +++ b/setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index c2c9d0d915..6f4f163f08 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index b2f14b1892..26efb37962 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/cpp/operator/test_act.cu b/tests/cpp/operator/test_act.cu index 32b068de58..b4280818a8 100644 --- a/tests/cpp/operator/test_act.cu +++ b/tests/cpp/operator/test_act.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast.cu b/tests/cpp/operator/test_cast.cu index 81c975b0a8..35d9dd2efd 100644 --- a/tests/cpp/operator/test_cast.cu +++ b/tests/cpp/operator/test_cast.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_current_scaling.cu b/tests/cpp/operator/test_cast_current_scaling.cu index 18325d6daf..4dd6cd2d58 100644 --- a/tests/cpp/operator/test_cast_current_scaling.cu +++ b/tests/cpp/operator/test_cast_current_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_dbias.cu b/tests/cpp/operator/test_cast_dbias.cu index 0f8ff2b6a3..18f07153c6 100644 --- a/tests/cpp/operator/test_cast_dbias.cu +++ b/tests/cpp/operator/test_cast_dbias.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_dbias_dgelu.cu b/tests/cpp/operator/test_cast_dbias_dgelu.cu index 572b4a02ad..8213e5665a 100644 --- a/tests/cpp/operator/test_cast_dbias_dgelu.cu +++ b/tests/cpp/operator/test_cast_dbias_dgelu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_float8blockwise.cu b/tests/cpp/operator/test_cast_float8blockwise.cu index fe4ae2d264..8e9da91d08 100644 --- a/tests/cpp/operator/test_cast_float8blockwise.cu +++ b/tests/cpp/operator/test_cast_float8blockwise.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_gated_swiglu.cu b/tests/cpp/operator/test_cast_gated_swiglu.cu index 35ae462106..298b978f2a 100644 --- a/tests/cpp/operator/test_cast_gated_swiglu.cu +++ b/tests/cpp/operator/test_cast_gated_swiglu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index 3800921446..b5e11c30e1 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_mxfp8_gated_swiglu.cu b/tests/cpp/operator/test_cast_mxfp8_gated_swiglu.cu index 512ee7e810..3ff0e8ae99 100644 --- a/tests/cpp/operator/test_cast_mxfp8_gated_swiglu.cu +++ b/tests/cpp/operator/test_cast_mxfp8_gated_swiglu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index afd7927da2..1904d03df7 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_transpose.cu b/tests/cpp/operator/test_cast_transpose.cu index 863570cd3d..44c78e4a09 100644 --- a/tests/cpp/operator/test_cast_transpose.cu +++ b/tests/cpp/operator/test_cast_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_transpose_current_scaling.cu b/tests/cpp/operator/test_cast_transpose_current_scaling.cu index e78137ca41..225d24317a 100644 --- a/tests/cpp/operator/test_cast_transpose_current_scaling.cu +++ b/tests/cpp/operator/test_cast_transpose_current_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_transpose_dbias.cu b/tests/cpp/operator/test_cast_transpose_dbias.cu index 0368bcf1a4..5b06b28327 100644 --- a/tests/cpp/operator/test_cast_transpose_dbias.cu +++ b/tests/cpp/operator/test_cast_transpose_dbias.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu b/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu index 15744fbeea..9a4a2fa080 100644 --- a/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu +++ b/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_cast_transpose_dgeglu.cu b/tests/cpp/operator/test_cast_transpose_dgeglu.cu index 0e75c41e62..a87c0c5a42 100644 --- a/tests/cpp/operator/test_cast_transpose_dgeglu.cu +++ b/tests/cpp/operator/test_cast_transpose_dgeglu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_causal_softmax.cu b/tests/cpp/operator/test_causal_softmax.cu index ab64ed5642..8ae63a81e1 100644 --- a/tests/cpp/operator/test_causal_softmax.cu +++ b/tests/cpp/operator/test_causal_softmax.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_dequantize_mxfp8.cu b/tests/cpp/operator/test_dequantize_mxfp8.cu index a7a993f1fa..a529f93d7c 100644 --- a/tests/cpp/operator/test_dequantize_mxfp8.cu +++ b/tests/cpp/operator/test_dequantize_mxfp8.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_memset.cu b/tests/cpp/operator/test_memset.cu index c6a9ac13be..00f9e7614c 100644 --- a/tests/cpp/operator/test_memset.cu +++ b/tests/cpp/operator/test_memset.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_multi_cast_transpose.cu b/tests/cpp/operator/test_multi_cast_transpose.cu index 0bbca55375..2bb35c4b89 100644 --- a/tests/cpp/operator/test_multi_cast_transpose.cu +++ b/tests/cpp/operator/test_multi_cast_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_multi_padding.cu b/tests/cpp/operator/test_multi_padding.cu index 742672d8f7..3ac48ff214 100644 --- a/tests/cpp/operator/test_multi_padding.cu +++ b/tests/cpp/operator/test_multi_padding.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_multi_unpadding.cu b/tests/cpp/operator/test_multi_unpadding.cu index ca685b9628..98ded4f636 100644 --- a/tests/cpp/operator/test_multi_unpadding.cu +++ b/tests/cpp/operator/test_multi_unpadding.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_normalization.cu b/tests/cpp/operator/test_normalization.cu index 20ad38ca24..db5d6be773 100644 --- a/tests/cpp/operator/test_normalization.cu +++ b/tests/cpp/operator/test_normalization.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_normalization.h b/tests/cpp/operator/test_normalization.h index 271345686e..16b4929741 100644 --- a/tests/cpp/operator/test_normalization.h +++ b/tests/cpp/operator/test_normalization.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_normalization_mxfp8.cu b/tests/cpp/operator/test_normalization_mxfp8.cu index 08d70eb724..10b33f8e2a 100644 --- a/tests/cpp/operator/test_normalization_mxfp8.cu +++ b/tests/cpp/operator/test_normalization_mxfp8.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_qdq.cu b/tests/cpp/operator/test_qdq.cu index 68b183f109..4e364fffa4 100644 --- a/tests/cpp/operator/test_qdq.cu +++ b/tests/cpp/operator/test_qdq.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_swap_first_dims.cu b/tests/cpp/operator/test_swap_first_dims.cu index 4c2cf415ff..7234f555ba 100644 --- a/tests/cpp/operator/test_swap_first_dims.cu +++ b/tests/cpp/operator/test_swap_first_dims.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index f6e0da057a..1660ff4e7f 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/operator/test_transpose.cu b/tests/cpp/operator/test_transpose.cu index c372cddd47..9c233adc4a 100644 --- a/tests/cpp/operator/test_transpose.cu +++ b/tests/cpp/operator/test_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index d70eb13536..ed961bfe96 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index b8993dfb62..42178fec40 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/util/CMakeLists.txt b/tests/cpp/util/CMakeLists.txt index 7540687089..6d70b7b84f 100644 --- a/tests/cpp/util/CMakeLists.txt +++ b/tests/cpp/util/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/cpp/util/test_nvrtc.cpp b/tests/cpp/util/test_nvrtc.cpp index e885140ce1..d41084449e 100644 --- a/tests/cpp/util/test_nvrtc.cpp +++ b/tests/cpp/util/test_nvrtc.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp/util/test_string.cpp b/tests/cpp/util/test_string.cpp index 6a9fe0d9a5..59631c0453 100644 --- a/tests/cpp/util/test_string.cpp +++ b/tests/cpp/util/test_string.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index ed3ddeb885..0d7258a81d 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index 0a20aa1cea..cdd6f9cf14 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py index cb5676d514..6b7520d147 100644 --- a/tests/jax/conftest.py +++ b/tests/jax/conftest.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """conftest for tests/jax""" diff --git a/tests/jax/distributed_test_base.py b/tests/jax/distributed_test_base.py index 137fa480dd..aa7a8fb8d5 100644 --- a/tests/jax/distributed_test_base.py +++ b/tests/jax/distributed_test_base.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import operator diff --git a/tests/jax/multi_process_launch.sh b/tests/jax/multi_process_launch.sh index d430e0f413..3cdca7f396 100644 --- a/tests/jax/multi_process_launch.sh +++ b/tests/jax/multi_process_launch.sh @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/pytest.ini b/tests/jax/pytest.ini index 70d4188c5f..490671a631 100644 --- a/tests/jax/pytest.ini +++ b/tests/jax/pytest.ini @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index c8bd9d47c3..80fcc68843 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_distributed_dense.py b/tests/jax/test_distributed_dense.py index 15b1463437..b8caf188d4 100644 --- a/tests/jax/test_distributed_dense.py +++ b/tests/jax/test_distributed_dense.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index 6b7e04f124..d0018543d1 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_distributed_helper.py b/tests/jax/test_distributed_helper.py index c9647c13cb..ef19ec9fe7 100644 --- a/tests/jax/test_distributed_helper.py +++ b/tests/jax/test_distributed_helper.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_distributed_layernorm.py b/tests/jax/test_distributed_layernorm.py index d551b73905..e9a2fa49e2 100644 --- a/tests/jax/test_distributed_layernorm.py +++ b/tests/jax/test_distributed_layernorm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_distributed_layernorm_mlp.py b/tests/jax/test_distributed_layernorm_mlp.py index 667840da2c..d214597cb3 100644 --- a/tests/jax/test_distributed_layernorm_mlp.py +++ b/tests/jax/test_distributed_layernorm_mlp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import re diff --git a/tests/jax/test_distributed_softmax.py b/tests/jax/test_distributed_softmax.py index 8cdd4c3f59..0665baa4e3 100644 --- a/tests/jax/test_distributed_softmax.py +++ b/tests/jax/test_distributed_softmax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_functions.py b/tests/jax/test_functions.py index 48a2fb4f88..6d250a481b 100644 --- a/tests/jax/test_functions.py +++ b/tests/jax/test_functions.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index f7267af5b8..ac1b7c3505 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Tests for fused attention""" diff --git a/tests/jax/test_layer.py b/tests/jax/test_layer.py index ca07e15742..8c16d162ed 100644 --- a/tests/jax/test_layer.py +++ b/tests/jax/test_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Test transformer_engine.jax.flax.TransformerLayer""" diff --git a/tests/jax/test_misc.py b/tests/jax/test_misc.py index 6db492921d..20cb271db9 100644 --- a/tests/jax/test_misc.py +++ b/tests/jax/test_misc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_multi_process_distributed_grouped_gemm.py b/tests/jax/test_multi_process_distributed_grouped_gemm.py index 31209d1bc9..94fed0859f 100644 --- a/tests/jax/test_multi_process_distributed_grouped_gemm.py +++ b/tests/jax/test_multi_process_distributed_grouped_gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index 9d1bcc820f..43f2553eed 100644 --- a/tests/jax/test_permutation.py +++ b/tests/jax/test_permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_recipe_characteristics.py b/tests/jax/test_recipe_characteristics.py index 5171a6c622..1a265ac2f3 100644 --- a/tests/jax/test_recipe_characteristics.py +++ b/tests/jax/test_recipe_characteristics.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_sanity_import.py b/tests/jax/test_sanity_import.py index 5e1bca2c9c..15ca7761c7 100644 --- a/tests/jax/test_sanity_import.py +++ b/tests/jax/test_sanity_import.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/jax/test_softmax.py b/tests/jax/test_softmax.py index 9dd03ea0fd..7af9613538 100644 --- a/tests/jax/test_softmax.py +++ b/tests/jax/test_softmax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Tests for the softmax primitives""" diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py index 071b8a73a2..6d969de0d3 100644 --- a/tests/jax/test_triton_custom_calls.py +++ b/tests/jax/test_triton_custom_calls.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Tests for Triton-based custom calls in TE JAX.""" diff --git a/tests/jax/utils.py b/tests/jax/utils.py index 7194e387c7..8055792308 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Utility for the TE layer tests""" diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index e58b2da3a8..3efb516b57 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 4aedcff1b8..eb7905bcd5 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import logging diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 2d4fe69e32..9480b8de70 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/attention/test_cp_utils.py b/tests/pytorch/attention/test_cp_utils.py index 0dd5ba601e..e5051aab36 100644 --- a/tests/pytorch/attention/test_cp_utils.py +++ b/tests/pytorch/attention/test_cp_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py index 864276a676..c662252f9e 100644 --- a/tests/pytorch/attention/test_kv_cache.py +++ b/tests/pytorch/attention/test_kv_cache.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/conftest.py b/tests/pytorch/debug/conftest.py index 20edc6aab7..26f601f893 100644 --- a/tests/pytorch/debug/conftest.py +++ b/tests/pytorch/debug/conftest.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import pytest diff --git a/tests/pytorch/debug/run_distributed.py b/tests/pytorch/debug/run_distributed.py index 358841943a..285ec7ba0c 100644 --- a/tests/pytorch/debug/run_distributed.py +++ b/tests/pytorch/debug/run_distributed.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/test_api_features.py b/tests/pytorch/debug/test_api_features.py index fbf619d481..5387634cb3 100644 --- a/tests/pytorch/debug/test_api_features.py +++ b/tests/pytorch/debug/test_api_features.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/test_config.py b/tests/pytorch/debug/test_config.py index 9b6bcd1cd5..bb1ea52fb5 100644 --- a/tests/pytorch/debug/test_config.py +++ b/tests/pytorch/debug/test_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import pathlib diff --git a/tests/pytorch/debug/test_distributed.py b/tests/pytorch/debug/test_distributed.py index ab5b60a139..a8debadae9 100644 --- a/tests/pytorch/debug/test_distributed.py +++ b/tests/pytorch/debug/test_distributed.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/test_log.py b/tests/pytorch/debug/test_log.py index 5456ab820b..7edc0cc90b 100644 --- a/tests/pytorch/debug/test_log.py +++ b/tests/pytorch/debug/test_log.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/test_numerics.py b/tests/pytorch/debug/test_numerics.py index ed8cdc1773..ab9a2d054a 100644 --- a/tests/pytorch/debug/test_numerics.py +++ b/tests/pytorch/debug/test_numerics.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/test_perf.py b/tests/pytorch/debug/test_perf.py index c8c9ae3c1f..0523492310 100644 --- a/tests/pytorch/debug/test_perf.py +++ b/tests/pytorch/debug/test_perf.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/test_sanity.py b/tests/pytorch/debug/test_sanity.py index 97be3003d5..aee5474e76 100644 --- a/tests/pytorch/debug/test_sanity.py +++ b/tests/pytorch/debug/test_sanity.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/debug/utils.py b/tests/pytorch/debug/utils.py index f03ee56b5f..cfa62483b7 100644 --- a/tests/pytorch/debug/utils.py +++ b/tests/pytorch/debug/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/run_fsdp2_model.py index c343299242..5df3468861 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/run_fsdp2_model.py @@ -1,6 +1,6 @@ #!/usr/bin/python3 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/run_gemm_with_overlap.py b/tests/pytorch/distributed/run_gemm_with_overlap.py index 073fa08117..96a7e43231 100644 --- a/tests/pytorch/distributed/run_gemm_with_overlap.py +++ b/tests/pytorch/distributed/run_gemm_with_overlap.py @@ -1,6 +1,6 @@ #!/usr/bin/python3 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py index b2bd6dd773..53c7a5e7cc 100644 --- a/tests/pytorch/distributed/run_layer_with_overlap.py +++ b/tests/pytorch/distributed/run_layer_with_overlap.py @@ -1,6 +1,6 @@ #!/usr/bin/python3 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py index 6cad80fde7..8e24e636e8 100644 --- a/tests/pytorch/distributed/run_numerics.py +++ b/tests/pytorch/distributed/run_numerics.py @@ -1,6 +1,6 @@ #!/usr/bin/python3 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index 3605b3c708..0f3d2cbbf0 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -1,6 +1,6 @@ #!/usr/bin/python3 -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 51b920eab5..8a434b2148 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 3f4848e105..95bf5aa05a 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import os diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py index 5844d81097..c484038938 100644 --- a/tests/pytorch/distributed/test_fusible_ops.py +++ b/tests/pytorch/distributed/test_fusible_ops.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index 61c813b8f2..603433e0da 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_numerics.py b/tests/pytorch/distributed/test_numerics.py index 05b3e54280..491678de14 100644 --- a/tests/pytorch/distributed/test_numerics.py +++ b/tests/pytorch/distributed/test_numerics.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_numerics_exact.py b/tests/pytorch/distributed/test_numerics_exact.py index 72aa786646..b63fea5d2f 100644 --- a/tests/pytorch/distributed/test_numerics_exact.py +++ b/tests/pytorch/distributed/test_numerics_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py index f7c0e1fe88..2e7a63e0a2 100644 --- a/tests/pytorch/distributed/test_sanity.py +++ b/tests/pytorch/distributed/test_sanity.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 91d6fc6ed1..e328e57758 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py index 8ec8a29d80..306d0627f5 100644 --- a/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py +++ b/tests/pytorch/layernorm_mlp/test_selective_activation_checkpoint.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 9f860551d0..911b7660dc 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index a29dcb4279..01a4a01205 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index 0292063ab9..a96fea3af0 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 2467c7e2e1..80ccb2f23d 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index 904dfc2eab..98be9a4f54 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py index 1407c1d8bf..b14eeb815b 100755 --- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/references/blockwise_fp8_gemm_reference.py b/tests/pytorch/references/blockwise_fp8_gemm_reference.py index 5aef986e37..c98277734f 100644 --- a/tests/pytorch/references/blockwise_fp8_gemm_reference.py +++ b/tests/pytorch/references/blockwise_fp8_gemm_reference.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/references/blockwise_quantizer_reference.py b/tests/pytorch/references/blockwise_quantizer_reference.py index 1ce7d3e427..f0bc2ba0fb 100644 --- a/tests/pytorch/references/blockwise_quantizer_reference.py +++ b/tests/pytorch/references/blockwise_quantizer_reference.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/references/quantize_scale_calc.py b/tests/pytorch/references/quantize_scale_calc.py index f36ddca3b2..a6ff425133 100644 --- a/tests/pytorch/references/quantize_scale_calc.py +++ b/tests/pytorch/references/quantize_scale_calc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/references/ref_per_tensor_cs.py b/tests/pytorch/references/ref_per_tensor_cs.py index 5e803f7ed5..c4a6d73d70 100644 --- a/tests/pytorch/references/ref_per_tensor_cs.py +++ b/tests/pytorch/references/ref_per_tensor_cs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_checkpoint.py b/tests/pytorch/test_checkpoint.py index 99a3af0d61..1383264fdc 100644 --- a/tests/pytorch/test_checkpoint.py +++ b/tests/pytorch/test_checkpoint.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index c5b4b48b67..385998a8c5 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_cpu_offloading_v1.py b/tests/pytorch/test_cpu_offloading_v1.py index 8a8e036304..153bceca7d 100644 --- a/tests/pytorch/test_cpu_offloading_v1.py +++ b/tests/pytorch/test_cpu_offloading_v1.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 3ddf33b16a..1b9e11792e 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 64f1c3d159..4de49115b3 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_deferred_init.py b/tests/pytorch/test_deferred_init.py index 4ce522495a..58c8543485 100644 --- a/tests/pytorch/test_deferred_init.py +++ b/tests/pytorch/test_deferred_init.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_float8_blockwise_gemm_exact.py b/tests/pytorch/test_float8_blockwise_gemm_exact.py index 9ae8a60699..a33f860422 100644 --- a/tests/pytorch/test_float8_blockwise_gemm_exact.py +++ b/tests/pytorch/test_float8_blockwise_gemm_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_float8_blockwise_scaling_exact.py b/tests/pytorch/test_float8_blockwise_scaling_exact.py index 153f0b7e04..ab386c0c2f 100644 --- a/tests/pytorch/test_float8_blockwise_scaling_exact.py +++ b/tests/pytorch/test_float8_blockwise_scaling_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index fd47b66c7e..99ab9c4984 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_float8blockwisetensor.py b/tests/pytorch/test_float8blockwisetensor.py index c59f8d8c6a..fe6db9aa41 100644 --- a/tests/pytorch/test_float8blockwisetensor.py +++ b/tests/pytorch/test_float8blockwisetensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index efef64a1e6..f70be45918 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_fused_rope.py b/tests/pytorch/test_fused_rope.py index 9e4ddbdad1..50624df9e0 100644 --- a/tests/pytorch/test_fused_rope.py +++ b/tests/pytorch/test_fused_rope.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. from typing import Callable, Tuple, Union, List diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index fa134ba4bd..f559362d82 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. import torch diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 735cc9b953..ce15dd1421 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_gqa.py b/tests/pytorch/test_gqa.py index 3ef4806182..1ad71c73d6 100644 --- a/tests/pytorch/test_gqa.py +++ b/tests/pytorch/test_gqa.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_hf_integration.py b/tests/pytorch/test_hf_integration.py index b014201c25..e3c8ed0a94 100644 --- a/tests/pytorch/test_hf_integration.py +++ b/tests/pytorch/test_hf_integration.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_jit.py b/tests/pytorch/test_jit.py index e670070bc0..3ec06cd4eb 100644 --- a/tests/pytorch/test_jit.py +++ b/tests/pytorch/test_jit.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_multi_tensor.py b/tests/pytorch/test_multi_tensor.py index 94012354db..b7caa094ae 100644 --- a/tests/pytorch/test_multi_tensor.py +++ b/tests/pytorch/test_multi_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index c30c7bcc76..abe2806e66 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index 2ce6eb82bb..50cd150c4e 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_parallel_cross_entropy.py b/tests/pytorch/test_parallel_cross_entropy.py index 7028dc993a..7b92672af7 100644 --- a/tests/pytorch/test_parallel_cross_entropy.py +++ b/tests/pytorch/test_parallel_cross_entropy.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_partial_cast.py b/tests/pytorch/test_partial_cast.py index cb0c4d75bd..bbb18503b1 100644 --- a/tests/pytorch/test_partial_cast.py +++ b/tests/pytorch/test_partial_cast.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_permutation.py b/tests/pytorch/test_permutation.py index 9a0cf6fb7c..be1ff30472 100644 --- a/tests/pytorch/test_permutation.py +++ b/tests/pytorch/test_permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_qk_norm.py b/tests/pytorch/test_qk_norm.py index d45ec283cc..873bd91863 100644 --- a/tests/pytorch/test_qk_norm.py +++ b/tests/pytorch/test_qk_norm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 0353944ed6..b2e8fca7cb 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index ea26f0b108..91d4b89013 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index fc4a5d6515..e9d24c1a8e 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/test_sanity_import.py b/tests/pytorch/test_sanity_import.py index 5657cf0d85..68136fae2f 100644 --- a/tests/pytorch/test_sanity_import.py +++ b/tests/pytorch/test_sanity_import.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index bdf469c59a..ca5fbc997a 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index e51f03e3d8..0175f04e2e 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 79948e28f7..5975efedaf 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 2d7932d5aa..02388d2e70 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/activation/activation_template.h b/transformer_engine/common/activation/activation_template.h index 7353c3e1d5..ffbffafd1a 100644 --- a/transformer_engine/common/activation/activation_template.h +++ b/transformer_engine/common/activation/activation_template.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index 4979023ef1..675341f7db 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/activation/relu.cu b/transformer_engine/common/activation/relu.cu index c0ef9fd65a..fd70e38c1a 100644 --- a/transformer_engine/common/activation/relu.cu +++ b/transformer_engine/common/activation/relu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/activation/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index 6957a91e61..cc812a17fa 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 73467d7275..de1a8864da 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh index b750142f5b..a5c8327cdf 100644 --- a/transformer_engine/common/cast/core/common.cuh +++ b/transformer_engine/common/cast/core/common.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index 9138fda040..81304981d3 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index f08b09317e..540912c5af 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 6d4454402c..8453b9a68b 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/fp8/dequantize_fp8.cuh b/transformer_engine/common/cast/fp8/dequantize_fp8.cuh index 2514758b5a..6a0eaf94fb 100644 --- a/transformer_engine/common/cast/fp8/dequantize_fp8.cuh +++ b/transformer_engine/common/cast/fp8/dequantize_fp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh index 225ef93ed9..6123d7130b 100644 --- a/transformer_engine/common/cast/fp8/gated_fp8.cuh +++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/fp8/quantize_fp8.cuh b/transformer_engine/common/cast/fp8/quantize_fp8.cuh index efc5015b75..96a42b494d 100644 --- a/transformer_engine/common/cast/fp8/quantize_fp8.cuh +++ b/transformer_engine/common/cast/fp8/quantize_fp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index c56ebe172c..ecdbb5c657 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 4f0e1b80f7..3f5c44120e 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index cbb46f3f28..d6aae78d25 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index 4a39e54a35..dd1b4fa40e 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh b/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh index 5073de5b11..5e68b3760c 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/state_counter.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh b/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh index 63fc1d6bd9..dc2d650e7c 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/swizzle.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index cff8464903..bdbe5cddc3 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 5307cad37f..38677a7075 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 28b47e32d2..1ceb08a9d0 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh index 83ad8fd40b..b4bccf2397 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 7322bf2655..455074e325 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 76f46298db..66a3da55dd 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 56369db27f..2a3c64e8dd 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.cc b/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.cc index 71ea00de3a..c26d0d1be0 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.cc +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.cc @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.h b/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.h index aa6021a190..985bc383b8 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.h +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/ipcsocket.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 6c7bed55ac..9c597be306 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu index 1dcd54d0d7..3d8848d95a 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h index 4d52fbb644..c8d7c87313 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/common.cu b/transformer_engine/common/common.cu index 666f57188d..0ec40dc01c 100644 --- a/transformer_engine/common/common.cu +++ b/transformer_engine/common/common.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 0e264eaae3..0bc9536844 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cudnn_utils.cpp b/transformer_engine/common/cudnn_utils.cpp index eaf6de680a..05ee35ccc7 100644 --- a/transformer_engine/common/cudnn_utils.cpp +++ b/transformer_engine/common/cudnn_utils.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/cudnn_utils.h b/transformer_engine/common/cudnn_utils.h index 0016ad7f55..0777d1e03d 100644 --- a/transformer_engine/common/cudnn_utils.h +++ b/transformer_engine/common/cudnn_utils.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/dropout/dropout.cu b/transformer_engine/common/dropout/dropout.cu index bab349161e..b20b76bbf6 100644 --- a/transformer_engine/common/dropout/dropout.cu +++ b/transformer_engine/common/dropout/dropout.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/context_parallel.cu b/transformer_engine/common/fused_attn/context_parallel.cu index 5921d97d52..cf1fffd94f 100644 --- a/transformer_engine/common/fused_attn/context_parallel.cu +++ b/transformer_engine/common/fused_attn/context_parallel.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/flash_attn.cu b/transformer_engine/common/fused_attn/flash_attn.cu index 59207d59a5..6c66746e62 100644 --- a/transformer_engine/common/fused_attn/flash_attn.cu +++ b/transformer_engine/common/fused_attn/flash_attn.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 611beb7b84..fde0d38921 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 efa4c78439..d3746fc042 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 872b798bb4..c34eae4e6e 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 1028df6452..336e3d5386 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 57b7afcf43..3b30c6e716 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 5d806290a9..3630041ccf 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index c2efa25829..a1a932fdf5 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/kv_cache.cu b/transformer_engine/common/fused_attn/kv_cache.cu index 3b78cab239..52b46a9774 100644 --- a/transformer_engine/common/fused_attn/kv_cache.cu +++ b/transformer_engine/common/fused_attn/kv_cache.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index df1eae0dd7..727aac447b 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 72047a73f2..7d23bb5c55 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_rope/fused_rope.cu b/transformer_engine/common/fused_rope/fused_rope.cu index 597a5d3c29..27dc11ab43 100644 --- a/transformer_engine/common/fused_rope/fused_rope.cu +++ b/transformer_engine/common/fused_rope/fused_rope.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 94082594f6..2aa2805fed 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 03d22942b5..7540b5c41d 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 03e972332a..2719c68c97 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_router/utils.h b/transformer_engine/common/fused_router/utils.h index b6f9d87bdc..4ae0b467b5 100644 --- a/transformer_engine/common/fused_router/utils.h +++ b/transformer_engine/common/fused_router/utils.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu index bbe722a8f5..6ea8017e07 100644 --- a/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu index 79318cd28b..27f86673c5 100644 --- a/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_masked_softmax.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu index 03cdd68279..431148cd1d 100644 --- a/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu +++ b/transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/gemm/config.cpp b/transformer_engine/common/gemm/config.cpp index cf211beaf9..ec2381dc1e 100644 --- a/transformer_engine/common/gemm/config.cpp +++ b/transformer_engine/common/gemm/config.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/gemm/config.h b/transformer_engine/common/gemm/config.h index 54ccf06a53..c7b528bf4b 100644 --- a/transformer_engine/common/gemm/config.h +++ b/transformer_engine/common/gemm/config.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 118bf19335..4b7d8179b0 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/gemm/cutlass_grouped_gemm.cu b/transformer_engine/common/gemm/cutlass_grouped_gemm.cu index 18736c4f54..ef720d1984 100644 --- a/transformer_engine/common/gemm/cutlass_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cutlass_grouped_gemm.cu @@ -1,5 +1,5 @@ /*************************************************************************************************** - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. **************************************************************************************************/ diff --git a/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh b/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh index 1add571325..eb99edc4d3 100644 --- a/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh +++ b/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh @@ -1,5 +1,5 @@ /*************************************************************************************************** - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. **************************************************************************************************/ diff --git a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh index b6f6799a49..bc46341e88 100644 --- a/transformer_engine/common/hadamard_transform/customized_pipeline.cuh +++ b/transformer_engine/common/hadamard_transform/customized_pipeline.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index ea5e22bbfb..5d45996dc8 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 6e071ec79f..85bb98f0f1 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 3932b328ae..c1fb87d048 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform.cu b/transformer_engine/common/hadamard_transform/hadamard_transform.cu index c01ce7b78f..de930aa2cb 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 11325041ae..a839be1701 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh b/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh index ad3bbf5cd7..f86061abb0 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_utils.cuh @@ -1,5 +1,5 @@ /************************************************************************* -* Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +* Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 4e48088586..55cd44d9de 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 19fbe431aa..5e67b7645e 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h b/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h index 649b5ced50..5bef067910 100644 --- a/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h +++ b/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm.h b/transformer_engine/common/include/transformer_engine/comm_gemm.h index 14cf56a002..06b56789a3 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm.h +++ b/transformer_engine/common/include/transformer_engine/comm_gemm.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h index cffc411a0d..6307eab14c 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h +++ b/transformer_engine/common/include/transformer_engine/comm_gemm_overlap.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/cudnn.h b/transformer_engine/common/include/transformer_engine/cudnn.h index 70acead631..ce44f87d9e 100644 --- a/transformer_engine/common/include/transformer_engine/cudnn.h +++ b/transformer_engine/common/include/transformer_engine/cudnn.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/dropout.h b/transformer_engine/common/include/transformer_engine/dropout.h index 6ba1ab9126..57866abcdb 100644 --- a/transformer_engine/common/include/transformer_engine/dropout.h +++ b/transformer_engine/common/include/transformer_engine/dropout.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 6622019280..16b4b8ff4d 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/fused_rope.h b/transformer_engine/common/include/transformer_engine/fused_rope.h index 19047f463b..aea5256a2c 100644 --- a/transformer_engine/common/include/transformer_engine/fused_rope.h +++ b/transformer_engine/common/include/transformer_engine/fused_rope.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 8cf4b222a5..1f026a703d 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 950014cc9b..38214f8872 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index 112cb9b54d..b6e9719aad 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/multi_stream.h b/transformer_engine/common/include/transformer_engine/multi_stream.h index e406a07867..013a424941 100644 --- a/transformer_engine/common/include/transformer_engine/multi_stream.h +++ b/transformer_engine/common/include/transformer_engine/multi_stream.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 03d35dc2ed..1bea4cb21f 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/normalization.h b/transformer_engine/common/include/transformer_engine/normalization.h index 651ae87b4c..7f5bd92fc1 100644 --- a/transformer_engine/common/include/transformer_engine/normalization.h +++ b/transformer_engine/common/include/transformer_engine/normalization.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/padding.h b/transformer_engine/common/include/transformer_engine/padding.h index 0783fc2b21..13775b65a5 100644 --- a/transformer_engine/common/include/transformer_engine/padding.h +++ b/transformer_engine/common/include/transformer_engine/padding.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/permutation.h b/transformer_engine/common/include/transformer_engine/permutation.h index 570eb02fb1..1fb1963512 100644 --- a/transformer_engine/common/include/transformer_engine/permutation.h +++ b/transformer_engine/common/include/transformer_engine/permutation.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index b1773a8db3..83d436db30 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/softmax.h b/transformer_engine/common/include/transformer_engine/softmax.h index 9f1c423172..e8883017a0 100644 --- a/transformer_engine/common/include/transformer_engine/softmax.h +++ b/transformer_engine/common/include/transformer_engine/softmax.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 624e71d1e3..b4489abdd8 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 19cb646be2..7fc9d78980 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/include/transformer_engine/transpose.h b/transformer_engine/common/include/transformer_engine/transpose.h index cc069ee3ec..cd73935abb 100644 --- a/transformer_engine/common/include/transformer_engine/transpose.h +++ b/transformer_engine/common/include/transformer_engine/transpose.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/multi_tensor/adam.cu b/transformer_engine/common/multi_tensor/adam.cu index 9dec2c178a..5d89179c44 100644 --- a/transformer_engine/common/multi_tensor/adam.cu +++ b/transformer_engine/common/multi_tensor/adam.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/multi_tensor/compute_scale.cu b/transformer_engine/common/multi_tensor/compute_scale.cu index 0ac9ab7371..66871ccfa4 100644 --- a/transformer_engine/common/multi_tensor/compute_scale.cu +++ b/transformer_engine/common/multi_tensor/compute_scale.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/multi_tensor/l2norm.cu b/transformer_engine/common/multi_tensor/l2norm.cu index cc66562af5..8a7f265d40 100644 --- a/transformer_engine/common/multi_tensor/l2norm.cu +++ b/transformer_engine/common/multi_tensor/l2norm.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh index b78612181b..3062ead551 100644 --- a/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh +++ b/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/multi_tensor/scale.cu b/transformer_engine/common/multi_tensor/scale.cu index ac457adb06..b3266200c4 100644 --- a/transformer_engine/common/multi_tensor/scale.cu +++ b/transformer_engine/common/multi_tensor/scale.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/multi_tensor/sgd.cu b/transformer_engine/common/multi_tensor/sgd.cu index 9235de3304..0159581d32 100644 --- a/transformer_engine/common/multi_tensor/sgd.cu +++ b/transformer_engine/common/multi_tensor/sgd.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/common.cpp b/transformer_engine/common/normalization/common.cpp index 70e814e806..852b418b39 100644 --- a/transformer_engine/common/normalization/common.cpp +++ b/transformer_engine/common/normalization/common.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index 37144052a9..79de2ac140 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/kernel_traits.h b/transformer_engine/common/normalization/kernel_traits.h index 78d9212de6..12fc095c38 100644 --- a/transformer_engine/common/normalization/kernel_traits.h +++ b/transformer_engine/common/normalization/kernel_traits.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/layernorm/ln_api.cpp b/transformer_engine/common/normalization/layernorm/ln_api.cpp index b83ae25f25..24664d363b 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh index b68e79cd98..c4b00b87c3 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_kernels.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu index 1eeb08415b..68aa0942c1 100644 --- a/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_bwd_semi_cuda_kernel.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu index 787c75ef8c..464df8d276 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh index 38c4096073..5a37cf46da 100644 --- a/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/layernorm/ln_fwd_kernels.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index ea6c972bf5..137df79bde 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh index 3f3cdd065b..d620ee5260 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu index 9bd56c4ec9..60238f256d 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu index 90b4f13405..5522fd5c6b 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh index 7fed7f123a..900fb58be2 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_fwd_kernels.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/nvshmem_api/CMakeLists.txt b/transformer_engine/common/nvshmem_api/CMakeLists.txt index 67136b1baa..1e72e42b0a 100644 --- a/transformer_engine/common/nvshmem_api/CMakeLists.txt +++ b/transformer_engine/common/nvshmem_api/CMakeLists.txt @@ -1,5 +1,5 @@ ########################################################################## -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. ########################################################################## diff --git a/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu b/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu index d5f6aeecce..efa7d0d53a 100644 --- a/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu +++ b/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.h b/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.h index c878e97af5..1f757bc270 100644 --- a/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.h +++ b/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/nvtx.h b/transformer_engine/common/nvtx.h index ada7a59092..f3ff10cf06 100644 --- a/transformer_engine/common/nvtx.h +++ b/transformer_engine/common/nvtx.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/permutation/permutation.cu b/transformer_engine/common/permutation/permutation.cu index d66298b692..fbba27941c 100644 --- a/transformer_engine/common/permutation/permutation.cu +++ b/transformer_engine/common/permutation/permutation.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 98e2a29df8..64ee2a5a16 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/recipe/current_scaling.cu b/transformer_engine/common/recipe/current_scaling.cu index ee2c845159..15ec1621bc 100644 --- a/transformer_engine/common/recipe/current_scaling.cu +++ b/transformer_engine/common/recipe/current_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/recipe/delayed_scaling.cu b/transformer_engine/common/recipe/delayed_scaling.cu index e1f9bcf644..a0da551d0f 100644 --- a/transformer_engine/common/recipe/delayed_scaling.cu +++ b/transformer_engine/common/recipe/delayed_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/recipe/fp8_block_scaling.cu b/transformer_engine/common/recipe/fp8_block_scaling.cu index 42a7b8d696..f69fd6c262 100644 --- a/transformer_engine/common/recipe/fp8_block_scaling.cu +++ b/transformer_engine/common/recipe/fp8_block_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index 8a7ecc6b01..534a796324 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 5ebc7ba4f3..682d8b53f5 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/recipe/recipe_common.cuh b/transformer_engine/common/recipe/recipe_common.cuh index 11f9bc1299..07839407a3 100644 --- a/transformer_engine/common/recipe/recipe_common.cuh +++ b/transformer_engine/common/recipe/recipe_common.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 2cb43e8f27..73647d5717 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/swizzle/swizzle_block_scaling.cu b/transformer_engine/common/swizzle/swizzle_block_scaling.cu index 4be85474af..c5ad1aed43 100644 --- a/transformer_engine/common/swizzle/swizzle_block_scaling.cu +++ b/transformer_engine/common/swizzle/swizzle_block_scaling.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 4a140b4376..370d9723cf 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/cast_transpose.cu b/transformer_engine/common/transpose/cast_transpose.cu index 648070c8d1..dd27fa83ee 100644 --- a/transformer_engine/common/transpose/cast_transpose.cu +++ b/transformer_engine/common/transpose/cast_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h index 89266f4bbc..66bf8e86df 100644 --- a/transformer_engine/common/transpose/cast_transpose.h +++ b/transformer_engine/common/transpose/cast_transpose.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/cast_transpose_fusion.cu b/transformer_engine/common/transpose/cast_transpose_fusion.cu index a04ab902ca..77c1322e7d 100644 --- a/transformer_engine/common/transpose/cast_transpose_fusion.cu +++ b/transformer_engine/common/transpose/cast_transpose_fusion.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/multi_cast_transpose.cu b/transformer_engine/common/transpose/multi_cast_transpose.cu index bf38565686..33e1c19d8f 100644 --- a/transformer_engine/common/transpose/multi_cast_transpose.cu +++ b/transformer_engine/common/transpose/multi_cast_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index 661cf339ae..c636a627a4 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu index fcf7a151c3..df869b4331 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ 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 b49a54fbdb..798c712fda 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/rtc/cast_transpose.cu b/transformer_engine/common/transpose/rtc/cast_transpose.cu index 952d70f38b..e40c463a6b 100644 --- a/transformer_engine/common/transpose/rtc/cast_transpose.cu +++ b/transformer_engine/common/transpose/rtc/cast_transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu b/transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu index 34359561aa..49b533ffbd 100644 --- a/transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu +++ b/transformer_engine/common/transpose/rtc/cast_transpose_fusion.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/rtc/swap_first_dims.cu b/transformer_engine/common/transpose/rtc/swap_first_dims.cu index 89a07697a6..0e045fabdb 100644 --- a/transformer_engine/common/transpose/rtc/swap_first_dims.cu +++ b/transformer_engine/common/transpose/rtc/swap_first_dims.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/rtc/transpose.cu b/transformer_engine/common/transpose/rtc/transpose.cu index 6d05c68106..fb4c9feba3 100644 --- a/transformer_engine/common/transpose/rtc/transpose.cu +++ b/transformer_engine/common/transpose/rtc/transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/swap_first_dims.cu b/transformer_engine/common/transpose/swap_first_dims.cu index 08249a8231..33346e5499 100644 --- a/transformer_engine/common/transpose/swap_first_dims.cu +++ b/transformer_engine/common/transpose/swap_first_dims.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/transpose.cu b/transformer_engine/common/transpose/transpose.cu index 9f0acd8071..a280df21a5 100644 --- a/transformer_engine/common/transpose/transpose.cu +++ b/transformer_engine/common/transpose/transpose.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/transpose/transpose_fusion.cu b/transformer_engine/common/transpose/transpose_fusion.cu index 75fa07a5b3..670fe6f92f 100644 --- a/transformer_engine/common/transpose/transpose_fusion.cu +++ b/transformer_engine/common/transpose/transpose_fusion.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/triton/__init__.py b/transformer_engine/common/triton/__init__.py index 76c9b98d0e..dd9011e50d 100644 --- a/transformer_engine/common/triton/__init__.py +++ b/transformer_engine/common/triton/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/triton/cross_entropy.py b/transformer_engine/common/triton/cross_entropy.py index 282b23cda8..bec2620467 100644 --- a/transformer_engine/common/triton/cross_entropy.py +++ b/transformer_engine/common/triton/cross_entropy.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/triton/pad.py b/transformer_engine/common/triton/pad.py index 8f15e7dcba..3c43e0c53a 100644 --- a/transformer_engine/common/triton/pad.py +++ b/transformer_engine/common/triton/pad.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index de30c7c532..e53b2a9455 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/common/util/cuda_driver.cpp b/transformer_engine/common/util/cuda_driver.cpp index 01e3edf57a..1e98528e52 100644 --- a/transformer_engine/common/util/cuda_driver.cpp +++ b/transformer_engine/common/util/cuda_driver.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/cuda_driver.h b/transformer_engine/common/util/cuda_driver.h index 3425e0af35..2715d8e4e4 100644 --- a/transformer_engine/common/util/cuda_driver.h +++ b/transformer_engine/common/util/cuda_driver.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/cuda_nvml.cpp b/transformer_engine/common/util/cuda_nvml.cpp index 0af9cd7411..25e1a53519 100644 --- a/transformer_engine/common/util/cuda_nvml.cpp +++ b/transformer_engine/common/util/cuda_nvml.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/cuda_nvml.h b/transformer_engine/common/util/cuda_nvml.h index 14131a3cdd..ad5a496253 100644 --- a/transformer_engine/common/util/cuda_nvml.h +++ b/transformer_engine/common/util/cuda_nvml.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 2e5ef8b8e1..f99900bac8 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/cuda_runtime.h b/transformer_engine/common/util/cuda_runtime.h index 6b999870dd..c696f6b57a 100644 --- a/transformer_engine/common/util/cuda_runtime.h +++ b/transformer_engine/common/util/cuda_runtime.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/curanddx.hpp b/transformer_engine/common/util/curanddx.hpp index 4d7c90a019..6dd0b57177 100644 --- a/transformer_engine/common/util/curanddx.hpp +++ b/transformer_engine/common/util/curanddx.hpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/handle_manager.h b/transformer_engine/common/util/handle_manager.h index adb2f55587..1a538eaff0 100644 --- a/transformer_engine/common/util/handle_manager.h +++ b/transformer_engine/common/util/handle_manager.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/logging.h b/transformer_engine/common/util/logging.h index c2ce684c4e..c542afa393 100644 --- a/transformer_engine/common/util/logging.h +++ b/transformer_engine/common/util/logging.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/math.h b/transformer_engine/common/util/math.h index 005a600670..05fe2f5398 100644 --- a/transformer_engine/common/util/math.h +++ b/transformer_engine/common/util/math.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/multi_stream.cpp b/transformer_engine/common/util/multi_stream.cpp index 70d7376afa..6b19f36741 100644 --- a/transformer_engine/common/util/multi_stream.cpp +++ b/transformer_engine/common/util/multi_stream.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/multi_stream.h b/transformer_engine/common/util/multi_stream.h index 26f2d19df8..c82af7e744 100644 --- a/transformer_engine/common/util/multi_stream.h +++ b/transformer_engine/common/util/multi_stream.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/padding.cu b/transformer_engine/common/util/padding.cu index 4e569e1674..1859d8a5cb 100644 --- a/transformer_engine/common/util/padding.cu +++ b/transformer_engine/common/util/padding.cu @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 7f296c9e38..c22fb33ffe 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index bce124e705..faa4e36809 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index f6e79c0cee..7925fdceea 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/rtc.h b/transformer_engine/common/util/rtc.h index 7de1e4d55c..65faf7bcc2 100644 --- a/transformer_engine/common/util/rtc.h +++ b/transformer_engine/common/util/rtc.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/shared_lib_wrapper.h b/transformer_engine/common/util/shared_lib_wrapper.h index 3ccc8239b8..e8abe68a2a 100644 --- a/transformer_engine/common/util/shared_lib_wrapper.h +++ b/transformer_engine/common/util/shared_lib_wrapper.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/string.h b/transformer_engine/common/util/string.h index 0064144102..28f825b036 100644 --- a/transformer_engine/common/util/string.h +++ b/transformer_engine/common/util/string.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/string_header.h.in b/transformer_engine/common/util/string_header.h.in index b9fa83a94f..6c373a5718 100644 --- a/transformer_engine/common/util/string_header.h.in +++ b/transformer_engine/common/util/string_header.h.in @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/system.h b/transformer_engine/common/util/system.h index 5636ab5095..90c984a46a 100644 --- a/transformer_engine/common/util/system.h +++ b/transformer_engine/common/util/system.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/util/vectorized_pointwise.h b/transformer_engine/common/util/vectorized_pointwise.h index dd6869e027..0aa2df7d26 100644 --- a/transformer_engine/common/util/vectorized_pointwise.h +++ b/transformer_engine/common/util/vectorized_pointwise.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 2d37e9c85a..26549191a3 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/common/utils.py b/transformer_engine/common/utils.py index a808e1571f..acbb1ca5fb 100644 --- a/transformer_engine/common/utils.py +++ b/transformer_engine/common/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """The utilities for Transformer Engine""" diff --git a/transformer_engine/debug/__init__.py b/transformer_engine/debug/__init__.py index 62f7f41728..446e192d86 100644 --- a/transformer_engine/debug/__init__.py +++ b/transformer_engine/debug/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/__init__.py b/transformer_engine/debug/features/__init__.py index 51a7cc6d1f..3ad59237ae 100644 --- a/transformer_engine/debug/features/__init__.py +++ b/transformer_engine/debug/features/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/_test_dummy_feature.py b/transformer_engine/debug/features/_test_dummy_feature.py index 4dee97b707..f74cd95e9d 100644 --- a/transformer_engine/debug/features/_test_dummy_feature.py +++ b/transformer_engine/debug/features/_test_dummy_feature.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/api.py b/transformer_engine/debug/features/api.py index 94fc6d129c..9c30f87c3b 100644 --- a/transformer_engine/debug/features/api.py +++ b/transformer_engine/debug/features/api.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/disable_fp8_gemm.py b/transformer_engine/debug/features/disable_fp8_gemm.py index ef2cccbe4a..befebb412c 100644 --- a/transformer_engine/debug/features/disable_fp8_gemm.py +++ b/transformer_engine/debug/features/disable_fp8_gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/disable_fp8_layer.py b/transformer_engine/debug/features/disable_fp8_layer.py index c3b0e4cca9..3839e5f2bf 100644 --- a/transformer_engine/debug/features/disable_fp8_layer.py +++ b/transformer_engine/debug/features/disable_fp8_layer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/fake_quant.py b/transformer_engine/debug/features/fake_quant.py index 58c7379b5b..f48b49b725 100644 --- a/transformer_engine/debug/features/fake_quant.py +++ b/transformer_engine/debug/features/fake_quant.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index d09fb10579..ffcc6b1ad4 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/log_tensor_stats.py b/transformer_engine/debug/features/log_tensor_stats.py index ff37e659a0..100fa64481 100644 --- a/transformer_engine/debug/features/log_tensor_stats.py +++ b/transformer_engine/debug/features/log_tensor_stats.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/per_tensor_scaling.py b/transformer_engine/debug/features/per_tensor_scaling.py index dd1f42cf06..a4bab4eaf5 100644 --- a/transformer_engine/debug/features/per_tensor_scaling.py +++ b/transformer_engine/debug/features/per_tensor_scaling.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/utils/__init__.py b/transformer_engine/debug/features/utils/__init__.py index aae2ec4e99..d691c1828c 100644 --- a/transformer_engine/debug/features/utils/__init__.py +++ b/transformer_engine/debug/features/utils/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py index b5b462f5a2..9ce56dd76d 100644 --- a/transformer_engine/debug/features/utils/stats_buffer.py +++ b/transformer_engine/debug/features/utils/stats_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/features/utils/stats_computation.py b/transformer_engine/debug/features/utils/stats_computation.py index 8c480441c5..46a48e2abf 100644 --- a/transformer_engine/debug/features/utils/stats_computation.py +++ b/transformer_engine/debug/features/utils/stats_computation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/pytorch/__init__.py b/transformer_engine/debug/pytorch/__init__.py index 8bdbe287de..731b1f0c1d 100644 --- a/transformer_engine/debug/pytorch/__init__.py +++ b/transformer_engine/debug/pytorch/__init__.py @@ -1,3 +1,3 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index caa5f5a7ee..7a8670f043 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/pytorch/debug_state.py b/transformer_engine/debug/pytorch/debug_state.py index c47e859bb3..63f856c931 100644 --- a/transformer_engine/debug/pytorch/debug_state.py +++ b/transformer_engine/debug/pytorch/debug_state.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/debug/pytorch/utils.py b/transformer_engine/debug/pytorch/utils.py index 18ed3556f1..ef125904a7 100644 --- a/transformer_engine/debug/pytorch/utils.py +++ b/transformer_engine/debug/pytorch/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/__init__.py b/transformer_engine/jax/__init__.py index 6259a7ad84..d0afc1ff25 100644 --- a/transformer_engine/jax/__init__.py +++ b/transformer_engine/jax/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Transformer Engine bindings for JAX. diff --git a/transformer_engine/jax/activation.py b/transformer_engine/jax/activation.py index daa3679c48..b2b90a10c9 100644 --- a/transformer_engine/jax/activation.py +++ b/transformer_engine/jax/activation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Activation functions for Transformer Engine in JAX. diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 09a29f4cb8..21db296c34 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX multi-head attention modules""" diff --git a/transformer_engine/jax/checkpoint_policies.py b/transformer_engine/jax/checkpoint_policies.py index a03db09b9e..7312eefb11 100644 --- a/transformer_engine/jax/checkpoint_policies.py +++ b/transformer_engine/jax/checkpoint_policies.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Checkpoint policies for Transformer Engine in JAX. diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py index c0285e157a..6a2f9b7378 100644 --- a/transformer_engine/jax/cpp_extensions/__init__.py +++ b/transformer_engine/jax/cpp_extensions/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Python interface for c++ extensions""" diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index c5fb85041e..573603ef3a 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE custom ops for activation""" diff --git a/transformer_engine/jax/cpp_extensions/amax.py b/transformer_engine/jax/cpp_extensions/amax.py index afc248a0ad..700ba9061c 100644 --- a/transformer_engine/jax/cpp_extensions/amax.py +++ b/transformer_engine/jax/cpp_extensions/amax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE custom ops for amax calculation""" diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index ef921c2762..0cdfcebf38 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE custom ops for attention""" diff --git a/transformer_engine/jax/cpp_extensions/base.py b/transformer_engine/jax/cpp_extensions/base.py index 61deab5b80..b26e01c0c7 100644 --- a/transformer_engine/jax/cpp_extensions/base.py +++ b/transformer_engine/jax/cpp_extensions/base.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE base custom ops""" diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 76a8b225ba..71f133bfc4 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX te modules""" diff --git a/transformer_engine/jax/cpp_extensions/misc.py b/transformer_engine/jax/cpp_extensions/misc.py index 225d577cd3..3b6d6b6342 100644 --- a/transformer_engine/jax/cpp_extensions/misc.py +++ b/transformer_engine/jax/cpp_extensions/misc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE miscellaneous for custom ops""" diff --git a/transformer_engine/jax/cpp_extensions/normalization.py b/transformer_engine/jax/cpp_extensions/normalization.py index 862780620e..70fdf4c474 100644 --- a/transformer_engine/jax/cpp_extensions/normalization.py +++ b/transformer_engine/jax/cpp_extensions/normalization.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE custom ops for normalization""" diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index b3f24e9337..1fcecb0e96 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE custom ops for quantization""" diff --git a/transformer_engine/jax/cpp_extensions/softmax.py b/transformer_engine/jax/cpp_extensions/softmax.py index 6d8b24b07d..ff30c9bba3 100644 --- a/transformer_engine/jax/cpp_extensions/softmax.py +++ b/transformer_engine/jax/cpp_extensions/softmax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX/TE custom ops for softmax""" diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 75d22fbf53..a83a1e0a80 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index 34ce29ae13..6c5a976344 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/amax.cpp b/transformer_engine/jax/csrc/extensions/amax.cpp index 46f167fcaf..5ffccaffb4 100644 --- a/transformer_engine/jax/csrc/extensions/amax.cpp +++ b/transformer_engine/jax/csrc/extensions/amax.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 79436fb8b7..540aeb8b2d 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp b/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp index 7082bfb035..87a889621e 100644 --- a/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp +++ b/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/cgemm_helper.h b/transformer_engine/jax/csrc/extensions/cgemm_helper.h index 84b2b81540..2b980e7ee4 100644 --- a/transformer_engine/jax/csrc/extensions/cgemm_helper.h +++ b/transformer_engine/jax/csrc/extensions/cgemm_helper.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/cublas.cpp b/transformer_engine/jax/csrc/extensions/cublas.cpp index 0d3397ce84..a9f29b0ffb 100644 --- a/transformer_engine/jax/csrc/extensions/cublas.cpp +++ b/transformer_engine/jax/csrc/extensions/cublas.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/cudnn.cpp b/transformer_engine/jax/csrc/extensions/cudnn.cpp index 48eab30851..92070433a8 100644 --- a/transformer_engine/jax/csrc/extensions/cudnn.cpp +++ b/transformer_engine/jax/csrc/extensions/cudnn.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/ffi.cpp b/transformer_engine/jax/csrc/extensions/ffi.cpp index a0425efda6..6bb2f18234 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.cpp +++ b/transformer_engine/jax/csrc/extensions/ffi.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/ffi.h b/transformer_engine/jax/csrc/extensions/ffi.h index d4f76a011a..f9d327102b 100644 --- a/transformer_engine/jax/csrc/extensions/ffi.h +++ b/transformer_engine/jax/csrc/extensions/ffi.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 6566ff1689..e25a67a401 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/misc.cpp b/transformer_engine/jax/csrc/extensions/misc.cpp index 176115ade9..7e72438e1a 100644 --- a/transformer_engine/jax/csrc/extensions/misc.cpp +++ b/transformer_engine/jax/csrc/extensions/misc.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/misc.h b/transformer_engine/jax/csrc/extensions/misc.h index 21b50c1af4..eb7be0a66a 100644 --- a/transformer_engine/jax/csrc/extensions/misc.h +++ b/transformer_engine/jax/csrc/extensions/misc.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/normalization.cpp b/transformer_engine/jax/csrc/extensions/normalization.cpp index b01e23c128..3361ddf64a 100644 --- a/transformer_engine/jax/csrc/extensions/normalization.cpp +++ b/transformer_engine/jax/csrc/extensions/normalization.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 9784565cc9..a5986404c9 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/quantization.cpp b/transformer_engine/jax/csrc/extensions/quantization.cpp index 1f7db84383..c5a766f7f2 100644 --- a/transformer_engine/jax/csrc/extensions/quantization.cpp +++ b/transformer_engine/jax/csrc/extensions/quantization.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/softmax.cpp b/transformer_engine/jax/csrc/extensions/softmax.cpp index ee3e5b35e8..2fdb8ea678 100644 --- a/transformer_engine/jax/csrc/extensions/softmax.cpp +++ b/transformer_engine/jax/csrc/extensions/softmax.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/utils.cpp b/transformer_engine/jax/csrc/extensions/utils.cpp index 3ba073737c..52ab2edf0f 100644 --- a/transformer_engine/jax/csrc/extensions/utils.cpp +++ b/transformer_engine/jax/csrc/extensions/utils.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/csrc/extensions/utils.h b/transformer_engine/jax/csrc/extensions/utils.h index 37acf6744e..c55c8d86ce 100644 --- a/transformer_engine/jax/csrc/extensions/utils.h +++ b/transformer_engine/jax/csrc/extensions/utils.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index c499b0651e..23d91f7db0 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Dense layer transformation operations for Transformer Engine in JAX. diff --git a/transformer_engine/jax/flax/__init__.py b/transformer_engine/jax/flax/__init__.py index d1a9cb47f8..dd7d2a47ba 100644 --- a/transformer_engine/jax/flax/__init__.py +++ b/transformer_engine/jax/flax/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Transformer Engine bindings for JAX""" diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index dcfb812896..3d82d8f0b4 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 1395976b9f..ad5a60e4c2 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/layernorm.py b/transformer_engine/jax/layernorm.py index 0f5c6aeef6..3f3f3802db 100644 --- a/transformer_engine/jax/layernorm.py +++ b/transformer_engine/jax/layernorm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Layer normalization operations for Transformer Engine in JAX. diff --git a/transformer_engine/jax/layernorm_dense.py b/transformer_engine/jax/layernorm_dense.py index 14726553f2..8c21496ffe 100644 --- a/transformer_engine/jax/layernorm_dense.py +++ b/transformer_engine/jax/layernorm_dense.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Fused Layer normalization and dense layer transformation operations for Transformer Engine in JAX. diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py index 47fed6c3a7..a8de32830b 100644 --- a/transformer_engine/jax/layernorm_mlp.py +++ b/transformer_engine/jax/layernorm_mlp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Multi-layer perceptron (MLP) operations with layer normalization for Transformer Engine in JAX. diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 32de0b1a3c..2e16e674cc 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/pyproject.toml b/transformer_engine/jax/pyproject.toml index ff0e356ed9..d3162ae96d 100755 --- a/transformer_engine/jax/pyproject.toml +++ b/transformer_engine/jax/pyproject.toml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/quantize/__init__.py b/transformer_engine/jax/quantize/__init__.py index 878067a783..4505611a48 100644 --- a/transformer_engine/jax/quantize/__init__.py +++ b/transformer_engine/jax/quantize/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/quantize/dequantizer.py b/transformer_engine/jax/quantize/dequantizer.py index 80ebc6b875..74787b9308 100644 --- a/transformer_engine/jax/quantize/dequantizer.py +++ b/transformer_engine/jax/quantize/dequantizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/quantize/device_utils.py b/transformer_engine/jax/quantize/device_utils.py index 9f5d2f4587..b9f0ee65f3 100644 --- a/transformer_engine/jax/quantize/device_utils.py +++ b/transformer_engine/jax/quantize/device_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/quantize/hadamard.py b/transformer_engine/jax/quantize/hadamard.py index 5f6f0ec2b5..1bad6be101 100644 --- a/transformer_engine/jax/quantize/hadamard.py +++ b/transformer_engine/jax/quantize/hadamard.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Randomized Hadamard Transform (RHT) utilities for JAX.""" diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index 6358edf468..7d81d71bc8 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/quantize/metadata.py b/transformer_engine/jax/quantize/metadata.py index a987643eb7..52367216c4 100644 --- a/transformer_engine/jax/quantize/metadata.py +++ b/transformer_engine/jax/quantize/metadata.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/quantize/misc.py b/transformer_engine/jax/quantize/misc.py index c1e169d005..b7841bfa4e 100644 --- a/transformer_engine/jax/quantize/misc.py +++ b/transformer_engine/jax/quantize/misc.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index 4edc187795..f5ca6aeaed 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/quantize/scaling_modes.py b/transformer_engine/jax/quantize/scaling_modes.py index eea27a35d8..61c3af178c 100644 --- a/transformer_engine/jax/quantize/scaling_modes.py +++ b/transformer_engine/jax/quantize/scaling_modes.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index 90f139c3da..c26cb8a531 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/setup.py b/transformer_engine/jax/setup.py index ccdbcdb529..2d25242825 100644 --- a/transformer_engine/jax/setup.py +++ b/transformer_engine/jax/setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index b4b8c42027..9b13412c14 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Sharding utilities for Transformer Engine in JAX. diff --git a/transformer_engine/jax/softmax.py b/transformer_engine/jax/softmax.py index 24fca6bc71..8302e7ccee 100644 --- a/transformer_engine/jax/softmax.py +++ b/transformer_engine/jax/softmax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """JAX softmax modules""" diff --git a/transformer_engine/jax/triton_extensions/__init__.py b/transformer_engine/jax/triton_extensions/__init__.py index 13a36421bf..c98254d7d3 100644 --- a/transformer_engine/jax/triton_extensions/__init__.py +++ b/transformer_engine/jax/triton_extensions/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index 01b15c5adc..849673fe31 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 41ce15303c..064b2843c6 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 9f4a9678eb..5e1eb6954b 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py index 67afd835d0..c4c2aa3e72 100644 --- a/transformer_engine/pytorch/attention/__init__.py +++ b/transformer_engine/pytorch/attention/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/__init__.py b/transformer_engine/pytorch/attention/dot_product_attention/__init__.py index 112a20d51c..941f94f105 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/__init__.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index c1ff46c75a..c726ed8849 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. 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 1bcff966b7..75b360e485 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. 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 f506035c1e..6e5a12a103 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 @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py index fd799957b4..74d9583ce5 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/softmax.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/softmax.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 10a06ed965..bf19388d7e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py index 4ae1bd09a1..08e50aad8b 100644 --- a/transformer_engine/pytorch/attention/inference.py +++ b/transformer_engine/pytorch/attention/inference.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index beb13b7f1e..f875fd1e0a 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index a32b2d3edb..77ad57ed8f 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index a1fae730c5..3cce4600d9 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/cpp_extensions/__init__.py b/transformer_engine/pytorch/cpp_extensions/__init__.py index 944d1849bf..bb6e921132 100644 --- a/transformer_engine/pytorch/cpp_extensions/__init__.py +++ b/transformer_engine/pytorch/cpp_extensions/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 88c223eb46..e226ef32d4 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index d4ff0b96d9..2a97e2ac71 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 58ed063066..d0b8d3474e 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/cpu_offload_v1.py b/transformer_engine/pytorch/cpu_offload_v1.py index e79e37b019..f92c436941 100644 --- a/transformer_engine/pytorch/cpu_offload_v1.py +++ b/transformer_engine/pytorch/cpu_offload_v1.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/cross_entropy.py b/transformer_engine/pytorch/cross_entropy.py index 30002cdbfd..733b9c10e1 100644 --- a/transformer_engine/pytorch/cross_entropy.py +++ b/transformer_engine/pytorch/cross_entropy.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index e054424dd4..fa6f142b68 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 978bee52dc..1e1e3326c4 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 80479dccf4..52ef02a347 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 14cc084c0c..9ea14e1af0 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp b/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp index d1dcf68c3d..4392fa4b43 100644 --- a/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp +++ b/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 2480d9aba9..b455e03757 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/bias.cpp b/transformer_engine/pytorch/csrc/extensions/bias.cpp index b0435d2723..c59e3c4f64 100644 --- a/transformer_engine/pytorch/csrc/extensions/bias.cpp +++ b/transformer_engine/pytorch/csrc/extensions/bias.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index aa9d800c7b..3bbc99b444 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp b/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp index 38947c5a9d..a126ab0d60 100644 --- a/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp +++ b/transformer_engine/pytorch/csrc/extensions/comm_gemm_overlap.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/dropout.cpp b/transformer_engine/pytorch/csrc/extensions/dropout.cpp index e6f29d0da7..bea8f3a7b5 100644 --- a/transformer_engine/pytorch/csrc/extensions/dropout.cpp +++ b/transformer_engine/pytorch/csrc/extensions/dropout.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp index 3be2ca9396..d6693a485e 100644 --- a/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/fp8_partial_cast.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 335052296f..07ddfbeb6f 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/misc.cpp b/transformer_engine/pytorch/csrc/extensions/misc.cpp index 2c0014a6b8..d667a61d44 100644 --- a/transformer_engine/pytorch/csrc/extensions/misc.cpp +++ b/transformer_engine/pytorch/csrc/extensions/misc.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp index acf04900e5..145e1d4b40 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/adam.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp index e60b001f6f..328970ffa8 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/compute_scale.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/l2norm.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/l2norm.cpp index d33a2520e3..b02cf1fbba 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/l2norm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/l2norm.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp index 2db936f84a..4bb83bfeed 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/sgd.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/sgd.cpp index 2c6a6b7c4c..a70fe12b56 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/sgd.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/sgd.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index 3c5c17fc6f..d7a07724c1 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp b/transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp index 9c31678ee5..ac68727ac8 100644 --- a/transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/nvshmem_comm.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/padding.cpp b/transformer_engine/pytorch/csrc/extensions/padding.cpp index d4b64a485c..6c66fda015 100644 --- a/transformer_engine/pytorch/csrc/extensions/padding.cpp +++ b/transformer_engine/pytorch/csrc/extensions/padding.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/permutation.cpp b/transformer_engine/pytorch/csrc/extensions/permutation.cpp index 97cf400851..226705b169 100644 --- a/transformer_engine/pytorch/csrc/extensions/permutation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/permutation.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index d0f450bc71..e73eca7861 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp index 63c26ee303..c02d2ec616 100644 --- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp +++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 9befe14f88..2ae0d648a1 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/softmax.cpp b/transformer_engine/pytorch/csrc/extensions/softmax.cpp index 2e0e482eb4..3bb6a5e7b3 100644 --- a/transformer_engine/pytorch/csrc/extensions/softmax.cpp +++ b/transformer_engine/pytorch/csrc/extensions/softmax.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 7dfdf99547..477d7c87e7 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/pybind.h b/transformer_engine/pytorch/csrc/pybind.h index 65665d01b6..25ffef0588 100644 --- a/transformer_engine/pytorch/csrc/pybind.h +++ b/transformer_engine/pytorch/csrc/pybind.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index fd748d1b21..a73efc008a 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 368e9dcdfa..16171054cb 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/util.cpp b/transformer_engine/pytorch/csrc/util.cpp index ce547d302e..96fd2ccb3a 100644 --- a/transformer_engine/pytorch/csrc/util.cpp +++ b/transformer_engine/pytorch/csrc/util.cpp @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/util.h b/transformer_engine/pytorch/csrc/util.h index 57eee86d2a..4d72db922e 100644 --- a/transformer_engine/pytorch/csrc/util.h +++ b/transformer_engine/pytorch/csrc/util.h @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ diff --git a/transformer_engine/pytorch/custom_recipes/__init__.py b/transformer_engine/pytorch/custom_recipes/__init__.py index 6e859ba5db..f115ffe743 100644 --- a/transformer_engine/pytorch/custom_recipes/__init__.py +++ b/transformer_engine/pytorch/custom_recipes/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/custom_recipes/gemm.py b/transformer_engine/pytorch/custom_recipes/gemm.py index cc98a8a57a..8f853ff093 100644 --- a/transformer_engine/pytorch/custom_recipes/gemm.py +++ b/transformer_engine/pytorch/custom_recipes/gemm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/custom_recipes/quantization.py b/transformer_engine/pytorch/custom_recipes/quantization.py index 876ca7fcb9..85920f5032 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization.py +++ b/transformer_engine/pytorch/custom_recipes/quantization.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py b/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py index 96cbca772c..5bdc537e4b 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py index b371ca4842..d00d0c8b94 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/custom_recipes/utils.py b/transformer_engine/pytorch/custom_recipes/utils.py index 20dc6f11b0..3e23661f14 100644 --- a/transformer_engine/pytorch/custom_recipes/utils.py +++ b/transformer_engine/pytorch/custom_recipes/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 9f589498a4..5497ee7967 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/export.py b/transformer_engine/pytorch/export.py index a86f8ee58c..89306fbe1e 100644 --- a/transformer_engine/pytorch/export.py +++ b/transformer_engine/pytorch/export.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/float8_tensor.py b/transformer_engine/pytorch/float8_tensor.py index eeafc23c70..45069adeef 100644 --- a/transformer_engine/pytorch/float8_tensor.py +++ b/transformer_engine/pytorch/float8_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/fp8.py b/transformer_engine/pytorch/fp8.py index f937b3de99..6bcf2d53c7 100644 --- a/transformer_engine/pytorch/fp8.py +++ b/transformer_engine/pytorch/fp8.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 1822c47d8b..f587ca9946 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index e9a65a72ff..5884188b7e 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/__init__.py b/transformer_engine/pytorch/module/__init__.py index ac682190c2..3cf15efc11 100644 --- a/transformer_engine/pytorch/module/__init__.py +++ b/transformer_engine/pytorch/module/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 6151ecafd3..88b58a353a 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index ab7cd9ab47..ad5cd04341 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/fp8_padding.py b/transformer_engine/pytorch/module/fp8_padding.py index d59a26ca33..8ac49c9bae 100644 --- a/transformer_engine/pytorch/module/fp8_padding.py +++ b/transformer_engine/pytorch/module/fp8_padding.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/fp8_unpadding.py b/transformer_engine/pytorch/module/fp8_unpadding.py index 6f9702a96f..c5d396837f 100644 --- a/transformer_engine/pytorch/module/fp8_unpadding.py +++ b/transformer_engine/pytorch/module/fp8_unpadding.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index c4d35a9c2c..d0a5618afb 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py index 52802c618c..d4f0a78ba2 100644 --- a/transformer_engine/pytorch/module/layernorm.py +++ b/transformer_engine/pytorch/module/layernorm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 667c199c49..13b94f2327 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 56e050fe88..ddb33f303c 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b65f7005eb..f3220d5860 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py index cac3e18220..ace4be31de 100644 --- a/transformer_engine/pytorch/module/rmsnorm.py +++ b/transformer_engine/pytorch/module/rmsnorm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/numerics_debug.py b/transformer_engine/pytorch/numerics_debug.py index 5a73f5b61b..45d9aacde3 100644 --- a/transformer_engine/pytorch/numerics_debug.py +++ b/transformer_engine/pytorch/numerics_debug.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/onnx_extensions.py b/transformer_engine/pytorch/onnx_extensions.py index 79f9a9fb47..4d3b90bf63 100644 --- a/transformer_engine/pytorch/onnx_extensions.py +++ b/transformer_engine/pytorch/onnx_extensions.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/__init__.py b/transformer_engine/pytorch/ops/__init__.py index 156c33210a..2b270ea3de 100644 --- a/transformer_engine/pytorch/ops/__init__.py +++ b/transformer_engine/pytorch/ops/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 103e537dd0..4520dbc313 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 28d49bf7b9..665ffe359c 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index a444facd0a..9d54e12dba 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/add_extra_input.py b/transformer_engine/pytorch/ops/basic/add_extra_input.py index 1fcfa0466a..47f2b6e248 100644 --- a/transformer_engine/pytorch/ops/basic/add_extra_input.py +++ b/transformer_engine/pytorch/ops/basic/add_extra_input.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/all_gather.py b/transformer_engine/pytorch/ops/basic/all_gather.py index fc768ad83b..4e5c192876 100644 --- a/transformer_engine/pytorch/ops/basic/all_gather.py +++ b/transformer_engine/pytorch/ops/basic/all_gather.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/all_reduce.py b/transformer_engine/pytorch/ops/basic/all_reduce.py index d9a253924c..f2e4b2481d 100644 --- a/transformer_engine/pytorch/ops/basic/all_reduce.py +++ b/transformer_engine/pytorch/ops/basic/all_reduce.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 9f09e6634b..2714d718fe 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 6910163825..8b60251088 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/constant_scale.py b/transformer_engine/pytorch/ops/basic/constant_scale.py index 4de70c0e9f..d4b3660acf 100644 --- a/transformer_engine/pytorch/ops/basic/constant_scale.py +++ b/transformer_engine/pytorch/ops/basic/constant_scale.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/dropout.py b/transformer_engine/pytorch/ops/basic/dropout.py index 38b2a59a73..8850604aad 100644 --- a/transformer_engine/pytorch/ops/basic/dropout.py +++ b/transformer_engine/pytorch/ops/basic/dropout.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/identity.py b/transformer_engine/pytorch/ops/basic/identity.py index 788b3aac8a..9e90bd98c0 100644 --- a/transformer_engine/pytorch/ops/basic/identity.py +++ b/transformer_engine/pytorch/ops/basic/identity.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/l2normalization.py b/transformer_engine/pytorch/ops/basic/l2normalization.py index ff4f923819..be155c9356 100644 --- a/transformer_engine/pytorch/ops/basic/l2normalization.py +++ b/transformer_engine/pytorch/ops/basic/l2normalization.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 3922f85cad..631f0fafc9 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 34228affc7..61caaaf65d 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/quantize.py b/transformer_engine/pytorch/ops/basic/quantize.py index 1278701a9b..d126b554b5 100644 --- a/transformer_engine/pytorch/ops/basic/quantize.py +++ b/transformer_engine/pytorch/ops/basic/quantize.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/reduce_scatter.py b/transformer_engine/pytorch/ops/basic/reduce_scatter.py index eabbb461bc..0169da2490 100644 --- a/transformer_engine/pytorch/ops/basic/reduce_scatter.py +++ b/transformer_engine/pytorch/ops/basic/reduce_scatter.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/reshape.py b/transformer_engine/pytorch/ops/basic/reshape.py index fcdb3b0bbe..f8ae86fecd 100644 --- a/transformer_engine/pytorch/ops/basic/reshape.py +++ b/transformer_engine/pytorch/ops/basic/reshape.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 316c292c53..3179d0a447 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index 21113c2127..f4568ff25d 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py index a33ef4acf8..d5b9ce0e96 100644 --- a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py +++ b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index 1df55b83a0..186619caae 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index 0c12e3ab3e..5e7339db85 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py index 39ee4ab2fa..f7f59e65c9 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index ca3d57ac98..1c5edfcfcb 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 8a0f77dd56..4efb33e037 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index fe93410707..25b40f76e3 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 5149aa1ffb..4943ffb1bd 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 517632d651..fe04aa1e0b 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index fecf28f0a9..bf7af48d03 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/linear.py b/transformer_engine/pytorch/ops/linear.py index d1e6382291..d5829b0c50 100644 --- a/transformer_engine/pytorch/ops/linear.py +++ b/transformer_engine/pytorch/ops/linear.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 421c92b823..47286dfced 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index 2afda58e47..a0db3cd2d0 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/optimizers/__init__.py b/transformer_engine/pytorch/optimizers/__init__.py index c76f75743d..792eab094a 100644 --- a/transformer_engine/pytorch/optimizers/__init__.py +++ b/transformer_engine/pytorch/optimizers/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index b5c87b4815..1995655c33 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/optimizers/fused_sgd.py b/transformer_engine/pytorch/optimizers/fused_sgd.py index 8a76ec5901..08e465e951 100644 --- a/transformer_engine/pytorch/optimizers/fused_sgd.py +++ b/transformer_engine/pytorch/optimizers/fused_sgd.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py index 64ec0a28da..a5cbd27337 100644 --- a/transformer_engine/pytorch/optimizers/multi_tensor_apply.py +++ b/transformer_engine/pytorch/optimizers/multi_tensor_apply.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index d15814585e..0c16b35e11 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/pyproject.toml b/transformer_engine/pytorch/pyproject.toml index e5a4549db2..0b42b0a8da 100755 --- a/transformer_engine/pytorch/pyproject.toml +++ b/transformer_engine/pytorch/pyproject.toml @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 3f5995230c..eba547afb0 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index c9a4467a82..3414581f7c 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index a6030dd9df..52d1d9d6ca 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """ diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py index 9719ccb35c..99f6a99efa 100644 --- a/transformer_engine/pytorch/setup.py +++ b/transformer_engine/pytorch/setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index ada624a902..cb199d24b5 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index 2214edbff2..ba3407e13b 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 01e03e5355..03c16ebbed 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 67df40c047..43cbdcf9e6 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 6dcf9ae79a..88081f51bf 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 0c244628d6..8b707af3b2 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/storage/__init__.py b/transformer_engine/pytorch/tensor/storage/__init__.py index 9cb228f3a7..d7a2719200 100644 --- a/transformer_engine/pytorch/tensor/storage/__init__.py +++ b/transformer_engine/pytorch/tensor/storage/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. """Storage for quantized tensors.""" diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 38d117b2a0..157981b4d7 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 8d12c30700..adf3ce8aea 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index e7840d2c43..3bdf80c55e 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 04ab092ee2..b9d568c9cd 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 94f761f2b0..05e2d22e9c 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/torch_version.py b/transformer_engine/pytorch/torch_version.py index ff1a0abb89..3e299af1fd 100644 --- a/transformer_engine/pytorch/torch_version.py +++ b/transformer_engine/pytorch/torch_version.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index b3ad8ccc55..9b9ccc5185 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/triton/__init__.py b/transformer_engine/pytorch/triton/__init__.py index 80766864d9..d86cededd7 100644 --- a/transformer_engine/pytorch/triton/__init__.py +++ b/transformer_engine/pytorch/triton/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index e6d0397673..b574d69e0f 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/triton/pad.py b/transformer_engine/pytorch/triton/pad.py index 790b8277b2..547bc27760 100644 --- a/transformer_engine/pytorch/triton/pad.py +++ b/transformer_engine/pytorch/triton/pad.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 27662e1b28..8dff6b0426 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 16e126493f..47af9fabe1 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. From 27dc83bf0efd0aec5d659ea6aaf805746a2cd012 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 2 Jan 2026 15:02:19 +0530 Subject: [PATCH 146/521] Document environment variables (#2552) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Document envvars Signed-off-by: Kirthi Shankar Sivamani * Add remaining envvars Signed-off-by: Kirthi Shankar Sivamani * More missing ones Signed-off-by: Kirthi Shankar Sivamani * Update docs/envvars.rst Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update docs/envvars.rst Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> --- docs/envvars.rst | 494 +++++++++++++++++++++++++++++++++++++++++++++++ docs/index.rst | 1 + 2 files changed, 495 insertions(+) create mode 100644 docs/envvars.rst diff --git a/docs/envvars.rst b/docs/envvars.rst new file mode 100644 index 0000000000..86b313b133 --- /dev/null +++ b/docs/envvars.rst @@ -0,0 +1,494 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +.. _environment_variables: + +Environment Variables +===================== + +This document describes the environment variables used by Transformer Engine. They provide an alternate method to alter Transformer Engine's behavior during build and runtime, but are less rigorously maintained compared to the API and may be subject to change. + +Build-Time Environment Variables +--------------------------------- + +These environment variables control the build and compilation process of Transformer Engine. + +Build Configuration +^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_BUILD_DEBUG + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable debug build mode. When set to ``1``, the build includes debug symbols (``-g``) and disables optimizations. + +.. envvar:: NVTE_BUILD_MAX_JOBS + + :Type: ``int`` + :Default: Maximum available + :Description: Number of parallel jobs to use during the build process. If not set, the system will use the maximum available parallel jobs. Also respects the standard ``MAX_JOBS`` environment variable. + +.. envvar:: NVTE_BUILD_THREADS_PER_JOB + + :Type: ``int`` + :Default: ``1`` + :Description: Number of threads to use per parallel build job. This is passed to the CUDA compiler via the ``--threads`` flag. + +.. envvar:: NVTE_FRAMEWORK + + :Type: ``str`` + :Default: Auto-detected + :Description: Comma-separated list of frameworks to build support for (``pytorch``, ``jax``, ``all``, or ``none``). If not specified, automatically detects installed frameworks. + +.. envvar:: NVTE_USE_CCACHE + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable ccache for faster recompilation. When set to ``1``, uses ccache as a compiler launcher for both C++ and CUDA compilation. + +.. envvar:: NVTE_CCACHE_BIN + + :Type: ``str`` + :Default: ``ccache`` + :Description: Path to the ccache binary. Only used when :envvar:`NVTE_USE_CCACHE` is set to ``1``. + +.. envvar:: NVTE_CMAKE_BUILD_DIR + + :Type: ``str`` + :Default: None + :Description: Path to the CMake build directory for incremental builds. If set, CMake will use this directory for build artifacts. + +.. envvar:: NVTE_RELEASE_BUILD + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable release build mode. When set to ``1``, prepares the build for distribution (e.g., PyPI wheel). This affects library installation paths and build tool management. + +.. envvar:: NVTE_PROJECT_BUILDING + + :Type: ``int`` (0 or 1) + :Default: Not set + :Description: Internal flag set to ``1`` during the build process to indicate that the project is being built. Not intended for external use. + +Optional Dependencies +^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_UB_WITH_MPI + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable MPI support for userbuffers. When set to ``1``, requires ``MPI_HOME`` to be set to the MPI installation directory. + +.. envvar:: NVTE_ENABLE_NVSHMEM + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable NVSHMEM support. When set to ``1``, requires ``NVSHMEM_HOME`` to be set to the NVSHMEM installation directory. + +.. envvar:: NVTE_BUILD_ACTIVATION_WITH_FAST_MATH + + :Type: CMake option + :Default: ``OFF`` + :Description: Compile activation kernels (GELU, ReLU, SwiGLU) with the ``--use_fast_math`` CUDA compiler flag for improved performance at the cost of some precision. + +CUDA Configuration +^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_CUDA_ARCHS + + :Type: ``str`` + :Default: Auto-detected based on CUDA version + :Description: Semicolon-separated list of CUDA compute architectures to compile for (e.g., ``"80;90"`` for A100 and H100, or ``"75;80;89;90"``). If not set, automatically determined based on the installed CUDA Toolkit version. CUDA 13.0+ defaults to ``"75;80;89;90;100;120"``, CUDA 12.8+ defaults to ``"70;80;89;90;100;120"``, and earlier versions default to ``"70;80;89;90"``. Setting this can significantly reduce build time and binary size by targeting only the GPU architectures you need. + +.. envvar:: NVTE_CUDA_INCLUDE_DIR + + :Type: ``str`` + :Default: Auto-detected + :Description: Path to CUDA include directory containing ``cuda_runtime.h``. If not set, Transformer Engine searches in common locations (``CUDA_HOME``, ``CUDA_DIR``, ``/usr/local/cuda``). This is used for NVRTC kernel compilation. + +Runtime Environment Variables +------------------------------ + +These environment variables control the behavior of Transformer Engine during execution. + +Attention Backend Selection +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_FLASH_ATTN + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable or disable FlashAttention backend for DotProductAttention. When set to ``0``, FlashAttention will not be used. + +.. envvar:: NVTE_FUSED_ATTN + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable or disable FusedAttention backend (cuDNN-based) for DotProductAttention. When set to ``0``, FusedAttention will not be used. + +.. envvar:: NVTE_UNFUSED_ATTN + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable or disable UnfusedDotProductAttention backend (native PyTorch). When set to ``0``, UnfusedDotProductAttention will not be used. + +.. envvar:: NVTE_FUSED_ATTN_BACKEND + + :Type: ``int`` (0, 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. + +.. envvar:: NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT + + :Type: ``int`` (0 or 1) + :Default: Auto-determined + :Description: Control workspace-related optimizations in FusedAttention. ``0`` disables optimizations, ``1`` enables them. These optimizations trade memory for performance. When unset, Transformer Engine determines the code path based on internal logic. For deterministic behavior with cuDNN ≥8.9.5 and <9.0.0, this is automatically set to ``1``. + +.. envvar:: NVTE_FUSED_ATTN_USE_FAv2_BWD + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: When using FusedAttention, use FlashAttention-2 implementation for the backward pass instead of the cuDNN implementation. This can be useful due to performance differences between various versions of flash-attn and FusedAttention. + +.. envvar:: NVTE_ALLOW_NONDETERMINISTIC_ALGO + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Allow non-deterministic algorithms for Transformer Engine execution. When set to ``0``, only deterministic algorithms are allowed. This is relevant for both PyTorch and JAX attention implementations. + +.. envvar:: NVTE_FUSED_RING_ATTENTION_USE_SCAN + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: **(JAX only)** Use scan loop for ring attention implementation. When set to ``1``, the fused ring attention will use a scan-based iteration approach. + +.. envvar:: NVTE_APPLY_QK_LAYER_SCALING + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Apply QK layer scaling in UnfusedDotProductAttention. This is an FP16 training trick required for certain GPT-like models. When set to ``1`` and a layer number is provided, the softmax scale is divided by the layer number, and the layer number is used as the softmax scale during the softmax operation. Only effective when using FP16 dtype and when the layer number is specified. + +Context Parallelism +^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_BATCH_MHA_P2P_COMM + + :Type: ``int`` (0 or 1) + :Default: ``0`` (or auto-enabled for pre-Blackwell GPUs with CP size 2) + :Description: Use batched P2P communication (``batch_isend_irecv``) for KV exchange in context parallel MultiheadAttention. When enabled, send and receive operations are batched together, which can improve communication efficiency. This is automatically enabled for devices with compute capability < 10.0 (pre-Blackwell GPUs) when context parallel size is 2. Setting this to ``1`` forces batched P2P communication regardless of device architecture. + +FP8 Configuration +^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_UNFUSED_FP8_UPDATE + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use unfused kernel for FP8 amax and scale updates. When set to ``1``, amax and scale updates are computed using separate unfused kernels instead of fused operations. + +.. envvar:: NVTE_FP8_DPA_BWD + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable FP8 in the backward pass of DotProductAttention. ``1`` = FP8 forward and backward, ``0`` = FP8 forward and FP16/BF16 backward. + +.. envvar:: NVTE_DPA_FP8CS_O_in_F16 + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: For Float8CurrentScaling in DotProductAttention, use FP16/BF16 for the output tensor in the backward pass. ``1`` = use F16/BF16 output in backward, ``0`` = use FP8 output in backward. + +.. envvar:: NVTE_DPA_FP8_RECIPE + + :Type: ``str`` + :Default: Empty (use same as linear layers) + :Description: Override FP8 recipe for DotProductAttention layers. Valid values: ``"F16"`` (disable FP8), ``"DelayedScaling"``, or ``"Float8CurrentScaling"``. This allows using different FP8 recipes for attention vs. linear layers. + +.. envvar:: NVTE_DPA_FP8_RECIPE_DPA + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable FP8 in DotProductAttention when using :envvar:`NVTE_DPA_FP8_RECIPE`. When set to ``1``, the DotProductAttention layer will use the FP8 recipe specified by :envvar:`NVTE_DPA_FP8_RECIPE`. This provides fine-grained control over which attention components use FP8. + +.. envvar:: NVTE_DPA_FP8_RECIPE_MHA + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable FP8 in MultiheadAttention (MHA) when using :envvar:`NVTE_DPA_FP8_RECIPE`. When set to ``1``, the MultiheadAttention QKV and output projection layers will use the FP8 recipe specified by :envvar:`NVTE_DPA_FP8_RECIPE`. This provides fine-grained control over which attention components use FP8. + +.. envvar:: NVTE_DPA_FP8_FORMAT + + :Type: ``str`` + :Default: ``"HYBRID"`` + :Description: FP8 format for DotProductAttention when switching recipes. Valid values: ``"HYBRID"``, ``"E4M3"``, ``"E5M2"``. Only used when :envvar:`NVTE_DPA_FP8_RECIPE` is set. + +.. envvar:: NVTE_DPA_FP8DS_AMAX_ALGO + + :Type: ``str`` + :Default: ``"most_recent"`` + :Description: Amax computation algorithm for DelayedScaling recipe in DotProductAttention. Valid values: ``"most_recent"``, ``"max"``. Only used when :envvar:`NVTE_DPA_FP8_RECIPE` is set to ``"DelayedScaling"``. + +.. envvar:: NVTE_DPA_FP8DS_AMAX_HISTLEN + + :Type: ``int`` + :Default: ``1`` + :Description: Amax history length for DelayedScaling recipe in DotProductAttention. Only used when :envvar:`NVTE_DPA_FP8_RECIPE` is set to ``"DelayedScaling"``. + +.. envvar:: NVTE_DPA_FP8DS_REDUCE_AMAX + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Reduce amax across distributed ranks for DelayedScaling recipe in DotProductAttention. Only used when :envvar:`NVTE_DPA_FP8_RECIPE` is set to ``"DelayedScaling"``. + +.. envvar:: NVTE_UnfusedDPA_Emulate_FP8 + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Allow FP8 emulation in UnfusedDotProductAttention. When set to ``1``, UnfusedDotProductAttention can emulate FP8 operations using FP16/BF16 computation. + +Kernel Configuration +^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_USE_FAST_MATH + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable fast math optimizations in runtime-compiled (NVRTC) kernels. This trades numerical accuracy for performance. These optimizations are experimental and inconsistently implemented. + +.. envvar:: NVTE_DISABLE_NVRTC + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Disable NVRTC (CUDA Runtime Compilation) support. When set to ``1``, runtime kernel compilation is disabled. This can be useful in environments where NVRTC is not available or not desired. + +.. envvar:: NVTE_USE_CUTLASS_GROUPED_GEMM + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Use CUTLASS implementation for grouped GEMM operations instead of cuBLAS. When set to ``1``, enables CUTLASS grouped GEMM kernels, which may provide better performance for certain workloads on Hopper (SM90) GPUs. + +.. envvar:: NVTE_CUTLASS_GROUPED_GEMM_WARN_FALLBACK + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Emit a warning when falling back from CUTLASS to cuBLAS for grouped GEMM operations. + +Torch Compilation and Fusion +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_TORCH_COMPILE + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable PyTorch 2.x ``torch.compile`` support for compatible Transformer Engine operations. When set to ``0``, disables compilation support and uses regular PyTorch eager mode. + +.. envvar:: NVTE_BIAS_GELU_NVFUSION + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable GELU fusion with bias using NVFusion in PyTorch. When set to ``0``, uses separate bias and GELU operations. + +.. envvar:: NVTE_BIAS_DROPOUT_FUSION + + :Type: ``int`` (0 or 1) + :Default: ``1`` + :Description: Enable fusion of bias and dropout operations. When set to ``0``, bias and dropout are computed separately. + +LayerNorm/RMSNorm SM Margins +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_FWD_LAYERNORM_SM_MARGIN + + :Type: ``int`` + :Default: ``0`` + :Description: Number of SMs (Streaming Multiprocessors) to reserve (not use) during forward LayerNorm/RMSNorm operations. This can be used to control resource allocation and overlap computation with communication. + +.. envvar:: NVTE_BWD_LAYERNORM_SM_MARGIN + + :Type: ``int`` + :Default: ``0`` + :Description: Number of SMs to reserve during backward LayerNorm/RMSNorm operations. + +.. envvar:: NVTE_INF_LAYERNORM_SM_MARGIN + + :Type: ``int`` + :Default: ``0`` + :Description: Number of SMs to reserve during inference LayerNorm/RMSNorm operations. + +GEMM Configuration +^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_EXT_MARGIN_SM + + :Type: ``int`` + :Default: Total SM count + :Description: External SM margin for GEMM operations. Specifies the number of SMs to use for GEMM operations. The actual number of SMs used is ``sm_count - NVTE_EXT_MARGIN_SM``. + +.. envvar:: NVTE_AG_P2P_MULTI_ATOMIC + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable multi-atomic mode for AllGather with atomic GEMM using P2P communication. When set to ``1``, uses ``userbuffers_sendrecv_multiatomic`` for communication during atomic GEMM overlap with AllGather operations. This disables copy engine (CE) usage and enables push mode for userbuffers. This is an advanced optimization for tensor-parallel communication-computation overlap. + +CPU Offloading +^^^^^^^^^^^^^^ + +.. envvar:: NVTE_CPU_OFFLOAD_V1 + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable legacy version of CPU offloading implementation. + +Debugging and Profiling +^^^^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_DEBUG + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable debug mode. When set to ``1``, enables verbose debug output and additional checks in attention operations. + +.. envvar:: NVTE_DEBUG_LEVEL + + :Type: ``int`` (0, 1, or 2) + :Default: ``0`` + :Description: Debug verbosity level. Higher values enable more verbose debug output. Only effective when :envvar:`NVTE_DEBUG` is set to ``1``. + +.. envvar:: NVTE_PRINT_LAYER_NUMBER + + :Type: ``int`` + :Default: ``1`` + :Description: Layer number to print debug information for during attention operations. + +.. envvar:: NVTE_PRINT_RANK + + :Type: ``int`` + :Default: ``0`` + :Description: Distributed rank to print debug information for during attention operations. + +.. envvar:: NVTE_NVTX_ENABLED + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable NVTX (NVIDIA Tools Extension) range profiling for Transformer Engine operations. When set to ``1``, NVTX markers are added to operations for profiling with NVIDIA Nsight Systems. + +.. envvar:: NVTE_DEBUG_NUMERICS + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: **(JAX only)** Enable verbose printing of tensor numerics for debugging purposes. + +Testing +^^^^^^^ + +.. envvar:: NVTE_TEST_NVINSPECT_ENABLED + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable NVInspect integration for testing. When set to ``1``, enables the NVInspect debugging API for numerical analysis during tests. + +.. envvar:: NVTE_TEST_NVINSPECT_CONFIG_FILE + + :Type: ``str`` + :Default: None + :Description: Path to NVInspect configuration file. Required when :envvar:`NVTE_TEST_NVINSPECT_ENABLED` is set to ``1``. + +.. envvar:: NVTE_TEST_NVINSPECT_FEATURE_DIRS + + :Type: ``str`` + :Default: None + :Description: Comma-separated list of directories containing NVInspect features. Required when :envvar:`NVTE_TEST_NVINSPECT_ENABLED` is set to ``1``. + +.. envvar:: NVTE_TEST_ARTIFACTS_DIR + + :Type: ``str`` + :Default: System temp directory + :Description: Directory for storing test artifacts (e.g., generated ONNX models). + +ONNX Export +^^^^^^^^^^^ + +.. envvar:: NVTE_ONNX_KVCACHE_MAX_SEQ_LEN + + :Type: ``int`` + :Default: ``128`` + :Description: Maximum sequence length for KV cache during ONNX export. This is used for attention masking in exported ONNX models. + +JAX-Specific Variables +^^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_JAX_CUSTOM_CALLS + + :Type: ``str`` + :Default: None + :Description: Control which JAX custom call primitives are enabled or disabled. Format: ``"true"`` (enable all), ``"false"`` (disable all), or comma-separated key-value pairs like ``"GemmPrimitive=false,DBiasQuantizePrimitive=true"``. This provides fine-grained control over which operations use custom CUDA kernels vs. JAX native implementations. + +.. envvar:: NVTE_JAX_CUSTOM_CALLS_RE + + :Type: ``str`` + :Default: None + :Description: **Deprecated** (use :envvar:`NVTE_JAX_CUSTOM_CALLS` instead). Regex pattern to match primitive names for enabling/disabling. Example: ``"DBiasQuantizePrimitive"`` or ``"^(?!DBiasQuantizePrimitive$).+$"``. + +.. envvar:: NVTE_JAX_UNITTEST_LEVEL + + :Type: ``str`` + :Default: None + :Description: Test level for JAX unit tests (``"L0"``, ``"L1"``, ``"L2"``). Used internally by the test suite. + +Examples +-------- + +Building with Debug Symbols +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: bash + + export NVTE_BUILD_DEBUG=1 + export NVTE_USE_CCACHE=1 + pip install -e . + +Using Specific Attention Backend +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: bash + + # Use only FlashAttention, disable FusedAttention + export NVTE_FLASH_ATTN=1 + export NVTE_FUSED_ATTN=0 + python train.py + +Configuring FP8 for Attention +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: bash + + # Use DelayedScaling for attention, CurrentScaling for linear layers + export NVTE_DPA_FP8_RECIPE="DelayedScaling" + export NVTE_DPA_FP8_FORMAT="HYBRID" + export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent" + export NVTE_DPA_FP8DS_AMAX_HISTLEN=1024 + python train.py + +Enable Profiling +^^^^^^^^^^^^^^^^ + +.. code-block:: bash + + # Enable NVTX markers for profiling + export NVTE_NVTX_ENABLED=1 + nsys profile --trace=nvtx,cuda python train.py + +JAX Custom Calls Control +^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. code-block:: bash + + # Disable all custom calls + export NVTE_JAX_CUSTOM_CALLS="false" + python train_jax.py + + # Disable specific primitives + export NVTE_JAX_CUSTOM_CALLS="GemmPrimitive=false,DBiasQuantizePrimitive=false" + python train_jax.py diff --git a/docs/index.rst b/docs/index.rst index 3f707d8904..99611cda99 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -56,5 +56,6 @@ Transformer Engine documentation api/c/index debug + envvars examples/attention/attention.ipynb examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb From c988548f72bbc271fe2ab7bad1046b91b577aa29 Mon Sep 17 00:00:00 2001 From: xiaoxi-wangfj <690912414@qq.com> Date: Sat, 3 Jan 2026 07:23:29 +0800 Subject: [PATCH 147/521] [PyTorch] Fix garbage initialized permuted_scale (#2547) Signed-off-by: xiaoxi-wangfj <690912414@qq.com> Co-authored-by: Teddy Do --- transformer_engine/pytorch/triton/permutation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 8dff6b0426..8c9003bb5f 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -165,7 +165,7 @@ def permute_with_mask_map( alloc((num_out_tokens,), dtype=probs.dtype, device="cuda") if probs is not None else None ) permuted_scale = ( - torch.empty((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda") + alloc((num_out_tokens, scale_hidden_dim), dtype=scale.dtype, device="cuda") if scale is not None else None ) From 4f364c8e6010823cc5e94bd75d32dc1e743bc999 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 5 Jan 2026 11:22:30 +0530 Subject: [PATCH 148/521] Fix out of bound ID passed to `cutlass::arch::NamedBarrier::sync` (#2554) Fix barrier ID Signed-off-by: Kirthi Shankar Sivamani --- .../group_row_cast_col_hadamard_transform_cast_fusion.cu | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 c1fb87d048..8b077f6f1f 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 @@ -983,7 +983,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); - int row_quant_barrier_id = 10; + // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8 + // in order to go over the reserved named barrier count. + constexpr int row_quant_barrier_id = 2; cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); From c90a9214091badd1234b2d9ca851bd97f8edb0f6 Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Sun, 4 Jan 2026 22:52:46 -0700 Subject: [PATCH 149/521] Add tests that reset_parameters doesn't change parameter initial value ranges (#2550) * Add tests for 2528 and 2529 Signed-off-by: Peter St. John * Update tests/pytorch/test_deferred_init.py Signed-off-by: Kirthi Shankar Sivamani * Update tests/pytorch/test_deferred_init.py Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Peter St. John Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_deferred_init.py | 43 ++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_deferred_init.py b/tests/pytorch/test_deferred_init.py index 58c8543485..f61bf22194 100644 --- a/tests/pytorch/test_deferred_init.py +++ b/tests/pytorch/test_deferred_init.py @@ -28,7 +28,6 @@ class TestDeferredInit: - @staticmethod def get_module_args(module): hidden_size = num_heads * head_dim @@ -82,3 +81,45 @@ def test_reset_parameters( "on CUDA device" ) del module + + @pytest.mark.parametrize("module_type", _core_modules) + def test_reset_parameters_doesnt_change_parameter_stats( + self, + module_type: torch.nn.Module, + ) -> None: + """Test for github issue #2528 and #2529 to ensure that reset_parameters() doesn't change + the parameter mean and std""" + args, kwargs = TestDeferredInit.get_module_args(module_type) + kwargs["device"] = "cuda" + module = module_type(*args, **kwargs) + + param_stats = { + name: {"mean": param.mean(), "std": param.std()} + for name, param in module.named_parameters() + } + + with torch.no_grad(): + module.reset_parameters() + + param_stats_after = { + name: {"mean": param.mean(), "std": param.std()} + for name, param in module.named_parameters() + } + + for name, stats in param_stats_after.items(): + torch.testing.assert_close( + stats["mean"], + param_stats[name]["mean"], + atol=1e-3, + rtol=1e-3, + msg=f"{name} mean changed after reset_parameters", + ) + torch.testing.assert_close( + stats["std"], + param_stats[name]["std"], + atol=1e-3, + rtol=1e-3, + msg=f"{name} std changed after reset_parameters", + ) + + del module From a976740700869369d69966fe39fbb3e507232ace Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 6 Jan 2026 10:03:22 +0100 Subject: [PATCH 150/521] [docs] Getting started refactor (#2534) * docs: Add comprehensive Getting Started guide with benchmarks - Add new Getting Started documentation with PyTorch and JAX tutorials - Include benchmark scripts demonstrating TE performance benefits - Add CSS styling for code output and tabs - Replace old quickstart notebooks with improved documentation - Add transformer layer diagram (SVG) - Update docs configuration and workflow Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * 2026 in copyright 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> --- .github/workflows/docs.yml | 2 +- docs/_static/css/output-style.css | 60 ++ docs/_static/css/rtabs.css | 43 + docs/conf.py | 3 + docs/examples/advanced_optimizations.ipynb | 2 +- docs/examples/quickstart.ipynb | 606 ------------- docs/examples/quickstart_jax.ipynb | 833 ------------------ docs/examples/quickstart_jax_utils.py | 4 - docs/examples/te_jax_integration.ipynb | 2 +- docs/getting_started.rst | 16 - docs/getting_started/getting_started_jax.out | 34 + docs/getting_started/getting_started_jax.py | 523 +++++++++++ .../getting_started_jax_summary.csv | 7 + .../getting_started_pytorch.out | 42 + .../getting_started_pytorch.py | 497 +++++++++++ .../getting_started_pytorch_summary.csv | 7 + .../getting_started_utils_jax.py | 76 ++ .../getting_started_utils_pytorch.py | 124 +++ docs/getting_started/index.rst | 566 ++++++++++++ docs/getting_started/transformer_layer.svg | 82 ++ docs/index.rst | 2 +- 21 files changed, 2068 insertions(+), 1463 deletions(-) create mode 100644 docs/_static/css/output-style.css create mode 100644 docs/_static/css/rtabs.css delete mode 100644 docs/examples/quickstart.ipynb delete mode 100644 docs/examples/quickstart_jax.ipynb delete mode 100644 docs/getting_started.rst create mode 100644 docs/getting_started/getting_started_jax.out create mode 100644 docs/getting_started/getting_started_jax.py create mode 100644 docs/getting_started/getting_started_jax_summary.csv create mode 100644 docs/getting_started/getting_started_pytorch.out create mode 100644 docs/getting_started/getting_started_pytorch.py create mode 100644 docs/getting_started/getting_started_pytorch_summary.csv create mode 100644 docs/getting_started/getting_started_utils_jax.py create mode 100644 docs/getting_started/getting_started_utils_pytorch.py create mode 100644 docs/getting_started/index.rst create mode 100644 docs/getting_started/transformer_layer.svg diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 388c822eeb..6fde0338a1 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -17,7 +17,7 @@ jobs: uses: actions/checkout@v3 - name: 'Install dependencies' run: | - pip install sphinx==8.1.3 sphinx_rtd_theme==3.0.1 nbsphinx==0.9.5 IPython ipython_genutils==0.2.0 ipywidgets==8.0.2 astroid==3.3.2 + pip install sphinx==8.1.3 sphinx_rtd_theme==3.0.1 nbsphinx==0.9.5 IPython ipython_genutils==0.2.0 ipywidgets==8.0.2 astroid==3.3.2 sphinx-tabs==3.4.7 pip install breathe==4.35.0 sphinx-autoapi==3.3.2 sudo apt-get install -y pandoc graphviz doxygen export GIT_SHA=$(git show-ref --hash HEAD) diff --git a/docs/_static/css/output-style.css b/docs/_static/css/output-style.css new file mode 100644 index 0000000000..864d8587a3 --- /dev/null +++ b/docs/_static/css/output-style.css @@ -0,0 +1,60 @@ +/* Custom styling for program output blocks */ + +.program-output { + background-color: #f8f9fa; + padding: 0; /* No padding at all */ + margin: 0; /* No margins at all */ + border-radius: 0; /* No rounded corners */ + font-family: 'Courier New', monospace; + font-size: 14px; + line-height: 1.5; + width: 100%; + max-width: 100%; +} + +.program-output pre { + margin: 0; + padding: 0; + background: transparent !important; + border: none !important; + color: #2c3e50; + width: 100%; +} + +.program-output .highlight { + background: transparent !important; + margin: 0; + width: 100%; +} + +/* Alternative lighter style */ +.output-block { + background-color: #fafbfc; + border: 1px solid #e1e4e8; + padding: 10px 14px; + margin: 10px 0; + border-radius: 3px; + font-family: 'SF Mono', 'Consolas', monospace; + font-size: 13px; + color: #24292e; +} + +/* Console-like output style */ +.console-output { + background-color: #1e1e1e; + border-left: 3px solid #76b900; + padding: 14px 18px; + margin: 12px 0; + border-radius: 5px; + font-family: 'Fira Code', 'Consolas', monospace; + font-size: 13px; + color: #d4d4d4; + box-shadow: 0 2px 4px rgba(0,0,0,0.1); +} + +.console-output pre { + margin: 0; + color: #d4d4d4; + background: transparent !important; +} + diff --git a/docs/_static/css/rtabs.css b/docs/_static/css/rtabs.css new file mode 100644 index 0000000000..7f4213ef99 --- /dev/null +++ b/docs/_static/css/rtabs.css @@ -0,0 +1,43 @@ +/* Custom styling for sphinx-tabs */ + +.sphinx-tabs { + margin-bottom: 1rem; +} + +.sphinx-tabs-tab { + background-color: #f4f4f4; + border: 1px solid #ccc; + border-bottom: none; + padding: 0.5rem 1rem; + margin-right: 0.5rem; + cursor: pointer; + font-weight: 500; + transition: background-color 0.2s; +} + +.sphinx-tabs-tab:hover { + background-color: #e0e0e0; +} + +.sphinx-tabs-tab[aria-selected="true"] { + background-color: #76b900; /* NVIDIA green */ + color: white; + border-color: #76b900; + margin-right: 0.5rem; +} + +.sphinx-tabs-panel { + border: 1px solid #ccc; + padding: 1rem; + background-color: #f9f9f9; +} + +/* Dark mode support for RTD theme */ +.rst-content .sphinx-tabs-tab { + color: #333; +} + +.rst-content .sphinx-tabs-tab[aria-selected="true"] { + color: white; +} + diff --git a/docs/conf.py b/docs/conf.py index 0734008137..43a7230666 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -58,6 +58,7 @@ "nbsphinx", "breathe", "autoapi.extension", + "sphinx_tabs.tabs", ] templates_path = ["_templates"] @@ -83,6 +84,8 @@ html_css_files = [ "css/nvidia_font.css", "css/nvidia_footer.css", + "css/rtabs.css", + "css/output-style.css", ] html_theme_options = { diff --git a/docs/examples/advanced_optimizations.ipynb b/docs/examples/advanced_optimizations.ipynb index 7c08bb6586..1b0694a05f 100644 --- a/docs/examples/advanced_optimizations.ipynb +++ b/docs/examples/advanced_optimizations.ipynb @@ -13,7 +13,7 @@ "id": "6dcbf25a", "metadata": {}, "source": [ - "This guide is a follow-up to the discussion in the [quickstart guide](quickstart.ipynb). We will focus on techniques to achieve maximum performance when training a basic GPT encoder layer. For convenience, we use some helper functions defined in [quickstart_utils.py](quickstart_utils.py). " + "This guide is a follow-up to the discussion in the [Getting Started guide](../getting_started/index.rst). We will focus on techniques to achieve maximum performance when training a basic GPT encoder layer. For convenience, we use some helper functions defined in [quickstart_utils.py](quickstart_utils.py). " ] }, { diff --git a/docs/examples/quickstart.ipynb b/docs/examples/quickstart.ipynb deleted file mode 100644 index 0ad2f4fee8..0000000000 --- a/docs/examples/quickstart.ipynb +++ /dev/null @@ -1,606 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "da9fd6a8", - "metadata": {}, - "source": [ - "# Getting Started\n", - "\n", - "## Overview\n", - "\n", - "Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, providing better performance with lower memory utilization in both training and inference. It provides support for 8-bit floating point (FP8) precision on Hopper GPUs, implements a collection of highly optimized building blocks for popular Transformer architectures, and exposes an automatic-mixed-precision-like API that can be used seamlessly with your PyTorch code. It also includes a framework-agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers.\n", - "\n", - "## Let's build a Transformer layer!\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We build a basic Transformer layer using regular PyTorch modules. This will be our baseline for later comparisons with Transformer Engine.\n", - "\n", - "
\n", - "\n", - "Let's start with creating a GPT encoder layer using plain PyTorch. Figure 1 shows the overall structure.\n", - "\n", - "
\n", - "\n", - "
Figure 1: Structure of a GPT encoder layer.
\n", - "
\n", - "\n", - "We construct the components as follows:\n", - "\n", - "- `LayerNorm`: `torch.nn.LayerNorm`\n", - "- `QKV Projection`: `torch.nn.Linear` (conceptually three `Linear` layers for Q, K, and V separately, but we fuse into a single `Linear` layer that is three times larger)\n", - "- `DotProductAttention`: `DotProductAttention` from [quickstart_utils.py](quickstart_utils.py)\n", - "- `Projection`: `torch.nn.Linear`\n", - "- `Dropout`: `torch.nn.Dropout`\n", - "- `MLP`: `BasicMLP` from [quickstart_utils.py](quickstart_utils.py)\n", - "\n", - "Over the course of this tutorial we will use a few modules and helper functions defined in [quickstart_utils.py](quickstart_utils.py). Putting it all together:" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "2be43d64", - "metadata": {}, - "outputs": [], - "source": [ - "import torch\n", - "import quickstart_utils as utils\n", - "\n", - "class BasicTransformerLayer(torch.nn.Module):\n", - " def __init__(\n", - " self,\n", - " hidden_size: int,\n", - " ffn_hidden_size: int,\n", - " num_attention_heads: int,\n", - " layernorm_eps: int = 1e-5,\n", - " attention_dropout: float = 0.1,\n", - " hidden_dropout: float = 0.1,\n", - " ):\n", - " super().__init__()\n", - " self.num_attention_heads = num_attention_heads\n", - " self.kv_channels = hidden_size // num_attention_heads\n", - " self.ln1 = torch.nn.LayerNorm(hidden_size, eps=layernorm_eps)\n", - " self.qkv_projection = torch.nn.Linear(hidden_size, 3 * hidden_size, bias=True)\n", - " self.attention = utils.DotProductAttention(\n", - " num_attention_heads=num_attention_heads,\n", - " kv_channels=self.kv_channels,\n", - " attention_dropout=attention_dropout,\n", - " )\n", - " self.projection = torch.nn.Linear(hidden_size, hidden_size, bias=True)\n", - " self.dropout = torch.nn.Dropout(hidden_dropout)\n", - " self.ln2 = torch.nn.LayerNorm(hidden_size, eps=layernorm_eps)\n", - " self.mlp = utils.BasicMLP(\n", - " hidden_size=hidden_size,\n", - " ffn_hidden_size=ffn_hidden_size,\n", - " ) \n", - " \n", - " def forward(\n", - " self, \n", - " x: torch.Tensor, \n", - " attention_mask: torch.Tensor\n", - " ) -> torch.Tensor:\n", - " res = x\n", - " x = self.ln1(x)\n", - " \n", - " # Fused QKV projection\n", - " qkv = self.qkv_projection(x)\n", - " qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels)\n", - " q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3)\n", - " \n", - " x = self.attention(q, k, v, attention_mask)\n", - " x = self.projection(x)\n", - " x = self.dropout(x)\n", - " x = res + x\n", - " res = x\n", - " x = self.ln2(x)\n", - " x = self.mlp(x)\n", - " \n", - " return x + res" - ] - }, - { - "cell_type": "markdown", - "id": "40724d1d", - "metadata": {}, - "source": [ - "That's it! We now have a simple Transformer layer. We can test it:" - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "a786f0ea", - "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 = torch.float16\n", - "\n", - "# Synthetic data\n", - "x = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)\n", - "dy = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype)" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "ffdbfb7a", - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "BasicTransformerLayer(\n", - " (ln1): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n", - " (qkv_projection): Linear(in_features=4096, out_features=12288, bias=True)\n", - " (attention): DotProductAttention(\n", - " (dropout): Dropout(p=0.1, inplace=False)\n", - " )\n", - " (projection): Linear(in_features=4096, out_features=4096, bias=True)\n", - " (dropout): Dropout(p=0.1, inplace=False)\n", - " (ln2): LayerNorm((4096,), eps=1e-05, elementwise_affine=True)\n", - " (mlp): BasicMLP(\n", - " (linear1): Linear(in_features=4096, out_features=16384, bias=True)\n", - " (linear2): Linear(in_features=16384, out_features=4096, bias=True)\n", - " )\n", - ")" - ] - }, - "execution_count": 3, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "basic_transformer = BasicTransformerLayer(\n", - " hidden_size,\n", - " ffn_hidden_size,\n", - " num_attention_heads,\n", - ")\n", - "basic_transformer.to(dtype=dtype).cuda()" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "0162ad40", - "metadata": {}, - "outputs": [], - "source": [ - "torch.manual_seed(1234)\n", - "y = basic_transformer(x, attention_mask=None)" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "65ae6dd6", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 43.0663916015625 ms\n" - ] - } - ], - "source": [ - "utils.speedometer(\n", - " basic_transformer,\n", - " x,\n", - " dy,\n", - " forward_kwargs = { \"attention_mask\": None },\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "43717e36", - "metadata": {}, - "source": [ - "## Meet Transformer Engine\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We modify the example Transformer layer to include the simplest TE modules: `Linear` and `LayerNorm`.\n", - "\n", - "
\n", - "\n", - "Now that we have a basic Transformer layer, let's use Transformer Engine to speed up the training. " - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "004d3c92", - "metadata": {}, - "outputs": [], - "source": [ - "import transformer_engine.pytorch as te" - ] - }, - { - "cell_type": "markdown", - "id": "1931f911", - "metadata": {}, - "source": [ - "TE provides a set of PyTorch modules that can be used to build Transformer layers. The simplest of the provided modules are the `Linear` and `LayerNorm` layers, which we can use instead of `torch.nn.Linear` and `torch.nn.LayerNorm`. Let's modify `BasicTransformerLayer`:" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "1f44db50", - "metadata": {}, - "outputs": [], - "source": [ - "class BasicTEMLP(torch.nn.Module):\n", - " def __init__(self,\n", - " hidden_size: int,\n", - " ffn_hidden_size: int) -> None:\n", - " super().__init__()\n", - " self.linear1 = te.Linear(hidden_size, ffn_hidden_size, bias=True)\n", - " self.linear2 = te.Linear(ffn_hidden_size, hidden_size, bias=True)\n", - "\n", - " def forward(self, x):\n", - " x = self.linear1(x)\n", - " x = torch.nn.functional.gelu(x, approximate='tanh')\n", - " x = self.linear2(x)\n", - " return x \n", - " \n", - "class BasicTETransformerLayer(torch.nn.Module):\n", - " def __init__(self,\n", - " hidden_size: int,\n", - " ffn_hidden_size: int,\n", - " num_attention_heads: int,\n", - " layernorm_eps: int = 1e-5,\n", - " attention_dropout: float = 0.1,\n", - " hidden_dropout: float = 0.1):\n", - " super().__init__()\n", - " self.num_attention_heads = num_attention_heads\n", - " self.kv_channels = hidden_size // num_attention_heads\n", - " self.ln1 = te.LayerNorm(hidden_size, eps=layernorm_eps)\n", - " self.qkv_projection = te.Linear(hidden_size, 3 * hidden_size, bias=True)\n", - " self.attention = utils.DotProductAttention(\n", - " num_attention_heads=num_attention_heads,\n", - " kv_channels=self.kv_channels,\n", - " attention_dropout=attention_dropout,\n", - " )\n", - " self.projection = te.Linear(hidden_size, hidden_size, bias=True)\n", - " self.dropout = torch.nn.Dropout(hidden_dropout)\n", - " self.ln2 = te.LayerNorm(hidden_size, eps=layernorm_eps)\n", - " self.mlp = BasicTEMLP(\n", - " hidden_size=hidden_size,\n", - " ffn_hidden_size=ffn_hidden_size,\n", - " )\n", - " \n", - " def forward(self, \n", - " x: torch.Tensor, \n", - " attention_mask: torch.Tensor):\n", - " res = x\n", - " x = self.ln1(x)\n", - " \n", - " # Fused QKV projection\n", - " qkv = self.qkv_projection(x)\n", - " qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels)\n", - " q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3)\n", - " \n", - " x = self.attention(q, k, v, attention_mask)\n", - " x = self.projection(x)\n", - " x = self.dropout(x)\n", - " x = res + x\n", - " res = x\n", - " x = self.ln2(x)\n", - " x = self.mlp(x)\n", - " \n", - " return x + res" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "916531e8", - "metadata": {}, - "outputs": [], - "source": [ - "basic_te_transformer = BasicTETransformerLayer(\n", - " hidden_size, \n", - " ffn_hidden_size, \n", - " num_attention_heads,\n", - ")\n", - "basic_te_transformer.to(dtype=dtype).cuda()\n", - "utils.share_parameters_with_basic_te_model(basic_te_transformer, basic_transformer)" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "3643fa54", - "metadata": {}, - "outputs": [], - "source": [ - "torch.manual_seed(1234)\n", - "y = basic_te_transformer(x, attention_mask=None)" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "10b92894", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 43.1413232421875 ms\n" - ] - } - ], - "source": [ - "utils.speedometer(\n", - " basic_te_transformer,\n", - " x,\n", - " dy,\n", - " forward_kwargs = { \"attention_mask\": None },\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "3f990226", - "metadata": {}, - "source": [ - "## Fused TE Modules\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We optimize the example Transformer layer with TE modules for fused operations.\n", - "\n", - "
\n", - "\n", - "The `Linear` layer is enough to build any Transformer model and it enables usage of Transformer Engine even for very custom Transformers. However, having more knowledge about the model allows for additional optimizations like kernel fusion, increasing the achievable speedup.\n", - "\n", - "Transformer Engine therefore provides coarser modules that span multiple layers:\n", - "\n", - "* `LayerNormLinear`\n", - "* `LayerNormMLP`\n", - "* `TransformerLayer`\n", - "\n", - "Building a third iteration of our Transformer layer with `LayerNormLinear` and `LayerNormMLP`:" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "c55eae1f", - "metadata": {}, - "outputs": [], - "source": [ - "class FusedTETransformerLayer(torch.nn.Module):\n", - " def __init__(self,\n", - " hidden_size: int,\n", - " ffn_hidden_size: int,\n", - " num_attention_heads: int,\n", - " layernorm_eps: int = 1e-5,\n", - " attention_dropout: float = 0.1,\n", - " hidden_dropout: float = 0.1):\n", - " super().__init__()\n", - " self.num_attention_heads = num_attention_heads\n", - " self.kv_channels = hidden_size // num_attention_heads\n", - " self.ln_qkv = te.LayerNormLinear(hidden_size, 3 * hidden_size, eps=layernorm_eps, bias=True)\n", - " self.attention = utils.DotProductAttention(\n", - " num_attention_heads=num_attention_heads,\n", - " kv_channels=self.kv_channels,\n", - " attention_dropout=attention_dropout,\n", - " )\n", - " self.projection = te.Linear(hidden_size, hidden_size, bias=True)\n", - " self.dropout = torch.nn.Dropout(hidden_dropout)\n", - " self.ln_mlp = te.LayerNormMLP(hidden_size, ffn_hidden_size, eps=layernorm_eps, bias=True)\n", - " \n", - " \n", - " def forward(self, \n", - " x: torch.Tensor, \n", - " attention_mask: torch.Tensor):\n", - " res = x\n", - " qkv = self.ln_qkv(x)\n", - " \n", - " # Split qkv into query, key and value\n", - " qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels)\n", - " q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3)\n", - " \n", - " x = self.attention(q, k, v, attention_mask)\n", - " x = self.projection(x)\n", - " x = self.dropout(x)\n", - " x = res + x\n", - " res = x\n", - " x = self.ln_mlp(x)\n", - " \n", - " return x + res" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "85949421", - "metadata": {}, - "outputs": [], - "source": [ - "fused_te_transformer = FusedTETransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads)\n", - "fused_te_transformer.to(dtype=dtype).cuda()\n", - "utils.share_parameters_with_fused_te_model(fused_te_transformer, basic_transformer)" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "2c263e71", - "metadata": {}, - "outputs": [], - "source": [ - "torch.manual_seed(1234)\n", - "y = fused_te_transformer(x, attention_mask=None)" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "id": "24e101bc", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 43.1981201171875 ms\n" - ] - } - ], - "source": [ - "utils.speedometer(\n", - " fused_te_transformer,\n", - " x,\n", - " dy,\n", - " forward_kwargs = { \"attention_mask\": None },\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "33f13c26", - "metadata": {}, - "source": [ - "Finally, the `TransformerLayer` module is convenient for creating standard Transformer architectures and it provides the highest degree of performance optimization:" - ] - }, - { - "cell_type": "code", - "execution_count": 15, - "id": "ec8c3685", - "metadata": {}, - "outputs": [], - "source": [ - "te_transformer = te.TransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads)\n", - "te_transformer.to(dtype=dtype).cuda()\n", - "utils.share_parameters_with_transformerlayer_te_model(te_transformer, basic_transformer)" - ] - }, - { - "cell_type": "code", - "execution_count": 16, - "id": "e48cd590", - "metadata": {}, - "outputs": [], - "source": [ - "torch.manual_seed(1234)\n", - "y = te_transformer(x, attention_mask=None)" - ] - }, - { - "cell_type": "code", - "execution_count": 17, - "id": "3ec3707d-e63f-4899-8308-b11c55b5caa4", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 39.99169921875 ms\n" - ] - } - ], - "source": [ - "utils.speedometer(\n", - " te_transformer,\n", - " x,\n", - " dy,\n", - " forward_kwargs = { \"attention_mask\": None },\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "4034c3eb-8958-49f2-85f6-30c94977d884", - "metadata": {}, - "source": [ - "## Enabling FP8\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We configure a TE module to perform compute in FP8.\n", - "\n", - "
\n", - "\n", - "Enabling FP8 support is very simple in Transformer Engine. We just need to wrap the modules within an [autocast](../api/pytorch.rst#transformer_engine.pytorch.autocast) context manager. Note that autocast should only be used to wrap the forward pass and must exit before starting a backward pass. See the [FP8 tutorial](fp8_primer.ipynb) for a detailed explanation of FP8 recipes and the supported options." - ] - }, - { - "cell_type": "code", - "execution_count": 18, - "id": "31256aa7-3d5e-425c-91ab-502b1326a748", - "metadata": {}, - "outputs": [], - "source": [ - "from transformer_engine.common.recipe import Format, DelayedScaling\n", - "\n", - "te_transformer = te.TransformerLayer(hidden_size, ffn_hidden_size, num_attention_heads)\n", - "te_transformer.to(dtype=dtype).cuda()\n", - "utils.share_parameters_with_transformerlayer_te_model(te_transformer, basic_transformer)\n", - "\n", - "fp8_format = Format.HYBRID\n", - "fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo=\"max\")\n", - "torch.manual_seed(1234)\n", - "with te.autocast(enabled=True, fp8_recipe=fp8_recipe):\n", - " y = te_transformer(x, attention_mask=None)" - ] - }, - { - "cell_type": "code", - "execution_count": 19, - "id": "793ebd2d-b84b-47bc-811a-7991df8500aa", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 28.61394775390625 ms\n" - ] - } - ], - "source": [ - "utils.speedometer(\n", - " te_transformer,\n", - " x,\n", - " dy,\n", - " forward_kwargs = { \"attention_mask\": None },\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe },\n", - ")" - ] - } - ], - "metadata": { - "language_info": { - "name": "python" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/examples/quickstart_jax.ipynb b/docs/examples/quickstart_jax.ipynb deleted file mode 100644 index 5b4e439c00..0000000000 --- a/docs/examples/quickstart_jax.ipynb +++ /dev/null @@ -1,833 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "962d87bb", - "metadata": {}, - "source": [ - "\n", - "\n", - "# Getting Started\n", - "\n", - "## Overview\n", - "\n", - "Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, providing better performance with lower memory utilization in both training and inference. It provides support for 8-bit floating point (FP8) precision on Hopper, Ada, as well as 8-bit and 4-bit floating point (NVFP4) precision on Blackwell GPUs, implements a collection of highly optimized building blocks for popular Transformer architectures, and exposes an automatic-mixed-precision-like API that can be used seamlessly with your JAX code. It also includes a framework-agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers.\n", - "\n", - "This guide shows how to start using Transformer Engine with JAX. Similar tutorial for pyTorch is available [here](quickstart.ipynb).\n", - "We recommend you to try understanding the basics of JAX first, using these resources:\n", - "\n", - "- Thinking in JAX: https://docs.jax.dev/en/latest/notebooks/thinking_in_jax.html\n", - "- JAX 101: https://docs.jax.dev/en/latest/jax-101.html\n", - "- Key concepts in JAX: https://docs.jax.dev/en/latest/key-concepts.html#jax-arrays-jax-array\n", - "- Flax 101: https://flax-linen.readthedocs.io/en/latest/guides/flax_fundamentals/index.html\n", - "\n", - "## Let's build a Transformer decoder layer!\n", - "_This is based upon the GPT decoder layer with causal masking, which prevents each position from attending to future positions._\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We build a basic Transformer layer using regular Flax modules. This will be our baseline for later comparisons with Transformer Engine.\n", - "\n", - "
\n", - "\n", - "Let's start with creating the transformer layer using plain [FLAX Linen](https://flax.readthedocs.io/en/stable/) . Figure 1 shows the overall structure.\n", - "\n", - "
\n", - "\n", - "
Figure 1: Structure of a GPT decoder layer.
\n", - "
\n", - "\n", - "We construct the components as follows:\n", - "\n", - "- `LayerNorm`: `nn.LayerNorm` (Flax)\n", - "- `QKV Projection`: `nn.Dense` (conceptually there are three seperate `Dense` layers for Q, K, and V separately, but we fuse them together into a single `Dense` layer that is three times larger)\n", - "- `DotProductAttention`: `nn.MuliheadDotProductAttention` (Flax)\n", - "- `Projection`: `nn.Dense` (Flax)\n", - "- `Dropout`: `nn.Dropout` (Flax)\n", - "- `MLP`: `FlaxMLP` implemented using `nn.Dense` and `nn.gelu`\n", - "\n", - "Over the course of this tutorial we will use a few modules and helper functions defined in [quickstart_jax_utils.py](quickstart_jax_utils.py). Putting it all together: \n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "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": 2, - "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", - "\n", - " @nn.compact\n", - " def __call__(self, x: jnp.ndarray) -> jnp.ndarray:\n", - " x = nn.Dense(features=self.ffn_hidden_size, use_bias=True)(x)\n", - " x = nn.gelu(x, approximate=True) # equivalent to tanh approximation\n", - " x = nn.Dense(features=self.hidden_size, use_bias=True)(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", - " \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)(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", - " 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)(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", - " )\n", - " x = mlp(x)\n", - " \n", - " return x + res\n" - ] - }, - { - "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: 19.258604049682617 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": "ccb16f31", - "metadata": {}, - "source": [ - "## Meet Transformer Engine\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "Now that we have a basic Transformer layer in Flax, let's use Transformer Engine to speed up the training. The following examples show how to use TE modules.\n", - "\n", - "
\n", - "\n", - "As a reminder, the FlaxTransformerLayer above used:\n", - "\n", - "- `nn.LayerNorm`: Flax LayerNorm\n", - "- `nn.Dense`: Flax Dense layer for QKV projection \n", - "- `nn.MultiheadDotProductAttention`: Flax MultiheadDotProductAttention\n", - "- `nn.Dense`: Flax Dense layer for projection\n", - "- `nn.Dropout`: Flax Dropout\n", - "- `FlaxMLP`: Custom MLP implemented from `nn.Dense`\n", - "\n", - "Below we show how to use Transformer Engine Flax modules for better performance:\n" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "bed20d6b", - "metadata": {}, - "outputs": [], - "source": [ - "import transformer_engine.jax as te\n", - "import transformer_engine.jax.flax as te_flax" - ] - }, - { - "cell_type": "markdown", - "id": "f28cb444", - "metadata": {}, - "source": [ - "TE provides a set of Flax Linen modules that can be used to build Transformer layers. The simplest of the provided modules are the `DenseGeneral ` and `LayerNorm` layers, which we can use instead of `flax.linen.Dense` and ` flax.linen.LayerNorm`. Let's modify our `FlaxTransformerLayer`:" - ] - }, - { - "cell_type": "code", - "execution_count": 8, - "id": "56105579", - "metadata": {}, - "outputs": [], - "source": [ - "class TEUnfusedMLP(nn.Module):\n", - " hidden_size : int\n", - " ffn_hidden_size: int\n", - "\n", - " @nn.compact\n", - " def __call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray:\n", - " x = te_flax.DenseGeneral(features=self.ffn_hidden_size, use_bias=True) (x)\n", - " x = x.reshape(*x.shape[:-1], 1, x.shape[-1])\n", - " x = te.activation.activation(x, activation_type=('gelu',))\n", - " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True) (x)\n", - " return x\n", - "\n", - "class TEUnfusedTransformerLayer(nn.Module):\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", - " use_te_attention: bool = True # True for TE attention, False for Flax attention\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", - " res = x\n", - " x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x)\n", - "\n", - " # Fused QKV projection\n", - " qkv = te_flax.DenseGeneral(features=3 * self.hidden_size, use_bias=True)(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", - " # Attention - either TE or Flax implementation\n", - " if self.use_te_attention:\n", - " # Use TE's DotProductAttention\n", - " attention = te_flax.DotProductAttention(\n", - " head_dim=self.kv_channels,\n", - " num_attention_heads=self.num_attention_heads,\n", - " num_gqa_groups=self.num_attention_heads, # No GQA\n", - " attention_dropout=self.attention_dropout,\n", - " attn_mask_type='causal',\n", - " )\n", - " x = attention(\n", - " q, k, v,\n", - " # Causal mask does not need an explicit instatiated mask as specialized kernels exist to handle it\n", - " sequence_descriptor=None, \n", - " deterministic=deterministic\n", - " )\n", - " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", - " x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))\n", - " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x)\n", - " x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic)\n", - " else:\n", - " # Use Flax's MultiHeadDotProductAttention\n", - " q_reshaped = q.reshape(q.shape[0], q.shape[1], self.hidden_size)\n", - " k_reshaped = k.reshape(k.shape[0], k.shape[1], self.hidden_size)\n", - " v_reshaped = v.reshape(v.shape[0], v.shape[1], self.hidden_size)\n", - "\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", - " attention = nn.MultiHeadDotProductAttention(\n", - " num_heads=self.num_attention_heads,\n", - " qkv_features=self.kv_channels,\n", - " dropout_rate=self.attention_dropout,\n", - " )\n", - " x = attention(q_reshaped, k_reshaped, v_reshaped, mask=attention_mask, deterministic=deterministic)\n", - "\n", - " x = res + x\n", - "\n", - " # Second residual connection\n", - " res = x\n", - " x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x)\n", - "\n", - " # MLP\n", - " mlp = TEUnfusedMLP(\n", - " hidden_size=self.hidden_size,\n", - " ffn_hidden_size=self.ffn_hidden_size\n", - " )\n", - "\n", - " x = mlp(x, deterministic=deterministic)\n", - "\n", - " return x + res" - ] - }, - { - "cell_type": "markdown", - "id": "a76911ac", - "metadata": {}, - "source": [ - "Testing performance of the model, using `DenseGeneral`, `LayerNorm` and activation from TE, while keeping Flax's `MultiHeadDotProductAttention` the same as the first simple Transformer in JAX implementation. To read more about this implementation from Flax, you can refer to this documentation: https://flax.readthedocs.io/en/latest/api_reference/flax.nnx/nn/attention.html" - ] - }, - { - "cell_type": "code", - "execution_count": 9, - "id": "4b67511f", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 16.003193855285645 ms\n" - ] - } - ], - "source": [ - "te_unfused_transformer_with_flax_MHA = TEUnfusedTransformerLayer(\n", - " hidden_size, \n", - " ffn_hidden_size, \n", - " num_attention_heads,\n", - " use_te_attention=False\n", - ")\n", - "\n", - "te_params = te_unfused_transformer_with_flax_MHA.init(key, x, attention_mask=None, deterministic=False)\n", - "\n", - "utils.speedometer(\n", - " model_apply_fn=te_unfused_transformer_with_flax_MHA.apply,\n", - " variables=te_params, # Ensure the correct `params` is passed\n", - " input=x,\n", - " output_grad=dy,\n", - " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " rngs={\"dropout\": dropout_key},\n", - ")\n" - ] - }, - { - "cell_type": "markdown", - "id": "0b230058", - "metadata": {}, - "source": [ - "Now, we move on to also replace the attention sub-layer with TE's `DotProductAttention` implementation" - ] - }, - { - "cell_type": "code", - "execution_count": 10, - "id": "5146cd99", - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n", - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 8.897695541381836 ms\n" - ] - } - ], - "source": [ - "te_unfused_transformer = TEUnfusedTransformerLayer(\n", - " hidden_size, \n", - " ffn_hidden_size, \n", - " num_attention_heads,\n", - ")\n", - "\n", - "te_params = te_unfused_transformer.init(key, x, attention_mask=None, deterministic=False)\n", - "\n", - "utils.speedometer(\n", - " model_apply_fn=te_unfused_transformer.apply,\n", - " variables=te_params, # Ensure the correct `params` is passed\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": "c9a101d3", - "metadata": {}, - "source": [ - "## Enabling Quantization (FP8 or FP4)\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We configure a TE module to perform compute in FP8.\n", - "\n", - "
\n", - "\n", - "Enabling FP8 support is very simple in Transformer Engine. We just need to wrap the modules within an [autocast](../api/jax.rst#transformer_engine.jax.fp8_autocast) context manager. See the [FP8 tutorial](fp8_primer.ipynb) for a detailed explanation of FP8 recipes and the supported options.\n", - "\n", - "
\n", - "\n", - "Important: FP8 Metadata Initialization\n", - "\n", - "When using FP8, the model **must be initialized within the `autocast` context**. This creates a special collection called `fp8_metas` that contains scaling factors and other metadata required for FP8 computation. If you initialize a model outside of `autocast` and then try to use it with FP8, you will get a `ScopeCollectionNotFound` error because the `fp8_metas` collection was never created.\n", - "\n", - "
" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "c2eee376", - "metadata": {}, - "outputs": [], - "source": [ - "from transformer_engine.common.recipe import Format, DelayedScaling\n", - "fp8_format = Format.HYBRID\n", - "fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo=\"max\")" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "de96827c", - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n", - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n", - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 5.651178359985352 ms\n" - ] - } - ], - "source": [ - "with te.autocast(enabled=True, recipe=fp8_recipe):\n", - " te_unfused_params = te_unfused_transformer.init(key, x, attention_mask=None, deterministic=False)\n", - "\n", - " # Example usage of forward \n", - " y = te_unfused_transformer.apply(te_unfused_params, x, attention_mask=None, deterministic=True)\n", - "\n", - "utils.speedometer(\n", - " model_apply_fn=te_unfused_transformer.apply,\n", - " variables=te_unfused_params, # Ensure the correct `params` is passed\n", - " input=x,\n", - " output_grad=dy,\n", - " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe},\n", - " rngs={\"dropout\": dropout_key},\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "3801b201", - "metadata": {}, - "source": [ - "\n", - "## Fused TE Modules\n", - "\n", - "
\n", - "\n", - "Summary\n", - " \n", - "We optimize the example Transformer layer with TE modules for fused operations.\n", - "\n", - "
\n", - "\n", - "The `DenseGeneral` layer is enough to build any Transformer model and it enables usage of the Transformer Engine even for very custom Transformers. However, having more knowledge about the model allows for additional optimizations such as kernel fusions in mixed-precision recipes, increasing the achievable speedup.\n", - "\n", - "Transformer Engine therefore provides coarser modules that span multiple layers:\n", - "\n", - "* `LayerNormDenseGeneral`\n", - "* `LayerNormMLP`\n", - "* `TransformerLayer`\n", - "\n", - "To see a complete list of all the functions TE Flax support, you can view it here: https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/jax.html#modules\n", - "\n", - "Building a third iteration of our Transformer layer with `LayerNormDenseGeneral` and `LayerNormMLP`:" - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "11203785", - "metadata": {}, - "outputs": [], - "source": [ - "class TEFusedTransformerLayer(nn.Module):\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", - "\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", - " res = x\n", - "\n", - " # Fused QKV projection\n", - " qkv,_ = te_flax.LayerNormDenseGeneral(features=3 * self.hidden_size, \n", - " epsilon=self.layernorm_eps, \n", - " use_bias=True, \n", - " return_layernorm_output=False)(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", - " # Attention using TE's DotProductAttention\n", - " attention = te_flax.DotProductAttention(\n", - " head_dim=self.kv_channels,\n", - " num_attention_heads=self.num_attention_heads,\n", - " num_gqa_groups=self.num_attention_heads, \n", - " attention_dropout=self.attention_dropout,\n", - " attn_mask_type='causal',\n", - " )\n", - " x = attention(q, k, v, sequence_descriptor=None, deterministic=deterministic)\n", - " # Reshape from [batch, seq_len, num_heads, head_dim] to [batch, seq_len, hidden_size]\n", - " x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3]))\n", - " x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x)\n", - " x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic)\n", - "\n", - " x = res + x\n", - "\n", - " # Second residual connection\n", - " res = x\n", - " x,_ = te_flax.LayerNormMLP(intermediate_dim=self.ffn_hidden_size, \n", - " epsilon=self.layernorm_eps,\n", - " use_bias=True,\n", - " activations=('gelu',),\n", - " intermediate_dropout_rate=0.0,\n", - " return_layernorm_output=False\n", - " )(x, deterministic=deterministic)\n", - "\n", - " return x + res" - ] - }, - { - "cell_type": "markdown", - "id": "334cff59", - "metadata": {}, - "source": [ - "Similar to the unnfused model, we also compare the performance of fused model when using Flax's MultiheadDotProductAttention implementation and TE's." - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "id": "6b0c705e", - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n", - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n", - "/mnt/jberchtold/ptyche-lustre-home/transformerengine/transformer_engine/jax/flax/transformer.py:626: UserWarning: transpose_batch_sequence defaults to False in DotProductAttention starting TransformerEngine v2.10\n", - " warnings.warn(\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 5.493879318237305 ms\n" - ] - } - ], - "source": [ - "te_fused_transformer = TEFusedTransformerLayer(\n", - " hidden_size, \n", - " ffn_hidden_size, \n", - " num_attention_heads\n", - ")\n", - "\n", - "with te.autocast(enabled=True, recipe=fp8_recipe):\n", - " te_fused_params = te_fused_transformer.init(key, x, attention_mask=None, deterministic=False)\n", - " # Example usage of forward \n", - " y = te_fused_transformer.apply(te_fused_params, x, attention_mask=None, deterministic=True)\n", - "\n", - "utils.speedometer(\n", - " model_apply_fn=te_fused_transformer.apply,\n", - " variables=te_fused_params,\n", - " input=x,\n", - " output_grad=dy,\n", - " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe},\n", - " rngs={\"dropout\": dropout_key},\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "a45c12c8", - "metadata": {}, - "source": [ - "Finally, the `TransformerLayer` module is convenient for creating standard Transformer architectures." - ] - }, - { - "cell_type": "code", - "execution_count": 15, - "id": "b2aaa8ef", - "metadata": {}, - "outputs": [], - "source": [ - "te_transformer = te_flax.TransformerLayer(\n", - " hidden_size=hidden_size,\n", - " mlp_hidden_size=ffn_hidden_size, \n", - " num_attention_heads=num_attention_heads,\n", - " mlp_activations=(\"gelu\",),\n", - " self_attn_mask_type='causal',\n", - " layernorm_epsilon=1e-5,\n", - " use_bias=True,\n", - " intermediate_dropout=0.0,\n", - " enable_relative_embedding=False,\n", - " self_attn_bias_type='no_bias',\n", - " hidden_dropout=0.0,\n", - ")\n", - "\n", - "with te.autocast(enabled=True, recipe=fp8_recipe):\n", - " te_transformer_params = te_transformer.init(key, x, deterministic=False)\n", - " y = te_transformer.apply(te_transformer_params, x, attention_mask=None, deterministic=True)" - ] - }, - { - "cell_type": "code", - "execution_count": 20, - "id": "b9cdbf22", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 5.334172248840332 ms\n" - ] - } - ], - "source": [ - "utils.speedometer(\n", - " model_apply_fn=te_transformer.apply,\n", - " model_init_fn=te_transformer.init,\n", - " variables=te_transformer_params,\n", - " input=x,\n", - " output_grad=dy,\n", - " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " autocast_kwargs = { \"enabled\": True, \"recipe\": fp8_recipe },\n", - " rngs={\"dropout\": dropout_key},\n", - ")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.12.3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/examples/quickstart_jax_utils.py b/docs/examples/quickstart_jax_utils.py index 3a2947f338..0c5ec5295e 100644 --- a/docs/examples/quickstart_jax_utils.py +++ b/docs/examples/quickstart_jax_utils.py @@ -5,13 +5,9 @@ import jax import jax.numpy as jnp import time -import math from typing import Callable, Any, Dict, Optional, Tuple -from flax import linen as nn import transformer_engine.jax as te -import transformer_engine.jax.flax as te_flax -from transformer_engine.jax.flax.transformer import DotProductAttention as TEDotProductAttention def speedometer( diff --git a/docs/examples/te_jax_integration.ipynb b/docs/examples/te_jax_integration.ipynb index 70647e421a..66d16ed52f 100644 --- a/docs/examples/te_jax_integration.ipynb +++ b/docs/examples/te_jax_integration.ipynb @@ -264,7 +264,7 @@ "id": "5e9310c9", "metadata": {}, "source": [ - "# Transformer Engine" + "## Transformer Engine" ] }, { diff --git a/docs/getting_started.rst b/docs/getting_started.rst deleted file mode 100644 index f5ab2ae695..0000000000 --- a/docs/getting_started.rst +++ /dev/null @@ -1,16 +0,0 @@ -.. - Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - - See LICENSE for license information. - -Getting Started -=============== - -Choose your framework to get started with Transformer Engine: - -.. toctree:: - :maxdepth: 1 - - PyTorch - JAX - diff --git a/docs/getting_started/getting_started_jax.out b/docs/getting_started/getting_started_jax.out new file mode 100644 index 0000000000..c11f3b1965 --- /dev/null +++ b/docs/getting_started/getting_started_jax.out @@ -0,0 +1,34 @@ +pyxis: importing docker image: gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-jax-py3-devel +pyxis: imported docker image: gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-jax-py3-devel +# BENCHMARK_BASELINE_OUTPUT_START +Baseline Flax: +Mean time: 86.580 ms +# BENCHMARK_BASELINE_OUTPUT_END + +# BENCHMARK_TE_UNFUSED_OUTPUT_START +TE Unfused: +Mean time: 42.252 ms +# BENCHMARK_TE_UNFUSED_OUTPUT_END + +# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_START +TE Unfused + TE Attention: +Mean time: 35.054 ms +# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_END + +# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_START +TE Unfused + TE Attention + FP8: +Mean time: 22.638 ms +# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_END + +# BENCHMARK_TE_FUSED_FP8_OUTPUT_START +TE Fused + TE Attention + FP8: +Mean time: 23.703 ms +# BENCHMARK_TE_FUSED_FP8_OUTPUT_END + +# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_START +TE TransformerLayer + FP8: +Mean time: 22.812 ms +# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_END + + +Summary written to getting_started_jax_summary.csv diff --git a/docs/getting_started/getting_started_jax.py b/docs/getting_started/getting_started_jax.py new file mode 100644 index 0000000000..88ea9f6dc8 --- /dev/null +++ b/docs/getting_started/getting_started_jax.py @@ -0,0 +1,523 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Getting Started with Transformer Engine - JAX Example +====================================================== + +This example shows how to build a Transformer decoder layer using JAX/Flax +and how to optimize it with Transformer Engine. +""" + +import jax +import jax.numpy as jnp +from flax import linen as nn +from typing import Optional + +import transformer_engine.jax as te +import transformer_engine.jax.flax as te_flax +from transformer_engine.jax.sharding import MeshResource +from transformer_engine.common.recipe import Format, DelayedScaling + +from getting_started_utils_jax import speedometer + + +# Configuration +hidden_size = 4096 +sequence_length = 2048 +batch_size = 8 +ffn_hidden_size = 16384 +num_attention_heads = 32 +dtype = jnp.bfloat16 + +# Create synthetic data +key = jax.random.PRNGKey(42) +x = jax.random.normal(key, (batch_size, sequence_length, hidden_size)).astype(dtype) +mesh_resource = MeshResource() + + +# ============================================================================= +# Baseline: Pure Flax Implementation +# ============================================================================= + + +# BASELINE_MLP_START +class FlaxMLP(nn.Module): + """Feed-forward network in Transformer layer. + Built with plain Flax modules. + """ + + hidden_size: int + ffn_hidden_size: int + + @nn.compact + def __call__(self, x: jnp.ndarray) -> jnp.ndarray: + x = nn.Dense(features=self.ffn_hidden_size, use_bias=True)(x) + x = nn.gelu(x, approximate=True) + x = nn.Dense(features=self.hidden_size, use_bias=True)(x) + return x + + +# BASELINE_MLP_END + + +# BASELINE_LAYER_START +class FlaxTransformerLayer(nn.Module): + """Basic Transformer layer using plain Flax modules.""" + + hidden_size: int + ffn_hidden_size: int + num_attention_heads: int + layernorm_eps: float = 1e-5 + attention_dropout: float = 0.1 + + def setup(self): + self.kv_channels = self.hidden_size // self.num_attention_heads + + @nn.compact + def __call__( + self, + x: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = False, + ) -> jnp.ndarray: + if attention_mask is None: + attention_mask = nn.make_causal_mask(x[..., 0], dtype=jnp.bool_) + + res = x + x = nn.LayerNorm(epsilon=self.layernorm_eps)(x) + + # Fused QKV projection + qkv = nn.Dense(features=3 * self.hidden_size, use_bias=True)(x) + qkv = qkv.reshape( + qkv.shape[0], qkv.shape[1], self.num_attention_heads, 3 * self.kv_channels + ) + q, k, v = jnp.split(qkv, 3, axis=3) + + dropout_rng = None + if not deterministic and self.attention_dropout > 0: + dropout_rng = self.make_rng("dropout") + + x = nn.dot_product_attention( + query=q, + key=k, + value=v, + mask=attention_mask, + dropout_rng=dropout_rng, + dropout_rate=self.attention_dropout, + deterministic=deterministic, + broadcast_dropout=True, + ) + + x = x.reshape(x.shape[0], x.shape[1], self.hidden_size) + x = nn.Dense(features=self.hidden_size, use_bias=True)(x) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + x = res + x + + res = x + x = nn.LayerNorm(epsilon=self.layernorm_eps)(x) + mlp = FlaxMLP(hidden_size=self.hidden_size, ffn_hidden_size=self.ffn_hidden_size) + x = mlp(x) + + return x + res + + +# BASELINE_LAYER_END + + +print("# BENCHMARK_BASELINE_OUTPUT_START") +# BENCHMARK_BASELINE_START +baseline = FlaxTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, +) +params = baseline.init(key, x, deterministic=False) + +print("Baseline Flax:") +time_baseline = speedometer( + baseline.apply, params, x, forward_kwargs={"deterministic": True}, label="baseline" +) +# BENCHMARK_BASELINE_END +print("# BENCHMARK_BASELINE_OUTPUT_END\n") + + +# ============================================================================= +# TE Unfused: Basic TE Modules +# ============================================================================= + + +# TE_UNFUSED_MLP_START +class TEUnfusedMLP(nn.Module): + """MLP using TE modules.""" + + hidden_size: int + ffn_hidden_size: int + + @nn.compact + def __call__(self, x: jnp.ndarray, deterministic: bool) -> jnp.ndarray: + x = te_flax.DenseGeneral(features=self.ffn_hidden_size, use_bias=True)(x) + x = x.reshape(*x.shape[:-1], 1, x.shape[-1]) + x = te.activation.activation(x, activation_type=("gelu",)) + x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x) + return x + + +# TE_UNFUSED_MLP_END + + +# TE_UNFUSED_LAYER_START +class TEUnfusedTransformerLayer(nn.Module): + """Transformer layer using basic TE modules (without TE attention).""" + + hidden_size: int + ffn_hidden_size: int + num_attention_heads: int + layernorm_eps: float = 1e-5 + attention_dropout: float = 0.1 + + def setup(self): + self.kv_channels = self.hidden_size // self.num_attention_heads + + @nn.compact + def __call__( + self, + x: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = False, + ) -> jnp.ndarray: + if attention_mask is None: + attention_mask = nn.make_causal_mask(x[..., 0], dtype=jnp.bool_) + + res = x + x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x) + + qkv = te_flax.DenseGeneral(features=3 * self.hidden_size, use_bias=True)(x) + qkv = qkv.reshape( + qkv.shape[0], qkv.shape[1], self.num_attention_heads, 3 * self.kv_channels + ) + q, k, v = jnp.split(qkv, 3, axis=3) + + dropout_rng = None + if not deterministic and self.attention_dropout > 0: + dropout_rng = self.make_rng("dropout") + + x = nn.dot_product_attention( + query=q, + key=k, + value=v, + mask=attention_mask, + dropout_rng=dropout_rng, + dropout_rate=self.attention_dropout, + deterministic=deterministic, + broadcast_dropout=True, + ) + + x = x.reshape(x.shape[0], x.shape[1], self.hidden_size) + x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + + x = res + x + + res = x + x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x) + mlp = TEUnfusedMLP(hidden_size=self.hidden_size, ffn_hidden_size=self.ffn_hidden_size) + x = mlp(x, deterministic=deterministic) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + + return x + res + + +# TE_UNFUSED_LAYER_END + + +print("# BENCHMARK_TE_UNFUSED_OUTPUT_START") +# BENCHMARK_TE_UNFUSED_START +te_unfused = TEUnfusedTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, +) +params = te_unfused.init(key, x, deterministic=False) + +print("TE Unfused:") +time_te_unfused = speedometer( + te_unfused.apply, params, x, forward_kwargs={"deterministic": True}, label="te_unfused" +) +# BENCHMARK_TE_UNFUSED_END +print("# BENCHMARK_TE_UNFUSED_OUTPUT_END\n") + + +# ============================================================================= +# TE Unfused + TE Attention +# ============================================================================= + + +# TE_UNFUSED_ATTN_LAYER_START +class TEUnfusedAttnTransformerLayer(nn.Module): + """Transformer layer using TE modules including TE DotProductAttention.""" + + hidden_size: int + ffn_hidden_size: int + num_attention_heads: int + layernorm_eps: float = 1e-5 + attention_dropout: float = 0.1 + + def setup(self): + self.kv_channels = self.hidden_size // self.num_attention_heads + + @nn.compact + def __call__( + self, + x: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = False, + ) -> jnp.ndarray: + res = x + x = te_flax.LayerNorm(epsilon=self.layernorm_eps, dtype=jnp.bfloat16)(x) + + qkv = te_flax.DenseGeneral( + features=3 * self.hidden_size, use_bias=True, dtype=jnp.bfloat16 + )(x) + qkv = qkv.reshape( + qkv.shape[0], qkv.shape[1], self.num_attention_heads, 3 * self.kv_channels + ) + q, k, v = jnp.split(qkv, 3, axis=3) + + attention = te_flax.DotProductAttention( + head_dim=self.kv_channels, + num_attention_heads=self.num_attention_heads, + num_gqa_groups=self.num_attention_heads, + attention_dropout=self.attention_dropout, + attn_mask_type="causal", + transpose_batch_sequence=False, + ) + x = attention(q, k, v, deterministic=deterministic) + x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3])) + x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True, dtype=jnp.bfloat16)(x) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + + x = res + x + + res = x + x = te_flax.LayerNorm(epsilon=self.layernorm_eps)(x) + mlp = TEUnfusedMLP(hidden_size=self.hidden_size, ffn_hidden_size=self.ffn_hidden_size) + x = mlp(x, deterministic=deterministic) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + + return x + res + + +# TE_UNFUSED_ATTN_LAYER_END + + +print("# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_START") +# BENCHMARK_TE_UNFUSED_ATTN_START +te_unfused_attn = TEUnfusedAttnTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, +) + +with te.autocast(enabled=False, mesh_resource=mesh_resource): + params = te_unfused_attn.init(key, x, deterministic=False) + +print("TE Unfused + TE Attention:") +time_te_unfused_attn = speedometer( + te_unfused_attn.apply, + params, + x, + forward_kwargs={"deterministic": True}, + autocast_kwargs={"enabled": False, "mesh_resource": mesh_resource}, + label="te_unfused_attn", +) +# BENCHMARK_TE_UNFUSED_ATTN_END +print("# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_END\n") + + +# ============================================================================= +# TE Unfused + FP8 +# ============================================================================= + +print("# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_START") +# BENCHMARK_TE_UNFUSED_FP8_START +recipe = DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=16, amax_compute_algo="max") + +te_unfused_fp8 = TEUnfusedAttnTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, +) + +with te.autocast(enabled=True, recipe=recipe, mesh_resource=mesh_resource): + params = te_unfused_fp8.init(key, x, deterministic=False) + +print("TE Unfused + TE Attention + FP8:") +time_te_unfused_fp8 = speedometer( + te_unfused_fp8.apply, + params, + x, + forward_kwargs={"deterministic": True}, + autocast_kwargs={"enabled": True, "recipe": recipe, "mesh_resource": mesh_resource}, + label="te_unfused_fp8", +) +# BENCHMARK_TE_UNFUSED_FP8_END +print("# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_END\n") + + +# ============================================================================= +# TE Fused + FP8: Optimized Modules with FP8 +# ============================================================================= + + +# TE_FUSED_LAYER_START +class TEFusedTransformerLayer(nn.Module): + """Transformer layer using fused TE modules for better performance.""" + + hidden_size: int + ffn_hidden_size: int + num_attention_heads: int + layernorm_eps: float = 1e-5 + attention_dropout: float = 0.1 + + def setup(self): + self.kv_channels = self.hidden_size // self.num_attention_heads + + @nn.compact + def __call__( + self, + x: jnp.ndarray, + attention_mask: Optional[jnp.ndarray] = None, + deterministic: bool = False, + ) -> jnp.ndarray: + res = x + + # Fused LayerNorm + QKV projection + qkv, _ = te_flax.LayerNormDenseGeneral( + features=3 * self.hidden_size, + epsilon=self.layernorm_eps, + use_bias=True, + return_layernorm_output=False, + )(x) + qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, self.num_attention_heads, self.kv_channels) + q, k, v = qkv[:, :, 0, :, :], qkv[:, :, 1, :, :], qkv[:, :, 2, :, :] + + attention = te_flax.DotProductAttention( + head_dim=self.kv_channels, + num_attention_heads=self.num_attention_heads, + num_gqa_groups=self.num_attention_heads, + attention_dropout=self.attention_dropout, + attn_mask_type="causal", + qkv_layout="bshd_bshd_bshd", + transpose_batch_sequence=False, + ) + x = attention(q, k, v, deterministic=deterministic) + x = x.reshape((x.shape[0], x.shape[1], x.shape[2] * x.shape[3])) + x = te_flax.DenseGeneral(features=self.hidden_size, use_bias=True)(x) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + + x = res + x + + res = x + # Fused LayerNorm + MLP + x, _ = te_flax.LayerNormMLP( + intermediate_dim=self.ffn_hidden_size, + epsilon=self.layernorm_eps, + use_bias=True, + activations=("gelu",), + intermediate_dropout_rate=0.0, + return_layernorm_output=False, + )(x, deterministic=deterministic) + x = nn.Dropout(rate=self.attention_dropout)(x, deterministic=deterministic) + + return x + res + + +# TE_FUSED_LAYER_END + + +print("# BENCHMARK_TE_FUSED_FP8_OUTPUT_START") +# BENCHMARK_TE_FUSED_FP8_START +te_fused_fp8 = TEFusedTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, +) + +with te.autocast(enabled=True, recipe=recipe, mesh_resource=mesh_resource): + params = te_fused_fp8.init(key, x, deterministic=False) + +print("TE Fused + TE Attention + FP8:") +time_te_fused_fp8 = speedometer( + te_fused_fp8.apply, + params, + x, + forward_kwargs={"deterministic": True}, + autocast_kwargs={"enabled": True, "recipe": recipe, "mesh_resource": mesh_resource}, + label="te_fused_fp8", +) +# BENCHMARK_TE_FUSED_FP8_END +print("# BENCHMARK_TE_FUSED_FP8_OUTPUT_END\n") + + +# ============================================================================= +# TE TransformerLayer + FP8: Ready-to-use Module +# ============================================================================= + +print("# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_START") +# BENCHMARK_TE_TRANSFORMER_LAYER_START +te_transformer_layer = te_flax.TransformerLayer( + hidden_size=hidden_size, + mlp_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + mlp_activations=("gelu",), + self_attn_mask_type="causal", + layernorm_epsilon=1e-5, + use_bias=True, + attention_dropout=0.0, + intermediate_dropout=0.0, + hidden_dropout=0.0, + enable_relative_embedding=False, + self_attn_bias_type="no_bias", + dtype=jnp.bfloat16, + transpose_batch_sequence=False, +) + +with te.autocast(enabled=True, recipe=recipe, mesh_resource=mesh_resource): + params = te_transformer_layer.init(key, x, deterministic=False) + +print("TE TransformerLayer + FP8:") +time_te_transformer_layer = speedometer( + te_transformer_layer.apply, + params, + x, + forward_kwargs={"deterministic": True}, + autocast_kwargs={"enabled": True, "recipe": recipe, "mesh_resource": mesh_resource}, + label="te_transformer_layer", +) +# BENCHMARK_TE_TRANSFORMER_LAYER_END +print("# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_END\n") + +# Write summary CSV for RST documentation +with open("getting_started_jax_summary.csv", "w") as f: + f.write("Implementation,Time (ms),Speedup\n") + f.write(f"Baseline Flax,{time_baseline:.2f},1.00x\n") + f.write(f"TE Unfused,{time_te_unfused:.2f},{time_baseline/time_te_unfused:.2f}x\n") + f.write( + "TE Unfused + TE" + f" Attention,{time_te_unfused_attn:.2f},{time_baseline/time_te_unfused_attn:.2f}x\n" + ) + f.write( + "TE Unfused + TE Attention +" + f" FP8,{time_te_unfused_fp8:.2f},{time_baseline/time_te_unfused_fp8:.2f}x\n" + ) + f.write( + "TE Fused + TE Attention +" + f" FP8,{time_te_fused_fp8:.2f},{time_baseline/time_te_fused_fp8:.2f}x\n" + ) + f.write( + "TE TransformerLayer +" + f" FP8,{time_te_transformer_layer:.2f},{time_baseline/time_te_transformer_layer:.2f}x\n" + ) +print("\nSummary written to getting_started_jax_summary.csv") diff --git a/docs/getting_started/getting_started_jax_summary.csv b/docs/getting_started/getting_started_jax_summary.csv new file mode 100644 index 0000000000..5b6a4249b3 --- /dev/null +++ b/docs/getting_started/getting_started_jax_summary.csv @@ -0,0 +1,7 @@ +Implementation,Time (ms),Speedup +Baseline Flax,86.58,1.00x +TE Unfused,42.25,2.05x +TE Unfused + TE Attention,35.05,2.47x +TE Unfused + TE Attention + FP8,22.64,3.82x +TE Fused + TE Attention + FP8,23.70,3.65x +TE TransformerLayer + FP8,22.81,3.80x diff --git a/docs/getting_started/getting_started_pytorch.out b/docs/getting_started/getting_started_pytorch.out new file mode 100644 index 0000000000..9b9387a8b2 --- /dev/null +++ b/docs/getting_started/getting_started_pytorch.out @@ -0,0 +1,42 @@ +pyxis: importing docker image: gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-pytorch-py3-devel-amd64 +pyxis: imported docker image: gitlab-master.nvidia.com/dl/transformerengine/transformerengine:main-pytorch-py3-devel-amd64 +/usr/local/lib/python3.12/dist-packages/torch/library.py:357: UserWarning: Warning only once for all operators, other operators may also be overridden. + Overriding a previously registered kernel for the same operator and the same dispatch key + operator: flash_attn::_flash_attn_backward(Tensor dout, Tensor q, Tensor k, Tensor v, Tensor out, Tensor softmax_lse, Tensor(a6!)? dq, Tensor(a7!)? dk, Tensor(a8!)? dv, float dropout_p, float softmax_scale, bool causal, SymInt window_size_left, SymInt window_size_right, float softcap, Tensor? alibi_slopes, bool deterministic, Tensor? rng_state=None) -> Tensor + registered at /usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py:926 + dispatch key: ADInplaceOrView + previous kernel: no debug info + new kernel: registered at /usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py:926 (Triggered internally at /opt/pytorch/pytorch/aten/src/ATen/core/dispatch/OperatorEntry.cpp:208.) + self.m.impl( +# BENCHMARK_BASELINE_OUTPUT_START +Baseline PyTorch: +Mean time: 48.280 ms +# BENCHMARK_BASELINE_OUTPUT_END + +# BENCHMARK_TE_UNFUSED_OUTPUT_START +TE Unfused: +Mean time: 49.342 ms +# BENCHMARK_TE_UNFUSED_OUTPUT_END + +# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_START +TE Unfused + TE Attention: +Mean time: 35.709 ms +# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_END + +# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_START +TE Unfused + TE Attention + FP8: +Mean time: 23.406 ms +# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_END + +# BENCHMARK_TE_FUSED_FP8_OUTPUT_START +TE Fused + TE Attention + FP8: +Mean time: 22.964 ms +# BENCHMARK_TE_FUSED_FP8_OUTPUT_END + +# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_START +TE TransformerLayer + FP8: +Mean time: 21.670 ms +# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_END + + +Summary written to getting_started_pytorch_summary.csv diff --git a/docs/getting_started/getting_started_pytorch.py b/docs/getting_started/getting_started_pytorch.py new file mode 100644 index 0000000000..bbffd300c8 --- /dev/null +++ b/docs/getting_started/getting_started_pytorch.py @@ -0,0 +1,497 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Getting Started with Transformer Engine - PyTorch Example +========================================================== + +This example shows how to build a Transformer layer using PyTorch +and how to optimize it with Transformer Engine. +""" + +from typing import Optional +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Format, DelayedScaling + +from getting_started_utils_pytorch import DotProductAttention, speedometer + + +# Configuration +hidden_size = 4096 +sequence_length = 2048 +batch_size = 8 +ffn_hidden_size = 16384 +num_attention_heads = 32 +dtype = torch.bfloat16 + +# Create synthetic data +x = torch.rand(sequence_length, batch_size, hidden_size).cuda().to(dtype=dtype) + + +# ============================================================================= +# Baseline: Pure PyTorch Implementation +# ============================================================================= + + +# BASELINE_MLP_START +class PyTorchMLP(torch.nn.Module): + """Feed-forward network in Transformer layer. + Built with plain PyTorch modules. + """ + + hidden_size: int + ffn_hidden_size: int + + def __init__(self, hidden_size: int, ffn_hidden_size: int) -> None: + super().__init__() + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + self.linear1 = torch.nn.Linear(hidden_size, ffn_hidden_size, bias=True) + self.linear2 = torch.nn.Linear(ffn_hidden_size, hidden_size, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear1(x) + x = torch.nn.functional.gelu(x, approximate="tanh") + x = self.linear2(x) + return x + + +# BASELINE_MLP_END + + +# BASELINE_LAYER_START +class PyTorchTransformerLayer(torch.nn.Module): + """Basic Transformer layer using plain PyTorch modules.""" + + def __init__( + self, + hidden_size: int, + ffn_hidden_size: int, + num_attention_heads: int, + layernorm_eps: float = 1e-5, + attention_dropout: float = 0.1, + hidden_dropout: float = 0.1, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.kv_channels = hidden_size // num_attention_heads + self.ln1 = torch.nn.LayerNorm(hidden_size, eps=layernorm_eps) + self.qkv_projection = torch.nn.Linear(hidden_size, 3 * hidden_size, bias=True) + self.attention = DotProductAttention( + num_attention_heads=num_attention_heads, + kv_channels=self.kv_channels, + attention_dropout=attention_dropout, + ) + self.projection = torch.nn.Linear(hidden_size, hidden_size, bias=True) + self.dropout = torch.nn.Dropout(hidden_dropout) + self.ln2 = torch.nn.LayerNorm(hidden_size, eps=layernorm_eps) + self.mlp = PyTorchMLP(hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size) + + def forward( + self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + res = x + x = self.ln1(x) + + # Fused QKV projection + qkv = self.qkv_projection(x) + qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels) + q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3) + + x = self.attention(q, k, v, attention_mask) + x = self.projection(x) + x = self.dropout(x) + x = res + x + + # Second residual connection + res = x + x = self.ln2(x) + x = self.mlp(x) + + return x + res + + +# BASELINE_LAYER_END + + +print("# BENCHMARK_BASELINE_OUTPUT_START") +# BENCHMARK_BASELINE_START +baseline = ( + PyTorchTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + ) + .to(dtype=dtype) + .cuda() +) + +print("Baseline PyTorch:") +time_baseline = speedometer(baseline, x, forward_kwargs={"attention_mask": None}, label="baseline") +# BENCHMARK_BASELINE_END +print("# BENCHMARK_BASELINE_OUTPUT_END\n") + + +# ============================================================================= +# TE Unfused: Basic TE Modules +# ============================================================================= + + +# TE_UNFUSED_MLP_START +class TEUnfusedMLP(torch.nn.Module): + """MLP using TE modules.""" + + hidden_size: int + ffn_hidden_size: int + + def __init__(self, hidden_size: int, ffn_hidden_size: int) -> None: + super().__init__() + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + self.linear1 = te.Linear(hidden_size, ffn_hidden_size, bias=True) + self.linear2 = te.Linear(ffn_hidden_size, hidden_size, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.linear1(x) + x = torch.nn.functional.gelu(x, approximate="tanh") + x = self.linear2(x) + return x + + +# TE_UNFUSED_MLP_END + + +# TE_UNFUSED_LAYER_START +class TEUnfusedTransformerLayer(torch.nn.Module): + """Transformer layer using basic TE modules.""" + + def __init__( + self, + hidden_size: int, + ffn_hidden_size: int, + num_attention_heads: int, + layernorm_eps: float = 1e-5, + attention_dropout: float = 0.1, + hidden_dropout: float = 0.1, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.kv_channels = hidden_size // num_attention_heads + self.ln1 = te.LayerNorm(hidden_size, eps=layernorm_eps) + self.qkv_projection = te.Linear(hidden_size, 3 * hidden_size, bias=True) + self.attention = DotProductAttention( + num_attention_heads=num_attention_heads, + kv_channels=self.kv_channels, + attention_dropout=attention_dropout, + ) + self.projection = te.Linear(hidden_size, hidden_size, bias=True) + self.dropout1 = torch.nn.Dropout(hidden_dropout) + self.ln2 = te.LayerNorm(hidden_size, eps=layernorm_eps) + self.mlp = TEUnfusedMLP(hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size) + self.dropout2 = torch.nn.Dropout(hidden_dropout) + + def forward( + self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + res = x + x = self.ln1(x) + + # Fused QKV projection + qkv = self.qkv_projection(x) + qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels) + q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3) + + x = self.attention(q, k, v, attention_mask) + x = self.projection(x) + x = self.dropout1(x) + x = res + x + + # Second residual connection + res = x + x = self.ln2(x) + x = self.mlp(x) + x = self.dropout2(x) + + return x + res + + +# TE_UNFUSED_LAYER_END + + +print("# BENCHMARK_TE_UNFUSED_OUTPUT_START") +# BENCHMARK_TE_UNFUSED_START +te_unfused = ( + TEUnfusedTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + ) + .to(dtype=dtype) + .cuda() +) + +print("TE Unfused:") +time_te_unfused = speedometer( + te_unfused, x, forward_kwargs={"attention_mask": None}, label="te_unfused" +) +# BENCHMARK_TE_UNFUSED_END +print("# BENCHMARK_TE_UNFUSED_OUTPUT_END\n") + + +# ============================================================================= +# TE Unfused + TE Attention +# ============================================================================= + + +# TE_UNFUSED_ATTN_LAYER_START +class TEUnfusedAttnTransformerLayer(torch.nn.Module): + """Transformer layer using TE modules including TE DotProductAttention.""" + + def __init__( + self, + hidden_size: int, + ffn_hidden_size: int, + num_attention_heads: int, + layernorm_eps: float = 1e-5, + attention_dropout: float = 0.1, + hidden_dropout: float = 0.1, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.kv_channels = hidden_size // num_attention_heads + self.ln1 = te.LayerNorm(hidden_size, eps=layernorm_eps) + self.qkv_projection = te.Linear(hidden_size, 3 * hidden_size, bias=True) + self.attention = te.DotProductAttention( + num_attention_heads=num_attention_heads, + kv_channels=self.kv_channels, + attention_dropout=attention_dropout, + attn_mask_type="causal", + ) + self.projection = te.Linear(hidden_size, hidden_size, bias=True) + self.dropout1 = torch.nn.Dropout(hidden_dropout) + self.ln2 = te.LayerNorm(hidden_size, eps=layernorm_eps) + self.mlp = TEUnfusedMLP(hidden_size=hidden_size, ffn_hidden_size=ffn_hidden_size) + self.dropout2 = torch.nn.Dropout(hidden_dropout) + + def forward( + self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + res = x + x = self.ln1(x) + + # Fused QKV projection + qkv = self.qkv_projection(x) + qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels) + q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3) + + x = self.attention(q, k, v, attention_mask) + x = self.projection(x) + x = self.dropout1(x) + x = res + x + + # Second residual connection + res = x + x = self.ln2(x) + x = self.mlp(x) + x = self.dropout2(x) + + return x + res + + +# TE_UNFUSED_ATTN_LAYER_END + + +print("# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_START") +# BENCHMARK_TE_UNFUSED_ATTN_START +te_unfused_attn = ( + TEUnfusedAttnTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + ) + .to(dtype=dtype) + .cuda() +) + +print("TE Unfused + TE Attention:") +time_te_unfused_attn = speedometer( + te_unfused_attn, x, forward_kwargs={"attention_mask": None}, label="te_unfused_attn" +) +# BENCHMARK_TE_UNFUSED_ATTN_END +print("# BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_END\n") + + +# ============================================================================= +# TE Unfused + FP8 +# ============================================================================= + +print("# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_START") +# BENCHMARK_TE_UNFUSED_FP8_START +recipe = DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=16, amax_compute_algo="max") + +te_unfused_fp8 = ( + TEUnfusedAttnTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + ) + .to(dtype=dtype) + .cuda() +) + +print("TE Unfused + TE Attention + FP8:") +time_te_unfused_fp8 = speedometer( + te_unfused_fp8, + x, + forward_kwargs={"attention_mask": None}, + autocast_kwargs={"enabled": True, "recipe": recipe}, + label="te_unfused_fp8", +) +# BENCHMARK_TE_UNFUSED_FP8_END +print("# BENCHMARK_TE_UNFUSED_FP8_OUTPUT_END\n") + + +# ============================================================================= +# TE Fused + FP8: Optimized Modules with FP8 +# ============================================================================= + + +# TE_FUSED_LAYER_START +class TEFusedTransformerLayer(torch.nn.Module): + """Transformer layer using fused TE modules for better performance.""" + + def __init__( + self, + hidden_size: int, + ffn_hidden_size: int, + num_attention_heads: int, + layernorm_eps: float = 1e-5, + attention_dropout: float = 0.1, + hidden_dropout: float = 0.1, + ): + super().__init__() + self.num_attention_heads = num_attention_heads + self.kv_channels = hidden_size // num_attention_heads + + # Fused LayerNorm + QKV projection + self.ln_qkv = te.LayerNormLinear(hidden_size, 3 * hidden_size, eps=layernorm_eps, bias=True) + self.attention = te.DotProductAttention( + num_attention_heads=num_attention_heads, + kv_channels=self.kv_channels, + attention_dropout=attention_dropout, + attn_mask_type="causal", + ) + self.projection = te.Linear(hidden_size, hidden_size, bias=True) + self.dropout1 = torch.nn.Dropout(hidden_dropout) + + # Fused LayerNorm + MLP + self.ln_mlp = te.LayerNormMLP(hidden_size, ffn_hidden_size, eps=layernorm_eps, bias=True) + self.dropout2 = torch.nn.Dropout(hidden_dropout) + + def forward( + self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None + ) -> torch.Tensor: + res = x + + # Fused LayerNorm + QKV projection + qkv = self.ln_qkv(x) + qkv = qkv.view(qkv.size(0), qkv.size(1), self.num_attention_heads, 3 * self.kv_channels) + q, k, v = torch.split(qkv, qkv.size(3) // 3, dim=3) + + x = self.attention(q, k, v, attention_mask) + x = self.projection(x) + x = self.dropout1(x) + x = res + x + + # Fused LayerNorm + MLP + res = x + x = self.ln_mlp(x) + x = self.dropout2(x) + + return x + res + + +# TE_FUSED_LAYER_END + + +print("# BENCHMARK_TE_FUSED_FP8_OUTPUT_START") +# BENCHMARK_TE_FUSED_FP8_START +te_fused_fp8 = ( + TEFusedTransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + ) + .to(dtype=dtype) + .cuda() +) + +print("TE Fused + TE Attention + FP8:") +time_te_fused_fp8 = speedometer( + te_fused_fp8, + x, + forward_kwargs={"attention_mask": None}, + autocast_kwargs={"enabled": True, "recipe": recipe}, + label="te_fused_fp8", +) +# BENCHMARK_TE_FUSED_FP8_END +print("# BENCHMARK_TE_FUSED_FP8_OUTPUT_END\n") + + +# ============================================================================= +# TE TransformerLayer + FP8: Ready-to-use Module +# ============================================================================= + +print("# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_START") +# BENCHMARK_TE_TRANSFORMER_LAYER_START +te_transformer_layer = ( + te.TransformerLayer( + hidden_size=hidden_size, + ffn_hidden_size=ffn_hidden_size, + num_attention_heads=num_attention_heads, + self_attn_mask_type="causal", + layernorm_epsilon=1e-5, + bias=True, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + .to(dtype=dtype) + .cuda() +) + +print("TE TransformerLayer + FP8:") +time_te_transformer_layer = speedometer( + te_transformer_layer, + x, + forward_kwargs={"attention_mask": None}, + autocast_kwargs={"enabled": True, "recipe": recipe}, + label="te_transformer_layer", +) +# BENCHMARK_TE_TRANSFORMER_LAYER_END +print("# BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_END\n") + + +# Write summary CSV for RST documentation +with open("getting_started_pytorch_summary.csv", "w") as f: + f.write("Implementation,Time (ms),Speedup\n") + f.write(f"Baseline PyTorch,{time_baseline:.2f},1.00x\n") + f.write(f"TE Unfused,{time_te_unfused:.2f},{time_baseline/time_te_unfused:.2f}x\n") + f.write( + "TE Unfused + TE" + f" Attention,{time_te_unfused_attn:.2f},{time_baseline/time_te_unfused_attn:.2f}x\n" + ) + f.write( + "TE Unfused + TE Attention +" + f" FP8,{time_te_unfused_fp8:.2f},{time_baseline/time_te_unfused_fp8:.2f}x\n" + ) + f.write( + "TE Fused + TE Attention +" + f" FP8,{time_te_fused_fp8:.2f},{time_baseline/time_te_fused_fp8:.2f}x\n" + ) + f.write( + "TE TransformerLayer +" + f" FP8,{time_te_transformer_layer:.2f},{time_baseline/time_te_transformer_layer:.2f}x\n" + ) +print("\nSummary written to getting_started_pytorch_summary.csv") diff --git a/docs/getting_started/getting_started_pytorch_summary.csv b/docs/getting_started/getting_started_pytorch_summary.csv new file mode 100644 index 0000000000..b3a5d7330e --- /dev/null +++ b/docs/getting_started/getting_started_pytorch_summary.csv @@ -0,0 +1,7 @@ +Implementation,Time (ms),Speedup +Baseline PyTorch,48.28,1.00x +TE Unfused,49.34,0.98x +TE Unfused + TE Attention,35.71,1.35x +TE Unfused + TE Attention + FP8,23.41,2.06x +TE Fused + TE Attention + FP8,22.96,2.10x +TE TransformerLayer + FP8,21.67,2.23x diff --git a/docs/getting_started/getting_started_utils_jax.py b/docs/getting_started/getting_started_utils_jax.py new file mode 100644 index 0000000000..e489395fc7 --- /dev/null +++ b/docs/getting_started/getting_started_utils_jax.py @@ -0,0 +1,76 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Utility functions for Getting Started with Transformer Engine - JAX +==================================================================== + +Helper classes and functions for the getting started examples. +""" + +import time +from typing import Callable, Any, Optional + +import jax +import jax.numpy as jnp +from flax import linen as nn +import transformer_engine.jax as te +from transformer_engine.jax.sharding import MeshResource + + +def speedometer( + apply_fn: Callable, + params: Any, + x: jnp.ndarray, + forward_kwargs: dict = {}, + autocast_kwargs: Optional[dict] = None, + timing_iters: int = 100, + warmup_iters: int = 10, + label: str = "benchmark", +) -> float: + """Measure average forward + backward pass time for a JAX module. + + Args: + apply_fn: JIT-compiled apply function + params: Model parameters + x: Input tensor + forward_kwargs: Additional kwargs for forward pass + autocast_kwargs: Kwargs for te.autocast context + timing_iters: Number of timing iterations + warmup_iters: Number of warmup iterations + label: Optional label for logging + + Returns: + Average time per iteration in milliseconds + """ + if autocast_kwargs is None: + autocast_kwargs = {"enabled": False} + else: + autocast_kwargs = dict(autocast_kwargs) + autocast_kwargs.setdefault("mesh_resource", MeshResource()) + + def loss_fn(params, x): + y = apply_fn(params, x, **forward_kwargs) + return jnp.sum(y) + + # JIT compile within autocast context + with te.autocast(**autocast_kwargs): + grad_fn = jax.jit(jax.value_and_grad(loss_fn)) + + # Warmup runs + for _ in range(warmup_iters): + loss, grads = grad_fn(params, x) + jax.block_until_ready((loss, grads)) + + # Timing runs + times = [] + for _ in range(timing_iters): + start = time.perf_counter() + loss, grads = grad_fn(params, x) + jax.block_until_ready((loss, grads)) + times.append(time.perf_counter() - start) + + avg_time = sum(times) / len(times) * 1000 + print(f"Mean time: {avg_time:.3f} ms") + return avg_time diff --git a/docs/getting_started/getting_started_utils_pytorch.py b/docs/getting_started/getting_started_utils_pytorch.py new file mode 100644 index 0000000000..c76e17645a --- /dev/null +++ b/docs/getting_started/getting_started_utils_pytorch.py @@ -0,0 +1,124 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Utility functions for Getting Started with Transformer Engine - PyTorch +======================================================================== + +Helper classes and functions for the getting started examples. +""" + +import math +from typing import Optional +import torch +import transformer_engine.pytorch as te + + +def speedometer( + module: torch.nn.Module, + x: torch.Tensor, + forward_kwargs: dict = {}, + autocast_kwargs: Optional[dict] = None, + timing_iters: int = 100, + warmup_iters: int = 10, + label: str = "benchmark", +) -> float: + """Measure average forward + backward pass time for a PyTorch module. + + Args: + module: PyTorch module to benchmark + x: Input tensor + forward_kwargs: Additional kwargs for forward pass + autocast_kwargs: Kwargs for te.autocast context + timing_iters: Number of timing iterations + warmup_iters: Number of warmup iterations + label: Optional label for logging + + Returns: + Average time per iteration in milliseconds + """ + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + if autocast_kwargs is None: + autocast_kwargs = {"enabled": False} + + # Warmup runs + torch.cuda.synchronize() + for _ in range(warmup_iters): + with te.autocast(**autocast_kwargs): + y = module(x, **forward_kwargs) + loss = y.sum() + loss.backward() + torch.cuda.synchronize() + + # Timing runs + start.record() + for _ in range(timing_iters): + with te.autocast(**autocast_kwargs): + y = module(x, **forward_kwargs) + loss = y.sum() + loss.backward() + end.record() + torch.cuda.synchronize() + + avg_time = start.elapsed_time(end) / timing_iters + print(f"Mean time: {avg_time:.3f} ms") + return avg_time + + +class DotProductAttention(torch.nn.Module): + """Attention operation in Transformer layer. + + Built with plain PyTorch modules. + """ + + def __init__( + self, + num_attention_heads: int, + kv_channels: int, + attention_dropout: float, + ) -> None: + super().__init__() + self.projection_size = kv_channels * num_attention_heads + self.hidden_size_per_attention_head = kv_channels + self.norm_factor = math.sqrt(self.hidden_size_per_attention_head) + self.dropout = torch.nn.Dropout(attention_dropout) + + def masked_softmax(self, inp: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor: + if mask is not None: + inp.masked_fill_(mask, -10000.0) + return torch.nn.Softmax(dim=-1)(inp) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + b = query.size(1) + np = query.size(2) + sq = query.size(0) + sk = key.size(0) + hn = value.size(3) + + query = query.view(sq, b * np, -1) + key = key.view(sk, b * np, -1) + + bmm1 = ( + torch.bmm(query.transpose(0, 1), key.transpose(0, 1).transpose(1, 2)) / self.norm_factor + ) + + attention_scores = bmm1.view(b, np, sq, sk) + attention_probs = self.masked_softmax(attention_scores, attention_mask) + attention_probs = self.dropout(attention_probs) + + value = value.view(sk, b * np, -1) + attention_probs = attention_probs.view(b * np, sq, -1) + context = torch.bmm(attention_probs, value.transpose(0, 1)) + context = context.view(b, np, sq, hn) + context = context.permute(2, 0, 1, 3).contiguous() + context = context.view(sq, b, self.projection_size) + + return context diff --git a/docs/getting_started/index.rst b/docs/getting_started/index.rst new file mode 100644 index 0000000000..9e10f82c14 --- /dev/null +++ b/docs/getting_started/index.rst @@ -0,0 +1,566 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Getting Started +=============== + +Overview +-------- + +Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, +providing better performance with lower memory utilization in both training and inference. +It provides support for 8-bit floating point (FP8) precision on Hopper and Ada GPUs, as well as +8-bit and 4-bit floating point (NVFP4) precision on Blackwell GPUs. + +TE implements a collection of highly optimized building blocks for popular Transformer +architectures and exposes an automatic-mixed-precision-like API that can be used seamlessly +with your deep learning code. + + +Currently two frameworks are supported: PyTorch and JAX. + +.. tabs:: + + .. tab:: PyTorch + + Basic knowledge of PyTorch is recommended: + + - `PyTorch Tutorials `_ + - `PyTorch Documentation `_ + + .. tab:: JAX + + We recommend understanding the basics of JAX first: + + - `Thinking in JAX `_ + - `JAX 101 `_ + - `Key concepts in JAX `_ + - `Flax 101 `_ + + +Baseline: Pure Framework Implementation +--------------------------------------- + +Let's build a Transformer decoder layer! + +We'll create a basic GPT-style layer with causal masking, +which prevents each position from attending to future positions. This will be our baseline +for later comparisons with Transformer Engine. + +.. raw:: html + :file: transformer_layer.svg + +.. raw:: html + +

Structure of a GPT decoder layer

+ +We construct the components as follows: + +.. tabs:: + + .. tab:: PyTorch + + * **LayerNorm**: ``torch.nn.LayerNorm`` + * **QKV Projection**: ``torch.nn.Linear`` (fused Q, K, V into single layer 3x larger) + * **DotProductAttention**: Custom implementation using ``torch.bmm`` + * **Projection**: ``torch.nn.Linear`` + * **Dropout**: ``torch.nn.Dropout`` + * **MLP**: Two ``torch.nn.Linear`` layers with ``torch.nn.functional.gelu`` activation + + .. tab:: JAX + + * **LayerNorm**: ``nn.LayerNorm`` + * **QKV Projection**: ``nn.Dense`` (fused Q, K, V into single layer 3x larger) + * **DotProductAttention**: ``nn.dot_product_attention`` + * **Projection**: ``nn.Dense`` + * **Dropout**: ``nn.Dropout`` + * **MLP**: Two ``nn.Dense`` layers with ``nn.gelu`` activation + +Putting it all together: + +.. tabs:: + + .. tab:: PyTorch + + First, define the MLP block: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BASELINE_MLP_START + :end-before: # BASELINE_MLP_END + + Now, putting it all together into a GPT decoder layer: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BASELINE_LAYER_START + :end-before: # BASELINE_LAYER_END + + Benchmark the baseline implementation: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BENCHMARK_BASELINE_START + :end-before: # BENCHMARK_BASELINE_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_pytorch.out + :language: text + :start-after: # BENCHMARK_BASELINE_OUTPUT_START + :end-before: # BENCHMARK_BASELINE_OUTPUT_END + + .. tab:: JAX + + First, define the MLP block: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BASELINE_MLP_START + :end-before: # BASELINE_MLP_END + + Now, putting it all together into a GPT decoder layer: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BASELINE_LAYER_START + :end-before: # BASELINE_LAYER_END + + Benchmark the baseline implementation: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BENCHMARK_BASELINE_START + :end-before: # BENCHMARK_BASELINE_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_jax.out + :language: text + :start-after: # BENCHMARK_BASELINE_OUTPUT_START + :end-before: # BENCHMARK_BASELINE_OUTPUT_END + + +TE Unfused: Basic TE Modules +---------------------------- + +Now let's replace the standard framework modules with TE equivalents. +This is the simplest way to start using Transformer Engine. + +.. tabs:: + + .. tab:: PyTorch + + Replace PyTorch modules with TE equivalents: + + .. code-block:: python + + import transformer_engine.pytorch as te + + Mapping: + + * ``torch.nn.Linear`` → ``te.Linear`` + * ``torch.nn.LayerNorm`` → ``te.LayerNorm`` + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # TE_UNFUSED_MLP_START + :end-before: # TE_UNFUSED_MLP_END + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # TE_UNFUSED_LAYER_START + :end-before: # TE_UNFUSED_LAYER_END + + Benchmark the TE unfused implementation: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BENCHMARK_TE_UNFUSED_START + :end-before: # BENCHMARK_TE_UNFUSED_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_pytorch.out + :language: text + :start-after: # BENCHMARK_TE_UNFUSED_OUTPUT_START + :end-before: # BENCHMARK_TE_UNFUSED_OUTPUT_END + + .. tab:: JAX + + Replace Flax modules with TE equivalents: + + .. code-block:: python + + import transformer_engine.jax as te + import transformer_engine.jax.flax as te_flax + + Mapping: + + * ``nn.Dense`` → ``te_flax.DenseGeneral`` + * ``nn.LayerNorm`` → ``te_flax.LayerNorm`` + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # TE_UNFUSED_MLP_START + :end-before: # TE_UNFUSED_MLP_END + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # TE_UNFUSED_LAYER_START + :end-before: # TE_UNFUSED_LAYER_END + + Benchmark the TE unfused implementation: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BENCHMARK_TE_UNFUSED_START + :end-before: # BENCHMARK_TE_UNFUSED_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_jax.out + :language: text + :start-after: # BENCHMARK_TE_UNFUSED_OUTPUT_START + :end-before: # BENCHMARK_TE_UNFUSED_OUTPUT_END + + +TE Unfused + TE Attention +------------------------- + +Now let's also replace the attention mechanism with TE's optimized ``DotProductAttention``. +TE's attention automatically selects the best available backend — for example, FlashAttention or cuDNN fused attention — based on your hardware and input configuration, +delivering optimal performance without manual tuning. + +.. tabs:: + + .. tab:: PyTorch + + Replace the custom attention with TE's optimized implementation: + + * Custom ``DotProductAttention`` → ``te.DotProductAttention`` + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # TE_UNFUSED_ATTN_LAYER_START + :end-before: # TE_UNFUSED_ATTN_LAYER_END + + Benchmark TE Unfused with TE Attention: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BENCHMARK_TE_UNFUSED_ATTN_START + :end-before: # BENCHMARK_TE_UNFUSED_ATTN_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_pytorch.out + :language: text + :start-after: # BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_START + :end-before: # BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_END + + .. tab:: JAX + + Replace Flax's attention with TE's optimized implementation: + + * ``nn.dot_product_attention`` → ``te_flax.DotProductAttention`` + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # TE_UNFUSED_ATTN_LAYER_START + :end-before: # TE_UNFUSED_ATTN_LAYER_END + + Benchmark TE Unfused with TE Attention: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BENCHMARK_TE_UNFUSED_ATTN_START + :end-before: # BENCHMARK_TE_UNFUSED_ATTN_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_jax.out + :language: text + :start-after: # BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_START + :end-before: # BENCHMARK_TE_UNFUSED_ATTN_OUTPUT_END + + +TE Unfused + TE Attention + FP8 +------------------------------- + +Now let's combine TE modules with TE Attention and enable FP8 precision. +Wrap your code within an ``autocast`` context manager to enable FP8. +This provides significant speedups on supported hardware (Hopper, Ada, Blackwell GPUs). + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + from transformer_engine.common.recipe import Format, DelayedScaling + + recipe = DelayedScaling( + fp8_format=Format.HYBRID, + amax_history_len=16, + amax_compute_algo="max" + ) + + with te.autocast(enabled=True, recipe=recipe): + y = te_unfused(x, attention_mask=None) + + .. note:: + + The ``autocast`` should only wrap the forward pass and must exit before + starting a backward pass. + + Benchmark TE Unfused with FP8: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BENCHMARK_TE_UNFUSED_FP8_START + :end-before: # BENCHMARK_TE_UNFUSED_FP8_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_pytorch.out + :language: text + :start-after: # BENCHMARK_TE_UNFUSED_FP8_OUTPUT_START + :end-before: # BENCHMARK_TE_UNFUSED_FP8_OUTPUT_END + + .. tab:: JAX + + .. code-block:: python + + from transformer_engine.common.recipe import Format, DelayedScaling + + recipe = DelayedScaling( + fp8_format=Format.HYBRID, + amax_history_len=16, + amax_compute_algo="max" + ) + + with te.autocast(enabled=True, recipe=recipe): + params = te_unfused.init(key, x, deterministic=False) + y = te_unfused.apply(params, x, deterministic=True) + + .. important:: + + When using FP8 in JAX, the model **must be initialized within the autocast context** + to create the ``fp8_metas`` collection. + + Benchmark TE Unfused with FP8: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BENCHMARK_TE_UNFUSED_FP8_START + :end-before: # BENCHMARK_TE_UNFUSED_FP8_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_jax.out + :language: text + :start-after: # BENCHMARK_TE_UNFUSED_FP8_OUTPUT_START + :end-before: # BENCHMARK_TE_UNFUSED_FP8_OUTPUT_END + + +TE Fused + TE Attention + FP8: Optimized Modules +------------------------------------------------ + +Fused modules use kernel fusion to combine multiple operations. +While speedups are modest on a single GPU, they scale better in multi-GPU setups. +Combined with TE Attention and FP8, this delivers peak performance. + +.. tabs:: + + .. tab:: PyTorch + + Fused modules available: + + * ``te.LayerNormLinear`` - fuses LayerNorm + Linear + * ``te.LayerNormMLP`` - fuses LayerNorm + MLP + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # TE_FUSED_LAYER_START + :end-before: # TE_FUSED_LAYER_END + + Benchmark TE Fused with FP8: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BENCHMARK_TE_FUSED_FP8_START + :end-before: # BENCHMARK_TE_FUSED_FP8_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_pytorch.out + :language: text + :start-after: # BENCHMARK_TE_FUSED_FP8_OUTPUT_START + :end-before: # BENCHMARK_TE_FUSED_FP8_OUTPUT_END + + .. tab:: JAX + + Fused modules available: + + * ``te_flax.LayerNormDenseGeneral`` - fuses LayerNorm + Dense + * ``te_flax.LayerNormMLP`` - fuses LayerNorm + MLP + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # TE_FUSED_LAYER_START + :end-before: # TE_FUSED_LAYER_END + + Benchmark TE Fused with FP8: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BENCHMARK_TE_FUSED_FP8_START + :end-before: # BENCHMARK_TE_FUSED_FP8_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_jax.out + :language: text + :start-after: # BENCHMARK_TE_FUSED_FP8_OUTPUT_START + :end-before: # BENCHMARK_TE_FUSED_FP8_OUTPUT_END + + +TE TransformerLayer + FP8: Ready-to-use Module +---------------------------------------------- + +For the simplest integration, Transformer Engine provides a ready-to-use ``TransformerLayer`` +module that includes all optimizations out of the box. + +.. tabs:: + + .. tab:: PyTorch + + Just use ``te.TransformerLayer`` - it handles everything for you: + + .. literalinclude:: getting_started_pytorch.py + :language: python + :start-after: # BENCHMARK_TE_TRANSFORMER_LAYER_START + :end-before: # BENCHMARK_TE_TRANSFORMER_LAYER_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_pytorch.out + :language: text + :start-after: # BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_START + :end-before: # BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_END + + .. tab:: JAX + + Just use ``te_flax.TransformerLayer`` - it handles everything for you: + + .. literalinclude:: getting_started_jax.py + :language: python + :start-after: # BENCHMARK_TE_TRANSFORMER_LAYER_START + :end-before: # BENCHMARK_TE_TRANSFORMER_LAYER_END + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: getting_started_jax.out + :language: text + :start-after: # BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_START + :end-before: # BENCHMARK_TE_TRANSFORMER_LAYER_OUTPUT_END + + +Benchmark Summary +----------------- + +The table below summarizes the performance improvements achieved with Transformer Engine +on an NVIDIA H100 GPU. Results may vary depending on hardware and configuration. While this +tutorial focuses on a simple single-GPU scenario, features like fused layers can provide +additional benefits in more complex setups such as multi-GPU training. + +.. tabs:: + + .. tab:: PyTorch + + .. csv-table:: + :header-rows: 1 + :widths: 40, 20, 20 + :file: getting_started_pytorch_summary.csv + + .. tab:: JAX + + .. csv-table:: + :header-rows: 1 + :widths: 40, 20, 20 + :file: getting_started_jax_summary.csv diff --git a/docs/getting_started/transformer_layer.svg b/docs/getting_started/transformer_layer.svg new file mode 100644 index 0000000000..28ba3dd386 --- /dev/null +++ b/docs/getting_started/transformer_layer.svg @@ -0,0 +1,82 @@ + + + + + + + + + + + + + + + + + + LayerNorm + + + + + + QKV Projection + + + + + + Dot Product + Attention + + + + + + Projection + + + + + + Dropout + + + + + + + + + + + + + + + + + LayerNorm + + + + + + MLP + + + + + + + + + + + + diff --git a/docs/index.rst b/docs/index.rst index 99611cda99..0edcb863b6 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -29,7 +29,7 @@ Transformer Engine documentation :caption: Getting Started installation - getting_started + getting_started/index faq .. toctree:: From df69100c3bbb34b1e5f756d69093003d4846199d Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 6 Jan 2026 10:49:27 -0800 Subject: [PATCH 151/521] [Common] Fix long compile time in padding.cu on arch 75 (#2562) * Fix long compile time in padding.cu Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/common/util/padding.cu | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/util/padding.cu b/transformer_engine/common/util/padding.cu index 1859d8a5cb..8359238289 100644 --- a/transformer_engine/common/util/padding.cu +++ b/transformer_engine/common/util/padding.cu @@ -94,7 +94,6 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP #pragma unroll for (int i2 = 0; i2 < nvec; ++i2) { const int row = tile_row + i1 * nvec + i2; - size_t row_offset = static_cast(row) * row_length; const int col = tile_col + j1 * nvec; Vec local_input; Vec local_output; @@ -102,7 +101,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - local_input.data.elt[j2] = input[row_offset + col + j2]; + local_input.data.elt[j2] = input[static_cast(row) * row_length + col + j2]; } } } @@ -113,14 +112,14 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - output[row_offset + col + j2] = local_output.data.elt[j2]; + output[static_cast(row) * row_length + col + j2] = local_output.data.elt[j2]; } } } else if (row < padded_num_rows) { // padding for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - output[row_offset + col + j2] = local_zero; + output[static_cast(row) * row_length + col + j2] = local_zero; } } } @@ -179,7 +178,6 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult #pragma unroll for (int i2 = 0; i2 < nvec; ++i2) { const int row = tile_row + i1 * nvec + i2; - size_t row_offset = static_cast(row) * row_length; const int col = tile_col + j1 * nvec; Vec local_input; Vec local_output; @@ -187,7 +185,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - local_input.data.elt[j2] = input[row_offset + col + j2]; + local_input.data.elt[j2] = input[static_cast(row) * row_length + col + j2]; } } } @@ -198,7 +196,7 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult if (row < num_rows) { for (int j2 = 0; j2 < nvec; ++j2) { if (col + j2 < row_length) { - output[row_offset + col + j2] = local_output.data.elt[j2]; + output[static_cast(row) * row_length + col + j2] = local_output.data.elt[j2]; } } } From 404a3ee04a9011e57e7ed852cc8dd86a3e21e1be Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 6 Jan 2026 12:01:57 -0800 Subject: [PATCH 152/521] [JAX] Fix test_layer to support fused attention and adjust test encoder tolerance to account for minor diff (#2563) Fix failing unit tests Signed-off-by: Jeremy Berchtold --- examples/jax/encoder/test_model_parallel_encoder.py | 8 ++++---- tests/jax/test_layer.py | 9 +++++++++ 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index b534db8576..f29cc4e0be 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -503,7 +503,7 @@ def test_te_delayed_scaling_fp8(self): self.args.use_fp8 = True self.args.fp8_recipe = "DelayedScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.361 and actual[1] > 0.84 + assert actual[0] < 0.362 and actual[1] > 0.84 @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) def test_te_mxfp8(self): @@ -535,7 +535,7 @@ def test_te_delayed_scaling_fp8_with_sp(self): self.args.use_fp8 = True self.args.fp8_recipe = "DelayedScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.361 and actual[1] > 0.84 + assert actual[0] < 0.362 and actual[1] > 0.84 @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) def test_te_mxfp8_with_sp(self): @@ -569,7 +569,7 @@ def test_te_delayed_scaling_fp8_shardy(self): self.args.use_fp8 = True self.args.fp8_recipe = "DelayedScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.361 and actual[1] > 0.84 + assert actual[0] < 0.362 and actual[1] > 0.84 @unittest.skipIf(not is_fp8_supported, fp8_reason) def test_te_delayed_scaling_fp8_with_sp_shardy(self): @@ -579,7 +579,7 @@ def test_te_delayed_scaling_fp8_with_sp_shardy(self): self.args.use_fp8 = True self.args.fp8_recipe = "DelayedScaling" actual = train_and_evaluate(self.args) - assert actual[0] < 0.361 and actual[1] > 0.84 + assert actual[0] < 0.362 and actual[1] > 0.84 @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) def test_te_mxfp8_shardy(self): diff --git a/tests/jax/test_layer.py b/tests/jax/test_layer.py index 8c16d162ed..0499d5cba7 100644 --- a/tests/jax/test_layer.py +++ b/tests/jax/test_layer.py @@ -430,6 +430,9 @@ class EncoderRunner(BaseRunner): "attention/DotProductAttention_0/_UnfusedDotProductAttention_0/softmax_offset": ( "attention/DotProductAttention_0/softmax_offset" ), + "attention/DotProductAttention_0/_FusedDotProductAttention_0/softmax_offset": ( + "attention/DotProductAttention_0/softmax_offset" + ), "mlp/wi_kernel": "mlp/wi/kernel", "mlp/wi_bias": "mlp/wi/bias", "mlp/wo_kernel": "mlp/wo/kernel", @@ -478,6 +481,9 @@ class DecoderRunner(BaseRunner): "encoder_decoder_attention/DotProductAttention_0/_UnfusedDotProductAttention_0/softmax_offset": ( "encoder_decoder_attention/DotProductAttention_0/softmax_offset" ), + "encoder_decoder_attention/DotProductAttention_0/_FusedDotProductAttention_0/softmax_offset": ( + "encoder_decoder_attention/DotProductAttention_0/softmax_offset" + ), "self_attention/qkv/scale": "pre_self_attention_layer_norm/scale", "self_attention/qkv/ln_bias": "pre_self_attention_layer_norm/ln_bias", "self_attention/query/scale": "pre_self_attention_layer_norm/scale", @@ -485,6 +491,9 @@ class DecoderRunner(BaseRunner): "self_attention/DotProductAttention_0/_UnfusedDotProductAttention_0/softmax_offset": ( "self_attention/DotProductAttention_0/softmax_offset" ), + "self_attention/DotProductAttention_0/_FusedDotProductAttention_0/softmax_offset": ( + "self_attention/DotProductAttention_0/softmax_offset" + ), "mlp/wi_kernel": "mlp/wi/kernel", "mlp/wi_bias": "mlp/wi/bias", "mlp/wo_kernel": "mlp/wo/kernel", From 702fc5eecadd4709e049c2e370fdf5607e79975f Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Tue, 6 Jan 2026 16:34:28 -0800 Subject: [PATCH 153/521] Fix 50% comparison mismatch in sort_chunks_by_index (#2566) * force initialization to int32 Signed-off-by: tdophung * address greptile comment Signed-off-by: tdophung --------- Signed-off-by: tdophung --- tests/jax/test_permutation.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index 43f2553eed..d61ea8eb75 100644 --- a/tests/jax/test_permutation.py +++ b/tests/jax/test_permutation.py @@ -97,7 +97,9 @@ def reference_make_row_id_map( # Compute total tokens per expert and expert offsets tokens_per_expert = jnp.sum(routing_map, axis=0) - expert_offsets = jnp.concatenate([jnp.array([0]), jnp.cumsum(tokens_per_expert)[:-1]]) + expert_offsets = jnp.concatenate( + [jnp.array([0], dtype=jnp.int32), jnp.cumsum(tokens_per_expert)[:-1].astype(jnp.int32)] + ) # Compute destination rows for all (token, expert) pairs # dest_row[i, j] = expert_offsets[j] + cumsum_per_expert[i, j] - 1 if routed, else -1 @@ -115,7 +117,9 @@ def reference_make_row_id_map( # Gather the sorted destination rows and expert indices using advanced indexing # Create indices for gathering - token_idx = jnp.broadcast_to(jnp.arange(num_tokens)[:, None], (num_tokens, num_experts)) + token_idx = jnp.broadcast_to( + jnp.arange(num_tokens, dtype=jnp.int32)[:, None], (num_tokens, num_experts) + ) sorted_dest_rows = dest_rows_all[token_idx, sorted_expert_indices] # Build row_id_map: [dest_row_0, ..., dest_row_{E-1}, expert_idx_0, ..., expert_idx_{E-1}, n_routed] @@ -373,11 +377,15 @@ def reference_make_chunk_sort_map( Row ID map for chunk sorting of shape [num_tokens,]. """ # Compute source chunk boundaries (cumulative sum of original split_sizes) - src_cumsum = jnp.concatenate([jnp.array([0]), jnp.cumsum(split_sizes)]) + src_cumsum = jnp.concatenate( + [jnp.array([0], dtype=jnp.int32), jnp.cumsum(split_sizes).astype(jnp.int32)] + ) # Compute destination chunk boundaries based on sorted order sorted_sizes = split_sizes[sorted_indices] - dest_cumsum = jnp.concatenate([jnp.array([0]), jnp.cumsum(sorted_sizes)]) + dest_cumsum = jnp.concatenate( + [jnp.array([0], dtype=jnp.int32), jnp.cumsum(sorted_sizes).astype(jnp.int32)] + ) # For each source chunk, compute its destination offset # inverse_indices[i] = position of chunk i in sorted order @@ -386,7 +394,7 @@ def reference_make_chunk_sort_map( # Create row_id_map: for each token position, compute its destination # First, figure out which chunk each position belongs to - position_indices = jnp.arange(num_tokens) + position_indices = jnp.arange(num_tokens, dtype=jnp.int32) # chunk_ids[i] = which chunk position i belongs to chunk_ids = jnp.searchsorted(src_cumsum[1:], position_indices, side="right") From de51c96b2b7c7ca4a856038945de4150fdb6af6c Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Wed, 7 Jan 2026 10:27:20 -0800 Subject: [PATCH 154/521] [NVFP4][MOE] Bug Fix for NVFP4 Grouped Quant (#2564) * fix Signed-off-by: Zhongbo Zhu * resolve review comments Signed-off-by: Zhongbo Zhu * Comment tweaks Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Zhongbo Zhu Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- ...cast_col_hadamard_transform_cast_fusion.cu | 33 ++++++++++++------- .../transformer_engine/hadamard_transform.h | 3 +- .../pytorch/csrc/extensions/cast.cpp | 8 ++++- 3 files changed, 30 insertions(+), 14 deletions(-) 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 8b077f6f1f..1ef1f81e82 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 @@ -1125,8 +1125,9 @@ template (tile_scheduler_workspace), 0, + sizeof(uint32_t), stream)); // Launch kernel cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream}; @@ -1308,8 +1308,6 @@ void group_row_col_rht_gemm_ntt_w_sfc(int packed_sequence_length, int hidden_siz tile_scheduler_workspace, mma, rng_state); NVTE_CHECK_CUDA(cudaGetLastError()); NVTE_CHECK(status == cutlass::Status::kSuccess, "Kernel launch failed."); - - NVTE_CHECK_CUDA(cudaFreeAsync(tile_scheduler_workspace, stream)); } } // namespace @@ -1318,7 +1316,8 @@ void group_row_col_rht_gemm_ntt_w_sfc(int packed_sequence_length, int hidden_siz void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector &output_list, const size_t *split_sections, size_t num_tensors, const Tensor &hadamard_matrix_, - QuantizationConfig &quant_config, cudaStream_t stream) { + QuantizationConfig &quant_config, Tensor &quant_workspace, + cudaStream_t stream) { NVTE_API_CALL(group_hadamard_transform_cast_fusion); using transformer_engine::detail::kMaxTensorsPerKernel; @@ -1399,6 +1398,12 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector(rng_state_tensor.data.dptr); } + uint32_t *tile_scheduler_workspace = nullptr; + NVTE_CHECK(quant_workspace.data.dptr != nullptr, "Quantization workspace must be provided."); + NVTE_CHECK(quant_workspace.data.buffer_size_bytes() >= sizeof(uint32_t), + "Quantization workspace must be at least 4 bytes."); + tile_scheduler_workspace = reinterpret_cast(quant_workspace.data.dptr); + // Template arguments using TA = cute::bfloat16_t; using TB = cute::bfloat16_t; @@ -1461,7 +1466,9 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector(rowwise_data_base_ptr), /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), /*args=*/kernel_args, - /*rng_state=*/rng_state, /*sm_count=*/sm_count, + /*rng_state=*/rng_state, + /*tile_scheduler_workspace=*/tile_scheduler_workspace, + /*sm_count=*/sm_count, /*stream=*/stream, /*k_tile_size=*/k_tile_size); } else { NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", @@ -1478,7 +1485,7 @@ void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETenso const size_t *split_sections, const size_t num_tensors, const NVTEQuantizationConfig quant_config, - cudaStream_t stream) { + NVTETensor quant_workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion); using namespace transformer_engine; NVTE_CHECK(num_tensors > 0, "Number of tensors should be greater than 0."); @@ -1489,6 +1496,8 @@ void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETenso output_list[i] = convertNVTETensorCheck(outputs[i]); } + Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace); + QuantizationConfig quant_config_cpp; if (quant_config != nullptr) { quant_config_cpp = *reinterpret_cast(quant_config); @@ -1497,5 +1506,5 @@ void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETenso // Call the multi-tensor Hadamard transform amax implementation. group_hadamard_transform_cast_fusion(*input_tensor, output_list, split_sections, num_tensors, *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, - stream); + *quant_workspace_tensor, stream); } diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index b6e9719aad..13103cc388 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -115,13 +115,14 @@ void nvte_group_hadamard_transform_cast_fusion_columnwise( * \param[in] split_sections Array specifying splits in dimension 0 for each output tensor. * \param[in] num_tensors Number of output tensors, must be > 0. * \param[in] quant_config Quantization configuration. + * \param[in] quant_workspace Workspace buffer. Must be at least 4 bytes. * \param[in] stream CUDA stream used for the operation. */ void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETensor* outputs, const NVTETensor hadamard_matrix, const size_t* split_sections, size_t num_tensors, const NVTEQuantizationConfig quant_config, - cudaStream_t stream); + NVTETensor quant_workspace, cudaStream_t stream); #ifdef __cplusplus } // extern "C" diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 3bbc99b444..4e5e5223f7 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -872,10 +872,16 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, auto rht_matrix_nvte = makeTransformerEngineTensor(quantizer.rht_matrix); if (all_aligned_token_dim) { + // allocate a tile scheduler workspace + auto tile_scheduler_workspace_torch = + at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32)); + auto nvte_tile_scheduler_workspace = + makeTransformerEngineTensor(tile_scheduler_workspace_torch); // call the fully-fused grouped kernel for rowwise quantization & colwise RHT quantization transpose nvte_group_hadamard_transform_cast_fusion( input.data(), reinterpret_cast(nvte_tensor_output_list.data()), - rht_matrix_nvte.data(), split_sections.data(), num_tensors, quant_config_list[0], stream); + rht_matrix_nvte.data(), split_sections.data(), num_tensors, quant_config_list[0], + nvte_tile_scheduler_workspace.data(), stream); } else { // Separate quantization for rowwise usage and columnwise usage // Rowwise quantization fusion with grouped version From 08dc786cbd24485dc3bed0607d99740021beb2b7 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 7 Jan 2026 15:30:03 -0800 Subject: [PATCH 155/521] Fix 50% comparison mismatch in sort_chunks_by_index (Cont.) (#2575) * force initialization to int32 Signed-off-by: tdophung * address greptile comment Signed-off-by: tdophung * del useless comments, add more restriction to int32 Signed-off-by: tdophung --------- Signed-off-by: tdophung --- tests/jax/test_permutation.py | 13 ++----------- 1 file changed, 2 insertions(+), 11 deletions(-) diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index d61ea8eb75..5bb59c6ed5 100644 --- a/tests/jax/test_permutation.py +++ b/tests/jax/test_permutation.py @@ -19,11 +19,6 @@ from utils import assert_allclose, pytest_parametrize_wrapper -# ============================================================================= -# Test parameter definitions with L0 (fast) and L2 (comprehensive) levels -# ============================================================================= - -# All dispatch/combine test cases ALL_DISPATCH_COMBINE_CASES = [ (128, 5, 128, 3), (1024, 8, 128, 8), @@ -35,7 +30,6 @@ "L2": ALL_DISPATCH_COMBINE_CASES, } -# All sort chunks test cases ALL_SORT_CHUNKS_CASES = [ (8, 4096, 1280), (64, 4096, 4096), @@ -46,7 +40,6 @@ "L2": ALL_SORT_CHUNKS_CASES, } -# All dispatch/combine with padding test cases ALL_DISPATCH_COMBINE_PADDING_CASES = [ (128, 5, 128, 3, 8), (1024, 8, 128, 8, 16), @@ -58,14 +51,12 @@ "L2": ALL_DISPATCH_COMBINE_PADDING_CASES, } -# Dtypes for testing ALL_DTYPES = [jnp.float32, jnp.bfloat16] DTYPES = { "L0": ALL_DTYPES, "L2": ALL_DTYPES, } -# With probs options ALL_WITH_PROBS = [True, False] WITH_PROBS = { "L0": [True], @@ -389,7 +380,7 @@ def reference_make_chunk_sort_map( # For each source chunk, compute its destination offset # inverse_indices[i] = position of chunk i in sorted order - inverse_indices = jnp.argsort(sorted_indices) + inverse_indices = jnp.argsort(sorted_indices).astype(jnp.int32) dest_offsets = dest_cumsum[inverse_indices] # Create row_id_map: for each token position, compute its destination @@ -397,7 +388,7 @@ def reference_make_chunk_sort_map( position_indices = jnp.arange(num_tokens, dtype=jnp.int32) # chunk_ids[i] = which chunk position i belongs to - chunk_ids = jnp.searchsorted(src_cumsum[1:], position_indices, side="right") + chunk_ids = jnp.searchsorted(src_cumsum[1:], position_indices, side="right").astype(jnp.int32) # within_chunk_offset[i] = position i's offset within its chunk within_chunk_offset = position_indices - src_cumsum[chunk_ids] From 5f828c25d90f1535db2895528117beccd37faaec Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 7 Jan 2026 17:23:12 -0800 Subject: [PATCH 156/521] Solve pytorch-triton and triton package contention (#2540) * Add triton version detection logic, and NVTE_USE_PYTORCH_TRITON knob for jax Signed-off-by: tdophung * change build requirements and installation to reflect new option Signed-off-by: tdophung * reduce boilerplate comments Signed-off-by: tdophung * format code Signed-off-by: tdophung * fix typo Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * make env var more precise Signed-off-by: tdophung * make env variables checking consitent Signed-off-by: tdophung --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- build_tools/jax.py | 25 ++- build_tools/pytorch.py | 14 +- .../jax/triton_extensions/__init__.py | 33 +++- .../jax/triton_extensions/utils.py | 162 +++++++++++++++++- 4 files changed, 228 insertions(+), 6 deletions(-) diff --git a/build_tools/jax.py b/build_tools/jax.py index 276c9943d6..f07c0a202f 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -19,8 +19,29 @@ def install_requirements() -> List[str]: def test_requirements() -> List[str]: - """Test dependencies for TE/JAX extensions.""" - return ["numpy", "triton"] + """Test dependencies for TE/JAX extensions. + + Triton Package Selection: + The triton package is selected based on NVTE_USE_PYTORCH_TRITON environment variable: + + Default (NVTE_USE_PYTORCH_TRITON unset or "0"): + Returns 'triton' - OpenAI's standard package from PyPI. + Install with: pip install triton + + NVTE_USE_PYTORCH_TRITON=1: + Returns 'pytorch-triton' - for mixed JAX+PyTorch environments. + Install with: pip install pytorch-triton --index-url https://download.pytorch.org/whl/cu121 + + Note: Do NOT install pytorch-triton from PyPI directly - that's a placeholder. + """ + use_pytorch_triton = bool(int(os.environ.get("NVTE_USE_PYTORCH_TRITON", "0"))) + + triton_package = "pytorch-triton" if use_pytorch_triton else "triton" + + return [ + "numpy", + triton_package, + ] def xla_path() -> str: diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index b4815a0942..98511e45cb 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -13,7 +13,17 @@ def install_requirements() -> List[str]: - """Install dependencies for TE/PyTorch extensions.""" + """Install dependencies for TE/PyTorch extensions. + + IMPORTANT - PyTorch Index Required for pytorch-triton: + These dependencies MUST be installed using PyTorch's package index: + + pip install pytorch-triton --index-url https://download.pytorch.org/whl/ + + - pytorch-triton is only available from PyTorch's index (not PyPI) + - The 'pytorch-triton' package on PyPI is a placeholder that will fail + - torch.compile() requires pytorch-triton, not OpenAI's 'triton' package + """ return [ "torch>=2.1", "einops", @@ -22,7 +32,7 @@ def install_requirements() -> List[str]: "packaging", "pydantic", "nvdlfw-inspect", - "triton", + "pytorch-triton", ] diff --git a/transformer_engine/jax/triton_extensions/__init__.py b/transformer_engine/jax/triton_extensions/__init__.py index c98254d7d3..d9708fde9f 100644 --- a/transformer_engine/jax/triton_extensions/__init__.py +++ b/transformer_engine/jax/triton_extensions/__init__.py @@ -9,7 +9,33 @@ IMPORTANT: This module requires Triton to be installed. If you don't have Triton, use transformer_engine.jax.cpp_extensions instead (CUDA/FFI based primitives). -Install Triton: pip install triton + +Triton Package Options: +----------------------- +There are two compatible Triton packages: + +1. Standard 'triton' from OpenAI (recommended for JAX-only environments): + pip install triton + +2. 'pytorch-triton' from PyTorch's index (for mixed JAX+PyTorch environments): + pip install torch --index-url https://download.pytorch.org/whl/cu121 + # pytorch-triton is automatically installed as a dependency + + Both packages work with JAX Triton kernels. The pytorch-triton package + has version format "X.Y.Z+" (e.g., "3.0.0+45fff310c8"). + +WARNING: Do NOT run 'pip install pytorch-triton' directly! The package on PyPI +is a placeholder that will fail with "RuntimeError: Should never be installed". +The real pytorch-triton only comes bundled with PyTorch from PyTorch's index. + + +Environment Variables: + NVTE_USE_PYTORCH_TRITON: If set to "1", acknowledge using pytorch-triton + for JAX Triton kernels (suppresses compatibility warnings). Set this + when both JAX and PyTorch are installed in the same environment. + + Example: + export NVTE_USE_PYTORCH_TRITON=1 Usage: @@ -23,6 +49,11 @@ def lowering(ctx, x, **kwargs): # Use permutation functions from transformer_engine.jax.triton_extensions import make_row_id_map, permute_with_mask_map + + # Check Triton package info + from transformer_engine.jax.triton_extensions import get_triton_info + info = get_triton_info() + print(f"Using Triton {info['version']} from {info['source']}") """ from .utils import * diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 064b2843c6..59fc5c60af 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -6,9 +6,33 @@ This module provides utility functions for integrating Triton kernels into JAX primitives. Triton is only imported when this module is used. + +Triton Package Compatibility: + There are two Triton packages that can be used: + + 1. 'triton' (from OpenAI/PyPI): Standard package, works with JAX out of the box. + Install with: pip install triton + + 2. 'pytorch-triton' (from PyTorch's index): Bundled with PyTorch, includes + PyTorch-specific patches. Version format: "3.0.0+" + + IMPORTANT: The 'pytorch-triton' package on PyPI (version 0.0.1) is a + placeholder that will NOT work. The real pytorch-triton is only available + from PyTorch's package index and is auto-installed with PyTorch: + pip install torch --index-url https://download.pytorch.org/whl/cu121 + + pytorch-triton has been tested to work with JAX Triton kernels. + +Environment Variables: + NVTE_USE_PYTORCH_TRITON: If set to "1", explicitly acknowledge using + pytorch-triton for JAX Triton kernels (suppresses warnings). This is + useful when both JAX and PyTorch are installed in the same environment. + Default is "0". """ import hashlib +import os +import warnings from typing import Any, Callable, Mapping import zlib @@ -17,6 +41,113 @@ import jax.numpy as jnp +# Placeholder package version on PyPI that should never be used +_PYTORCH_TRITON_PLACEHOLDER_VERSION = "0.0.1" + + +def _detect_triton_package(): + """Detect which Triton package is installed and validate compatibility. + + Returns: + tuple: (triton_version: str or None, is_pytorch_triton: bool, is_placeholder: bool) + + The function detects: + - None: Triton not installed + - Standard triton from OpenAI (versions like "3.1.0") + - Real pytorch-triton from PyTorch's index (versions like "3.0.0+45fff310c8") + - Placeholder pytorch-triton from PyPI (version "0.0.1" - broken, raises RuntimeError) + """ + try: + import triton + + triton_version = getattr(triton, "__version__", "unknown") + except ImportError: + return None, False, False + except RuntimeError as e: + # The placeholder pytorch-triton package from PyPI raises: + # RuntimeError: "Should never be installed" + if "Should never be installed" in str(e): + return _PYTORCH_TRITON_PLACEHOLDER_VERSION, False, True + raise + + # Check for placeholder package (version 0.0.1 from PyPI) + is_placeholder = triton_version == _PYTORCH_TRITON_PLACEHOLDER_VERSION + + # Real pytorch-triton versions have a commit SHA suffix like "3.0.0+45fff310c8" + is_pytorch_triton = "+" in triton_version and len(triton_version.split("+")[-1]) >= 8 + + return triton_version, is_pytorch_triton, is_placeholder + + +def _check_triton_compatibility(): + """Check Triton package compatibility and emit warnings if necessary. + + This function handles the case where both JAX and PyTorch may be installed, + each expecting different Triton packages: + - JAX typically uses the standard 'triton' package from OpenAI + - PyTorch uses 'pytorch-triton' which is versioned with commit SHAs + + The NVTE_USE_PYTORCH_TRITON environment variable can be used to explicitly + acknowledge using pytorch-triton with JAX (suppresses warnings). + + Raises: + ImportError: If triton is not installed or the placeholder package is detected. + """ + triton_version, is_pytorch_triton, is_placeholder = _detect_triton_package() + + # Handle placeholder package from PyPI + if is_placeholder: + raise ImportError( + "Detected the placeholder 'pytorch-triton' package (version 0.0.1) from PyPI.\n" + "This is NOT a functional Triton installation.\n\n" + "The placeholder package exists to prevent namespace conflicts. To fix this:\n\n" + "Option 1 - Use standard Triton (recommended for JAX-only environments):\n" + " pip uninstall pytorch-triton triton\n" + " pip install triton\n\n" + "Option 2 - Use real pytorch-triton (for mixed JAX+PyTorch environments):\n" + " pip uninstall pytorch-triton triton\n" + " pip install torch --index-url https://download.pytorch.org/whl/cu121\n" + " # pytorch-triton is automatically installed as a torch dependency\n\n" + "Note: Do NOT run 'pip install pytorch-triton' directly - this installs\n" + "the broken placeholder. The real pytorch-triton only comes from PyTorch's index." + ) + + if triton_version is None: + raise ImportError( + "Triton is required for transformer_engine.jax.triton_extensions.\n\n" + "Option 1 - Install standard Triton (recommended for JAX-only):\n" + " pip install triton\n\n" + "Option 2 - Install PyTorch with pytorch-triton (for mixed environments):\n" + " pip install torch --index-url https://download.pytorch.org/whl/cu121\n\n" + "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead." + ) + + use_pytorch_triton_explicit = bool(int(os.environ.get("NVTE_USE_PYTORCH_TRITON", "0"))) + + if is_pytorch_triton: + if use_pytorch_triton_explicit: + # User explicitly opted in - just log info (no warning) + pass # Silent acknowledgment, no warning needed + else: + # pytorch-triton detected but user didn't explicitly opt in + warnings.warn( + f"Detected pytorch-triton package (version {triton_version}) instead of the" + " standard 'triton' package from OpenAI. This typically happens when PyTorch is" + " installed alongside JAX.\n\npytorch-triton is compatible with JAX Triton" + " kernels. To suppress this warning, set:\n export" + " NVTE_USE_PYTORCH_TRITON=1\n\nAlternatively, for a JAX-only environment:\n - Use" + " separate virtual environments for JAX and PyTorch, or\n - Use" + " transformer_engine.jax.cpp_extensions instead (CUDA-based, no Triton needed)", + category=UserWarning, + stacklevel=3, + ) + + return triton_version, is_pytorch_triton + + +# Perform compatibility check and get triton info +_TRITON_VERSION, _IS_PYTORCH_TRITON = _check_triton_compatibility() + try: from jax._src.lib import gpu_triton from triton.compiler import compiler as tc @@ -30,12 +161,41 @@ ) from e -__all__ = ["triton_call_lowering"] +__all__ = ["triton_call_lowering", "get_triton_info"] # Triton kernel cache (module-level, shared across all kernels) _TRITON_KERNEL_CACHE = {} +def get_triton_info(): + """Get information about the installed Triton package. + + Returns: + dict: Dictionary containing: + - version (str): Triton version string (e.g., "3.1.0" or "3.0.0+45fff310c8") + - is_pytorch_triton (bool): True if using real pytorch-triton from PyTorch's index + - is_openai_triton (bool): True if using standard triton from OpenAI/PyPI + - env_acknowledged (bool): True if NVTE_USE_PYTORCH_TRITON=1 is set + - source (str): "pytorch" or "openai" indicating the package source + + Example: + from transformer_engine.jax.triton_extensions import get_triton_info + info = get_triton_info() + print(f"Triton version: {info['version']} (from {info['source']})") + if info['is_pytorch_triton']: + print("Using pytorch-triton - compatible with both PyTorch and JAX") + """ + env_acknowledged = bool(int(os.environ.get("NVTE_USE_PYTORCH_TRITON", "0"))) + + return { + "version": _TRITON_VERSION, + "is_pytorch_triton": _IS_PYTORCH_TRITON, + "is_openai_triton": not _IS_PYTORCH_TRITON, + "env_acknowledged": env_acknowledged and _IS_PYTORCH_TRITON, + "source": "pytorch" if _IS_PYTORCH_TRITON else "openai", + } + + def get_triton_dtype(aval): """Convert JAX dtype to Triton type string. From 5f0e3b935a18595e8915f9068d9c2181407d96f9 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 8 Jan 2026 16:30:34 -0800 Subject: [PATCH 157/521] [JAX] Refactor and trim TE JAX Attn testing (#2542) * Pick a leaner set of combinations for TE JAX CP attn tests such that only those cp,dp,tp combinations are picked where cp*dp*tp is equal to num gpus Signed-off-by: Kshitij Lakhani * Consolidate the test cases run for different B,S,H,D and QKV layout Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Code and comments 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 * Make FP16 + GQA test cross attn instead of self attn to generalize the test Signed-off-by: Kshitij Lakhani --------- Signed-off-by: Kshitij Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/distributed_test_base.py | 8 +- tests/jax/test_distributed_fused_attn.py | 2 +- tests/jax/test_fused_attn.py | 169 ++++++++++++++++++++--- tests/jax/utils.py | 7 + 4 files changed, 162 insertions(+), 24 deletions(-) diff --git a/tests/jax/distributed_test_base.py b/tests/jax/distributed_test_base.py index aa7a8fb8d5..6d963f5c7b 100644 --- a/tests/jax/distributed_test_base.py +++ b/tests/jax/distributed_test_base.py @@ -12,7 +12,7 @@ from transformer_engine.jax.sharding import MeshResource -from utils import assert_allclose, is_devices_enough +from utils import assert_allclose, is_devices_enough, is_devices_equal def generate_configs(): @@ -49,7 +49,11 @@ def generate_context_parallel_configs_for_attn(): TP_sizes = (1, 2) for dp, cp, tp in product(DP_sizes, CP_sizes, TP_sizes): ndev = cp * tp * dp - if is_devices_enough(ndev): + # Run only those dp,cp,tp combinations which require exactly ndev GPUs. + # For e.g., if num_GPUs is 8 and ndev=8 , all the dp,cp,tp combinations fulfilling ndev = cp * tp * dp are picked. + # However, if num_GPUs is 8 and ndev=4, then all the dp,cp,tp combinations fulfilling ndev = cp * tp * dp are ignored. + # To explicitly pick combinations associated with ndev=4, one can set CUDA_VISIBLE_DEVICES=0,1,2,3, thereby forcing num_GPUs to 4 instead of 8. + if is_devices_equal(ndev): # Do not run cp1 case in L1 as that is already covered in TestDistributedSelfAttn and TestDistributedCrossAttn (as these do not have any cp combinations) if cp != 1: configsL1.append( diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index d0018543d1..d5ebe9f261 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -334,7 +334,7 @@ def test_cross_attn( class TestDistributedContextParallelSelfAttn: - + # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests def impl_test_context_parallel_attn( self, device_count, diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index ac1b7c3505..a0aee50430 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1068,41 +1068,70 @@ def check_dqkv(primitive, reference, pad, idx): ], ) @pytest.mark.parametrize( - "qkv_layout", - [ - pytest.param(QKVLayout.BS3HD, id="QKV_PACKED"), - pytest.param(QKVLayout.BSHD_BS2HD, id="KV_PACKED"), - pytest.param(QKVLayout.BSHD_BSHD_BSHD, id="SEPARATE"), - pytest.param(QKVLayout.T3HD, id="RAGGED_QKV_PACKED"), - pytest.param(QKVLayout.THD_T2HD, id="RAGGED_KV_PACKED"), - pytest.param(QKVLayout.THD_THD_THD, id="RAGGED_SEPARATE"), - ], -) -@pytest.mark.parametrize( - "b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype", + "b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype, qkv_layout", [ + # large data size + bf16 + qkv packed pytest.param( - 2, 2048, 2048, 12, 12, 64, 64, jnp.bfloat16, id="2-2048-2048-12-12-64-64-BF16-SELF" + 2, + 2048, + 2048, + 12, + 12, + 64, + 64, + jnp.bfloat16, + QKVLayout.BS3HD, + id="2-2048-2048-12-12-64-64-BF16-SELF-QKV_PACKED", ), pytest.param( 2, - 512, - 1024, + 2048, + 2048, 12, 12, 64, 64, jnp.bfloat16, - id="2-512-1024-12-12-64-64-BF16-CROSS", + QKVLayout.T3HD, + id="2-2048-2048-12-12-64-64-BF16-SELF-RAGGED_QKV_PACKED", ), + # mid data size + bf16 + cross attn + kv packed pytest.param( - 2, 2048, 2048, 12, 6, 64, 64, jnp.bfloat16, id="2-2048-2048-12-6-64-64-BF16-GQA" + 2, + 512, + 1024, + 12, + 12, + 64, + 64, + jnp.bfloat16, + QKVLayout.BSHD_BS2HD, + id="2-512-1024-12-12-64-64-BF16-CROSS-KV_PACKED", ), pytest.param( - 4, 128, 128, 16, 16, 64, 64, jnp.float16, id="4-128-128-16-16-64-64-FP16-SELF" + 2, + 512, + 1024, + 12, + 12, + 64, + 64, + jnp.bfloat16, + QKVLayout.THD_T2HD, + id="2-512-1024-12-12-64-64-BF16-CROSS-RAGGED_KV_PACKED", ), + # large data size + bf16 + cross attn + diff hidden v dim + qkv separate pytest.param( - 4, 128, 128, 16, 16, 64, 32, jnp.float16, id="4-128-128-16-16-64-32-FP16-SELF" + 2, + 2048, + 1024, + 12, + 12, + 64, + 32, + jnp.bfloat16, + QKVLayout.BSHD_BSHD_BSHD, + id="2-2048-1024-12-12-64-32-BF16-CROSS-SEPARATE", ), pytest.param( 2, @@ -1113,10 +1142,108 @@ def check_dqkv(primitive, reference, pad, idx): 64, 32, jnp.bfloat16, - id="2-2048-1024-12-12-64-32-BF16-CROSS", + QKVLayout.THD_THD_THD, + id="2-2048-1024-12-12-64-32-BF16-CROSS-RAGGED_SEPARATE", + ), + # large data size + bf16 + gqa + kv packed + pytest.param( + 2, + 2048, + 2048, + 12, + 6, + 64, + 64, + jnp.bfloat16, + QKVLayout.BSHD_BS2HD, + id="2-2048-2048-12-6-64-64-BF16-GQA-KV_PACKED", + ), + pytest.param( + 2, + 2048, + 2048, + 12, + 6, + 64, + 64, + jnp.bfloat16, + QKVLayout.THD_T2HD, + id="2-2048-2048-12-6-64-64-BF16-GQA-RAGGED_KV_PACKED", ), + # small data size + fp16 + diff hidden v dim + qkv packed pytest.param( - 2, 2048, 2048, 12, 6, 128, 64, jnp.float16, id="2-2048-2048-12-6-128-64-FP16-GQA" + 4, + 128, + 128, + 16, + 16, + 64, + 32, + jnp.float16, + QKVLayout.BS3HD, + id="4-128-128-16-16-64-32-FP16-SELF-QKV_PACKED", + ), + pytest.param( + 4, + 128, + 128, + 16, + 16, + 64, + 32, + jnp.float16, + QKVLayout.T3HD, + id="4-128-128-16-16-64-32-FP16-SELF-RAGGED_QKV_PACKED", + ), + # small data size + fp16 + kv packed + pytest.param( + 4, + 128, + 128, + 16, + 16, + 64, + 64, + jnp.float16, + QKVLayout.BSHD_BS2HD, + id="4-128-128-16-16-64-64-FP16-SELF-KV_PACKED", + ), + pytest.param( + 4, + 128, + 128, + 16, + 16, + 64, + 64, + jnp.float16, + QKVLayout.THD_T2HD, + id="4-128-128-16-16-64-64-FP16-SELF-RAGGED_KV_PACKED", + ), + # large data size + fp16 + cross attn + gqa + diff hidden v dim + qkv separate + pytest.param( + 2, + 1024, + 2048, + 12, + 6, + 128, + 64, + jnp.float16, + QKVLayout.BSHD_BSHD_BSHD, + id="2-1024-2048-12-6-128-64-FP16-CROSS-GQA-SEPARATE", + ), + pytest.param( + 2, + 1024, + 2048, + 12, + 6, + 128, + 64, + jnp.float16, + QKVLayout.THD_THD_THD, + id="2-1024-2048-12-6-128-64-FP16-CROSS-GQA-RAGGED_SEPARATE", ), ], ) diff --git a/tests/jax/utils.py b/tests/jax/utils.py index 8055792308..c22b0a6063 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -47,6 +47,13 @@ def is_devices_enough(required): return len(jax.devices()) >= required +def is_devices_equal(required): + """ + Check if the available GPUs is exactly equal + """ + return len(jax.devices()) == required + + def _generate_drop_path_shape(shape: Sequence[int], batch_dim: int) -> Sequence[int]: # Generate broadcast dims for drop_path. drop_path_shape = list(range(0, len(shape))) From 32f403fd8bcf7f74559cf3526201bd4d68960c2b Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Fri, 9 Jan 2026 13:13:20 -0800 Subject: [PATCH 158/521] Update list of authorized CI users (#2581) Update list of CI users Signed-off-by: Tim Moon --- .github/workflows/trigger-ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml index 13ea45b070..c56601ae98 100644 --- a/.github/workflows/trigger-ci.yml +++ b/.github/workflows/trigger-ci.yml @@ -56,8 +56,8 @@ jobs: || github.actor == 'vcherepanov-nv' || github.actor == 'tdophung' || github.actor == 'vthumbe1503' - || github.actor == 'janekb04' || github.actor == 'shengfangd' + || github.actor == 'kainzhong' ) steps: - name: Check if comment is issued by authorized person From 2f8ae81c3b78db38f5ace8735eedb66269159c91 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Fri, 9 Jan 2026 19:05:28 -0800 Subject: [PATCH 159/521] Debug doc generation (#2576) Debug Doxygen and LaTeX warnings Signed-off-by: Tim Moon --- .gitignore | 1 + .../common/include/transformer_engine/cast.h | 5 +++-- .../transformer_engine/cast_transpose_noop.h | 4 ++-- .../include/transformer_engine/fused_attn.h | 6 ++--- .../common/include/transformer_engine/gemm.h | 2 ++ .../include/transformer_engine/multi_tensor.h | 2 +- .../transformer_engine/normalization.h | 14 +++++++++--- .../include/transformer_engine/recipe.h | 4 ++-- .../include/transformer_engine/swizzle.h | 1 + .../transformer_engine/transformer_engine.h | 22 ++++++++++++------- .../include/transformer_engine/transpose.h | 10 ++++----- 11 files changed, 44 insertions(+), 27 deletions(-) diff --git a/.gitignore b/.gitignore index 74acd6ad7f..7a86041a1e 100644 --- a/.gitignore +++ b/.gitignore @@ -41,3 +41,4 @@ compile_commands.json .nfs tensor_dumps/ artifacts/ +*.DS_Store diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 5e67b7645e..576494a4de 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -130,7 +130,7 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, * \param[in] stream CUDA stream used for the operation. */ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor dbias, - NVTETensor workplace, cudaStream_t stream); + NVTETensor workspace, cudaStream_t stream); /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. @@ -263,7 +263,8 @@ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t str * * \param[in] inputs List of input tensors to be cast. * \param[in,out] outputs List of output quantized tensors. - * \param[in] quant_config (Optional) Quantization configurations. + * \param[in] quant_config (Optional) Quantization configurations. + * \param[in] num_tensors Number of input and output tensors. * \param[in] stream CUDA stream used for the operation. */ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, diff --git a/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h b/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h index 5bef067910..d21a28e521 100644 --- a/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h +++ b/transformer_engine/common/include/transformer_engine/cast_transpose_noop.h @@ -4,8 +4,8 @@ * See LICENSE for license information. ************************************************************************/ -/*! \file transpose_with_noop.h - * \brief Functions handling transposes with no-op. +/*! \file cast_transpose_noop.h + * \brief Transpose functions with no-op flag. */ #ifndef TRANSFORMER_ENGINE_CAST_TRANSPOSE_WITH_NOOP_H_ diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 16b4b8ff4d..dd70ccf8df 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -409,7 +409,6 @@ void nvte_fused_attn_bwd_qkvpacked( * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). - * \param[in] deterministic Whether to execute with deterministic behaviours. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ @@ -673,7 +672,7 @@ void nvte_populate_rng_state_async(NVTETensor rng_state_dst, const NVTETensor se * \param[in] len batch_size x sequence_length. * \param[in] stream CUDA stream used for this operation. */ -uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlen, NVTETensor workspace, size_t len, +uint32_t nvte_get_runtime_num_segments(NVTETensor cu_seqlens, NVTETensor workspace, size_t len, cudaStream_t stream); /*! \brief Set the seed and offset for RNG state. @@ -830,8 +829,7 @@ void nvte_convert_thd_to_bshd(NVTETensor tensor, NVTETensor cu_seqlens, NVTETens * \param[in] tensor Input tensor. * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. * \param[out] new_tensor Output tensor. - * \param[in] b Batch size. - * \param[in] max_seq_len Maximum sequence length. + * \param[in] t Packed sequence length. * \param[in] stream CUDA stream used for this operation. */ void nvte_convert_bshd_to_thd(NVTETensor tensor, NVTETensor cu_seqlens, NVTETensor new_tensor, diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 38214f8872..8a51b54fdb 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -255,9 +255,11 @@ class MatmulConfigWrapper { MatmulConfigWrapper(const MatmulConfigWrapper &) = delete; MatmulConfigWrapper &operator=(const MatmulConfigWrapper &) = delete; + /*! \brief Move constructor. */ MatmulConfigWrapper(MatmulConfigWrapper &&other) : config_{other.config_} { other.config_ = nullptr; } + /*! \brief Move-assignment operator. */ MatmulConfigWrapper &operator=(MatmulConfigWrapper &&other) { if (config_ != nullptr) { nvte_destroy_matmul_config(config_); diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 1bea4cb21f..303801a88a 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -288,7 +288,7 @@ void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor ** * and populate the amax of the corresponding output tensor. * * \param[in] input Input tensor. - * \param[in,out] amaxes Array of output tensors. Only the amax is updated. + * \param[in,out] outputs Array of output tensors. Only the amax is updated. * \param[in] split_sections Size of each tensor split along dimension 0. * \param[in] num_tensors Number of tensor splits. * \param[in] stream CUDA stream used for the operation. diff --git a/transformer_engine/common/include/transformer_engine/normalization.h b/transformer_engine/common/include/transformer_engine/normalization.h index 7f5bd92fc1..29b98ca54f 100644 --- a/transformer_engine/common/include/transformer_engine/normalization.h +++ b/transformer_engine/common/include/transformer_engine/normalization.h @@ -163,11 +163,16 @@ void nvte_rmsnorm_bwd_add(const NVTETensor dz, const NVTETensor x, const NVTETen NVTETensor dgamma, NVTETensor workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream); -/*! \brief Helper to enable cuDNN backend for normalization +/*! \brief Set whether to enable cuDNN backend for normalization forward. * - * \param[in] bool Enable if True + * \param[in] enable Whether to enable cuDNN backend. */ void nvte_enable_cudnn_norm_fwd(bool enable); + +/*! \brief Set whether to enable cuDNN backend for normalization backward. + * + * \param[in] enable Whether to enable cuDNN backend. + */ void nvte_enable_cudnn_norm_bwd(bool enable); /*! \brief Control whether norm computes `gamma += 1.0` for zero-centered gamma @@ -176,11 +181,14 @@ void nvte_enable_cudnn_norm_bwd(bool enable); * Currently this only applies to the CuDNN backend. If CuDNN is not used, * this setting has no effect. * - * \param[in] bool Enable if True + * \param[in] enable Whether to enable zero-centered gamma. */ void nvte_enable_zero_centered_gamma_in_weight_dtype(bool enable); +#ifdef __cplusplus +/*! \brief Normalization function type */ enum class NVTE_Norm_Type { LayerNorm, RMSNorm }; +#endif #ifdef __cplusplus } // extern "C" diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index 83d436db30..c0cec8a3b9 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -23,7 +23,7 @@ extern "C" { * the last, the last entry shifts to the second to last) and the * first entry is set to zero. The scaling factor is estimated so the * FP8 tensor's maximum absolute value is - * @f$ 2^{-\text{margin}} \text{max}_\text{fp8\_dtype} @f$. + * @f$ 2^{-margin} \max_{fp8\_dtype} @f$. * * \param[in] amax_history History of maximum absolute values. * Shape: [history_length, num_scales] @@ -54,7 +54,7 @@ void nvte_delayed_scaling_recipe_amax_and_scale_update( * the last, the last entry shifts to the second to last) and the * first entry is set to zero. The scaling factor is estimated so the * FP8 tensor's maximum absolute value is - * @f$ 2^{-\text{margin}} \text{max}_\text{fp8\_dtype} @f$. + * @f$ 2^{-margin} \max_{fp8\_dtype} @f$. * * \param[in] amax_reduction_buffer The contiguous buffer used for amax reduction. * Shape: [num_scales * num_tensors] diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index b4489abdd8..4e3544d3c7 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -34,6 +34,7 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud * * \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: diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 7fc9d78980..fd0125c8d0 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -51,8 +51,11 @@ struct NVTEShape { * It does not own the memory it points to. */ struct NVTEBasicTensor { + /*! Pointer to data buffer. */ void *data_ptr; + /*! Data type. */ NVTEDType dtype; + /*! Tensor shape. */ NVTEShape shape; }; @@ -144,7 +147,7 @@ void *nvte_tensor_columnwise_data(const NVTETensor tensor); * * \param[data] Pointer to start of shape array. If NULL, the shape * will be filled with zeros. - * \param[data] Number of dimensions (must be <= 14) + * \param[ndim] Number of dimensions (must be <= 14) * * \return A shape. The shape will own its own copy of the data. */ @@ -177,7 +180,7 @@ size_t nvte_tensor_ndims(const NVTETensor tensor); /*! \brief Get the size of a specific tensor dimension. * * \param[in] tensor Tensor. - * \param[in] size_t Dimension index. + * \param[in] dim Dimension index. * * \return Size of the tensor at the specified dimension. */ @@ -258,8 +261,7 @@ NVTEShape nvte_tensor_scale_inv_shape(const NVTETensor tensor); /*! \brief Reset tensor value to zero. * * \param[in] tensor Tensor. - * - * \return A scale_inv shape of the input tensor. + * \param[in] stream CUDA stream to use for the operation. */ void nvte_zero_tensor(const NVTETensor tensor, cudaStream_t stream); @@ -539,7 +541,7 @@ enum class DType { /*! \brief Check if TE datatype is FP8 * * Return true if TE datatype is FP8 - * \param[in] DType TE Datatype of interest + * \param[in] t TE Datatype of interest */ inline bool is_fp8_dtype(const DType t) { return t == DType::kFloat8E4M3 || t == DType::kFloat8E5M2; @@ -548,14 +550,14 @@ inline bool is_fp8_dtype(const DType t) { /*! \brief Check if TE datatype is FP4 * * Return true if TE datatype is FP4 - * \param[in] DType TE Datatype of interest + * \param[in] t TE Datatype of interest */ inline bool is_fp4_dtype(const DType t) { return t == DType::kFloat4E2M1; } /*! \brief Check if TE datatype is high precision (FP32, FP16, BF16) * * Return true if TE datatype is high precision - * \param[in] DType TE Datatype of interest + * \param[in] t TE Datatype of interest */ inline bool is_high_precision_dtype(const DType t) { return t == DType::kFloat32 || t == DType::kBFloat16 || t == DType::kFloat16; @@ -579,6 +581,7 @@ class TensorWrapper { * \param[in] scale_dptr Pointer to the scale value. * \param[in] scale_inv_shape Shape of scale_inv * \param[in] scale_inv_dptr Pointer to the inverse of scale value. + * \param[in] scaling_mode Tensor data format. */ TensorWrapper(void *dptr, const NVTEShape &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, float *scale_inv_dptr = nullptr, @@ -615,6 +618,7 @@ class TensorWrapper { * \param[in] scale_dptr Pointer to the scale value. * \param[in] scale_inv_shape Shape of scale_inv * \param[in] scale_inv_dptr Pointer to the inverse of scale value. + * \param[in] scaling_mode Tensor data format. */ TensorWrapper(void *dptr, const std::vector &shape, const DType dtype, float *amax_dptr = nullptr, float *scale_dptr = nullptr, @@ -766,7 +770,7 @@ class TensorWrapper { /*! \brief Get the size of this TensorWrapper in the given dimension. * - * \param[in] size_t Dimension index. + * \param[in] dim Dimension index. * * \return Size of this TensorWrapper in given dimension. */ @@ -935,9 +939,11 @@ class QuantizationConfigWrapper { QuantizationConfigWrapper(const QuantizationConfigWrapper &) = delete; QuantizationConfigWrapper &operator=(const QuantizationConfigWrapper &) = delete; + /*! \brief Move constructor. */ QuantizationConfigWrapper(QuantizationConfigWrapper &&other) : config_{other.config_} { other.config_ = nullptr; } + /*! \brief Move-assignment operator. */ QuantizationConfigWrapper &operator=(QuantizationConfigWrapper &&other) { if (config_ != nullptr) { nvte_destroy_quantization_config(config_); diff --git a/transformer_engine/common/include/transformer_engine/transpose.h b/transformer_engine/common/include/transformer_engine/transpose.h index cd73935abb..5f9a8fe149 100644 --- a/transformer_engine/common/include/transformer_engine/transpose.h +++ b/transformer_engine/common/include/transformer_engine/transpose.h @@ -231,7 +231,7 @@ void nvte_cast_transpose_dbias_dsrelu(const NVTETensor input, const NVTETensor a * - columnwise data of `output` is equal to `transpose(cast(dact(input)))` * * \param[in] input Input tensor of shape [N, H]. - * \param[in] gated_act_input Tensor used as input to the forward of + * \param[in] act_input Tensor used as input to the forward of * gated activation operation. * Shape [N, H * 2]. * \param[in,out] output Result of the cast. @@ -250,7 +250,7 @@ void nvte_dgeglu_cast_transpose(const NVTETensor input, const NVTETensor act_inp * - columnwise data of `output` is equal to `transpose(cast(dact(input)))` * * \param[in] input Input tensor of shape [N, H]. - * \param[in] gated_act_input Tensor used as input to the forward of + * \param[in] act_input Tensor used as input to the forward of * gated activation operation. * Shape [N, H * 2]. * \param[in,out] output Result of the cast. @@ -269,7 +269,7 @@ void nvte_dswiglu_cast_transpose(const NVTETensor input, const NVTETensor act_in * - columnwise data of `output` is equal to `transpose(cast(dact(input)))` * * \param[in] input Input tensor of shape [N, H]. - * \param[in] gated_act_input Tensor used as input to the forward of + * \param[in] act_input Tensor used as input to the forward of * gated activation operation. * Shape [N, H * 2]. * \param[in,out] output Result of the cast. @@ -288,7 +288,7 @@ void nvte_dreglu_cast_transpose(const NVTETensor input, const NVTETensor act_inp * - columnwise data of `output` is equal to `transpose(cast(dact(input)))` * * \param[in] input Input tensor of shape [N, H]. - * \param[in] gated_act_input Tensor used as input to the forward of + * \param[in] act_input Tensor used as input to the forward of * gated activation operation. * Shape [N, H * 2]. * \param[in,out] output Result of the cast. @@ -307,7 +307,7 @@ void nvte_dqgeglu_cast_transpose(const NVTETensor input, const NVTETensor act_in * - columnwise data of `output` is equal to `transpose(cast(dact(input)))` * * \param[in] input Input tensor of shape [N, H]. - * \param[in] gated_act_input Tensor used as input to the forward of + * \param[in] act_input Tensor used as input to the forward of * gated activation operation. * Shape [N, H * 2]. * \param[in,out] output Result of the cast. From fe8fad59c9add7f8aa841fed2a0b4087b931856f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 13 Jan 2026 11:03:24 +0100 Subject: [PATCH 160/521] [PyTorch] Bunch of fixes for cpu offloading (#2535) * code drop Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix 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> --- tests/pytorch/test_cpu_offloading.py | 11 ++- transformer_engine/pytorch/cpu_offload.py | 72 +++++++++++++++---- transformer_engine/pytorch/csrc/extensions.h | 3 +- .../pytorch/csrc/extensions/cast.cpp | 37 +++++----- .../pytorch/csrc/extensions/pybind.cpp | 2 +- .../pytorch/module/grouped_linear.py | 7 +- transformer_engine/pytorch/module/linear.py | 3 +- .../pytorch/optimizers/fused_adam.py | 7 +- .../pytorch/quantized_tensor.py | 14 ---- 9 files changed, 103 insertions(+), 53 deletions(-) diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 385998a8c5..7da8dcf863 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -54,9 +54,13 @@ class Utils: + # Tensor used for simulating long-running GPU work in long_job() tensor1 = torch.randn((1024, 1024), device="cuda", dtype=torch.bfloat16) - _B = 64 - _S = 256 + # Test tensor dimensions: _B x _S x _D = 128 x 512 x 256 = 16,777,216 elements + # This exceeds the 256K element threshold for offloading (cpu_offload.py line 443). + # For quantized tensors, scale_inv tensors (~524K elements for block scaling) also exceed threshold. + _B = 128 + _S = 512 _H = 4 _D = 256 @@ -395,6 +399,9 @@ def test_multiple_tensor_offload(self, recipe): offload_synchronizer.push_tensor(x1) offload_synchronizer.push_tensor(x1) offload_synchronizer.push_tensor(x1) + # Verify x1 is not corrupted after pushing (important for QuantizedTensor) + if recipe is not None: + x1.dequantize() # Should not raise - tensor should still be valid offload_synchronizer.fwd_step() # Only one copy of tensor on cpu is allocated. assert Utils.get_cpu_memory_mb() == pytest.approx(init_cpu_memory + 1 * x_size, 0.1) diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index d0b8d3474e..05219b7b18 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -19,6 +19,7 @@ from .quantized_tensor import ( restore_from_saved, prepare_for_saving, + QuantizedTensor, ) @@ -255,6 +256,8 @@ def start_offload(self): Start offloading of tensors. Puts copy from GPU to CPU tasks on offload stream. Before each copy event, the offload stream waits for the event signalling that the tensor is ready to be offloaded. This event is recorded in the start_offload or push_tensor call. + + Note: tensor_list only contains regular tensors (QuantizedTensors are decomposed in push_tensor). """ self._validate_state(func_name="start_offload", allowed_states=["not_offloaded"]) self.state = "offload_started" @@ -275,19 +278,18 @@ def start_offload(self): with torch.cuda.stream(self.offload_stream): if allocate_cpu_buffers: - # empty_like is defined also for QuantizedTensors offloaded_tensor = torch.empty_like( tensor, device=torch.device("cpu"), pin_memory=True ) self.cpu_tensor_group.tensor_list.append(offloaded_tensor) else: - assert self.cpu_tensor_group.tensor_list[tensor_id].shape == tensor.shape, ( + offloaded_tensor = self.cpu_tensor_group.tensor_list[tensor_id] + assert offloaded_tensor.shape == tensor.shape, ( "CPU buffer shape does not match the offloaded tensor shape:" - f" {self.cpu_tensor_group.tensor_list[tensor_id].shape} != {tensor.shape} " - " Make sure that tensor shaped do not change between" + f" {offloaded_tensor.shape} != {tensor.shape} " + "Make sure that tensor shapes do not change between" " iterations if retain_pinned_cpu_buffers is True." ) - offloaded_tensor = self.cpu_tensor_group.tensor_list[tensor_id] offloaded_tensor.copy_(tensor, non_blocking=True) # aux is a dictionary that contains auxiliary data like information which tensors were deduplicated, @@ -318,6 +320,9 @@ def start_reload(self): """ Start reloading of tensors. It allocates new tensors on GPU and puts copy from CPU tasks on offload stream. + + Note: tensor_list only contains regular tensors (QuantizedTensors are decomposed in push_tensor + and reconstructed in pop_tensor). """ self._validate_state(func_name="start_reload", allowed_states=["offload_finished"]) self.state = "reload_started" @@ -330,7 +335,6 @@ def start_reload(self): # cannot move tensors from pool of one stream to another without # calling cudaFree and cudaMalloc again. - # empty_like is defined also for QuantizedTensors. reloaded_tensor = torch.empty_like(tensor, device=torch.device("cuda")) self.offload_stream.wait_stream(torch.cuda.current_stream()) @@ -347,16 +351,29 @@ def start_reload(self): self.bwd_gpu_tensor_group ) - def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor: + def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor | tuple[list, list]: """ It is called when a tensor is saved for backward pass. If tensor is offloaded, returns int representing the index of the tensor in the offloaded tensor group. If tensor is not offloaded, returns the tensor itself. + For QuantizedTensor, returns (list of push results for each component, tensor_objs) tuple. """ self._validate_state(func_name="push_tensor", allowed_states=["not_offloaded"]) if self._check_if_offload(tensor): + # For QuantizedTensor: decompose into component tensors, push each one recursively + if isinstance(tensor, QuantizedTensor): + # Make a copy because prepare_for_saving modifies the object (sets fields to None) + tensor_copy = tensor.detach() + # Inline prepare_for_saving logic - QuantizedTensor is a torch.Tensor subclass, + # so the generic prepare_for_saving would not call tensor.prepare_for_saving() + saved_tensors, tensor_obj = tensor_copy.prepare_for_saving() + push_results = [ + self.push_tensor(t) if t is not None else None for t in saved_tensors + ] + return (push_results, [tensor_obj]) + self.fwd_gpu_tensor_group.tensor_list.append(tensor) # The group is processed and offloaded at the end of the forward pass of current layer. # To enable offloading of tensors faster we use self.offload_stream and record @@ -370,23 +387,39 @@ def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor: return len(self.fwd_gpu_tensor_group.tensor_list) - 1 return tensor - def pop_tensor(self, tensor_or_tensor_id: torch.Tensor | int) -> torch.Tensor: + def pop_tensor( + self, tensor_or_tensor_id: torch.Tensor | int | tuple[list, list] + ) -> torch.Tensor: """ It is called when a tensor is used in backward pass. Returns the tensor. If tensor was offloaded/reloaded, wait for the reload of a tensor to finish. + For QuantizedTensor (tuple input), reconstructs from component tensors. """ self._validate_state( func_name="pop_tensor", allowed_states=["not_offloaded", "reload_started"] ) - # 1. tensor not offloaded + # 1. tensor not offloaded (regular tensor returned as-is from push) if isinstance(tensor_or_tensor_id, torch.Tensor): return tensor_or_tensor_id - # 2. the layer was not offloaded at all + + # 2. QuantizedTensor case: tuple of (push_results, tensor_objs) + if isinstance(tensor_or_tensor_id, tuple): + push_results, tensor_objs = tensor_or_tensor_id + # Recursively pop each component + reloaded_tensors = [ + self.pop_tensor(pr) if pr is not None else None for pr in push_results + ] + # Inline restore_from_saved - tensor_objs[0] is the QuantizedTensor copy + tensor_obj = tensor_objs[0] + tensor_obj.restore_from_saved(reloaded_tensors) + return tensor_obj + + # 3. Regular tensor index case if self.state == "not_offloaded": return self.fwd_gpu_tensor_group.tensor_list[tensor_or_tensor_id] - # 3. the layer was offloaded + # 4. the layer was offloaded assert self.state == "reload_started" # wait for the tensor to be reloaded torch.cuda.current_stream().wait_event( @@ -406,6 +439,10 @@ def _check_if_offload(self, t: torch.Tensor) -> bool: """ Check if tensor needs to be offloaded. """ + # Only offload tensors with at least 256k elements (~1MB for float32) + if t.numel() < 256 * 1024: + return False + if ( not isinstance(t, torch.nn.Parameter) and not getattr(t, "_TE_do_not_offload", False) @@ -418,7 +455,6 @@ def _check_if_offload(self, t: torch.Tensor) -> bool: " this tensor will be skipped." ) return False - return True return False @@ -488,11 +524,13 @@ def bwd_step(self, layer_num: int): self.previous_bwd_layer_id = layer_num self.current_layer_id = layer_num - def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor: + def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor | tuple[list, list]: """Default push tensor method""" return self.layer_states[self.num_of_fwds].push_tensor(tensor) - def pop_tensor(self, tensor_or_tensor_id: torch.Tensor | int) -> torch.Tensor: + def pop_tensor( + self, tensor_or_tensor_id: torch.Tensor | int | tuple[list, list] + ) -> torch.Tensor: """Default pop tensor method""" return self.layer_states[self.current_layer_id].pop_tensor(tensor_or_tensor_id) @@ -592,6 +630,12 @@ def bwd_step(self, layer_num: int): for layer in self.start_reload_map[layer_num]: self.layer_states[layer].start_reload() + def push_tensor(self, tensor: torch.Tensor) -> int | torch.Tensor | tuple[list, list]: + """Push tensor - skip processing if layer won't be offloaded to reduce CPU overhead.""" + if not self.offload_layer_map.get(self.num_of_fwds, False): + return tensor + return self.layer_states[self.num_of_fwds].push_tensor(tensor) + class ManualOffloadSynchronizer(OffloadSynchronizer): """ diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 52ef02a347..60b931abfd 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -254,7 +254,8 @@ std::vector multi_tensor_quantize(const std::vector &ten std::vector split_quantize(const at::Tensor &tensor, const std::vector &split_sections, - std::vector quantizer_list); + std::vector quantizer_list, + bool disable_bulk_allocation = false); /*************************************************************************************************** * Bias gradient fusions diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 4e5e5223f7..ac06841879 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1095,7 +1095,8 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, std::vector split_quantize(const at::Tensor &tensor, const std::vector &split_sections, - std::vector quantizer_list) { + std::vector quantizer_list, + bool disable_bulk_allocation) { init_extension(); // Check number of tensors @@ -1147,22 +1148,24 @@ std::vector split_quantize(const at::Tensor &tensor, enum class QuantizationMethod { UNFUSED, FUSED_NVFP4 }; AllocationMethod allocation_method = AllocationMethod::UNFUSED; QuantizationMethod quantization_method = QuantizationMethod::UNFUSED; - if (std::all_of(quantizer_list.begin(), quantizer_list.end(), - [](const py::handle &quantizer) -> bool { - return detail::IsFloat8BlockwiseQuantizers(quantizer.ptr()); - })) { - allocation_method = AllocationMethod::BULK_FP8_BLOCKWISE; - } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(), - [](const py::handle &quantizer) -> bool { - return detail::IsMXFP8Quantizers(quantizer.ptr()); - })) { - allocation_method = AllocationMethod::BULK_MXFP8; - } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(), - [](const py::handle &quantizer) -> bool { - return detail::IsNVFP4Quantizers(quantizer.ptr()); - })) { - allocation_method = AllocationMethod::BULK_NVFP4; - quantization_method = QuantizationMethod::FUSED_NVFP4; + if (!disable_bulk_allocation) { + if (std::all_of(quantizer_list.begin(), quantizer_list.end(), + [](const py::handle &quantizer) -> bool { + return detail::IsFloat8BlockwiseQuantizers(quantizer.ptr()); + })) { + allocation_method = AllocationMethod::BULK_FP8_BLOCKWISE; + } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(), + [](const py::handle &quantizer) -> bool { + return detail::IsMXFP8Quantizers(quantizer.ptr()); + })) { + allocation_method = AllocationMethod::BULK_MXFP8; + } else if (std::all_of(quantizer_list.begin(), quantizer_list.end(), + [](const py::handle &quantizer) -> bool { + return detail::IsNVFP4Quantizers(quantizer.ptr()); + })) { + allocation_method = AllocationMethod::BULK_NVFP4; + quantization_method = QuantizationMethod::FUSED_NVFP4; + } } // Allocate output tensors diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index e73eca7861..c5c8905294 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -248,7 +248,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Multi-tensor quantize", py::arg("tensor_list"), py::arg("quantizer_list")); 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("quantizer_list"), py::arg("disable_bulk_allocation") = false); m.def("te_general_grouped_gemm", &transformer_engine::pytorch::te_general_grouped_gemm, "Grouped GEMM"); m.def("fp8_transpose", &transformer_engine::pytorch::fp8_transpose, "Transpose with FP8 I/O", diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index d0a5618afb..1e6f0b00ab 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -143,7 +143,12 @@ def forward( inp_view = inp.reshape(-1, in_features) inputmats: list if fp8 and not debug: - inputmats = tex.split_quantize(inp_view, m_splits, input_quantizers) + # Disable bulk allocation when CPU offloading is active: offloading skips small + # tensors (like scales), but bulk allocation shares storage across all tensors, + # so if scales can't be offloaded, nothing in the group can be offloaded. + inputmats = tex.split_quantize( + inp_view, m_splits, input_quantizers, disable_bulk_allocation=cpu_offloading + ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( inp_view, input_quantizers, m_splits, activation_dtype diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index f3220d5860..b8349f84a0 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -428,7 +428,8 @@ def forward( # weights if weights are externally touched outside this module ctx.weight_object = weight - mark_not_offload(weight, weightmat, bias) + if cpu_offloading: + mark_not_offload(weight, weightmat, bias) # TODO(ksivamani): Check memory usage tensors_to_save, tensor_objects = prepare_for_saving( saved_inputmat, diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 1995655c33..495056d652 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -14,6 +14,7 @@ from torch.distributed._tensor import DTensor import transformer_engine_torch as tex from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor, Float8Quantizer +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from .multi_tensor_apply import multi_tensor_applier @@ -372,10 +373,12 @@ def _initialize_state( store_param_remainders (bool): Store only trailing remainder bits. """ dtype = self.name_to_dtype_map[state_name] + # Handle QuantizedTensor by dequantizing first + param_for_empty = param.dequantize() if isinstance(param, QuantizedTensor) else param if store_param_remainders: - data = torch.zeros(param.shape, dtype=torch.int16, device=param.device) + data = torch.zeros_like(param_for_empty, dtype=torch.int16) else: - data = torch.empty(param.shape, dtype=dtype, device=param.device) + data = torch.empty_like(param_for_empty, dtype=dtype) if zero_buffer: data.zero_() diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 3414581f7c..ac827e794a 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -20,11 +20,6 @@ _stride_from_shape, ) -_quantized_tensor_cpu_supported_ops = ( - torch.ops.aten.empty_like.default, - torch.ops.aten.copy_.default, -) - class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -539,15 +534,6 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): if kwargs is None: kwargs = {} - def check_if_cpu(arg): - if isinstance(cls, QuantizedTensor) and arg.device.type == "cpu": - assert ( - func in _quantized_tensor_cpu_supported_ops - ), f"QuantizedTensor on CPU does not support this operation: {func}" - return arg - - args = tree_map(check_if_cpu, args) - # Do not force the QuantizedTensor type on the returned tensor return torch._C._disabled_torch_function_impl(func, types, args, kwargs) From 69636a08171d162a223ae3a01e4b36902b64dba1 Mon Sep 17 00:00:00 2001 From: Victor Oliveira Date: Tue, 13 Jan 2026 11:34:03 -0800 Subject: [PATCH 161/521] ONNX: Fix FP8 quantization for the second MLP in LayerNormMLP (#2577) ONNX: Fix FP8 quantization for the second MLP in LayernormMLP Signed-off-by: Victor Oliveira --- .../pytorch/module/layernorm_mlp.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ddb33f303c..4256028c8b 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2243,14 +2243,23 @@ def onnx_forward( assert not TEDebugState.debug_enabled, "Debug mode is not supported in ONNX export" assert_warmed_up(self) + + # Get quantizers ( fc1_input_quantizer, fc1_weight_quantizer, + _, + _, + _, + _, fc2_input_quantizer, fc2_weight_quantizer, - output_quantizer, - *_, + fc2_output_quantizer, + _, + _, + _, ) = self._get_quantizers(False, is_grad_enabled) + inp_dtype = inp.dtype fc1_weight, fc2_weight = self._get_weight_tensors() @@ -2324,7 +2333,7 @@ def _clamped_swiglu(x, limit, alpha): fc2_out = onnx_gemm(fc2_weight, act_out, fc2_bias) - if output_quantizer is not None: + if fc2_output_quantizer is not None: raise NotImplementedError("ONNX export of quantized output is not supported") if self.return_layernorm_output: From bd007993dd7403253c6fd9084ebc409f9ea7adad Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 14 Jan 2026 15:58:07 -0800 Subject: [PATCH 162/521] Revert adding pytorch-triton as a build requirement (#2592) * Remove pyhtorch-triton as a requirement and remove auto-fetching pytorch-triton as it is a placeeholder in pyPI Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix docstring Signed-off-by: tdophung --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- build_tools/pytorch.py | 23 ++--------------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index 98511e45cb..fdfdee9b1c 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -13,27 +13,8 @@ def install_requirements() -> List[str]: - """Install dependencies for TE/PyTorch extensions. - - IMPORTANT - PyTorch Index Required for pytorch-triton: - These dependencies MUST be installed using PyTorch's package index: - - pip install pytorch-triton --index-url https://download.pytorch.org/whl/ - - - pytorch-triton is only available from PyTorch's index (not PyPI) - - The 'pytorch-triton' package on PyPI is a placeholder that will fail - - torch.compile() requires pytorch-triton, not OpenAI's 'triton' package - """ - return [ - "torch>=2.1", - "einops", - "onnxscript", - "onnx", - "packaging", - "pydantic", - "nvdlfw-inspect", - "pytorch-triton", - ] + """Install dependencies for TE/PyTorch extensions.""" + return ["torch>=2.1", "einops", "onnxscript", "onnx", "packaging", "pydantic", "nvdlfw-inspect"] def test_requirements() -> List[str]: From fcfa0c3c8105aa1225ecad61b38e55b1d583fb86 Mon Sep 17 00:00:00 2001 From: Hongbin Liu Date: Thu, 15 Jan 2026 09:29:48 +0800 Subject: [PATCH 163/521] (Bug fix) Fix accuracy issue for blockwise scaling+E8 scale on Blackwell (#2589) * bug fix Signed-off-by: hongbinl * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/common/swizzle/swizzle_block_scaling.cu Mask to 8 bits to prevent potential bit overlap Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Hongbin Liu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/common/swizzle/swizzle_block_scaling.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Hongbin Liu * fix bug in 2d too Signed-off-by: Hongbin Liu --------- Signed-off-by: hongbinl Signed-off-by: Hongbin Liu Signed-off-by: Hongbin Liu 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> --- .../common/swizzle/swizzle_block_scaling.cu | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/swizzle/swizzle_block_scaling.cu b/transformer_engine/common/swizzle/swizzle_block_scaling.cu index c5ad1aed43..37993787a5 100644 --- a/transformer_engine/common/swizzle/swizzle_block_scaling.cu +++ b/transformer_engine/common/swizzle/swizzle_block_scaling.cu @@ -113,7 +113,8 @@ void __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE) } // pack the exponent bits of the scaling factors - uint32_t packed_exponents = (sf.x >> 23) | (sf.y >> 15) | (sf.z >> 7) | (sf.w << 1); + uint32_t packed_exponents = ((sf.x >> 23) & 0xFF) | (((sf.y >> 23) & 0xFF) << 8) | + (((sf.z >> 23) & 0xFF) << 16) | (((sf.w >> 23) & 0xFF) << 24); // partially swizzle the scaling factors constexpr uint32_t ACTIVE_MASK = 0xFFFFFFFF; // no divergent branches @@ -198,8 +199,9 @@ void __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE) uint32_t sf = *reinterpret_cast(warp_src); // broadcast it to four scaling factors for 1x32 tiles - sf = (sf << 1) | (sf >> 7); - sf = sf | (sf >> 16); + // extract and broadcast the exponent byte to four bytes for E8M0 format + uint32_t exp_byte = (sf >> 23) & 0xFF; + sf = exp_byte | (exp_byte << 8) | (exp_byte << 16) | (exp_byte << 24); // broadcast it to sixteen scaling factors for 1x32 tiles const uint4 sf4{sf, sf, sf, sf}; From 4df43dbe8422599edd3c21806a04b356bfe8227c Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Wed, 14 Jan 2026 19:52:11 -0600 Subject: [PATCH 164/521] docs: Update README Latest News section (#2583) * Move older news to Previous Signed-off-by: Santosh Bhavani * Add Nov 2025 news entries Signed-off-by: Santosh Bhavani --------- Signed-off-by: Santosh Bhavani --- README.rst | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/README.rst b/README.rst index 9241e26cdd..55be0e583f 100644 --- a/README.rst +++ b/README.rst @@ -13,23 +13,14 @@ Transformer Engine Latest News =========== +* [11/2025] `NVIDIA Blackwell Architecture Sweeps MLPerf Training v5.1 Benchmarks `_ +* [11/2025] `Scale Biology Transformer Models with PyTorch and NVIDIA BioNeMo Recipes `_ +* [11/2025] `FP8 Training of Large-Scale RL Models `_ * [09/2025] `Pretraining Large Language Models with NVFP4 `_ * [09/2025] `Native FP8 Mixed Precision Training for Ling 2.0, Open Sourced! `_ * [09/2025] `Faster Training Throughput in FP8 Precision with NVIDIA NeMo `_ * [08/2025] `How we built DeepL's next-generation LLMs with FP8 for training and inference `_ * [08/2025] `NVFP4 Trains with Precision of 16-bit and Speed and Efficiency of 4-bit `_ -* [06/2025] `Floating Point 8: An Introduction to Efficient, Lower-Precision AI Training `_ -* [05/2025] `Advanced Optimization Strategies for LLM Training on NVIDIA Grace Hopper `_ -* [03/2025] `Stable and Scalable FP8 Deep Learning Training on Blackwell | GTC 2025 `_ -* [03/2025] `Measure and Improve AI Workload Performance with NVIDIA DGX Cloud Benchmarking `_ - -.. image:: docs/examples/comparison-fp8-bf16-training-nvidia-dgx-cloud-benchmarking-performance-explorer.jpg - :width: 600 - :alt: Comparison of FP8 versus BF16 training, as seen in NVIDIA DGX Cloud Benchmarking Performance Explorer - -* [02/2025] `Understanding the Language of Life's Biomolecules Across Evolution at a New Scale with Evo 2 `_ -* [02/2025] `NVIDIA DGX Cloud Introduces Ready-To-Use Templates to Benchmark AI Platform Performance `_ -* [01/2025] `Continued Pretraining of State-of-the-Art LLMs for Sovereign AI and Regulated Industries with iGenius and NVIDIA DGX Cloud `_ `Previous News <#previous-news>`_ @@ -425,6 +416,18 @@ Videos Previous News ============= +* [06/2025] `Floating Point 8: An Introduction to Efficient, Lower-Precision AI Training `_ +* [05/2025] `Advanced Optimization Strategies for LLM Training on NVIDIA Grace Hopper `_ +* [03/2025] `Stable and Scalable FP8 Deep Learning Training on Blackwell | GTC 2025 `_ +* [03/2025] `Measure and Improve AI Workload Performance with NVIDIA DGX Cloud Benchmarking `_ + +.. image:: docs/examples/comparison-fp8-bf16-training-nvidia-dgx-cloud-benchmarking-performance-explorer.jpg + :width: 600 + :alt: Comparison of FP8 versus BF16 training, as seen in NVIDIA DGX Cloud Benchmarking Performance Explorer + +* [02/2025] `Understanding the Language of Life's Biomolecules Across Evolution at a New Scale with Evo 2 `_ +* [02/2025] `NVIDIA DGX Cloud Introduces Ready-To-Use Templates to Benchmark AI Platform Performance `_ +* [01/2025] `Continued Pretraining of State-of-the-Art LLMs for Sovereign AI and Regulated Industries with iGenius and NVIDIA DGX Cloud `_ * [11/2024] `Developing a 172B LLM with Strong Japanese Capabilities Using NVIDIA Megatron-LM `_ * [11/2024] `How FP8 boosts LLM training by 18% on Amazon SageMaker P5 instances `_ * [11/2024] `Efficiently train models with large sequence lengths using Amazon SageMaker model parallel `_ From 2236292a4cf2e6f276ba8444bac1125a3313f394 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 15 Jan 2026 11:36:20 -0800 Subject: [PATCH 165/521] [JAX] Disable fused attention in encoder tests for determinism (#2601) disable fused attention in encoder tests for determinism Signed-off-by: Jeremy Berchtold --- examples/jax/encoder/test_model_parallel_encoder.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index f29cc4e0be..02937bc394 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Encoder training on multi-GPU with tesnor parallelism""" import argparse +import os import unittest from functools import partial @@ -489,6 +490,9 @@ class TestEncoder(unittest.TestCase): def setUp(self): """Run 5 epochs for testing""" + # TODO(jberchtold): Remove once fused attention from cuDNN supports determinism on Blackwell + if "NVTE_FUSED_ATTN" not in os.environ: + os.environ["NVTE_FUSED_ATTN"] = "0" self.args = encoder_parser(["--epochs", "5"]) @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") From 6cbdb0423b8b0ce8895a2dfc833c78e4e84c4aca Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 15 Jan 2026 14:18:16 -0800 Subject: [PATCH 166/521] [JAX] Install Cmake in TE/JAX build Github Action (#2603) * install cmake in jax build github action Signed-off-by: Jeremy Berchtold * Update build.yml Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --------- Signed-off-by: Jeremy Berchtold Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- .github/workflows/build.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 427b2f27fa..d80564274e 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -92,7 +92,7 @@ jobs: options: --user root steps: - name: 'Dependencies' - run: pip install pybind11[global] + run: pip install cmake==3.21.0 pybind11[global] - name: 'Checkout' uses: actions/checkout@v3 with: @@ -144,7 +144,7 @@ jobs: - name: 'Dependencies' run: | docker exec builder bash -c '\ - pip install pybind11[global] einops onnxscript && \ + pip install cmake==3.21.0 pybind11[global] einops onnxscript && \ pip install torch --no-cache-dir --index-url https://download.pytorch.org/whl/cu130 ' - name: 'Build' From 6a34b6574fa6c29d9d07fdcddf9812cbb1488878 Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Thu, 15 Jan 2026 14:41:55 -0800 Subject: [PATCH 167/521] fix: enable opt for cutlass sources to avoid infinite compile time (#2595) Signed-off-by: Kaining Zhong --- transformer_engine/common/CMakeLists.txt | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 5975efedaf..a83cbe3e30 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -230,12 +230,24 @@ add_library(transformer_engine SHARED ${transformer_engine_SOURCES}) target_include_directories(transformer_engine PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/include") -# CUTLASS kernels require SM90a and cause hang in debug build +# Grouped GEMM kernels require SM90a set_property( SOURCE gemm/cutlass_grouped_gemm.cu APPEND PROPERTY - COMPILE_OPTIONS "--generate-code=arch=compute_90a,code=sm_90a;-g0") + COMPILE_OPTIONS "--generate-code=arch=compute_90a,code=sm_90a") + +# CUTLASS kernels could cause hang in debug build +set(CUTLASS_KERNEL_SOURCES + gemm/cutlass_grouped_gemm.cu + hadamard_transform/group_hadamard_transform_cast_fusion.cu + hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu + hadamard_transform/hadamard_transform_cast_fusion.cu) +set_property( + SOURCE ${CUTLASS_KERNEL_SOURCES} + APPEND + PROPERTY + COMPILE_OPTIONS "-g0;-dopt=on") # Configure dependencies target_link_libraries(transformer_engine PUBLIC From a652730fd38b1415ba048cde44e84f87589209b5 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 16 Jan 2026 15:24:06 -0800 Subject: [PATCH 168/521] [JAX] Custom partitioning for Permutation primitives (#2591) * initial impl, not tested Signed-off-by: tdophung * consolidate different unpermute primitives with with_pad and with_merging_probs booleans. Implement partitioning for all permutation primitives Signed-off-by: tdophung * Add distributed test for non-padding permutation Signed-off-by: tdophung * fix issues in distributed test for padding permutation. Make common kernel zero intiialize output permuted scales, permuted probs and output tokens Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert zeroing in triton common kernel as it is a race condition. Instead, add extra input (aliased wiuth output) buffer to inner primitive of permutation on jax side to pass in zero intitiated buffers done with jnp zeros Signed-off-by: tdophung * fix utils to handle input output aliasing in autotuned kernels Signed-off-by: tdophung * Clean up comments, and add more comments explaining input output alias in utils Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint and greptile comment Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix issues that lint fixing introduced Signed-off-by: tdophung --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_distributed_permutation.py | 597 +++++++++ .../common/triton/permutation.py | 16 + transformer_engine/jax/permutation.py | 19 +- .../jax/triton_extensions/permutation.py | 1115 +++++++++++++---- .../jax/triton_extensions/utils.py | 40 +- .../pytorch/triton/permutation.py | 14 +- 6 files changed, 1515 insertions(+), 286 deletions(-) create mode 100644 tests/jax/test_distributed_permutation.py diff --git a/tests/jax/test_distributed_permutation.py b/tests/jax/test_distributed_permutation.py new file mode 100644 index 0000000000..5b6d8fec47 --- /dev/null +++ b/tests/jax/test_distributed_permutation.py @@ -0,0 +1,597 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for distributed/sharded execution of MoE permutation primitives. + +Testing Strategy: +================= +MoE permutation is data-dependent - the destination index for each token depends +on how many tokens before it are routed to the same expert. This means: + +1. We CANNOT compare sharded output against global reference directly +2. Instead, we verify that each GPU's LOCAL output is correct according to its + LOCAL routing (which produces LOCAL row_id_map with LOCAL indices) + +For data-parallel MoE without expert parallelism: +- Each GPU has ALL experts replicated +- Each GPU processes a subset of tokens (sharded on token/batch dimension) +- Each GPU computes its own local row_id_map from its local routing_map slice +- Each GPU's output is local and doesn't need to match global output + +These tests verify: +1. Local token_dispatch: sharded input -> local row_id_map -> local permute (forward + backward) +2. Local roundtrip: dispatch + combine recovers original input (forward + backward) +""" + +import pytest + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding, PartitionSpec + +from distributed_test_base import generate_configs +from utils import assert_allclose, pytest_parametrize_wrapper + +# High-level API with VJP support +from transformer_engine.jax.permutation import ( + token_dispatch, + token_combine, +) + +# Reference implementations from test_permutation.py +from test_permutation import ( + reference_make_row_id_map, + _reference_permute_impl, + _reference_unpermute_impl, + reference_token_combine, +) + +# Dispatch/combine test cases: (num_tokens, num_experts, hidden_size, topk) +# topk = number of experts each token is routed to +# Includes small, medium-large, and largest stress test cases. +ALL_DISPATCH_COMBINE_CASES = [ + (128, 4, 64, 2), + (4096, 32, 1280, 2), + (4096, 256, 4096, 6), +] +DISPATCH_COMBINE_CASES = { + "L0": ALL_DISPATCH_COMBINE_CASES[0:1], + "L2": ALL_DISPATCH_COMBINE_CASES, +} + +# Dispatch/combine with padding test cases: (num_tokens, num_experts, hidden_size, topk, align_size) +ALL_DISPATCH_COMBINE_PADDING_CASES = [ + (128, 4, 64, 2, 8), + (4096, 32, 1280, 2, 128), + (4096, 256, 4096, 6, 16), +] +DISPATCH_COMBINE_PADDING_CASES = { + "L0": ALL_DISPATCH_COMBINE_PADDING_CASES[0:1], + "L2": ALL_DISPATCH_COMBINE_PADDING_CASES, +} + +# Dtypes for testing +ALL_DTYPES = [jnp.float32, jnp.bfloat16] +DTYPES = { + "L0": [jnp.float32], + "L2": ALL_DTYPES, +} + + +class TestDistributedPermutation: + """Test distributed/sharded execution of MoE permutation primitives. + + These tests validate that custom partitioning produces correct LOCAL results + when inputs are sharded across multiple devices. + + Key insight: With data-parallel MoE, each GPU independently processes its + local tokens. The row_id_map is generated locally and contains LOCAL indices. + We verify correctness by comparing each shard's output against the reference + implementation run on that shard's local data. + """ + + @staticmethod + def compute_padded_output_size( + num_tokens: int, + num_experts: int, + topk: int, + align_size: int, + num_dp_devices: int, + ) -> int: + """Compute global_num_out_tokens for distributed padding tests. + + Each device processes local_num_tokens tokens. We compute the worst-case + padded output size per device, then multiply by num_dp_devices to get + a global size that ensures global / num_dp >= local_worst. + """ + local_num_tokens = num_tokens // num_dp_devices + local_raw_out = local_num_tokens * topk + local_worst = ((local_raw_out + num_experts * (align_size - 1)) // align_size) * align_size + return local_worst * num_dp_devices + + @staticmethod + def generate_routing_map( + num_tokens: int, + num_experts: int, + topk: int = 2, # Number of experts each token is routed to (max 1s per row). + key: jax.Array = None, + ): + if key is None: + key = jax.random.PRNGKey(0) + + routing_map = jnp.zeros((num_tokens, num_experts), dtype=jnp.int32) + for token_idx in range(num_tokens): + key, subkey = jax.random.split(key) + expert_indices = jax.random.choice(subkey, num_experts, shape=(topk,), replace=False) + routing_map = routing_map.at[token_idx, expert_indices].set(1) + + return routing_map + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,hidden_size,topk", + DISPATCH_COMBINE_CASES, + ) + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("use_shardy", [False, True]) + def test_local_token_dispatch( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + hidden_size, + topk, + dtype, + use_shardy, + ): + """ + Test token_dispatch with sharded inputs. + + Verifies that sharded execution produces the same result as chunk-wise + reference execution. The sharded primitive: + 1. Receives global num_out_tokens (partition function divides it) + 2. Each GPU operates on its local shard independently + 3. Results are gathered (concatenated) across GPUs + + Output ordering: [GPU0_expert0, GPU0_expert1, ... | GPU1_expert0, ...] + + The reference processes each chunk independently and concatenates, + matching the sharded execution's output ordering. + Tests both forward pass (output values) and backward pass (gradients). + """ + jax.config.update("jax_use_shardy_partitioner", use_shardy) + key = jax.random.PRNGKey(42) + + # Generate global inputs + key, inp_key, prob_key = jax.random.split(key, 3) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + routing_map = self.generate_routing_map(num_tokens, num_experts, topk, key) + probs = jax.random.uniform( + prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 + ) + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + # Shard on token (batch) dimension + dp_axis = mesh_resource.dp_resource + sharded_pspec = PartitionSpec(dp_axis, None) + + # Compute num_out_tokens as concrete values + # Global num_out_tokens is passed to token_dispatch (partition function divides it) + # Local num_out_tokens is used for reference implementation + num_dp_devices = mesh.shape[dp_axis] if dp_axis else 1 + global_num_out_tokens = num_tokens * topk + local_num_tokens = num_tokens // num_dp_devices + local_num_out_tokens = local_num_tokens * topk + + with mesh: + inp_sharding = NamedSharding(mesh, sharded_pspec) + routing_sharding = NamedSharding(mesh, sharded_pspec) + probs_sharding = NamedSharding(mesh, sharded_pspec) + + # Shard the inputs + inp_sharded = jax.device_put(inp, inp_sharding) + routing_sharded = jax.device_put(routing_map, routing_sharding) + probs_sharded = jax.device_put(probs, probs_sharding) + + # ================================================================ + # Forward pass test + # ================================================================ + @jax.jit + def target_dispatch(x, rm, p): + # Pass global num_out_tokens - partition function divides it + out, perm_probs, rid_map, _, _ = token_dispatch( + x, rm, global_num_out_tokens, probs=p + ) + return out, perm_probs, rid_map + + # Reference: process each GPU's shard independently, then concatenate + # This matches how the sharded primitive operates: + # - Each GPU processes its local shard + # - Results are gathered (concatenated) across GPUs + # Output ordering: [GPU0_exp0, GPU0_exp1, ... | GPU1_exp0, GPU1_exp1, ...] + inp_shards = jnp.reshape(inp, (num_dp_devices, local_num_tokens, hidden_size)) + routing_shards = jnp.reshape( + routing_map, (num_dp_devices, local_num_tokens, num_experts) + ) + probs_shards = jnp.reshape(probs, (num_dp_devices, local_num_tokens, num_experts)) + + ref_outputs = [] + ref_perm_probs_list = [] + ref_rid_maps = [] + for i in range(num_dp_devices): + shard_rid_map = reference_make_row_id_map(routing_shards[i]) + shard_out, shard_perm_probs = _reference_permute_impl( + inp_shards[i], shard_rid_map, probs_shards[i], local_num_out_tokens + ) + ref_outputs.append(shard_out) + ref_perm_probs_list.append(shard_perm_probs) + ref_rid_maps.append(shard_rid_map) + + # Concatenate like all_gather would + ref_out = jnp.concatenate(ref_outputs, axis=0) + ref_perm_probs = jnp.concatenate(ref_perm_probs_list, axis=0) + ref_rid_map = jnp.concatenate(ref_rid_maps, axis=0) + + # Run target on sharded inputs + target_out, target_perm_probs, target_rid_map = target_dispatch( + inp_sharded, routing_sharded, probs_sharded + ) + + # Compare forward outputs + assert_allclose(jax.device_get(target_out), ref_out, dtype=dtype) + assert_allclose(jax.device_get(target_perm_probs), ref_perm_probs, dtype=dtype) + + # Verify row_id_map n_routed column matches routing_map sum + target_rid_map_np = jax.device_get(target_rid_map) + assert jnp.array_equal( + target_rid_map_np[:, -1], ref_rid_map[:, -1] + ), "n_routed column mismatch" + + # Sanity checks + target_out_np = jax.device_get(target_out) + target_perm_probs_np = jax.device_get(target_perm_probs) + assert not np.any(np.isnan(target_out_np)), "Output contains NaN" + assert not np.any(np.isnan(target_perm_probs_np)), "Permuted probs contain NaN" + assert np.all(target_perm_probs_np >= 0), "Permuted probs contain negative values" + + # ================================================================ + # Backward pass test (gradients) + # ================================================================ + def target_loss(x, rm, p): + out, perm_probs, _, _, _ = token_dispatch(x, rm, global_num_out_tokens, probs=p) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + # Reference loss: process chunks independently and sum + def ref_chunk_loss(inp_chunk, routing_chunk, probs_chunk): + rid_map = reference_make_row_id_map(routing_chunk) + out, perm_probs = _reference_permute_impl( + inp_chunk, rid_map, probs_chunk, local_num_out_tokens + ) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + target_grad_fn = jax.jit(jax.grad(target_loss, argnums=(0, 2))) + ref_chunk_grad_fn = jax.jit(jax.grad(ref_chunk_loss, argnums=(0, 2))) + + target_inp_grad, target_probs_grad = target_grad_fn( + inp_sharded, routing_sharded, probs_sharded + ) + + # Compute reference gradients per chunk, then concatenate + ref_inp_grads = [] + ref_probs_grads = [] + for i in range(num_dp_devices): + chunk_inp_grad, chunk_probs_grad = ref_chunk_grad_fn( + inp_shards[i], routing_shards[i], probs_shards[i] + ) + ref_inp_grads.append(chunk_inp_grad) + ref_probs_grads.append(chunk_probs_grad) + + ref_inp_grad = jnp.concatenate(ref_inp_grads, axis=0) + ref_probs_grad = jnp.concatenate(ref_probs_grads, axis=0) + + assert_allclose(jax.device_get(target_inp_grad), ref_inp_grad, dtype=dtype) + assert_allclose(jax.device_get(target_probs_grad), ref_probs_grad, dtype=dtype) + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,hidden_size,topk", + DISPATCH_COMBINE_CASES, + ) + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("use_shardy", [False, True]) + def test_local_roundtrip( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + hidden_size, + topk, + dtype, + use_shardy, + ): + """ + Test roundtrip: token_dispatch followed by token_combine with sharded inputs. + + Each GPU: + 1. Gets a shard of the input and routing_map + 2. Performs local dispatch (permute) + 3. Performs local combine (unpermute) + 4. With uniform merging probs, should recover original input + + Tests both forward pass and backward pass (gradient should be 2*x). + """ + jax.config.update("jax_use_shardy_partitioner", use_shardy) + key = jax.random.PRNGKey(42) + + # Generate global inputs + key, inp_key = jax.random.split(key, 2) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + routing_map = self.generate_routing_map(num_tokens, num_experts, topk, key) + + # Uniform merging probs for perfect roundtrip + uniform_merging_probs = routing_map.astype(dtype) / jnp.maximum( + jnp.sum(routing_map, axis=1, keepdims=True), 1.0 + ) + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + dp_axis = mesh_resource.dp_resource + sharded_pspec = PartitionSpec(dp_axis, None) + + # Compute num_out_tokens as concrete value + # Global num_out_tokens is passed to token_dispatch (partition function divides it) + global_num_out_tokens = num_tokens * topk + + with mesh: + inp_sharding = NamedSharding(mesh, sharded_pspec) + routing_sharding = NamedSharding(mesh, sharded_pspec) + merging_sharding = NamedSharding(mesh, sharded_pspec) + + inp_sharded = jax.device_put(inp, inp_sharding) + routing_sharded = jax.device_put(routing_map, routing_sharding) + merging_sharded = jax.device_put(uniform_merging_probs, merging_sharding) + + # ================================================================ + # Forward pass test + # ================================================================ + @jax.jit + def roundtrip(x, rm, mprobs): + dispatched, _, rid_map, _, _ = token_dispatch(x, rm, global_num_out_tokens) + return token_combine(dispatched, rid_map, mprobs) + + roundtrip_out = roundtrip(inp_sharded, routing_sharded, merging_sharded) + + # Should recover original input + assert_allclose(jax.device_get(roundtrip_out), jax.device_get(inp_sharded), dtype=dtype) + + # ================================================================ + # Backward pass test (gradients) + # ================================================================ + def roundtrip_loss(x, rm, mprobs): + dispatched, _, rid_map, _, _ = token_dispatch(x, rm, global_num_out_tokens) + combined = token_combine(dispatched, rid_map, mprobs) + return jnp.sum(combined**2) + + # With uniform merging probs, roundtrip is identity, so gradient should be 2*x + grad_fn = jax.jit(jax.grad(roundtrip_loss, argnums=0)) + computed_grad = grad_fn(inp_sharded, routing_sharded, merging_sharded) + + expected_grad = 2.0 * inp_sharded + + assert_allclose( + jax.device_get(computed_grad), jax.device_get(expected_grad), dtype=dtype + ) + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,hidden_size,topk,align_size", + DISPATCH_COMBINE_PADDING_CASES, + ) + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("use_shardy", [False, True]) + def test_local_token_dispatch_with_padding( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + hidden_size, + topk, + align_size, + dtype, + use_shardy, + ): + """ + Test token_dispatch with padding using sharded inputs. + + Tests both forward pass (output values) and backward pass (gradients). + """ + jax.config.update("jax_use_shardy_partitioner", use_shardy) + key = jax.random.PRNGKey(42) + + # Generate global inputs + key, inp_key, prob_key = jax.random.split(key, 3) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + routing_map = self.generate_routing_map(num_tokens, num_experts, topk, key) + probs = jax.random.uniform( + prob_key, (num_tokens, num_experts), dtype=dtype, minval=0.1, maxval=1.0 + ) + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + dp_axis = mesh_resource.dp_resource + sharded_pspec = PartitionSpec(dp_axis, None) + num_dp_devices = mesh.shape[dp_axis] if dp_axis else 1 + + # For padding + sharding, we need to account for per-shard padding overhead. + # Each shard needs E*(A-1) extra space for worst-case padding. + # Compute global_num_out_tokens such that global / num_dp >= local_worst. + global_num_out_tokens = self.compute_padded_output_size( + num_tokens, num_experts, topk, align_size, num_dp_devices + ) + + with mesh: + inp_sharding = NamedSharding(mesh, sharded_pspec) + routing_sharding = NamedSharding(mesh, sharded_pspec) + probs_sharding = NamedSharding(mesh, sharded_pspec) + + inp_sharded = jax.device_put(inp, inp_sharding) + routing_sharded = jax.device_put(routing_map, routing_sharding) + probs_sharded = jax.device_put(probs, probs_sharding) + + # ================================================================ + # Forward pass test + # ================================================================ + @jax.jit + def dispatch_with_padding(x, rm, p): + out, perm_probs, rid_map, pad_offsets, _ = token_dispatch( + x, rm, global_num_out_tokens, probs=p, align_size=align_size + ) + return out, perm_probs, rid_map, pad_offsets + + out, perm_probs, rid_map, pad_offsets = dispatch_with_padding( + inp_sharded, routing_sharded, probs_sharded + ) + + # Sanity checks + out_np = jax.device_get(out) + perm_probs_np = jax.device_get(perm_probs) + assert not np.any(np.isnan(out_np)), "Output contains NaN" + assert not np.any(np.isnan(perm_probs_np)), "Permuted probs contain NaN" + assert np.all(perm_probs_np >= 0), "Permuted probs contain negative values" + + # ================================================================ + # Backward pass test (gradients) + # ================================================================ + def loss_with_padding(x, rm, p): + out, perm_probs, _, _, _ = token_dispatch( + x, rm, global_num_out_tokens, probs=p, align_size=align_size + ) + return jnp.sum(out**2) + jnp.sum(perm_probs**2) + + grad_fn = jax.jit(jax.grad(loss_with_padding, argnums=(0, 2))) + inp_grad, probs_grad = grad_fn(inp_sharded, routing_sharded, probs_sharded) + + # Gradients should not contain NaN + assert not np.any(np.isnan(jax.device_get(inp_grad))), "Input gradient contains NaN" + assert not np.any(np.isnan(jax.device_get(probs_grad))), "Probs gradient contains NaN" + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,hidden_size,topk,align_size", + DISPATCH_COMBINE_PADDING_CASES, + ) + @pytest_parametrize_wrapper("dtype", DTYPES) + @pytest_parametrize_wrapper("use_shardy", [False, True]) + def test_local_roundtrip_with_padding( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + hidden_size, + topk, + align_size, + dtype, + use_shardy, + ): + """ + Test roundtrip with padding/alignment using sharded inputs. + + With uniform merging probs, should recover original input. + Tests both forward pass and backward pass. + """ + jax.config.update("jax_use_shardy_partitioner", use_shardy) + key = jax.random.PRNGKey(42) + + # Generate inputs + key, inp_key = jax.random.split(key, 2) + inp = jax.random.uniform( + inp_key, (num_tokens, hidden_size), dtype=dtype, minval=-1.0, maxval=1.0 + ) + routing_map = self.generate_routing_map(num_tokens, num_experts, topk, key) + + # Uniform merging probs + uniform_merging_probs = routing_map.astype(dtype) / jnp.maximum( + jnp.sum(routing_map, axis=1, keepdims=True), 1.0 + ) + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + dp_axis = mesh_resource.dp_resource + sharded_pspec = PartitionSpec(dp_axis, None) + num_dp_devices = mesh.shape[dp_axis] if dp_axis else 1 + + # For padding + sharding, we need to account for per-shard padding overhead. + # Each shard needs E*(A-1) extra space for worst-case padding. + # Compute global_num_out_tokens such that global / num_dp >= local_worst. + global_num_out_tokens = self.compute_padded_output_size( + num_tokens, num_experts, topk, align_size, num_dp_devices + ) + + with mesh: + inp_sharding = NamedSharding(mesh, sharded_pspec) + routing_sharding = NamedSharding(mesh, sharded_pspec) + merging_sharding = NamedSharding(mesh, sharded_pspec) + + inp_sharded = jax.device_put(inp, inp_sharding) + routing_sharded = jax.device_put(routing_map, routing_sharding) + merging_sharded = jax.device_put(uniform_merging_probs, merging_sharding) + + # ================================================================ + # Forward pass test + # ================================================================ + @jax.jit + def roundtrip_with_padding(x, rm, mprobs): + dispatched, _, rid_map, pad_offsets, _ = token_dispatch( + x, rm, global_num_out_tokens, align_size=align_size + ) + return token_combine(dispatched, rid_map, mprobs, pad_offsets) + + roundtrip_out = roundtrip_with_padding(inp_sharded, routing_sharded, merging_sharded) + + # Should recover original input + assert_allclose(jax.device_get(roundtrip_out), jax.device_get(inp_sharded), dtype=dtype) + + # ================================================================ + # Backward pass test (gradients) + # ================================================================ + def roundtrip_loss_with_padding(x, rm, mprobs): + dispatched, _, rid_map, pad_offsets, _ = token_dispatch( + x, rm, global_num_out_tokens, align_size=align_size + ) + combined = token_combine(dispatched, rid_map, mprobs, pad_offsets) + return jnp.sum(combined**2) + + # With uniform merging probs, roundtrip is identity, so gradient should be 2*x + grad_fn = jax.jit(jax.grad(roundtrip_loss_with_padding, argnums=0)) + computed_grad = grad_fn(inp_sharded, routing_sharded, merging_sharded) + + expected_grad = 2.0 * inp_sharded + + assert_allclose( + jax.device_get(computed_grad), jax.device_get(expected_grad), dtype=dtype + ) diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index e53b2a9455..4602f41cfd 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -201,8 +201,15 @@ def _permute_kernel( scale_ptr, permuted_scale_ptr, pad_offsets_ptr, + # Pre-allocated output buffers for JAX input_output_aliases. + # These are aliased to output_ptr/permuted_probs_ptr in JAX, so they point to the same memory. + # In PyTorch, pass the same tensors as output_ptr/permuted_probs_ptr. + output_buf_ptr, # pylint: disable=unused-argument + permuted_probs_buf_ptr, # pylint: disable=unused-argument # sizes scale_hidden_dim, + num_tokens, # pylint: disable=unused-argument + num_out_tokens, # pylint: disable=unused-argument # strides stride_row_id_map_token, stride_row_id_map_expert, @@ -228,12 +235,17 @@ def _permute_kernel( FUSION_PAD: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): + # Note: When FUSION_PAD=True, output buffers should be pre-zeroed by the caller + # to ensure padding positions contain zeros. + # PyTorch: Use torch.zeros() for output buffer allocation + # JAX: Pre-zeroed buffers should be passed (when input_output_aliases works) expert_idx = 0 pid_t = tl.program_id(0) pid_h = tl.program_id(1) cur_off = pid_h * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = cur_off < hidden_size + src_row = pid_t.to(tl.int64) input_off = src_row * stride_input_token + cur_off * stride_input_hidden inp = tl.load(input_ptr + input_off, mask=mask) @@ -306,6 +318,10 @@ def _unpermute_kernel( merging_probs_ptr, permuted_probs_ptr, pad_offsets_ptr, + # Dummy parameters for JAX input_output_aliases compatibility (matches _permute_kernel signature pattern) + # These are unused in the unpermute kernel but maintain consistency with the permute kernel. + output_buf_ptr, # pylint: disable=unused-argument + unpermuted_probs_buf_ptr, # pylint: disable=unused-argument # strides stride_row_id_map_token, stride_row_id_map_expert, diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 2e16e674cc..405d5f7661 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -137,7 +137,7 @@ def token_dispatch( ) -@partial(jax.custom_vjp, nondiff_argnums=(1, 3, 4, 5, 6)) +@partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5, 6)) def _token_dispatch( inp: jnp.ndarray, routing_map: jnp.ndarray, @@ -240,6 +240,7 @@ def _token_dispatch_fwd_rule( num_experts, worst_case_out_tokens, hidden_size, + align_size=align_size, ) else: # No padding @@ -268,7 +269,6 @@ def _token_dispatch_fwd_rule( def _token_dispatch_bwd_rule( - _routing_map: jnp.ndarray, _num_out_tokens: int, _worst_case_out_tokens: int, _align_size: Optional[int], @@ -281,8 +281,12 @@ def _token_dispatch_bwd_rule( Optional[jnp.ndarray], Optional[jnp.ndarray], ], -) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: - """Backward pass rule for token_dispatch.""" +) -> Tuple[jnp.ndarray, None, Optional[jnp.ndarray]]: + """Backward pass rule for token_dispatch. + + Returns gradients for (inp, routing_map, probs). + routing_map gradient is None since it's a discrete routing decision. + """ row_id_map, pad_offsets, num_tokens, num_experts, hidden_size, with_probs = residuals output_grad, permuted_probs_grad, _, _, _ = g # Ignore row_id_map, pad_offsets, target grads @@ -309,7 +313,9 @@ def _token_dispatch_bwd_rule( hidden_size, ) - return inp_grad, probs_grad if with_probs else None + # Return gradients for (inp, routing_map, probs) + # routing_map is non-differentiable (discrete routing), so return None + return inp_grad, None, probs_grad if with_probs else None _token_dispatch.defvjp(_token_dispatch_fwd_rule, _token_dispatch_bwd_rule) @@ -497,6 +503,8 @@ def _token_combine_bwd_rule( else: # Simple case: just permute gradients back if pad_offsets is not None: + # Note: align_size uses default (128) since buffer sizes are already + # determined from forward pass (stored in residuals as num_out_tokens) inp_grad, _ = permute_with_mask_map_and_pad( output_grad, row_id_map, @@ -506,6 +514,7 @@ def _token_combine_bwd_rule( num_experts, num_out_tokens, hidden_size, + align_size=128, # Default, sizes already computed in forward ) # The permute kernel only writes to positions that tokens map to. # Padded positions may contain uninitialized (NaN) values - replace with zeros. diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index 849673fe31..bd8bd8ff13 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -8,9 +8,13 @@ import jax import jax.numpy as jnp +from jax.sharding import PartitionSpec +from jax.experimental.custom_partitioning import SdyShardingRule import triton from transformer_engine.jax.cpp_extensions.base import BasePrimitive, register_primitive +from transformer_engine.jax.cpp_extensions.misc import get_padded_spec, NamedSharding +from transformer_engine.jax.sharding import get_mesh_axis_size from transformer_engine.common.triton.permutation import ( _row_id_map_pass_1_kernel, _row_id_map_pass_2_kernel, @@ -93,7 +97,6 @@ def impl(routing_map, num_tokens, num_experts, block_size): @staticmethod def lowering(ctx, routing_map, *, num_tokens, num_experts, block_size): """MLIR lowering using triton_call_lowering.""" - # Compute strides routing_stride_token = num_experts routing_stride_expert = 1 row_id_stride_token = num_experts * 2 + 1 @@ -101,11 +104,10 @@ def lowering(ctx, routing_map, *, num_tokens, num_experts, block_size): grid = (num_experts, triton.cdiv(num_tokens, block_size)) - # All scalar arguments must be passed as constexprs return triton_call_lowering( ctx, _row_id_map_pass_1_kernel, - routing_map, # Only tensor arguments here + routing_map, grid=grid, constexprs={ "num_tokens": num_tokens, @@ -117,6 +119,76 @@ def lowering(ctx, routing_map, *, num_tokens, num_experts, block_size): }, ) + @staticmethod + def infer_sharding_from_operands( + num_tokens, num_experts, block_size, mesh, arg_infos, result_infos + ): + """Infer output sharding from input sharding.""" + del num_tokens, num_experts, block_size, result_infos + routing_map_spec = get_padded_spec(arg_infos[0]) + # row_id_map has same token dimension sharding as routing_map + # Shape: (num_tokens, num_experts * 2 + 1) + row_id_map_sharding = NamedSharding( + mesh, + PartitionSpec(routing_map_spec[0], None), + desc="RowIdMapPass1.row_id_map_sharding", + ) + # Workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + workspace_sharding = NamedSharding( + mesh, + PartitionSpec(None, None), + desc="RowIdMapPass1.workspace_sharding", + ) + return [row_id_map_sharding, workspace_sharding] + + @staticmethod + def partition(num_tokens, num_experts, block_size, mesh, arg_infos, result_infos): + """Row id map 1st pass partition.""" + del num_tokens, result_infos + routing_map_spec = get_padded_spec(arg_infos[0]) + + # Input sharding + arg_shardings = (arg_infos[0].sharding,) + + # Output shardings + row_id_map_sharding = NamedSharding( + mesh, + PartitionSpec(routing_map_spec[0], None), + desc="RowIdMapPass1.row_id_map_sharding", + ) + workspace_sharding = NamedSharding( + mesh, + PartitionSpec(None, None), + desc="RowIdMapPass1.workspace_sharding", + ) + out_shardings = [row_id_map_sharding, workspace_sharding] + + def sharded_impl(routing_map): + # Each shard processes its local tokens + local_num_tokens = routing_map.shape[0] + return RowIdMapPass1Primitive.impl( + routing_map, + num_tokens=local_num_tokens, + num_experts=num_experts, + block_size=block_size, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule(num_tokens, num_experts, block_size, mesh, value_types, result_types): + """Shardy sharding rule for this primitive.""" + del num_tokens, num_experts, block_size, mesh, value_types, result_types + prefix = "RowIdMapPass1" + # routing_map shape: (num_tokens, num_experts) + input_spec = (f"{prefix}_tokens", f"{prefix}_experts") + # row_id_map shape: (num_tokens, num_experts * 2 + 1) + # Note: row_id_cols != experts since it's num_experts * 2 + 1 + row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_row_id_cols") + # workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + workspace_spec = (f"{prefix}_experts", f"{prefix}_ws_blocks") + return SdyShardingRule((input_spec,), (row_id_map_spec, workspace_spec)) + register_primitive(RowIdMapPass1Primitive) @@ -185,6 +257,69 @@ def lowering(ctx, row_id_map, workspace, *, num_tokens, num_experts, block_size) }, ) + @staticmethod + def infer_sharding_from_operands( + num_tokens, num_experts, block_size, mesh, arg_infos, result_infos + ): + """Infer output sharding from input sharding.""" + del num_tokens, num_experts, block_size, result_infos + row_id_map_spec = get_padded_spec(arg_infos[0]) + # Output has same sharding as input (in-place operation) + row_id_map_sharding = NamedSharding( + mesh, + PartitionSpec(*row_id_map_spec), + desc="RowIdMapPass2.row_id_map_sharding", + ) + workspace_sharding = NamedSharding( + mesh, + PartitionSpec(None, None), + desc="RowIdMapPass2.workspace_sharding", + ) + return [row_id_map_sharding, workspace_sharding] + + @staticmethod + def partition(num_tokens, num_experts, block_size, mesh, arg_infos, result_infos): + """Partition the primitive for distributed execution.""" + del num_tokens, result_infos + row_id_map_spec = get_padded_spec(arg_infos[0]) + + # Input shardings + arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding) + + # Output shardings (same as inputs for in-place operation) + row_id_map_sharding = NamedSharding( + mesh, + PartitionSpec(*row_id_map_spec), + desc="RowIdMapPass2.row_id_map_sharding", + ) + workspace_sharding = NamedSharding( + mesh, + PartitionSpec(None, None), + desc="RowIdMapPass2.workspace_sharding", + ) + out_shardings = [row_id_map_sharding, workspace_sharding] + + def sharded_impl(row_id_map, workspace): + local_num_tokens = row_id_map.shape[0] + return RowIdMapPass2Primitive.impl( + row_id_map, + workspace, + num_tokens=local_num_tokens, + num_experts=num_experts, + block_size=block_size, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule(num_tokens, num_experts, block_size, mesh, value_types, result_types): + """Shardy sharding rule for this primitive.""" + del num_tokens, num_experts, block_size, mesh, value_types, result_types + prefix = "RowIdMapPass2" + row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_cols") + workspace_spec = (f"{prefix}_ws_experts", f"{prefix}_ws_blocks") + return SdyShardingRule((row_id_map_spec, workspace_spec), (row_id_map_spec, workspace_spec)) + register_primitive(RowIdMapPass2Primitive) @@ -240,6 +375,52 @@ def lowering(ctx, row_id_map, *, num_tokens, num_experts): }, ) + @staticmethod + def infer_sharding_from_operands(num_tokens, num_experts, mesh, arg_infos, result_infos): + """Infer output sharding from input sharding.""" + del num_tokens, num_experts, result_infos + row_id_map_spec = get_padded_spec(arg_infos[0]) + # Output has same sharding as input (in-place operation) + return NamedSharding( + mesh, + PartitionSpec(*row_id_map_spec), + desc="RowIdMapPass3.row_id_map_sharding", + ) + + @staticmethod + def partition(num_tokens, num_experts, mesh, arg_infos, result_infos): + """Partition the primitive for distributed execution.""" + del num_tokens, result_infos + row_id_map_spec = get_padded_spec(arg_infos[0]) + + # Input sharding + arg_shardings = (arg_infos[0].sharding,) + + # Output sharding (same as input for in-place operation) + out_sharding = NamedSharding( + mesh, + PartitionSpec(*row_id_map_spec), + desc="RowIdMapPass3.row_id_map_sharding", + ) + + def sharded_impl(row_id_map): + local_num_tokens = row_id_map.shape[0] + return RowIdMapPass3Primitive.impl( + row_id_map, + num_tokens=local_num_tokens, + num_experts=num_experts, + ) + + return mesh, sharded_impl, out_sharding, arg_shardings + + @staticmethod + def shardy_sharding_rule(num_tokens, num_experts, mesh, value_types, result_types): + """Shardy sharding rule for this primitive.""" + del num_tokens, num_experts, mesh, value_types, result_types + prefix = "RowIdMapPass3" + row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_cols") + return SdyShardingRule((row_id_map_spec,), (row_id_map_spec,)) + register_primitive(RowIdMapPass3Primitive) @@ -251,8 +432,12 @@ class PermuteWithMaskMapPrimitive(BasePrimitive): name = "te_permute_with_mask_map_triton" multiple_results = True - # scale, permuted_scale are dummy inputs (not used when PERMUTE_SCALE=False) - # pad_offsets can be shape (0,) when not doing padding, or (num_experts,) when padding + # Outer primitive has 6 tensor inputs: inp, row_id_map, probs, scale, permuted_scale, pad_offsets + # Static args for outer primitive: num_tokens, num_experts, num_out_tokens, hidden_size, + # with_probs, with_pad, align_size + # Inner primitive adds output_buf, permuted_probs_buf) + + # impl_static_args is for the outer primitive's impl() which has 6 tensor inputs. impl_static_args = ( 6, 7, @@ -260,7 +445,8 @@ class PermuteWithMaskMapPrimitive(BasePrimitive): 9, 10, 11, - ) # num_tokens, num_experts, num_out_tokens, hidden_size, with_probs, with_pad + 12, + ) inner_primitive = None outer_primitive = None @@ -272,6 +458,8 @@ def abstract( scale_aval, # dummy, same shape as inp permuted_scale_aval, # dummy, same shape as inp pad_offsets_aval, + output_buf_aval=None, # Pre-zeroed output buffer (inner primitive only) + permuted_probs_buf_aval=None, # Pre-zeroed permuted_probs buffer (inner primitive only) *, num_tokens, num_experts, @@ -279,10 +467,12 @@ def abstract( hidden_size, with_probs, with_pad, + align_size, ): """Shape/dtype inference for permute.""" del row_id_map_aval, scale_aval, permuted_scale_aval, pad_offsets_aval - del num_tokens, num_experts, with_pad + del num_tokens, num_experts, with_pad, align_size + del output_buf_aval, permuted_probs_buf_aval # Used for input_output_aliases only output_shape = (num_out_tokens, hidden_size) output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) @@ -308,9 +498,29 @@ def impl( hidden_size, with_probs, with_pad, + align_size, # align_size is only used for sharding, but must be passed since abstract() requires it ): """Forward to inner primitive.""" + assert PermuteWithMaskMapPrimitive.inner_primitive is not None + + # Create pre-zeroed output buffers for the inner primitive. + # When with_pad=True, this ensures padding positions contain zeros. + # These buffers are aliased to the outputs via input_output_aliases in the lowering. + if with_pad: + output_buf = jnp.zeros((num_out_tokens, hidden_size), dtype=inp.dtype) + if with_probs: + permuted_probs_buf = jnp.zeros((num_out_tokens,), dtype=probs.dtype) + else: + permuted_probs_buf = jnp.zeros((0,), dtype=inp.dtype) + else: + # When not padding, use empty buffers (kernel ignores them, lowering skips aliasing) + output_buf = jnp.empty((num_out_tokens, hidden_size), dtype=inp.dtype) + if with_probs: + permuted_probs_buf = jnp.empty((num_out_tokens,), dtype=probs.dtype) + else: + permuted_probs_buf = jnp.empty((0,), dtype=inp.dtype) + return PermuteWithMaskMapPrimitive.inner_primitive.bind( inp, row_id_map, @@ -318,12 +528,15 @@ def impl( scale, permuted_scale, pad_offsets, + output_buf, + permuted_probs_buf, num_tokens=num_tokens, num_experts=num_experts, num_out_tokens=num_out_tokens, hidden_size=hidden_size, with_probs=with_probs, with_pad=with_pad, + align_size=align_size, ) @staticmethod @@ -335,6 +548,8 @@ def lowering( scale, permuted_scale, pad_offsets, + output_buf, # Pre-zeroed output buffer (for input_output_aliases) + permuted_probs_buf, # Pre-zeroed permuted_probs buffer (for input_output_aliases) *, num_tokens, num_experts, @@ -342,9 +557,10 @@ def lowering( hidden_size, with_probs, with_pad, + align_size, ): """MLIR lowering using triton_call_lowering.""" - del num_out_tokens + del align_size inp_stride_token = hidden_size inp_stride_hidden = 1 output_stride_token = hidden_size @@ -371,6 +587,18 @@ def lowering( 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. + # Input indices: 0=inp, 1=row_id_map, 2=probs, 3=scale, 4=permuted_scale, + # 5=pad_offsets, 6=output_buf, 7=permuted_probs_buf + # Output indices: 0=output, 1=permuted_probs + if with_pad: + input_output_aliases = {6: 0} + if with_probs: + input_output_aliases[7] = 1 + else: + input_output_aliases = None + return triton_call_lowering( ctx, _permute_kernel, @@ -380,9 +608,14 @@ def lowering( scale, permuted_scale, pad_offsets, + output_buf, + permuted_probs_buf, grid=grid, + input_output_aliases=input_output_aliases, constexprs={ "scale_hidden_dim": 0, + "num_tokens": num_tokens, + "num_out_tokens": num_out_tokens, "stride_row_id_map_token": row_id_stride_token, "stride_row_id_map_expert": row_id_stride_expert, "stride_input_token": inp_stride_token, @@ -405,24 +638,242 @@ def lowering( }, ) + @staticmethod + def infer_sharding_from_operands( + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + with_probs, + with_pad, + align_size, + mesh, + arg_infos, + result_infos, + ): + """Infer output sharding from input sharding. + + For batch-dimension partitioning: + - Input (num_tokens, hidden_size) is sharded on token dim + - Output (num_out_tokens, hidden_size) gets same token dim sharding + - Permuted probs (num_out_tokens,) gets same token dim sharding + """ + del align_size # Used only in partition + del num_tokens, num_experts, num_out_tokens, hidden_size, with_pad, result_infos + inp_spec = get_padded_spec(arg_infos[0]) + # Output has same sharding pattern: (token_shard, None) + output_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0], None), + desc="PermuteWithMaskMap.output_sharding", + ) + if with_probs: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0]), + desc="PermuteWithMaskMap.permuted_probs_sharding", + ) + else: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="PermuteWithMaskMap.permuted_probs_sharding_empty", + ) + return [output_sharding, permuted_probs_sharding] + + @staticmethod + def partition( + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + with_probs, + with_pad, + align_size, + mesh, + arg_infos, + result_infos, + ): + """Partition the primitive for distributed execution. + + For batch-dimension partitioning, each GPU processes its local tokens + independently. The row_id_map contains local destination indices, + so no inter-GPU communication is needed. + """ + del num_tokens, result_infos + inp_spec = get_padded_spec(arg_infos[0]) + + # Input shardings - preserve original shardings + arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + + # Output shardings + output_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0], None), + desc="PermuteWithMaskMap.output_sharding", + ) + if with_probs: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0]), + desc="PermuteWithMaskMap.permuted_probs_sharding", + ) + else: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="PermuteWithMaskMap.permuted_probs_sharding_empty", + ) + out_shardings = [output_sharding, permuted_probs_sharding] + + # Get number of data parallel devices from the batch sharding axis + batch_axis = inp_spec[0] + if batch_axis is not None: + num_dp_devices = get_mesh_axis_size(batch_axis, mesh) + else: + num_dp_devices = 1 + + def sharded_impl(inp, row_id_map, probs, scale, permuted_scale, pad_offsets): + # Each shard processes its local tokens independently (data parallelism) + local_num_tokens = inp.shape[0] + + # ========================================================================= + # MoE Permutation Sharding (data parallelism, no expert parallelism) + # ========================================================================= + # Each GPU has ALL experts and processes its local batch of tokens. + # + # TopK bounds output: each token goes to at most topK experts, so: + # global_num_out_tokens = global_num_in_tokens * topK + # local_num_out_tokens = local_num_in_tokens * topK + # = global_num_out_tokens / num_dp_devices + # + # E = num_experts + # A = align_size for padding to group gemm size in cuBLAS + # With padding (align_size != 128, which is the default/no-op value): + # The global num_out_tokens passed here is already worst_case_out_tokens. + # We need to recalculate local worst-case from local raw tokens. + # local_raw_out_tokens = global_raw_out_tokens / num_dp_devices + # local_worst_case = ((local_raw_out + E*(A-1)) // A) * A + # + # Local permute produces output ordered by expert: [E0 | E1 | ... | EN] + # where each expert section contains tokens routed to that expert. + # + # Global assembly (if needed) should be done outside this primitive. + + # ========================================================================= + # Output size calculation + # ========================================================================= + # For both padding and non-padding cases, use simple division. + # The global num_out_tokens is already the worst-case buffer size. + # + # IMPORTANT for padding + sharding: + # Padding overhead is per-shard (each shard needs E*(A-1) extra space). + # The caller must account for this by passing a sufficiently large + # global num_out_tokens such that: global_worst / num_dp >= local_worst + # where local_worst = ((local_raw + E*(A-1)) // A) * A + + local_num_out_tokens = num_out_tokens // num_dp_devices + + # Local permute - output stays sharded on this GPU + local_output, local_permuted_probs = PermuteWithMaskMapPrimitive.impl( + inp, + row_id_map, + probs, + scale, + permuted_scale, + pad_offsets, + num_tokens=local_num_tokens, + num_experts=num_experts, + num_out_tokens=local_num_out_tokens, + hidden_size=hidden_size, + with_probs=with_probs, + with_pad=with_pad, + align_size=align_size, + ) + + return local_output, local_permuted_probs + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule( + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + with_probs, + with_pad, + align_size, + mesh, + value_types, + result_types, + ): + """Shardy sharding rule for this primitive.""" + del ( + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + align_size, + mesh, + value_types, + result_types, + ) + prefix = "PermuteWithMaskMap" + # inp: (num_tokens, hidden_size) + inp_spec = (f"{prefix}_tokens", f"{prefix}_hidden") + # row_id_map: (num_tokens, num_experts * 2 + 1) + row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_row_id_cols") + # probs: (num_tokens, num_experts) or (0,) + probs_spec = ( + (f"{prefix}_tokens", f"{prefix}_experts") if with_probs else (f"{prefix}_empty",) + ) + # scale: (num_tokens, hidden_size) - same shape as inp, permuted together + scale_spec = (f"{prefix}_tokens", f"{prefix}_hidden") + # permuted_scale: (num_out_tokens, hidden_size) - same shape as output + permuted_scale_spec = (f"{prefix}_out_tokens", f"{prefix}_hidden") + # pad_offsets: (num_experts,) or (0,) - uses same experts factor as probs + pad_offsets_spec = (f"{prefix}_experts",) if with_pad else (f"{prefix}_pad_empty",) + # output: (num_out_tokens, hidden_size) + output_spec = (f"{prefix}_out_tokens", f"{prefix}_hidden") + # permuted_probs: (num_out_tokens,) or (0,) + permuted_probs_spec = (f"{prefix}_out_tokens",) if with_probs else (f"{prefix}_empty2",) + + return SdyShardingRule( + ( + inp_spec, + row_id_map_spec, + probs_spec, + scale_spec, + permuted_scale_spec, + pad_offsets_spec, + ), + (output_spec, permuted_probs_spec), + ) + register_primitive(PermuteWithMaskMapPrimitive) class UnpermuteWithMaskMapPrimitive(BasePrimitive): """ - Unpermute the input tensor based on the row_id_map. + Unpermute the input tensor based on the row_id_map, optionally with fused unpadding. """ name = "te_unpermute_with_mask_map_triton" multiple_results = True + # Outer primitive has 5 tensor inputs: inp, row_id_map, merging_probs, permuted_probs, pad_offsets + # Static args for outer primitive: num_tokens, num_experts, hidden_size, + # with_merging_probs, with_probs, with_unpad + # Inner primitive has adds output_buf, unpermuted_probs_buf impl_static_args = ( 5, 6, 7, 8, 9, - ) # num_tokens, num_experts, hidden_size, with_merging_probs, with_probs + 10, + ) inner_primitive = None outer_primitive = None @@ -432,16 +883,20 @@ def abstract( row_id_map_aval, merging_probs_aval, permuted_probs_aval, - pad_offsets_aval, # dummy, not used when FUSION_UNPAD=False + pad_offsets_aval, + output_buf_aval=None, # Dummy (inner primitive only) + unpermuted_probs_buf_aval=None, # Dummy (inner primitive only) *, num_tokens, num_experts, hidden_size, with_merging_probs, with_probs, + with_unpad, ): """Shape/dtype inference for unpermute.""" - del row_id_map_aval, merging_probs_aval, with_merging_probs, pad_offsets_aval + del row_id_map_aval, merging_probs_aval, with_merging_probs, pad_offsets_aval, with_unpad + del output_buf_aval, unpermuted_probs_buf_aval output_shape = (num_tokens, hidden_size) output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) @@ -468,20 +923,33 @@ def impl( hidden_size, with_merging_probs, with_probs, + with_unpad, ): """Forward to inner primitive.""" assert UnpermuteWithMaskMapPrimitive.inner_primitive is not None + + # Create dummy buffers for kernel signature consistency with _permute_kernel. + # These are not used for pre-zeroing since unpermute writes to all output positions. + output_buf = jnp.empty((num_tokens, hidden_size), dtype=inp.dtype) + if with_probs: + unpermuted_probs_buf = jnp.empty((num_tokens, num_experts), dtype=permuted_probs.dtype) + else: + unpermuted_probs_buf = jnp.empty((0,), dtype=inp.dtype) + return UnpermuteWithMaskMapPrimitive.inner_primitive.bind( inp, row_id_map, merging_probs, permuted_probs, pad_offsets, + output_buf, + unpermuted_probs_buf, num_tokens=num_tokens, num_experts=num_experts, hidden_size=hidden_size, with_merging_probs=with_merging_probs, with_probs=with_probs, + with_unpad=with_unpad, ) @staticmethod @@ -492,12 +960,15 @@ def lowering( merging_probs, permuted_probs, pad_offsets, + output_buf, # Dummy for kernel signature consistency + unpermuted_probs_buf, # Dummy for kernel signature consistency *, num_tokens, num_experts, hidden_size, with_merging_probs, with_probs, + with_unpad, ): """MLIR lowering using triton_call_lowering.""" # Compute strides @@ -523,7 +994,6 @@ def lowering( block_size = _get_min_block_size(_unpermute_kernel) grid = (num_tokens, triton.cdiv(hidden_size, block_size)) - # Pass all 5 inputs including pad_offsets (even though FUSION_UNPAD=False) return triton_call_lowering( ctx, _unpermute_kernel, @@ -532,6 +1002,8 @@ def lowering( merging_probs, permuted_probs, pad_offsets, + output_buf, + unpermuted_probs_buf, grid=grid, constexprs={ "stride_row_id_map_token": row_id_stride_token, @@ -550,174 +1022,170 @@ def lowering( "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), "WITH_MERGING_PROBS": with_merging_probs, "PERMUTE_PROBS": with_probs, - "FUSION_UNPAD": False, + "FUSION_UNPAD": with_unpad, "BLOCK_SIZE": block_size, }, ) - -register_primitive(UnpermuteWithMaskMapPrimitive) - - -class UnpermuteWithMaskMapAndUnpadPrimitive(BasePrimitive): - """ - Unpermute the input tensor based on the row_id_map with fused unpadding. - """ - - name = "te_unpermute_with_mask_map_and_unpad_triton" - multiple_results = True - impl_static_args = ( - 5, - 6, - 7, - 8, - 9, - ) # num_tokens, num_experts, hidden_size, with_merging_probs, with_probs - inner_primitive = None - outer_primitive = None - @staticmethod - def abstract( - inp_aval, - row_id_map_aval, - merging_probs_aval, - permuted_probs_aval, - pad_offsets_aval, - *, + def infer_sharding_from_operands( num_tokens, num_experts, hidden_size, with_merging_probs, with_probs, + with_unpad, + mesh, + arg_infos, + result_infos, ): - """Shape/dtype inference for unpermute with unpadding.""" - del row_id_map_aval, merging_probs_aval, with_merging_probs, pad_offsets_aval - - output_shape = (num_tokens, hidden_size) - output_aval = jax.core.ShapedArray(output_shape, inp_aval.dtype) - + """Infer output sharding from input sharding. + + For batch-dimension partitioning: + - row_id_map (num_tokens, num_experts*2+1) is sharded on token dim + - Output (num_tokens, hidden_size) gets same token dim sharding + """ + del num_tokens, num_experts, hidden_size, with_merging_probs, with_unpad, result_infos + row_id_map_spec = get_padded_spec(arg_infos[1]) + # Output has same token dimension sharding as row_id_map + output_sharding = NamedSharding( + mesh, + PartitionSpec(row_id_map_spec[0], None), + desc="UnpermuteWithMaskMap.output_sharding", + ) if with_probs: - unpermuted_probs_shape = (num_tokens, num_experts) - unpermuted_probs_aval = jax.core.ShapedArray( - unpermuted_probs_shape, permuted_probs_aval.dtype + unpermuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(row_id_map_spec[0], None), + desc="UnpermuteWithMaskMap.unpermuted_probs_sharding", ) else: - unpermuted_probs_aval = jax.core.ShapedArray((0,), inp_aval.dtype) - - return output_aval, unpermuted_probs_aval + unpermuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="UnpermuteWithMaskMap.unpermuted_probs_sharding_empty", + ) + return [output_sharding, unpermuted_probs_sharding] @staticmethod - def impl( - inp, - row_id_map, - merging_probs, - permuted_probs, - pad_offsets, + def partition( num_tokens, num_experts, hidden_size, with_merging_probs, with_probs, + with_unpad, + mesh, + arg_infos, + result_infos, ): - """Forward to inner primitive.""" - assert UnpermuteWithMaskMapAndUnpadPrimitive.inner_primitive is not None - return UnpermuteWithMaskMapAndUnpadPrimitive.inner_primitive.bind( - inp, - row_id_map, - merging_probs, - permuted_probs, - pad_offsets, - num_tokens=num_tokens, - num_experts=num_experts, - hidden_size=hidden_size, - with_merging_probs=with_merging_probs, - with_probs=with_probs, + """Partition the primitive for distributed execution.""" + del num_tokens, result_infos + row_id_map_spec = get_padded_spec(arg_infos[1]) + + # Input shardings - preserve original shardings + arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + + # Output shardings + output_sharding = NamedSharding( + mesh, + PartitionSpec(row_id_map_spec[0], None), + desc="UnpermuteWithMaskMap.output_sharding", ) + if with_probs: + unpermuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(row_id_map_spec[0], None), + desc="UnpermuteWithMaskMap.unpermuted_probs_sharding", + ) + else: + unpermuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="UnpermuteWithMaskMap.unpermuted_probs_sharding_empty", + ) + out_shardings = [output_sharding, unpermuted_probs_sharding] + + def sharded_impl(inp, row_id_map, merging_probs, permuted_probs, pad_offsets): + # Each shard processes its local tokens + local_num_tokens = row_id_map.shape[0] + return UnpermuteWithMaskMapPrimitive.impl( + inp, + row_id_map, + merging_probs, + permuted_probs, + pad_offsets, + num_tokens=local_num_tokens, + num_experts=num_experts, + hidden_size=hidden_size, # hidden_size is not sharded + with_merging_probs=with_merging_probs, + with_probs=with_probs, + with_unpad=with_unpad, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings @staticmethod - def lowering( - ctx, - inp, - row_id_map, - merging_probs, - permuted_probs, - pad_offsets, - *, + def shardy_sharding_rule( num_tokens, num_experts, hidden_size, with_merging_probs, with_probs, + with_unpad, + mesh, + value_types, + result_types, ): - """MLIR lowering using triton_call_lowering.""" - # Compute strides - inp_stride_token = hidden_size - inp_stride_hidden = 1 - output_stride_token = hidden_size - output_stride_hidden = 1 - row_id_stride_token = num_experts * 2 + 1 - row_id_stride_expert = 1 - - if with_merging_probs: - merging_probs_stride_token = num_experts - merging_probs_stride_expert = 1 - else: - merging_probs_stride_token = 0 - merging_probs_stride_expert = 0 - - permuted_probs_stride_token = 1 - unpermuted_probs_stride_token = num_experts - unpermuted_probs_stride_expert = 1 - - # Grid - use minimum BLOCK_SIZE from autotune configs - block_size = _get_min_block_size(_unpermute_kernel) - grid = (num_tokens, triton.cdiv(hidden_size, block_size)) + """Shardy sharding rule for this primitive.""" + del num_tokens, num_experts, hidden_size, mesh, value_types, result_types + prefix = "UnpermuteWithMaskMap" + # inp: (num_out_tokens, hidden_size) + inp_spec = (f"{prefix}_out_tokens", f"{prefix}_hidden") + # row_id_map: (num_tokens, num_experts * 2 + 1) + row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_row_id_cols") + # merging_probs: (num_tokens, num_experts) or (0,) + merging_probs_spec = ( + (f"{prefix}_tokens", f"{prefix}_experts") + if with_merging_probs + else (f"{prefix}_empty",) + ) + # permuted_probs: (num_out_tokens,) or (0,) + permuted_probs_spec = (f"{prefix}_out_tokens",) if with_probs else (f"{prefix}_empty2",) + # pad_offsets: (num_experts,) when with_unpad=True, or dummy (0,) otherwise + pad_offsets_spec = (f"{prefix}_experts",) if with_unpad else (f"{prefix}_pad_empty",) + # output: (num_tokens, hidden_size) + output_spec = (f"{prefix}_tokens", f"{prefix}_hidden") + # unpermuted_probs: (num_tokens, num_experts) or (0,) + unpermuted_probs_spec = ( + (f"{prefix}_tokens", f"{prefix}_experts") if with_probs else (f"{prefix}_empty3",) + ) - return triton_call_lowering( - ctx, - _unpermute_kernel, - inp, - row_id_map, - merging_probs, - permuted_probs, - pad_offsets, - grid=grid, - constexprs={ - "stride_row_id_map_token": row_id_stride_token, - "stride_row_id_map_expert": row_id_stride_expert, - "stride_input_token": inp_stride_token, - "stride_input_hidden": inp_stride_hidden, - "stride_output_token": output_stride_token, - "stride_output_hidden": output_stride_hidden, - "stride_merging_probs_token": merging_probs_stride_token, - "stride_merging_probs_expert": merging_probs_stride_expert, - "stride_permuted_probs_token": permuted_probs_stride_token, - "stride_unpermuted_probs_token": unpermuted_probs_stride_token, - "stride_unpermuted_probs_expert": unpermuted_probs_stride_expert, - "num_experts": num_experts, - "hidden_size": hidden_size, - "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), - "WITH_MERGING_PROBS": with_merging_probs, - "PERMUTE_PROBS": with_probs, - "FUSION_UNPAD": True, - "BLOCK_SIZE": block_size, - }, + return SdyShardingRule( + (inp_spec, row_id_map_spec, merging_probs_spec, permuted_probs_spec, pad_offsets_spec), + (output_spec, unpermuted_probs_spec), ) -register_primitive(UnpermuteWithMaskMapAndUnpadPrimitive) +register_primitive(UnpermuteWithMaskMapPrimitive) class UnpermuteBwdWithMergingProbsPrimitive(BasePrimitive): """ - Backward pass for unpermute with merging probabilities. + Backward pass for unpermute with merging probabilities, optionally with fused unpadding. This kernel computes gradients for both the input and merging_probs. """ name = "te_unpermute_bwd_with_merging_probs_triton" multiple_results = True - impl_static_args = (5, 6, 7, 8) # num_tokens, num_experts, num_out_tokens, hidden_size + impl_static_args = ( + 5, + 6, + 7, + 8, + 9, + ) # num_tokens, num_experts, num_out_tokens, hidden_size, with_unpad inner_primitive = None outer_primitive = None @@ -727,15 +1195,16 @@ def abstract( fwd_input_aval, merging_probs_aval, row_id_map_aval, - pad_offsets_aval, # dummy, not used when FUSION_UNPAD=False + pad_offsets_aval, *, num_tokens, num_experts, num_out_tokens, hidden_size, + with_unpad, ): """Shape/dtype inference for unpermute backward with merging probs.""" - del fwd_input_aval, row_id_map_aval, pad_offsets_aval + del fwd_input_aval, row_id_map_aval, pad_offsets_aval, with_unpad # fwd_input_grad has same shape as fwd_input fwd_input_grad_shape = (num_out_tokens, hidden_size) @@ -760,6 +1229,7 @@ def impl( num_experts, num_out_tokens, hidden_size, + with_unpad, ): """Forward to inner primitive.""" assert UnpermuteBwdWithMergingProbsPrimitive.inner_primitive is not None @@ -773,6 +1243,7 @@ def impl( num_experts=num_experts, num_out_tokens=num_out_tokens, hidden_size=hidden_size, + with_unpad=with_unpad, ) @staticmethod @@ -788,6 +1259,7 @@ def lowering( num_experts, num_out_tokens, hidden_size, + with_unpad, ): """MLIR lowering using triton_call_lowering.""" del num_out_tokens @@ -812,7 +1284,6 @@ def lowering( # Get min block size from autotune configs for consistency block_size = _get_min_block_size(_unpermute_bwd_with_merging_probs_kernel) - # Pass all 5 inputs including pad_offsets (even though FUSION_UNPAD=False) return triton_call_lowering( ctx, _unpermute_bwd_with_merging_probs_kernel, @@ -838,152 +1309,126 @@ def lowering( "num_experts": num_experts, "hidden_size": hidden_size, "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), - "FUSION_UNPAD": False, + "FUSION_UNPAD": with_unpad, "BLOCK_SIZE": block_size, }, ) - -register_primitive(UnpermuteBwdWithMergingProbsPrimitive) - - -class UnpermuteBwdWithMergingProbsAndUnpadPrimitive(BasePrimitive): - """ - Backward pass for unpermute with merging probabilities and fused unpadding. - - This kernel computes gradients for both the input and merging_probs, - while handling padded outputs. - """ - - name = "te_unpermute_bwd_with_merging_probs_and_unpad_triton" - multiple_results = True - impl_static_args = (5, 6, 7, 8) # num_tokens, num_experts, num_out_tokens, hidden_size - inner_primitive = None - outer_primitive = None - @staticmethod - def abstract( - fwd_output_grad_aval, - fwd_input_aval, - merging_probs_aval, - row_id_map_aval, - pad_offsets_aval, - *, + def infer_sharding_from_operands( num_tokens, num_experts, num_out_tokens, hidden_size, + with_unpad, + mesh, + arg_infos, + result_infos, ): - """Shape/dtype inference for unpermute backward with merging probs and unpadding.""" - del fwd_input_aval, row_id_map_aval, pad_offsets_aval - - # fwd_input_grad has same shape as fwd_input - fwd_input_grad_shape = (num_out_tokens, hidden_size) - fwd_input_grad_aval = jax.core.ShapedArray(fwd_input_grad_shape, fwd_output_grad_aval.dtype) - - # merging_probs_grad has same shape as merging_probs - merging_probs_grad_shape = (num_tokens, num_experts) - merging_probs_grad_aval = jax.core.ShapedArray( - merging_probs_grad_shape, merging_probs_aval.dtype + """Infer output sharding from input sharding.""" + del num_tokens, num_experts, num_out_tokens, hidden_size, with_unpad, result_infos + fwd_output_grad_spec = get_padded_spec(arg_infos[0]) + merging_probs_spec = get_padded_spec(arg_infos[2]) + # fwd_input_grad has same token sharding as fwd_output_grad + fwd_input_grad_sharding = NamedSharding( + mesh, + PartitionSpec(fwd_output_grad_spec[0], None), + desc="UnpermuteBwdWithMergingProbs.fwd_input_grad_sharding", ) - - return fwd_input_grad_aval, merging_probs_grad_aval + # merging_probs_grad has same sharding as merging_probs + merging_probs_grad_sharding = NamedSharding( + mesh, + PartitionSpec(merging_probs_spec[0], None), + desc="UnpermuteBwdWithMergingProbs.merging_probs_grad_sharding", + ) + return [fwd_input_grad_sharding, merging_probs_grad_sharding] @staticmethod - def impl( - fwd_output_grad, - fwd_input, - merging_probs, - row_id_map, - pad_offsets, + def partition( num_tokens, num_experts, num_out_tokens, hidden_size, + with_unpad, + mesh, + arg_infos, + result_infos, ): - """Forward to inner primitive.""" - assert UnpermuteBwdWithMergingProbsAndUnpadPrimitive.inner_primitive is not None - return UnpermuteBwdWithMergingProbsAndUnpadPrimitive.inner_primitive.bind( - fwd_output_grad, - fwd_input, - merging_probs, - row_id_map, - pad_offsets, - num_tokens=num_tokens, - num_experts=num_experts, - num_out_tokens=num_out_tokens, - hidden_size=hidden_size, + """Partition the primitive for distributed execution.""" + del num_tokens, num_out_tokens, result_infos + fwd_output_grad_spec = get_padded_spec(arg_infos[0]) + merging_probs_spec = get_padded_spec(arg_infos[2]) + + arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + + fwd_input_grad_sharding = NamedSharding( + mesh, + PartitionSpec(fwd_output_grad_spec[0], None), + desc="UnpermuteBwdWithMergingProbs.fwd_input_grad_sharding", + ) + merging_probs_grad_sharding = NamedSharding( + mesh, + PartitionSpec(merging_probs_spec[0], None), + desc="UnpermuteBwdWithMergingProbs.merging_probs_grad_sharding", ) + out_shardings = [fwd_input_grad_sharding, merging_probs_grad_sharding] + + def sharded_impl(fwd_output_grad, fwd_input, merging_probs, row_id_map, pad_offsets): + local_num_tokens = row_id_map.shape[0] + # NOTE: local_num_out_tokens is obtained from the actual tensor shape, + # which reflects the data-dependent output size from the forward pass. + local_num_out_tokens = fwd_input.shape[0] + return UnpermuteBwdWithMergingProbsPrimitive.impl( + fwd_output_grad, + fwd_input, + merging_probs, + row_id_map, + pad_offsets, + num_tokens=local_num_tokens, + num_experts=num_experts, + num_out_tokens=local_num_out_tokens, + hidden_size=hidden_size, # hidden_size is not sharded + with_unpad=with_unpad, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings @staticmethod - def lowering( - ctx, - fwd_output_grad, - fwd_input, - merging_probs, - row_id_map, - pad_offsets, - *, + def shardy_sharding_rule( num_tokens, num_experts, num_out_tokens, hidden_size, + with_unpad, + mesh, + value_types, + result_types, ): - """MLIR lowering using triton_call_lowering.""" - del num_out_tokens - - # Compute strides - row_id_stride_token = num_experts * 2 + 1 - row_id_stride_expert = 1 - fwd_output_grad_stride_token = hidden_size - fwd_output_grad_stride_hidden = 1 - fwd_input_grad_stride_token = hidden_size - fwd_input_grad_stride_hidden = 1 - fwd_input_stride_token = hidden_size - fwd_input_stride_hidden = 1 - merging_probs_stride_token = num_experts - merging_probs_stride_expert = 1 - merging_probs_grad_stride_token = num_experts - merging_probs_grad_stride_expert = 1 - - # Grid - one program per token - grid = (num_tokens,) - - # Get min block size from autotune configs for consistency - block_size = _get_min_block_size(_unpermute_bwd_with_merging_probs_kernel) - - return triton_call_lowering( - ctx, - _unpermute_bwd_with_merging_probs_kernel, - fwd_output_grad, - fwd_input, - merging_probs, - row_id_map, - pad_offsets, - grid=grid, - constexprs={ - "stride_row_id_map_token": row_id_stride_token, - "stride_row_id_map_expert": row_id_stride_expert, - "stride_fwd_output_grad_token": fwd_output_grad_stride_token, - "stride_fwd_output_grad_hidden": fwd_output_grad_stride_hidden, - "stride_fwd_input_grad_token": fwd_input_grad_stride_token, - "stride_fwd_input_grad_hidden": fwd_input_grad_stride_hidden, - "stride_fwd_input_token": fwd_input_stride_token, - "stride_fwd_input_hidden": fwd_input_stride_hidden, - "stride_merging_probs_token": merging_probs_stride_token, - "stride_merging_probs_expert": merging_probs_stride_expert, - "stride_merging_probs_grad_token": merging_probs_grad_stride_token, - "stride_merging_probs_grad_expert": merging_probs_grad_stride_expert, - "num_experts": num_experts, - "hidden_size": hidden_size, - "PROBS_LOAD_WIDTH": triton.next_power_of_2(num_experts), - "FUSION_UNPAD": True, - "BLOCK_SIZE": block_size, - }, + """Shardy sharding rule for this primitive.""" + del num_tokens, num_experts, num_out_tokens, hidden_size, mesh, value_types, result_types + prefix = "UnpermuteBwdWithMergingProbs" + fwd_output_grad_spec = (f"{prefix}_tokens", f"{prefix}_hidden") + fwd_input_spec = (f"{prefix}_out_tokens", f"{prefix}_hidden") + merging_probs_spec = (f"{prefix}_tokens", f"{prefix}_experts") + row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_row_id_cols") + # pad_offsets: (num_experts,) when with_unpad=True, or dummy (0,) otherwise + pad_offsets_spec = (f"{prefix}_experts",) if with_unpad else (f"{prefix}_pad_empty",) + fwd_input_grad_spec = (f"{prefix}_out_tokens", f"{prefix}_hidden") + merging_probs_grad_spec = (f"{prefix}_tokens", f"{prefix}_experts") + + return SdyShardingRule( + ( + fwd_output_grad_spec, + fwd_input_spec, + merging_probs_spec, + row_id_map_spec, + pad_offsets_spec, + ), + (fwd_input_grad_spec, merging_probs_grad_spec), ) -register_primitive(UnpermuteBwdWithMergingProbsAndUnpadPrimitive) +register_primitive(UnpermuteBwdWithMergingProbsPrimitive) def unpermute_bwd_with_merging_probs( @@ -1027,7 +1472,7 @@ def unpermute_bwd_with_merging_probs( merging_probs_grad : jnp.ndarray Gradient w.r.t. merging_probs of shape `[num_tokens, num_experts]`. """ - # Create dummy pad_offsets (not used when FUSION_UNPAD=False, but required by kernel signature) + # Create dummy pad_offsets (not used when with_unpad=False, but required by kernel signature) dummy_pad_offsets = jnp.zeros((0,), dtype=jnp.int32) # Pass arguments in kernel order: fwd_output_grad, fwd_input, merging_probs, row_id_map, pad_offsets return UnpermuteBwdWithMergingProbsPrimitive.outer_primitive.bind( @@ -1040,6 +1485,7 @@ def unpermute_bwd_with_merging_probs( num_experts=num_experts, num_out_tokens=num_out_tokens, hidden_size=hidden_size, + with_unpad=False, ) @@ -1088,7 +1534,7 @@ def unpermute_bwd_with_merging_probs_and_unpad( merging_probs_grad : jnp.ndarray Gradient w.r.t. merging_probs of shape `[num_tokens, num_experts]`. """ - return UnpermuteBwdWithMergingProbsAndUnpadPrimitive.outer_primitive.bind( + return UnpermuteBwdWithMergingProbsPrimitive.outer_primitive.bind( fwd_output_grad, fwd_input, merging_probs, @@ -1098,6 +1544,7 @@ def unpermute_bwd_with_merging_probs_and_unpad( num_experts=num_experts, num_out_tokens=num_out_tokens, hidden_size=hidden_size, + with_unpad=True, ) @@ -1147,6 +1594,54 @@ def lowering(ctx, split_sizes, sorted_indices, *, num_tokens, num_splits): }, ) + @staticmethod + def infer_sharding_from_operands(num_tokens, num_splits, mesh, arg_infos, result_infos): + """Infer output sharding from input sharding.""" + del num_tokens, num_splits, result_infos, arg_infos + # row_id_map is replicated since split_sizes and sorted_indices are typically small + return NamedSharding( + mesh, + PartitionSpec(None), + desc="MakeChunkSortMap.row_id_map_sharding", + ) + + @staticmethod + def partition(num_tokens, num_splits, mesh, arg_infos, result_infos): + """Partition the primitive for distributed execution.""" + del result_infos + + arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + + out_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="MakeChunkSortMap.row_id_map_sharding", + ) + + def sharded_impl(split_sizes, sorted_indices): + return MakeChunkSortMapPrimitive.impl( + split_sizes, + sorted_indices, + num_tokens=num_tokens, + num_splits=num_splits, + ) + + return mesh, sharded_impl, out_sharding, arg_shardings + + @staticmethod + def shardy_sharding_rule(num_tokens, num_splits, mesh, value_types, result_types): + """Shardy sharding rule for this primitive.""" + del num_tokens, num_splits, mesh, value_types, result_types + prefix = "MakeChunkSortMap" + split_sizes_spec = (f"{prefix}_splits",) + sorted_indices_spec = (f"{prefix}_splits",) + row_id_map_spec = (f"{prefix}_tokens",) + + return SdyShardingRule( + (split_sizes_spec, sorted_indices_spec), + (row_id_map_spec,), + ) + register_primitive(MakeChunkSortMapPrimitive) @@ -1228,6 +1723,91 @@ def lowering(ctx, inp, row_id_map, probs, *, num_tokens, hidden_size, is_forward }, ) + @staticmethod + def infer_sharding_from_operands( + num_tokens, hidden_size, is_forward, with_probs, mesh, arg_infos, result_infos + ): + """Infer output sharding from input sharding.""" + del num_tokens, hidden_size, is_forward, result_infos + inp_spec = get_padded_spec(arg_infos[0]) + output_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0], None), + desc="SortChunksByMap.output_sharding", + ) + if with_probs: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0]), + desc="SortChunksByMap.permuted_probs_sharding", + ) + else: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="SortChunksByMap.permuted_probs_sharding_empty", + ) + return [output_sharding, permuted_probs_sharding] + + @staticmethod + def partition(num_tokens, hidden_size, is_forward, with_probs, mesh, arg_infos, result_infos): + """Partition the primitive for distributed execution.""" + del num_tokens, result_infos + inp_spec = get_padded_spec(arg_infos[0]) + + arg_shardings = tuple(arg_i.sharding for arg_i in arg_infos) + + output_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0], None), + desc="SortChunksByMap.output_sharding", + ) + if with_probs: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(inp_spec[0]), + desc="SortChunksByMap.permuted_probs_sharding", + ) + else: + permuted_probs_sharding = NamedSharding( + mesh, + PartitionSpec(None), + desc="SortChunksByMap.permuted_probs_sharding_empty", + ) + out_shardings = [output_sharding, permuted_probs_sharding] + + def sharded_impl(inp, row_id_map, probs): + local_num_tokens = inp.shape[0] + return SortChunksByMapPrimitive.impl( + inp, + row_id_map, + probs, + num_tokens=local_num_tokens, + hidden_size=hidden_size, # hidden_size is not sharded + is_forward=is_forward, + with_probs=with_probs, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule( + num_tokens, hidden_size, is_forward, with_probs, mesh, value_types, result_types + ): + """Shardy sharding rule for this primitive.""" + del num_tokens, hidden_size, is_forward, mesh, value_types, result_types + prefix = "SortChunksByMap" + inp_spec = (f"{prefix}_tokens", f"{prefix}_hidden") + row_id_map_spec = (f"{prefix}_tokens",) + probs_spec = (f"{prefix}_tokens",) if with_probs else (f"{prefix}_empty",) + output_spec = (f"{prefix}_tokens", f"{prefix}_hidden") + permuted_probs_spec = (f"{prefix}_tokens",) if with_probs else (f"{prefix}_empty2",) + + return SdyShardingRule( + (inp_spec, row_id_map_spec, probs_spec), + (output_spec, permuted_probs_spec), + ) + register_primitive(SortChunksByMapPrimitive) @@ -1356,6 +1936,7 @@ def permute_with_mask_map( hidden_size=hidden_size, with_probs=with_probs, with_pad=False, + align_size=128, # Default value, no-op for non-padding case ) if not with_probs: @@ -1373,6 +1954,7 @@ def permute_with_mask_map_and_pad( num_experts: int, num_out_tokens: int, hidden_size: int, + align_size: int = 128, ) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: """ Permute the input tensor based on the row_id_map with fused padding. @@ -1395,13 +1977,18 @@ def permute_with_mask_map_and_pad( Number of tokens in the permuted tensor (including padding). hidden_size : int Hidden size of the input tensor. + align_size : int + Alignment size for padding (default: 128). Used for distributed sharding + to correctly compute local buffer sizes. Returns ------- output : jnp.ndarray Permuted and padded output tensor of shape `[num_out_tokens, hidden_size]`. + Padding positions are zero-filled. permuted_probs : Optional[jnp.ndarray] Permuted probabilities if probs was provided, None otherwise. + Padding positions are zero-filled. """ with_probs = probs is not None @@ -1426,8 +2013,14 @@ def permute_with_mask_map_and_pad( hidden_size=hidden_size, with_probs=with_probs, with_pad=True, + align_size=align_size, ) + # Note: Zero-filling of padding positions is handled by pre-zeroing the output + # buffers in impl() using jnp.zeros(), then aliasing them to the kernel's outputs + # via input_output_aliases. The kernel only writes to valid positions, leaving + # padding positions at zero. + if not with_probs: permuted_probs = None @@ -1479,7 +2072,7 @@ def unpermute_with_mask_map( merging_probs = jnp.zeros((0,), dtype=inp.dtype) if not with_probs: permuted_probs = jnp.zeros((0,), dtype=inp.dtype) - # Create dummy pad_offsets (not used when FUSION_UNPAD=False, but required by kernel signature) + # Create dummy pad_offsets (not used when with_unpad=False, but required by kernel signature) dummy_pad_offsets = jnp.zeros((0,), dtype=jnp.int32) output, unpermuted_probs = UnpermuteWithMaskMapPrimitive.outer_primitive.bind( @@ -1493,6 +2086,7 @@ def unpermute_with_mask_map( hidden_size=hidden_size, with_merging_probs=with_merging_probs, with_probs=with_probs, + with_unpad=False, ) if not with_probs: @@ -1550,7 +2144,7 @@ def unpermute_with_mask_map_and_unpad( if not with_probs: permuted_probs = jnp.zeros((0,), dtype=inp.dtype) - output, unpermuted_probs = UnpermuteWithMaskMapAndUnpadPrimitive.outer_primitive.bind( + output, unpermuted_probs = UnpermuteWithMaskMapPrimitive.outer_primitive.bind( inp, row_id_map, merging_probs, @@ -1561,6 +2155,7 @@ def unpermute_with_mask_map_and_unpad( hidden_size=hidden_size, with_merging_probs=with_merging_probs, with_probs=with_probs, + with_unpad=True, ) if not with_probs: diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 59fc5c60af..6ea4092cbc 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -409,7 +409,8 @@ def lowering(ctx, x, *, block_size): kernel_constexprs = constexprs if constexprs is not None else {} # Handle autotuned kernels - compile all configs - if isinstance(kernel_fn, autotuner.Autotuner): + is_autotuned = isinstance(kernel_fn, autotuner.Autotuner) + if is_autotuned: # Compile all configs for runtime selection kernel_calls = [] actual_kernel_fn = kernel_fn.fn @@ -450,24 +451,23 @@ def lowering(ctx, x, *, block_size): kernel_calls.append((config_call, str(config))) - # Create autotuned kernel call - # Convert input_output_aliases to format with sizes - if input_output_aliases is None: - input_output_aliases = {} - - input_output_aliases_with_sizes = tuple( - ( - input_idx, - output_idx, - ctx.avals_in[input_idx].size * ctx.avals_in[input_idx].dtype.itemsize, - ) - for input_idx, output_idx in input_output_aliases.items() - ) - + # IMPORTANT: We pass an empty tuple for input_output_aliases_with_sizes. + # + # Background: + # 1. jax.ffi.ffi_lowering(operand_output_aliases=...) is a HINT to XLA that an + # output can reuse an input's buffer. XLA may or may not honor this. + # 2. TritonAutotunedKernelCall's input_output_aliases_with_sizes triggers + # save/restore logic during autotuning (see jaxlib/gpu/triton_kernels.cc:630-701). + # + # The problem: The save phase (triton_kernels.cc:632) only saves if buffers[input_idx] == buffers[output_idx], + # but the restore phase (triton_kernels.cc:697-700) unconditionally iterates over all aliases and tries + # to access input_copies[input_idx]. If XLA didn't actually alias the buffers, input_copies[input_idx] doesn't exist, creating an empty vector whose .data() returns nullptr, causing CUDA_ERROR_INVALID_VALUE during the restore memcpy. + # + # WAR: Don't pass aliases to TritonAutotunedKernelCall. kernel_call = gpu_triton.TritonAutotunedKernelCall( f"{actual_kernel_fn.__name__}_autotuned", kernel_calls, - input_output_aliases_with_sizes, + (), # Empty to avoid buggy save/restore in jaxlib/gpu/triton_kernels.cc ) else: @@ -498,15 +498,17 @@ def lowering(ctx, x, *, block_size): serialized_metadata = b"" call_proto = kernel_call.to_proto(actual_kernel_fn.__name__, serialized_metadata) - if input_output_aliases is None: - input_output_aliases = {} + if input_output_aliases: + ffi_operand_output_aliases = input_output_aliases + 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=input_output_aliases, + operand_output_aliases=ffi_operand_output_aliases, ) return rule(ctx, *array_args) diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 8c9003bb5f..6b5de9ab0f 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -157,8 +157,8 @@ def permute_with_mask_map( scale_hidden_dim : int Hidden size of the scale tensor. """ - # Use torch.zeros when pad_offsets is provided to ensure padding regions are zeroed, - # since the kernel doesn't write to padding positions. + # Use torch.zeros when pad_offsets is provided to ensure padding regions are zeroed. + # The kernel writes only to valid positions, leaving padding positions at zero. alloc = torch.zeros if pad_offsets is not None else torch.empty output = alloc((num_out_tokens, hidden_size), dtype=inp.dtype, device="cuda") permuted_probs = ( @@ -178,7 +178,13 @@ def permute_with_mask_map( scale, permuted_scale, pad_offsets, + # Pass output buffers as input parameters (for JAX input_output_aliases compatibility). + # In PyTorch, these point to the same memory as the output pointers below. + output, + permuted_probs, scale_hidden_dim, + num_tokens, + num_out_tokens, row_id_map.stride(0), row_id_map.stride(1), inp.stride(0), @@ -252,6 +258,10 @@ def unpermute_with_mask_map( merging_probs, permuted_probs, pad_offsets, + # Dummy buffer parameters for kernel signature consistency with _permute_kernel. + # These are unused in unpermute but maintain consistent interface. + output, # output_buf_ptr (unused, passed for signature consistency) + unpermuted_probs, # unpermuted_probs_buf_ptr (unused, passed for signature consistency) row_id_map.stride(0), row_id_map.stride(1), inp.stride(0), From 99df881061ba6949081fdd8f00dccd0f617c6594 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Fri, 16 Jan 2026 16:49:30 -0800 Subject: [PATCH 169/521] Add logic for block-scaled tensors with GEMM swizzled scales (#2486) * Add general C API for setting tensor params Signed-off-by: Tim Moon * Implement general accessors for NVTETensor Signed-off-by: Tim Moon * Refactor tex swizzling to skip if scales are already swizzled Signed-off-by: Tim Moon * Add checks for non-swizzled scales in MXFP8 and NVFP4 kernels Signed-off-by: Tim Moon * Support pre-swizzled scales in MXFP8Tensor Signed-off-by: Tim Moon * Add tex function to swizzle MXFP8 scales Signed-off-by: Tim Moon * Fix bug in inplace swizzle function Signed-off-by: Tim Moon * Tweak comments to use "compact/swizzled format" Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MXFP8 quantize kernel with pre-swizzled scales Signed-off-by: Tim Moon * Expose pre-swizzled scales in modules Signed-off-by: Tim Moon * Fix bug in multi-swizzle Signed-off-by: Tim Moon * Support MXFP8 gated activations with swizzled scales Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add PyTorch infrastructure for pre-swizzled NVFP4 tensors Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Deprecate DSv3-specific quantization logic in C API Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove support for DSv3 compact data from quantizer Signed-off-by: Tim Moon * Remove DSv3 compact data format from core lib Signed-off-by: Tim Moon * Fix bug in FP8 all-gather Signed-off-by: Tim Moon * Fix linter warnings Signed-off-by: Tim Moon * Update JAX to use new swizzled scale API Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Review suggestion from @greptile-apps Signed-off-by: Tim Moon * Review suggestions from @greptile-apps Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update C++ swizzle test with swizzled scales API Signed-off-by: Tim Moon * Return default tensor params when querying params for invalid NVTETensor Signed-off-by: Tim Moon * Debug DSv3 FP8 test failures Signed-off-by: Tim Moon * Debug Userbuffers test failures Signed-off-by: Tim Moon * Make sure gated activations populate FP8 transpose if needed Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Review suggestions from @greptile-apps Signed-off-by: Tim Moon * Disable pre-swizzling with debug quantizer Signed-off-by: Tim Moon * Review suggestion from @greptile-apps Signed-off-by: Tim Moon * Fix merge conflicts and review suggestions Update copyright years. Tweak comments. Fix various complaints from @greptile-apps. Signed-off-by: Tim Moon * Use explicitly sized types in config accessors Miscellaneous review suggestions from @ptrendx. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make util header for function that compute swizzled scale index Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestions 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> * Update expected error message in FP8 block-scaling test Signed-off-by: Tim Moon * Review suggestion from @yaox12 Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon 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: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/cpp/operator/test_swizzle.cu | 1 + tests/cpp/test_common.h | 4 + .../test_float8_blockwise_gemm_exact.py | 2 +- .../test_float8_blockwise_scaling_exact.py | 120 ------ tests/pytorch/test_float8blockwisetensor.py | 15 +- tests/pytorch/test_fusible_ops.py | 145 +++++-- .../common/cast/dispatch/gated.cuh | 29 ++ .../common/cast/dispatch/quantize.cuh | 22 +- .../common/cast/mxfp8/dequantize_mxfp8.cuh | 1 + .../common/cast/mxfp8/gated_mxfp8.cuh | 280 ++++++++----- .../common/cast/mxfp8/quantize_mxfp8.cuh | 370 +++++++++------- .../common/cast/mxfp8/swizzle.cuh | 45 ++ .../common/cast/nvfp4/dequantize_nvfp4.cuh | 1 + .../common/cast/nvfp4/quantize_nvfp4.cuh | 11 +- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 1 + .../comm_gemm_overlap/comm_gemm_overlap.cpp | 10 + transformer_engine/common/common.h | 30 +- transformer_engine/common/gemm/config.cpp | 28 +- transformer_engine/common/gemm/config.h | 8 +- .../common/gemm/cublaslt_gemm.cu | 37 +- .../hadamard_transform_cast_fusion.cu | 1 + .../common/include/transformer_engine/gemm.h | 41 +- .../include/transformer_engine/swizzle.h | 7 +- .../transformer_engine/transformer_engine.h | 137 +++--- .../common/normalization/layernorm/ln_api.cpp | 5 + .../normalization/rmsnorm/rmsnorm_api.cpp | 5 + .../common/recipe/mxfp8_scaling.cu | 2 +- transformer_engine/common/swizzle/swizzle.cu | 9 + .../common/swizzle/swizzle_block_scaling.cu | 15 +- .../common/transformer_engine.cpp | 199 ++++++++- .../common/transpose/cast_transpose.h | 5 +- .../quantize_transpose_square_blockwise.cu | 12 +- .../common/transpose/transpose.cu | 12 +- .../common/transpose/transpose.h | 20 + transformer_engine/common/util/ptx.cuh | 3 + .../common/util/pybind_helper.h | 3 +- .../debug/pytorch/debug_quantization.py | 11 +- .../jax/csrc/extensions/gemm.cpp | 16 +- .../pytorch/cpp_extensions/gemm.py | 10 +- transformer_engine/pytorch/csrc/common.cpp | 4 +- transformer_engine/pytorch/csrc/common.h | 17 +- transformer_engine/pytorch/csrc/extensions.h | 17 +- .../pytorch/csrc/extensions/cast.cpp | 22 +- .../pytorch/csrc/extensions/gemm.cpp | 13 +- .../pytorch/csrc/extensions/normalization.cpp | 42 +- .../pytorch/csrc/extensions/pybind.cpp | 2 + .../pytorch/csrc/extensions/swizzle.cpp | 394 ++++++++++++++++++ transformer_engine/pytorch/csrc/quantizer.cpp | 171 +++----- .../pytorch/csrc/type_converters.cpp | 22 +- transformer_engine/pytorch/csrc/util.cpp | 263 ------------ transformer_engine/pytorch/csrc/util.h | 41 +- transformer_engine/pytorch/distributed.py | 238 +++++------ transformer_engine/pytorch/module/base.py | 3 + .../pytorch/module/grouped_linear.py | 20 +- .../pytorch/module/layernorm_linear.py | 23 +- .../pytorch/module/layernorm_mlp.py | 31 +- transformer_engine/pytorch/module/linear.py | 28 +- .../pytorch/ops/basic/basic_linear.py | 10 +- .../ops/fused/userbuffers_backward_linear.py | 1 + transformer_engine/pytorch/permutation.py | 2 + .../pytorch/quantized_tensor.py | 17 +- .../pytorch/tensor/float8_blockwise_tensor.py | 211 ++++------ .../pytorch/tensor/float8_tensor.py | 1 + .../pytorch/tensor/mxfp8_tensor.py | 129 +++--- .../pytorch/tensor/nvfp4_tensor.py | 13 + .../float8_blockwise_tensor_storage.py | 60 +-- .../tensor/storage/mxfp8_tensor_storage.py | 24 +- .../tensor/storage/nvfp4_tensor_storage.py | 22 +- 68 files changed, 2012 insertions(+), 1502 deletions(-) create mode 100644 transformer_engine/common/cast/mxfp8/swizzle.cuh create mode 100644 transformer_engine/common/transpose/transpose.h create mode 100644 transformer_engine/pytorch/csrc/extensions/swizzle.cpp delete mode 100644 transformer_engine/pytorch/csrc/util.cpp diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 1660ff4e7f..694b348a9b 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -85,6 +85,7 @@ void performTestSwizzle1D(const int num_tiles_M, const int num_tiles_K, bool row std::vector scaling_mode = {SF_MODE_X, SF_MODE_Y, 0}; Tensor input("input", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); Tensor output("output", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); + output.set_with_gemm_swizzled_scales(true); fillUniform(&input); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 42178fec40..b528a79b4f 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -286,6 +286,10 @@ class Tensor { tensor_.set_amax(nullptr, DType::kFloat32, tensor_.defaultShape); } + void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales){ + tensor_.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + } + void to_cpu() const; void from_cpu() const; void set_scale(float scale); diff --git a/tests/pytorch/test_float8_blockwise_gemm_exact.py b/tests/pytorch/test_float8_blockwise_gemm_exact.py index a33f860422..eff571b5cd 100644 --- a/tests/pytorch/test_float8_blockwise_gemm_exact.py +++ b/tests/pytorch/test_float8_blockwise_gemm_exact.py @@ -884,7 +884,7 @@ def test_illegal_2D_by_2D_enforced( is_w_1d_scaled, ) -> None: # 2D block quantization by 2D block quantization is not supported. - expected_err_msg = "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling supported" + expected_err_msg = "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling GEMM is supported" cublas_gemm_test_constraint_enforced( x_dtype, w_dtype, diff --git a/tests/pytorch/test_float8_blockwise_scaling_exact.py b/tests/pytorch/test_float8_blockwise_scaling_exact.py index ab386c0c2f..09f3986ad0 100644 --- a/tests/pytorch/test_float8_blockwise_scaling_exact.py +++ b/tests/pytorch/test_float8_blockwise_scaling_exact.py @@ -87,126 +87,6 @@ def initialize_for_many_scales( return result -@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) -@pytest.mark.parametrize( - "M, N", - [ - # full tile cases - (128, 128), - (256, 256), - (256, 1024), - (1024, 256), - # Padding required cases - (256, 272), - (303, 300), - (305, 256), - # Some larger tiles. - (2000, 2000), - (2048, 2000), - (2000, 1024), - (2048, 1024), - ], -) -@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize("quant_dtype", [torch.float8_e4m3fn, torch.float8_e5m2], ids=str) -@pytest.mark.parametrize("eps", [0], ids=["eps_0"]) -@pytest.mark.parametrize("pow_2_scales", [True], ids=["pow2scales"]) -def test_quantization_1D_block_tiling_with_compact_data_and_scales( - x_dtype: torch.dtype, - M: int, - N: int, - quant_dtype: torch.dtype, - eps: float, - pow_2_scales: bool, -) -> None: - te_dtype = TE_DType[quant_dtype] - tile_size = (1, 128) - # This test runs a comparison of the ref class versus the class using - # CUDA kernels to quantize. They should quantize identically for pixels - # that are not DC values in the scale factor shape. - ref_quantizer = BlockwiseQuantizerReference() - sut_quantizer = Float8BlockQuantizer( - fp8_dtype=te_dtype, - rowwise=True, - columnwise=True, - amax_epsilon=eps, - force_pow_2_scales=pow_2_scales, - block_scaling_dim=1, - all_gather_usage=True, - ) - - # Setup device and random seed - device = "cuda" - seed = 0 - torch.manual_seed(seed) - torch.cuda.manual_seed(seed) - - # Input - x = initialize_for_many_scales((M, N), tile_size, dtype=x_dtype, device=device) - - x_fp8_sut = sut_quantizer.make_empty((M, N), dtype=x_dtype, device=device, requires_grad=False) - x_fp8_sut = sut_quantizer.update_quantized(x, x_fp8_sut) - x_fp8_sut_cpp_alloc = sut_quantizer(x) - - assert x_fp8_sut._rowwise_data is not None - qx: torch.Tensor = x_fp8_sut._rowwise_data.view(dtype=quant_dtype) - assert x_fp8_sut._rowwise_scale_inv is not None - sx: torch.Tensor = x_fp8_sut._rowwise_scale_inv - qx_t = x_fp8_sut._columnwise_data - sx_t = x_fp8_sut._columnwise_scale_inv - - qresult_ref = ref_quantizer.quantize( - x, - quant_dtype=quant_dtype, - return_transpose=True, - eps=eps, - pow_2_scales=pow_2_scales, - quant_tile_shape=tile_size, - munge_scale_shapes=False, - ) - qx_ref, sx_ref, qx_t_ref, sx_t_ref = ( - qresult_ref.data, - qresult_ref.scale, - qresult_ref.data_t, - qresult_ref.scale_t, - ) - - # match the reference quantize transpose output with the columnwise non-transpose method - qx_t_ref = qx_t_ref.transpose(-1, -2).contiguous() - sx_t_ref = sx_t_ref.transpose(-1, -2).contiguous() - - # Check - torch.testing.assert_close(qx.float(), qx_ref.float(), atol=0.0, rtol=0.0) - torch.testing.assert_close(sx, sx_ref, atol=0.0, rtol=0.0) - assert qx_t is not None - qx_t = qx_t.view(dtype=quant_dtype) - assert qx_t_ref is not None - assert sx_t is not None - assert sx_t_ref is not None - torch.testing.assert_close(qx_t.float(), qx_t_ref.float(), atol=0.0, rtol=0.0) - torch.testing.assert_close(sx_t, sx_t_ref, atol=0.0, rtol=0.0) - - # check that the C++ and Python allocators are equivalent - torch.testing.assert_close( - x_fp8_sut._rowwise_data, x_fp8_sut_cpp_alloc._rowwise_data, atol=0.0, rtol=0.0 - ) - torch.testing.assert_close( - x_fp8_sut._rowwise_scale_inv, x_fp8_sut_cpp_alloc._rowwise_scale_inv, atol=0.0, rtol=0.0 - ) - torch.testing.assert_close( - x_fp8_sut._columnwise_data, x_fp8_sut_cpp_alloc._columnwise_data, atol=0.0, rtol=0.0 - ) - torch.testing.assert_close( - x_fp8_sut._columnwise_scale_inv, - x_fp8_sut_cpp_alloc._columnwise_scale_inv, - atol=0.0, - rtol=0.0, - ) - - # check if the fp8 output between C++ and Python are the same - assert x_fp8_sut._data_format == x_fp8_sut_cpp_alloc._data_format - - def check_quantization_block_tiling_versus_reference( x_dtype: torch.dtype, M: int, diff --git a/tests/pytorch/test_float8blockwisetensor.py b/tests/pytorch/test_float8blockwisetensor.py index fe6db9aa41..7add4ee5ab 100644 --- a/tests/pytorch/test_float8blockwisetensor.py +++ b/tests/pytorch/test_float8blockwisetensor.py @@ -175,16 +175,12 @@ def test_quantize_dequantize_columnwise_only( ) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) @pytest.mark.parametrize("dq_columnwise", [True, False]) - @pytest.mark.parametrize("all_gather_usage", [True, False]) def test_quantize_dequantize_dims( self, dims: DimsType, block_scaling_dim: int, dq_columnwise: bool, - all_gather_usage: bool, ) -> None: - if all_gather_usage and block_scaling_dim != 1: - pytest.skip("all_gather_usage only implemented for 1D block quantization.") atol = _tols[tex.DType.kFloat8E4M3]["atol"] rtol = _tols[tex.DType.kFloat8E4M3]["rtol"] quantizer = Float8BlockQuantizer( @@ -192,7 +188,6 @@ def test_quantize_dequantize_dims( rowwise=True, columnwise=dq_columnwise, block_scaling_dim=block_scaling_dim, - all_gather_usage=all_gather_usage, ) self._test_quantize_dequantize( quantizer=quantizer, @@ -218,7 +213,6 @@ def test_quantize_dequantize_compact_format( rowwise=True, columnwise=dq_columnwise, block_scaling_dim=block_scaling_dim, - all_gather_usage=(block_scaling_dim == 1), ) self._test_quantize_dequantize( quantizer=quantizer, @@ -283,13 +277,8 @@ def test_data_accessors(self, dims: DimsType, block_scaling_dim: int) -> None: @pytest.mark.parametrize("dims", [[256, 512], [250, 500]]) @pytest.mark.parametrize("block_scaling_dim", [1, 2]) - @pytest.mark.parametrize("all_gather_usage", [True, False]) - def test_serialization( - self, dims: DimsType, block_scaling_dim: int, all_gather_usage: bool - ) -> None: + def test_serialization(self, dims: DimsType, block_scaling_dim: int) -> None: """Test serialization of Float8BlockwiseQTensor""" - if all_gather_usage and block_scaling_dim != 1: - pytest.skip("all_gather_usage only implemented for 1D block quantization.") device = "cuda" dtype = torch.bfloat16 x_hp = torch.rand(_to_list(dims), dtype=dtype, device=device) @@ -298,7 +287,6 @@ def test_serialization( rowwise=True, columnwise=True, block_scaling_dim=block_scaling_dim, - all_gather_usage=all_gather_usage, ) # Create FP8 tensor @@ -322,7 +310,6 @@ def test_serialization( assert x_fp8_loaded._is_2D_scaled == x_fp8._is_2D_scaled assert x_fp8_loaded.dtype == x_fp8.dtype assert x_fp8_loaded._fp8_dtype == x_fp8._fp8_dtype - assert x_fp8_loaded._data_format == x_fp8._data_format # Test that dequantized values match x_fp8_dequant = x_fp8.dequantize() diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index ce15dd1421..7183e30e71 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2737,7 +2737,11 @@ def test_linear( # Check that original and loaded model match exactly tols = {"rtol": 0, "atol": 0} for param_load, param_save in zip(model_load.parameters(), model_save.parameters()): - torch.testing.assert_close(param_load, param_save, **tols) + torch.testing.assert_close( # Force dequantization by casting to FP64 + param_load.to(dtype=torch.float64, device="cpu"), + param_save.to(dtype=torch.float64, device="cpu"), + **tols, + ) torch.testing.assert_close(param_load.grad, param_save.grad, **tols) for y_load, y_save in zip(ys_load, ys_save): torch.testing.assert_close(y_load, y_save, **tols) @@ -2754,7 +2758,6 @@ def setup_class(cls) -> None: @pytest.mark.parametrize("requires_grad", (False, True)) @pytest.mark.parametrize("bias", (False, True)) - @pytest.mark.parametrize("normalization", ("LayerNorm", "RMSNorm")) @pytest.mark.parametrize("quantized_compute", (False, True)) @pytest.mark.parametrize("quantized_weight", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @@ -2764,25 +2767,18 @@ def test_layernorm_mlp( *, requires_grad: bool, bias: bool, - normalization: str, quantized_compute: bool, quantized_weight: bool, dtype: torch.dtype, quantization: Optional[str], device: torch.device = "cuda", - hidden_size: int = 32, - sequence_length: int = 512, + hidden_size: int = 256, + sequence_length: int = 48, batch_size: int = 4, - ffn_hidden_size: int = 64, + ffn_hidden_size: int = 384, layernorm_epsilon: float = 1e-5, ) -> None: - """ - LayerNorm/RMSNorm + Linear + GELU + Linear - - Note that this test checks only if the module runs - as when chaining multiple modules it is hard to validate - numerical accuracy. - """ + """LayerNorm/RMSNorm + Linear + SwiGLU + Linear""" # Make input shape in_shape = (sequence_length, batch_size, hidden_size) @@ -2798,38 +2794,90 @@ def test_layernorm_mlp( pytest.skip("Quantization scheme is not used") # Random data - _, x_test = make_reference_and_test_tensors( + x_ref, x_test = make_reference_and_test_tensors( in_shape, quantization=quantization, test_dtype=dtype, test_device=device, requires_grad=requires_grad, ) - _, dy_test = make_reference_and_test_tensors( + norm_w_ref, norm_w_test = make_reference_and_test_tensors( + hidden_size, + test_dtype=dtype, + test_device=device, + ) + norm_b_ref, norm_b_test = make_reference_and_test_tensors( + hidden_size, + test_dtype=dtype, + test_device=device, + ) + w1_ref, w1_test = make_reference_and_test_tensors( + (ffn_hidden_size, hidden_size), + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + w2_ref, w2_test = make_reference_and_test_tensors( + (hidden_size, ffn_hidden_size // 2), + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + 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, + test_dtype=dtype, + test_device=device, + ) + b2_ref, b2_test = make_reference_and_test_tensors( + hidden_size, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( in_shape, quantization=quantization, test_dtype=dtype, test_device=device, 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 + x = torch.nn.functional.layer_norm( + x, + (hidden_size,), + weight=norm_w_ref, + bias=norm_b_ref, + eps=layernorm_epsilon, + ) + x = torch.nn.functional.linear(x, w1_ref, bias=b1_ref) + x1, x2 = x.chunk(2, dim=-1) + x = torch.nn.functional.silu(x1) * x2 + x = torch.nn.functional.linear(x, w2_ref, bias=b2_ref) + y_ref = x + y_ref.backward(dy_ref) - # Implementation with fusible operations + # Construct operations recipe = make_recipe(quantization) with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): - if normalization == "LayerNorm": - norm = te_ops.LayerNorm( - hidden_size, - eps=layernorm_epsilon, - device=device, - dtype=dtype, - ) - else: - norm = te_ops.RMSNorm( - hidden_size, - eps=layernorm_epsilon, - device=device, - dtype=dtype, - ) + norm = te_ops.LayerNorm( + hidden_size, + eps=layernorm_epsilon, + device=device, + dtype=dtype, + ) ffn1 = te_ops.Linear( hidden_size, ffn_hidden_size, @@ -2837,15 +2885,48 @@ def test_layernorm_mlp( device=device, dtype=dtype, ) - act = te_ops.GELU() + act = te_ops.SwiGLU() ffn2 = te_ops.Linear( - ffn_hidden_size, + ffn_hidden_size // 2, hidden_size, bias=bias, device=device, dtype=dtype, ) + + # Copy weights + with torch.no_grad(): + norm.weight.copy_(norm_w_test) + norm.bias.copy_(norm_b_test) + ffn1.weight.copy_(w1_test) + ffn2.weight.copy_(w2_test) + if bias: + ffn1.bias.copy_(b1_test) + ffn2.bias.copy_(b2_test) + del norm_w_test, norm_b_test, w1_test, b1_test, w2_test, b2_test + + # Fuse ops and perform forward and backward pass forward = te_ops.Sequential(norm, ffn1, act, ffn2) with te.autocast(enabled=quantized_compute, recipe=recipe): 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) + 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) diff --git a/transformer_engine/common/cast/dispatch/gated.cuh b/transformer_engine/common/cast/dispatch/gated.cuh index 540912c5af..06e8f0e306 100644 --- a/transformer_engine/common/cast/dispatch/gated.cuh +++ b/transformer_engine/common/cast/dispatch/gated.cuh @@ -14,6 +14,7 @@ #include #include "../../common.h" +#include "../../transpose/transpose.h" #include "../../utils.cuh" #include "../fp8/gated_fp8.cuh" #include "../mxfp8/gated_mxfp8.cuh" @@ -53,6 +54,20 @@ void quantize_gated_fwd_helper(const NVTETensor nvte_input, NVTETensor nvte_outp } else { fp8::cast_gated_fwd(input, output, p, stream); } + if (is_fp8_dtype(output->dtype()) && output->has_columnwise_data()) { + // FP8 kernel only populates row-wise data, so perform + // transpose separately if needed + Tensor transpose_in, transpose_out, dummy; + transpose_in.scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + transpose_in.data.dptr = output->data.dptr; + transpose_in.data.shape = {output->flat_first_dim(), output->flat_last_dim()}; + transpose_in.data.dtype = output->data.dtype; + transpose_out.scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + transpose_out.data.dptr = output->columnwise_data.dptr; + transpose_out.data.shape = {output->flat_last_dim(), output->flat_first_dim()}; + transpose_out.data.dtype = output->data.dtype; + detail::transpose(transpose_in, /*noop=*/dummy, &transpose_out, stream); + } break; } case NVTE_MXFP8_1D_SCALING: { @@ -129,6 +144,20 @@ void quantize_gated_bwd_helper(const NVTETensor nvte_grad, const NVTETensor nvte } else { fp8::cast_gated_bwd(gated_input, grad, output, p, stream); } + if (is_fp8_dtype(output->dtype()) && output->has_columnwise_data()) { + // FP8 kernel only populates row-wise data, so perform + // transpose separately if needed + Tensor transpose_in, transpose_out, dummy; + transpose_in.scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + transpose_in.data.dptr = output->data.dptr; + transpose_in.data.shape = {output->flat_first_dim(), output->flat_last_dim()}; + transpose_in.data.dtype = output->data.dtype; + transpose_out.scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + transpose_out.data.dptr = output->columnwise_data.dptr; + transpose_out.data.shape = {output->flat_last_dim(), output->flat_first_dim()}; + transpose_out.data.dtype = output->data.dtype; + detail::transpose(transpose_in, /*noop=*/dummy, &transpose_out, stream); + } break; } case NVTE_MXFP8_1D_SCALING: { diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 8453b9a68b..a02e7f4f07 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -150,17 +150,10 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, FP8BlockwiseRowwiseOption rowwise_option = FP8BlockwiseRowwiseOption::NONE; FP8BlockwiseColumnwiseOption columnwise_option = FP8BlockwiseColumnwiseOption::NONE; if (output_tensor->has_data()) { - bool rowwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT); - rowwise_option = rowwise_compact ? FP8BlockwiseRowwiseOption::ROWWISE_COMPACT - : FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; + rowwise_option = FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; } if (output_tensor->has_columnwise_data()) { - bool columnwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT); - columnwise_option = columnwise_compact - ? FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT - : FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; + columnwise_option = FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; } quantize_transpose_vector_blockwise( input_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, @@ -298,17 +291,10 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens FP8BlockwiseRowwiseOption rowwise_option = FP8BlockwiseRowwiseOption::NONE; FP8BlockwiseColumnwiseOption columnwise_option = FP8BlockwiseColumnwiseOption::NONE; if (output_tensor->has_data()) { - bool rowwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT); - rowwise_option = rowwise_compact ? FP8BlockwiseRowwiseOption::ROWWISE_COMPACT - : FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; + rowwise_option = FP8BlockwiseRowwiseOption::ROWWISE_GEMM_READY; } if (output_tensor->has_columnwise_data()) { - bool columnwise_compact = (quant_config_cpp.float8_block_scale_tensor_format == - Float8BlockScaleTensorFormat::COMPACT); - columnwise_option = columnwise_compact - ? FP8BlockwiseColumnwiseOption::COLUMNWISE_COMPACT - : FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; + columnwise_option = FP8BlockwiseColumnwiseOption::COLUMNWISE_GEMM_READY; } quantize_transpose_vector_blockwise( grad_tensor->data, output_tensor->scale_inv, output_tensor->columnwise_scale_inv, diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index ecdbb5c657..f8fecaa4e1 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -239,6 +239,7 @@ 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."); diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 3f5c44120e..dc9a190e1f 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_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 { @@ -51,7 +52,8 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 template + bool ROWWISE_SCALING, bool COLWISE_SCALING, bool WITH_GEMM_SWIZZLED_SCALES, + size_t THREADS_PER_CHUNK> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_gated_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_grad, const __grid_constant__ CUtensorMap tensor_map_input_act, @@ -68,6 +70,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) using IType2 = typename ptx::FPx2; using OType2 = typename ptx::FPx2; + using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; + constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; static_assert(STAGES >= 1); @@ -355,14 +359,17 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // 2. Compute E8M0 scaling factor const e8m0_t biased_exponent_act = ptx::float_to_e8m0(thread_amax_act * Quantized_Limits::max_norm_rcp); - const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; const size_t global_scales_offset_X = scales_offset_X_colwise; - const size_t scale_idx = - global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = gemm_swizzled_scale_idx(global_scales_offset_X, global_scales_offset_Y, + DIVUP(rows, static_cast(128))); + } else { + scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + } const bool row_out_of_bounds_colwise = (row_base_colwise + stage_offset_Y) >= rows; const bool out_of_bounds_colwise = row_out_of_bounds_colwise || col_out_of_bounds_colwise; - if (tid_Y_colwise == 0 && (!out_of_bounds_colwise)) { scales_colwise[scale_idx] = biased_exponent_act; } @@ -374,8 +381,14 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const e8m0_t biased_exponent_gate = ptx::float_to_e8m0(thread_amax_gate * Quantized_Limits::max_norm_rcp); - // const size_t scale_idx_gate = scale_idx + scale_stride_colwise / 2; - const size_t scale_idx_gate = scale_idx + gate_scale_idx_offset_colwise; + size_t scale_idx_gate; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx_gate = gemm_swizzled_scale_idx( + global_scales_offset_X + gate_scale_idx_offset_colwise, global_scales_offset_Y, + DIVUP(rows, static_cast(128))); + } else { + scale_idx_gate = scale_idx + gate_scale_idx_offset_colwise; + } if (tid_Y_colwise == 0 && (!out_of_bounds_colwise)) { scales_colwise[scale_idx_gate] = biased_exponent_gate; } @@ -557,7 +570,14 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) ptx::float_to_e8m0(thread_amax_act * Quantized_Limits::max_norm_rcp); const size_t stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const size_t stage_scales_offset_X = scales_offset_X_rowwise; - const size_t scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + const size_t output_cols = (IS_BWD ? 2 : 1) * cols; + scale_idx = gemm_swizzled_scale_idx(stage_scales_offset_Y, stage_scales_offset_X, + DIVUP(output_cols, static_cast(128))); + } else { + scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + } const bool row_out_of_bounds_rowwise = (row_base_rowwise + stage_offset_Y) >= rows; const bool out_of_bounds_rowwise = row_out_of_bounds_rowwise || col_out_of_bounds_rowwise; if (!out_of_bounds_rowwise) { @@ -573,7 +593,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if constexpr (IS_BWD) { const e8m0_t biased_exponent_gate = ptx::float_to_e8m0(thread_amax_gate * Quantized_Limits::max_norm_rcp); - const size_t scale_idx_gate = scale_idx + gate_scale_idx_offset_rowwise; + + size_t scale_idx_gate; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + const size_t output_cols = (IS_BWD ? 2 : 1) * cols; + scale_idx_gate = gemm_swizzled_scale_idx( + stage_scales_offset_Y, stage_scales_offset_X + gate_scale_idx_offset_rowwise, + DIVUP(output_cols, static_cast(128))); + } else { + scale_idx_gate = scale_idx + gate_scale_idx_offset_rowwise; + } if (!out_of_bounds_rowwise) { scales_rowwise[scale_idx_gate] = biased_exponent_gate; } @@ -667,7 +696,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) parity ^= 1; destroy_barriers(mbar, is_master_thread); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} +} // NOLINT(readability/fn_size) + } // namespace gated_kernel template has_data(); const bool USE_COLWISE_SCALING = output->has_columnwise_data(); + const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; if (USE_ROWWISE_SCALING) { NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); @@ -722,113 +753,140 @@ void quantize_gated(const Tensor &gated_input, const Tensor &grad, Tensor *outpu gated_input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, - alignas(64) CUtensorMap tensor_map_grad{}; - alignas(64) CUtensorMap tensor_map_input_act{}; - alignas(64) CUtensorMap tensor_map_input_gate{}; - alignas(64) CUtensorMap tensor_map_output_act_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_gate_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_act_colwise{}; - alignas(64) CUtensorMap tensor_map_output_gate_colwise{}; + alignas(64) CUtensorMap tensor_map_grad{}; + alignas(64) CUtensorMap tensor_map_input_act{}; + alignas(64) CUtensorMap tensor_map_input_gate{}; + alignas(64) CUtensorMap tensor_map_output_act_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_gate_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_act_colwise{}; + alignas(64) CUtensorMap tensor_map_output_gate_colwise{}; - constexpr size_t input_type_bit_size = TypeInfo::size; - constexpr size_t output_type_bit_size = TypeInfo::size; + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; - if constexpr (IS_BWD) { - create_2D_tensor_map(tensor_map_grad, grad.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, - cols, 0, input_type_bit_size); - } + if constexpr (IS_BWD) { + create_2D_tensor_map(tensor_map_grad, grad.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, + cols, 0, input_type_bit_size); + } - const uint32_t tensor_stride_elems = output_cols; - create_2D_tensor_map(tensor_map_input_act, gated_input.data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols * 2, 0, input_type_bit_size); - create_2D_tensor_map(tensor_map_input_gate, gated_input.data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols * 2, cols, input_type_bit_size); - - if (USE_ROWWISE_SCALING) { - create_2D_tensor_map(tensor_map_output_act_rowwise, output->data, rows, cols, - BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, 0, - output_type_bit_size); - create_2D_tensor_map(tensor_map_output_gate_rowwise, output->data, rows, cols, - BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, cols, - output_type_bit_size); - } + const uint32_t tensor_stride_elems = output_cols; + create_2D_tensor_map(tensor_map_input_act, gated_input.data, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols * 2, 0, input_type_bit_size); + create_2D_tensor_map(tensor_map_input_gate, gated_input.data, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols * 2, cols, input_type_bit_size); + + if (USE_ROWWISE_SCALING) { + create_2D_tensor_map(tensor_map_output_act_rowwise, output->data, rows, cols, + BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, 0, + output_type_bit_size); + create_2D_tensor_map(tensor_map_output_gate_rowwise, output->data, rows, cols, + BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, cols, + output_type_bit_size); + } - if (USE_COLWISE_SCALING) { - create_2D_tensor_map(tensor_map_output_act_colwise, output->columnwise_data, rows, cols, - BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, 0, - output_type_bit_size); - create_2D_tensor_map(tensor_map_output_gate_colwise, output->columnwise_data, rows, - cols, BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, cols, - output_type_bit_size); - } + if (USE_COLWISE_SCALING) { + create_2D_tensor_map(tensor_map_output_act_colwise, output->columnwise_data, rows, + cols, BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, 0, + output_type_bit_size); + create_2D_tensor_map(tensor_map_output_gate_colwise, output->columnwise_data, rows, + cols, BUFF_DIM_Y, BUFF_DIM_X, tensor_stride_elems, cols, + output_type_bit_size); + } - const size_t buff_elems_total = BUFFS_NUM * BUFF_DIM_Y * BUFF_DIM_X; - const size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; - const size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; - const size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); - const size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); - - const size_t grad_mem = (IS_BWD ? buff_size_aligned_in : 0); - const size_t in_act_mem = buff_size_aligned_in; - const size_t in_gate_mem = buff_size_aligned_in; - const size_t in_mem = grad_mem + in_act_mem + in_gate_mem; - - const size_t out_act_mem = buff_size_aligned_out; - const size_t out_gate_mem = (IS_BWD ? buff_size_aligned_out : 0); - size_t out_mem = out_act_mem + out_gate_mem; - - if (USE_ROWWISE_SCALING && USE_COLWISE_SCALING) { out_mem *= 2; } - - const size_t shmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - - switch (scaling_type) { - case ScalingType::ROWWISE: { - auto kernel = - quantize_gated_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - - kernel<<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, - tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, - tensor_map_output_act_colwise, tensor_map_output_gate_colwise, scales_rowwise_ptr, - scales_colwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise, p); - break; - } - case ScalingType::COLWISE: { - auto kernel = - quantize_gated_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - - kernel<<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, - tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, - tensor_map_output_act_colwise, tensor_map_output_gate_colwise, scales_rowwise_ptr, - scales_colwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise, p); - break; - } - case ScalingType::BIDIMENSIONAL: { - auto kernel = - quantize_gated_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); - - kernel<<>>( - tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, - tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, - tensor_map_output_act_colwise, tensor_map_output_gate_colwise, scales_rowwise_ptr, - scales_colwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise, p); - break; - } - } NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) - ); // NOLINT(*) + const size_t buff_elems_total = BUFFS_NUM * BUFF_DIM_Y * BUFF_DIM_X; + const size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; + const size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; + const size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + const size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + const size_t grad_mem = (IS_BWD ? buff_size_aligned_in : 0); + const size_t in_act_mem = buff_size_aligned_in; + const size_t in_gate_mem = buff_size_aligned_in; + const size_t in_mem = grad_mem + in_act_mem + in_gate_mem; + + const size_t out_act_mem = buff_size_aligned_out; + const size_t out_gate_mem = (IS_BWD ? buff_size_aligned_out : 0); + size_t out_mem = out_act_mem + out_gate_mem; + + if (USE_ROWWISE_SCALING && USE_COLWISE_SCALING) { out_mem *= 2; } + + const size_t shmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; + + // Zero out swizzled scales if padding is needed + /// TODO (tmoon) Handle this within the cast kernel + if (with_gemm_swizzled_scales) { + constexpr size_t TILE_DIM_X = 128; // Tile dim in data buffer + constexpr size_t TILE_DIM_Y = 128; + if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) { + if (USE_ROWWISE_SCALING) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, + output->scale_inv.buffer_size_bytes(), stream)); + } + if (USE_COLWISE_SCALING) { + NVTE_CHECK_CUDA( + cudaMemsetAsync(output->columnwise_scale_inv.dptr, 0, + output->columnwise_scale_inv.buffer_size_bytes(), stream)); + } + } + } + + switch (scaling_type) { + case ScalingType::ROWWISE: { + auto kernel = + quantize_gated_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); + + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, + tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, + tensor_map_output_act_colwise, tensor_map_output_gate_colwise, + scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise, p); + break; + } + case ScalingType::COLWISE: { + auto kernel = + quantize_gated_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); + + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, + tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, + tensor_map_output_act_colwise, tensor_map_output_gate_colwise, + scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise, p); + break; + } + case ScalingType::BIDIMENSIONAL: { + auto kernel = + quantize_gated_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem_size)); + + kernel<<>>( + tensor_map_grad, tensor_map_input_act, tensor_map_input_gate, + tensor_map_output_act_rowwise, tensor_map_output_gate_rowwise, + tensor_map_output_act_colwise, tensor_map_output_gate_colwise, + scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise, p); + break; + } + } NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) } } // namespace mxfp8 diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index d6aae78d25..70a68132ad 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -22,6 +22,7 @@ #include "../../utils.cuh" #include "../core/common.cuh" #include "specialized/quantize_mxfp8.cuh" +#include "swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -43,7 +44,8 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 template + bool COLWISE_SCALING, bool WITH_GEMM_SWIZZLED_SCALES, size_t CHUNK_DIM_Y, + size_t CHUNK_DIM_X, size_t THREADS_PER_CHUNK> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_act_input, @@ -60,6 +62,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) using IType2 = typename ptx::FPx2; using OType2 = typename ptx::FPx2; + using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; + if constexpr (NO_ACTIVATIONS) { if (noop != nullptr && noop[0] == 1.0f) { return; @@ -106,7 +110,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; const size_t 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 rowwise_scale_is_within_bounds = SCALE_DIM_X * scales_offset_X_rowwise < cols; // helps resolving bank conflicts in shmem const int thread_lane = threadIdx.x % THREADS_PER_WARP; @@ -263,11 +267,15 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // 2. Compute E8M0 scaling factor const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; const size_t global_scales_offset_X = scales_offset_X_colwise; - const size_t scale_idx = - global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = gemm_swizzled_scale_idx(global_scales_offset_X, global_scales_offset_Y, + DIVUP(rows, static_cast(128))); + } else { + scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + } scales_colwise[scale_idx] = biased_exponent; const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); @@ -411,7 +419,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const int stage_scales_offset_X = scales_offset_X_rowwise; - const int scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = gemm_swizzled_scale_idx(stage_scales_offset_Y, stage_scales_offset_X, + DIVUP(cols, static_cast(128))); + } 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; } @@ -550,7 +564,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, bool use_colwise_scaling = output->has_columnwise_data(); NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); - if (use_rowwise_scaling) { NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); } @@ -560,17 +573,21 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, } CheckNoopTensor(*noop, "cast_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(); - constexpr bool CAST_DBIAS_ONLY = IS_DBIAS && (!IS_DACT) && (!IS_ACT); - + // Tensor chunk handled by each CUDA block constexpr size_t CHUNK_DIM_Y = CAST_DBIAS_ONLY ? 128 : 64; constexpr size_t CHUNK_DIM_X = CAST_DBIAS_ONLY ? 128 : 64; - constexpr size_t THREADS_PER_CHUNK = CAST_DBIAS_ONLY ? 128 : 64; + // CUDA block config + constexpr size_t THREADS_PER_CHUNK = CAST_DBIAS_ONLY ? 128 : 64; constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; + constexpr size_t BUFF_DIM_Y = THREADS_Y; constexpr size_t BUFF_DIM_X = CHUNK_DIM_X; @@ -579,6 +596,8 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, const dim3 grid(blocks_X, blocks_Y); const size_t block_size = THREADS_PER_CHUNK; + const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + const size_t scale_stride_rowwise = use_rowwise_scaling ? output->scale_inv.shape[1] : 1; const size_t scale_stride_colwise = use_colwise_scaling ? output->columnwise_scale_inv.shape[1] : 1; @@ -619,168 +638,195 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( output->dtype(), OType, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + + if (specialized::hasSpec() && + !WITH_GEMM_SWIZZLED_SCALES) { + switch (scaling_type) { + case ScalingType::ROWWISE: { + using traits = specialized::CastTraits; + auto kernel = specialized::quantize_mxfp8_kernel_cast_only; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + traits::smem); + + dim3 block(traits::threadLayout::num, traits::warpLayout::N, + traits::warpLayout::M); + dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN, + (rows + traits::blockDimM - 1) / traits::blockDimM); + kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output->data.dptr), + scales_rowwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); + + 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); + // 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; + create_2D_tensor_map(tensor_map_input, input.data, rows, cols, + traits::blockIterDim::M, traits::blockIterDim::N, + /*stride_elems=*/cols, + /*offset_elems=*/0, input_type_bit_size, + traits::input_swizzle_pattern); + + alignas(64) CUtensorMap tensor_map_rowwise_output{}; + alignas(64) CUtensorMap tensor_map_colwise_output{}; + constexpr size_t output_type_bit_size = TypeInfo::size; + create_2D_tensor_map(tensor_map_rowwise_output, output->data, rows, cols, + traits::blockIterDim::M, traits::blockIterDim::N, + /*stride_elems=*/cols, + /*offset_elems=*/0, output_type_bit_size, + traits::output_swizzle_pattern); + create_2D_tensor_map(tensor_map_colwise_output, output->columnwise_data, rows, + cols, traits::blockIterDim::M, traits::blockIterDim::N, + cols, 0, output_type_bit_size, + traits::output_swizzle_pattern); + + dim3 block(traits::rowThreadLayout::num, traits::numWarps); + dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N, + (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); + kernel<<>>( + tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, + scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + + break; + } + default: { + NVTE_ERROR("Invalid scaling type."); + } + } + return; + } - if (specialized::hasSpec()) { - switch (scaling_type) { - case ScalingType::ROWWISE: { - using traits = specialized::CastTraits; - auto kernel = specialized::quantize_mxfp8_kernel_cast_only; + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_act_input{}; + alignas(64) CUtensorMap tensor_map_output_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_colwise{}; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - traits::smem); + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; - dim3 block(traits::threadLayout::num, traits::warpLayout::N, traits::warpLayout::M); - dim3 grid((cols + traits::blockDimN - 1) / traits::blockDimN, - (rows + traits::blockDimM - 1) / traits::blockDimM); - kernel<<>>( - reinterpret_cast(input.data.dptr), - reinterpret_cast(output->data.dptr), - scales_rowwise_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); + create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, + cols, 0, input_type_bit_size); - 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); - // 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; - create_2D_tensor_map(tensor_map_input, input.data, rows, cols, - traits::blockIterDim::M, traits::blockIterDim::N, - /*stride_elems=*/cols, - /*offset_elems=*/0, input_type_bit_size, - traits::input_swizzle_pattern); - - alignas(64) CUtensorMap tensor_map_rowwise_output{}; - alignas(64) CUtensorMap tensor_map_colwise_output{}; - constexpr size_t output_type_bit_size = TypeInfo::size; - create_2D_tensor_map(tensor_map_rowwise_output, output->data, rows, cols, - traits::blockIterDim::M, traits::blockIterDim::N, - /*stride_elems=*/cols, - /*offset_elems=*/0, output_type_bit_size, - traits::output_swizzle_pattern); - create_2D_tensor_map(tensor_map_colwise_output, output->columnwise_data, rows, cols, - traits::blockIterDim::M, traits::blockIterDim::N, cols, 0, - output_type_bit_size, traits::output_swizzle_pattern); - - dim3 block(traits::rowThreadLayout::num, traits::numWarps); - dim3 grid((cols + traits::blockDIM::N - 1) / traits::blockDIM::N, - (rows + traits::blockDIM::M - 1) / traits::blockDIM::M); - kernel<<>>( - tensor_map_input, tensor_map_rowwise_output, tensor_map_colwise_output, - scales_rowwise_ptr, scales_colwise_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - - break; - } - default: { - NVTE_ERROR("Invalid scaling type."); + if constexpr (IS_DACT) { + create_2D_tensor_map(tensor_map_act_input, act_input->data, rows, cols, BUFF_DIM_Y, + BUFF_DIM_X, cols, 0, input_type_bit_size); } - } - return; - } - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_act_input{}; - alignas(64) CUtensorMap tensor_map_output_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_colwise{}; - - constexpr size_t input_type_bit_size = TypeInfo::size; - constexpr size_t output_type_bit_size = TypeInfo::size; - - create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, - cols, 0, input_type_bit_size); - if constexpr (IS_DACT) { - create_2D_tensor_map(tensor_map_act_input, act_input->data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols, 0, input_type_bit_size); - } + if (use_rowwise_scaling) { + create_2D_tensor_map(tensor_map_output_rowwise, output->data, rows, cols, + BUFF_DIM_Y, BUFF_DIM_X, cols, 0, output_type_bit_size); + } - if (use_rowwise_scaling) { - create_2D_tensor_map(tensor_map_output_rowwise, output->data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols, 0, output_type_bit_size); - } + 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, output_type_bit_size); + } - 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, output_type_bit_size); - } + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; + constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); + const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); + const size_t out_mem = out_rowwise_mem + out_colwise_mem; + + const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; + + // Zero out swizzled scales if padding is needed + /// TODO (tmoon) Handle this within the cast kernel + if (with_gemm_swizzled_scales) { + constexpr size_t TILE_DIM_X = 128; // Tile dim in data buffer + constexpr size_t TILE_DIM_Y = 128; + if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) { + if (use_rowwise_scaling) { + NVTE_CHECK_CUDA(cudaMemsetAsync(output->scale_inv.dptr, 0, + output->scale_inv.buffer_size_bytes(), stream)); + } + if (use_colwise_scaling) { + NVTE_CHECK_CUDA( + cudaMemsetAsync(output->columnwise_scale_inv.dptr, 0, + output->columnwise_scale_inv.buffer_size_bytes(), stream)); + } + } + } - constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; - constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; - constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; - constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; - constexpr size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); - - constexpr size_t elt_input_mem = buff_size_aligned_in; - constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); - constexpr size_t in_mem = elt_input_mem + act_input_mem; - - const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); - const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); - const size_t out_mem = out_rowwise_mem + out_colwise_mem; - - const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - - switch (scaling_type) { - case ScalingType::ROWWISE: { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - 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: { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - 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: { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - 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; - } - } + switch (scaling_type) { + case ScalingType::ROWWISE: { + auto kernel = quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + 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: { + auto kernel = quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + 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: { + auto kernel = quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + 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; + } + } - if constexpr (IS_DBIAS) { - common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); - }); // NOLINT(*) - ); // NOLINT(*) + if constexpr (IS_DBIAS) { + common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); + }); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) } } // namespace mxfp8 diff --git a/transformer_engine/common/cast/mxfp8/swizzle.cuh b/transformer_engine/common/cast/mxfp8/swizzle.cuh new file mode 100644 index 0000000000..7648e3f5cb --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/swizzle.cuh @@ -0,0 +1,45 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file swizzle.cuh + * \brief Helper function for GEMM-swizzled scales + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_SWIZZLE_CUH_ +#define TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_SWIZZLE_CUH_ + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace swizzle { + +/*! \brief Convert compact scale indices into GEMM swizzled scale index + * + * MXFP8 GEMM expects scaling factors to be in a "swizzled" order + * (https://docs.nvidia.com/cuda/cublas/#d-block-scaling-factors-layout). + * This function converts indices from "compact" order (i.e. matching + * the FP8 data) to swizzled order. + * + */ +__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_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; + const size_t idx_in_tile_X = j % TILE_DIM_X; + const size_t idx_in_tile_Y = i % TILE_DIM_Y; + size_t idx = (tile_idx_Y * num_tiles_X + tile_idx_X) * TILE_SIZE; + idx += (idx_in_tile_Y % 32) * 16 + (idx_in_tile_Y / 32) * 4 + idx_in_tile_X; + return idx; +} + +} // namespace swizzle +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_CAST_MXFP8_SWIZZLE_CUH_ diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 38677a7075..ccdc4c93e3 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -80,6 +80,7 @@ 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."); diff --git a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh index b4bccf2397..e7854ffde3 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh @@ -142,17 +142,10 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) 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 / SCALE_DIM_X) * sizeof(fp8e4m3); - constexpr size_t buff_size_mxfp8_scales = - (CHUNK_DIM_Y / SCALE_DIM_Y) * CHUNK_DIM_X * sizeof(fp8e8m0); - 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); - constexpr size_t out_mem_rowwise_scales = (ROWWISE_SCALING ? buff_size_nvfp4_scales : 0); - constexpr size_t out_mem_colwise_scales = (COLWISE_SCALING ? buff_size_mxfp8_scales : 0); extern __shared__ char dynamic_shmem[]; uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); @@ -167,8 +160,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) 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); - e8m0_t *out_colwise_scales_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + (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; @@ -557,6 +549,7 @@ inline void quantize(const Tensor &input, const Tensor *noop, Tensor *output, cu 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) { diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 455074e325..5da9cc5a5b 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -1179,6 +1179,7 @@ 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(!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."); NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), 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 2a3c64e8dd..aad2ec0686 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -172,7 +172,17 @@ CommOverlapCore::~CommOverlapCore() { TensorWrapper CommOverlapCore::get_tensor_chunk(const TensorWrapper &source, size_t chunk_offset, const std::vector &chunk_shape) { + // Check tensor format const auto scaling_mode = source.scaling_mode(); + NVTE_CHECK(scaling_mode == NVTE_DELAYED_TENSOR_SCALING || scaling_mode == NVTE_MXFP8_1D_SCALING, + "Unsupported tensor format (", to_string(scaling_mode), ")."); + if (scaling_mode == NVTE_MXFP8_1D_SCALING) { + uint8_t has_swizzled_scales = false; + nvte_get_tensor_param_v2(source.data(), NVTETensorParam::kNVTEWithGEMMSwizzledScales, + &has_swizzled_scales, sizeof(has_swizzled_scales), nullptr); + NVTE_CHECK(has_swizzled_scales, + "Expected MXFP8 tensor to have scales in GEMM swizzled format."); + } // Tensor dimensions std::vector shape = shape_to_vector(source.shape()); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 0bc9536844..970b7aef6c 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -133,6 +133,23 @@ struct Tensor { NVTEScalingMode scaling_mode; NVTETensor nvte_tensor; + /*! \brief Whether scaling factors are in format expected by GEMM + * + * Only meaningful for MXFP8 and NVFP4. + */ + bool with_gemm_swizzled_scales = false; + + /*! Map from NVTETensorParam to parameter sizes */ + static constexpr size_t attr_sizes[] = { + sizeof(NVTEBasicTensor), // kNVTERowwiseData + sizeof(NVTEBasicTensor), // kNVTEColumnwiseData + sizeof(NVTEBasicTensor), // kNVTEScale + sizeof(NVTEBasicTensor), // kNVTEAmax + sizeof(NVTEBasicTensor), // kNVTERowwiseScaleInv + sizeof(NVTEBasicTensor), // kNVTEColumnwiseScaleInv + sizeof(NVTEBasicTensor), // kNVTEColumnwiseAmax + sizeof(uint8_t) // kNVTEWithGEMMSwizzledScales + }; Tensor() : scaling_mode{NVTE_DELAYED_TENSOR_SCALING}, nvte_tensor{0} {} @@ -146,6 +163,7 @@ struct Tensor { scale_inv.clear(); columnwise_scale_inv.clear(); scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + with_gemm_swizzled_scales = false; } explicit operator NVTETensor() const noexcept { return nvte_tensor; } @@ -389,22 +407,20 @@ struct QuantizationConfig { bool force_pow_2_scales = false; float amax_epsilon = 0.0f; NVTETensor noop_tensor = nullptr; - Float8BlockScaleTensorFormat float8_block_scale_tensor_format = - Float8BlockScaleTensorFormat::GEMM_READY; NVTETensor rng_state = nullptr; bool nvfp4_2d_quantization = false; bool stochastic_rounding = false; bool use_fast_math = false; static constexpr size_t attr_sizes[] = { - sizeof(bool), // force_pow_2_scales + sizeof(uint8_t), // force_pow_2_scales sizeof(float), // amax_epsilon sizeof(NVTETensor), // noop_tensor - sizeof(Float8BlockScaleTensorFormat), // float8_block_scale_tensor_format + sizeof(Float8BlockScaleTensorFormat), // (deprecated) sizeof(NVTETensor), // rng_seed and offset - sizeof(bool), // nvfp4_2d_quantization - sizeof(bool), // stochastic_rounding - sizeof(bool) // use_fast_math + sizeof(uint8_t), // nvfp4_2d_quantization + sizeof(uint8_t), // stochastic_rounding + sizeof(uint8_t) // use_fast_math }; }; diff --git a/transformer_engine/common/gemm/config.cpp b/transformer_engine/common/gemm/config.cpp index ec2381dc1e..2532e96bb8 100644 --- a/transformer_engine/common/gemm/config.cpp +++ b/transformer_engine/common/gemm/config.cpp @@ -36,6 +36,12 @@ void nvte_get_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, " bytes)"); + // bool size is implementation-dependent, so we explicitly specify + // uint8_t in the user-facing API. + auto bool_to_uint8 = [](bool in, void *out) { + *reinterpret_cast(out) = static_cast(in); + }; + // Write to buffer NVTE_CHECK(config != nullptr, "Invalid NVTEMatmulConfig (got NULL)"); const auto &config_ = *reinterpret_cast(config); @@ -47,19 +53,19 @@ void nvte_get_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA std::memcpy(buf, &config_.dbias_tensor, attr_size); break; case kNVTEMatmulConfigWithGELUEpilogue: - std::memcpy(buf, &config_.with_gelu_epilogue, attr_size); + bool_to_uint8(config_.with_gelu_epilogue, buf); break; case kNVTEMatmulConfigWithDGELUEpilogue: - std::memcpy(buf, &config_.with_dgelu_epilogue, attr_size); + bool_to_uint8(config_.with_dgelu_epilogue, buf); break; case kNVTEMatmulConfigEpilogueAuxTensor: std::memcpy(buf, &config_.epilogue_aux_tensor, attr_size); break; case kNVTEMatmulConfigUseSplitAccumulator: - std::memcpy(buf, &config_.use_split_accumulator, attr_size); + bool_to_uint8(config_.use_split_accumulator, buf); break; case kNVTEMatmulConfigSMCount: - std::memcpy(buf, &config_.sm_count, attr_size); + *reinterpret_cast(buf) = static_cast(config_.sm_count); break; default: NVTE_ERROR("Unsupported NVTEMatmulConfigAttribute (got ", static_cast(attr), ")"); @@ -79,6 +85,12 @@ void nvte_set_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA " bytes)"); NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + // bool size is implementation-dependent, so we explicitly specify + // uint8_t in the user-facing API. + auto uint8_to_bool = [](const void *in, bool &out) { + out = static_cast(*reinterpret_cast(in)); + }; + // Read from buffer NVTE_CHECK(config != nullptr, "Invalid NVTEMatmulConfig (got NULL)"); auto &config_ = *reinterpret_cast(config); @@ -90,19 +102,19 @@ void nvte_set_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA std::memcpy(&config_.dbias_tensor, buf, attr_size); break; case kNVTEMatmulConfigWithGELUEpilogue: - std::memcpy(&config_.with_gelu_epilogue, buf, attr_size); + uint8_to_bool(buf, config_.with_gelu_epilogue); break; case kNVTEMatmulConfigWithDGELUEpilogue: - std::memcpy(&config_.with_dgelu_epilogue, buf, attr_size); + uint8_to_bool(buf, config_.with_dgelu_epilogue); break; case kNVTEMatmulConfigEpilogueAuxTensor: std::memcpy(&config_.epilogue_aux_tensor, buf, attr_size); break; case kNVTEMatmulConfigUseSplitAccumulator: - std::memcpy(&config_.use_split_accumulator, buf, attr_size); + uint8_to_bool(buf, config_.use_split_accumulator); break; case kNVTEMatmulConfigSMCount: - std::memcpy(&config_.sm_count, buf, attr_size); + config_.sm_count = static_cast(*reinterpret_cast(buf)); break; default: NVTE_ERROR("Unsupported NVTEMatmulConfigAttribute (got ", static_cast(attr), ")"); diff --git a/transformer_engine/common/gemm/config.h b/transformer_engine/common/gemm/config.h index c7b528bf4b..86a617b5fe 100644 --- a/transformer_engine/common/gemm/config.h +++ b/transformer_engine/common/gemm/config.h @@ -23,11 +23,11 @@ struct MatmulConfig { static constexpr size_t attr_sizes[] = { sizeof(NVTETensor), // bias_tensor sizeof(NVTETensor), // dbias_tensor - sizeof(bool), // with_gelu_epilogue - sizeof(bool), // with_dgelu_epilogue + sizeof(uint8_t), // with_gelu_epilogue + sizeof(uint8_t), // with_dgelu_epilogue sizeof(NVTETensor), // epilogue_aux_tensor - sizeof(bool), // use_split_accumulator - sizeof(int) // sm_count + sizeof(uint8_t), // use_split_accumulator + sizeof(int32_t) // sm_count }; }; diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 4b7d8179b0..02faad40d3 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -503,6 +503,14 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #if CUBLAS_VERSION >= 120800 NVTE_CHECK(cublas_version() >= 120800, "MXFP8 requires cuBLAS 12.8+, but run-time cuBLAS version is ", cublas_version()); + + // Check that scales are in expected format + NVTE_CHECK(inputA->with_gemm_swizzled_scales, + "MXFP8 scales are not in format expected by GEMM"); + NVTE_CHECK(inputB->with_gemm_swizzled_scales, + "MXFP8 scales are not in format expected by GEMM"); + + // Configure cuBLAS scales fp8e8m0 *A_scale_inverse = reinterpret_cast(param.A_scale_inv); fp8e8m0 *B_scale_inverse = reinterpret_cast(param.B_scale_inv); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, @@ -513,6 +521,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, &B_scale_inverse, sizeof(B_scale_inverse))); scaling_mode_a = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0; scaling_mode_b = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0; + // Workaround for heuristic cache bug in cublasLt. This separates the MXFP8 cache key from non-block scaling. // CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE is unused for block scaling so it's safe to set. if (cublas_version() <= 120803) { @@ -529,17 +538,22 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #if CUBLAS_VERSION >= 120800 NVTE_CHECK(cublas_version() >= 120800, "FP4 requires cuBLAS 12.8+, but run-time cuBLAS version is ", cublas_version()); - // make sure alpha beta computation dtype remains fp32 by CUBLASLT_MATMUL_DESC_SCALE_TYPE - cublasDataType_t scale_type = CUDA_R_32F; + + // Check that scales are in expected format + NVTE_CHECK(inputA->with_gemm_swizzled_scales, + "NVFP4 block scales are not in format expected by GEMM"); + NVTE_CHECK(inputB->with_gemm_swizzled_scales, + "NVFP4 block scales are not in format expected by GEMM"); + + // alpha and beta are device pointers to FP32 + const cublasDataType_t scale_type = CUDA_R_32F; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type))); - - // Set pointer mode: alpha and beta are both device pointers - // https://docs.nvidia.com/cuda/cublas/#cublasltpointermode-t - cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE; + const cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode))); + // Configure cuBLAS scales fp8e4m3 *A_scale_inverse = reinterpret_cast(param.A_scale_inv); fp8e4m3 *B_scale_inverse = reinterpret_cast(param.B_scale_inv); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, @@ -561,6 +575,14 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK(cublas_version() >= 120900, "FP8 block scaling requires cuBLAS 12.9+, but run-time cuBLAS version is ", cublas_version()); + + // Check that matrix formats are valid + NVTE_CHECK((!(inputA->scaling_mode == NVTE_BLOCK_SCALING_2D && + inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)), + "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling GEMM is supported, " + "but got 2D by 2D"); + + // Configure cuBLAS scales float *A_scale_inverse = reinterpret_cast(param.A_scale_inv); float *B_scale_inverse = reinterpret_cast(param.B_scale_inv); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, @@ -569,9 +591,6 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &B_scale_inverse, sizeof(B_scale_inverse))); - NVTE_CHECK((!(inputA->scaling_mode == NVTE_BLOCK_SCALING_2D && - inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)), - "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling supported, but got 2D by 2D"); scaling_mode_a = inputA->scaling_mode == NVTE_BLOCK_SCALING_1D ? CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F : CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F; 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 a839be1701..0696deaaa7 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -739,6 +739,7 @@ void hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &out NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16, "Input tensor must be BF16 tensor, but dtype is ", to_string(input_.dtype()), "."); NVTE_CHECK(input_.dim() >= 2, "Input must be a 2D tensor."); + NVTE_CHECK(!output_.with_gemm_swizzled_scales, "Output must have scales in compact format."); const SimpleTensor &input = input_.data; SimpleTensor &global_amax = output_.amax; SimpleTensor &output_t = output_.data; diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 8a51b54fdb..b304ed34be 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -57,24 +57,24 @@ NVTEMatmulConfig nvte_create_matmul_config(); /*! \brief Query an option in matrix multiplication configuration. * - * \param[in] config Matrix multiplication configuration. - * \param[in] attr Option type. - * \param[out] buf Memory address to write option value. Ignored if - * NULL. - * \param[in] size_in_bytes Size of buf. - * \param[out] size_written Number of bytes that have been written to - * buf. If buf is NULL, then the number of - * bytes that would have been written. + * \param[in] config Matrix multiplication configuration. + * \param[in] attr Option type. + * \param[out] buf Memory address to write option value to. + * Ignored if NULL. + * \param[in] size_in_bytes Size of buf. + * \param[out] size_written Number of bytes that have been written to + * buf. If buf is NULL, then the number of + * bytes that would have been written. */ void nvte_get_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written); /*! \brief Set an option in matrix multiplication configuration. * - * \param[in] config Matrix multiplication configuration. - * \param[in] attr Option type. - * \param[out] buf Memory address to read option value. - * \param[in] size_in_bytes Size of buf. + * \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. */ void nvte_set_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigAttribute attr, const void *buf, size_t size_in_bytes); @@ -296,14 +296,15 @@ class MatmulConfigWrapper { /*! \brief Set whether to compute GELU in GEMM epilogue. */ void set_with_gelu_epilogue(bool with_gelu_epilogue) { - nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigWithGELUEpilogue, - &with_gelu_epilogue, sizeof(bool)); + const auto val = static_cast(with_gelu_epilogue); + nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigWithGELUEpilogue, &val, sizeof(val)); } /*! \brief Set whether to compute GELU backward in GEMM epilogue. */ void set_with_dgelu_epilogue(bool with_dgelu_epilogue) { - nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigWithDGELUEpilogue, - &with_dgelu_epilogue, sizeof(bool)); + const auto val = static_cast(with_dgelu_epilogue); + nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigWithDGELUEpilogue, &val, + sizeof(val)); } /*! \brief Set auxilliary tensor for GEMM epilogue. */ @@ -314,13 +315,15 @@ class MatmulConfigWrapper { /*! \brief Set whether to use split accumulator for FP8 GEMM. */ void set_use_split_accumulator(bool use_split_accumulator) { - nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigUseSplitAccumulator, - &use_split_accumulator, sizeof(bool)); + const auto val = static_cast(use_split_accumulator); + nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigUseSplitAccumulator, &val, + sizeof(val)); } /*! \brief Set number of streaming multiprocessors to use in GEMM kernel. */ void set_sm_count(int sm_count) { - nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigSMCount, &sm_count, sizeof(int)); + const auto val = static_cast(sm_count); + nvte_set_matmul_config_attribute(config_, kNVTEMatmulConfigSMCount, &val, sizeof(val)); } private: diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 4e3544d3c7..5e420b2d42 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -4,8 +4,8 @@ * See LICENSE for license information. ************************************************************************/ -/*! \file cast.h - * \brief Functions to cast to/from FP8. +/*! \file swizzle.h + * \brief Functions to convert scaling factors into format expected by GEMM. */ #ifndef TRANSFORMER_ENGINE_SWIZZLE_H_ @@ -47,7 +47,7 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen /*! \brief Swizzling FP8 block scaling scaling factors into mxfp8 interleaved layout for GEMM * - * \param[in] input Input FP8 block scaling tensor with GEMM_READY scale_inv. + * \param[in] input Input FP8 block-scaled tensor. * \param[in,out] output Output mxfp8 tensor which hosts swizzled scale_inv. * \param[in] stream CUDA stream used for the operation. * @@ -57,7 +57,6 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen * Requirements: * - input is an FP8 block scaling tensor * - input has rowwise usage - * - input.scale_inv is in GEMM_READY format * - output is an MXFP8 tensor * - output has rowwise usage * - output.scale_inv has appropriate shape diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index fd0125c8d0..ae41f238a4 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -13,6 +13,7 @@ #include #include +#include #ifdef __cplusplus extern "C" { @@ -63,13 +64,14 @@ struct NVTEBasicTensor { * \brief Indicates the kind of the tensor parameter to set/get. */ enum NVTETensorParam { - kNVTERowwiseData = 0, /*!< Data usable in rowwise manner */ - kNVTEColumnwiseData = 1, /*!< Data usable in columnwise manner */ - kNVTEScale = 2, /*!< Scale tensor */ - kNVTEAmax = 3, /*!< Amax tensor */ - kNVTERowwiseScaleInv = 4, /*!< Scale inverse tensor for decoding Rowwise Data */ - kNVTEColumnwiseScaleInv = 5, /*!< Scale inverse tensor for decoding Columnwise Data */ - kNVTEColumnwiseAmax = 6, /*!< Columnwise Amax tensor */ + kNVTERowwiseData = 0, /*!< Data usable in rowwise manner */ + kNVTEColumnwiseData = 1, /*!< Data usable in columnwise manner */ + kNVTEScale = 2, /*!< Scale tensor */ + kNVTEAmax = 3, /*!< Amax tensor */ + kNVTERowwiseScaleInv = 4, /*!< Scale inverse tensor for decoding Rowwise Data */ + 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 */ kNVTENumTensorParams }; @@ -266,6 +268,8 @@ NVTEShape nvte_tensor_scale_inv_shape(const NVTETensor tensor); void nvte_zero_tensor(const NVTETensor tensor, cudaStream_t stream); /*! \brief Set a parameter of the tensor. + * + * \warning Deprecated in favor of nvte_set_tensor_param_v2. * * \param[in/out] tensor Tensor. * \param[in] param_name The parameter to be set. @@ -275,12 +279,38 @@ void nvte_set_tensor_param(NVTETensor *tensor, NVTETensorParam param_name, const NVTEBasicTensor *param); /*! \brief Get a value of the parameter of the tensor. + * + * \warning Deprecated in favor of nvte_set_tensor_param_v2. * * \param[in] tensor Tensor. * \param[in] param_name The parameter to be set. */ NVTEBasicTensor nvte_get_tensor_param(const NVTETensor tensor, NVTETensorParam param_name); +/*! \brief Set a tensor parameter. + * + * \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. + */ +void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const void *buf, + size_t size_in_bytes); + +/*! \brief Query a tensor parameter. + * + * \param[in] tensor Tensor. + * \param[in] param Tensor parameter type. + * \param[out] buf Memory address to write parameter value. + * Ignored if NULL. + * \param[in] size_in_bytes Size of buf. + * \param[out] size_written Number of bytes that have been written to + * buf. If buf is NULL, then the number of + * bytes that would have been written. + */ +void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, void *buf, + size_t size_in_bytes, size_t *size_written); + /*! \brief Get the granularity of scaling of this tensor. * * \param[in] tensor Tensor. @@ -326,12 +356,7 @@ enum NVTEQuantizationConfigAttribute { conditional early even when captured in a static CUDA graph. */ kNVTEQuantizationConfigNoopTensor = 2, - /*! Data format for an FP8 block-scaled tensor - * - * This is not the right design since the tensor format is a - * property of the tensor, not the quantization. This enum will - * likely be refactored away in the future. - */ + /*! \warning Deprecated */ kNVTEQuantizationConfigFloat8BlockScaleTensorFormat = 3, /*! RNG state (NVTETensor with 2 elements - seed and offset */ kNVTEQuantizationConfigRNGState = 4, @@ -355,14 +380,14 @@ NVTEQuantizationConfig nvte_create_quantization_config(); /*! \brief Query an option in quantization config. * - * \param[in] config Quantization config. - * \param[in] attr Option type. - * \param[out] buf Memory address to write option value. Ignored if - * NULL. - * \param[in] size_in_bytes Size of buf. - * \param[out] size_written Number of bytes that have been written to - * buf. If buf is NULL, then the number of - * bytes that would have been written. + * \param[in] config Quantization config. + * \param[in] attr Option type. + * \param[out] buf Memory address to write option value. + * Ignored if NULL. + * \param[in] size_in_bytes Size of buf. + * \param[out] size_written Number of bytes that have been written to + * buf. If buf is NULL, then the number of + * bytes that would have been written. */ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, NVTEQuantizationConfigAttribute attr, void *buf, @@ -370,10 +395,10 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, /*! \brief Set an option in quantization config. * - * \param[in] config Quantization config. - * \param[in] attr Option type. - * \param[out] buf Memory address to read option value. - * \param[in] size_in_bytes Size of buf. + * \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. */ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, NVTEQuantizationConfigAttribute attr, const void *buf, @@ -589,20 +614,20 @@ class TensorWrapper { const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING) { tensor_ = nvte_create_tensor(scaling_mode); NVTEBasicTensor data = {dptr, static_cast(dtype), shape}; - nvte_set_tensor_param(&tensor_, kNVTERowwiseData, &data); + nvte_set_tensor_param_v2(tensor_, kNVTERowwiseData, &data, sizeof(data)); NVTEBasicTensor amax = {amax_dptr, kNVTEFloat32, amax_dptr != nullptr ? defaultShape : emptyShape}; - nvte_set_tensor_param(&tensor_, kNVTEAmax, &amax); + nvte_set_tensor_param_v2(tensor_, kNVTEAmax, &amax, sizeof(amax)); NVTEBasicTensor scale = {scale_dptr, kNVTEFloat32, scale_dptr != nullptr ? defaultShape : emptyShape}; - nvte_set_tensor_param(&tensor_, kNVTEScale, &scale); + nvte_set_tensor_param_v2(tensor_, kNVTEScale, &scale, sizeof(scale)); if (scale_inv_dptr == nullptr && scale_inv_shape.ndim == defaultShape.ndim && scale_inv_shape.ndim == 1 && scale_inv_shape.data[0] == defaultShape.data[0]) { // Scale-inv pointer has not been provided and shape matches default scale_inv_shape = emptyShape; } NVTEBasicTensor scale_inv = {scale_inv_dptr, kNVTEFloat32, scale_inv_shape}; - nvte_set_tensor_param(&tensor_, kNVTERowwiseScaleInv, &scale_inv); + nvte_set_tensor_param_v2(tensor_, kNVTERowwiseScaleInv, &scale_inv, sizeof(scale_inv)); } /*! \brief Constructs new TensorWrapper. @@ -673,7 +698,7 @@ class TensorWrapper { const ShapeType &shape) noexcept { NVTEShape nvte_shape = this->convertShape(shape); NVTEBasicTensor data = {dptr, static_cast(type), nvte_shape}; - nvte_set_tensor_param(&tensor_, param, &data); + nvte_set_tensor_param_v2(tensor_, param, &data, sizeof(data)); return *this; } @@ -712,10 +737,17 @@ class TensorWrapper { return set_parameter(kNVTEColumnwiseAmax, dptr, type, shape); } + void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales) { + const auto val = static_cast(with_gemm_swizzled_scales); + nvte_set_tensor_param_v2(tensor_, kNVTEWithGEMMSwizzledScales, &val, sizeof(val)); + } + // Parameter getters NVTEBasicTensor get_parameter(const NVTETensorParam param) const noexcept { - return nvte_get_tensor_param(tensor_, param); + NVTEBasicTensor ret; + nvte_get_tensor_param_v2(tensor_, param, &ret, sizeof(ret), nullptr); + return ret; } NVTEBasicTensor get_rowwise_data() const noexcept { return get_parameter(kNVTERowwiseData); } @@ -740,6 +772,12 @@ class TensorWrapper { return get_parameter(kNVTEColumnwiseAmax); } + bool get_with_gemm_swizzled_scales() const { + uint8_t val = 0; + nvte_get_tensor_param_v2(tensor_, kNVTEWithGEMMSwizzledScales, &val, sizeof(val), nullptr); + return static_cast(val); + } + /*! \brief Get an underlying NVTETensor. * * \return NVTETensor held by this TensorWrapper. @@ -919,15 +957,8 @@ class TensorWrapper { NVTETensor tensor_ = nullptr; }; -/*! \enum Float8BlockScaleTensorFormat - * \brief Data format for an FP8 block-scaled tensor - */ -enum class Float8BlockScaleTensorFormat { - /*! FP8 data is transposed if needed and scales are swizzled */ - GEMM_READY = 0, - /*! FP8 data is untransposed and scales are not swizzled or padded */ - COMPACT = 1 -}; +/*! \warning Deprecated */ +enum class Float8BlockScaleTensorFormat { GEMM_READY = 0, COMPACT = 1, INVALID }; /*! \struct QuantizationConfigWrapper * \brief C++ wrapper for NVTEQuantizationConfigWrapper. @@ -968,8 +999,9 @@ class QuantizationConfigWrapper { /*! \brief Set whether to force power of 2 scales */ void set_force_pow_2_scales(bool force_pow_2_scales) { - nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigForcePow2Scales, - &force_pow_2_scales, sizeof(bool)); + const auto val = static_cast(force_pow_2_scales); + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigForcePow2Scales, &val, + sizeof(val)); } /*! \brief Set small value to add to amax */ @@ -984,12 +1016,8 @@ class QuantizationConfigWrapper { sizeof(NVTETensor)); } - /*! \brief Set FP8 block-scaled tensor format */ - void set_float8_block_scale_tensor_format(Float8BlockScaleTensorFormat format) { - nvte_set_quantization_config_attribute(config_, - kNVTEQuantizationConfigFloat8BlockScaleTensorFormat, - &format, sizeof(Float8BlockScaleTensorFormat)); - } + /*! \warning Deprecated */ + void set_float8_block_scale_tensor_format(Float8BlockScaleTensorFormat format) {} /*! \brief Set stochastic rounding state */ void set_rng_state(NVTETensor rng_state) { @@ -999,20 +1027,23 @@ class QuantizationConfigWrapper { /*! \brief Set whether to use 2D block scaling for NVFP4 */ void set_nvfp4_2d_quantization(bool nvfp4_2d_quantization) { + const auto val = static_cast(nvfp4_2d_quantization); nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigNVFP42DQuantization, - &nvfp4_2d_quantization, sizeof(bool)); + &val, sizeof(val)); } /*! \brief Set whether to use stochastic rounding */ void set_stochastic_rounding(bool stochastic_rounding) { - nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigStochasticRounding, - &stochastic_rounding, sizeof(bool)); + const auto val = static_cast(stochastic_rounding); + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigStochasticRounding, &val, + sizeof(val)); } /*! \brief Set whether to enable fast math operations */ void set_use_fast_math(bool use_fast_math) { - nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigUseFastMath, - &use_fast_math, sizeof(bool)); + const auto val = static_cast(use_fast_math); + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigUseFastMath, &val, + sizeof(val)); } private: diff --git a/transformer_engine/common/normalization/layernorm/ln_api.cpp b/transformer_engine/common/normalization/layernorm/ln_api.cpp index 24664d363b..7bd5a1bbd0 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -27,10 +27,15 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size const float epsilon, Tensor* z, Tensor* mu, Tensor* rsigma, Tensor* workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream) { + // Check for unsupported configurations if (is_fp8_dtype(z->data.dtype) && !is_delayed_tensor_scaling(z->scaling_mode) && !is_mxfp8_scaling(z->scaling_mode)) { NVTE_ERROR("Not implemented scaling mode: " + to_string(z->scaling_mode) + "."); } + if (is_mxfp8_scaling(z->scaling_mode)) { + NVTE_CHECK(!z->with_gemm_swizzled_scales, + "MXFP8 output must have scales in compact format, not swizzled for GEMM."); + } NVTE_CHECK(x.data.shape.size() == 2, "x must be 2D tensor."); NVTE_CHECK(gamma.data.shape == beta.data.shape, "Gamma and Beta must have the same shape."); diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index 137df79bde..6f6656534a 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -23,10 +23,15 @@ using namespace normalization; void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tensor *z, Tensor *rsigma, Tensor *workspace, const int multiprocessorCount, const bool zero_centered_gamma, cudaStream_t stream) { + // Check for unsupported configurations if (is_fp8_dtype(z->data.dtype) && !is_delayed_tensor_scaling(z->scaling_mode) && !is_mxfp8_scaling(z->scaling_mode)) { NVTE_ERROR("Not implemented scaling mode: " + to_string(z->scaling_mode) + "."); } + if (is_mxfp8_scaling(z->scaling_mode)) { + NVTE_CHECK(!z->with_gemm_swizzled_scales, + "MXFP8 output must have scales in compact format, not swizzled for GEMM."); + } NVTE_CHECK(x.data.shape.size() == 2, "x must be 2D tensor."); diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index 534a796324..5a6490c042 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -110,7 +110,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) OType *output_rowwise_minus_offset = output_rowwise - start_offset; OType *output_colwise_minus_offset = output_colwise - start_offset; int warp_idx = threadIdx.x / 32; - int lane_idx = threadIdx.x % 32; + // int lane_idx = threadIdx.x % 32; int c = blockIdx.x * kColsPerTile + threadIdx.x; int r = blockIdx.y * kRowsPerTile; diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 73647d5717..4425c4e9f7 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -340,6 +340,10 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s // Check tensors CheckInputTensor(*input, "scaling_factor_input"); CheckInputTensor(*output, "scaling_factor_output"); + NVTE_CHECK(!input->with_gemm_swizzled_scales, + "Expected input tensor with scales in compact format."); + NVTE_CHECK(output->with_gemm_swizzled_scales, + "Expected output tensor with scales in GEMM swizzled format."); switch (scaling_mode) { case NVTE_MXFP8_1D_SCALING: NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP8, got ", @@ -656,6 +660,11 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, NVTE_CHECK( (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)), "Not implemented scaling mode " + to_string(scaling_mode) + "."); + NVTE_CHECK(!input[i]->with_gemm_swizzled_scales, + "Expected input tensors with scales in compact format."); + NVTE_CHECK(output[i]->with_gemm_swizzled_scales, + "Expected output tensors with scales in GEMM swizzled format."); + // 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) + "]"); diff --git a/transformer_engine/common/swizzle/swizzle_block_scaling.cu b/transformer_engine/common/swizzle/swizzle_block_scaling.cu index 37993787a5..90bc3985a4 100644 --- a/transformer_engine/common/swizzle/swizzle_block_scaling.cu +++ b/transformer_engine/common/swizzle/swizzle_block_scaling.cu @@ -98,7 +98,8 @@ void __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE) // calculate this warp's input base pointer constexpr uint32_t in_x_stride = WARP_SIZE * sizeof(uint4); - const void* const warp_src = in + in_tile_y * in_y_stride + in_tile_x * in_x_stride; + const void* const warp_src = + (reinterpret_cast(in) + in_tile_y * in_y_stride + in_tile_x * in_x_stride); // load scaling factors for this lane's initial four 1x128 tiles uint4 sf; @@ -129,7 +130,8 @@ void __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE) // store them cooperatively for 512 1x32 tiles in a 128x128 tile constexpr uint32_t out_x_stride = 512; - void* const warp_dst = out + out_tile_y * out_y_stride + out_tile_x * out_x_stride; + void* const warp_dst = + (reinterpret_cast(out) + out_tile_y * out_y_stride + out_tile_x * out_x_stride); reinterpret_cast(warp_dst)[lane] = sf; } @@ -193,7 +195,8 @@ void __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE) // calculate this warp's input base pointer constexpr uint32_t in_x_stride = sizeof(float); - const void* const warp_src = in + in_tile_y * in_y_stride + in_tile_x * in_x_stride; + const void* const warp_src = + (reinterpret_cast(in) + in_tile_y * in_y_stride + in_tile_x * in_x_stride); // load scaling factor for this warp's 128x128 tile uint32_t sf = *reinterpret_cast(warp_src); @@ -208,7 +211,8 @@ void __global__ __launch_bounds__(WARPS_X_PER_TB* WARPS_Y_PER_TB* WARP_SIZE) // store it cooperatively for 512 1x32 tiles in a 128x128 tile constexpr uint32_t out_x_stride = 512; - void* const warp_dst = out + out_tile_y * out_y_stride + out_tile_x * out_x_stride; + void* const warp_dst = + (reinterpret_cast(out) + out_tile_y * out_y_stride + out_tile_x * out_x_stride); reinterpret_cast(warp_dst)[lane] = sf4; } @@ -261,6 +265,9 @@ void swizzle_block_scaling_to_mxfp8_scaling_factors(const Tensor* input, Tensor* NVTE_CHECK(output->scale_inv.dtype == DType::kFloat8E8M0, "Output must have E8M0 scaling factors"); + NVTE_CHECK(output->with_gemm_swizzled_scales, + "Expected output tensor with scales in GEMM swizzled format."); + NVTE_CHECK(input->data.dptr != nullptr, "Input must have rowwise data"); NVTE_CHECK(output->data.dptr == input->data.dptr, "Output must share data with input"); NVTE_CHECK(input->scale_inv.dptr != nullptr, "Input must have rowwise scaling factors"); diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 370d9723cf..6880dd560a 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -6,12 +6,16 @@ #include +#include #include #include #include #include #include +#include +#include #include +#include #include "common.h" #include "common/util/cuda_runtime.h" @@ -778,7 +782,8 @@ void nvte_set_tensor_param(NVTETensor *tensor, NVTETensorParam param_name, t->columnwise_amax = *param; break; default: - NVTE_ERROR("Unknown tensor parameter!"); + NVTE_ERROR("Unsupported tensor parameter (", static_cast(param_name), + "). Consider using nvte_set_tensor_param_v2 instead."); } } @@ -803,7 +808,148 @@ NVTEBasicTensor nvte_get_tensor_param(const NVTETensor tensor, NVTETensorParam p case kNVTEColumnwiseAmax: return t.columnwise_amax; default: - NVTE_ERROR("Unknown tensor parameter!"); + NVTE_ERROR("Unsupported tensor parameter (", static_cast(param_name), + "). Consider using nvte_set_tensor_param_v2 instead."); + } +} + +void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const void *buf, + size_t size_in_bytes) { + // Check attribute and buffer + NVTE_CHECK(param < kNVTENumTensorParams, "Invalid NVTETensorParam (got ", static_cast(param), + ")"); + NVTE_CHECK(tensor != nullptr, "Tensor pointer can't be NULL."); + auto &t = *transformer_engine::convertNVTETensorCheck(tensor); + const auto &attr_size = transformer_engine::Tensor::attr_sizes[param]; + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for tensor parameter " + "(parameter ", + static_cast(param), " needs ", attr_size, " bytes, but buffer has ", + size_in_bytes, " bytes)"); + NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + + // Read from buffer + switch (param) { + case kNVTERowwiseData: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.data = *basic_tensor; + break; + } + case kNVTEColumnwiseData: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.columnwise_data = *basic_tensor; + break; + } + case kNVTEScale: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.scale = *basic_tensor; + break; + } + case kNVTEAmax: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.amax = *basic_tensor; + break; + } + case kNVTERowwiseScaleInv: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.scale_inv = *basic_tensor; + break; + } + case kNVTEColumnwiseScaleInv: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.columnwise_scale_inv = *basic_tensor; + break; + } + case kNVTEColumnwiseAmax: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.columnwise_amax = *basic_tensor; + break; + } + case kNVTEWithGEMMSwizzledScales: + t.with_gemm_swizzled_scales = static_cast(*reinterpret_cast(buf)); + break; + default: + NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); + } +} + +void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, void *buf, + size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + + // Check param + NVTE_CHECK(param < kNVTENumTensorParams, "Invalid NVTETensorParam (got ", static_cast(param), + ")"); + + // Write attribute size if provided + const auto &attr_size = Tensor::attr_sizes[param]; + if (size_written != nullptr) { + *size_written = attr_size; + } + + // Return immediately if buffer is not provided + if (buf == nullptr) { + return; + } + + // Check buffer size + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for tensor parameter " + "(parameter ", + static_cast(param), " needs ", attr_size, " bytes, but buffer has ", + size_in_bytes, " bytes)"); + + // Get C++ tensor + const Tensor *t = convertNVTETensor(tensor); + std::optional dummy; + if (t == nullptr) { + // Make dummy tensor if provided tensor is invalid + dummy.emplace(); + t = &(*dummy); + } + + // Write to buffer + switch (param) { + case kNVTERowwiseData: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->data); + break; + } + case kNVTEColumnwiseData: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->columnwise_data); + break; + } + case kNVTEScale: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->scale); + break; + } + case kNVTEAmax: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->amax); + break; + } + case kNVTERowwiseScaleInv: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->scale_inv); + break; + } + case kNVTEColumnwiseScaleInv: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->columnwise_scale_inv); + break; + } + case kNVTEColumnwiseAmax: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->columnwise_amax); + break; + } + case kNVTEWithGEMMSwizzledScales: + *reinterpret_cast(buf) = static_cast(t->with_gemm_swizzled_scales); + break; + default: + NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } } @@ -854,10 +1000,12 @@ NVTEQuantizationConfig nvte_create_quantization_config() { void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, NVTEQuantizationConfigAttribute attr, void *buf, size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + // Write attribute size NVTE_CHECK(attr < kNVTEQuantizationConfigNumAttributes, "Invalid NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); - const auto &attr_size = transformer_engine::QuantizationConfig::attr_sizes[attr]; + const auto &attr_size = QuantizationConfig::attr_sizes[attr]; if (size_written != nullptr) { *size_written = attr_size; } @@ -874,12 +1022,18 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, " bytes)"); + // bool size is implementation-dependent, so we explicitly specify + // uint8_t in the user-facing API. + auto bool_to_uint8 = [](bool in, void *out) { + *reinterpret_cast(out) = static_cast(in); + }; + // Write to buffer NVTE_CHECK(config != nullptr, "Invalid NVTEQuantizationConfig (got NULL)"); - const auto &config_ = *reinterpret_cast(config); + const auto &config_ = *reinterpret_cast(config); switch (attr) { case kNVTEQuantizationConfigForcePow2Scales: - std::memcpy(buf, &config_.force_pow_2_scales, attr_size); + bool_to_uint8(config_.force_pow_2_scales, buf); break; case kNVTEQuantizationConfigAmaxEpsilon: std::memcpy(buf, &config_.amax_epsilon, attr_size); @@ -887,20 +1041,23 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigNoopTensor: std::memcpy(buf, &config_.noop_tensor, attr_size); break; - case kNVTEQuantizationConfigFloat8BlockScaleTensorFormat: - std::memcpy(buf, &config_.float8_block_scale_tensor_format, attr_size); + case kNVTEQuantizationConfigFloat8BlockScaleTensorFormat: { + // Deprecated + const auto invalid = Float8BlockScaleTensorFormat::INVALID; + std::memcpy(buf, &invalid, attr_size); break; + } case kNVTEQuantizationConfigRNGState: std::memcpy(buf, &config_.rng_state, attr_size); break; case kNVTEQuantizationConfigNVFP42DQuantization: - std::memcpy(buf, &config_.nvfp4_2d_quantization, attr_size); + bool_to_uint8(config_.nvfp4_2d_quantization, buf); break; case kNVTEQuantizationConfigStochasticRounding: - std::memcpy(buf, &config_.stochastic_rounding, attr_size); + bool_to_uint8(config_.stochastic_rounding, buf); break; case kNVTEQuantizationConfigUseFastMath: - std::memcpy(buf, &config_.use_fast_math, attr_size); + bool_to_uint8(config_.use_fast_math, buf); break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); @@ -910,10 +1067,12 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, NVTEQuantizationConfigAttribute attr, const void *buf, size_t size_in_bytes) { + using namespace transformer_engine; + // Check attribute and buffer NVTE_CHECK(attr < kNVTEQuantizationConfigNumAttributes, "Invalid NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); - const auto &attr_size = transformer_engine::QuantizationConfig::attr_sizes[attr]; + const auto &attr_size = QuantizationConfig::attr_sizes[attr]; NVTE_CHECK(size_in_bytes >= attr_size, "Buffer is too small for quantization config attribute " "(attribute ", @@ -921,12 +1080,18 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, " bytes)"); NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + // bool size is implementation-dependent, so we explicitly specify + // uint8_t in the user-facing API. + auto uint8_to_bool = [](const void *in, bool &out) { + out = static_cast(*reinterpret_cast(in)); + }; + // Read from buffer NVTE_CHECK(config != nullptr, "Invalid NVTEQuantizationConfig (got NULL)"); - auto &config_ = *reinterpret_cast(config); + auto &config_ = *reinterpret_cast(config); switch (attr) { case kNVTEQuantizationConfigForcePow2Scales: - std::memcpy(&config_.force_pow_2_scales, buf, attr_size); + uint8_to_bool(buf, config_.force_pow_2_scales); break; case kNVTEQuantizationConfigAmaxEpsilon: std::memcpy(&config_.amax_epsilon, buf, attr_size); @@ -935,19 +1100,19 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, std::memcpy(&config_.noop_tensor, buf, attr_size); break; case kNVTEQuantizationConfigFloat8BlockScaleTensorFormat: - std::memcpy(&config_.float8_block_scale_tensor_format, buf, attr_size); + // Deprecated break; case kNVTEQuantizationConfigRNGState: std::memcpy(&config_.rng_state, buf, attr_size); break; case kNVTEQuantizationConfigNVFP42DQuantization: - std::memcpy(&config_.nvfp4_2d_quantization, buf, attr_size); + uint8_to_bool(buf, config_.nvfp4_2d_quantization); break; case kNVTEQuantizationConfigStochasticRounding: - std::memcpy(&config_.stochastic_rounding, buf, attr_size); + uint8_to_bool(buf, config_.stochastic_rounding); break; case kNVTEQuantizationConfigUseFastMath: - std::memcpy(&config_.use_fast_math, buf, attr_size); + uint8_to_bool(buf, config_.use_fast_math); break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h index 66bf8e86df..a5ec2306b1 100644 --- a/transformer_engine/common/transpose/cast_transpose.h +++ b/transformer_engine/common/transpose/cast_transpose.h @@ -36,7 +36,7 @@ enum class FP8BlockwiseRowwiseOption { NONE, // Rowwise data, scales in GEMM format ROWWISE_GEMM_READY, - // Rowwise data, scales in compact format, needs extra processing (padding, transposing) before GEMM + // Deprecated ROWWISE_COMPACT }; @@ -50,8 +50,7 @@ enum class FP8BlockwiseColumnwiseOption { // On Hopper sm90, GEMM_READY means that columnwise quantization also fuses transpose op // On higher sm versions with TN,NT,NN fp8 gemm, GEMM_READY doesn't fuse transpose COLUMNWISE_GEMM_READY, - // Columnwise data in original shape - // Scales in compact format, needs extra processing (padding, transposing) before GEMM + // Deprecated COLUMNWISE_COMPACT }; diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index c636a627a4..0e286009a5 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -492,7 +492,7 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor } NVTE_CHECK(input.shape == output.shape, "Input and output must have the same shape."); - const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u; + const size_t row_length = input.shape.size() > 0 ? input.shape.back() : 1; size_t num_rows = 1; for (size_t i = 0; (i < input.shape.size() - 1) && (input.shape.size() > 0); ++i) { num_rows *= input.shape.at(i); @@ -511,12 +511,14 @@ void quantize_transpose_square_blockwise(const SimpleTensor& input, SimpleTensor const float* noop_ptr = reinterpret_cast(noop_tensor.dptr); if (return_transpose) { - NVTE_CHECK(output_t.shape.size() == input.shape.size(), - "output_t must have same number of dimensions as input."); + NVTE_CHECK(output_t.shape.size() == input.shape.size(), "input (shape=", input.shape, + ") and output_t (shape=", output_t.shape, ") have incompatible dims."); if (output_t.shape.size() > 0) { - NVTE_CHECK(output_t.shape[0] == row_length, "Wrong dimension 0 of output_t."); + NVTE_CHECK(output_t.shape.front() == input.shape.back(), "input (shape=", input.shape, + ") and output_t (shape=", output_t.shape, ") have incompatible dims."); for (size_t i = 1; i < output_t.shape.size(); ++i) { - NVTE_CHECK(output_t.shape.at(i) == input.shape.at(i - 1), "Wrong dimension in output_t"); + NVTE_CHECK(output_t.shape[i] == input.shape[i - 1], "input (shape=", input.shape, + ") and output_t (shape=", output_t.shape, ") have incompatible dims."); } } NVTE_CHECK(output.dtype == output_t.dtype, "output and output_t need to have the same type."); diff --git a/transformer_engine/common/transpose/transpose.cu b/transformer_engine/common/transpose/transpose.cu index a280df21a5..49f1333024 100644 --- a/transformer_engine/common/transpose/transpose.cu +++ b/transformer_engine/common/transpose/transpose.cu @@ -14,8 +14,10 @@ #include "../util/rtc.h" #include "../util/string.h" #include "../utils.cuh" +#include "./transpose.h" namespace transformer_engine { +namespace detail { namespace { @@ -203,7 +205,8 @@ void transpose(const Tensor &input, const Tensor &noop, Tensor *output_, cudaStr NVTE_CHECK(input.data.dptr != nullptr, "Input is not allocated."); NVTE_CHECK(output.data.dptr != nullptr, "Output is not allocated."); - NVTE_CHECK(input.data.dtype == output.data.dtype, "Input and output type must match."); + NVTE_CHECK(input.data.dtype == output.data.dtype, "Input (dtype=", to_string(input.data.dtype), + ") and output (dtype=", to_string(output.data.dtype), ") do not match."); if (noop.data.dptr != nullptr) { NVTE_CHECK(noop.numel() == 1, "Expected 1 element, ", "but found ", noop.numel(), "."); @@ -283,19 +286,20 @@ void transpose(const Tensor &input, const Tensor &noop, Tensor *output_, cudaStr }); // NOLINT(*) } +} // namespace detail } // namespace transformer_engine void nvte_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_transpose); using namespace transformer_engine; auto noop = Tensor(); - transpose(*convertNVTETensorCheck(input), noop, convertNVTETensor(output), stream); + detail::transpose(*convertNVTETensorCheck(input), noop, convertNVTETensor(output), stream); } void nvte_transpose_with_noop(const NVTETensor input, const NVTETensor noop, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_transpose_with_noop); using namespace transformer_engine; - transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(noop), - convertNVTETensor(output), stream); + detail::transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(noop), + convertNVTETensor(output), stream); } diff --git a/transformer_engine/common/transpose/transpose.h b/transformer_engine/common/transpose/transpose.h new file mode 100644 index 0000000000..36246f4abd --- /dev/null +++ b/transformer_engine/common/transpose/transpose.h @@ -0,0 +1,20 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_COMMON_TRANSPOSE_TRANSPOSE_H_ +#define TRANSFORMER_ENGINE_COMMON_TRANSPOSE_TRANSPOSE_H_ + +#include "../common.h" + +namespace transformer_engine { +namespace detail { + +void transpose(const Tensor &input, const Tensor &noop, Tensor *output_, cudaStream_t stream); + +} // namespace detail +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_TRANSPOSE_TRANSPOSE_H_ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index c22fb33ffe..4cdd8297a8 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -840,6 +840,7 @@ __device__ __forceinline__ int32_t elect_one_sync(uint32_t mask = 0xFFFFFFFFu) { return pred; #else NVTE_DEVICE_ERROR("elect_one_sync is only supported on SM 10.0+."); + return 0; #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -891,6 +892,7 @@ __device__ __forceinline__ bf16 get_amax(bf16 a, bf16 b) { return r; #else NVTE_DEVICE_ERROR("get_amax is only supported on SM 10.0+."); + return 0.f; #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } @@ -903,6 +905,7 @@ __device__ __forceinline__ fp16 get_amax(fp16 a, fp16 b) { return r; #else NVTE_DEVICE_ERROR("get_amax is only supported on SM 10.0+."); + return 0.f; #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index faa4e36809..6adba23a8f 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -83,7 +83,8 @@ pybind11::enum_( \ m, "Float8BlockScaleTensorFormat", pybind11::module_local()) \ .value("GEMM_READY", transformer_engine::Float8BlockScaleTensorFormat::GEMM_READY) \ - .value("COMPACT", transformer_engine::Float8BlockScaleTensorFormat::COMPACT); \ + .value("COMPACT", transformer_engine::Float8BlockScaleTensorFormat::COMPACT) \ + .value("INVALID", transformer_engine::Float8BlockScaleTensorFormat::INVALID); \ pybind11::enum_(m, "CommOverlapType", \ pybind11::module_local()) \ .value("RS", transformer_engine::CommOverlapType::RS) \ diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 7a8670f043..29a108c75f 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -62,12 +62,17 @@ def __init__( self.tp_group = tp_group # used in inspect_tensor calls self.iteration = TEDebugState.get_iteration() - # .internal = True is slightly faster, but results - # in errors when caching the weights. - # Setting .internal = False is safer. + # Configure parent quantizer if parent_quantizer is not None: + # .internal = True is slightly faster, but results + # in errors when caching the weights. + # Setting .internal = False is safer. parent_quantizer.internal = False + # .optimize_for_gemm = True is not supported because debug + # quantizers perform non-GEMM operations. + parent_quantizer.optimize_for_gemm = False + self.rowwise_gemm_name, self.columnwise_gemm_name = _tensor_to_gemm_names_map[tensor_name] # next iteration when this quantizer will call any API diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index e25a67a401..4303682bfb 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -65,23 +65,33 @@ std::tuple> xla_buffer_to_nvte_gemm_operand( NVTE_CHECK(typeToSize(scale_dtype) == 1, "Inverse scale factors need to have an 8-bit data type."); } - if (!is_nvfp4) { + if (scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING) { + // Assume MXFP8 scales are already swizzled if (rowwise) { input.set_rowwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); } else { input.set_columnwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); } - } else { // Swizzle for NVFP4 + input.set_with_gemm_swizzled_scales(true); + } else if (is_nvfp4) { // Swizzle for NVFP4 NVTE_CHECK(rowwise, "NVFP4 GEMM expects rowwise for both LHS and RHS"); input.set_rowwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); // Create tensor to hold swizzled scale factor TensorWrapper output(get_nvte_scaling_mode(scaling_mode)); output.set_rowwise_data(buffer.untyped_data(), input_dtype, input_shape); output.set_rowwise_scale_inv(swizzle_scale_ptr, scale_dtype, scale_shape); + output.set_with_gemm_swizzled_scales(true); // Launch swizzle kernel nvte_swizzle_scaling_factors(input.data(), output.data(), stream); // Set swizzled scales into the input tensor input.set_rowwise_scale_inv(swizzle_scale_ptr, scale_dtype, scale_shape); + input.set_with_gemm_swizzled_scales(true); + } else { // Tensor scaling + if (rowwise) { + input.set_rowwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); + } else { + input.set_columnwise_scale_inv(scale_inv.untyped_data(), scale_dtype, scale_shape); + } } } @@ -669,6 +679,7 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_swizzle_i.set_rowwise_scale_inv(lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); lhs_i.set_rowwise_scale_inv(swizzled_lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); } + lhs_i.set_with_gemm_swizzled_scales(true); if (rhs_use_colwise) { rhs_swizzle_i.set_columnwise_data(rhs_vptr, rhs_dtype, rhs_shape_i); rhs_swizzle_i.set_columnwise_scale_inv(rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); @@ -678,6 +689,7 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type rhs_swizzle_i.set_rowwise_scale_inv(rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); rhs_i.set_rowwise_scale_inv(swizzled_rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); } + rhs_i.set_with_gemm_swizzled_scales(true); if (!is_empty_gemm) { lhs_swizzle_wrapper_list.push_back(std::move(lhs_swizzle_i)); diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 2a97e2ac71..406e7075f7 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -164,17 +164,9 @@ def general_gemm( bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype] if isinstance(A, Float8BlockwiseQTensorStorage) or isinstance(B, Float8BlockwiseQTensorStorage): - # There is not use_split_accumulator == False - # implementation for Float8BlockwiseQTensorStorage GEMM + # FP8 block-scaling requires split accumulator use_split_accumulator = True - # Check that data format is supported - if ( - A._data_format != tex.Float8BlockScaleTensorFormat.GEMM_READY - or B._data_format != tex.Float8BlockScaleTensorFormat.GEMM_READY - ): - raise RuntimeError("GEMM with Float8BlockwiseQTensor requires GEMM_READY format") - args = ( A, transa, # transa diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index fa6f142b68..645dbb48d2 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -301,11 +301,13 @@ std::vector convertShape(const NVTEShape& shape) { return std::vector(shape.data, shape.data + shape.ndim); } -size_t roundup(const size_t value, const size_t multiple) { +size_t roundup(size_t value, size_t multiple) { assert(multiple > 0); return ((value + multiple - 1) / multiple) * multiple; } +size_t ceildiv(size_t numer, size_t denom) { return (numer + denom - 1) / denom; } + void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr) { NVTE_SCOPED_GIL_RELEASE({ nvte_extract_seed_and_offset(rng_state_ptr, arg.captured_, arg.seed_.ptr, arg.seed_.val, diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 1e1e3326c4..bc22e03097 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -120,6 +120,7 @@ class Quantizer { bool rowwise_usage = true; bool columnwise_usage = true; bool internal = false; + bool optimize_for_gemm = false; py::handle quantizer; protected: @@ -231,8 +232,6 @@ class Float8BlockQuantizer : public Quantizer { bool force_pow_2_scales = false; // Amax within quantization tile has a floor of epsilon. float amax_epsilon = 0.0; - // Whether quantized tensor will be used in an all-gather - bool all_gather_usage = false; private: int block_scaling_dim = 2; @@ -358,11 +357,12 @@ inline size_t typeToNumBits(transformer_engine::DType t) { case transformer_engine::DType::kByte: case transformer_engine::DType::kFloat8E4M3: case transformer_engine::DType::kFloat8E5M2: + case transformer_engine::DType::kFloat8E8M0: return 8; case transformer_engine::DType::kFloat4E2M1: return 4; default: - NVTE_ERROR("Invalid type"); + NVTE_ERROR("Invalid type (", static_cast(t), ")."); } } @@ -386,8 +386,10 @@ inline at::ScalarType GetATenDType(transformer_engine::DType t) { return at::kFloat8_e4m3fn; case transformer_engine::DType::kFloat8E5M2: return at::kFloat8_e5m2; + case transformer_engine::DType::kFloat8E8M0: + return at::kByte; // e8m0 dtype requires PyTorch 2.7.0+ default: - NVTE_ERROR("Invalid type"); + NVTE_ERROR("Invalid type (", static_cast(t), ")."); } } @@ -414,8 +416,7 @@ inline transformer_engine::DType GetTransformerEngineDType(at::ScalarType t) { case torch::kInt64: return transformer_engine::DType::kInt64; default: - std::cout << "Type: " << static_cast(t) << std::endl; - NVTE_ERROR("Invalid type"); + NVTE_ERROR("Invalid type (", static_cast(t), ")."); } } @@ -477,7 +478,9 @@ void* getDataPtr(at::Tensor tensor, int offset = 0); std::vector convertShape(const NVTEShape& shape); -size_t roundup(const size_t value, const size_t multiple); +size_t roundup(size_t value, size_t multiple); + +size_t ceildiv(size_t numer, size_t denom); NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape); diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 60b931abfd..9dc0d1f37b 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -7,7 +7,12 @@ #ifndef TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_ #define TRANSFORMER_ENGINE_PYTORCH_CSRC_EXTENSIONS_H_ +#include #include +#include +#include +#include +#include #include "common.h" @@ -78,11 +83,6 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, int64_t window_size_right, bool return_max_logit, bool cuda_graph); -std::pair quantizer_helper(py::handle quantizer, - const std::vector &shape, DType dtype, - bool create_hp_tensor_for_cs, - std::optional data); - 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, @@ -475,6 +475,13 @@ void fused_multi_row_padding(at::Tensor input, at::Tensor output, void fused_multi_row_unpadding(at::Tensor input, at::Tensor output, std::vector input_row_list, std::vector unpadded_input_row_list); + +/*************************************************************************************************** + * Scale swizzling for GEMM + **************************************************************************************************/ + +void inplace_swizzle_scale_for_gemm(py::handle &tensor); + /*************************************************************************************************** * NVSHMEM APIs **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index ac06841879..5c9d0f5b07 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -327,9 +327,9 @@ std::tuple, std::vector> bulk_allocate_fp (columnwise_usage ? py::cast(columnwise_scale_list[i]) : py::none()); // Construct Python tensor - tensor_py_list.emplace_back(Float8BlockwiseQTensorClass( - rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, fp8_dtype, - quantizer_py_list[i], is_2D_scaled, Float8BlockScaleTensorFormat::GEMM_READY)); + tensor_py_list.emplace_back( + Float8BlockwiseQTensorClass(rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, + fp8_dtype, quantizer_py_list[i], is_2D_scaled)); // Construct C++ tensor tensor_cpp_list.emplace_back(makeTransformerEngineTensor( @@ -365,6 +365,8 @@ std::tuple, std::vector> bulk_allocate_mx const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); 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; @@ -475,8 +477,8 @@ std::tuple, std::vector> bulk_allocate_mx // Construct Python tensor tensor_py_list.emplace_back(MXFP8TensorClass(rowwise_data, rowwise_scale, columnwise_data, - columnwise_scale, fp8_dtype, - quantizer_py_list[i])); + columnwise_scale, fp8_dtype, quantizer_py_list[i], + with_gemm_swizzled_scales)); // Construct C++ tensor tensor_cpp_list.emplace_back(makeTransformerEngineTensor( @@ -488,6 +490,7 @@ std::tuple, std::vector> bulk_allocate_mx columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, rowwise_usage ? rowwise_scale_shapes[i] : std::vector{0}, columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode)); + tensor_cpp_list.back().set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); } return retval; @@ -517,6 +520,7 @@ std::tuple, std::vector, bool> bulk_alloc const auto columnwise_usage = quantizer_cpp_list[0]->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; constexpr size_t scale_elem_size = 1; // Helper function to construct tensor view @@ -675,9 +679,9 @@ 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])); + 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)); // Construct C++ tensor // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, @@ -693,6 +697,7 @@ std::tuple, std::vector, bool> bulk_alloc columnwise_usage ? columnwise_scale_list[i].data_ptr() : nullptr, 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); // Set the amax rowwise and amax columnwise if available if (rowwise_usage) { @@ -703,6 +708,7 @@ std::tuple, std::vector, bool> bulk_alloc tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, std::vector{1}); } + tensor_cpp_list.emplace_back(std::move(tensor_wrapper)); } } diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 07ddfbeb6f..d75b0f14c7 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -240,9 +240,12 @@ std::vector gemm(py::handle A, bool transa, py::handle B, bool trans auto main_stream = at::cuda::getCurrentCUDAStream(); if (A_tensor.numel() != 0 && B_tensor.numel() != 0) { // Optionally swizzle the scaling factors - swizzled_scale_inverses_list.emplace_back(std::move(swizzle_scaling_factors(A_tensor, transa))); - swizzled_scale_inverses_list.emplace_back( - std::move(swizzle_scaling_factors(B_tensor, !transb))); + auto [A_row_scales, A_col_scales] = swizzle_scales_for_gemm(A_tensor, transa, !transa); + auto [B_row_scales, B_col_scales] = swizzle_scales_for_gemm(B_tensor, !transb, transb); + swizzled_scale_inverses_list.emplace_back(std::move(A_row_scales)); + swizzled_scale_inverses_list.emplace_back(std::move(A_col_scales)); + swizzled_scale_inverses_list.emplace_back(std::move(B_row_scales)); + swizzled_scale_inverses_list.emplace_back(std::move(B_col_scales)); // Emulate the FP8 block scaling recipe with MXFP8 on Blackwell and newer // as it is not natively supported by cublasLt @@ -501,9 +504,9 @@ std::optional> te_general_grouped_gemm( // Optionally swizzle the scaling factors swizzled_scale_inverses_list.emplace_back( - multi_tensor_swizzle_scaling_factors(te_A_wrappers, transa)); + multi_tensor_swizzle_scales_for_gemm(te_A_wrappers, transa, !transa)); swizzled_scale_inverses_list.emplace_back( - multi_tensor_swizzle_scaling_factors(te_B_wrappers, !transb)); + multi_tensor_swizzle_scales_for_gemm(te_B_wrappers, !transb, transb)); // Emulate the FP8 block scaling recipe with MXFP8 on Blackwell and newer // as it is not natively supported by cublasLt diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index d7a07724c1..3214c3a9db 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -89,14 +89,8 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe TensorWrapper mu_nvte = makeTransformerEngineTensor(mu_py); TensorWrapper rsigma_nvte = makeTransformerEngineTensor(rsigma_py); - // Output tensor + // Quantizer auto quantizer_cpp = convert_quantizer(quantizer); - TensorWrapper out_nvte; - if (out.is_none()) { - std::tie(out_nvte, out) = quantizer_cpp->create_tensor(shape, out_dtype); - } else { - out_nvte = makeTransformerEngineTensor(out, quantizer); - } // Choose implementation enum class Impl { @@ -135,6 +129,19 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe } } + // Output tensor + TensorWrapper out_nvte; + if (out.is_none()) { + if (impl == Impl::FULLY_FUSED) { + // FP8 has no special logic to optimize for GEMM, MXFP8 cuDNN + // kernel does not support GEMM swizzled scales + quantizer_cpp->optimize_for_gemm = false; + } + std::tie(out_nvte, out) = quantizer_cpp->create_tensor(shape, out_dtype); + } else { + out_nvte = makeTransformerEngineTensor(out, quantizer); + } + // Construct unquantized output tensor if needed TensorWrapper unquantized_out_nvte; py::object unquantized_out; @@ -318,14 +325,8 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w at::Tensor rsigma_py = at::empty({static_cast(outer_size)}, at::CUDA(at::kFloat)); TensorWrapper rsigma_nvte = makeTransformerEngineTensor(rsigma_py); - // Output tensor + // Quantizer auto quantizer_cpp = convert_quantizer(quantizer); - TensorWrapper out_nvte; - if (out.is_none()) { - std::tie(out_nvte, out) = quantizer_cpp->create_tensor(shape, out_dtype); - } else { - out_nvte = makeTransformerEngineTensor(out, quantizer); - } // Choose implementation enum class Impl { @@ -364,6 +365,19 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w } } + // Output tensor + TensorWrapper out_nvte; + if (out.is_none()) { + if (impl == Impl::FULLY_FUSED) { + // FP8 has no special logic to optimize for GEMM, MXFP8 cuDNN + // kernel does not support GEMM swizzled scales + quantizer_cpp->optimize_for_gemm = false; + } + std::tie(out_nvte, out) = quantizer_cpp->create_tensor(shape, out_dtype); + } else { + out_nvte = makeTransformerEngineTensor(out, quantizer); + } + // Construct unquantized output tensor if needed TensorWrapper unquantized_out_nvte; py::object unquantized_out; diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c5c8905294..79dd9ea5ce 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -290,6 +290,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Fused Multi-tensor padding", py::call_guard()); m.def("fused_multi_row_unpadding", &transformer_engine::pytorch::fused_multi_row_unpadding, "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"); // attention kernels m.def("fa_prepare_fwd", &transformer_engine::pytorch::fa_prepare_fwd, diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp new file mode 100644 index 0000000000..a4750d9aa0 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -0,0 +1,394 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include "common.h" +#include "common/common.h" +#include "extensions.h" +#include "pybind.h" +#include "util.h" + +namespace transformer_engine { +namespace pytorch { + +namespace { + +void reset_tensor_data(transformer_engine::TensorWrapper &tensor, bool rowwise, bool columnwise) { + NVTEShape shape; + shape.ndim = 1; + shape.data[0] = 0; + const transformer_engine::DType dtype = transformer_engine::DType::kFloat32; + if (rowwise) { + tensor.set_rowwise_data(nullptr, dtype, shape); + tensor.set_rowwise_scale_inv(nullptr, dtype, shape); + } + if (columnwise) { + tensor.set_columnwise_data(nullptr, dtype, shape); + tensor.set_columnwise_scale_inv(nullptr, dtype, shape); + } +} + +} // namespace + +std::tuple, std::optional> swizzle_scales_for_gemm( + transformer_engine::TensorWrapper &tensor, bool rowwise_usage, bool columnwise_usage) { + // Return early if scale swizzling is not required + const auto scaling_mode = tensor.scaling_mode(); + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + case NVTE_NVFP4_1D_SCALING: + // Tensor format requires scale swizzling + break; + case NVTE_INVALID_SCALING: + NVTE_ERROR("Invalid scaling mode for swizzling scaling factors."); + default: + // Tensor format does not require scale swizzling for GEMM + return {std::nullopt, std::nullopt}; + } + + // Return early if scales are already swizzled + if (tensor.get_with_gemm_swizzled_scales()) { + return {std::nullopt, std::nullopt}; + } + + // CUDA stream + auto stream = at::cuda::getCurrentCUDAStream(); + + // Swizzle row-wise scales if needed + std::optional rowwise_scales_pyt; + if (rowwise_usage) { + // Buffer for unswizzled scales + const auto input_scales_nvte = tensor.get_rowwise_scale_inv(); + void *input_scales_dptr = input_scales_nvte.data_ptr; + const NVTEShape input_scales_shape = input_scales_nvte.shape; + const auto scales_dtype = static_cast(input_scales_nvte.dtype); + + // Allocate buffer for swizzled scales + const NVTEShape output_scales_shape = input_scales_shape; + rowwise_scales_pyt = allocateSpace(input_scales_shape, scales_dtype, false); + void *output_scales_dptr = getDataPtr(*rowwise_scales_pyt); + + // Initialize TE tensors with scales + const auto data_nvte = tensor.get_rowwise_data(); + const auto data_dtype = static_cast(data_nvte.dtype); + TensorWrapper input_nvte(scaling_mode); + input_nvte.set_rowwise_data(nullptr, data_dtype, data_nvte.shape); + input_nvte.set_rowwise_scale_inv(input_scales_dptr, scales_dtype, input_scales_shape); + TensorWrapper output_nvte(scaling_mode); + output_nvte.set_rowwise_data(nullptr, data_dtype, data_nvte.shape); + output_nvte.set_rowwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape); + output_nvte.set_with_gemm_swizzled_scales(true); + + // Launch kernel + NVTE_SCOPED_GIL_RELEASE( + { nvte_swizzle_scaling_factors(input_nvte.data(), output_nvte.data(), stream); }); + + // Update tensor with swizzled scales + tensor.set_rowwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape); + } + + // Swizzle column-wise scales if needed + std::optional columnwise_scales_pyt; + if (columnwise_usage) { + // Buffer for unswizzled scales + const auto input_scales_nvte = tensor.get_columnwise_scale_inv(); + void *input_scales_dptr = input_scales_nvte.data_ptr; + const NVTEShape input_scales_shape = input_scales_nvte.shape; + const auto scales_dtype = static_cast(input_scales_nvte.dtype); + + // Allocate buffer for swizzled scales + const NVTEShape output_scales_shape = input_scales_shape; + columnwise_scales_pyt = allocateSpace(input_scales_shape, scales_dtype, false); + void *output_scales_dptr = getDataPtr(*columnwise_scales_pyt); + + // Initialize TE tensors with scales + const auto data_nvte = tensor.get_columnwise_data(); + const auto data_dtype = static_cast(data_nvte.dtype); + TensorWrapper input_nvte(scaling_mode); + input_nvte.set_columnwise_data(nullptr, data_dtype, data_nvte.shape); + input_nvte.set_columnwise_scale_inv(input_scales_dptr, scales_dtype, input_scales_shape); + TensorWrapper output_nvte(scaling_mode); + output_nvte.set_columnwise_data(nullptr, data_dtype, data_nvte.shape); + output_nvte.set_columnwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape); + output_nvte.set_with_gemm_swizzled_scales(true); + + // Launch kernel + NVTE_SCOPED_GIL_RELEASE( + { nvte_swizzle_scaling_factors(input_nvte.data(), output_nvte.data(), stream); }); + + // Update tensor with swizzled scales + tensor.set_columnwise_scale_inv(output_scales_dptr, scales_dtype, output_scales_shape); + } + + // Update tensor + reset_tensor_data(tensor, !rowwise_usage, !columnwise_usage); + tensor.set_with_gemm_swizzled_scales(true); + + return {std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)}; +} + +std::optional multi_tensor_swizzle_scales_for_gemm( + std::vector &tensors, bool rowwise_usage, + bool columnwise_usage) { + // Checks and trivial cases + NVTE_CHECK(rowwise_usage != columnwise_usage, + "Expect exactly one of rowwise_usage=", rowwise_usage, + " and columnwise_usage=", columnwise_usage, "."); + if (tensors.empty()) { + return std::nullopt; + } + const auto scaling_mode = tensors.front().scaling_mode(); + for (const auto &tensor : tensors) { + NVTE_CHECK(tensor.scaling_mode() == scaling_mode, "Tensors have different scaling modes"); + } + + // Return early if scale swizzling is not required + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + case NVTE_NVFP4_1D_SCALING: + // Tensor format requires scale swizzling + break; + case NVTE_INVALID_SCALING: + NVTE_ERROR("Invalid scaling mode for swizzling scaling factors."); + default: + // Tensor format does not require scale swizzling for GEMM + return std::nullopt; + } + + // Filter out tensors that already have swizzled scales + std::vector tensors_needing_swizzle; + for (auto &tensor : tensors) { + if (!tensor.get_with_gemm_swizzled_scales()) { + tensors_needing_swizzle.push_back(&tensor); + } + } + if (tensors_needing_swizzle.empty()) { + return std::nullopt; + } + + // Determine buffer size needed for swizzled scales + std::vector output_scales_offsets; + size_t output_scales_bytes = 0; + for (auto &tensor : tensors_needing_swizzle) { + const auto scales_nvte = + (rowwise_usage ? tensor->get_rowwise_scale_inv() : tensor->get_columnwise_scale_inv()); + const auto &shape = scales_nvte.shape; + const auto dtype = static_cast(scales_nvte.dtype); + const auto dtype_bits = transformer_engine::pytorch::typeToNumBits(dtype); + const auto size = product(shape, 0, shape.ndim); + output_scales_bytes = roundup(output_scales_bytes, 16); // align to 16B + output_scales_offsets.push_back(output_scales_bytes); + output_scales_bytes += ceildiv(size * dtype_bits, 8); + } + + // Allocate buffer for swizzled scales + auto output_scales_pyt = allocateSpace(std::vector{output_scales_bytes}, + 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) { + 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); + 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); + } 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); + } + } + + // 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({ + nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte_raw.data(), outputs_nvte_raw.data(), + inputs_nvte_raw.size(), + at::cuda::getCurrentCUDAStream()); + }); + + // Update tensors with swizzled scales + for (size_t i = 0; i < tensors_needing_swizzle.size(); ++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); + } 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); + } + } + + return std::move(output_scales_pyt); +} + +at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapper &input, + bool rowwise) { + // Check input tensor + const NVTEScalingMode scaling_mode = input.scaling_mode(); + NVTE_CHECK(scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D, + "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]; + } + } + NVTEShape data_shape{}; + data_shape.data[0] = data_flat_first_dim; + data_shape.data[1] = data_flat_last_dim; + data_shape.ndim = 2; + + // Recreate input tensor with rowwise usage + transformer_engine::TensorWrapper input_cu(scaling_mode); + input_cu.set_rowwise_data(data.data_ptr, input.dtype(), data_shape); + const NVTEBasicTensor scale_inv = + rowwise ? input.get_rowwise_scale_inv() : input.get_columnwise_scale_inv(); + input_cu.set_rowwise_scale_inv( + scale_inv.data_ptr, static_cast(scale_inv.dtype), scale_inv.shape); + + // Create output tensor + transformer_engine::TensorWrapper output_cu(NVTE_MXFP8_1D_SCALING); + output_cu.set_rowwise_data(data.data_ptr, input.dtype(), data_shape); + // Output swizzled mxfp8 scaling factor dimensions + const size_t swizzled_scale_inv_first_dim = ceildiv(data_flat_first_dim, 128) * 128; + const size_t swizzled_scale_inv_last_dim = ceildiv(data_flat_last_dim, 128) * 4; + // Allocate memory for swizzled mxfp8 scaling factors + at::Tensor swizzled_scale_inv = + allocateSpace(std::vector{swizzled_scale_inv_first_dim, swizzled_scale_inv_last_dim}, + transformer_engine::DType::kByte, false); + // Set rowwise scaling factors on output + void *const swizzled_scale_inv_dptr = getDataPtr(swizzled_scale_inv, 0); + NVTEShape swizzled_scale_inv_shape{}; + swizzled_scale_inv_shape.data[0] = swizzled_scale_inv_first_dim; + swizzled_scale_inv_shape.data[1] = swizzled_scale_inv_last_dim; + swizzled_scale_inv_shape.ndim = 2; + output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, transformer_engine::DType::kFloat8E8M0, + swizzled_scale_inv_shape); + output_cu.set_with_gemm_swizzled_scales(true); + + // Convert scaling factors from FP8 block scaling GEMM_READY format to mxfp8 swizzled format + NVTE_SCOPED_GIL_RELEASE({ + nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(input_cu.data(), output_cu.data(), + at::cuda::getCurrentCUDAStream()); + }); + + // Set the input tensor to be the converted mxfp8 tensor and return the swizzled scaling factor + // for it to be kept alive during the GEMM + input = std::move(output_cu); + return swizzled_scale_inv; +} + +void inplace_swizzle_scale_for_gemm(py::handle &tensor) { + // Convert Python tensor to C++ tensor + auto tensor_nvte = makeTransformerEngineTensor(tensor, py::none()); + + // Return early if scale swizzling is not required + const auto scaling_mode = tensor_nvte.scaling_mode(); + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + case NVTE_NVFP4_1D_SCALING: + // Tensor format requires scale swizzling + break; + case NVTE_INVALID_SCALING: + NVTE_ERROR("Invalid scaling mode for swizzling scaling factors."); + default: + // Tensor format does not require scale swizzling for GEMM + return; + } + + // Return early if scales are already swizzled + if (tensor_nvte.get_with_gemm_swizzled_scales()) { + return; + } + + // Check what scaling factors the tensor contains + auto is_empty = [](const NVTEBasicTensor &t) -> bool { + return t.shape.ndim == 1 && t.shape.data[0] == 0; + }; + const bool has_rowwise_scales = !is_empty(tensor_nvte.get_rowwise_scale_inv()); + const bool has_columnwise_scales = !is_empty(tensor_nvte.get_columnwise_scale_inv()); + + // Swizzle scaling factors + auto [rowwise_scales, columnwise_scales] = + swizzle_scales_for_gemm(tensor_nvte, has_rowwise_scales, has_columnwise_scales); + + // Update Python tensor with swizzled scales + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + if (has_rowwise_scales) { + tensor.attr("_rowwise_scale_inv") = rowwise_scales; + } + if (has_columnwise_scales) { + tensor.attr("_columnwise_scale_inv") = columnwise_scales; + } + tensor.attr("_with_gemm_swizzled_scales") = true; + break; + case NVTE_NVFP4_1D_SCALING: + if (has_rowwise_scales) { + tensor.attr("_rowwise_scale_inv") = rowwise_scales; + } + if (has_columnwise_scales) { + tensor.attr("_columnwise_scale_inv") = columnwise_scales; + } + tensor.attr("_with_gemm_swizzled_scales") = true; + break; + default: + NVTE_ERROR("Invalid scaling mode for swizzling scaling factors."); + } +} + +} // namespace pytorch +} // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index a73efc008a..1c968e276d 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -52,10 +52,12 @@ Quantizer::Quantizer(const py::handle& quantizer) { this->rowwise_usage = true; this->columnwise_usage = true; this->internal = false; + this->optimize_for_gemm = false; } else { this->rowwise_usage = quantizer.attr("rowwise_usage").cast(); this->columnwise_usage = quantizer.attr("columnwise_usage").cast(); this->internal = quantizer.attr("internal").cast(); + this->optimize_for_gemm = quantizer.attr("optimize_for_gemm").cast(); this->quantizer = quantizer; } } @@ -555,7 +557,6 @@ Float8BlockQuantizer::Float8BlockQuantizer(const py::handle& quantizer) : Quanti this->amax_epsilon = quantizer.attr("amax_epsilon").cast(); NVTE_CHECK(this->block_scaling_dim == 1 || this->block_scaling_dim == 2, "Unsupported block scaling dim."); - this->all_gather_usage = quantizer.attr("all_gather_usage").cast(); } void Float8BlockQuantizer::set_quantization_params(TensorWrapper* tensor) const {} @@ -575,10 +576,6 @@ std::pair Float8BlockQuantizer::create_tensor( opts = opts.dtype(torch::kUInt8).device(torch::kCUDA); scale_opts = scale_opts.dtype(torch::kFloat32).device(torch::kCUDA); - Float8BlockScaleTensorFormat data_format = - (all_gather_usage ? Float8BlockScaleTensorFormat::COMPACT - : Float8BlockScaleTensorFormat::GEMM_READY); - if (rowwise_usage) { data_rowwise = at::empty(torch_shape, opts); auto scale_shape = get_scale_shape(shape, false); @@ -597,21 +594,13 @@ std::pair Float8BlockQuantizer::create_tensor( NVTE_CHECK(torch_shape.size() == shape.size(), "Shape expected to match torch shape. Shape ", columnwise_shape, " torch shape: ", torch_columnwise_shape); if (torch_shape.size() > 0) { - if (!all_gather_usage) { - torch_columnwise_shape.reserve(torch_shape.size()); - columnwise_shape.reserve(shape.size()); - torch_columnwise_shape.push_back(torch_shape[torch_shape.size() - 1]); - columnwise_shape.push_back(shape[shape.size() - 1]); - for (size_t i = 0; i < torch_shape.size() - 1; ++i) { - torch_columnwise_shape.push_back(torch_shape[i]); - columnwise_shape.push_back(shape[i]); - } - } else { - // assert we are doing 1D scaling - NVTE_CHECK(block_scaling_dim == 1, - "Compact columnwise format is not supported for 128x128 2D block scaling."); - torch_columnwise_shape = torch_shape; - columnwise_shape = shape; + torch_columnwise_shape.reserve(torch_shape.size()); + columnwise_shape.reserve(shape.size()); + torch_columnwise_shape.push_back(torch_shape[torch_shape.size() - 1]); + columnwise_shape.push_back(shape[shape.size() - 1]); + for (size_t i = 0; i < torch_shape.size() - 1; ++i) { + torch_columnwise_shape.push_back(torch_shape[i]); + columnwise_shape.push_back(shape[i]); } } auto scale_shape = get_scale_shape(shape, true); @@ -635,7 +624,7 @@ std::pair Float8BlockQuantizer::create_tensor( "rowwise_data"_a = data_rowwise, "columnwise_data"_a = data_colwise, "rowwise_scale_inv"_a = scale_inv_rowwise, "columnwise_scale_inv"_a = scale_inv_colwise, "fp8_dtype"_a = this->dtype, "quantizer"_a = this->quantizer, - "is_2D_scaled"_a = (block_scaling_dim == 2), "data_format"_a = data_format); + "is_2D_scaled"_a = (block_scaling_dim == 2)); } else { py::handle Float8BlockwiseQTensorClass( reinterpret_cast(Float8BlockwiseQTensorPythonClass)); @@ -643,8 +632,7 @@ std::pair Float8BlockQuantizer::create_tensor( "shape"_a = torch_shape, "dtype"_a = GetATenDType(dtype), "rowwise_data"_a = data_rowwise, "columnwise_data"_a = data_colwise, "rowwise_scale_inv"_a = scale_inv_rowwise, "columnwise_scale_inv"_a = scale_inv_colwise, "fp8_dtype"_a = this->dtype, - "quantizer"_a = this->quantizer, "is_2D_scaled"_a = (block_scaling_dim == 2), - "data_format"_a = data_format); + "quantizer"_a = this->quantizer, "is_2D_scaled"_a = (block_scaling_dim == 2)); } return {std::move(tensor), std::move(ret)}; @@ -654,6 +642,7 @@ std::pair Float8BlockQuantizer::convert_and_update_te py::object tensor) const { const DType dtype = tensor.attr("_fp8_dtype").cast(); bool is_2D_scaled = tensor.attr("_is_2D_scaled").cast(); + const bool with_gemm_swizzled_scales = true; // Extract buffers from Python tensor auto get_tensor = [&tensor](const char* name) -> std::optional { @@ -675,13 +664,10 @@ std::pair Float8BlockQuantizer::convert_and_update_te opts = opts.dtype(torch::kUInt8).device(torch::kCUDA); scale_opts = scale_opts.dtype(torch::kFloat32).device(torch::kCUDA); - auto get_columnwise_shape = [&columnwise_data](bool all_gather_usage) -> std::vector { + auto get_columnwise_shape = [&columnwise_data]() -> std::vector { if (!columnwise_data) { return std::vector(); } - if (all_gather_usage) { - return getTensorShape(*columnwise_data); - } std::vector shape = getTensorShape(*columnwise_data); std::vector shape_transposed(shape.size()); for (size_t i = 0; i + 1 < shape.size(); ++i) { @@ -696,12 +682,12 @@ std::pair Float8BlockQuantizer::convert_and_update_te if (rowwise_data) { shape = getTensorShape(*rowwise_data); if (columnwise_data) { - auto expected_shape = get_columnwise_shape(all_gather_usage); + auto expected_shape = get_columnwise_shape(); NVTE_CHECK(shape == expected_shape, "BlockwiseFP8 row-wise data (shape=", shape, ") and column-wise data (shape=", expected_shape, ") do not match"); } } else { - shape = get_columnwise_shape(all_gather_usage); + shape = get_columnwise_shape(); } std::vector torch_shape; for (auto s : shape) { @@ -738,21 +724,13 @@ std::pair Float8BlockQuantizer::convert_and_update_te std::vector columnwise_shape; std::vector torch_columnwise_shape; if (torch_shape.size() > 0) { - if (!all_gather_usage) { - torch_columnwise_shape.reserve(torch_shape.size()); - columnwise_shape.reserve(shape.size()); - torch_columnwise_shape.push_back(torch_shape[torch_shape.size() - 1]); - columnwise_shape.push_back(shape[shape.size() - 1]); - for (size_t i = 0; i < torch_shape.size() - 1; ++i) { - torch_columnwise_shape.push_back(torch_shape[i]); - columnwise_shape.push_back(shape[i]); - } - } else { - // assert we are doing 1D scaling - NVTE_CHECK(block_scaling_dim == 1, - "Compact columnwise format is not supported for 128x128 2D block scaling."); - torch_columnwise_shape = torch_shape; - columnwise_shape = shape; + torch_columnwise_shape.reserve(torch_shape.size()); + columnwise_shape.reserve(shape.size()); + torch_columnwise_shape.push_back(torch_shape[torch_shape.size() - 1]); + columnwise_shape.push_back(shape[shape.size() - 1]); + for (size_t i = 0; i < torch_shape.size() - 1; ++i) { + torch_columnwise_shape.push_back(torch_shape[i]); + columnwise_shape.push_back(shape[i]); } } if (!columnwise_data) { @@ -798,6 +776,7 @@ std::pair Float8BlockQuantizer::convert_and_update_te const auto scale_inv_colwise_shape = getTensorShape(scale_inv_colwise); ret.set_columnwise_scale_inv(scale_inv_colwise_dptr, DType::kFloat32, scale_inv_colwise_shape); } + ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); set_quantization_params(&ret); return {std::move(ret), std::move(tensor)}; } @@ -813,9 +792,6 @@ void Float8BlockQuantizer::quantize(const TensorWrapper& input, TensorWrapper& o } quant_config.set_force_pow_2_scales(force_pow_2_scales); quant_config.set_amax_epsilon(amax_epsilon); - if (all_gather_usage) { - quant_config.set_float8_block_scale_tensor_format(Float8BlockScaleTensorFormat::COMPACT); - } NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, at::cuda::getCurrentCUDAStream()); }); @@ -832,10 +808,6 @@ std::vector Float8BlockQuantizer::get_scale_shape(const std::vector scale_shape; bool rowwise_usage = !columnwise; @@ -845,26 +817,17 @@ std::vector Float8BlockQuantizer::get_scale_shape(const std::vector Float8BlockQuantizer::get_scale_shape(const std::vector MXFP8Quantizer::create_tensor(const std::ve DType dtype) const { using namespace pybind11::literals; + // Scaling factor format + const bool with_gemm_swizzled_scales = this->optimize_for_gemm; + // Tensor dimensions const std::vector shape_int64(shape.begin(), shape.end()); size_t flat_first_dim = 1; @@ -951,19 +909,17 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve py::object out_py; if (internal) { py::handle MXFP8TensorClass(reinterpret_cast(MXFP8TensorStoragePythonClass)); - out_py = MXFP8TensorClass("rowwise_data"_a = rowwise_data_py, - "columnwise_data"_a = columnwise_data_py, - "rowwise_scale_inv"_a = rowwise_scale_inv_py, - "columnwise_scale_inv"_a = columnwise_scale_inv_py, - "fp8_dtype"_a = this->dtype, "quantizer"_a = this->quantizer); + out_py = MXFP8TensorClass(rowwise_data_py, rowwise_scale_inv_py, columnwise_data_py, + columnwise_scale_inv_py, this->dtype, this->quantizer, + with_gemm_swizzled_scales); } else { py::handle MXFP8TensorClass(reinterpret_cast(MXFP8TensorPythonClass)); - out_py = MXFP8TensorClass("shape"_a = shape_int64, "dtype"_a = GetATenDType(dtype), - "rowwise_data"_a = rowwise_data_py, - "columnwise_data"_a = columnwise_data_py, - "rowwise_scale_inv"_a = rowwise_scale_inv_py, - "columnwise_scale_inv"_a = columnwise_scale_inv_py, - "fp8_dtype"_a = this->dtype, "quantizer"_a = this->quantizer); + out_py = MXFP8TensorClass( + "shape"_a = shape_int64, "dtype"_a = GetATenDType(dtype), + "rowwise_data"_a = rowwise_data_py, "columnwise_data"_a = columnwise_data_py, + "rowwise_scale_inv"_a = rowwise_scale_inv_py, + "columnwise_scale_inv"_a = columnwise_scale_inv_py, "fp8_dtype"_a = this->dtype, + "quantizer"_a = this->quantizer, "with_gemm_swizzled_scales"_a = with_gemm_swizzled_scales); } // Construct C++ MXFP8 tensor @@ -978,6 +934,7 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve out_cpp.set_columnwise_scale_inv(columnwise_scale_inv_tensor.data_ptr(), DType::kFloat8E8M0, columnwise_scale_inv_shape); } + out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(out_py)}; @@ -987,6 +944,9 @@ std::pair MXFP8Quantizer::convert_and_update_tensor( py::object tensor) const { NVTE_CHECK(detail::IsMXFP8Tensor(tensor.ptr()), "MXFP8Quantizer must output to MXFP8Tensor."); + // Scaling factor format + const bool with_gemm_swizzled_scales = this->optimize_for_gemm; + // Extract buffers from Python tensor auto get_tensor = [&tensor](const char* name) -> std::optional { auto attr_py = tensor.attr(name); @@ -1070,6 +1030,7 @@ std::pair MXFP8Quantizer::convert_and_update_tensor( // Coerce other attrs tensor.attr("_fp8_dtype") = dtype; + tensor.attr("_with_gemm_swizzled_scales") = with_gemm_swizzled_scales; // Construct C++ MXFP8 tensor TensorWrapper out_cpp(NVTE_MXFP8_1D_SCALING); @@ -1083,6 +1044,7 @@ std::pair MXFP8Quantizer::convert_and_update_tensor( out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E8M0, getTensorShape(*columnwise_scale_inv)); } + out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(tensor)}; @@ -1173,6 +1135,9 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve DType dtype) const { using namespace pybind11::literals; + // Scaling factor format + const bool with_gemm_swizzled_scales = false; /// TODO (tmoon) self->optimize_for_gemm + // Tensor dimensions const std::vector shape_int64(shape.begin(), shape.end()); size_t flat_first_dim = 1; @@ -1235,12 +1200,9 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve py::object out_py; if (internal) { py::handle NVFP4TensorClass(reinterpret_cast(NVFP4TensorStoragePythonClass)); - out_py = NVFP4TensorClass( - "rowwise_data"_a = rowwise_data_py, "columnwise_data"_a = columnwise_data_py, - "rowwise_scale_inv"_a = rowwise_scale_inv_py, - "columnwise_scale_inv"_a = columnwise_scale_inv_py, "amax_rowwise"_a = amax_rowwise_py, - "amax_columnwise"_a = amax_columnwise_py, "fp4_dtype"_a = this->dtype, - "quantizer"_a = this->quantizer); + out_py = NVFP4TensorClass(rowwise_data_py, rowwise_scale_inv_py, columnwise_data_py, + columnwise_scale_inv_py, amax_rowwise_py, amax_columnwise_py, + this->dtype, this->quantizer, with_gemm_swizzled_scales); } else { py::handle NVFP4TensorClass(reinterpret_cast(NVFP4TensorPythonClass)); out_py = NVFP4TensorClass( @@ -1249,7 +1211,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve "rowwise_scale_inv"_a = rowwise_scale_inv_py, "columnwise_scale_inv"_a = columnwise_scale_inv_py, "amax_rowwise"_a = amax_rowwise_py, "amax_columnwise"_a = amax_columnwise_py, "fp4_dtype"_a = this->dtype, - "quantizer"_a = this->quantizer); + "quantizer"_a = this->quantizer, "with_gemm_swizzled_scales"_a = with_gemm_swizzled_scales); } // Construct C++ tensor @@ -1272,6 +1234,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve out_cpp.set_columnwise_amax(amax_columnwise.data_ptr(), DType::kFloat32, std::vector{1}); } + out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(out_py)}; @@ -1301,6 +1264,9 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( py::object tensor) const { NVTE_CHECK(detail::IsNVFP4Tensor(tensor.ptr()), "NVFP4Quantizer must output to IsNVFP4Tensor."); + // Scaling factor format + const bool with_gemm_swizzled_scales = false; // TODO (tmoon) Enable with optimize_for_gemm + // Extract buffers from Python tensor auto get_tensor = [&tensor](const char* name) -> std::optional { auto attr_py = tensor.attr(name); @@ -1438,6 +1404,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( out_cpp.set_columnwise_amax(amax_columnwise->data_ptr(), DType::kFloat32, std::vector{1}); } + out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(tensor)}; diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 16171054cb..3f998bb66f 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -55,8 +55,9 @@ TensorWrapper NVTETensorFromFloat8Tensor(py::handle tensor, Quantizer *quantizer TensorWrapper NVTETensorFromMXFP8Tensor(py::handle tensor, Quantizer *quantizer) { auto ret = TensorWrapper(NVTE_MXFP8_1D_SCALING); - bool rowwise_usage = !(tensor.attr("_rowwise_data").is_none()); - bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); + 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(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for MXFP8 Tensor."); @@ -78,6 +79,9 @@ TensorWrapper NVTETensorFromMXFP8Tensor(py::handle tensor, Quantizer *quantizer) getTensorShape(scale_inv)); } + // Scale layout + ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + // Quantizer state quantizer->set_quantization_params(&ret); @@ -93,6 +97,7 @@ TensorWrapper NVTETensorFromFloat8BlockwiseQTensor(py::handle tensor, Quantizer auto ret = TensorWrapper(is_2D_scaled ? NVTE_BLOCK_SCALING_2D : NVTE_BLOCK_SCALING_1D); + // Row-wise data if (rowwise_usage) { const at::Tensor &data_rowwise = tensor.attr("_rowwise_data").cast(); const at::Tensor &scale_inv_rowwise = tensor.attr("_rowwise_scale_inv").cast(); @@ -102,6 +107,8 @@ TensorWrapper NVTETensorFromFloat8BlockwiseQTensor(py::handle tensor, Quantizer const auto scale_inv_rowwise_shape = getTensorShape(scale_inv_rowwise); ret.set_rowwise_scale_inv(scale_inv_rowwise_dptr, DType::kFloat32, scale_inv_rowwise_shape); } + + // Column-wise data if (columnwise_usage) { const at::Tensor &data_colwise = tensor.attr("_columnwise_data").cast(); const at::Tensor &scale_inv_colwise = tensor.attr("_columnwise_scale_inv").cast(); @@ -112,7 +119,10 @@ TensorWrapper NVTETensorFromFloat8BlockwiseQTensor(py::handle tensor, Quantizer const auto scale_inv_colwise_shape = getTensorShape(scale_inv_colwise); ret.set_columnwise_scale_inv(scale_inv_colwise_dptr, DType::kFloat32, scale_inv_colwise_shape); } + + // Quantizer state quantizer->set_quantization_params(&ret); + return ret; } @@ -121,8 +131,9 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) auto ret = TensorWrapper(NVTE_NVFP4_1D_SCALING); - bool rowwise_usage = !(tensor.attr("_rowwise_data").is_none()); - bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); + 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(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for NVFP4 Tensor."); @@ -150,6 +161,9 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) getTensorShape(amax_columnwise)); } + // Scale layout + ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + // Quantizer state quantizer->set_quantization_params(&ret); diff --git a/transformer_engine/pytorch/csrc/util.cpp b/transformer_engine/pytorch/csrc/util.cpp deleted file mode 100644 index 96fd2ccb3a..0000000000 --- a/transformer_engine/pytorch/csrc/util.cpp +++ /dev/null @@ -1,263 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include "util.h" - -#include "common.h" -#include "common/common.h" - -std::optional swizzle_scaling_factors(transformer_engine::TensorWrapper& input, - bool rowwise) { - using namespace transformer_engine::pytorch; - - if (input.scaling_mode() == NVTE_INVALID_SCALING) { - NVTE_ERROR("Invalid scaling mode for swizzle."); - } else if (input.scaling_mode() != NVTE_MXFP8_1D_SCALING && - input.scaling_mode() != NVTE_NVFP4_1D_SCALING) { - return std::nullopt; - } - - NVTE_CHECK(input.element_size_bits() == 4 || input.element_size_bits() == 8, - "4-bit or 8-bit input required for swizzling scaling factors."); - - const auto nvfp4 = input.scaling_mode() == NVTE_NVFP4_1D_SCALING; - - NVTEBasicTensor scale_inv; - NVTEShape nvte_input_shape; - if (rowwise) { - nvte_input_shape = input.shape(); - scale_inv = input.get_rowwise_scale_inv(); - } else { - nvte_input_shape = input.get_columnwise_data().shape; - scale_inv = input.get_columnwise_scale_inv(); - } - - auto input_shape = nvte_shape_to_vector(nvte_input_shape); - auto scale_inv_shape = nvte_shape_to_vector(scale_inv.shape); - - NVTE_CHECK(input_shape.size() >= 2, "Wrong ndims for swizzle input shape."); - - // Allocate memory for swizzled output. - auto options = at::TensorOptions().dtype(torch::kByte).device(torch::kCUDA); - std::vector scale_inv_shape_int; - for (size_t i = 0; i < scale_inv_shape.size(); ++i) { - scale_inv_shape_int.push_back(static_cast(scale_inv_shape[i])); - } - auto swizzled_scale_inv = at::empty(scale_inv_shape_int, options); - void* scale_inv_dptr = scale_inv.data_ptr; - void* swizzled_scale_inv_dptr = getDataPtr(swizzled_scale_inv, 0); - - transformer_engine::TensorWrapper input_cu(input.scaling_mode()); - transformer_engine::TensorWrapper output_cu(input.scaling_mode()); - - const auto input_dtype = - (nvfp4) ? transformer_engine::DType::kFloat4E2M1 : transformer_engine::DType::kFloat8E4M3; - const auto scale_inv_dtype = - (nvfp4) ? transformer_engine::DType::kFloat8E4M3 : transformer_engine::DType::kFloat8E8M0; - - if (rowwise) { - input_cu.set_rowwise_data(input.dptr(), input_dtype, input_shape); - input_cu.set_rowwise_scale_inv(scale_inv_dptr, scale_inv_dtype, scale_inv_shape); - output_cu.set_rowwise_data(input.dptr(), input_dtype, input_shape); - output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, scale_inv_shape); - } else { - input_cu.set_columnwise_data(input.columnwise_dptr(), input_dtype, input_shape); - input_cu.set_columnwise_scale_inv(scale_inv_dptr, scale_inv_dtype, scale_inv_shape); - output_cu.set_columnwise_data(input.columnwise_dptr(), input_dtype, input_shape); - output_cu.set_columnwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, scale_inv_shape); - } - - // Launch kernel - nvte_swizzle_scaling_factors(input_cu.data(), output_cu.data(), at::cuda::getCurrentCUDAStream()); - - if (rowwise) { - input.set_rowwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, scale_inv_shape); - } else { - input.set_columnwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, scale_inv_shape); - } - - return swizzled_scale_inv; -} - -std::optional multi_tensor_swizzle_scaling_factors( - std::vector& tensors, bool rowwise) { - using namespace transformer_engine::pytorch; - - if (tensors.empty()) { - return std::nullopt; - } - - bool all_same_scaling_mode = std::all_of( - tensors.cbegin(), tensors.cend(), [&tensors](const transformer_engine::TensorWrapper& val) { - return val.scaling_mode() == tensors.front().scaling_mode(); - }); - NVTE_CHECK(all_same_scaling_mode, "Scaling mode of the input tensors must be the same."); - - if (tensors.front().scaling_mode() == NVTE_INVALID_SCALING) { - NVTE_ERROR("Invalid scaling mode for swizzle."); - } else if (tensors.front().scaling_mode() != NVTE_MXFP8_1D_SCALING && - tensors.front().scaling_mode() != NVTE_NVFP4_1D_SCALING) { - return std::nullopt; - } - - const auto scaling_mode = tensors.front().scaling_mode(); - const auto nvfp4 = scaling_mode == NVTE_NVFP4_1D_SCALING; - - std::vector wrappers; - std::vector input_tensors, output_tensors; - - // Collect scale_inv shapes and calculate buffer size and offsets for scale_invs - std::vector> scale_inv_shapes; - std::vector scale_inv_dptrs; - size_t buffer_size = 0; - std::vector scale_inv_offsets; - constexpr size_t scale_elem_size = 1; - for (auto& tensor : tensors) { - NVTEBasicTensor scale_inv; - if (rowwise) { - scale_inv = tensor.get_rowwise_scale_inv(); - } else { - scale_inv = tensor.get_columnwise_scale_inv(); - } - auto scale_inv_shape = nvte_shape_to_vector(scale_inv.shape); - buffer_size = roundup(buffer_size, 16); // align to 16B - scale_inv_offsets.push_back(buffer_size); - buffer_size += product(scale_inv_shape) * scale_elem_size; - scale_inv_shapes.emplace_back(scale_inv_shape); - scale_inv_dptrs.push_back(scale_inv.data_ptr); - } - - // Allocate full buffer - auto buffer = at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8)); - - const auto input_dtype = - (nvfp4) ? transformer_engine::DType::kFloat4E2M1 : transformer_engine::DType::kFloat8E4M3; - const auto scale_inv_dtype = - (nvfp4) ? transformer_engine::DType::kFloat8E4M3 : transformer_engine::DType::kFloat8E8M0; - - for (size_t i = 0; i < tensors.size(); ++i) { - auto& tensor = tensors[i]; - void* scale_inv_dptr = scale_inv_dptrs[i]; - void* swizzled_scale_inv_dptr = getDataPtr(buffer, scale_inv_offsets[i]); - - // Empty tensors don't require scale swizzling - if (tensor.numel() == 0) { - continue; - } - - // Tensor shape - NVTEShape nvte_input_shape; - if (rowwise) { - nvte_input_shape = tensor.shape(); - } else { - nvte_input_shape = tensor.get_columnwise_data().shape; - } - - auto input_shape = nvte_shape_to_vector(nvte_input_shape); - // Reconstruct input only to avoid swizzling both directions if not needed. - // Use any 8 bit type, it's irrelevant. - transformer_engine::TensorWrapper input_cu(scaling_mode); - transformer_engine::TensorWrapper output_cu(scaling_mode); - if (rowwise) { - input_cu.set_rowwise_data(tensor.dptr(), input_dtype, input_shape); - input_cu.set_rowwise_scale_inv(scale_inv_dptr, scale_inv_dtype, scale_inv_shapes[i]); - output_cu.set_rowwise_data(tensor.dptr(), input_dtype, input_shape); - output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, - scale_inv_shapes[i]); - // Set the swizzled scaling factor to the original tensor. - tensor.set_rowwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, scale_inv_shapes[i]); - } else { - input_cu.set_columnwise_data(tensor.columnwise_dptr(), input_dtype, input_shape); - input_cu.set_columnwise_scale_inv(scale_inv_dptr, scale_inv_dtype, scale_inv_shapes[i]); - output_cu.set_columnwise_data(tensor.columnwise_dptr(), input_dtype, input_shape); - output_cu.set_columnwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, - scale_inv_shapes[i]); - // Set the swizzled scaling factor to the original tensor. - tensor.set_columnwise_scale_inv(swizzled_scale_inv_dptr, scale_inv_dtype, - scale_inv_shapes[i]); - } - - input_tensors.emplace_back(input_cu.data()); - output_tensors.emplace_back(output_cu.data()); - wrappers.emplace_back(std::move(input_cu)); - wrappers.emplace_back(std::move(output_cu)); - } - - // Launch kernel - nvte_multi_tensor_swizzle_scaling_factors(input_tensors.data(), output_tensors.data(), - input_tensors.size(), at::cuda::getCurrentCUDAStream()); - - return buffer; -} - -at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapper& input, - bool rowwise) { - using namespace transformer_engine::pytorch; - using transformer_engine::DIVUP; - - // Check input tensor - const NVTEScalingMode scaling_mode = input.scaling_mode(); - NVTE_CHECK(scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D, - "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]; - } - } - NVTEShape data_shape{}; - data_shape.data[0] = data_flat_first_dim; - data_shape.data[1] = data_flat_last_dim; - data_shape.ndim = 2; - - // Recreate input tensor with rowwise usage - transformer_engine::TensorWrapper input_cu(scaling_mode); - input_cu.set_rowwise_data(data.data_ptr, input.dtype(), data_shape); - const NVTEBasicTensor scale_inv = - rowwise ? input.get_rowwise_scale_inv() : input.get_columnwise_scale_inv(); - input_cu.set_rowwise_scale_inv( - scale_inv.data_ptr, static_cast(scale_inv.dtype), scale_inv.shape); - - // Create output tensor - transformer_engine::TensorWrapper output_cu(NVTE_MXFP8_1D_SCALING); - output_cu.set_rowwise_data(data.data_ptr, input.dtype(), data_shape); - // Output swizzled mxfp8 scaling factor dimensions - const size_t swizzled_scale_inv_first_dim = DIVUP(data_flat_first_dim, 128) * 128; - const size_t swizzled_scale_inv_last_dim = DIVUP(data_flat_last_dim, 128) * 4; - // Allocate memory for swizzled mxfp8 scaling factors - const auto options = at::TensorOptions().dtype(torch::kByte).device(torch::kCUDA); - at::Tensor swizzled_scale_inv = at::empty( - std::vector{swizzled_scale_inv_first_dim, swizzled_scale_inv_last_dim}, options); - // Set rowwise scaling factors on output - void* const swizzled_scale_inv_dptr = getDataPtr(swizzled_scale_inv, 0); - NVTEShape swizzled_scale_inv_shape{}; - swizzled_scale_inv_shape.data[0] = swizzled_scale_inv_first_dim; - swizzled_scale_inv_shape.data[1] = swizzled_scale_inv_last_dim; - swizzled_scale_inv_shape.ndim = 2; - output_cu.set_rowwise_scale_inv(swizzled_scale_inv_dptr, transformer_engine::DType::kFloat8E8M0, - swizzled_scale_inv_shape); - - // Convert scaling factors from FP8 block scaling GEMM_READY format to mxfp8 swizzled format - nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(input_cu.data(), output_cu.data(), - at::cuda::getCurrentCUDAStream()); - - // Set the input tensor to be the converted mxfp8 tensor and return the swizzled scaling factor - // for it to be kept alive during the GEMM - input = std::move(output_cu); - return swizzled_scale_inv; -} diff --git a/transformer_engine/pytorch/csrc/util.h b/transformer_engine/pytorch/csrc/util.h index 4d72db922e..8988c18261 100644 --- a/transformer_engine/pytorch/csrc/util.h +++ b/transformer_engine/pytorch/csrc/util.h @@ -10,33 +10,44 @@ #include #include +#include +#include #include "transformer_engine/transformer_engine.h" -/*! \brief Swizzle the scaling factor of the input tensor. +namespace transformer_engine { +namespace pytorch { + +/*! \brief Convert tensor block scales into GEMM swizzled format. * - * The returned swizzled scaling factor tensor should be kept alive during the GEMM. + * The returned swizzled scales should be kept alive during the GEMM. */ -std::optional swizzle_scaling_factors(transformer_engine::TensorWrapper &input, - bool rowwise); +std::tuple, std::optional> swizzle_scales_for_gemm( + TensorWrapper& tensor, bool rowwise_usage, bool columnwise_usage); -/*! \brief Swizzle the scaling factor of the input tensors. +/*! \brief Convert multiple tensor block scales into GEMM swizzled format. * - * The returned swizzled scaling factor tensors should be kept alive during the GEMMs. + * The returned swizzled scales should be kept alive during the GEMMs. */ -std::optional multi_tensor_swizzle_scaling_factors( - std::vector &inputs, bool rowwise); +std::optional multi_tensor_swizzle_scales_for_gemm(std::vector& tensors, + bool rowwise_usage, + bool columnwise_usage); /*! \brief Convert a block scaling tensor to an mxfp8 tensor in-place. * - * If rowwise==false, the columnwise data will be reinterpreted as rowwise data to avoid - * transposing it in memory. Due to differences in how block scaling and mxfp8 store data, - * this requires the calling code to treat the output tensor as having been tranposed in this case. + * If rowwise==false, the columnwise data will be reinterpreted as + * rowwise data to avoid transposing it in memory. Due to differences + * in how block scaling and mxfp8 store data, this requires the + * calling code to treat the output tensor as having been transposed + * in this case. * - * Returns the swizzled scaling factor of the converted mxfp8 tensor. - * The returned swizzled scaling factor tensor should be kept alive during the GEMM. + * Returns the swizzled scaling factor of the converted mxfp8 tensor. + * The returned swizzled scaling factor tensor should be kept alive + * during the GEMM. */ -at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapper &input, - bool rowwise); +at::Tensor convert_block_scaling_to_mxfp8_tensor(TensorWrapper& input, bool rowwise); + +} // namespace pytorch +} // namespace transformer_engine #endif // TRANSFORMER_ENGINE_PYTORCH_CSRC_UTIL_H_ diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 5497ee7967..004a04ab4c 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -48,7 +48,7 @@ from .tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage from .tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from .tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage -from ..debug.pytorch.debug_quantization import DebugQuantizedTensor, DebugQuantizer +from ..debug.pytorch.debug_quantization import DebugQuantizedTensor __all__ = ["checkpoint", "CudaRNGStatesTracker"] @@ -930,6 +930,34 @@ def reduce_scatter_along_first_dim( return output, handle +@dataclass +class _AsyncHandle: + """Handle for asynchronous collectives.""" + + async_handle: torch.distributed.Work + post_process_function: Optional[Callable] = None + post_process_function_args: Optional[Tuple[Any, ...]] = None + post_process_function_kwargs: Optional[Dict[str, Any]] = None + _synchronized: bool = False + + def wait(self) -> None: + """Synchronize the asynchronous communicaton. + + Perform post-processing if needed. + + """ + if self._synchronized: + return + self.async_handle.wait() + if self.post_process_function is not None: + args = self.post_process_function_args + args = () if args is None else args + kwargs = self.post_process_function_kwargs + kwargs = {} if kwargs is None else kwargs + self.post_process_function(*args, **kwargs) + self._synchronized = True + + def _all_gather_fp8( inp: torch.Tensor, process_group: dist_group_type, @@ -1020,73 +1048,7 @@ def _all_gather_fp8( return out, handle -def _get_quantizer_format(quantizer: Quantizer) -> Optional[bool]: - """Get quantizer format.""" - if isinstance(quantizer, DebugQuantizer): - quantizer = quantizer.parent_quantizer - if isinstance(quantizer, Float8BlockQuantizer): - return quantizer.all_gather_usage - return None - - -def _set_quantizer_format(quantizer: Quantizer, compact: bool = False) -> None: - """Make quantizer compact""" - _quantizer = quantizer - if isinstance(quantizer, DebugQuantizer): - _quantizer = quantizer.parent_quantizer - if isinstance(_quantizer, Float8BlockQuantizer): - _quantizer.all_gather_usage = compact - - -def _post_process_fp8_blockwise_gather( - out: Float8BlockwiseQTensorStorage, - quantizer: Float8BlockQuantizer, - handle: Optional[torch.distributed.Work] = None, -) -> Float8BlockwiseQTensorStorage: - """Post-process FP8 blockwise gather.""" - if handle is not None: - handle.wait() - handle = None - - if out._is_gemm_ready_format(): - return out - - needs_columnwise_data_transpose = quantizer is not None and quantizer.columnwise_usage - need_rowwise_scale_transpose = quantizer is not None and quantizer.rowwise_usage - - # CuBLAS requires transpose of the scale inv tensor, suppose orig input is 256x1024 - # columnwise compact format means doing 128x1 quantization of it - # so quantized tensor is 256x1024, scale inv is 2x1024 - # If we were doing GEMM_READY format, then it's equivalent to do 1x128 quantization - # on a transposed 1024x256 tensor, so scale inv is 1024x2, cublas requries 2x1024 - # Thereforce, it turns out we don't need to transpose the scale inv, only columnwise data - if needs_columnwise_data_transpose: - out._transpose_columnwise_data() - if need_rowwise_scale_transpose: - out._rowwise_scale_inv = out._rowwise_scale_inv.transpose(-2, -1).contiguous() - out._data_format = tex.Float8BlockScaleTensorFormat.GEMM_READY - return out - - -@dataclass -class _FP8BlockwiseAllGatherAsyncHandle: - """Handle for asynchronous FP8 blockwise all-gather.""" - - tensor: Float8BlockwiseQTensorStorage - quantizer: Float8BlockQuantizer - async_handle: torch.distributed.Work - _synchronized: bool = False - - def wait(self) -> None: - """Wait for the async operation to complete and post-process the tensor.""" - if self._synchronized: - return - self.async_handle.wait() - _post_process_fp8_blockwise_gather(self.tensor, self.quantizer) - self._synchronized = True - - -def _all_gather_fp8_blockwise( +def _start_all_gather_fp8_blockwise( inp: torch.Tensor, process_group: dist_group_type, *, @@ -1125,44 +1087,25 @@ def _all_gather_fp8_blockwise( ) world_size = get_distributed_world_size(process_group) - # Check that quantizer is valid - if quantizer is not None and not isinstance(quantizer, Float8BlockQuantizer): - raise ValueError(f"Got non-FP8 blockwise quantizer ({quantizer.__class__.__name__})") - if not (quantizer.block_scaling_dim == 1 and quantizer.block_len == 128): - raise NotImplementedError("Only 1D blockwise quantization is supported for allgather") - # Output tensor dims if out_shape is None: out_shape = list(inp.size()) out_shape[0] *= world_size - # Doing BF16 gather for now as baseline because it's simpler - if ( - not isinstance(inp, Float8BlockwiseQTensorStorage) - and quantizer is not None - and not quantizer.is_quantizable(inp) - ): - out = torch.empty( - out_shape, - dtype=dtype, - device=device, - memory_format=torch.contiguous_format, - ) + # Check that quantizer is valid + if quantizer is None: + raise ValueError("Quantizer is missing") + if not isinstance(quantizer, Float8BlockQuantizer): + raise ValueError(f"Got non-FP8 blockwise quantizer ({quantizer.__class__.__name__})") + + # Fall back to high-precision all-gather if FP8 is not supported + if not quantizer.is_quantizable(inp) or quantizer.block_scaling_dim != 1: + out = torch.empty(out_shape, dtype=dtype, device=device) torch.distributed.all_gather_into_tensor(out, inp, group=process_group, async_op=False) - orig_all_gather_usage = quantizer.all_gather_usage - quantizer.all_gather_usage = False out = quantizer(out) - quantizer.all_gather_usage = orig_all_gather_usage return out, None - # Implementation of fp8 gather needs to account for: - # * Getting columnwise data as a transpose of how it is stored for GEMMS. - # * Gathering non GEMM swizzled scales. - - # Cast input tensor to Float8BlockwiseQTensor with required data - # Set to compact usage in case the quantizer is not correctly configured - orig_all_gather_usage = quantizer.all_gather_usage - quantizer.all_gather_usage = True + # Quantize input tensor if needed if not isinstance(inp, Float8BlockwiseQTensorStorage): inp = quantizer(inp) elif (quantizer.rowwise_usage and inp._rowwise_data is None) or ( @@ -1177,14 +1120,9 @@ def _all_gather_fp8_blockwise( # Construct Float8BlockwiseQTensor output tensor out = quantizer.make_empty(out_shape, dtype=dtype, device=device) - quantizer.all_gather_usage = orig_all_gather_usage - - # Begin to do network communication, need to make sure compact format - if inp._data_format != tex.Float8BlockScaleTensorFormat.COMPACT: - raise RuntimeError( - "All-gather with FP8 block-wise quantized tensor requires compact data format, " - f"but found data_format={inp._data_format}" - ) + # Temporary buffers for all-gathering transposed buffers + interleaved_rowwise_scale_inv = None + interleaved_columnwise_data = None # Coalesce NCCL collectives with torch.distributed._coalescing_manager( @@ -1193,11 +1131,17 @@ def _all_gather_fp8_blockwise( async_ops=async_op, ) as coalescing_manager: - # Gather Float8BlockwiseQTensor data for row-wise usage + # Gather row-wise data if quantizer.rowwise_usage: - # Launch all-gathers + scale_inv_shape = list(inp._rowwise_scale_inv.size()) + scale_inv_shape[0] *= world_size + interleaved_rowwise_scale_inv = torch.empty( + scale_inv_shape, + dtype=inp._rowwise_scale_inv.dtype, + device=device, + ) torch.distributed.all_gather_into_tensor( - out._rowwise_scale_inv, + interleaved_rowwise_scale_inv, inp._rowwise_scale_inv, group=process_group, ) @@ -1207,36 +1151,73 @@ def _all_gather_fp8_blockwise( group=process_group, ) - # Gather Float8BlockwiseQTensor data for column-wise usage + # Column-wise data if quantizer.columnwise_usage: - # Launch all-gathers + data_shape = list(inp._columnwise_data.size()) + data_shape[0] *= world_size + interleaved_columnwise_data = torch.empty( + data_shape, + dtype=inp._columnwise_data.dtype, + device=device, + ) torch.distributed.all_gather_into_tensor( out._columnwise_scale_inv, inp._columnwise_scale_inv, group=process_group, ) torch.distributed.all_gather_into_tensor( - out._columnwise_data, + interleaved_columnwise_data, inp._columnwise_data, group=process_group, ) - handle = coalescing_manager if async_op else None - - # Unlike MXFP8, this fp8 blockwise tensor primarily works with Hopper - # This means that we need to transpose the gathered columnwise data - # Example usage is grad_output tensor, ie. dY in linear backward - # We want to gather two FP8 tensors (rowwise and columnwise) along dim0 - # and then transpose the columnwise data to match the rowwise data - # Make sure FP8 transpose is populated if needed - + # Finalize communication if needed + async_handle = None if async_op: - handle = _FP8BlockwiseAllGatherAsyncHandle(out, quantizer, handle) + async_handle = _AsyncHandle( + coalescing_manager, + post_process_function=_finish_all_gather_fp8_blockwise, + post_process_function_args=( + out, + world_size, + interleaved_rowwise_scale_inv, + interleaved_columnwise_data, + ), + ) else: - # if it's a sync op, we need to do the transpose here as post processing step - _post_process_fp8_blockwise_gather(out, quantizer, handle) + _finish_all_gather_fp8_blockwise( + out, + world_size, + interleaved_rowwise_scale_inv, + interleaved_columnwise_data, + ) - return out, handle + return out, async_handle + + +def _finish_all_gather_fp8_blockwise( + out: Float8BlockwiseQTensorStorage, + world_size: int, + interleaved_rowwise_scale_inv: Optional[torch.Tensor], + interleaved_columnwise_data: Optional[torch.Tensor], +) -> Float8BlockwiseQTensorStorage: + """Post-process FP8 blockwise gather.""" + + # Fix interleaving in row-wise scales + if interleaved_rowwise_scale_inv is not None: + dim0 = out._rowwise_scale_inv.size(0) + view_in = interleaved_rowwise_scale_inv.view(world_size, dim0, -1) + view_out = out._rowwise_scale_inv.view(dim0, world_size, -1) + tex.swap_first_dims(view_in, out=view_out) + + # Fix interleaving in column-wise data + if interleaved_columnwise_data is not None: + dim0 = out._columnwise_data.size(0) + view_in = interleaved_columnwise_data.view(world_size, dim0, -1) + view_out = out._columnwise_data.view(dim0, world_size, -1) + tex.swap_first_dims(view_in, out=view_out) + + return out def _swap_first_dims(tensor: torch.Tensor, world_size: int): @@ -1250,7 +1231,7 @@ def _swap_first_dims(tensor: torch.Tensor, world_size: int): """ shape = tensor.shape - assert tensor.ndim >= 2, "Wrong number of dimensions for fixing interleave." + assert len(shape) >= 2, "Wrong number of dimensions for fixing interleave." first_dim = shape[0] flattened_trailing = math.prod(shape[1:]) assert first_dim % world_size == 0, "Wrong dimensions for fixing interleave." @@ -1681,7 +1662,7 @@ def gather_along_first_dim( if isinstance(inp, Float8BlockwiseQTensorStorage) or isinstance( quantizer, Float8BlockQuantizer ): - return _all_gather_fp8_blockwise( + return _start_all_gather_fp8_blockwise( inp, process_group, async_op=async_op, @@ -1719,10 +1700,6 @@ def gather_along_first_dim( ) if isinstance(inp, QuantizedTensorStorage): inp = inp.dequantize() - # Falling back to high-precision all-gather for Float8BlockQuantizer - # means that it should directly output GEMM_READY format - compact = _get_quantizer_format(quantizer) - _set_quantizer_format(quantizer, compact=False) out = torch.empty( out_shape, dtype=inp.dtype, @@ -1731,7 +1708,6 @@ def gather_along_first_dim( ) torch.distributed.all_gather_into_tensor(out, inp, group=process_group) out = quantizer(out) - _set_quantizer_format(quantizer, compact=compact) return out, None # Dequantize quantized tensor if not supported diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index ad5cd04341..875d245a8f 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -560,6 +560,8 @@ def fill_userbuffers_buffer_for_all_gather( "Userbuffers requires MXFP8 tensor dims that are divisible by 128, " f"but got MXFP8 tensor with shape={tuple(local_shape)}" ) + if local_tensor._with_gemm_swizzled_scales: + raise ValueError("Userbuffers assumes MXFP8 tensors have unswizzled scales") local_scale_inv = ( local_tensor._rowwise_scale_inv if with_rowwise_data @@ -592,6 +594,7 @@ def fill_userbuffers_buffer_for_all_gather( columnwise_scale_inv=columnwise_scale_inv, fp8_dtype=local_tensor._fp8_dtype, quantizer=quantizer, + with_gemm_swizzled_scales=False, ) return global_tensor, local_tensor diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 1e6f0b00ab..e6e69b3e4a 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -720,13 +720,9 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" super().set_meta_tensor(fwd, recipe) - # customize quantizers based on each recipe & layer configs + # Recipe-specific quantizer configuration recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): - assert not self.tp_size > 1, ( - "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " - "Because the TP communication is handled outside of this module." - ) self._customize_quantizers_float8_current_scaling(fwd, recipe) def reset_parameters(self, defer_init=False): @@ -879,9 +875,12 @@ def backward_dw(self): def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: """Customize quantizers based on current scaling recipe + linear.""" - assert ( - recipe.float8_current_scaling() - ), "current scaling recipe quantizer customization here" + + assert not self.tp_size > 1, ( + "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " + "Because the TP communication is handled outside of this module." + ) + if fwd: for i in range(self.num_gemms): # set configs about amax epsilon and power_2_scale @@ -954,9 +953,9 @@ def _get_quantizers(self): ] for i in range(self.num_gemms) ] - # TODO: use internal after #1638 is merged. # pylint: disable=fixme for i in range(self.num_gemms): - input_quantizers[i].internal = False + input_quantizers[i].internal = True + input_quantizers[i].optimize_for_gemm = True if torch.is_grad_enabled(): grad_output_quantizers = [ self.quantizers["scaling_bwd"][ @@ -966,6 +965,7 @@ 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 return ( input_quantizers, weight_quantizers, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 13b94f2327..ca30ef9567 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -64,7 +64,6 @@ restore_from_saved, ) from ...debug.pytorch.debug_state import TEDebugState -from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..cpu_offload import ( is_cpu_offload_enabled, @@ -253,8 +252,6 @@ def forward( if fp8 or debug: ln_out = input_quantizer(ln_out) input_quantizer.set_usage(rowwise=True, columnwise=False) - if isinstance(input_quantizer, Float8BlockQuantizer): - input_quantizer.all_gather_usage = False ln_out_total = input_quantizer(ln_out_total) else: quantizer = None @@ -1409,15 +1406,12 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" super().set_meta_tensor(fwd, recipe) - # customize quantizers based on each recipe & layer configs + # Recipe-specific quantizer configuration recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) - elif recipe.float8_block_scaling(): - self._customize_quantizers_float8_blockwise_scaling(fwd, recipe) elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) - # elif other recipes (mxfp8, etc) def reset_layer_norm_parameters(self) -> None: """Init LN params""" @@ -1619,12 +1613,16 @@ def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): output_quantizer = None input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] input_quantizer.internal = True + if not (self.parallel_mode == "column" and self.sequence_parallel): + input_quantizer.optimize_for_gemm = True (weight_quantizer,) = self._get_weight_quantizers() if fp8_output: output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] if is_grad_enabled: grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] grad_output_quantizer.internal = True + if not (self.parallel_mode == "row" and self.sequence_parallel): + grad_output_quantizer.optimize_for_gemm = True if fp8_grad: grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] @@ -1808,14 +1806,3 @@ def _get_weight_quantizers(self) -> List[Quantizer]: weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] weight_quantizer.internal = True return [weight_quantizer] - - def _customize_quantizers_float8_blockwise_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on blockwise scaling recipe + layernorm_linear.""" - assert ( - recipe.float8_block_scaling() - ), "blockwise scaling recipe quantizer customization here" - if fwd: - if self.sequence_parallel and self.parallel_mode == "column": - self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT - ].all_gather_usage = True diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 4256028c8b..35e4522138 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -431,8 +431,6 @@ def _forward( if fp8 or debug: ln_out = fc1_input_quantizer(ln_out) fc1_input_quantizer.set_usage(rowwise=True, columnwise=False) - if isinstance(fc1_input_quantizer, Float8BlockQuantizer): - fc1_input_quantizer.all_gather_usage = False ln_out_total = fc1_input_quantizer(ln_out_total) else: quantizer = None @@ -1964,15 +1962,12 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" super().set_meta_tensor(fwd, recipe) - # customize quantizers based on each recipe & layer configs + # Recipe-specific quantizer configuration recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) - elif recipe.float8_block_scaling(): - self._customize_quantizers_float8_blockwise_scaling(fwd, recipe) elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) - # elif for other recipes (mxfp8, etc.) def reset_layer_norm_parameters(self) -> None: """Init LN params""" @@ -2193,6 +2188,8 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): if self.fp8 or self.fp8_calibration: fc1_input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] fc1_input_quantizer.internal = True + if not self.sequence_parallel: + fc1_input_quantizer.optimize_for_gemm = True fc2_input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM2_INPUT] fc2_input_quantizer.set_usage( rowwise=True, @@ -2201,7 +2198,8 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): (MXFP8Quantizer, Float8BlockQuantizer, NVFP4Quantizer), ), ) - fc1_input_quantizer.internal = True + fc2_input_quantizer.internal = True + fc2_input_quantizer.optimize_for_gemm = True if fp8_output: fc2_output_quantizer = self.quantizers["scaling_fwd"][ tex.FP8FwdTensors.GEMM2_OUTPUT @@ -2211,10 +2209,13 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): tex.FP8BwdTensors.GRAD_OUTPUT2 ] fc2_grad_output_quantizer.internal = True + if not self.sequence_parallel: + fc2_grad_output_quantizer.optimize_for_gemm = True fc1_grad_output_quantizer = self.quantizers["scaling_bwd"][ tex.FP8BwdTensors.GRAD_OUTPUT1 ] fc1_grad_output_quantizer.internal = True + fc1_grad_output_quantizer.optimize_for_gemm = True return ( fc1_input_quantizer, @@ -2467,22 +2468,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]: fc2_weight_quantizer.internal = True return [fc1_weight_quantizer, fc2_weight_quantizer] - def _customize_quantizers_float8_blockwise_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on blockwise scaling recipe + layernorm_mlp.""" - assert ( - recipe.float8_block_scaling() - ), "blockwise scaling recipe quantizer customization here" - if fwd: - if self.sequence_parallel and self.set_parallel_mode: - self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT - ].all_gather_usage = True - else: - if self.sequence_parallel and self.set_parallel_mode: - self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 - ].all_gather_usage = True - def backward_dw(self): """ Execute the delayed weight gradient computation. diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b8349f84a0..38104604d8 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1313,15 +1313,12 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" super().set_meta_tensor(fwd, recipe) - # customize quantizers based on each recipe & layer configs + # Recipe-specific quantizer configuration recipe = FP8GlobalStateManager.get_fp8_recipe() if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) - elif recipe.float8_block_scaling(): - self._customize_quantizers_float8_blockwise_scaling(fwd, recipe) elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) - # elif for other recipes (mxfp8, etc.) def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) @@ -1489,12 +1486,16 @@ def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): output_quantizer = None input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] input_quantizer.internal = True + if not (self.parallel_mode == "column" and self.sequence_parallel): + input_quantizer.optimize_for_gemm = True (weight_quantizer,) = self._get_weight_quantizers() if fp8_output: output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] if is_grad_enabled: grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] grad_output_quantizer.internal = True + if not (self.parallel_mode == "row" and self.sequence_parallel): + grad_output_quantizer.optimize_for_gemm = True if fp8_grad: grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] return ( @@ -1669,22 +1670,3 @@ def _get_weight_quantizers(self) -> List[Quantizer]: weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] weight_quantizer.internal = True return [weight_quantizer] - - def _customize_quantizers_float8_blockwise_scaling(self, fwd: bool, recipe: Recipe) -> None: - """Customize quantizers based on blockwise scaling recipe + linear.""" - assert ( - recipe.float8_block_scaling() - ), "blockwise scaling recipe quantizer customization here" - - if fwd: - if self.sequence_parallel and self.parallel_mode == "column": - # set compact for inp tensor X - self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT - ].all_gather_usage = True - else: - if self.sequence_parallel and self.parallel_mode == "row": - # set compact for grad_output tensor dY - self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 - ].all_gather_usage = True diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 2714d718fe..e640f3ffb1 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -342,15 +342,21 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: super().reset_recipe_state(recipe=recipe) - # Input/grad output quantizers use internal tensors + # Configure input/grad output tensor + # 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) grad_output_quantizer = self.get_quantizer("backward", 0) if input_quantizer is not None: input_quantizer.internal = True + if not (self.tensor_parallel_mode == "column" and self.sequence_parallel): + input_quantizer.optimize_for_gemm = True if grad_output_quantizer is not None: grad_output_quantizer.internal = True + if not (self.tensor_parallel_mode == "row" and self.sequence_parallel): + grad_output_quantizer.optimize_for_gemm = True - # Handle weight quantizer + # Configure weight quantizer # Note: This function may be called in base class constructor, # before any basic linear attrs have been set. weight_quantizer = self.get_quantizer("forward", 1) diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 4943ffb1bd..6c889ba047 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -292,6 +292,7 @@ def _functional_backward( rowwise=True, columnwise=with_columnwise, ) + grad_output_quantizer.optimize_for_gemm = False dy_local = grad_output_quantizer(dy_local) else: dy_local = maybe_dequantize(dy_local, dtype) diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index 0c16b35e11..5beeed1262 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -294,6 +294,7 @@ def forward( columnwise_scale_inv=None, quantizer=None, requires_grad=output.requires_grad, + with_gemm_swizzled_scales=False, ) ctx.save_for_backward(row_id_map, pad_offsets) @@ -504,6 +505,7 @@ def backward(ctx, unpermuted_act_grad): columnwise_scale_inv=None, quantizer=None, requires_grad=act_grad.requires_grad, + with_gemm_swizzled_scales=False, ) if not ctx.needs_input_grad[2]: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index ac827e794a..0a6ad61ff0 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -199,10 +199,21 @@ class Quantizer(abc.ABC): """ internal: bool + """Whether to solely optimize for matrix multiplication + + The resulting quantized tensors are not guaranteed to support any + operation other than matrix multiplication. Use with care since + this is likely to break communication, checkpointing, and many + other features. + + """ + optimize_for_gemm: bool + def __init__(self, *, rowwise: bool, columnwise: bool) -> None: self.rowwise_usage = rowwise self.columnwise_usage = columnwise self.internal = False + self.optimize_for_gemm = False def __repr__(self): return ( @@ -314,7 +325,11 @@ def supports_only_rowwise_all_gather(self) -> bool: return False def is_quantizable(self, inp: torch.Tensor) -> bool: # pylint: disable=unused-argument - """Returns whether or not given tensor can be quantized""" + """Whether tensor supports quantized all-gather + + Consider a less misleading function name. + + """ return True def get_usages(self) -> Dict[str, bool]: diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 03c16ebbed..ecafb6ddfc 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -4,14 +4,14 @@ """Tensor class with FP8 data quantized with NxN tiles""" from __future__ import annotations -from typing import Optional, Tuple, Iterable, Union - +from collections.abc import Iterable import math +from typing import Any, Optional, Tuple, Union + import torch + import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType -from transformer_engine_torch import Float8BlockScaleTensorFormat - from transformer_engine.common.recipe import Float8BlockScaling, Recipe from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ..quantized_tensor import QuantizedTensor, Quantizer @@ -35,8 +35,6 @@ class Float8BlockQuantizer(Quantizer): amax_epsilon: float force_pow_2_scales: bool block_scaling_dim: int - # Whether to produce tensors that will be used in all-gather - all_gather_usage: bool def __init__( self, @@ -47,7 +45,6 @@ def __init__( amax_epsilon: float = 0.0, force_pow_2_scales: bool = True, block_scaling_dim: int = 2, - all_gather_usage: bool = False, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = fp8_dtype @@ -55,7 +52,6 @@ def __init__( self.force_pow_2_scales = force_pow_2_scales self.amax_epsilon = amax_epsilon self.block_scaling_dim = block_scaling_dim - self.all_gather_usage = all_gather_usage def copy(self) -> Float8BlockQuantizer: """Create shallow copy""" @@ -65,11 +61,11 @@ def copy(self) -> Float8BlockQuantizer: rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, block_scaling_dim=self.block_scaling_dim, - all_gather_usage=self.all_gather_usage, amax_epsilon=self.amax_epsilon, force_pow_2_scales=self.force_pow_2_scales, ) quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm return quantizer @@ -123,103 +119,86 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: return tex.quantize(tensor, self) def get_scale_shape(self, shape: Iterable[int], columnwise: bool) -> Tuple[int, int]: - """Calculate the shape of the scaling tensor for blockwise quantization. + """Scaling tensor shape. - This method determines the shape of the scaling tensor needed for blockwise quantization, - taking into account the input tensor shape and whether columnwise scaling is used. - The scales are padded to multiples of 4 on the inner dimension for compatibility with GEMM. + This method determines the shape of the scaling tensor based + on the quantizer configuration. The scales are padded to + multiples of 4 for compatibility with GEMM. Parameters ---------- shape : Iterable[int] - Shape of the input tensor to be quantized + Logical tensor shape. columnwise : bool - Whether to use columnwise scaling (True) or rowwise scaling (False) + Whether the data is scaled column-wise (True) or row-wise (False). Returns ------- Tuple[int, int] - Shape of the scaling tensor as (outer_dim, inner_dim) - For 2D tensors: - - If columnwise: (roundup(K/blocksize), round_to_multiple(roundup(M/blocksize), 4)) - - If rowwise: (roundup(M/blocksize), round_to_multiple(roundup(K/blocksize), 4)) - For 1D tensors: - - If columnwise: (roundup(M/blocksize), round_to_multiple(K, 4)) - - If rowwise: (roundup(K/blocksize), round_to_multiple(M, 4)) + Scaling tensor shape. + """ - M, K = 1, 1 - for i in range(len(shape) - 1): - M *= shape[i] - if len(shape) > 0: - K = shape[-1] - # 2D 128x128 quantization block scaling - # CuBLAS requries 128x128 scaling factor to be padded - # currently rowwise and columnwise format option doesn't apply to 2D scaling + + # Flatten tensor to 2D + dim0 = math.prod(shape[:-1]) + dim1 = shape[-1] if shape else 1 + + # Check block dims + if self.block_scaling_dim not in (1, 2): + raise RuntimeError( + "Only 1D or 2D blocks are supported, " + f"but got block_scaling_dim={self.block_scaling_dim}" + ) + + # 128x128 block scaling if self.block_scaling_dim == 2: + scale_dim0 = (dim0 + self.block_len - 1) // self.block_len + scale_dim1 = (dim1 + self.block_len - 1) // self.block_len if columnwise: - outer = math.ceil(K / self.block_len) - inner = round_up_to_nearest_multiple(math.ceil(M / self.block_len), 4) - return (outer, inner) - # rowwise - outer = math.ceil(M / self.block_len) - inner = round_up_to_nearest_multiple(math.ceil(K / self.block_len), 4) - return (outer, inner) - # 1D 1x128 quantization block scaling - # CuBLAS requries 1x128 scaling factor to be padded and transposed - assert self.block_scaling_dim == 1, "Only 1D or 2D blocks supported" + return (scale_dim1, round_up_to_nearest_multiple(scale_dim0, 4)) + return (scale_dim0, round_up_to_nearest_multiple(scale_dim1, 4)) + + # 1x128 block scaling if columnwise: - columnwise_compact = self.all_gather_usage - outer = math.ceil(M / self.block_len) - inner = round_up_to_nearest_multiple(K, 4) if not columnwise_compact else K - # GEMM READY case: scaling factor is [outer, inner], already transposed here for CuBLAS - # for COMPACT case, since we apply 1x128 scaling here without transposing columnwise data, scaling factor is also [outer, inner] - # so no need to swap inner outer here - return (outer, inner) - # rowwise - rowwise_compact = self.all_gather_usage - outer = math.ceil(K / self.block_len) - inner = round_up_to_nearest_multiple(M, 4) if not rowwise_compact else M - # GEMM READY case: scaling factor is [outer, inner], already transposed here for CuBLAS need - # for COMPACT case, since we apply 128x1 scaling, scaling block applies to inner dim, so we need to swap outer and inner here - return (outer, inner) if not rowwise_compact else (inner, outer) + return ( + (dim0 + self.block_len - 1) // self.block_len, + round_up_to_nearest_multiple(dim1, 4), + ) + return ( + (dim1 + self.block_len - 1) // self.block_len, + round_up_to_nearest_multiple(dim0, 4), + ) def get_columnwise_shape(self, shape: Iterable[int]) -> Tuple[int, ...]: - """Calculate the shape of a tensor after columnwise permutation. + """Column-wise data shape - This method rearranges the dimensions of a tensor to be columnwise, - moving the last dimension to the front and keeping the order of other dimensions. + GEMMs expect that the column-wise data is transposed relative + to the logical tensor shape. Parameters ---------- shape : Iterable[int] - Original shape of the tensor + Logical tensor shape. Returns ------- Tuple[int, ...] - New shape with dimensions rearranged for columnwise layout. - For a shape (d1, d2, ..., dn), returns (dn, d1, d2, ..., dn-1). - Returns empty tuple for empty input shape. + Column-wise data shape. """ - if len(shape) == 0: - return tuple() - # currently columnwise format option only applies to 1D quantizer - # for 2D scaling, columnwise format should always be GEMM_READY_DATA_AND_SCALES - # since currently 2D scaling only applies to module weights - if self.block_scaling_dim == 1 and self.all_gather_usage: - return shape - colwise_shape = [shape[-1]] - for i in range(len(shape) - 1): - colwise_shape.append(shape[i]) + colwise_shape = [] + if shape: + colwise_shape.append(shape[-1]) + colwise_shape.extend(shape[:-1]) return tuple(colwise_shape) def is_quantizable(self, inp: torch.Tensor) -> bool: """Returns whether or not given inp can be quantized""" - if inp.ndim < 2: + shape = inp.size() + if len(shape) < 2: return False - if inp.shape[-1] % self.block_len != 0: + if shape[-1] % self.block_len != 0: return False - if math.prod(inp.shape[:-1]) % self.block_len != 0: + if math.prod(shape[:-1]) % self.block_len != 0: return False return True @@ -233,44 +212,36 @@ def make_empty( pin_memory: bool = False, ) -> Float8BlockwiseQTensor: """Construct quantized tensor with uninitialized data""" - if device is None: - device = torch.device("cuda") - data_format = ( - tex.Float8BlockScaleTensorFormat.COMPACT - if self.all_gather_usage - else tex.Float8BlockScaleTensorFormat.GEMM_READY - ) + tensor_kwargs = { + "device": torch.device("cuda") if device is None else device, + "pin_memory": pin_memory, + } - # Allocate FP8 data - data = None - scale_inv = None + # Allocate buffers for row-scaled data + rowwise_data = None + rowwise_scale_inv = None if self.rowwise_usage: - data = torch.empty(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, + 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, - device=device, - pin_memory=pin_memory, + **tensor_kwargs, ) - # Allocate FP8 data transpose if needed + # 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, - device=device, - pin_memory=pin_memory, + **tensor_kwargs, ) - columnwise_scale_shape = self.get_scale_shape(shape, columnwise=True) columnwise_scale_inv = torch.empty( - columnwise_scale_shape, + self.get_scale_shape(shape, columnwise=True), dtype=torch.float32, - device=device, - pin_memory=pin_memory, + **tensor_kwargs, ) # Construct FP8 tensor @@ -278,13 +249,12 @@ def make_empty( shape=shape, dtype=dtype, fp8_dtype=self.dtype, - rowwise_data=data, - rowwise_scale_inv=scale_inv, + 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, - data_format=data_format, requires_grad=requires_grad, ) @@ -334,7 +304,6 @@ def __new__( fp8_dtype: TE_DType, quantizer: Quantizer, is_2D_scaled: bool, - data_format: tex.Float8BlockScaleTensorFormat = Float8BlockScaleTensorFormat.GEMM_READY, **kwargs, ): instance = super().__new__( @@ -346,7 +315,6 @@ def __new__( fp8_dtype, quantizer, is_2D_scaled, - data_format, *args, **kwargs, ) @@ -357,8 +325,7 @@ def __repr__(self, *, tensor_contents=None): return ( f"Float8BlockwiseQTensor(fp8_dtype={self._fp8_dtype}," f" is_2D_scaled={self._is_2D_scaled}," - f" data={self.dequantize(dtype=self.dtype)})," - f" data_format={self._data_format}" + f" data={self.dequantize(dtype=self.dtype)})" ) def quantize_( @@ -509,7 +476,7 @@ def _make_in_reduce_ex( dtype: torch.dtype, quantizer: Quantizer, is_2D_scaled: bool, - data_format: tex.Float8BlockScaleTensorFormat, + data_format: Any = None, # pylint: disable=unused-argument ) -> Float8BlockwiseQTensor: """Build Float8BlockwiseQTensor, for use in __reduce__ @@ -527,7 +494,6 @@ def _make_in_reduce_ex( dtype=dtype, quantizer=quantizer, is_2D_scaled=is_2D_scaled, - data_format=data_format, ) def __reduce_ex__(self, protocol: int) -> tuple: @@ -544,7 +510,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: self.dtype, self._quantizer, self._is_2D_scaled, - self._data_format, + None, # data_format ), ) @@ -570,7 +536,6 @@ def _set_from_tensor(dst: Float8BlockwiseQTensor, src: Float8BlockwiseQTensor): dst._fp8_dtype = src._fp8_dtype dst._rowwise_scale_inv = src._rowwise_scale_inv dst._columnwise_scale_inv = src._columnwise_scale_inv - dst._data_format = src._data_format # Check that tensor dimensions match if ( @@ -618,13 +583,6 @@ def forward( ) -> Float8BlockwiseQTensor: # pylint: disable=missing-function-docstring - # Check for invalid configurations - if not tensor._is_gemm_ready_format(): - raise NotImplementedError( - "View is only supported with GEMM_READY data format, " - f"but found data_format={tensor._data_format}" - ) - # Return input tensor if shape is not provided ctx.shape = tensor.shape if shape is None: @@ -693,14 +651,6 @@ def backward( # pylint: disable=missing-function-docstring if isinstance(grad, Float8BlockwiseQTensor): - - # Check for invalid configurations - if not grad._is_gemm_ready_format(): - raise NotImplementedError( - "View is only supported with GEMM_READY data format, " - f"but found data_format={grad._data_format}" - ) - new_data = ( grad._rowwise_data.view(*ctx.shape) if grad._rowwise_data is not None else None ) @@ -740,13 +690,6 @@ def forward( ) -> Float8BlockwiseQTensor: # pylint: disable=missing-function-docstring - # Check for invalid configurations - if not tensor._is_gemm_ready_format(): - raise NotImplementedError( - "Reshape is only supported with GEMM_READY data format, " - f"but found data_format={tensor._data_format}" - ) - # Return input tensor if shape is not provided ctx.shape = tensor.shape if shape is None: @@ -814,14 +757,6 @@ def backward( # pylint: disable=missing-function-docstring if isinstance(grad, Float8BlockwiseQTensor): - - # Check for invalid configurations - if not grad._is_gemm_ready_format(): - raise NotImplementedError( - "Reshape is only supported with GEMM_READY data format, " - f"but found data_format={grad._data_format}" - ) - new_rowwise_data = None new_columnwise_data = None if grad._rowwise_data is not None: diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 43cbdcf9e6..3aeace0a77 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -293,6 +293,7 @@ def copy(self) -> Float8CurrentScalingQuantizer: amax=self.amax, ) quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm return quantizer diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 88081f51bf..8dd2255d89 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -54,6 +54,7 @@ def copy(self) -> MXFP8Quantizer: columnwise=self.columnwise_usage, ) quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm return quantizer @@ -156,6 +157,7 @@ def make_empty( 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: @@ -179,6 +181,7 @@ def create_tensor_from_data( columnwise_scale_inv=None, fp8_dtype=fp8_dtype, quantizer=self, + with_gemm_swizzled_scales=False, ) def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: @@ -188,6 +191,10 @@ def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: return self.create_tensor_from_data(data, scale_inv, fake_dtype=torch.float32) def onnx_dequantize(self, tensor: Union[MXFP8TensorStorage, MXFP8Tensor]) -> torch.Tensor: + if tensor._with_gemm_swizzled_scales: + raise NotImplementedError( + "ONNX MXFP8 dequantization is only supported with scales in compact format." + ) return torch.ops.tex.mxfp8_dequantize(tensor._rowwise_data, tensor._rowwise_scale_inv) def _get_compatible_recipe(self) -> Union[type[Recipe], None]: @@ -229,9 +236,10 @@ def __new__( columnwise_scale_inv: Optional[torch.Tensor], fp8_dtype: TE_DType, quantizer: Optional[Quantizer], + with_gemm_swizzled_scales: bool, **kwargs, ): - instance = super().__new__( + return super().__new__( cls, rowwise_data, rowwise_scale_inv, @@ -239,10 +247,10 @@ def __new__( columnwise_scale_inv, fp8_dtype, quantizer, + with_gemm_swizzled_scales, *args, **kwargs, ) - return instance def __repr__(self, *, tensor_contents=None): return f"MXFP8Tensor(fp8_dtype={self._fp8_dtype}, data={self.dequantize(dtype=self.dtype)})" @@ -334,39 +342,44 @@ def contiguous( @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): - # View op if func == aten.view.default: tensor = args[0] - data = tensor._rowwise_data - out_data = data.__torch_dispatch__( - func, - types, - [data] + list(args[1:]), - kwargs, - ) - out_shape = out_data.size() + shape = args[1] + if len(shape) < 2 or shape[-1] != tensor.size(-1): + raise ValueError( + f"Attempted to make view with size={tuple(shape)} " + f"from MXFP8 tensor with shape={tuple(tensor.size())}." + ) + rowwise_data_view = None + columnwise_data_view = None + if tensor._rowwise_data is not None: + rowwise_data_view = tensor._rowwise_data.view(shape) + if tensor._columnwise_data is not None: + columnwise_data_view = tensor._columnwise_data.view(shape) return MXFP8Tensor( - shape=out_shape, + shape=shape, dtype=tensor.dtype, - rowwise_data=out_data, + rowwise_data=rowwise_data_view, rowwise_scale_inv=tensor._rowwise_scale_inv, - columnwise_data=tensor._columnwise_data, + columnwise_data=columnwise_data_view, columnwise_scale_inv=tensor._columnwise_scale_inv, quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) if func == torch.ops.aten.copy_.default: dst, src = args[0], args[1] if isinstance(src, MXFP8Tensor) and isinstance(dst, MXFP8Tensor): - # Booleans to check if src has all the usages that dst needs to respect dst quantizer usages. - # If not, default to base class behavior. - rowwise_matches = src._rowwise_data is not None or dst._rowwise_data is None - columnwise_matches = ( - src._columnwise_data is not None or dst._columnwise_data is None - ) - if rowwise_matches and columnwise_matches: + if src._rowwise_data is None and dst._rowwise_data is not None: + pass + elif src._columnwise_data is None and dst._columnwise_data is not None: + pass + elif src._with_gemm_swizzled_scales != dst._with_gemm_swizzled_scales: + pass + else: + # src and dst match, so we can directly copy data if dst._rowwise_data is not None: dst._rowwise_data.copy_(src._rowwise_data.detach(), *args[2:], **kwargs) dst._rowwise_scale_inv.copy_( @@ -381,26 +394,25 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) return dst - # FSDP2 related functions. if func == aten.split.Tensor: - # This is called if entire model is initialized on CUDA device and - # then splitted. Finally the shard needed by the process is used - # and other splitted shards are discarded. + # With FSDP2, this is called if entire model is + # initialized on CUDA device and then splitted. Finally + # the shard needed by the process is used and other + # splitted shards are discarded. + tensor = args[0] + split_size = args[1] if "dim" in kwargs: dim_to_split = kwargs["dim"] else: dim_to_split = args[2] if len(args) > 2 else 0 - tensor = args[0] - split_size = args[1] - dim0_size = tensor.size(0) - dimlast_size = math.prod(tensor.shape[1:]) + + # Fall back to high-precision if split is non-trivial if ( - dim0_size % split_size != 0 - or dim_to_split != 0 + dim_to_split != 0 + or tensor.size(0) % split_size != 0 or split_size % MXFP8_BLOCK_SCALING_SIZE != 0 - or dimlast_size % MXFP8_BLOCK_SCALING_SIZE != 0 + or tensor._with_gemm_swizzled_scales ): - # Handle splitting by dequantizing and splitting the hp tensor return super().__torch_dispatch__(func, types, args, kwargs) out_data = [] @@ -460,28 +472,26 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, + with_gemm_swizzled_scales=False, ) for splitted_tensor_data in zip(*out_data) ] + if func == torch.ops.aten.as_strided.default: # Applied on unsharded param in FSDP2. In our case, this should be a no-op # This is needed for the case where some MXFP8 shards need padding i.e dimension 0 # of the unsharded param is not a multiple of the world size. If that is the case, # we down the dequantization route and weights are allgathered in high precision. # If weight doesnt need padding, this is just a no-op. + tensor = args[0] shape = args[1] strides = args[2] - tensor = args[0] if ( - len(shape) != 2 - or len(strides) != 2 - or strides[1] != 1 - or shape[0] != tensor.shape[0] - or shape[1] != tensor.shape[1] + len(shape) == len(strides) == 2 + and tuple(strides) == (shape[-1], 1) + and tuple(shape) == tuple(tensor.size()) ): - return super().__torch_dispatch__(func, types, args, kwargs) - - return MXFP8Tensor.make_like(tensor) + return MXFP8Tensor.make_like(tensor) if func == aten.slice.Tensor: # FSDP2 needed function. @@ -489,19 +499,12 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): # of the unsharded param is not a multiple of the world size. If that is the case, # we down the dequantization route and weights are allgathered in high precision instead. # If sharded weight doesnt have padding, this is just a no-op. + tensor = args[0] dim = args[1] start = args[2] length = args[3] - tensor = args[0] - if ( - dim != 0 - or length != tensor.shape[0] - or start != 0 - or length % MXFP8_BLOCK_SCALING_SIZE != 0 - or start % MXFP8_BLOCK_SCALING_SIZE != 0 - ): - return super().__torch_dispatch__(func, types, args, kwargs) - return MXFP8Tensor.make_like(tensor) + if start == 0 and length == tensor.size(dim): + return MXFP8Tensor.make_like(tensor) if func == aten.new_zeros.default: rowwise_data = None @@ -558,7 +561,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): quantizer=tensor._quantizer, requires_grad=False, fp8_dtype=tensor._fp8_dtype, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) + # Default case return super().__torch_dispatch__(func, types, args, kwargs) @@ -584,19 +589,24 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m # pylint: disable=unused-argument from transformer_engine.pytorch.distributed import _get_module_fsdp_state + # Get FSDP state fsdp_state = _get_module_fsdp_state(module) reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward + # Remove padding from scale inverses before allgather # Rowwise scale_inv should be divisible by [128,4], columnwise by [4, 128] rowwise_scale_inv = self._rowwise_scale_inv columnwise_scale_inv = self._columnwise_scale_inv shape = self.shape + if self._with_gemm_swizzled_scales: + raise NotImplementedError( + "FSDP2 is only supported for MXFP8Tensors with compact scales" + ) if rowwise_scale_inv is not None: # Remove padding from rowwise scale_inv flattened_in_shape0 = math.prod(shape[:-1]) if rowwise_scale_inv.size(0) != flattened_in_shape0: rowwise_scale_inv = rowwise_scale_inv[:flattened_in_shape0] - if columnwise_scale_inv is not None: # Remove padding from columnwise scale_inv flattened_in_shape0 = math.prod(shape[:-1]) // MXFP8_BLOCK_SCALING_SIZE @@ -681,7 +691,7 @@ def fsdp_post_all_gather( out._columnwise_data = columnwise_data out._columnwise_scale_inv = columnwise_scale_inv else: - # We ll be here when post all gather is called the first time. + # We'll be here when post all gather is called the first time. # MXFP8Tensor constructor makes a copy of the quantizer to # save as its own quantizer. For the consequent iterations, # the same quantizer is used. Copy is needed in the first iteration, @@ -696,6 +706,7 @@ def fsdp_post_all_gather( dtype=param_dtype, shape=rowwise_data.shape if rowwise_data is not None else columnwise_data.shape, quantizer=self._quantizer, + with_gemm_swizzled_scales=False, ) out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) return out, all_gather_outputs @@ -711,6 +722,7 @@ def _make_in_reduce_ex( dtype: torch.dtype, shape: torch.shape, quantizer: Optional[Quantizer] = None, + with_gemm_swizzled_scales: bool = False, ) -> MXFP8Tensor: """Build MXFP8Tensor, for use in __reduce__ @@ -727,6 +739,7 @@ def _make_in_reduce_ex( dtype=dtype, shape=shape, quantizer=quantizer, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, ) def __reduce_ex__(self, protocol: int) -> tuple: @@ -742,6 +755,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: self.dtype, self.shape, self._quantizer, + self._with_gemm_swizzled_scales, ), ) @@ -763,7 +777,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: if not devices_match(new_device, tensor.device): tensor = tensor.to(device=new_device) - # Just copy FP8 data if other tensor is MXFP8Tensor + # Just copy data if other tensor is MXFP8Tensor if isinstance(tensor, MXFP8Tensor): if ( # pylint: disable=too-many-boolean-expressions self.size() != tensor.size() @@ -791,6 +805,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: self._fp8_dtype = tensor._fp8_dtype self._rowwise_scale_inv = tensor._rowwise_scale_inv self._columnwise_scale_inv = tensor._columnwise_scale_inv + self._with_gemm_swizzled_scales = tensor._with_gemm_swizzled_scales return # Quantize to FP8 @@ -862,6 +877,7 @@ def forward( columnwise_scale_inv=tensor._columnwise_scale_inv, fp8_dtype=tensor._fp8_dtype, quantizer=tensor._quantizer, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) @staticmethod @@ -888,6 +904,7 @@ def backward( columnwise_scale_inv=grad._columnwise_scale_inv, fp8_dtype=grad._fp8_dtype, quantizer=grad._quantizer, + with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, ) return dgrad, None return grad.view(ctx.shape), None @@ -948,6 +965,7 @@ def forward( columnwise_scale_inv=tensor._columnwise_scale_inv, fp8_dtype=tensor._fp8_dtype, quantizer=tensor._quantizer, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) @staticmethod @@ -973,6 +991,7 @@ def backward( columnwise_scale_inv=grad._columnwise_scale_inv, fp8_dtype=grad._fp8_dtype, quantizer=grad._quantizer, + with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, ) return dgrad, None return grad.view(ctx.shape), None diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 8b707af3b2..101cf78a8f 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -193,6 +193,7 @@ def copy(self) -> NVFP4Quantizer: stochastic_rounding=self.stochastic_rounding, ) quantizer.internal = self.internal + quantizer.optimize_for_gemm = self.optimize_for_gemm quantizer.rht_matrix = self.rht_matrix quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t @@ -359,6 +360,7 @@ def make_empty( fp4_dtype=self.dtype, quantizer=self, requires_grad=requires_grad, + with_gemm_swizzled_scales=False, ) def calibrate(self, tensor: torch.Tensor) -> None: @@ -418,6 +420,7 @@ def __new__( amax_columnwise: Optional[torch.Tensor], fp4_dtype: TE_DType, quantizer: Quantizer, + with_gemm_swizzled_scales: bool, **kwargs, ): instance = super().__new__( @@ -430,6 +433,7 @@ def __new__( amax_columnwise, fp4_dtype, quantizer, + with_gemm_swizzled_scales, *args, **kwargs, ) @@ -592,6 +596,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): amax_columnwise=amax_columnwise, quantizer=tensor._quantizer, requires_grad=tensor.requires_grad, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) # Default case @@ -610,6 +615,7 @@ def _make_in_reduce_ex( fp4_dtype: TE_DType, dtype: torch.dtype, quantizer: Quantizer, + with_gemm_swizzled_scales: bool = False, ) -> NVFP4Tensor: """Build NVFP4Tensor, for use in __reduce__ @@ -629,6 +635,7 @@ def _make_in_reduce_ex( amax_columnwise=amax_columnwise, quantizer=quantizer, requires_grad=False, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, ) def __reduce_ex__(self, protocol: int) -> tuple: @@ -646,6 +653,7 @@ def __reduce_ex__(self, protocol: int) -> tuple: self._fp4_dtype, self.dtype, self._quantizer, + self._with_gemm_swizzled_scales, ), ) @@ -696,6 +704,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: self._columnwise_scale_inv = tensor._columnwise_scale_inv self._amax_rowwise = tensor._amax_rowwise self._amax_columnwise = tensor._amax_columnwise + self._with_gemm_swizzled_scales = tensor._with_gemm_swizzled_scales return # Quantize to FP8 @@ -782,6 +791,7 @@ def forward( quantizer=tensor._quantizer, fp4_dtype=tensor._fp4_dtype, requires_grad=tensor.requires_grad, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) @staticmethod @@ -823,6 +833,7 @@ def backward( quantizer=grad._quantizer, fp4_dtype=grad._fp4_dtype, requires_grad=grad.requires_grad, + with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, ) return dgrad, None return grad.view(ctx.shape), None @@ -902,6 +913,7 @@ def forward( quantizer=tensor._quantizer, fp4_dtype=tensor._fp4_dtype, requires_grad=tensor.requires_grad, + with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, ) @staticmethod @@ -943,6 +955,7 @@ def backward( quantizer=grad._quantizer, fp4_dtype=grad._fp4_dtype, requires_grad=grad.requires_grad, + with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, ) return dgrad, None return grad.view(ctx.shape), None diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 157981b4d7..278d7dc039 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -11,7 +11,6 @@ import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType -from transformer_engine_torch import Float8BlockScaleTensorFormat from ...quantized_tensor import QuantizedTensorStorage, Quantizer @@ -36,7 +35,6 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage): _rowwise_scale_inv: Optional[torch.Tensor] _columnwise_scale_inv: Optional[torch.Tensor] _is_2D_scaled: bool - _data_format: Float8BlockScaleTensorFormat def __new__( cls, @@ -47,7 +45,6 @@ def __new__( fp8_dtype: TE_DType, quantizer: Quantizer, is_2D_scaled: bool, - data_format: Float8BlockScaleTensorFormat, *args, **kwargs, ): @@ -62,7 +59,6 @@ def __new__( instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv instance._is_2D_scaled = is_2D_scaled - instance._data_format = data_format return instance @@ -87,13 +83,8 @@ def get_metadata(self) -> Dict[str, Any]: "fp8_dtype": self._fp8_dtype, "quantizer": self._quantizer, "is_2D_scaled": self._is_2D_scaled, - "data_format": self._data_format, } - def _is_gemm_ready_format(self) -> bool: - """Whether data is in GEMM_READY format""" - return self._data_format == Float8BlockScaleTensorFormat.GEMM_READY - def prepare_for_saving( self, ) -> Tuple[list[Optional[torch.Tensor]], Float8BlockwiseQTensorStorage]: @@ -153,36 +144,18 @@ def _dequantize_vectorwise(self, *, dtype: torch.dtype = torch.float32) -> torch for i in range(len(q.shape) - 1): q_M *= q.shape[i] inner_q_dimension_tiled = True - if self._is_gemm_ready_format(): - scales_tiled_dim, scales_untiled_dim = scale_inv.shape - inner_scale_dimension_tiled = False - scales_are_compact = False - else: - scales_untiled_dim, scales_tiled_dim = scale_inv.shape - inner_scale_dimension_tiled = True - scales_are_compact = True + scales_tiled_dim, scales_untiled_dim = scale_inv.shape else: assert self._columnwise_data is not None, "No data to dequantize" q = self._columnwise_data scale_inv = self._columnwise_scale_inv scales_tiled_dim, scales_untiled_dim = scale_inv.shape - inner_scale_dimension_tiled = False - if self._is_gemm_ready_format(): - inner_q_dimension_tiled = True - transpose_output = True - if len(q.shape) >= 1: - q_M = q.shape[0] - for i in range(1, len(q.shape)): - q_K *= q.shape[i] - scales_are_compact = False - else: - inner_q_dimension_tiled = False - transpose_output = False - if len(q.shape) >= 1: - q_K = q.shape[-1] - for i in range(len(q.shape) - 1): - q_M *= q.shape[i] - scales_are_compact = True + inner_q_dimension_tiled = True + transpose_output = True + if len(q.shape) >= 1: + q_M = q.shape[0] + for i in range(1, len(q.shape)): + q_K *= q.shape[i] orig_shape = q.shape q = q.reshape(q_M, q_K) @@ -202,15 +175,10 @@ def _dequantize_vectorwise(self, *, dtype: torch.dtype = torch.float32) -> torch ).contiguous() padded_M, padded_K = q.shape q_tiled = q.reshape(scales_tiled_dim, block_len, q_K) - if not scales_are_compact and scales_untiled_dim > q_M: + if scales_untiled_dim > q_M: # untiled scale dimension is 4 element aligned. scale_inv = scale_inv[:, :q_M].contiguous() - if scales_are_compact and inner_scale_dimension_tiled: - dq_scale = scale_inv.contiguous().reshape(q_M, scales_tiled_dim, 1) - elif scales_are_compact and not inner_scale_dimension_tiled: - dq_scale = scale_inv.contiguous().reshape(scales_tiled_dim, 1, q_K) - else: - dq_scale = scale_inv.transpose(-2, -1).contiguous().reshape(q_M, scales_tiled_dim, 1) + dq_scale = scale_inv.transpose(-2, -1).contiguous().reshape(q_M, scales_tiled_dim, 1) torch_q_dtype = TE_DType_To_Torch[self._fp8_dtype] result = q_tiled.view(torch_q_dtype).to(torch.float32) * dq_scale if padded_M != q_M or padded_K != q_K: @@ -233,12 +201,6 @@ def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: if not self._is_2D_scaled: return self._dequantize_vectorwise(dtype=dtype) - if not self._is_gemm_ready_format(): - raise NotImplementedError( - "Dequantize is only supported with GEMM_READY data format, " - f"but found _data_format={self._data_format}" - ) - def format_scale_as_logical_shape(q_K, scales, block_len): # The GEMM for 2D blocks required padding in the scales. derived_scale_k_shape = math.ceil(q_K / block_len) @@ -304,8 +266,6 @@ def size(self, *args, **kwargs): if self._rowwise_data is not None: return self._rowwise_data.size(*args, **kwargs) dims = list(self._columnwise_data.size(*args, **kwargs)) - if not self._is_gemm_ready_format(): # compact format - return torch.Size(dims) reordered = [] for i in range(1, len(dims)): reordered.append(dims[i]) @@ -366,7 +326,7 @@ def __repr__(self): return ( "Float8BlockwiseQTensorStorage(" f"fp8_dtype={self._fp8_dtype}, " - f"{descriptor}_scaled_data={data}" + f"{descriptor}_scaled_data={data})" ) def update_usage( diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 3bdf80c55e..1951731c75 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -57,13 +57,23 @@ class MXFP8TensorStorage(QuantizedTensorStorage): """ + # Row-scaled FP8 data _rowwise_data: Optional[torch.Tensor] + # Column-scaled FP8 data _columnwise_data: Optional[torch.Tensor] - _quantizer: Optional[Quantizer] - _fp8_dtype: TE_DType + # Scaling factors for row-scaled FP8 data _rowwise_scale_inv: torch.Tensor + # Scaling factors for column-scaled FP8 data _columnwise_scale_inv: torch.Tensor + # Builder class for casting to MXFP8 + _quantizer: Optional[Quantizer] + # FP8 data type + _fp8_dtype: TE_DType + # Whether scaling factors are in the swizzled format expected by + # GEMM + _with_gemm_swizzled_scales: bool + def __new__( cls, rowwise_data: Optional[torch.Tensor], @@ -72,6 +82,7 @@ def __new__( columnwise_scale_inv: Optional[torch.Tensor], fp8_dtype: TE_DType, quantizer: Optional[Quantizer], + with_gemm_swizzled_scales: bool, *args, **kwargs, ): @@ -81,10 +92,11 @@ def __new__( instance = super().__new__(cls, *args, **kwargs) instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data - instance._quantizer = quantizer.copy() if quantizer is not None else None - instance._fp8_dtype = fp8_dtype instance._rowwise_scale_inv = rowwise_scale_inv instance._columnwise_scale_inv = columnwise_scale_inv + instance._quantizer = quantizer.copy() if quantizer is not None else None + instance._fp8_dtype = fp8_dtype + instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales return instance @@ -108,6 +120,7 @@ def get_metadata(self) -> Dict[str, Any]: "columnwise_scale_inv": self._columnwise_scale_inv, "fp8_dtype": self._fp8_dtype, "quantizer": self._quantizer, + "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, } def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorStorage]: @@ -197,6 +210,7 @@ def view(self, shape: torch.Size): columnwise_scale_inv=self._columnwise_scale_inv, fp8_dtype=self._fp8_dtype, quantizer=self._quantizer, + with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, ) def __repr__(self): @@ -255,7 +269,7 @@ def update_usage( self._columnwise_data = None self._columnwise_scale_inv = None - def get_usages(self) -> Tuple[bool, bool]: + def get_usages(self) -> Dict[str, bool]: """Get the usage of the tensor""" return { "rowwise": self._rowwise_data is not None, diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index b9d568c9cd..b064d711ce 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -71,15 +71,29 @@ class NVFP4TensorStorage(QuantizedTensorStorage): """ + # Row-scaled FP4 data _rowwise_data: Optional[torch.Tensor] + # Column-scaled FP4 data _columnwise_data: Optional[torch.Tensor] - _quantizer: Optional[Quantizer] + # Block scaling factors for row-scaled FP4 data _rowwise_scale_inv: torch.Tensor + # Block scaling factors for column-scaled FP4 data _columnwise_scale_inv: torch.Tensor - _fp4_dtype: TE_DType + # Input absolute maximum value (used to compute tensor scale for + # row-scaled FP4 data) _amax_rowwise: torch.Tensor + # Input absolute maximum value (used to compute tensor scale for + # column-scaled FP4 data) _amax_columnwise: torch.Tensor + # Builder class for casting to MXFP8 + _quantizer: Optional[Quantizer] + # FP4 data type + _fp4_dtype: TE_DType + # Whether scaling factors are in the swizzled format expected by + # GEMM + _with_gemm_swizzled_scales: bool + def __new__( cls, rowwise_data: Optional[torch.Tensor], @@ -90,6 +104,7 @@ def __new__( amax_columnwise: torch.Tensor, fp4_dtype: TE_DType, quantizer: Optional[Quantizer], + with_gemm_swizzled_scales: bool, *args, **kwargs, ): @@ -104,6 +119,7 @@ def __new__( instance._columnwise_scale_inv = columnwise_scale_inv instance._amax_rowwise = amax_rowwise instance._amax_columnwise = amax_columnwise + instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales return instance @@ -131,6 +147,7 @@ def get_metadata(self) -> Dict[str, Any]: "amax_columnwise": self._amax_columnwise, "fp4_dtype": self._fp4_dtype, "quantizer": self._quantizer, + "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, } def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], NVFP4TensorStorage]: @@ -248,6 +265,7 @@ def view(self, shape: torch.Size): amax_columnwise=self._amax_columnwise, quantizer=self._quantizer, fp4_dtype=self._fp4_dtype, + with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, ) def __repr__(self): From dfdd38201e447b77d265979d43371943c461674a Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Tue, 20 Jan 2026 09:14:51 -0800 Subject: [PATCH 170/521] Changed VERSION to 2.13.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 d5e1cb2914..90cc92ea66 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.12.0.dev0 +2.13.0.dev0 From 27fc168e2166e2338a540e3fcd59f9c8491f3b10 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 20 Jan 2026 13:59:29 -0800 Subject: [PATCH 171/521] [Common] Enable determinism for cuDNN >= 9.18.1 on Blackwell (#2584) * update FE to 1.17 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add determinism flag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add determinism to test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add determinism to qa/ Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * move bias/dbias/versioning/dropout logic to C API 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 qa/L0_pytorch_unittest/test.sh make .xml file specific to deterministic tests in qa/ Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add determinism to Jax extension Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add determinism to Jax 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 * Update tests/jax/test_fused_attn.py fix typo Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Update transformer_engine/common/fused_attn/fused_attn.cpp fix indentation Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> 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 the AI fixes 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 Jax extension call Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes based on comments Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix selection logic and fwd arg 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 version check in Jax test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix pytorch CI failures Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix Jax CI failures 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 non-/determinism logic and CI 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 formatting 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 transformer_engine/common/fused_attn/fused_attn.cpp fix and/or logic Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update to 9.18.1 for requirement 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 * reduce Jax CI tests for determinism 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> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- 3rdparty/cudnn-frontend | 2 +- qa/L0_jax_unittest/test.sh | 1 + qa/L0_pytorch_unittest/test.sh | 1 + tests/jax/test_fused_attn.py | 212 +++++++++++++++++- tests/pytorch/attention/test_attention.py | 41 +++- .../common/fused_attn/fused_attn.cpp | 41 ++-- .../include/transformer_engine/fused_attn.h | 3 +- .../jax/cpp_extensions/attention.py | 21 +- transformer_engine/jax/csrc/extensions.h | 2 +- .../jax/csrc/extensions/attention.cpp | 8 +- .../attention/dot_product_attention/utils.py | 5 +- transformer_engine/pytorch/csrc/extensions.h | 2 +- .../pytorch/csrc/extensions/attention.cpp | 4 +- 13 files changed, 299 insertions(+), 44 deletions(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 0258951d4d..b372d39879 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 0258951d4d512f4714eb1574496f4d57669b1b93 +Subproject commit b372d39879d44c91a8d5b342022e74802b6a8da2 diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index ee9ce130aa..3453e35d2c 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -29,6 +29,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" mkdir -p "$XML_LOG_DIR" 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_*" +NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_fused_attn_with_determinism.xml $TE_PATH/tests/jax/test_fused_attn.py -k "TestFusedAttnWithDeterminism" || test_fail "tests/jax/test_fused_attn.py" pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements" 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" diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 21eed28367..a13dfada79 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -45,6 +45,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_e python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading_v1.xml $TE_PATH/tests/pytorch/test_cpu_offloading_v1.py || test_fail "test_cpu_offloading_v1.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py" +NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint 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" diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index a0aee50430..f9946e1f7f 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. """Tests for fused attention""" +import os from enum import Enum, auto from dataclasses import dataclass, field from functools import partial @@ -49,6 +50,9 @@ from distributed_test_base import assert_equal_collectives from utils import assert_allclose, print_debug_tensor_stats +# Get determinism +_deterministic = not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + @pytest.fixture(autouse=True, scope="module") def init(): @@ -413,15 +417,25 @@ def _check_configs(self): pytest.skip( "seqlen_q > seqlen_kv is not supported with sliding window attention in cuDNN" ) - # TODO(KshitijLakhani): Set the upper limit for skipping this test when cuDNN adds support - if ( - get_device_compute_capability(0) >= 100 - and self.dropout_prob == 0.1 - and self.attn_bias_type is not AttnBiasType.NO_BIAS - ): - pytest.skip( - "For sm100+, bprop kernel support for dropout + determinism (bias) is not supported" - ) + + if get_device_compute_capability(0) >= 100 and self.is_training: + if FusedAttnHelper.is_non_deterministic_allowed() and ( + (self.dropout_prob != 0.0 and self.attn_bias_type != AttnBiasType.NO_BIAS) + or get_cudnn_version() < 90700 + ): + pytest.skip( + "For sm100+, non-deterministic bprop (cuDNN 9.7+) does not support bias with" + " dropout" + ) + if not FusedAttnHelper.is_non_deterministic_allowed() and ( + self.dropout_prob != 0.0 + or self.attn_bias_type != AttnBiasType.NO_BIAS + or get_cudnn_version() < 91801 + ): + pytest.skip( + "For sm100+, deterministic bprop (cuDNN 9.18.1+) does not support bias or" + " dropout" + ) # Test the MLA case where head dims for qk differ from head dims for v, only if the tensors # are provided in BSHD_BSHD_BSHD or THD_THD_THD formats if self.head_dim_qk != self.head_dim_v and not self.qkv_layout.is_separate(): @@ -1269,6 +1283,7 @@ def check_dqkv(primitive, reference, pad, idx): pytest.param(SeqDescFormat.SegmentIDs, id="SegmentIDs"), ], ) +@pytest.mark.skipif(_deterministic, reason="Test non-determinism only") class TestFusedAttn: """ Fused attention tester @@ -1392,3 +1407,182 @@ def test_backward( seq_desc_format, ) runner.test_backward() + + +@pytest.mark.parametrize( + "attn_mask_type", + [ + pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"), + pytest.param(AttnMaskType.PADDING_MASK, id="PADDING"), + pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL"), + pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL"), + pytest.param( + AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK, id="PADDING_CAUSAL_BOTTOM_RIGHT" + ), + ], +) +@pytest.mark.parametrize( + "softmax_type", + [ + pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), + ], +) +@pytest.mark.parametrize( + "b, s_q, s_kv, h_q, h_kv, d_qk, d_v, dtype, qkv_layout", + [ + # large data size + fp16 + cross attn + gqa + diff hidden v dim + qkv separate + pytest.param( + 2, + 1024, + 2048, + 12, + 6, + 128, + 64, + jnp.bfloat16, + QKVLayout.BSHD_BSHD_BSHD, + id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-SEPARATE", + ), + pytest.param( + 2, + 1024, + 2048, + 12, + 6, + 128, + 64, + jnp.bfloat16, + QKVLayout.THD_THD_THD, + id="2-1024-2048-12-6-128-64-BF16-CROSS-GQA-RAGGED_SEPARATE", + ), + ], +) +@pytest.mark.parametrize( + "dropout_prob", + [ + pytest.param(0.0, id="DROP_0.0"), + ], +) +@pytest.mark.parametrize( + "swa", + [ + pytest.param(False, id="NO_SWA"), + ], +) +@pytest.mark.parametrize( + "seq_desc_format", + [ + pytest.param(SeqDescFormat.Seqlens, id="Seqlens"), + ], +) +@pytest.mark.skipif(not _deterministic, reason="Test determinism only") +class TestFusedAttnWithDeterminism: + """ + Fused attention tester with determinism + """ + + @staticmethod + @pytest.mark.parametrize( + "is_training", + [ + pytest.param(True, id="TRAINING"), + ], + ) + @pytest.mark.parametrize( + "attn_bias_type, bias_shape", + [ + pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), + pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape._1HSS, id="POST_SCALE_BIAS-1HSS"), + ], + ) + def _test_forward( + b, + s_q, + s_kv, + h_q, + h_kv, + d_qk, + d_v, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + dtype, + is_training, + qkv_layout, + bias_shape, + swa, + seq_desc_format, + ): + """ + Test forward with parameterized configs + This test is not intended to run automatically during CI as it is time-consuming + It is kept for development and debugging + """ + TestFusedAttn._test_forward( + b, + s_q, + s_kv, + h_q, + h_kv, + d_qk, + d_v, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + dtype, + is_training, + qkv_layout, + bias_shape, + swa, + seq_desc_format, + ) + + @staticmethod + @pytest.mark.parametrize( + "attn_bias_type, bias_shape", + [ + pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), + pytest.param(AttnBiasType.POST_SCALE_BIAS, BiasShape._1HSS, id="POST_SCALE_BIAS-1HSS"), + ], + ) + def test_backward( + b, + s_q, + s_kv, + h_q, + h_kv, + d_qk, + d_v, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + dtype, + qkv_layout, + bias_shape, + swa, + seq_desc_format, + ): + """ + Test backward with parameterized configs + """ + TestFusedAttn.test_backward( + b, + s_q, + s_kv, + h_q, + h_kv, + d_qk, + d_v, + attn_bias_type, + attn_mask_type, + softmax_type, + dropout_prob, + dtype, + qkv_layout, + bias_shape, + swa, + seq_desc_format, + ) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index eb7905bcd5..9111d3511c 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -72,6 +72,14 @@ f" sm{device_compute_capability[0] * 10 + device_compute_capability[1]}" ) + +# Get determinism +_deterministic = ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() +) + + # Reset RNG seed and states seed = 1234 reset_rng_states() @@ -160,6 +168,7 @@ def test_dot_product_attention( qkv_layout=qkv_layout, pad_between_seqs=pad_between_seqs, is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported: @@ -170,6 +179,7 @@ def test_dot_product_attention( qkv_layout=qkv_layout, pad_between_seqs=pad_between_seqs, is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -886,11 +896,14 @@ def _run_dot_product_attention( reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" if backend == "FlashAttention": os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT"] = "1" if workspace_opt else "0" + if backend == "UnfusedDotProductAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True # Create seqlens @@ -1292,6 +1305,7 @@ def test_transformer_layer( qkv_format.replace("hd", "h3d") if fused_qkv_params else qkv_format.replace("hd", "3hd") ), is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported: @@ -1305,6 +1319,7 @@ def test_transformer_layer( else qkv_format.replace("hd", "3hd") ), is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends @@ -1432,10 +1447,13 @@ def _run_transformer_layer( reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" if backend == "FlashAttention": os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" + if backend == "UnfusedDotProductAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True # Create input tensor @@ -1629,6 +1647,7 @@ def test_dpa_fp8_extra_state(model, dtype): qkv_dtype=torch.float8_e4m3fn, qkv_layout="sb3hd", is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not fused_attn_supported and not flash_attn_supported: @@ -1819,6 +1838,7 @@ def test_mha_fp8_vs_f16( fp8=True, fp8_meta=fp8_meta, is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: @@ -1830,6 +1850,7 @@ def test_mha_fp8_vs_f16( qkv_dtype=dtype, qkv_layout=qkv_format.replace("hd", "h3d"), is_training=is_training, + deterministic=_deterministic, ) _, fused_attn_supported_f16, _ = available_backends if not fused_attn_supported_f16: @@ -1838,6 +1859,7 @@ def test_mha_fp8_vs_f16( if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") flash_attn_fwd_fp8, param_names, flash_attn_bwd_fp8 = _run_mha_fp8_vs_f16( @@ -1847,6 +1869,7 @@ def test_mha_fp8_vs_f16( if fused_attn_supported_fp8: os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = True") fused_attn_fwd_fp8, param_names, fused_attn_bwd_fp8 = _run_mha_fp8_vs_f16( @@ -1856,6 +1879,7 @@ def test_mha_fp8_vs_f16( if fused_attn_supported_f16: os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" _attention_backends["backend_selection_requires_update"] = True logging.info("[test_mha_fp8_vs_f16]: run with fp8_mha = False") fused_attn_fwd_f16, param_names, fused_attn_bwd_f16 = _run_mha_fp8_vs_f16( @@ -2068,6 +2092,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal fp8=True, fp8_meta=fp8_meta, is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if flash_attn_supported + fused_attn_supported < 1: @@ -2078,6 +2103,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal qkv_dtype=dtype, qkv_layout=qkv_layout, is_training=is_training, + deterministic=_deterministic, ) _, fused_attn_supported, _ = available_backends if not fused_attn_supported: @@ -2088,6 +2114,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FlashAttention)") flash_attn_fwd_fp8, flash_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( @@ -2097,6 +2124,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal if unfused_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (UnfusedDotProductAttention)") unfused_attn_fwd_fp8, unfused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( @@ -2105,6 +2133,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" _attention_backends["backend_selection_requires_update"] = True logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FusedAttention)") fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( @@ -2113,6 +2142,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" if config.dropout_p == 0.0: # test cuDNN FP8 dropout: need a FP16/BF16 reference on Blackwell logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = False (FusedAttention)") @@ -2367,13 +2397,16 @@ def test_custom_mha_fp8_vs_f16(dtype, model): qkv_dtype=torch.float8_e4m3fn, qkv_layout="t3hd" if cudnn_frontend_version == 0 else "bs3hd", is_training=is_training, + deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends if not (fused_attn_backends and unfused_attn_supported): pytest.skip("Not enough backends to run this test with.") fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_custom_mha_fp8(dtype, config, "FusedAttention") - unfused_attn_fwd_f16, unfused_attn_bwd_f16 = _run_ref_mha_f16(dtype, config, "UnfusedAttention") + unfused_attn_fwd_f16, unfused_attn_bwd_f16 = _run_ref_mha_f16( + dtype, config, "UnfusedDotProductAttention" + ) atol = 5e-1 rtol = 5e-1 @@ -2406,10 +2439,13 @@ def _run_custom_mha_fp8(dtype, config, backend): reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" if backend == "FlashAttention": os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" + if backend == "UnfusedDotProductAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True inp = 0.0001 * torch.randint( @@ -2460,10 +2496,13 @@ def _run_ref_mha_f16(dtype, config, backend): os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" if backend == "FlashAttention": os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" + if backend == "UnfusedDotProductAttention": + os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True inp = torch.load("qkv.pt").to(device="cuda") diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index fde0d38921..415bfae063 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -206,7 +206,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph) { + int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { using namespace transformer_engine; NVTE_Fused_Attn_Backend backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; const int device_id = cuda::current_device(); @@ -440,7 +440,16 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // 9.13.1+: vanilla, off-by-one, learnable (cudnn_runtime_version >= 91301 || (cudnn_runtime_version < 91301 && - softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX))) { + softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX)) && + // determinism on Blackwell + // pre-9.18.1: fwd: deterministic; bwd: non-deterministic + // 9.18.1+: fwd: deterministic; bwd: non-deterministic/deterministic + (sm_arch_ < 100 || + (sm_arch_ >= 100 && (!is_training || + (is_training && !deterministic && + (dropout == 0.0 || bias_type == NVTE_Bias_Type::NVTE_NO_BIAS)) || + (is_training && deterministic && cudnn_runtime_version >= 91801 && + 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)) { @@ -553,7 +562,7 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h, h, max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, return_max_logit, - cuda_graph); + cuda_graph, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -595,7 +604,8 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, wkspace, stream, handle); #else NVTE_ERROR( - "cuDNN 8.9.0 is required for BF16/FP16 fused attention with arbitrary sequence length. \n"); + "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) @@ -669,7 +679,8 @@ void nvte_fused_attn_bwd_qkvpacked( NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h, h, - max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, false, cuda_graph); + max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, false, cuda_graph, + deterministic); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -855,7 +866,7 @@ void nvte_fused_attn_fwd_kvpacked( NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, - return_max_logit, cuda_graph); + return_max_logit, cuda_graph, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -897,7 +908,8 @@ void nvte_fused_attn_fwd_kvpacked( input_page_table_v, input_rng_state, wkspace, stream, handle); #else NVTE_ERROR( - "cuDNN 8.9.3 is required for BF16/FP16 fused attention with arbitrary sequence length. \n"); + "cuDNN 8.9.3 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) @@ -982,10 +994,10 @@ void nvte_fused_attn_bwd_kvpacked( const NVTEDType Q_type = static_cast(input_Q->data.dtype); const NVTEDType KV_type = static_cast(input_KV->data.dtype); - NVTE_Fused_Attn_Backend fused_attention_backend = - nvte_get_fused_attn_backend(true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, - softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, - d, window_size_left, window_size_right, false, cuda_graph); + NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( + true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, + h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, false, + cuda_graph, deterministic); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -1166,7 +1178,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, 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); + return_max_logit, cuda_graph, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) @@ -1189,7 +1201,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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"); + "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) @@ -1262,7 +1275,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false, - cuda_graph); + cuda_graph, deterministic); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { #if (CUDNN_VERSION >= 8901) diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index dd70ccf8df..0fabb81aef 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -208,13 +208,14 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] window_size_right Sliding window size (the right half). * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. + * \param[in] deterministic Whether determinism is required or not. */ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bool is_training, NVTEDType q_dtype, NVTEDType kv_dtype, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph); + int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic); /*! \brief Compute dot product attention with packed QKV input. * diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 0cdfcebf38..ee10115aa1 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -144,6 +144,7 @@ def get_fused_attn_backend(self): self.head_dim_v, self.window_size[0], self.window_size[1], + not self.is_non_deterministic_allowed(), ) @staticmethod @@ -3563,13 +3564,21 @@ def fused_attn_bwd( softmax_offset, (None, HEAD_AXES, None, None) ) - # TODO(KshitijLakhani): Add a check for cuDNN version when determinism does get supported on - # sm100+ compute_capabilities = get_all_device_compute_capability() - if any(x >= 100 for x in compute_capabilities): - assert not ( - attn_bias_type != AttnBiasType.NO_BIAS and dropout_probability != 0 - ), "For sm100+, bprop kernel support for dropout + determinism (bias) is not supported" + if any(x >= 100 for x in compute_capabilities) and is_training: + assert ( + FusedAttnHelper.is_non_deterministic_allowed() + and get_cudnn_version() >= (9, 7, 0) + and (attn_bias_type == AttnBiasType.NO_BIAS or dropout_probability == 0.0) + ) or ( + not FusedAttnHelper.is_non_deterministic_allowed() + and get_cudnn_version() >= (9, 18, 1) + and attn_bias_type == AttnBiasType.NO_BIAS + and dropout_probability == 0.0 + ), ( + "For sm100+, non-deterministic bprop (cuDNN 9.7+) does not support bias with dropout," + " and deterministic bprop (cuDNN 9.18.1+) does not support bias or dropout" + ) fused_config = _FusedAttnConfig( attn_bias_type=attn_bias_type, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index a83a1e0a80..5f93392633 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -113,7 +113,7 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right); + int64_t window_size_right, bool deterministic); pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 540aeb8b2d..4fe8e728a3 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -16,12 +16,12 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right) { + int64_t window_size_right, bool deterministic) { auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false); + false, false, deterministic); return backend; } @@ -266,7 +266,7 @@ static void FusedAttnForwardImpl( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false); + false, false, deterministic); nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); /* Auxiliary tensors (to be propagated to the backward pass later) */ @@ -522,7 +522,7 @@ static void FusedAttnBackwardImpl( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false); + false, false, deterministic); PrepareFusedAttnBackwardAuxTensors(&aux_input_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, softmax_aux, rng_state, bias, softmax_offset); diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index bf19388d7e..cb74a15e77 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -994,6 +994,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt window_size[1], return_max_logit, cuda_graph, + deterministic, ) if fused_attention_backend == FusedAttnBackend["No_Backend"]: logger.debug("Disabling FusedAttention as no backend supports the provided input") @@ -1064,10 +1065,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention for determinism reasons with post_scale_bias") use_fused_attention = False fused_attention_backend = None - if is_training and device_compute_capability >= (10, 0): - logger.debug("Disabling FusedAttention for determinism reasons on Blackwell") - use_fused_attention = False - fused_attention_backend = None # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9dc0d1f37b..591c89f83f 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -81,7 +81,7 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float p_dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph); + int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic); std::vector fused_attn_fwd( size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index b455e03757..be645d91b9 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -45,12 +45,12 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, float p_dropout, 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, int64_t window_size_left, - int64_t window_size_right, bool return_max_logit, bool cuda_graph) { + int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic) { NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, attn_mask_type, softmax_type, p_dropout, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, window_size_left, window_size_right, - return_max_logit, cuda_graph); + return_max_logit, cuda_graph, deterministic); return fused_attention_backend; } From fbb16f4a71393fe7188f9f198438be6f5265e29b Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Wed, 21 Jan 2026 16:05:27 +0100 Subject: [PATCH 172/521] [Common] Tuned NVFP4 cast kernel (#2412) * Implemented persistent nvfp4 kernel Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix FP4 guard in ptx Signed-off-by: Oleg Goncharov * Fix Signed-off-by: Oleg Goncharov * Fix in ptx. reduxf32 guard 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 * Fixes per PR review Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixes per PR review. Added parameter to turn off the persistency Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Modified reference CPU implementation in C++ unit tests to match GPU (numerical truncation). Tightened the numerical tolerance Signed-off-by: Oleg Goncharov * Disabled persistency by default, as non-persistent kernel is more performant when inputs are large Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use the tuned kernel also for the rowwise only quantization Signed-off-by: Oleg Goncharov * Fixed typo Signed-off-by: Oleg Goncharov * Addressed comments from the PR review Signed-off-by: Oleg Goncharov * Resolved conflicts Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Macros renaming 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 | 130 +-- .../common/cast/core/common.cuh | 6 + .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 7 + .../quantize_transpose_nvfp4_tuned_1D.cuh | 789 ++++++++++++++++++ transformer_engine/common/util/ptx.cuh | 306 +++++++ 5 files changed, 1184 insertions(+), 54 deletions(-) create mode 100644 transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 1904d03df7..c4df8759f2 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -54,12 +54,16 @@ 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) { +float compute_global_encode_scaling_factor_FP4(const float global_amax, const bool use_fast_math) { constexpr float fp8_max = 448.0f; // 448.0f; constexpr float fp4_max = 6.0f; // 6.0f; float global_encode_scale = fp8_max * fp4_max / global_amax; - // If scale is infinity, return max value of float32 - global_encode_scale = fminf(global_encode_scale, Numeric_Traits::maxNorm); + // If scale is infinity, return the max normalized value + const float max_norm_clamp = use_fast_math + ? Numeric_Traits::maxNorm + : Numeric_Traits::maxNorm; + + global_encode_scale = fminf(global_encode_scale, max_norm_clamp); // If global amax is 0 or infinity, return 1 if (global_amax == 0.0f || global_encode_scale == 0.0f) { return 1.0f; @@ -76,10 +80,11 @@ void quantize_nvfp4_1d(float (*OP)(const float), const size_t rows, const size_t cols, const size_t scales_stride, - const float global_amax) { + const float global_amax, + const bool use_fast_math) { // Compute a global encoding/decoding scaling factor for all S_dec_b - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax); + const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math); constexpr size_t block_size_X = 16; const size_t blocks_X = divide_round_up(cols, block_size_X); @@ -114,14 +119,20 @@ void quantize_nvfp4_1d(float (*OP)(const float), const float S_dec_b = block_amax / 6.0f; // Scale & Store per-block decoding scaling factor - const float S_dec_b_fp8 = S_dec_b * S_enc; + const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); + 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_fp8 == 0 ? 0.f : S_enc / S_dec_b_fp8; + const float S_enc_b_fp8 = S_dec_b_fp32 == 0.f ? 0.f : S_enc / S_dec_b_fp32; const size_t scale_idx = i * scales_stride + block_X; - scales[scale_idx] = static_cast(S_dec_b_fp8); - const float scale_reciprocal = S_enc_b_fp8; + scales[scale_idx] = S_dec_b_fp8; + + float scale_reciprocal = S_enc_b_fp8; + if (use_fast_math) { + // Numerical truncation to match GPU implementation, if mixed precision FMA instruction is used + scale_reciprocal = static_cast(static_cast(scale_reciprocal)); + } for (size_t j = j_min; j < j_max; j += 2) { const int idx_pair = (i * cols + j) / 2; @@ -136,7 +147,7 @@ void quantize_nvfp4_1d(float (*OP)(const float), fp4e2m1x2 casted_to_e2m1_pair(scaled_elt_pair); output[idx_pair] = casted_to_e2m1_pair; - // const double2 truncated_pair = cvt_fp4x2_to_double2(casted_to_e2m1_pair); + const double2 truncated_pair = cvt_fp4x2_to_double2(casted_to_e2m1_pair); } } } @@ -149,9 +160,10 @@ void compute_2d_mathematical_scales(float (*OP)(const float), const size_t rows, const size_t cols, const float global_amax, - std::vector>& math_scales) { + std::vector>& math_scales, + const bool use_fast_math) { - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax); + const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math); 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); @@ -195,13 +207,14 @@ void quantize_nvfp4_2d(float (*OP)(const float), const size_t rows, const size_t cols, const size_t scales_stride, - const float global_amax) { + const float global_amax, + const bool use_fast_math) { // Step 1: Compute mathematical 8x8 scaling factors std::vector> math_scales; - compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales); + compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math); - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax); + const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math); 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); @@ -282,11 +295,12 @@ void quantize_nvfp4(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_2d_quantization = false) { if (use_2d_quantization) { - quantize_nvfp4_2d(OP, input, output, scales, rows, cols, scales_stride, global_amax); + quantize_nvfp4_2d(OP, input, output, scales, rows, cols, scales_stride, global_amax, use_fast_math); } else { - quantize_nvfp4_1d(OP, input, output, scales, rows, cols, scales_stride, global_amax); + quantize_nvfp4_1d(OP, input, output, scales, rows, cols, scales_stride, global_amax, use_fast_math); } } @@ -302,6 +316,7 @@ void compute_ref(float (*OP)(const float), 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 input_t = create_transpose(input, rows, cols); @@ -309,7 +324,7 @@ 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, global_amax, math_scales); + compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; @@ -336,12 +351,16 @@ 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); // scales already filled - quantize_nvfp4_2d(OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, global_amax); // scales_t already filled + quantize_nvfp4_2d(OP, input, output, nullptr, rows, cols, scales_stride, global_amax, + use_fast_math); // scales already filled + quantize_nvfp4_2d(OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, global_amax, + use_fast_math); // scales_t already filled } else { - quantize_nvfp4(OP, input, output, scales, rows, cols, scales_stride, global_amax, use_2d_quantization); - quantize_nvfp4(OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, global_amax, use_2d_quantization); + 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); } } @@ -349,6 +368,8 @@ void compare_nvfp4_tensors(const std::string& name, const fp4e2m1 *test_data, const fp4e2m1 *ref_data, const int rows, const int cols, double atol = 1e-5, double rtol = 1e-8) { + constexpr int max_mismatches_to_print = 3; + std::vector mismatch_messages; size_t total_mismatches = 0; @@ -362,29 +383,16 @@ void compare_nvfp4_tensors(const std::string& name, const double t = (k == 0 ? test_data_pair.x : test_data_pair.y); const double r = (k == 0 ? ref_data_pair.x : ref_data_pair.y); - bool mismatch = fabs(t - r) > atol && (r == 0 || fabs((t - r) / r) > rtol); - /* For Float32 the floating point comparison is enough to error out */ - bool assertion = false; - if (mismatch && !assertion) { - /* Check if it is just a failure of round to nearest choosing different - side of the real value */ - const double mean = (t + r) / 2; - const double mean_p = mean >= 0 ? mean * (1 + 1e-6) : mean * (1 - 1e-6); - const double mean_m = mean >= 0 ? mean * (1 - 1e-6) : mean * (1 + 1e-6); - const double cast_mean_p = static_cast(static_cast(mean_p)); - const double cast_mean_m = static_cast(static_cast(mean_m)); - assertion = !(cast_mean_m == std::min(t,r) && cast_mean_p == std::max(t,r)); - } - if (assertion) { + const bool mismatch = fabs(t - r) > (atol + fabs(r) * rtol); + if (mismatch) { total_mismatches++; - std::string msg = "Mismatch at place (" + std::to_string(idx + k) + "): " + - std::to_string(t) + " vs " + std::to_string(r) + - " (abs_diff: " + std::to_string(fabs(t - r)) + - ", rel_diff: " + std::to_string(r == 0 ? 0.0 : fabs((t - r) / r)) + ")"; - mismatch_messages.push_back(msg); - // Optional: limit number of detailed messages to avoid overwhelming output - if (mismatch_messages.size() <= 100) { + if (total_mismatches <= max_mismatches_to_print) { + std::string msg = "Mismatch at place (" + std::to_string(idx + k) + "): " + + std::to_string(t) + " vs " + std::to_string(r) + + " (abs_diff: " + std::to_string(fabs(t - r)) + + ", rel_diff: " + std::to_string(r == 0 ? 0.0 : fabs((t - r) / r)) + ")"; + mismatch_messages.push_back(msg); std::cout << "Error in tensor " << name << ": " << msg << std::endl; } } @@ -400,8 +408,9 @@ void compare_nvfp4_tensors(const std::string& name, std::cout << "STATUS: FAILED for output" << std::endl; std::cout << "Total mismatches found: " << total_mismatches << std::endl; std::cout << "Mismatch rate: " << (100.0 * total_mismatches) / (rows * cols) << "%" << std::endl; - if (mismatch_messages.size() > 100) { - std::cout << "... and " << (mismatch_messages.size() - 100) << " more mismatches (showing first 100)" << std::endl; + if (mismatch_messages.size() > max_mismatches_to_print) { + std::cout << "... and " << (mismatch_messages.size() - max_mismatches_to_print) + << " more mismatches (showing first " << max_mismatches_to_print << ")" << std::endl; } std::cout << "============================" << std::endl; @@ -519,7 +528,8 @@ void compareResults_nvfp4(const Tensor &test, template void performTest(float (*OP)(const float), - const std::vector& shape) { + const std::vector& shape, + const bool use_fast_math) { using namespace test; DType itype = TypeInfo::dtype; @@ -580,15 +590,16 @@ void performTest(float (*OP)(const float), cols, scales_stride, scales_stride_t, + use_fast_math, use_2d_quantization); - - QuantizationConfigWrapper quant_config; - // 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(); + + 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()); @@ -619,8 +630,8 @@ void performTest(float (*OP)(const float), } ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); - const double atol = 0.05; - const double rtol = 0.1; + 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); @@ -666,12 +677,18 @@ std::vector Activation_types = { ActivationType::Identity }; +std::vector use_fast_nvfp4_scaling_vec = { + false, + true +}; + } // namespace class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam , - transformer_engine::DType>> {}; + transformer_engine::DType, + bool>> {}; TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { // Skip tests for pre-Blackwell architectures @@ -685,6 +702,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { const ActivationType Act_type = std::get<0>(GetParam()); 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()); // Skip tests if the input tensor is 1D if (tensor_dims.size() < 2) { @@ -702,7 +720,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { } TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, - performTest(OP, tensor_dims); + performTest(OP, tensor_dims, use_fast_math); ); } @@ -724,7 +742,8 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Combine( ::testing::ValuesIn(Activation_types), ::testing::ValuesIn(tensor_dims), - ::testing::Values(DType::kBFloat16)), + ::testing::Values(DType::kBFloat16), + ::testing::ValuesIn(use_fast_nvfp4_scaling_vec)), [](const testing::TestParamInfo& info) { std::string name = to_string(std::get<0>(info.param)); const auto& shape = std::get<1>(info.param); @@ -732,5 +751,8 @@ INSTANTIATE_TEST_SUITE_P( 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; }); diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh index a5c8327cdf..0997b01f7e 100644 --- a/transformer_engine/common/cast/core/common.cuh +++ b/transformer_engine/common/cast/core/common.cuh @@ -35,6 +35,12 @@ inline bool dimensions_supported_by_TMA(const Tensor *const t) { return cols % alignment_requirement == 0; } +__device__ __forceinline__ unsigned char *align_smem_ptr_per_TMA_requirements(unsigned char *p) { + size_t addr = reinterpret_cast(p); + addr = (addr + TMA_SHMEM_ALIGNMENT - 1) & ~(TMA_SHMEM_ALIGNMENT - 1); + return reinterpret_cast(addr); +} + namespace kernel { constexpr size_t THREADS_PER_BLOCK = 256; diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 5da9cc5a5b..99776db281 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -21,6 +21,7 @@ #include "../../util/ptx.cuh" #include "../../utils.cuh" #include "core_nvfp4.cuh" +#include "specialized/quantize_transpose_nvfp4_tuned_1D.cuh" namespace transformer_engine { namespace dispatch { @@ -1159,6 +1160,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, #if FP4_TYPE_SUPPORTED using namespace quantize_transpose_kernel; using namespace ptx; + bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to @@ -1166,6 +1168,11 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, // TODO(Frank): Is there a better way to do this? bool return_transpose = output->has_columnwise_data(); + if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { + quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); + return; + } + constexpr bool COMPUTE_ACTIVATIONS = false; using ParamOP = Empty; constexpr float (*OP)(float, const ParamOP &) = nullptr; 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 new file mode 100644 index 0000000000..af1b01d6b2 --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -0,0 +1,789 @@ +/************************************************************************* + * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_transpose_nvfp4_tuned_1D.cuh + * \brief Tuned kernel to cast to NVFP4 and transpose. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_TUNED_1D_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_TRANSPOSE_NVFP4_TUNED_1D_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_transpose_tuned_kernel { + +using namespace quantization_and_transposition_SF; +using namespace core; +using namespace ptx; + +#if FP4_TYPE_SUPPORTED + +struct TunableConfig { + static constexpr int CHUNK_DIM_Y = 128; + static constexpr int CHUNK_DIM_X = 128; + static constexpr int PREFETCH_STAGES = 1; + static constexpr bool PERSISTENT = false; +}; + +constexpr int SCALE_DIM = 16; // NVFP4 block (x16 elts) +constexpr int THREADS_NUM = 128; +constexpr int ELTS_PER_THREAD = 16; +constexpr int TILE_DIM_Y = 64; +constexpr int TILE_DIM_X = 64; + +static_assert(ELTS_PER_THREAD == SCALE_DIM && "Hardcoded and fixed parameter\0"); + +static_assert((THREADS_NUM * ELTS_PER_THREAD <= TILE_DIM_Y * TILE_DIM_X) && + "Unbalanced threads workload\0"); + +static_assert((TunableConfig::CHUNK_DIM_Y % TILE_DIM_Y == 0) && + "Chunk size Y must be evenly divisible by the tile size Y\0"); +static_assert((TunableConfig::CHUNK_DIM_X % TILE_DIM_X == 0) && + "Chunk size X must be evenly divisible by the tile size X\0"); + +static_assert((TILE_DIM_Y % SCALE_DIM == 0) && + "Tile size Y must be evenly divisible by the scale dim\0"); +static_assert((TILE_DIM_X % SCALE_DIM == 0) && + "Tile size X must be evenly divisible by the scale dim\0"); + +constexpr int TILES_Y = TunableConfig::CHUNK_DIM_Y / TILE_DIM_Y; +constexpr int TILES_X = TunableConfig::CHUNK_DIM_X / TILE_DIM_X; + +constexpr int THREADS_PER_SCALE_ROWWISE = SCALE_DIM / ELTS_PER_THREAD; + +constexpr int SCALES_PER_CHUNK_Y = TunableConfig::CHUNK_DIM_Y / SCALE_DIM; +constexpr int SCALES_PER_CHUNK_X = TunableConfig::CHUNK_DIM_X / SCALE_DIM; + +constexpr int SCALES_PER_TILE_Y = TILE_DIM_Y / SCALE_DIM; +constexpr int SCALES_PER_TILE_X = TILE_DIM_X / SCALE_DIM; + +constexpr int STAGES_Y = TILES_Y; +constexpr int STAGES_X = TILES_X; +constexpr int STAGES = STAGES_Y * STAGES_X; + +constexpr int BUFFS_NUM = TunableConfig::PREFETCH_STAGES + 1; +constexpr int BUFFS_NUM_IN = BUFFS_NUM; +constexpr int BUFFS_NUM_OUT = BUFFS_NUM; +constexpr int BUFFS_NUM_OUT_TR = 2; +constexpr int BUFF_DIM_Y = TILE_DIM_Y; +constexpr int BUFF_DIM_X = TILE_DIM_X; +constexpr int BUFF_SIZE = BUFF_DIM_Y * BUFF_DIM_X; +constexpr int BUFF_SIZE_TOTAL = BUFF_SIZE * BUFFS_NUM; + +// Input buffer (BF16) +constexpr int BUFF_IN_DIM_Y = BUFF_DIM_Y; +constexpr int BUFF_IN_DIM_X = BUFF_DIM_X; +constexpr int BUFF_IN_SIZE = BUFF_IN_DIM_Y * BUFF_IN_DIM_X; +constexpr int BUFF_IN_ELTS_NUM = BUFF_IN_DIM_Y * BUFF_IN_DIM_X; + +// Output buffer (NVFP4) +constexpr int BUFF_OUT_DIM_Y = BUFF_DIM_Y; +constexpr int BUFF_OUT_DIM_X = (BUFF_DIM_X * 4) / 8; +constexpr int BUFF_OUT_SIZE = BUFF_OUT_DIM_Y * BUFF_OUT_DIM_X; + +// Output transpose buffer (NVFP4) +constexpr int BUFF_OUT_TR_DIM_Y = BUFF_DIM_X; +constexpr int BUFF_OUT_TR_DIM_X = (BUFF_DIM_Y * 4) / 8; +constexpr int BUFF_OUT_TR_SIZE = BUFF_OUT_TR_DIM_Y * BUFF_OUT_TR_DIM_X; + +// Manual swizzling parameters to reduce SHMEM bank conflicts +constexpr int PACK_SIZE = 8; +constexpr int WAVES = ELTS_PER_THREAD / PACK_SIZE; + +constexpr int THREADS_X_ROWWISE = TILE_DIM_X / ELTS_PER_THREAD; +constexpr int THREADS_Y_ROWWISE = THREADS_NUM / THREADS_X_ROWWISE; + +constexpr int THREADS_X_TR = TILE_DIM_X / 2; +constexpr int THREADS_Y_TR = THREADS_NUM / THREADS_X_TR; + +constexpr int ITERATIONS_NORMAL = BUFF_DIM_Y / THREADS_Y_ROWWISE; +constexpr int ITERATIONS_TR = SCALES_PER_TILE_Y / THREADS_Y_TR; +static_assert(ITERATIONS_TR >= 1 && "Number of transpose iterations should be >=1\0"); +static_assert((SCALES_PER_TILE_Y % THREADS_Y_TR == 0) && + "Partial transpose iterations are not supported\0"); + +constexpr int BUFF_OUT_IT_OFFSET = BUFF_OUT_TR_DIM_X / ITERATIONS_TR / STAGES; + +static_assert(BUFF_DIM_Y >= SCALE_DIM && + "Number of buffer rows must be greater or equal to the size of the columwise " + "scaling block\0"); +static_assert(TunableConfig::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"); + +// Number of 4-bit elements that span 32 banks (4-byte each) of shared memory +constexpr int TOTAL_BANKS_WIDTH = (32 * 4 * 8) / 4; // 256 + +// Number of threads (rowwise scaling) that span 32 banks (4-byte banks) of shared memory +constexpr int THREADS_PER_BANK = TOTAL_BANKS_WIDTH / ELTS_PER_THREAD; + +using IType = bf16; +using IType2 = typename ptx::FPx2; +using IType3D = IType[BUFFS_NUM_IN][BUFF_IN_DIM_Y][BUFF_IN_DIM_X]; +using IType2x3D = IType2[BUFFS_NUM_IN][BUFF_IN_DIM_Y][BUFF_IN_DIM_X / 2]; +using OType2x3D = fp4e2m1x2[BUFFS_NUM_OUT][BUFF_OUT_DIM_Y][BUFF_OUT_DIM_X]; +using OType2xt3D = fp4e2m1x2[BUFFS_NUM_OUT_TR][BUFF_OUT_TR_DIM_Y][BUFF_OUT_TR_DIM_X]; +using ScalesType2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_Y][SCALES_PER_CHUNK_X]; +using ScalesTypeTr2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_X][SCALES_PER_CHUNK_Y]; +using RNG_t = typename transformer_engine::curanddx::detail::philox4x32_native_state<10>; + +template +struct SCALING_COEFFICIENT_TYPE {}; +template <> +struct SCALING_COEFFICIENT_TYPE { + using type = float; +}; +template <> +struct SCALING_COEFFICIENT_TYPE { + using type = bf16; +}; + +__device__ __forceinline__ float get_amax_of_pair(const IType2 pair) { + return static_cast(__hmax(__habs(pair.x), __habs(pair.y))); +} + +// Compute "correct" per-block encoding scaling factor +template +__device__ __forceinline__ SF_TYPE +compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const float S_enc) { + constexpr float float_max = detail::TypeExtrema::max; + const float scale_rcp = fminf(S_enc / static_cast(S_dec_block), float_max); + return static_cast(scale_rcp); +} + +template +__device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_ptr, + fp4e2m1x2 *__restrict__ sOut_tr_ptr, + nvfp4_scale_t *__restrict__ sSFcolwise_ptr, + const float S_enc_colwise, const int stage_Y, + const int stage_X, const int buff_in, + const int buff_out_tr, RNG_t &rng, + uint4 &random_uint4, int &rnd_idx) { + using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; + + const auto &sIn2x = *reinterpret_cast(sIn_ptr); + auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); + auto &sSFcolwise = *reinterpret_cast(sSFcolwise_ptr); + + const int warp = threadIdx.x / THREADS_PER_WARP; + const int thread_lane = threadIdx.x % THREADS_PER_WARP; + + const int tid_Y_colwise = (thread_lane % 4 + warp) % 4; + const int tid_X_colwise = thread_lane; + + const int thread_offset_Y_colwise = tid_Y_colwise * SCALE_DIM; + const int thread_offset_X_colwise = tid_X_colwise * 2; + + const int in_thread_offset_Y = thread_offset_Y_colwise; + const int in_thread_offset_X = thread_offset_X_colwise / 2; + + const int out_tr_thread_offset_Y = thread_offset_X_colwise; + const int out_tr_thread_offset_X = thread_offset_Y_colwise / 2; + + const int scale_tr_offset_Y = (stage_X * TILE_DIM_X) + 2 * tid_X_colwise; + const int scale_tr_offset_X = (stage_Y * SCALES_PER_TILE_Y) + tid_Y_colwise; + + __align__(8) IType rIn[2][SCALE_DIM]; + // Read (cache) a pair of input elements (S2R). Find NVFP4-block AMAX + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + const IType2 elt_pair = + ptx::ld_shared_b32(&sIn2x[buff_in][in_thread_offset_Y + i][in_thread_offset_X]); + rIn[0][i] = elt_pair.x; + rIn[1][i] = elt_pair.y; + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, elt_pair); + } + const float block_amax[2] = {static_cast(__habs(thread_amax_2x.x)), + static_cast(__habs(thread_amax_2x.y))}; +#pragma unroll + for (int w = 0; w < 2; ++w) { + const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax[w], S_enc_colwise); + + // Store scaling factors to SMEM buffer (R2S) + sSFcolwise[scale_tr_offset_Y + w][scale_tr_offset_X] = S_dec_b_fp8; + + const scaling_coeff_type SFcoefficient = + compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_colwise); + + // Scale elements + __align__(8) uint32_t rOut[SCALE_DIM / 8]; +#pragma unroll + for (int e = 0; e < SCALE_DIM / 8; ++e) { + const uint64_t elts03 = *reinterpret_cast(&rIn[w][8 * e]); + const uint64_t elts47 = *reinterpret_cast(&rIn[w][8 * e + 4]); + if constexpr (USE_STOCHASTIC_ROUNDING) { + const uint32_t rbits03 = core::get_rbits(rng, random_uint4, rnd_idx); + const uint32_t rbits47 = core::get_rbits(rng, random_uint4, rnd_idx); + rOut[e] = ptx::mul_cvt_bf16_to_fp4_8x_stochastic_rounding( + elts03, elts47, SFcoefficient, rbits03, rbits47); + } else { + rOut[e] = ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest(elts03, elts47, + SFcoefficient); + } + } + uint64_t &out_pack_16x = *reinterpret_cast(rOut); + ptx::st_shared_b64(&sOut_tr[buff_out_tr][out_tr_thread_offset_Y + w][out_tr_thread_offset_X], + out_pack_16x); + } +} + +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) { + using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; + + const auto &sIn = *reinterpret_cast(sIn_ptr); + auto &sOut = *reinterpret_cast(sOut_ptr); + auto &sSFrowwise = *reinterpret_cast(sSFrowwise_ptr); + + const int thread_lane = threadIdx.x % THREADS_PER_WARP; + const int bank_group = thread_lane / THREADS_PER_BANK; + + const int tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; + const int tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; + + const int thread_offset_Y_rowwise = tid_Y_rowwise; + const int thread_offset_X_rowwise = tid_X_rowwise * ELTS_PER_THREAD; + + const int SF_thread_offset_rowwise_Y = tid_Y_rowwise; + const int SF_thread_offset_rowwise_X = tid_X_rowwise / THREADS_PER_SCALE_ROWWISE; + + const bool SF_storing_thread = (tid_X_rowwise % THREADS_PER_SCALE_ROWWISE == 0); + + const int stage_rowwise_scales_offset_Y = SF_thread_offset_rowwise_Y + stage_Y * TILE_DIM_Y; + const int stage_rowwise_scales_offset_X = + SF_thread_offset_rowwise_X + stage_X * SCALES_PER_TILE_X; +#pragma unroll + for (int it = 0; it < ITERATIONS_NORMAL; ++it) { + const int it_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; + + __align__(16) IType2 rIn[WAVES][PACK_SIZE / 2]; + + // Read (cache) input elements (S2R). Find NVFP4-block AMAX + 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) % ELTS_PER_THREAD; + const int swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + + // Load elements + __uint128_t &elts_8x = *reinterpret_cast<__uint128_t *>(&rIn[w]); + elts_8x = ptx::ld_shared_b128(&sIn[buff_in][it_offset_Y_rowwise][swizzled_thread_idx]); +#pragma unroll + for (int e = 0; e < PACK_SIZE / 2; ++e) { + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, rIn[w][e]); + } + } + 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); + + // Store scaling factors to SMEM buffer (R2S) + if (SF_storing_thread) { + const int scales_offset_Y = stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE; + const int scales_offset_X = stage_rowwise_scales_offset_X; + sSFrowwise[scales_offset_Y][scales_offset_X] = S_dec_b_fp8; + } + +// Scale elements +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const uint64_t elts03 = *reinterpret_cast(&rIn[w][0]); + const uint64_t elts47 = *reinterpret_cast(&rIn[w][2]); + + uint32_t out_x8; + if constexpr (USE_STOCHASTIC_ROUNDING) { + const uint32_t rbits03 = core::get_rbits(rng, random_uint4, rnd_idx); + const uint32_t rbits47 = core::get_rbits(rng, random_uint4, rnd_idx); + out_x8 = ptx::mul_cvt_bf16_to_fp4_8x_stochastic_rounding( + elts03, elts47, SFcoefficient, rbits03, rbits47); + } else { + out_x8 = ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest(elts03, elts47, + SFcoefficient); + } + + const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % ELTS_PER_THREAD; + const int swizzled_idx = (swizzled_group_idx + thread_offset_X_rowwise) / 2; + ptx::st_shared_b32(&sOut[buff_out][it_offset_Y_rowwise][swizzled_idx], out_x8); + } + } +} + +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, + const __grid_constant__ CUtensorMap tensor_map_output_t, nvfp4_scale_t *const scales_ptr, + nvfp4_scale_t *const scales_t_ptr, const float *noop, const float *const amax_rowwise_ptr, + const float *const amax_colwise_ptr, const size_t rows, const size_t cols, + const size_t scale_stride, const size_t scale_stride_t, const size_t *rng_state) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + const size_t rng_sequence = + threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; + RNG_t rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; + // Index of the random number. It increments each time when used and resets to 0 if reaches 4x + int rnd_idx = 0; + + const bool leading_thread = (threadIdx.x == 0); + + constexpr int buff_elems = BUFF_DIM_Y * BUFF_IN_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + + constexpr int in_mem = buff_size_aligned_in; + + constexpr int out_mem_rowwise_data = buff_size_aligned_out; + constexpr int out_mem_colwise_data = RETURN_TRANSPOSE ? buff_size_aligned_out_t : 0; + constexpr int out_mem_rowwise_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + + // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned + extern __shared__ unsigned char dynamic_shmem[]; + unsigned char *dshmem = common::align_smem_ptr_per_TMA_requirements(dynamic_shmem); + + IType *sIn_ptr = reinterpret_cast(dshmem); + fp4e2m1x2 *sOut_ptr = reinterpret_cast(dshmem + in_mem); + fp4e2m1x2 *sOut_tr_ptr = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); + + auto &sIn = *reinterpret_cast(sIn_ptr); + auto &sOut = *reinterpret_cast(sOut_ptr); + auto &sOut_tr = *reinterpret_cast(sOut_tr_ptr); + + nvfp4_scale_t *sSFrowwise_ptr = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); + nvfp4_scale_t *sSFcolwise_ptr = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + + auto &sSFrowwise = *reinterpret_cast(sSFrowwise_ptr); + auto &sSFcolwise = *reinterpret_cast(sSFcolwise_ptr); + + constexpr int shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + // Compute a global encoding/decoding scaling factors for all S_dec_b + const float S_enc_rowwise = + (amax_rowwise_ptr == nullptr) + ? 1.0f + : core::compute_global_encode_scaling_factor_FP4(*amax_rowwise_ptr); + + const float S_enc_colwise = + (amax_colwise_ptr == nullptr) + ? S_enc_rowwise + : core::compute_global_encode_scaling_factor_FP4(*amax_colwise_ptr); + + __shared__ uint64_t workID_mbar; + __shared__ __uint128_t workID_response; + constexpr uint32_t workID_response_size = sizeof(workID_response); + static_assert(workID_response_size == 16); + + __shared__ uint64_t IN_buff_readable_mbar[BUFFS_NUM]; + + // Coordinates of the first chunk (CTA) to process + int32_t ctaid_X = blockIdx.x; + int32_t ctaid_Y = blockIdx.y; + + // Initialize shared memory barriers with the number of threads participating in them + if (leading_thread) { +#pragma unroll + for (int buff = 0; buff < BUFFS_NUM; ++buff) { + ptx::mbarrier_init(&IN_buff_readable_mbar[buff], 1); + } + ptx::mbarrier_init(&workID_mbar, 1); + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + bool job_finished = false; + int buff_in = 0; + int buff_out = 0; + int buff_out_tr = 0; + int IN_buff_readable_parity[BUFFS_NUM] = {0, 0}; + int ctaid_parity = 0; + +// Prefetch input data only when processing the first chunk, +// which enables the one-iteration overlap throughout the entire kernel life +#pragma unroll + for (int stage = 0; stage < TunableConfig::PREFETCH_STAGES; ++stage) { + const int buff_in = stage; + const int stage_Y = stage / STAGES_X; + const int stage_X = stage % STAGES_X; + + const int stage_offset_Y = stage_Y * TILE_DIM_Y; + const int stage_offset_X = stage_X * TILE_DIM_X; + + const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + const int block_offset_X = ctaid_X * TunableConfig::CHUNK_DIM_X; + + const int global_offset_Y = block_offset_Y + stage_offset_Y; + const int global_offset_X = block_offset_X + stage_offset_X; + + uint64_t *barrier = &IN_buff_readable_mbar[buff_in]; + if (leading_thread) { + uint64_t *dst = reinterpret_cast(&sIn[buff_in]); + const uint64_t *src = reinterpret_cast(&tensor_map_input); + + // Arrive on the barrier and tell how many bytes are expected to come in + ptx::mbarrier_arrive_expect_tx(barrier, shmem_buff_size); + + // Initiate bulk tensor copy + ptx::cp_async_bulk_tensor_2d_global_to_shared(dst, src, global_offset_X, global_offset_Y, + barrier); + } + } + + while (!job_finished) { + const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + const int block_offset_X = ctaid_X * TunableConfig::CHUNK_DIM_X; + + const int block_offset_Y_tr = ctaid_X * TunableConfig::CHUNK_DIM_X; + const int block_offset_X_tr = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + + const int chunk_rows = rows - block_offset_Y; + const int chunk_cols = cols - block_offset_X; + + const int scales_block_offset_Y_rowwise = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + const int scales_block_offset_X_rowwise = ctaid_X * SCALES_PER_CHUNK_X; + const int scales_block_offset_Y_tr = ctaid_X * TunableConfig::CHUNK_DIM_X; + const int scales_block_offset_X_tr = ctaid_Y * SCALES_PER_CHUNK_Y; + + if constexpr (TunableConfig::PERSISTENT) { + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx_cta_relaxed_shared_cta(&workID_mbar, workID_response_size); + ptx::try_cancel_cta(&workID_mbar, &workID_response); + } + } + +#pragma unroll + for (int stage = 0; stage < STAGES; ++stage) { + const int stage_Y = stage / STAGES_X; + const int stage_X = stage % STAGES_X; + + const int stage_offset_Y = stage_Y * TILE_DIM_Y; + const int stage_offset_X = stage_X * TILE_DIM_X; + + if (stage == STAGES - TunableConfig::PREFETCH_STAGES) { + if constexpr (TunableConfig::PERSISTENT) { + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&workID_mbar, ctaid_parity); + ptx::get_cancelled_cta_id_2D(&workID_response, ctaid_X, ctaid_Y); + ctaid_parity ^= 1; + } else { + ctaid_X = -1; + ctaid_Y = -1; + } + if (ctaid_X == -1 && ctaid_Y == -1) { + job_finished = true; + } + } + + // Prefetch next stage Input data + if (!job_finished || (stage < STAGES - TunableConfig::PREFETCH_STAGES)) { + const int next_prefetch_buff = (buff_in + TunableConfig::PREFETCH_STAGES) % BUFFS_NUM; + const int next_prefetch_stage = (stage + TunableConfig::PREFETCH_STAGES) % STAGES; + const int next_prefetch_stage_Y = next_prefetch_stage / STAGES_X; + const int next_prefetch_stage_X = next_prefetch_stage % STAGES_X; + + const int next_prefetch_stage_offset_Y = next_prefetch_stage_Y * TILE_DIM_Y; + const int next_prefetch_stage_offset_X = next_prefetch_stage_X * TILE_DIM_X; + + // Offsets change, because coordinates of the next "to-be-prefetched" CTA do also chage + const int block_offset_Y = ctaid_Y * TunableConfig::CHUNK_DIM_Y; + const int block_offset_X = ctaid_X * TunableConfig::CHUNK_DIM_X; + + const int global_offset_Y = block_offset_Y + next_prefetch_stage_offset_Y; + const int global_offset_X = block_offset_X + next_prefetch_stage_offset_X; + + uint64_t *barrier = &IN_buff_readable_mbar[next_prefetch_buff]; + if (leading_thread) { + uint64_t *dst = reinterpret_cast(&sIn[next_prefetch_buff]); + const uint64_t *src = reinterpret_cast(&tensor_map_input); + + // Arrive on the barrier and tell how many bytes are expected to come in + ptx::mbarrier_arrive_expect_tx(barrier, shmem_buff_size); + + // Initiate bulk tensor copy + ptx::cp_async_bulk_tensor_2d_global_to_shared(dst, src, global_offset_X, global_offset_Y, + barrier); + } + ptx::fence_proxy_async_shared_cta(); + } + + // Wait for the data to have arrived + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&IN_buff_readable_mbar[buff_in], + IN_buff_readable_parity[buff_in]); + IN_buff_readable_parity[buff_in] ^= 1; + + // Wait for TMA transfer to have finished reading shared memory + // I.e. the OUT buffer is ready to be written to + ptx::cp_async_bulk_wait_group_read(); + + // NVFP4 Quantization + rowwise_scaling( + sIn_ptr, sOut_ptr, sSFrowwise_ptr, S_enc_rowwise, stage_Y, stage_X, buff_in, buff_out, + rng, random_uint4, rnd_idx); + + if constexpr (RETURN_TRANSPOSE) { + colwise_scaling( + sIn_ptr, sOut_tr_ptr, sSFcolwise_ptr, S_enc_colwise, stage_Y, stage_X, buff_in, + buff_out_tr, rng, random_uint4, rnd_idx); + } + + // 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 (leading_thread) { + const int global_offset_Y = block_offset_Y + stage_offset_Y; + const int global_offset_X = block_offset_X + stage_offset_X; + const int global_offset_Y_tr = block_offset_Y_tr + stage_offset_X; + const int global_offset_X_tr = block_offset_X_tr + stage_offset_Y; + + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), global_offset_X, + global_offset_Y, reinterpret_cast(&sOut[buff_out])); + + if constexpr (RETURN_TRANSPOSE) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output_t), global_offset_X_tr, + global_offset_Y_tr, reinterpret_cast(&sOut_tr[buff_out_tr])); + } + + // Create a "bulk async-group" out of the previous bulk copy operation + ptx::cp_async_bulk_commit_group(); + } + + buff_in = (buff_in + 1) % BUFFS_NUM_IN; + buff_out = (buff_out + 1) % BUFFS_NUM_OUT; + buff_out_tr = (buff_out_tr + 1) % BUFFS_NUM_OUT_TR; + } // end of stages + + // Vectorized store of scaling factors (S2G) + { + // Rowwise + { + using ScalesVec = Vec; + // number of scales in X dimension of this chunk + const int count = min(SCALES_PER_CHUNK_X, chunk_cols / SCALE_DIM); + + for (size_t row = threadIdx.x; row < TunableConfig::CHUNK_DIM_Y; row += THREADS_NUM) { + const size_t row_global = scales_block_offset_Y_rowwise + row; + if (row_global < rows) { + ScalesVec &scales_vec = *reinterpret_cast(sSFrowwise[row]); + const size_t scale_idx_global = + row_global * scale_stride + scales_block_offset_X_rowwise; + scales_vec.store_to_elts(&scales_ptr[scale_idx_global], 0, count); + } + } + } + + // Colwise + if constexpr (RETURN_TRANSPOSE) { + using ScalesVec = Vec; + // number of scales in Y dimension of this chunk + const int count = min(SCALES_PER_CHUNK_Y, chunk_rows / SCALE_DIM); + + for (size_t row_tr = threadIdx.x; row_tr < TunableConfig::CHUNK_DIM_X; + row_tr += THREADS_NUM) { + const size_t row_tr_global = scales_block_offset_Y_tr + row_tr; + if (row_tr_global < cols) { + ScalesVec &scales_vec = *reinterpret_cast(sSFcolwise[row_tr]); + const size_t scale_idx_global = + row_tr_global * scale_stride_t + scales_block_offset_X_tr; + scales_vec.store_to_elts(&scales_t_ptr[scale_idx_global], 0, count); + } + } + } + + if (!job_finished) { + // Ensures all reads from SFs buffer have completed and it's ready to be reused + __syncthreads(); + } + } + } + + if (leading_thread) { +#pragma unroll + for (int buff = 0; buff < BUFFS_NUM; ++buff) { + ptx::mbarrier_invalid(&IN_buff_readable_mbar[buff]); + } + ptx::mbarrier_invalid(&workID_mbar); + } +#else + NVTE_DEVICE_ERROR("sm_100 or higher is required."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +#endif // FP4_TYPE_SUPPORTED +} // namespace quantize_transpose_tuned_kernel + +inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace quantize_transpose_tuned_kernel; + using namespace ptx; + + 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; + + // If transposed output is allocated, return the transposed data + // Otherwise, it's not necesary to return the transposed data. + const bool return_transpose = output->has_columnwise_data(); + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output", false); + + NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); + 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"); + + if (return_transpose) { + NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), + "Transposed output must have FP4 type."); + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, + "Transposed scaling tensor must be allocated"); + } + + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + + NVTE_CHECK(rows % 32 == 0, + "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA + NVTE_CHECK(cols % 32 == 0, + "Number of tensor cols must be a multiple of 32"); // 16B alignment for TMA + + const int blocks_Y = DIVUP(rows, static_cast(TunableConfig::CHUNK_DIM_Y)); + const int blocks_X = DIVUP(cols, static_cast(TunableConfig::CHUNK_DIM_X)); + const dim3 grid(blocks_X, blocks_Y); + const int block_size = THREADS_NUM; + + const size_t scale_stride = output->scale_inv.shape[1]; + const size_t scale_stride_transpose = + return_transpose ? output->columnwise_scale_inv.shape[1] : 0; + + nvfp4_scale_t *const scales_ptr = reinterpret_cast(output->scale_inv.dptr); + nvfp4_scale_t *const scales_transpose_ptr = + reinterpret_cast(output->columnwise_scale_inv.dptr); + + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + const float *const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); + const float *const amax_colwise_ptr = + reinterpret_cast(output->columnwise_amax.dptr); + + const NVTETensor rng_state_tensor = (quant_config != nullptr) ? quant_config->rng_state : nullptr; + const size_t *rng_state = nullptr; + if (rng_state_tensor != nullptr) { + 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}, + "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); + } + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_output{}; + alignas(64) CUtensorMap tensor_map_output_transpose{}; + + 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, output->data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, + 4); + if (return_transpose) { + create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, + BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); + } + + constexpr int buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr int buff_elems_total_in = BUFFS_NUM_IN * buff_elems; + constexpr int buff_size_aligned_in = + DIVUP_TO_MULTIPLE(buff_elems_total_in * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT * BUFF_OUT_SIZE, TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_aligned_out_t = + DIVUP_TO_MULTIPLE(BUFFS_NUM_OUT_TR * BUFF_OUT_TR_SIZE, TMA_SHMEM_ALIGNMENT); + + constexpr int buff_size_scales = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + constexpr int buff_size_scales_transpose = DIVUP_TO_MULTIPLE( + TunableConfig::CHUNK_DIM_X * SCALES_PER_CHUNK_Y * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT); + + const int in_mem = buff_size_aligned_in; + + const int out_data_mem = buff_size_aligned_out; + const int out_data_transpose_mem = return_transpose ? buff_size_aligned_out_t : 0; + const int out_scales_mem = buff_size_scales; + const int out_scales_transpose_mem = return_transpose ? buff_size_scales_transpose : 0; + + const int out_mem = out_data_mem + out_data_transpose_mem; + + const int dshmem_size = + in_mem + out_mem + out_scales_transpose_mem + out_scales_mem + TMA_SHMEM_ALIGNMENT; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + 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); + }););); +#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_TRANSPOSE_NVFP4_TUNED_1D_CUH_ diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 4cdd8297a8..9bcf6e2289 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -164,6 +164,18 @@ __device__ __forceinline__ void mbarrier_arrive_expect_tx(uint64_t *mbar, const #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +__device__ __forceinline__ void mbarrier_arrive_expect_tx_cta_relaxed_shared_cta( + uint64_t *mbar, const uint32_t tx_count) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile("mbarrier.arrive.expect_tx.relaxed.cta.shared::cta.b64 _, [%0], %1;" ::"r"(mbar_ptr), + "r"(tx_count)); +#else + NVTE_DEVICE_ERROR( + "mbarrier_arrive_expect_tx_cta_relaxed_shared_cta is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + __device__ __forceinline__ void fence_mbarrier_init_release_cluster() { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) asm volatile("fence.mbarrier_init.release.cluster;"); @@ -243,6 +255,75 @@ __device__ __forceinline__ void mbarrier_wait_parity(uint64_t *mbar, const uint3 #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +__device__ __forceinline__ void mbarrier_wait_parity_acquire_cta_shared_cta(uint64_t *mbar, + uint32_t phase_parity) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + asm volatile( + "{\n\t" + ".reg .b64 r1; \n\t" + ".reg .pred waitComplete; \n\t" // predicate representing if barrier condition is met + "WAIT: \n\t" // loop around barrier wait + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 waitComplete, [%0], %1; \n\t" + "@waitComplete bra DONE; \n\t" // mbarrier conditions are met + "bra WAIT; \n\t" // just a time-out, try again + "DONE: \n\t" + "}\n\t" + : + : "r"(mbar_ptr), "r"(phase_parity) + : "memory"); +#else + NVTE_DEVICE_ERROR("mbarrier_wait_parity_acquire_cta_shared_cta is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void try_cancel_cta(uint64_t *mbar, __uint128_t *response_data_ptr) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + uint32_t mbar_ptr = __cvta_generic_to_shared(mbar); + uint32_t workID_response = __cvta_generic_to_shared(response_data_ptr); + asm volatile( + "clusterlaunchcontrol.try_cancel.async.mbarrier::complete_tx::bytes.multicast::cluster::" + "all.b128 " + "[%0], [%1];" ::"r"(workID_response), + "r"(mbar_ptr)); + } else { + NVTE_DEVICE_ERROR( + "Cluster Launch Control PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } +} + +__device__ __forceinline__ void get_cancelled_cta_id_2D(__uint128_t *response_data_ptr, + int32_t &ctaid_X, int32_t &ctaid_Y) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + uint32_t workID_response = __cvta_generic_to_shared(response_data_ptr); + asm volatile( + "{\n\t" + ".reg .s32 x_ctaid; \n\t" + ".reg .s32 y_ctaid; \n\t" + "mov .s32 x_ctaid, -1; \n\t" + "mov .s32 y_ctaid, -1; \n\t" + ".reg.b128 try_cancel_response; \n\t" + "ld.shared.b128 try_cancel_response, [%2]; \n\t" + ".reg .pred P1; \n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 P1, try_cancel_response; \n\t" + "@P1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {x_ctaid, y_ctaid, _, " + "_}, try_cancel_response; \n\t" + "mov .s32 %0, x_ctaid; \n\t" + "mov .s32 %1, y_ctaid; \n\t" + "}\n\t" + : "=r"(ctaid_X), "=r"(ctaid_Y) + : "r"(workID_response) + : "memory"); + } else { + NVTE_DEVICE_ERROR( + "Cluster Launch Control PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } +} + constexpr uint32_t FP32_MANTISSA_BITS = 23; constexpr uint32_t FP32_EXPONENT_BIAS = 127; @@ -657,6 +738,179 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x(const float2 in01, c return mul_cvt_fp32_to_fp4_4x_with_rn(in01, in23, scale, rbits); } } + +template +__device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_round_to_nearest( + const uint64_t in03, const uint64_t in47, const SCALING_COEFFICIENT_TYPE scaling_coefficient) { + uint32_t out_8x = 0; + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + if constexpr (std::is_same::value) { + asm volatile( + "{\n" + ".reg.f32 zero; \n\t" + "mov.b32 zero, 0; \n\t" + ".reg.b16 scaling_coeff; \n\t" + "mov.b16 scaling_coeff, %3; \n\t" + ".reg.b16 v0_h, v1_h, v2_h, v3_h, v4_h, v5_h, v6_h, v7_h; \n\t" + "mov.b64 {v0_h, v1_h, v2_h, v3_h}, %1; \n\t" + "mov.b64 {v4_h, v5_h, v6_h, v7_h}, %2; \n\t" + + ".reg.f32 v0, v1, v2, v3, v4, v5, v6, v7; \n\t" + "fma.rn.f32.bf16 v0, v0_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v1, v1_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v2, v2_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v3, v3_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v4, v4_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v5, v5_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v6, v6_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v7, v7_h, scaling_coeff, zero; \n\t" + + ".reg.b8 f0, f1, f2, f3; \n\t" + // Elements reordered to match e2m1x4 packing order (v1,v0) + "cvt.rn.satfinite.e2m1x2.f32 f0, v1, v0;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v3, v2;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f2, v5, v4;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f3, v7, v6;\n\t" + "mov.b32 %0, {f0, f1, f2, f3};\n" + "}" + : "=r"(out_8x) + : "l"(in03), "l"(in47), "h"(reinterpret_cast(scaling_coefficient))); + } else if constexpr (std::is_same::value) { + asm volatile( + "{\n" + ".reg.b64 scaling_coeff_2x; \n\t" + "mov.b64 scaling_coeff_2x, {%3, %3}; \n\t" + ".reg.b16 v0_bf16, v1_bf16, v2_bf16, v3_bf16, v4_bf16, v5_bf16, v6_bf16, v7_bf16; \n\t" + "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16}, %1; \n\t" + "mov.b64 {v4_bf16, v5_bf16, v6_bf16, v7_bf16}, %2; \n\t" + + ".reg.b32 v0, v1, v2, v3, v4, v5, v6, v7; \n\t" + "cvt.f32.bf16 v0, v0_bf16; \n\t" + "cvt.f32.bf16 v1, v1_bf16; \n\t" + "cvt.f32.bf16 v2, v2_bf16; \n\t" + "cvt.f32.bf16 v3, v3_bf16; \n\t" + "cvt.f32.bf16 v4, v4_bf16; \n\t" + "cvt.f32.bf16 v5, v5_bf16; \n\t" + "cvt.f32.bf16 v6, v6_bf16; \n\t" + "cvt.f32.bf16 v7, v7_bf16; \n\t" + + ".reg.b64 v01, v23, v45, v67; \n\t" + "mov.b64 v01, {v0, v1}; \n\t" + "mov.b64 v23, {v2, v3}; \n\t" + "mov.b64 v45, {v4, v5}; \n\t" + "mov.b64 v67, {v6, v7}; \n\t" + "mul.f32x2 v01, v01, scaling_coeff_2x; \n\t" + "mul.f32x2 v23, v23, scaling_coeff_2x; \n\t" + "mul.f32x2 v45, v45, scaling_coeff_2x; \n\t" + "mul.f32x2 v67, v67, scaling_coeff_2x; \n\t" + // Elements reordered to match the packing order (v1,v0) + "mov.b64 {v1, v0}, v01; \n\t" + "mov.b64 {v3, v2}, v23; \n\t" + "mov.b64 {v5, v4}, v45; \n\t" + "mov.b64 {v7, v6}, v67; \n\t" + + ".reg.b8 f0, f1, f2, f3; \n\t" + "cvt.rn.satfinite.e2m1x2.f32 f0, v0, v1;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f1, v2, v3;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f2, v4, v5;\n\t" + "cvt.rn.satfinite.e2m1x2.f32 f3, v6, v7;\n\t" + "mov.b32 %0, {f0, f1, f2, f3};\n\t" + "}" + : "=r"(out_8x) + : "l"(in03), "l"(in47), "f"(scaling_coefficient)); + } else { + NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); + } + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return out_8x; +} + +template +__device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding( + const uint64_t in03, const uint64_t in47, const SCALING_COEFFICIENT_TYPE scaling_coefficient, + const uint32_t rbits03, const uint32_t rbits47) { + uint32_t out_8x = 0; + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + if constexpr (std::is_same::value) { + asm volatile( + "{\n" + ".reg.f32 zero; \n\t" + "mov.b32 zero, 0; \n\t" + ".reg.b16 scaling_coeff; \n\t" + "mov.b16 scaling_coeff, %3; \n\t" + ".reg.b16 v0_h, v1_h, v2_h, v3_h, v4_h, v5_h, v6_h, v7_h; \n\t" + "mov.b64 {v0_h, v1_h, v2_h, v3_h}, %1; \n\t" + "mov.b64 {v4_h, v5_h, v6_h, v7_h}, %2; \n\t" + + ".reg.f32 v0, v1, v2, v3, v4, v5, v6, v7; \n\t" + "fma.rn.f32.bf16 v0, v0_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v1, v1_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v2, v2_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v3, v3_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v4, v4_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v5, v5_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v6, v6_h, scaling_coeff, zero; \n\t" + "fma.rn.f32.bf16 v7, v7_h, scaling_coeff, zero; \n\t" + + ".reg.b16 b03, b47; \n\t" + // Elements reordered to match e2m1x4 packing order (v3,v2,v1,v0) + "cvt.rs.satfinite.e2m1x4.f32 b03, {v3, v2, v1, v0}, %4; \n\t" + "cvt.rs.satfinite.e2m1x4.f32 b47, {v7, v6, v5, v4}, %5; \n\t" + "mov.b32 %0, {b03, b47};\n" + "}" + : "=r"(out_8x) + : "l"(in03), "l"(in47), "h"(reinterpret_cast(scaling_coefficient)), + "r"(rbits03), "r"(rbits47)); + } else if constexpr (std::is_same::value) { + asm volatile( + "{\n" + ".reg.b16 v0_bf16, v1_bf16, v2_bf16, v3_bf16, v4_bf16, v5_bf16, v6_bf16, v7_bf16; \n\t" + "mov.b64 {v0_bf16, v1_bf16, v2_bf16, v3_bf16}, %1; \n\t" + "mov.b64 {v4_bf16, v5_bf16, v6_bf16, v7_bf16}, %2; \n\t" + + ".reg.b32 v0, v1, v2, v3, v4, v5, v6, v7; \n\t" + "cvt.f32.bf16 v0, v0_bf16; \n\t" + "cvt.f32.bf16 v1, v1_bf16; \n\t" + "cvt.f32.bf16 v2, v2_bf16; \n\t" + "cvt.f32.bf16 v3, v3_bf16; \n\t" + "cvt.f32.bf16 v4, v4_bf16; \n\t" + "cvt.f32.bf16 v5, v5_bf16; \n\t" + "cvt.f32.bf16 v6, v6_bf16; \n\t" + "cvt.f32.bf16 v7, v7_bf16; \n\t" + + "mul.f32 v0, v0, %3; \n\t" + "mul.f32 v1, v1, %3; \n\t" + "mul.f32 v2, v2, %3; \n\t" + "mul.f32 v3, v3, %3; \n\t" + "mul.f32 v4, v4, %3; \n\t" + "mul.f32 v5, v5, %3; \n\t" + "mul.f32 v6, v6, %3; \n\t" + "mul.f32 v7, v7, %3; \n\t" + ".reg.b16 b03, b47; \n\t" + // Elements reordered to match e2m1x4 packing order (v3,v2,v1,v0) + "cvt.rs.satfinite.e2m1x4.f32 b03, {v3, v2, v1, v0}, %4; \n\t" + "cvt.rs.satfinite.e2m1x4.f32 b47, {v7, v6, v5, v4}, %5; \n\t" + "mov.b32 %0, {b03, b47};\n" + "}" + : "=r"(out_8x) + : "l"(in03), "l"(in47), "f"(scaling_coefficient), "r"(rbits03), "r"(rbits47)); + } else { + NVTE_DEVICE_ERROR("Not supported scaling coefficient type."); + } + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return out_8x; +} + #endif // FP4_TYPE_SUPPORTED // SIMD like "Fused" cast + multiplication (x2) @@ -1508,6 +1762,58 @@ __device__ __forceinline__ floatx4 up_cast(const bf16x4 &in) { return out; } +// Loads single BF16/FP16 element from shared memory state space +__device__ __forceinline__ bf16 ld_shared_b16(const bf16 *__restrict__ src_smem) { + const uint32_t src_smem_ptr = __cvta_generic_to_shared(src_smem); + bf16 dst; + asm volatile("ld.shared.b16 %0, [%1];" + : "=h"(reinterpret_cast(dst)) + : "r"(src_smem_ptr)); + return dst; +} + +// Loads pair of BF16/FP16 values from shared memory state space +__device__ __forceinline__ bf16x2 ld_shared_b32(const bf16x2 *__restrict__ src_smem) { + const uint32_t src_smem_ptr = __cvta_generic_to_shared(src_smem); + bf16x2 dst; + asm volatile("ld.shared.b32 %0, [%1];" + : "=r"(reinterpret_cast(dst)) + : "r"(src_smem_ptr)); + return dst; +} + +// Loads 8x BF16 values from shared memory state space +__device__ __forceinline__ __uint128_t ld_shared_b128(const bf16 *__restrict__ src_smem) { + uint64_t elts03, elts47; + const uint32_t src_smem_ptr = __cvta_generic_to_shared(src_smem); + asm volatile( + "{\n\t" + ".reg.b128 xy; \n\t" + "ld.shared.b128 xy, [%2]; \n\t" + "mov.b128 {%0, %1}, xy; \n" + "}\n" + : "=l"(elts03), "=l"(elts47) + : "r"(src_smem_ptr)); + return (static_cast<__uint128_t>(elts47) << 64) | static_cast<__uint128_t>(elts03); +} + +#if FP4_TYPE_SUPPORTED +// Vectorized store of x8 FP4 elements into shared memory state space +__device__ __forceinline__ void st_shared_b32(fp4e2m1x2 *__restrict__ dst_smem, + uint32_t fp4_pack_x8) { + const uint32_t dst_smem_ptr = __cvta_generic_to_shared(dst_smem); + asm volatile("st.shared.b32 [%0], %1;" : : "r"(dst_smem_ptr), "r"(fp4_pack_x8)); +} +#endif + +// Vectorized store of x16 FP4 elements into shared memory state space +#if FP4_TYPE_SUPPORTED +__device__ __forceinline__ void st_shared_b64(fp4e2m1x2 *__restrict__ dst_smem, + uint64_t fp4_pack_x16) { + const uint32_t dst_smem_ptr = __cvta_generic_to_shared(dst_smem); + asm volatile("st.shared.b64 [%0], %1;" : : "r"(dst_smem_ptr), "l"(fp4_pack_x16)); +} +#endif } // namespace ptx namespace { From 36f4e4517c7e58f73f816a41378357bcf6872c8f Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Wed, 21 Jan 2026 19:10:09 +0100 Subject: [PATCH 173/521] Fixed the year to 2026 (#2611) Signed-off-by: Oleg Goncharov --- .../nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 af1b01d6b2..4119001686 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 @@ -1,5 +1,5 @@ /************************************************************************* - * Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * * See LICENSE for license information. ************************************************************************/ From 605786f4c3f2c2c5a90a913a627bf9b54ccd3a68 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 21 Jan 2026 14:25:09 -0800 Subject: [PATCH 174/521] [pyTorch] CPU performance optimizations (#2439) * PoC of the changes Signed-off-by: Przemek Tredak * Early exit from the Free function for the empty tensor Signed-off-by: Przemek Tredak * Use the proper function for nvtx range Signed-off-by: Przemek Tredak * Only do mark_not_offload when the cpu_offloading is enabled Signed-off-by: Przemek Tredak * First pass on making the setattr issue not come back Signed-off-by: Przemek Tredak * Actually add pytest.ini Signed-off-by: Przemek Tredak * Changes to __init__ Signed-off-by: Przemek Tredak * A different way Signed-off-by: Przemek Tredak * WAR the fact that it is not possible to set __setattr__ dynamically Signed-off-by: Przemek Tredak * Simpler solution and fixes Signed-off-by: Przemek Tredak * Fix for the inference mode DPA Signed-off-by: Przemek Tredak * Start of debugging debug tools Signed-off-by: Przemek Tredak * More fixes in debug Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Speculative moving the validate_name to the constructor 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 * Making the debug tools names saner Signed-off-by: Przemek Tredak * Change the setattr usage in the tensor parallel group setting Signed-off-by: Przemek Tredak * Adding try/finally - it does not seem to impact the time in observable way Signed-off-by: Przemek Tredak * Fixing lint issues and the thunder test Signed-off-by: Przemek Tredak * Fix 1 of the debug tests Signed-off-by: Przemek Tredak * Removed the warning and enforcement in the CI Signed-off-by: Przemek Tredak * try-finally in the context manager Signed-off-by: Przemek Tredak * Fixing the debug tests Signed-off-by: Przemek Tredak Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Przemek Tredak Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 2 +- tests/pytorch/debug/test_sanity.py | 25 ++- .../common/transformer_engine.cpp | 4 +- .../dot_product_attention.py | 21 +- .../pytorch/attention/multi_head_attention.py | 5 +- transformer_engine/pytorch/distributed.py | 8 +- transformer_engine/pytorch/module/base.py | 188 ++++++++++-------- .../pytorch/module/grouped_linear.py | 9 +- .../pytorch/module/layernorm_linear.py | 13 +- .../pytorch/module/layernorm_mlp.py | 13 +- transformer_engine/pytorch/module/linear.py | 14 +- transformer_engine/pytorch/transformer.py | 11 +- 12 files changed, 170 insertions(+), 143 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 9111d3511c..6fe0ffdaee 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2790,7 +2790,7 @@ def forward( cu_seqlens, max_s, ) -> torch.Tensor: - with self.prepare_forward(inp, num_gemms=3) as inp: + with self.prepare_forward_ctx(inp, num_gemms=3) as inp: out = _custom_mha_fp8.apply( inp, self.qkv_weight, diff --git a/tests/pytorch/debug/test_sanity.py b/tests/pytorch/debug/test_sanity.py index aee5474e76..2bc4b35590 100644 --- a/tests/pytorch/debug/test_sanity.py +++ b/tests/pytorch/debug/test_sanity.py @@ -30,10 +30,17 @@ stats: [min, max, mean, std, l1_norm, l2_norm, cur_amax, dynamic_range] start_step : 0 end_step: 1 +""", + "log_fp8": """log_fp8: + layers: + layer_types: [linear] + enabled: + True + transformer_engine: LogFp8TensorStats: enabled: True tensors: [activation, gradient, weight] - stats: [underflows, overflows] + stats: [underflows%] start_step : 0 end_step: 1 """, @@ -46,22 +53,26 @@ FakeQuant: enabled: True gemms: [fprop, dgrad, wgrad] + tensors: [activation, weight, gradient] quant_format: FP8E5M2 """, } +# Configs that require FP8 to be enabled +fp8_required_configs = {"log_fp8"} + def _get_model(model_key): if model_key == "linear": - return te.Linear(D, D) + return te.Linear(D, D, name="layer") if model_key == "layernorm_linear": - return te.LayerNormLinear(D, D) + return te.LayerNormLinear(D, D, name="layer") if model_key == "layernorm_mlp": - return te.LayerNormMLP(D, D, D) + return te.LayerNormMLP(D, D, D, name="layer") if model_key == "mha_attention": - return te.MultiheadAttention(D, H) + return te.MultiheadAttention(D, H, name="layer") if model_key == "transformer_layer": - return te.TransformerLayer(D, D, H) + return te.TransformerLayer(D, D, H, name="layer") def _run_forward_backward(model, fp8): @@ -95,4 +106,6 @@ def _run_test(model_key, fp8, config, feature_dirs, config_file, log_dir): def test_sanity_debug(model_key, fp8, config_key, feature_dirs): if fp8 and not fp8_available: pytest.skip(reason_for_no_fp8) + if not fp8 and config_key in fp8_required_configs: + pytest.skip(f"Config '{config_key}' requires FP8") _run_test(model_key, fp8, configs[config_key], feature_dirs) diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 6880dd560a..06971443dd 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -454,9 +454,9 @@ class TensorAllocator { } void Free(NVTETensor t) { - std::lock_guard lock(mutex); 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 @@ -564,9 +564,9 @@ class GroupedTensorAllocator { } void Free(NVTEGroupedTensor t) { - std::lock_guard lock(mutex); uintptr_t index = reinterpret_cast(t); if (index == 0) return; + std::lock_guard lock(mutex); NVTE_CHECK(index <= memory.size(), "Invalid grouped tensor."); free_list.push_back(index); // Clean up 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 6e5a12a103..51ffbc2e48 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 @@ -676,9 +676,9 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # assume attention uses the same fp8_group as GEMMs fp8_group = FP8GlobalStateManager.get_fp8_group() - self.fp8_parameters = FP8GlobalStateManager.with_fp8_parameters() - self.fp8 = FP8GlobalStateManager.is_fp8_enabled() - self.fp8_calibration = FP8GlobalStateManager.is_fp8_calibration() + self.fast_setattr("fp8_parameters", FP8GlobalStateManager.with_fp8_parameters()) + self.fast_setattr("fp8", FP8GlobalStateManager.is_fp8_enabled()) + self.fast_setattr("fp8_calibration", FP8GlobalStateManager.is_fp8_calibration()) fp8_enabled = self.fp8 or self.fp8_calibration self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration if self.fp8_parameters or fp8_enabled: @@ -703,7 +703,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: ) else: # If fp8 isn't enabled, turn off and return. - self.fp8_initialized = False + self.fast_setattr("fp8_initialized", False) return if self.fp8_parameters and not self.fp8_initialized: @@ -721,7 +721,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Allocate scales and amaxes self.init_fp8_meta_tensors(fp8_recipes) - self.fp8_initialized = True + self.fast_setattr("fp8_initialized", True) self.fp8_meta["recipe"] = fp8_recipe_dpa if fp8_recipe != fp8_recipe_dpa: @@ -1000,7 +1000,7 @@ def forward( cases. It is ignored for other backends and when context parallelism is enabled. """ - with self.prepare_forward( + with self.prepare_forward_ctx( query_layer, num_gemms=3, allow_non_contiguous=True, @@ -1145,10 +1145,11 @@ def forward( if attn_mask_type == "padding_causal": attn_mask_type = attn_mask_type + "_bottom_right" - self.attention_type = "cross" - self.flash_attention.attention_type = self.attention_type - self.fused_attention.attention_type = self.attention_type - self.unfused_attention.attention_type = self.attention_type + if self.attention_type != "cross": + self.fast_setattr("attention_type", "cross") + self.flash_attention.attention_type = self.attention_type + self.fused_attention.attention_type = self.attention_type + self.unfused_attention.attention_type = self.attention_type query_layer, key_layer, value_layer = [ x.contiguous() if not x.is_contiguous() else x diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index f875fd1e0a..d813e7c8f1 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -8,7 +8,6 @@ from typing import Callable, List, Optional, Tuple, Union import torch -from transformer_engine.debug.pytorch.debug_state import TEDebugState from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule @@ -335,6 +334,7 @@ def __init__( self.hidden_size_kv = self.hidden_size_per_attention_head * self.num_gqa_groups self.name = name + TransformerEngineBaseModule._validate_name(self) common_gemm_kwargs = { "fuse_wgrad_accumulation": fuse_wgrad_accumulation, @@ -739,9 +739,6 @@ def forward( core_attention_bias_type in AttnBiasTypes ), f"core_attention_bias_type {core_attention_bias_type} is not supported!" - if TEDebugState.debug_enabled: - TransformerEngineBaseModule._validate_name(self) - # ================================================= # Pre-allocate memory for key-value cache for inference # ================================================= diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 004a04ab4c..f269e21b8c 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -729,8 +729,8 @@ def checkpoint( if isinstance(function, TransformerEngineBaseModule): # If this TE module is FSDP-wrapped, clear its FSDP group information because there's no need # to scatter/gather activations that we will recompute anyway. - setattr(function, "fsdp_wrapped", False) - setattr(function, "fsdp_group", None) + function.fast_setattr("fsdp_wrapped", False) + function.fast_setattr("fsdp_group", None) # Otherwise discard unused te.utils.checkpoint.checkpoint() arguments # and execute TE's own checkpointing @@ -2022,7 +2022,7 @@ def prepare_te_modules_for_fsdp(fsdp_root: torch.nn.Module) -> None: ) root_state = _get_module_fsdp_state(fsdp_root) assert root_state is not None, "Root module does not have a valid _FSDPState." - setattr(fsdp_root.module, "fsdp_group", root_state.process_group) + fsdp_root.module.fast_setattr("fsdp_group", root_state.process_group) # Iterate through all FSDP-wrapped submodules and inject FSDP information into TE modules fsdp_states, fsdp_modules = _get_fsdp_states_with_modules(fsdp_root) @@ -2033,7 +2033,7 @@ def prepare_te_modules_for_fsdp(fsdp_root: torch.nn.Module) -> None: "TE modules with primary weights in FP8 cannot be FSDP-wrapped. " "Please initialize your model without the te.quantized_model_init(...) context." ) - setattr(fsdp_module.module, "fsdp_group", state.process_group) + fsdp_module.module.fast_setattr("fsdp_group", state.process_group) class FullyShardedDataParallel(FSDP): diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 875d245a8f..841cdf04ca 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -10,9 +10,8 @@ import warnings from enum import Enum from abc import ABC, abstractmethod -from typing import Any, Dict, Generator, List, Optional, Set, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union from contextlib import contextmanager -import logging from types import MethodType import torch @@ -50,6 +49,8 @@ is_non_tn_fp8_gemm_supported, torch_get_autocast_gpu_dtype, get_nvtx_range_context, + nvtx_range_push, + nvtx_range_pop, ) from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from ...common.recipe import DelayedScaling, Recipe @@ -605,10 +606,10 @@ def fill_userbuffers_buffer_for_all_gather( class TransformerEngineBaseModule(torch.nn.Module, ABC): """Base TE module.""" - def __init__(self) -> None: + def __init__(self, name: Optional[str] = None) -> None: super().__init__() assert torch.cuda.is_available(), "TransformerEngine needs CUDA." - self.name = None + self.name = name self.next_iter_when_debug_should_be_run = 0 self.fp8_initialized = False self.fp8 = False @@ -633,26 +634,22 @@ def __init__(self) -> None: if not TEDebugState.debug_enabled: TEDebugState.initialize() + self._validate_name() - # Names of attributes that can be set quickly (see __setattr__ - # method) - _fast_setattr_names: Set[str] = { - "activation_dtype", - "fp8", - "fp8_initialized", - "fp8_calibration", - "fp8_parameters", - } + def fast_setattr(self, name: str, value: Any) -> None: + """ + Fast version of the Module's set attribute function. + Should be used for regular attributes, but not properties nor parameters/buffers. + """ + self.__dict__[name] = value - def __setattr__(self, name: str, value: Any) -> None: - if name in TransformerEngineBaseModule._fast_setattr_names: - # torch.nn.Module has a custom __setattr__ that handles - # modules, parameters, and buffers. This is unnecessary - # overhead when setting plain attrs. - self.__dict__[name] = value - else: - # Default case - super().__setattr__(name, value) + def module_setattr(self, name: str, value: Any) -> None: + """ + Regular version of the Module's set attribute function. + Should be used only when the fast version cannot be used - for the properties, + parameters and buffers. + """ + super().__setattr__(name, value) def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> None: """ @@ -773,7 +770,7 @@ def init_fp8_meta_tensors(self, recipe: Recipe) -> None: self.set_meta_tensor(True, recipe) self.set_meta_tensor(False, recipe) - self.fp8_meta_tensors_initialized = True + self.fast_setattr("fp8_meta_tensors_initialized", True) def get_fp8_meta_tensors(self) -> None: """Get scales and amaxes.""" @@ -930,7 +927,7 @@ def set_activation_dtype(self, inp: torch.Tensor) -> None: """Get activation data type for AMP.""" # Native AMP (`torch.autocast`) gets highest priority if torch.is_autocast_enabled(): - self.activation_dtype = torch_get_autocast_gpu_dtype() + self.fast_setattr("activation_dtype", torch_get_autocast_gpu_dtype()) return # All checks after this have already been performed once, thus skip @@ -945,7 +942,7 @@ def set_activation_dtype(self, inp: torch.Tensor) -> None: "Data types for parameters must match when outside of autocasted region. " f" Found input dtype: {dtype} and {name!r} dtype: {param.dtype}" ) - self.activation_dtype = dtype + self.fast_setattr("activation_dtype", dtype) def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None: """ @@ -957,8 +954,8 @@ def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> N tp_group : ProcessGroup, default = None tensor parallel process group. """ - self.tp_group = tp_group - self.tp_group_initialized = True + self.fast_setattr("tp_group", tp_group) + self.fast_setattr("tp_group_initialized", True) def _get_fp8_params(self) -> Union[List[torch.Tensor], None]: """returns the FP8 weights.""" @@ -974,48 +971,51 @@ def _get_fp8_params(self) -> Union[List[torch.Tensor], None]: # assume FP8 execution. def init_fp8_metadata(self, num_gemms: int = 1) -> None: """Initialize fp8 related metadata and tensors during fprop.""" - _original_recipe = self.fp8_meta.get("recipe", None) - - self.fp8_parameters = FP8GlobalStateManager.with_fp8_parameters() - self.fp8 = FP8GlobalStateManager.is_fp8_enabled() - self.fp8_calibration = FP8GlobalStateManager.is_fp8_calibration() - fp8_enabled = self.fp8 or self.fp8_calibration - self.fp8_meta["fp8_checkpoint"] = self.fp8 or self.fp8_calibration - - if self.fp8_parameters or fp8_enabled: - if ( - self.fp8_initialized - and FP8GlobalStateManager.get_fp8_recipe() == self.fp8_meta["recipe"] - ): + meta = self.fp8_meta + + fp8 = FP8GlobalStateManager.is_fp8_enabled() + fp8_parameters = FP8GlobalStateManager.with_fp8_parameters() + fp8_calibration = FP8GlobalStateManager.is_fp8_calibration() + self.fast_setattr("fp8_parameters", fp8_parameters) + self.fast_setattr("fp8", fp8) + self.fast_setattr("fp8_calibration", fp8_calibration) + fp8_enabled = fp8 or fp8_calibration + meta["fp8_checkpoint"] = fp8_enabled + + _original_recipe = None + + if fp8_parameters or fp8_enabled: + _original_recipe = meta.get("recipe", None) + if self.fp8_initialized and FP8GlobalStateManager.get_fp8_recipe() == _original_recipe: # FP8 init has already been run and recipe is the same, don't do anything. return - self.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() else: # If fp8 isn't enabled, turn off and return. - self.fp8_initialized = False + self.fast_setattr("fp8_initialized", False) return - if self.fp8_parameters and not self.fp8_initialized: - self.fp8_meta["num_gemms"] = num_gemms - self.init_fp8_meta_tensors(self.fp8_meta["recipe"]) + if fp8_parameters and not self.fp8_initialized: + meta["num_gemms"] = num_gemms + self.init_fp8_meta_tensors(meta["recipe"]) if fp8_enabled: # Set FP8 and other FP8 metadata - self.fp8_meta["num_gemms"] = num_gemms - self.fp8_meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() + meta["num_gemms"] = num_gemms + meta["fp8_group"] = FP8GlobalStateManager.get_fp8_group() # Set FP8_MAX per tensor according to recipe - if hasattr(self.fp8_meta["recipe"], "fp8_format"): - self.fp8_meta["fp8_max_fwd"] = self.fp8_meta["recipe"].fp8_format.value.max_fwd - self.fp8_meta["fp8_max_bwd"] = self.fp8_meta["recipe"].fp8_format.value.max_bwd + if hasattr(meta["recipe"], "fp8_format"): + meta["fp8_max_fwd"] = meta["recipe"].fp8_format.value.max_fwd + meta["fp8_max_bwd"] = meta["recipe"].fp8_format.value.max_bwd # Allocate scales and amaxes - self.init_fp8_meta_tensors(self.fp8_meta["recipe"]) - self.fp8_initialized = True + self.init_fp8_meta_tensors(meta["recipe"]) + self.fast_setattr("fp8_initialized", True) - self.fp8_meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() + meta["recipe"] = FP8GlobalStateManager.get_fp8_recipe() - _current_recipe = self.fp8_meta["recipe"] + _current_recipe = meta["recipe"] if _original_recipe is not None and not ( issubclass(_current_recipe.__class__, _original_recipe.__class__) or issubclass(_original_recipe.__class__, _current_recipe.__class__) @@ -1028,22 +1028,18 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # Clear cached workspaces as they were created with the old recipe/quantizer type self._fp8_workspaces.clear() - @contextmanager def prepare_forward( self, inp: torch.Tensor, num_gemms: int = 1, allow_non_contiguous: bool = False, allow_different_data_and_param_types: bool = False, - ) -> Generator[torch.Tensor, None, None]: - """Checks and prep for FWD. - The context manager is needed because there isn't a way for a module to know - if it's the last FP8 module in the forward autocast. It is useful - to setup the forward aggregated amax reduction for every module - just in case. The autocast exit will pick up the most recent one. - """ - self.allow_different_data_and_param_types = allow_different_data_and_param_types - self.forwarded_at_least_once = True + ) -> torch.Tensor: + """Checks and prepares for FWD execution.""" + self.fast_setattr( + "allow_different_data_and_param_types", allow_different_data_and_param_types + ) + self.fast_setattr("forwarded_at_least_once", True) # Activation recomputation is used and this is the second forward phase. if self.fp8 and in_fp8_activation_recompute_phase(): @@ -1074,13 +1070,37 @@ def prepare_forward( if self.training and is_fp8_activation_recompute_enabled(): FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute(self.fp8_meta) - with get_nvtx_range_context(self.__class__.__name__ + " forward"): - if not allow_non_contiguous and not inp.is_contiguous(): - inp = inp.contiguous() - yield inp + nvtx_range_push(self.__class__.__name__ + " forward") + if not allow_non_contiguous and not inp.is_contiguous(): + inp = inp.contiguous() + return inp + 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() if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase(): FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta) + nvtx_range_pop() + + @contextmanager + def prepare_forward_ctx( + self, + inp: torch.Tensor, + num_gemms: int = 1, + allow_non_contiguous: bool = False, + allow_different_data_and_param_types: bool = False, + ) -> Generator[torch.Tensor, None, None]: + """Checks and prepares for FWD execution.""" + inp = self.prepare_forward( + inp, num_gemms, allow_non_contiguous, allow_different_data_and_param_types + ) + try: + yield inp + finally: + self.end_forward() def set_nccl_overlap_warning_if_tp(self) -> None: """When using TP, the NCCL communication needs to be scheduled @@ -1315,9 +1335,9 @@ def clear(self): # Update the parameter based on its type if not is_dtensor: - setattr(self, name, param) + self.module_setattr(name, param) else: - setattr(self, name, dtensor_param) + self.module_setattr(name, dtensor_param) @abstractmethod def forward(self): @@ -1516,7 +1536,6 @@ def is_debug_iter(self) -> bool: debug = TEDebugState.debug_enabled if not debug: return False - self._validate_name() # If layer is run first time in new iteration, # we need to check if the debug should be enabled for this layer - @@ -1530,14 +1549,14 @@ def is_debug_iter(self) -> bool: debug = False else: debug = TEDebugState.get_iteration() >= self.next_iter_when_debug_should_be_run - self.debug_last_iteration = TEDebugState.get_iteration() - self.debug_enabled_in_this_iteration = debug + self.fast_setattr("debug_last_iteration", TEDebugState.get_iteration()) + self.fast_setattr("debug_enabled_in_this_iteration", debug) else: # If this is the same iteration as previous invocation of the module, # we use the debug value from the first invocation in the iteration. debug = self.debug_enabled_in_this_iteration - self.debug_last_iteration = TEDebugState.get_iteration() + self.fast_setattr("debug_last_iteration", TEDebugState.get_iteration()) if self.wgrad_store is not None: if debug and self.wgrad_store.delay_wgrad_compute(): @@ -1553,7 +1572,9 @@ def no_debug_features_active(self, quantizers): # Sometimes features inform that they will not be enabled for particular layer # for multiple next iterations. - self.next_iter_when_debug_should_be_run = next_iter_when_debug_should_be_run(quantizers) + self.fast_setattr( + "next_iter_when_debug_should_be_run", next_iter_when_debug_should_be_run(quantizers) + ) if not run_current: return True @@ -1565,22 +1586,13 @@ def no_debug_features_active(self, quantizers): def _validate_name(self): """ Validate name passed to the module. - This is invoked in the forward() method as module names are assigned after Model is initialized in Megatron-LM. - If no name is assigned, it creates a default name with layer count as the variable. + It creates a default name with layer count as the variable + which may be changed by the user of the module. """ if self.name is not None: return - assert TEDebugState.debug_enabled - import nvdlfw_inspect.api as debug_api - - if self.name is None: - debug_api.log_message( - "Names are not provided to debug modules. ", - "Creating and using generic names. Pass names to debug modules for better" - " insight. ", - level=logging.WARNING, - ) - self.name = f"Layer_{TEDebugState.get_layer_count()}" + + self.name = f"Layer_{TEDebugState.get_layer_count()}" def _check_weight_tensor_recipe_correspondence(self) -> None: """ diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index e6e69b3e4a..c9ceb714e3 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -614,7 +614,7 @@ def __init__( save_original_input: bool = False, name: Optional[str] = None, ) -> None: - super().__init__() + super().__init__(name) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.num_gemms = num_gemms @@ -633,7 +633,6 @@ def __init__( ), "GroupedLinear doesn't support Userbuffer overlap." self.get_rng_state_tracker = get_rng_state_tracker self.rng_tracker_name = rng_tracker_name - self.name = name self.wgrad_store = WeightGradStore(delay_wgrad_compute) @@ -789,7 +788,8 @@ def forward( is_grad_enabled = torch.is_grad_enabled() - with self.prepare_forward(inp, num_gemms=self.num_gemms) as inp: + inp = self.prepare_forward(inp, num_gemms=self.num_gemms) + try: weight_tensors = self._get_weight_tensors() bias_tensors = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] @@ -844,6 +844,9 @@ def forward( ) out = linear_fn(*autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors) + finally: + self.end_forward() + if self.return_bias: return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors] return out diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index ca30ef9567..702916696b 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1158,9 +1158,9 @@ def __init__( ub_name: Optional[str] = None, delay_wgrad_compute: bool = False, symmetric_ar_type: Optional[str] = None, - name: str = None, + name: Optional[str] = None, ) -> None: - super().__init__() + super().__init__(name) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.in_features = in_features @@ -1179,7 +1179,6 @@ def __init__( self.symmetric_ar_type = symmetric_ar_type self.wgrad_store = WeightGradStore(delay_wgrad_compute, ub_bulk_wgrad) - self.name = name if tp_group is None: self.tp_size = tp_size @@ -1508,10 +1507,11 @@ def forward( ).is_fp8_ubuf(): fp8_grad = True - with self.prepare_forward( + inp = self.prepare_forward( inp, allow_non_contiguous=False # removed .contiguous from inside the layer - ) as inp: + ) + try: # Get concatenated weight and bias tensors weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() @@ -1590,6 +1590,9 @@ def forward( non_tensor_args, ) + finally: + self.end_forward() + if self.return_layernorm_output: out, ln_out = out diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 35e4522138..bec6744518 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -1787,7 +1787,7 @@ def __init__( zero_centered_gamma: bool = False, device: Union[torch.device, str] = "cuda", ub_overlap_ag: bool = False, - name: str = None, + name: Optional[str] = None, ub_overlap_rs: bool = False, ub_overlap_rs_dgrad: bool = False, ub_bulk_dgrad: bool = False, @@ -1796,7 +1796,7 @@ def __init__( symmetric_ar_type: Optional[str] = None, checkpoint: bool = False, ) -> None: - super().__init__() + super().__init__(name) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.fuse_wgrad_accumulation = fuse_wgrad_accumulation @@ -1827,7 +1827,6 @@ def __init__( for use_fp8 in [False, True] ) ) - self.name = name self.wgrad_store = WeightGradStore(delay_wgrad_compute, ub_bulk_wgrad) @@ -2047,8 +2046,9 @@ def forward( if get_ub("fc2_fprop", FP8GlobalStateManager.is_fp8_enabled()).is_fp8_ubuf(): fp8_output = True - with self.prepare_forward(inp, num_gemms=2) as inp: + inp = self.prepare_forward(inp, num_gemms=2) + try: quantizers = ( self._get_quantizers(fp8_output, is_grad_enabled) if not debug @@ -2087,7 +2087,7 @@ def forward( # Disable bias_gelu_nvfusion for determinism checkpointing in non-reentrant mode if self.bias_gelu_nvfusion and not use_reentrant_activation_recompute(): - self.bias_gelu_nvfusion = False + self.fast_setattr("bias_gelu_nvfusion", False) if is_grad_enabled: fwd_fn = _LayerNormMLP.apply @@ -2157,6 +2157,9 @@ def forward( non_tensor_args, ) + finally: + self.end_forward() + if self.return_layernorm_output: out, ln_out = out diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 38104604d8..23ad8cacb0 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -428,8 +428,8 @@ def forward( # weights if weights are externally touched outside this module ctx.weight_object = weight - if cpu_offloading: mark_not_offload(weight, weightmat, bias) + # TODO(ksivamani): Check memory usage tensors_to_save, tensor_objects = prepare_for_saving( saved_inputmat, @@ -1098,7 +1098,7 @@ def __init__( save_original_input: bool = False, name: Optional[str] = None, ) -> None: - super().__init__() + super().__init__(name) params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.in_features = in_features @@ -1111,7 +1111,6 @@ def __init__( self.rng_tracker_name = rng_tracker_name self.symmetric_ar_type = symmetric_ar_type self.save_original_input = save_original_input - self.name = name self.wgrad_store = WeightGradStore(delay_wgrad_compute, ub_bulk_wgrad) @@ -1395,11 +1394,8 @@ def forward( ).is_fp8_ubuf(): fp8_grad = True - with self.prepare_forward( - inp, - allow_non_contiguous=isinstance(inp, QuantizedTensor), - ) as inp: - + inp = self.prepare_forward(inp, allow_non_contiguous=isinstance(inp, QuantizedTensor)) + try: weight_tensor, bias_tensor = self._get_weight_and_bias_tensors() quantizers = ( @@ -1470,6 +1466,8 @@ def forward( bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None, non_tensor_args, ) + finally: + self.end_forward() if self.gemm_bias_unfused_add: out = out + cast_if_needed(bias_tensor, self.activation_dtype) diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 9b9ccc5185..7c3125a165 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -12,7 +12,6 @@ from transformer_engine.pytorch.torch_version import torch_version from transformer_engine.pytorch.module import LayerNormMLP, LayerNorm, RMSNorm -from transformer_engine.debug.pytorch.debug_state import TEDebugState from transformer_engine.pytorch.attention.multi_head_attention import MultiheadAttention from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.jit import ( @@ -398,6 +397,7 @@ def __init__( self.softmax_type = softmax_type self.name = name + TransformerEngineBaseModule._validate_name(self) attention_args = ( hidden_size, @@ -446,7 +446,7 @@ def __init__( qk_norm_type=qk_norm_type, qk_norm_eps=qk_norm_eps, qk_norm_before_rope=qk_norm_before_rope, - name=name + ".self_attention" if name is not None else None, + name=self.name + ".self_attention" if self.name is not None else None, ) if layer_type == "decoder": @@ -463,7 +463,7 @@ def __init__( qk_norm_type=qk_norm_type, qk_norm_eps=qk_norm_eps, qk_norm_before_rope=qk_norm_before_rope, - name=name + ".inter_attention" if name is not None else None, + name=self.name + ".inter_attention" if self.name is not None else None, ) # LayerNorm -> activation(Linear + Bias) -> Linear @@ -499,7 +499,7 @@ def __init__( activation_params=activation_params, normalization=normalization, device=device, - name=name + ".layernorm_mlp" if name is not None else None, + name=self.name + ".layernorm_mlp" if self.name is not None else None, ) self.hidden_dropout = hidden_dropout @@ -768,9 +768,6 @@ def forward( enc_dec_attn_mask[i].dtype == torch.bool for i in range(len(enc_dec_attn_mask)) ), "Encoder-decoder attention mask must be boolean tensor(s)" - if TEDebugState.debug_enabled: - TransformerEngineBaseModule._validate_name(self) - # For AMP if torch.is_autocast_enabled(): hidden_states = cast_if_needed(hidden_states, torch_get_autocast_gpu_dtype()) From 8bf37f0e09f0b111f31750a629bc4fd78e0c12d9 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 21 Jan 2026 17:06:10 -0800 Subject: [PATCH 175/521] [JAX] Fix cb.CUDAOptions usage for Triton 3.6.0 (#2610) * Fix cb.CUDAOptions usage for Triton 3.6.0 Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update utils.py Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update utils.py Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update utils.py Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --------- 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> --- transformer_engine/jax/triton_extensions/utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 6ea4092cbc..2627a08929 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -36,6 +36,8 @@ from typing import Any, Callable, Mapping import zlib +from packaging import version + from jax import core import jax import jax.numpy as jnp @@ -274,13 +276,16 @@ def compile_triton( return _TRITON_KERNEL_CACHE[cache_key] # Compile kernel + cuda_option_kwargs = {} + if version.parse(_TRITON_VERSION) < version.parse("3.6.0"): + cuda_option_kwargs["cluster_dims"] = (1, 1, 1) options = cb.CUDAOptions( num_warps=num_warps, num_stages=num_stages, num_ctas=num_ctas, - cluster_dims=(1, 1, 1), debug=False, enable_fp_fusion=enable_fp_fusion, + **cuda_option_kwargs, ) # Mark constants as constexpr in signature @@ -303,8 +308,6 @@ def compile_triton( # Create kernel object for JAX # From jax/jaxlib/gpu/triton_kernels.cc: - from packaging import version - if version.parse(jax.__version__) >= version.parse("0.8.2"): kernel = gpu_triton.TritonKernel( compiled.name, # arg0: kernel_name (str) From 3d46bf61e3bb336f15ef063b5d72fc3454eb53c2 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 22 Jan 2026 09:45:47 -0800 Subject: [PATCH 176/521] Permutation to always return group_size/tokens_per_expert (#2613) return tokens_per_experts always Signed-off-by: tdophung --- transformer_engine/jax/permutation.py | 38 +++++++++++++++++---------- 1 file changed, 24 insertions(+), 14 deletions(-) diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 405d5f7661..438511fa55 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -52,7 +52,7 @@ def token_dispatch( Optional[jnp.ndarray], jnp.ndarray, Optional[jnp.ndarray], - Optional[jnp.ndarray], + jnp.ndarray, ]: """ Dispatch tokens to experts based on routing map. @@ -101,9 +101,11 @@ def token_dispatch( pad_offsets : Optional[jnp.ndarray] Per-expert cumulative padding offsets of shape [num_experts] when using padding, None otherwise. Pass this to token_combine when unpadding is needed. - target_tokens_per_expert : Optional[jnp.ndarray] - Aligned token counts per expert of shape [num_experts] when using padding, - None otherwise. + tokens_per_expert : jnp.ndarray + Token counts per expert of shape [num_experts]: + - Without padding: actual token counts (sum of routing_map columns) + - With padding: aligned token counts (ceil(actual / align_size) * align_size) + This gives the effective number of tokens per expert in the output buffer. Note ---- @@ -151,10 +153,10 @@ def _token_dispatch( Optional[jnp.ndarray], jnp.ndarray, Optional[jnp.ndarray], - Optional[jnp.ndarray], + jnp.ndarray, ]: """Internal token_dispatch with custom VJP.""" - (output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert), _ = ( + (output, permuted_probs, row_id_map, pad_offsets, tokens_per_expert), _ = ( _token_dispatch_fwd_rule( inp, routing_map, @@ -165,7 +167,7 @@ def _token_dispatch( use_padding, ) ) - return output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert + return output, permuted_probs, row_id_map, pad_offsets, tokens_per_expert def _token_dispatch_fwd_rule( @@ -182,7 +184,7 @@ def _token_dispatch_fwd_rule( Optional[jnp.ndarray], jnp.ndarray, Optional[jnp.ndarray], - Optional[jnp.ndarray], + jnp.ndarray, ], Tuple[jnp.ndarray, Optional[jnp.ndarray], int, int, int, bool], ]: @@ -212,11 +214,11 @@ def _token_dispatch_fwd_rule( with_probs = probs is not None - if use_padding: - # Compute tokens_per_expert internally from routing_map - # This can be a traced value since output shape uses worst_case_out_tokens - tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) + # Compute tokens_per_expert from routing_map (actual counts) + # This is well-optimized by XLA as a simple column-wise reduction + tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) + if use_padding: # Calculate aligned token counts per expert target_tokens_per_expert = (jnp.ceil(tokens_per_expert / align_size) * align_size).astype( jnp.int32 @@ -242,10 +244,12 @@ def _token_dispatch_fwd_rule( hidden_size, align_size=align_size, ) + + # Return aligned counts when using padding + out_tokens_per_expert = target_tokens_per_expert else: # No padding pad_offsets = None - target_tokens_per_expert = None output, permuted_probs = permute_with_mask_map( inp, @@ -257,14 +261,20 @@ def _token_dispatch_fwd_rule( hidden_size, ) + # Return actual counts when not using padding + out_tokens_per_expert = tokens_per_expert + # Return (primals, residuals) + # out_tokens_per_expert is: + # - target_tokens_per_expert (aligned) when using padding + # - tokens_per_expert (actual) when not using padding residuals = (row_id_map, pad_offsets, num_tokens, num_experts, hidden_size, with_probs) return ( output, permuted_probs, row_id_map, pad_offsets, - target_tokens_per_expert, + out_tokens_per_expert, ), residuals From 0f0e229b57a2a2357a33b89d274b2bb8b4f955a2 Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Thu, 22 Jan 2026 12:00:04 -0800 Subject: [PATCH 177/521] [PyT] Update THD sink attention logic for cudnn >=9.18.0 (#2568) * Update THD sink attention logic for newer cudnn versions THD Sink attention is supported in 9.18.0 Signed-off-by: Chen Cui * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update thd sink attention logic for cp>1 Signed-off-by: Chen Cui * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add unit test for thd + sink attention Signed-off-by: Chen Cui * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address comments Signed-off-by: Chen Cui * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * do not skip thd cp sink attention test Signed-off-by: Chen Cui * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * disable deterministic mode for sink attention Signed-off-by: Chen Cui --------- Signed-off-by: Chen Cui 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> --- tests/pytorch/attention/test_attention.py | 9 ++++++ .../attention/test_attention_with_cp.py | 9 ++++-- .../dot_product_attention/context_parallel.py | 18 ++++++----- .../attention/dot_product_attention/utils.py | 31 ++++++++++--------- 4 files changed, 42 insertions(+), 25 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6fe0ffdaee..65ca74c484 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -429,6 +429,15 @@ def test_dpa_softmax(dtype, model_configs, model): ) +@pytest.mark.skipif(get_cudnn_version() < (9, 18, 0), reason="cuDNN 9.18.0+ is required.") +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@pytest.mark.parametrize("model_configs", [model_configs_softmax]) +@pytest.mark.parametrize("model", model_configs_softmax.keys()) +def test_dpa_softmax_thd(dtype, model_configs, model): + """Test DotProductAttention module with different softmax types""" + test_dot_product_attention(dtype, model_configs, model, True, True, "thd_thd_thd", False, False) + + model_configs_mla = { # test: ModelConfig(b, sq, hq, dqk) "mla_1_0": ModelConfig(8, 128, 16, 64, head_dim_v=128), diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 9480b8de70..06ed6e5723 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -283,9 +283,14 @@ def test_cp_with_fused_attention( pytest.skip( "CP implementation only supports cp_comm_type=a2a for non-vanilla softmax types!" ) - if config.softmax_type != "vanilla" and qkv_format == "thd": + if ( + get_cudnn_version() < (9, 18, 0) + and config.softmax_type != "vanilla" + and qkv_format == "thd" + ): pytest.skip( - "CP implementation does not support qkv_format=thd for non-vanilla softmax types!" + "Unless cudnn version >= 9.18.0, CP implementation does not support qkv_format=thd for" + " non-vanilla softmax types!" ) dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16} 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 75b360e485..a5931188dc 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -4026,28 +4026,30 @@ def attn_forward_func_with_cp( assert not sliding_window_attn or cp_comm_type in [ "a2a", "all_gather", - ], "Context parallelism does not support sliding window attention with {cp_comm_type=}!" + ], 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", - ], "Context parallelism does not support MLA with {cp_comm_type=}!" + ], 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" - ), "Context parallelism does not support {softmax_type=} with FP8 attention!" + ), f"Context parallelism does not support {softmax_type=} with FP8 attention!" assert ( softmax_type == "vanilla" or use_fused_attention - ), "Context parallelism only supports {softmax_type=} with FusedAttention backend!" + ), f"Context parallelism only supports {softmax_type=} with FusedAttention backend!" assert ( softmax_type == "vanilla" or cp_comm_type == "a2a" - ), "Context parallelism only supports {softmax_type=} with cp_comm_type = 'a2a'!" - assert ( - softmax_type == "vanilla" or qkv_format != "thd" - ), "Context parallelism does not support {softmax_type=} with qkv_format = 'thd'!" + ), f"Context parallelism only supports {softmax_type=} with cp_comm_type = 'a2a'!" + if get_cudnn_version() < (9, 18, 0): + assert softmax_type == "vanilla" or qkv_format != "thd", ( + f"Before cuDNN 9.18.0, context parallelism does not support {softmax_type=} with" + " qkv_format = 'thd'!" + ) args = [ is_training, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index cb74a15e77..fcac740cc3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -716,22 +716,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_unfused_attention = False if qkv_format == "thd": - logger.debug( - "Disabling FusedAttention for softmax_type = %s and qkv_format = thd", softmax_type - ) - use_fused_attention = False - logger.debug( - "Disabling UnfusedDotProductAttention for softmax_type = %s and qkv_format = thd", - softmax_type, - ) - use_unfused_attention = False + if 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: - logger.debug( - "Disabling UnfusedDotProductAttention for context parallelism with softmax_type" - " = %s", - softmax_type, - ) - use_unfused_attention = False if cp_comm_type != "a2a": logger.debug( "Disabling FusedAttention for context parallelism with softmax_type = %s and" @@ -1049,6 +1041,15 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_flash_attention_2 = False if use_fused_attention and deterministic: + if softmax_type != "vanilla": + logger.debug( + "Disabling FusedAttention for determinism reasons with softmax_type = %s. " + "Sink attention (off-by-one and learnable softmax) requires " + "NVTE_ALLOW_NONDETERMINISTIC_ALGO=1", + softmax_type, + ) + use_fused_attention = False + fused_attention_backend = None if fused_attention_backend == FusedAttnBackend["FP8"] and is_training: logger.debug("Disabling FusedAttention for determinism reasons with FP8") use_fused_attention = False From c6a92a4dced73ffabdd41d77bf3bfa2eb67f6f1c Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Thu, 22 Jan 2026 12:07:20 -0800 Subject: [PATCH 178/521] Add support for SWA (left, right) with FusedAttention (#2477) * SWA (left, right) with FusedAttention changes cherry-picked from https://github.com/NVIDIA/TransformerEngine/pull/1369 Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test_kv_cache failures Signed-off-by: Sudhakar Singh * remove unnecessary comments Signed-off-by: Sudhakar Singh * fix some more filter issues, address feedback Signed-off-by: Sudhakar Singh * fix for local test case failures - `bottom_right_diagonal` should be calculated in `fused_attn_fwd` call as well Signed-off-by: Sudhakar Singh * make conditions more accurate Signed-off-by: Sudhakar Singh * add cp tests to test swa (left, right) Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove dead code and make conditions better Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * feedback form Charlene Signed-off-by: Sudhakar Singh * small er Signed-off-by: Sudhakar Singh * plumb `bottom_right_diagonal` through jax Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * plumb `bottom_right_diagonal` through jax Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add missing fields Signed-off-by: Sudhakar Singh * use proper mask type in CP 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> --- tests/pytorch/attention/test_attention.py | 7 +- .../attention/test_attention_with_cp.py | 15 ++- tests/pytorch/utils.py | 14 +-- .../common/fused_attn/fused_attn.cpp | 97 ++++++++++--------- .../fused_attn_f16_arbitrary_seqlen.cu | 75 ++++++++------ .../fused_attn_f16_arbitrary_seqlen.h | 19 ++-- .../common/fused_attn/fused_attn_fp8.cu | 2 + transformer_engine/common/fused_attn/utils.h | 12 ++- .../include/transformer_engine/fused_attn.h | 63 ++++++------ .../jax/cpp_extensions/attention.py | 22 ++++- transformer_engine/jax/csrc/extensions.h | 4 +- .../jax/csrc/extensions/attention.cpp | 62 ++++++------ .../dot_product_attention/backends.py | 21 +++- .../dot_product_attention.py | 43 ++++++-- .../attention/dot_product_attention/utils.py | 82 +++++++++------- .../pytorch/attention/multi_head_attention.py | 26 +++++ .../pytorch/cpp_extensions/fused_attn.py | 22 +++++ transformer_engine/pytorch/csrc/extensions.h | 15 +-- .../pytorch/csrc/extensions/attention.cpp | 51 +++++----- transformer_engine/pytorch/transformer.py | 55 ++++++++++- 20 files changed, 474 insertions(+), 233 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 65ca74c484..bd0ac41974 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -153,6 +153,7 @@ def test_dot_product_attention( if config.window_size == (-1, -1) and swa: config.window_size = [2, 2] + 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 qkv_format == "thd" and "padding" not in config.attn_mask_type: @@ -171,6 +172,7 @@ def test_dot_product_attention( deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends + if not fused_attn_supported: is_training = False available_backends, _, fused_attn_backends = get_available_attention_backends( @@ -701,9 +703,10 @@ def test_dpa_bias_shapes(dtype, model_configs, model): @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_swa]) @pytest.mark.parametrize("model", model_configs_swa.keys()) -def test_dpa_sliding_window(dtype, model_configs, model): +@pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "sbhd_sbhd_sbhd"]) +def test_dpa_sliding_window(dtype, model_configs, model, qkv_layout): """Test DotProductAttention module with sliding window attention""" - test_dot_product_attention(dtype, model_configs, model, False, True, None, True, False) + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, True, False) model_configs_alibi_slopes = { diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 06ed6e5723..836598087b 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -147,7 +147,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): 2, 4096, 12, 128, attn_mask_type="causal", attn_bias_type="post_scale_bias" ), # MHA "cp_1_3": ModelConfig(2, 4096, 12, 128, attn_bias_type="post_scale_bias"), # MHA - "cp_1_4": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", window_size=(512, 0)), # MHA + "cp_1_4": 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_2": ModelConfig( @@ -163,7 +163,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): 2, 4096, 12, 128, num_gqa_groups=2, attn_bias_type="post_scale_bias" ), # GQA "cp_2_4": ModelConfig( - 2, 4096, 12, 128, num_gqa_groups=2, attn_mask_type="causal", window_size=(512, 0) + 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 @@ -187,7 +187,16 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): qkv_formats = ["bshd", "sbhd", "thd"] cp_comm_types = ["p2p", "all_gather", "a2a", "a2a+p2p"] if test_essential: - configs = ["cp_1_0", "cp_1_1", "cp_1_4", "cp_2_0", "cp_2_2", "cp_3_2", "cp_4_2"] + configs = [ + "cp_1_0", + "cp_1_1", + "cp_1_4", + "cp_2_0", + "cp_2_2", + "cp_2_4", + "cp_3_2", + "cp_4_2", + ] model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs} dtypes = ["bf16", "fp8"] qkv_formats = ["sbhd", "thd"] diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index ca5fbc997a..b6a84a8e2b 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -353,11 +353,11 @@ def test(): backends = {0: "F16_max512_seqlen", 1: "F16_arbitrary_seqlen", 2: "FP8"} if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() - with logging_context(highest_level=AttentionLogging._log_level): - for i in range(3): - os.environ["NVTE_FUSED_ATTN_BACKEND"] = str(i) - _attention_backends["backend_selection_requires_update"] = True - available_backends, flash_attention_backend, fused_attention_backend = test() - if fused_attention_backend == FusedAttnBackend[backends[i]]: - fused_attn_backends.append(fused_attention_backend) + + for i in range(3): + os.environ["NVTE_FUSED_ATTN_BACKEND"] = str(i) + _attention_backends["backend_selection_requires_update"] = True + available_backends, flash_attention_backend, fused_attention_backend = test() + if fused_attention_backend == FusedAttnBackend[backends[i]]: + fused_attn_backends.append(fused_attention_backend) return available_backends, flash_attention_backend, fused_attn_backends diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 415bfae063..4f8367aac7 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -406,9 +406,11 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (window_size_right == -1 || window_size_right == 0)) || // 9.2: SWA (left, 0) + top-left diagonal + {bshd, sbhd} (cudnn_runtime_version >= 90200 && - ((window_size_left == -1 && (window_size_right == -1 || window_size_right == 0)) || - ((window_size_left >= 0 || window_size_left == -1) && window_size_right == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + ((window_size_left == -1 && window_size_right == -1 && + attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK) || + ((window_size_left == -1 || window_size_left >= 0) && window_size_right == 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 && max_seqlen_q == max_seqlen_kv)) && max_seqlen_q <= max_seqlen_kv && dropout == 0.0 && @@ -418,12 +420,14 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // 9.6: SWA (left, 0) + top-left/bottom-right diagonal + {bshd, sbhd, thd} (cudnn_runtime_version >= 90600 && ((window_size_left == -1 && (window_size_right == -1 || window_size_right == 0)) || - ((window_size_left >= 0 || window_size_left == -1) && window_size_right == 0 && + ((window_size_left >= 0 || window_size_left == -1) && + (window_size_right >= 0 || window_size_right == -1) && ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK && // TODO(cyang): fix bug for BRCM + cross-attention on sm100 (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_PADDING_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && (sm_arch_ < 100 || (sm_arch_ >= 100 && ((max_seqlen_q == max_seqlen_kv && @@ -515,16 +519,14 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // NVTE fused attention FWD with packed QKV // DEPRECATED: This API is deprecated. // Please use nvte_fused_attn_fwd with separate Q, K, V tensors instead. -void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, - const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, - NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, - const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, - size_t max_seqlen, 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - NVTETensor workspace, cudaStream_t stream) { +void nvte_fused_attn_fwd_qkvpacked( + const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, + NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, + const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, + 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_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) { NVTE_API_CALL(nvte_flash_attn_fwd_qkvpacked); using namespace transformer_engine; @@ -598,10 +600,10 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, fused_attn_arbitrary_seqlen_fwd( b, h, h, max_seqlen, max_seqlen, d, d, t, t, 0, 0, 0, 0, 0, 0, is_training, return_max_logit, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, &Q_view, &K_view, &V_view, input_Bias, - input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, input_cu_seqlens, - input_cu_seqlens_padded, input_cu_seqlens_padded, nullptr, nullptr, input_rng_state, - wkspace, stream, handle); + window_size_left, window_size_right, bottom_right_diagonal, &Q_view, &K_view, &V_view, + input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, + input_cu_seqlens, input_cu_seqlens_padded, input_cu_seqlens_padded, nullptr, nullptr, + input_rng_state, wkspace, stream, handle); #else NVTE_ERROR( "cuDNN 8.9.0 is required for BF16/FP16 fused attention with arbitrary sequence length. " @@ -639,8 +641,8 @@ void nvte_fused_attn_bwd_qkvpacked( NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, size_t max_seqlen, 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 deterministic, bool cuda_graph, - NVTETensor workspace, cudaStream_t stream) { + 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_qkvpacked); using namespace transformer_engine; @@ -736,10 +738,11 @@ void nvte_fused_attn_bwd_qkvpacked( fused_attn_arbitrary_seqlen_bwd( b, h, h, max_seqlen, max_seqlen, d, d, t, t, attn_scale, dropout, qkv_layout, bias_type, - attn_mask_type, softmax_type, window_size_left, window_size_right, deterministic, &Q_view, - &K_view, &V_view, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, &dQ_view, - &dK_view, &dV_view, output_dBias, output_dSoftmaxOffset, input_cu_seqlens, input_cu_seqlens, - input_cu_seqlens_padded, input_cu_seqlens_padded, input_rng_state, wkspace, stream, handle); + attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, + deterministic, &Q_view, &K_view, &V_view, input_O, input_dO, input_Bias, + input_SoftmaxOffset, output_S, &dQ_view, &dK_view, &dV_view, output_dBias, + output_dSoftmaxOffset, input_cu_seqlens, input_cu_seqlens, input_cu_seqlens_padded, + input_cu_seqlens_padded, input_rng_state, wkspace, stream, handle); #else const char *err_msg = "cuDNN 8.9.0 is required for BF16/FP16 fused attention " @@ -790,7 +793,8 @@ void nvte_fused_attn_fwd_kvpacked( 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, NVTETensor workspace, cudaStream_t stream) { + int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, + cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd_kvpacked); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -902,10 +906,10 @@ void nvte_fused_attn_fwd_kvpacked( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, 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, input_Q, &K_view, &V_view, 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); + window_size_left, window_size_right, bottom_right_diagonal, input_Q, &K_view, &V_view, + 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.3 is required for BF16/FP16 fused attention with arbitrary sequence length. " @@ -945,8 +949,8 @@ void nvte_fused_attn_bwd_kvpacked( 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 deterministic, bool cuda_graph, NVTETensor workspace, - cudaStream_t stream) { + 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_kvpacked); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -1052,11 +1056,11 @@ void nvte_fused_attn_bwd_kvpacked( fused_attn_arbitrary_seqlen_bwd( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, t_q, t_kv, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, deterministic, - input_Q, &K_view, &V_view, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, - output_dQ, &dK_view, &dV_view, 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); + bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, deterministic, input_Q, &K_view, &V_view, input_O, input_dO, + input_Bias, input_SoftmaxOffset, output_S, output_dQ, &dK_view, &dV_view, 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.3 is required for BF16/FP16 fused attention " @@ -1106,8 +1110,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso bool return_max_logit, bool cuda_graph, 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, NVTETensor workspace, - cudaStream_t stream) { + int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_fwd); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -1195,10 +1199,10 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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, 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); + 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. " @@ -1228,8 +1232,9 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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 deterministic, - bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { + 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); @@ -1302,8 +1307,8 @@ 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, - 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, + 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 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 d3746fc042..53023361e4 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 @@ -55,10 +55,10 @@ void fused_attn_arbitrary_seqlen_fwd_impl( int64_t max_pages_per_seq_v, int64_t bias_b, int64_t bias_h, 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, 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, + 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) { using namespace transformer_engine; @@ -75,6 +75,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( if (is_bottom_right && s_q == s_kv && !is_padding) { is_causal = true; is_bottom_right = false; + bottom_right_diagonal = false; } bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (is_training && dropout_probability != 0.0f); @@ -129,6 +130,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, true, tensorType, cudnn_frontend::DataType_t::NOT_SET, @@ -254,9 +256,16 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_causal_mask_bottom_right(is_bottom_right) .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 >= 90200 && window_size_left != -1) { sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); } + if (cudnn_runtime_version >= 90600 && window_size_right != -1) { + sdpa_options.set_diagonal_band_right_bound(window_size_right); + } sdpa_options.set_alibi_mask(is_alibi); @@ -542,13 +551,14 @@ void fused_attn_arbitrary_seqlen_bwd_impl( int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t bias_b, int64_t bias_h, 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 deterministic, void *devPtrQ, - void *devPtrKTranspose, void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, - void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrdQ, void *devPtrdK, void *devPtrdV, - void *devPtrdO, void *devPtrdBias, void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, - void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, - void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, - void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + 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, void *devPtrSoftmaxOffset, + void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, void *devPtrdBias, + void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, void *devPtrDropoutOffset, + void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrSeqOffsetsQ, + void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, void *workspace, + size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -563,6 +573,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (is_bottom_right && s_q == s_kv && !is_padding) { is_causal = true; is_bottom_right = false; + bottom_right_diagonal = false; } bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (dropout_probability != 0.0f); @@ -621,6 +632,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, deterministic, tensorType, cudnn_frontend::DataType_t::NOT_SET, @@ -781,9 +793,17 @@ void fused_attn_arbitrary_seqlen_bwd_impl( sdpa_backward_options.set_max_total_seq_len_kv(s_kv); } + 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 >= 90200 && window_size_left != -1) { sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); } + if (cudnn_runtime_version >= 90600 && window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } if (cudnn_runtime_version >= 90000) { sdpa_backward_options.set_deterministic_algorithm(deterministic); @@ -1044,8 +1064,8 @@ void fused_attn_arbitrary_seqlen_fwd( 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, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, const Tensor *input_Bias, + 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, @@ -1180,11 +1200,11 @@ void fused_attn_arbitrary_seqlen_fwd( 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, is_training, return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, 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); + 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) { @@ -1206,13 +1226,14 @@ void fused_attn_arbitrary_seqlen_bwd( 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_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 deterministic, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_S, - Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, Tensor *output_dBias, - Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { + 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, + const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, + const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, + Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto QKV_type = input_Q->data.dtype; void *devPtrQ = input_Q->data.dptr; @@ -1273,8 +1294,8 @@ void fused_attn_arbitrary_seqlen_bwd( 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, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, - deterministic, devPtrQ, devPtrK, devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, - devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdO, devPtrdBias, + 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); 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 c34eae4e6e..4dd7f3d1da 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 @@ -25,8 +25,8 @@ void fused_attn_arbitrary_seqlen_fwd( 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, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, const Tensor *input_Bias, + 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, @@ -37,13 +37,14 @@ void fused_attn_arbitrary_seqlen_bwd( 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_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 deterministic, const Tensor *input_Q, - const Tensor *input_K, const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_S, - Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, Tensor *output_dBias, - Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, - Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + 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, + const Tensor *input_O, const Tensor *input_dO, const Tensor *input_Bias, + const Tensor *input_SoftmaxOffset, Tensor *output_S, Tensor *output_dQ, Tensor *output_dK, + Tensor *output_dV, Tensor *output_dBias, Tensor *output_dSoftmaxOffset, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, + 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 diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 3630041ccf..f886ec77f4 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1707,6 +1707,7 @@ void fused_attn_fp8_fwd_impl_v1( 0, 0, true, + true, qkv_tensor_type, o_tensor_type, cudnn_frontend::DataType_t::NOT_SET, @@ -2035,6 +2036,7 @@ void fused_attn_fp8_bwd_impl_v1( NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX, 0, 0, + true, false, qkv_tensor_type, o_tensor_type, diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 7d23bb5c55..fdfc4abe82 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -110,6 +110,7 @@ struct FADescriptor_v1 { NVTE_Softmax_Type softmax_type; std::int64_t window_size_left; std::int64_t window_size_right; + bool bottom_right_diagonal; bool deterministic; cudnn_frontend::DataType_t qkv_tensor_type; cudnn_frontend::DataType_t o_tensor_type; @@ -121,15 +122,16 @@ struct FADescriptor_v1 { 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, attnScale, isTraining, dropoutProbability, layout, mask_type, softmax_type, - window_size_left, window_size_right, deterministic, bias_type, qkv_tensor_type, - o_tensor_type, do_tensor_type, dqkv_tensor_type, generate_max_sum_exp) < + window_size_left, window_size_right, bottom_right_diagonal, deterministic, + bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type, dqkv_tensor_type, + generate_max_sum_exp) < 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.attnScale, rhs.isTraining, rhs.dropoutProbability, rhs.layout, rhs.mask_type, rhs.softmax_type, - rhs.window_size_left, rhs.window_size_right, rhs.deterministic, rhs.bias_type, - rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, - rhs.dqkv_tensor_type, rhs.generate_max_sum_exp); + 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.generate_max_sum_exp); } }; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 0fabb81aef..cddd3d7506 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -270,22 +270,21 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ [[deprecated( "nvte_fused_attn_fwd_qkvpacked() is deprecated. Please use nvte_fused_attn_fwd() with separate " "Q, K, V tensors instead.")]] -void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, - const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, - NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, - const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, - size_t max_seqlen, 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - NVTETensor workspace, cudaStream_t stream); +void nvte_fused_attn_fwd_qkvpacked( + const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, + NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, + const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, + 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_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); /*! \brief Compute the backward of the dot product attention with packed QKV input. * @@ -333,6 +332,7 @@ void nvte_fused_attn_fwd_qkvpacked(const NVTETensor QKV, const NVTETensor Bias, * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] deterministic Whether to execute with deterministic behaviours. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. @@ -347,8 +347,8 @@ void nvte_fused_attn_bwd_qkvpacked( NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, size_t max_seqlen, 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 deterministic, bool cuda_graph, - NVTETensor workspace, cudaStream_t stream); + 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 Compute dot product attention with packed KV input. * @@ -410,6 +410,7 @@ void nvte_fused_attn_bwd_qkvpacked( * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ @@ -425,7 +426,8 @@ void nvte_fused_attn_fwd_kvpacked( 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, - int64_t window_size_right, NVTETensor workspace, cudaStream_t stream); + int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, + cudaStream_t stream); /*! \brief Compute the backward of the dot product attention with packed KV input. * @@ -479,6 +481,7 @@ void nvte_fused_attn_fwd_kvpacked( * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] deterministic Whether to execute with deterministic behaviours. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. @@ -495,8 +498,8 @@ void nvte_fused_attn_bwd_kvpacked( 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 deterministic, bool cuda_graph, NVTETensor workspace, - cudaStream_t stream); + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, bool cuda_graph, + NVTETensor workspace, cudaStream_t stream); /*! \brief Compute dot product attention with separate Q, K and V. * @@ -560,19 +563,23 @@ void nvte_fused_attn_bwd_kvpacked( * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] workspace Workspace tensor. * \param[in] stream CUDA stream used for this operation. */ -void nvte_fused_attn_fwd( - const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, - const NVTETensor SoftmaxOffset, NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, - const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, - const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, - const NVTETensor page_table_k, 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_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, NVTETensor workspace, cudaStream_t stream); +void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, + const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, + NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, + const NVTETensor cu_seqlens_q, const NVTETensor cu_seqlens_kv, + const NVTETensor cu_seqlens_q_padded, + const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, + 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_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); /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * @@ -629,6 +636,7 @@ void nvte_fused_attn_fwd( * \param[in] softmax_type Attention softmax type. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). + * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. * \param[in] deterministic Whether to execute with deterministic behaviours. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] workspace Workspace tensor. @@ -644,8 +652,9 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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 deterministic, - bool cuda_graph, NVTETensor workspace, cudaStream_t stream); + 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. * diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index ee10115aa1..e5d75e1501 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -70,6 +70,7 @@ "is_training", "max_segments_per_seq", "window_size", + "bottom_right_diagonal", "context_parallel_load_balanced", "cp_axis", "cp_striped_window_size", @@ -91,6 +92,7 @@ class _FusedAttnConfig: is_training: bool max_segments_per_seq: int window_size: Tuple[int, int] + bottom_right_diagonal: bool context_parallel_load_balanced: bool cp_axis: str cp_striped_window_size: Tuple[int, int] # Only for CP + Ring P2P + THD + SWA @@ -371,6 +373,11 @@ def abstract( *bias_batch_shape, bias_heads, _, _ = bias_aval.shape bias_batch = reduce(operator.mul, bias_batch_shape) + bottom_right_diagonal = config.attn_mask_type in [ + AttnMaskType.CAUSAL_BOTTOM_RIGHT_MASK, + AttnMaskType.PADDING_CAUSAL_BOTTOM_RIGHT_MASK, + ] + # do a dummy kernel call here to get workspace buffer shapes/dtypes that XLA needs to # prepare for the active fused-attn backend input_batch = reduce(operator.mul, batch_shape) @@ -395,6 +402,7 @@ def abstract( config.max_segments_per_seq, config.window_size[0], config.window_size[1], + bottom_right_diagonal, ) wkspace_aval = q_aval.update( shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) @@ -503,6 +511,7 @@ def lowering( deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), window_size_left=window_size_left, window_size_right=window_size_right, + bottom_right_diagonal=config.bottom_right_diagonal, softmax_type=int(config.softmax_type.value), ) @@ -813,6 +822,7 @@ def abstract( config.max_segments_per_seq, config.window_size[0], config.window_size[1], + config.bottom_right_diagonal, ) dq_aval = q_aval.update(shape=q_aval.shape, dtype=q_dtype) @@ -948,6 +958,7 @@ def lowering( deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), window_size_left=window_size_left, window_size_right=window_size_right, + bottom_right_diagonal=config.bottom_right_diagonal, softmax_type=int(config.softmax_type.value), ) @@ -1357,9 +1368,10 @@ def get_adjusted_max_segments_per_seq(self, max_seqlen, cp_size): def get_step_config(self) -> _FusedAttnConfig: """Returns a _FusedAttnConfig for single CP step call to fused attention.""" + adjusted_mask = self.get_adjusted_mask() return _FusedAttnConfig( attn_bias_type=self.config.attn_bias_type, - attn_mask_type=self.get_adjusted_mask(), + attn_mask_type=adjusted_mask, softmax_type=self.config.softmax_type, qkv_layout=self.config.qkv_layout, scaling_factor=self.config.scaling_factor, @@ -1367,6 +1379,7 @@ def get_step_config(self) -> _FusedAttnConfig: is_training=self.config.is_training, max_segments_per_seq=self.config.max_segments_per_seq, window_size=self.config.window_size, + bottom_right_diagonal=adjusted_mask.is_bottom_right(), context_parallel_load_balanced=self.config.context_parallel_load_balanced, cp_axis=self.config.cp_axis, cp_striped_window_size=None, @@ -1375,9 +1388,10 @@ def get_step_config(self) -> _FusedAttnConfig: def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: """Returns a _FusedAttnConfig for single CP step call (made via a striped AG primitive) to fused attention.""" + adjusted_mask = self.get_adjusted_mask() return _FusedAttnConfig( attn_bias_type=self.config.attn_bias_type, - attn_mask_type=self.get_adjusted_mask(), + attn_mask_type=adjusted_mask, softmax_type=self.config.softmax_type, qkv_layout=self.config.qkv_layout, scaling_factor=self.config.scaling_factor, @@ -1385,6 +1399,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: is_training=self.config.is_training, max_segments_per_seq=self.get_adjusted_max_segments_per_seq(max_seqlen, cp_size), window_size=self.config.window_size, + bottom_right_diagonal=adjusted_mask.is_bottom_right(), context_parallel_load_balanced=self.config.context_parallel_load_balanced, cp_axis=self.config.cp_axis, cp_striped_window_size=None, @@ -2430,6 +2445,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: is_training=self.config.is_training, max_segments_per_seq=self.config.max_segments_per_seq, window_size=self.config.window_size, + bottom_right_diagonal=attn_mask_type.is_bottom_right(), context_parallel_load_balanced=self.config.context_parallel_load_balanced, cp_axis=self.config.cp_axis, cp_striped_window_size=None, @@ -3418,6 +3434,7 @@ def fused_attn_fwd( is_training=is_training, max_segments_per_seq=max_segments_per_seq, window_size=(-1, -1) if window_size is None else window_size, + bottom_right_diagonal=attn_mask_type.is_bottom_right(), context_parallel_load_balanced=context_parallel_causal_load_balanced, cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, @@ -3590,6 +3607,7 @@ def fused_attn_bwd( is_training=is_training, max_segments_per_seq=max_segments_per_seq, window_size=(-1, -1) if window_size is None else window_size, + bottom_right_diagonal=attn_mask_type.is_bottom_right(), context_parallel_load_balanced=context_parallel_causal_load_balanced, cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 5f93392633..3fd086e257 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -121,7 +121,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, 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); + int64_t window_size_right, bool bottom_right_diagonal); pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, @@ -129,7 +129,7 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, 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); + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal); // GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmHandler); diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 4fe8e728a3..92e67ac191 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -144,7 +144,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, 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) { + 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 q_tensor = TensorWrapper(nullptr, q_shape, dtype); auto k_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim}; @@ -192,7 +192,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( 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, query_workspace_tensor.data(), nullptr); + window_size_left, window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), + nullptr); } nvte_tensor_pack_destroy(&aux_output_tensors); @@ -237,7 +238,7 @@ static void FusedAttnForwardImpl( size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, - int64_t window_size_left, int64_t window_size_right) { + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ @@ -328,7 +329,7 @@ static void FusedAttnForwardImpl( 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, workspace_tensor.data(), stream); + window_size_left, window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); nvte_tensor_pack_destroy(&aux_output_tensors); } @@ -346,6 +347,7 @@ static void FusedAttnForwardImpl( size_t max_segments_per_seq = get_attr_value(attrs, "max_segments_per_seq"); \ auto window_size_left = get_attr_value(attrs, "window_size_left"); \ auto window_size_right = get_attr_value(attrs, "window_size_right"); \ + bool bottom_right_diagonal = get_attr_value(attrs, "bottom_right_diagonal"); \ float scaling_factor = get_attr_value(attrs, "scaling_factor"); \ float dropout_probability = get_attr_value(attrs, "dropout_probability"); \ NVTE_Bias_Type bias_type = \ @@ -384,7 +386,7 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty input_batch, bias_batch, q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, softmax_type, qkv_layout, dtype, wkspace_dtype, - is_training, deterministic, window_size_left, window_size_right); + is_training, deterministic, window_size_left, window_size_right, bottom_right_diagonal); return ffi_with_cuda_error_check(); } @@ -415,7 +417,7 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, 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) { + 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 q_tensor = TensorWrapper(nullptr, q_shape, dtype); auto dq_tensor = TensorWrapper(nullptr, q_shape, dtype); @@ -467,17 +469,18 @@ 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, 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, 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); @@ -496,7 +499,7 @@ static void FusedAttnBackwardImpl( size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, - int64_t window_size_left, int64_t window_size_right) { + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ @@ -593,16 +596,17 @@ 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, 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, 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); } @@ -631,7 +635,7 @@ Error_Type FusedAttnBackwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_T q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, deterministic, window_size_left, - window_size_right); + window_size_right, bottom_right_diagonal); return ffi_with_cuda_error_check(); } diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index c726ed8849..ef7fa0dcc0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -261,6 +261,7 @@ def forward( attn_mask_type: str = "causal", attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, alibi_slopes: Optional[torch.Tensor] = None, @@ -346,6 +347,11 @@ def forward( attention_mask=attention_mask, window_size=window_size, attention_type=self.attention_type, + bottom_right_alignment=( + attn_mask_type not in ["causal", "padding_causal"] + if bottom_right_diagonal is None + else bottom_right_diagonal + ), ) ) @@ -449,7 +455,11 @@ def forward( actual_seqlens_q=actual_seqlens_q if "padding" in attn_mask_type else None, actual_seqlens_kv=actual_seqlens_kv if "padding" in attn_mask_type else None, alibi_slopes=alibi_slopes, - bottom_right_alignment=attn_mask_type not in ["causal", "padding_causal"], + bottom_right_alignment=( + attn_mask_type not in ["causal", "padding_causal"] + if bottom_right_diagonal is None + else bottom_right_diagonal + ), ) matmul_result = torch.baddbmm( matmul_result, @@ -1110,6 +1120,7 @@ def forward( attn_mask_type, softmax_type, window_size, + bottom_right_diagonal, rng_gen, fused_attention_backend, use_FAv2_bwd, @@ -1213,6 +1224,7 @@ def forward( attn_mask_type, softmax_type, window_size, + bottom_right_diagonal, rng_gen, softmax_offset, cuda_graph=is_graph_capturing(), @@ -1290,6 +1302,7 @@ def forward( attn_mask_type, softmax_type, window_size, + bottom_right_diagonal, rng_gen, softmax_offset, return_max_logit, @@ -1377,6 +1390,7 @@ def forward( ctx.attn_mask_type = attn_mask_type ctx.softmax_type = softmax_type ctx.window_size = window_size + ctx.bottom_right_diagonal = bottom_right_diagonal ctx.fused_attention_backend = ( fused_attention_backend if ctx.fp8 else FusedAttnBackend["F16_arbitrary_seqlen"] ) @@ -1527,6 +1541,7 @@ def backward(ctx, d_out, *_args): ctx.attn_mask_type, ctx.softmax_type, ctx.window_size, + ctx.bottom_right_diagonal, ctx.deterministic, is_graph_capturing(), ) @@ -1592,6 +1607,7 @@ def backward(ctx, d_out, *_args): ctx.attn_mask_type, ctx.softmax_type, ctx.window_size, + ctx.bottom_right_diagonal, ctx.deterministic, is_graph_capturing(), ) @@ -1631,6 +1647,7 @@ def backward(ctx, d_out, *_args): None, None, None, + None, d_softmax_offset, None, None, @@ -1728,6 +1745,7 @@ def forward( attn_mask_type: str = "causal", attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, fused_attention_backend: tex.NVTE_Fused_Attn_Backend = tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, @@ -1935,6 +1953,7 @@ def forward( attn_mask_type, self.softmax_type, window_size, + bottom_right_diagonal, None, # rng_gen fused_attention_backend, use_FAv2_bwd, 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 51ffbc2e48..5a554d86ec 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 @@ -228,6 +228,11 @@ class DotProductAttention(TransformerEngineBaseModule): map to ``window_size = (-1, 0)`` and Transformer Engine distinguishes them based on ``attn_mask_type``. Similar to :attr:`attn_mask_type`, ``window_size`` can be overridden by :attr:`window_size` in ``forward`` as well. + bottom_right_diagonal: Optional[bool], default = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the encoder. + If `None`, it will be set to `False` for `attn_mask_type` = + {'causal', 'padding_causal'} and `True` for other mask types. attention_type : str, default = "self" type of attention, either ``"self"`` and ``"cross"``. layer_number : int, default = None @@ -324,6 +329,7 @@ def __init__( qkv_format: str = "sbhd", attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, sequence_parallel: bool = False, tp_size: int = 1, get_rng_state_tracker: Optional[Callable] = None, @@ -350,6 +356,7 @@ def __init__( attn_mask_type = "padding_causal" self.attn_mask_type = attn_mask_type self.window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size) + self.bottom_right_diagonal = bottom_right_diagonal if tp_group is None: self.tp_size = tp_size if tp_size == 1: @@ -811,6 +818,7 @@ def forward( max_seqlen_kv: int = None, attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, checkpoint_core_attention: bool = False, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, @@ -963,6 +971,16 @@ def forward( causal masks are aligned to the bottom right corner. window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention. + bottom_right_diagonal: Optional[bool], default = None + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the encoder. + If `None`, it will be set to `False` for `attn_mask_type` = + {'causal', 'padding_causal'} and `True` for other mask types. + Note: This parameter will be automatically overridden based on the + `attn_mask_type` - it will be forced to `False` for 'causal' and + 'padding_causal' mask types, and forced to `True` for mask types + containing 'bottom_right' (e.g., 'causal_bottom_right', + 'padding_causal_bottom_right'), regardless of the explicitly passed value. checkpoint_core_attention : bool, default = False If true, forward activations for attention are recomputed during the backward pass in order to save memory that would @@ -1081,6 +1099,15 @@ def forward( if window_size is None: window_size = self.window_size window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size) + if bottom_right_diagonal is None: + bottom_right_diagonal = self.bottom_right_diagonal + if attn_mask_type in {"causal", "padding_causal"}: + bottom_right_diagonal = False + if bottom_right_diagonal is None or attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + }: + bottom_right_diagonal = True # checks for qkv_format if qkv_format is None: @@ -1144,6 +1171,8 @@ def forward( assert "padding" in attn_mask_type, "KV caching requires padding mask!" if attn_mask_type == "padding_causal": attn_mask_type = attn_mask_type + "_bottom_right" + # since attention mask is changed, set `bottom_right_diagonal` to True + bottom_right_diagonal = True if self.attention_type != "cross": self.fast_setattr("attention_type", "cross") @@ -1257,7 +1286,6 @@ def forward( if self.layer_number == 1: _alibi_cache["_alibi_slopes_require_update"] = True _alibi_cache["_alibi_bias_require_update"] = True - bottom_right_alignment = (attn_mask_type not in ["causal", "padding_causal"],) if core_attention_bias_type == "alibi": assert ( core_attention_bias is None @@ -1266,7 +1294,7 @@ def forward( _alibi_cache["_num_heads"] != query_layer.shape[-2] or _alibi_cache["_max_seqlen_q"] != max_seqlen_q or _alibi_cache["_max_seqlen_kv"] != max_seqlen_kv - or _alibi_cache["_bottom_right_alignment"] != bottom_right_alignment + or _alibi_cache["_bottom_right_alignment"] != bottom_right_diagonal or _alibi_cache["_alibi_slopes"] is None ): _alibi_cache["_alibi_slopes_require_update"] = True @@ -1323,6 +1351,7 @@ def forward( head_dim_v=head_dim_v, attn_mask_type=attn_mask_type, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, alibi_slopes_shape=alibi_slopes.shape if alibi_slopes is not None else None, core_attention_bias_type=core_attention_bias_type, core_attention_bias_shape=core_attention_bias_shape, @@ -1446,9 +1475,7 @@ def forward( if use_fused_attention: fu_core_attention_bias_type = core_attention_bias_type fu_core_attention_bias = core_attention_bias - if core_attention_bias_type == "alibi" and ( - alibi_slopes is not None or max_seqlen_q != max_seqlen_kv - ): + if core_attention_bias_type == "alibi" and (alibi_slopes is not None): fu_core_attention_bias_type = "post_scale_bias" _, fu_core_attention_bias = dpa_utils.get_alibi( _alibi_cache, @@ -1457,7 +1484,7 @@ def forward( max_seqlen_kv, alibi_slopes=alibi_slopes, bias_dtype=query_layer.dtype, - bottom_right_alignment=attn_mask_type not in ["causal", "padding_causal"], + bottom_right_alignment=bottom_right_diagonal, ) if checkpoint_core_attention: return self._checkpointed_attention_forward( @@ -1475,6 +1502,7 @@ def forward( attn_mask_type=attn_mask_type, attention_mask=attention_mask, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, fused_attention_backend=fused_attention_backend, core_attention_bias_type=fu_core_attention_bias_type, core_attention_bias=fu_core_attention_bias, @@ -1505,6 +1533,7 @@ def forward( attn_mask_type=attn_mask_type, attention_mask=attention_mask, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, fused_attention_backend=fused_attention_backend, core_attention_bias_type=fu_core_attention_bias_type, core_attention_bias=fu_core_attention_bias, @@ -1539,6 +1568,7 @@ def forward( attn_mask_type=attn_mask_type, attention_mask=attention_mask, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, alibi_slopes=alibi_slopes, @@ -1562,6 +1592,7 @@ def forward( attn_mask_type=attn_mask_type, attention_mask=attention_mask, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, alibi_slopes=alibi_slopes, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index fcac740cc3..56e6f093d1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -200,6 +200,9 @@ class AttentionParams: `causal_bottom_right`, `padding_causal_bottom_right`, `arbitrary`} window_size : Tuple[int, int], default = None Sliding window attention size. + bottom_right_diagonal: bool, default = `None` + Whether to align sliding window and ALiBi diagonal to the bottom right corner + of the softmax matrix. alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias @@ -249,6 +252,7 @@ class AttentionParams: head_dim_v: int = 64 attn_mask_type: str = "no_mask" window_size: Union[Tuple[int, int], None] = None + bottom_right_diagonal: bool = True alibi_slopes_shape: Union[torch.Size, List, None] = None core_attention_bias_type: str = "no_bias" core_attention_bias_shape: str = "1hss" @@ -325,6 +329,7 @@ def get_attention_backend( head_dim_v = attention_params.head_dim_v attn_mask_type = attention_params.attn_mask_type window_size = attention_params.window_size + bottom_right_diagonal = attention_params.bottom_right_diagonal alibi_slopes_shape = attention_params.alibi_slopes_shape core_attention_bias_type = attention_params.core_attention_bias_type core_attention_bias_shape = attention_params.core_attention_bias_shape @@ -859,39 +864,43 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # backend | window_size | diagonal alignment # --------------------------------------------------------------------------------- # FlashAttention | (-1, -1) or (>=0, >=0) | bottom right - # FusedAttention | (-1, 0) or (>=0, 0) | top left - # UnfusedDotProductAttention | (-1, -1) or (>=0, >=0) | both; + # FusedAttention | (-1, 0) or (>=0, >=0) | top left, bottom right + # UnfusedDotProductAttention | (-1, -1) or (>=0, >=0) | top left, bottom right # | | converts window_size to an 'arbitrary' mask if window_size is None: window_size = check_set_window_size(attn_mask_type, window_size) - else: - 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): - logger.debug( - "Disabling FusedAttention as it does not support sliding window attention" - " for FP8" - ) - use_fused_attention = False - elif window_size[1] != 0 or attention_dropout != 0.0: - logger.debug( - "Disabling FusedAttention as it only supports sliding window attention " - "with (left, 0) and no dropout" - ) - use_fused_attention = False - elif max_seqlen_q > max_seqlen_kv: - logger.debug( - "Disabling FusedAttention as it does not support sliding window attention " - "with s_q > s_kv for cross-attention" - ) - use_fused_attention = False - if use_flash_attention_2 and (window_size[0] != -1 or window_size[1] not in [-1, 0]): - if not FlashAttentionUtils.is_installed: - FlashAttentionUtils.version_required = PkgVersion("2.3") - elif not FlashAttentionUtils.v2_3_plus: - logger.debug( - "Disabling FlashAttention as sliding window attention requires flash-attn 2.3+" - ) - use_flash_attention_2 = False + 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): + logger.debug( + "Disabling FusedAttention as it does not support sliding window attention for FP8" + ) + use_fused_attention = False + elif attention_dropout != 0.0: + logger.debug( + "Disabling FusedAttention as it only supports sliding window attention " + "without dropout" + ) + use_fused_attention = False + elif max_seqlen_q > max_seqlen_kv: + logger.debug( + "Disabling FusedAttention as it does not support sliding window attention " + "with s_q > s_kv for cross-attention" + ) + use_fused_attention = False + if use_flash_attention_2 and (window_size[0] != -1 or window_size[1] not in [-1, 0]): + if not FlashAttentionUtils.is_installed: + FlashAttentionUtils.version_required = PkgVersion("2.3") + elif not FlashAttentionUtils.v2_3_plus: + logger.debug( + "Disabling FlashAttention as sliding window attention requires flash-attn 2.3+" + ) + use_flash_attention_2 = False + elif not bottom_right_diagonal and max_seqlen_q != max_seqlen_kv: + logger.debug( + "Disabling FlashAttention as it only supports sliding window with bottom right" + " diagonal alignment for cross-attention" + ) + use_flash_attention = False # Filter: Attention bias # backend | bias types | ALiBi diagonal alignment @@ -913,6 +922,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt elif not FlashAttentionUtils.v2_4_plus: logger.debug("Disabling FlashAttention as ALiBi requires flash-attn 2.4+") use_flash_attention_2 = False + elif not bottom_right_diagonal and max_seqlen_q != max_seqlen_kv: + logger.debug( + "Disabling FlashAttention as it only supports ALiBi with bottom right diagonal" + " alignment for cross-attention" + ) + use_flash_attention = False if ( core_attention_bias_type not in ["no_bias", "alibi"] @@ -930,13 +945,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if ( use_fused_attention and core_attention_bias_type == "alibi" - and (alibi_slopes_shape is not None or max_seqlen_q != max_seqlen_kv) + and (alibi_slopes_shape is not None) ): fu_core_attention_bias_type = "post_scale_bias" fu_core_attention_bias_requires_grad = False - if alibi_slopes_shape is None: - fu_core_attention_bias_shape = "1hss" - elif len(alibi_slopes_shape) == 1 and alibi_slopes_shape[0] == num_heads: + + if len(alibi_slopes_shape) == 1 and alibi_slopes_shape[0] == num_heads: fu_core_attention_bias_shape = "1hss" elif ( len(alibi_slopes_shape) == 2 diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index d813e7c8f1..01c4955d78 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -31,6 +31,7 @@ from transformer_engine.pytorch.attention.dot_product_attention import DotProductAttention from transformer_engine.pytorch.attention.inference import InferenceParams from transformer_engine.pytorch.attention.rope import apply_rotary_pos_emb +from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils from transformer_engine.pytorch.cpu_offload import start_offload, is_cpu_offload_enabled @@ -92,6 +93,11 @@ class MultiheadAttention(torch.nn.Module): map to ``window_size = (-1, 0)`` and Transformer Engine distinguishes them based on ``attn_mask_type``. Similar to :attr:`attn_mask_type`, ``window_size`` can be overridden by :attr:`window_size` in :meth:`forward` as well. + bottom_right_diagonal: Optional[bool], default = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the encoder. + If `None`, it will be set to `False` for `attn_mask_type` = + {`causal`, `padding_causal`} and `True` for other mask types. num_gqa_groups : int, default = None number of GQA groups in the transformer layer. Grouped Query Attention is described in @@ -247,6 +253,7 @@ def __init__( layer_number: Optional[int] = None, attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, tp_group: Optional[dist_group_type] = None, tp_size: int = 1, num_gqa_groups: Optional[int] = None, @@ -285,6 +292,7 @@ def __init__( self.qkv_format = qkv_format self.attn_mask_type = attn_mask_type self.window_size = window_size + self.bottom_right_diagonal = bottom_right_diagonal self.layer_number = 1 if layer_number is None else layer_number self.input_layernorm = input_layernorm self.attention_type = attention_type @@ -621,6 +629,7 @@ def forward( encoder_output: Optional[torch.Tensor] = None, attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[InferenceParams] = None, @@ -667,6 +676,11 @@ def forward( aligned to the bottom right corner. window_size: Optional[Tuple[int, int]], default = None sliding window size for local attention. + bottom_right_diagonal: Optional[bool], default = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the encoder. + If `None`, it will be set to `False` for `attn_mask_type` = + {`causal`, `padding_causal`} and `True` for other mask types. encoder_output : Optional[torch.Tensor], default = None Output of the encoder block to be fed into the decoder block if using ``layer_type="decoder"``. @@ -731,6 +745,17 @@ def forward( if window_size is None: window_size = self.window_size + window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size) + if bottom_right_diagonal is None: + bottom_right_diagonal = self.bottom_right_diagonal + if attn_mask_type in {"causal", "padding_causal"}: + bottom_right_diagonal = False + if bottom_right_diagonal is None or attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + }: + bottom_right_diagonal = True + if "padding" in attn_mask_type and attention_mask is not None: for mask in attention_mask: assert mask.dtype == torch.bool, "Attention mask must be in boolean type!" @@ -1001,6 +1026,7 @@ def forward( attention_mask=attention_mask, attn_mask_type=attn_mask_type, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, checkpoint_core_attention=checkpoint_core_attention, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index e226ef32d4..101e5b2525 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -137,6 +137,7 @@ def fused_attn_fwd( attn_mask_type: str = "padding", softmax_type: str = "vanilla", window_size: Tuple[int, int] = (-1, -1), + bottom_right_diagonal: bool = None, rng_gen: torch.Generator = None, softmax_offset: torch.Tensor = None, return_max_logit: bool = False, @@ -212,6 +213,9 @@ def fused_attn_fwd( in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding window and causal mask specifically. + bottom_right_diagonal: bool, default = None + whether to align sliding window and ALiBi diagonal to the top left (False) or + bottom right (True) corner of the softmax matrix. rng_gen : torch.Generator, default = None random number generator; if None, uses the default CUDA generator from PyTorch; otherwise, uses rng_gen @@ -255,6 +259,12 @@ def fused_attn_fwd( max_logit : if return_max_logit = True, shape [h] and same data type as O; otherwise None """ + if bottom_right_diagonal is None: + bottom_right_diagonal = attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + } + if attn_scale is None: d = q.size(-1) attn_scale = 1.0 / math.sqrt(d) @@ -306,6 +316,7 @@ def fused_attn_fwd( AttnMaskType[attn_mask_type], SoftmaxType[softmax_type], window_size, + bottom_right_diagonal, cu_seqlens_q, cu_seqlens_kv, q, @@ -370,6 +381,7 @@ def fused_attn_bwd( attn_mask_type: str = "padding", softmax_type: str = "vanilla", window_size: Tuple[int, int] = (-1, -1), + bottom_right_diagonal: bool = None, deterministic: bool = False, cuda_graph: bool = False, ) -> Tuple[Union[torch.Tensor, None], ...]: @@ -442,6 +454,9 @@ def fused_attn_bwd( in [i + seqlen_k - seqlen_q - window_size[0], i + seqlen_k - seqlen_q + window_size[1]] inclusive. Special cases (-1, -1) and (-1, 0) mean no sliding window and causal mask specifically. + bottom_right_diagonal: bool, default = None + whether to align sliding window and ALiBi diagonal to the top left (False) or + bottom right (True) corner of the softmax matrix. deterministic : bool, default = False whether to execute the backward pass with deterministic behaviours. cuda_graph : bool, default = False @@ -462,6 +477,12 @@ def fused_attn_bwd( gradient tensor of softmax offset of shape [1, h_q, 1, 1]. See softmax_type in DotProductAttention for details. """ + if bottom_right_diagonal is None: + bottom_right_diagonal = attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + } + if attn_scale is None: d = q.size(-1) attn_scale = 1.0 / math.sqrt(d) @@ -500,6 +521,7 @@ def fused_attn_bwd( AttnMaskType[attn_mask_type], SoftmaxType[softmax_type], window_size, + bottom_right_diagonal, deterministic, cu_seqlens_q, cu_seqlens_kv, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 591c89f83f..f7cf32eaf6 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -87,9 +87,10 @@ 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, 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::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, py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, @@ -99,10 +100,10 @@ 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_Softmax_Type softmax_type, const std::vector window_size, 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, + 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 std::vector Aux_CTX_Tensors, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index be645d91b9..bf62db8c33 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -100,9 +100,10 @@ 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, 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::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, py::handle s_quantizer, py::handle o_quantizer, const std::optional Bias, @@ -235,7 +236,7 @@ std::vector fused_attn_fwd( 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], workspace.data(), + softmax_type, window_size[0], window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); }); @@ -295,7 +296,7 @@ std::vector fused_attn_fwd( 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], workspace.data(), + softmax_type, window_size[0], window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); }); @@ -310,10 +311,10 @@ 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_Softmax_Type softmax_type, const std::vector window_size, 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, + 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 std::vector Aux_CTX_Tensors, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, @@ -532,14 +533,14 @@ std::vector fused_attn_bwd( // populate tensors with appropriate shapes and dtypes NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd(te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), - te_S.data(), te_dP.data(), &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], window_size[1], deterministic, cuda_graph, - workspace.data(), at::cuda::getCurrentCUDAStream()); + nvte_fused_attn_bwd( + te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), + &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], + window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace @@ -549,14 +550,14 @@ std::vector fused_attn_bwd( // execute kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_fused_attn_bwd(te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), - te_S.data(), te_dP.data(), &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], window_size[1], deterministic, cuda_graph, - workspace.data(), at::cuda::getCurrentCUDAStream()); + nvte_fused_attn_bwd( + te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), + &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], + window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), + at::cuda::getCurrentCUDAStream()); }); // destroy tensor wrappers diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 7c3125a165..fdb3869199 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -34,7 +34,7 @@ from transformer_engine.pytorch.distributed import get_distributed_world_size from transformer_engine.pytorch.export import is_in_onnx_export_mode from transformer_engine.pytorch.module.base import TransformerEngineBaseModule - +import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils warnings.filterwarnings("module", category=DeprecationWarning, module="transformer") @@ -148,11 +148,21 @@ class TransformerLayer(torch.nn.Module): distinguishes them based on :attr:`self_attn_mask_type` or :attr:`enc_dec_attn_mask_type`. Similar to :attr:`self_attn_mask_type`, :attr:`window_size` can be overridden by :attr:`window_size` in :meth:`forward` as well. + bottom_right_diagonal: Optional[bool], default = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the encoder. + If `None`, it will be set to `False` for `self_attn_mask_type` = + {`causal`, `padding_causal`} and `True` for other mask types. enc_dec_attn_mask_type : {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'}, default = "no_mask" type of attention mask passed into softmax operation for decoder. enc_dec_window_size : Optional[Tuple[int, int]], default = None sliding window size for local attention in decoder. + enc_dec_bottom_right_diagonal: Optional[bool], default = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the decoder. + If `None`, it will be set to `False` for `enc_dec_attn_mask_type` = + {`causal`, `padding_causal`} and `True` for other mask types. zero_centered_gamma : bool, default = False if set to ``True``, gamma parameter in LayerNorm is initialized to 0 and the LayerNorm formula changes to @@ -301,7 +311,9 @@ def __init__( kv_channels: Optional[int] = None, self_attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, enc_dec_attn_mask_type: str = "no_mask", + enc_dec_bottom_right_diagonal: Optional[bool] = None, enc_dec_window_size: Optional[Tuple[int, int]] = None, tp_group: Optional[dist_group_type] = None, tp_size: int = 1, @@ -343,8 +355,10 @@ def __init__( self.self_attn_mask_type = self_attn_mask_type self.window_size = window_size + self.bottom_right_diagonal = bottom_right_diagonal self.enc_dec_attn_mask_type = enc_dec_attn_mask_type self.enc_dec_window_size = enc_dec_window_size + self.enc_dec_bottom_right_diagonal = enc_dec_bottom_right_diagonal params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype ub_bulk_wgrad = ub_tp_comm_overlap and ub_bulk_wgrad ub_bulk_dgrad = ub_tp_comm_overlap and ub_bulk_dgrad @@ -606,10 +620,12 @@ def forward( attention_mask: Optional[torch.Tensor] = None, self_attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, + bottom_right_diagonal: Optional[bool] = None, encoder_output: Optional[torch.Tensor] = None, enc_dec_attn_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, enc_dec_attn_mask_type: Optional[str] = None, enc_dec_window_size: Optional[Tuple[int, int]] = None, + enc_dec_bottom_right_diagonal: Optional[bool] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[InferenceParams] = None, @@ -654,6 +670,11 @@ def forward( causal masks are aligned to the bottom right corner. window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention in encoder. + bottom_right_diagonal: Optional[bool] = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the encoder. + If `None`, it will be set to `False` for `self_attn_mask_type` = + {`causal`, `padding_causal`} and `True` for other mask types. encoder_output : Optional[torch.Tensor], default = None Output of the encoder block to be fed into the decoder block if using :attr:`layer_type` = ``"decoder"``. @@ -670,6 +691,11 @@ def forward( Type of attention mask passed into softmax operation for decoder. enc_dec_window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention in decoder. + enc_dec_bottom_right_diagonal: Optional[bool] = `None` + Align sliding window and ALiBi diagonal to the top left (`False`) + or bottom right (`True`) corner of the softmax matrix in the decoder. + If `None`, it will be set to `False` for `enc_dec_attn_mask_type` = + {`causal`, `padding_causal`} and `True` for other mask types. is_first_microbatch : {True, False, None}, default = None During training using either gradient accumulation or pipeline parallelism a minibatch of data is further split @@ -736,10 +762,35 @@ def forward( self_attn_mask_type = self.self_attn_mask_type if window_size is None: window_size = self.window_size + window_size = dpa_utils.check_set_window_size(self_attn_mask_type, window_size) + if enc_dec_attn_mask_type is None: enc_dec_attn_mask_type = self.enc_dec_attn_mask_type if enc_dec_window_size is None: enc_dec_window_size = self.enc_dec_window_size + enc_dec_window_size = dpa_utils.check_set_window_size( + enc_dec_attn_mask_type, enc_dec_window_size + ) + + if bottom_right_diagonal is None: + bottom_right_diagonal = self.bottom_right_diagonal + if self_attn_mask_type in {"causal", "padding_causal"}: + bottom_right_diagonal = False + if bottom_right_diagonal is None or self_attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + }: + bottom_right_diagonal = True + + if enc_dec_bottom_right_diagonal is None: + enc_dec_bottom_right_diagonal = self.enc_dec_bottom_right_diagonal + if enc_dec_attn_mask_type in {"causal", "padding_causal"}: + enc_dec_bottom_right_diagonal = False + if enc_dec_bottom_right_diagonal is None or enc_dec_attn_mask_type in { + "causal_bottom_right", + "padding_causal_bottom_right", + }: + enc_dec_bottom_right_diagonal = True assert ( self_attn_mask_type in AttnMaskTypes @@ -778,6 +829,7 @@ def forward( attention_mask=attention_mask, attn_mask_type=self_attn_mask_type, window_size=window_size, + bottom_right_diagonal=bottom_right_diagonal, inference_params=inference_params, is_first_microbatch=is_first_microbatch, checkpoint_core_attention=checkpoint_core_attention, @@ -813,6 +865,7 @@ def forward( attention_mask=enc_dec_attn_mask, attn_mask_type=enc_dec_attn_mask_type, window_size=enc_dec_window_size, + bottom_right_diagonal=enc_dec_bottom_right_diagonal, encoder_output=encoder_output, inference_params=inference_params, is_first_microbatch=is_first_microbatch, From 52ee5ea06737f1e1604d154b943aa51fab9b0f3d Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 22 Jan 2026 17:14:33 -0800 Subject: [PATCH 179/521] Fix bugs in permutation custom partitioning (#2617) * Use correct block size for workspace in row id map creation, also shard workspace correctly based on 2nd dim of routing_map/row_id map Signed-off-by: DoubleCheeseCheetos * reduce size of largest test case on single_GPU scenario to fit on L40 and A100 in CI line up Signed-off-by: tdophung --------- Signed-off-by: DoubleCheeseCheetos Signed-off-by: tdophung Co-authored-by: DoubleCheeseCheetos --- tests/jax/test_permutation.py | 4 +-- .../jax/triton_extensions/permutation.py | 29 ++++++++++++------- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index 5bb59c6ed5..138a817240 100644 --- a/tests/jax/test_permutation.py +++ b/tests/jax/test_permutation.py @@ -23,7 +23,7 @@ (128, 5, 128, 3), (1024, 8, 128, 8), (4096, 32, 1280, 2), - (4096, 256, 4096, 6), + (4096, 64, 4096, 6), ] DISPATCH_COMBINE_CASES = { "L0": ALL_DISPATCH_COMBINE_CASES[0:2], @@ -44,7 +44,7 @@ (128, 5, 128, 3, 8), (1024, 8, 128, 8, 16), (4096, 32, 1280, 2, 128), - (4096, 256, 4096, 6, 16), + (4096, 64, 4096, 6, 16), ] DISPATCH_COMBINE_PADDING_CASES = { "L0": ALL_DISPATCH_COMBINE_PADDING_CASES[0:2], diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index bd8bd8ff13..0c80f9f18c 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -65,8 +65,6 @@ class RowIdMapPass1Primitive(BasePrimitive): @staticmethod def abstract(routing_map_aval, *, num_tokens, num_experts, block_size): """Shape/dtype inference for pass 1.""" - del block_size # Only affects grid, not output shape - assert routing_map_aval.shape == ( num_tokens, num_experts, @@ -75,7 +73,7 @@ def abstract(routing_map_aval, *, num_tokens, num_experts, block_size): row_id_map_shape = (num_tokens, num_experts * 2 + 1) workspace_shape = ( num_experts, - triton.cdiv(num_tokens, DEFAULT_BLOCK_SIZE), + triton.cdiv(num_tokens, block_size), ) return ( @@ -134,9 +132,10 @@ def infer_sharding_from_operands( desc="RowIdMapPass1.row_id_map_sharding", ) # Workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + # Second dim depends on num_tokens, so it must be sharded on the same axis as tokens workspace_sharding = NamedSharding( mesh, - PartitionSpec(None, None), + PartitionSpec(None, routing_map_spec[0]), desc="RowIdMapPass1.workspace_sharding", ) return [row_id_map_sharding, workspace_sharding] @@ -156,9 +155,11 @@ def partition(num_tokens, num_experts, block_size, mesh, arg_infos, result_infos PartitionSpec(routing_map_spec[0], None), desc="RowIdMapPass1.row_id_map_sharding", ) + # Workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + # Second dim depends on num_tokens, so it must be sharded on the same axis as tokens workspace_sharding = NamedSharding( mesh, - PartitionSpec(None, None), + PartitionSpec(None, routing_map_spec[0]), desc="RowIdMapPass1.workspace_sharding", ) out_shardings = [row_id_map_sharding, workspace_sharding] @@ -186,7 +187,8 @@ def shardy_sharding_rule(num_tokens, num_experts, block_size, mesh, value_types, # Note: row_id_cols != experts since it's num_experts * 2 + 1 row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_row_id_cols") # workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) - workspace_spec = (f"{prefix}_experts", f"{prefix}_ws_blocks") + # Second dim depends on num_tokens, so use same factor to ensure same sharding + workspace_spec = (f"{prefix}_experts", f"{prefix}_tokens") return SdyShardingRule((input_spec,), (row_id_map_spec, workspace_spec)) @@ -208,10 +210,9 @@ class RowIdMapPass2Primitive(BasePrimitive): def abstract(row_id_map_aval, workspace_aval, *, num_tokens, num_experts, block_size): """Shape/dtype inference for pass 2 (in-place operation).""" del row_id_map_aval, workspace_aval - del block_size row_id_map_shape = (num_tokens, num_experts * 2 + 1) - workspace_shape = (num_experts, triton.cdiv(num_tokens, DEFAULT_BLOCK_SIZE)) + workspace_shape = (num_experts, triton.cdiv(num_tokens, block_size)) return ( jax.core.ShapedArray(row_id_map_shape, jnp.int32), @@ -270,9 +271,11 @@ def infer_sharding_from_operands( PartitionSpec(*row_id_map_spec), desc="RowIdMapPass2.row_id_map_sharding", ) + # Workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + # Second dim depends on num_tokens, so it must be sharded on the same axis as tokens workspace_sharding = NamedSharding( mesh, - PartitionSpec(None, None), + PartitionSpec(None, row_id_map_spec[0]), desc="RowIdMapPass2.workspace_sharding", ) return [row_id_map_sharding, workspace_sharding] @@ -292,9 +295,11 @@ def partition(num_tokens, num_experts, block_size, mesh, arg_infos, result_infos PartitionSpec(*row_id_map_spec), desc="RowIdMapPass2.row_id_map_sharding", ) + # Workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + # Second dim depends on num_tokens, so it must be sharded on the same axis as tokens workspace_sharding = NamedSharding( mesh, - PartitionSpec(None, None), + PartitionSpec(None, row_id_map_spec[0]), desc="RowIdMapPass2.workspace_sharding", ) out_shardings = [row_id_map_sharding, workspace_sharding] @@ -317,7 +322,9 @@ def shardy_sharding_rule(num_tokens, num_experts, block_size, mesh, value_types, del num_tokens, num_experts, block_size, mesh, value_types, result_types prefix = "RowIdMapPass2" row_id_map_spec = (f"{prefix}_tokens", f"{prefix}_cols") - workspace_spec = (f"{prefix}_ws_experts", f"{prefix}_ws_blocks") + # workspace shape: (num_experts, cdiv(num_tokens, BLOCK_SIZE)) + # Second dim depends on num_tokens, so use same factor to ensure same sharding + workspace_spec = (f"{prefix}_ws_experts", f"{prefix}_tokens") return SdyShardingRule((row_id_map_spec, workspace_spec), (row_id_map_spec, workspace_spec)) From a0a89a8eb6e173fb73d4f10dbc94e3c8a1609927 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Fri, 23 Jan 2026 06:37:48 +0100 Subject: [PATCH 180/521] [Common] Disabled the tuned NVFP4 kernels (#2615) * Disabled the tuned NVFP4 kernels Signed-off-by: Oleg Goncharov * Disabled fast math in cpp tests Signed-off-by: Oleg Goncharov --------- Signed-off-by: Oleg Goncharov --- tests/cpp/operator/test_cast_nvfp4_transpose.cu | 7 +------ .../common/cast/nvfp4/quantize_transpose_nvfp4.cuh | 8 ++++---- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index c4df8759f2..d8d495d61f 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -677,11 +677,6 @@ std::vector Activation_types = { ActivationType::Identity }; -std::vector use_fast_nvfp4_scaling_vec = { - false, - true -}; - } // namespace class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam @@ -743,7 +738,7 @@ INSTANTIATE_TEST_SUITE_P( ::testing::ValuesIn(Activation_types), ::testing::ValuesIn(tensor_dims), ::testing::Values(DType::kBFloat16), - ::testing::ValuesIn(use_fast_nvfp4_scaling_vec)), + ::testing::Values(false)), [](const testing::TestParamInfo& info) { std::string name = to_string(std::get<0>(info.param)); const auto& shape = std::get<1>(info.param); diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 99776db281..61c6ba9cef 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -1168,10 +1168,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, // TODO(Frank): Is there a better way to do this? bool return_transpose = output->has_columnwise_data(); - if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { - quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); - return; - } + // if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { + // quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); + // return; + // } constexpr bool COMPUTE_ACTIVATIONS = false; using ParamOP = Empty; From 72592763737e2f54145013be40d81d5101590e3c Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Sat, 24 Jan 2026 16:05:05 -0800 Subject: [PATCH 181/521] [PyTorch] Support user-defined op fusions (#2597) * Expose option for custom op fusions Refactor fusion functions to remove index bookkeeping. Refactor fused ops to use consistent operation order. Signed-off-by: Tim Moon * Add tests for custom ops Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix linter warnings and numerical test failures Signed-off-by: Tim Moon * Tweak pattern matching logic with fixed window sizes Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use TF32 tols in fused op tests Signed-off-by: Tim Moon * Review suggestion from @greptile-apps Signed-off-by: Tim Moon * Backpropagate fixes from #2622 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 | 322 +++++++++++++++++- transformer_engine/pytorch/ops/__init__.py | 10 +- .../pytorch/ops/fused/__init__.py | 60 ++-- .../ops/fused/backward_activation_bias.py | 121 +++---- .../pytorch/ops/fused/backward_add_rmsnorm.py | 104 +++--- .../pytorch/ops/fused/backward_linear_add.py | 117 ++++--- .../ops/fused/backward_linear_scale.py | 109 +++--- .../fused/forward_linear_bias_activation.py | 117 +++---- .../ops/fused/forward_linear_bias_add.py | 119 +++---- .../ops/fused/forward_linear_scale_add.py | 128 ++++--- .../ops/fused/userbuffers_backward_linear.py | 157 ++++----- .../ops/fused/userbuffers_forward_linear.py | 144 ++++---- transformer_engine/pytorch/ops/fuser.py | 168 ++++++--- 13 files changed, 996 insertions(+), 680 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 7183e30e71..a23de29e02 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -2329,13 +2329,13 @@ def test_backward_activation_bias( backward_ops = model._module_groups[0]._backward_ops if with_quantization: assert len(backward_ops) == 2 - assert isinstance(backward_ops[0][0], BackwardActivationBias) - assert isinstance(backward_ops[1][0], te_ops.Quantize) + assert isinstance(backward_ops[0][0], te_ops.Quantize) + assert isinstance(backward_ops[1][0], BackwardActivationBias) else: assert len(backward_ops) == 3 - assert isinstance(backward_ops[0][0], act_type) + assert isinstance(backward_ops[0][0], te_ops.Quantize) assert isinstance(backward_ops[1][0], te_ops.Bias) - assert isinstance(backward_ops[2][0], te_ops.Quantize) + assert isinstance(backward_ops[2][0], act_type) # Expected numerical error tols = dtype_tols(dtype) @@ -2930,3 +2930,317 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: 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) + + +class TestCustomOps: + """Test with ops that are defined externally""" + + def test_custom_basic_op( + self, + *, + shape: Iterable[int] = (7, 5), + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ) -> None: + """Custom basic op""" + + class CustomScaleOp(te.ops.BasicOperation): + """Custom op that applies a learnable scale""" + + def __init__(self) -> None: + super().__init__() + self.scale: torch.nn.Parameter + scale = torch.ones((), dtype=dtype, device=device) + scale = torch.nn.Parameter(scale) + self.register_parameter("scale", scale) + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + ctx.save_for_backward(self.scale, input_) + return self.scale * input_ + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> torch.Tensor: + ( + scale, + input_, + ) = ctx.saved_tensors + grad_scale = torch.inner(input_.reshape(-1), grad_output.reshape(-1)) + grad_scale = grad_scale.reshape(()) + grad_input = scale * grad_output + return grad_input, (grad_scale,) + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + ) + w_ref, w_test = make_reference_and_test_tensors( + (), + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + y_ref = w_ref * x_ref + y_ref.backward(dy_ref) + + # Implementation with fusible operation + op = CustomScaleOp() + forward = te.ops.Sequential(te.ops.Identity(), op, te.ops.Identity()) + with torch.no_grad(): + op.scale.copy_(w_test) + del w_test + y_test = forward(x_test) + y_test.backward(dy_test) + + # Check results + tols = dtype_tols(dtype) + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dw_test = op.scale.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(dx_test, x_ref.grad, **tols) + torch.testing.assert_close(dw_test, w_ref.grad, **tols) + + def test_custom_forward_fused_op( + self, + *, + shape: Iterable[int] = (7, 11), + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ): + """Custom fused op in forward pass""" + + class CustomForwardLinearSiLU(te.ops.FusedOperation): + """Custom fused op for GEMM + SiLU""" + + _enabled = True + + def __init__(self, *, linear, silu) -> None: + super().__init__((linear, silu)) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + **unused, + ) -> torch.Tensor: + weight = self.basic_ops[0].weight + dtype = weight.dtype + device = weight.device + + # Perform compute on CPU, because why not? + x = input_.cpu() + w = weight.cpu() + y = torch.matmul(x, w.T) + z = torch.nn.functional.silu(y) + out = z.to(device=device) + + # Save state for linear backward + linear_op_ctx = basic_op_ctxs[0] + linear_op_ctx.save_for_backward(input_, weight) + linear_op_ctx.with_quantized_compute = False + linear_op_ctx.input_quantizer = None + linear_op_ctx.weight_quantizer = None + linear_op_ctx.grad_output_quantizer = None + linear_op_ctx.grad_input_quantizer = None + linear_op_ctx.dtype = dtype + linear_op_ctx.input_requires_grad = True + linear_op_ctx.weight_requires_grad = True + + # Save state for SiLU backward + silu_op_ctx = basic_op_ctxs[1] + silu_op_ctx.save_for_backward(y.to(device=device)) + silu_op_ctx.dtype = dtype + silu_op_ctx.prev_op_grad_output_quantizer = None + + return out, [(), ()] + + @staticmethod + def fuse_ops( + ops: list[FusibleOperation], + **unused, + ) -> list[FusibleOperation]: + """Apply fusion the first time this function is called""" + if CustomForwardLinearSiLU._enabled: + CustomForwardLinearSiLU._enabled = False + op = CustomForwardLinearSiLU(linear=ops[0], silu=ops[1]) + return [op] + ops[2:] + return ops + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + ) + w_ref, w_test = make_reference_and_test_tensors( + (shape[-1], shape[-1]), + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + y_ref = torch.nn.functional.linear(x_ref, w_ref) + y_ref = torch.nn.functional.silu(y_ref) + y_ref.backward(dy_ref) + + # Implementation with fusible operation + te.ops.register_forward_fusion(CustomForwardLinearSiLU.fuse_ops) + model = te.ops.Sequential( + te.ops.Linear(shape[-1], shape[-1], bias=False), + te.ops.SiLU(), + ) + with torch.no_grad(): + model[0].weight.copy_(w_test) + del w_test + y_test = model(x_test) + y_test.backward(dy_test) + + # Check that forward operations have been fused + forward_ops = model._module_groups[0]._forward_ops + assert len(forward_ops) == 1 + assert isinstance(forward_ops[0][0], CustomForwardLinearSiLU) + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + + # Check results + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dw_test = model[0].weight.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(dx_test, x_ref.grad, **tols) + torch.testing.assert_close(dw_test, w_ref.grad, **tols) + + def test_custom_backward_fused_op( + self, + *, + shape: Iterable[int] = (13, 5), + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ): + """Custom fused op in backward pass""" + + class CustomBackwardLinearScale(te.ops.FusedOperation): + """Custom fused op for backward linear + scale""" + + _enabled: bool = True + + def __init__(self, *, scale, linear) -> None: + super().__init__((scale, linear)) + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + **unused, + ) -> torch.Tensor: + + # Load state from linear forward + linear_op_ctx = basic_op_ctxs[1] + x, w = linear_op_ctx.saved_tensors + dtype = linear_op_ctx.dtype + device = w.device + + # Perform compute in FP64 and apply scale before dgrad + # GEMM instead of after + scale = self.basic_ops[0].scale + dy = grad_output.double() + x = x.double() + w = w.double() + dx = torch.matmul(dy, scale * w) + dw = torch.matmul(dy.T, x) + dx = dx.to(dtype=dtype) + dw = dw.to(dtype=dtype) + + return dx, [(), (dw,)], [(), ()] + + @staticmethod + def fuse_ops( + ops: list[FusibleOperation], + **unused, + ) -> list[FusibleOperation]: + """Apply fusion the first time this function is called""" + if CustomBackwardLinearScale._enabled: + CustomBackwardLinearScale._enabled = False + op = CustomBackwardLinearScale(scale=ops[0], linear=ops[1]) + return [op] + ops[2:] + return ops + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + ) + w_ref, w_test = make_reference_and_test_tensors( + (shape[-1], shape[-1]), + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + scale = 1.234 + + # Plain PyTorch implementation + y_ref = torch.nn.functional.linear(scale * x_ref, w_ref) + y_ref.backward(dy_ref) + + # Implementation with fusible operation + te.ops.register_backward_fusion(CustomBackwardLinearScale.fuse_ops, prepend=True) + model = te.ops.Sequential( + te.ops.ConstantScale(scale), + te.ops.Linear(shape[-1], shape[-1], bias=False), + ) + with torch.no_grad(): + model[1].weight.copy_(w_test) + del w_test + y_test = model(x_test) + y_test.backward(dy_test) + + # Check that forward operations have been fused + backward_ops = model._module_groups[0]._backward_ops + assert len(backward_ops) == 1 + assert isinstance(backward_ops[0][0], CustomBackwardLinearScale) + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + + # Check results + y_test = y_test.to(dtype=torch.float64, device="cpu") + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + dw_test = model[1].weight.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) + torch.testing.assert_close(dx_test, x_ref.grad, **tols) + torch.testing.assert_close(dw_test, w_ref.grad, **tols) diff --git a/transformer_engine/pytorch/ops/__init__.py b/transformer_engine/pytorch/ops/__init__.py index 2b270ea3de..99f51a9c7a 100644 --- a/transformer_engine/pytorch/ops/__init__.py +++ b/transformer_engine/pytorch/ops/__init__.py @@ -8,7 +8,9 @@ """ -from transformer_engine.pytorch.ops.basic import * -from transformer_engine.pytorch.ops.linear import Linear -from transformer_engine.pytorch.ops.op import FusibleOperation -from transformer_engine.pytorch.ops.sequential import Sequential +from .basic import * +from .fuser import register_backward_fusion, register_forward_fusion +from .linear import Linear +from .op import BasicOperation, FusedOperation, FusibleOperation +from .sequential import Sequential +from . import fused diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index f4568ff25d..19608894e0 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -4,39 +4,27 @@ """Compound tensor operation supported by the operation fuser.""" -from .backward_activation_bias import ( - BackwardActivationBias, - fuse_backward_activation_bias, -) -from .backward_add_rmsnorm import ( - BackwardAddRMSNorm, - fuse_backward_add_rmsnorm, -) -from .backward_linear_add import ( - BackwardLinearAdd, - fuse_backward_linear_add, -) -from .backward_linear_scale import ( - BackwardLinearScale, - fuse_backward_linear_scale, -) -from .forward_linear_bias_activation import ( - ForwardLinearBiasActivation, - fuse_forward_linear_bias_activation, -) -from .forward_linear_bias_add import ( - ForwardLinearBiasAdd, - fuse_forward_linear_bias_add, -) -from .forward_linear_scale_add import ( - ForwardLinearScaleAdd, - fuse_forward_linear_scale_add, -) -from .userbuffers_backward_linear import ( - UserbuffersBackwardLinear, - fuse_userbuffers_backward_linear, -) -from .userbuffers_forward_linear import ( - UserbuffersForwardLinear, - fuse_userbuffers_forward_linear, -) +from ..fuser import register_backward_fusion, register_forward_fusion +from .backward_activation_bias import BackwardActivationBias +from .backward_add_rmsnorm import BackwardAddRMSNorm +from .backward_linear_add import BackwardLinearAdd +from .backward_linear_scale import BackwardLinearScale +from .forward_linear_bias_activation import ForwardLinearBiasActivation +from .forward_linear_bias_add import ForwardLinearBiasAdd +from .forward_linear_scale_add import ForwardLinearScaleAdd +from .userbuffers_backward_linear import UserbuffersBackwardLinear +from .userbuffers_forward_linear import UserbuffersForwardLinear + + +# Register forward fusions +register_forward_fusion(UserbuffersForwardLinear.fuse_forward_ops) +register_forward_fusion(ForwardLinearBiasAdd.fuse_forward_ops) +register_forward_fusion(ForwardLinearBiasActivation.fuse_forward_ops) +register_forward_fusion(ForwardLinearScaleAdd.fuse_forward_ops) + +# Register backward fusions +register_backward_fusion(UserbuffersBackwardLinear.fuse_backward_ops) +register_backward_fusion(BackwardLinearAdd.fuse_backward_ops) +register_backward_fusion(BackwardLinearScale.fuse_backward_ops) +register_backward_fusion(BackwardActivationBias.fuse_backward_ops) +register_backward_fusion(BackwardAddRMSNorm.fuse_backward_ops) diff --git a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py index d5b9ce0e96..4ab082d32b 100644 --- a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py +++ b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py @@ -53,8 +53,8 @@ def fuser_backward( ]: # Get basic operation contexts - activation_op_ctx = basic_op_ctxs[0] - bias_op_ctx = basic_op_ctxs[1] + bias_op_ctx = basic_op_ctxs[0] + activation_op_ctx = basic_op_ctxs[1] # Saved tensors from forward pass (act_input,) = activation_op_ctx.saved_tensors @@ -79,68 +79,59 @@ def fuser_backward( # Clear activation input tensor clear_tensor_data(act_input) - return dx, [(), (db,)], [(), ()] + return dx, [(db,), ()], [(), ()] - -def fuse_backward_activation_bias( - ops: list[tuple[FusibleOperation, list[int]]], - recipe: Optional[Recipe], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fused backward dact + dbias + quantize - - Parameters - ---------- - ops : list of tuples - Backward pass operations and the indices of the corresponding - basic operations. - recipe : Recipe, optional - Used quantization recipe - - Returns - ------- - ops : list of tuples - Updated backward pass operations - - """ - - # Check if recipe supports bias activation fusion - if recipe is None: - return ops - - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 3: + @staticmethod + def fuse_backward_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for backward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Backward pass operations. + recipe : Recipe, optional + Quantization recipe. + + Returns + ------- + ops : list of FusibleOperation + Updated backward pass operations + + """ + + # Check if recipe supports bias activation fusion + if recipe is None: + return ops + + # Scan through ops, fusing if possible + out = [] + window, ops = ops[:3], ops[3:] + while len(window) == 3: + if ( + isinstance(window[2], _fusible_activations) + and isinstance(window[1], Bias) + and window[0].get_grad_output_quantizer() is not None + ): + # Construct fused op if window matches pattern + op = BackwardActivationBias(bias=window[1], activation=window[2]) + window = [window[0], op] + else: + # Shift window if window doesn't match pattern + out.extend(window[:-2]) + window = window[-2:] + + # Adjust window to expected size + out.extend(window[:-3]) + window = window[-3:] + while ops and len(window) < 3: + window.append(ops[0]) + ops = ops[1:] + + # Return list of ops out.extend(window) - - # Check if first op is a supported activation - window, ops = ops[:1], ops[1:] - op, _ = window[0] - if not isinstance(op, _fusible_activations): - continue - - # Check if second op is bias - op, _ = ops[0] - if not isinstance(op, Bias): - continue - - # Check if third op has a grad input quantizer - op, _ = ops[1] - if not op.num_quantizers("backward") > 0: - continue - - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = BackwardActivationBias( - activation=window[0][0], - bias=window[1][0], - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py index 186619caae..a3c81e60c8 100644 --- a/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py +++ b/transformer_engine/pytorch/ops/fused/backward_add_rmsnorm.py @@ -42,7 +42,7 @@ def fuser_backward( # Get basic operations rmsnorm_op = self.basic_ops[1] - rmsnorm_op_ctx = basic_op_ctxs[0] + rmsnorm_op_ctx = basic_op_ctxs[1] # Saved tensors from forward pass x, rstdevs = rmsnorm_op_ctx.saved_tensors @@ -53,7 +53,7 @@ def fuser_backward( # Check input tensors dtype = rmsnorm_op_ctx.dtype - extra_grad = basic_op_grad_extra_outputs[1][0] + extra_grad = basic_op_grad_extra_outputs[0][0] dy = maybe_dequantize(grad_output.contiguous(), dtype).view(x.size()) w = maybe_dequantize(rmsnorm_op.weight, dtype).view((inner_dim,)) add = maybe_dequantize(extra_grad.contiguous(), dtype).view(x.size()) @@ -77,57 +77,51 @@ def fuser_backward( grad_input = dx.view(grad_output.size()) grad_weight = dw.view(weight_dims) - return grad_input, [(grad_weight,), ()], [(), ()] - - -def fuse_backward_add_rmsnorm( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fused backward RMNorm + add - - Parameters - ---------- - ops : list of tuples - Backward pass operations and the indices of the corresponding - basic operations. - - Returns - ------- - ops : list of tuples - Updated backward pass operations - - """ - - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 2: + return grad_input, [(), (grad_weight,)], [(), ()] + + @staticmethod + def fuse_backward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for backward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Backward pass operations. + + Returns + ------- + ops : list of FusibleOperation + Updated backward pass operations + + """ + + # Scan through ops, fusing if possible + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + if ( + isinstance(window[0], MakeExtraOutput) + and isinstance(window[1], RMSNorm) + and not window[0]._in_place + ): + # Construct fused op if window matches pattern + op = BackwardAddRMSNorm(add=window[0], rmsnorm=window[1]) + window = [op] + else: + # Shift window if window doesn't match pattern + out.extend(window[:-1]) + window = window[-1:] + + # Adjust window to expected size + out.extend(window[:-2]) + window = window[-2:] + while ops and len(window) < 2: + window.append(ops[0]) + ops = ops[1:] + + # Return list of ops out.extend(window) - - # Check if first op is linear - window, ops = ops[:1], ops[1:] - op, _ = window[0] - if not isinstance(op, RMSNorm): - continue - - # Check if second op is "make extra output" - op, _ = ops[0] - if not isinstance(op, MakeExtraOutput): - continue - if op._in_place: - continue - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = BackwardAddRMSNorm( - rmsnorm=window[0][0], - add=window[1][0], - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index 5e7339db85..c06e212e87 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -45,7 +45,7 @@ def fuser_backward( # Get basic operations linear_op = self.basic_ops[1] - linear_op_ctx = basic_op_ctxs[0] + linear_op_ctx = basic_op_ctxs[1] # Saved tensors from forward pass (x_local, w) = linear_op_ctx.saved_tensors @@ -71,7 +71,7 @@ def fuser_backward( accumulate_into_main_grad = False # Linear backward pass - grad_input = basic_op_grad_extra_outputs[1][0] + grad_input = basic_op_grad_extra_outputs[0][0] grad_input, grad_weight = BasicLinear._functional_backward( grad_output=grad_output, input=x_local, @@ -109,61 +109,60 @@ def fuser_backward( zero=getattr(weight_param, "zero_out_wgrad", False), ) - return grad_input, [(grad_weight,), ()], [(), ()] - - -def fuse_backward_linear_add( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fused backward dgrad GEMM + add - - Parameters - ---------- - ops : list of tuples - Backward pass operations and the indices of the corresponding - basic operations. - - Returns - ------- - ops : list of tuples - Updated backward pass operations - - """ - - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 2: + return grad_input, [(), (grad_weight,)], [(), ()] + + @staticmethod + def fuse_backward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for backward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Backward pass operations. + + Returns + ------- + ops : list of FusibleOperation + Updated backward pass operations + + """ + + # Scan through ops, fusing if possible + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + + # Check if window matches pattern + matches_pattern = True + if not (isinstance(window[0], MakeExtraOutput) and isinstance(window[1], BasicLinear)): + matches_pattern = False + elif not window[0]._in_place: + # Fused op accumulates grad input in-place + matches_pattern = False + elif window[1].tensor_parallel_mode == "column": + # Column tensor-parallelism requires communication + # after the dgrad GEMM + matches_pattern = False + + if matches_pattern: + # Construct fused op if window matches pattern + op = BackwardLinearAdd(backward_add=window[0], linear=window[1]) + window = [op] + else: + # Shift window if window doesn't match pattern + out.extend(window[:-1]) + window = window[-1:] + + # Adjust window to expected size + out.extend(window[:-2]) + window = window[-2:] + while ops and len(window) < 2: + window.append(ops[0]) + ops = ops[1:] + + # Return list of ops out.extend(window) - - # Check if first op is linear - window, ops = ops[:1], ops[1:] - op, _ = window[0] - if not isinstance(op, BasicLinear): - continue - if op.tensor_parallel_mode == "column": - # Row tensor-parallelism requires communication after the - # GEMM - continue - - # Check if second op is "make extra output" - op, _ = ops[0] - if not isinstance(op, MakeExtraOutput): - continue - if not op._in_place: - continue - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = BackwardLinearAdd( - linear=window[0][0], - backward_add=window[1][0], - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py index f7f59e65c9..709073e6f8 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py @@ -45,7 +45,7 @@ def fuser_backward( # Get basic operations linear_op = self.basic_ops[0] - linear_op_ctx = basic_op_ctxs[1] + linear_op_ctx = basic_op_ctxs[0] scale_op = self.basic_ops[1] # Saved tensors from forward pass @@ -109,58 +109,57 @@ def fuser_backward( zero=getattr(weight_param, "zero_out_wgrad", False), ) - return grad_input, [(), (grad_weight,)], [(), ()] - - -def fuse_backward_linear_scale( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fused backward dgrad GEMM + constant scale - - Parameters - ---------- - ops : list of tuples - Backward pass operations and the indices of the corresponding - basic operations. - - Returns - ------- - ops : list of tuples - Updated backward pass operations - - """ - - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 2: + return grad_input, [(grad_weight,), ()], [(), ()] + + @staticmethod + def fuse_backward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for backward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Backward pass operations. + + Returns + ------- + ops : list of FusibleOperation + Updated backward pass operations + + """ + + # Scan through ops, fusing if possible + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + + # Check if window matches pattern + matches_pattern = True + if not (isinstance(window[0], BasicLinear) and isinstance(window[1], ConstantScale)): + matches_pattern = False + elif window[0].tensor_parallel_mode == "column": + # Column tensor-parallelism requires communication + # after the dgrad GEMM + matches_pattern = False + + if matches_pattern: + # Construct fused op if window matches pattern + op = BackwardLinearScale(linear=window[0], scale=window[1]) + window = [op] + else: + # Shift window if window doesn't match pattern + out.extend(window[:-1]) + window = window[-1:] + + # Adjust window to expected size + out.extend(window[:-2]) + window = window[-2:] + while ops and len(window) < 2: + window.append(ops[0]) + ops = ops[1:] + + # Return list of ops out.extend(window) - - # Check if first op is constant scale - window, ops = ops[:1], ops[1:] - op, _ = window[0] - if not isinstance(op, ConstantScale): - continue - - # Check if second op is linear - op, _ = ops[0] - if not isinstance(op, BasicLinear): - continue - if op.tensor_parallel_mode == "column": - # Column tensor-parallelism requires communication after the dgrad GEMM - continue - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = BackwardLinearScale( - scale=window[0][0], - linear=window[1][0], - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index 1c5edfcfcb..dfc11a19e7 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -134,62 +134,63 @@ def fuser_forward( return output, [() for _ in range(len(self.basic_ops))] - -def fuse_forward_linear_bias_activation( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fuse forward GEMM + bias + activation - - Parameters - ---------- - ops : list of tuples - Forward pass operations and the indices of the corresponding - basic operations. - - Returns - ------- - ops : list of tuples - Updated forward pass operations - - """ - - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 2: + @staticmethod + def fuse_forward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for forward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Forward pass operations. + + Returns + ------- + ops : list of FusibleOperation + Updated forward pass operations + + """ + + # Scan through ops, fusing if possible + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + + # Check if window matches pattern + matches_pattern = True + if not (isinstance(window[0], BasicLinear) and isinstance(window[1], Bias)): + matches_pattern = False + elif window[0].tensor_parallel_mode == "row": + # Row tensor-parallelism requires communication after + # the GEMM + matches_pattern = False + elif window[0].weight.dtype not in (torch.float16, torch.bfloat16): + # cuBLAS only supports fused GEMM+bias+activation with + # FP16 and BF16 output + matches_pattern = False + + if matches_pattern: + # Construct fused op if window matches pattern + op = ForwardLinearBiasActivation( + linear=window[0], + bias=window[1], + activation=None, + ) + window = [op] + else: + # Shift window if window doesn't match pattern + out.extend(window[:-1]) + window = window[-1:] + + # Adjust window to expected size + out.extend(window[:-2]) + window = window[-2:] + while ops and len(window) < 2: + window.append(ops[0]) + ops = ops[1:] + + # Return list of ops out.extend(window) - - # Check if first op is linear - window, ops = ops[:1], ops[1:] - op1, _ = window[0] - if not isinstance(op1, BasicLinear): - continue - if op1.tensor_parallel_mode == "row": - # Row tensor-parallelism requires communication after the - # GEMM - continue - if op1.weight.dtype not in (torch.float16, torch.bfloat16): - # cuBLAS only supports fused GEMM+bias+activation with - # FP16 and BF16 output - continue - - # Check if second op is bias - op2, _ = ops[0] - if not isinstance(op2, Bias): - continue - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = ForwardLinearBiasActivation( - linear=window[0][0], - bias=window[1][0], - activation=None, - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 4efb33e037..2dfc0566b7 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -131,72 +131,63 @@ def fuser_forward( return output, [() for _ in range(len(self.basic_ops))] + @staticmethod + def fuse_forward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for forward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Forward pass operations. + + Returns + ------- + ops : list of FusibleOperation + Updated forward pass operations + + """ + + # Scan through ops, fusing if possible + out = [] + window = [] + while ops: + + # Shift window + out.extend(window) + window = [ops[0]] + ops = ops[1:] -def fuse_forward_linear_bias_add( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fuse forward GEMM + bias + add - - Parameters - ---------- - ops : list of tuples - Forward pass operations and the indices of the corresponding - basic operations. + # Check if first op is linear + if not isinstance(window[0], BasicLinear): + continue + if window[0].tensor_parallel_mode == "row": + # Row tensor-parallelism requires communication after + # the GEMM + continue + linear = window[0] - Returns - ------- - ops : list of tuples - Updated forward pass operations + # Check if next op is bias + bias = None + if ops and isinstance(ops[0], Bias): + window.append(ops[0]) + ops = ops[1:] + bias = window[-1] + + # Check if next op is in-place add extra input + if ops and isinstance(ops[0], AddExtraInput) and ops[0]._in_place: + window.append(ops[0]) + ops = ops[1:] + add = window[-1] + else: + continue - """ + # Replace window with fused op + op = ForwardLinearBiasAdd(linear=linear, bias=bias, add=add) + window = [op] - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 2: + # Return list of ops out.extend(window) - - # Check if first op is linear - window, ops = ops[:1], ops[1:] - op, _ = window[0] - if not isinstance(op, BasicLinear): - continue - if op.tensor_parallel_mode == "row": - # Row tensor-parallelism requires communication after the - # GEMM - continue - linear = op - op, _ = ops[0] - - # Check if next op is bias - bias = None - if isinstance(op, Bias): - bias = op - window.extend(ops[:1]) - ops = ops[1:] - if len(ops) == 0: - continue - op, _ = ops[0] - - # Check if next op is in-place add extra input - if not isinstance(op, AddExtraInput): - continue - if not op._in_place: - continue - add = op - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = ForwardLinearBiasAdd( - linear=linear, - bias=bias, - add=add, - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index 25b40f76e3..ae4bdd4b19 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -110,70 +110,66 @@ def fuser_forward( return output, [() for _ in range(len(self.basic_ops))] - -def fuse_forward_linear_scale_add( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Fuse forward GEMM + scale + add - - Parameters - ---------- - ops : list of tuples - Forward pass operations and the indices of the corresponding - basic operations. - - Returns - ------- - ops : list of tuples - Updated forward pass operations - - """ - - # Scan through ops, fusing if possible - out = [] - window = [] - while len(ops) >= 3: + @staticmethod + def fuse_forward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for forward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Forward pass operations. + + Returns + ------- + ops : list of FusibleOperation + Updated forward pass operations + + """ + + # Scan through ops, fusing if possible + out = [] + window, ops = ops[:3], ops[3:] + while len(window) == 3: + + # Check if window matches pattern + matches_pattern = True + if not ( + isinstance(window[0], BasicLinear) + and isinstance(window[1], ConstantScale) + and isinstance(window[2], AddExtraInput) + ): + matches_pattern = False + elif window[0].tensor_parallel_mode == "row": + # Row tensor-parallelism requires communication after + # the GEMM + matches_pattern = False + elif not window[2]._in_place: + # Fused op accumulates output in-place + matches_pattern = False + + if matches_pattern: + # Construct fused op if window matches pattern + op = ForwardLinearScaleAdd( + linear=window[0], + scale=window[1], + add=window[2], + ) + window = [op] + else: + # Shift window if window doesn't match pattern + out.extend(window[:-2]) + window = window[-2:] + + # Adjust window to expected size + out.extend(window[:-3]) + window = window[-3:] + while ops and len(window) < 3: + window.append(ops[0]) + ops = ops[1:] + + # Return list of ops out.extend(window) - - # Check if first op is linear - window, ops = ops[:1], ops[1:] - op, _ = window[0] - if not isinstance(op, BasicLinear): - continue - if op.tensor_parallel_mode == "row": - # Row tensor-parallelism requires communication after the - # GEMM - continue - linear = op - op, _ = ops[0] - - # Check if next op is constant scale - if not isinstance(op, ConstantScale): - continue - scale = op - window.extend(ops[:1]) - ops = ops[1:] - op, _ = ops[0] - - # Check if next op is in-place add extra input - if not isinstance(op, AddExtraInput): - continue - if not op._in_place: - continue - add = op - window.extend(ops[:1]) - ops = ops[1:] - - # Replace window with fused op - op = ForwardLinearScaleAdd( - linear=linear, - scale=scale, - add=add, - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out.extend(window) - out.extend(ops) - return out + return out diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 6c889ba047..90ade030c8 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -503,7 +503,7 @@ def fuser_backward( # Get basic operations idx = self._op_idxs["linear"] linear_op = self.basic_ops[idx] - linear_op_ctx = basic_op_ctxs[-1] + linear_op_ctx = basic_op_ctxs[0] bias_op = None if self._op_idxs["bias"] is not None: idx = self._op_idxs["bias"] @@ -578,99 +578,84 @@ def fuser_backward( grad_params[self._op_idxs["linear"]] = (grad_weight,) if bias_op is not None: grad_params[self._op_idxs["bias"]] = (grad_bias,) - grad_params.reverse() grad_extra_inputs = [() for _ in range(len(self.basic_ops))] return grad_input, grad_params, grad_extra_inputs + @staticmethod + def fuse_backward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for backward pass. -def fuse_userbuffers_backward_linear( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Substitute linear operations with Userbuffers implementation + Parameters + ---------- + ops : list of FusibleOperation + Backward pass operations. + recipe : Recipe, optional + Quantization recipe. - Parameters - ---------- - ops : list of tuples - Backward pass operations and the indices of the corresponding - basic operations. + Returns + ------- + ops : list of FusibleOperation + Updated backward pass operations - Returns - ------- - ops : list of tuples - Updated backward pass operations + """ - """ + # Return immediately if environment is not distributed + if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1: + return ops - # Return immediately if environment is not distributed - if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1: - return ops - - # Sliding window in list of ops - window = [] - - def peek_next_op() -> Optional[FusibleOperation]: - """Get next op in list of ops""" - nonlocal ops - if not ops: - return None - return ops[-1][0] - - def pop_next_op() -> FusibleOperation: - """Remove next op from list of ops and add to sliding window""" - nonlocal ops, window - window.insert(0, ops[-1]) - ops = ops[:-1] - return window[0][0] - - # Scan through ops in reverse order, fusing if possible - out_reversed = [] - while ops: - out_reversed.extend(reversed(window)) - window.clear() - - # Check if next op is linear - next_op = pop_next_op() - if not isinstance(next_op, BasicLinear): - continue - linear = next_op - if linear._userbuffers_options is None: - continue - - # Check if next op is bias - bias = None - if linear.tensor_parallel_mode != "row" and isinstance(peek_next_op(), Bias): - bias = pop_next_op() - - # Check if next op is reduce-scatter - reduce_scatter = None - if linear.tensor_parallel_mode is None and isinstance(peek_next_op(), ReduceScatter): - reduce_scatter = pop_next_op() - - # Check for invalid combinations - if reduce_scatter is None: - if linear.tensor_parallel_mode is None: - continue - if linear.tensor_parallel_size == 1: - continue - if linear.tensor_parallel_mode == "row" and bias is not None: - continue - else: - if linear.tensor_parallel_mode is not None: + # Scan through ops, fusing if possible + out = [] + window = [] + while ops: + + # Shift window + out.extend(window) + window, ops = ops[:1], ops[1:] + + # Check if first op is linear + if not isinstance(window[0], BasicLinear): continue - if reduce_scatter.process_group_size == 1: + linear = window[0] + if linear._userbuffers_options is None: continue - # Replace window with fused op - op = UserbuffersBackwardLinear( - linear=linear, - bias=bias, - reduce_scatter=reduce_scatter, - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] - - # Return list of ops - out_reversed.extend(reversed(window)) - out = out_reversed - out.reverse() - return out + # Check if next op is bias + bias = None + if linear.tensor_parallel_mode != "row" and ops and isinstance(ops[0], Bias): + bias, ops = ops[0], ops[1:] + window.append(bias) + + # Check if next op is reduce-scatter + reduce_scatter = None + if linear.tensor_parallel_mode is None and ops and isinstance(ops[0], ReduceScatter): + reduce_scatter, ops = ops[0], ops[1:] + window.append(reduce_scatter) + + # Check for invalid combinations + if reduce_scatter is None: + if linear.tensor_parallel_mode is None: + continue + if linear.tensor_parallel_size == 1: + continue + if linear.tensor_parallel_mode == "row" and bias is not None: + continue + else: + if linear.tensor_parallel_mode is not None: + continue + if reduce_scatter.process_group_size == 1: + continue + + # Replace window with fused op + op = UserbuffersBackwardLinear( + linear=linear, + bias=bias, + reduce_scatter=reduce_scatter, + ) + window = [op] + + # Return list of ops + out.extend(window) + return out diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index fe04aa1e0b..6ef9bf083b 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -369,93 +369,79 @@ def fuser_forward( return output, [() for _ in range(len(self.basic_ops))] + @staticmethod + def fuse_forward_ops( + ops: list[FusibleOperation], + **unused, # pylint: disable=unused-argument + ) -> list[FusibleOperation]: + """Apply operation fusion for forward pass. -def fuse_userbuffers_forward_linear( - ops: list[tuple[FusibleOperation, list[int]]], -) -> list[tuple[FusibleOperation, list[int]]]: - """Substitute linear operations with Userbuffers implementation - - Parameters - ---------- - ops : list of tuples - Forward pass operations and the indices of the corresponding - basic operations. - - Returns - ------- - ops : list of tuples - Updated forward pass operations + Parameters + ---------- + ops : list of FusibleOperation + Forward pass operations. - """ + Returns + ------- + ops : list of FusibleOperation + Updated forward pass operations - # Return immediately if environment is not distributed - if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1: - return ops - - # Sliding window in list of ops - window = [] - - def peek_next_op() -> Optional[FusibleOperation]: - """Get next op in list of ops""" - nonlocal ops - if not ops: - return None - return ops[0][0] - - def pop_next_op() -> FusibleOperation: - """Remove next op from list of ops and add to sliding window""" - nonlocal ops, window - window.append(ops[0]) - ops = ops[1:] - return window[-1][0] - - # Scan through ops, fusing if possible - out = [] - while ops: - out.extend(window) - window.clear() + """ - # Check if next op is linear - next_op = pop_next_op() - if not isinstance(next_op, BasicLinear): - continue - linear = next_op - if linear._userbuffers_options is None: - continue + # Return immediately if environment is not distributed + if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1: + return ops - # Check if next op is bias - bias = None - if linear.tensor_parallel_mode != "row" and isinstance(peek_next_op(), Bias): - bias = pop_next_op() + # Scan through ops, fusing if possible + out = [] + window = [] + while ops: - # Check if next op is reduce-scatter - reduce_scatter = None - if linear.tensor_parallel_mode is None and isinstance(peek_next_op(), ReduceScatter): - reduce_scatter = pop_next_op() + # Shift window + out.extend(window) + window, ops = ops[:1], ops[1:] - # Check for invalid combinations - if reduce_scatter is None: - if linear.tensor_parallel_mode is None: - continue - if linear.tensor_parallel_size == 1: - continue - if linear.tensor_parallel_mode == "row" and bias is not None: - continue - else: - if linear.tensor_parallel_mode is not None: + # Check if first op is linear + if not isinstance(window[0], BasicLinear): continue - if reduce_scatter.process_group_size == 1: + linear = window[0] + if linear._userbuffers_options is None: continue - # Replace window with fused op - op = UserbuffersForwardLinear( - linear=linear, - bias=bias, - reduce_scatter=reduce_scatter, - ) - basic_op_idxs = [basic_op_idxs[0] for _, basic_op_idxs in window] - window = [(op, basic_op_idxs)] + # Check if next op is bias + bias = None + if linear.tensor_parallel_mode != "row" and ops and isinstance(ops[0], Bias): + bias, ops = ops[0], ops[1:] + window.append(bias) + + # Check if next op is reduce-scatter + reduce_scatter = None + if linear.tensor_parallel_mode is None and ops and isinstance(ops[0], ReduceScatter): + reduce_scatter, ops = ops[0], ops[1:] + window.append(reduce_scatter) + + # Check for invalid combinations + if reduce_scatter is None: + if linear.tensor_parallel_mode is None: + continue + if linear.tensor_parallel_size == 1: + continue + if linear.tensor_parallel_mode == "row" and bias is not None: + continue + else: + if linear.tensor_parallel_mode is not None: + continue + if reduce_scatter.process_group_size == 1: + continue + + # Replace window with fused op + op = UserbuffersForwardLinear( + linear=linear, + bias=bias, + reduce_scatter=reduce_scatter, + ) + window = [op] - # Return list of ops - out.extend(window) - return out + # Return list of ops + out.extend(window) + return out diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index bf7af48d03..7fe6ea37ed 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -5,33 +5,20 @@ """Manager class for a pipeline of fusible operations.""" from __future__ import annotations -from collections.abc import Callable, Iterable -from typing import Any, Optional +from collections.abc import Callable, Iterable, Sequence import itertools +from typing import Any, Optional, TypeAlias import torch -from transformer_engine.pytorch.quantization import FP8GlobalStateManager, Recipe, DelayedScaling -from transformer_engine.pytorch.ops.op import ( +from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling +from ..quantized_tensor import prepare_for_saving, restore_from_saved +from .op import ( BasicOperation, FusibleOperation, + FusedOperation, OperationContext, ) -from transformer_engine.pytorch.ops.fused import ( - fuse_backward_activation_bias, - fuse_backward_add_rmsnorm, - fuse_backward_linear_add, - fuse_backward_linear_scale, - fuse_forward_linear_bias_activation, - fuse_forward_linear_bias_add, - fuse_forward_linear_scale_add, - fuse_userbuffers_backward_linear, - fuse_userbuffers_forward_linear, -) -from transformer_engine.pytorch.quantized_tensor import ( - prepare_for_saving, - restore_from_saved, -) def _split_tuple(t: tuple, idx: int) -> tuple[tuple, tuple]: @@ -57,6 +44,12 @@ def _is_graph_capturing() -> bool: return _is_graph_capturing_function() +# Type alias for a function that may perform operation fusion +OperationFusionFunction: TypeAlias = ( + "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]" +) + + class _OperationFuserAutogradFunction(torch.autograd.Function): """Autograd function for a pipeline of operations @@ -241,7 +234,7 @@ def backward( dx = grad_output grad_params = [None for _ in range(len(basic_ops))] grad_extra_inputs = [None for _ in range(len(basic_ops))] - for op, basic_op_idxs in backward_ops: + for op, basic_op_idxs in reversed(backward_ops): # Stop if no more gradients are required if all(not basic_op_ctxs[idx].requires_grad for idx in basic_op_idxs): @@ -315,6 +308,10 @@ class OperationFuser: """ + # Functions to perform operation fusion + forward_fusion_functions: list[OperationFusionFunction] = [] + backward_fusion_functions: list[OperationFusionFunction] = [] + def __init__( self, ops: list[FusibleOperation], @@ -334,7 +331,7 @@ def __init__( self._basic_op_num_extra_inputs: list[int] = list(op.num_extra_inputs for op in basic_ops) self.num_extra_inputs: int = sum(self._basic_op_num_extra_inputs) - # Ops for forward and backward pass, will be populated in fuse_ops + # Ops for forward and backward pass, will be populated in maybe_fuse_ops self._forward_ops: list[tuple[FusibleOperation, list[int]]] self._backward_ops: list[tuple[FusibleOperation, list[int]]] @@ -349,31 +346,48 @@ def __init__( self._flat_basic_op_params = sum(self._basic_op_params, []) @classmethod - def _fuse_forward_ops( - cls, - ops: list[tuple[FusibleOperation, list[int]]], - recipe: Optional[Recipe], # pylint: disable=unused-argument - ) -> list[tuple[FusibleOperation, list[int]]]: - """Attempt to fuse operations in forward pass""" - ops = fuse_userbuffers_forward_linear(ops) - ops = fuse_forward_linear_bias_add(ops) - ops = fuse_forward_linear_bias_activation(ops) - ops = fuse_forward_linear_scale_add(ops) - return ops - - @classmethod - def _fuse_backward_ops( + def _fuse_ops( cls, - ops: list[tuple[FusibleOperation, list[int]]], + basic_ops: Sequence[BasicOperation], + fusion_funcs: Iterable[OperationFusionFunction], recipe: Optional[Recipe], ) -> list[tuple[FusibleOperation, list[int]]]: - """Attempt to fuse operations in backward pass""" - ops = fuse_userbuffers_backward_linear(ops) - ops = fuse_backward_linear_add(ops) - ops = fuse_backward_linear_scale(ops) - ops = fuse_backward_activation_bias(ops, recipe) - ops = fuse_backward_add_rmsnorm(ops) - return ops + """Apply operation fusions""" + + # Apply op fusions + fused_ops = list(basic_ops) + for func in fusion_funcs: + fused_ops = func(fused_ops, recipe=recipe) + + def raise_mismatch_error() -> None: + """Throw error indicating invalid op fusion""" + raise RuntimeError( + "Found mismatch after fusing operations " + f"(basic_ops={[o.__class__.__name__ for o in basic_ops]}, " + f"fused_ops={[o.__class__.__name__ for o in fused_ops]})" + ) + + # Determine basic op indices corresponding to each op + out = [] + idx = 0 + for op in fused_ops: + if isinstance(op, FusedOperation): + idxs = [] + for basic_op in op.basic_ops: + if basic_op is not basic_ops[idx]: + raise_mismatch_error() + idxs.append(idx) + idx += 1 + out.append((op, idxs)) + else: + if op is not basic_ops[idx]: + raise_mismatch_error() + out.append((op, [idx])) + idx += 1 + if idx != len(basic_ops): + raise_mismatch_error() + + return out def maybe_fuse_ops( self, @@ -424,12 +438,16 @@ def maybe_fuse_ops( op.pre_first_fuser_forward() # Prepare basic op lists for fusions - forward_ops = [(op, [idx]) for idx, op in enumerate(self._basic_ops)] - backward_ops = list(reversed(forward_ops[first_op_requiring_backward:])) - - # Fuse ops - self._forward_ops = self._fuse_forward_ops(forward_ops, recipe) - self._backward_ops = self._fuse_backward_ops(backward_ops, recipe) + self._forward_ops = OperationFuser._fuse_ops( + self._basic_ops, + OperationFuser.forward_fusion_functions, + recipe=recipe, + ) + self._backward_ops = OperationFuser._fuse_ops( + self._basic_ops, + OperationFuser.backward_fusion_functions, + recipe=recipe, + ) # Save current fusion params self.recipe_type, self.first_op_requiring_backward = fusion_params @@ -491,3 +509,55 @@ def __call__( *extra_inputs, ) return forward_func(*args) + + +def register_forward_fusion( + op_fusion_func: OperationFusionFunction, + prepend: bool = False, +) -> None: + """Register function to perform operation fusion for forward pass. + + The fusion function should have the following signature: + + func(ops, *, recipe) -> updated ops + + Parameters + ---------- + op_fusion_func: function + Function that takes a list of operations and may substitute + them with fused operations. + prepend: bool, default = ``False`` + Whether the operation fuser should apply this fusion function + first. The default is to apply it last. + + """ + if prepend: + OperationFuser.forward_fusion_functions.insert(0, op_fusion_func) + else: + OperationFuser.forward_fusion_functions.append(op_fusion_func) + + +def register_backward_fusion( + op_fusion_func: OperationFusionFunction, + prepend: bool = False, +) -> None: + """Register function to perform operation fusion for backward pass. + + The fusion function should have the following signature: + + func(ops, *, recipe) -> updated ops + + Parameters + ---------- + op_fusion_func: function + Function that takes a list of operations and may substitute + them with fused operations. + prepend: bool, default = ``False`` + Whether the operation fuser should apply this fusion function + first. The default is to apply it last. + + """ + if prepend: + OperationFuser.backward_fusion_functions.insert(0, op_fusion_func) + else: + OperationFuser.backward_fusion_functions.append(op_fusion_func) From 2dbfbc743ea20fda3549f7c6019f9b99cc672455 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Mon, 26 Jan 2026 12:39:32 -0600 Subject: [PATCH 182/521] fix(examples): te_llama compatibility with transformers >= 4.57 (#2572) * fix(examples): te_llama compatibility with HuggingFace transformers >= 4.57 The te_llama.py example was failing with HuggingFace transformers 4.57+ due to API changes in how decoder layer outputs are handled. Changes: - Handle case where hidden_states is passed as a tuple (older HF versions) - Return tensor directly instead of wrapped in tuple (HF 4.57+ expects this) - Fix regex pattern to use raw string (fixes SyntaxWarning) Error fixed: AttributeError: 'tuple' object has no attribute 'contiguous' Tested with: - transformer_engine 2.5.0 - transformers 4.57.3 - PyTorch container nvcr.io/nvidia/pytorch:25.08-py3 Signed-off-by: Santosh Bhavani * docs(te_llama): add requirements.txt Signed-off-by: Santosh Bhavani * fix(docs): add missing notebook output names Signed-off-by: Santosh Bhavani --------- Signed-off-by: Santosh Bhavani --- docs/examples/te_llama/requirements.txt | 5 + docs/examples/te_llama/te_llama.py | 15 +- ...tutorial_accelerate_hf_llama_with_te.ipynb | 1535 +++++++++-------- 3 files changed, 793 insertions(+), 762 deletions(-) create mode 100644 docs/examples/te_llama/requirements.txt diff --git a/docs/examples/te_llama/requirements.txt b/docs/examples/te_llama/requirements.txt new file mode 100644 index 0000000000..093849001b --- /dev/null +++ b/docs/examples/te_llama/requirements.txt @@ -0,0 +1,5 @@ +transformers==4.57.0 +accelerate==1.10.0 +peft==0.15.2 +datasets==4.0.0 +sentencepiece==0.2.1 diff --git a/docs/examples/te_llama/te_llama.py b/docs/examples/te_llama/te_llama.py index b2d4d183ab..6dfa9b67bb 100644 --- a/docs/examples/te_llama/te_llama.py +++ b/docs/examples/te_llama/te_llama.py @@ -72,10 +72,15 @@ def forward(self, hidden_states, *args, attention_mask, **kwargs): forward pass of the `TransformerLayer`. Also, make sure the output format matches the output of the HF's `LlamaDecoderLayer`. """ - return ( - super().forward( - hidden_states, attention_mask=attention_mask, rotary_pos_emb=self.te_rope_emb - ), + # Handle case where hidden_states might be a tuple (from previous layer output) + # This can happen with older versions of HuggingFace transformers + if isinstance(hidden_states, tuple): + hidden_states = hidden_states[0] + + # Return tensor directly for HuggingFace transformers >= 4.57 + # (older versions wrapped output in tuple and extracted with layer_outputs[0]) + return super().forward( + hidden_states, attention_mask=attention_mask, rotary_pos_emb=self.te_rope_emb ) @@ -162,7 +167,7 @@ def replace_params(hf_state_dict, te_state_dict, config): # collect all layer prefixes to update all_layer_prefixes = set() for param_key in hf_state_dict.keys(): - layer_prefix_pat = "model.layers.\d+." + layer_prefix_pat = r"model.layers.\d+." m = re.match(layer_prefix_pat, param_key) if m is not None: all_layer_prefixes.add(m.group()) diff --git a/docs/examples/te_llama/tutorial_accelerate_hf_llama_with_te.ipynb b/docs/examples/te_llama/tutorial_accelerate_hf_llama_with_te.ipynb index 00499cff5f..ac9252ff15 100644 --- a/docs/examples/te_llama/tutorial_accelerate_hf_llama_with_te.ipynb +++ b/docs/examples/te_llama/tutorial_accelerate_hf_llama_with_te.ipynb @@ -1,763 +1,784 @@ { - "cells": [ - { - "cell_type": "markdown", - "id": "6a5b2993", - "metadata": {}, - "source": [ - "# Accelerating Hugging Face Llama 2 and 3 Fine-Tuning with Transformer Engine\n", - "\n", - "
\n", - "\n", - "Goal\n", - "\n", - "This tutorial showcases how to accelerate finetuning a full [Llama 2](https://huggingface.co/meta-llama/Llama-2-7b-hf) or [Llama 3](https://huggingface.co/meta-llama/Meta-Llama-3-8B) models from Hugging Face by using `TransformerLayer` from the [Transformer Engine library](https://github.com/NVIDIA/TransformerEngine) in `BF16` and `FP8` precisions.\n", - "\n", - "
\n" - ] - }, - { - "cell_type": "markdown", - "id": "331f476a", - "metadata": {}, - "source": [ - "## Dependencies for this tutorial\n", - "\n", - "Following files and media are necessary to effectively run this tutorial:\n", - "\n", - "1. `te_llama.py`\n", - " - This file contains the code to load a Hugging Face Llama 2 or Llama 3 checkpoint in Transformer Engine's `TransformerLayer` instead of Hugging Face's `LlamaDecoderLayer`. This is used in the following two sections of the tutorial - \"Improvement 1\" and \"Improvement 2\".\n", - "2. `utils.py`\n", - " - This file contains the code related to dataloading, hyperparameters, setting up model/optimizers/accelerator, model training and other miscellaneous tasks like restarting the jupyter notebook from within the cell. \n", - "3. `media/`\n", - " - This directory contains the images used in the following tutorial.\n", - "\n", - "These packages are necessary to run this tutorial:\n", - "`pytorch`, `transformer_engine`, `accelerate`, `transformers`, `peft`, `datasets`.\n", - "\n", - "\n", - "
\n", - "\n", - "Note on running the tutorial with Llama 3 weights\n", - "\n", - "This tutorial shows the cell outputs when run with Llama 2 7B weights. It can be run with Llama 3 8B weights simply by providing the directory with those weights (in Hugging Face format) instead of Llama 2 7B weights. These two models are almost identical, the biggest difference being the model dimension (the smallest Llama 3 model has 8B parameters, whereas the smallest Llama 2 has 7B), which enables this tutorial to work for both of them.\n", - "\n", - "
\n" - ] - }, - { - "cell_type": "markdown", - "id": "44abae4f", - "metadata": {}, - "source": [ - "## Table of contents\n", - "1. From \"Transformer\" to \"Llama\"\n", - "2. Hugging Face's `LlamaModel`\n", - " - Hugging Face's `LlamaDecoderLayer`\n", - "3. [Baseline] Running HF `LlamaModel` (Precision: `BF16`)\n", - "6. [Improvement 1] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `BF16`)\n", - " - Transformer Engine's `TransformerLayer`\n", - " - `TransformerLayer` options explained\n", - " - Mapping weights from HF's `LlamaDecoderLayer` to TE's `TransformerLayer`\n", - "7. [Improvement 2] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `FP8`)\n", - "8. Conclusion" - ] - }, - { - "cell_type": "markdown", - "id": "e37e2cc1", - "metadata": {}, - "source": [ - "## From \"Transformer\" to \"Llama\" \n", - "\n", - "
\n", - "\n", - "
Fig 1: Llama visualized as a transformer. (generated with [Nvidia's AI-foundation models](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/sdxl))
\n", - "
\n", - "\n", - "A flashback:\n", - "\n", - "- 2017: [\"Attention Is All You Need\"](https://arxiv.org/abs/1706.03762) paper introduced pioneering \"Transformer\" architecture and changed the NLP field forever.\n", - "- 2018-2020: Emergence of GPT model series that showed causal decoder architectures are great fit for pretraining, few-shot and zero-shot learning.\n", - "- Fast forward to 2023-2024: Following GPT-3/GPT-4 success stories, researchers and companies raced to produce the next best pretrained model that could further be finetuned for application-specific use-cases.\n", - "- February 2023: Meta releases [Llama 2](https://llama.meta.com/llama2) models (Large Language Model Meta AI). \n", - " - These models range from 7B to 70B parameters.\n", - " - LLaMA 2 was pretrained on 2 trillion tokens.\n", - "- April 2024: Meta releases [Llama 3](https://llama.meta.com/llama3) models.\n", - " - These models range from 8B to 70B parameters.\n", - " - LLaMA 3 was pretrained on 15 trillion tokens.\n", - "\n", - "For more information on Llama 2 consider reading the [Huggingface tutorial](https://huggingface.co/blog/llama2). As a quick summary, here are some of the important differences b/w the conventional transformer decoder architecture vs Llama 2 architecture:\n", - "\n", - "1. Decoder only model (causal language modeling and next word prediction)\n", - "2. RMSNorm in place of the LayerNorm\n", - "3. SwiGLU activation function\n", - "4. RoPE as positional embeddings \n", - "5. Grouped Query Attention for the 70B model\n", - "6. Trained on 4K context length\n", - "\n", - "Hugging Face also released a [tutorial about Llama 3](https://huggingface.co/blog/llama3). The key points are:\n", - "\n", - "1. Use of bigger tokenizer - 128256 vs 32K.\n", - "2. Grouped Query Attention is used also by smaller 8B model.\n", - "3. The context length increased to 8K for all models.\n", - "3. Llama 3 was trained on 8x more data than Llama 2.\n", - "\n", - "
\n", - "\n", - "
Fig 2: Comparing GPT and Llama architectures.
\n", - "
" - ] - }, - { - "cell_type": "markdown", - "id": "a110de1a", - "metadata": {}, - "source": [ - "## Hugging Face's `LlamaModel`\n", - "Hugging Face provides an open-source implementation of `Llama` model in [modeling_llama.py](https://github.com/huggingface/transformers/blob/3d2900e829ab16757632f9dde891f1947cfc4be0/src/transformers/models/llama/modeling_llama.py#L4).\n", - "\n", - "Here's a block diagram that shows how Llama model is implemented in the Hugging Face repo. Notice the modular encapsulated form and `LlamaDecoderLayer` at the core of the model implementation.\n", - "\n", - "
\n", - "\n", - "
Fig 3: Causal Llama Model Block Diagram.
\n", - "
\n", - "\n", - "The above diagram translates to the following text output of the model in PyTorch. Notice that the core of the model has 32 `LlamaDecoderLayer`s. \n", - "\n", - "```\n", - "LlamaForCausalLM(\n", - " (model): LlamaModel(\n", - " (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n", - " (layers): ModuleList(\n", - " (0-31): 32 x LlamaDecoderLayer(\n", - " (self_attn): LlamaFlashAttention2(\n", - " (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (rotary_emb): LlamaRotaryEmbedding()\n", - " )\n", - " (mlp): LlamaMLP(\n", - " (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", - " (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", - " (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n", - " (act_fn): SiLU()\n", - " )\n", - " (input_layernorm): LlamaRMSNorm()\n", - " (post_attention_layernorm): LlamaRMSNorm()\n", - " )\n", - " )\n", - " (norm): LlamaRMSNorm()\n", - " )\n", - " (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n", - ")\n", - "```\n", - "\n", - "#### Hugging Face's `LlamaDecoderLayer`\n", - "\n", - "Let's take a closer look at `LlamaDecoderLayer`. It is composed of `input_layernorm`, `self_attn`, `post_attention_layernorm` and `mlp` modules. Each module has associated weights as shown in the diagram.\n", - "\n", - "
\n", - "\n", - "
Fig 4: Causal Llama Model Block Diagram (with simplified illustration of the [LlamaDecoderLayer](https://github.com/huggingface/transformers/blob/e770f0316d2a9b787c9d1440f204fcb65e176682/src/transformers/models/llama/modeling_llama.py#L695)).
\n", - "
\n", - "\n", - "##### Self_Attn Layer\n", - "For simplicity in the block diagram illustration of the \"self_attn\" box, we omit the \"Grouped Query Attention\" operation and only showcase the modules which have associated weights.\n", - " \n", - "##### MLP Layer\n", - "\n", - "SwiGLU is an activation defined as follows in the [modeling_llama.py](https://github.com/huggingface/transformers/blob/7c4995f93d8d24aae05e1e43279c96dce736e5c8/src/transformers/models/llama/modeling_llama.py#L236) file in the Hugging Face github repo:\n", - "```\n", - "\"\"\"\n", - "1. `self.up_proj`, `self.gate_proj` and `self.down_proj` are \"Linear\" layers\n", - "2. `self.act_fn` is a \"Swish\" function\n", - "\n", - "\"\"\"\n", - "down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))\n", - "```\n", - "It requires a set of 3 weights as compared to 2 weights in conventional \"MLP\" layers e.g. in the traditional transformer or GPT architectures. This is also illustrated in the following figure:\n", - "\n", - "
\n", - "\n", - "
Fig 5: A look inside the feedforward layer with swiglu activation function.
\n", - "
" - ] - }, - { - "cell_type": "markdown", - "id": "c9529229", - "metadata": {}, - "source": [ - "## [Baseline] Running HF `LlamaModel` (Precision: `BF16`)\n", - "\n", - "Llama 2 weights are loaded into the Hugging Face native implementation `LlamaForCausalLM` (refer to [modeling_llama.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py)). \n", - "\n", - "For this and other subsequent runs, the `batch_size` is `8`. The `LlamaDecoderLayer` is left unchanged in the baseline as follows:\n", - "\n", - "
\n", - "\n", - "
Fig 6: Revisiting \"LlamaDecoderLayer\".
\n", - "
\n", - "\n", - "
\n", - "Note\n", - "\n", - "The baseline implementation will be run in `BF16` precision.\n", - "\n", - "
" - ] - }, - { - "cell_type": "markdown", - "id": "b38eb3ac", - "metadata": {}, - "source": [ - "
\n", - "\n", - "Note\n", - " \n", - "This tutorial loads and trains a Llama 3 8B or a Llama 2 7B model which takes up most of the GPU memory and therefore, we need to restart the jupyter notebook each time before running the following sections. A small utility method `restart_jupyter_notebook` is defined in the accompanying `utils.py` file. This function restarts the jupyter notebook so that the GPU memory is flushed before the model is loaded again from the checkpoint in order to avoid running into OOM (Out Of Memory) errors.\n", - "\n", - "If the utility doesn't work, comment this line `restart_jupyter_notebook()` in the following cell and manually restart the jupyter notebook before running the cell. Repeat the same for other sections in this tutorial.\n", - "\n", - "
\n" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "2e9d7a8c", - "metadata": {}, - "outputs": [ + "cells": [ { - "name": "stdout", - "output_type": "stream", - "text": [ - "10 finetuning steps complete!\n", - "Average time taken per step: 248 milliseconds\n" - ] - } - ], - "source": [ - "# Restart the notebook (to flush the GPU memory)\n", - "from utils import restart_jupyter_notebook\n", - "restart_jupyter_notebook()\n", - "\n", - "\n", - "# Import necessary packages, methods and variables\n", - "from utils import *\n", - "\n", - "\n", - "# Provide Huggingface Access Token\n", - "hyperparams.hf_access_token = \"\"\n", - "assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n", - "\n", - "# Provide a directory to cache weights in to avoid downloading them every time.\n", - "# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n", - "hyperparams.weights_cache_dir = \"\"\n", - "\n", - "# For Llama 2, uncomment this line (also set by default)\n", - "hyperparams.model_name = \"meta-llama/Llama-2-7b-hf\"\n", - "\n", - "# For Llama 3, uncomment this line\n", - "# hyperparams.model_name = \"meta-llama/Meta-Llama-3-8B\"\n", - "\n", - "hyperparams.mixed_precision = \"bf16\"\n", - "\n", - "\n", - "# Init the model and accelerator wrapper\n", - "model = init_baseline_model(hyperparams)\n", - "accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)\n", - "\n", - "\n", - "# Finetune the model\n", - "finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)" - ] - }, - { - "cell_type": "markdown", - "id": "4035ccb7", - "metadata": {}, - "source": [ - "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 (or ms per batch) | Speedup (over baseline) |\n", - "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", - "| HF (baseline) | BF16 | 248 | 1 |" - ] - }, - { - "cell_type": "markdown", - "id": "3db90dff", - "metadata": {}, - "source": [ - "## [Improvement 1] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `BF16`)\n", - "\n", - "In addition to basic layers like `Linear` and `LayerNorm`, Transformer Engine offers larger modules like `MultiheadAttention` (combines \"LayerNorm\" and \"Self Attention\") and `LayerNormMLP` (combines \"LayerNorm\" and \"MLP\") that could replace their counterparts in the `LlamaDecoderLayer` and potentially provide a speedup. Transformer Engine also offers a full `TransformerLayer` (which further combines `MultiheadAttention` and `LayerNormMLP` layers) which could replace `LlamaDecoderLayer` and provide a speedup (with careful mapping of the weights since the name of the weights are different for those two layers). Let's take a closer look at Transformer Engine's `TransformerLayer`. \n", - "\n", - "#### Transformer Engine's `TransformerLayer`\n", - "\n", - "At a higher level, TE's `TransformerLayer` could be visualized as an apt replacement for the `LlamaDecoderLayer`. But the internals of the `TransformerLayer` are organized a bit differently. \n", - "\n", - "
\n", - "\n", - "
Fig 7: Transformer Engine's `TransformerLayer`
\n", - "
\n", - "\n", - "Just like Hugging Face's `LlamaDecoderLayer`, Transformer Engine's `TransformerLayer` encapsulates `self_attention` (as `MultiheadAttention`) and `mlp` (as `LayerNormMLP`). A major difference is that the two `Norm`s are included in the `MultiheadAttention` and `LayerNormMLP` layers as shown in the following output prompt:\n", - "\n", - "```\n", - "TransformerLayer(\n", - " (self_attention): MultiheadAttention(\n", - " (layernorm_qkv): LayerNormLinear()\n", - " (core_attention): DotProductAttention()\n", - " (proj): Linear()\n", - " )\n", - " (layernorm_mlp): LayerNormMLP()\n", - ")\n", - "```\n", - "\n", - "Another difference is that Transformer Engine implements an efficient version of feedforward layer with SwiGLU in which the weights from the `up_proj` and `gate_proj` modules are merged together and SwiGLU is applied using a custom fused kernel. This is done so that only one big and efficient Matrix Multiplication operation is issued to the GPU instead of two smaller ones.\n", - "\n", - "
\n", - "\n", - "
Fig 8: Abstract illustration of the SwiGLU implementation in Transformer Engine.
\n", - "
\n", - "\n", - "#### `TransformerLayer` options explained\n", - "\n", - "
\n", - "\n", - "Note\n", - " \n", - "Here, we go over some of the options in `TransformerLayer` that are needed for the tutorial. For a complete list of options, refer the [TransformerLayer API documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html?highlight=transformerlayer#transformer_engine.pytorch.TransformerLayer).\n", - "\n", - "
\n", - "\n", - "In the accompanying `te_llama.py` file, `TELlamaDecoderLayer` is defined as a wrapper over TE's `TransformerLayer` with a few needed options that make `TransformerLayer` a plug-in replacement for the HF's `LlamaDecoderLayer`.\n", - "\n", - "```\n", - "class TELlamaDecoderLayer(te.pytorch.TransformerLayer):\n", - " def __init__(self, config):\n", - " super().__init__(\n", - " config.hidden_size,\n", - " config.intermediate_size,\n", - " config.num_attention_heads,\n", - " bias=False,\n", - " layernorm_epsilon=config.rms_norm_eps,\n", - " hidden_dropout=0,\n", - " attention_dropout=0,\n", - " fuse_qkv_params=False,\n", - " normalization=\"RMSNorm\",\n", - " activation=\"swiglu\",\n", - " attn_input_format=\"bshd\",\n", - " num_gqa_groups=config.num_key_value_heads,\n", - " )\n", - " te_rope = RotaryPositionEmbedding(config.hidden_size//config.num_attention_heads)\n", - " self.te_rope_emb = te_rope(max_seq_len=config.max_position_embeddings).cuda()\n", - "```\n", - "\n", - "Here's a list summarizing each option briefly:\n", - "\n", - "1. `hidden_size`: size of each input sample.\n", - "2. `ffn_hidden_size`: intermediate size to which samples are projected.\n", - "3. `num_attention_heads`: number of attention heads in the transformer layer.\n", - "4. `bias`: switch to add additive biases to the submodule layers.\n", - "5. `layernorm_epsilon`: a value added to the denominator of layer normalization for numerical stability. Default is `1e-5`.\n", - "6. `hidden_dropout`: dropout probability for the dropout op after FC2 layer (fully connected layer no. 2). Default is `0.1`.\n", - "7. `attention_dropout`: dropout probability for the dropout op during multi-head attention. Default is `0.1`. \n", - "8. `fuse_qkv_params`: if set to True, TransformerLayer module exposes a single fused parameter for query-key-value. This enables optimizations such as QKV fusion without concatentations/splits and also enables the argument fuse_wgrad_accumulation.\n", - "9. `normalization`: type of normalization applied. Default is `LayerNorm`.\n", - "10. `activation`: type of activation used in the MLP block. Default is `gelu`.\n", - "11. `attn_input_format`: controls whether the dimensions of the intermediate hidden states is 'batch first' ('bshd') or 'sequence first' ('sbhd'). `s` stands for the sequence length, `b` batch size, `h` the number of heads, `d` head size. Note that these formats are very closely related to the `qkv_format` in the `MultiHeadAttention` and `DotProductAttention` modules.\n", - "12. `num_gqa_groups`: number of GQA groups in the transformer layer. Grouped Query Attention is described in [this paper](https://arxiv.org/pdf/2305.13245.pdf). This only affects the keys and values, not the querys. GQA-1 is equivalent to Multi-Query Attention ([MQA](https://arxiv.org/pdf/1911.02150.pdf)), while GQA-H is equivalent to MultiHead Attention, i.e. `num_gqa_groups = num_attention_heads`.\n", - "\n", - "\n", - "Further, note that `RotaryPositionEmbedding` is defined as part of the `TELlamaDecoderLayer` (wrapper around TE's `TransformerLayer`) itself since it expects this rope cache if RoPE is used in the model. \n", - "\n", - "Let's revisit how `LlamaDecoderLayer`s form the core of the decoder layer stack in HF's llama implementation:\n", - "```\n", - "ModuleList(\n", - " (0-31): 32 x LlamaDecoderLayer(\n", - " (self_attn): LlamaAttention(\n", - " (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", - " (rotary_emb): LlamaRotaryEmbedding()\n", - " )\n", - " (mlp): LlamaMLP(\n", - " (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", - " (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", - " (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n", - " (act_fn): SiLU()\n", - " )\n", - " (input_layernorm): LlamaRMSNorm()\n", - " (post_attention_layernorm): LlamaRMSNorm()\n", - " )\n", - ")\n", - "```\n", - "\n", - "A major portion of the Hugging Face model implementation (32 `LlamaDecoderLayer` layers) could be potentially replaced with Transformer Engine's `TransformerLayer` layers. Let's see how it is made possible.\n", - "\n", - "\n", - "#### Mapping weights from HF's `LlamaDecoderLayer` to TE's `TransformerLayer`\n", - "\n", - "Refer the accompanying file `te_llama.py` which provides a reference to create a Llama 2 model with TE's `TransformerLayer` after replacing HF's `LlamaDecoderLayer`.\n", - "\n", - "Briefly, following pieces of code are put together:\n", - "\n", - "1. `TELlamaDecoderLayer` is added as a wrapper for `TransformerLayer`. \n", - "```\n", - "class TELlamaDecoderLayer(te.pytorch.TransformerLayer):\n", - " \"\"\"\n", - " Wrapper class over TE's `TransformerLayer`. This makes the wrapper very\n", - " similar to HF's `LlamaDecoderLayer` and easier to replace it in the code.\n", - "\n", - " Args:\n", - " config: LlamaConfig\n", - " args: positional args (for compatibility with `LlamaDecoderLayer`)\n", - " kwargs: keyword args (for compatibility with `LlamaDecoderLayer`)\n", - " \"\"\"\n", - " def __init__(self, config, *args, **kwargs):\n", - " super().__init__(\n", - " hidden_size=config.hidden_size,\n", - " ffn_hidden_size=config.intermediate_size,\n", - " num_attention_heads=config.num_attention_heads,\n", - " bias=False,\n", - " layernorm_epsilon=config.rms_norm_eps,\n", - " hidden_dropout=0,\n", - " attention_dropout=0,\n", - " fuse_qkv_params=False,\n", - " normalization=\"RMSNorm\",\n", - " activation=\"swiglu\",\n", - " attn_input_format=\"bshd\",\n", - " )\n", - " te_rope = RotaryPositionEmbedding(config.hidden_size//config.num_attention_heads)\n", - " self.te_rope_emb = te_rope(max_seq_len=config.max_position_embeddings).cuda()\n", - "\n", - " def forward(self,\n", - " hidden_states,\n", - " *args,\n", - " attention_mask,\n", - " **kwargs):\n", - " \"\"\"\n", - " Custom forward to make sure we only pass relevant arguments to the\n", - " forward pass of the `TransformerLayer`. Also, make sure the output\n", - " format matches the output of the HF's `LlamaDecoderLayer`.\n", - " \"\"\"\n", - " return (super().forward(hidden_states, attention_mask=attention_mask, rotary_pos_emb=self.te_rope_emb),)\n", - "```\n", - "\n", - "2. Before creating a `LlamaForCausalLM`, `replace_decoder` context manager is used to monkey-patch `LlamaDecoderLayer` with `TELlamaDecoderLayer`.\n", - "\n", - "```\n", - "@contextmanager\n", - "def replace_decoder(te_decoder_cls):\n", - " \"\"\"\n", - " Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.\n", - " \"\"\"\n", - " original_llama_decoder_cls = transformers.models.llama.modeling_llama.LlamaDecoderLayer\n", - " transformers.models.llama.modeling_llama.LlamaDecoderLayer = te_decoder_cls\n", - " try:\n", - " yield\n", - " finally:\n", - " transformers.models.llama.modeling_llama.LlamaDecoderLayer = original_llama_decoder_cls\n", - ".\n", - ".\n", - ".\n", - "class TELlamaForCausalLM:\n", - " \"\"\"\n", - " Causal LM created with `LlamaModel`. The underlying `LlamaDecoderLayer`\n", - " class is monkey-patched with `TELlamaDecoderLayer` class before\n", - " initializing the causal LM with `LlamaForCausalLM`.\n", - "\n", - " Args:\n", - " config: LlamaConfig\n", - " \"\"\"\n", - "\n", - " def __new__(cls, config: LlamaConfig):\n", - " with replace_decoder(te_decoder_cls=TELlamaDecoderLayer):\n", - " llama_for_causal_lm = LlamaForCausalLM(config)\n", - " return llama_for_causal_lm\n", - ".\n", - ".\n", - ".\n", - "```\n", - "\n", - "3. A custom `pretrained_from_local` method is added that copies the weights from the checkpoint (which is meant for HF Llama implementation) to the modified `TELlamaForCausalLM` by carefully mapping the weights from the `LlamaDecoderLayer` (HF) to `TransformerLayer` (TE). The method `replace_params` maps and copies apt weights from `LlamaDecoderLayer` to the `TransformerLayer`. Refer to the following diagram for more details.\n", - "\n", - "```\n", - "def replace_params(hf_state_dict, te_state_dict):\n", - " # collect all layer prefixes to update\n", - " all_layer_prefixes = set()\n", - " for param_key in hf_state_dict.keys():\n", - " layer_prefix_pat = 'model.layers.\\d+.'\n", - " m = re.match(layer_prefix_pat, param_key)\n", - " if m is not None:\n", - " all_layer_prefixes.add(m.group())\n", - "\n", - " for layer_prefix in all_layer_prefixes:\n", - " # When loading weights into models with less number of layers, skip the\n", - " # copy if the corresponding layer doesn't exist in TE model\n", - " if layer_prefix + 'self_attention.layernorm_qkv.layer_norm_weight' in te_state_dict:\n", - " te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.layer_norm_weight'].data[:] = hf_state_dict[layer_prefix + 'input_layernorm.weight'].data[:]\n", - "\n", - " if layer_prefix + 'self_attention.layernorm_qkv.query_weight' in te_state_dict:\n", - " te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.query_weight'].data[:] = hf_state_dict[layer_prefix + 'self_attn.q_proj.weight'].data[:]\n", - "\n", - " if layer_prefix + 'self_attention.layernorm_qkv.key_weight' in te_state_dict:\n", - " te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.key_weight'].data[:] = hf_state_dict[layer_prefix + 'self_attn.k_proj.weight'].data[:]\n", - " .\n", - " .\n", - " .\n", - "\n", - " return all_layer_prefixes\n", - "```\n", - "\n", - "The following figure shows how the weights get mapped from the HF's `LlamaDecoderLayer` to TE's `TransformerLayer`.\n", - "\n", - "
\n", - "\n", - "
Fig 9: Replace `LlamaDecoderLayer` with `TransformerLayer`.
\n", - "
\n", - "\n", - "After initializing the modified Llama model this way, the core decoder layers get changed to `TELlamaDecoderLayer` (wrapper around `TransformerLayer`) as shown in the following output:\n", - "```\n", - "ModuleList(\n", - " (0-31): 32 x TELlamaDecoderLayer(\n", - " (self_attention): MultiheadAttention(\n", - " (layernorm_qkv): LayerNormLinear()\n", - " (core_attention): DotProductAttention(\n", - " (flash_attention): FlashAttention()\n", - " (fused_attention): FusedAttention()\n", - " (unfused_attention): UnfusedDotProductAttention(\n", - " (scale_mask_softmax): FusedScaleMaskSoftmax()\n", - " (attention_dropout): Dropout(p=0, inplace=False)\n", - " )\n", - " )\n", - " (proj): Linear()\n", - " )\n", - " (layernorm_mlp): LayerNormMLP()\n", - " )\n", - ")\n", - "```\n", - "\n", - "In summary, the model gets changed as follows with a large chunk of the implementation (core decoder layers) coming from Transformer Engine.\n", - "\n", - "
\n", - "\n", - "
Fig 10: Language model after the HF's `LlamaDecoderLayer`s are replaced with TE's `TransformerLayer`s.
\n", - "
\n", - "\n", - "\n", - "
\n", - "Note\n", - "\n", - "Let's first run this \"TELlama\" implementation in `BF16` precision.\n", - "
" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "bdb34b91", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Accelerating Hugging Face Llama 2 and 3 Fine-Tuning with Transformer Engine\n", + "\n", + "
\n", + "\n", + "Goal\n", + "\n", + "This tutorial showcases how to accelerate finetuning a full [Llama 2](https://huggingface.co/meta-llama/Llama-2-7b-hf) or [Llama 3](https://huggingface.co/meta-llama/Meta-Llama-3-8B) models from Hugging Face by using `TransformerLayer` from the [Transformer Engine library](https://github.com/NVIDIA/TransformerEngine) in `BF16` and `FP8` precisions.\n", + "\n", + "
\n" + ], + "id": "6a5b2993" + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "10 finetuning steps complete!\n", - "Average time taken per step: 185 milliseconds\n" - ] - } - ], - "source": [ - "# Restart the notebook (to flush the GPU memory)\n", - "from utils import restart_jupyter_notebook\n", - "restart_jupyter_notebook()\n", - "\n", - "\n", - "# Import necessary packages, methods and variables\n", - "from utils import *\n", - "\n", - "\n", - "# Provide Huggingface Access Token\n", - "hyperparams.hf_access_token = \"\"\n", - "assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n", - "\n", - "# Provide a directory to cache weights in to avoid downloading them every time.\n", - "# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n", - "hyperparams.weights_cache_dir = \"\"\n", - "\n", - "# For Llama 2, uncomment this line (also set by default)\n", - "hyperparams.model_name = \"meta-llama/Llama-2-7b-hf\"\n", - "\n", - "# For Llama 3, uncomment this line\n", - "# hyperparams.model_name = \"meta-llama/Meta-Llama-3-8B\"\n", - "\n", - "hyperparams.mixed_precision = \"bf16\"\n", - "\n", - "\n", - "# Init the model and accelerator wrapper\n", - "model = init_te_llama_model(hyperparams)\n", - "accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)\n", - "\n", - "\n", - "# Finetune the model\n", - "finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)" - ] - }, - { - "cell_type": "markdown", - "id": "0c9fbd65", - "metadata": {}, - "source": [ - "Compared to the \"baseline\" implementation, we see that using Transformer Engine's `TransformerLayer` in place of Huggging Face's `LlamaDecoderLayer` gives a speedup of **34%** even when using only BF16 precision!\n", - "\n", - "| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n", - "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", - "| HF (baseline) | BF16 | 248 | 1 |\n", - "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 185 | 1.34 |" - ] - }, - { - "cell_type": "markdown", - "id": "98cd8efb", - "metadata": {}, - "source": [ - "## [Improvement 2] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `FP8`)\n", - "\n", - "Now that most of the HF Llama model implementation (`LlamaDecoderLayer`s) has been swapped with Transformer Engine implementation (`TELlamaDecoderLayer` or `TransformerLayer`), let's see how finetuning in `FP8` precision helps improve performance.\n", - "\n", - "#### How to run the model in `FP8` precision\n", - "\n", - "After the substitution, the model can be run in `FP8` precision by the following change over the previous BF16 runs. (For more information, refer the corresponding `wrap_with_accelerator` function in the accompanying `utils.py` file).\n", - "\n", - "```\n", - "# Specify the `FP8RecipeKwargs` (additional argument required to run in `fp8` precision)\n", - "fp8_kwarg_handler = [FP8RecipeKwargs(backend=\"te\")]\n", - "\n", - "# Pass the `FP8RecipeKwargs` to the `Accelerator` init call\n", - "accelerator = Accelerator(\n", - " ...\n", - " kwargs_handlers=fp8_kwarg_handler\n", - ")\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "772c6f22", - "metadata": {}, - "outputs": [ + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Dependencies for this tutorial\n", + "\n", + "Following files and media are necessary to effectively run this tutorial:\n", + "\n", + "1. `te_llama.py`\n", + " - This file contains the code to load a Hugging Face Llama 2 or Llama 3 checkpoint in Transformer Engine's `TransformerLayer` instead of Hugging Face's `LlamaDecoderLayer`. This is used in the following two sections of the tutorial - \"Improvement 1\" and \"Improvement 2\".\n", + "2. `utils.py`\n", + " - This file contains the code related to dataloading, hyperparameters, setting up model/optimizers/accelerator, model training and other miscellaneous tasks like restarting the jupyter notebook from within the cell. \n", + "3. `requirements.txt`\n", + " - This file contains the necessary Python packages for this tutorial.\n", + "4. `media/`\n", + " - This directory contains the images used in the following tutorial.\n", + "\n", + "\n", + "
\n", + "\n", + "Note on running the tutorial with Llama 3 weights\n", + "\n", + "This tutorial shows the cell outputs when run with Llama 2 7B weights. It can be run with Llama 3 8B weights simply by providing the directory with those weights (in Hugging Face format) instead of Llama 2 7B weights. These two models are almost identical, the biggest difference being the model dimension (the smallest Llama 3 model has 8B parameters, whereas the smallest Llama 2 has 7B), which enables this tutorial to work for both of them.\n", + "\n", + "
\n", + "" + ], + "id": "331f476a" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Setup\n", + "\n", + "Install the required Python packages using the following command:" + ], + "id": "b56526b3" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Uncomment and run this cell when running the tutorial for the first time\n", + "# %pip install -r requirements.txt" + ], + "id": "099697e2", + "execution_count": null, + "outputs": [] + }, { - "name": "stdout", - "output_type": "stream", - "text": [ - "10 finetuning steps complete!\n", - "Average time taken per step: 160 milliseconds\n" - ] + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Table of contents\n", + "1. From \"Transformer\" to \"Llama\"\n", + "2. Hugging Face's `LlamaModel`\n", + " - Hugging Face's `LlamaDecoderLayer`\n", + "3. [Baseline] Running HF `LlamaModel` (Precision: `BF16`)\n", + "6. [Improvement 1] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `BF16`)\n", + " - Transformer Engine's `TransformerLayer`\n", + " - `TransformerLayer` options explained\n", + " - Mapping weights from HF's `LlamaDecoderLayer` to TE's `TransformerLayer`\n", + "7. [Improvement 2] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `FP8`)\n", + "8. Conclusion" + ], + "id": "44abae4f" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## From \"Transformer\" to \"Llama\" \n", + "\n", + "
\n", + "\n", + "
Fig 1: Llama visualized as a transformer. (generated with [Nvidia's AI-foundation models](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/ai-foundation/models/sdxl))
\n", + "
\n", + "\n", + "A flashback:\n", + "\n", + "- 2017: [\"Attention Is All You Need\"](https://arxiv.org/abs/1706.03762) paper introduced pioneering \"Transformer\" architecture and changed the NLP field forever.\n", + "- 2018-2020: Emergence of GPT model series that showed causal decoder architectures are great fit for pretraining, few-shot and zero-shot learning.\n", + "- Fast forward to 2023-2024: Following GPT-3/GPT-4 success stories, researchers and companies raced to produce the next best pretrained model that could further be finetuned for application-specific use-cases.\n", + "- February 2023: Meta releases [Llama 2](https://llama.meta.com/llama2) models (Large Language Model Meta AI). \n", + " - These models range from 7B to 70B parameters.\n", + " - LLaMA 2 was pretrained on 2 trillion tokens.\n", + "- April 2024: Meta releases [Llama 3](https://llama.meta.com/llama3) models.\n", + " - These models range from 8B to 70B parameters.\n", + " - LLaMA 3 was pretrained on 15 trillion tokens.\n", + "\n", + "For more information on Llama 2 consider reading the [Huggingface tutorial](https://huggingface.co/blog/llama2). As a quick summary, here are some of the important differences b/w the conventional transformer decoder architecture vs Llama 2 architecture:\n", + "\n", + "1. Decoder only model (causal language modeling and next word prediction)\n", + "2. RMSNorm in place of the LayerNorm\n", + "3. SwiGLU activation function\n", + "4. RoPE as positional embeddings \n", + "5. Grouped Query Attention for the 70B model\n", + "6. Trained on 4K context length\n", + "\n", + "Hugging Face also released a [tutorial about Llama 3](https://huggingface.co/blog/llama3). The key points are:\n", + "\n", + "1. Use of bigger tokenizer - 128256 vs 32K.\n", + "2. Grouped Query Attention is used also by smaller 8B model.\n", + "3. The context length increased to 8K for all models.\n", + "3. Llama 3 was trained on 8x more data than Llama 2.\n", + "\n", + "
\n", + "\n", + "
Fig 2: Comparing GPT and Llama architectures.
\n", + "
" + ], + "id": "e37e2cc1" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Hugging Face's `LlamaModel`\n", + "Hugging Face provides an open-source implementation of `Llama` model in [modeling_llama.py](https://github.com/huggingface/transformers/blob/3d2900e829ab16757632f9dde891f1947cfc4be0/src/transformers/models/llama/modeling_llama.py#L4).\n", + "\n", + "Here's a block diagram that shows how Llama model is implemented in the Hugging Face repo. Notice the modular encapsulated form and `LlamaDecoderLayer` at the core of the model implementation.\n", + "\n", + "
\n", + "\n", + "
Fig 3: Causal Llama Model Block Diagram.
\n", + "
\n", + "\n", + "The above diagram translates to the following text output of the model in PyTorch. Notice that the core of the model has 32 `LlamaDecoderLayer`s. \n", + "\n", + "```\n", + "LlamaForCausalLM(\n", + " (model): LlamaModel(\n", + " (embed_tokens): Embedding(32000, 4096, padding_idx=0)\n", + " (layers): ModuleList(\n", + " (0-31): 32 x LlamaDecoderLayer(\n", + " (self_attn): LlamaFlashAttention2(\n", + " (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (rotary_emb): LlamaRotaryEmbedding()\n", + " )\n", + " (mlp): LlamaMLP(\n", + " (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", + " (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", + " (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n", + " (act_fn): SiLU()\n", + " )\n", + " (input_layernorm): LlamaRMSNorm()\n", + " (post_attention_layernorm): LlamaRMSNorm()\n", + " )\n", + " )\n", + " (norm): LlamaRMSNorm()\n", + " )\n", + " (lm_head): Linear(in_features=4096, out_features=32000, bias=False)\n", + ")\n", + "```\n", + "\n", + "### Hugging Face's `LlamaDecoderLayer`\n", + "\n", + "Let's take a closer look at `LlamaDecoderLayer`. It is composed of `input_layernorm`, `self_attn`, `post_attention_layernorm` and `mlp` modules. Each module has associated weights as shown in the diagram.\n", + "\n", + "
\n", + "\n", + "
Fig 4: Causal Llama Model Block Diagram (with simplified illustration of the [LlamaDecoderLayer](https://github.com/huggingface/transformers/blob/e770f0316d2a9b787c9d1440f204fcb65e176682/src/transformers/models/llama/modeling_llama.py#L695)).
\n", + "
\n", + "\n", + "#### Self_Attn Layer\n", + "For simplicity in the block diagram illustration of the \"self_attn\" box, we omit the \"Grouped Query Attention\" operation and only showcase the modules which have associated weights.\n", + " \n", + "#### MLP Layer\n", + "\n", + "SwiGLU is an activation defined as follows in the [modeling_llama.py](https://github.com/huggingface/transformers/blob/7c4995f93d8d24aae05e1e43279c96dce736e5c8/src/transformers/models/llama/modeling_llama.py#L236) file in the Hugging Face github repo:\n", + "```\n", + "\"\"\"\n", + "1. `self.up_proj`, `self.gate_proj` and `self.down_proj` are \"Linear\" layers\n", + "2. `self.act_fn` is a \"Swish\" function\n", + "\n", + "\"\"\"\n", + "down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))\n", + "```\n", + "It requires a set of 3 weights as compared to 2 weights in conventional \"MLP\" layers e.g. in the traditional transformer or GPT architectures. This is also illustrated in the following figure:\n", + "\n", + "
\n", + "\n", + "
Fig 5: A look inside the feedforward layer with swiglu activation function.
\n", + "
" + ], + "id": "a110de1a" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Baseline] Running HF `LlamaModel` (Precision: `BF16`)\n", + "\n", + "Llama 2 weights are loaded into the Hugging Face native implementation `LlamaForCausalLM` (refer to [modeling_llama.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py)). \n", + "\n", + "For this and other subsequent runs, the `batch_size` is `8`. The `LlamaDecoderLayer` is left unchanged in the baseline as follows:\n", + "\n", + "
\n", + "\n", + "
Fig 6: Revisiting \"LlamaDecoderLayer\".
\n", + "
\n", + "\n", + "
\n", + "Note\n", + "\n", + "The baseline implementation will be run in `BF16` precision.\n", + "\n", + "
" + ], + "id": "c9529229" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "
\n", + "\n", + "Note\n", + " \n", + "This tutorial loads and trains a Llama 3 8B or a Llama 2 7B model which takes up most of the GPU memory and therefore, we need to restart the jupyter notebook each time before running the following sections. A small utility method `restart_jupyter_notebook` is defined in the accompanying `utils.py` file. This function restarts the jupyter notebook so that the GPU memory is flushed before the model is loaded again from the checkpoint in order to avoid running into OOM (Out Of Memory) errors.\n", + "\n", + "If the utility doesn't work, comment this line `restart_jupyter_notebook()` in the following cell and manually restart the jupyter notebook before running the cell. Repeat the same for other sections in this tutorial.\n", + "\n", + "
\n" + ], + "id": "b38eb3ac" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Restart the notebook (to flush the GPU memory)\n", + "from utils import restart_jupyter_notebook\n", + "restart_jupyter_notebook()\n", + "\n", + "\n", + "# Import necessary packages, methods and variables\n", + "from utils import *\n", + "\n", + "\n", + "# Provide Huggingface Access Token\n", + "hyperparams.hf_access_token = \"\"\n", + "assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n", + "\n", + "# Provide a directory to cache weights in to avoid downloading them every time.\n", + "# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n", + "hyperparams.weights_cache_dir = \"\"\n", + "\n", + "# For Llama 2, uncomment this line (also set by default)\n", + "hyperparams.model_name = \"meta-llama/Llama-2-7b-hf\"\n", + "\n", + "# For Llama 3, uncomment this line\n", + "# hyperparams.model_name = \"meta-llama/Meta-Llama-3-8B\"\n", + "\n", + "hyperparams.mixed_precision = \"bf16\"\n", + "\n", + "\n", + "# Init the model and accelerator wrapper\n", + "model = init_baseline_model(hyperparams)\n", + "accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)\n", + "\n", + "\n", + "# Finetune the model\n", + "finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)" + ], + "execution_count": 1, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "10 finetuning steps complete!\n", + "Average time taken per step: 248 milliseconds\n" + ] + } + ], + "id": "2e9d7a8c" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "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 (or ms per batch) | Speedup (over baseline) |\n", + "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", + "| HF (baseline) | BF16 | 248 | 1 |" + ], + "id": "4035ccb7" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Improvement 1] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `BF16`)\n", + "\n", + "In addition to basic layers like `Linear` and `LayerNorm`, Transformer Engine offers larger modules like `MultiheadAttention` (combines \"LayerNorm\" and \"Self Attention\") and `LayerNormMLP` (combines \"LayerNorm\" and \"MLP\") that could replace their counterparts in the `LlamaDecoderLayer` and potentially provide a speedup. Transformer Engine also offers a full `TransformerLayer` (which further combines `MultiheadAttention` and `LayerNormMLP` layers) which could replace `LlamaDecoderLayer` and provide a speedup (with careful mapping of the weights since the name of the weights are different for those two layers). Let's take a closer look at Transformer Engine's `TransformerLayer`. \n", + "\n", + "### Transformer Engine's `TransformerLayer`\n", + "\n", + "At a higher level, TE's `TransformerLayer` could be visualized as an apt replacement for the `LlamaDecoderLayer`. But the internals of the `TransformerLayer` are organized a bit differently. \n", + "\n", + "
\n", + "\n", + "
Fig 7: Transformer Engine's `TransformerLayer`
\n", + "
\n", + "\n", + "Just like Hugging Face's `LlamaDecoderLayer`, Transformer Engine's `TransformerLayer` encapsulates `self_attention` (as `MultiheadAttention`) and `mlp` (as `LayerNormMLP`). A major difference is that the two `Norm`s are included in the `MultiheadAttention` and `LayerNormMLP` layers as shown in the following output prompt:\n", + "\n", + "```\n", + "TransformerLayer(\n", + " (self_attention): MultiheadAttention(\n", + " (layernorm_qkv): LayerNormLinear()\n", + " (core_attention): DotProductAttention()\n", + " (proj): Linear()\n", + " )\n", + " (layernorm_mlp): LayerNormMLP()\n", + ")\n", + "```\n", + "\n", + "Another difference is that Transformer Engine implements an efficient version of feedforward layer with SwiGLU in which the weights from the `up_proj` and `gate_proj` modules are merged together and SwiGLU is applied using a custom fused kernel. This is done so that only one big and efficient Matrix Multiplication operation is issued to the GPU instead of two smaller ones.\n", + "\n", + "
\n", + "\n", + "
Fig 8: Abstract illustration of the SwiGLU implementation in Transformer Engine.
\n", + "
\n", + "\n", + "### `TransformerLayer` options explained\n", + "\n", + "
\n", + "\n", + "Note\n", + " \n", + "Here, we go over some of the options in `TransformerLayer` that are needed for the tutorial. For a complete list of options, refer the [TransformerLayer API documentation](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/api/pytorch.html?highlight=transformerlayer#transformer_engine.pytorch.TransformerLayer).\n", + "\n", + "
\n", + "\n", + "In the accompanying `te_llama.py` file, `TELlamaDecoderLayer` is defined as a wrapper over TE's `TransformerLayer` with a few needed options that make `TransformerLayer` a plug-in replacement for the HF's `LlamaDecoderLayer`.\n", + "\n", + "```\n", + "class TELlamaDecoderLayer(te.pytorch.TransformerLayer):\n", + " def __init__(self, config):\n", + " super().__init__(\n", + " config.hidden_size,\n", + " config.intermediate_size,\n", + " config.num_attention_heads,\n", + " bias=False,\n", + " layernorm_epsilon=config.rms_norm_eps,\n", + " hidden_dropout=0,\n", + " attention_dropout=0,\n", + " fuse_qkv_params=False,\n", + " normalization=\"RMSNorm\",\n", + " activation=\"swiglu\",\n", + " attn_input_format=\"bshd\",\n", + " num_gqa_groups=config.num_key_value_heads,\n", + " )\n", + " te_rope = RotaryPositionEmbedding(config.hidden_size//config.num_attention_heads)\n", + " self.te_rope_emb = te_rope(max_seq_len=config.max_position_embeddings).cuda()\n", + "```\n", + "\n", + "Here's a list summarizing each option briefly:\n", + "\n", + "1. `hidden_size`: size of each input sample.\n", + "2. `ffn_hidden_size`: intermediate size to which samples are projected.\n", + "3. `num_attention_heads`: number of attention heads in the transformer layer.\n", + "4. `bias`: switch to add additive biases to the submodule layers.\n", + "5. `layernorm_epsilon`: a value added to the denominator of layer normalization for numerical stability. Default is `1e-5`.\n", + "6. `hidden_dropout`: dropout probability for the dropout op after FC2 layer (fully connected layer no. 2). Default is `0.1`.\n", + "7. `attention_dropout`: dropout probability for the dropout op during multi-head attention. Default is `0.1`. \n", + "8. `fuse_qkv_params`: if set to True, TransformerLayer module exposes a single fused parameter for query-key-value. This enables optimizations such as QKV fusion without concatentations/splits and also enables the argument fuse_wgrad_accumulation.\n", + "9. `normalization`: type of normalization applied. Default is `LayerNorm`.\n", + "10. `activation`: type of activation used in the MLP block. Default is `gelu`.\n", + "11. `attn_input_format`: controls whether the dimensions of the intermediate hidden states is 'batch first' ('bshd') or 'sequence first' ('sbhd'). `s` stands for the sequence length, `b` batch size, `h` the number of heads, `d` head size. Note that these formats are very closely related to the `qkv_format` in the `MultiHeadAttention` and `DotProductAttention` modules.\n", + "12. `num_gqa_groups`: number of GQA groups in the transformer layer. Grouped Query Attention is described in [this paper](https://arxiv.org/pdf/2305.13245.pdf). This only affects the keys and values, not the querys. GQA-1 is equivalent to Multi-Query Attention ([MQA](https://arxiv.org/pdf/1911.02150.pdf)), while GQA-H is equivalent to MultiHead Attention, i.e. `num_gqa_groups = num_attention_heads`.\n", + "\n", + "\n", + "Further, note that `RotaryPositionEmbedding` is defined as part of the `TELlamaDecoderLayer` (wrapper around TE's `TransformerLayer`) itself since it expects this rope cache if RoPE is used in the model. \n", + "\n", + "Let's revisit how `LlamaDecoderLayer`s form the core of the decoder layer stack in HF's llama implementation:\n", + "```\n", + "ModuleList(\n", + " (0-31): 32 x LlamaDecoderLayer(\n", + " (self_attn): LlamaAttention(\n", + " (q_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (k_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (v_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n", + " (rotary_emb): LlamaRotaryEmbedding()\n", + " )\n", + " (mlp): LlamaMLP(\n", + " (gate_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", + " (up_proj): Linear(in_features=4096, out_features=11008, bias=False)\n", + " (down_proj): Linear(in_features=11008, out_features=4096, bias=False)\n", + " (act_fn): SiLU()\n", + " )\n", + " (input_layernorm): LlamaRMSNorm()\n", + " (post_attention_layernorm): LlamaRMSNorm()\n", + " )\n", + ")\n", + "```\n", + "\n", + "A major portion of the Hugging Face model implementation (32 `LlamaDecoderLayer` layers) could be potentially replaced with Transformer Engine's `TransformerLayer` layers. Let's see how it is made possible.\n", + "\n", + "\n", + "### Mapping weights from HF's `LlamaDecoderLayer` to TE's `TransformerLayer`\n", + "\n", + "Refer the accompanying file `te_llama.py` which provides a reference to create a Llama 2 model with TE's `TransformerLayer` after replacing HF's `LlamaDecoderLayer`.\n", + "\n", + "Briefly, following pieces of code are put together:\n", + "\n", + "1. `TELlamaDecoderLayer` is added as a wrapper for `TransformerLayer`. \n", + "```\n", + "class TELlamaDecoderLayer(te.pytorch.TransformerLayer):\n", + " \"\"\"\n", + " Wrapper class over TE's `TransformerLayer`. This makes the wrapper very\n", + " similar to HF's `LlamaDecoderLayer` and easier to replace it in the code.\n", + "\n", + " Args:\n", + " config: LlamaConfig\n", + " args: positional args (for compatibility with `LlamaDecoderLayer`)\n", + " kwargs: keyword args (for compatibility with `LlamaDecoderLayer`)\n", + " \"\"\"\n", + " def __init__(self, config, *args, **kwargs):\n", + " super().__init__(\n", + " hidden_size=config.hidden_size,\n", + " ffn_hidden_size=config.intermediate_size,\n", + " num_attention_heads=config.num_attention_heads,\n", + " bias=False,\n", + " layernorm_epsilon=config.rms_norm_eps,\n", + " hidden_dropout=0,\n", + " attention_dropout=0,\n", + " fuse_qkv_params=False,\n", + " normalization=\"RMSNorm\",\n", + " activation=\"swiglu\",\n", + " attn_input_format=\"bshd\",\n", + " )\n", + " te_rope = RotaryPositionEmbedding(config.hidden_size//config.num_attention_heads)\n", + " self.te_rope_emb = te_rope(max_seq_len=config.max_position_embeddings).cuda()\n", + "\n", + " def forward(self,\n", + " hidden_states,\n", + " *args,\n", + " attention_mask,\n", + " **kwargs):\n", + " \"\"\"\n", + " Custom forward to make sure we only pass relevant arguments to the\n", + " forward pass of the `TransformerLayer`. Also, make sure the output\n", + " format matches the output of the HF's `LlamaDecoderLayer`.\n", + " \"\"\"\n", + " return (super().forward(hidden_states, attention_mask=attention_mask, rotary_pos_emb=self.te_rope_emb),)\n", + "```\n", + "\n", + "2. Before creating a `LlamaForCausalLM`, `replace_decoder` context manager is used to monkey-patch `LlamaDecoderLayer` with `TELlamaDecoderLayer`.\n", + "\n", + "```\n", + "@contextmanager\n", + "def replace_decoder(te_decoder_cls):\n", + " \"\"\"\n", + " Replace `LlamaDecoderLayer` with custom `TELlamaDecoderLayer`.\n", + " \"\"\"\n", + " original_llama_decoder_cls = transformers.models.llama.modeling_llama.LlamaDecoderLayer\n", + " transformers.models.llama.modeling_llama.LlamaDecoderLayer = te_decoder_cls\n", + " try:\n", + " yield\n", + " finally:\n", + " transformers.models.llama.modeling_llama.LlamaDecoderLayer = original_llama_decoder_cls\n", + ".\n", + ".\n", + ".\n", + "class TELlamaForCausalLM:\n", + " \"\"\"\n", + " Causal LM created with `LlamaModel`. The underlying `LlamaDecoderLayer`\n", + " class is monkey-patched with `TELlamaDecoderLayer` class before\n", + " initializing the causal LM with `LlamaForCausalLM`.\n", + "\n", + " Args:\n", + " config: LlamaConfig\n", + " \"\"\"\n", + "\n", + " def __new__(cls, config: LlamaConfig):\n", + " with replace_decoder(te_decoder_cls=TELlamaDecoderLayer):\n", + " llama_for_causal_lm = LlamaForCausalLM(config)\n", + " return llama_for_causal_lm\n", + ".\n", + ".\n", + ".\n", + "```\n", + "\n", + "3. A custom `pretrained_from_local` method is added that copies the weights from the checkpoint (which is meant for HF Llama implementation) to the modified `TELlamaForCausalLM` by carefully mapping the weights from the `LlamaDecoderLayer` (HF) to `TransformerLayer` (TE). The method `replace_params` maps and copies apt weights from `LlamaDecoderLayer` to the `TransformerLayer`. Refer to the following diagram for more details.\n", + "\n", + "```\n", + "def replace_params(hf_state_dict, te_state_dict):\n", + " # collect all layer prefixes to update\n", + " all_layer_prefixes = set()\n", + " for param_key in hf_state_dict.keys():\n", + " layer_prefix_pat = 'model.layers.\\d+.'\n", + " m = re.match(layer_prefix_pat, param_key)\n", + " if m is not None:\n", + " all_layer_prefixes.add(m.group())\n", + "\n", + " for layer_prefix in all_layer_prefixes:\n", + " # When loading weights into models with less number of layers, skip the\n", + " # copy if the corresponding layer doesn't exist in TE model\n", + " if layer_prefix + 'self_attention.layernorm_qkv.layer_norm_weight' in te_state_dict:\n", + " te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.layer_norm_weight'].data[:] = hf_state_dict[layer_prefix + 'input_layernorm.weight'].data[:]\n", + "\n", + " if layer_prefix + 'self_attention.layernorm_qkv.query_weight' in te_state_dict:\n", + " te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.query_weight'].data[:] = hf_state_dict[layer_prefix + 'self_attn.q_proj.weight'].data[:]\n", + "\n", + " if layer_prefix + 'self_attention.layernorm_qkv.key_weight' in te_state_dict:\n", + " te_state_dict[layer_prefix + 'self_attention.layernorm_qkv.key_weight'].data[:] = hf_state_dict[layer_prefix + 'self_attn.k_proj.weight'].data[:]\n", + " .\n", + " .\n", + " .\n", + "\n", + " return all_layer_prefixes\n", + "```\n", + "\n", + "The following figure shows how the weights get mapped from the HF's `LlamaDecoderLayer` to TE's `TransformerLayer`.\n", + "\n", + "
\n", + "\n", + "
Fig 9: Replace `LlamaDecoderLayer` with `TransformerLayer`.
\n", + "
\n", + "\n", + "After initializing the modified Llama model this way, the core decoder layers get changed to `TELlamaDecoderLayer` (wrapper around `TransformerLayer`) as shown in the following output:\n", + "```\n", + "ModuleList(\n", + " (0-31): 32 x TELlamaDecoderLayer(\n", + " (self_attention): MultiheadAttention(\n", + " (layernorm_qkv): LayerNormLinear()\n", + " (core_attention): DotProductAttention(\n", + " (flash_attention): FlashAttention()\n", + " (fused_attention): FusedAttention()\n", + " (unfused_attention): UnfusedDotProductAttention(\n", + " (scale_mask_softmax): FusedScaleMaskSoftmax()\n", + " (attention_dropout): Dropout(p=0, inplace=False)\n", + " )\n", + " )\n", + " (proj): Linear()\n", + " )\n", + " (layernorm_mlp): LayerNormMLP()\n", + " )\n", + ")\n", + "```\n", + "\n", + "In summary, the model gets changed as follows with a large chunk of the implementation (core decoder layers) coming from Transformer Engine.\n", + "\n", + "
\n", + "\n", + "
Fig 10: Language model after the HF's `LlamaDecoderLayer`s are replaced with TE's `TransformerLayer`s.
\n", + "
\n", + "\n", + "\n", + "
\n", + "Note\n", + "\n", + "Let's first run this \"TELlama\" implementation in `BF16` precision.\n", + "
" + ], + "id": "3db90dff" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Restart the notebook (to flush the GPU memory)\n", + "from utils import restart_jupyter_notebook\n", + "restart_jupyter_notebook()\n", + "\n", + "\n", + "# Import necessary packages, methods and variables\n", + "from utils import *\n", + "\n", + "\n", + "# Provide Huggingface Access Token\n", + "hyperparams.hf_access_token = \"\"\n", + "assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n", + "\n", + "# Provide a directory to cache weights in to avoid downloading them every time.\n", + "# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n", + "hyperparams.weights_cache_dir = \"\"\n", + "\n", + "# For Llama 2, uncomment this line (also set by default)\n", + "hyperparams.model_name = \"meta-llama/Llama-2-7b-hf\"\n", + "\n", + "# For Llama 3, uncomment this line\n", + "# hyperparams.model_name = \"meta-llama/Meta-Llama-3-8B\"\n", + "\n", + "hyperparams.mixed_precision = \"bf16\"\n", + "\n", + "\n", + "# Init the model and accelerator wrapper\n", + "model = init_te_llama_model(hyperparams)\n", + "accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)\n", + "\n", + "\n", + "# Finetune the model\n", + "finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)" + ], + "execution_count": 1, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "10 finetuning steps complete!\n", + "Average time taken per step: 185 milliseconds\n" + ] + } + ], + "id": "bdb34b91" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Compared to the \"baseline\" implementation, we see that using Transformer Engine's `TransformerLayer` in place of Huggging Face's `LlamaDecoderLayer` gives a speedup of **34%** even when using only BF16 precision!\n", + "\n", + "| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n", + "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", + "| HF (baseline) | BF16 | 248 | 1 |\n", + "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 185 | 1.34 |" + ], + "id": "0c9fbd65" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Improvement 2] Replace HF's `LlamaDecoderLayer` with TE's `TransformerLayer` (Precision: `FP8`)\n", + "\n", + "Now that most of the HF Llama model implementation (`LlamaDecoderLayer`s) has been swapped with Transformer Engine implementation (`TELlamaDecoderLayer` or `TransformerLayer`), let's see how finetuning in `FP8` precision helps improve performance.\n", + "\n", + "### How to run the model in `FP8` precision\n", + "\n", + "After the substitution, the model can be run in `FP8` precision by the following change over the previous BF16 runs. (For more information, refer the corresponding `wrap_with_accelerator` function in the accompanying `utils.py` file).\n", + "\n", + "```\n", + "# Specify the `FP8RecipeKwargs` (additional argument required to run in `fp8` precision)\n", + "fp8_kwarg_handler = [FP8RecipeKwargs(backend=\"te\")]\n", + "\n", + "# Pass the `FP8RecipeKwargs` to the `Accelerator` init call\n", + "accelerator = Accelerator(\n", + " ...\n", + " kwargs_handlers=fp8_kwarg_handler\n", + ")\n", + "```" + ], + "id": "98cd8efb" + }, + { + "cell_type": "code", + "metadata": {}, + "source": [ + "# Restart the notebook (to flush the GPU memory)\n", + "from utils import restart_jupyter_notebook\n", + "restart_jupyter_notebook()\n", + "\n", + "\n", + "# Import necessary packages, methods and variables\n", + "from utils import *\n", + "\n", + "\n", + "# Provide Huggingface Access Token\n", + "hyperparams.hf_access_token = \"\"\n", + "assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n", + "\n", + "# Provide a directory to cache weights in to avoid downloading them every time.\n", + "# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n", + "hyperparams.weights_cache_dir = \"\"\n", + "\n", + "# For Llama 2, uncomment this line (also set by default)\n", + "hyperparams.model_name = \"meta-llama/Llama-2-7b-hf\"\n", + "\n", + "# For Llama 3, uncomment this line\n", + "# hyperparams.model_name = \"meta-llama/Meta-Llama-3-8B\"\n", + "\n", + "hyperparams.mixed_precision = \"fp8\"\n", + "\n", + "\n", + "# Init the model and accelerator wrapper\n", + "model = init_te_llama_model(hyperparams)\n", + "accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)\n", + "\n", + "\n", + "# Finetune the model\n", + "finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)" + ], + "execution_count": 1, + "outputs": [ + { + "output_type": "stream", + "name": "stdout", + "text": [ + "10 finetuning steps complete!\n", + "Average time taken per step: 160 milliseconds\n" + ] + } + ], + "id": "772c6f22" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n", + "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", + "| HF (baseline) | BF16 | 248 | 1 |\n", + "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 185 | 1.34 |\n", + "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | FP8 | 160 | 1.55 |\n", + "\n", + "\n", + "After turning on FP8 precision, we get even more speedup of **55%** (with Llama 2 7B)!\n", + "\n", + "### Llama 3 performance results\n", + "Running the same tutorial with **Llama 3 8B** yields the following performance numbers:\n", + "\n", + "| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n", + "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", + "| HF (baseline) | BF16 | 270 | 1 |\n", + "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 217 | 1.24 |\n", + "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | FP8 | 185 | 1.46 |\n", + "\n", + "For Llama 3 8B, we get the most speedup of **46%** with FP8 precision!\n", + "\n" + ], + "id": "e7cf9c3a" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Conclusion\n", + "\n", + "Using `TransformerLayer` module from Transformer Engine as a substitute for Hugging Face's `LlamaDecoderLayer` provides a speedup over Hugging Face's native Llama 2 and Llama 3 implementations. This needs careful initialization of the model such that the model weights (which are meant for `LlamaDecoderLayer`) are correctly mapped to their counterparts in TE's `TransformerLayer`. Even with `BF16` precision, `TransformerLayer` provides a speedup over the baseline implementation. With `FP8` precision, the speed up is even more pronounced!" + ], + "id": "95d6c42b" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.12" } - ], - "source": [ - "# Restart the notebook (to flush the GPU memory)\n", - "from utils import restart_jupyter_notebook\n", - "restart_jupyter_notebook()\n", - "\n", - "\n", - "# Import necessary packages, methods and variables\n", - "from utils import *\n", - "\n", - "\n", - "# Provide Huggingface Access Token\n", - "hyperparams.hf_access_token = \"\"\n", - "assert hyperparams.hf_access_token, \"Provide a HF API Access Token!\"\n", - "\n", - "# Provide a directory to cache weights in to avoid downloading them every time.\n", - "# (By default, weights are cached in `~/.cache/huggingface/hub/models`)\n", - "hyperparams.weights_cache_dir = \"\"\n", - "\n", - "# For Llama 2, uncomment this line (also set by default)\n", - "hyperparams.model_name = \"meta-llama/Llama-2-7b-hf\"\n", - "\n", - "# For Llama 3, uncomment this line\n", - "# hyperparams.model_name = \"meta-llama/Meta-Llama-3-8B\"\n", - "\n", - "hyperparams.mixed_precision = \"fp8\"\n", - "\n", - "\n", - "# Init the model and accelerator wrapper\n", - "model = init_te_llama_model(hyperparams)\n", - "accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator(model, hyperparams)\n", - "\n", - "\n", - "# Finetune the model\n", - "finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler)" - ] - }, - { - "cell_type": "markdown", - "id": "e7cf9c3a", - "metadata": {}, - "source": [ - "| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n", - "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", - "| HF (baseline) | BF16 | 248 | 1 |\n", - "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 185 | 1.34 |\n", - "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | FP8 | 160 | 1.55 |\n", - "\n", - "\n", - "After turning on FP8 precision, we get even more speedup of **55%** (with Llama 2 7B)!\n", - "\n", - "#### Llama 3 performance results\n", - "Running the same tutorial with **Llama 3 8B** yields the following performance numbers:\n", - "\n", - "| Models | Precision | Step Time (or ms per batch) | Speedup (over baseline) |\n", - "|-------------------------------------------------------------|-----------|-----------------------------|-------------------------|\n", - "| HF (baseline) | BF16 | 270 | 1 |\n", - "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | BF16 | 217 | 1.24 |\n", - "| TE (replace `LlamaDecoderLayer` with `TE.TransformerLayer`) | FP8 | 185 | 1.46 |\n", - "\n", - "For Llama 3 8B, we get the most speedup of **46%** with FP8 precision!\n", - "\n" - ] - }, - { - "cell_type": "markdown", - "id": "95d6c42b", - "metadata": {}, - "source": [ - "## Conclusion\n", - "\n", - "Using `TransformerLayer` module from Transformer Engine as a substitute for Hugging Face's `LlamaDecoderLayer` provides a speedup over Hugging Face's native Llama 2 and Llama 3 implementations. This needs careful initialization of the model such that the model weights (which are meant for `LlamaDecoderLayer`) are correctly mapped to their counterparts in TE's `TransformerLayer`. Even with `BF16` precision, `TransformerLayer` provides a speedup over the baseline implementation. With `FP8` precision, the speed up is even more pronounced!" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.12" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file From 2104e4c1b3dcc7249e9e4de9252bf4fa59a03030 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 27 Jan 2026 12:10:19 -0800 Subject: [PATCH 183/521] [JAX] Use "nyu-mll/glue" instead of "glue" for encoder datasets to fix 404 error (#2625) * Use "nyu-mll/glue" instead of "glue" for encoder datasets to fix 404 error Signed-off-by: Jeremy Berchtold * rename mnist dataset path Signed-off-by: Jeremy Berchtold * add dataset manifest Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold --- examples/jax/datasets.txt | 3 +++ examples/jax/encoder/test_model_parallel_encoder.py | 4 ++-- examples/jax/encoder/test_multigpu_encoder.py | 4 ++-- examples/jax/encoder/test_multiprocessing_encoder.py | 4 ++-- examples/jax/encoder/test_single_gpu_encoder.py | 4 ++-- examples/jax/mnist/test_single_gpu_mnist.py | 4 ++-- 6 files changed, 13 insertions(+), 10 deletions(-) create mode 100644 examples/jax/datasets.txt diff --git a/examples/jax/datasets.txt b/examples/jax/datasets.txt new file mode 100644 index 0000000000..fd3f5bc41e --- /dev/null +++ b/examples/jax/datasets.txt @@ -0,0 +1,3 @@ +# Datasets used by TE encoder tests. Pull these to pre-emptively cache datasets +ylecun/mnist +nyu-mll/glue \ No newline at end of file diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index 02937bc394..73b93798a0 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -219,11 +219,11 @@ def get_datasets(max_seq_len): vocab = {} word_id = 0 - train_ds = load_dataset("glue", "cola", split="train") + train_ds = load_dataset("nyu-mll/glue", "cola", split="train") train_ds.set_format(type="np") train_ds, vocab, word_id = data_preprocess(train_ds, vocab, word_id, max_seq_len) - test_ds = load_dataset("glue", "cola", split="validation") + test_ds = load_dataset("nyu-mll/glue", "cola", split="validation") test_ds.set_format(type="np") test_ds, vocab, word_id = data_preprocess(test_ds, vocab, word_id, max_seq_len) return train_ds, test_ds, word_id diff --git a/examples/jax/encoder/test_multigpu_encoder.py b/examples/jax/encoder/test_multigpu_encoder.py index 98184ccd75..22a89cc0a9 100644 --- a/examples/jax/encoder/test_multigpu_encoder.py +++ b/examples/jax/encoder/test_multigpu_encoder.py @@ -197,11 +197,11 @@ def get_datasets(max_seq_len): vocab = {} word_id = 0 - train_ds = load_dataset("glue", "cola", split="train") + train_ds = load_dataset("nyu-mll/glue", "cola", split="train") train_ds.set_format(type="np") train_ds, vocab, word_id = data_preprocess(train_ds, vocab, word_id, max_seq_len) - test_ds = load_dataset("glue", "cola", split="validation") + test_ds = load_dataset("nyu-mll/glue", "cola", split="validation") test_ds.set_format(type="np") test_ds, vocab, word_id = data_preprocess(test_ds, vocab, word_id, max_seq_len) return train_ds, test_ds, word_id diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index 327540521c..0166b60acd 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -307,11 +307,11 @@ def get_datasets(max_seq_len): vocab = {} word_id = 0 - train_ds = load_dataset("glue", "cola", split="train") + train_ds = load_dataset("nyu-mll/glue", "cola", split="train") train_ds.set_format(type="np") train_ds, vocab, word_id = data_preprocess(train_ds, vocab, word_id, max_seq_len) - test_ds = load_dataset("glue", "cola", split="validation") + test_ds = load_dataset("nyu-mll/glue", "cola", split="validation") test_ds.set_format(type="np") test_ds, vocab, word_id = data_preprocess(test_ds, vocab, word_id, max_seq_len) return train_ds, test_ds, word_id diff --git a/examples/jax/encoder/test_single_gpu_encoder.py b/examples/jax/encoder/test_single_gpu_encoder.py index 82c7fed38e..6d67296bd2 100644 --- a/examples/jax/encoder/test_single_gpu_encoder.py +++ b/examples/jax/encoder/test_single_gpu_encoder.py @@ -195,11 +195,11 @@ def get_datasets(max_seq_len): vocab = {} word_id = 0 - train_ds = load_dataset("glue", "cola", split="train") + train_ds = load_dataset("nyu-mll/glue", "cola", split="train") train_ds.set_format(type="np") train_ds, vocab, word_id = data_preprocess(train_ds, vocab, word_id, max_seq_len) - test_ds = load_dataset("glue", "cola", split="validation") + test_ds = load_dataset("nyu-mll/glue", "cola", split="validation") test_ds.set_format(type="np") test_ds, vocab, word_id = data_preprocess(test_ds, vocab, word_id, max_seq_len) return train_ds, test_ds, word_id diff --git a/examples/jax/mnist/test_single_gpu_mnist.py b/examples/jax/mnist/test_single_gpu_mnist.py index 0c76d51c37..ef85f4a7ab 100644 --- a/examples/jax/mnist/test_single_gpu_mnist.py +++ b/examples/jax/mnist/test_single_gpu_mnist.py @@ -146,7 +146,7 @@ def eval_model(state, test_ds, batch_size, var_collect): def get_datasets(): """Load MNIST train and test datasets into memory.""" - train_ds = load_dataset("mnist", split="train", trust_remote_code=True) + train_ds = load_dataset("ylecun/mnist", split="train", trust_remote_code=True) train_ds.set_format(type="np") batch_size = train_ds["image"].shape[0] shape = (batch_size, IMAGE_H, IMAGE_W, IMAGE_C) @@ -154,7 +154,7 @@ def get_datasets(): "image": train_ds["image"].astype(np.float32).reshape(shape) / 255.0, "label": train_ds["label"], } - test_ds = load_dataset("mnist", split="test", trust_remote_code=True) + test_ds = load_dataset("ylecun/mnist", split="test", trust_remote_code=True) test_ds.set_format(type="np") batch_size = test_ds["image"].shape[0] shape = (batch_size, IMAGE_H, IMAGE_W, IMAGE_C) From f04b094c20d075374baa8b6dcbc566247b32bd36 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, 28 Jan 2026 01:27:17 +0100 Subject: [PATCH 184/521] [PyTorch] ONNX test fix + export for FP8 attention (#2598) * jjit bug fix Signed-off-by: Pawel Gadzinski * fix' Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix 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> --- qa/L1_pytorch_onnx_unittest/test.sh | 3 +- tests/pytorch/test_onnx_export.py | 22 +++++++-- .../dot_product_attention/backends.py | 46 +++++++++++++++++++ .../dot_product_attention.py | 4 +- .../attention/dot_product_attention/utils.py | 4 +- transformer_engine/pytorch/jit.py | 34 ++++++++++---- 6 files changed, 97 insertions(+), 16 deletions(-) diff --git a/qa/L1_pytorch_onnx_unittest/test.sh b/qa/L1_pytorch_onnx_unittest/test.sh index b3a520e129..6f9ff54e48 100644 --- a/qa/L1_pytorch_onnx_unittest/test.sh +++ b/qa/L1_pytorch_onnx_unittest/test.sh @@ -6,4 +6,5 @@ : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_onnx_export.xml $TE_PATH/tests/pytorch/test_onnx_export.py +# NVTE_UnfusedDPA_Emulate_FP8=1 enables FP8 attention emulation when no native backend is available +NVTE_UnfusedDPA_Emulate_FP8=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_onnx_export.xml $TE_PATH/tests/pytorch/test_onnx_export.py diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index 50cd150c4e..9aea3bc274 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -713,6 +713,14 @@ def test_export_layernorm_mlp_activation(seed_default_rng, activation): _test_export_layernorm_mlp(activation=activation) +# Quantization recipes with fp8_dpa=True for attention emulation export test +dpa_quantization_recipes = [None] # None = no quantization +if fp8_available: + dpa_quantization_recipes.append(recipe.DelayedScaling(fp8_dpa=True)) + dpa_quantization_recipes.append(recipe.Float8CurrentScaling(fp8_dpa=True)) + + +@pytest.mark.parametrize("fp8_recipe", dpa_quantization_recipes) @pytest.mark.parametrize( "precision, use_mask, attn_mask_type", [ @@ -730,6 +738,7 @@ def test_export_core_attention( precision: torch.dtype, use_mask: bool, attn_mask_type: str, + fp8_recipe: recipe.Recipe, ): # Set dimensions (these are arbitrary). seq_len, batch_size, num_attention_heads, kv_channels = (64, 4, 1, 64) @@ -749,22 +758,25 @@ def test_export_core_attention( mask_str = get_attn_mask_str(use_mask, attn_mask_type) high_prec_str = dtype2str(precision) - fname = f"te.core_attention{mask_str}{high_prec_str}.onnx" + fp8_str = "_fp8_dpa" if fp8_recipe is not None else "" + fname = f"te.core_attention{fp8_str}{mask_str}{high_prec_str}.onnx" + + is_fp8 = fp8_recipe is not None model = te.attention.DotProductAttention( num_attention_heads=num_attention_heads, kv_channels=kv_channels, - attention_dropout=0.5, qkv_format=qkv_format, attn_mask_type=attn_mask_type, ).to(device="cuda") - do_export(model, inp, fname, input_names=input_names, fp8_recipe=None) - te_outputs = te_infer(model, inp, is_fp8=False, fp8_recipe=None) + do_export(model, inp, fname, input_names=input_names, fp8_recipe=fp8_recipe) + te_outputs = te_infer(model, inp, is_fp8=is_fp8, fp8_recipe=fp8_recipe) serialize_inputs_outputs(fname, inp, te_outputs, input_names=input_names) if precision in (torch.bfloat16,): return + atol = 5e-1 if is_fp8 else 1e-2 validate_result( - fname, inp, model, is_fp8=True, atol=1e-2, input_names=input_names, te_outputs=te_outputs + fname, inp, model, is_fp8=True, atol=atol, input_names=input_names, te_outputs=te_outputs ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index ef7fa0dcc0..aa6c063951 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -164,6 +164,11 @@ class FP8EmulationFunc(torch.autograd.Function): @staticmethod def forward(ctx, tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layout): # pylint: disable=missing-function-docstring + if is_in_onnx_export_mode(): + return FP8EmulationFunc.onnx_forward( + tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layout + ) + if quantizer_name == "QKV_quantizer": query_layer, key_layer, value_layer = [ x.contiguous() for x in [tensor1, tensor2, tensor3] @@ -202,6 +207,47 @@ def backward(ctx, grad1, grad2, grad3): tensors = grad1, grad2, grad3 return tensors[0], tensors[1], tensors[2], None, None, None + @staticmethod + def onnx_forward(tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layout=None): + """ + ONNX-compatible forward for FP8 emulation using operations with defined ONNX translations. + """ + # pylint: disable=unused-argument + is_qkv_quantizer = quantizer_name == "QKV_quantizer" + assert isinstance( + quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer) + ), "ONNX FP8 emulation path supports only Float8 quantizers." + + if is_qkv_quantizer: + # Flatten + concatenate + quantize + split. Equivalent to combine_and_quantize Case 3. + orig_dtype = tensor1.dtype + shapes = [tensor1.shape, tensor2.shape, tensor3.shape] + numels = [tensor1.numel(), tensor2.numel(), tensor3.numel()] + + # Flatten and concatenate + combined = torch.cat( + [tensor1.reshape(-1), tensor2.reshape(-1), tensor3.reshape(-1)], dim=0 + ) + + # Quantize + dequantize combined tensor using quantizer's ONNX methods + combined_fp8 = quantizer.onnx_quantize(combined) + out = quantizer.onnx_dequantize(combined_fp8).to(orig_dtype) + + # Split back + out1 = out[: numels[0]].reshape(shapes[0]) + out2 = out[numels[0] : numels[0] + numels[1]].reshape(shapes[1]) + out3 = out[numels[0] + numels[1] :].reshape(shapes[2]) + + return out1, out2, out3 + if quantizer_name in ["S_quantizer", "O_quantizer"]: + # Emulate FP8 on single tensor using quantizer's ONNX methods + orig_dtype = tensor1.dtype + t_fp8 = quantizer.onnx_quantize(tensor1) + out = quantizer.onnx_dequantize(t_fp8).to(orig_dtype) + return out, tensor2, tensor3 + # Pass-through + return tensor1, tensor2, tensor3 + class UnfusedDotProductAttention(torch.nn.Module): """Parallel attention w/o QKV and Proj Gemms 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 5a554d86ec..5d830dca33 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 @@ -1552,7 +1552,9 @@ def forward( ) if use_unfused_attention: - allow_emulation = os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" + allow_emulation = ( + os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode() + ) if checkpoint_core_attention: return self._checkpointed_attention_forward( self.unfused_attention, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 56e6f093d1..0c5a519813 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -479,7 +479,9 @@ def get_attention_backend( logger.debug("Disabling FlashAttention 3 for FP8 training") use_flash_attention_3 = False if use_unfused_attention: - allow_emulation = os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" + allow_emulation = ( + os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode() + ) if not allow_emulation: logger.debug("Disabling UnfusedDotProductAttention for FP8 attention") use_unfused_attention = False diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py index 5884188b7e..1b93b8254c 100644 --- a/transformer_engine/pytorch/jit.py +++ b/transformer_engine/pytorch/jit.py @@ -46,17 +46,35 @@ def wrapper(*args, **kwargs): # Decorator to disable Torch Dynamo # See: https://github.com/NVIDIA/TransformerEngine/issues/308 -no_torch_dynamo = lambda recursive=True: lambda func: func if torch.__version__ >= "2": import torch._dynamo - if torch.__version__ >= "2.1": - no_torch_dynamo = lambda recursive=True: lambda f: ( - f if is_in_onnx_export_mode() else torch._dynamo.disable(f, recursive=recursive) - ) - else: - # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True - no_torch_dynamo = lambda recursive=True: torch._dynamo.disable + def no_torch_dynamo(recursive=True): + """Decorator to disable Torch Dynamo, except during ONNX export.""" + + def decorator(f): + # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True + disabled_f = ( + torch._dynamo.disable(f, recursive=recursive) + if torch.__version__ >= "2.1" + else torch._dynamo.disable(f) + ) + + @wraps(f) + def wrapper(*args, **kwargs): + if is_in_onnx_export_mode(): + return f(*args, **kwargs) + return disabled_f(*args, **kwargs) + + return wrapper + + return decorator + +else: + # Fallback for PyTorch < 2.0: no-op decorator + def no_torch_dynamo(recursive=True): # pylint: disable=unused-argument + """No-op decorator for PyTorch < 2.0.""" + return lambda func: func def set_jit_fusion_options() -> None: From b9f4013143c4c3f565bfdf375263b218328217c8 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, 28 Jan 2026 01:28:30 +0100 Subject: [PATCH 185/521] [common] Add support for cuBLASLt GEMM for GroupedTensor (#2502) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add FP8 scale support and fix alignment for grouped GEMM - Add FP8 scale_inv pointer handling in nvte_grouped_gemm for proper FP8 GEMM - Fix random padding in tests to ensure 16-byte alignment for all dtypes - Reorder GroupedGemmSetupWorkspace members for natural alignment - Remove debug prints Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Grouped GEMM: code cleanup and NULL C support - Remove unused alignment parameter from GroupedGemmSetupWorkspace::from_buffers - Simplify select_grouped_operand by removing dead code branches - Add GroupedOperandSelection.tensor field to avoid passing tensor separately - Extract set_fp8_scale_pointers and init_matrix_layouts helpers - Add safety check for FP8 on Hopper column-wise fallback - Support NULL C tensor when beta=0 (uses D as placeholder) - Remove unused get_scale_inv() from test - Add use_null_c test parameter and test case - Fix documentation: alpha/beta are single element tensors only Signed-off-by: Piotr Gadzinski Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Grouped GEMM: per-matrix alpha/beta support - Change alpha/beta from single values to per-matrix arrays - Validate alpha/beta have exactly num_tensors elements - Update kernel to index alpha_ptr[idx] and beta_ptr[idx] - Move alpha/beta validation to validate_grouped_gemm_inputs - Update tests to use per-matrix alpha/beta arrays - Update documentation Signed-off-by: Piotr Gadzinski Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix alpha/beta numel - use SimpleTensor::numel() Signed-off-by: Piotr Gadzinski Signed-off-by: Pawel Gadzinski * Refactor: move grouped GEMM to separate file and cleanup API Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * Require Blackwell (SM100) and cuBLAS 13.1+ for grouped GEMM Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/common/gemm/config.h Co-authored-by: Przemyslaw Tredak Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * changed Signed-off-by: Pawel Gadzinski * suggestions Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactored hopper tensor selection Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Piotr Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak --- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_grouped_gemm.cu | 308 +++++++++ tests/cpp/test_common.cu | 163 +++++ tests/cpp/test_common.h | 54 ++ transformer_engine/common/CMakeLists.txt | 1 + transformer_engine/common/gemm/config.cpp | 103 +++ transformer_engine/common/gemm/config.h | 19 + .../common/gemm/cublaslt_gemm.cu | 35 +- .../common/gemm/cublaslt_grouped_gemm.cu | 645 ++++++++++++++++++ .../common/include/transformer_engine/gemm.h | 171 +++++ .../common/util/cuda_runtime.cpp | 8 + transformer_engine/common/util/cuda_runtime.h | 6 + 12 files changed, 1494 insertions(+), 20 deletions(-) create mode 100644 tests/cpp/operator/test_grouped_gemm.cu create mode 100644 transformer_engine/common/gemm/cublaslt_grouped_gemm.cu diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 26efb37962..08a683949b 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -30,6 +30,7 @@ add_executable(test_operator test_causal_softmax.cu test_swizzle.cu test_swap_first_dims.cu + test_grouped_gemm.cu ../test_common.cu) # Find required packages diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu new file mode 100644 index 0000000000..35c4375cbe --- /dev/null +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -0,0 +1,308 @@ +/************************************************************************* + * 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 +#include +#include + +#include "../test_common.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum class InputCase { + kFP8Current, + kBF16, +}; + +enum class ShapeCase { + kAllSame, + kSameFirst, + kSameLast, + kAllDifferent, +}; + +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: 6 pointer arrays (A, B, C, D, alpha, beta) + 6 int arrays (a_rows, a_cols, b_rows, b_cols, d_rows, d_cols) + size_t size = 6 * ptr_bytes + 6 * int_bytes; + const size_t alignment = 256; + size = ((size + alignment - 1) / alignment) * alignment; + return size; +} + +Tensor make_fp8_operand(const std::string& name, const std::vector& shape) { + Tensor input_fp32(name + "_fp32", shape, DType::kFloat32); + fillUniform(&input_fp32); + + Tensor fp8(name, shape, TypeInfo::dtype, true, true, NVTE_DELAYED_TENSOR_SCALING); + + nvte_compute_amax(input_fp32.data(), fp8.data(), 0); + QuantizationConfigWrapper config; + nvte_compute_scale_from_amax(fp8.data(), config, 0); + nvte_quantize(input_fp32.data(), fp8.data(), 0); + return fp8; +} + +Tensor make_bf16_operand(const std::string& name, const std::vector& shape) { + Tensor t(name, shape, DType::kBFloat16); + const size_t numel = shape[0] * shape[1]; + std::vector<__nv_bfloat16> ones(numel, __float2bfloat16(1.0f)); + NVTE_CHECK_CUDA(cudaMemcpy(t.rowwise_dptr(), ones.data(), + numel * sizeof(__nv_bfloat16), cudaMemcpyHostToDevice)); + return t; +} + +struct TestParams { + InputCase input_case; + bool transa; + bool transb; + ShapeCase shape_case; + bool use_null_c = false; // When true, pass nullptr for C (valid when beta=0) +}; + +// Returns a vector of (M, N, K) tuples for each GEMM in the group. +// M - number of rows in output D +// N - number of columns in output D +// K - reduction dimension shared between A and B +std::vector> make_shapes(ShapeCase scase) { + switch (scase) { + case ShapeCase::kAllSame: + return {{64, 64, 32}, {64, 64, 32}, {64, 64, 32}}; + case ShapeCase::kSameFirst: + // Same M (first dim), varying N and K + return {{64, 80, 32}, {64, 96, 48}, {64, 112, 64}}; + case ShapeCase::kSameLast: + // Same N (last dim), varying M and K + return {{64, 80, 32}, {80, 80, 48}, {96, 80, 64}}; + case ShapeCase::kAllDifferent: + default: + return {{64, 96, 32}, {80, 112, 48}, {96, 128, 64}}; + } +} + +void run_grouped_gemm_case(const TestParams& params) { +#if CUBLAS_VERSION < 130100 + GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.1+, but compile-time cuBLAS version is " + << CUBLAS_VERSION << "."; +#else + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP() << "Grouped GEMM requires Blackwell (SM100) or newer."; + } + + 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); + + 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{M, K} + : std::vector{K, M}; + const std::vector b_shape = params.transb ? std::vector{K, N} + : std::vector{N, 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; + } + } + 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; + 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]); + } + + 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()); + + std::vector C_tensors; + std::vector 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{static_cast(M), static_cast(N)}, + DType::kBFloat16)); + } + D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i), + std::vector{static_cast(M), static_cast(N)}, + DType::kBFloat16)); + NVTE_CHECK_CUDA(cudaMemset(D_group_tensors.back().rowwise_dptr(), 0, bytes(D_group_tensors.back().rowwise_shape(), D_group_tensors.back().dtype()))); + } + + 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]); + } + D_views.push_back(&D_group_tensors[i]); + } + + std::optional grouped_C; + if (!params.use_null_c) { + grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING); + } + 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); + + 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 >= 130100 +} + +class GroupedGemmTest : public ::testing::TestWithParam {}; + +TEST_P(GroupedGemmTest, CompareWithMultiTensorGemm) { + run_grouped_gemm_case(GetParam()); +} + +std::string MakeGroupedGemmTestName(const testing::TestParamInfo& info) { + constexpr const char* kInputNames[] = {"FP8Current", "BF16"}; + constexpr const char* kShapeNames[] = {"AllSame", "SameM", "SameN", "AllDiff"}; + 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; +} + +// TestParams: {input_case, transa, transb, shape_case, use_null_c} +const std::vector kTestParams = { + // Basic tests + {InputCase::kFP8Current, true, false, ShapeCase::kAllDifferent, false}, + {InputCase::kFP8Current, false, true, ShapeCase::kAllDifferent, false}, + {InputCase::kFP8Current, false, false, ShapeCase::kAllSame, false}, + {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}, + // Test NULL C (valid when beta=0) + {InputCase::kBF16, false, false, ShapeCase::kAllSame, true}, +}; + +INSTANTIATE_TEST_SUITE_P(OperatorTest, + GroupedGemmTest, + ::testing::ValuesIn(kTestParams), + MakeGroupedGemmTestName); + +} // namespace diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index ed961bfe96..af99d9c42f 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -1057,4 +1058,166 @@ std::array get_scale_tensor_dims(const size_t rows, return {unpadded_blocks_Y, unpadded_blocks_X, blocks_Y, blocks_X}; } +GroupedBuffers build_grouped_tensor(const std::vector& tensors, + const NVTEScalingMode scaling_mode) { + NVTE_CHECK(!tensors.empty(), "No tensors provided for grouped tensor build."); + const NVTEShape shape = tensors[0]->rowwise_shape(); + const DType dtype = tensors[0]->dtype(); + const size_t num_tensors = tensors.size(); + const size_t elem_size = typeToNumBits(dtype) / 8; + GroupedBuffers grouped; + grouped.elem_size = elem_size; + grouped.num_tensors = num_tensors; + grouped.dtype = dtype; + grouped.scaling_mode = scaling_mode; + grouped.tensor_bytes.resize(num_tensors); + grouped.offsets_host.resize(num_tensors, 0); + + std::vector first_dims(num_tensors); + std::vector last_dims(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + const auto s = tensors[i]->rowwise_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 bool same_first = std::all_of(first_dims.begin(), first_dims.end(), + [&](int64_t v) { return v == first_dims[0]; }); + 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); + auto random_padding = [&]() -> int64_t { + // Random padding ensuring 16-byte alignment regardless of element size + // 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 + return dist(gen) * static_cast(align_elements); + }; + + auto numel = [&](size_t idx) -> int64_t { + return first_dims[idx] * last_dims[idx]; + }; + + const bool need_offsets = !same_first || !same_last; + if (need_offsets) { + offsets[0] = 0; + for (size_t i = 1; i < num_tensors; ++i) { + offsets[i] = offsets[i - 1] + numel(i - 1) + random_padding(); + } + } else { + for (size_t i = 0; i < num_tensors; ++i) { + offsets[i] = static_cast(i) * numel(0); + } + } + grouped.offsets_host = offsets; + + int64_t logical_first = 0; + int64_t logical_last = 0; + if (same_first && same_last) { + logical_first = first_dims[0] * static_cast(num_tensors); + logical_last = last_dims[0]; + } else if (same_first && !same_last) { + logical_first = first_dims[0]; + logical_last = std::accumulate(last_dims.begin(), last_dims.end(), int64_t{0}); + } else if (!same_first && same_last) { + logical_first = std::accumulate(first_dims.begin(), first_dims.end(), int64_t{0}); + logical_last = last_dims[0]; + } else { + logical_first = 1; + logical_last = 0; + for (size_t i = 0; i < num_tensors; ++i) { + logical_last += first_dims[i] * last_dims[i]; + } + } + size_t logical_data[2] = {static_cast(logical_first), + static_cast(logical_last)}; + grouped.logical_shape = nvte_make_shape(logical_data, 2); + grouped.handle.reset(nvte_create_grouped_tensor(scaling_mode, num_tensors, grouped.logical_shape)); + + const int64_t last_idx = static_cast(num_tensors - 1); + 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; + + 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, + tensors[i]->rowwise_dptr(), + grouped.tensor_bytes[i], + cudaMemcpyDeviceToDevice)); + } + + NVTEBasicTensor data_tensor{grouped.data.get(), static_cast(dtype), grouped.logical_shape}; + NVTEGroupedTensor h = grouped.handle.get(); + nvte_set_grouped_tensor_param(&h, kNVTEGroupedRowwiseData, &data_tensor); + + const bool include_columnwise = isFp8Type(dtype) || isFp4Type(dtype); + if (include_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, + tensors[i]->columnwise_dptr(), + grouped.tensor_bytes[i], + cudaMemcpyDeviceToDevice)); + } + NVTEBasicTensor col_tensor{grouped.columnwise_data.get(), + static_cast(dtype), + grouped.logical_shape}; + nvte_set_grouped_tensor_param(&h, kNVTEGroupedColumnwiseData, &col_tensor); + } + + if (!same_first) { + grouped.first_dims_dev = cuda_alloc(num_tensors * sizeof(int64_t)); + NVTE_CHECK_CUDA(cudaMemcpy(grouped.first_dims_dev.get(), first_dims.data(), + num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTEShape fd_shape = nvte_make_shape(&num_tensors, 1); + NVTEBasicTensor fd_tensor{grouped.first_dims_dev.get(), kNVTEInt64, fd_shape}; + nvte_set_grouped_tensor_param(&h, kNVTEGroupedFirstDims, &fd_tensor); + } + + if (!same_last) { + grouped.last_dims_dev = cuda_alloc(num_tensors * sizeof(int64_t)); + NVTE_CHECK_CUDA(cudaMemcpy(grouped.last_dims_dev.get(), last_dims.data(), + num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTEShape ld_shape = nvte_make_shape(&num_tensors, 1); + NVTEBasicTensor ld_tensor{grouped.last_dims_dev.get(), kNVTEInt64, ld_shape}; + nvte_set_grouped_tensor_param(&h, kNVTEGroupedLastDims, &ld_tensor); + } + + if (!same_first || !same_last) { + grouped.offsets_dev = cuda_alloc(num_tensors * 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); + NVTEBasicTensor off_tensor{grouped.offsets_dev.get(), kNVTEInt64, off_shape}; + nvte_set_grouped_tensor_param(&h, kNVTEGroupedTensorOffsets, &off_tensor); + } + + if (isFp8Type(dtype)) { + std::vector scale_inv_cpu(num_tensors, 1.f); + for (size_t i = 0; i < num_tensors; ++i) { + tensors[i]->to_cpu(); + scale_inv_cpu[i] = tensors[i]->rowwise_cpu_scale_inv_ptr()[0]; + } + grouped.scale_inv = cuda_alloc(sizeof(float) * num_tensors); + NVTE_CHECK_CUDA(cudaMemcpy(grouped.scale_inv.get(), scale_inv_cpu.data(), + sizeof(float) * num_tensors, cudaMemcpyHostToDevice)); + NVTEShape scale_shape = nvte_make_shape(&num_tensors, 1); + NVTEBasicTensor scale_tensor{grouped.scale_inv.get(), kNVTEFloat32, scale_shape}; + nvte_set_grouped_tensor_param(&h, kNVTEGroupedRowwiseScaleInv, &scale_tensor); + nvte_set_grouped_tensor_param(&h, kNVTEGroupedColumnwiseScaleInv, &scale_tensor); + } + + return grouped; +} + } // namespace test diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index b528a79b4f..082677c978 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -504,6 +504,60 @@ 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); } +}; +struct GroupedTensorDeleter { + void operator()(NVTEGroupedTensor h) const { if (h) nvte_destroy_grouped_tensor(h); } +}; + +template +using CudaPtr = std::unique_ptr; +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. +struct GroupedBuffers { + GroupedTensorHandle handle; + CudaPtr<> data; + CudaPtr<> scale_inv; + CudaPtr first_dims_dev; + CudaPtr last_dims_dev; + CudaPtr offsets_dev; + CudaPtr<> columnwise_data; + NVTEShape logical_shape{}; + std::vector offsets_host; + std::vector tensor_bytes; + size_t num_tensors{0}; + size_t elem_size{0}; + DType dtype{DType::kFloat32}; + NVTEScalingMode scaling_mode{NVTE_DELAYED_TENSOR_SCALING}; + + GroupedBuffers() = default; + GroupedBuffers(const GroupedBuffers&) = delete; + GroupedBuffers& operator=(const GroupedBuffers&) = delete; + GroupedBuffers(GroupedBuffers&&) = default; + GroupedBuffers& operator=(GroupedBuffers&&) = default; + ~GroupedBuffers() = default; + + // Convenience accessors for raw pointers + NVTEGroupedTensor get_handle() const { return handle.get(); } + void* get_data() const { return data.get(); } +}; + +GroupedBuffers build_grouped_tensor(const std::vector& tensors, + const NVTEScalingMode scaling_mode); + } // namespace test #if FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a83cbe3e30..efe958f844 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -144,6 +144,7 @@ list(APPEND transformer_engine_cuda_sources fused_attn/fused_attn_fp8.cu fused_attn/utils.cu gemm/cublaslt_gemm.cu + gemm/cublaslt_grouped_gemm.cu normalization/layernorm/ln_bwd_semi_cuda_kernel.cu normalization/layernorm/ln_fwd_cuda_kernel.cu normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu diff --git a/transformer_engine/common/gemm/config.cpp b/transformer_engine/common/gemm/config.cpp index 2532e96bb8..286fc0cc96 100644 --- a/transformer_engine/common/gemm/config.cpp +++ b/transformer_engine/common/gemm/config.cpp @@ -126,3 +126,106 @@ void nvte_destroy_matmul_config(NVTEMatmulConfig config) { delete reinterpret_cast(config); } } + +NVTEGroupedMatmulConfig nvte_create_grouped_matmul_config() { + return new transformer_engine::GroupedMatmulConfig; +} + +void nvte_get_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, + NVTEGroupedMatmulConfigAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written) { + // Write attribute size + NVTE_CHECK(attr < kNVTEGroupedMatmulConfigNumAttributes, + "Invalid NVTEGroupedMatmulConfigAttribute (got ", static_cast(attr), ")"); + NVTE_CHECK(size_written != nullptr, "Invalid size_written (got NULL)"); + const auto &attr_size = transformer_engine::GroupedMatmulConfig::attr_sizes[attr]; + *size_written = attr_size; + + // Return immediately if buffer is not provided + if (buf == nullptr) { + return; + } + + // Check buffer size + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for grouped matmul config attribute " + "(attribute ", + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, + " bytes)"); + + // Write to buffer + NVTE_CHECK(config != nullptr, "Invalid NVTEGroupedMatmulConfig (got NULL)"); + const auto &config_ = *reinterpret_cast(config); + switch (attr) { + case kNVTEGroupedMatmulConfigAvgM: { + int64_t val = config_.avg_m.value_or(0); + std::memcpy(buf, &val, attr_size); + break; + } + case kNVTEGroupedMatmulConfigAvgN: { + int64_t val = config_.avg_n.value_or(0); + std::memcpy(buf, &val, attr_size); + break; + } + case kNVTEGroupedMatmulConfigAvgK: { + int64_t val = config_.avg_k.value_or(0); + std::memcpy(buf, &val, attr_size); + break; + } + case kNVTEGroupedMatmulConfigSMCount: + std::memcpy(buf, &config_.sm_count, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEGroupedMatmulConfigAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_set_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, + NVTEGroupedMatmulConfigAttribute attr, + const void *buf, size_t size_in_bytes) { + // Check attribute and buffer + NVTE_CHECK(attr < kNVTEGroupedMatmulConfigNumAttributes, + "Invalid NVTEGroupedMatmulConfigAttribute (got ", static_cast(attr), ")"); + const auto &attr_size = transformer_engine::GroupedMatmulConfig::attr_sizes[attr]; + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for grouped matmul config attribute " + "(attribute ", + static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, + " bytes)"); + NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + + // Read from buffer + NVTE_CHECK(config != nullptr, "Invalid NVTEGroupedMatmulConfig (got NULL)"); + auto &config_ = *reinterpret_cast(config); + switch (attr) { + case kNVTEGroupedMatmulConfigAvgM: { + int64_t val; + std::memcpy(&val, buf, attr_size); + config_.avg_m = val; + break; + } + case kNVTEGroupedMatmulConfigAvgN: { + int64_t val; + std::memcpy(&val, buf, attr_size); + config_.avg_n = val; + break; + } + case kNVTEGroupedMatmulConfigAvgK: { + int64_t val; + std::memcpy(&val, buf, attr_size); + config_.avg_k = val; + break; + } + case kNVTEGroupedMatmulConfigSMCount: + std::memcpy(&config_.sm_count, buf, attr_size); + break; + default: + NVTE_ERROR("Unsupported NVTEGroupedMatmulConfigAttribute (got ", static_cast(attr), ")"); + } +} + +void nvte_destroy_grouped_matmul_config(NVTEGroupedMatmulConfig config) { + if (config != nullptr) { + delete reinterpret_cast(config); + } +} diff --git a/transformer_engine/common/gemm/config.h b/transformer_engine/common/gemm/config.h index 86a617b5fe..ad38e88334 100644 --- a/transformer_engine/common/gemm/config.h +++ b/transformer_engine/common/gemm/config.h @@ -9,6 +9,9 @@ #include +#include +#include + namespace transformer_engine { struct MatmulConfig { @@ -31,6 +34,22 @@ struct MatmulConfig { }; }; +struct GroupedMatmulConfig { + // Average dimension hints for cuBLASLt algorithm selection heuristics. + // nullopt means "not set" - compute automatically from tensor shapes. + std::optional avg_m; + std::optional avg_n; + std::optional avg_k; + + // Number of streaming multiprocessors to use in GEMM kernel + int sm_count = 0; + + // Note: API transfers the value type, not std::optional + static constexpr size_t attr_sizes[] = {sizeof(decltype(avg_m)::value_type), + sizeof(decltype(avg_n)::value_type), + sizeof(decltype(avg_k)::value_type), sizeof(sm_count)}; +}; + } // namespace transformer_engine #endif // TRANSFORMER_ENGINE_GEMM_CONFIG_H_ diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 02faad40d3..e4e97abd91 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -302,13 +302,6 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla return ret; } -/* cuBLAS version number at run-time */ -size_t cublas_version() { - // Cache version to avoid cuBLAS logging overhead - static size_t version = cublasLtGetVersion(); - return version; -} - } // namespace namespace transformer_engine { @@ -501,8 +494,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #endif // CUBLAS_VERSION >= 120800 } else if (mxfp8_gemm) { #if CUBLAS_VERSION >= 120800 - NVTE_CHECK(cublas_version() >= 120800, - "MXFP8 requires cuBLAS 12.8+, but run-time cuBLAS version is ", cublas_version()); + NVTE_CHECK(cuda::cublas_version() >= 120800, + "MXFP8 requires cuBLAS 12.8+, but run-time cuBLAS version is ", + cuda::cublas_version()); // Check that scales are in expected format NVTE_CHECK(inputA->with_gemm_swizzled_scales, @@ -524,7 +518,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, // Workaround for heuristic cache bug in cublasLt. This separates the MXFP8 cache key from non-block scaling. // CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE is unused for block scaling so it's safe to set. - if (cublas_version() <= 120803) { + if (cuda::cublas_version() <= 120803) { const int64_t dummy_a_vec_stride = 1; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE, &dummy_a_vec_stride, @@ -536,8 +530,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #endif // CUBLAS_VERSION >= 120800 } else if (use_fp4) { // NVFP4 GEMM #if CUBLAS_VERSION >= 120800 - NVTE_CHECK(cublas_version() >= 120800, - "FP4 requires cuBLAS 12.8+, but run-time cuBLAS version is ", cublas_version()); + NVTE_CHECK(cuda::cublas_version() >= 120800, + "FP4 requires cuBLAS 12.8+, but run-time cuBLAS version is ", + cuda::cublas_version()); // Check that scales are in expected format NVTE_CHECK(inputA->with_gemm_swizzled_scales, @@ -572,9 +567,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, (inputB->scaling_mode == NVTE_BLOCK_SCALING_1D || inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)) { #if CUBLAS_VERSION >= 120900 - NVTE_CHECK(cublas_version() >= 120900, + NVTE_CHECK(cuda::cublas_version() >= 120900, "FP8 block scaling requires cuBLAS 12.9+, but run-time cuBLAS version is ", - cublas_version()); + cuda::cublas_version()); // Check that matrix formats are valid NVTE_CHECK((!(inputA->scaling_mode == NVTE_BLOCK_SCALING_2D && @@ -607,7 +602,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, } #if CUBLAS_VERSION >= 120800 - if (cublas_version() >= 120800) { + if (cuda::cublas_version() >= 120800) { NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scaling_mode_a, sizeof(scaling_mode_a))); @@ -624,7 +619,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_AMAX_D_POINTER, &D_amax, sizeof(D_amax))); #if CUBLAS_VERSION >= 120800 - if (cublas_version() >= 120800) { + if (cuda::cublas_version() >= 120800) { // NOTE: In all current cases where FP8 output is supported, the input is // scaled identically to the output. NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, @@ -711,9 +706,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK(cuda::cudart_version() >= 12020 && cuda::cudart_version() < 13000, "Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but run-time CUDA version is ", cuda::cudart_version()); - NVTE_CHECK(cublas_version() >= 120205 && cublas_version() < 130000, + NVTE_CHECK(cuda::cublas_version() >= 120205 && cuda::cublas_version() < 130000, "Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but run-time cuBLAS version is ", - cublas_version()); + cuda::cublas_version()); if (m_split == 0) m_split = 1; if (n_split == 0) n_split = 1; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( @@ -939,9 +934,9 @@ void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor "Atomic GEMM requires CUDA version >=12.2.0 and <13.0.0, but run-time CUDA version is ", transformer_engine::cuda::cudart_version()); NVTE_CHECK( - cublas_version() >= 120205 && cublas_version() < 130000, + cuda::cublas_version() >= 120205 && cuda::cublas_version() < 130000, "Atomic GEMM requires cuBLAS version >=12.2.5 and <13.0.0, but run-time cuBLAS version is ", - cublas_version()); + cuda::cublas_version()); const Tensor *inputA = convertNVTETensorCheck(A); const Tensor *inputB = convertNVTETensorCheck(B); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu new file mode 100644 index 0000000000..a1206474ea --- /dev/null +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -0,0 +1,645 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include + +#include "../common.h" +#include "../util/cuda_runtime.h" +#include "../util/handle_manager.h" +#include "../util/logging.h" +#include "./config.h" + +namespace { + +inline void CreateCublasHandle(cublasLtHandle_t *handle) { + NVTE_CHECK_CUBLAS(cublasLtCreate(handle)); +} + +} // namespace + +#if CUBLAS_VERSION >= 130100 + +namespace { + +// Helper struct to pass per-tensor shape/offset info (pointer or uniform value) +struct TensorShapeInfo { + const int64_t *first_dims; // nullptr if uniform + const int64_t *last_dims; // nullptr if uniform + const int64_t *offsets; // nullptr if need to compute + int64_t uniform_first; // used if first_dims == nullptr + int64_t uniform_last; // used if last_dims == nullptr + + // Create from GroupedTensor + static TensorShapeInfo from_tensor(const transformer_engine::GroupedTensor *t) { + const bool has_first = t->first_dims.has_data(); + const bool has_last = t->last_dims.has_data(); + // When per-tensor dims are not provided, we must be in the uniform-shape case. + NVTE_CHECK(has_first || t->all_same_first_dim(), + "GroupedTensor is missing first_dims for varying shapes"); + NVTE_CHECK(has_last || t->all_same_last_dim(), + "GroupedTensor is missing last_dims for varying shapes"); + + const int64_t *first_ptr = + has_first ? static_cast(t->first_dims.dptr) : nullptr; + const int64_t *last_ptr = has_last ? static_cast(t->last_dims.dptr) : nullptr; + + const int64_t uniform_first = has_first ? 0 : static_cast(t->get_common_first_dim()); + const int64_t uniform_last = has_last ? 0 : static_cast(t->get_common_last_dim()); + + return {first_ptr, last_ptr, + t->tensor_offsets.has_data() ? static_cast(t->tensor_offsets.dptr) + : nullptr, + uniform_first, uniform_last}; + } + + // Create for C tensor (uses D's dimensions, only has offsets) + static TensorShapeInfo create_shape_info_for_C(const transformer_engine::GroupedTensor *C, + const transformer_engine::GroupedTensor *D) { + const bool has_first = D->first_dims.has_data(); + const bool has_last = D->last_dims.has_data(); + NVTE_CHECK(has_first || D->all_same_first_dim(), + "GroupedTensor D is missing first_dims for varying shapes"); + NVTE_CHECK(has_last || D->all_same_last_dim(), + "GroupedTensor D is missing last_dims for varying shapes"); + + const int64_t *first_ptr = + has_first ? static_cast(D->first_dims.dptr) : nullptr; + const int64_t *last_ptr = has_last ? static_cast(D->last_dims.dptr) : nullptr; + const int64_t uniform_first = has_first ? 0 : static_cast(D->get_common_first_dim()); + const int64_t uniform_last = has_last ? 0 : static_cast(D->get_common_last_dim()); + + return {first_ptr, last_ptr, + C->tensor_offsets.has_data() ? static_cast(C->tensor_offsets.dptr) + : nullptr, + uniform_first, uniform_last}; + } +}; + +// Helper functions to compute average dimensions from logical_shape for heuristics +// These are hints for cuBLASLt algorithm selection, don't need to be exact +inline int64_t compute_avg_first_dim(const transformer_engine::GroupedTensor *t) { + // logical_shape[0] is either num_tensors*M (uniform) or sum_of_M (varying first) + // In both cases, dividing by num_tensors gives the average + return static_cast(t->logical_shape.data[0]) / static_cast(t->num_tensors); +} + +inline int64_t compute_avg_last_dim(const transformer_engine::GroupedTensor *t) { + if (t->all_same_last_dim()) { + // logical_shape[1] is the common N + return static_cast(t->logical_shape.data[1]); + } + // When varying, logical_shape[1] should be sum of last dims if provided; otherwise fallback to avg via division. + return static_cast(t->logical_shape.data[1]) / static_cast(t->num_tensors); +} + +// Workspace layout for grouped GEMM +struct GroupedGemmSetupWorkspace { + void **A_ptrs; + void **B_ptrs; + void **C_ptrs; + void **D_ptrs; + float **alpha_ptrs; + float **beta_ptrs; + // 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 (8-byte aligned), then int arrays (4-byte aligned) + static GroupedGemmSetupWorkspace from_buffers(char *setup_ws_ptr, size_t num_tensors) { + GroupedGemmSetupWorkspace ws; + size_t offset = 0; + const size_t ptr_size = num_tensors * sizeof(void *); + const size_t int_size = num_tensors * sizeof(int); + + // Pointer arrays first (all 8-byte aligned) + ws.A_ptrs = reinterpret_cast(setup_ws_ptr + offset); + offset += ptr_size; + ws.B_ptrs = reinterpret_cast(setup_ws_ptr + offset); + offset += ptr_size; + ws.C_ptrs = reinterpret_cast(setup_ws_ptr + offset); + offset += ptr_size; + ws.D_ptrs = reinterpret_cast(setup_ws_ptr + offset); + offset += ptr_size; + ws.alpha_ptrs = reinterpret_cast(setup_ws_ptr + offset); + offset += ptr_size; + ws.beta_ptrs = reinterpret_cast(setup_ws_ptr + offset); + offset += ptr_size; + + // Int arrays for storage dimensions (4-byte aligned) + 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); + + 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); + // Layout: 6 ptr arrays, then 6 int arrays + size_t size = 6 * ptr_size + 6 * int_size; + size = ((size + alignment - 1) / alignment) * alignment; + return size; + } +}; + +// ----------------------------------------------------------------------------- +// Helper routines to keep nvte_grouped_gemm readable +// ----------------------------------------------------------------------------- +inline void validate_grouped_gemm_inputs(const transformer_engine::GroupedTensor *inputA, + const transformer_engine::GroupedTensor *inputB, + const transformer_engine::GroupedTensor *inputC, + const transformer_engine::GroupedTensor *outputD, + const transformer_engine::Tensor *alpha_tensor, + const transformer_engine::Tensor *beta_tensor) { + const size_t num_tensors = inputA->num_tensors; + NVTE_CHECK(num_tensors >= 1, "Grouped GEMM: number of tensors must be at least 1"); + NVTE_CHECK(inputB->num_tensors == num_tensors, + "Grouped GEMM: A and B must have the same number of tensors"); + // C can be NULL (will use D as C when beta=0) + if (inputC != nullptr) { + NVTE_CHECK(inputC->num_tensors == num_tensors, + "Grouped GEMM: A and C must have the same number of tensors"); + } + NVTE_CHECK(outputD->num_tensors == num_tensors, + "Grouped GEMM: A and D must have the same number of tensors"); + + // Validate alpha/beta have per-matrix values + 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); + + auto is_fp8_or_16bit = [](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; + }; + 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_fp8_or_16bit(inputA->dtype()) && is_fp8_or_16bit(inputB->dtype()), + "Grouped GEMM inputs must be FP8, BF16, or FP16."); + // Only check C dtype if C is provided + if (inputC != nullptr) { + NVTE_CHECK(is_output_dtype(inputC->dtype()), "Grouped GEMM: C must be BF16, FP16, or FP32."); + } + NVTE_CHECK(is_output_dtype(outputD->dtype()), "Grouped GEMM: D must be BF16, FP16, or FP32."); + NVTE_CHECK(inputA->has_data() || inputA->has_columnwise_data(), + "Grouped GEMM: A tensor is missing both row-wise and column-wise data"); + NVTE_CHECK(inputB->has_data() || inputB->has_columnwise_data(), + "Grouped GEMM: B tensor is missing both row-wise and column-wise data"); +} + +// 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. +struct GroupedOperandSelection { + TensorShapeInfo shape; // Shape info with dims already swapped for columnwise if needed + char *dptr = nullptr; + void *scale_inv = nullptr; + transformer_engine::DType dtype = transformer_engine::DType::kNumTypes; + bool trans = false; +}; + +// 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) { + 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(), + "GroupedTensor is missing first_dims for varying shapes"); + NVTE_CHECK(has_last || t->all_same_last_dim(), + "GroupedTensor is missing last_dims for varying shapes"); + + const int64_t *first_ptr = has_first ? static_cast(t->first_dims.dptr) : nullptr; + const int64_t *last_ptr = has_last ? static_cast(t->last_dims.dptr) : nullptr; + const int64_t uniform_first = has_first ? 0 : static_cast(t->get_common_first_dim()); + const int64_t uniform_last = has_last ? 0 : static_cast(t->get_common_last_dim()); + + 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}; +} + +inline GroupedOperandSelection select_grouped_operand(const transformer_engine::GroupedTensor *t, + bool trans, bool is_A) { + using namespace transformer_engine; + const bool has_row = t->has_data(); + const bool has_col = t->has_columnwise_data(); + NVTE_CHECK(has_row || has_col, + "Grouped GEMM operand is missing both row-wise and column-wise data"); + + // Currently only unquantized data and tensor-scaled FP8 are supported. + const auto sm = t->scaling_mode; + NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING, + "Grouped GEMM is only supported with unquantized data and tensor-scaled FP8 data"); + + const DType row_dtype = t->data.dtype; + const DType col_dtype = t->columnwise_data.dtype; + GroupedOperandSelection sel; + sel.trans = trans; + + 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(); + + // Helper to select columnwise storage (swaps dims in shape) + auto use_columnwise = [&]() { + sel.dptr = static_cast(t->columnwise_data.dptr); + sel.scale_inv = t->columnwise_scale_inv.dptr; + sel.dtype = col_dtype; + sel.shape = create_shape_info(t, /*swap_dims=*/true); + }; + + // Helper to select row-wise storage + auto use_rowwise = [&]() { + sel.dptr = static_cast(t->data.dptr); + sel.scale_inv = t->scale_inv.dptr; + sel.dtype = row_dtype; + sel.shape = create_shape_info(t, /*swap_dims=*/false); + }; + + // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed. + if (is_fp8 && !non_tn_fp8_ok) { + if (is_A) { + if (!sel.trans) { + NVTE_CHECK(has_col, "Grouped GEMM: A is missing column-wise data needed for FP8 TN layout"); + use_columnwise(); + sel.trans = true; // using pre-transposed storage + return sel; + } + } else { // B + if (sel.trans) { + NVTE_CHECK(has_col, "Grouped GEMM: B is missing column-wise data needed for FP8 TN layout"); + use_columnwise(); + sel.trans = false; // using pre-transposed storage + return sel; + } + } + } + + // If only column-wise data is available, mirror the transpose flag (pre-transposed storage). + if (!has_row && has_col) { + // On Hopper FP8, this would break TN requirement - should have been handled above + NVTE_CHECK( + !is_fp8 || non_tn_fp8_ok, + "Grouped GEMM: FP8 on Hopper requires row-wise data for this transpose configuration"); + use_columnwise(); + sel.trans = !trans; // flip transpose for pre-transposed storage + return sel; + } + + // Default: use row-wise data + use_rowwise(); + return sel; +} + +inline void *validate_and_get_workspace_ptr(transformer_engine::Tensor *ws, size_t required_size, + const char *workspace_name) { + NVTE_CHECK(ws != nullptr, workspace_name, " tensor is null."); + const size_t provided_size = get_buffer_size_bytes(ws->data.numel(), ws->data.dtype); + NVTE_CHECK(provided_size >= required_size, "Grouped GEMM: Insufficient ", workspace_name, + ". Required: ", required_size, " bytes, Available: ", provided_size, " bytes."); + return ws->data.dptr; +} + +inline void init_matrix_layouts(cublasLtMatrixLayoutOpaque_t &descA, + cublasLtMatrixLayoutOpaque_t &descB, + cublasLtMatrixLayoutOpaque_t &descC, + cublasLtMatrixLayoutOpaque_t &descD, + const GroupedGemmSetupWorkspace &ws, + const GroupedOperandSelection &A_sel, + const GroupedOperandSelection &B_sel, + const transformer_engine::GroupedTensor *D, size_t num_tensors) { + const cudaDataType_t A_type = get_cuda_dtype(A_sel.dtype); + const cudaDataType_t B_type = get_cuda_dtype(B_sel.dtype); + const cudaDataType_t D_type = get_cuda_dtype(D->dtype()); + + // Storage dimensions computed by kernel, leading dimension = rows + NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descA, A_type, num_tensors, ws.a_rows, + ws.a_cols, ws.a_rows)); + NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descB, B_type, num_tensors, ws.b_rows, + ws.b_cols, ws.b_rows)); + NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descC, D_type, num_tensors, ws.d_rows, + ws.d_cols, ws.d_rows)); + NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descD, D_type, num_tensors, ws.d_rows, + ws.d_cols, ws.d_rows)); +} + +inline void init_matmul_desc(cublasLtMatmulDescOpaque_t &matmulDesc, cublasOperation_t op_A, + cublasOperation_t op_B) { + NVTE_CHECK_CUBLAS(cublasLtMatmulDescInit(&matmulDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); + + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &op_A, + sizeof(op_A))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_TRANSB, &op_B, + sizeof(op_B))); + + cublasLtPointerMode_t pointer_mode = CUBLASLT_POINTER_MODE_DEVICE; + 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))); +} + +inline void set_fp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, + const GroupedOperandSelection &A_sel, + const GroupedOperandSelection &B_sel) { + const bool is_fp8_a = is_fp8_dtype(A_sel.dtype); + const bool is_fp8_b = is_fp8_dtype(B_sel.dtype); + if (!is_fp8_a && !is_fp8_b) return; + + if (is_fp8_a) { + void *a_scale_inv = A_sel.scale_inv; + NVTE_CHECK(a_scale_inv != nullptr, "FP8 grouped GEMM: A scale_inv is required"); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( + &matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale_inv, sizeof(a_scale_inv))); + } + if (is_fp8_b) { + void *b_scale_inv = B_sel.scale_inv; + NVTE_CHECK(b_scale_inv != nullptr, "FP8 grouped GEMM: B scale_inv is required"); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( + &matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale_inv, sizeof(b_scale_inv))); + } +} + +// Constants for grouped GEMM workspace (declared early for use in heuristics) +static constexpr size_t kGroupedGemmAlignment = 256; +static constexpr size_t kGroupedGemmCublasWorkspaceSize = 32ull * 1024 * 1024; // 32 MiB + +inline cublasLtMatmulAlgo_t select_grouped_gemm_algo(cublasLtHandle_t handle, + cublasLtMatmulDescOpaque_t &matmulDesc, + cublasLtMatrixLayoutOpaque_t &descA, + cublasLtMatrixLayoutOpaque_t &descB, + cublasLtMatrixLayoutOpaque_t &descC, + cublasLtMatrixLayoutOpaque_t &descD, + int64_t avg_m, int64_t avg_n, int64_t avg_k) { + cublasLtMatmulPreferenceOpaque_t preference; + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceInit(&preference)); + NVTE_CHECK_CUBLAS( + cublasLtMatmulPreferenceSetAttribute(&preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, + &kGroupedGemmCublasWorkspaceSize, sizeof(size_t))); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + &preference, CUBLASLT_MATMUL_PREF_GROUPED_DESC_D_AVERAGE_ROWS, &avg_m, sizeof(int64_t))); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + &preference, CUBLASLT_MATMUL_PREF_GROUPED_DESC_D_AVERAGE_COLS, &avg_n, sizeof(int64_t))); + NVTE_CHECK_CUBLAS(cublasLtMatmulPreferenceSetAttribute( + &preference, CUBLASLT_MATMUL_PREF_GROUPED_AVERAGE_REDUCTION_DIM, &avg_k, sizeof(int64_t))); + + cublasLtMatmulHeuristicResult_t heuristicResult; + int returnedResults = 0; + auto status = cublasLtMatmulAlgoGetHeuristic(handle, &matmulDesc, &descA, &descB, &descC, &descD, + &preference, 1, &heuristicResult, &returnedResults); + NVTE_CHECK(status != CUBLAS_STATUS_NOT_SUPPORTED, + "Unable to find suitable cuBLAS grouped GEMM algorithm"); + NVTE_CHECK_CUBLAS(status); + NVTE_CHECK(returnedResults > 0, "No suitable algorithm found for grouped GEMM"); + return heuristicResult.algo; +} + +// Single kernel that sets up all GEMM parameters. +// Rationale: cuBLASLt grouped matmul API needs flat arrays of pointers and per-matrix dimensions, +// but NVTEGroupedTensor stores a single contiguous buffer + optional per-tensor offsets/shapes. +// We bridge the mismatch on GPU by computing per-group pointers and storage dims in one kernel. +__global__ void setup_grouped_gemm_kernel( + // Output arrays + void **A_ptrs, void **B_ptrs, void **C_ptrs, void **D_ptrs, int *a_rows, int *a_cols, + int *b_rows, int *b_cols, int *d_rows, int *d_cols, float **alpha_ptrs, float **beta_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, + size_t num_tensors) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_tensors) return; + + // Get dimensions for this tensor (from array or uniform value) + int64_t a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first; + int64_t a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last; + int64_t b_first = B_meta.first_dims ? B_meta.first_dims[idx] : B_meta.uniform_first; + int64_t b_last = B_meta.last_dims ? B_meta.last_dims[idx] : B_meta.uniform_last; + int64_t d_first = D_meta.first_dims ? D_meta.first_dims[idx] : D_meta.uniform_first; + int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last; + + // Compute offsets (from array or compute from uniform dims) + int64_t a_offset = + A_meta.offsets ? A_meta.offsets[idx] : (idx * A_meta.uniform_first * A_meta.uniform_last); + int64_t b_offset = + B_meta.offsets ? B_meta.offsets[idx] : (idx * B_meta.uniform_first * B_meta.uniform_last); + int64_t c_offset = + C_meta.offsets ? C_meta.offsets[idx] : (idx * C_meta.uniform_first * C_meta.uniform_last); + int64_t d_offset = + D_meta.offsets ? D_meta.offsets[idx] : (idx * D_meta.uniform_first * D_meta.uniform_last); + + // Compute data pointers + A_ptrs[idx] = a_base + a_offset * a_elem_size; + B_ptrs[idx] = b_base + b_offset * b_elem_size; + C_ptrs[idx] = c_base + c_offset * c_elem_size; + D_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); + // For OUTPUTS (D, C): cuBLAS writes in column-major, so rows=first (M), cols=last (N). + d_rows[idx] = static_cast(d_first); + d_cols[idx] = static_cast(d_last); + + // Fill alpha/beta pointers (per-matrix) + alpha_ptrs[idx] = alpha_ptr + idx; + beta_ptrs[idx] = beta_ptr + idx; +} + +// Launch the setup kernel to populate workspace arrays +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) { + // Use shape info from selection (already accounts for columnwise dimension swap) + TensorShapeInfo A_meta = A_sel.shape; + TensorShapeInfo B_meta = B_sel.shape; + TensorShapeInfo C_meta = TensorShapeInfo::create_shape_info_for_C(C, D); + TensorShapeInfo D_meta = TensorShapeInfo::from_tensor(D); + + char *c_base = static_cast(C->data.dptr); + char *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 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; + + 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, 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), + num_tensors); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +inline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) { + return GroupedGemmSetupWorkspace::required_setup_size(num_tensors, kGroupedGemmAlignment); +} + +} // namespace + +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, + NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_grouped_gemm); + using namespace transformer_engine; + + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.1+ + const int current_device = cuda::current_device(); + NVTE_CHECK(cuda::sm_arch(current_device) >= 100, + "nvte_grouped_gemm requires Blackwell (SM100) or newer architecture."); + NVTE_CHECK(cuda::cublas_version() >= 130100, + "nvte_grouped_gemm requires cuBLAS 13.1+, but run-time cuBLAS version is ", + cuda::cublas_version()); + + // Convert to internal types + const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A); + const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B); + const GroupedTensor *inputC_raw = convertNVTEGroupedTensor(C); // Can be NULL + GroupedTensor *outputD = convertNVTEGroupedTensorCheck(D); + const Tensor *alpha_tensor = convertNVTETensorCheck(alpha); + const Tensor *beta_tensor = convertNVTETensorCheck(beta); + Tensor *wspace_setup = convertNVTETensor(workspace_setup); + Tensor *wspace_cublas = convertNVTETensor(workspace_cublas); + + // Parse config (if provided) + GroupedMatmulConfig config_; + if (config != nullptr) { + config_ = *reinterpret_cast(config); + } + + // Validate inputs and num_tensors + validate_grouped_gemm_inputs(inputA, inputB, inputC_raw, outputD, alpha_tensor, beta_tensor); + + // 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 = inputA->num_tensors; + + // Select operand storage (row-wise vs column-wise) and adjust transpose flags to + // mirror the non-grouped GEMM logic for FP8 layout constraints. + const auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true); + const auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false); + + // Workspaces: setup (pointer arrays) and cuBLAS + const size_t setup_workspace_size = grouped_gemm_setup_workspace_size(num_tensors); + const size_t cublas_workspace_size = kGroupedGemmCublasWorkspaceSize; + + void *setup_workspace_ptr = validate_and_get_workspace_ptr(wspace_setup, setup_workspace_size, + "Grouped GEMM setup workspace"); + void *cublas_workspace_ptr = validate_and_get_workspace_ptr(wspace_cublas, cublas_workspace_size, + "Grouped GEMM cuBLAS workspace"); + + auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers( + static_cast(setup_workspace_ptr), num_tensors); + launch_grouped_gemm_setup(setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, + beta_tensor, num_tensors, stream); + + // Get cuBLAS handle + using cublasHandleManager = detail::HandleManager; + cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); + + // Setup cuBLAS operations + cublasOperation_t op_A = A_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N; + cublasOperation_t op_B = B_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N; + + // Create grouped matrix layouts + cublasLtMatrixLayoutOpaque_t descA, descB, descC, descD; + init_matrix_layouts(descA, descB, descC, descD, setup_workspace, A_sel, B_sel, outputD, + num_tensors); + + // Create matmul descriptor + cublasLtMatmulDescOpaque_t matmulDesc; + init_matmul_desc(matmulDesc, op_A, op_B); + set_fp8_scale_pointers(matmulDesc, A_sel, B_sel); + + // 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 = + config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); + + // Heuristic selection + cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo(handle, matmulDesc, descA, descB, descC, + descD, avg_m_val, avg_n_val, avg_k_val); + + // Execute the grouped GEMM + 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)); +} + +#else // CUBLAS_VERSION < 130100 + +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, + NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, + cudaStream_t stream) { + NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.1+, but compile-time cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); +} + +#endif // CUBLAS_VERSION >= 130100 diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index b304ed34be..1afc9828e8 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -11,6 +11,8 @@ #ifndef TRANSFORMER_ENGINE_GEMM_H_ #define TRANSFORMER_ENGINE_GEMM_H_ +#include + #include "transformer_engine.h" #ifdef __cplusplus @@ -20,6 +22,9 @@ extern "C" { /*! \brief Configuration for matrix multiplication. */ typedef void *NVTEMatmulConfig; +/*! \brief Configuration for grouped matrix multiplication. */ +typedef void *NVTEGroupedMatmulConfig; + /*! \enum NVTEMatmulConfigAttribute * \brief Type of option for matrix multiplication. */ @@ -52,6 +57,36 @@ enum NVTEMatmulConfigAttribute { kNVTEMatmulConfigNumAttributes }; +/*! \enum NVTEGroupedMatmulConfigAttribute + * \brief Type of option for grouped matrix multiplication. + */ +enum NVTEGroupedMatmulConfigAttribute { + /*! Average M dimension hint + * + * Optional hint for average M dimension across all matrices in the group. + * Used by cuBLASLt for algorithm selection heuristics. If not set, + * computed automatically from D's logical shape. + */ + kNVTEGroupedMatmulConfigAvgM = 0, + /*! Average N dimension hint + * + * Optional hint for average N dimension across all matrices in the group. + * Used by cuBLASLt for algorithm selection heuristics. If not set, + * computed automatically from D's logical shape. + */ + kNVTEGroupedMatmulConfigAvgN = 1, + /*! Average K (reduction) dimension hint + * + * Optional hint for average K dimension across all matrices in the group. + * Used by cuBLASLt for algorithm selection heuristics. If not set, + * computed automatically from A's logical shape. + */ + kNVTEGroupedMatmulConfigAvgK = 2, + /*! Number of streaming multiprocessors to use in GEMM kernel. */ + kNVTEGroupedMatmulConfigSMCount = 3, + kNVTEGroupedMatmulConfigNumAttributes +}; + /*! \brief Create a matrix multiplication configuration. */ NVTEMatmulConfig nvte_create_matmul_config(); @@ -82,6 +117,38 @@ void nvte_set_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA /*! \brief Destroy a matrix multiplication configuration. */ void nvte_destroy_matmul_config(NVTEMatmulConfig config); +/*! \brief Create a grouped matrix multiplication configuration. */ +NVTEGroupedMatmulConfig nvte_create_grouped_matmul_config(); + +/*! \brief Query an option in grouped matrix multiplication configuration. + * + * \param[in] config Grouped matrix multiplication configuration. + * \param[in] attr Option type. + * \param[out] buf Memory address to write option value. Ignored if + * NULL. + * \param[in] size_in_bytes Size of buf. + * \param[out] size_written Number of bytes that have been written to + * buf. If buf is NULL, then the number of + * bytes that would have been written. + */ +void nvte_get_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, + NVTEGroupedMatmulConfigAttribute attr, void *buf, + size_t size_in_bytes, size_t *size_written); + +/*! \brief Set an option in grouped matrix multiplication configuration. + * + * \param[in] config Grouped matrix multiplication configuration. + * \param[in] attr Option type. + * \param[out] buf Memory address to read option value. + * \param[in] size_in_bytes Size of buf. + */ +void nvte_set_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, + NVTEGroupedMatmulConfigAttribute attr, + const void *buf, size_t size_in_bytes); + +/*! \brief Destroy a grouped matrix multiplication configuration. */ +void nvte_destroy_grouped_matmul_config(NVTEGroupedMatmulConfig config); + /*! \brief Compute matrix multiplication of 2 matrices, potentially fused with other operations (deprecated). * * This has been deprecated in favor of nvte_cublas_gemm_v2. @@ -228,6 +295,46 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor bool transa, bool transb, bool grad, NVTETensor *workspace, 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.1+ (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.1+ (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, + NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus @@ -331,6 +438,70 @@ class MatmulConfigWrapper { NVTEMatmulConfig config_ = nullptr; }; +/*! \struct GroupedMatmulConfigWrapper + * \brief C++ wrapper for NVTEGroupedMatmulConfig. + */ +class GroupedMatmulConfigWrapper { + public: + GroupedMatmulConfigWrapper() : config_{nvte_create_grouped_matmul_config()} {} + + GroupedMatmulConfigWrapper(const GroupedMatmulConfigWrapper &) = delete; + GroupedMatmulConfigWrapper &operator=(const GroupedMatmulConfigWrapper &) = delete; + + GroupedMatmulConfigWrapper(GroupedMatmulConfigWrapper &&other) : config_{other.config_} { + other.config_ = nullptr; + } + GroupedMatmulConfigWrapper &operator=(GroupedMatmulConfigWrapper &&other) { + if (config_ != nullptr) { + nvte_destroy_grouped_matmul_config(config_); + } + config_ = other.config_; + other.config_ = nullptr; + return *this; + } + + ~GroupedMatmulConfigWrapper() { + if (config_ != nullptr) { + nvte_destroy_grouped_matmul_config(config_); + config_ = nullptr; + } + } + + /*! \brief Get the underlying NVTEGroupedMatmulConfig. + * + * \return NVTEGroupedMatmulConfig held by this GroupedMatmulConfigWrapper. + */ + operator NVTEGroupedMatmulConfig() const noexcept { return config_; } + + /*! \brief Set average M dimension hint for algorithm selection. */ + void set_avg_m(int64_t avg_m) { + nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigAvgM, &avg_m, + sizeof(int64_t)); + } + + /*! \brief Set average N dimension hint for algorithm selection. */ + void set_avg_n(int64_t avg_n) { + nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigAvgN, &avg_n, + sizeof(int64_t)); + } + + /*! \brief Set average K dimension hint for algorithm selection. */ + void set_avg_k(int64_t avg_k) { + nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigAvgK, &avg_k, + sizeof(int64_t)); + } + + /*! \brief Set number of streaming multiprocessors to use. */ + void set_sm_count(int sm_count) { + nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigSMCount, &sm_count, + sizeof(int)); + } + + private: + /*! \brief Wrapped NVTEGroupedMatmulConfig. */ + NVTEGroupedMatmulConfig config_ = nullptr; +}; + } // namespace transformer_engine #endif // __cplusplus diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index f99900bac8..4b43940a51 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -6,6 +6,8 @@ #include "../util/cuda_runtime.h" +#include + #include #include @@ -210,6 +212,12 @@ int cudart_version() { return version; } +size_t cublas_version() { + // Cache version to avoid cuBLAS logging overhead + static size_t version = cublasLtGetVersion(); + return version; +} + } // namespace cuda } // namespace transformer_engine diff --git a/transformer_engine/common/util/cuda_runtime.h b/transformer_engine/common/util/cuda_runtime.h index c696f6b57a..f0aa239622 100644 --- a/transformer_engine/common/util/cuda_runtime.h +++ b/transformer_engine/common/util/cuda_runtime.h @@ -73,6 +73,12 @@ const std::string &include_directory(bool required = false); */ int cudart_version(); +/* \brief cuBLAS version number at run-time + * + * Versions may differ between compile-time and run-time. + */ +size_t cublas_version(); + } // namespace cuda } // namespace transformer_engine From f8cca8b95de7e5637bc52f1b55113b38a6ad0774 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, 29 Jan 2026 14:37:57 -0800 Subject: [PATCH 186/521] [Pytorch] Fix wheel test (#2635) fix wheel Signed-off-by: Pawel Gadzinski --- qa/L0_pytorch_wheel/test.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/qa/L0_pytorch_wheel/test.sh b/qa/L0_pytorch_wheel/test.sh index fcf1a52b9c..fe4aab456e 100644 --- a/qa/L0_pytorch_wheel/test.sh +++ b/qa/L0_pytorch_wheel/test.sh @@ -28,11 +28,11 @@ WHL_BASE="transformer_engine-${VERSION}" # Core wheel. NVTE_RELEASE_BUILD=1 pip3 wheel --no-build-isolation -vvv --wheel-dir ./dist . || error_exit "Failed to setup bdist_wheel" -wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl" +python3 -m wheel unpack dist/${WHL_BASE}-* || error_exit "Failed to unpack dist/${WHL_BASE}-*.whl" sed -i "s/Name: transformer-engine/Name: transformer-engine-cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" sed -i "s/Name: transformer_engine/Name: transformer_engine_cu12/g" "transformer_engine-${VERSION}/transformer_engine-${VERSION}.dist-info/METADATA" mv "${WHL_BASE}/${WHL_BASE}.dist-info" "${WHL_BASE}/transformer_engine_cu12-${VERSION}.dist-info" || error_exit "Failed to move ${WHL_BASE}.dist-info to transformer_engine_cu12-${VERSION}.dist-info" -wheel pack ${WHL_BASE} || error_exit "Failed to pack ${WHL_BASE}" +python3 -m wheel pack ${WHL_BASE} || error_exit "Failed to pack ${WHL_BASE}" rm dist/*.whl || error_exit "Failed to remove dist/*.whl" mv *.whl dist/ || error_exit "Failed to move *.whl to dist/" NVTE_RELEASE_BUILD=1 NVTE_BUILD_METAPACKAGE=1 pip3 wheel --no-build-isolation --no-deps -vvv --wheel-dir ./dist . || error_exit "Failed to setup metapackage" From c3769cb799d633af1ea741bfdaf5f05a3584fd09 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Fri, 30 Jan 2026 10:17:59 -0800 Subject: [PATCH 187/521] Fix minimum version of cublas for grouped gemm (#2631) * version change Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * ifx Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski --- tests/cpp/operator/test_grouped_gemm.cu | 6 ++-- .../common/gemm/cublaslt_gemm.cu | 33 ++++++++++--------- .../common/gemm/cublaslt_grouped_gemm.cu | 20 +++++------ .../common/include/transformer_engine/gemm.h | 4 +-- 4 files changed, 33 insertions(+), 30 deletions(-) diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu index 35c4375cbe..a694052b15 100644 --- a/tests/cpp/operator/test_grouped_gemm.cu +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -102,8 +102,8 @@ std::vector> make_shapes(ShapeCase scase) { } void run_grouped_gemm_case(const TestParams& params) { -#if CUBLAS_VERSION < 130100 - GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.1+, but compile-time cuBLAS version is " +#if CUBLAS_VERSION < 130200 + GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.2+, but compile-time cuBLAS version is " << CUBLAS_VERSION << "."; #else if (getDeviceComputeCapability() < blackwellComputeCapability) { @@ -267,7 +267,7 @@ void run_grouped_gemm_case(const TestParams& params) { atol, rtol); } -#endif // CUBLAS_VERSION >= 130100 +#endif // CUBLAS_VERSION >= 130200 } class GroupedGemmTest : public ::testing::TestWithParam {}; diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index e4e97abd91..c58c3cb47a 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -494,9 +494,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #endif // CUBLAS_VERSION >= 120800 } else if (mxfp8_gemm) { #if CUBLAS_VERSION >= 120800 - NVTE_CHECK(cuda::cublas_version() >= 120800, + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800, "MXFP8 requires cuBLAS 12.8+, but run-time cuBLAS version is ", - cuda::cublas_version()); + transformer_engine::cuda::cublas_version()); // Check that scales are in expected format NVTE_CHECK(inputA->with_gemm_swizzled_scales, @@ -518,7 +518,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, // Workaround for heuristic cache bug in cublasLt. This separates the MXFP8 cache key from non-block scaling. // CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE is unused for block scaling so it's safe to set. - if (cuda::cublas_version() <= 120803) { + if (transformer_engine::cuda::cublas_version() <= 120803) { const int64_t dummy_a_vec_stride = 1; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE, &dummy_a_vec_stride, @@ -530,9 +530,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, #endif // CUBLAS_VERSION >= 120800 } else if (use_fp4) { // NVFP4 GEMM #if CUBLAS_VERSION >= 120800 - NVTE_CHECK(cuda::cublas_version() >= 120800, + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120800, "FP4 requires cuBLAS 12.8+, but run-time cuBLAS version is ", - cuda::cublas_version()); + transformer_engine::cuda::cublas_version()); // Check that scales are in expected format NVTE_CHECK(inputA->with_gemm_swizzled_scales, @@ -567,9 +567,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, (inputB->scaling_mode == NVTE_BLOCK_SCALING_1D || inputB->scaling_mode == NVTE_BLOCK_SCALING_2D)) { #if CUBLAS_VERSION >= 120900 - NVTE_CHECK(cuda::cublas_version() >= 120900, + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120900, "FP8 block scaling requires cuBLAS 12.9+, but run-time cuBLAS version is ", - cuda::cublas_version()); + transformer_engine::cuda::cublas_version()); // Check that matrix formats are valid NVTE_CHECK((!(inputA->scaling_mode == NVTE_BLOCK_SCALING_2D && @@ -602,7 +602,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, } #if CUBLAS_VERSION >= 120800 - if (cuda::cublas_version() >= 120800) { + if (transformer_engine::cuda::cublas_version() >= 120800) { NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scaling_mode_a, sizeof(scaling_mode_a))); @@ -619,7 +619,7 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( operationDesc, CUBLASLT_MATMUL_DESC_AMAX_D_POINTER, &D_amax, sizeof(D_amax))); #if CUBLAS_VERSION >= 120800 - if (cuda::cublas_version() >= 120800) { + if (transformer_engine::cuda::cublas_version() >= 120800) { // NOTE: In all current cases where FP8 output is supported, the input is // scaled identically to the output. NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(operationDesc, @@ -703,12 +703,14 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, "Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but compile-time cuBLAS version is ", CUBLAS_VERSION); #else - NVTE_CHECK(cuda::cudart_version() >= 12020 && cuda::cudart_version() < 13000, + NVTE_CHECK(transformer_engine::cuda::cudart_version() >= 12020 && + transformer_engine::cuda::cudart_version() < 13000, "Atomic GEMM requires CUDA >=12.2.0 and <13.0.0, but run-time CUDA version is ", - cuda::cudart_version()); - NVTE_CHECK(cuda::cublas_version() >= 120205 && cuda::cublas_version() < 130000, + transformer_engine::cuda::cudart_version()); + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 120205 && + transformer_engine::cuda::cublas_version() < 130000, "Atomic GEMM requires cuBLAS >=12.2.5 and <13.0.0, but run-time cuBLAS version is ", - cuda::cublas_version()); + transformer_engine::cuda::cublas_version()); if (m_split == 0) m_split = 1; if (n_split == 0) n_split = 1; NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( @@ -934,9 +936,10 @@ void nvte_cublas_atomic_gemm(const NVTETensor A, const NVTETensor B, NVTETensor "Atomic GEMM requires CUDA version >=12.2.0 and <13.0.0, but run-time CUDA version is ", transformer_engine::cuda::cudart_version()); NVTE_CHECK( - cuda::cublas_version() >= 120205 && cuda::cublas_version() < 130000, + transformer_engine::cuda::cublas_version() >= 120205 && + transformer_engine::cuda::cublas_version() < 130000, "Atomic GEMM requires cuBLAS version >=12.2.5 and <13.0.0, but run-time cuBLAS version is ", - cuda::cublas_version()); + transformer_engine::cuda::cublas_version()); const Tensor *inputA = convertNVTETensorCheck(A); const Tensor *inputB = convertNVTETensorCheck(B); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index a1206474ea..b3e216dc4f 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -26,7 +26,7 @@ inline void CreateCublasHandle(cublasLtHandle_t *handle) { } // namespace -#if CUBLAS_VERSION >= 130100 +#if CUBLAS_VERSION >= 130200 namespace { @@ -543,13 +543,13 @@ 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.1+ - const int current_device = cuda::current_device(); - NVTE_CHECK(cuda::sm_arch(current_device) >= 100, + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.2+ + const int current_device = transformer_engine::cuda::current_device(); + NVTE_CHECK(transformer_engine::cuda::sm_arch(current_device) >= 100, "nvte_grouped_gemm requires Blackwell (SM100) or newer architecture."); - NVTE_CHECK(cuda::cublas_version() >= 130100, - "nvte_grouped_gemm requires cuBLAS 13.1+, but run-time cuBLAS version is ", - cuda::cublas_version()); + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 130200, + "nvte_grouped_gemm requires cuBLAS 13.2+, but run-time cuBLAS version is ", + transformer_engine::cuda::cublas_version()); // Convert to internal types const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A); @@ -631,15 +631,15 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT kGroupedGemmCublasWorkspaceSize, stream)); } -#else // CUBLAS_VERSION < 130100 +#else // CUBLAS_VERSION < 130200 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, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.1+, but compile-time cuBLAS version is ", + NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.2+, but compile-time cuBLAS version is ", CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); } -#endif // CUBLAS_VERSION >= 130100 +#endif // CUBLAS_VERSION >= 130200 diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 1afc9828e8..7403448722 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -299,7 +299,7 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor /* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ /*! \brief Grouped matrix multiplication: D = alpha * op(A) @ op(B) + beta * C * - * \note Requires cuBLAS 13.1+ (CUDA 13.1+) and Blackwell (SM100) or newer GPU architecture. + * \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. * @@ -322,7 +322,7 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor * \param[in] stream CUDA stream for the operation. * * Requirements: - * - cuBLAS 13.1+ (CUDA 13.1+) + * - 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] From 3ceb248e01a2c0dc1215fe0f46ebc235f289ba0d 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, 2 Feb 2026 09:11:35 -0800 Subject: [PATCH 188/521] More detailed documentation for recipes (#2343) * Code drop: Update recipes documentation and remove custom recipes from low precision training Signed-off-by: Pawel Gadzinski * Fix SVG css import path for diagrams Signed-off-by: Pawel Gadzinski * Refactor low_precision_training docs: remove optimizers, fix imports, add GPU checks Changes: - Remove optimizer code from all recipe examples (keep only forward/backward) - Fix Format imports (use Format.E4M3 instead of string 'E4M3') - Fix params_dtype for PyTorch examples (add params_dtype=torch.bfloat16) - Add GPU capability assertions before START blocks for blockwise/mxfp8/nvfp4 - Fix JAX imports (Float8CurrentScaling from common.recipe, NVFP4BlockScaling) - Add global_shard_guard for TransformerLayer examples in JAX - Fix fused_layers_jax.py return tuple unpacking - Update memory_usage JAX examples with dynamic GPU measurement - Remove memory_usage_3_jax (JAX doesn't support FP8 weight storage) - Update performance_considerations.rst for JAX differences - Delete unused .out files and fp8_autocast_jax.py Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix JAX memory usage .out files with correct output Signed-off-by: Pawel Gadzinski * responded to comments Signed-off-by: Pawel Gadzinski * applied suggestions form greptile Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * year change Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * jax compute capability fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix 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> --- docs/_static/css/diagram-colors.css | 134 +++++ docs/_static/css/sphinx_tabs.css | 45 ++ docs/_static/css/svg-responsive.css | 72 +++ docs/_templates/layout.html | 4 + docs/conf.py | 5 +- .../fp8_blockwise_scaling.rst | 254 ++++++++++ .../img/blockwise_swizzle_flow.svg | 146 ++++++ .../img/combined_scaling.svg | 342 +++++++++++++ .../img/transpose_handling.svg | 347 +++++++++++++ .../pytorch_blockwise_scaling_example.py | 37 ++ .../fp8_current_scaling.rst | 180 +++++++ .../img/fp8_cast_process.svg | 55 ++ .../img/fp8_current_scaling_all_gather.svg | 78 +++ .../fp8_current_scaling/img/fp8_formats.svg | 164 ++++++ .../img/fp8_scaling_concept.svg | 112 +++++ .../jax_current_scaling_example.py | 33 ++ .../pytorch_current_scaling_example.py | 29 ++ .../fp8_delayed_scaling.rst | 163 ++++++ .../img/scaling_comparison.svg | 82 +++ ...jax_delayed_scaling_distributed_example.py | 15 + .../jax_delayed_scaling_example.py | 39 ++ ...rch_delayed_scaling_distributed_example.py | 18 + .../pytorch_delayed_scaling_example.py | 37 ++ .../features/low_precision_training/index.rst | 17 + .../introduction/autocast_jax.py | 83 +++ .../introduction/autocast_pytorch.py | 69 +++ .../introduction/bf16_fp16_training_jax.py | 39 ++ .../bf16_fp16_training_pytorch.py | 52 ++ .../introduction/img/fp8_linear_flow.svg | 172 +++++++ .../img/fp_formats_comparison.svg | 183 +++++++ .../img/master_weights_approaches.svg | 112 +++++ .../img/mixed_precision_operations.svg | 105 ++++ .../introduction/introduction.rst | 285 +++++++++++ .../mxfp8/img/fp8_1d_scaling.svg | 177 +++++++ .../mxfp8/img/mxfp8_row_col.svg | 266 ++++++++++ .../img/mxfp8_scale_linearize_and_swizzle.svg | 190 +++++++ .../mxfp8/img/mxfp8_swizzle_both_tensors.svg | 101 ++++ .../mxfp8/img/mxfp8_tensor_scaling_layout.svg | 63 +++ .../mxfp8/jax_mxfp8_example.py | 39 ++ .../low_precision_training/mxfp8/mxfp8.rst | 213 ++++++++ .../mxfp8/pytorch_mxfp8_example.py | 34 ++ .../nvfp4/img/nvfp4_all_gather.svg | 118 +++++ .../nvfp4/img/nvfp4_hierarchical_scaling.svg | 186 +++++++ .../nvfp4/img/nvfp4_row_col.svg | 208 ++++++++ .../nvfp4/img/nvfp4_vs_fp8.svg | 91 ++++ .../low_precision_training/nvfp4/img/rht.svg | 138 +++++ .../nvfp4/img/stochastic_rounding.svg | 95 ++++ .../nvfp4/jax_nvfp4_example.py | 43 ++ .../low_precision_training/nvfp4/nvfp4.rst | 275 ++++++++++ .../nvfp4/pytorch_nvfp4_example.py | 35 ++ .../fused_layers_jax.py | 41 ++ .../fused_layers_pytorch.py | 37 ++ .../img/fused_layers.svg | 120 +++++ .../img/gemm_access_pattern.svg | 214 ++++++++ .../img/hopper_vs_blackwell_layout.svg | 122 +++++ .../img/sequence_parallel_quantization.svg | 159 ++++++ .../img/transpose_fusion.svg | 181 +++++++ .../memory_usage_1_jax.out | 9 + .../memory_usage_1_jax.py | 45 ++ .../memory_usage_1_pytorch.out | 4 + .../memory_usage_1_pytorch.py | 38 ++ .../memory_usage_2_jax.out | 10 + .../memory_usage_2_jax.py | 48 ++ .../memory_usage_2_pytorch.out | 4 + .../memory_usage_2_pytorch.py | 39 ++ .../memory_usage_3_pytorch.out | 4 + .../memory_usage_3_pytorch.py | 44 ++ .../performance_considerations.rst | 473 ++++++++++++++++++ .../save_original_input_pytorch.out | 4 + .../save_original_input_pytorch.py | 51 ++ docs/index.rst | 8 + 71 files changed, 7434 insertions(+), 1 deletion(-) create mode 100644 docs/_static/css/diagram-colors.css create mode 100644 docs/_static/css/sphinx_tabs.css create mode 100644 docs/_static/css/svg-responsive.css create mode 100644 docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst create mode 100644 docs/features/low_precision_training/fp8_blockwise_scaling/img/blockwise_swizzle_flow.svg create mode 100644 docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg create mode 100644 docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg create mode 100644 docs/features/low_precision_training/fp8_blockwise_scaling/pytorch_blockwise_scaling_example.py create mode 100644 docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst create mode 100644 docs/features/low_precision_training/fp8_current_scaling/img/fp8_cast_process.svg create mode 100644 docs/features/low_precision_training/fp8_current_scaling/img/fp8_current_scaling_all_gather.svg create mode 100644 docs/features/low_precision_training/fp8_current_scaling/img/fp8_formats.svg create mode 100644 docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg create mode 100644 docs/features/low_precision_training/fp8_current_scaling/jax_current_scaling_example.py create mode 100644 docs/features/low_precision_training/fp8_current_scaling/pytorch_current_scaling_example.py create mode 100644 docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst create mode 100644 docs/features/low_precision_training/fp8_delayed_scaling/img/scaling_comparison.svg create mode 100644 docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_distributed_example.py create mode 100644 docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_example.py create mode 100644 docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py create mode 100644 docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_example.py create mode 100644 docs/features/low_precision_training/index.rst create mode 100644 docs/features/low_precision_training/introduction/autocast_jax.py create mode 100644 docs/features/low_precision_training/introduction/autocast_pytorch.py create mode 100644 docs/features/low_precision_training/introduction/bf16_fp16_training_jax.py create mode 100644 docs/features/low_precision_training/introduction/bf16_fp16_training_pytorch.py create mode 100644 docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg create mode 100644 docs/features/low_precision_training/introduction/img/fp_formats_comparison.svg create mode 100644 docs/features/low_precision_training/introduction/img/master_weights_approaches.svg create mode 100644 docs/features/low_precision_training/introduction/img/mixed_precision_operations.svg create mode 100644 docs/features/low_precision_training/introduction/introduction.rst create mode 100644 docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg create mode 100644 docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg create mode 100644 docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg create mode 100644 docs/features/low_precision_training/mxfp8/img/mxfp8_swizzle_both_tensors.svg create mode 100644 docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg create mode 100644 docs/features/low_precision_training/mxfp8/jax_mxfp8_example.py create mode 100644 docs/features/low_precision_training/mxfp8/mxfp8.rst create mode 100644 docs/features/low_precision_training/mxfp8/pytorch_mxfp8_example.py create mode 100644 docs/features/low_precision_training/nvfp4/img/nvfp4_all_gather.svg create mode 100644 docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg create mode 100644 docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg create mode 100644 docs/features/low_precision_training/nvfp4/img/nvfp4_vs_fp8.svg create mode 100644 docs/features/low_precision_training/nvfp4/img/rht.svg create mode 100644 docs/features/low_precision_training/nvfp4/img/stochastic_rounding.svg create mode 100644 docs/features/low_precision_training/nvfp4/jax_nvfp4_example.py create mode 100644 docs/features/low_precision_training/nvfp4/nvfp4.rst create mode 100644 docs/features/low_precision_training/nvfp4/pytorch_nvfp4_example.py create mode 100644 docs/features/low_precision_training/performance_considerations/fused_layers_jax.py create mode 100644 docs/features/low_precision_training/performance_considerations/fused_layers_pytorch.py create mode 100644 docs/features/low_precision_training/performance_considerations/img/fused_layers.svg create mode 100644 docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg create mode 100644 docs/features/low_precision_training/performance_considerations/img/hopper_vs_blackwell_layout.svg create mode 100644 docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg create mode 100644 docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.out create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.py create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.out create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.py create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.out create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.py create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.out create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.py create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.out create mode 100644 docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.py create mode 100644 docs/features/low_precision_training/performance_considerations/performance_considerations.rst create mode 100644 docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.out create mode 100644 docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.py diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css new file mode 100644 index 0000000000..96a2a8a6dc --- /dev/null +++ b/docs/_static/css/diagram-colors.css @@ -0,0 +1,134 @@ +/* Diagram color definitions for Transformer Engine documentation */ + +/* High precision (BF16/FP16) elements */ +.hp { + fill: #ede7f6; + stroke: #673ab7; + stroke-width: 2; +} + +/* FP8 precision elements */ +.fp8 { + fill: #fff8e1; + stroke: #ffa726; + stroke-width: 2; +} + +/* GEMM/computation operations */ +.gemm { + fill: #ffe0b2; + stroke: #fb8c00; + stroke-width: 2.5; +} + +/* Quantization operations */ +.quantize { + fill: #e8f5e9; + stroke: #66bb6a; + stroke-width: 2; +} + +/* Amax computation operations */ +.amax { + fill: #e1f5fe; + stroke: #039be5; + stroke-width: 2; +} + +/* Text styles */ +.text { + font-family: 'Segoe UI', Arial, sans-serif; + font-size: 14px; + text-anchor: middle; + fill: #212121; +} + +.small-text { + font-family: 'Segoe UI', Arial, sans-serif; + font-size: 14px; + text-anchor: middle; + fill: #757575; +} + +.label { + font-family: 'Segoe UI', Arial, sans-serif; + font-size: 14px; + text-anchor: middle; + fill: #424242; +} + +.title { + font-family: 'Segoe UI', Arial, sans-serif; + font-size: 18px; + font-weight: 600; + text-anchor: middle; + fill: #212121; +} + +.section-title { + font-family: 'Segoe UI', Arial, sans-serif; + font-size: 15px; + font-weight: 600; + text-anchor: middle; +} + +/* Arrows */ +/* Note: marker-end references #arrowhead marker which must be defined in each SVG's section */ +.arrow { + stroke: #616161; + stroke-width: 2; + fill: none; + marker-end: url(#arrowhead); +} + +/* Additional box and element styles */ +.box-blue { + fill: #e3f2fd; + stroke: #1976d2; + stroke-width: 2; +} + +.box-orange { + fill: #fff3e0; + stroke: #f57c00; + stroke-width: 2; +} + +.box-green { + fill: #c8e6c9; + stroke: #388e3c; + stroke-width: 2; +} + +.box-dashed { + stroke-dasharray: 5,5; +} + +/* LayerNorm specific */ +.layernorm { + fill: #b3e5fc; + stroke: #0277bd; + stroke-width: 2.5; +} + +/* Fused layers */ +.fused { + fill: #b2dfdb; + stroke: #00695c; + stroke-width: 3; +} + +/* Generic computation blocks */ +.computation { + fill: #f5f5f5; + stroke: #757575; + stroke-width: 2; +} + +/* FP32 precision (alternative red) */ +.fp32 { + fill: #ffcdd2; + stroke: #d32f2f; + stroke-width: 2.5; +} + diff --git a/docs/_static/css/sphinx_tabs.css b/docs/_static/css/sphinx_tabs.css new file mode 100644 index 0000000000..c3e524e0e9 --- /dev/null +++ b/docs/_static/css/sphinx_tabs.css @@ -0,0 +1,45 @@ +/* Custom styling for sphinx-tabs */ + +.sphinx-tabs { + margin-bottom: 1rem; +} + +.sphinx-tabs-tab { + background-color: #f4f4f4; + border: 1px solid #ccc; + border-bottom: none; + padding: 0.5rem 1rem; + margin-right: 0.5rem; + cursor: pointer; + font-weight: 500; + transition: background-color 0.2s; +} + +.sphinx-tabs-tab:hover { + background-color: #e0e0e0; +} + +.sphinx-tabs-tab[aria-selected="true"] { + background-color: #76b900; /* NVIDIA green */ + color: white; + border-color: #76b900; + margin-right: 0.5rem; +} + +.sphinx-tabs-panel { + border: 1px solid #ccc; + padding: 1rem; + background-color: #f9f9f9; +} + +/* Dark mode support for RTD theme */ +.rst-content .sphinx-tabs-tab { + color: #333; +} + +.rst-content .sphinx-tabs-tab[aria-selected="true"] { + color: white; +} + + + diff --git a/docs/_static/css/svg-responsive.css b/docs/_static/css/svg-responsive.css new file mode 100644 index 0000000000..3ffe14eb14 --- /dev/null +++ b/docs/_static/css/svg-responsive.css @@ -0,0 +1,72 @@ +/* Responsive styling for SVG images */ + +/* Make all SVG images responsive */ +.document svg, +.document object[type="image/svg+xml"], +.rst-content svg { + max-width: 100%; + height: auto; + display: block; + margin: 1em auto; +} + +/* For raw HTML embedded SVGs */ +.document .raw-html svg { + max-width: 100%; + height: auto; + width: 100%; +} + +/* Ensure container doesn't overflow */ +.document .raw-html { + max-width: 100%; + overflow-x: auto; +} + +/* Figure containers with captions */ +.svg-figure { + text-align: center; + margin: 20px auto; +} + +.svg-figure img { + display: block; + margin: 0 auto; + height: auto; +} + +/* Different width classes for figures */ +.svg-figure.width-70 img { + width: 70%; + max-width: 100%; +} + +.svg-figure.width-80 img { + width: 80%; + max-width: 100%; +} + +.svg-figure.width-90 img { + width: 90%; + max-width: 100%; +} + +.svg-figure.width-100 img { + width: 100%; +} + +/* Figure captions */ +.svg-caption { + font-style: italic; + margin-top: 10px; + color: #555; + font-size: 0.95em; + line-height: 1.4; +} + + + + + + + diff --git a/docs/_templates/layout.html b/docs/_templates/layout.html index f94e526f57..99ae0702a8 100644 --- a/docs/_templates/layout.html +++ b/docs/_templates/layout.html @@ -67,6 +67,10 @@ overflow: visible !important; } + .quant { + background-color: yellow !important; + } + + + + + + + + + + + + + Input Tensor + + FP32/BF16 + + + + + + + + Quantize + + + + + + + FP8 (Compact) + + + + + FP32 Scales + + + + FP8 Data + + + + + + + + All-Gather + + + + + + + Swizzle + + + + + + + FP8 (GEMM Ready) + + + + + Swizzled Scales + + + + FP8 Data + + + + + + + + GEMM + + + + + + + + + + Input Tensor + + FP32/BF16 + + + + + + + + Quantize + + + Swizzle + + + + + + + FP8 (GEMM Ready) + + + + + Swizzled Scales + + + + FP8 Data + + + + + + + + GEMM + + diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg b/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg new file mode 100644 index 0000000000..dbf6999aef --- /dev/null +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/img/combined_scaling.svg @@ -0,0 +1,342 @@ + + + + + + + + + + Delayed/Current FP8 Scaling + (Single scaling factor per tensor) + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + 1 scaling factor + + + + + Blockwise FP8 Scaling – 1 dimension + (One scaling factor per 128 elements) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + Scaling factors (one per block) + + + + + Blockwise FP8 Scaling – 2 dimensions + (One scaling factor per 128x128 block of elements) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + + + + + + + + + + + + + + + Scaling factors (1 per 2D block) + + diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg b/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg new file mode 100644 index 0000000000..e9a3b7b7d1 --- /dev/null +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/img/transpose_handling.svg @@ -0,0 +1,347 @@ + + + + + + + 1D Blockwise Scaling + + + + Rowwise Quantization + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + Columnwise Quantization + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + + + … + + + + + + + + + + + + + + + + + + + + + + + + + + 2D Blockwise Scaling + + + + Rowwise Quantization + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + + + + + + + + + + + + + Columnwise Quantization + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + + … + … + + + … + + + + + + + + + + + + + + + + + + + diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/pytorch_blockwise_scaling_example.py b/docs/features/low_precision_training/fp8_blockwise_scaling/pytorch_blockwise_scaling_example.py new file mode 100644 index 0000000000..5100fc1a1d --- /dev/null +++ b/docs/features/low_precision_training/fp8_blockwise_scaling/pytorch_blockwise_scaling_example.py @@ -0,0 +1,37 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Check for Hopper or newer GPU +major, minor = torch.cuda.get_device_capability() +assert major >= 9, f"FP8 Blockwise Scaling requires SM90 (Hopper) or later, got SM{major}{minor}" + +# START_BLOCKWISE_SCALING_EXAMPLE + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Float8BlockScaling + +# Create FP8 Blockwise Scaling recipe +recipe = Float8BlockScaling( + fp8_format=te.common.recipe.Format.E4M3, # E4M3 or HYBRID (default: E4M3) + x_block_scaling_dim=1, # 1D scaling for activations (default: 1) + w_block_scaling_dim=2, # 2D scaling for weights (default: 2) + grad_block_scaling_dim=1, # 1D scaling for gradients (default: 1) +) + +# Create a linear layer with bfloat16 parameters +layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + +# Forward and backward pass +inp = torch.randn(32, 128, 1024, dtype=torch.bfloat16, device="cuda") + +with te.autocast(enabled=True, recipe=recipe): + output = layer(inp) + loss = output.sum() + +loss.backward() + +# END_BLOCKWISE_SCALING_EXAMPLE diff --git a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst new file mode 100644 index 0000000000..a4830a3fd5 --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst @@ -0,0 +1,180 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +FP8 Current Scaling +=================================== + +FP8 current scaling recipe is the simplest low precision recipe provided by Transformer Engine. +To understand how this recipe works, we first need to examine what the FP8 data type is and how it differs from other floating point formats. + + +FP8 data type +------------- + +The FP8 datatype, introduced in Hopper architecture, is actually 2 distinct datatypes, useful in different parts of the training of neural networks: + +* E4M3 -- consists of 1 sign bit, 4 exponent bits and 3 bits of mantissa. It can store values up to +/-448 and ``nan``. +* E5M2 -- consists of 1 sign bit, 5 exponent bits and 2 bits of mantissa. It can store values up to +/-57344, +/- ``inf`` and ``nan``. The tradeoff of the increased dynamic range is lower precision of the stored values. + +.. raw:: html + :file: img/fp8_formats.svg + +*Figure 1: Structure of the floating point datatypes. All of the values shown (in FP16, BF16, FP8 E4M3 and FP8 E5M2) are the closest representations of value 0.3952.* + + +**E4M3 and E5M2 usage in training** + +By default, Transformer Engine uses a hybrid approach: + +* *Forward pass* - activations and weights require more precision, so E4M3 datatype is used to store them. +* *Backward pass* - gradients are less susceptible to precision loss but require higher dynamic range, so E5M2 datatype is preferred. + +The user can configure this behavior via the ``fp8_format`` parameter of the recipe. + + +Scaling factors +--------------- + + +Limited dynamic range of FP8 datatype is insufficient for many tensors. +To address this, values in the tensor are scaled. FP8 Current Scaling recipe uses one **FP32** scale factor per tensor. The representation of a tensor element ``x`` in FP8 precision is given by: + +.. code-block:: python + + x = x_fp8 * s + +where + +* ``x_fp8`` is the FP8 value (E4M3 or E5M2), +* ``s`` is a global **FP32** scaling factor applied to the entire tensor. + +**FP8 Current Scaling quantization** + +Let's take a closer look at how quantization to FP8 with scaling factor is implemented in +the FP8 Current Scaling recipe. + +.. raw:: html + :file: img/fp8_scaling_concept.svg + +*Figure 3: Quantization to FP8 consists of amax (absolute maximum) computation, scaling to fit the FP8 range and casting to the respective FP8 format.* + +Quantization to FP8 consists of 3 steps: + +1. Computation of the absolute maximum value of the tensor - we refer to it as ``amax``. +2. Applying the scaling factor of ``fp8_max / amax`` to the tensor, to fit it into the FP8 range +3. Casting into the respective FP8 format using *Round To Nearest Even (RTNE)*. Values round to the nearest representable FP8 value. When exactly halfway between two values, rounds to the one with even mantissa to minimize systematic bias. + +**Performance analysis** + +Quantization is a memory-bound operation that requires reading the tensor twice: + +* First read: compute ``amax`` across all elements. +* Second read: apply the scaling factor and cast to FP8. + +This is a significant overhead compared to other recipes, which typically require only a single memory read. + +.. raw:: html + :file: img/fp8_cast_process.svg + +*Figure 4: FP8 quantization with current scaling recipe - two tensor reads are needed, one to compute amax and one to apply the scaling factor and cast to FP8.* + + +Transpose handling +------------------ + + + +*Ada and Hopper* + +On Ada and Hopper, the backward pass requires a transposed FP8 tensor. +The columnwise layout is physically different from the rowwise layout, so a transpose operation is needed. +All 3 options from :ref:`Performance Considerations Transpose handling section ` are supported. + +*Blackwell and later* + +Blackwell hardware supports multiple GEMM layouts natively, eliminating the need for explicit transposes. +The rowwise and columnwise tensors share the same physical memory layout. + +.. figure:: ../performance_considerations/img/hopper_vs_blackwell_layout.svg + :align: center + :alt: Comparison of rowwise and columnwise tensor layouts on Blackwell vs Hopper + + *Figure 6: On Blackwell, rowwise and columnwise usages share the same memory layout. On Hopper, columnwise usage requires a physical transpose.* + + +Distributed training +-------------------- + +**Quantized all-gather** + +FP8 all-gather is supported on all architectures (Ada and later). + +**Amax reduction** + +Tensors that are gathered across nodes (e.g. input and gradient in sequence parallelism) require amax synchronization before quantization. +Each node computes its local ``amax``, then a reduction produces the global maximum across all nodes. +All nodes use this synchronized amax to compute identical scaling factors, enabling quantized all-gather. + +.. raw:: html + :file: img/fp8_current_scaling_all_gather.svg + +*Figure 7: Quantization and all-gather flow for FP8 current scaling showing amax computation and synchronization.* + + +Supported devices +----------------- + +Ada and later (SM 8.9+) + +Examples +-------- + +Here's how to use FP8 Current Scaling recipe in PyTorch and JAX: + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM89 (Ada) or later +
+ + .. literalinclude:: pytorch_current_scaling_example.py + :language: python + :start-after: # START_CURRENT_SCALING_EXAMPLE + :end-before: # END_CURRENT_SCALING_EXAMPLE + + .. tab:: JAX + + .. raw:: html + +
+ Requires SM89 (Ada) or later +
+ + .. literalinclude:: jax_current_scaling_example.py + :language: python + :start-after: # START_CURRENT_SCALING_EXAMPLE + :end-before: # END_CURRENT_SCALING_EXAMPLE + + +---- + +Developer Notes +--------------- + +This section contains implementation details that may be useful for developers +but are not required for using FP8 Current Scaling in practice. + +All-gather of columnwise tensors +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +On Blackwell and later, rowwise and columnwise tensors share the same memory layout, +so all-gather of columnwise tensors is directly supported. + +For Hopper and Ada, all-gather of transposed FP8 tensors is not supported. +The rowwise tensor is gathered first, then transposed to columnwise format. \ No newline at end of file diff --git a/docs/features/low_precision_training/fp8_current_scaling/img/fp8_cast_process.svg b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_cast_process.svg new file mode 100644 index 0000000000..294fca318b --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_cast_process.svg @@ -0,0 +1,55 @@ + + + + + + + + + + + FP8 quantization + + + + High Precision + Tensor + + + + + + + Quantize + + + + Compute amax + 1 tensor read + + + + + + + Apply Scale + + Cast + 1 tensor read + + + + + + + FP8 + Tensor + + diff --git a/docs/features/low_precision_training/fp8_current_scaling/img/fp8_current_scaling_all_gather.svg b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_current_scaling_all_gather.svg new file mode 100644 index 0000000000..f984e1dd31 --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_current_scaling_all_gather.svg @@ -0,0 +1,78 @@ + + + + + + + + + + + Quantization + all gather for FP8 current scaling + + + + High Precision + Tensor + + + + + + + Compute + Amax + + + + + + + Synchronize + Amax + + + + + + + Scale + + Cast + + + + + + + FP8 + Tensor + + + + + + + All-Gather + + + + + + + FP8 Gathered + Tensor + + + diff --git a/docs/features/low_precision_training/fp8_current_scaling/img/fp8_formats.svg b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_formats.svg new file mode 100644 index 0000000000..bf86a29a6c --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_formats.svg @@ -0,0 +1,164 @@ + + + + + + + sign + exponent + mantissa + + + FP16 + + + + 0 + + + + 0 + + 1 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + 0 + + 1 + + 0 + + 1 + + 0 + + 0 + + 1 + + 1 + + = 0.395264 + + + + BF16 + + + + 0 + + + + 0 + + 1 + + 1 + + 1 + + 1 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + 0 + + 1 + + 0 + + 1 + + 0 + + = 0.394531 + + + + FP8 E4M3 + + + + 0 + + + + 0 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + 1 + + = 0.40625 + + + + FP8 E5M2 + + + + 0 + + + + 0 + + 1 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + = 0.375 + + + diff --git a/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg new file mode 100644 index 0000000000..9442b4e4aa --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/img/fp8_scaling_concept.svg @@ -0,0 +1,112 @@ + + + + + Original Tensor Values + + + + + + + 0 + + + + + + + + + + + + + + + + + amax + + + + + + Original range + + + + + + Scaled Values (fit FP8 range) + + + + + + + 0 + + + + + + FP8 range + + + + - FP8 range max + + + + + + + + + + + + + + Cast to FP8 (quantized values) + + + + + + + 0 + + + + + + FP8 range + + + + + + + + + + + + + + + + + + + diff --git a/docs/features/low_precision_training/fp8_current_scaling/jax_current_scaling_example.py b/docs/features/low_precision_training/fp8_current_scaling/jax_current_scaling_example.py new file mode 100644 index 0000000000..107b13c53b --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/jax_current_scaling_example.py @@ -0,0 +1,33 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_CURRENT_SCALING_EXAMPLE + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import DenseGeneral +from transformer_engine.common.recipe import Float8CurrentScaling, Format + +# Create FP8 Current Scaling recipe +# Available formats: +# - Format.HYBRID (default) -- E4M3 for forward pass, E5M2 for backward pass +# - Format.E4M3 -- E4M3 for both forward and backward pass +recipe = Float8CurrentScaling(fp8_format=Format.HYBRID) + +with te.autocast(enabled=True, recipe=recipe): + # Create and initialize layer + layer = DenseGeneral(features=1024) + key = jax.random.PRNGKey(0) + x = jax.random.normal(key, (32, 128, 1024), dtype=jnp.bfloat16) + var_collect = layer.init(key, x) + + # Forward and backward pass + def loss_fn(var_collect): + output = layer.apply(var_collect, x) + return output.sum() + + loss, grads = jax.value_and_grad(loss_fn)(var_collect) + +# END_CURRENT_SCALING_EXAMPLE diff --git a/docs/features/low_precision_training/fp8_current_scaling/pytorch_current_scaling_example.py b/docs/features/low_precision_training/fp8_current_scaling/pytorch_current_scaling_example.py new file mode 100644 index 0000000000..7ac1271890 --- /dev/null +++ b/docs/features/low_precision_training/fp8_current_scaling/pytorch_current_scaling_example.py @@ -0,0 +1,29 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_CURRENT_SCALING_EXAMPLE + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Float8CurrentScaling, Format + +# Create FP8 Current Scaling recipe +# Available formats: +# - Format.HYBRID (default) -- E4M3 for forward pass, E5M2 for backward pass +# - Format.E4M3 -- E4M3 for both forward and backward pass +recipe = Float8CurrentScaling(fp8_format=Format.HYBRID) + +# Create a simple linear layer with bfloat16 parameters +layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + +# Forward and backward pass +inp = torch.randn(32, 128, 1024, dtype=torch.bfloat16, device="cuda") + +with te.autocast(enabled=True, recipe=recipe): + output = layer(inp) + loss = output.sum() + +loss.backward() + +# END_CURRENT_SCALING_EXAMPLE diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst new file mode 100644 index 0000000000..9d05305eda --- /dev/null +++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst @@ -0,0 +1,163 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +FP8 Delayed Scaling +=================================== + +FP8 Delayed Scaling recipe estimates scaling factors from historical amax values rather than computing them +for each tensor. Compared to Current Scaling recipe, +this reduces tensor reads per quantization from two to one, +improving memory efficiency. + +Both this and :doc:`FP8 Current Scaling <../fp8_current_scaling/fp8_current_scaling>` recipe use +the same FP8 formats (E4M3/E5M2) with one FP32 scaling factor per tensor. +Reading the FP8 Current Scaling documentation first is recommended. + +Quantization with delayed scaling factors +----------------------------------------- + +FP8 Current Scaling requires two tensor reads per quantization: one to compute amax, +one to cast. FP8 Delayed Scaling eliminates the first read by predicting the scaling factor +from historical amax values - hence *delayed* (using past values) versus *current* (using present values). + +The quantization process works as follows: + +1. **Compute scaling factor from history** (no tensor read needed): + The scaling factor is derived from stored ``amax_history`` using the formula: + + ``scaling_factor = FP8_MAX / amax`` + + where ``amax`` is computed from history using either ``max`` (maximum over window, default) or ``most_recent`` algorithm. + +2. **Quantize the tensor** (one tensor read): + Apply the scaling factor and cast to FP8. Values exceeding FP8 range are clipped. + +3. **Update history**: + Record the actual amax from this quantization for future iterations. + +Each module maintains an ``amax_history`` tensor of configurable length (``amax_history_len``) +for each quantized tensor. + +.. raw:: html + :file: img/scaling_comparison.svg + +*Figure 1. Comparison of FP8 Current Scaling and FP8 Delayed Scaling quantization processes.* + +Amax History Management +----------------------- + +The ``amax_history`` buffer acts as a sliding window of recent amax values. +Position 0 serves as a staging area for the current amax, while positions 1 to N-1 +store the history from oldest to newest. Each quantization writes the observed amax +to position 0, and after the pass completes, the history is rotated: + +.. code-block:: text + + Before rotation: [amax_N, amax_1, amax_2, ..., amax_N-1] (amax_N = current, amax_1 = oldest) + After rotation: [0, amax_2, ..., amax_N-1, amax_N] (amax_1 dropped, amax_N appended) + +The scaling factor is computed **before** the rotation, so it uses all ``amax_history_len`` values. +Position 0 serves as a staging area — it is zeroed after the scale update, ready for the next iteration's amax. + +The implementation differs between PyTorch and JAX: + +.. tabs:: + + .. tab:: PyTorch + + Each module creates two ``amax_history`` tensors, initialized to zero: + + - Forward: shape ``(amax_history_len, num_gemms * 3)`` — three FP8 tensors per GEMM (input, weight, output) + - Backward: shape ``(amax_history_len, num_gemms * 2)`` — two FP8 tensors per GEMM (grad_output, grad_input) + + When the autocast context exits, a single CUDA kernel processes all tensors at once — + performing amax reduction across GPUs and history rotation. This batched approach + minimizes kernel launch overhead compared to updating each tensor separately. + + .. tab:: JAX + + Each quantizer maintains its own ``amax_history`` with shape ``(amax_history_len,)`` + and updates independently. + +Here's how to use FP8 Delayed Scaling in PyTorch and JAX: + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM89 (Ada) or later +
+ + .. literalinclude:: pytorch_delayed_scaling_example.py + :language: python + :start-after: # START_DELAYED_SCALING_EXAMPLE + :end-before: # END_DELAYED_SCALING_EXAMPLE + + .. tab:: JAX + + .. raw:: html + +
+ Requires SM89 (Ada) or later +
+ + .. literalinclude:: jax_delayed_scaling_example.py + :language: python + :start-after: # START_DELAYED_SCALING_EXAMPLE + :end-before: # END_DELAYED_SCALING_EXAMPLE + + +Distributed Training +-------------------- + +FP8 Delayed Scaling uses the same data formats as FP8 Current Scaling - quantized all-gather is supported. +However, amax reduction works slightly differently in different frameworks. + +.. tabs:: + + .. tab:: PyTorch + + Amax reduction is controlled by two parameters: + + - ``reduce_amax`` in recipe: enables/disables reduction (required for SP and CP) + - ``amax_reduction_group`` in ``autocast``: specifies the process group for reduction + + We recommend reducing amax across all GPUs where the tensor is sharded, + including data parallel ranks. + + .. literalinclude:: pytorch_delayed_scaling_distributed_example.py + :language: python + :start-after: # START_AMAX_REDUCTION_EXAMPLE + :end-before: # END_AMAX_REDUCTION_EXAMPLE + + In data parallel training, some modules may not execute on certain ranks + (e.g., MoE experts that receive no tokens). This is handled as follows: + + - **First iteration**: All modules must execute on all ranks to register + their ``amax_history`` tensors in the global buffer. Mismatched registration + would cause the ``all_reduce`` to hang due to different tensor sizes across ranks. + - **Subsequent iterations**: The ``autocast`` context must be entered and exited + on all ranks (this triggers the collective reduction). Individual modules can be + skipped - if no rank executes a module, its history is not rotated and scale + remains unchanged. + + + .. tab:: JAX + + Amax reduction is always enabled and managed automatically. + Reduction scope: all parallelism axes except pipeline parallelism (TP, SP, DP/FSDP). + + .. literalinclude:: jax_delayed_scaling_distributed_example.py + :language: python + :start-after: # START_AMAX_REDUCTION_EXAMPLE + :end-before: # END_AMAX_REDUCTION_EXAMPLE + +Supported devices +----------------- + +Ada and later (SM 8.9+) \ No newline at end of file diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/img/scaling_comparison.svg b/docs/features/low_precision_training/fp8_delayed_scaling/img/scaling_comparison.svg new file mode 100644 index 0000000000..aff4ba0da3 --- /dev/null +++ b/docs/features/low_precision_training/fp8_delayed_scaling/img/scaling_comparison.svg @@ -0,0 +1,82 @@ + + + + + + + + + + + Current Scaling + + + + Tensor + + + + + + + Amax Computation + + + + + + + Quantization + (uses tensor + amax) + + + + + + + FP8 Tensor + + + + Delayed Scaling + + + + Tensor + + + + amax history + + + + read amax + + + + Quantization + (uses tensor + amax from history) + (updates amax history) + + + + update amax + + + + + + + FP8 Tensor + + + diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_distributed_example.py b/docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_distributed_example.py new file mode 100644 index 0000000000..f354ddaf77 --- /dev/null +++ b/docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_distributed_example.py @@ -0,0 +1,15 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_AMAX_REDUCTION_EXAMPLE +import transformer_engine.jax as te +from transformer_engine.common.recipe import DelayedScaling + +# Amax reduction scope is managed internally +recipe = DelayedScaling(reduce_amax=True) # Must be True in JAX + +with te.autocast(enabled=True, recipe=recipe): + output = layer.apply(params, inp) + +# END_AMAX_REDUCTION_EXAMPLE diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_example.py b/docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_example.py new file mode 100644 index 0000000000..5971117686 --- /dev/null +++ b/docs/features/low_precision_training/fp8_delayed_scaling/jax_delayed_scaling_example.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +from transformer_engine.jax.quantize import get_device_compute_capability + +# Requires Ada (SM89) or newer for FP8 support +assert get_device_compute_capability() >= 89, "This example requires SM89 (Ada) or newer" + +# START_DELAYED_SCALING_EXAMPLE + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import DenseGeneral +from transformer_engine.common.recipe import DelayedScaling + +# Create FP8 Delayed Scaling recipe +recipe = DelayedScaling( + margin=0, # Margin for scaling factor computation (default: 0) + amax_history_len=1024, # Length of amax history window (default: 1024) + amax_compute_algo="max", # How to compute amax from history (default: "max") +) + +with te.autocast(enabled=True, recipe=recipe): + # Initialize layer and data + layer = DenseGeneral(features=1024) + key = jax.random.PRNGKey(0) + x = jax.random.normal(key, (32, 128, 1024), dtype=jnp.bfloat16) + var_collect = layer.init(key, x) + + # Forward and backward pass + def loss_fn(var_collect): + output = layer.apply(var_collect, x) + return output.sum() + + loss, grads = jax.value_and_grad(loss_fn)(var_collect) + +# END_DELAYED_SCALING_EXAMPLE diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py new file mode 100644 index 0000000000..863b71e8c6 --- /dev/null +++ b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_distributed_example.py @@ -0,0 +1,18 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_AMAX_REDUCTION_EXAMPLE +import torch.distributed as dist +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import DelayedScaling + +# Create process group for amax reduction (e.g., all 8 GPUs) +amax_reduction_group = dist.new_group(ranks=[0, 1, 2, 3, 4, 5, 6, 7]) + +recipe = DelayedScaling(reduce_amax=True) + +with te.autocast(recipe=recipe, amax_reduction_group=amax_reduction_group): + output = model(inp) + +# END_AMAX_REDUCTION_EXAMPLE diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_example.py b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_example.py new file mode 100644 index 0000000000..45d244f47d --- /dev/null +++ b/docs/features/low_precision_training/fp8_delayed_scaling/pytorch_delayed_scaling_example.py @@ -0,0 +1,37 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or newer for FP8 support +assert torch.cuda.get_device_capability()[0] >= 9 or ( + torch.cuda.get_device_capability()[0] == 8 and torch.cuda.get_device_capability()[1] >= 9 +), "This example requires SM89 (Ada) or newer" + +# START_DELAYED_SCALING_EXAMPLE + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import DelayedScaling + +# Create FP8 Delayed Scaling recipe +recipe = DelayedScaling( + margin=0, # Margin for scaling factor computation (default: 0) + amax_history_len=1024, # Length of amax history window (default: 1024) + amax_compute_algo="max", # How to compute amax from history (default: "max") +) + +# Create a linear layer with bfloat16 parameters +layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + +# Forward and backward pass +inp = torch.randn(32, 128, 1024, dtype=torch.bfloat16, device="cuda") + +with te.autocast(enabled=True, recipe=recipe): + output = layer(inp) + loss = output.sum() + +loss.backward() + +# END_DELAYED_SCALING_EXAMPLE diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst new file mode 100644 index 0000000000..8b392d2bbb --- /dev/null +++ b/docs/features/low_precision_training/index.rst @@ -0,0 +1,17 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Low precision training +=================================== + +.. toctree:: + + introduction/introduction.rst + performance_considerations/performance_considerations.rst + fp8_current_scaling/fp8_current_scaling.rst + fp8_delayed_scaling/fp8_delayed_scaling.rst + fp8_blockwise_scaling/fp8_blockwise_scaling.rst + mxfp8/mxfp8.rst + nvfp4/nvfp4.rst \ No newline at end of file diff --git a/docs/features/low_precision_training/introduction/autocast_jax.py b/docs/features/low_precision_training/introduction/autocast_jax.py new file mode 100644 index 0000000000..0abb670064 --- /dev/null +++ b/docs/features/low_precision_training/introduction/autocast_jax.py @@ -0,0 +1,83 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +from transformer_engine.jax.quantize import get_device_compute_capability + +# Requires Ada (SM89) or newer for FP8 support +assert get_device_compute_capability() >= 89, "This example requires SM89 (Ada) or newer" + +# START_AUTOCAST_BASIC + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import TransformerLayer +from transformer_engine.common.recipe import DelayedScaling, Format + +# Set up recipe +recipe = DelayedScaling() + +# Model initialization must happen inside autocast +with te.autocast(enabled=True, recipe=recipe): + layer = TransformerLayer( + hidden_size=1024, + mlp_hidden_size=4096, + num_attention_heads=16, + ) + + init_key, dropout_key = jax.random.split(jax.random.PRNGKey(0)) + x = jax.random.normal(init_key, (32, 128, 1024), dtype=jnp.bfloat16) + var_collect = layer.init({"params": init_key, "dropout": dropout_key}, x) + + # Forward and backward pass (both inside autocast for JAX) + def loss_fn(var_collect): + output = layer.apply(var_collect, x, rngs={"dropout": dropout_key}) + return output.sum() + + loss, grads = jax.value_and_grad(loss_fn)(var_collect) + +# END_AUTOCAST_BASIC + + +# START_AUTOCAST_SEQUENTIAL + +encoder_recipe = DelayedScaling(fp8_format=Format.E4M3) +decoder_recipe = DelayedScaling(fp8_format=Format.HYBRID) + +with te.autocast(enabled=True, recipe=encoder_recipe): + encoder = TransformerLayer(hidden_size=1024, mlp_hidden_size=4096, num_attention_heads=16) + encoder_var_collect = encoder.init({"params": init_key, "dropout": dropout_key}, x) + hidden = encoder.apply(encoder_var_collect, x, rngs={"dropout": dropout_key}) + +with te.autocast(enabled=True, recipe=decoder_recipe): + decoder = TransformerLayer(hidden_size=1024, mlp_hidden_size=4096, num_attention_heads=16) + decoder_var_collect = decoder.init({"params": init_key, "dropout": dropout_key}, hidden) + output = decoder.apply(decoder_var_collect, hidden, rngs={"dropout": dropout_key}) + +# END_AUTOCAST_SEQUENTIAL + + +# START_AUTOCAST_NESTED + +outer_recipe = DelayedScaling(fp8_format=Format.E4M3) +inner_recipe = DelayedScaling(fp8_format=Format.HYBRID) + +with te.autocast(enabled=True, recipe=outer_recipe): + # layer1 uses outer_recipe + layer1 = TransformerLayer(hidden_size=1024, mlp_hidden_size=4096, num_attention_heads=16) + var_collect1 = layer1.init({"params": init_key, "dropout": dropout_key}, x) + hidden = layer1.apply(var_collect1, x, rngs={"dropout": dropout_key}) + + with te.autocast(enabled=True, recipe=inner_recipe): + # layer2 uses inner_recipe (overrides outer) + layer2 = TransformerLayer(hidden_size=1024, mlp_hidden_size=4096, num_attention_heads=16) + var_collect2 = layer2.init({"params": init_key, "dropout": dropout_key}, hidden) + hidden = layer2.apply(var_collect2, hidden, rngs={"dropout": dropout_key}) + + # layer3 uses outer_recipe again + layer3 = TransformerLayer(hidden_size=1024, mlp_hidden_size=4096, num_attention_heads=16) + var_collect3 = layer3.init({"params": init_key, "dropout": dropout_key}, hidden) + output = layer3.apply(var_collect3, hidden, rngs={"dropout": dropout_key}) + +# END_AUTOCAST_NESTED diff --git a/docs/features/low_precision_training/introduction/autocast_pytorch.py b/docs/features/low_precision_training/introduction/autocast_pytorch.py new file mode 100644 index 0000000000..2c1528ff9e --- /dev/null +++ b/docs/features/low_precision_training/introduction/autocast_pytorch.py @@ -0,0 +1,69 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or newer for FP8 support +assert torch.cuda.get_device_capability()[0] >= 9 or ( + torch.cuda.get_device_capability()[0] == 8 and torch.cuda.get_device_capability()[1] >= 9 +), "This example requires SM89 (Ada) or newer" + +# START_AUTOCAST_BASIC + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import DelayedScaling, Format + +recipe = DelayedScaling() +layer = te.Linear(1024, 1024) +inp = torch.randn(32, 1024, dtype=torch.float32, device="cuda") + +with te.autocast(enabled=True, recipe=recipe): + output = layer(inp) + +# .backward() is called outside of autocast +loss = output.sum() +loss.backward() + +# END_AUTOCAST_BASIC + + +# START_AUTOCAST_SEQUENTIAL + +encoder_recipe = DelayedScaling(fp8_format=Format.E4M3) +decoder_recipe = DelayedScaling(fp8_format=Format.HYBRID) + +encoder = te.Linear(1024, 1024) +decoder = te.Linear(1024, 1024) + +with te.autocast(enabled=True, recipe=encoder_recipe): + hidden = encoder(inp) + +with te.autocast(enabled=True, recipe=decoder_recipe): + output = decoder(hidden) + +# END_AUTOCAST_SEQUENTIAL + + +# START_AUTOCAST_NESTED + +outer_recipe = DelayedScaling(fp8_format=Format.E4M3) +inner_recipe = DelayedScaling(fp8_format=Format.HYBRID) + +layer1 = te.Linear(1024, 1024) +layer2 = te.Linear(1024, 1024) +layer3 = te.Linear(1024, 1024) + +with te.autocast(enabled=True, recipe=outer_recipe): + # layer1 uses outer_recipe + x = layer1(inp) + + with te.autocast(enabled=True, recipe=inner_recipe): + # layer2 uses inner_recipe (overrides outer) + x = layer2(x) + + # layer3 uses outer_recipe again + output = layer3(x) + +# END_AUTOCAST_NESTED diff --git a/docs/features/low_precision_training/introduction/bf16_fp16_training_jax.py b/docs/features/low_precision_training/introduction/bf16_fp16_training_jax.py new file mode 100644 index 0000000000..a3c9c2ae45 --- /dev/null +++ b/docs/features/low_precision_training/introduction/bf16_fp16_training_jax.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_BF16_FP16_TRAINING + +import jax +import jax.numpy as jnp +from transformer_engine.jax.flax import TransformerLayer + + +def run_forward_backward(params_dtype, compute_dtype): + # Create TransformerLayer + layer = TransformerLayer( + hidden_size=1024, + mlp_hidden_size=4096, + num_attention_heads=16, + dtype=params_dtype, + ) + + # Initialize parameters and optimizer + init_key, dropout_key = jax.random.split(jax.random.PRNGKey(0)) + x = jax.random.normal(init_key, (32, 128, 1024), dtype=compute_dtype) + var_collect = layer.init({"params": init_key, "dropout": dropout_key}, x) + + # Forward and backward pass + def loss_fn(var_collect): + output = layer.apply(var_collect, x, rngs={"dropout": dropout_key}) + assert output.dtype == compute_dtype + return output.sum() + + loss, grads = jax.value_and_grad(loss_fn)(var_collect) + + +run_forward_backward(jnp.float32, jnp.float32) # high precision training +run_forward_backward(jnp.float32, jnp.bfloat16) # bfloat16 training with master weights in FP32 +run_forward_backward(jnp.bfloat16, jnp.bfloat16) # bfloat16 training with weights in BF16 + +# END_BF16_FP16_TRAINING diff --git a/docs/features/low_precision_training/introduction/bf16_fp16_training_pytorch.py b/docs/features/low_precision_training/introduction/bf16_fp16_training_pytorch.py new file mode 100644 index 0000000000..4eb6ce1f84 --- /dev/null +++ b/docs/features/low_precision_training/introduction/bf16_fp16_training_pytorch.py @@ -0,0 +1,52 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_BF16_FP16_TRAINING + +import torch +import transformer_engine.pytorch as te +from contextlib import nullcontext + + +def run_forward_backward(params_dtype, autocast_precision, grad_scaler_enabled): + if grad_scaler_enabled: + grad_scaler = torch.amp.GradScaler("cuda") + + layer = te.TransformerLayer( + hidden_size=1024, + ffn_hidden_size=4096, + num_attention_heads=16, + params_dtype=params_dtype, + ) + x = torch.randn(32, 128, 1024, dtype=params_dtype, device="cuda") + + autocast_ctx = ( + torch.autocast(device_type="cuda", dtype=autocast_precision) + if autocast_precision is not None + else nullcontext() + ) + with autocast_ctx: + output = layer(x) + assert ( + output.dtype == autocast_precision if autocast_precision is not None else params_dtype + ) + loss = output.sum() + if grad_scaler_enabled: + grad_scaler.scale(loss).backward() + else: + loss.backward() + + +run_forward_backward(torch.float32, torch.float32, False) # high precision training +run_forward_backward( + torch.float32, torch.bfloat16, False +) # bfloat16 training with master weights in FP32 +run_forward_backward( + torch.float32, torch.float16, True +) # fp16 training with master weights in FP32, needs loss scaling +run_forward_backward( + torch.bfloat16, torch.bfloat16, False +) # bfloat16 training with weights in BF16 + +# END_BF16_FP16_TRAINING diff --git a/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg b/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg new file mode 100644 index 0000000000..e1861ebc1c --- /dev/null +++ b/docs/features/low_precision_training/introduction/img/fp8_linear_flow.svg @@ -0,0 +1,172 @@ + + + + + + + + + + + FP8 Linear Layer – Forward and Backward Pass + + + Forward Pass + + + + InputT + + + + Input + + + + + + + Quantize + + + + + + + + + Input + + + + N + + + + Weight + + + + + + + Quantize + + + + + + + + + Weight + + + + WeightT + + + + T + + + + FP8 GEMM + (TN) + + + + + + + Output + + + + + + Backward Pass + + + + WeightT + + + + Output grad. + + + + + + + Quantize + + + + + + + + + Output grad. + + + + Output grad.T + + + + FP8 GEMM + (TN) + + + + Input grad. + + + + FP8 GEMM + (TN) + + + + Weight grad. + + + + InputT + + + + + N + + + T + + + + + + N + + + T + + + + + + + + Higher Precision (FP32/BF16/FP16) + + + + Lower Precision (FP8, MXFP8 etc.) + + + diff --git a/docs/features/low_precision_training/introduction/img/fp_formats_comparison.svg b/docs/features/low_precision_training/introduction/img/fp_formats_comparison.svg new file mode 100644 index 0000000000..a6c46b364d --- /dev/null +++ b/docs/features/low_precision_training/introduction/img/fp_formats_comparison.svg @@ -0,0 +1,183 @@ + + + + + + + sign + exponent + mantissa + + + FP32 + + + + 0 + + + + 0 + + 1 + + 1 + + 1 + + 1 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + 0 + + 1 + + 0 + + 1 + + 0 + + 0 + + 1 + + 0 + + 1 + + 0 + + 1 + + 1 + + 1 + + 1 + + 0 + + 1 + + 0 + + 1 + + 0 + + 0 + + 0 + + = 0.3952 + + + + BF16 + + + + 0 + + + + 0 + + 1 + + 1 + + 1 + + 1 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + 0 + + 1 + + 0 + + 1 + + 0 + + ≈ 0.3945 + + + + FP16 + + + + 0 + + + + 0 + + 1 + + 1 + + 0 + + 1 + + + + 1 + + 0 + + 0 + + 1 + + 0 + + 1 + + 0 + + 0 + + 1 + + 0 + + ≈ 0.3950 + + diff --git a/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg b/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg new file mode 100644 index 0000000000..b231fefd90 --- /dev/null +++ b/docs/features/low_precision_training/introduction/img/master_weights_approaches.svg @@ -0,0 +1,112 @@ + + + + + + + + + + + Master Weights Storage Approaches + + + + + + + Low Precision Weights + (no master weights) + + + + Model + + Weights (BF16/FP16) + + + + + + + Forward/Backward + + + + + + + Optimizer + + State (FP32) + + + Master Weights in Model + + + + Model + + Weights (FP32) + + + + + cast to BF16/FP16 + + + + Forward/Backward + + + + + + + Optimizer + + State (FP32) + + + Master Weights in Optimizer + + + + cast to BF16/FP16 + + + + + + + Model + + Weights (BF16/FP16) + + + + + + + Forward/Backward + + + + + + + Optimizer + + State (FP32) + + Master (FP32) + + + + + diff --git a/docs/features/low_precision_training/introduction/img/mixed_precision_operations.svg b/docs/features/low_precision_training/introduction/img/mixed_precision_operations.svg new file mode 100644 index 0000000000..7a61759184 --- /dev/null +++ b/docs/features/low_precision_training/introduction/img/mixed_precision_operations.svg @@ -0,0 +1,105 @@ + + + + + + + + + + + Transformer Layer – default precision of operation in low precision recipe + + + + Input + + + + + Layer Norm + + + + + QKV Linear + + + + + QK^T + + + + + Softmax + + + + + + Scores * V + + + + + Output Linear + + + + + Dropout + Add + + + + + + Layer Norm + + + + + FFN Linear 1 + + + + + GELU + + + + + FFN Linear 2 + + + + + Output + + + Memory State: + + + + Parameters + + + + Gradients + + + + + + Higher Precision (FP32/BF16/FP16) + + + + Lower Precision (FP8, MXFP8 etc.) + + + diff --git a/docs/features/low_precision_training/introduction/introduction.rst b/docs/features/low_precision_training/introduction/introduction.rst new file mode 100644 index 0000000000..760a63b0b1 --- /dev/null +++ b/docs/features/low_precision_training/introduction/introduction.rst @@ -0,0 +1,285 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Introduction +=================================== + +Transformer Engine accelerates deep learning on NVIDIA GPUs in several ways, +with low precision training being one of the most important. +This chapter introduces mixed precision training and FP8 support. + + +Training in BF16/FP16 +--------------------- + +Deep learning traditionally uses 32-bit floating-point (FP32) numbers. +NVIDIA GPUs support lower precision formats—FP16 since Pascal, BF16 since Ampere—which offer higher throughput and lower memory usage. +Let's compare these formats. + +.. raw:: html + :file: img/fp_formats_comparison.svg + +*Figure 1: Comparison of FP32, BF16, and FP16 floating-point formats showing bit allocation for sign, exponent, and mantissa.* + +The key differences between these formats are: + +* **FP32** (32 bits total): 1 sign bit + 8 exponent bits + 23 mantissa bits – standard single-precision format +* **BF16** (16 bits total): 1 sign bit + 8 exponent bits + 7 mantissa bits – maintains FP32's exponent range but has reduced precision +* **FP16** (16 bits total): 1 sign bit + 5 exponent bits + 10 mantissa bits – reduced range but higher precision than BF16 + +BF16's advantage is that it shares the same exponent range as FP32, +making it easier to convert between the two formats without overflow/underflow issues. +FP16 offers better precision for smaller values but has a limited dynamic range, +which results in the need to perform loss scaling to avoid overflow/underflow—see `this paper on loss scaling `__ for more details. + +**Mixed precision** + +Not all operations should be run in reduced precision to preserve accuracy. +Modern deep learning frameworks use *mixed precision training*, +where different operations use different precisions based on their numerical properties: + +* Matrix multiplications are compute-heavy and remain numerically stable at lower precision, making them ideal candidates for acceleration. +* Operations like layer normalization and softmax can work with low precision inputs and outputs, but may use high precision internally or for their weights. +* Operations like loss computation and exponentiation need high precision throughout. + +**Master weights** + +Another consideration in mixed precision training is how to store the model weights. +Lower precision formats like FP16 and BF16 have limited representational granularity, +which becomes problematic during gradient updates. +When a small gradient is added to a not so small weight stored in low precision, +the result may round back to the original value if the update falls below the format's precision threshold. +Moreover, some elements of the gradient itself can be too small to be represented in low precision, +especially after the accumulation from multiple GPUs in the data parallel training setting. + +The solution is to maintain *master weights* in FP32. +During training, weights are cast to lower precision for forward and backward passes, +but the gradient updates are applied to the full-precision master copy. +This ensures that even small gradients accumulate correctly over time. + +There are two common software approaches to storing master weights: + +* *In the optimizer*: + The model holds low-precision weights, + while the optimizer maintains FP32 copies alongside momentum and other state. + During each step, + the optimizer updates its FP32 copy and casts the result back to the model's low-precision weights. + + This approach makes it easier to shard master weights together with other optimizer state, for example in ZeRO optimizer. + + Since the casting happens only during the optimizer step, this approach is also faster when optimizer runs less frequently than the model, e.g. when performing gradient accumulation or pipeline parallel training. + +* *In the model*: + The model stores weights directly in FP32, + and they are cast to lower precision on-the-fly during forward and backward passes. + This approach works seamlessly with any standard optimizer, requiring no special support. + +.. raw:: html + :file: img/master_weights_approaches.svg + +*Figure 2: Three approaches to weight storage—low precision only (no master weights), master weights stored in the model, and master weights stored in the optimizer.* + +.. tabs:: + + .. tab:: PyTorch + + The PyTorch API of Transformer Engine provides several mechanisms to control precision: + + * **Weight precision**: Use the ``params_dtype`` argument in any TE layer constructor. + * **Computation precision**: Use the ``torch.autocast`` context manager. When enabled, inputs are cast to the autocast dtype before computation. + * **Input dtype**: When ``torch.autocast`` is not used, the input tensor's dtype determines the computation precision. In this case, inputs and parameters must have matching dtypes. + + .. literalinclude:: bf16_fp16_training_pytorch.py + :language: python + :start-after: # START_BF16_FP16_TRAINING + :end-before: # END_BF16_FP16_TRAINING + + + .. tab:: JAX + + The JAX API of Transformer Engine provides two mechanisms to control precision: + + * **Weight precision**: Use the ``dtype`` argument in any TE layer constructor. + * **Computation precision**: Determined by the dtype of the input tensor. + + For training with master weights in FP32 and computation in BF16, + cast the input tensor to BF16 before passing it to the layer. + + .. literalinclude:: bf16_fp16_training_jax.py + :language: python + :start-after: # START_BF16_FP16_TRAINING + :end-before: # END_BF16_FP16_TRAINING + + + +Lower precisions +---------------- + +Transformer Engine's primary feature is supporting even lower precision than BF16/FP16, such as FP8, MXFP8, NVFP4, etc. +The logic of these precisions is more complicated than the logic of BF16/FP16 – they require scaling factors to +properly represent the full range of values in the tensor. Sometimes it is one scaling factor per tensor, +sometimes it is one scaling factor per block of values. A precision format combined with the logic for training +is called **a recipe**. + +In this section we present common logic for all the recipes. Each one of them is described in more detail in a separate section later. +Let's now see how we can train in lower precisions in supported frameworks. + +.. tabs:: + + .. tab:: PyTorch + + The PyTorch API of Transformer Engine provides an ``autocast`` context manager to control precision. + It's similar to the ``torch.autocast`` context manager, but tailored for low precision training. + The most important argument is the ``recipe`` argument, which accepts objects inheriting from + :class:`~transformer_engine.common.recipe.Recipe`. + + Forward computations need to be performed inside the ``autocast`` context manager, + while the ``.backward()`` call should be outside of it (it inherits the setting from the + corresponding forward pass). + + Here is a basic example: + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada or newer) +
+ + .. literalinclude:: autocast_pytorch.py + :language: python + :start-after: # START_AUTOCAST_BASIC + :end-before: # END_AUTOCAST_BASIC + + You can use multiple recipes in the same model in the following ways: + + **Sequential contexts** – apply different recipes to different parts of your model: + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada or newer) +
+ + .. literalinclude:: autocast_pytorch.py + :language: python + :start-after: # START_AUTOCAST_SEQUENTIAL + :end-before: # END_AUTOCAST_SEQUENTIAL + + **Nested contexts** – the inner context overrides the outer one for its scope: + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada or newer) +
+ + .. literalinclude:: autocast_pytorch.py + :language: python + :start-after: # START_AUTOCAST_NESTED + :end-before: # END_AUTOCAST_NESTED + + + .. tab:: JAX + + The JAX API of Transformer Engine provides an ``autocast`` context manager similar to PyTorch. + The key difference is that in JAX, model initialization must happen inside the ``autocast`` context + to properly capture quantization metadata in the parameter tree. + + Here is a basic example: + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada or newer) +
+ + .. literalinclude:: autocast_jax.py + :language: python + :start-after: # START_AUTOCAST_BASIC + :end-before: # END_AUTOCAST_BASIC + + You can use multiple recipes in the same model in the following ways: + + **Sequential contexts** – apply different recipes to different parts of your model: + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada or newer) +
+ + .. literalinclude:: autocast_jax.py + :language: python + :start-after: # START_AUTOCAST_SEQUENTIAL + :end-before: # END_AUTOCAST_SEQUENTIAL + + **Nested contexts** – the inner context overrides the outer one for its scope: + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada or newer) +
+ + .. literalinclude:: autocast_jax.py + :language: python + :start-after: # START_AUTOCAST_NESTED + :end-before: # END_AUTOCAST_NESTED + + .. note:: + Python context managers like ``autocast`` may interact unexpectedly with JAX's JIT compilation. + For finer-grained control, consider passing the recipe directly to TE modules instead. + See the `TE JAX Integration notebook `_ + for details. + +**Mixed precision with 8- or 4-bit precisions** + +From now on, we will refer to FP8/MXFP8/NVFP4 etc. as *low precision* +and to FP32/BF16/FP16 as *high precision*. This terminology will be +used throughout the rest of the documentation. + +Not all operations run in low precision: + +- **Linear operations**: run in low precision. +- **Attention computations**: run in high precision by default (some recipes allow low precision as an option). +- **Other operations** (layer normalization, softmax, etc.): run in high precision. + +Within high-precision operations, there are two categories: + +- **Configurable precision**: most operations run in parameter precision (FP32/BF16/FP16) or the precision specified by ``torch.autocast``. +- **Fixed FP32 precision**: some operations, or parts of operations—such as the division in layernorm—always run in FP32, regardless of other settings. + +.. raw:: html + :file: img/mixed_precision_operations.svg + +*Figure 3: Default precision of operations in a TransformerLayer forward pass. Only linear operations are in lower precision. Dot product attention is shown as three separate operations (QK^T, Softmax, Scores * V), though in practice these may be fused into a single kernel.* + +**Linear layer data flow** + +Let's see how data flow of a linear layer works by default on a single H100 GPU with FP8 precision: + +H100 (Hopper) architecture natively supports FP8 Matrix Multiplication only in **TN** layout (Transpose-NoTranspose), +so GEMM with tensors ``A`` and ``B`` returns ``B * A^T``. + +*Forward pass* + +* Input is quantized to FP8 – both ``input`` and ``input^T`` quantized versions are created. +* Weights are stored in high precision and quantized to low precision before the GEMM – both ``weight`` and ``weight^T`` quantized versions are created. +* FP8 GEMM with layout **TN** is run with ``weight`` and ``input`` tensors, +* Outputs – ``input * weight^T`` tensor – are returned in high precision. + +*Backward pass* + +* Output gradients are quantized to FP8 – both ``output_grad`` and ``output_grad^T`` quantized versions are created. +* FP8 GEMM with layout **TN** is performed with ``weight^T`` and ``output_grad`` tensors to compute input gradients. +* FP8 GEMM with layout **TN** is performed with ``input^T`` and ``output_grad^T`` tensors to compute weight gradients. +* Input gradients – ``output_grad * weight`` tensor – are returned in high precision. +* Weight gradients – ``output_grad^T * input`` tensor – are returned in high precision. + + +.. raw:: html + :file: img/fp8_linear_flow.svg + +*Figure 4: Forward pass of a Linear layer with low precision data flow.* diff --git a/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg b/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg new file mode 100644 index 0000000000..30f16d9a71 --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/img/fp8_1d_scaling.svg @@ -0,0 +1,177 @@ + + + + + + + + MXFP8 + (One scaling factor per 32 elements) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + E8M0 scaling factors (one per 32 elements) + + + diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg new file mode 100644 index 0000000000..42ea0308bb --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_row_col.svg @@ -0,0 +1,266 @@ + + + + + + + Rowwise (1x32 blocks) + + + + Data + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + + + Scales + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + + Columnwise (32x1 blocks) + + + + Data + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + + + Scales + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + + diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg new file mode 100644 index 0000000000..6e4ed44d56 --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_scale_linearize_and_swizzle.svg @@ -0,0 +1,190 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + 1 + 2 + 3 + + K + + + 1 + K + + + 2 + K + + + 3 + + 2K + + + 1 + 2K + + + 1 + 2K + + + 3 + + + … + … + … + + + ⋮ + ⋮ + ⋱ + + + 128x4 + + + + + + + + + + + + 1 + + + 2 + + + … + + + K + 1 + + + K + 2 + + + … + + + 1x512 + + + + + + + 128 4-bit elements + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 0 + 1 + 2 + 3 + 4 + 5 + 6 + 7 + ... + + + + + + + + + + + + + + + + + + + + + + + 0 + 32 + 64 + 96 + 1 + 33 + 65 + 97 + ... + + + + diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_swizzle_both_tensors.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_swizzle_both_tensors.svg new file mode 100644 index 0000000000..d8489ecc4f --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_swizzle_both_tensors.svg @@ -0,0 +1,101 @@ + + + + + + + + + + + + + + + + Input Tensor + + FP32/BF16 + + + + + + + + Quantize + + + + + + + MXFP8 Tensor + + + + + Scales + + + + FP8 Data + + + + + + + + Communication + (All-Gather) + (Optional) + + + + + + + Swizzle + + + + + + + MXFP8 Tensor + + + + + Swizzle Scales + + + + FP8 Data + + + + + + + + GEMM + + diff --git a/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg b/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg new file mode 100644 index 0000000000..3b81ff0a36 --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/img/mxfp8_tensor_scaling_layout.svg @@ -0,0 +1,63 @@ + + + + + + + + + + + + + + + + + + … + … + … + + + ⋮ + ⋮ + ⋮ + ⋱ + + + FP8 Tensor (128×128 blocks) + + + + + + + + + + + + + + + + + … + … + … + + + ⋮ + ⋮ + ⋱ + + + Scaling Factors (128×4 blocks) + diff --git a/docs/features/low_precision_training/mxfp8/jax_mxfp8_example.py b/docs/features/low_precision_training/mxfp8/jax_mxfp8_example.py new file mode 100644 index 0000000000..96ef1a2573 --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/jax_mxfp8_example.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Check for Blackwell or newer GPU +from transformer_engine.jax.quantize import get_device_compute_capability + +assert ( + get_device_compute_capability() >= 100 +), f"MXFP8 requires SM100 (Blackwell) or later, got SM{get_device_compute_capability()}" + +# START_MXFP8_EXAMPLE + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import DenseGeneral +from transformer_engine.common.recipe import MXFP8BlockScaling, Format + +# Create MXFP8 recipe +recipe = MXFP8BlockScaling( + fp8_format=Format.E4M3, # FP8 format (default: E4M3, E5M2 not supported) +) + +with te.autocast(enabled=True, recipe=recipe): + # Initialize layer and data + layer = DenseGeneral(features=1024) + key = jax.random.PRNGKey(0) + x = jax.random.normal(key, (32, 128, 1024), dtype=jnp.bfloat16) + var_collect = layer.init(key, x) + + # Forward and backward pass + def loss_fn(var_collect): + output = layer.apply(var_collect, x) + return output.sum() + + loss, grads = jax.value_and_grad(loss_fn)(var_collect) + +# END_MXFP8_EXAMPLE diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst new file mode 100644 index 0000000000..f8f8f48b0d --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst @@ -0,0 +1,213 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +MXFP8 +===== + + +MXFP8 (Microscaling FP8) is an enhanced FP8 blockwise scaling recipe that leverages native hardware +acceleration on Blackwell GPUs (SM 10.0+). By using one scaling factor per 32 consecutive values +(rather than 128), MXFP8 delivers finer-grained quantization with improved numerical precision. + + + +Data Format +----------- + +The representation of an FP8 tensor element ``x`` in MXFP8 precision is given by: + +.. code-block:: python + + x = x_fp8 * s_block + +where + +* ``x_fp8`` is the FP8 value in E4M3 format, +* ``s_block`` is a local **E8M0** scaling factor shared by a block of 32 elements. + E8M0 is an 8-bit format with 8 exponent bits and 0 mantissa bits, representing only powers of 2. + + +**FP8 format** + +Like FP8 Blockwise Scaling, E4M3 is used by default for both forward and backward passes. +The finer-grained scaling provides sufficient dynamic range without requiring the E5M2 format. +The ``fp8_format`` parameter also supports ``HYBRID`` mode (E4M3 for forward, E5M2 for backward). +Pure E5M2 training is not supported. + + +**Block size** + +Block size is 32. +Blocks are one-dimensional, containing 32 consecutive values. No 2D scaling is performed. + +There are some assumptions on the dimensions of the tensor: + +* the tensor must have at least 2 dimensions, +* the last dimension must be divisible by 32, +* the product of all dimensions except the last must be divisible by 32. + + +**Scaling factors** + +Scaling factors are stored as E8M0 (8 exponent bits, 0 mantissa bits), which inherently represents +powers of 2. This differs from FP8 Blockwise Scaling, which uses 32-bit floating point numbers +optionally constrained to powers of 2. Note that FP32 also has 8 exponent bits, so the representable +ranges are the same when the power-of-2 constraint is enabled. + +Each block's scaling factor is computed through the following steps: + +1. Find the maximum absolute value (``amax_block``) across all 32 elements in the block. +2. Compute the E8M0 biased exponent: ``e = float_to_e8m0(amax_block / max_fp8)``, where ``max_fp8 = 448`` + (the maximum representable value in E4M3 format). + + Since E8M0 and FP32 share the same exponent bias (127), ``float_to_e8m0`` simply extracts + the 8-bit exponent from the FP32 representation, rounding up if the mantissa is non-zero. + +3. The scaling factor is ``s_block = 2^(e - 127)``. + +This ensures that the largest value in each block fits within the FP8 representable range without overflow. + + +.. raw:: html + :file: img/fp8_1d_scaling.svg + +*Figure 1. MXFP8 uses one E8M0 scaling factor per 32 consecutive elements, providing fine-grained +quantization and compact scaling factor representation.* + + +Handling transposes +------------------- + +Blackwell architecture supports multiple FP8 GEMM layouts (TN, NT, NN), so columnwise usage +does not require explicit transposition. However, rowwise and columnwise quantizations are different: + +- *Rowwise* - 1 scaling factor per 32 consecutive elements along a row (1×32 blocks). +- *Columnwise* - 1 scaling factor per 32 consecutive elements along a column (32×1 blocks). + +Since the scaling factor blocks have different orientations, rowwise and columnwise MXFP8 tensors +are numerically different — one cannot derive one from the other. Both must be quantized +independently from the full-precision data. + +.. raw:: html + :file: img/mxfp8_row_col.svg + +*Figure 2. MXFP8 rowwise vs columnwise quantization layout.* + + +Distributed training +-------------------- + +**Scale synchronization** + +The blockwise scaled tensor does not need any scale synchronization among the nodes. +This is because each scaling factor is local to its 32-element block, +unlike :doc:`FP8 Current <../fp8_current_scaling/fp8_current_scaling>`/:doc:`Delayed Scaling <../fp8_delayed_scaling/fp8_delayed_scaling>` where a single global scale applies to the entire tensor, even when sharded. + +**Quantized all-gather** + +MXFP8 all-gather is supported. + + +Examples +-------- + +Here's how to use MXFP8 recipe in PyTorch and JAX: + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: pytorch_mxfp8_example.py + :language: python + :start-after: # START_MXFP8_EXAMPLE + :end-before: # END_MXFP8_EXAMPLE + + .. tab:: JAX + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: jax_mxfp8_example.py + :language: python + :start-after: # START_MXFP8_EXAMPLE + :end-before: # END_MXFP8_EXAMPLE + + +Supported devices +----------------- + +SM 10.0, SM 10.3 + + +---- + +Developer Notes +--------------- + +This section contains implementation details that may be useful for developers +but are not required for using MXFP8 in practice. + +Swizzling scaling factors +^^^^^^^^^^^^^^^^^^^^^^^^^ + +Like :doc:`FP8 Blockwise Scaling <../fp8_blockwise_scaling/fp8_blockwise_scaling>`, MXFP8 uses different data layouts for communication and computation. +MXFP8 GEMMs require scaling factors in a specific hardware layout +(see `cuBLAS documentation `__). +The conversion to this GEMM-ready layout is called *swizzling*. When no communication is needed, +swizzling can be fused with quantization. When communication is required, swizzled scaling factors +cannot be communicated across devices, so Transformer Engine performs swizzling after communication, +just before each GEMM operation. + +.. raw:: html + :file: img/mxfp8_swizzle_both_tensors.svg + +*Figure 3. MXFP8 swizzling process: standard scaling factors are rearranged into the hardware-required layout.* + + +Blackwell Tensor Cores compute matrix multiplications using ``128x128`` tiles. +Scaling factors are stored in row-major order, but to process a tile, we need a ``128x4`` vertical +slice of scaling factors. In row-major storage, these vertical slices are scattered in memory +with gaps between each row. The hardware requires them to be stored contiguously. + +.. raw:: html + :file: img/mxfp8_tensor_scaling_layout.svg + +*Figure 4. FP8 tensor (left) is divided into 128x128 tiles. Each tile requires a 128x4 block of scaling factors (right). These vertical blocks are not contiguous in memory.* + +Swizzling transforms the layout to meet hardware requirements by: + +1. **Linearizing** the ``128x4`` blocks so they are stored contiguously one after another. +2. **Permuting** the 4-byte elements within each block. + +Specifically, if we index the 128 4-byte elements in a scaling factor block as :math:`0, 1, \dots, 127`, the hardware expects them in the following interleaved order: + +.. code-block:: text + + 0, 32, 64, 96, 1, 33, 65, 97, ..., k, 32 + k, 64 + k, 96 + k, ..., 31, 63, 95, 127 + + +.. raw:: html + :file: img/mxfp8_scale_linearize_and_swizzle.svg + +*Figure 5. Linearization and swizzling of scaling factors. The 2D grid of scaling factors is first flattened into a contiguous sequence of blocks (top), then the rows within each block are interleaved to match the hardware access pattern (bottom).* + +For columnwise scaling factors, the process is analogous but with ``4x128`` horizontal blocks instead of ``128x4`` vertical blocks. + +All-gather of columnwise tensors +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +All-gather of columnwise tensors is supported and necessary because: + +- columnwise quantized tensors cannot be computed from rowwise quantized ones, +- gathering high-precision tensors is avoided in most cases for performance reasons. \ No newline at end of file diff --git a/docs/features/low_precision_training/mxfp8/pytorch_mxfp8_example.py b/docs/features/low_precision_training/mxfp8/pytorch_mxfp8_example.py new file mode 100644 index 0000000000..3cc70137b5 --- /dev/null +++ b/docs/features/low_precision_training/mxfp8/pytorch_mxfp8_example.py @@ -0,0 +1,34 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Check for Blackwell or newer GPU +major, minor = torch.cuda.get_device_capability() +assert major >= 10, f"MXFP8 requires SM100 (Blackwell) or later, got SM{major}{minor}" + +# START_MXFP8_EXAMPLE + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import MXFP8BlockScaling, Format + +# Create MXFP8 recipe +recipe = MXFP8BlockScaling( + fp8_format=Format.E4M3, # E4M3 (default) or HYBRID; pure E5M2 not supported +) + +# Create a linear layer with bfloat16 parameters +layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + +# Forward and backward pass +inp = torch.randn(32, 128, 1024, dtype=torch.bfloat16, device="cuda") + +with te.autocast(enabled=True, recipe=recipe): + output = layer(inp) + loss = output.sum() + +loss.backward() + +# END_MXFP8_EXAMPLE diff --git a/docs/features/low_precision_training/nvfp4/img/nvfp4_all_gather.svg b/docs/features/low_precision_training/nvfp4/img/nvfp4_all_gather.svg new file mode 100644 index 0000000000..3e215551a7 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/img/nvfp4_all_gather.svg @@ -0,0 +1,118 @@ + + + + + + + + + + + Quantization + All-Gather for NVFP4 + + + + High Precision + Tensor + + + + + + + Compute + Amax + + + + + + + Synchronize + Amax + + + + + + + Compute + s_global + + + + + + + Scale + Cast + (s_block, + s_global) + + + + + + + NVFP4 + Tensor + + + + + + + All-Gather + + + + + + + NVFP4 Gathered + Tensor + + + diff --git a/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg b/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg new file mode 100644 index 0000000000..05e67b7889 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/img/nvfp4_hierarchical_scaling.svg @@ -0,0 +1,186 @@ + + + + + + + + NVFP4 Hierarchical Scaling + (Block scaling + Global scaling) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + + + … + … + + + … + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + … + … + … + + + E4M3 scaling factors (one per 16 elements) + + + + + Global Scale (FP32) + (one per tensor) + + + + + + \ No newline at end of file diff --git a/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg b/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg new file mode 100644 index 0000000000..30363d6ce2 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/img/nvfp4_row_col.svg @@ -0,0 +1,208 @@ + + + + + + + Rowwise (1×16 blocks) + + + + Data [A, B] + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + ⋮ + ⋮ + + + + + + + s_block [A, B/16] + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + ⋮ + ⋮ + + + + + + + s_global + + + + + Columnwise (16×1 blocks) — transposed storage + + + + Data [B, A] + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + ⋮ + ⋮ + + + + + + + s_block [B, A/16] + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + … + … + ⋮ + ⋮ + + + + + + + s_global + + + diff --git a/docs/features/low_precision_training/nvfp4/img/nvfp4_vs_fp8.svg b/docs/features/low_precision_training/nvfp4/img/nvfp4_vs_fp8.svg new file mode 100644 index 0000000000..68f6bf9039 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/img/nvfp4_vs_fp8.svg @@ -0,0 +1,91 @@ + + + + + + + FP8 E4M3 + + + + 0 + + + + 1 + + 0 + + 0 + + 0 + + + + 1 + + 1 + + 1 + + (1 sign, 4 exp, 3 mantissa) + + + + FP8 E5M2 + + + + 0 + + + + 1 + + 0 + + 0 + + 0 + + 0 + + + + 1 + + 1 + + (1 sign, 5 exp, 2 mantissa) + + + + NVFP4 + + + + 0 + + + + 1 + + 0 + + + + 1 + + (1 sign, 2 exp, 1 mantissa) + + + + diff --git a/docs/features/low_precision_training/nvfp4/img/rht.svg b/docs/features/low_precision_training/nvfp4/img/rht.svg new file mode 100644 index 0000000000..0250c27ae5 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/img/rht.svg @@ -0,0 +1,138 @@ + + + + + + + + + + + + Random Hadamard Transform for WGRAD GEMM + + + + + + + Without RHT + + + + + Activations + + + + + + + Quantize + + + + + + + WGRAD + GEMM + + + + + Output Grad + + + + + + + Quantize + + + + + + + + + + Weight Grad + + + + + With RHT + + + + + Activations + + + + + + + RHT + + + + + + + Quantize + + + + + + + WGRAD + GEMM + + + + + Output Grad + + + + + + + RHT + + + + + + + Quantize + + + + + + + + + + Weight Grad + + + diff --git a/docs/features/low_precision_training/nvfp4/img/stochastic_rounding.svg b/docs/features/low_precision_training/nvfp4/img/stochastic_rounding.svg new file mode 100644 index 0000000000..eb745f6e84 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/img/stochastic_rounding.svg @@ -0,0 +1,95 @@ + + + + + + + + + + + + Round to Nearest + + + + + + + v₁ + + + + v₂ + + + + x + + + + + Round to v₁ + + + 100% + + + Round to v₂ + + + 0% + + + + + + + Stochastic Rounding + + + + + + + v₁ + + + + v₂ + + + + x + + + + + Round to v₁ + + + 60% + + + Round to v₂ + + + 40% + + + + + diff --git a/docs/features/low_precision_training/nvfp4/jax_nvfp4_example.py b/docs/features/low_precision_training/nvfp4/jax_nvfp4_example.py new file mode 100644 index 0000000000..6c94f31345 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/jax_nvfp4_example.py @@ -0,0 +1,43 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Check for Blackwell or newer GPU +from transformer_engine.jax.quantize import get_device_compute_capability + +assert ( + get_device_compute_capability() >= 100 +), f"NVFP4 requires SM100 (Blackwell) or later, got SM{get_device_compute_capability()}" + +# START_NVFP4_EXAMPLE + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import DenseGeneral +from transformer_engine.common.recipe import NVFP4BlockScaling + +# Define NVFP4 recipe +# 2D weight quantization and RHT are enabled by default +recipe = NVFP4BlockScaling() +# To disable features, use: +# recipe = NVFP4BlockScaling(disable_rht=True, disable_2d_quantization=True) + +with te.autocast(enabled=True, recipe=recipe): + # Initialize layer and data + layer = DenseGeneral(features=1024) + key, sr_key = jax.random.split(jax.random.PRNGKey(0)) + x = jax.random.normal(key, (32, 128, 1024), dtype=jnp.bfloat16) + + # NVFP4 requires sr_rng for stochastic rounding + rngs = {"sr_rng": sr_key} + var_collect = layer.init({"params": key, "sr_rng": sr_key}, x) + + # Forward and backward pass + def loss_fn(var_collect): + output = layer.apply(var_collect, x, rngs=rngs) + return output.sum() + + loss, grads = jax.value_and_grad(loss_fn)(var_collect) + +# END_NVFP4_EXAMPLE diff --git a/docs/features/low_precision_training/nvfp4/nvfp4.rst b/docs/features/low_precision_training/nvfp4/nvfp4.rst new file mode 100644 index 0000000000..0415963a71 --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/nvfp4.rst @@ -0,0 +1,275 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +NVFP4 +=================================== + +NVFP4 is the first 4-bit recipe introduced in Transformer Engine – +please refer to the `NVFP4 paper `__ for more details. +It is a more complex recipe than the previous ones – apart from the new data format, +it introduces multiple features which help training stability. + +Data Format +---------------------- + +The NVFP4 datatype consists of 1 sign bit, 2 exponent bits, and 1 mantissa bit (E2M1). +It can represent values of magnitude up to +/- 6. +NVFP4 uses a hierarchical block scaling approach where multiple scaling factors are combined to recover the high precision value. + +.. raw:: html + :file: img/nvfp4_vs_fp8.svg + +*Figure 1. Bit layout comparison between standard FP8 formats (E4M3 and E5M2) and NVFP4 (E2M1).* + + +The representation of an NVFP4 tensor element ``x`` is given by: + +.. code-block:: python + + x = x_e2m1 * s_block * s_global + +where + +* ``x_e2m1`` is the 4-bit value, +* ``s_block`` is a local **FP8 E4M3** scaling factor shared by a block of 16 consecutive elements, +* ``s_global`` is a global **FP32** scaling factor applied to the entire tensor. + +**Scaling Factor Computation** + +The scaling factors are computed as follows: + +1. Global scaling factor (``s_global``): + +.. code-block:: python + + s_global = global_amax / (fp8_max * fp4_max) + # where: + # - global_amax: maximum absolute value across the entire tensor + # - fp8_max: maximum representable value in FP8 E4M3 (448.0) + # - fp4_max: maximum representable value in NVFP4 E2M1 (6.0) + +2. Block scaling factor (``s_block``): + +.. code-block:: python + + s_block = (block_amax / fp4_max) / s_global + # where: + # - block_amax: maximum absolute value within the block + # - fp4_max: maximum representable value in NVFP4 E2M1 (6.0) + # - s_block is stored in FP8 E4M3 format + + +.. raw:: html + :file: img/nvfp4_hierarchical_scaling.svg + +*Figure 2. NVFP4 hierarchical scaling structure showing the combination of block-level and global scaling factors.* + +This hierarchical structure uses fine-grained block scaling to handle the tensor's dynamic range, +while the FP4 values represent the block-level dynamic range. The global scaling factor +aligns values to the representable range of the E4M3 × E2M1 combination. + +**2D weight scaling** + +NVFP4 can be: + +* 1 dimensional - each block of 16 consecutive elements shares a scaling factor, +* 2 dimensional - each block of 16x16 elements shares a scaling factor. + +By default, NVFP4 uses 2D scaling for weights and 1D scaling for activations and gradients. +Set ``disable_2d_quantization=True`` in the recipe configuration to force 1D scaling for weights as well (activations and gradients always use 1D). +The motivation for using 2D scaling for weights is to ensure that rowwise and columnwise +quantized tensors are numerically equivalent. +Please refer to the `NVFP4 paper `__ for more details. + + +Stochastic Rounding +------------------- + +Stochastic rounding is applied when casting scaled values to NVFP4 format. Instead of deterministic rounding +(always rounding to nearest even value), each scaled value is probabilistically rounded to one of the two +nearest representable NVFP4 values. The probability of rounding to a given value is inversely proportional to +the distance to that value, which ensures that the expected value of the quantized +tensor equals the original value, eliminating systematic quantization bias during training. +Stochastic rounding is hardware-accelerated using native GPU instructions introduced with the +Blackwell architecture. + +.. raw:: html + :file: img/stochastic_rounding.svg + +*Figure 3. Stochastic rounding illustration. Given a value* ``x`` *to be quantized, and the two nearest +representable NVFP4 values* ``v1`` *(lower) and* ``v2`` *(higher), deterministic rounding always +rounds to the nearest value, while stochastic rounding probabilistically rounds to either value. +If* ``x`` *is 40% of the way from* ``v1`` *to* ``v2``, *there is a 60% chance of rounding to* ``v1`` +*and a 40% chance of rounding to* ``v2``. + +Stochastic rounding is enabled only for gradients. It can be disabled by setting +``disable_stochastic_rounding=True`` in the recipe configuration. + + +Random Hadamard Transform +-------------------------- + +Random Hadamard Transform (RHT) applies an orthogonal rotation to the tensor **before quantization**, +smoothing outliers in the tensor distributions and making them easier to represent accurately in NVFP4. +RHT is applied to columnwise quantization of inputs and gradients, which are operands +for the **wgrad GEMM**. This GEMM is particularly sensitive +to quantization errors, hence the additional outlier smoothing. +RHT is supported only for BF16 inputs/gradients. + +The transform is defined as: + +.. math:: + + x' = x H + +where :math:`H` is the RHT matrix defined below. The quantization scale factor is computed +from the rotated tensor :math:`x'`. + +**Hadamard matrix** + +The :math:`d \times d` Hadamard matrix has elements :math:`\pm 1` and satisfies :math:`H_d H_d^T = d I`. +When normalized by :math:`1/\sqrt{d}`, the matrix becomes orthogonal and can be applied +to both operands of a matrix multiplication: + +.. math:: + + C = (AH)(H^T B) = AB + +where the transforms cancel within the dot-product since :math:`H H^T = I`. + +**Sign matrix** + +In the RHT implementation, a :math:`d`-dimensional diagonal sign matrix :math:`S_d` is applied +together with the Hadamard matrix: + +.. math:: + + H = \frac{1}{\sqrt{d}} S_d H_d + +where diagonal entries of :math:`S_d` are :math:`\{-1, 1\}` and flip the signs of different rows of :math:`H_d`. +As described in the paper, a single random sign vector is shared across all linear layers throughout training. +In the implementation, this vector is fixed and the RHT matrix is computed once at initialization and cached. + +**Tiled implementation** + +The Hadamard transform is performed in a tiled approach along the last dimension of the tensor. +For an :math:`m \times k` tensor, the data is reshaped to :math:`(mk/d) \times d` +and multiplied by the :math:`d \times d` matrix :math:`H`. In this implementation, :math:`d = 16`. + + +.. raw:: html + :file: img/rht.svg + +*Figure 4. WGRAD GEMM pipeline comparison: without RHT (left) and with RHT applied (right).* + +Handling transposes +------------------- + +Like :doc:`MXFP8 <../mxfp8/mxfp8>`, NVFP4 requires both rowwise and columnwise quantized tensors +for different GEMM operands. Unlike MXFP8 which supports multiple layouts (TN, NT, NN), +**NVFP4 GEMM only supports the TN layout**. + +NVFP4 stores columnwise data and scaling factors in a **transposed layout**: + +- **Rowwise**: data ``[A, B]`` with 1×16 horizontal blocks, ``scales`` shape ``[A, B/16]`` +- **Columnwise**: data ``[B, A]`` (transposed) with 1×16 horizontal blocks, ``scales`` shape ``[B, A/16]`` + +Scale tensors are padded for hardware alignment: first dimension to a multiple of 128, +second dimension to a multiple of 4 (e.g. rowwise: ``[roundup(A, 128), roundup(B/16, 4)]``). + +.. raw:: html + :file: img/nvfp4_row_col.svg + +*Figure 5. NVFP4 rowwise vs columnwise quantization layout. Unlike MXFP8, columnwise scales are stored transposed.* + + +Distributed training +-------------------- + +**Amax reduction** + +Block scaling factors (``s_block``) do not require synchronization between nodes, +as each scaling factor is local to its block of 16 elements. +However, the global scaling factor (``s_global``) requires amax synchronization for gathered tensors. +For tensors that are gathered (e.g., input and gradient in sequence parallelism), +amax reduction is performed before quantization. +If before synchronization there was ``amax_1`` on node 1, +``amax_2`` on node 2, etc., after synchronization there will be ``max(amax_1, amax_2, ...)`` on all nodes. + +**Quantized all-gather** + +NVFP4 all-gather is supported. + +.. raw:: html + :file: img/nvfp4_all_gather.svg + +*Figure 6. Quantization and all-gather flow for NVFP4 showing amax synchronization and hierarchical scaling.* + +Examples +-------- + +Here's how to use NVFP4 recipe in PyTorch and JAX. The examples show how to configure features like 2D weight quantization and Random Hadamard Transform (RHT): + +.. tabs:: + + .. tab:: PyTorch + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: pytorch_nvfp4_example.py + :language: python + :start-after: # START_NVFP4_EXAMPLE + :end-before: # END_NVFP4_EXAMPLE + + .. tab:: JAX + + .. raw:: html + +
+ Requires SM100 (Blackwell) or later +
+ + .. literalinclude:: jax_nvfp4_example.py + :language: python + :start-after: # START_NVFP4_EXAMPLE + :end-before: # END_NVFP4_EXAMPLE + + +Supported devices +----------------- + +* **Training**: SM 10.0, SM 10.3 +* **Inference**: SM 10.0+ + + +---- + +Developer Notes +--------------- + +This section contains implementation details that may be useful for developers +but are not required for using NVFP4 in practice. + +Swizzling scaling factors +^^^^^^^^^^^^^^^^^^^^^^^^^ + +NVFP4 requires swizzling of block scaling factors (``s_block``) before GEMM operations, +similar to :doc:`MXFP8 <../mxfp8/mxfp8>`. Key differences: + +- Block size is 16 (vs 32 for MXFP8) +- Both rowwise and columnwise scaling factors are swizzled, but thanks to the transposed + columnwise layout, a single rowwise swizzle kernel handles both cases. +- Scaling factors are stored as FP8 E4M3 (vs E8M0 for MXFP8) + +All-gather of columnwise tensors +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +All-gather of columnwise tensors is supported. To enable quantized all-gather, +all nodes must use the same ``s_global``, which is computed from the synchronized global amax. +This is automatically enabled for column-parallel and row-parallel linear layers. diff --git a/docs/features/low_precision_training/nvfp4/pytorch_nvfp4_example.py b/docs/features/low_precision_training/nvfp4/pytorch_nvfp4_example.py new file mode 100644 index 0000000000..07b680defa --- /dev/null +++ b/docs/features/low_precision_training/nvfp4/pytorch_nvfp4_example.py @@ -0,0 +1,35 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Check for Blackwell or newer GPU +major, minor = torch.cuda.get_device_capability() +assert major >= 10, f"NVFP4 requires SM100 (Blackwell) or later, got SM{major}{minor}" + +# START_NVFP4_EXAMPLE + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import NVFP4BlockScaling + +# Define NVFP4 recipe +# 2D weight quantization and RHT are enabled by default +recipe = NVFP4BlockScaling() +# To disable features, use: +# recipe = NVFP4BlockScaling(disable_rht=True, disable_2d_quantization=True) + +# Create a linear layer with bfloat16 parameters +layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + +# Forward and backward pass +inp = torch.randn(32, 128, 1024, dtype=torch.bfloat16, device="cuda") + +with te.autocast(enabled=True, recipe=recipe): + output = layer(inp) + loss = output.sum() + +loss.backward() + +# END_NVFP4_EXAMPLE diff --git a/docs/features/low_precision_training/performance_considerations/fused_layers_jax.py b/docs/features/low_precision_training/performance_considerations/fused_layers_jax.py new file mode 100644 index 0000000000..4f2f39ca34 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/fused_layers_jax.py @@ -0,0 +1,41 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ + +# START_FUSED_LAYERS + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import LayerNorm, DenseGeneral, LayerNormDenseGeneral +from transformer_engine.common.recipe import DelayedScaling + +key = jax.random.PRNGKey(0) +x = jax.random.normal(key, (32, 128, 1024), dtype=jnp.bfloat16) + +# Example 1: Separate LayerNorm and DenseGeneral layers +layer_norm = LayerNorm() +dense = DenseGeneral(features=1024) + +# Initialize parameters +ln_params = layer_norm.init(key, x) +dense_params = dense.init(key, x) + +# Two separate operations +normalized = layer_norm.apply(ln_params, x) +output_separate = dense.apply(dense_params, normalized) + +# Example 2: Fused LayerNormDenseGeneral layer +fused_layer = LayerNormDenseGeneral(features=1024) + +# Initialize and apply with FP8 autocast +recipe = DelayedScaling() +with te.autocast(enabled=True, recipe=recipe): + fused_params = fused_layer.init(key, x) + output_fused, _ = fused_layer.apply(fused_params, x) # Returns (output, ln_output) + +# The fused layer is more efficient as it combines LayerNorm and quantization + +# END_FUSED_LAYERS diff --git a/docs/features/low_precision_training/performance_considerations/fused_layers_pytorch.py b/docs/features/low_precision_training/performance_considerations/fused_layers_pytorch.py new file mode 100644 index 0000000000..2108f45a08 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/fused_layers_pytorch.py @@ -0,0 +1,37 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ +cc = torch.cuda.get_device_capability() +assert cc[0] == 8 and cc[1] >= 9 or cc[0] == 9, "This example requires SM89 (Ada) or SM90 (Hopper)" + +# START_FUSED_LAYERS + +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import DelayedScaling + +# Example 1: Separate LayerNorm and Linear layers +layer_norm = te.LayerNorm(1024) +linear = te.Linear(1024, 1024) + +inp = torch.randn(32, 128, 1024, dtype=torch.bfloat16, device="cuda") + +# Two separate operations: LayerNorm produces FP32, then Linear quantizes it +normalized = layer_norm(inp) +output_separate = linear(normalized) + +# Example 2: Fused LayerNormLinear layer +fused_layer = te.LayerNormLinear(1024, 1024, params_dtype=torch.bfloat16) + +# Single operation: LayerNorm output is directly quantized +recipe = DelayedScaling() +with te.autocast(enabled=True, recipe=recipe): + output_fused = fused_layer(inp) + +# The fused layer is more efficient as it avoids redundant quantization + +# END_FUSED_LAYERS diff --git a/docs/features/low_precision_training/performance_considerations/img/fused_layers.svg b/docs/features/low_precision_training/performance_considerations/img/fused_layers.svg new file mode 100644 index 0000000000..8b7ffb5b50 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/img/fused_layers.svg @@ -0,0 +1,120 @@ + + + + + + + + + + + LayerNorm + Linear: Separate vs Fused + + + + + + Scenario 1: Separate Layers + + + + Input + + + + + + + LayerNorm + + + + + + + Output + + + + + + + Linear + + + + Quantize + + + + + + + FP8 tensor + + + + + + + ... + + + + + + + Output + + + + Scenario 2: Fused Layer + + + + Input + + + + + + + LayerNormLinear + + + + + LayerNorm + Quantize + + + + + + + FP8 tensor + + + + + + + ... + + + + + + + Output + + diff --git a/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg b/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg new file mode 100644 index 0000000000..fa720427e7 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/img/gemm_access_pattern.svg @@ -0,0 +1,214 @@ + + + + + + + + + + NN GEMM + + + + A + + + + + + + + + + + + + + + + + + + + + + + + + + + rowwise + + + + + B + + + + + + + + + + + + + + + + + + + + + + + + + + + columnwise + + + + + A×B + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + TN GEMM + + + + A + + + + + + + + + + + + + + + + + + + + + + + + + + + rowwise + + + + + B + + + + + + + + + + + + + + + + + + + + + + + + + + + rowwise + + + + + A×BT + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/docs/features/low_precision_training/performance_considerations/img/hopper_vs_blackwell_layout.svg b/docs/features/low_precision_training/performance_considerations/img/hopper_vs_blackwell_layout.svg new file mode 100644 index 0000000000..6f9bc4d5a1 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/img/hopper_vs_blackwell_layout.svg @@ -0,0 +1,122 @@ + + + + + + + + FP8 tensor on Hopper + + + + rowwise + + + 0 + + 1 + + 2 + + 3 + + + 4 + + 5 + + 6 + + 7 + + + 8 + + 9 + + 10 + + 11 + + + + + columnwise + + + 0 + + 4 + + 8 + + + 1 + + 5 + + 9 + + + 2 + + 6 + + 10 + + + 3 + + 7 + + 11 + + + + + + + + FP8 tensor on Blackwell + + + + rowwise and columnwise + + + 0 + + 1 + + 2 + + 3 + + + 4 + + 5 + + 6 + + 7 + + + 8 + + 9 + + 10 + + 11 + + + diff --git a/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg b/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg new file mode 100644 index 0000000000..5b61ac2478 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/img/sequence_parallel_quantization.svg @@ -0,0 +1,159 @@ + + + + + + + + + + + All-Gather of Quantized Tensors (one scenario) + + + Input Tensor quantized all-gather + + + FWD: + + + + High Precision + Tensor + + + + + + + Quantize + + + + + + + Rowwise + Quantized + + + + + + + All-Gather + + + + + + ... + + + BWD: + + + + + + + Columnwise + Quantized + + + + + + + All-Gather + + + + + + ... + + + + + + Gradient Tensor quantized all-gather + + + BWD: + + + + High Precision + Tensor + + + + + + + Quantize + + + + + + + Col. Quantized + + + + + + + Row. Quantized + + + + + + + + + + All-Gather + + + + + + ... + + + + + High Precision (FP32/BF16/FP16) + + + Lower Precision (FP8, etc.) + + + Quantization + + + All-Gather + + + + diff --git a/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg b/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg new file mode 100644 index 0000000000..194b1237e1 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/img/transpose_fusion.svg @@ -0,0 +1,181 @@ + + + + + + + + + + + Option 1: Quantize both usages in forward + + + FORWARD: + + + + High Precision + Tensor + + + + + + + Quantize + + + + + + + Quantized + Rowwise + + + BACKWARD: + + + + + + + Quantized + Columnwise + + + + + + Option 2: Separate Quantizations (quantize when needed) + + + FORWARD: + + + + High Precision + Tensor + + + + + + + Quantize + + + + + + + Quantized + Rowwise + + + + + + BACKWARD: + + + + High Precision + Tensor + + + + + + + Quantize + + + + + + + Quantized + Columnwise + + + + + + Option 3: Convert Rowwise to Columnwise in Backward (reuse saved tensor) + + + FORWARD: + + + + High Precision + Tensor + + + + + + + Quantize + + + + + + + Quantized + Rowwise + + + + + + BACKWARD: + + + + Quantized + Rowwise + + + + + + + Make + Columnwise + + + + + + + Quantized + Columnwise + + + + + High Precision (FP32/BF16/FP16) + + + Lower Precision (FP8, etc.) + + + Quantization / Make Columnwise + + + + diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.out b/docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.out new file mode 100644 index 0000000000..717769b1ed --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.out @@ -0,0 +1,9 @@ +# START_MEMORY_USAGE_1 +Tensors in memory: + Shape: (1024, 1024), Dtype: bfloat16, Size: 2048.0 KB + Shape: (1024, 1024), Dtype: bfloat16, Size: 2048.0 KB + Total from all live arrays: 4.00 MB +# END_MEMORY_USAGE_1 +Processing events... +Generated: + No reports were generated diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.py b/docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.py new file mode 100644 index 0000000000..8c1250575e --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_1_jax.py @@ -0,0 +1,45 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ + +print("# START_MEMORY_USAGE_1") + +import jax +import jax.numpy as jnp +from transformer_engine.jax.flax import DenseGeneral + + +key = jax.random.PRNGKey(0) +jax.clear_caches() + + +# Initialize layer with BF16 parameters +layer = DenseGeneral(features=1024, dtype=jnp.bfloat16) +x = jax.random.normal(key, (1024, 1024), dtype=jnp.bfloat16) +var_collect = layer.init(key, x) + + +@jax.jit +def loss_fn(var_collect, x): + output = layer.apply(var_collect, x) + return output.sum() + + +# Trace the backward pass - this allocates saved tensors +_, backward_fn = jax.vjp(loss_fn, var_collect, x) + + +del x + +print("Tensors in memory:") +total_bytes = 0 +for arr in jax.live_arrays(): + total_bytes += arr.nbytes + if arr.nbytes > 200000: # do not count small tensors + print(f" Shape: {arr.shape}, Dtype: {arr.dtype}, Size: {arr.nbytes / 1024:.1f} KB") +print(f" Total from all live arrays: {total_bytes / (1024**2):.2f} MB") + + +print("# END_MEMORY_USAGE_1") diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.out b/docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.out new file mode 100644 index 0000000000..b00749241d --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.out @@ -0,0 +1,4 @@ + +# START_MEMORY_USAGE_1 +Memory usage after forward pass: 6.00 MB +# END_MEMORY_USAGE_1 diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.py b/docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.py new file mode 100644 index 0000000000..dd4ce24471 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_1_pytorch.py @@ -0,0 +1,38 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ +cc = torch.cuda.get_device_capability() +assert cc[0] == 8 and cc[1] >= 9 or cc[0] == 9, "This example requires SM89 (Ada) or SM90 (Hopper)" + +print("# START_MEMORY_USAGE_1") +import torch +import transformer_engine.pytorch as te + + +def measure_memory(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + init_memory = torch.cuda.memory_allocated() + layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + + inp = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda") + out = layer(inp) + del inp # Input is saved by model for backward, not by user script + + mem_after_forward = torch.cuda.memory_allocated() - init_memory + return mem_after_forward + + +# Warmup run +measure_memory() + +# Actual measurement +mem_after_forward = measure_memory() +print(f"Memory usage after forward pass: {mem_after_forward/1024**2:.2f} MB") +# END_MEMORY_USAGE_1 +print("# END_MEMORY_USAGE_1") diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.out b/docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.out new file mode 100644 index 0000000000..ab720b57a8 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.out @@ -0,0 +1,10 @@ +# START_MEMORY_USAGE_2 +Tensors in memory: + Shape: (1024, 1024), Dtype: float8_e4m3fn, Size: 1024.0 KB + Shape: (1024, 1024), Dtype: float8_e4m3fn, Size: 1024.0 KB + Shape: (1024, 1024), Dtype: bfloat16, Size: 2048.0 KB + Total from all live arrays: 4.02 MB +# END_MEMORY_USAGE_2 +Processing events... +Generated: + No reports were generated diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.py b/docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.py new file mode 100644 index 0000000000..3baa55bb8a --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_2_jax.py @@ -0,0 +1,48 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ + +print("# START_MEMORY_USAGE_2") + +import jax +import jax.numpy as jnp +import transformer_engine.jax as te +from transformer_engine.jax.flax import DenseGeneral +from transformer_engine.common.recipe import DelayedScaling + + +key = jax.random.PRNGKey(0) +recipe = DelayedScaling() +jax.clear_caches() + + +# Initialize layer with BF16 parameters (outside autocast) +layer = DenseGeneral(features=1024, dtype=jnp.bfloat16) +x = jax.random.normal(key, (1024, 1024), dtype=jnp.bfloat16) + + +# Forward and backward pass with FP8 compute +with te.autocast(enabled=True, recipe=recipe): + var_collect = layer.init(key, x) + + @jax.jit + def loss_fn(var_collect, x): + output = layer.apply(var_collect, x) + return output.sum() + + # Trace the backward pass - this allocates saved tensors + _, backward_fn = jax.vjp(loss_fn, var_collect, x) + +del x + +print("Tensors in memory:") +total_bytes = 0 +for arr in jax.live_arrays(): + total_bytes += arr.nbytes + if arr.nbytes > 200000: # do not count small tensors + print(f" Shape: {arr.shape}, Dtype: {arr.dtype}, Size: {arr.nbytes / 1024:.1f} KB") +print(f" Total from all live arrays: {total_bytes / (1024**2):.2f} MB") + +print("# END_MEMORY_USAGE_2") diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.out b/docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.out new file mode 100644 index 0000000000..cc1e402581 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.out @@ -0,0 +1,4 @@ + +# START_MEMORY_USAGE_2 +Memory after forward pass: 6.02 MB +# END_MEMORY_USAGE_2 diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.py b/docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.py new file mode 100644 index 0000000000..5c247177d8 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_2_pytorch.py @@ -0,0 +1,39 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ +cc = torch.cuda.get_device_capability() +assert cc[0] == 8 and cc[1] >= 9 or cc[0] == 9, "This example requires SM89 (Ada) or SM90 (Hopper)" + +print("# START_MEMORY_USAGE_2") +import torch +import transformer_engine.pytorch as te + + +def measure_memory(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + init_memory = torch.cuda.memory_allocated() + layer = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + + inp = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda") + with te.autocast(enabled=True): + out = layer(inp) + del inp # Input is saved by model for backward, not by user script + + mem_after_forward = torch.cuda.memory_allocated() - init_memory + return mem_after_forward + + +# Warmup run +measure_memory() + +# Actual measurement +mem_after_forward = measure_memory() +print(f"Memory after forward pass: {mem_after_forward/1024**2:.2f} MB") +# END_MEMORY_USAGE_2 +print("# END_MEMORY_USAGE_2") diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.out b/docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.out new file mode 100644 index 0000000000..ea4d0dc891 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.out @@ -0,0 +1,4 @@ + +# START_MEMORY_USAGE_3 +Memory after forward pass: 3.02 MB +# END_MEMORY_USAGE_3 diff --git a/docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.py b/docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.py new file mode 100644 index 0000000000..ce6905ce4b --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/memory_usage_3_pytorch.py @@ -0,0 +1,44 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ +cc = torch.cuda.get_device_capability() +assert cc[0] == 8 and cc[1] >= 9 or cc[0] == 9, "This example requires SM89 (Ada) or SM90 (Hopper)" + +print("# START_MEMORY_USAGE_3") +import torch +import transformer_engine.pytorch as te + + +def measure_memory(): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + init_memory = torch.cuda.memory_allocated() + + # FP8 inference with FP8 weights + with te.quantized_model_init(enabled=True), torch.no_grad(): + layer_fp8 = te.Linear(1024, 1024, params_dtype=torch.bfloat16) + + with torch.no_grad(): + inp = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda") + with te.autocast(enabled=True): + out = layer_fp8(inp) + del inp # Input is not saved by model for backward in inference + + mem_after_forward = torch.cuda.memory_allocated() - init_memory + + return mem_after_forward + + +# Warmup run +measure_memory() + +# Actual measurement +mem_after_forward = measure_memory() +print(f"Memory after forward pass: {mem_after_forward/1024**2:.2f} MB") +# END_MEMORY_USAGE_3 +print("# END_MEMORY_USAGE_3") diff --git a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst new file mode 100644 index 0000000000..a495af56c1 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst @@ -0,0 +1,473 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Performance Considerations +=================================== + +.. _handling_transposes: + +Handling transposes +------------------- + +In the last chapter we demonstrated that for FP8 on Hopper architecture, +some tensors need to be physically transposed in memory to perform needed GEMMs. +Dealing with transposes in Transformer low precision training is a bit tricky. +Let's start by introducing the concept of *tensor usages*. + +**Tensor usages** + +Each quantized tensor may have two usages: + +- *rowwise usage* -- which is used for matrix multiplication, when the consecutive elements in row are accessed, +- *columnwise usage* -- which is used for matrix multiplication, when the consecutive elements in column are accessed, + +To understand what access of consecutive elements means, let's consider two matrices ``A`` and ``B`` +and analyze how their elements are accessed during multiplication. + +For NN (non-transposed, non-transposed) multiplication ``C = A * B``, the formula is ``C_ij = sum_k(A_ik * B_kj)``. +To compute element ``C_ij``, we iterate over the i-th row of ``A`` (elements ``A_i0, A_i1, ...``) +and the j-th column of ``B`` (elements ``B_0j, B_1j, ...``). Thus, ``A`` is accessed rowwise +and ``B`` is accessed columnwise. + +For NT (non-transposed, transposed) multiplication ``C = A * B^T``, the formula changes to ``C_ij = sum_k(A_ik * B_jk)``. +Now we iterate over the i-th row of ``A`` and the j-th row of ``B`` (elements ``B_j0, B_j1, ...``). +Both tensors are accessed rowwise. + +The figure below illustrates these access patterns: + +.. figure:: img/gemm_access_pattern.svg + :align: center + :width: 60% + :alt: Matrix multiplication access pattern showing rowwise access for first tensor and columnwise access for second tensor + + Figure 1: Access patterns in matrix multiplication for matrices in ``A * B`` and ``A * B^T`` operations. + +Based on the visualization above, we can derive general rules for when each matrix +is accessed in rowwise or columnwise fashion. The key insight is that: + +- The **first tensor** in a matrix multiplication is accessed along its rows (rowwise) when non-transposed, + or along its columns (columnwise) when transposed. +- The **second tensor** follows the opposite pattern: columnwise when non-transposed, rowwise when transposed. + +.. table:: Table 1: Summary of tensor access patterns based on transpose state. + :align: center + + +------------------+--------------+---------------+ + | | First tensor | Second tensor | + +------------------+--------------+---------------+ + | Non-transposed | rowwise | columnwise | + +------------------+--------------+---------------+ + | Transposed | columnwise | rowwise | + +------------------+--------------+---------------+ + +**Input, weight and output gradient usages** + +Now let's apply these rules to a Linear layer. During training, a Linear layer performs +three GEMM operations: one in the forward pass and two in the backward pass. + + +.. table:: Table 2: Tensor access patterns for GEMM operations in a Linear layer during training. + :align: center + + +-------------------+-------------------------------------+---------------------------+---------------------------+ + | GEMM | Formula | First tensor usage | Second tensor usage | + +===================+=====================================+===========================+===========================+ + | Forward | ``output = input * weight^T`` | input: rowwise | weight: rowwise | + +-------------------+-------------------------------------+---------------------------+---------------------------+ + | Weight gradient | ``wgrad = output_grad^T * input`` | output_grad: columnwise | input: columnwise | + +-------------------+-------------------------------------+---------------------------+---------------------------+ + | Input gradient | ``dgrad = output_grad * weight`` | output_grad: rowwise | weight: columnwise | + +-------------------+-------------------------------------+---------------------------+---------------------------+ + +An important observation is that the **forward pass uses only rowwise tensors** - both input +and weight are accessed rowwise. + +The backward pass introduces columnwise access. For weight gradient, both output gradient and input +are accessed columnwise. For input gradient, output gradient is rowwise while weight is columnwise. + +As a result, each tensor (input, weight, output gradient) needs both rowwise and columnwise +usages during training. This has implications for memory layout and transpose operations. + + +**Architecture differences** + +The physical memory layout requirements for rowwise and columnwise usages differ between architectures +and recipes. For FP8 tensors: + +- *Hopper*: cannot efficiently access elements in columnwise fashion, so columnwise tensors need to be physically transposed in memory. Note that higher precision formats (BF16/FP16) do not have this limitation. +- *Blackwell*: supports columnwise access natively, so no transpose is needed. + +We will see that for most of the recipes and devices, rowwise usage and columnwise usage need different tensors. +Thus by *rowwise tensor* and *columnwise tensor* we mean tensors that are used in rowwise and columnwise usages respectively. + +.. figure:: img/hopper_vs_blackwell_layout.svg + :align: center + :alt: Comparison of rowwise and columnwise tensor layouts on Blackwell vs Hopper + + Figure 2: On Blackwell, rowwise and columnwise usages share the same memory layout. + On Hopper, columnwise usage requires a physical transpose. + +**Quantization fusions** + +This section is relevant only for recipes for which columnwise tensors +are different from rowwise tensors. + +Note that performing rowwise and columnwise quantization at the same time +enables some fusions, which usually lead to better performance. +We showcase 3 example scenarios of producing quantized tensors in rowwise and columnwise usages, +TE will use best possible fusion for given recipe and TE module configuration: + +1. *Computation of quantized tensor in both rowwise and columnwise usages in a single kernel in forward pass*. + + This is the fastest one, + but since the columnwise usage is saved for backward pass, it may lead to increased memory usage, + if the high precision tensor also needs to be saved for backward - for example if it is the attention output which is saved anyway. + +2. *Computation of quantized tensor in rowwise usage in forward pass and fused quantization to produce columnwise usage in backward pass*. + + This is usually slower than the previous one, since high precision tensor needs to be read twice. + It is used for example when high precision tensor is gathered both in forward and in backward + and quantized tensor gather is not implemented for such recipe. + +3. *Computation of quantized tensor in rowwise usage in forward pass and transpose to columnwise usage in backward pass*. + + It is more memory efficient than Option 1, but not all recipes can utilize it (otherwise + the quantization accuracy would drop due to double quantization errors). + +Transformer Engine chooses the best possible fusion internally taking the recipe and the operation into account. + +.. raw:: html + :file: img/transpose_fusion.svg + +*Figure 3: Three scenarios of producing quantized tensors in rowwise and columnwise usages.* + + + +Memory usage +------------ + +This section discusses memory usage in low precision training. +Contrary to intuition, FP8 training does not always reduce memory compared to BF16/FP16. + +*Master weights* + +Transformer Engine by default stores weights in high precision and quantizes them to low precision before each GEMM. +Moreover, one can specify which high precision should be used to store the weights in the +model (FP32/BF16/FP16) -- or choose not to store high precision weights in the model at all. +There are multiple scenarios to consider, three of them are listed below: + +1. model weights are in FP32, quantized to low precision before each GEMM, +2. model weights are in BF16/FP16, quantized to low precision before each GEMM, master weights in optimizer are in FP32. +3. model weights are stored directly in low precision, and master weights in optimizer are in FP32. + +Note that each of these scenarios may have different memory footprint. + +*Activations saved for backward* + +Unlike weights, activations do not require a high precision copy for optimizer updates. +As shown in Table 2, the input needs rowwise usage in forward and columnwise usage +for weight gradient computation in backward — so it must be saved between passes. + +The memory impact depends on which scenario from Figure 3 is used. +Additionally, on architectures where rowwise and columnwise usage tensors share the same memory layout +(e.g., FP8 on Blackwell, as shown in Figure 2), a single quantized tensor serves both usages, +reducing memory overhead compared to architectures requiring separate tensors. + +Output gradients, on the other hand, are computed during backward and do not need to be saved — +both rowwise and columnwise usages are produced on the fly as needed. + +The FP8 examples below are analyzed on Hopper (SM90) or Ada (SM89) architecture, where rowwise +and columnwise tensors require separate memory layouts. + +.. tabs:: + + .. tab:: PyTorch + + **1. Baseline: high precision forward pass** + + Let's start with a forward pass in higher precision to establish a baseline. + + .. raw:: html + +
+ Needs to be run on SM89 (Ada) or SM90 (Hopper) +
+ + .. literalinclude:: memory_usage_1_pytorch.py + :language: python + :start-after: # START_MEMORY_USAGE_1 + :end-before: # END_MEMORY_USAGE_1 + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: memory_usage_1_pytorch.out + :language: text + :start-after: # START_MEMORY_USAGE_1 + :end-before: # END_MEMORY_USAGE_1 + + Layer size is ``1024 * 1024 * 2 (2 bytes per parameter) = 2MB``. + Memory after forward pass is ``2 MB (weight) + 2 MB (input saved for backward) + 2 MB (output) = 6 MB``. + + **2. FP8 training with model weights in BF16** + + Now let's see the memory usage in FP8 training with high precision weights. + + .. raw:: html + +
+ Needs to be run on SM89 (Ada) or SM90 (Hopper) +
+ + .. literalinclude:: memory_usage_2_pytorch.py + :language: python + :start-after: # START_MEMORY_USAGE_2 + :end-before: # END_MEMORY_USAGE_2 + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: memory_usage_2_pytorch.out + :language: text + :start-after: # START_MEMORY_USAGE_2 + :end-before: # END_MEMORY_USAGE_2 + + Total memory usage is ``2 MB (weight) + 1 MB (weight in FP8) + 1 MB (input in FP8 saved for backward) + 2 MB (output) = 6 MB``. + + **3. FP8 inference with model weights stored directly in low precision** + + For inference scenarios, model weights can be stored directly in low precision. Since we are only + performing forward passes without gradient updates, master weights in high precision are not needed. + + .. raw:: html + +
+ Needs to be run on SM89 (Ada) or SM90 (Hopper) +
+ + .. literalinclude:: memory_usage_3_pytorch.py + :language: python + :start-after: # START_MEMORY_USAGE_3 + :end-before: # END_MEMORY_USAGE_3 + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: memory_usage_3_pytorch.out + :language: text + :start-after: # START_MEMORY_USAGE_3 + :end-before: # END_MEMORY_USAGE_3 + + Total memory usage is ``1 MB (weight in FP8) + 2 MB (output) = 3 MB``. + This is lower than the BF16 baseline (6 MB) since no copies are saved for backward in inference mode. + + **4. Saving original input instead of quantized** + + By default, TE saves the columnwise quantized input for the backward pass (needed for weight gradient). + However, when the high precision input is already being saved (e.g., for a residual connection), + keeping an additional quantized copy wastes memory. + + The ``save_original_input=True`` option tells the layer to reference the original high precision input + instead of caching a separate quantized copy. The input is re-quantized during backward when needed. + Below is an example with a residual block where input is kept for the addition: + + .. raw:: html + +
+ Needs to be run on SM89 (Ada) or SM90 (Hopper) +
+ + .. literalinclude:: save_original_input_pytorch.py + :language: python + :start-after: # START_SAVE_ORIGINAL_INPUT + :end-before: # END_SAVE_ORIGINAL_INPUT + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: save_original_input_pytorch.out + :language: text + :start-after: # START_SAVE_ORIGINAL_INPUT + :end-before: # END_SAVE_ORIGINAL_INPUT + + .. tab:: JAX + + **1. Baseline: high precision forward pass** + + Let's start with a forward pass in higher precision to establish a baseline. + + .. raw:: html + +
+ Needs to be run on SM89 (Ada) or SM90 (Hopper) +
+ + .. literalinclude:: memory_usage_1_jax.py + :language: python + :start-after: # START_MEMORY_USAGE_1 + :end-before: # END_MEMORY_USAGE_1 + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: memory_usage_1_jax.out + :language: text + :start-after: # START_MEMORY_USAGE_1 + :end-before: # END_MEMORY_USAGE_1 + + Layer size is ``1024 * 1024 * 2 (2 bytes per parameter) = 2MB``. + Memory after forward pass is ``2 MB (weight) + 2 MB (input saved for backward) = 4 MB``. + + **2. FP8 training with master weights in BF16** + + Now let's see the memory usage in FP8 training with high precision weights. + + .. raw:: html + +
+ Needs to be run on SM89 (Ada) or SM90 (Hopper) +
+ + .. literalinclude:: memory_usage_2_jax.py + :language: python + :start-after: # START_MEMORY_USAGE_2 + :end-before: # END_MEMORY_USAGE_2 + + .. raw:: html + +
+ Output: +
+ + .. container:: program-output + + .. literalinclude:: memory_usage_2_jax.out + :language: text + :start-after: # START_MEMORY_USAGE_2 + :end-before: # END_MEMORY_USAGE_2 + + Memory after forward pass is ``2 MB (weight in BF16) + 1 MB (input in FP8) + 1 MB (weight in FP8) = 4 MB``. + +Fused layers +------------ + + +Transformer Engine provides fused layers such as ``LayerNormLinear`` (``LayerNormDenseGeneral`` in JAX) and ``LayerNormMLP`` +that enable kernel fusion optimizations. One key optimization is fusing layer normalization +with quantization. + +In a typical Transformer architecture, LayerNorm precedes a Linear layer. Without fusion, +the LayerNorm outputs in high precision, and the Linear layer must then quantize this input before +performing the GEMM — adding overhead. With ``LayerNormLinear``, these operations are fused +into a single kernel: the LayerNorm output is quantized directly, eliminating the separate +quantization step and reducing memory movement. + + +.. raw:: html + :file: img/fused_layers.svg + +*Figure 4: Comparison of separate LayerNorm and Linear layers versus fused LayerNormLinear layer, showing reduced quantization overhead.* + + +Let's see how we can use fused layers in different frameworks. + +.. tabs:: + + .. tab:: PyTorch + + In PyTorch, Transformer Engine provides fused layers like ``LayerNormLinear`` and ``LayerNormMLP``. + These layers combine normalization and linear operations with optimized quantization. + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada, Hopper, Blackwell, or newer) +
+ + .. literalinclude:: fused_layers_pytorch.py + :language: python + :start-after: # START_FUSED_LAYERS + :end-before: # END_FUSED_LAYERS + + The fused ``LayerNormLinear`` layer is particularly efficient in FP8 training because + it avoids an intermediate quantization step. The LayerNorm output is directly quantized + for the GEMM operation, reducing memory movement and improving performance. + + .. tab:: JAX + + In JAX, Transformer Engine provides fused layers like ``LayerNormDenseGeneral`` and ``LayerNormMLP``. + These layers combine normalization and dense operations with optimized quantization. + + .. raw:: html + +
+ Needs to be run on SM89+ (Ada, Hopper, Blackwell, or newer) +
+ + .. literalinclude:: fused_layers_jax.py + :language: python + :start-after: # START_FUSED_LAYERS + :end-before: # END_FUSED_LAYERS + + The fused ``LayerNormDenseGeneral`` layer is particularly efficient in FP8 training because + it avoids an intermediate quantization step. The LayerNorm output is directly quantized + for the GEMM operation, reducing memory movement and improving performance. + + +Distributed training +-------------------- + +Transformer Engine handles collective operations internally, so users typically don't need to manage +the interaction between communication and low precision computation. + +Recall that each Linear layer involves six tensors: weight, input, output, and their gradients. +Of these, output and gradients are returned in high precision, and weights are generally not +communicated (except in FSDP, which is outside the scope of this section). This leaves two +tensors where low precision communication matters: **input** and **output gradient**. + +For sequence parallelism, TE supports all-gather of quantized tensors. This provides several benefits: + +1. *Reduced memory usage* — no need to store high precision tensors for backward pass. +2. *Reduced communication* — smaller tensors mean less data to transfer. +3. *Parallelized quantization* — quantization work is distributed across GPUs. + +Support varies by recipe — for example, columnwise quantized all-gather is not available +for all configurations. + +The figure below illustrates one possible all-gather scenario for input and output gradient tensors. +Actual behavior depends on the recipe and module configuration. + +.. raw:: html + :file: img/sequence_parallel_quantization.svg + +*Figure 5: All-gather of quantized tensors for input and gradient tensors. +This is one possible scenario — actual behavior varies depending on the recipe and module configuration.* + + diff --git a/docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.out b/docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.out new file mode 100644 index 0000000000..21227220f8 --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.out @@ -0,0 +1,4 @@ +# START_SAVE_ORIGINAL_INPUT +save_original_input=False: 25.0 MB +save_original_input=True: 24.0 MB +# END_SAVE_ORIGINAL_INPUT diff --git a/docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.py b/docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.py new file mode 100644 index 0000000000..c9efa7107e --- /dev/null +++ b/docs/features/low_precision_training/performance_considerations/save_original_input_pytorch.py @@ -0,0 +1,51 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch + +# Requires Ada (SM89) or Hopper (SM90), different results on Blackwell+ +cc = torch.cuda.get_device_capability() +assert cc[0] == 8 and cc[1] >= 9 or cc[0] == 9, "This example requires SM89 (Ada) or SM90 (Hopper)" + +print("# START_SAVE_ORIGINAL_INPUT") +# START_SAVE_ORIGINAL_INPUT +import torch +import transformer_engine.pytorch as te +from transformer_engine.common.recipe import Float8CurrentScaling + +recipe = Float8CurrentScaling() + + +def residual_block(layer, inp): + """Residual connection: input is saved for addition after linear.""" + out = layer(inp) + return out + inp # inp must be kept for this addition + + +def measure_memory(use_save_original): + torch.cuda.empty_cache() + torch.cuda.reset_peak_memory_stats() + + layer = te.Linear( + 1024, 1024, params_dtype=torch.bfloat16, save_original_input=use_save_original + ) + inp = torch.randn(1024, 1024, dtype=torch.bfloat16, device="cuda", requires_grad=True) + + with te.autocast(enabled=True, recipe=recipe): + out = residual_block(layer, inp) + out.sum().backward() + + return torch.cuda.max_memory_allocated() / 1024**2 + + +# Warmup runs +measure_memory(False) +measure_memory(True) + +# Actual measurements +for use_save_original in [False, True]: + peak = measure_memory(use_save_original) + print(f"save_original_input={use_save_original}: {peak:.1f} MB") +# END_SAVE_ORIGINAL_INPUT +print("# END_SAVE_ORIGINAL_INPUT") diff --git a/docs/index.rst b/docs/index.rst index 0edcb863b6..336cd2d47f 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -39,6 +39,14 @@ Transformer Engine documentation api/common api/framework + +.. toctree:: + :hidden: + :caption: Features + + features/low_precision_training/index.rst + + .. toctree:: :hidden: :caption: Examples and Tutorials From 94ba75d7f470a412a5d4f3cea3728792f49b38f7 Mon Sep 17 00:00:00 2001 From: Vadim Markovtsev Date: Tue, 3 Feb 2026 02:00:07 +0100 Subject: [PATCH 189/521] Support building with headers from nvidia wheels (#2623) * Support building with headers from nvidia wheels There are two changes: 1. `import nvidia` returns a namespace package with `__file__` equal to `None` 2. Add the way to force headers from nvidia wheels. Without that envvar, it's practically impossible with CUDA installed system-wide. I successfully built the package with torch using the following `uv` configuration: ``` [tool.uv.extra-build-dependencies] "transformer-engine-torch" = [ "ninja", "nvidia-cuda-crt==13.0.88", "nvidia-cuda-cccl==13.0.85", { requirement = "torch", match-runtime = true }, { requirement = "pytorch-triton", match-runtime = true }, { requirement = "nvidia-cusolver", match-runtime = true }, { requirement = "nvidia-curand", match-runtime = true }, { requirement = "nvidia-cublas", match-runtime = true }, { requirement = "nvidia-cusparse", match-runtime = true }, { requirement = "nvidia-cudnn-cu13", match-runtime = true }, { requirement = "nvidia-nvtx", match-runtime = true }, { requirement = "nvidia-cuda-nvrtc", match-runtime = true }, { requirement = "nvidia-cuda-runtime", match-runtime = true }, ] ``` Signed-off-by: Vadim Markovtsev * Apply suggestion from @ksivaman Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Vadim Markovtsev Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Kirthi Shankar Sivamani --- build_tools/utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index 8a52440310..885901068a 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -228,9 +228,10 @@ def nvcc_path() -> Tuple[str, str]: def get_cuda_include_dirs() -> Tuple[str, str]: """Returns the CUDA header directory.""" + force_wheels = bool(int(os.getenv("NVTE_BUILD_USE_NVIDIA_WHEELS", "0"))) # If cuda is installed via toolkit, all necessary headers # are bundled inside the top level cuda directory. - if cuda_toolkit_include_path() is not None: + if not force_wheels and cuda_toolkit_include_path() is not None: return [cuda_toolkit_include_path()] # Use pip wheels to include all headers. @@ -239,7 +240,10 @@ def get_cuda_include_dirs() -> Tuple[str, str]: except ModuleNotFoundError as e: raise RuntimeError("CUDA not found.") - cuda_root = Path(nvidia.__file__).parent + if nvidia.__file__ is not None: + cuda_root = Path(nvidia.__file__).parent + else: + cuda_root = Path(nvidia.__path__[0]) # namespace return [ subdir / "include" for subdir in cuda_root.iterdir() From 29b84c168ebbc151990e06d7a147532273837376 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Tue, 3 Feb 2026 02:49:34 +0100 Subject: [PATCH 190/521] [Common] Fix NVFP4 tuned-kernel numerics (#2639) * Fixed scaling-factor computation for FP32 to match the reference implementation. Signed-off-by: Oleg Goncharov * Uncommented the tuned kernel path 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> --- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 8 +++---- .../quantize_transpose_nvfp4_tuned_1D.cuh | 21 ++++++++++++++++--- 2 files changed, 22 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 61c6ba9cef..99776db281 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -1168,10 +1168,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, // TODO(Frank): Is there a better way to do this? bool return_transpose = output->has_columnwise_data(); - // if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { - // quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); - // return; - // } + if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { + quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); + return; + } constexpr bool COMPUTE_ACTIVATIONS = false; using ParamOP = Empty; 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 4119001686..061a88fd6d 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 @@ -163,9 +163,24 @@ __device__ __forceinline__ float get_amax_of_pair(const IType2 pair) { template __device__ __forceinline__ SF_TYPE compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const float S_enc) { - constexpr float float_max = detail::TypeExtrema::max; - const float scale_rcp = fminf(S_enc / static_cast(S_dec_block), float_max); - return static_cast(scale_rcp); + NVTE_DEVICE_ERROR("Unsupported scaling-factor type. Only FP32 and BF16 are supported."); +} + +template <> +__device__ __forceinline__ float compute_nvfp4_scaling_coefficient( + const nvfp4_scale_t S_dec_block, const float S_enc) { + const float S_dec = 1.0f / S_enc; + const float scale_rcp = + fminf(1.0f / (static_cast(S_dec_block) * S_dec), detail::TypeExtrema::max); + return scale_rcp; +} + +template <> +__device__ __forceinline__ bf16 +compute_nvfp4_scaling_coefficient(const nvfp4_scale_t S_dec_block, const float S_enc) { + const float scale_rcp = + fminf(S_enc / (static_cast(S_dec_block)), detail::TypeExtrema::max); + return static_cast(scale_rcp); } template From 74faf7ec422229bcecf9e079d96b429da071e7b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 3 Feb 2026 06:29:19 +0100 Subject: [PATCH 191/521] [PyTorch Debug] NVFP4 debug stats support (#2296) * init Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * year update in license Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/debug/1_getting_started.rst | 2 +- docs/debug/3_api_features.rst | 7 +- tests/pytorch/debug/test_log.py | 120 ++++++++++ .../debug/features/disable_fp8_gemm.py | 44 ++-- .../debug/features/disable_fp8_layer.py | 63 +++-- .../features/disable_quantization_gemm.py | 59 +++++ .../features/disable_quantization_layer.py | 61 +++++ .../debug/features/log_fp8_tensor_stats.py | 28 ++- .../debug/features/log_nvfp4_tensor_stats.py | 225 ++++++++++++++++++ .../debug/features/utils/stats_computation.py | 62 +++++ .../debug/pytorch/debug_quantization.py | 52 ++-- 11 files changed, 636 insertions(+), 87 deletions(-) create mode 100644 transformer_engine/debug/features/disable_quantization_gemm.py create mode 100644 transformer_engine/debug/features/disable_quantization_layer.py create mode 100644 transformer_engine/debug/features/log_nvfp4_tensor_stats.py diff --git a/docs/debug/1_getting_started.rst b/docs/debug/1_getting_started.rst index de72b2242d..cce2616998 100644 --- a/docs/debug/1_getting_started.rst +++ b/docs/debug/1_getting_started.rst @@ -15,7 +15,7 @@ Transformer Engine provides a set of precision debug tools which allow you to ea - log the statistics for each of the tensors in every matrix multiply (GEMM) operation, - run selected GEMMs in higher precision, - run current scaling - with one scaling factor per tensor - for particular GEMMs, -- test new precisions and integrate them with FP8 training, +- test new precisions and integrate them with quantized training (FP8, NVFP4, etc.), - ... and many more. There are 4 things one needs to do to use Transformer Engine debug features: diff --git a/docs/debug/3_api_features.rst b/docs/debug/3_api_features.rst index fc48371a3f..a973a0b4fe 100644 --- a/docs/debug/3_api_features.rst +++ b/docs/debug/3_api_features.rst @@ -8,7 +8,10 @@ Debug features .. autoapiclass:: transformer_engine.debug.features.log_tensor_stats.LogTensorStats .. autoapiclass:: transformer_engine.debug.features.log_fp8_tensor_stats.LogFp8TensorStats -.. autoapiclass:: transformer_engine.debug.features.disable_fp8_gemm.DisableFP8GEMM -.. autoapiclass:: transformer_engine.debug.features.disable_fp8_layer.DisableFP8Layer +.. autoapiclass:: transformer_engine.debug.features.log_nvfp4_tensor_stats.LogNvfp4TensorStats +.. autoapiclass:: transformer_engine.debug.features.disable_quantization_gemm.DisableQuantizationGEMM +.. autoapiclass:: transformer_engine.debug.features.disable_quantization_layer.DisableQuantizationLayer .. autoapiclass:: transformer_engine.debug.features.per_tensor_scaling.PerTensorScaling .. autoapiclass:: transformer_engine.debug.features.fake_quant.FakeQuant +.. autoapiclass:: transformer_engine.debug.features.disable_fp8_gemm.DisableFP8GEMM +.. autoapiclass:: transformer_engine.debug.features.disable_fp8_layer.DisableFP8Layer \ No newline at end of file diff --git a/tests/pytorch/debug/test_log.py b/tests/pytorch/debug/test_log.py index 7edc0cc90b..5d6fc41ac7 100644 --- a/tests/pytorch/debug/test_log.py +++ b/tests/pytorch/debug/test_log.py @@ -15,6 +15,7 @@ is_fp8_available, is_mxfp8_available, is_fp8_block_scaling_available, + is_nvfp4_available, ) from transformer_engine.pytorch.quantization import RecipeState from transformer_engine.debug.pytorch.debug_state import TEDebugState @@ -29,6 +30,7 @@ fp8_block_scaling_available, reason_for_no_fp8_block_scaling = is_fp8_block_scaling_available( return_reason=True ) +nvfp4_available, reason_for_no_nvfp4 = is_nvfp4_available(return_reason=True) LOG_QUANTIZED_CONFIG_BASE = """ log: @@ -363,6 +365,124 @@ def test_log_every_3_or_5_layers(layer, configs_dir, feature_dirs): TEDebugState._reset() +# NVFP4 tests +LOG_NVFP4_CONFIG_BASE = """ +log: + layers: + layer_name_regex_pattern: .* + enabled: + True + transformer_engine: + LogNvfp4TensorStats: + enabled: True + stats: [ + {stats} + ] + tensors: [activation, gradient, weight] + freq: 2 + start_step: 0 + end_step: 10 +""" + + +def test_nvfp4_numeric(feature_dirs): + """Test that NVFP4 underflows% and MSE stats are computed correctly with known values.""" + if not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + + log_nvfp4_config = LOG_NVFP4_CONFIG_BASE.format(stats="underflows%, mse") + + with debug_session(log_nvfp4_config, feature_dirs) as log_dir: + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + from transformer_engine.pytorch.quantization import RecipeState + + recipe_state = RecipeState.create( + recipe.NVFP4BlockScaling(), + mode="forward", + num_quantizers=3, + ) + + # Create test tensor with known distribution + torch.manual_seed(42) + tensor = torch.randn(128, 128, dtype=torch.bfloat16).cuda() + # Add some small values that should underflow to zero in FP4 + tensor[0, :16] = 0.0001 + + quantizer = recipe_state.make_quantizers()[0] + quantized_tensor = quantizer(tensor) + + debug_api.transformer_engine.inspect_tensor( + layer_name="test_layer", + tensor_name="activation", + iteration=0, + tp_group=None, + tensor=tensor, + quantizer=quantizer, + rowwise_quantized_tensor=quantized_tensor, + columnwise_quantized_tensor=quantized_tensor, + ) + debug_api.step() + + dequantized_tensor = quantized_tensor.dequantize() + output = read_log(log_dir) + + # Validate both stats are present + assert "nvfp4_underflows%" in output, "underflows% stat missing" + assert "nvfp4_mse" in output, "mse stat missing" + + # Extract values and validate numerics + underflows_value = None + mse_value = None + + for line in output.splitlines(): + if "nvfp4_underflows%" in line and "value=" in line: + underflows_value = float(line.split("value=")[1].split()[0]) + if "nvfp4_mse" in line and "value=" in line: + mse_value = float(line.split("value=")[1].split()[0]) + + # Compute expected underflows: non-zero elements that became zero after quantization + orig_nonzero_mask = tensor != 0 + dequant_zero_mask = dequantized_tensor == 0 + expected_underflows = ( + (orig_nonzero_mask & dequant_zero_mask).sum().float() / tensor.numel() * 100 + ) + + # Allow some tolerance + assert underflows_value == pytest.approx(expected_underflows.cpu().item(), abs=1e-4) + + # Compute expected MSE + expected_mse = torch.nn.functional.mse_loss( + dequantized_tensor.float(), tensor.float(), reduction="mean" + ) + + assert mse_value == pytest.approx(expected_mse.cpu().item(), abs=1e-4) + + +def test_fp8_stats_allows_nvfp4_with_recipe_prefix(feature_dirs): + """Test that LogFp8TensorStats allows recipe-prefixed stats with NVFP4 for what-if analysis.""" + if not nvfp4_available: + pytest.skip(reason_for_no_nvfp4) + + # Use recipe-prefixed stat with NVFP4 - should work (computes MXFP8 separately) + log_fp8_config = LOG_QUANTIZED_CONFIG_BASE.format(stats="mxfp8_mse") + + with debug_session(log_fp8_config, feature_dirs) as log_dir: + model = te.Linear(128, 128, params_dtype=torch.bfloat16) + inp = torch.randn(128, 128, dtype=torch.bfloat16).cuda() + + # Should work - recipe-prefixed stats compute MXFP8 separately for comparison + for _ in range(2): + with te.autocast(recipe=recipe.NVFP4BlockScaling()): + output = model(inp) + loss = output.sum() + loss.backward() + debug_api.step() + + output = read_log(log_dir) + # Should have logged MXFP8 MSE stat (what-if scenario) + assert "mxfp8_mse" in output + + def test_log_grouped_gemm(feature_dirs): if not fp8_available: pytest.skip(reason_for_no_fp8) diff --git a/transformer_engine/debug/features/disable_fp8_gemm.py b/transformer_engine/debug/features/disable_fp8_gemm.py index befebb412c..9bbb7ef4ad 100644 --- a/transformer_engine/debug/features/disable_fp8_gemm.py +++ b/transformer_engine/debug/features/disable_fp8_gemm.py @@ -2,17 +2,28 @@ # # See LICENSE for license information. -"""DisableFP8GEMM Feature support for nvidia-dlframework-inspect""" +"""DisableFP8GEMM Feature support for nvidia-dlframework-inspect -from nvdlfw_inspect.registry import Registry, api_method -from transformer_engine.debug.features.api import TEConfigAPIMapper +DEPRECATED: This is a backward compatibility alias for DisableQuantizationGEMM. +New code should use DisableQuantizationGEMM instead, which works with all quantization formats. +""" + +import warnings + +from nvdlfw_inspect.registry import Registry +from transformer_engine.debug.features.disable_quantization_gemm import DisableQuantizationGEMM @Registry.register_feature(namespace="transformer_engine") -class DisableFP8GEMM(TEConfigAPIMapper): +class DisableFP8GEMM(DisableQuantizationGEMM): """ GEMM operations are executed in higher precision, even when FP8 autocast is enabled. + .. deprecated:: + Use :class:`DisableQuantizationGEMM` instead. This class is maintained for + backward compatibility only. DisableQuantizationGEMM works with all quantization + formats (FP8, NVFP4, etc.), not just FP8. + Parameters ---------- @@ -32,22 +43,17 @@ class DisableFP8GEMM(TEConfigAPIMapper): layers: layer_types: [fc1] transformer_engine: - DisableFP8GEMM: + DisableFP8GEMM: # Deprecated: use DisableQuantizationGEMM enabled: True gemms: [dgrad, wgrad] """ - @api_method - def fp8_gemm_enabled( - self, config, layer_name: str, gemm: str, iteration: int - ): # pylint: disable=unused-argument - """API call responsible for choice between high-precision and FP8 GEMM execution.""" - - for key in config: - if key != "gemm": - raise ValueError(f'[NVTORCH INSPECT ERROR] Unexpected key in config: "{key}".') - - # If this feature is invoked, then FP8 GEMM is disabled. - # If not, then default behaviour in TransformerEngineAPI - # is that fp8_gemm() API call returns True. - return False, iteration + 1 + def __init__(self, *args, **kwargs): + warnings.warn( + "DisableFP8GEMM is deprecated. " + "Use DisableQuantizationGEMM instead, which works with all quantization " + "formats (FP8, NVFP4, etc.).", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(*args, **kwargs) diff --git a/transformer_engine/debug/features/disable_fp8_layer.py b/transformer_engine/debug/features/disable_fp8_layer.py index 3839e5f2bf..5ae03ef456 100644 --- a/transformer_engine/debug/features/disable_fp8_layer.py +++ b/transformer_engine/debug/features/disable_fp8_layer.py @@ -2,17 +2,27 @@ # # See LICENSE for license information. -"""DisableFP8Layer Feature support for nvidia-dlframework-inspect""" +"""DisableFP8Layer Feature support for nvidia-dlframework-inspect -import nvdlfw_inspect.api as debug_api -from nvdlfw_inspect.registry import Registry, api_method +DEPRECATED: This is a backward compatibility alias for DisableQuantizationLayer. +New code should use DisableQuantizationLayer instead, which works with all quantization formats. +""" + +import warnings + +from nvdlfw_inspect.registry import Registry +from transformer_engine.debug.features.disable_quantization_layer import DisableQuantizationLayer @Registry.register_feature(namespace="transformer_engine") -class DisableFP8Layer: +class DisableFP8Layer(DisableQuantizationLayer): """ Disables all FP8 GEMMs in the layer. + .. deprecated:: + Use :class:`DisableQuantizationLayer` instead. This class is maintained for + backward compatibility only. DisableQuantizationLayer works with all quantization + formats (FP8, NVFP4, etc.), not just FP8. Example ------- @@ -20,36 +30,19 @@ class DisableFP8Layer: example_disable_fp8_layer: enabled: True - layers: - layer_types: [fc1] - transformer_engine: - DisableFP8Layer: - enabled: True + layers: + layer_types: [fc1] + transformer_engine: + DisableFP8Layer: # Deprecated: use DisableQuantizationLayer + enabled: True """ - @api_method - def fp8_gemm_enabled( - self, config, layer_name: str, gemm: str, iteration: int - ): # pylint: disable=unused-argument - """API call responsible for selecting between high-precision and FP8 GEMM execution.""" - for key in config: - if key not in ["enabled", "gemm"]: - raise ValueError(f'[NVTORCH INSPECT ERROR] Unexpected key in config: "{key}".') - # If FP8 training, disable FP8 for the selected layers if this feature is enabled in config. - debug_api.log_message("FP8 Disabled", layer_name) - - # If this feature is invoked, then FP8 GEMM is disabled. - # If not, then default behavior in TransformerEngineAPI - # is that fp8_gemm() API call returns True. - return False, iteration + 1 - - def parse_config_and_api(self, config, **_kwargs): - """Determines whether to run the API - DisableFP8Layer is the only feature provided by the Transformer Engine - which does not inherit from TEConfigAPIMapper - this mapper is primarly responsible for - parsing gemms and tensors fields from the config, which are not needed for this feature. - - Explanation of the parse_config_and_api can be found in the - nvidia-dlframework-inspect documentation. - """ - return config["enabled"], None + def __init__(self, *args, **kwargs): + warnings.warn( + "DisableFP8Layer is deprecated. " + "Use DisableQuantizationLayer instead, which works with all quantization " + "formats (FP8, NVFP4, etc.).", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(*args, **kwargs) diff --git a/transformer_engine/debug/features/disable_quantization_gemm.py b/transformer_engine/debug/features/disable_quantization_gemm.py new file mode 100644 index 0000000000..932c2f83dd --- /dev/null +++ b/transformer_engine/debug/features/disable_quantization_gemm.py @@ -0,0 +1,59 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DisableQuantizationGEMM Feature support for nvidia-dlframework-inspect""" + +from nvdlfw_inspect.registry import Registry, api_method +from transformer_engine.debug.features.api import TEConfigAPIMapper + + +@Registry.register_feature(namespace="transformer_engine") +class DisableQuantizationGEMM(TEConfigAPIMapper): + """ + Disables specific GEMM operations from using quantization, forcing high-precision execution. + + Works with any quantization format (FP8, NVFP4, etc.). + + Parameters + ---------- + + gemms: List[str] + list of gemms to disable quantization for + + - fprop + - dgrad + - wgrad + + Example + ------- + .. code-block:: yaml + + example_disable_quantization_gemm: + enabled: True + layers: + layer_types: [fc1] + transformer_engine: + DisableQuantizationGEMM: + enabled: True + gemms: [dgrad, wgrad] + """ + + @api_method + def fp8_gemm_enabled( + self, config, layer_name: str, gemm: str, iteration: int + ): # pylint: disable=unused-argument + """API call responsible for choice between high-precision and quantized GEMM execution. + + Note: Method name kept as 'fp8_gemm_enabled' for backward compatibility with the debug API, + but it applies to all quantization formats (FP8, NVFP4, etc.). + """ + + for key in config: + if key != "gemm": + raise ValueError(f'[NVTORCH INSPECT ERROR] Unexpected key in config: "{key}".') + + # If this feature is invoked, then quantized GEMM is disabled (returns to high precision). + # If not, then default behavior in TransformerEngineAPI + # is that fp8_gemm() API call returns True. + return False, iteration + 1 diff --git a/transformer_engine/debug/features/disable_quantization_layer.py b/transformer_engine/debug/features/disable_quantization_layer.py new file mode 100644 index 0000000000..081e310ed2 --- /dev/null +++ b/transformer_engine/debug/features/disable_quantization_layer.py @@ -0,0 +1,61 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DisableQuantizationLayer Feature support for nvidia-dlframework-inspect""" + +import nvdlfw_inspect.api as debug_api +from nvdlfw_inspect.registry import Registry, api_method + + +@Registry.register_feature(namespace="transformer_engine") +class DisableQuantizationLayer: + """ + Disables all quantized GEMMs in the layer, forcing high-precision execution. + + Works with any quantization format (FP8, NVFP4, etc.). + + Example + ------- + .. code-block:: yaml + + example_disable_quantization_layer: + enabled: True + layers: + layer_types: [fc1] + transformer_engine: + DisableQuantizationLayer: + enabled: True + """ + + @api_method + def fp8_gemm_enabled( + self, config, layer_name: str, gemm: str, iteration: int + ): # pylint: disable=unused-argument + """API call responsible for selecting between high-precision and quantized GEMM execution. + + Note: Method name kept as 'fp8_gemm_enabled' for backward compatibility with the debug API, + but it applies to all quantization formats (FP8, NVFP4, etc.). + """ + for key in config: + if key not in ["enabled", "gemm"]: + raise ValueError(f'[NVTORCH INSPECT ERROR] Unexpected key in config: "{key}".') + # If quantized training, disable quantization for the selected layers if this feature is enabled. + debug_api.log_message("Quantization Disabled", layer_name) + + # If this feature is invoked, then quantized GEMM is disabled (returns to high precision). + # If not, then default behavior in TransformerEngineAPI + # is that fp8_gemm() API call returns True. + return False, iteration + 1 + + def parse_config_and_api(self, config, **_kwargs): + """Determines whether to run the API. + + DisableQuantizationLayer is the only feature provided by the Transformer Engine + which does not inherit from TEConfigAPIMapper - this mapper is primarily responsible for + parsing gemms and tensors fields from the config, which are not needed for this feature. + + Explanation of the parse_config_and_api can be found in the + nvidia-dlframework-inspect documentation. + """ + return config["enabled"], None diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index ffcc6b1ad4..df42fb1376 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -9,12 +9,13 @@ import torch import nvdlfw_inspect.api as debug_api - +import transformer_engine_torch as tex from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats from nvdlfw_inspect.registry import Registry, api_method from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS +from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Quantizer, @@ -22,7 +23,14 @@ ) from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer -from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter + +try: + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + _nvfp4_available = True +except ImportError: + _nvfp4_available = False + NVFP4Quantizer = None ALL_RECIPE_NAMES = ["fp8_delayed_scaling", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"] @@ -39,6 +47,8 @@ def _get_recipe_name(quantizer: Optional[Quantizer]): return "mxfp8" if isinstance(quantizer, Float8BlockQuantizer): return "fp8_block_scaling" + if _nvfp4_available and isinstance(quantizer, NVFP4Quantizer): + return "nvfp4" raise ValueError(f"Unsupported quantizer type: {type(quantizer)}") @@ -164,6 +174,16 @@ def check_if_stat_is_supported(self, stat: str, current_recipe: str): if recipe_from_stat != "" and recipe_from_stat not in ALL_RECIPE_NAMES: raise ValueError(f"Stat {stat} contains an unsupported recipe name: {recipe_from_stat}") + # Block any NVFP4 stats in LogFp8TensorStats (FP8-specific logic won't work) + # But allow recipe-prefixed FP8 stats like "mxfp8_underflows%" even with NVFP4 quantizer + if recipe_from_stat == "nvfp4": + raise ValueError( + f"[NVTORCH INSPECT ERROR] Cannot compute NVFP4 stats '{stat}' in LogFp8TensorStats." + " FP8-specific statistics do not work with NVFP4. Use LogNvfp4TensorStats for" + " NVFP4-specific stats, or use FP8 recipe-prefixed stats (e.g.," + " 'mxfp8_underflows%', 'fp8_block_scaling_mse') for what-if FP8 comparisons." + ) + if recipe_from_stat in ["fp8_delayed_scaling", "fp8_current_scaling"] and columnwise: raise ValueError( f"Stat {stat} is not supported. Columnwise tensor statistics are not supported for" @@ -189,6 +209,7 @@ def check_if_stat_is_supported(self, stat: str, current_recipe: str): def get_recipe_from_stat(self, stat: str, default_recipe: str = ""): """Returns the recipe name from the stat string.""" + columnwise_stat = stat.endswith("_columnwise") for recipe_name in ALL_RECIPE_NAMES: if recipe_name in stat: @@ -213,7 +234,7 @@ def update_aux_dict( Yields the aux_dict. Needs to clean after usage, because it possibly change the usage of the quantized tensor. """ - fp8_dtype = None + fp8_dtype = tex.DType.kFloat8E4M3 if recipe_name in ["fp8_delayed_scaling", "fp8_current_scaling", "fp8_block_scaling"]: assert isinstance( quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer, Float8BlockQuantizer) @@ -282,6 +303,7 @@ def inspect_tensor( ), "[NVTORCH INSPECT ERROR] LogFp8TensorStats cannot be run without low-precision recipe." quantized_tensor = rowwise_quantized_tensor + assert isinstance( quantized_tensor, QuantizedTensor ), "[NVTORCH INSPECT ERROR] LogFp8TensorStats quantized_tensor must be a QuantizedTensor." diff --git a/transformer_engine/debug/features/log_nvfp4_tensor_stats.py b/transformer_engine/debug/features/log_nvfp4_tensor_stats.py new file mode 100644 index 0000000000..ec2b3c38d3 --- /dev/null +++ b/transformer_engine/debug/features/log_nvfp4_tensor_stats.py @@ -0,0 +1,225 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""LogNvfp4TensorStats Feature support for nvidia-dlframework-inspect""" + +from typing import Dict, Optional +from contextlib import contextmanager + +import torch +import nvdlfw_inspect.api as debug_api + +from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats +from nvdlfw_inspect.registry import Registry, api_method + +from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS +from transformer_engine.pytorch.tensor import Quantizer, QuantizedTensor +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer +from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter +from transformer_engine.pytorch.tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage + + +@Registry.register_feature(namespace="transformer_engine") +class LogNvfp4TensorStats(BaseLogTensorStats): + """Logs statistics of NVFP4 quantized tensors. + + In distributed runs each rank first computes its local statistics; the values + are gathered the next time `debug_api.step()` is called. Remember to call + `debug_api.step()` every training step so the logs are flushed. + + The feature is micro-batch aware: if several forward/backward passes occur + between successive `debug_api.step()` calls, statistics are accumulated for all + tensors except weights. + + Collecting NVFP4 statistics is expensive. Choosing a larger `freq` reduces the + overhead, and if the feature is skipped for a step the additional cost is + minimal. When no other debug feature is active, the layer runs at normal + Transformer Engine speed. + + Parameters + ---------- + + stats: List[str] + List of statistics to collect. Available stats: + - underflows% - percentage of non-zero elements clipped to 0 (from packed FP4 data) + - mse - mean squared error = sum((quantized_tensor - original_tensor)**2) / num_elements + + tensors/tensors_struct: List[str] + list of tensors to log + - activation, + - gradient, + - weight, + + freq: Optional[int], default = 1 + frequency of logging stats, stats will be logged every `freq` steps + start_step: Optional[int], default = None + start step of logging stats + end_step: Optional[int], default = None + end step of logging stats + start_end_list: Optional[list([int, int])], default = None + non-overlapping list of (start, end) pairs in incremental order. If not None, will ignore start_step and end_step + + Example + ------- + .. code-block:: yaml + + example_nvfp4_tensor_stat_collection: + enabled: True + layers: + layer_types: [layernorm_linear] + transformer_engine: + LogNvfp4TensorStats: + enabled: True + tensors_struct: + - tensor: activation + stats: [underflows%, mse] + freq: 1 + - tensor: gradient + stats: [underflows%, mse] + freq: 5 + start_step: 0 + end_step: 80 + """ + + def check_if_stat_is_supported(self, stat: str): + """Returns True if stat is supported, raises ValueError otherwise.""" + supported_stats = [ + "underflows%", + "mse", + ] + if stat not in supported_stats: + raise ValueError( + f"Stat {stat} is not supported for NVFP4. Supported stats: {supported_stats}" + ) + return True + + def get_stat_with_prefix(self, stat: str) -> str: + """Add nvfp4_ prefix to stat name for use in stats_computation.""" + return f"nvfp4_{stat}" + + @contextmanager + def update_aux_dict( + self, + aux_dict: Dict, + quantized_tensor: QuantizedTensor, + quantizer: Quantizer, # pylint: disable=unused-argument + original_tensor: torch.Tensor, + ): + """ + Updates the aux_dict with the quantized tensor and additional NVFP4-specific data. + Yields the aux_dict. + """ + aux_dict = { + "nvfp4": quantized_tensor, + "original_tensor": original_tensor, + } + + try: + yield aux_dict + finally: + pass + + @api_method + def inspect_tensor_enabled( + self, config: Dict, layer_name: str, tensor_name: str, iteration: int + ): # pylint: disable=unused-argument + """API call used to determine whether to run inspect_tensor() in the forward.""" + run_current, next_iter = next_enabled_iter( + config.get("start_step", None), + config.get("end_step", None), + config.get("start_end_list", None), + config.get("freq", 1), + iteration, + ) + STATS_BUFFERS.layers_to_next_iter[layer_name] = next_iter + return run_current, next_iter + + @api_method + def inspect_tensor( + self, + config: Dict, + layer_name: str, + tensor_name: str, + iteration: int, + tp_group, + tensor: torch.Tensor, + rowwise_quantized_tensor: Optional[QuantizedTensor] = None, + columnwise_quantized_tensor: Optional[QuantizedTensor] = None, + quantizer: Optional[Quantizer] = None, + ): + """ + API call used to collect the data about the tensor after process_tensor()/quantization. + """ + assert rowwise_quantized_tensor is columnwise_quantized_tensor + assert ( + quantizer is not None + ), "[NVTORCH INSPECT ERROR] LogNvfp4TensorStats cannot be run without NVFP4 quantizer." + + quantized_tensor = rowwise_quantized_tensor + + # Ensure we're working with NVFP4 tensors + if not isinstance(quantizer, NVFP4Quantizer): + raise ValueError( + "[NVTORCH INSPECT ERROR] LogNvfp4TensorStats requires NVFP4Quantizer, " + f"but got {type(quantizer).__name__}" + ) + + assert isinstance(quantized_tensor, NVFP4TensorStorage), ( + "[NVTORCH INSPECT ERROR] LogNvfp4TensorStats quantized_tensor must be a" + " NVFP4TensorStorage." + ) + + for stat in config["stats"]: + self.check_if_stat_is_supported(stat) + + start_step = config.get("start_step", None) + end_step = config.get("end_step", None) + start_end_list = config.get("start_end_list", None) + if start_end_list is not None: + start_end_list = tuple(tuple(int(x) for x in interval) for interval in start_end_list) + + options = ( + start_step, + end_step, + start_end_list, + "nvfp4", + ) + + skip_reduction, reduction_group, reduce_within_microbatch = get_reduction_params( + tensor_name, tp_group + ) + + # Add nvfp4_ prefix to all stats for internal use + prefixed_stats = [self.get_stat_with_prefix(stat) for stat in config["stats"]] + + STATS_BUFFERS.try_add_buffer( + layer_name=layer_name, + tensor_name=tensor_name, + stats=prefixed_stats, + options=options, + reduction_group=reduction_group, + reduce_within_microbatch=reduce_within_microbatch, + ) + + with self.update_aux_dict( + aux_dict={}, + quantized_tensor=quantized_tensor, + quantizer=quantizer, + original_tensor=tensor, + ) as aux_dict: + STATS_BUFFERS.feed( + layer_name, + tensor_name, + options, + tensor, + iteration, + skip_reduction, + aux_dict=aux_dict, + ) + + debug_api.log_message( + f"Feature={self.__class__.__name__}, API=inspect_tensor: {tensor_name}", + layer_name, + extra_cachable_args=(tensor_name,), + ) diff --git a/transformer_engine/debug/features/utils/stats_computation.py b/transformer_engine/debug/features/utils/stats_computation.py index 46a48e2abf..b0002ffee6 100644 --- a/transformer_engine/debug/features/utils/stats_computation.py +++ b/transformer_engine/debug/features/utils/stats_computation.py @@ -443,3 +443,65 @@ def add_max_blockwise_dynamic_range_stats( add_underflows_stats(_recipe_name, _columnwise) add_scale_inv_stats(_recipe_name, _columnwise) add_mse_stats(_recipe_name, _columnwise) + + +# NVFP4-specific statistics + + +def count_nonzero_nvfp4(fp4_data: torch.Tensor) -> torch.Tensor: + """Count the number of non-zero elements in the FP4 data. + + FP4 data is stored as 2 4-bit values per byte (uint8). + We need to unpack and count non-zeros. + """ + # Each byte contains two FP4 values + # Value 0 in FP4 E2M1 format is represented as 0 (and also 8 for -0.0) + zero_vals = torch.tensor([0, 8], device=fp4_data.device, dtype=torch.uint8) + + # Extract first and second nibbles + first_nibble = fp4_data % 16 + second_nibble = fp4_data // 16 + + # Count zeros + first_zeros = torch.isin(first_nibble, zero_vals).sum() + second_zeros = torch.isin(second_nibble, zero_vals).sum() + + total_elements = fp4_data.numel() * 2 + return total_elements - first_zeros - second_zeros + + +def add_nvfp4_underflows_stats(): + """Register underflow stats for NVFP4. + + Computes underflows by counting zeros in packed FP4 data vs original tensor. + """ + stat_num = "nvfp4_underflows_num" + stat_pct = "nvfp4_underflows%" + + stats_to_num[stat_num] = len(stats_to_num) + stats_to_num[stat_pct] = len(stats_to_num) + + # Count non-zeros in original vs FP4 packed data + STATS[stat_num] = ( + lambda x, aux_dict: x.count_nonzero() + - count_nonzero_nvfp4(aux_dict["nvfp4"]._rowwise_data), + lambda buffers, _sn=stat_num: sum(_get(buffers, _sn)), + ) + STATS[stat_pct] = ( + lambda x, aux_dict: ( + x.count_nonzero() - count_nonzero_nvfp4(aux_dict["nvfp4"]._rowwise_data) + ) + / aux_dict["nvfp4"].numel() + * 100, + lambda buffers, _sn_num=stat_num: 100 + * sum(_get(buffers, _sn_num)) + / sum(_get(buffers, "numel")), + ) + + DEPENDENCIES[stat_num] = {stat_num} + DEPENDENCIES[stat_pct] = {stat_num, "numel"} + + +# Register NVFP4 stats +add_nvfp4_underflows_stats() +add_mse_stats("nvfp4") # Reuse existing MSE function diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 29a108c75f..455079143b 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -36,7 +36,7 @@ } API_CALL_MODIFY = "modify_tensor()" -STANDARD_FP8_QUANTIZE = "FP8 Quantize" +STANDARD_QUANTIZE = "Quantize" HIGH_PRECISION = "High Precision" @@ -88,7 +88,7 @@ def __init__( # inspect_tensor*_enabled are bool fields, # indicating whether some feature will need to run inspect_tensor_* calls. # - # *_tensor_plan are one of [API_CALL_MODIFY, STANDARD_FP8_QUANTIZE, HIGH_PRECISION] + # *_tensor_plan are one of [API_CALL_MODIFY, STANDARD_QUANTIZE, HIGH_PRECISION] # determining what will happen when the quantizer is used for that tensor. self.output_tensor = tensor_name in ["output", "wgrad", "dgrad"] if self.output_tensor: @@ -170,7 +170,7 @@ def get_enabled_look_at_tensors(self): def get_tensors_plan(self): """ Returns (rowwise_plan, columnwise_plan). Each element of the tuple is one of - API_CALL_MODIFY, STANDARD_FP8_QUANTIZE, or HIGH_PRECISION, indicating the behavior + API_CALL_MODIFY, STANDARD_QUANTIZE, or HIGH_PRECISION, indicating the behavior of this quantizer with respect to these tensors. """ import nvdlfw_inspect.api as debug_api @@ -191,16 +191,16 @@ def get_tensors_plan(self): rowwise_plan = API_CALL_MODIFY else: if self.parent_quantizer is not None: - fp8_quantize = self.process_enabled_api_call( - debug_api.transformer_engine.fp8_gemm_enabled( + quantize_enabled = self.process_enabled_api_call( + debug_api.transformer_engine.fp8_gemm_enabled( # API name kept for compatibility layer_name=self.layer_name, gemm=self.rowwise_gemm_name, iteration=self.iteration, ) ) - if fp8_quantize: - rowwise_plan = STANDARD_FP8_QUANTIZE + if quantize_enabled: + rowwise_plan = STANDARD_QUANTIZE if rowwise_plan is None: rowwise_plan = HIGH_PRECISION @@ -218,16 +218,16 @@ def get_tensors_plan(self): columnwise_plan = API_CALL_MODIFY else: if self.parent_quantizer is not None: - fp8_quantize = self.process_enabled_api_call( - debug_api.transformer_engine.fp8_gemm_enabled( + quantize_enabled = self.process_enabled_api_call( + debug_api.transformer_engine.fp8_gemm_enabled( # API name kept for compatibility layer_name=self.layer_name, gemm=self.columnwise_gemm_name, iteration=self.iteration, ) ) - if fp8_quantize: - columnwise_plan = STANDARD_FP8_QUANTIZE + if quantize_enabled: + columnwise_plan = STANDARD_QUANTIZE if columnwise_plan is None: columnwise_plan = HIGH_PRECISION @@ -278,7 +278,7 @@ def _call_inspect_tensor_api( del args["quantizer"] if ( - self.rowwise_tensor_plan in [API_CALL_MODIFY, STANDARD_FP8_QUANTIZE] + self.rowwise_tensor_plan in [API_CALL_MODIFY, STANDARD_QUANTIZE] and self.inspect_tensor_postquantize_enabled_rowwise ): args["tensor"] = rowwise_gemm_tensor @@ -286,7 +286,7 @@ def _call_inspect_tensor_api( debug_api.transformer_engine.inspect_tensor_postquantize(**args) if ( - self.columnwise_tensor_plan in [API_CALL_MODIFY, STANDARD_FP8_QUANTIZE] + self.columnwise_tensor_plan in [API_CALL_MODIFY, STANDARD_QUANTIZE] and self.inspect_tensor_postquantize_enabled_columnwise ): args["tensor"] = columnwise_gemm_tensor @@ -317,14 +317,14 @@ def quantize( self.parent_quantizer.set_usage(rowwise=True) rowwise_gemm_tensor, columnwise_gemm_tensor = None, None - if STANDARD_FP8_QUANTIZE in [self.rowwise_tensor_plan, self.columnwise_tensor_plan]: + if STANDARD_QUANTIZE in [self.rowwise_tensor_plan, self.columnwise_tensor_plan]: quantized_tensor = self.parent_quantizer(tensor) - # if both rowwise_tensor_plan and columnwise_tensor_plan need to be in fp8, + # if both rowwise_tensor_plan and columnwise_tensor_plan need to be quantized, # one tensor with columnwise=True and rowwise=True is computed # and both rowwise_tensor_plan and columnwise_tensor_plan point to it. - if self.rowwise_tensor_plan == STANDARD_FP8_QUANTIZE: + if self.rowwise_tensor_plan == STANDARD_QUANTIZE: rowwise_gemm_tensor = quantized_tensor - if self.columnwise_tensor_plan == STANDARD_FP8_QUANTIZE: + if self.columnwise_tensor_plan == STANDARD_QUANTIZE: columnwise_gemm_tensor = quantized_tensor # 2. modify_tensor() is called, if it is used. @@ -379,7 +379,7 @@ def process_gemm_output(self, tensor: torch.Tensor): """This call is invoked after the gemm to inspect and modify the output tensor.""" import nvdlfw_inspect.api as debug_api - assert self.parent_quantizer is None, "FP8 output is not supported for debug=True." + assert self.parent_quantizer is None, "Quantized output is not supported for debug=True." assert self.output_tensor tensor_to_gemm = {"output": "fprop", "wgrad": "wgrad", "dgrad": "dgrad"} if self.rowwise_tensor_plan == API_CALL_MODIFY: @@ -420,9 +420,9 @@ def any_feature_enabled(self) -> bool: ): return True if self.parent_quantizer is not None: - if self.rowwise_tensor_plan != STANDARD_FP8_QUANTIZE: + if self.rowwise_tensor_plan != STANDARD_QUANTIZE: return True - if self.columnwise_tensor_plan != STANDARD_FP8_QUANTIZE: + if self.columnwise_tensor_plan != STANDARD_QUANTIZE: return True return False @@ -446,7 +446,7 @@ def update_quantized( if self.parent_quantizer is not None: if ( dst.rowwise_gemm_tensor is not None - and self.rowwise_tensor_plan == STANDARD_FP8_QUANTIZE + and self.rowwise_tensor_plan == STANDARD_QUANTIZE ): if hasattr(dst.rowwise_gemm_tensor, "quantize_"): dst.rowwise_gemm_tensor.quantize_(src, noop_flag=None) @@ -455,7 +455,7 @@ def update_quantized( updated_rowwise_gemm = True if ( dst.columnwise_gemm_tensor is not None - and self.columnwise_tensor_plan == STANDARD_FP8_QUANTIZE + and self.columnwise_tensor_plan == STANDARD_QUANTIZE and not updated_rowwise_gemm ): if hasattr(dst.columnwise_gemm_tensor, "quantize_"): @@ -540,14 +540,12 @@ def _update_parent_quantizer_usage(self): """ Updates the usage of the parent quantizer. """ - rowwise_gemm_quantize = ( - self.rowwise_usage and self.rowwise_tensor_plan == STANDARD_FP8_QUANTIZE - ) + rowwise_gemm_quantize = self.rowwise_usage and self.rowwise_tensor_plan == STANDARD_QUANTIZE columnwise_gemm_quantize = ( - self.columnwise_usage and self.columnwise_tensor_plan == STANDARD_FP8_QUANTIZE + self.columnwise_usage and self.columnwise_tensor_plan == STANDARD_QUANTIZE ) - if STANDARD_FP8_QUANTIZE in [self.rowwise_tensor_plan, self.columnwise_tensor_plan]: + if STANDARD_QUANTIZE in [self.rowwise_tensor_plan, self.columnwise_tensor_plan]: self.parent_quantizer.set_usage( rowwise=rowwise_gemm_quantize, columnwise=columnwise_gemm_quantize, From 59f6f3876767d07045152bfae07b5dd4c54e1725 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 3 Feb 2026 18:17:10 -0800 Subject: [PATCH 192/521] [JAX] Update JAX container in readme (#2648) * Update README.rst Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update README.rst Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update README.rst Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --------- Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- README.rst | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/README.rst b/README.rst index 55be0e583f..cd55a2e18f 100644 --- a/README.rst +++ b/README.rst @@ -175,15 +175,22 @@ For example to use the NGC PyTorch container interactively, .. code-block:: bash - docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:25.08-py3 + docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:26.01-py3 For example to use the NGC JAX container interactively, .. code-block:: bash - docker run --gpus all -it --rm nvcr.io/nvidia/jax:25.08-py3 + docker run --gpus all -it --rm nvcr.io/nvidia/jax:26.01-py3 -Where 25.08 (corresponding to August 2025 release) is the container version. +Where 26.01 (corresponding to January 2026 release) is the container version. + +We recommend updating to the latest NGC container available here: + +* https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch +* https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax + +If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. If you would like to use examples from TE main branch and are running into import errors, please try the latest pip package or building from source, although NGC containers are recommended for ease-of-use for most users. **Benefits of using NGC containers:** From 71971e33dc01a015c943c8c3b73800eb047b8a78 Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Thu, 5 Feb 2026 16:59:30 -0800 Subject: [PATCH 193/521] Fix exp2f_rcp to properly handle nan and 0xFE cases (#2647) Signed-off-by: Kaining Zhong --- tests/cpp/test_common.h | 10 +++++++--- transformer_engine/common/util/ptx.cuh | 10 +++++++--- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 082677c978..5bb6400629 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -425,10 +425,14 @@ inline fp8e8m0 float_to_e8m0(float val) { } inline float exp2f_rcp(fp8e8m0 biased_exp) { - if (biased_exp == 0) { - return 1.0f; + int32_t int_val = 0; + if (biased_exp == 255) { + int_val = 0x7fffffff; + } else if (biased_exp == 254) { + int_val = 0x00400000; + } else { + int_val = (254 - biased_exp) << FP32_MANTISSA_BITS; // 127 - (biased_exp - 127) } - int32_t int_val = (254 - biased_exp) << FP32_MANTISSA_BITS; // 127 - (biased_exp - 127) float fp32_val = *reinterpret_cast(&int_val); return fp32_val; } diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 9bcf6e2289..5367d7e781 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -328,9 +328,13 @@ constexpr uint32_t FP32_MANTISSA_BITS = 23; constexpr uint32_t FP32_EXPONENT_BIAS = 127; __device__ __forceinline__ float exp2f_rcp(e8m0_t biased_exp) { - return (biased_exp == 0) ? 1 - : __int_as_float((254 - biased_exp) - << FP32_MANTISSA_BITS); // 127 - (biased_exp - 127) + // Handle the special case of NaN. + if (biased_exp == 255) return __int_as_float(0x7fffffff); + // Handle the special case where the unbiased exponent is 127, so the reciprocal is 2^-127 which needs the first bit of + // the mantissa to be 1, which can't be obtained by shifting `FP32_MANTISSA_BITS` bits to the left. + if (biased_exp == 254) return __int_as_float(0x00400000); + // Fast calculation when the unbiased exp is in [-126, 126], and only the exponent part is used to express the reciprocal. + return __int_as_float((254 - biased_exp) << FP32_MANTISSA_BITS); } __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { From 739394721bda555e11637598b2e5aa109b94afbf Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Fri, 6 Feb 2026 20:46:04 +0100 Subject: [PATCH 194/521] [Common] MXFP8 kernel for grouped tensors (#2586) * Rebased to main Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed the year to 2026 Signed-off-by: Oleg Goncharov * Added compilation guards Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added BWD pass Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added dbias and dact tests. Refactoring. Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added grouped MXFP8 DACT and ACT API and tests Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed a typo Signed-off-by: Oleg Goncharov * Fixes per the review Signed-off-by: Oleg Goncharov * More fixes from the review Signed-off-by: Oleg Goncharov * Fixes per the review Signed-off-by: Oleg Goncharov * Relaxed requirement for last dim from mod128 to mod32 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 * Added alignment checks when tensor descriptors are modified 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: vthumbe1503 --- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_cast_mxfp8_grouped.cu | 820 +++++++++++++++ transformer_engine/common/activation/gelu.cu | 73 ++ transformer_engine/common/activation/relu.cu | 73 ++ .../common/activation/swiglu.cu | 36 + transformer_engine/common/cast/cast.cu | 22 + .../common/cast/dispatch/quantize.cuh | 84 ++ .../cast/mxfp8/group_quantize_mxfp8.cuh | 962 ++++++++++++++++++ .../include/transformer_engine/activation.h | 110 ++ .../common/include/transformer_engine/cast.h | 146 +++ 10 files changed, 2327 insertions(+) create mode 100644 tests/cpp/operator/test_cast_mxfp8_grouped.cu create mode 100644 transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 08a683949b..56880a428d 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -11,6 +11,7 @@ add_executable(test_operator test_cast_mxfp8_gated_swiglu.cu test_qdq.cu test_cast_mxfp8.cu + test_cast_mxfp8_grouped.cu test_cast_nvfp4_transpose.cu test_cast_float8blockwise.cu test_dequantize_mxfp8.cu diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu new file mode 100644 index 0000000000..8b084ca452 --- /dev/null +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -0,0 +1,820 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include +#include +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum ProcessingMethod { + CAST_ONLY, + CAST_DBIAS, + CAST_DBIAS_DACT, + CAST_DACT, + CAST_ACT +}; + +enum ActivationKind { + Identity, + GeLU, + SiLU, + ReLU, + QGeLU, + SReLU +}; + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + +template +void compute_ref(const ProcessingMethod processing_method, + float (*OP)(const float), + const bool rowwise, + const bool colwise, + const InputType* input, + const InputType* grad, + OutputType* output_rowwise, + OutputType* output_colwise, + fp8e8m0* output_scales_rowwise, + fp8e8m0* output_scales_colwise, + InputType* output_dbias, + const size_t rows, + const size_t cols, + const size_t scales_stride_rowwise, + const size_t scales_stride_colwise, + const bool is_single_tensor) +{ + const size_t tile_size_Y = 32; + const size_t tile_size_X = 32; + const size_t tiles_num_Y = (rows + tile_size_Y - 1) / tile_size_Y; + const size_t tiles_num_X = (cols + tile_size_X - 1) / tile_size_X; + + std::vector output_dbias_fp32(cols, 0); + #pragma omp parallel proc_bind(spread) + { + // Buffers to cache intermediate computations + std::vector cache_buffer(tile_size_Y * tile_size_X); + + std::vector thread_dbias(cols, 0); + #pragma omp for schedule(static) + for (size_t t = 0; t < tiles_num_Y * tiles_num_X; ++t) { + const size_t tile_Y = t / tiles_num_X; + const size_t tile_X = t % tiles_num_X; + const size_t tile_offset_Y = tile_Y * tile_size_Y; + const size_t tile_offset_X = tile_X * tile_size_X; + + const size_t i_min = tile_offset_Y; + const size_t i_max = std::min(i_min + tile_size_Y, rows); + + const size_t j_min = tile_offset_X; + const size_t j_max = std::min(j_min + tile_size_X, cols); + + // Cache computations + for (size_t i = i_min; i < i_max; ++i) { + for (size_t j = j_min; j < j_max; ++j) { + + const size_t idx = i * cols + j; + const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); + + float elt = static_cast(input[idx]); + if (processing_method == ProcessingMethod::CAST_DBIAS) { + // grad is the input + elt = static_cast(grad[idx]); + } + if (processing_method != ProcessingMethod::CAST_ONLY + && processing_method != ProcessingMethod::CAST_DBIAS) { + elt = OP(elt); + } + if (processing_method == ProcessingMethod::CAST_DACT || + processing_method == ProcessingMethod::CAST_DBIAS_DACT) { + elt *= static_cast(grad[idx]); + } + thread_dbias[j] += elt; + + // Numerical truncation: after downcast to InputType (BF16/FP16), upcast it back to FP32 + elt = static_cast(static_cast(elt)); + + cache_buffer[cache_idx] = elt; + if (isinf(elt) || isnan(elt)) { + continue; + } + } + } + + if (rowwise) { + for (size_t i = i_min; i < i_max; ++i) { + float block_amax = 0.0f; + + for (size_t j = j_min; j < j_max; ++j) { + const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); + block_amax = std::max(block_amax, std::abs(cache_buffer[cache_idx])); + } + + const fp8e8m0 biased_exponent = float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); + const size_t scale_idx = i * scales_stride_rowwise + tile_X; + output_scales_rowwise[scale_idx] = biased_exponent; + const float scale_reciprocal = exp2f_rcp(biased_exponent); + + for (size_t j = j_min; j < j_max; ++j) { + const size_t idx = i * cols + j; + const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); + output_rowwise[idx] = static_cast(cache_buffer[cache_idx] * scale_reciprocal); + } + } + } + if (colwise) { + for (size_t j = j_min; j < j_max; ++j) { + float block_amax = 0.0f; + + for (size_t i = i_min; i < i_max; ++i) { + const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); + block_amax = std::max(block_amax, std::abs(cache_buffer[cache_idx])); + } + + const fp8e8m0 biased_exponent = float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); + const size_t scale_idx = tile_Y * scales_stride_colwise + j; + output_scales_colwise[scale_idx] = biased_exponent; + const float scale_reciprocal = exp2f_rcp(biased_exponent); + + for (size_t i = i_min; i < i_max; ++i) { + const size_t idx = i * cols + j; + const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); + output_colwise[idx] = static_cast(cache_buffer[cache_idx] * scale_reciprocal); + } + } + } + } + #pragma omp critical + { + for (size_t j = 0; j < cols; ++j) { + output_dbias_fp32[j] += thread_dbias[j]; + } + } + } + + if (is_single_tensor) { + for (size_t j = 0; j < cols; ++j) { + output_dbias[j] = static_cast(output_dbias_fp32[j]); + } + } +} + +template +void compare_scaled_elts(const std::string &name, + const T* ref_data, + const T* test_data, + const size_t rows, + const size_t cols, + const bool rowwise, + const size_t tolerable_mismatches_limit = 0, + const double atol = 1e-5, + const double rtol = 1e-8) { + size_t mismatches_num = 0; + int first_mismatch_idx = -1; + + for (size_t i = 0; i < rows * cols; ++i) { + double t = static_cast(test_data[i]); + double r = static_cast(ref_data[i]); + bool mismatch = fabs(t - r) > atol && (r == 0 || fabs((t - r) / r) > rtol); + /* For Float32 the floating point comparison is enough to error out */ + bool assertion = false; + if (mismatch && !assertion) { + /* Check if it is just a failure of round to nearest choosing different + side of the real value */ + const double mean = (t + r) / 2; + const double mean_p = mean >= 0 ? mean * (1 + 1e-6) : mean * (1 - 1e-6); + const double mean_m = mean >= 0 ? mean * (1 - 1e-6) : mean * (1 + 1e-6); + const double cast_mean_p = static_cast(static_cast(mean_p)); + const double cast_mean_m = static_cast(static_cast(mean_m)); + assertion = !(cast_mean_m == std::min(t,r) && cast_mean_p == std::max(t,r)); + } + std::string direction = rowwise ? "rowwise" : "columnwise"; + if (assertion) { + mismatches_num++; + if (first_mismatch_idx == -1) { + first_mismatch_idx = i; + } + } + if (mismatches_num > tolerable_mismatches_limit) { + const double first_mismatch_t = static_cast(test_data[first_mismatch_idx]); + const double first_mismatch_r = static_cast(ref_data[first_mismatch_idx]); + + GTEST_FAIL() << mismatches_num << " mismatche(s) which is more than tolerable mismatch limit of " + << tolerable_mismatches_limit << "." << std::endl + << "Error in tensor " << name << " in " + << direction << " direction." << std::endl + << "First mismatch at place " << first_mismatch_idx + << " (" << std::to_string(first_mismatch_idx) << "): " + << first_mismatch_t << " vs " << first_mismatch_r; + } + } +} + +/** + * Scaling along single dimension (either rows or columns) + * Produces one set of output data and the corresponding data of the fused operation (dbias): + * 1) Scaled rows + row-wise scaling factors + * OR + * 2) Scaled columns + column-wise scaling factors + */ +template +void performTest(const ProcessingMethod processing_method, + float (*OP)(const float), + const ShapeRepresentation shape_rep, + const size_t num_tensors, + const std::vector& logical_shape_vec, + const std::vector& first_dims_h, + const std::vector& last_dims_h, + const std::vector& offsets_h, + const bool rowwise, + const bool colwise) { + using namespace test; + + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + const size_t rows = logical_shape_vec[0]; + const size_t cols = logical_shape_vec[1]; + + size_t elts_num = 0; + size_t rowwise_sfs_num = 0; + size_t colwise_sfs_num = 0; + + std::vector rowwise_scales_first_dim(num_tensors, 0); + std::vector rowwise_scales_last_dim(num_tensors, 0); + std::vector rowwise_scales_offset(num_tensors + 1, 0); + std::vector colwise_scales_first_dim(num_tensors, 0); + std::vector colwise_scales_last_dim(num_tensors, 0); + std::vector colwise_scales_offset(num_tensors + 1, 0); + + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; + const size_t elts = M * K; + elts_num += elts; + + const size_t unpadded_rowwise_blocks_Y = M; + const size_t unpadded_rowwise_blocks_X = divide_round_up(K, 32); + const size_t unpadded_colwise_blocks_Y = divide_round_up(M, 32); + const size_t unpadded_colwise_blocks_X = K; + + rowwise_scales_first_dim[t] = round_up_to_nearest_multiple(unpadded_rowwise_blocks_Y, 128); + rowwise_scales_last_dim[t] = round_up_to_nearest_multiple(unpadded_rowwise_blocks_X, 4); + colwise_scales_first_dim[t] = round_up_to_nearest_multiple(unpadded_colwise_blocks_Y, 4); + colwise_scales_last_dim[t] = round_up_to_nearest_multiple(unpadded_colwise_blocks_X, 128); + + const size_t rowwise_sfs = rowwise_scales_first_dim[t] * rowwise_scales_last_dim[t]; + const size_t colwise_sfs = colwise_scales_first_dim[t] * colwise_scales_last_dim[t]; + + rowwise_sfs_num += rowwise_sfs; + colwise_sfs_num += colwise_sfs; + + rowwise_scales_offset[t+1] = rowwise_sfs_num; + colwise_scales_offset[t+1] = colwise_sfs_num; + } + + const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS) || (shape_rep == VARYING_FIRST_DIM); + + std::vector scales_rowwise_shape = {rowwise_sfs_num}; + std::vector scales_colwise_shape = {colwise_sfs_num}; + + std::mt19937 gen; + std::uniform_real_distribution<> dis(-2.0, 1.0); + + std::vector in_data(elts_num); + std::vector grad_data(elts_num); + + std::vector out_data_rowwise_h(rowwise ? elts_num : 0); + std::vector out_data_colwise_h(colwise ? elts_num : 0); + std::vector out_scales_rowwise_h(rowwise ? rowwise_sfs_num : 0); + std::vector out_scales_colwise_h(colwise ? colwise_sfs_num : 0); + + std::vector out_data_rowwise_ref(rowwise ? elts_num : 0); + std::vector out_data_colwise_ref(colwise ? elts_num : 0); + std::vector out_scales_rowwise_ref(rowwise ? rowwise_sfs_num : 0); + std::vector out_scales_colwise_ref(colwise ? colwise_sfs_num : 0); + + std::vector ref_output_dbias(is_single_tensor ? cols : 0); + + for (size_t i = 0; i < elts_num; ++i) { + const float val = dis(gen); + grad_data[i] = static_cast(val); + in_data[i] = static_cast(val); + } + + const OutputType zero_elt = static_cast(0.0f); + const fp8e8m0 zero_SF = static_cast(0.0f); + if (rowwise) { + std::fill(out_data_rowwise_h.begin(), out_data_rowwise_h.end(), zero_elt); + std::fill(out_data_rowwise_ref.begin(), out_data_rowwise_ref.end(), zero_elt); + std::fill(out_scales_rowwise_h.begin(), out_scales_rowwise_h.end(), zero_SF); + std::fill(out_scales_rowwise_ref.begin(), out_scales_rowwise_ref.end(), zero_SF); + } + if (colwise) { + std::fill(out_data_colwise_h.begin(), out_data_colwise_h.end(), zero_elt); + std::fill(out_data_colwise_ref.begin(), out_data_colwise_ref.end(), zero_elt); + std::fill(out_scales_colwise_h.begin(), out_scales_colwise_h.end(), zero_SF); + std::fill(out_scales_colwise_ref.begin(), out_scales_colwise_ref.end(), zero_SF); + } + + const size_t in_data_size = elts_num * sizeof(InputType); + const size_t out_data_size = elts_num * sizeof(OutputType); + const size_t rowwise_scales_size = rowwise_sfs_num * sizeof(fp8e8m0); + const size_t colwise_scales_size = colwise_sfs_num * sizeof(fp8e8m0); + + const size_t first_dims_size = num_tensors * sizeof(size_t); + const size_t last_dims_size = num_tensors * sizeof(size_t); + const size_t offsets_size = (num_tensors + 1) * sizeof(size_t); + + InputType* grad_data_d; + InputType* in_data_d; + OutputType* out_data_rowwise_d; + OutputType* out_data_colwise_d; + fp8e8m0* out_scales_rowwise_d; + fp8e8m0* out_scales_colwise_d; + size_t* first_dims_d; + size_t* last_dims_d; + size_t* offsets_d; + + cudaMalloc((void**)&grad_data_d, in_data_size); + cudaMalloc((void**)&in_data_d, in_data_size); + cudaMalloc((void**)&first_dims_d, first_dims_size); + cudaMalloc((void**)&last_dims_d, last_dims_size); + cudaMalloc((void**)&offsets_d, offsets_size); + + cudaMemcpy(grad_data_d, grad_data.data(), in_data_size, cudaMemcpyHostToDevice); + cudaMemcpy(in_data_d, in_data.data(), in_data_size, cudaMemcpyHostToDevice); + cudaMemcpy(first_dims_d, first_dims_h.data(), first_dims_size, cudaMemcpyHostToDevice); + cudaMemcpy(last_dims_d, last_dims_h.data(), last_dims_size, cudaMemcpyHostToDevice); + cudaMemcpy(offsets_d, offsets_h.data(), offsets_size, cudaMemcpyHostToDevice); + + NVTEShape logical_shape_ = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); + + NVTEShape first_dims_shape_; + NVTEShape last_dims_shape_; + NVTEShape offsets_shape_; + + first_dims_shape_.ndim = 1; + last_dims_shape_.ndim = 1; + offsets_shape_.ndim = 1; + + first_dims_shape_.data[0] = num_tensors; + last_dims_shape_.data[0] = num_tensors; + offsets_shape_.data[0] = num_tensors + 1; + + NVTEGroupedTensor grad_group_tensor = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape_); + NVTEGroupedTensor in_group_tensor = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape_); + NVTEGroupedTensor out_group_tensor = nvte_create_grouped_tensor(NVTE_MXFP8_1D_SCALING, num_tensors, logical_shape_); + + NVTEBasicTensor grad_data_tensor = {grad_data_d, static_cast(itype), logical_shape_}; + NVTEBasicTensor in_data_tensor = {in_data_d, static_cast(itype), logical_shape_}; + nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &in_data_tensor); + nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &grad_data_tensor); + + if ((shape_rep == VARYING_FIRST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { + NVTEBasicTensor first_dims_tensor = {first_dims_d, kNVTEInt64, first_dims_shape_}; + nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, &first_dims_tensor); + nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, &first_dims_tensor); + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, &first_dims_tensor); + } + + if ((shape_rep == VARYING_LAST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { + NVTEBasicTensor last_dims_tensor = {last_dims_d, kNVTEInt64, last_dims_shape_}; + nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor); + nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor); + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor); + } + + if (shape_rep != SAME_BOTH_DIMS) { + NVTEBasicTensor offsets_tensor = {offsets_d, kNVTEInt64, offsets_shape_}; + nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, &offsets_tensor); + nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, &offsets_tensor); + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, &offsets_tensor); + } + + if (rowwise) { + cudaMalloc((void**)&out_data_rowwise_d, out_data_size); + cudaMalloc((void**)&out_scales_rowwise_d, rowwise_scales_size); + cudaMemset(out_data_rowwise_d, 0, out_data_size); + cudaMemset(out_scales_rowwise_d, 0, rowwise_scales_size); + NVTEBasicTensor out_data_rowwise_tensor = {out_data_rowwise_d, static_cast(otype), logical_shape_}; + NVTEShape scales_rowwise_shape_ = nvte_make_shape(scales_rowwise_shape.data(), scales_rowwise_shape.size()); + NVTEBasicTensor out_scales_rowwise_tensor = {out_scales_rowwise_d, NVTEDType::kNVTEFloat8E8M0, scales_rowwise_shape_}; + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &out_data_rowwise_tensor); + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseScaleInv, &out_scales_rowwise_tensor); + } + + if (colwise) { + cudaMalloc((void**)&out_data_colwise_d, out_data_size); + cudaMalloc((void**)&out_scales_colwise_d, colwise_scales_size); + cudaMemset(out_data_colwise_d, 0, out_data_size); + cudaMemset(out_scales_colwise_d, 0, colwise_scales_size); + NVTEBasicTensor out_data_colwise_tensor = {out_data_colwise_d, static_cast(otype), logical_shape_}; + NVTEShape scales_colwise_shape_ = nvte_make_shape(scales_colwise_shape.data(), scales_colwise_shape.size()); + NVTEBasicTensor out_scales_colwise_tensor = {out_scales_colwise_d, NVTEDType::kNVTEFloat8E8M0, scales_colwise_shape_}; + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedColumnwiseData, &out_data_colwise_tensor); + nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedColumnwiseScaleInv, &out_scales_colwise_tensor); + } + + Tensor output_dbias("output_dbias", std::vector{ cols }, itype); + + // Reference (CPU) + if (is_single_tensor) { + + const size_t unpadded_rowwise_blocks_X = divide_round_up(cols, 32); + const size_t unpadded_colwise_blocks_X = cols; + + const size_t scales_stride_rowwise = round_up_to_nearest_multiple(unpadded_rowwise_blocks_X, 4); + const size_t scales_stride_colwise = round_up_to_nearest_multiple(unpadded_colwise_blocks_X, 128); + + compute_ref( + processing_method, OP, rowwise, colwise, in_data.data(), grad_data.data(), + out_data_rowwise_ref.data(), out_data_colwise_ref.data(), + out_scales_rowwise_ref.data(), out_scales_colwise_ref.data(), + ref_output_dbias.data(), rows, cols, + scales_stride_rowwise, + scales_stride_colwise, + is_single_tensor); + } else { + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; + + const size_t scales_stride_rowwise = rowwise_scales_last_dim[t]; + const size_t scales_stride_colwise = colwise_scales_last_dim[t]; + const size_t data_offset = offsets_h[t]; + const size_t rowwise_sfs_offset = rowwise_scales_offset[t]; + const size_t colwise_sfs_offset = colwise_scales_offset[t]; + + const InputType* const grad_ptr = grad_data.data() + data_offset; + const InputType* const in_ptr = in_data.data() + data_offset; + OutputType* const out_data_rowwise_ptr = out_data_rowwise_ref.data() + data_offset; + OutputType* const out_data_colwise_ptr = out_data_colwise_ref.data() + data_offset; + fp8e8m0* const out_scales_rowwise_ptr = out_scales_rowwise_ref.data() + rowwise_sfs_offset; + fp8e8m0* const out_scales_colwise_ptr = out_scales_colwise_ref.data() + colwise_sfs_offset; + + compute_ref( + processing_method, OP, rowwise, colwise, in_ptr, grad_ptr, + out_data_rowwise_ptr, out_data_colwise_ptr, + out_scales_rowwise_ptr, out_scales_colwise_ptr, + ref_output_dbias.data(), M, K, + scales_stride_rowwise, + scales_stride_colwise, + is_single_tensor); + } + } + + // GPU + Tensor workspace; + switch (processing_method) { + case ProcessingMethod::CAST_ONLY: { + nvte_group_quantize(in_group_tensor, out_group_tensor, 0); + break; + } + case ProcessingMethod::CAST_DBIAS: { + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias.data(), workspace.data(), 0); + workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias.data(), workspace.data(), 0); + break; + } + case ProcessingMethod::CAST_DBIAS_DACT: { + auto nvte_group_quantize_dbias_dact = &nvte_group_quantize_dbias_dgelu; + if (OP == &dsilu) { nvte_group_quantize_dbias_dact = &nvte_group_quantize_dbias_dsilu; } + else if (OP == &drelu) { nvte_group_quantize_dbias_dact = &nvte_group_quantize_dbias_drelu; } + else if (OP == &dqgelu) { nvte_group_quantize_dbias_dact = &nvte_group_quantize_dbias_dqgelu; } + else if (OP == &dsrelu) { nvte_group_quantize_dbias_dact = &nvte_group_quantize_dbias_dsrelu; } + + nvte_group_quantize_dbias_dact(grad_group_tensor, in_group_tensor, out_group_tensor, + output_dbias.data(), workspace.data(), 0); + workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); + nvte_group_quantize_dbias_dact(grad_group_tensor, in_group_tensor, out_group_tensor, + output_dbias.data(), workspace.data(), 0); + break; + } + case ProcessingMethod::CAST_ACT: { + auto nvte_group_act = &nvte_group_gelu; + if (OP == &silu) { nvte_group_act = &nvte_group_silu; } + else if (OP == &relu) { nvte_group_act = &nvte_group_relu; } + else if (OP == &qgelu) { nvte_group_act = &nvte_group_qgelu; } + else if (OP == &srelu) { nvte_group_act = &nvte_group_srelu; } + nvte_group_act(in_group_tensor, out_group_tensor, 0); + break; + } + case ProcessingMethod::CAST_DACT: { + auto nvte_group_dact = &nvte_group_dgelu; + if (OP == &dsilu) { nvte_group_dact = &nvte_group_dsilu; } + else if (OP == &drelu) { nvte_group_dact = &nvte_group_drelu; } + else if (OP == &dqgelu) { nvte_group_dact = &nvte_group_dqgelu; } + else if (OP == &dsrelu) { nvte_group_dact = &nvte_group_dsrelu; } + nvte_group_dact(grad_group_tensor, in_group_tensor, out_group_tensor, 0); + break; + } + } + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + auto [atol, rtol] = getTolerances(otype); + const size_t scale_diff_abs_tolerance = 0; + const double abs_tolerable_mismatches_limit = 0.0; + const double rel_tolerable_mismatches_limit = 0.0; + + if (rowwise) { + cudaMemcpy(out_data_rowwise_h.data(), out_data_rowwise_d, out_data_size, cudaMemcpyDeviceToHost); + cudaMemcpy(out_scales_rowwise_h.data(), out_scales_rowwise_d, rowwise_scales_size, cudaMemcpyDeviceToHost); + + size_t mismatches_scales = 0; + compare_scaling_factors("rowwise_scales", out_scales_rowwise_h.data(), out_scales_rowwise_ref.data(), + 1, rowwise_sfs_num, rowwise_sfs_num, mismatches_scales, scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, rel_tolerable_mismatches_limit); + + const size_t mismatches_elts = 32 * mismatches_scales; + + compare_scaled_elts("rowwise_output", out_data_rowwise_ref.data(), + out_data_rowwise_h.data(), rows, cols, true, mismatches_elts); + } + + if (colwise) { + cudaMemcpy(out_data_colwise_h.data(), out_data_colwise_d, out_data_size, cudaMemcpyDeviceToHost); + cudaMemcpy(out_scales_colwise_h.data(), out_scales_colwise_d, colwise_scales_size, cudaMemcpyDeviceToHost); + + size_t mismatches_scales = 0; + compare_scaling_factors("colwise_scales", out_scales_colwise_h.data(), out_scales_colwise_ref.data(), + 1, colwise_sfs_num, colwise_sfs_num, mismatches_scales, scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, rel_tolerable_mismatches_limit); + + const size_t mismatches_elts = 32 * mismatches_scales; + + compare_scaled_elts("colwise_output", out_data_colwise_ref.data(), + out_data_colwise_h.data(), rows, cols, false, mismatches_elts); + } + + if (processing_method == ProcessingMethod::CAST_DBIAS + || processing_method == ProcessingMethod::CAST_DBIAS_DACT) + { + auto [atol_dbias, rtol_dbias] = getTolerances(itype); + if (itype == DType::kFloat32) { + atol_dbias = 1e-4; + rtol_dbias *= sqrt(static_cast(rows)) ; + } else { + rtol_dbias *= 4; + } + compareResults("output_dbias", output_dbias, ref_output_dbias.data(), true, atol_dbias, rtol_dbias); + } + + cudaFree(grad_data_d); + cudaFree(in_data_d); + cudaFree(first_dims_d); + cudaFree(last_dims_d); + cudaFree(offsets_d); + if (rowwise) { + cudaFree(out_data_rowwise_d); + cudaFree(out_scales_rowwise_d); + } + if (colwise) { + cudaFree(out_data_colwise_d); + cudaFree(out_scales_colwise_d); + } +} + +std::vector processing_methods = { + ProcessingMethod::CAST_ONLY, + ProcessingMethod::CAST_DBIAS, + ProcessingMethod::CAST_DBIAS_DACT, + ProcessingMethod::CAST_DACT, + ProcessingMethod::CAST_ACT, +}; + +std::vector activation_kinds = { + ActivationKind::Identity, + ActivationKind::GeLU, + // ActivationKind::SiLU, + // ActivationKind::ReLU, + // ActivationKind::QGeLU, + // ActivationKind::SReLU, +}; + +enum ScalingDirection { + ROWWISE = 0, + COLWISE = 1, + BOTH = 2 +}; + +std::vector scaling_directions = { + ScalingDirection::ROWWISE, + ScalingDirection::COLWISE, + ScalingDirection::BOTH, +}; + +// {shape_representation, num_tensors, [logical_shape_M, logical_shape_K], [M_i], [K_i]} +std::vector> input_config = { + {SAME_BOTH_DIMS, 1, 128,128}, + {SAME_BOTH_DIMS, 2, 256,128}, + {VARYING_FIRST_DIM, 2, 512,128, 128,384}, + {VARYING_FIRST_DIM, 2, 384,160, 128,256}, + {VARYING_FIRST_DIM, 5, 4096,512, 128,256,384,1024,2304}, + {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, + {VARYING_BOTH_DIMS, 2, 1,(128*128)+(256*256), 128,256, 128,256}, + {VARYING_BOTH_DIMS, 2, 1,(256*128)+(512*640), 256,512, 128,640}, +}; + +} // namespace + +class GroupedFusedCastMXFP8TestSuite : public ::testing::TestWithParam + , // Config + transformer_engine::DType, // InputType + transformer_engine::DType // OutputType + >> {}; + +TEST_P(GroupedFusedCastMXFP8TestSuite, Test) { + // Skip tests for pre-Blackwell architectures + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const ProcessingMethod processing_method = std::get<0>(GetParam()); + const ActivationKind activation = std::get<1>(GetParam()); + const ScalingDirection scaling_direction = std::get<2>(GetParam()); + const std::vector input_config = std::get<3>(GetParam()); + const DType input_type = std::get<4>(GetParam()); + const DType output_type = std::get<5>(GetParam()); + + const ShapeRepresentation shape_rep = static_cast(input_config[0]); + const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS) || (shape_rep == VARYING_FIRST_DIM); + + const size_t num_tensors = input_config[1]; + const std::vector logical_shape = {input_config[2], input_config[3]}; + std::vector first_dims(num_tensors); + std::vector last_dims(num_tensors); + std::vector offsets(num_tensors + 1, 0); + for (size_t t = 0; t < num_tensors; ++t) { + switch (shape_rep) { + case SAME_BOTH_DIMS: { + first_dims[t] = logical_shape[0] / num_tensors; + last_dims[t] = logical_shape[1]; + break; + } + case VARYING_FIRST_DIM: { + first_dims[t] = input_config[t + 4]; + last_dims[t] = logical_shape[1]; + break; + } + case VARYING_LAST_DIM: { + first_dims[t] = logical_shape[0]; + last_dims[t] = input_config[t + 4]; + break; + } + case VARYING_BOTH_DIMS: { + first_dims[t] = input_config[t + 4]; + last_dims[t] = input_config[t + (4 + num_tensors)]; + break; + } + } + offsets[t+1] = offsets[t] + first_dims[t] * last_dims[t]; + // Skips tests if tensor shape is not as required by the kernel + if ((first_dims[t] % 128 != 0) || (last_dims[t] % 32 != 0)) { + GTEST_SKIP(); + } + } + // Skips DBias tests if last dimension of tensors variates + if ((processing_method == ProcessingMethod::CAST_DBIAS || processing_method == ProcessingMethod::CAST_DBIAS_DACT) + && !is_single_tensor) { + GTEST_SKIP(); + } + + // Skips non Act tests if the Activation type is not an identity + if ((processing_method == ProcessingMethod::CAST_ONLY || processing_method == ProcessingMethod::CAST_DBIAS) + && activation != ActivationKind::Identity) { + GTEST_SKIP(); + } + // Skips Act tests if the Activation is an identity + if ((processing_method == ProcessingMethod::CAST_DBIAS_DACT + || processing_method == ProcessingMethod::CAST_DACT + || processing_method == ProcessingMethod::CAST_ACT) && (activation == ActivationKind::Identity)) { + GTEST_SKIP(); + } + + bool rowwise = false; + bool colwise = false; + switch (scaling_direction) { + case ScalingDirection::ROWWISE: rowwise = true; break; + case ScalingDirection::COLWISE: colwise = true; break; + case ScalingDirection::BOTH: rowwise = true; colwise = true; break; + } + + auto OP = &identity; + + if (processing_method == ProcessingMethod::CAST_ACT) { + switch (activation) { + case ActivationKind::GeLU: OP = &gelu; break; + case ActivationKind::SiLU: OP = &silu; break; + case ActivationKind::ReLU: OP = &relu; break; + case ActivationKind::QGeLU: OP = &qgelu; break; + case ActivationKind::SReLU: OP = &srelu; break; + } + } else if (processing_method == ProcessingMethod::CAST_DACT + || processing_method == ProcessingMethod::CAST_DBIAS_DACT) { + switch (activation) { + case ActivationKind::GeLU: OP = &dgelu; break; + case ActivationKind::SiLU: OP = &dsilu; break; + case ActivationKind::ReLU: OP = &drelu; break; + case ActivationKind::QGeLU: OP = &dqgelu; break; + case ActivationKind::SReLU: OP = &dsrelu; break; + } + } + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(output_type, OutputType, + performTest(processing_method, OP, shape_rep, num_tensors, + logical_shape, first_dims, last_dims, offsets, + rowwise, colwise); + ); + ); +} + +std::string to_string(const ProcessingMethod method) { + switch (method) { + case ProcessingMethod::CAST_ONLY: return "CAST_ONLY"; + case ProcessingMethod::CAST_DBIAS: return "CAST_DBIAS"; + case ProcessingMethod::CAST_DBIAS_DACT: return "CAST_DBIAS_DACT"; + case ProcessingMethod::CAST_DACT: return "CAST_DACT"; + case ProcessingMethod::CAST_ACT: return "CAST_ACT"; + default: return ""; + } +} + +std::string to_string(const ActivationKind activation) { + switch (activation) { + case ActivationKind::Identity: return "Identity"; + case ActivationKind::GeLU: return "GeLU"; + case ActivationKind::SiLU: return "SiLU"; + case ActivationKind::ReLU: return "ReLU"; + case ActivationKind::QGeLU: return "QGeLU"; + case ActivationKind::SReLU: return "SReLU"; + default: return ""; + } +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + GroupedFusedCastMXFP8TestSuite, + ::testing::Combine( + ::testing::ValuesIn(processing_methods), + ::testing::ValuesIn(activation_kinds), + ::testing::ValuesIn(scaling_directions), + ::testing::ValuesIn(input_config), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)), + [](const testing::TestParamInfo& info) { + const ProcessingMethod method = std::get<0>(info.param); + std::string name = to_string(method); + name += "X" + to_string(std::get<1>(info.param)); + + switch (std::get<2>(info.param)) { + case ScalingDirection::ROWWISE: name += "_ROWWISE_"; break; + case ScalingDirection::COLWISE: name += "_COLWISE_"; break; + case ScalingDirection::BOTH: name += "_BIDIMENSIONAL_"; break; + } + + const std::vector input = std::get<3>(info.param); + + switch(static_cast(input[0])) { + case ShapeRepresentation::SAME_BOTH_DIMS: name += "SAME_BOTH_DIMS"; break; + case ShapeRepresentation::VARYING_FIRST_DIM: name += "VARYING_FIRST_DIM"; break; + case ShapeRepresentation::VARYING_LAST_DIM: name += "VARYING_LAST_DIM"; break; + case ShapeRepresentation::VARYING_BOTH_DIMS: name += "VARYING_BOTH_DIMS"; break; + }; + + name += "_N_" + std::to_string(input[1]); + + name += "_SHAPE_" + + std::to_string(input[2]) + + "X" + std::to_string(input[3]); + + name += "_" + test::typeName(std::get<4>(info.param)) + + "_" + test::typeName(std::get<5>(info.param)); + return name; + }); diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index 675341f7db..d209ea8d47 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -13,6 +13,14 @@ 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); @@ -20,6 +28,20 @@ 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; + NVTETensor 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) { @@ -33,6 +55,20 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor activati input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTETensor 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; @@ -54,6 +90,15 @@ 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); @@ -61,6 +106,20 @@ 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; + NVTETensor 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) { @@ -74,6 +133,20 @@ void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor activat input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTETensor 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/relu.cu b/transformer_engine/common/activation/relu.cu index fd70e38c1a..b6f758caf6 100644 --- a/transformer_engine/common/activation/relu.cu +++ b/transformer_engine/common/activation/relu.cu @@ -13,6 +13,14 @@ 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); @@ -20,6 +28,20 @@ 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; + NVTETensor 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) { @@ -33,6 +55,20 @@ void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor activati input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTETensor 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; @@ -54,6 +90,15 @@ 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); @@ -61,6 +106,20 @@ 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; + NVTETensor 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) { @@ -74,6 +133,20 @@ void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor activat input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTETensor 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/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index cc812a17fa..77d5b6867f 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -13,6 +13,14 @@ 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); @@ -20,6 +28,20 @@ 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; + NVTETensor 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) { @@ -33,6 +55,20 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor activati input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_group_quantize_dbias_dsilu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTETensor 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/cast/cast.cu b/transformer_engine/common/cast/cast.cu index de1a8864da..582172a88e 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -26,6 +26,15 @@ 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, + 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, nullptr, stream); +} + void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor noop, cudaStream_t stream) { NVTE_API_CALL(nvte_quantize_noop); @@ -60,6 +69,19 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d input, activation_input, output, dbias, workspace, nullptr, stream); } +void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, + NVTETensor 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; diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index a02e7f4f07..b83df1dedf 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -18,6 +18,7 @@ #include "../../util/vectorized_pointwise.h" #include "../core/common.cuh" #include "../fp8/quantize_fp8.cuh" +#include "../mxfp8/group_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" #include "../nvfp4/quantize_nvfp4.cuh" @@ -371,6 +372,89 @@ void group_quantize_fwd_helper(const NVTETensor input, NVTETensor *outputs, } } +template +void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + using namespace detail; + + NVTEScalingMode scaling_mode = nvte_grouped_tensor_scaling_mode(output); + + const NVTEGroupedTensor activation = nullptr; + NVTETensor dbias = nullptr; + NVTETensor workspace = nullptr; + + const GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); + const GroupedTensor *activations_tensor = convertNVTEGroupedTensor(activation); + Tensor *dbias_tensor = convertNVTETensor(dbias); + Tensor *workspace_tensor = convertNVTETensor(workspace); + + // Quantization config + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Noop flag + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + // Dispatch to quantization kernel depending on data format + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + mxfp8::group_quantize( + input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, + workspace_tensor, stream); + break; + } + default: + NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); + } +} + +template +void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, NVTETensor dbias, NVTETensor workspace, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + using namespace detail; + + NVTEScalingMode scaling_mode = nvte_grouped_tensor_scaling_mode(output); + + const GroupedTensor *grad_tensor = convertNVTEGroupedTensorCheck(grad); + const GroupedTensor *input_tensor = convertNVTEGroupedTensor(input); + GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); + Tensor *dbias_tensor = convertNVTETensor(dbias); + Tensor *workspace_tensor = convertNVTETensor(workspace); + + // Quantization config + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + // Noop flag + Tensor dummy_tensor; + Tensor *noop_tensor = &dummy_tensor; + if (quant_config_cpp.noop_tensor != nullptr) { + noop_tensor = convertNVTETensorCheck(quant_config_cpp.noop_tensor); + } + + // Dispatch to quantization kernel depending on data format + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + mxfp8::group_quantize( + grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, + stream); + break; + } + default: + NVTE_ERROR("Not implemented scaling mode: " + to_string(scaling_mode) + "."); + } +} + } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh new file mode 100644 index 0000000000..7801a2064d --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -0,0 +1,962 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_quantize_mxfp8.cuh + * \brief CUDA kernels to quantize grouped tensors to MXFP8. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_QUANTIZE_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_GROUP_QUANTIZE_MXFP8_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "../core/common.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace group_quantize_kernel { + +constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; +__device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; +__device__ alignas(128) CUtensorMap g_tensor_maps_act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; +__device__ alignas(128) CUtensorMap g_tensor_maps_output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; +__device__ alignas(128) CUtensorMap g_tensor_maps_output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + +constexpr size_t SCALE_DIM_Y = 32; +constexpr size_t SCALE_DIM_X = 32; + +constexpr size_t BUFFS_NUM = 2; +constexpr size_t PACK_SIZE = 4; +constexpr size_t WAVES = SCALE_DIM_X / PACK_SIZE; + +constexpr size_t CHUNK_DIM_Y = 128; +constexpr size_t CHUNK_DIM_X = 128; +constexpr size_t THREADS_PER_CHUNK = 128; + +constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; + +constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; +constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; + +constexpr size_t BUFF_DIM_Y = THREADS_Y; +constexpr size_t BUFF_DIM_X = CHUNK_DIM_X; +constexpr size_t BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; +static_assert(BUFF_DIM_Y == 32); + +constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; +static_assert(STAGES >= 1); + +// Number of 1-byte elements that span 32 banks (4-byte each) of shared memory +constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4) / 1; // 128 + +// 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; // 4 = 128 / 32 + +__device__ __forceinline__ size_t get_current_tensor_id( + const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr) { + if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { + const size_t current_row = current_offset / last_logical_dim; + const size_t rows_per_tensor = first_logical_dim / num_tensors; + return current_row / rows_per_tensor; + } else { + size_t low = 1; + size_t hi = num_tensors; // [low, hi] + + while (low < hi) { + const size_t mid = low + (hi - low) / 2; + const size_t mid_offset = static_cast(offsets_ptr[mid]); + + if (mid_offset <= current_offset) { + low = mid + 1; + } else { + hi = mid; + } + } + return low - 1; + } +} + +__device__ __forceinline__ size_t get_tensor_rows_num( + const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t first_logical_dim, + const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { + size_t rows_num = 0; + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + case ShapeRepresentation::VARYING_LAST_DIM: + rows_num = first_logical_dim; + break; + case ShapeRepresentation::VARYING_FIRST_DIM: + case ShapeRepresentation::VARYING_BOTH_DIMS: + rows_num = static_cast(first_dims_ptr[tensor_id]); + break; + } + return rows_num; +} + +__device__ __forceinline__ size_t get_tensor_cols_num( + const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t last_logical_dim, + const int64_t *const __restrict__ last_dims_ptr) { + size_t cols_num = 0; + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + case ShapeRepresentation::VARYING_FIRST_DIM: + cols_num = last_logical_dim; + break; + case ShapeRepresentation::VARYING_LAST_DIM: + case ShapeRepresentation::VARYING_BOTH_DIMS: + cols_num = static_cast(last_dims_ptr[tensor_id]); + break; + } + return cols_num; +} + +// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index +__device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, + CUtensorMap *global_tensor_map, + const uintptr_t global_data_ptr, + const size_t global_dim_Y, + const size_t global_dim_X, + const size_t data_type_size_bytes) { + __shared__ CUtensorMap shared_tensor_map; + shared_tensor_map = base_tensor_map; // Copy the base tensor map into shmem + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; + if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Shape not supported, as data stride must be 16B aligned."); + } + if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); + } + + asm volatile( + "{\n\t" + ".reg.b64 tensor_map_ptr; \n\t" + "mov.b64 tensor_map_ptr, %0; \n\t" + "tensormap.replace.tile.global_address.b1024.b64 [tensor_map_ptr], %1; \n\t" + "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 1, %2; \n\t" // DIM Y + "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 0, %3; \n\t" // DIM X + "tensormap.replace.tile.global_stride.b1024.b64 [tensor_map_ptr], 0, %4; \n" + "}\n" ::"l"(reinterpret_cast(&shared_tensor_map)), + "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), + "r"(static_cast(global_dim_X)), "l"(static_cast(global_stride_bytes)) + : "memory"); + *global_tensor_map = shared_tensor_map; + } else { + NVTE_DEVICE_ERROR( + "tensormap.replace is architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } +} + +template +__global__ void update_tma_descriptors( + const __grid_constant__ CUtensorMap base_tensor_map_input, + const __grid_constant__ CUtensorMap base_tensor_map_act_input, + const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise, + const __grid_constant__ CUtensorMap base_tensor_map_output_colwise, + const IType *const __restrict__ input_data_ptr, + const IType *const __restrict__ act_input_data_ptr, + const OType *const __restrict__ output_rowwise_data_ptr, + const OType *const __restrict__ output_colwise_data_ptr, const ShapeRepresentation shape_rep, + const size_t num_tensors, const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, const bool rowwise, const bool colwise, + const bool compute_dactivations) { + const bool leading_thread = (threadIdx.x == 0); + const size_t tensor_id = blockIdx.x; + + const size_t rows = + get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + + const size_t offset_elts = offsets_ptr[tensor_id]; + + if (leading_thread && (tensor_id < num_tensors)) { + { + const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id], + global_data_ptr, rows, cols, sizeof(IType)); + } + if (compute_dactivations) { + const uintptr_t global_data_ptr = + reinterpret_cast(act_input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_act_input, &g_tensor_maps_act_input[tensor_id], + global_data_ptr, rows, cols, sizeof(IType)); + } + if (rowwise) { + const uintptr_t global_data_ptr = + reinterpret_cast(output_rowwise_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output_rowwise, + &g_tensor_maps_output_rowwise[tensor_id], global_data_ptr, rows, cols, + sizeof(OType)); + } + if (colwise) { + const uintptr_t global_data_ptr = + reinterpret_cast(output_colwise_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output_colwise, + &g_tensor_maps_output_colwise[tensor_id], global_data_ptr, rows, cols, + sizeof(OType)); + } + } +} + +__device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map)); +#else + NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +template +__global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel( + const __grid_constant__ CUtensorMap tensor_map_input_static, + const __grid_constant__ CUtensorMap tensor_map_act_input_static, + const __grid_constant__ CUtensorMap tensor_map_output_rowwise_static, + const __grid_constant__ CUtensorMap tensor_map_output_colwise_static, + const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t first_logical_dim, + const size_t last_logical_dim, const int64_t *const __restrict__ offsets_ptr, + const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, e8m0_t *const __restrict__ scales_rowwise_ptr, + e8m0_t *const __restrict__ scales_colwise_ptr, const float *__restrict__ noop, + float *const __restrict__ dbias_workspace, float *const __restrict__ amax_ptr) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; + constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; + + using IType2 = typename ptx::FPx2; + using OType2 = typename ptx::FPx2; + + if constexpr (NO_ACTIVATIONS) { + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + } + + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; + + const size_t block_global_offset = blockIdx.x * ELTS_PER_CHUNK; + + const size_t tensor_id = get_current_tensor_id(shape_rep, num_tensors, block_global_offset, + first_logical_dim, last_logical_dim, offsets_ptr); + + const size_t rows = + get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + + const size_t scale_stride_rowwise = DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(32)), 4); + const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, 128); + + const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); + + // grouped tensor can be treated as continuous tensor for MXFP8 + const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[tensor_id]); + + const CUtensorMap &tensor_map_input = + is_single_tensor ? tensor_map_input_static : g_tensor_maps_input[tensor_id]; + const CUtensorMap &tensor_map_act_input = + is_single_tensor ? tensor_map_act_input_static : g_tensor_maps_act_input[tensor_id]; + const CUtensorMap &tensor_map_output_rowwise = + is_single_tensor ? tensor_map_output_rowwise_static : g_tensor_maps_output_rowwise[tensor_id]; + const CUtensorMap &tensor_map_output_colwise = + is_single_tensor ? tensor_map_output_colwise_static : g_tensor_maps_output_colwise[tensor_id]; + + const bool leading_thread = (threadIdx.x == 0); + + if (leading_thread && (!is_single_tensor)) { + fence_acquire_tensormap(&tensor_map_input); + if constexpr (COMPUTE_ACTIVATIONS) { + fence_acquire_tensormap(&tensor_map_act_input); + } + if constexpr (ROWWISE_SCALING) { + fence_acquire_tensormap(&tensor_map_output_rowwise); + } + if constexpr (COLWISE_SCALING) { + fence_acquire_tensormap(&tensor_map_output_colwise); + } + } + + const size_t blocks_X_num_in_current_tensor = DIVUP(cols, static_cast(128)); + const size_t block_id_in_current_tensor = + is_single_tensor ? blockIdx.x : (blockIdx.x - tensor_base / ELTS_PER_CHUNK); + + const size_t block_id_Y = block_id_in_current_tensor / blocks_X_num_in_current_tensor; + const size_t block_id_X = block_id_in_current_tensor % blocks_X_num_in_current_tensor; + + const size_t block_offset_Y = block_id_Y * CHUNK_DIM_Y; + const size_t block_offset_X = block_id_X * CHUNK_DIM_X; + + e8m0_t *const scales_rowwise = + scales_rowwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_X); + e8m0_t *const scales_colwise = + scales_colwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_Y); + + const size_t scales_block_offset_Y_rowwise = block_id_Y * CHUNK_DIM_Y; + const size_t scales_block_offset_X_rowwise = block_id_X * CHUNK_DIM_X / SCALE_DIM_X; + const size_t scales_block_offset_Y_colwise = block_id_Y * CHUNK_DIM_Y / SCALE_DIM_Y; + const size_t scales_block_offset_X_colwise = block_id_X * CHUNK_DIM_X; + + const size_t tid_Y_rowwise = threadIdx.x / THREADS_X; + const size_t tid_X_rowwise = threadIdx.x % THREADS_X; + const size_t tid_Y_colwise = 0; + const size_t tid_X_colwise = threadIdx.x; + + const size_t thread_offset_Y_rowwise = tid_Y_rowwise; + const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM_X; + + const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; + const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; + const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; + const size_t scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; + + // 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_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 = + DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0); + + // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned + extern __shared__ unsigned char dynamic_shmem[]; + unsigned char *dshmem = common::align_smem_ptr_per_TMA_requirements(dynamic_shmem); + + // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned + IType *in_sh = reinterpret_cast(dshmem); + IType *act_in_sh = reinterpret_cast(dshmem + elt_input_mem); + + OType *out_rowwise_data_sh = reinterpret_cast(dshmem + in_mem); + OType *out_colwise_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise); + IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer + + constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; + + float partial_dbias_colwise = 0.0f; + float thread_dbias_rowwise[SCALE_DIM_X]; + if constexpr (IS_DBIAS) { +#pragma unroll + for (int j = 0; j < SCALE_DIM_X; ++j) { + thread_dbias_rowwise[j] = 0.0f; + } + } + + float block_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, leading_thread); + + int parity = 0; + + if constexpr (IS_DACT) { + copy_2d_to_sharedx2(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, &act_in_sh[0], + &tensor_map_act_input, block_offset_X, block_offset_Y, shmem_buff_size, + &mbar[0], leading_thread); + } else { + copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, shmem_buff_size, + &mbar[0], leading_thread); + } + +#pragma unroll + for (int stage = 0; stage < STAGES; ++stage) { + const size_t buff = stage % BUFFS_NUM; + const size_t next_stage = stage + 1; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + + 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 size_t next_buff = next_stage % BUFFS_NUM; + const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_buff_offset = next_buff * BUFF_DIM; + if constexpr (IS_DACT) { + copy_2d_to_sharedx2(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, + global_offset_Y, &act_in_sh[next_buff_offset], &tensor_map_act_input, + global_offset_X, global_offset_Y, shmem_buff_size, &mbar[next_stage], + leading_thread); + } else { + copy_2d_to_shared(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, + global_offset_Y, shmem_buff_size, &mbar[next_stage], leading_thread); + } + } + + ptx::fence_proxy_async_shared_cta(); + + // Wait for the data to have arrived + ptx::mbarrier_wait_parity(&mbar[stage], parity); + + float thread_amax = 0.0f; + if constexpr (COLWISE_SCALING) { + const size_t shmem_offset_base_colwise = buff * BUFF_DIM + tid_X_colwise; + thread_amax = 0.0f; + float in_compute_colwise[BUFF_DIM_Y]; + IType in_colwise_IType[BUFF_DIM_Y]; + + // 1. Read/Compute elements. Find MXFP8-block AMAX + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + IType thread_amax_f16 = static_cast(0.0f); +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; + in_colwise_IType[i] = in_sh[shmem_offset_colwise]; + thread_amax_f16 = __hmax(thread_amax_f16, __habs(in_colwise_IType[i])); + } + thread_amax = static_cast(thread_amax_f16); + } else { +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; + + float elt = static_cast(in_sh[shmem_offset_colwise]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); + elt *= OP(act_in_elt, {}); + } + if constexpr (IS_DBIAS) { + partial_dbias_colwise += 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); + } + thread_amax = fmaxf(thread_amax, fabsf(elt)); + in_compute_colwise[i] = elt; + } + } + + // 2. Compute E8M0 scaling factor + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + + const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; + const size_t global_scales_offset_X = scales_offset_X_colwise; + const size_t scale_idx = + global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + scales_colwise[scale_idx] = biased_exponent; + + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; + +// 3. Scale elements +#pragma unroll + for (int i = 0; i < SCALE_DIM_Y; ++i) { + float in; + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + in = static_cast(in_colwise_IType[i]); + } else { + in = in_compute_colwise[i]; + } + const float scaled_out = in * block_scale_inverse; + + const size_t shmem_offset_elt = shmem_offset_base_colwise + i * BUFF_DIM_X; + out_colwise_data_sh[shmem_offset_elt] = static_cast(scaled_out); + } + } + + if constexpr (ROWWISE_SCALING) { + const size_t shmem_offset_base_rowwise = + buff * BUFF_DIM + thread_offset_Y_rowwise * BUFF_DIM_X; + thread_amax = 0.0f; + float in_compute_rowwise[SCALE_DIM_X]; + Vec in_cached[WAVES]; + + // used as an IType container for BF16/FP16 --> MXFP8 CAST ONLY + Vec in_IType[WAVES]; + + // 1. Read/Compute elements. Find MXFP8-block AMAX + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + 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]); + } + } + thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; + + // 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 constexpr (std::is_same_v) { +#pragma unroll + for (int e = 0; e < PACK_SIZE; ++e) { + thread_amax = fmaxf(thread_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) { + thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; + + Vec in; + Vec act_in; + + in.load_from(&in_sh[shmem_offset_rowwise]); + if constexpr (IS_DACT) { + act_in.load_from(&act_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 (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + float act_in_elt = static_cast(act_in.data.elt[e]); + elt *= OP(act_in_elt, {}); + } + + // If DBIAS was computed in the 1st pass (COLWISE) then no need to compute it again + if constexpr (IS_DBIAS && (!COLWISE_SCALING)) { + thread_dbias_rowwise[j] += 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)); + } + thread_amax = fmaxf(thread_amax, fabsf(elt)); + in_compute_rowwise[j] = elt; + } + } + } + + // 2. Compute E8M0 scaling factor + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; + const int stage_scales_offset_X = scales_offset_X_rowwise; + const int scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + scales_rowwise[scale_idx] = biased_exponent; + + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; + +// 3. Scale elements +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + Vec out; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 2; ++e) { + IType2 in; + OType2 &out_pair = reinterpret_cast(out.data.elt[e]); + if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { + in = in_IType[w].data.elt[e]; + } else if constexpr (IS_CACHED_ACT_OP) { + in.x = in_cached[w].data.elt[2 * e]; + in.y = in_cached[w].data.elt[2 * e + 1]; + } else { + const int j = w * PACK_SIZE + 2 * e; + in.x = in_compute_rowwise[j]; + in.y = in_compute_rowwise[j + 1]; + } + ptx::mul_cvt_2x(out_pair, in, block_scale_inverse_2x); + } + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; + const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_idx; + out.store_to(&out_rowwise_data_sh[shmem_offset_rowwise]); + } + } + + __builtin_assume(block_amax >= 0); + __builtin_assume(thread_amax >= 0); + block_amax = fmaxf(block_amax, thread_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 (leading_thread) { + const int global_offset_Y = block_offset_Y + stage_offset_Y; + const int global_offset_X = block_offset_X; + const int buff_offset = buff * BUFF_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])); + } + 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])); + } + + // Create a "bulk async-group" out of the previous bulk copy operation. + ptx::cp_async_bulk_commit_group(); + } + } + + parity ^= 1; + + if constexpr (IS_DBIAS) { + if (is_single_tensor) { + float thread_partial_dbias = 0.0f; + if constexpr (COLWISE_SCALING) { + thread_partial_dbias = partial_dbias_colwise; + } else { + // Reusing dshmem (in_sh) as dbias buffer [HEIGHT x WIDTH] + // HEIGHT = THREADS_Y + // WIDTH = THREADS_X * (SCALE_DIM_X + 1) + // Added extra 1-element padding per thread_X to reduce bank conflicts + float *partial_dbias_rowwise = reinterpret_cast(dshmem); + + constexpr int DBIAS_BUFF_WIDTH = THREADS_X * (SCALE_DIM_X + 1); + + const int shmem_thread_offset = + tid_Y_rowwise * DBIAS_BUFF_WIDTH + tid_X_rowwise * (SCALE_DIM_X + 1); +#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_group_offset = shmem_thread_offset + swizzled_group_idx; +#pragma unroll + for (int e = 0; e < PACK_SIZE; ++e) { + const int j = w * PACK_SIZE + e; + const int shmem_elt_idx = swizzled_group_offset + e; + partial_dbias_rowwise[shmem_elt_idx] = thread_dbias_rowwise[j]; + } + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < THREADS_Y; ++i) { + // Add extra element offset per MXFP8 scaling block [1x32] + const int scaling_block = threadIdx.x / SCALE_DIM_X; + thread_partial_dbias += + partial_dbias_rowwise[i * DBIAS_BUFF_WIDTH + threadIdx.x + scaling_block]; + } + } + const int dbias_stride = cols; + const int dbias_offset_Y = block_id_Y; + const int dbias_offset_X = block_id_X * CHUNK_DIM_X + threadIdx.x; + const int dbias_idx = dbias_offset_Y * dbias_stride + dbias_offset_X; + const bool col_out_of_bounds_dbias = (dbias_offset_X >= cols); + if (!col_out_of_bounds_dbias) { + dbias_workspace[dbias_idx] = thread_partial_dbias; + } + } + } + + if (amax_ptr != nullptr) { + const int warp_id = threadIdx.x / THREADS_PER_WARP; + // Reduce the amax over the block + block_amax = reduce_max(block_amax, warp_id); + } + + if (leading_thread && amax_ptr != nullptr) { + atomicMaxFloat(amax_ptr, block_amax); + } + + destroy_barriers(mbar, leading_thread); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +} // namespace group_quantize_kernel + +template +void group_quantize(const GroupedTensor *input, const GroupedTensor *activations, + const Tensor *noop, GroupedTensor *output, Tensor *dbias, Tensor *workspace, + cudaStream_t stream) { + using namespace group_quantize_kernel; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + + const bool use_rowwise_scaling = output->has_data(); + const bool use_colwise_scaling = output->has_columnwise_data(); + NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling, + "Either rowwise or columnwise output data need to be allocated."); + + ScalingType scaling_type = ScalingType::BIDIMENSIONAL; + if (!use_colwise_scaling) { + scaling_type = ScalingType::ROWWISE; + } else if (!use_rowwise_scaling) { + scaling_type = ScalingType::COLWISE; + } + + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (output->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else if (output->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (output->varying_both_dims()) { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + // Treat a grouped tensor with const last dims as a single tensor + const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); + + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must be same."); + NVTE_CHECK(input->has_data(), "Cannot quantize tensor without rowwise data."); + NVTE_CHECK(is_fp8_dtype(output->dtype()), "Output must have FP8 type."); + + if (IS_DACT) { + NVTE_CHECK(activations->has_data(), "Activations tensor must have data."); + NVTE_CHECK(input->num_tensors == activations->num_tensors, + "Number of grad and activations tensors must be same."); + NVTE_CHECK(input->dtype() == activations->dtype(), + "Grad and activations tensors must have the same type."); + } + + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t last_logical_dim = input->logical_shape.data[1]; + const size_t elts_total = first_logical_dim * last_logical_dim; + + const size_t num_tensors = input->num_tensors; + + size_t blocks = 0; + + if (is_single_tensor) { + const size_t blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); + const size_t blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); + blocks = blocks_Y * blocks_X; + } else { + NVTE_CHECK(num_tensors < MAX_SUPPORTED_TENSOR_DESCRIPTORS, + "Number of tensors in a group is larger than " + "the MAX number of supported descriptors (64)."); + // Only full tiles supported + NVTE_CHECK(last_logical_dim % CHUNK_DIM_X == 0, + "Last dimension of a grouped tensor should be divisible by 128."); + blocks = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); + } + const dim3 grid(blocks); + const size_t block_size = THREADS_PER_CHUNK; + + // Logical shape of a tensor with varying all dims is [1, M*K] + if (shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS) { + NVTE_CHECK(first_logical_dim % 128 == 0, + "First dimension of a grouped tensor should be divisible by 128."); + } + + const int64_t *const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + + float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; + float *const amax_ptr = reinterpret_cast(output->amax.dptr); + const float *noop_ptr = reinterpret_cast(noop->data.dptr); + + e8m0_t *const scales_rowwise_ptr = reinterpret_cast(output->scale_inv.dptr); + e8m0_t *const scales_colwise_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + + if (use_rowwise_scaling) { + NVTE_CHECK(scales_rowwise_ptr != nullptr, "Scaling tensor must be allocated"); + } + if (use_colwise_scaling) { + NVTE_CHECK(scales_colwise_ptr != nullptr, "Columnwise scaling tensor must be allocated"); + } + + const size_t dbias_rows = DIVUP(first_logical_dim, CHUNK_DIM_Y); + const size_t dbias_cols = last_logical_dim; + if constexpr (IS_DBIAS) { + NVTE_CHECK(is_single_tensor, + "DBias is only supported for tensors with the const last dimension."); + NVTE_CHECK(dbias->data.dtype == input->dtype(), + "DBias must have the same type as input_tensor."); + NVTE_CHECK(dbias->data.shape == std::vector{last_logical_dim}, "Wrong shape of DBias."); + NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor."); + + if (workspace->data.dptr == nullptr) { + workspace->data.shape = {dbias_rows, dbias_cols}; + workspace->data.dtype = DType::kFloat32; + return; + } + } + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + input->dtype(), IType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + output->dtype(), OType, + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_act_input{}; + alignas(64) CUtensorMap tensor_map_output_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_colwise{}; + + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; + + create_2D_tensor_map(tensor_map_input, input->data, first_logical_dim, last_logical_dim, + BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, input_type_bit_size); + + if constexpr (IS_DACT) { + create_2D_tensor_map(tensor_map_act_input, activations->data, first_logical_dim, + last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, + input_type_bit_size); + } + + if (use_rowwise_scaling) { + create_2D_tensor_map(tensor_map_output_rowwise, output->data, first_logical_dim, + last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, + output_type_bit_size); + } + + if (use_colwise_scaling) { + create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, + first_logical_dim, last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, + last_logical_dim, 0, output_type_bit_size); + } + + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; + constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); + const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); + const size_t out_mem = out_rowwise_mem + out_colwise_mem; + + const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; + + auto kernel = group_quantize_mxfp8_kernel; + switch (scaling_type) { + case ScalingType::ROWWISE: { + kernel = group_quantize_mxfp8_kernel; + break; + } + case ScalingType::COLWISE: { + kernel = group_quantize_mxfp8_kernel; + break; + } + case ScalingType::BIDIMENSIONAL: { + kernel = group_quantize_mxfp8_kernel; + break; + } + } + + // Update tensor descriptors before launching the kernel + if (!is_single_tensor) { + const IType *const input_dptr = reinterpret_cast(input->data.dptr); + + const IType *const act_input_dptr = + IS_DACT ? reinterpret_cast(activations->data.dptr) : nullptr; + + OType *const output_rowwise_dptr = + use_rowwise_scaling ? reinterpret_cast(output->data.dptr) : nullptr; + + OType *const output_colwise_dptr = + use_colwise_scaling ? reinterpret_cast(output->columnwise_data.dptr) + : nullptr; + update_tma_descriptors<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, input_dptr, act_input_dptr, output_rowwise_dptr, + output_colwise_dptr, shape_rep, num_tensors, first_logical_dim, last_logical_dim, + offsets_ptr, first_dims_ptr, last_dims_ptr, use_rowwise_scaling, + use_colwise_scaling, IS_DACT); + } + + NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, shape_rep, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_rowwise_ptr, + scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr); + + if constexpr (IS_DBIAS) { + common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); + } + + NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) + ); // NOLINT(*) +} + +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GROUP_QUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 55cd44d9de..4c9eed3365 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -52,6 +52,16 @@ enum class NVTE_Activation_Type { */ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the GeLU activation of the grouped input. + * If the scaling mode of the grouped 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 grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_gelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); + /*! \brief Computes the SiLU 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. @@ -62,6 +72,16 @@ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); */ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the SiLU activation of the grouped input. + * If the scaling mode of the grouped 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 grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); + /*! \brief Computes the 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. @@ -72,6 +92,16 @@ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream); */ void nvte_relu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the ReLU activation of the grouped input. + * If the scaling mode of the grouped 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 grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_relu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); + /*! \brief Computes the Quick GeLU 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. @@ -82,6 +112,16 @@ void nvte_relu(const NVTETensor input, NVTETensor output, cudaStream_t stream); */ void nvte_qgelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the Quick GeLU activation of the grouped input. + * If the scaling mode of the grouped 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 grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_qgelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); + /*! \brief Computes the Squared 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. @@ -92,6 +132,16 @@ void nvte_qgelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); */ void nvte_srelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the Squared ReLU activation of the grouped input. + * If the scaling mode of the grouped 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 grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); + /*! \brief Computes the GeLU 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. @@ -104,6 +154,18 @@ void nvte_srelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the GeLU activation gradient of the grouped input. + * If the scaling mode of the output grouped 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 grouped gradient. + * \param[in] input Input grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTETensor output, cudaStream_t stream); + /*! \brief Computes the SiLU 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. @@ -116,6 +178,18 @@ void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the SiLU activation gradient of the grouped input. + * If the scaling mode of the output grouped 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 grouped gradient. + * \param[in] input Input grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTETensor output, cudaStream_t stream); + /*! \brief Computes the 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. @@ -128,6 +202,18 @@ void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the ReLU activation gradient of the grouped input. + * If the scaling mode of the output grouped 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 grouped gradient. + * \param[in] input Input grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_drelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTETensor output, cudaStream_t stream); + /*! \brief Computes the Quick GeLU 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. @@ -140,6 +226,18 @@ void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the Quick GeLU activation gradient of the grouped input. + * If the scaling mode of the output grouped 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 grouped gradient. + * \param[in] input Input grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_dqgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTETensor output, cudaStream_t stream); + /*! \brief Computes the Squared 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. @@ -152,6 +250,18 @@ void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu void nvte_dsrelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the Squared ReLU activation gradient of the grouped input. + * If the scaling mode of the output grouped 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 grouped gradient. + * \param[in] input Input grouped tensor for activation. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTETensor output, cudaStream_t stream); + /*! \brief Computes the gated GeLU 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. diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 576494a4de..04712d3003 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -89,6 +89,17 @@ extern "C" { */ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Casts input grouped tensor to MXFP8. + * The type of quantized tensor in the output depends on the scaling mode of the output + * tensor. See file level comments. + * + * \param[in] input Input grouped tensor to be cast. + * \param[in,out] output Output grouped MXFP8 tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream); + /*! \brief Casts input tensor to FP8/MXFP8/BlockwiseFP8, providing the option to immediately exit the kernel * based on the value of the 'noop' tensor. * The type of quantized tensor in the output depends on the scaling mode of the output @@ -132,6 +143,26 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Casts input grouped tensor to MXFP8. Additionally, reduces the input along columns. + * 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. + * + * This function produces 2 results: + * - `output` is equal to `cast(dact(input))` + * - `dbias` is equal to `reduce(dact(input), dim=1)` + * + * Calling this function with the workspace being 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] input Input grouped tensor to be cast. + * \param[in,out] output Output grouped FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, @@ -155,6 +186,29 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Computes backward of GeLU operation on the grouped input, then casts to FP8/MXFP8. + * Additionally, reduces the result of the GeLU backward along columns. + * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * This function produces 2 results: + * - `output` is equal to `cast(dact(input))` + * - `dbias` is equal to `reduce(dact(input), dim=1)` + * + * Calling this function with the workspace being 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] input Input grouped tensor to be cast. + * \param[in] act_input Activation input grouped tensor. + * \param[in,out] output Output grouped FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor act_input, NVTEGroupedTensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + /*! \brief Computes backward of SiLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the SiLU backward along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, @@ -178,6 +232,29 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Computes backward of SiLU operation on the grouped input, then casts to FP8/MXFP8. + * Additionally, reduces the result of the SiLU backward along columns. + * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * This function produces 2 results: + * - `output` is equal to `cast(dact(input))` + * - `dbias` is equal to `reduce(dact(input), dim=1)` + * + * Calling this function with the workspace being 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] input Input grouped tensor to be cast. + * \param[in] act_input Activation input grouped tensor. + * \param[in,out] output Output grouped FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_dbias_dsilu(const NVTEGroupedTensor input, + const NVTEGroupedTensor act_input, NVTEGroupedTensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + /*! \brief Computes backward of ReLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the ReLU backward along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, @@ -201,6 +278,29 @@ void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor act_inpu NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Computes backward of ReLU operation on the grouped input, then casts to FP8/MXFP8. + * Additionally, reduces the result of the ReLU backward along columns. + * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * This function produces 2 results: + * - `output` is equal to `cast(dact(input))` + * - `dbias` is equal to `reduce(dact(input), dim=1)` + * + * Calling this function with the workspace being 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] input Input grouped tensor to be cast. + * \param[in] act_input Activation input grouped tensor. + * \param[in,out] output Output grouped FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor act_input, NVTEGroupedTensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + /*! \brief Computes backward of Quick GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the Quick GeLU backward along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, @@ -224,6 +324,29 @@ void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor act_inp NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Computes backward of Quick GeLU operation on the grouped input, then casts to FP8/MXFP8. + * Additionally, reduces the result of the Quick GeLU backward along columns. + * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * This function produces 2 results: + * - `output` is equal to `cast(dact(input))` + * - `dbias` is equal to `reduce(dact(input), dim=1)` + * + * Calling this function with the workspace being 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] input Input grouped tensor to be cast. + * \param[in] act_input Activation input grouped tensor. + * \param[in,out] output Output grouped FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor act_input, NVTEGroupedTensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + /*! \brief Computes backward of Squared ReLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the Squared ReLU backward along columns. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, @@ -247,6 +370,29 @@ void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor act_inp NVTETensor output, NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); +/*! \brief Computes backward of Squared ReLU operation on the grouped input, then casts to FP8/MXFP8. + * Additionally, reduces the result of the Squared ReLU backward along columns. + * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * This function produces 2 results: + * - `output` is equal to `cast(dact(input))` + * - `dbias` is equal to `reduce(dact(input), dim=1)` + * + * Calling this function with the workspace being 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] input Input grouped tensor to be cast. + * \param[in] act_input Activation input grouped tensor. + * \param[in,out] output Output grouped FP8/MXFP8 tensor. + * \param[out] dbias Result of the reduction of the input along columns. + * \param[out] workspace Workspace tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor act_input, NVTEGroupedTensor output, + NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + /*! \brief Casts input tensor from reduced to higher precision. * If the scaling mode of the input tensor is set to NVTE_MXFP8_1D_SCALING, * the block dequantization (MXFP8) of the specified shape of the block will be used. From dccf67e7b3b70cc3d30a2e071e51a83c31d17e31 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 6 Feb 2026 16:35:09 -0800 Subject: [PATCH 195/521] [Common] Bucket batch size with higher granularity for THD (#2653) bucket max_b with more granularity when >512 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- transformer_engine/common/fused_attn/utils.cu | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index 727aac447b..a897b09330 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -535,11 +535,13 @@ size_t get_max_batch_size(size_t batch_size) { // batch size is expected to be 10s-100s // b = 1, ..., 32 -> max_b = 32 // b = 33, ..., 512 -> max_b = next power of 2 - // otherwise -> max_b = b + // b = 513, ... -> max_b = increment by 512 if (log2_b <= 5) { max_b = 32; } else if (log2_b <= 9) { max_b = pow(2, log2_b); + } else { + max_b = (batch_size + 511) / 512 * 512; } return max_b; } From c1a0c9746f8be057cbb3ac1bf18547e824b4d5fb Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Mon, 9 Feb 2026 09:48:12 -0800 Subject: [PATCH 196/521] [PyTorch][Core][JAX] Expand troubleshooting docs (#2602) * expand troubleshooting docs Signed-off-by: Jeremy Berchtold * Update README.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update README.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update README.rst Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --------- Signed-off-by: Jeremy Berchtold Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- README.rst | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/README.rst b/README.rst index cd55a2e18f..3cc5f81293 100644 --- a/README.rst +++ b/README.rst @@ -315,6 +315,37 @@ Troubleshooting cd transformer_engine pip install -v -v -v --no-build-isolation . +**Problems using UV or Virtual Environments:** + +1. **Import Error:** + + * **Symptoms:** Cannot import ``transformer_engine`` + * **Solution:** Ensure your UV environment is active and that you have used ``uv pip install --no-build-isolation `` instead of a regular pip install to your system environment. + +2. **cuDNN Sublibrary Loading Failed:** + + * **Symptoms:** Errors at runtime with ``CUDNN_STATUS_SUBLIBRARY_LOADING_FAILED`` + * **Solution:** This can occur when TE is built against the container's system installation of cuDNN, but pip packages inside the virtual environment pull in pip packages for ``nvidia-cudnn-cu12/cu13``. To resolve this, when building TE from source please specify the following environment variables to point to the cuDNN in your virtual environment. + + + .. code-block:: bash + + export CUDNN_PATH=$(pwd)/.venv/lib/python3.12/site-packages/nvidia/cudnn + export CUDNN_HOME=$CUDNN_PATH + export LD_LIBRARY_PATH=$CUDNN_PATH/lib:$LD_LIBRARY_PATH + +3. **Building Wheels:** + + * **Symptoms:** Regular TE installs work correctly but UV wheel builds fail at runtime. + * **Solution:** Ensure that ``uv build --wheel --no-build-isolation -v`` is used during the wheel build as well as the pip installation of the wheel. Use ``-v`` for verbose output to verify that TE is not pulling in a mismatching version of PyTorch or JAX that differs from the UV environment's version. + +**JAX-specific Common Issues and Solutions:** + +1. **FFI Issues:** + + * **Symptoms:** ``No registered implementation for custom call to for platform CUDA`` + * **Solution:** Ensure ``--no-build-isolation`` is used during installation. If pre-building wheels, ensure that the wheel is both built and installed with ``--no-build-isolation``. See "Problems using UV or Virtual Environments" above if using UV. + .. troubleshooting-end-marker-do-not-remove Breaking Changes From b84124301f4b7033e3743fc7b509a456233da5ad 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, 9 Feb 2026 20:42:38 +0100 Subject: [PATCH 197/521] [PyTorch Debug] Skip logging stats if unsupported (#2652) fix Signed-off-by: Pawel Gadzinski --- .../debug/features/log_fp8_tensor_stats.py | 24 +++++++++---- .../debug/features/log_nvfp4_tensor_stats.py | 36 ++++++++++++------- 2 files changed, 42 insertions(+), 18 deletions(-) diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index df42fb1376..108b33fd86 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -6,6 +6,7 @@ from typing import Dict, Optional, List, Tuple from contextlib import contextmanager +import warnings import torch import nvdlfw_inspect.api as debug_api @@ -298,15 +299,26 @@ def inspect_tensor( API call used to collect the data about the tensor after process_tensor()/quantization. """ assert rowwise_quantized_tensor is columnwise_quantized_tensor - assert ( - quantizer is not None - ), "[NVTORCH INSPECT ERROR] LogFp8TensorStats cannot be run without low-precision recipe." + + # Skip logging if quantizer is None (layer runs in high precision) + if quantizer is None: + warnings.warn( + f"[LogFp8TensorStats] Skipping stats collection for layer '{layer_name}', " + f"tensor '{tensor_name}': layer runs in high precision (no quantizer)." + ) + return quantized_tensor = rowwise_quantized_tensor - assert isinstance( - quantized_tensor, QuantizedTensor - ), "[NVTORCH INSPECT ERROR] LogFp8TensorStats quantized_tensor must be a QuantizedTensor." + # Skip logging if quantized_tensor is not a QuantizedTensor (incompatible precision) + if not isinstance(quantized_tensor, QuantizedTensor): + warnings.warn( + f"[LogFp8TensorStats] Skipping stats collection for layer '{layer_name}', " + f"tensor '{tensor_name}': incompatible precision " + f"(expected QuantizedTensor, got {type(quantized_tensor).__name__})." + ) + return + recipe_name = _get_recipe_name(quantizer) for stat in config["stats"]: diff --git a/transformer_engine/debug/features/log_nvfp4_tensor_stats.py b/transformer_engine/debug/features/log_nvfp4_tensor_stats.py index ec2b3c38d3..18ac8619f3 100644 --- a/transformer_engine/debug/features/log_nvfp4_tensor_stats.py +++ b/transformer_engine/debug/features/log_nvfp4_tensor_stats.py @@ -6,6 +6,7 @@ from typing import Dict, Optional from contextlib import contextmanager +import warnings import torch import nvdlfw_inspect.api as debug_api @@ -152,23 +153,34 @@ def inspect_tensor( API call used to collect the data about the tensor after process_tensor()/quantization. """ assert rowwise_quantized_tensor is columnwise_quantized_tensor - assert ( - quantizer is not None - ), "[NVTORCH INSPECT ERROR] LogNvfp4TensorStats cannot be run without NVFP4 quantizer." + + # Skip logging if quantizer is None (layer runs in high precision) + if quantizer is None: + warnings.warn( + f"[LogNvfp4TensorStats] Skipping stats collection for layer '{layer_name}', " + f"tensor '{tensor_name}': layer runs in high precision (no quantizer)." + ) + return quantized_tensor = rowwise_quantized_tensor - # Ensure we're working with NVFP4 tensors + # Skip logging if not NVFP4 quantizer (incompatible precision) if not isinstance(quantizer, NVFP4Quantizer): - raise ValueError( - "[NVTORCH INSPECT ERROR] LogNvfp4TensorStats requires NVFP4Quantizer, " - f"but got {type(quantizer).__name__}" + warnings.warn( + f"[LogNvfp4TensorStats] Skipping stats collection for layer '{layer_name}', " + f"tensor '{tensor_name}': incompatible precision " + f"(expected NVFP4Quantizer, got {type(quantizer).__name__})." ) - - assert isinstance(quantized_tensor, NVFP4TensorStorage), ( - "[NVTORCH INSPECT ERROR] LogNvfp4TensorStats quantized_tensor must be a" - " NVFP4TensorStorage." - ) + return + + # Skip logging if quantized tensor is not NVFP4TensorStorage (incompatible precision) + if not isinstance(quantized_tensor, NVFP4TensorStorage): + warnings.warn( + f"[LogNvfp4TensorStats] Skipping stats collection for layer '{layer_name}', " + f"tensor '{tensor_name}': incompatible precision " + f"(expected NVFP4TensorStorage, got {type(quantized_tensor).__name__})." + ) + return for stat in config["stats"]: self.check_if_stat_is_supported(stat) From 2894e4931cfafe767018a94bade958f046ad635d Mon Sep 17 00:00:00 2001 From: Pingtian Li <158665726+Wohox@users.noreply.github.com> Date: Tue, 10 Feb 2026 07:27:03 +0800 Subject: [PATCH 198/521] [Pytorch] Add get_backward_dw_params api for TE module (#2614) * add grad reduce api for cuda graph hook Signed-off-by: Pingtian Li * fix code consistency Signed-off-by: Pingtian Li --------- Signed-off-by: Pingtian Li Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/pytorch/graph.py | 10 ++++++++++ transformer_engine/pytorch/module/base.py | 10 ++++++++-- transformer_engine/pytorch/module/grouped_linear.py | 3 +-- transformer_engine/pytorch/module/layernorm_mlp.py | 3 +-- 4 files changed, 20 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index f587ca9946..37fff943d6 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -853,12 +853,22 @@ def functionalized(*user_args, **user_kwargs): return functionalized def make_graphed_attribute_functions(graph_idx): + # Get te modules for current graph + te_modules = visited_te_modules.get(graph_idx, set()) # Attach backward_dw as an attribute to the graphed callable. def backward_dw(): if need_bwd_dw_graph.get(graph_idx, False): bwd_dw_graphs[graph_idx].replay() + # Trigger the grad accumulation hook for wgrad graphs. + for module in te_modules: + if ( + isinstance(module, TransformerEngineBaseModule) + and module.need_backward_dw() + ): + module._trigger_wgrad_accumulation_and_reduce_hooks() + # Attach reset as an attribute to the graphed callable. def reset(): fwd_graphs[graph_idx].reset() diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 841cdf04ca..09b12afa21 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1526,8 +1526,14 @@ def backward_dw(self): bias_tensor.grad = bgrad.to(bias_tensor.dtype) del wgrad del bgrad - for wgrad_accumulation_and_reduce_hook in self.wgrad_accumulation_and_reduce_hooks: - wgrad_accumulation_and_reduce_hook() + self._trigger_wgrad_accumulation_and_reduce_hooks() + + def _trigger_wgrad_accumulation_and_reduce_hooks(self): + """ + Trigger the wgrad accumulation and reduce hooks. + """ + for wgrad_accumulation_and_reduce_hook in self.wgrad_accumulation_and_reduce_hooks: + wgrad_accumulation_and_reduce_hook() def is_debug_iter(self) -> bool: """ diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index c9ceb714e3..6cb685a3f6 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -873,8 +873,7 @@ def backward_dw(self): del grad_biases_ del wgrad_list del tensor_list - for wgrad_accumulation_and_reduce_hook in self.wgrad_accumulation_and_reduce_hooks: - wgrad_accumulation_and_reduce_hook() + self._trigger_wgrad_accumulation_and_reduce_hooks() def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: """Customize quantizers based on current scaling recipe + linear.""" diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index bec6744518..fb88764b89 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2506,5 +2506,4 @@ def backward_dw(self): del fc2_wgrad del fc1_wgrad del fc1_bias_grad - for wgrad_accumulation_and_reduce_hook in self.wgrad_accumulation_and_reduce_hooks: - wgrad_accumulation_and_reduce_hook() + self._trigger_wgrad_accumulation_and_reduce_hooks() From b09ff7e9eb645f14719e631527fcf787079be00a Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Tue, 10 Feb 2026 09:13:01 -0800 Subject: [PATCH 199/521] [pyTorch] Fix the compilation warnings (#2663) * Fix the compilation warnings for the PyTorch extension Signed-off-by: Przemek Tredak * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak --------- Signed-off-by: Przemek Tredak Signed-off-by: Przemyslaw Tredak Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- transformer_engine/common/common.h | 7 ++++--- .../fused_attn/fused_attn_f16_arbitrary_seqlen.cu | 1 - .../common/fused_attn/fused_attn_fp8.cu | 2 +- transformer_engine/pytorch/csrc/extensions.h | 2 ++ .../pytorch/csrc/extensions/pybind.cpp | 12 ++++++++---- 5 files changed, 15 insertions(+), 9 deletions(-) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 970b7aef6c..99a2985d5e 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -313,6 +313,9 @@ struct GroupedTensor { SimpleTensor columnwise_amax; SimpleTensor scale; // for FP8-DS only + NVTEScalingMode scaling_mode; + size_t num_tensors; + // Shape information (OPTIONAL - empty if dimension is uniform across all tensors) // first_dims[i] = first dimension of tensor i (empty if all tensors have same first dim) // last_dims[i] = last dimension of tensor i (empty if all tensors have same last dim) @@ -330,8 +333,6 @@ struct GroupedTensor { // Always 2D with positive dimensions NVTEShape logical_shape; - NVTEScalingMode scaling_mode; - size_t num_tensors; NVTEGroupedTensor nvte_tensor; GroupedTensor(NVTEScalingMode scaling_mode, size_t num_tensors) @@ -342,12 +343,12 @@ struct GroupedTensor { amax(), columnwise_amax(), scale(), + scaling_mode(scaling_mode), num_tensors(num_tensors), first_dims(nullptr, std::vector{0}, DType::kInt64), last_dims(nullptr, std::vector{0}, DType::kInt64), tensor_offsets(nullptr, std::vector{0}, DType::kInt64), logical_shape(nvte_make_shape(nullptr, 1)), - scaling_mode(scaling_mode), nvte_tensor(0) {} explicit operator NVTEGroupedTensor() const noexcept { return nvte_tensor; } 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 53023361e4..d13ed97de1 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 @@ -250,7 +250,6 @@ void fused_attn_arbitrary_seqlen_fwd_impl( fe::graph::SDPA_attributes sdpa_options; sdpa_options = fe::graph::SDPA_attributes() .set_name("flash_attention") - .set_is_inference(false) .set_generate_stats(generate_stats) .set_causal_mask(is_causal) .set_causal_mask_bottom_right(is_bottom_right) diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index f886ec77f4..fe859b0b22 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1810,7 +1810,7 @@ void fused_attn_fp8_fwd_impl_v1( fe::graph::SDPA_fp8_attributes sdpa_options; sdpa_options = fe::graph::SDPA_fp8_attributes() .set_name("sdpa_fp8") - .set_is_inference(false) + .set_generate_stats(true) .set_causal_mask(is_causal) .set_attn_scale(attn_scale); diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index f7cf32eaf6..e0ea3d6b78 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -548,6 +548,7 @@ class CommOverlap : torch::CustomClassHolder, public transformer_engine::CommOve ~CommOverlap() {} + using transformer_engine::CommOverlapCore::copy_into_buffer; void copy_into_buffer(const at::Tensor &input, bool local_chunk = false); at::Tensor get_buffer(bool local_chunk = false, @@ -569,6 +570,7 @@ class CommOverlapP2P : torch::CustomClassHolder, public transformer_engine::Comm ~CommOverlapP2P() {} + using transformer_engine::CommOverlapP2PBase::copy_into_buffer; void copy_into_buffer(const at::Tensor &input, bool local_chunk = false); at::Tensor get_buffer(bool local_chunk = false, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 79dd9ea5ce..1e907d9bc0 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -492,8 +492,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("comm_cga_size") = 2, py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 16, py::arg("set_sm_margin") = true, py::arg("atomic_gemm") = false, py::arg("rs_overlap_first_gemm") = false) - .def("copy_into_buffer", &CommOverlap::copy_into_buffer, py::arg("input"), - py::arg("local_chunk") = false) + .def("copy_into_buffer", + static_cast( + &CommOverlap::copy_into_buffer), + py::arg("input"), py::arg("local_chunk") = false) .def("get_buffer", &CommOverlap::get_buffer, py::arg("local_chunk") = false, py::arg("shape") = std::nullopt) .def("get_communication_stream", &CommOverlap::get_communication_stream); @@ -510,8 +512,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("gemm_priority") = 0, py::arg("comm_priority") = 0, py::arg("num_comm_sm") = 1, py::arg("set_sm_margin") = false, py::arg("atomic_gemm") = false, py::arg("use_ce") = true, py::arg("aggregate") = false) - .def("copy_into_buffer", &CommOverlapP2P::copy_into_buffer, py::arg("input"), - py::arg("local_chunk") = false) + .def("copy_into_buffer", + static_cast( + &CommOverlapP2P::copy_into_buffer), + py::arg("input"), py::arg("local_chunk") = false) .def("get_buffer", &CommOverlapP2P::get_buffer, py::arg("local_chunk") = false, py::arg("shape") = std::nullopt) .def("get_communication_stream", &CommOverlapP2P::get_communication_stream); From 01ac7f8e10ad1618f8a04a7c3df7582416dc3d66 Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Tue, 10 Feb 2026 13:27:05 -0800 Subject: [PATCH 200/521] [Pytorch] Make test script generate checkpoints if they don't exist (#2650) Signed-off-by: Kaining Zhong --- qa/L0_pytorch_unittest/test.sh | 6 +++++- tests/pytorch/test_checkpoint.py | 2 +- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a13dfada79..cd2d85c91c 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -48,7 +48,11 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention.xml $T NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" -NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint 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" +export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint +if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then + python3 $TE_PATH/tests/pytorch/test_checkpoint.py --save-checkpoint all || error_exit "Failed to generate checkpoint files" +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" diff --git a/tests/pytorch/test_checkpoint.py b/tests/pytorch/test_checkpoint.py index 1383264fdc..0427886b84 100644 --- a/tests/pytorch/test_checkpoint.py +++ b/tests/pytorch/test_checkpoint.py @@ -101,7 +101,7 @@ def _save_checkpoint(name: str, checkpoint_dir: Optional[pathlib.Path] = None) - # Path to save checkpoint if checkpoint_dir is None: checkpoint_dir = TestLoadCheckpoint._checkpoint_dir() - checkpoint_dir.mkdir(exist_ok=True) + checkpoint_dir.mkdir(parents=True, exist_ok=True) checkpoint_file = checkpoint_dir / f"{name}.pt" # Create module and save checkpoint From 8d15258573eca4f14fec5477378a5aae2a4e3af3 Mon Sep 17 00:00:00 2001 From: Faradawn Yang <73060648+faradawn@users.noreply.github.com> Date: Tue, 10 Feb 2026 16:20:52 -0800 Subject: [PATCH 201/521] Fix Broken Quickstart Links (#2641) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix broken link of quickstart guide Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * Update README.rst Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * moved getting started guide to first and moved jax out of pytorch section Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * Update README.rst 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> --------- Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- README.rst | 2 +- examples/README.md | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/README.rst b/README.rst index 3cc5f81293..5a6721b04c 100644 --- a/README.rst +++ b/README.rst @@ -137,7 +137,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 `Quickstart Notebook `_. +For a more comprehensive tutorial, check out our `Getting Started Guide `_. .. overview-end-marker-do-not-remove diff --git a/examples/README.md b/examples/README.md index 004d1631f1..782dc42f58 100644 --- a/examples/README.md +++ b/examples/README.md @@ -23,8 +23,6 @@ Additionally, we offer [Jupyter notebook tutorials](https://github.com/NVIDIA/Tr - **FP8 Weight Caching**: Avoiding redundant FP8 casting during multiple gradient accumulation steps to improve efficiency. - [Introduction to FP8](https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/fp8_primer.ipynb) - Overview of FP8 datatypes (E4M3, E5M2), mixed precision training, delayed scaling strategies, and code examples for FP8 configuration and usage. -- [TE Quickstart](https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/quickstart.ipynb) - - Introduction to TE, building a Transformer Layer using PyTorch, and instructions on integrating TE modules like Linear and LayerNorm. - [Basic MNIST Example](https://github.com/NVIDIA/TransformerEngine/tree/main/examples/pytorch/mnist) # JAX @@ -34,7 +32,9 @@ Additionally, we offer [Jupyter notebook tutorials](https://github.com/NVIDIA/Tr - Model Parallelism: Divide a model across multiple GPUs for parallel training. - Multiprocessing with Model Parallelism: Multiprocessing for model parallelism, including multi-node support and hardware affinity setup. - [Basic MNIST Example](https://github.com/NVIDIA/TransformerEngine/tree/main/examples/jax/mnist) - +- [TE JAX Integration Tutorial](https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/te_jax_integration.ipynb) + - Introduction to integrating TE into an existing JAX model framework, building a Transformer Layer, and instructions on integrating TE modules like Linear and LayerNorm. + # Third party - [Hugging Face Accelerate + TE](https://github.com/huggingface/accelerate/tree/main/benchmarks/fp8/transformer_engine) - Scripts for training with Accelerate and TE. Supports single GPU, and multi-GPU via DDP, FSDP, and DeepSpeed ZeRO 1-3. From 8ebb47e5997d4fbfe3c9fddfa5b40e6a268368d8 Mon Sep 17 00:00:00 2001 From: Lifu Zhang Date: Tue, 10 Feb 2026 16:31:30 -0800 Subject: [PATCH 202/521] Fix on TE to support Mcore Vision Encoder CUDA Graph (#2657) * Fix on TE to support Mcore Vision Encoder CUDA Graph Signed-off-by: Lifu Zhang * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactoring code Signed-off-by: Lifu Zhang * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Lifu Zhang Co-authored-by: Lifu Zhang Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/graph.py | 29 ++++++++++++++++++----------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 37fff943d6..f4b1fb23ae 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -451,11 +451,12 @@ def hook_fn( if is_training: inputs = tuple(i for i in static_input_surface if i.requires_grad) with _none_grad_context_wrapper(inputs): + outputs_requiring_grad = tuple( + o for o in outputs if o is not None and o.requires_grad + ) torch.autograd.backward( - tuple(o for o in outputs if o.requires_grad), - grad_tensors=tuple( - torch.empty_like(o) for o in outputs if o.requires_grad - ), + outputs_requiring_grad, + grad_tensors=tuple(torch.empty_like(o) for o in outputs_requiring_grad), ) grad_inputs = tuple(input.grad for input in inputs) @@ -616,19 +617,22 @@ def hook_fn( # Note for _reuse_graph_input_output_buffers: grad output is only used # within backward, so we can reuse the same static buffers every time. static_grad_outputs_keys = tuple( - (o.shape, o.dtype, o.layout) for o in static_outputs if o.requires_grad + (o.shape, o.dtype, o.layout) + for o in static_outputs + if o is not None and o.requires_grad ) if static_grad_outputs_keys in static_grad_outputs_dict: static_grad_outputs = static_grad_outputs_dict[static_grad_outputs_keys] else: static_grad_outputs = tuple( - torch.empty_like(o) if o.requires_grad else None + torch.empty_like(o) if o is not None and o.requires_grad else None for o in static_outputs ) static_grad_outputs_dict[static_grad_outputs_keys] = static_grad_outputs else: static_grad_outputs = tuple( - torch.empty_like(o) if o.requires_grad else None for o in static_outputs + torch.empty_like(o) if o is not None and o.requires_grad else None + for o in static_outputs ) if is_training: inputs = tuple(i for i in static_input_surface if i.requires_grad) @@ -636,7 +640,9 @@ def hook_fn( bwd_graph, pool=mempool ): torch.autograd.backward( - tuple(o for o in static_outputs if o.requires_grad), + tuple( + o for o in static_outputs if o is not None and o.requires_grad + ), grad_tensors=tuple(o for o in static_grad_outputs if o is not None), retain_graph=retain_graph_in_backward, ) @@ -719,7 +725,8 @@ def hook_fn( ): # For now, assumes all static_outputs require grad static_grad_outputs = tuple( - torch.empty_like(o) if o.requires_grad else None for o in static_outputs + torch.empty_like(o) if o is not None and o.requires_grad else None + for o in static_outputs ) if is_training: inputs = tuple(i for i in static_input_surface if i.requires_grad) @@ -727,7 +734,7 @@ def hook_fn( bwd_graph, pool=mempool ): torch.autograd.backward( - tuple(o for o in static_outputs if o.requires_grad), + tuple(o for o in static_outputs if o is not None and o.requires_grad), grad_tensors=tuple(o for o in static_grad_outputs if o is not None), retain_graph=retain_graph_in_backward, ) @@ -794,7 +801,7 @@ def forward(ctx, skip_fp8_weight_update, *inputs): # Replay forward graph fwd_graph.replay() assert isinstance(static_outputs, tuple) - return tuple(o.detach() for o in static_outputs) + return tuple(o.detach() if o is not None else o for o in static_outputs) @staticmethod @torch.autograd.function.once_differentiable From ac81c85b56ec9cdb2cfd213394feb7460eb8fa44 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 11 Feb 2026 11:09:40 +0530 Subject: [PATCH 203/521] [PyTorch] Python `GroupedTensor` (#2654) * PyTorch-Python GroupedTensor Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/pytorch/tensor/storage/grouped_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Remove mxfp8 gq test Signed-off-by: Kirthi Shankar Sivamani * Fix recipe tests and FP8 weights Signed-off-by: Kirthi Shankar Sivamani * Fix device test Signed-off-by: Kirthi Shankar Sivamani * Disable grouped weights for unsupported recipes Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/pytorch/test_grouped_tensor.py | 385 +++++++ tests/pytorch/test_sanity.py | 132 ++- transformer_engine/common/recipe/__init__.py | 35 +- .../pytorch/module/grouped_linear.py | 93 +- .../pytorch/quantized_tensor.py | 18 +- .../pytorch/tensor/float8_tensor.py | 19 +- .../pytorch/tensor/mxfp8_tensor.py | 45 +- .../pytorch/tensor/nvfp4_tensor.py | 5 +- .../pytorch/tensor/storage/__init__.py | 1 + .../float8_blockwise_tensor_storage.py | 18 + .../tensor/storage/float8_tensor_storage.py | 18 + .../pytorch/tensor/storage/grouped_tensor.py | 942 ++++++++++++++++++ .../tensor/storage/mxfp8_tensor_storage.py | 18 + .../tensor/storage/nvfp4_tensor_storage.py | 20 + 14 files changed, 1714 insertions(+), 35 deletions(-) create mode 100644 tests/pytorch/test_grouped_tensor.py create mode 100644 transformer_engine/pytorch/tensor/storage/grouped_tensor.py diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py new file mode 100644 index 0000000000..318009c669 --- /dev/null +++ b/tests/pytorch/test_grouped_tensor.py @@ -0,0 +1,385 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for GroupedTensor class""" + +from typing import List, Tuple +import pytest +import torch +import transformer_engine.pytorch as te +from transformer_engine.pytorch.tensor.storage.grouped_tensor import GroupedTensor +from transformer_engine.pytorch import ( + Quantizer, + Float8Quantizer, + Float8CurrentScalingQuantizer, + Float8BlockQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, +) +from transformer_engine.pytorch.constants import TE_DType_To_Torch +import transformer_engine_torch as tex + +# Check available recipes +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +_quantization_params = [ + pytest.param( + "fp8_delayed_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + ), + pytest.param( + "fp8_current_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + ), + pytest.param( + "fp8_blockwise", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, reason=reason_for_no_fp8_block_scaling + ), + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + ), +] + + +def make_quantizer(quantization: str, num_tensors: int, shape: List[Tuple[int, int]]) -> Quantizer: + """Create quantizers for given quantization scheme""" + + if quantization == "fp8_delayed_scaling": + quantizer = Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device="cuda"), + amax=torch.zeros(1, dtype=torch.float32, device="cuda"), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + elif quantization == "fp8_current_scaling": + quantizer = Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + quantizer.set_usage(rowwise=True, columnwise=False) + elif quantization == "fp8_blockwise": + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + 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 == "nvfp4": + quantizer = NVFP4Quantizer( + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=False, + ) + else: + raise ValueError(f"Unknown quantization scheme: {quantization}") + + quantizer.internal = False + + return quantizer + + +def _get_rowwise_data_tensor(qtensor, quantization: str) -> torch.Tensor: + if quantization in ("fp8_delayed_scaling", "fp8_current_scaling"): + return qtensor._data + if quantization in ("fp8_blockwise", "mxfp8", "nvfp4"): + return qtensor._rowwise_data + raise ValueError(f"Unknown quantization scheme: {quantization}") + + +def _rowwise_offset_bytes(numel: int, quantization: str) -> int: + if quantization == "nvfp4": + return numel // 2 + return numel + + +class TestGroupedTensor: + @staticmethod + def setup_class(cls) -> None: + # Configure RNG + seed = 1234 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + def test_basic_construction_all_same_shape(self) -> None: + """Test GroupedTensor construction with all tensors having same shape""" + num_tensors = 4 + shape = [(256, 512) for _ in range(num_tensors)] + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=None, + device="cuda", + dtype=torch.float32, + ) + + assert grouped_tensor.num_tensors == num_tensors + assert grouped_tensor.all_same_shape() + assert grouped_tensor.all_same_first_dim() + assert grouped_tensor.all_same_last_dim() + assert grouped_tensor.logical_shape == (num_tensors * 256, 512) + assert grouped_tensor.get_common_first_dim() == 256 + assert grouped_tensor.get_common_last_dim() == 512 + assert grouped_tensor.has_data() + + def test_basic_construction_varying_first_dim(self) -> None: + """Test GroupedTensor construction with varying first dimension""" + num_tensors = 3 + shape = [(128, 512), (256, 512), (384, 512)] + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=None, + device="cuda", + dtype=torch.float32, + ) + + assert grouped_tensor.num_tensors == num_tensors + assert not grouped_tensor.all_same_shape() + assert not grouped_tensor.all_same_first_dim() + assert grouped_tensor.all_same_last_dim() + assert grouped_tensor.get_common_last_dim() == shape[0][1] + assert grouped_tensor.logical_shape == ( + sum(v for v, _ in shape), + shape[0][1], + ) # sum of first dims + + def test_split_into_quantized_tensors_no_quantization(self) -> None: + """Test split_into_quantized_tensors for unquantized tensors""" + num_tensors = 3 + shape = [(256, 512) for _ in range(num_tensors)] + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=None, + device="cuda", + dtype=torch.float32, + ) + + # Get the original data pointer + original_data_ptr = grouped_tensor.data.data_ptr() + + # Split into tensors + tensors = grouped_tensor.split_into_quantized_tensors() + + assert len(tensors) == num_tensors + + # Verify each tensor has correct shape and shares storage + for i, tensor in enumerate(tensors): + assert tensor.shape == shape[i] + assert isinstance(tensor, torch.Tensor) + assert not hasattr(tensor, "_data") # Not a quantized tensor + + # Verify data pointer is within the original grouped tensor storage + # The tensor should be a view of the original data + assert tensor.data_ptr() >= original_data_ptr + + # Calculate expected offset + expected_offset = i * (shape[i][0] * shape[i][1]) * tensor.element_size() + assert tensor.data_ptr() == original_data_ptr + expected_offset + + @pytest.mark.parametrize("quantization", _quantization_params) + def test_split_into_quantized_tensors_quantized(self, quantization: str) -> None: + """Test split_into_quantized_tensors for quantized tensors""" + num_tensors = 3 + shape = [(512, 512) for _ in range(num_tensors)] + quantizers = make_quantizer(quantization, num_tensors, shape) + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=quantizers, + device="cuda", + ) + + # Get the original data pointer + original_data_ptr = grouped_tensor.data.data_ptr() + + # Split into tensors + tensors = grouped_tensor.split_into_quantized_tensors() + + assert len(tensors) == num_tensors + + # Verify each tensor shares storage with the grouped tensor + for i, tensor in enumerate(tensors): + rowwise_data = _get_rowwise_data_tensor(tensor, quantization) + assert rowwise_data is not None + assert rowwise_data.data_ptr() >= original_data_ptr + numel = shape[i][0] * shape[i][1] + expected_offset = _rowwise_offset_bytes(i * numel, quantization) + assert rowwise_data.data_ptr() == original_data_ptr + expected_offset + + def test_split_varying_shapes(self) -> None: + """Test split_into_quantized_tensors with varying shapes""" + num_tensors = 3 + shape = [(128, 512), (256, 512), (384, 512)] + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=None, + device="cuda", + dtype=torch.float32, + ) + + original_data_ptr = grouped_tensor.data.data_ptr() + tensors = grouped_tensor.split_into_quantized_tensors() + + assert len(tensors) == num_tensors + + # Verify shapes and storage + cumulative_offset = 0 + for i, tensor in enumerate(tensors): + assert tensor.shape == shape[i] + expected_offset = cumulative_offset * tensor.element_size() + assert tensor.data_ptr() == original_data_ptr + expected_offset + cumulative_offset += shape[i][0] * shape[i][1] + + @pytest.mark.parametrize("quantization", _quantization_params) + def test_quantize_inplace(self, quantization: str) -> None: + """Test that quantize is done in-place for all recipes""" + num_tensors = 3 + shape = [(512, 512) for _ in range(num_tensors)] + quantizers = make_quantizer(quantization, num_tensors, shape) + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=quantizers, + device="cuda", + ) + + # Get original data pointers before quantization + original_data_ptr = grouped_tensor.data.data_ptr() + original_scale_inv_ptr = grouped_tensor.scale_inv.data_ptr() + original_scale_ptr = ( + grouped_tensor.scale.data_ptr() if grouped_tensor.scale is not None else None + ) + + # Create input tensors + input_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in shape] + + # Quantize in place + quantized_tensors = grouped_tensor.quantize(input_tensors) + + # Verify data pointers haven't changed (in-place operation) + assert grouped_tensor.data.data_ptr() == original_data_ptr + assert grouped_tensor.scale_inv.data_ptr() == original_scale_inv_ptr + if original_scale_ptr is not None: + assert grouped_tensor.scale.data_ptr() == original_scale_ptr + + # Verify returned tensors point to the same storage + for i, qtensor in enumerate(quantized_tensors): + rowwise_data = _get_rowwise_data_tensor(qtensor, quantization) + numel = shape[i][0] * shape[i][1] + expected_offset = _rowwise_offset_bytes(i * numel, quantization) + assert rowwise_data.data_ptr() == original_data_ptr + expected_offset + + @pytest.mark.parametrize("quantization", _quantization_params) + def test_quantize_varying_shapes(self, quantization: str) -> None: + """Test quantize with varying shapes""" + num_tensors = 3 + shape = [(256, 512), (512, 512), (768, 512)] + quantizers = make_quantizer(quantization, num_tensors, shape) + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=quantizers, + device="cuda", + ) + + # Get original data pointers + original_data_ptr = grouped_tensor.data.data_ptr() + + # Create input tensors with varying shapes + input_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in shape] + + # Quantize in place + quantized_tensors = grouped_tensor.quantize(input_tensors) + + # Verify data pointer hasn't changed + assert grouped_tensor.data.data_ptr() == original_data_ptr + + # Verify each tensor points to correct location + cumulative_numel = 0 + for qtensor, tensor_shape in zip(quantized_tensors, shape): + rowwise_data = _get_rowwise_data_tensor(qtensor, quantization) + expected_offset = _rowwise_offset_bytes(cumulative_numel, quantization) + assert rowwise_data.data_ptr() == original_data_ptr + expected_offset + cumulative_numel += tensor_shape[0] * tensor_shape[1] + + @pytest.mark.parametrize("quantization", _quantization_params) + def test_static_quantize_method(self, quantization: str) -> None: + """Test the static quantize method""" + num_tensors = 3 + shape = [(512, 512) for _ in range(num_tensors)] + quantizers = make_quantizer(quantization, num_tensors, shape) + + # Create input tensors + input_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in shape] + + # Use static quantize method + grouped_tensor = GroupedTensor.create_and_quantize( + tensors=input_tensors, + quantizer=quantizers, + device="cuda", + ) + + # Verify the grouped tensor was created correctly + assert grouped_tensor.num_tensors == num_tensors + assert grouped_tensor.has_data() + + # Verify quantized_tensors were created and point to same storage + assert grouped_tensor.quantized_tensors is not None + assert len(grouped_tensor.quantized_tensors) == num_tensors + + original_data_ptr = grouped_tensor.data.data_ptr() + for i, qtensor in enumerate(grouped_tensor.quantized_tensors): + rowwise_data = _get_rowwise_data_tensor(qtensor, quantization) + numel = shape[i][0] * shape[i][1] + expected_offset = _rowwise_offset_bytes(i * numel, quantization) + assert rowwise_data.data_ptr() == original_data_ptr + expected_offset + + def test_clear(self) -> None: + """Test clear method""" + num_tensors = 3 + shape = [(256, 512) for _ in range(num_tensors)] + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_tensors, + shape=shape, + quantizer=None, + device="cuda", + dtype=torch.float32, + ) + + assert grouped_tensor.has_data() + assert grouped_tensor.num_tensors == num_tensors + + grouped_tensor.clear() + + assert not grouped_tensor.has_data() + assert grouped_tensor.num_tensors == 0 + assert grouped_tensor.data is None + assert grouped_tensor.logical_shape == (0, 0) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index e9d24c1a8e..b94cbdcd96 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -from typing import Optional +from typing import Optional, List import torch import pytest @@ -137,6 +137,117 @@ def reset_global_fp8_state(): FP8GlobalStateManager.reset() +def check_grouped_tensor_pointers_helper(tensors, num_elems_in_byte=1, tensor_name="tensor"): + """ + Verify that tensors are stored in contiguous memory. + + Args: + tensors: List or iterable of tensors to check + num_elems_in_byte: Number of elements packed per byte (1 for normal, 2 for NVFP4) + tensor_name: Name to use in error messages + """ + tensor_list = list(tensors) + if len(tensor_list) < 2: + return # Nothing to check + + for i in range(1, len(tensor_list)): + prev_tensor = tensor_list[i - 1] + curr_tensor = tensor_list[i] + + # Calculate expected offset based on previous tensor size + prev_numel = prev_tensor.numel() + expected_offset = (prev_numel // num_elems_in_byte) * prev_tensor.element_size() + + # Verify current tensor's data pointer is correctly offset + expected_ptr = prev_tensor.data_ptr() + expected_offset + actual_ptr = curr_tensor.data_ptr() + + assert ( + actual_ptr == expected_ptr + ), f"{tensor_name} {i} data pointer mismatch: expected {expected_ptr}, got {actual_ptr}" + + +def check_grouped_tensor_pointers( + weights: List[torch.Tensor], fp8_recipe: Optional[recipe.Recipe] = None +): + """ + Verify that the pointers of the weights are in contiguous memory for GroupedTensor. + TODO(ksivaman): This check can be made way more efficient but for now leaving the brute force approach. + """ + + num_elems_in_a_data_byte = 1 if fp8_recipe is None else 2 if fp8_recipe.nvfp4() else 1 + + # Check data. + if hasattr(weights[0], "_data") and weights[0]._data is not None: + data_tensors = [w._data for w in weights] + check_grouped_tensor_pointers_helper(data_tensors, num_elems_in_byte=1, tensor_name="data") + + # Check transpose. + if hasattr(weights[0], "_transpose") and weights[0]._transpose is not None: + transpose_tensors = [w._transpose for w in weights] + check_grouped_tensor_pointers_helper( + transpose_tensors, num_elems_in_byte=1, tensor_name="transpose" + ) + + # Check scale_inv. + if hasattr(weights[0], "_scale_inv") and weights[0]._scale_inv is not None: + scale_inv_tensors = [w._scale_inv for w in weights] + check_grouped_tensor_pointers_helper( + scale_inv_tensors, num_elems_in_byte=1, tensor_name="scale_inv" + ) + + # Check rowwise scale_inv. + if hasattr(weights[0], "_rowwise_scale_inv") and weights[0]._rowwise_scale_inv is not None: + scale_inv_tensors = [w._rowwise_scale_inv for w in weights] + check_grouped_tensor_pointers_helper( + scale_inv_tensors, num_elems_in_byte=1, tensor_name="rowwise_scale_inv" + ) + + # Check columnwise scale_inv. + if ( + hasattr(weights[0], "_columnwise_scale_inv") + and weights[0]._columnwise_scale_inv is not None + ): + columnwise_scale_inv_tensors = [w._columnwise_scale_inv for w in weights] + check_grouped_tensor_pointers_helper( + columnwise_scale_inv_tensors, + num_elems_in_byte=1, + tensor_name="columnwise scale_inv", + ) + + # Check rowwise amax. + if hasattr(weights[0], "_rowwise_amax") and weights[0]._rowwise_amax is not None: + rowwise_amax_tensors = [w._rowwise_amax for w in weights] + check_grouped_tensor_pointers_helper( + rowwise_amax_tensors, num_elems_in_byte=1, tensor_name="rowwise amax" + ) + + # Check columnwise amax. + if hasattr(weights[0], "_columnwise_amax") and weights[0]._columnwise_amax is not None: + columnwise_amax_tensors = [w._columnwise_amax for w in weights] + check_grouped_tensor_pointers_helper( + columnwise_amax_tensors, num_elems_in_byte=1, tensor_name="columnwise amax" + ) + + # Check rowwise data. + if hasattr(weights[0], "_rowwise_data") and weights[0]._rowwise_data is not None: + rowwise_data_tensors = [w._rowwise_data for w in weights] + check_grouped_tensor_pointers_helper( + rowwise_data_tensors, + num_elems_in_byte=num_elems_in_a_data_byte, + tensor_name="rowwise data", + ) + + # Check columnwise data. + if hasattr(weights[0], "_columnwise_data") and weights[0]._columnwise_data is not None: + columnwise_data_tensors = [w._columnwise_data for w in weights] + check_grouped_tensor_pointers_helper( + columnwise_data_tensors, + num_elems_in_byte=num_elems_in_a_data_byte, + tensor_name="columnwise data", + ) + + def _test_sanity_e2e_amp(block, dtype, config, fp8_recipe, skip_wgrad): te_inp_hidden_states = torch.randn( (config.max_seqlen_q, config.batch_size, config.hidden_size), @@ -495,9 +606,18 @@ def test_sanity_grouped_linear( use_fp8 = fp8_recipe is not None with quantized_model_init(enabled=use_fp8 and fp8_model_params, recipe=fp8_recipe): te_grouped_linear = GroupedLinear( - num_gemms, config.hidden_size, ffn_hidden_size, bias=use_bias, params_dtype=dtype + num_gemms, + config.hidden_size, + ffn_hidden_size, + bias=use_bias, + params_dtype=dtype, ).cuda() + # Verify that weights are stored in contiguous GroupedTensor storage. + weights = [getattr(te_grouped_linear, f"weight{i}") for i in range(num_gemms)] + if fp8_recipe is None or not (fp8_recipe.delayed() or fp8_recipe.float8_current_scaling()): + check_grouped_tensor_pointers(weights, fp8_recipe) + inp_hidden_states = torch.randn( num_tokens, config.hidden_size, dtype=dtype, requires_grad=True ).cuda() @@ -956,7 +1076,13 @@ def test_replace_raw_data_for_float8tensor(): random_bf16_data = torch.randn(fp8_tensor.shape, dtype=torch.bfloat16, device="cuda") fp8_quantizer.update_quantized(random_bf16_data, fp8_tensor) - attrs_to_check = ["_quantizer", "_fp8_dtype", "_scale_inv", "_transpose", "_transpose_invalid"] + attrs_to_check = [ + "_quantizer", + "_fp8_dtype", + "_scale_inv", + "_transpose", + "_transpose_invalid", + ] attrs = {} for attr in attrs_to_check: attrs[attr] = getattr(fp8_tensor, attr) diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 64ee2a5a16..18577b0eb4 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -88,33 +88,40 @@ class Recipe: Base recipe class. """ - def nvfp4(self): + @classmethod + def nvfp4(cls): """Whether the given recipe is NVFP4 1D block scaling.""" - return isinstance(self, NVFP4BlockScaling) + return issubclass(cls, NVFP4BlockScaling) - def mxfp8(self): + @classmethod + def mxfp8(cls): """Whether the given recipe is MXFP8 block scaling.""" - return isinstance(self, MXFP8BlockScaling) + return issubclass(cls, MXFP8BlockScaling) - def delayed(self): + @classmethod + def delayed(cls): """Whether the given recipe is delayed scaling.""" - return isinstance(self, DelayedScaling) + return issubclass(cls, DelayedScaling) - def float8_current_scaling(self): + @classmethod + def float8_current_scaling(cls): """Whether the given recipe is (per-tensor) current scaling.""" - return isinstance(self, Float8CurrentScaling) + return issubclass(cls, Float8CurrentScaling) - def float8_per_tensor_scaling(self): + @classmethod + def float8_per_tensor_scaling(cls): """Whether the given recipe is per-tensor scaling.""" - return isinstance(self, (DelayedScaling, Float8CurrentScaling)) + return issubclass(cls, (DelayedScaling, Float8CurrentScaling)) - def float8_block_scaling(self): + @classmethod + def float8_block_scaling(cls): """Whether the given recipe is float8 blockwise scaling.""" - return isinstance(self, Float8BlockScaling) + return issubclass(cls, Float8BlockScaling) - def custom(self): + @classmethod + def custom(cls): """Whether the given recipe is custom.""" - return isinstance(self, CustomRecipe) + return issubclass(cls, CustomRecipe) @dataclass() diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 6cb685a3f6..b6596bc2e9 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -13,6 +13,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe +from transformer_engine.pytorch.tensor.storage.grouped_tensor import GroupedTensor from .base import ( get_dummy_wgrad, TransformerEngineBaseModule, @@ -147,7 +148,10 @@ def forward( # tensors (like scales), but bulk allocation shares storage across all tensors, # so if scales can't be offloaded, nothing in the group can be offloaded. inputmats = tex.split_quantize( - inp_view, m_splits, input_quantizers, disable_bulk_allocation=cpu_offloading + inp_view, + m_splits, + input_quantizers, + disable_bulk_allocation=cpu_offloading, ) elif debug: inputmats = DebugQuantizer.multi_tensor_quantize( @@ -365,7 +369,10 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], for i in range(ctx.num_gemms): grad_biases[i] = grad_output_mats[i].sum(dim=0) grad_output = DebugQuantizer.multi_tensor_quantize( - grad_output_view, ctx.grad_output_quantizers, ctx.m_splits, ctx.activation_dtype + grad_output_view, + ctx.grad_output_quantizers, + ctx.m_splits, + ctx.activation_dtype, ) else: # Only split grad output. Grad bias is fused with @@ -436,7 +443,8 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if ctx.input_quantizers[0] is not None: for input_quantizer in ctx.input_quantizers: if isinstance( - input_quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer) + input_quantizer, + (Float8Quantizer, Float8CurrentScalingQuantizer), ): input_quantizer.set_usage(rowwise=True, columnwise=True) else: @@ -446,7 +454,10 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], inputmats = tex.split_quantize(inp_view, ctx.m_splits, ctx.input_quantizers) elif ctx.debug: inputmats = DebugQuantizer.multi_tensor_quantize( - inp_view, ctx.input_quantizers, ctx.m_splits, ctx.activation_dtype + inp_view, + ctx.input_quantizers, + ctx.m_splits, + ctx.activation_dtype, ) else: inputmats = torch.split( @@ -616,7 +627,7 @@ def __init__( ) -> None: super().__init__(name) - params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype + self.params_dtype = torch.get_default_dtype() if params_dtype is None else params_dtype self.num_gemms = num_gemms self.in_features = in_features self.out_features = out_features @@ -631,12 +642,19 @@ def __init__( assert ( not ub_overlap_rs and not ub_overlap_ag ), "GroupedLinear doesn't support Userbuffer overlap." + self.init_method = init_method self.get_rng_state_tracker = get_rng_state_tracker self.rng_tracker_name = rng_tracker_name self.wgrad_store = WeightGradStore(delay_wgrad_compute) - self._offsets = {"input": 0, "weight": 1, "output": 2, "grad_output": 0, "grad_input": 1} + self._offsets = { + "input": 0, + "weight": 1, + "output": 2, + "grad_output": 0, + "grad_input": 1, + } self._num_fp8_tensors_per_gemm = { "fwd": 3, "bwd": 2, @@ -678,7 +696,7 @@ def __init__( self.out_features, self.in_features, device=device, - dtype=params_dtype, + dtype=self.params_dtype, ), ), init_fn=init_method, @@ -694,13 +712,13 @@ def __init__( torch.empty( self.out_features, device=device, - dtype=params_dtype, + dtype=self.params_dtype, ), ), init_fn=init_method_constant(0.0), ) else: - bias = torch.Tensor().to(dtype=params_dtype, device=device) + bias = torch.Tensor().to(dtype=self.params_dtype, device=device) setattr(self, f"bias{i}", bias) if self.primary_weights_in_fp8: @@ -724,8 +742,61 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) + def make_grouped_weights(self, defer_init=False) -> None: + """ + Convert parameters into a GroupedTensor and re-register them as parameters. + """ + + if defer_init: + return + + weight_quantizers = self._get_weight_quantizers() + recipe = ( + weight_quantizers[0]._get_compatible_recipe() + if weight_quantizers and weight_quantizers[0] is not None + else None + ) + if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): + self.set_tensor_parallel_attributes(defer_init=defer_init) + return + + weights = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] + + # Create the weight storage. + grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=self.num_gemms, + shape=[(self.out_features, self.in_features)] * self.num_gemms, + quantizer=weight_quantizers[0], + dtype=self.params_dtype, + device=weights[0].device, + ) + + # Copy existing params into storage. + with torch.no_grad(): + for i in range(self.num_gemms): + if self.primary_weights_in_fp8: + grouped_weights.quantized_tensors[i].copy_from_storage(weights[i]) + else: + grouped_weights.quantized_tensors[i].copy_(weights[i]) + + # Re-register the grouped weights as parameters. + for i in range(self.num_gemms): + self.register_parameter( + f"weight{i}", + torch.nn.Parameter(grouped_weights.quantized_tensors[i]), + init_fn=self.init_method, + get_rng_state_tracker=self.get_rng_state_tracker, + fp8_meta_index=self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"], + ) + + self.set_tensor_parallel_attributes(defer_init=defer_init) + def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) + self.make_grouped_weights(defer_init=defer_init) + + def set_tensor_parallel_attributes(self, defer_init=False) -> None: + """Set attributes needed for TP""" if not defer_init: # Set parallelism attributes for linear weights @@ -925,7 +996,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" - if not self.fp8 and not self.fp8_calibration: + if not self.fp8 and not self.fp8_calibration and not self.primary_weights_in_fp8: return [None] * self.num_gemms weight_quantizers = [ self.quantizers["scaling_fwd"][ @@ -934,7 +1005,7 @@ def _get_weight_quantizers(self) -> List[Quantizer]: for i in range(self.num_gemms) ] for i in range(self.num_gemms): - weight_quantizers[i].internal = True + weight_quantizers[i].internal = not self.primary_weights_in_fp8 return weight_quantizers def _get_quantizers(self): diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 0a6ad61ff0..d78677bc83 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -69,7 +69,9 @@ def get_usages(self) -> Dict[str, bool]: f"{self.__class__.__name__} class does not implement get_usages function" ) - def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]: + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]: """Prepare the tensor base for saving for backward""" raise NotImplementedError( f"{self.__class__.__name__} class does not implement prepare_for_saving function" @@ -115,11 +117,18 @@ def update_quantizer(self, quantizer: Quantizer): warnings.warn("Quantizer is being updated, this may affect model behavior") self._quantizer = quantizer + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data from another QuantizedTensorStorage.""" + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement copy_from_storage function" + ) + def prepare_for_saving( *tensors: Union[torch.Tensor, QuantizedTensorStorage], ) -> Tuple[ - list[Optional[Union[torch.Tensor, torch.nn.Parameter]]], list[Optional[QuantizedTensorStorage]] + list[Optional[Union[torch.Tensor, torch.nn.Parameter]]], + list[Optional[QuantizedTensorStorage]], ]: """Prepare tensors for saving. Needed because save_for_backward accepts only torch.Tensor/torch.nn.Parameter types, while we want to be able to save @@ -144,7 +153,10 @@ def restore_from_saved( return_saved_tensors: bool = False, ) -> ( list[Optional[torch.Tensor | QuantizedTensorStorage]] - | tuple[list[Optional[torch.Tensor | QuantizedTensorStorage]], list[Optional[torch.Tensor]]] + | tuple[ + list[Optional[torch.Tensor | QuantizedTensorStorage]], + list[Optional[torch.Tensor]], + ] ): """Recombine the tensor data and metadata during backward pass.""" tensor_objects = [] diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 3aeace0a77..55bca49af3 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -11,7 +11,11 @@ import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType -from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling, Recipe +from transformer_engine.common.recipe import ( + DelayedScaling, + Float8CurrentScaling, + Recipe, +) from ..utils import canonicalize_process_group, devices_match from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func from ..quantized_tensor import QuantizedTensor, Quantizer @@ -154,6 +158,10 @@ def calibrate(self, tensor: torch.Tensor) -> None: amin, amax = tensor.aminmax() self.amax.copy_(torch.max(-amin, amax)) + def get_columnwise_shape(self, rowwise_data_shape: Iterable[int]) -> Tuple[int, ...]: + """Calculate the shape of the columnwise data for Float8 1D blockwise quantization.""" + return [rowwise_data_shape[-1]] + list(rowwise_data_shape[:-1]) + def create_tensor_from_data( self, data: torch.Tensor, @@ -408,6 +416,10 @@ def create_tensor_from_data( quantizer=self, ) + def get_columnwise_shape(self, rowwise_data_shape: Iterable[int]) -> Tuple[int, ...]: + """Calculate the shape of the columnwise data for Float8 1D blockwise quantization.""" + return [rowwise_data_shape[-1]] + list(rowwise_data_shape[:-1]) + def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: """Function using primitives with ONNX defined translations.""" if tensor.dtype != torch.float32: @@ -769,7 +781,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): kwargs, ) return Float8Tensor.make_like( - tensor, data=func_out, data_transpose=func_transposed_out, shape=func_out.shape + tensor, + data=func_out, + data_transpose=func_transposed_out, + shape=func_out.shape, ) if func == torch.ops.aten.detach.default: diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 8dd2255d89..41d6c87f2b 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -164,6 +164,49 @@ def calibrate(self, tensor: torch.Tensor) -> None: # TODO(ksivamani): No calibration needed for mxfp8? pass + def get_scale_shape( + self, + shape: Iterable[int], + columnwise: bool, + ) -> Tuple[int, int]: + """Calculate the shape of the scaling tensor for MXFP8 1D blockwise quantization. + + This method determines the shape of the scaling tensor needed for blockwise quantization, + taking into account the input tensor shape and whether columnwise scaling is used. + + Parameters + ---------- + shape : Iterable[int] + Shape of the input tensor to be quantized + columnwise : bool + Whether to use columnwise scaling (True) or rowwise scaling (False) + + Returns + ------- + Tuple[int, int] + Shape of the scaling tensor as (outer_dim, inner_dim) + For MXFP8 1D blockwise quantization, blocksize is 32 + Swizzle kernel will be performed before GEMM to suit the need of CuBLAS. + CuBLAS doc: https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout + """ + if columnwise: + # Columnwise: scale_inv shape is [prod(shape[:-1]) // BLOCK_SIZE, shape[-1]] + # with padding to multiples of [4, 128] + return ( + round_up_to_nearest_multiple(math.prod(shape[:-1]) // MXFP8_BLOCK_SCALING_SIZE, 4), + round_up_to_nearest_multiple(shape[-1], 128), + ) + # Rowwise: scale_inv shape is [prod(shape[:-1]), shape[-1] // BLOCK_SIZE] + # with padding to multiples of [128, 4] + return ( + round_up_to_nearest_multiple(math.prod(shape[:-1]), 128), + round_up_to_nearest_multiple(shape[-1] // MXFP8_BLOCK_SCALING_SIZE, 4), + ) + + def get_columnwise_shape(self, rowwise_data_shape: Tuple[int, ...]) -> Tuple[int, ...]: + """Calculate the shape of the columnwise data for MXFP8 1D blockwise quantization.""" + return rowwise_data_shape + def create_tensor_from_data( self, data: torch.Tensor, @@ -704,7 +747,7 @@ def fsdp_post_all_gather( columnwise_scale_inv=columnwise_scale_inv, fp8_dtype=fp8_dtype, dtype=param_dtype, - shape=rowwise_data.shape if rowwise_data is not None else columnwise_data.shape, + shape=(rowwise_data.shape if rowwise_data is not None else columnwise_data.shape), quantizer=self._quantizer, with_gemm_swizzled_scales=False, ) diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 101cf78a8f..66f986a900 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -341,7 +341,10 @@ def make_empty( ) 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 + 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 diff --git a/transformer_engine/pytorch/tensor/storage/__init__.py b/transformer_engine/pytorch/tensor/storage/__init__.py index d7a2719200..7c8a014c1d 100644 --- a/transformer_engine/pytorch/tensor/storage/__init__.py +++ b/transformer_engine/pytorch/tensor/storage/__init__.py @@ -7,3 +7,4 @@ from .mxfp8_tensor_storage import MXFP8TensorStorage # noqa: F401 from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage # noqa: F401 from .nvfp4_tensor_storage import NVFP4TensorStorage # noqa: F401 +from .grouped_tensor import GroupedTensor # noqa: F401 diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 278d7dc039..4cd6d19cd8 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -73,6 +73,24 @@ def clear(self): if t is not None: t.data = _empty_tensor() + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data buffers from another Float8BlockwiseQTensorStorage.""" + if not isinstance(src, Float8BlockwiseQTensorStorage): + raise TypeError("copy_from_storage expects Float8BlockwiseQTensorStorage") + if self._fp8_dtype != src._fp8_dtype: + raise RuntimeError("FP8 dtype mismatch in copy_from_storage") + if self._is_2D_scaled != src._is_2D_scaled: + raise RuntimeError("Scale layout 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: + dst.copy_(src_tensor) + + _copy_optional(self._rowwise_data, src._rowwise_data) + _copy_optional(self._columnwise_data, src._columnwise_data) + _copy_optional(self._rowwise_scale_inv, src._rowwise_scale_inv) + _copy_optional(self._columnwise_scale_inv, src._columnwise_scale_inv) + def get_metadata(self) -> Dict[str, Any]: """Get this tensor's metadata.""" return { diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index adf3ce8aea..9adb86c453 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -104,6 +104,24 @@ def clear(self): t.data = _empty_tensor() self._transpose_invalid = True + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data buffers from another Float8TensorStorage.""" + if not isinstance(src, Float8TensorStorage): + raise TypeError("copy_from_storage expects Float8TensorStorage") + if self._fp8_dtype != src._fp8_dtype: + raise RuntimeError("FP8 dtype 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: + dst.copy_(src_tensor) + + _copy_optional(self._data, src._data) + _copy_optional(self._transpose, src._transpose) + _copy_optional(self._scale_inv, src._scale_inv) + def get_metadata(self) -> Dict[str, Any]: """Get this tensor's metadata.""" return { diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor.py new file mode 100644 index 0000000000..dad4d1d0ea --- /dev/null +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor.py @@ -0,0 +1,942 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped tensor class for handling collections of tensors with different shapes""" +from __future__ import annotations +from typing import Optional, Tuple, List, Union +import math + +import torch + +from ...quantized_tensor import QuantizedTensorStorage, Quantizer + +from ..mxfp8_tensor import MXFP8Tensor +from ..nvfp4_tensor import NVFP4Tensor +from ..float8_tensor import Float8Tensor +from ..float8_blockwise_tensor import Float8BlockwiseQTensor +from .float8_tensor_storage import Float8TensorStorage +from .mxfp8_tensor_storage import MXFP8TensorStorage +from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage +from .nvfp4_tensor_storage import NVFP4TensorStorage + + +class GroupedTensor: + """ + EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. + + Grouped tensor is a collection of tensors with different shapes but the same dtype and scaling mode. + + Shape Representation: + - logical_shape: 2D shape representing the conceptual layout, i.e. the shape when member tensors + are flattened to 2D and stacked together (REQUIRED) + + When all_same_shape(): [num_tensors * M, N] where each tensor is (M, N) + + When varying_first_dim(): [~sum_of_first_dims, N] where N is common + + When varying_last_dim(): [M, ~sum_of_last_dims] where M is common + + When varying_both_dims(): [1, total_elements] (fully flattened) + + - first_dims and last_dims are OPTIONAL (None if dimension is uniform) + + None first_dims: all tensors have the same first dimension + + None last_dims: all tensors have the same last dimension + + Both None: all tensors have identical shapes + + Both set: each tensor has unique shape (first_dims[i], last_dims[i]) + + Data Layout: + - ALL data fields are stored as 1D flattened arrays (data, columnwise_data, scale_inv, etc.) + - logical_shape provides the conceptual 2D interpretation + - All data is stored on device in contiguous layout + + Note: This structure is used only for combined storage of multiple tensors with the same dtype and scaling mode. + """ + + def __init__( + self, + num_tensors: int, + shape: List[Tuple[int, int]], + quantizer: Optional[Quantizer] = None, + dtype: Optional[torch.dtype] = None, + data: Optional[torch.Tensor] = None, + columnwise_data: Optional[torch.Tensor] = None, + scale_inv: Optional[torch.Tensor] = None, + columnwise_scale_inv: Optional[torch.Tensor] = None, + amax: Optional[torch.Tensor] = None, + columnwise_amax: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + first_dims: Optional[torch.Tensor] = None, + last_dims: Optional[torch.Tensor] = None, + tensor_offsets: Optional[torch.Tensor] = None, + offsets: Optional[List[int]] = None, + scale_inv_offsets: Optional[List[int]] = None, + columnwise_scale_inv_offsets: Optional[List[int]] = None, + logical_shape: Optional[Tuple[int, int]] = None, + ) -> None: + """ + Initialize a GroupedTensor. + + Args: + num_tensors: Number of tensors in the group + shape: 2D shape of each tensor (len num_tensors) + quantizer: Quantizer for the grouped tensor + data: Row-wise data buffer (1D flattened) + columnwise_data: Column-wise data buffer (1D flattened) + scale_inv: Row-wise scale inverse buffer + columnwise_scale_inv: Column-wise scale inverse buffer + amax: Row-wise amax buffer + columnwise_amax: Column-wise amax buffer + scale: Scale buffer (for FP8-DS only) + first_dims: Device tensor of int64 array of length num_tensors (or None if uniform) + last_dims: Device tensor of int64 array of length num_tensors (or None if uniform) + tensor_offsets: Device tensor of int64 array of length num_tensors (or None if uniform) + offsets: Vector of integer offsets for each tensor. + logical_shape: 2D tuple representing conceptual shape + """ + self.num_tensors = num_tensors + self.quantizer = quantizer + self.shape = shape + self.dtype = ( + dtype if dtype is not None else torch.float32 + ) # Default to float32 if not provided + + # Data buffers + self.data = data + self.columnwise_data = columnwise_data + self.scale_inv = scale_inv + self.columnwise_scale_inv = columnwise_scale_inv + self.amax = amax + self.columnwise_amax = columnwise_amax + self.scale = scale + + # For convenient indexing for python GroupedTensor API. + self.scale_inv_offsets = scale_inv_offsets + self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets + + # Shape information (OPTIONAL - None if dimension is uniform across all tensors) + # first_dims[i] = first dimension of tensor i (None if all tensors have same first dim) + # last_dims[i] = last dimension of tensor i (None if all tensors have same last dim) + self.first_dims = ( + first_dims # Device pointer to int64_t array of length num_tensors (or None) + ) + self.last_dims = ( + last_dims # Device pointer to int64_t array of length num_tensors (or None) + ) + + # Offsets for indexing into contiguous 1D layout (OPTIONAL - not needed if all_same_shape()) + # tensor_offsets[i] = element offset to start of tensor i (cumulative sum of numel for tensors 0..i-1) + # Usage: tensor_i_ptr = data.data_ptr() + tensor_offsets[i] * element_size + # If None and all_same_shape(): offset[i] = i * M * N (where M, N are common dimensions) + self.tensor_offsets = ( + tensor_offsets # Device pointer to int64_t array of length num_tensors (or None) + ) + self.offsets = offsets # Vector of integer offsets for each tensor. + + # Logical shape: conceptual 2D shape of the grouped data (REQUIRED) + # Represents how the 1D flattened data should be interpreted as 2D + # Always 2D with positive dimensions + self.logical_shape = logical_shape if logical_shape is not None else (0, 0) + + # Hold a reference to the quantized tensors that occupy same storage as the GroupedTensor. + # Used as a convenience. + self.quantized_tensors = None + + def has_data(self) -> bool: + """ + Check if the tensor has row-wise data. + + Returns: + True if data buffer is initialized, False otherwise + """ + return self.data is not None + + def has_columnwise_data(self) -> bool: + """ + Check if the tensor has column-wise data. + + Returns: + True if columnwise_data buffer is initialized, False otherwise + """ + return self.columnwise_data is not None + + def all_same_first_dim(self) -> bool: + """ + Check if all tensors in the group have the same first dimension. + + Returns: + True if first dimension is uniform across all tensors + """ + return self.first_dims is None + + def all_same_last_dim(self) -> bool: + """ + Check if all tensors in the group have the same last dimension. + + Returns: + True if last dimension is uniform across all tensors + """ + return self.last_dims is None + + def all_same_shape(self) -> bool: + """ + Check if all tensors in the group have identical shapes. + + Returns: + True if all tensors have the same shape + """ + return self.first_dims is None and self.last_dims is None + + def varying_both_dims(self) -> bool: + """ + Check if both dimensions vary across tensors. + + Returns: + True if both first and last dimensions vary + """ + return self.first_dims is not None and self.last_dims is not None + + def get_common_first_dim(self) -> int: + """ + Get the common first dimension when all tensors share it. + + Returns: + The common first dimension + + Raises: + RuntimeError: If first dimension varies across tensors or logical_shape is not 2D + """ + if not self.all_same_first_dim(): + raise RuntimeError("First dim varies across tensors") + if len(self.logical_shape) != 2: + raise RuntimeError("Logical shape must be 2D") + + if self.all_same_shape(): + # When both dims are uniform: logical_shape = [num_tensors * M, N] + return self.logical_shape[0] // self.num_tensors + # When varying last dims but not first dim: logical_shape = [M, sum_of_last_dims] + return self.logical_shape[0] + + def get_common_last_dim(self) -> int: + """ + Get the common last dimension when all tensors share it. + + Returns: + The common last dimension + + Raises: + RuntimeError: If last dimension varies across tensors or logical_shape is not 2D + """ + if not self.all_same_last_dim(): + raise RuntimeError("Last dim varies across tensors") + if len(self.logical_shape) != 2: + raise RuntimeError("Logical shape must be 2D") + + # For both uniform and varying first dim cases: logical_shape[1] is the common last dim + return self.logical_shape[1] + + def get_dtype(self) -> torch.dtype: + """ + Get the high precision data type of the tensor. + + Returns: + The high precision dtype of the data buffer + """ + + return self.dtype + + def clear(self) -> None: + """ + Reset tensor data and clear all buffers. + """ + self.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.logical_shape = (0, 0) + self.num_tensors = 0 + self.quantizer = None + self.quantized_tensors = None + self.offsets = None + self.scale_inv_offsets = None + self.columnwise_scale_inv_offsets = None + + def __repr__(self) -> str: + """String representation of the GroupedTensor.""" + return ( + f"GroupedTensor(num_tensors={self.num_tensors}, " + f"shape={self.shape}, " + f"logical_shape={self.logical_shape}, " + f"dtype={self.get_dtype()})" + ) + + def __str__(self) -> str: + """User-friendly string representation.""" + shape_info = [] + if self.all_same_shape(): + shape_info.append("uniform shape") + else: + if not self.all_same_first_dim(): + shape_info.append("varying first dim") + if not self.all_same_last_dim(): + shape_info.append("varying last dim") + + return ( + f"GroupedTensor with {self.num_tensors} tensors " + f"({', '.join(shape_info) if shape_info else 'uniform'}), " + f"logical_shape={self.logical_shape}, " + f"dtype={self.get_dtype()}" + ) + + @staticmethod + def make_grouped_tensor_with_shapes( + num_tensors: int, + shape: List[Tuple[int, int]], + quantizer: Optional[Quantizer] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> GroupedTensor: + """ + Create a GroupedTensor for storing multiple weight tensors of the same shape. + + Args: + num_tensors: Number of tensors + shape: 2D shape of each tensor (len num_tensors) + quantizer: Quantizer for each tensor + device: Device to allocate tensors on, defaults to current cuda device + dtype: Data type of the tensor (for high precision case) + + Returns: + A GroupedTensor. + """ + + # First dim + first_dim_list = [s[0] for s in shape] + uniform_first_dim = all(first_dim_list[0] == x for x in first_dim_list) + logical_first_dim = sum(first_dim_list) + if uniform_first_dim: + first_dims = None + else: + first_dims = torch.tensor([s[0] for s in shape], dtype=torch.int64, device=device) + + # Last dim + last_dim_list = [s[1] for s in shape] + logical_last_dim = last_dim_list[0] + assert all(logical_last_dim == x for x in last_dim_list), "Last dims should be uniform" + + return GroupedTensor.make_grouped_tensor( + num_tensors=num_tensors, + first_dims=first_dims, + last_dims=None, + logical_first_dim=logical_first_dim, + logical_last_dim=logical_last_dim, + quantizer=quantizer, + device=device, + dtype=dtype, + ) + + @staticmethod + def make_grouped_tensor( + num_tensors: int, + first_dims: Optional[torch.Tensor], + last_dims: Optional[torch.Tensor], + logical_first_dim: int, + logical_last_dim: int, + quantizer: Optional[Quantizer] = None, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ) -> GroupedTensor: + """ + Create a GroupedTensor for storing multiple weight tensors of the same shape. + + Args: + num_tensors: Number of tensors + first_dims: Device tensor of int64 array of length num_tensors (or None if uniform) + last_dims: Device tensor of int64 array of length num_tensors (or None if uniform) + logical_first_dim: Logical first dimension + logical_last_dim: Logical last dimension + quantizer: Quantizer for each tensor + Used to figure out the recipe and what to allocate. + device: Device to allocate tensors on, defaults to current cuda device + dtype: Data type of the tensor (for high precision case) + + Returns: + A GroupedTensor. + """ + + # Set device + if device is None: + device = torch.cuda.current_device() + + # Shape patterns and validation. + all_same_first = first_dims is None + all_same_last = last_dims is None + + assert all_same_last, "Last dim must be uniform for GroupedTensor" + assert logical_first_dim > 0, "Logical first dim must be positive for GroupedTensor" + assert logical_last_dim > 0, "Logical last dim must be positive for GroupedTensor" + + # assert ( + # logical_first_dim % 128 == 0 + # ), "Logical first dim must be divisible by 128" + # assert logical_last_dim % 128 == 0, "Logical last dim must be divisible by 128" + + # Calculate tensor offsets (cumulative element offsets) + tensor_offsets = None + offsets = None + shape = [] + if not all_same_first: + # Need explicit offsets for non-uniform shapes + # Offsets are based on number of elements and not pointers. + # Kernels need to calculate precise pointers based on size of elements. + + # TODO(ksivaman): Single kernel + remove the host offset calculation. + tensor_offsets = torch.cat( + [ + torch.zeros(1, device=first_dims.device, dtype=first_dims.dtype), + torch.cumsum(first_dims * logical_last_dim, dim=0), + ] + ) + offsets = tensor_offsets.tolist() + first_dims_list = first_dims.tolist() + for i in range(num_tensors): + shape.append((first_dims_list[i], logical_last_dim)) + else: + offsets = [ + i * logical_first_dim * logical_last_dim // num_tensors + for i in range(num_tensors + 1) + ] + for i in range(num_tensors): + shape.append((logical_first_dim // num_tensors, logical_last_dim)) + + # Calculate logical shape based + logical_shape = (logical_first_dim, logical_last_dim) + + no_quantization = quantizer is None + + rowwise_usage = quantizer.rowwise_usage if not no_quantization else True + columnwise_usage = quantizer.columnwise_usage if not no_quantization else False + + # Calculate total elements across all tensors + total_elements = logical_first_dim * logical_last_dim + + data = None + columnwise_data = None + scale_inv = None + columnwise_scale_inv = None + amax = None + columnwise_amax = None + scale = None + scale_inv_offsets = None + columnwise_scale_inv_offsets = None + if no_quantization: + assert dtype is not None, "dtype must be provided for unquantized GroupedTensor" + if rowwise_usage: + # Allocate rowwise data buffer (1D flattened, uint8) + data = torch.empty(total_elements, dtype=dtype, device=device) + + if columnwise_usage: + # Allocate columnwise data buffer (1D flattened, uint8) + columnwise_data = torch.empty(total_elements, dtype=dtype, device=device) + elif quantizer._get_compatible_recipe().mxfp8(): + if rowwise_usage: + # Allocate rowwise data buffer (1D flattened, uint8) + data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Scale inverse buffer for MXFP8 - complex shape based on block scaling + # For grouped tensors, we need to calculate scale_inv size for all tensors + total_scale_elements = 0 + scale_inv_offsets = [0] + for i, s in enumerate(shape): + scale_inv_shape = quantizer.get_scale_shape(s, False) + scale_elements = math.prod(scale_inv_shape) + total_scale_elements += scale_elements + if i < num_tensors - 1: + scale_inv_offsets.append(total_scale_elements) + scale_inv = torch.empty(total_scale_elements, dtype=torch.uint8, device=device) + + if columnwise_usage: + # Allocate columnwise data buffer (1D flattened, uint8) + columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Columnwise scale inverse buffer + 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) + columnwise_scale_elements = math.prod(scale_inv_shape) + total_columnwise_scale_elements += columnwise_scale_elements + if i < num_tensors - 1: + columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) + columnwise_scale_inv = torch.empty( + total_columnwise_scale_elements, dtype=torch.uint8, device=device + ) + elif quantizer._get_compatible_recipe().delayed(): + if rowwise_usage: + # Allocate rowwise data buffer (1D flattened, uint8) + data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Scale inverse - one per tensor + scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 + scale_inv_offsets = list(range(num_tensors)) + + if columnwise_usage: + # Allocate columnwise data buffer (1D flattened, uint8) + columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Columnwise scale inverse - one per tensor + columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 + columnwise_scale_inv_offsets = list(range(num_tensors)) + + # Amax buffer for delayed scaling - one per tensor + amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + elif quantizer._get_compatible_recipe().nvfp4(): + + if rowwise_usage: + # Allocate rowwise data buffer (1D flattened, uint8, but FP4 packs 2 values per byte) + data = torch.empty((total_elements) // 2, dtype=torch.uint8, device=device) + # Scale inverse buffer for NVFP4 - complex shape based on block scaling + # For simplicity, calculate total scale elements needed + total_scale_elements = 0 + scale_inv_offsets = [0] + for i, s in enumerate(shape): + scale_inv_shape = quantizer.get_scale_shape(s, False) + total_scale_elements += math.prod(scale_inv_shape) + if i < num_tensors - 1: + 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) + + if columnwise_usage: + # Allocate columnwise data buffer (1D flattened, uint8, FP4 packed) + columnwise_data = torch.empty( + (total_elements) // 2, dtype=torch.uint8, device=device + ) + # Columnwise scale inverse buffer + total_columnwise_scale_elements = 0 + columnwise_scale_inv_offsets = [0] + for i, s in enumerate(shape): + columnwise_scale_inv_shape = quantizer.get_scale_shape(s, True) + total_columnwise_scale_elements += math.prod(columnwise_scale_inv_shape) + if i < num_tensors - 1: + columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) + 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: + # Allocate rowwise data buffer (1D flattened, uint8) + data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Scale inverse - size depends on block configuration + # For simplicity, calculate total scale elements needed + total_scale_elements = 0 + scale_inv_offsets = [0] + for i, s in enumerate(shape): + scale_inv_shape = quantizer.get_scale_shape(s, False) + total_scale_elements += math.prod(scale_inv_shape) + if i < num_tensors - 1: + scale_inv_offsets.append(total_scale_elements) + scale_inv = torch.empty(total_scale_elements, dtype=torch.float32, device=device) + + if columnwise_usage: + # Allocate columnwise data buffer (1D flattened, uint8) + columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Columnwise scale inverse + total_columnwise_scale_elements = 0 + columnwise_scale_inv_offsets = [0] + for i, s in enumerate(shape): + columnwise_scale_inv_shape = quantizer.get_scale_shape(s, True) + total_columnwise_scale_elements += math.prod(columnwise_scale_inv_shape) + if i < num_tensors - 1: + columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) + columnwise_scale_inv = torch.empty( + total_columnwise_scale_elements, dtype=torch.float32, device=device + ) + elif quantizer._get_compatible_recipe().float8_current_scaling(): + # Current scaling - per-tensor scaling computed on the fly + if rowwise_usage: + # Allocate rowwise data buffer (1D flattened, uint8) + data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Scale inverse - one per tensor + scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 + scale_inv_offsets = list(range(num_tensors)) + + if columnwise_usage: + # Allocate columnwise data buffer (1D flattened, uint8) + columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) + # Columnwise scale inverse - one per tensor + columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 + columnwise_scale_inv_offsets = list(range(num_tensors)) + + # Scale and amax buffers for current scaling - one per tensor + scale = torch.empty(num_tensors, dtype=torch.float32, device=device) + amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + else: + raise ValueError(f"Unsupported quantizer for GroupedTensor: {quantizer}") + + grouped_tensor = GroupedTensor( + num_tensors=num_tensors, + shape=shape, + dtype=dtype, + quantizer=quantizer, + data=data, + columnwise_data=columnwise_data, + scale_inv=scale_inv, + columnwise_scale_inv=columnwise_scale_inv, + amax=amax, + columnwise_amax=columnwise_amax, + scale=scale, + first_dims=first_dims, + last_dims=last_dims, + tensor_offsets=tensor_offsets, + offsets=offsets, + scale_inv_offsets=scale_inv_offsets, + columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, + logical_shape=logical_shape, + ) + + grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() + return grouped_tensor + + def split_into_quantized_tensors( + self, + ) -> List[Union[QuantizedTensorStorage, torch.Tensor]]: + """ + Split the GroupedTensor into a list of `num_tensors` + quantized tensors based on the quantizer. No additional memory allocation is performed, + so the tensors returned are the same as the ones used to create the GroupedTensor. + + If quantizer is None, returns normal torch tensors. + If quantizer.internal is True, returns QuantizedTensorStorage. + Otherwise, returns QuantizedTensor. + + TODO(ksivaman): Block cases where any dims are varying. This is needed only + to expose the weights as separate parameters. + """ + + result = [] + + no_quantization = self.quantizer is None + + # Case 1: No quantization - return regular torch tensors + if no_quantization: + for i in range(self.num_tensors): + # Get tensor shape + tensor_shape = self.shape[i] + + # Get tensor data slice + if self.offsets is not None: + start_offset = self.offsets[i] + numel = tensor_shape[0] * tensor_shape[1] + end_offset = start_offset + numel + + if self.has_data(): + tensor_data = self.data[start_offset:end_offset].view(tensor_shape) + result.append(tensor_data) + elif self.has_columnwise_data(): + tensor_data = self.columnwise_data[start_offset:end_offset].view( + tensor_shape + ) + result.append(tensor_data) + else: + raise RuntimeError("GroupedTensor has no data to split") + else: + # All same shape case + numel = tensor_shape[0] * tensor_shape[1] + start_offset = i * numel + end_offset = start_offset + numel + + if self.has_data(): + tensor_data = self.data[start_offset:end_offset].view(tensor_shape) + result.append(tensor_data) + elif self.has_columnwise_data(): + tensor_data = self.columnwise_data[start_offset:end_offset].view( + tensor_shape + ) + result.append(tensor_data) + else: + raise RuntimeError("GroupedTensor has no data to split") + + return result + + # Case 2: Quantized tensors + recipe = self.quantizer._get_compatible_recipe() + + for i in range(self.num_tensors): + # Get tensor shape + tensor_shape = self.shape[i] + numel = tensor_shape[0] * tensor_shape[1] + + # Get data offsets + if self.offsets is not None: + data_start = self.offsets[i] + data_end = data_start + numel + else: + # All same shape + data_start = i * numel + data_end = data_start + numel + + # Special shape handling for NVFP4. + nvfp4 = self.quantizer._get_compatible_recipe().nvfp4() + if nvfp4: + data_start = data_start // 2 + data_end = data_end // 2 + + # Extract rowwise and columnwise data + rowwise_data = None + columnwise_data = None + + if self.has_data(): + if nvfp4: + rowwise_tensor_shape = self.quantizer.convert_shape_for_fp4(tensor_shape) + else: + rowwise_tensor_shape = tensor_shape + rowwise_data = self.data[data_start:data_end].view(rowwise_tensor_shape) + + if self.has_columnwise_data(): + columnwise_tensor_shape = self.quantizer.get_columnwise_shape(tensor_shape) + if nvfp4: + columnwise_tensor_shape = self.quantizer.convert_shape_for_fp4( + columnwise_tensor_shape + ) + columnwise_data = self.columnwise_data[data_start:data_end].view( + columnwise_tensor_shape + ) + + # MXFP8 format + if recipe.mxfp8(): + # Extract scale_inv data + rowwise_scale_inv = None + columnwise_scale_inv = None + + if self.scale_inv is not None and self.scale_inv_offsets is not None: + scale_start = self.scale_inv_offsets[i] + if i < self.num_tensors - 1: + scale_end = self.scale_inv_offsets[i + 1] + else: + scale_end = self.scale_inv.numel() + + # Calculate expected scale shape for MXFP8 + scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + rowwise_scale_inv = self.scale_inv[scale_start:scale_end].view(scale_shape) + + if ( + self.columnwise_scale_inv is not None + and self.columnwise_scale_inv_offsets is not None + ): + cscale_start = self.columnwise_scale_inv_offsets[i] + if i < self.num_tensors - 1: + cscale_end = self.columnwise_scale_inv_offsets[i + 1] + else: + cscale_end = self.columnwise_scale_inv.numel() + + cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( + cscale_shape + ) + + if self.quantizer.internal: + mxfp8_tensor_class = MXFP8TensorStorage + else: + mxfp8_tensor_class = MXFP8Tensor + tensor = mxfp8_tensor_class( + shape=tensor_shape, + dtype=self.dtype, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=self.quantizer.dtype, + quantizer=self.quantizer, + with_gemm_swizzled_scales=self.quantizer.optimize_for_gemm, + ) + result.append(tensor) + + # Delayed scaling or current scaling (both use Float8TensorStorage) + elif recipe.delayed() or recipe.float8_current_scaling(): + # Scale inverse - one per tensor + scale_inv = None + if self.scale_inv is not None: + scale_inv = self.scale_inv[i : i + 1] + + if self.quantizer.internal: + float8_tensor_class = Float8TensorStorage + else: + float8_tensor_class = Float8Tensor + + tensor = float8_tensor_class( + shape=tensor_shape, + dtype=self.dtype, + data=rowwise_data, + fp8_scale_inv=scale_inv, + fp8_dtype=self.quantizer.dtype, + quantizer=self.quantizer, + data_transpose=columnwise_data, + ) + result.append(tensor) + + # Float8 block scaling + elif recipe.float8_block_scaling(): + # Extract scale_inv data + rowwise_scale_inv = None + columnwise_scale_inv = None + + if self.scale_inv is not None and self.scale_inv_offsets is not None: + scale_start = self.scale_inv_offsets[i] + if i < self.num_tensors - 1: + scale_end = self.scale_inv_offsets[i + 1] + else: + scale_end = self.scale_inv.numel() + + # Get scale shape from quantizer + scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + rowwise_scale_inv = self.scale_inv[scale_start:scale_end].view(scale_shape) + + if ( + self.columnwise_scale_inv is not None + and self.columnwise_scale_inv_offsets is not None + ): + cscale_start = self.columnwise_scale_inv_offsets[i] + if i < self.num_tensors - 1: + cscale_end = self.columnwise_scale_inv_offsets[i + 1] + else: + cscale_end = self.columnwise_scale_inv.numel() + + # Get columnwise scale shape from quantizer + cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( + cscale_shape + ) + + # Compute is_2D_scaled and data_format from quantizer attributes + is_2D_scaled = self.quantizer.block_scaling_dim == 2 + + if self.quantizer.internal: + float8_blockwise_q_tensor_class = Float8BlockwiseQTensorStorage + else: + float8_blockwise_q_tensor_class = Float8BlockwiseQTensor + + tensor = float8_blockwise_q_tensor_class( + shape=tensor_shape, + dtype=self.dtype, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + fp8_dtype=self.quantizer.dtype, + quantizer=self.quantizer, + is_2D_scaled=is_2D_scaled, + ) + result.append(tensor) + + # NVFP4 format + elif recipe.nvfp4(): + # Extract scale_inv data + rowwise_scale_inv = None + columnwise_scale_inv = None + amax_rowwise = None + amax_columnwise = None + + if self.scale_inv is not None and self.scale_inv_offsets is not None: + scale_start = self.scale_inv_offsets[i] + if i < self.num_tensors - 1: + scale_end = self.scale_inv_offsets[i + 1] + else: + scale_end = self.scale_inv.numel() + + # Get scale shape from quantizer + scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + rowwise_scale_inv = self.scale_inv[scale_start:scale_end].view(scale_shape) + + if ( + self.columnwise_scale_inv is not None + and self.columnwise_scale_inv_offsets is not None + ): + cscale_start = self.columnwise_scale_inv_offsets[i] + if i < self.num_tensors - 1: + cscale_end = self.columnwise_scale_inv_offsets[i + 1] + else: + cscale_end = self.columnwise_scale_inv.numel() + + # Get columnwise scale shape from quantizer + cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( + cscale_shape + ) + + # Extract amax - one per tensor + if self.amax is not None: + amax_rowwise = self.amax[i : i + 1] + + if self.columnwise_amax is not None: + amax_columnwise = self.columnwise_amax[i : i + 1] + + if self.quantizer.internal: + nvfp4_tensor_class = NVFP4TensorStorage + else: + nvfp4_tensor_class = NVFP4Tensor + + tensor = nvfp4_tensor_class( + shape=tensor_shape, + dtype=self.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, + fp4_dtype=self.quantizer.dtype, + quantizer=self.quantizer, + with_gemm_swizzled_scales=self.quantizer.optimize_for_gemm, + ) + result.append(tensor) + + else: + raise ValueError(f"Unsupported quantization recipe: {recipe}") + + return result + + @staticmethod + def create_and_quantize( + tensors: int, + quantizer: None | Quantizer, + *, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + noop_flag: Optional[torch.Tensor] = None, + ) -> Tuple[QuantizedTensorStorage, ...]: + """ + Quantize given tensors into quantized tensors with underlying + storage allocated in a GroupedTensor. + """ + + grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=len(tensors), + shape=[t.shape for t in tensors], + quantizer=quantizer, + device=device, + dtype=dtype, + ) + + grouped_tensor.quantize(tensors, noop_flag=noop_flag) + + return grouped_tensor + + def quantize( + self, + tensors: List[torch.Tensor], + noop_flag: Optional[torch.Tensor] = None, + ) -> Tuple[QuantizedTensorStorage, ...]: + """ + Quantize the GroupedTensor inplace. + """ + + quantized_tensors = self.split_into_quantized_tensors() + for i in range(self.num_tensors): + self.quantizer.update_quantized(tensors[i], quantized_tensors[i], noop_flag=noop_flag) + return quantized_tensors diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 1951731c75..5c8510488f 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -111,6 +111,24 @@ def clear(self): if t is not None: t.data = _empty_tensor() + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data buffers from another MXFP8TensorStorage.""" + if not isinstance(src, MXFP8TensorStorage): + raise TypeError("copy_from_storage expects MXFP8TensorStorage") + if self._fp8_dtype != src._fp8_dtype: + raise RuntimeError("FP8 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") + + def _copy_optional(dst: Optional[torch.Tensor], src_tensor: Optional[torch.Tensor]): + if dst is not None and src_tensor is not None: + dst.copy_(src_tensor) + + _copy_optional(self._rowwise_data, src._rowwise_data) + _copy_optional(self._columnwise_data, src._columnwise_data) + _copy_optional(self._rowwise_scale_inv, src._rowwise_scale_inv) + _copy_optional(self._columnwise_scale_inv, src._columnwise_scale_inv) + def get_metadata(self) -> Dict[str, Any]: """Get this tensor's metadata.""" return { diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index b064d711ce..8be23d0c19 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -136,6 +136,26 @@ def clear(self): if t is not None: t.data = _empty_tensor() + def copy_from_storage(self, src: QuantizedTensorStorage) -> None: + """Copy data buffers from another NVFP4TensorStorage.""" + if not isinstance(src, NVFP4TensorStorage): + raise TypeError("copy_from_storage expects NVFP4TensorStorage") + if self._fp4_dtype != src._fp4_dtype: + 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") + + def _copy_optional(dst: Optional[torch.Tensor], src_tensor: Optional[torch.Tensor]): + if dst is not None and src_tensor is not None: + dst.copy_(src_tensor) + + _copy_optional(self._rowwise_data, src._rowwise_data) + _copy_optional(self._columnwise_data, src._columnwise_data) + _copy_optional(self._rowwise_scale_inv, src._rowwise_scale_inv) + _copy_optional(self._columnwise_scale_inv, src._columnwise_scale_inv) + _copy_optional(self._amax_rowwise, src._amax_rowwise) + _copy_optional(self._amax_columnwise, src._amax_columnwise) + def get_metadata(self) -> Dict[str, Any]: """Get this tensor's metadata.""" return { From 402ea54b49caac987941b7d3e79285cf61d98d8e Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 12 Feb 2026 02:35:39 +0530 Subject: [PATCH 204/521] [C] NVFP4 quantization for `GroupedTensor` (#2655) * NVFP4 GroupedQuantize Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Zhongbo Zhu Co-authored-by: Zhongbo Zhu * fix fp4 Signed-off-by: Zhongbo Zhu * Remove unnecessary file Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Zhongbo Zhu Co-authored-by: Zhongbo Zhu --- transformer_engine/common/CMakeLists.txt | 2 + .../graph_safe_group_hadamard_transform.cu | 586 +++++++ ...cast_col_hadamard_transform_cast_fusion.cu | 1513 +++++++++++++++++ .../transformer_engine/hadamard_transform.h | 34 + .../include/transformer_engine/multi_tensor.h | 11 + 5 files changed, 2146 insertions(+) create mode 100644 transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu create mode 100644 transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index efe958f844..f0968c62ee 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -173,10 +173,12 @@ list(APPEND transformer_engine_cuda_arch_specific_sources cast/cast.cu gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu + hadamard_transform/graph_safe_group_hadamard_transform.cu hadamard_transform/hadamard_transform.cu hadamard_transform/hadamard_transform_cast_fusion.cu hadamard_transform/group_hadamard_transform_cast_fusion.cu hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu + hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu multi_tensor/compute_scale.cu recipe/mxfp8_scaling.cu transpose/quantize_transpose_square_blockwise.cu diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu new file mode 100644 index 0000000000..986229aabf --- /dev/null +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -0,0 +1,586 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "common/common.h" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" +#include "hadamard_transform_utils.cuh" + +namespace transformer_engine { +namespace { + +constexpr int kMaxTensorsPerKernel = 64; +constexpr int kThreadsPerWarp = 32; + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + +__device__ __forceinline__ size_t get_current_tensor_id( + const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t* const __restrict__ offsets_ptr) { + if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { + const size_t current_row = current_offset / last_logical_dim; + const size_t rows_per_tensor = first_logical_dim / num_tensors; + return current_row / rows_per_tensor; + } else { + // upper_bound(offsets, current_offset) - 1 in range i in [0..num_tensors) + size_t low = 0; + size_t hi = num_tensors; // half-open [low, hi) + + while (low < hi) { + const size_t mid = low + (hi - low) / 2; + const size_t mid_offset = static_cast(offsets_ptr[mid]); + + if (mid_offset <= current_offset) { + low = mid + 1; + } else { + hi = mid; + } + } + + // low = first index where offsets[low] > current_offset (or low == num_tensors) + // id = low - 1, but need to evaluate if current_offset < offsets[0] + return (low == 0) ? 0 : (low - 1); + } +} + +template +__device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4], + IType* in_sh_ptr, uint32_t& local_pre_rht_amax_reg, + uint32_t& local_amax_reg, + uint32_t& local_amax_t_reg) { + uint32_t a_frag[4]; // A matrix fragment + uint32_t c_frag[4]; // Result fragment + + int warp_id = threadIdx.x / kThreadsPerWarp; + int local_rank = (threadIdx.x % kThreadsPerWarp); + + int ld_row_idx = local_rank % kHadamardDimension; + int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; + int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); + + uint32_t temp_amax_reg; + uint32_t temp_amax_t_reg; + + if (kReturnIdentityAmax) { + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + + mma_m16_n16_k16_b16_b16_b16_noacc( + a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag_i[0], b_frag_i[1], b_frag_i[2], + b_frag_i[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_reg); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(local_amax_reg) + : "r"(local_amax_reg), "r"(temp_amax_reg)); + } + + if (kReturnTransposedAmax) { + // TODO(Frank): This is not efficient, since we could directly load the + // matrix in transposed layout. + if (!kReturnIdentityAmax) { + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } + + matrix_transpose_m8_n8_b16_inplace(a_frag[0]); + matrix_transpose_m8_n8_b16_inplace(a_frag[1]); + matrix_transpose_m8_n8_b16_inplace(a_frag[2]); + matrix_transpose_m8_n8_b16_inplace(a_frag[3]); + + mma_m16_n16_k16_b16_b16_b16_noacc( + a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], + b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(local_amax_t_reg) + : "r"(local_amax_t_reg), "r"(temp_amax_t_reg)); + } + + if (kReturnPreRhtAmax) { + if (!kReturnIdentityAmax && !kReturnTransposedAmax) { + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } + + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(a_frag[0]) + : "r"(a_frag[0]), "r"(a_frag[1])); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(a_frag[2]) + : "r"(a_frag[2]), "r"(a_frag[3])); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(a_frag[0]) + : "r"(a_frag[0]), "r"(a_frag[2])); + asm volatile("max.xorsign.abs.bf16x2 %0, %1, %2;\n\t" + : "=r"(local_pre_rht_amax_reg) + : "r"(a_frag[0]), "r"(local_pre_rht_amax_reg)); + } +} + +template +__device__ __host__ constexpr int NextPowerOf2() { + static_assert(kN > 0, "kN must be > 0"); + // Round up to the next power of 2 by counting leading zeros. + return 1 << (32 - __builtin_clz(kN - 1)); +} + +template +__device__ __forceinline__ void ReduceMax(const float pre_rht_amax, const float identity_amax, + const float transpose_amax, float* staging_for_pre_rht, + float* staging_for_identity, float* staging_for_transpose, + float* output_pre_rht_amax_ptr, + float* output_identity_amax_ptr, + float* output_transpose_amax_ptr, const int warpid) { + // intra-warp reduction + constexpr int kWarpSize = 32; + int local_rank = threadIdx.x % 32; + float warp_pre_rht_amax = kReturnPreRhtAmax ? warp_reduce_max(pre_rht_amax) : 0.0f; + float warp_identity_amax = kReturnIdentityAmax ? warp_reduce_max(identity_amax) : 0.0f; + float warp_transpose_amax = + kReturnTransposedAmax ? warp_reduce_max(transpose_amax) : 0.0f; + + // inter-warp reduction + if (threadIdx.x % 32 == 0) { + if (kReturnPreRhtAmax) { + staging_for_pre_rht[warpid] = warp_pre_rht_amax; + } + if (kReturnIdentityAmax) { + staging_for_identity[warpid] = warp_identity_amax; + } + if (kReturnTransposedAmax) { + staging_for_transpose[warpid] = warp_transpose_amax; + } + } + __syncthreads(); + constexpr int kNumWarpsPow2 = NextPowerOf2(); + if (warpid == 0) { + if (kReturnIdentityAmax) { + float identity_accum = local_rank < kNumWarps ? staging_for_identity[local_rank] : 0.0f; + identity_accum = warp_reduce_max(identity_accum); + if (local_rank == 0) { + atomicMaxFloat(output_identity_amax_ptr, identity_accum); + } + } + } + if (warpid == 1) { + if (kReturnTransposedAmax) { + float transpose_accum = local_rank < kNumWarps ? staging_for_transpose[local_rank] : 0.0f; + transpose_accum = warp_reduce_max(transpose_accum); + if (local_rank == 0) { + atomicMaxFloat(output_transpose_amax_ptr, transpose_accum); + } + } + } + if (warpid == 2) { + if (kReturnPreRhtAmax) { + float pre_rht_accum = local_rank < kNumWarps ? staging_for_pre_rht[local_rank] : 0.0f; + pre_rht_accum = warp_reduce_max(pre_rht_accum); + if (local_rank == 0) { + atomicMaxFloat(output_pre_rht_amax_ptr, pre_rht_accum); + } + } + } +} + +__global__ void GraphSafeMultiZeroAmaxKernel(const size_t num_tensors, float* amax_rowwise_ptr, + float* amax_colwise_ptr) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + // Assign each thread a range for rowwise and colwise independently + if (amax_rowwise_ptr != nullptr) { + for (int i = tid; i < num_tensors; i += stride) { + amax_rowwise_ptr[i] = 0.f; + } + } + if (amax_colwise_ptr != nullptr) { + for (int i = tid; i < num_tensors; i += stride) { + amax_colwise_ptr[i] = 0.f; + } + } +} + +__global__ void GraphSafeMultiAmaxMemcpyD2DKernelPreRHT(const size_t num_tensors, + float* amax_rowwise_ptr, + float* amax_colwise_ptr) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int stride = blockDim.x * gridDim.x; + + if (amax_rowwise_ptr != nullptr && amax_colwise_ptr != nullptr) { + for (; tid < num_tensors; tid += stride) { + float* output_pre_rht_amax_ptr = amax_rowwise_ptr + tid; + float* output_transpose_amax_ptr = amax_colwise_ptr + tid; + *output_transpose_amax_ptr = *output_pre_rht_amax_ptr; + } + } +} + +template +__global__ void GraphSafeGroupHadamardAmaxTmaKernel( + const __grid_constant__ CUtensorMap tensor_map_input, uint16_t random_sign_mask, + uint16_t random_sign_mask_t, const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t* const __restrict__ offsets_ptr, const int64_t* const __restrict__ first_dims_ptr, + float* const __restrict__ amax_rowwise_ptr, float* const __restrict__ amax_colwise_ptr) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + + float* output_pre_rht_amax_ptr; + float* output_identity_amax_ptr = nullptr; + float* output_transpose_amax_ptr; + + // calculate the global offset to get tensor id + size_t global_offset = blockIdx.y * CHUNK_DIM_Y * last_logical_dim; + int tensor_id = get_current_tensor_id(shape_rep, num_tensors, global_offset, first_logical_dim, + last_logical_dim, offsets_ptr); + output_pre_rht_amax_ptr = static_cast(amax_rowwise_ptr) + tensor_id; + output_transpose_amax_ptr = static_cast(amax_colwise_ptr) + tensor_id; + + static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y && CHUNK_DIM_Y % BUFF_DIM_Y == 0); + static_assert(CHUNK_DIM_X >= BUFF_DIM_X && CHUNK_DIM_X % BUFF_DIM_X == 0); + + constexpr size_t STAGES_Y = CHUNK_DIM_Y / BUFF_DIM_Y; + constexpr size_t STAGES_X = CHUNK_DIM_X / BUFF_DIM_X; + + constexpr int kNumWarps = (THREADS_PER_CHUNK * THREADS_PER_Y) / kThreadsPerWarp; + + const int input_block_offset_Y = blockIdx.y * CHUNK_DIM_Y; + const int input_block_offset_X = blockIdx.x * CHUNK_DIM_X; + + extern __shared__ __align__(128) 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! + uint8_t* dshmem = reinterpret_cast((base_shmem_ptr + 127) & ~127ULL); + + // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned + constexpr size_t in_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType); + IType* in_sh_0 = reinterpret_cast(dshmem); + dshmem += in_buff_size; + IType* in_sh_1 = reinterpret_cast(dshmem); + dshmem += in_buff_size; + + IType* in_shs[2] = {in_sh_0, in_sh_1}; + + constexpr int shmem_buff_size = BUFF_DIM_X * BUFF_DIM_Y * sizeof(IType); + + const bool is_master_thread = (threadIdx.x == 0 && threadIdx.y == 0); + + // Initialize shared memory barrier with the number of threads participating in the barrier. +#pragma nv_diag_suppress static_var_with_dynamic_init + uint64_t* mbar = reinterpret_cast(dshmem); + dshmem += sizeof(uint64_t) * (STAGES_X * STAGES_Y); + + float* max_staging_identity = reinterpret_cast(dshmem); + dshmem += sizeof(float) * kNumWarps; + float* max_staging_transpose = reinterpret_cast(dshmem); + dshmem += sizeof(float) * kNumWarps; + float* max_staging_pre_rht = reinterpret_cast(dshmem); + dshmem += sizeof(float) * kNumWarps; + + initialize_barriers(mbar, + is_master_thread); + + copy_2d_to_shared(in_shs[0], reinterpret_cast(&tensor_map_input), + input_block_offset_X, input_block_offset_Y, shmem_buff_size, &mbar[0], + is_master_thread); + + uint32_t had_frag_i[4]; + uint32_t had_frag_t[4]; + get_hadamard_matrix_fragment( + had_frag_i, random_sign_mask, had_frag_t, random_sign_mask_t); + + float local_pre_rht_amax = 0.0; + float local_amax = 0.0; + float local_amax_t = 0.0; + uint32_t local_pre_rht_amax_reg = *reinterpret_cast(&local_pre_rht_amax); + uint32_t local_amax_reg = *reinterpret_cast(&local_amax); + uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t); + + for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) { + for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) { + int stage = STAGES_X * stage_y + stage_x; + + const int next_stage = stage + 1; + const int next_stage_x = stage_x + 1 == STAGES_X ? 0 : stage_x + 1; + const int next_stage_y = stage_x + 1 == STAGES_X ? stage_y + 1 : stage_y; + + if (next_stage < STAGES_X * STAGES_Y) { + const int input_global_offset_Y = input_block_offset_Y + next_stage_y * BUFF_DIM_Y; + const int input_global_offset_X = input_block_offset_X + next_stage_x * BUFF_DIM_X; + + copy_2d_to_shared(in_shs[next_stage % 2], // ping-pong + reinterpret_cast(&tensor_map_input), input_global_offset_X, + input_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); + + const size_t compute_stage_x_num = + BUFF_DIM_X / (kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)); + const size_t compute_stage_y_num = BUFF_DIM_Y / (kHadamardDimension * THREADS_PER_Y); + + const size_t in_row_stride = BUFF_DIM_X; + + IType* in_sh_ptr = in_shs[stage % 2]; + +#pragma unroll + for (size_t compute_stage_y = 0; compute_stage_y < compute_stage_y_num; compute_stage_y++) { + const int row_idx_offset = (compute_stage_y * kHadamardDimension * THREADS_PER_Y + + threadIdx.y * kHadamardDimension); + const int in_row_offset = row_idx_offset * in_row_stride; + +#pragma unroll + for (size_t compute_stage_x = 0; compute_stage_x < compute_stage_x_num; compute_stage_x++) { + ComputeKernel( + had_frag_i, had_frag_t, + in_sh_ptr + in_row_offset + + (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)), + local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); + } + + // Ensure all threads have finished their computation before new data over-writes the shared + // memory. + __syncthreads(); + } + } + } + + const int warpid = (threadIdx.x + threadIdx.y * blockDim.x) / kThreadsPerWarp; + + if constexpr (kReturnPreRhtAmax) { + unpack_max_of_packed_bf16(local_pre_rht_amax_reg, local_pre_rht_amax); + } + if constexpr (kReturnIdentityAmax) { + unpack_max_of_packed_bf16(local_amax_reg, local_amax); + } + if constexpr (kReturnTransposedAmax) { + unpack_max_of_packed_bf16(local_amax_t_reg, local_amax_t); + } + + ReduceMax( + local_pre_rht_amax, local_amax, local_amax_t, max_staging_pre_rht, max_staging_identity, + max_staging_transpose, output_pre_rht_amax_ptr, output_identity_amax_ptr, + output_transpose_amax_ptr, warpid); + + destroy_barriers(mbar, is_master_thread); +#else + NVTE_DEVICE_ERROR("Kernel is only supported on SM 10.0+."); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +} // namespace + +// broadcast_pre_rht_amax: when it's true, hadamard transform will be disabled +// if at this time, the amax buffers for output expects both amax_rowwise and amax_colwise +// then call MultiAmaxMemcpyD2DKernelPreRHT to D2D copy the amax values +void group_hadamard_transform_amax_graph_safe(const GroupedTensor* input, GroupedTensor* output, + uint16_t random_sign_mask, + uint16_t random_sign_mask_t, + bool broadcast_pre_rht_amax, cudaStream_t stream) { + NVTE_API_CALL(group_hadamard_transform_amax_graph_safe); +#if CUDA_VERSION >= 12080 + + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must be same."); + NVTE_CHECK(input->has_data(), "Cannot quantize tensor without rowwise data."); + + checkCuDriverContext(stream); + + bool all_return_pre_rht_amax = output->has_data(); + // there is no rowwise RHT transform in current recipe + bool all_return_identity_amax = false; + bool all_return_transposed_amax = output->has_columnwise_data(); + + NVTE_CHECK(all_return_pre_rht_amax || all_return_identity_amax || all_return_transposed_amax, + "At least one of return_pre_rht_amax, return_identity_amax, or return_transposed_amax " + "must be true"); + + if (broadcast_pre_rht_amax) { + NVTE_CHECK(all_return_pre_rht_amax, + "broadcast_pre_rht_amax is only supported when we compute pre-RHT amax"); + // if all_return_identity_amax and all_return_transposed_amax both are false, there is no need to broadcast anything + broadcast_pre_rht_amax &= (all_return_identity_amax || all_return_transposed_amax); + } + + const size_t num_tensors = input->num_tensors; + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t last_logical_dim = input->logical_shape.data[1]; + // const size_t elts_total = first_logical_dim * last_logical_dim; + NVTE_CHECK(first_logical_dim % 128 == 0, + "First dimension of a grouped tensor should be divisible by 128."); + NVTE_CHECK(last_logical_dim % 128 == 0, + "Last dimension of a grouped tensor should be divisible by 128."); + + float* const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); + float* const amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr); + + const int64_t* const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); + const int64_t* const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); + // const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + + // some sanity checks + if (all_return_pre_rht_amax) { + NVTE_CHECK(amax_rowwise_ptr != nullptr, "Amax rowwise pointer should not be nullptr."); + } + if (all_return_transposed_amax) { + NVTE_CHECK(amax_colwise_ptr != nullptr, "Amax columnwise pointer should not be nullptr."); + } + + // Multi zero out multiple amaxes if needed + dim3 block_setup_amax(kMaxTensorsPerKernel); + dim3 grid_setup_amax(1); + GraphSafeMultiZeroAmaxKernel<<>>( + num_tensors, amax_rowwise_ptr, amax_colwise_ptr); + NVTE_CHECK_CUDA(cudaGetLastError()); + + using IType = bf16; + constexpr int kHadamardDimension = 16; + + // four (1x4) 64x64 sub-tiles for ping-pong overlap + constexpr uint64_t kChunkBlockXSmall = 256; + constexpr uint64_t kChunkBlockYSmall = 64; + constexpr uint64_t kBuffDimX = 64; + constexpr uint64_t kBuffDimY = 64; + + alignas(64) CUtensorMap tensor_map_input{}; + + create_2D_tensor_map( + /*tensorMap=*/tensor_map_input, + /*tensor=*/input->data, + /*globalY=*/first_logical_dim, + /*globalX=*/last_logical_dim, + /*shmemY=*/kBuffDimY, + /*shmemX=*/kBuffDimX, + /*stride_elems=*/last_logical_dim, + /*offset_elems=*/0, + /*type_num_bits=*/sizeof(IType) * 8, + /*swizzle=*/CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B); + + constexpr uint64_t kThreadBlockX = 4; + constexpr uint64_t kThreadBlockY = 1; + constexpr uint64_t kNumWarps = kThreadBlockX * kThreadBlockY; + + dim3 block(kThreadBlockX * kThreadsPerWarp, kThreadBlockY); + dim3 grid(DIVUP(last_logical_dim, kChunkBlockXSmall), + DIVUP(first_logical_dim, kChunkBlockYSmall)); + + ShapeRepresentation shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + if (output->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else if (output->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (output->varying_both_dims()) { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + shape_rep == ShapeRepresentation::VARYING_FIRST_DIM); + + NVTE_CHECK(is_const_last_dim, + "Currently we only support const last dimension for graph safe hadamard transform."); + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + (all_return_transposed_amax && !broadcast_pre_rht_amax), kReturnTransposedAmax, + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + (all_return_identity_amax && !broadcast_pre_rht_amax), kReturnIdentityAmax, + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + all_return_pre_rht_amax, kReturnPreRhtAmax, + + // *2 for ping-pong + size_t in_sh_size = kBuffDimX * kBuffDimY * 2 * sizeof(IType); + size_t mbar_size = sizeof(uint64_t) * (kChunkBlockXSmall / kBuffDimX) * + (kChunkBlockYSmall / kBuffDimY); + size_t shmem_bytes = in_sh_size + mbar_size + kNumWarps * sizeof(float) * 3; + // Add padding in case shmem ptr is not aligned to 128 bytes. + shmem_bytes = (shmem_bytes + 128); + + auto kernel = GraphSafeGroupHadamardAmaxTmaKernel< + IType, kHadamardDimension, kChunkBlockYSmall, kChunkBlockXSmall, kBuffDimY, + kBuffDimX, kThreadBlockX * kThreadsPerWarp, kThreadBlockY, kReturnPreRhtAmax, + kReturnIdentityAmax, kReturnTransposedAmax>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + shmem_bytes); + + kernel<<>>( + tensor_map_input, random_sign_mask, random_sign_mask_t, shape_rep, num_tensors, + first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, + amax_rowwise_ptr, amax_colwise_ptr); + if (broadcast_pre_rht_amax) { + GraphSafeMultiAmaxMemcpyD2DKernelPreRHT<<>>(num_tensors, amax_rowwise_ptr, + amax_colwise_ptr); + }))); + + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("Hadamard transform requires CUDA 12.8+, but compile-time CUDA version is ", + CUDA_VERSION); +#endif // CUDA_VERSION >= 12080 +} + +} // namespace transformer_engine + +void nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input, + NVTEGroupedTensor output, int random_sign_mask, + int random_sign_mask_t, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_hadamard_transform_amax_graph_safe); + using namespace transformer_engine; + + GroupedTensor* input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor* output_tensor = convertNVTEGroupedTensorCheck(output); + + if (input_tensor->num_tensors == 0) { + return; + } + + // Call the group tensor Hadamard transform amax implementation. + group_hadamard_transform_amax_graph_safe( + input_tensor, output_tensor, static_cast(random_sign_mask), + static_cast(random_sign_mask_t), false, stream); +} + +// Grouped-tensor amax without doing hadamard transform +void nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_amax_graph_safe); + using namespace transformer_engine; + + GroupedTensor* input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor* output_tensor = convertNVTEGroupedTensorCheck(output); + + if (input_tensor->num_tensors == 0) { + return; + } + + group_hadamard_transform_amax_graph_safe(input_tensor, output_tensor, 0, 0, true, stream); +} 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 new file mode 100644 index 0000000000..030dddfce4 --- /dev/null +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -0,0 +1,1513 @@ +/************************************************************************* + * 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 "common/common.h" +#include "common/util/cuda_runtime.h" +#include "common/util/curanddx.hpp" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" +#include "customized_pipeline.cuh" +#include "cutlass/arch/barrier.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/cluster_launch.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cutlass/fast_math.h" +#include "cutlass/float8.h" +#include "cutlass/float_subbyte.h" +#include "cutlass/gemm/collective/builders/sm100_common.inl" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/platform/platform.h" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/print_error.hpp" + +namespace transformer_engine { +namespace detail { +namespace { + +using namespace cute; + +// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor +using cute::Tensor; + +constexpr int kMaxTensorsPerKernel = 64; +constexpr int kNVFP4BlockSize = 16; + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + +__device__ __forceinline__ size_t get_current_tensor_id( + const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr) { + if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { + const size_t current_row = current_offset / last_logical_dim; + const size_t rows_per_tensor = first_logical_dim / num_tensors; + return current_row / rows_per_tensor; + } else { + // upper_bound(offsets, current_offset) - 1 in range i in [0..num_tensors) + size_t low = 0; + size_t hi = num_tensors; // half-open [low, hi) + + while (low < hi) { + const size_t mid = low + (hi - low) / 2; + const size_t mid_offset = static_cast(offsets_ptr[mid]); + + if (mid_offset <= current_offset) { + low = mid + 1; + } else { + hi = mid; + } + } + + // low = first index where offsets[low] > current_offset (or low == num_tensors) + // id = low - 1, but need to evaluate if current_offset < offsets[0] + return (low == 0) ? 0 : (low - 1); + } +} + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverterBase( + cutlass::Array const &input, cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + auto output_ptr = reinterpret_cast(&output); + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + "cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\n" + "cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\n" + "}" + : "=h"(output_ptr[0]), "=h"(output_ptr[1]) + : "f"(input[0]), "f"(input[1]), "f"(input[2]), "f"(input[3]), "f"(input[4]), "f"(input[5]), + "f"(input[6]), "f"(input[7]), "r"(rbits[0]), "r"(rbits[1])); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return output; +} + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverter( + cutlass::Array const &input, cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + cutlass::Array *result_ptr = + reinterpret_cast *>(&output); + cutlass::Array const *source_ptr = + reinterpret_cast const *>(&input); + cutlass::Array const *rbits_ptr = + reinterpret_cast const *>(&rbits); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < 2; i++) { + result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]); + } + return output; +} + +template +struct SharedStorage { + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_; + using AtomThrShapeMNK = cute::Shape<_1, _1, _1>; + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage; + + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = + cutlass::detail::CustomizedPipelineTmaUmmaAsync, + AtomThrShapeMNK>; + using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage; + using SchedPipeline = cutlass::PipelineCLCFetchAsync; + using SchedPipelineStorage = typename SchedPipeline::SharedStorage; + using SchedThrottlePipeline = cutlass::PipelineAsync; + using SchedThrottlePipelineStorage = typename SchedThrottlePipeline::SharedStorage; + + struct TensorStorage : cute::aligned_struct<128, _1> { + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + } tensors; + + alignas(16) AccumulatorPipelineStorage accumulator; + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) cute::uint64_t tma_barrier[1]; + alignas(16) SchedPipelineStorage sched; + alignas(16) SchedThrottlePipelineStorage sched_throttle; + alignas(16) int32_t atomic_tile_id[SchedulerPipelineStageCount_]; + alignas(16) float global_a_amax[kMaxTensorsPerKernel]; + alignas(16) float global_d_amax[kMaxTensorsPerKernel]; + uint32_t atomic_tile_counter[SchedulerPipelineStageCount_]; + uint32_t tmem_base_ptr; +}; + +// Main RHT GEMM kernel entry -- highly templated for flexible architecture/config support +template +__launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_graph_safe( + MShape M, NShape packed_N, KShape K, ClusterShape cluster_shape, ClusterTileShape cluster_tile, + TA const *A, AStride dA, ASmemLayout sAlayout, CUTE_GRID_CONSTANT TmaLoadA const tma_load_a, + TB const *B, BStride dB, BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, + TQA *QA, QAStride dQA, TSFA *SFA, TSFALayout sfa_layout, TQA *QA_COLWISE, TSFA *SFA_COLWISE, + float *amax_rowwise, float *amax_colwise, const int64_t *offsets, const int64_t *first_dims, + size_t num_tensors, ShapeRepresentation shape_rep, uint32_t *tile_scheduler_workspace, + TiledMMA mma, const size_t *rng_state) { + using namespace cute; + + // Abort immediately if compilation is not supported + constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; + if constexpr (!is_blackwell_arch) { + NVTE_DEVICE_ERROR( + "group_row_col_rht_gemm_device_graph_safe is only supported on Blackwell " + "with architecture-specific compilation. " + "Try recompiling with sm_100a or similar."); + return; + } + static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, + "group_row_col_rht_gemm_device_graph_safe must generate row-wise " + "and/or column-wise output."); +#if !defined(CUTLASS_ARCH_CLC_ENABLED) + CUTLASS_NOT_IMPLEMENTED(); + return; +#endif + + using X = Underscore; + // Accumulator data type for main computation + using ElementAccumulator = float; + static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); + static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; + static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; + static constexpr bool kEnableRowQuant = kEnableRowQuant_; + static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_; + static constexpr bool kUseFastMath = kUseFastMath_; + + // Constant for RHT tensor processing (tile size etc) + static int constexpr RhtTensorSize = 16; + + // Get the total number of tokens to process + // Note that here M is the hidden size, which is the last logical dimension of the input tensor x + // The kernel is designed in column major, so M is the hidden size + size_t sum_token_dims = offsets[num_tensors] / M; + + // Transaction bytes for TMA transfer on RHT tensor blocks + static int constexpr kTmaRhtTensorTransactionBytes = + cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; + + // Mainloop pipeline stage calculation, vectorization parameters for scaling factors + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + static int constexpr SFVecSize = 16; + // Swizzle output layout for scaling factor arrays + using SwizzledSFALayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + + // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling + using MainloopPipeline = + cutlass::detail::CustomizedPipelineTmaUmmaAsync; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + using SchedPipeline = cutlass::PipelineCLCFetchAsync; + using SchedPipelineState = typename SchedPipeline::PipelineState; + using SchedThrottlePipeline = cutlass::PipelineAsync; + using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState; + + static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, "ClusterShape must be Shape<_1,_1,_1>"); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static int constexpr VectorSize = RhtTensorSize; + + // Compile-time safety: static shapes required for shared memory layouts + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + // CUTE_STATIC_ASSERT(is_static::value); + + auto cluster_size = size<0>(cluster_shape); + auto mainloop_tiler = Shape<_128, _16, _128>{}; + auto epilogue_tiler = Shape<_128, _128, _128>{}; + + static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler); + + struct TileScheduler { + uint32_t tiles_in_m = 0; + uint32_t tiles_in_n = 0; + uint32_t linear_idx = 0; + uint32_t next_linear_idx = 0; + uint32_t start_idx = 0; + uint32_t tile_m_idx = 0; + uint32_t tile_n_idx = 0; + int k_tile_max = 0; + uint32_t *atomic_tile_index_; + uint32_t *smem_tile_counter; + uint32_t atomic_offset; + cutlass::FastDivmodU64 divmod_tiles_in_m; + + CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, + uint32_t *atomic_tile_index, uint32_t *smem_tile_counter) + : tiles_in_m(tiles_m), + tiles_in_n(tiles_n), + linear_idx(blockIdx.x), + next_linear_idx(blockIdx.x), + start_idx(blockIdx.x), + k_tile_max(kmax), + atomic_tile_index_(atomic_tile_index), + smem_tile_counter(smem_tile_counter), + atomic_offset(gridDim.x), + divmod_tiles_in_m(uint64_t(tiles_m)) { + update_tile_idx(); + } + CUTLASS_DEVICE void update_tile_idx() { + uint64_t q, r; + divmod_tiles_in_m(q, r, uint64_t(linear_idx)); + tile_m_idx = static_cast(r); + tile_n_idx = static_cast(q) * uint32_t(k_tile_max); + } + CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; } + CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; } + CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; } + + CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; } + + CUTLASS_DEVICE bool is_valid() const { + return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()), + cute::make_coord(tiles_in_m, tiles_in_n)); + } + + CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; } + + CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; } + + // Fetch a new tile_id using atomics. + CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) { + uint32_t tile_id_counter = 0; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p atom.global.add.u32 %0, [%1], 1; \n\t" + "}" + : "=r"(tile_id_counter) + : "l"(atomic_tile_index_), "r"(pred)); + + return tile_id_counter; + } + + CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_consumer_state) { + sched_pipeline.consumer_wait(sched_pipeline_consumer_state); + next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()]; + cutlass::arch::fence_view_async_shared(); + sched_pipeline.consumer_release(sched_pipeline_consumer_state); + return; + } + + CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_producer_state) { + uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state); + // Wait for clcID buffer to become empty with a flipped phase + sched_pipeline.producer_acquire(sched_pipeline_producer_state); + auto is_leading_thread = cute::elect_one_sync(); + uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset; + uint32_t smem_addr = + cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]); + if (is_leading_thread) { + cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0); + } + + ++sched_pipeline_producer_state; + return sched_pipeline_producer_state; + } + + CUTLASS_DEVICE auto update_work_tile_info() { + linear_idx = next_linear_idx; + update_tile_idx(); + return; + } + }; + + // Allocate and alias shared memory to the kernel's shared storage type + extern __shared__ char shared_memory[]; + using SharedStorage = + SharedStorage; + SharedStorage &shared_storage = *reinterpret_cast(shared_memory); + + // Compute the number of tiles in M and N after tiling and assign scheduler + uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); + uint32_t tiles_in_n = uint32_t(size(ceil_div(sum_token_dims, size<2>(epilogue_tiler)))); + + TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace, + shared_storage.atomic_tile_counter); + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + + // Shapes for accumulated tiles in mainloop and epilogue + auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{}); + auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{}); + + // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended + auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); + auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); + + // Number of threads assigned for various epilogue roles depending on quantization settings + static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; + static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; + static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0; + static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0; + static int constexpr NumSchedThreads = 32; + static int constexpr NumMainloopLoadThreads = 32; + static int constexpr NumEpilogueThreads = + NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + NumMmaThreadCount + NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + // warp assignment + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_sched_warp = (warp_idx == 2); + bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); + bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + + MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, + cluster_shape, cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = + cutlass::make_producer_start_state(); + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + using AccumulatorPipelineInitBarriers = cute::bool_constant; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = + cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_col_quant_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = + size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; + accumulator_pipeline_params.initializing_warp = 1; + AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, accumulator_pipeline_params, + cluster_shape, AccumulatorPipelineInitBarriers{}, + cute::true_type{}); // Delay mask calculation + typename SchedPipeline::Params sched_pipeline_params; + if (is_sched_warp) { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer; + } else { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer; + } + sched_pipeline_params.producer_blockid = 0; + sched_pipeline_params.producer_arv_count = 1; + sched_pipeline_params.consumer_arv_count = + NumSchedThreads + + cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); + sched_pipeline_params.transaction_bytes = sizeof(uint32_t); + sched_pipeline_params.initializing_warp = 3; + SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape); + SchedPipelineState sched_pipeline_consumer_state; + SchedPipelineState sched_pipeline_producer_state = + cutlass::make_producer_start_state(); + + typename SchedThrottlePipeline::Params sched_throttle_pipeline_params; + if (is_dma_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer; + } + if (is_sched_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer; + } + sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; + sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; + sched_throttle_pipeline_params.dst_blockid = 0; + sched_throttle_pipeline_params.initializing_warp = 4; + + SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle, + sched_throttle_pipeline_params); + SchedThrottlePipelineState sched_pipeline_throttle_consumer_state; + SchedThrottlePipelineState sched_pipeline_throttle_producer_state = + cutlass::make_producer_start_state(); + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + + // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer + if (is_dma_warp) { + // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access). + cutlass::arch::warpgroup_reg_dealloc<32>(); + // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory. + Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N)); + Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); + + // Partition tensors for tiling according to the mainloop and cluster tilers. + Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gB_nk = + local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) + + // Shared memory tensors for pipeline + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + // Determine warp/tile positioning + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Partition global to local fragments for A and B + Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) + Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = + tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = + tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); + + auto [tBgB, tBsB] = + tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + if constexpr (kEnableRHTColQuant) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], + kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), + tBsB(_, 0)); + } + } + + do { + // is_first_wave indicates whether this scheduler wave is the first among a group. + bool is_first_wave = scheduler.is_first_wave(); + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_, scheduler.tile_m(), _); + int k_tile = 0; + + sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state); + sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state); + ++sched_pipeline_throttle_producer_state; + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { + int k_tile_idx_n = scheduler.tile_n_base() + k_tile; + ++k_tile; + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType *tma_barrier = + mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + if (cute::elect_one_sync()) { + copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), + tAsA(_, write_stage)); + } + } + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + // scheduler.advance(); + } while (scheduler.is_valid()); + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform. + cutlass::arch::warpgroup_reg_dealloc<32>(); + if constexpr (kEnableRHTColQuant) { + // Setup shared memory fragments for A and B tiles. + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Allocate "fragments" -- these are actually umma smem descriptors + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + // Wait until the B (Hadamard) tensor copy is complete + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + + auto barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_, _, _, read_stage); + auto tCrB_nk = tCrB(_, _, 0, 0); + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) { + int accumulator_k_block = + accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; + int tCrA_k_block = k_block * EpilogueUnrollFactor; + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < EpilogueUnrollFactor; i++) { + auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i); + gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; + } + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); + barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + } else if (is_sched_warp) { + // Scheduler warp manages tile assignment and pipeline progress for warps + cutlass::arch::warpgroup_reg_dealloc<32>(); + do { + sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state); + sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state); + ++sched_pipeline_throttle_consumer_state; + sched_pipeline_producer_state = + scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } else if (is_epilogue_col_quant_warp) { + // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage, + // and writing result tensors/scales to global memory. + cutlass::arch::warpgroup_reg_alloc<192>(); + if constexpr (kEnableRHTColQuant) { + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + auto acc_epilogue_pipelined_shape = + append(acc_shape_epilogue, Int{}); + auto bulk_tmem_epilogue_layout = make_layout( + acc_epilogue_pipelined_shape, + make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler))); + auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); + + // Use 256-bit fragments for aligned bulk stores + static int constexpr FragmentSize = 256 / sizeof_bits_v; + + // Wait for TMEM allocation for this pipeline to finish + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; + // g2s load all global_d_amax + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) { + shared_storage.global_d_amax[g] = __ldg(reinterpret_cast(amax_colwise + g)); + } + + size_t rng_seed = 0; + size_t rng_offset = 0; + // Setup RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + // TODO(zhongbo): double check the logic here + int group_idx = get_current_tensor_id(shape_rep, num_tensors, + (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, + packed_N, M, offsets); + + // Determine quantization scale factor layouts/output splits for this group + TSFDLayout sfd_layout; + int cur_N = static_cast(first_dims[group_idx]); + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // Build output tensors for columns and their quant scales + // TODO(zhongbo): double check the logic here + Tensor mD = make_tensor(cute::subbyte_iterator(reinterpret_cast( + reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)), + make_shape(M, cur_N), DStride{}); // (M,packed_N) + Tensor gD_mn = + local_tile(mD, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + // for every tensor [x, y] row major, x y both a multiple of 128 + // both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3 + Tensor mSFD = make_tensor( + make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) + + offsets[group_idx] / kNVFP4BlockSize)), + sfd_layout); + Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + + // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); + auto tiled_r2g = + make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); + auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); + + cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float c_global_amax_val = shared_storage.global_d_amax[group_idx]; + float global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + float global_decode_scale = 1.0f / global_encode_scale; + + // Scaling factor for fast math path + float global_encode_scale_multiplier = 1.0f; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + + do { + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); + ++k_tile) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + // TODO(zhongbo): double check the logic here + int cur_group_idx = get_current_tensor_id(shape_rep, num_tensors, + global_tile_n_offset * M, packed_N, M, offsets); + + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + c_global_amax_val = shared_storage.global_d_amax[group_idx]; + // update amax + global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + // TODO(zhongbo): double check the logic here + cur_N = first_dims[group_idx]; + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = + tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = + make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // update tensor + mD = make_tensor(cute::subbyte_iterator(reinterpret_cast( + reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)), + make_shape(M, cur_N), DStride{}); + gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + mSFD = make_tensor( + make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) + + offsets[group_idx] / kNVFP4BlockSize)), + sfd_layout); + gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + } + int group_start_offset = offsets[group_idx] / M; + int local_tile_n_idx = + (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler); + Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx); + + Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx); + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); + Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + + Tensor tTR_rAcc = + make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDrD = make_tensor(shape(tDgD)); + Tensor tTR_rAcc_frag = + recast>(coalesce(tTR_rAcc)); + Tensor tDrD_frag = recast>(coalesce(tDrD)); + + Tensor src = thr_r2g.retile_S(tDrD); + Tensor dst = thr_r2g.retile_D(tDgD); + + Tensor tDgSFD_view = make_tensor( + tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); + Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); + Tensor tDrSFD = make_tensor(shape(tDgSFD)); + + static int constexpr NumVecs = size(tDgD) / VectorSize; + Tensor tD_rRowSFD_frg = recast>(tDrSFD); + + // Compute amax and quantization scales for this tile + cutlass::maximum_absolute_value_reduction, + true> + amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // Copy from TMEM to registers + copy(tiled_t2r, tDtAcc, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with + // unfused kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); + } + + auto compute_frgs = reinterpret_cast *>( + tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); + } + + if constexpr (kUseFastMath) { + // Fast math: multiply with precomputed reciprocal + pvscales = cutlass::multiplies>{}( + vec_maxs, global_encode_scale_multiplier); + } else { + // Accurate math: perform division + pvscales = + cutlass::divides>{}(vec_maxs, fp4_max); + pvscales = cutlass::multiplies>{}( + pvscales, global_encode_scale); + } + auto pvscales_cvted = + cutlass::NumericArrayConverter{}(pvscales); + + tD_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}( + tD_rRowSFD_frg(_0{})); + auto qpvscale_scaled = cutlass::multiplies>{}( + qpvscale_ups, global_decode_scale); + cutlass::Array acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = + cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides>{}( + 1.0, qpvscale_scaled); + } + + // Prepare stochastic rounding random state if enabled + uint4 random_uint4 = uint4{0, 0, 0, 0}; + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + // "Prefetch" a stochastic rounding state for the first tile + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + // Apply round/quantize to each fragment, with or without stochastic rounding + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], cutlass::platform::numeric_limits::max()); + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale)); + } + } + + // Write quantized FP4 tile and dequant scale to gmem + copy(tiled_r2g, src, dst); + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + } else if (is_epilogue_row_quant_warp) { + // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage. + cutlass::arch::warpgroup_reg_alloc<136>(); + if constexpr (kEnableRowQuant) { + using S2RVectorType = uint128_t; + + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % 256; + size_t rng_seed = 0; + size_t rng_offset = 0; + // g2s load all global_a_amax for all groups/tensors + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) { + shared_storage.global_a_amax[g] = __ldg(reinterpret_cast(amax_rowwise + g)); + } + // RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + // Input/output tensors/partitions for row quant warp + Tensor mQA = + make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA)); + Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); + + Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _), + Step<_1, X, _1>{}); // (BLK_M,BLK_N) + // Swizzled shared memory A tile, with layout + Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>( + coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) + + // Set up layouts for partitioning – tile-by-warp, with vector granularity + using S2RWarpLayout = Layout>; + using WarpGroupLayout = Layout>; + using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); + using S2RValLayout = Layout, _1>>; + using S2RAtomA = Copy_Atom; + using R2GAtomQA = Copy_Atom; + using R2GAtomSFA = Copy_Atom; + auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{}); + + auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); + auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); + auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx); + Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) + + // Allocate temporary register tensors for copying quantization => output + Tensor tQArA = make_tensor_like( + make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) + Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn); + Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); + + Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); + Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); + + // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8 + // in order to go over the reserved named barrier count. + constexpr int row_quant_barrier_id = 2; + cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); + + int group_idx = get_current_tensor_id(shape_rep, num_tensors, + (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, + packed_N, M, offsets); + float a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + + float global_decode_scale = 1.0f / global_encode_scale; + float global_encode_scale_multiplier = 1.0f; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + auto sfa_converter = cutlass::NumericConverter{}; + do { + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + int cur_group_idx = get_current_tensor_id(shape_rep, num_tensors, + global_tile_n_offset * M, packed_N, M, offsets); + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Update group quantization parameters/scaling + global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + if constexpr (kUseFastMath) { + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } + } + + auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state); + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); + cutlass::arch::fence_view_async_shared(); + mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state); + ++mainloop_pipe_consumer_state; + ++k_tile; + + // static int constexpr NumVecs = size(tQArA) / VectorSize; + cutlass::maximum_absolute_value_reduction, + true> + amax_reduction; + auto compute_frgs = reinterpret_cast *>(tQArA.data()); + auto output_frgs = + reinterpret_cast *>(raw_pointer_cast(tQArQA.data())); + Tensor amax = + make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); + Tensor pvscales = make_tensor_like(amax); + transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + + tiles_in_m * tiles_in_n * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) { + auto amax_view = group_modes<1, rank(amax)>(amax); + auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales); + auto compute_frgs_up = + cutlass::NumericArrayConverter{}( + compute_frgs[v]); + amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up); + if constexpr (kUseFastMath) { + // Fast math: multiply with precomputed reciprocal + pvscales_view(_0{}, v) = cutlass::multiplies{}( + amax_view(_0{}, v), global_encode_scale_multiplier); + } else { + // Accurate math: perform division + pvscales_view(_0{}, v) = + cutlass::divides{}(amax_view(_0{}, v), fp4_max); + pvscales_view(_0{}, v) = cutlass::multiplies{}( + pvscales_view(_0{}, v), global_encode_scale); + } + filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v)); + auto qpvscale_ups = + cutlass::NumericConverter{}(filter(tQArSFA)(v)); + auto qpvscale_scaled = + cutlass::multiplies{}(qpvscale_ups, global_decode_scale); + ElementAccumulator acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = + cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + } + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales, cutlass::platform::numeric_limits::max()); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = + cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale)); + } + } + copy(tiled_r2g_QA, tQArQA, tQAgQA_mn); + copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn)); + } + // scheduler.advance(); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + + } else { + cutlass::arch::warpgroup_reg_dealloc<32>(); + } +} // NOLINT(readability/fn_size) + +template +void group_row_col_rht_gemm_ntt_w_sfc_graph_safe( + int packed_sequence_length, int hidden_size, size_t num_tensors, ShapeRepresentation shape_rep, + TA const *A, TB const *B, TQA *QA, TSFA *SFA, TQA *QA_COLWISE, TSFA *SFA_COLWISE, + float *amax_rowwise, float *amax_colwise, const int64_t *offsets, const int64_t *first_dims, + const size_t *rng_state, uint32_t *tile_scheduler_workspace, uint32_t sm_count, + cudaStream_t stream, int k_tile_size = 1024) { + using namespace cute; + static int constexpr SFVecSize = 16; + static int constexpr RhtTensorSize = 16; + + static_assert(RhtTensorSize == 16, "RhtTensorSize must be 16"); + using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0), + make_stride(make_stride(_0{}, _1{}), 0))); + using LinearSFDLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)), + make_stride(0, make_stride(_0{}, _1{})))); + + using SwizzledSFALayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFALayout = decltype(tile_to_shape( + SwizzledSFALayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{})); + using SwizzledSFDLayout = decltype(tile_to_shape( + SwizzledSFDLayoutAtom{}, make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{})); + + using SFALayout = cute::conditional_t; + using SFDLayout = cute::conditional_t; + SFALayout sfa_layout; + SFDLayout sfd_layout; + + if constexpr (kEnableSwizzleSFOutput) { + sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{}, + make_shape(hidden_size, packed_sequence_length), Step<_1, _2>{}); + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, + make_shape(hidden_size, packed_sequence_length), Step<_2, _1>{}); + } else { + sfa_layout = make_layout( + make_shape(make_shape(Int{}, hidden_size / SFVecSize), packed_sequence_length), + make_stride(make_stride(_0{}, _1{}), hidden_size / SFVecSize)); + sfd_layout = make_layout( + make_shape(hidden_size, make_shape(Int{}, packed_sequence_length / SFVecSize)), + make_stride(packed_sequence_length / SFVecSize, make_stride(_0{}, _1{}))); + } + + // Define shapes (dynamic) + auto M = hidden_size; + auto N = packed_sequence_length; + Tensor tensorA = make_tensor(A, make_shape(hidden_size, packed_sequence_length), LayoutLeft{}); + Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{}); + Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, packed_sequence_length), LayoutLeft{}); + Tensor tensorSFA = make_tensor(SFA, sfa_layout); + + // Define strides (from tensors) + auto dA = stride(tensorA); // (dM,dK) + auto dB = stride(tensorB); // (dN,dK) + auto dD = LayoutRight{}; // (dM,dN) + auto dQA = stride(tensorQA); // (dM,dK) + using ClusterShape = Shape<_1, _1, _1>; + auto cluster_shape = ClusterShape{}; + auto cluster_tile_shape = Shape<_128, Int, Int>{}; + auto cluster_tile_mainloop = Shape<_128, Int, _128>{}; + + // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles + static int constexpr EpilogueUnrollFactor = + size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape); + // Construct the MMA + auto mma = make_tiled_mma( + SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape), + UMMA::Major::MN, UMMA::Major::MN>{}, + Layout>{}); + + // Assert that the TiledMMA uses all CTAs in the CGA. + CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma)); + CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma))); + + // Determine the A and B shapes + auto mma_shape_B = + partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape))); + + using TiledMma = decltype(mma); + using AtomThrID = typename TiledMma::AtomThrID; + + using SmemShape_M = decltype(shape_div( + shape<0>(cluster_tile_shape), + shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{})))); + using SmemShape_N = decltype(shape_div( + shape<1>(cluster_tile_shape), + shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{})))); + using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape)); + + using SmemLayoutAtomB = + decltype(cutlass::gemm::collective::detail::sm100_smem_selector()); + + auto mma_shape_A = partition_shape_A( + mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop))); + using SmemShape_M_A = + decltype(shape_div(shape<0>(cluster_tile_mainloop), + shape_div(shape<0>(cluster_tile_mainloop), + size<0>(cluster_tile_mainloop) / size(AtomThrID{})))); + using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop)); + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>()); + + static uint32_t constexpr TotalTmemRows = 128; + static uint32_t constexpr Sm100TmemCapacityColumns = 512; + static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns; + static uint32_t constexpr AccumulatorPipelineStageCount = + TotalTmem / (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape)); + + // Define the smem layouts (static) + // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory + constexpr int SchedulerPipelineStageCount = 4; + static int constexpr MainloopPipelineBytes = sizeof( + typename cutlass::detail::CustomizedPipelineTmaUmmaAsync<1, Shape<_1, _1, _1>, + Shape<_1, _1, _1>>::SharedStorage); + + static int constexpr SchedulerWorkspaceBytes = sizeof(int) * SchedulerPipelineStageCount; + static int constexpr SchedulerThrottlePipelineBytes = + sizeof(typename cutlass::PipelineAsync::SharedStorage); + static int constexpr SchedulerPipelineBytes = + sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage); + + static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier); + static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB); + static int constexpr AccPipelineBytes = sizeof( + typename cutlass::PipelineUmmaAsync>::SharedStorage); + static int constexpr TmemBasePtrsBytes = sizeof(uint32_t); + static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes + static int constexpr kBytesPerStage = + cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes; + static int constexpr kReservedBytes = SchedulerWorkspaceBytes + SchedulerThrottlePipelineBytes + + SchedulerPipelineBytes + TmemBasePtrsBytes + + TmemDeallocBytes + BTensorBytes + + AccPipelineBytes; // Reserve for barriers and other uses + static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage; + auto sP = Int{}; // SMEM pipelines + + auto sA = UMMA::tile_to_mma_shape(SmemLayoutAtomA{}, append(mma_shape_A, sP), + Step<_2, _1, _3>{}); // (MMA,MMA_M,MMA_K,PIPE) + auto sB = UMMA::tile_to_mma_shape(SmemLayoutAtomB{}, + append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1) + auto sD = Layout<_1>{}; // XXX Dummy + + auto tma_load_a = + make_tma_copy_A_sm100(SM90_TMA_LOAD{}, tensorA, sA(_, _, _, 0), cluster_tile_mainloop, mma); + auto tma_load_b = + make_tma_copy_B_sm100(SM90_TMA_LOAD{}, tensorB, sB(_, _, _, 0), cluster_tile_shape, mma); + + // Assert checks on tile sizes -- no predication + assert(M % size<0>(cluster_tile_shape) == 0); + assert(N % size<1>(cluster_tile_shape) == 0); + + dim3 dimBlock(512); + dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape)); + dim3 dimGrid(sm_count, 1, 1); + + int smem_size = sizeof( + SharedStorage); + + auto *kernel_ptr = &group_row_col_rht_gemm_device_graph_safe< + decltype(M), decltype(N), decltype(k_tile_size), decltype(cluster_shape), + decltype(cluster_tile_shape), TA, decltype(dA), decltype(sA), decltype(tma_load_a), TB, + decltype(dB), decltype(sB), decltype(tma_load_b), TD, decltype(dD), decltype(sD), TSFD, + decltype(sfd_layout), TQA, decltype(dQA), TSFA, decltype(sfa_layout), decltype(mma), + AccumulatorPipelineStageCount, SchedulerPipelineStageCount, kEnableStochasticRounding, + kEnableRHTColQuant, kEnableRowQuant, kEnableSwizzleSFOutput, kUseFastMath>; + + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + // Set workspace and set to zero + NVTE_CHECK_CUDA(cudaMemsetAsync(reinterpret_cast(tile_scheduler_workspace), 0, + sizeof(uint32_t), stream)); + + // Launch kernel + cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream}; + cutlass::Status status = cutlass::launch_kernel_on_cluster( + params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, A, dA, + sA, tma_load_a, B, dB, sB, tma_load_b, QA, dQA, SFA, sfa_layout, QA_COLWISE, SFA_COLWISE, + amax_rowwise, amax_colwise, offsets, first_dims, num_tensors, shape_rep, + tile_scheduler_workspace, mma, rng_state); + NVTE_CHECK_CUDA(cudaGetLastError()); + NVTE_CHECK(status == cutlass::Status::kSuccess, "Kernel launch failed."); +} + +} // namespace +} // namespace detail + +void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, + GroupedTensor *output, + const Tensor &hadamard_matrix_, + QuantizationConfig &quant_config, + Tensor &quant_workspace, cudaStream_t stream) { + NVTE_API_CALL(group_hadamard_transform_cast_fusion_graph_safe); + + using transformer_engine::detail::kMaxTensorsPerKernel; + using transformer_engine::detail::ShapeRepresentation; + + void *input_base_ptr = reinterpret_cast(input->data.dptr); + // TODO(zhongbo): add input sanity checks here + + bool all_has_row_quant = output->has_data(); + bool all_has_col_quant = output->has_columnwise_data(); + + // Stochastic rounding config + const bool use_stochastic_rounding = quant_config.stochastic_rounding; + const size_t *rng_state = nullptr; + if (use_stochastic_rounding) { + NVTE_CHECK(quant_config.rng_state != nullptr, + "Enabled stochastic rounding without providing RNG state"); + 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}, + "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); + rng_state = reinterpret_cast(rng_state_tensor.data.dptr); + } + + uint32_t *tile_scheduler_workspace = nullptr; + NVTE_CHECK(quant_workspace.data.dptr != nullptr, "Quantization workspace must be provided."); + NVTE_CHECK(quant_workspace.data.buffer_size_bytes() >= sizeof(uint32_t), + "Quantization workspace must be at least 4 bytes."); + tile_scheduler_workspace = reinterpret_cast(quant_workspace.data.dptr); + + // Template arguments + using TA = cute::bfloat16_t; + using TB = cute::bfloat16_t; + using TD = cutlass::float_e2m1_t; + using TSFD = cutlass::float_ue4m3_t; + using TQA = TD; + using TSFA = TSFD; + + checkCuDriverContext(stream); + + // Check Hadamard matrix + constexpr int kHadamardDimension = 16; + + NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16, + "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(), "."); + const size_t hadamard_dimension = hadamard_matrix.shape[0]; + + const size_t num_tensors = input->num_tensors; + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t last_logical_dim = input->logical_shape.data[1]; + // const size_t elts_total = first_logical_dim * last_logical_dim; + NVTE_CHECK(first_logical_dim % 128 == 0, + "First dimension of a grouped tensor should be divisible by 128."); + NVTE_CHECK(last_logical_dim % 128 == 0, + "Last dimension of a grouped tensor should be divisible by 128."); + NVTE_CHECK(num_tensors <= kMaxTensorsPerKernel, + "Number of tensors should be less than or equal to ", kMaxTensorsPerKernel); + + ShapeRepresentation shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + if (output->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (output->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else if (output->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (output->varying_both_dims()) { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + TQA *const rowwise_data_base_ptr = reinterpret_cast(output->data.dptr); + TSFA *const rowwise_scale_inv_base_ptr = reinterpret_cast(output->scale_inv.dptr); + TQA *const colwise_data_base_ptr = reinterpret_cast(output->columnwise_data.dptr); + TSFA *const colwise_scale_inv_base_ptr = + reinterpret_cast(output->columnwise_scale_inv.dptr); + float *const amax_rowwise_base_ptr = reinterpret_cast(output->amax.dptr); + float *const amax_colwise_base_ptr = reinterpret_cast(output->columnwise_amax.dptr); + + const int64_t *const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); + // const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + + const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + shape_rep == ShapeRepresentation::VARYING_FIRST_DIM); + NVTE_CHECK(is_const_last_dim, + "Currently we only support const last dimension for graph safe hadamard transform."); + + auto sm_count = transformer_engine::cuda::sm_count(); + + int k_tile_size = 1024; + + const bool use_swizzle_sf_output = false; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_stochastic_rounding, kEnableStochasticRounding, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + all_has_col_quant, kEnableRhtColQuant, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + all_has_row_quant, kEnableRowQuant, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_swizzle_sf_output, kEnableSwizzleSFOutput, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + + if constexpr (kEnableRhtColQuant || kEnableRowQuant) { + detail::group_row_col_rht_gemm_ntt_w_sfc_graph_safe< + kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, + kEnableSwizzleSFOutput, TA, TB, TQA, TSFA, TD, TSFD, kUseFastMath>( + /*packed_sequence_length=*/first_logical_dim, + /*hidden_size=*/last_logical_dim, + /*num_tensors=*/num_tensors, + /*shape_rep=*/shape_rep, + /*A=*/reinterpret_cast(input_base_ptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*QA=*/reinterpret_cast(rowwise_data_base_ptr), + /*SFA=*/reinterpret_cast(rowwise_scale_inv_base_ptr), + /*QA_COLWISE=*/reinterpret_cast(colwise_data_base_ptr), + /*SFA_COLWISE=*/reinterpret_cast(colwise_scale_inv_base_ptr), + /*amax_rowwise=*/reinterpret_cast(amax_rowwise_base_ptr), + /*amax_colwise=*/reinterpret_cast(amax_colwise_base_ptr), + /*offsets=*/offsets_ptr, + /*first_dims=*/first_dims_ptr, + /*rng_state=*/rng_state, + /*tile_scheduler_workspace=*/tile_scheduler_workspace, + /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size); + } else { + NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", + kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, ")."); + } + + ););););); +} + +} // namespace transformer_engine + +void nvte_group_hadamard_transform_cast_fusion_graph_safe( + const NVTEGroupedTensor input, NVTEGroupedTensor output, const NVTETensor hadamard_matrix, + const NVTEQuantizationConfig quant_config, NVTETensor quant_workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_hadamard_transform_cast_fusion_graph_safe); + using namespace transformer_engine; + + GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); + GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); + + Tensor *quant_workspace_tensor = convertNVTETensorCheck(quant_workspace); + + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + + if (input_tensor->num_tensors == 0) { + return; + } + + // Call the multi-tensor Hadamard transform amax implementation. + group_hadamard_transform_cast_fusion_graph_safe( + input_tensor, output_tensor, *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, + *quant_workspace_tensor, stream); +} diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index 13103cc388..bee939f0cd 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -86,6 +86,24 @@ void nvte_group_hadamard_transform_amax(const NVTETensor input, NVTETensor* outp int random_sign_mask, int random_sign_mask_t, cudaStream_t stream); +/*! \brief Grouped-tensor amax with Hadamard transform (graph safe, device-managed grouping). + * + * This function is experimental and the API is not stable. + * + * This API assumes that the split info (grouping of tensors) is on device and unknown to the host; + * therefore, this is a graph safe API and the grouped-tensor argument is passed as a single device structure. + * + * \param[in] input NVTEGroupedTensor representing grouped input tensors. + * \param[in,out] output NVTEGroupedTensor for output amax (row/col). Only the row-wise and + * column-wise amaxes are updated. + * \param[in] random_sign_mask 16-bit sign mask for RHT. + * \param[in] random_sign_mask_t 16-bit sign mask for transposed RHT. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_hadamard_transform_amax_graph_safe(const NVTEGroupedTensor input, + NVTEGroupedTensor output, int random_sign_mask, + int random_sign_mask_t, cudaStream_t stream); + /*! * \brief Perform the grouped-tensor columnwise Hadamard transform cast fusion operation. * @@ -124,6 +142,22 @@ void nvte_group_hadamard_transform_cast_fusion(const NVTETensor input, NVTETenso const NVTEQuantizationConfig quant_config, NVTETensor quant_workspace, cudaStream_t stream); +/*! + * \brief Perform the grouped-tensor Hadamard transform cast fusion operation in graph-safe mode. + * + * This function is experimental and the API is not stable. Group_ prefix means contiguous input concatenated. + * + * \param[in] input NVTEGroupedTensor representing grouped input tensors. + * \param[in,out] output NVTEGroupedTensor for output (row/column-wise quantized results). + * \param[in] hadamard_matrix Hadamard matrix to use for transformation. + * \param[in] quant_config Quantization configuration. + * \param[in] quant_workspace Workspace buffer. Must be at least 4 bytes. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_hadamard_transform_cast_fusion_graph_safe( + const NVTEGroupedTensor input, NVTEGroupedTensor output, const NVTETensor hadamard_matrix, + const NVTEQuantizationConfig quant_config, NVTETensor quant_workspace, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index 303801a88a..b5eadcf678 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -296,6 +296,17 @@ void nvte_multi_tensor_compute_scale_inv_e8m0_cuda(int chunk_size, NVTETensor ** void nvte_group_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, cudaStream_t stream); +/*! \brief Grouped-tensor amax without doing hadamard transform. + * + * This function is experimental and the API is not stable. + * + * \param[in] input NVTEGroupedTensor Input tensor. + * \param[in,out] output NVTEGroupedTensor Output tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_amax_graph_safe(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif From c4175fcad075f31043dd11e459db97316bf64839 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Wed, 11 Feb 2026 20:50:29 -0600 Subject: [PATCH 205/521] fix(build): Handle namespace packages for PyPI CUDA detection (#2580) fix: handle nvidia namespace packages where __file__ is None Signed-off-by: Santosh Bhavani Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/common/__init__.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index 02388d2e70..40933f17a9 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -245,11 +245,13 @@ def _nvidia_cudart_include_dir() -> str: return "" # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia" - # above doesn't through. However, they don't set "__file__" attribute. - if nvidia.__file__ is None: - return "" + # above doesn't throw. However, they don't set "__file__" attribute. + if nvidia.__file__ is not None: + nvidia_root = Path(nvidia.__file__).parent + else: + nvidia_root = Path(nvidia.__path__[0]) # namespace package - include_dir = Path(nvidia.__file__).parent / "cuda_runtime" + include_dir = nvidia_root / "cuda_runtime" return str(include_dir) if include_dir.exists() else "" From 93d51c82e6c0241c97394d0293b04ada4dbeade9 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Thu, 12 Feb 2026 03:54:24 +0100 Subject: [PATCH 206/521] [Common] Fuse pre-swizzling into grouped MXFP8 quantization kernel (#2630) * Added GEMM-ready preswizzling option 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: Kirthi Shankar Sivamani --- .../cast/mxfp8/group_quantize_mxfp8.cuh | 244 ++++++++++-------- transformer_engine/common/common.h | 7 + 2 files changed, 142 insertions(+), 109 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 7801a2064d..a29a09836e 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -21,6 +21,7 @@ #include "../../util/ptx.cuh" #include "../../utils.cuh" #include "../core/common.cuh" +#include "swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -231,7 +232,7 @@ __device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tenso template + bool COLWISE_SCALING, bool WITH_GEMM_SWIZZLED_SCALES> __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel( const __grid_constant__ CUtensorMap tensor_map_input_static, const __grid_constant__ CUtensorMap tensor_map_act_input_static, @@ -250,6 +251,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel using IType2 = typename ptx::FPx2; using OType2 = typename ptx::FPx2; + using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; + if constexpr (NO_ACTIVATIONS) { if (noop != nullptr && noop[0] == 1.0f) { return; @@ -475,8 +478,14 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; const size_t global_scales_offset_X = scales_offset_X_colwise; - const size_t scale_idx = - global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + + size_t scale_idx = 0; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = gemm_swizzled_scale_idx(global_scales_offset_X, global_scales_offset_Y, + DIVUP(rows, static_cast(128))); + } else { + scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + } scales_colwise[scale_idx] = biased_exponent; const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); @@ -602,7 +611,14 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const int stage_scales_offset_X = scales_offset_X_rowwise; - const int scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + + size_t scale_idx = 0; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = gemm_swizzled_scale_idx(stage_scales_offset_Y, stage_scales_offset_X, + DIVUP(cols, static_cast(128))); + } else { + scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + } scales_rowwise[scale_idx] = biased_exponent; const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); @@ -803,6 +819,8 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations const dim3 grid(blocks); const size_t block_size = THREADS_PER_CHUNK; + const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + // Logical shape of a tensor with varying all dims is [1, M*K] if (shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS) { NVTE_CHECK(first_logical_dim % 128 == 0, @@ -848,111 +866,119 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations input->dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( output->dtype(), OType, - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_act_input{}; - alignas(64) CUtensorMap tensor_map_output_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_colwise{}; - - constexpr size_t input_type_bit_size = TypeInfo::size; - constexpr size_t output_type_bit_size = TypeInfo::size; - - create_2D_tensor_map(tensor_map_input, input->data, first_logical_dim, last_logical_dim, - BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, input_type_bit_size); - - if constexpr (IS_DACT) { - create_2D_tensor_map(tensor_map_act_input, activations->data, first_logical_dim, - last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, - input_type_bit_size); - } - - if (use_rowwise_scaling) { - create_2D_tensor_map(tensor_map_output_rowwise, output->data, first_logical_dim, - last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, - output_type_bit_size); - } - - if (use_colwise_scaling) { - create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, - first_logical_dim, last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, - last_logical_dim, 0, output_type_bit_size); - } - - constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; - constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; - constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; - constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; - constexpr size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); - - constexpr size_t elt_input_mem = buff_size_aligned_in; - constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); - constexpr size_t in_mem = elt_input_mem + act_input_mem; - - const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); - const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); - const size_t out_mem = out_rowwise_mem + out_colwise_mem; - - const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - - auto kernel = group_quantize_mxfp8_kernel; - switch (scaling_type) { - case ScalingType::ROWWISE: { - kernel = group_quantize_mxfp8_kernel; - break; - } - case ScalingType::COLWISE: { - kernel = group_quantize_mxfp8_kernel; - break; - } - case ScalingType::BIDIMENSIONAL: { - kernel = group_quantize_mxfp8_kernel; - break; - } - } - - // Update tensor descriptors before launching the kernel - if (!is_single_tensor) { - const IType *const input_dptr = reinterpret_cast(input->data.dptr); - - const IType *const act_input_dptr = - IS_DACT ? reinterpret_cast(activations->data.dptr) : nullptr; - - OType *const output_rowwise_dptr = - use_rowwise_scaling ? reinterpret_cast(output->data.dptr) : nullptr; - - OType *const output_colwise_dptr = - use_colwise_scaling ? reinterpret_cast(output->columnwise_data.dptr) - : nullptr; - update_tma_descriptors<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, input_dptr, act_input_dptr, output_rowwise_dptr, - output_colwise_dptr, shape_rep, num_tensors, first_logical_dim, last_logical_dim, - offsets_ptr, first_dims_ptr, last_dims_ptr, use_rowwise_scaling, - use_colwise_scaling, IS_DACT); - } - - NVTE_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, shape_rep, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_rowwise_ptr, - scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr); - - if constexpr (IS_DBIAS) { - common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); - } - - NVTE_CHECK_CUDA(cudaGetLastError());); // 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_act_input{}; + alignas(64) CUtensorMap tensor_map_output_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_colwise{}; + + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; + + create_2D_tensor_map(tensor_map_input, input->data, first_logical_dim, + last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, + input_type_bit_size); + + if constexpr (IS_DACT) { + create_2D_tensor_map(tensor_map_act_input, activations->data, first_logical_dim, + last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, + input_type_bit_size); + } + + if (use_rowwise_scaling) { + create_2D_tensor_map(tensor_map_output_rowwise, output->data, first_logical_dim, + last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, + output_type_bit_size); + } + + if (use_colwise_scaling) { + create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, + first_logical_dim, last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, + last_logical_dim, 0, output_type_bit_size); + } + + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; + constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); + const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); + const size_t out_mem = out_rowwise_mem + out_colwise_mem; + + const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; + + auto kernel = + group_quantize_mxfp8_kernel; + switch (scaling_type) { + case ScalingType::ROWWISE: { + kernel = + group_quantize_mxfp8_kernel; + break; + } + case ScalingType::COLWISE: { + kernel = + group_quantize_mxfp8_kernel; + break; + } + case ScalingType::BIDIMENSIONAL: { + kernel = + group_quantize_mxfp8_kernel; + break; + } + } + + // Update tensor descriptors before launching the kernel + if (!is_single_tensor) { + const IType *const input_dptr = reinterpret_cast(input->data.dptr); + + const IType *const act_input_dptr = + IS_DACT ? reinterpret_cast(activations->data.dptr) : nullptr; + + OType *const output_rowwise_dptr = + use_rowwise_scaling ? reinterpret_cast(output->data.dptr) : nullptr; + + OType *const output_colwise_dptr = + use_colwise_scaling ? reinterpret_cast(output->columnwise_data.dptr) + : nullptr; + update_tma_descriptors<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, input_dptr, act_input_dptr, output_rowwise_dptr, + output_colwise_dptr, shape_rep, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, + use_rowwise_scaling, use_colwise_scaling, IS_DACT); + } + + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, shape_rep, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_rowwise_ptr, + scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr); + + if constexpr (IS_DBIAS) { + common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); + } + + NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) } } // namespace mxfp8 diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 99a2985d5e..2d7f0e7e8c 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -335,6 +335,12 @@ struct GroupedTensor { NVTEGroupedTensor nvte_tensor; + /*! \brief Whether scaling factors are in format expected by GEMM + * + * Only meaningful for MXFP8 and NVFP4. + */ + bool with_gemm_swizzled_scales = false; + GroupedTensor(NVTEScalingMode scaling_mode, size_t num_tensors) : data(), columnwise_data(), @@ -401,6 +407,7 @@ struct GroupedTensor { num_tensors = 0; scaling_mode = NVTE_DELAYED_TENSOR_SCALING; nvte_tensor = 0; + with_gemm_swizzled_scales = false; } }; From 3774aa37e2aa777c0f3bf3dca752d64f961725e1 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Wed, 11 Feb 2026 21:56:49 -0800 Subject: [PATCH 207/521] [PyTorch] Add ops for MoE grouped MLP (#2664) * Add ops for MoE grouped MLP Signed-off-by: Tim Moon * Move testing utility functions to util submodule Signed-off-by: Tim Moon * Tweak docs Signed-off-by: Tim Moon * Change order of tensor compatibility checks in noop_cat Review suggestion from @ptrendx. Signed-off-by: Tim Moon * Add support for GLU interleaving in clamped SwiGLU Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon --- tests/pytorch/test_fusible_ops.py | 494 +++++++++++- tests/pytorch/utils.py | 47 +- transformer_engine/pytorch/module/_common.py | 25 +- .../pytorch/ops/basic/__init__.py | 4 +- .../pytorch/ops/basic/activation.py | 75 -- .../pytorch/ops/basic/grouped_linear.py | 702 ++++++++++++++++++ .../pytorch/ops/basic/swiglu.py | 498 +++++++++++++ 7 files changed, 1744 insertions(+), 101 deletions(-) create mode 100644 transformer_engine/pytorch/ops/basic/grouped_linear.py create mode 100644 transformer_engine/pytorch/ops/basic/swiglu.py diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index a23de29e02..2c1320e262 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -5,8 +5,10 @@ from __future__ import annotations from collections.abc import Iterable +import functools import io import math +import random from typing import Optional import pytest @@ -36,7 +38,14 @@ import transformer_engine_torch as tex # Import utility functions -from utils import dtype_tols, make_recipe, quantization_tols, reset_rng_states +from utils import ( + assert_close, + assert_close_grads, + dtype_tols, + make_recipe, + quantization_tols, + reset_rng_states, +) # Check for supported quantization schemes fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) @@ -107,6 +116,9 @@ def maybe_skip_quantization( @torch.no_grad() def make_reference_and_test_tensors( shape: int | Iterable[int], + *, + min: float = 0.0, + max: float = 1.0, quantization: Optional[str] = None, ref_dtype: torch.dtype = torch.float64, ref_device: torch.device = "cpu", @@ -127,7 +139,8 @@ def make_reference_and_test_tensors( """ # Random reference tensor - ref = torch.rand(shape, dtype=ref_dtype, device=ref_device) + ref = torch.empty(shape, dtype=ref_dtype, device=ref_device) + ref.uniform_(min, max) # Construct test tensor from reference tensor test = ref.to(device=test_device, dtype=test_dtype) @@ -1680,6 +1693,7 @@ def test_swiglu( quantization: Optional[str], quantize_forward: bool, quantize_backward: bool, + glu_interleave_size: Optional[int] = None, ): # Tensor dimensions @@ -1706,7 +1720,17 @@ def test_swiglu( ) # Plain PyTorch implementation - x1, x2 = x_ref.chunk(2, dim=-1) + x = x_ref + if glu_interleave_size is not None: + x = x.reshape( + *in_shape[:-1], + in_shape[-1] // (2 * glu_interleave_size), + 2, + glu_interleave_size, + ) + x = x.transpose(-3, -2) + x = x.reshape(in_shape) + x1, x2 = x.chunk(2, dim=-1) y_ref = torch.nn.functional.silu(x1) * x2 y_ref.backward(dy_ref) @@ -1714,7 +1738,7 @@ def test_swiglu( recipe = make_recipe(quantization) forward = te_ops.Sequential( te_ops.Quantize(forward=False, backward=quantize_backward), - te_ops.SwiGLU(), + te_ops.SwiGLU(glu_interleave_size=glu_interleave_size), te_ops.Quantize(forward=quantize_forward, backward=False), ) with te.autocast(enabled=quantized_compute, recipe=recipe): @@ -1727,10 +1751,19 @@ def test_swiglu( tols = quantization_tols(quantization) # Check results - y_test = y_test.to(dtype=torch.float64, device="cpu") - dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") - torch.testing.assert_close(y_test, y_ref, **tols) - torch.testing.assert_close(dx_test, x_ref.grad, **tols) + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + + def test_interleaved_swiglu(self): + """SwiGLU with block interleaved input format""" + self.test_swiglu( + out_shape=(32, 192), + dtype=torch.float32, + quantization=None, + quantize_forward=False, + quantize_backward=False, + glu_interleave_size=32, + ) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _quantization_list) @@ -1740,6 +1773,7 @@ def test_clamped_swiglu( self, *, out_shape: Iterable[int] = (32, 32), + glu_interleave_size: Optional[int] = None, dtype: torch.dtype, device: torch.device = "cuda", quantization: Optional[str], @@ -1748,7 +1782,7 @@ def test_clamped_swiglu( limit: float = 0.75, alpha: float = 1.702, ): - # Test SwiGLU variant used in GPT OSS. + """SwiGLU variant used in GPT-OSS""" # Tensor dimensions in_shape = list(out_shape) in_shape[-1] *= 2 @@ -1773,7 +1807,17 @@ def test_clamped_swiglu( ) # Plain PyTorch implementation - x_glu, x_linear = x_ref.chunk(2, dim=-1) + x = x_ref + if glu_interleave_size is not None: + x = x.reshape( + *in_shape[:-1], + in_shape[-1] // (2 * glu_interleave_size), + 2, + glu_interleave_size, + ) + x = x.transpose(-3, -2) + x = x.reshape(in_shape) + 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) @@ -1785,7 +1829,11 @@ def test_clamped_swiglu( forward = te_ops.Sequential( te_ops.Quantize(forward=False, backward=quantize_backward), - te_ops.ClampedSwiGLU(limit=limit, alpha=alpha), + te_ops.ClampedSwiGLU( + limit=limit, + alpha=alpha, + glu_interleave_size=glu_interleave_size, + ), te_ops.Quantize(forward=quantize_forward, backward=False), ) with te.autocast(enabled=quantized_compute, recipe=recipe): @@ -1801,10 +1849,19 @@ def test_clamped_swiglu( tols = dtype_tols(tex.DType.kFloat8E4M3) # Check results - y_test = y_test.to(dtype=torch.float64, device="cpu") - dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") - torch.testing.assert_close(y_test, y_ref, **tols) - torch.testing.assert_close(dx_test, x_ref.grad, **tols) + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + + def test_interleaved_clamped_swiglu(self): + """GPT-OSS SwiGLU with block interleaved input format""" + self.test_clamped_swiglu( + out_shape=(32, 192), + dtype=torch.float32, + quantization=None, + quantize_forward=False, + quantize_backward=False, + glu_interleave_size=32, + ) @pytest.mark.parametrize("scale", (1, 0, -2.5, 3.5)) @pytest.mark.parametrize("shape", ((), (1, 13), (4, 4, 2))) @@ -1924,6 +1981,231 @@ def test_dropout( abs(z_score) < 2.5758 ), f"Number of zeros is outside 99% confidence interval ({prob=}, {prob_observed=})" + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("dtype", _dtypes) + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantized_compute", (False, True)) + @pytest.mark.parametrize("quantized_weight", (False, True)) + @pytest.mark.parametrize("input_requires_grad", (False, True)) + @pytest.mark.parametrize("weight_requires_grad", (False, True)) + def test_grouped_linear( + self, + *, + group_size: int = 4, + bias: bool, + weight_shape: tuple[int, int] = (128, 128), + split_alignment: int = 128, + dtype: torch.dtype, + device: torch.device = "cuda", + quantization: Optional[str], + quantized_compute: bool, + quantized_weight: bool, + input_requires_grad: bool, + weight_requires_grad: bool, + ) -> None: + """Grouped GEMM""" + + # Split sizes + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + # Make input and weight shapes consistent + out_features, in_features = weight_shape + in_shape = (split_sizes.sum().item(), in_features) + out_shape = (in_shape[0], 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_compute or quantized_weight): + pytest.skip("Quantization scheme is not specified") + if quantization is not None and not (quantized_compute or quantized_weight): + 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") + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + requires_grad=input_requires_grad, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + out_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + ws_ref, ws_test = [], [] + bs_ref, bs_test = [], [] + for _ in range(group_size): + w_ref, w_test = make_reference_and_test_tensors( + (out_features, in_features), + quantization=quantization, + test_dtype=dtype, + test_device=device, + requires_grad=weight_requires_grad, + ) + b_ref, b_test = None, None + if bias: + b_ref, b_test = make_reference_and_test_tensors( + out_features, + test_dtype=dtype, + test_device=device, + requires_grad=weight_requires_grad, + ) + ws_ref.append(w_ref) + ws_test.append(w_test) + bs_ref.append(b_ref) + bs_test.append(b_test) + + # Plain PyTorch implementation + xs_ref = torch.split(x_ref, split_sizes.tolist()) + ys_ref = [] + for x, w, b in zip(xs_ref, ws_ref, bs_ref): + ys_ref.append(torch.nn.functional.linear(x, w, bias=b)) + y_ref = torch.cat(ys_ref) + if input_requires_grad or weight_requires_grad: + y_ref.backward(dy_ref) + + # Construct fusible operation + 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, + ) + with torch.no_grad(): + for group_idx in range(group_size): + getattr(op, f"weight{group_idx}").copy_(ws_test[group_idx]) + if bias: + 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) + + # Forward and backward pass with op + with te.autocast(enabled=quantized_compute, recipe=recipe): + y_test = op(x_test, split_sizes) + if input_requires_grad or weight_requires_grad: + y_test.backward(dy_test) + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + if quantized_compute: + tols = quantization_tols(quantization) + + # Check results + y_test = y_test.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(y_test, y_ref, **tols) + if input_requires_grad: + dx_test = x_test.grad.to(dtype=torch.float64, device="cpu") + 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 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) + else: + assert w_test.grad is None + if bias: + 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("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_swiglu( + self, + *, + in_shape: Iterable[int], + glu_interleave_size: Optional[int] = None, + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + input_requires_grad: bool, + scales_requires_grad: bool, + ) -> None: + """SwiGLU with post-scale""" + + # Tensor dims + out_shape = list(in_shape) + out_shape[-1] //= 2 + + # 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( + out_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + x = x_ref + if glu_interleave_size is not None: + x = x.reshape( + -1, + in_shape[-1] // (2 * glu_interleave_size), + 2, + glu_interleave_size, + ) + x = x.transpose(1, 2) + x = x.reshape(in_shape) + x1, x2 = x.chunk(2, dim=-1) + y = torch.nn.functional.silu(x1) * x2 + 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.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + 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( + in_shape=(32, 192), + glu_interleave_size=32, + input_requires_grad=True, + scales_requires_grad=True, + ) + class TestFusedOps: """Tests for fused operations""" @@ -2931,6 +3213,188 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: 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) + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("dtype", _dtypes) + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("glu_interleave_size", (None, 32)) + def test_grouped_mlp( + self, + *, + group_size: int = 4, + bias: bool, + hidden_size: int = 256, + dtype: torch.dtype, + quantization: Optional[str], + device: torch.device = "cuda", + split_alignment: int = 256, + glu_interleave_size: Optional[int], + ) -> None: + """GroupedLinear + ScaledSwiGLU + GroupedLinear""" + + # Split sizes + split_sizes = [split_alignment * i for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + + # Make input shape + in_shape = (split_sizes.sum().item(), hidden_size) + out_shape = in_shape + + # Skip invalid configurations + with_quantization = quantization is not None + maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + if with_quantization and dtype not in (torch.bfloat16, torch.float16): + pytest.skip("Quantized group GEMM is only supported with BF16/FP16") + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + min=-0.25, + max=0.25, + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + out_shape, + min=-0.25, + max=0.25, + quantization=quantization, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + probs_ref, probs_test = make_reference_and_test_tensors( + (in_shape[0],), + test_dtype=dtype, + test_device=device, + ) + fc1_ws_ref, fc1_ws_test = [], [] + fc1_bs_ref, fc1_bs_test = [], [] + fc2_ws_ref, fc2_ws_test = [], [] + 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), + min=-0.25, + max=0.25, + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + fc2_w_ref, fc2_w_test = make_reference_and_test_tensors( + (hidden_size, hidden_size), + min=-0.25, + max=0.25, + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + fc1_b_ref, fc1_b_test = None, None + fc2_b_ref, fc2_b_test = None, None + if bias: + fc1_b_ref, fc1_b_test = make_reference_and_test_tensors( + (2 * hidden_size,), + min=-0.5, + max=0.5, + test_dtype=dtype, + test_device=device, + ) + fc2_b_ref, fc2_b_test = make_reference_and_test_tensors( + (hidden_size,), + min=-0.5, + max=0.5, + test_dtype=dtype, + test_device=device, + ) + fc1_ws_ref.append(fc1_w_ref) + fc1_bs_ref.append(fc1_b_ref) + fc1_ws_test.append(fc1_w_test) + fc1_bs_test.append(fc1_b_test) + fc2_ws_ref.append(fc2_w_ref) + fc2_bs_ref.append(fc2_b_ref) + 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]) + if glu_interleave_size is not None: + x = x.reshape( + -1, + 2 * hidden_size // (2 * glu_interleave_size), + 2, + glu_interleave_size, + ) + x = x.transpose(1, 2) + x = x.reshape(-1, 2 * hidden_size) + x1, x2 = x.chunk(2, dim=-1) + x = torch.nn.functional.silu(x1) * x2 + x = x * probs[group_idx].unsqueeze(-1) + x = torch.nn.functional.linear(x, fc2_ws_ref[group_idx], bias=fc2_bs_ref[group_idx]) + ys.append(x) + y_ref = torch.cat(ys) + y_ref.backward(dy_ref) + + # Construct operations + recipe = make_recipe(quantization) + with te.quantized_model_init(enabled=with_quantization, recipe=recipe): + fc1 = te_ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=bias, + device=device, + dtype=dtype, + ) + fc2 = te_ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device=device, + dtype=dtype, + ) + module = te_ops.Sequential( + fc1, + te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), + fc2, + ) + + # Copy weights + with torch.no_grad(): + for group_idx in range(group_size): + getattr(fc1, f"weight{group_idx}").copy_(fc1_ws_test[group_idx]) + getattr(fc2, f"weight{group_idx}").copy_(fc2_ws_test[group_idx]) + if bias: + getattr(fc1, f"bias{group_idx}").copy_(fc1_bs_test[group_idx]) + getattr(fc2, f"bias{group_idx}").copy_(fc2_bs_test[group_idx]) + del fc1_ws_test, fc1_bs_test, fc2_ws_test, fc2_bs_test + + # Fuse ops and perform forward and backward pass + with te.autocast(enabled=with_quantization, recipe=recipe): + y_test = module(x_test, split_sizes, probs_test, split_sizes) + y_test.backward(dy_test) + + # Loose tols for sanity checking + tols = {"rtol": 0.125, "atol": 0.25} + if quantization == "nvfp4": + tols = {"rtol": 0.25, "atol": 0.5} + + # Check values + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + assert_close_grads(probs_test, probs_ref, **tols) + for group_idx in range(group_size): + assert_close_grads(getattr(fc2, f"weight{group_idx}"), fc2_ws_ref[group_idx], **tols) + assert_close_grads(getattr(fc2, f"bias{group_idx}"), fc2_bs_ref[group_idx], **tols) + assert_close_grads(getattr(fc1, f"weight{group_idx}"), fc1_ws_ref[group_idx], **tols) + assert_close_grads(getattr(fc1, f"bias{group_idx}"), fc1_bs_ref[group_idx], **tols) + class TestCustomOps: """Test with ops that are defined externally""" diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index b6a84a8e2b..c54295d478 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -15,7 +15,7 @@ import transformer_engine import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch import InferenceParams +from transformer_engine.pytorch import InferenceParams, QuantizedTensor from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends from transformer_engine.pytorch.attention.dot_product_attention.utils import ( get_attention_backend, @@ -361,3 +361,48 @@ def test(): if fused_attention_backend == FusedAttnBackend[backends[i]]: fused_attn_backends.append(fused_attention_backend) return available_backends, flash_attention_backend, fused_attn_backends + + +@torch.no_grad +def assert_close( + actual: Optional[torch.Tensor], + expected: Optional[torch.Tensor], + *, + check_device: bool = False, + check_dtype: bool = False, + check_layout: bool = False, + **kwargs, +) -> None: + """Assert that two tensors are close. + + This function is a wrapper around torch.testing.assert_close. It + changes the defaults for device and dtype checks (useful when the + reference implementation is computed in high precision on CPU) and + it can handle quantized tensors. + + """ + if isinstance(actual, QuantizedTensor): + actual = actual.dequantize() + if isinstance(expected, QuantizedTensor): + expected = expected.dequantize() + torch.testing.assert_close( + actual, + expected, + check_device=check_device, + check_dtype=check_dtype, + check_layout=check_layout, + **kwargs, + ) + + +def assert_close_grads( + actual: Optional[torch.Tensor], + expected: Optional[torch.Tensor], + **kwargs, +) -> None: + """Assert that two tensors have close gradients.""" + if actual is None and expected is None: + return + assert actual is not None + assert expected is not None + assert_close(actual.grad, expected.grad, **kwargs) diff --git a/transformer_engine/pytorch/module/_common.py b/transformer_engine/pytorch/module/_common.py index 88b58a353a..bf5a230e84 100644 --- a/transformer_engine/pytorch/module/_common.py +++ b/transformer_engine/pytorch/module/_common.py @@ -77,6 +77,8 @@ def forward( # Check first tensor if not tensors: raise ValueError("Attempted to concatenate 0 tensors") + + # Check concat dim num_dims = tensors[0].dim() if not -num_dims <= dim < num_dims: raise ValueError( @@ -109,11 +111,24 @@ def forward( ctx.dim = dim ctx.split_ranges = split_ranges - # Out-of-place concatenation if needed + # Tensor properties from first tensor dtype = tensors[0].dtype device = tensors[0].device strides = tensors[0].stride() data_ptr_stride = strides[dim] * tensors[0].element_size() + + # Out-of-place concatenation when view tensors have different storage + # Note: This works around an edge case with the split_quantize + # function, which might allocate a buffer and construct + # subviews. However, in order to reduce CPU overheads, these + # views are configured manually outside of PyTorch. PyTorch + # doesn't know these views share the same memory, and it + # blocks us from reconstructing the full tensor because it + # thinks we are accessing out-of-bounds memory. + if tensors[0].untyped_storage().nbytes() < out_shape[dim] * data_ptr_stride: + return torch.cat(tensors, dim=dim) + + # Out-of-place concatenation if tensor properties do not match data_ptr = tensors[0].data_ptr() + tensors[0].size(dim) * data_ptr_stride for tensor in tensors[1:]: if ( @@ -126,13 +141,7 @@ def forward( data_ptr += tensor.size(dim) * data_ptr_stride # No-op concatenation - out = tensors[0].new() - out.set_( - tensors[0].untyped_storage(), - tensors[0].storage_offset(), - out_shape, - strides, - ) + out = tensors[0].as_strided(out_shape, strides) out.requires_grad = any(tensor.requires_grad for tensor in tensors) return out diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 665ffe359c..32da121cce 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -14,8 +14,6 @@ SReLU, SReGLU, SiLU, - SwiGLU, - ClampedSwiGLU, ) from .add_extra_input import AddExtraInput from .all_gather import AllGather @@ -24,6 +22,7 @@ from .bias import Bias from .constant_scale import ConstantScale from .dropout import Dropout +from .grouped_linear import GroupedLinear from .identity import Identity from .l2normalization import L2Normalization from .layer_norm import LayerNorm @@ -32,3 +31,4 @@ from .reduce_scatter import ReduceScatter from .reshape import Reshape from .rmsnorm import RMSNorm +from .swiglu import ClampedSwiGLU, ScaledSwiGLU, SwiGLU diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 9d54e12dba..2f1debdf5e 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -27,8 +27,6 @@ "SReLU", "SReGLU", "SiLU", - "SwiGLU", - "ClampedSwiGLU", ] @@ -355,76 +353,3 @@ def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: return tex.dsilu(*args, **kwargs) - - -class SwiGLU(_ActivationOperation): - r"""Swish gated linear unit - - The input tensor is split into chunks :math:`a` and :math:`b` - along the last dimension and the following is computed: - - .. math:: - - \text{GEGLU}(a,b) = \text{SiLU}(a) * b - - where - - .. math:: - - \text{SiLU}(x) = x \sigma(x) = \frac{x}{1+\exp(-x)} - - .. warning:: - - Transformer Engine's gated activations and PyTorch's GLU - activation follow opposite conventions for :math:`a` and - :math:`b`. Transformer Engine applies the gating function to - the first half of the input tensor, while PyTorch applies it to - the second half. - - The Sigmoid Linear Unit (SiLU) gating function is also known as - the swish function. See - `GLU Variants Improve Transformer`__ - and `Gaussian Error Linear Units (GELUs)`__. - - """ - - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: - return tex.swiglu(*args, **kwargs) - - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: - return tex.dswiglu(*args, **kwargs) - - -class ClampedSwiGLU(_ActivationOperation): - r"""GPT-OSS - Implementation based on `GPT-OSS`__. - - 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. - - .. warning:: The input tensor is chunked along the last dimension to get gates/pre-activations which is differnt - from GPT OSS implementation where the gates/pre-activations are assumed to be interleaved in the input tensor. - - Parameters - ---------- - limit : float - The clamp limit. - alpha : float - The scaling factor for the sigmoid function used in the activation. - cache_quantized_input : bool, default = False - Quantize input tensor when caching for use in the backward pass. - """ - - def __init__( - self, *, limit: float = 7.0, alpha: float = 1.702, cache_quantized_input: bool = False - ): - super().__init__(cache_quantized_input=cache_quantized_input) - self.limit = limit - self.alpha = alpha - - def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: - return tex.clamped_swiglu(*args, limit=self.limit, alpha=self.alpha, **kwargs) - - def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: - return tex.clamped_dswiglu(*args, limit=self.limit, alpha=self.alpha, **kwargs) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py new file mode 100644 index 0000000000..eb8a67600d --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -0,0 +1,702 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible operation for grouped linear layer.""" + +from __future__ import annotations +from collections.abc import Callable, Iterable, Sequence +import contextlib +import math +from typing import Any, Optional + +import torch + +import transformer_engine_torch as tex +from ...cpp_extensions import general_grouped_gemm +from ...distributed import CudaRNGStatesTracker +from ...module.base import ( + _2X_ACC_FPROP, + _2X_ACC_DGRAD, + _2X_ACC_WGRAD, + get_dummy_wgrad, +) +from ...quantization import FP8GlobalStateManager, Recipe +from ...tensor import MXFP8Quantizer, MXFP8Tensor, Quantizer +from ...utils import ( + canonicalize_device, + canonicalize_dtype, + clear_tensor_data, + devices_match, + round_up_to_nearest_multiple, +) +from .._common import is_quantized_tensor, maybe_dequantize +from ..op import BasicOperation, OperationContext + + +class GroupedLinear(BasicOperation): + r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i`` + + This feature is experimental and subject to change. + + This is equivalent to splitting the input tensor along its first + dimension, applying a separate ``torch.nn.Linear`` to each split, + and concatenating along the first dimension. + + Parameters + ---------- + num_groups : int + Number of linear transformations. + in_features : int + Inner dimension of input tensor. + out_features : int + Inner dimension of output tensor. + bias : bool, default = ``True`` + Apply additive bias. + device : torch.device, default = default CUDA device + Tensor device. + dtype : torch.dtype, default = default dtype + Tensor datatype. + rng_state_tracker_function : callable + Function that returns ``CudaRNGStatesTracker``, which is used + for model-parallel weight initialization. + accumulate_into_main_grad : bool, default = ``False`` + Whether to directly accumulate weight gradients into the + weight's ``main_grad`` attribute instead of relying on PyTorch + autograd. The weight's ``main_grad`` must be set externally + and there is no guarantee that `grad` will be set or be + meaningful. This is primarily intented to integrate with + Megatron-LM. This argument along with weight tensor having + attribute ``overwrite_main_grad`` set to True will overwrite + ``main_grad`` instead of accumulating. + + """ + + # Operation expects input split sizes + num_extra_inputs: int = 1 + + def __init__( + self, + num_groups: int, + in_features: int, + out_features: int, + *, + bias: bool = True, + device: Optional[torch.device | str] = None, + dtype: Optional[torch.dtype] = None, + rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]] = None, + accumulate_into_main_grad: bool = False, + ) -> None: + super().__init__() + + # Weight tensor dimensions + self.num_groups: int = num_groups + self.in_features: int = in_features + self.out_features: int = out_features + if self.num_groups <= 0: + raise ValueError(f"Invalid number of groups ({self.num_groups})") + if self.in_features <= 0: + raise ValueError(f"Invalid input size ({self.in_features})") + if self.out_features <= 0: + raise ValueError(f"Invalid output size ({self.out_features})") + + # Weight tensor attributes + device = canonicalize_device(device) + dtype = canonicalize_dtype(dtype) + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise ValueError(f"Supported dtypes are float32, float16, bfloat16 (got {dtype})") + + # Initialize recipe state if needed for natively quantized weight + self._with_quantized_weight: bool = FP8GlobalStateManager.with_fp8_parameters() + if self._with_quantized_weight: + self.reset_recipe_state(recipe=FP8GlobalStateManager.get_fp8_recipe()) + + # RNG state tracker + self._rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]] + self._rng_state_tracker_function = rng_state_tracker_function + + # Register weights + self.weight0: torch.nn.Parameter + for group_idx in range(self.num_groups): + weight_tensor = torch.empty( + self.out_features, + self.in_features, + device="meta", + dtype=dtype, + ) + self.register_parameter( + f"weight{group_idx}", + torch.nn.Parameter(weight_tensor), + ) + + # Register biases + self.bias0: Optional[torch.nn.Parameter] + for group_idx in range(self.num_groups): + bias_tensor = None + if bias: + bias_tensor = torch.empty( + self.out_features, + device="meta", + dtype=dtype, + ) + bias_tensor = torch.nn.Parameter(bias_tensor) + self.register_parameter(f"bias{group_idx}", bias_tensor) + + # Initialize weights if needed + if device.type != "meta": + self.reset_parameters() + + # Whether to accumulate weight gradient into main_grad + self._accumulate_into_main_grad: bool = accumulate_into_main_grad + + def num_quantizers(self, mode: str) -> int: + if mode == "forward": + return 2 * self.num_groups + if mode == "backward": + return self.num_groups + return 0 + + @property + def has_bias(self) -> bool: + """Whether an additive bias is being applied""" + return self.bias0 is not None + + def reset_parameters(self) -> None: + """Initialize parameter buffers and values""" + + # Parameter device + device = self.weight0.device + if device.type == "meta": + device = canonicalize_device(None) + + # Initialize weight values + # Note: Allocate a single buffer in order to support grouped + # GEMM kernels that expect a single weight buffer. + packed_weights = torch.empty( + self.num_groups, + self.out_features, + self.in_features, + dtype=self.weight0.dtype, + device=device, + ) + weights = [packed_weights[idx] for idx in range(self.num_groups)] + for weight in weights: + init_context = contextlib.nullcontext() + if self._rng_state_tracker_function is not None: + init_context = self._rng_state_tracker_function().fork() + with init_context: + torch.nn.init.kaiming_uniform_(weight, a=math.sqrt(5)) + + # Quantize weights if needed + if self._with_quantized_weight: + + # Configure quantizers + quantizers = [ + self.get_quantizer("forward", 2 * idx + 1) for idx in range(self.num_groups) + ] + with_rowwise_usage = True + with_columnwise_usage = torch.is_grad_enabled() + for quantizer in quantizers: + if quantizer is None: + raise RuntimeError( + "Tried to quantize weight with deferred initialization " + "due to meta device, but no quantizer was available. " + "This is most likely because the weight was initialized " + "within quantized_model_init, but the forward pass was not " + "performed within autocast." + ) + quantizer.set_usage( + rowwise=with_rowwise_usage, + columnwise=with_columnwise_usage, + ) + quantizer.internal = False + + # Quantize weights + weights = self._quantize_weights(weights, quantizers) + + # Register weights + for group_idx, weight in enumerate(weights): + if not isinstance(weight, torch.nn.Parameter): + weight = torch.nn.Parameter(weight) + setattr(self, f"weight{group_idx}", weight) + + # Initialize biases if needed + if self.bias0 is not None: + packed_biases = torch.zeros( + self.num_groups, + self.out_features, + dtype=self.bias0.dtype, + device=device, + ) + for group_idx in range(self.num_groups): + bias = torch.nn.Parameter(packed_biases[group_idx]) + setattr(self, f"bias{group_idx}", bias) + + def _quantize_weights( + self, + weights: Sequence[torch.Tensor], + quantizers: Sequence[Quantizer], + ) -> Sequence[torch.Tensor]: + """Construct quantized weight tensors.""" + + # Manually construct MXFP8 weights + if isinstance(quantizers[0], MXFP8Quantizer): + return self._quantize_weights_mxfp8(weights, quantizers) + + # Use quantizers to construct quantized weights + with torch.no_grad(): + return [quantizer(weight) for quantizer, weight in zip(quantizers, weights)] + + def _quantize_weights_mxfp8( + self, + weights: Sequence[torch.Tensor], + quantizers: Sequence[Quantizer], + ) -> Sequence[MXFP8Tensor]: + """Construct MXFP8 weight tensors. + + Instead of allocating separate buffers for each weight tensor, + this function constructs large buffers and assigns subviews to + each tensor. This is intended to support grouped GEMM kernels + that expect packed buffers. + + """ + + # Tensor dimensions + num_groups = len(weights) + out_features, in_features = weights[0].size() + packed_shape = (num_groups, out_features, in_features) + unpacked_shape = (out_features, in_features) + + # Tensor attributes + device = weights[0].device + dtype = weights[0].dtype + requires_grad = torch.is_grad_enabled() + with_rowwise_usage = quantizers[0].rowwise_usage + with_columnwise_usage = quantizers[0].columnwise_usage + + # Construct packed buffers + rowwise_data = [None] * num_groups + rowwise_scales = [None] * num_groups + columnwise_data = [None] * num_groups + columnwise_scales = [None] * num_groups + if with_rowwise_usage: + scale_shape = ( + num_groups, + round_up_to_nearest_multiple(out_features, 128), + round_up_to_nearest_multiple(in_features // 32, 4), + ) + packed_data = torch.empty(packed_shape, dtype=torch.uint8, device=device) + packed_scales = torch.empty(scale_shape, dtype=torch.uint8, device=device) + rowwise_data = [packed_data[idx] for idx in range(num_groups)] + rowwise_scales = [packed_scales[idx] for idx in range(num_groups)] + if with_columnwise_usage: + scale_shape = ( + num_groups, + round_up_to_nearest_multiple(out_features // 32, 4), + round_up_to_nearest_multiple(in_features, 128), + ) + packed_data = torch.empty(packed_shape, dtype=torch.uint8, device=device) + packed_scales = torch.empty(scale_shape, dtype=torch.uint8, device=device) + columnwise_data = [packed_data[idx] for idx in range(num_groups)] + columnwise_scales = [packed_scales[idx] for idx in range(num_groups)] + + # Construct MXFP8 tensors and cast to MXFP8 + out = [] + with torch.no_grad(): + for group_idx in range(num_groups): + weight = MXFP8Tensor( + shape=unpacked_shape, + dtype=dtype, + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise_data=rowwise_data[group_idx], + rowwise_scale_inv=rowwise_scales[group_idx], + columnwise_data=columnwise_data[group_idx], + columnwise_scale_inv=columnwise_scales[group_idx], + quantizer=quantizers[group_idx], + requires_grad=requires_grad, + with_gemm_swizzled_scales=False, + ) + weight.copy_(weights[group_idx]) + out.append(weight) + + return out + + def pre_first_fuser_forward(self) -> None: + super().pre_first_fuser_forward() + + # Initialize params if needed + if any(param.device.type == "meta" for param in self.parameters()): + self.reset_parameters() + + # Check that weights are consistent + dtype = self.weight0.dtype + device = self.weight0.device + weight_requires_grad = self.weight0.requires_grad + weight_tensor_type = type(self.weight0.data) + for group_idx in range(self.num_groups): + weight = getattr(self, f"weight{group_idx}") + if weight.dtype != dtype: + raise RuntimeError( + f"Weight {group_idx} has invalid dtype (expected {dtype}, got {weight.dtype})." + ) + if not devices_match(weight.device, device): + raise RuntimeError( + f"Weight {group_idx} has invalid device " + f"(expected {device}, got {weight.device})." + ) + if weight.requires_grad != weight_requires_grad: + raise RuntimeError( + f"Weight {group_idx} has requires_grad={weight.requires_grad}, " + f"but expected requires_grad={weight_requires_grad}." + ) + if type(weight.data) != weight_tensor_type: # pylint: disable=unidiomatic-typecheck + raise RuntimeError( + f"Weight {group_idx} has invalid tensor type " + f"(expected {weight_tensor_type.__name__}, " + f"got {type(weight.data).__name__})." + ) + + # Check that biases are consistent + for group_idx in range(self.num_groups): + bias = getattr(self, f"bias{group_idx}") + if self.has_bias: + if bias is None: + raise RuntimeError(f"Expected biases, but bias {group_idx} is uninitialized") + if bias.dtype != dtype: + raise RuntimeError( + f"Bias {group_idx} has invalid dtype (expected {dtype}, got {bias.dtype})." + ) + if not devices_match(bias.device, device): + raise RuntimeError( + f"Bias {group_idx} has invalid device " + f"(expected {device}, got {bias.device})." + ) + if bias.requires_grad != weight_requires_grad: + raise RuntimeError( + f"Bias {group_idx} has requires_grad={bias.requires_grad}, " + f"but expected requires_grad={weight_requires_grad}." + ) + else: + if bias is not None: + raise RuntimeError(f"Expected no biases, but bias {group_idx} is initialized") + + def pre_fuser_forward(self, *, requires_grad: bool) -> None: + super().pre_fuser_forward(requires_grad=requires_grad) + if FP8GlobalStateManager.is_fp8_enabled(): + # Assume weights have consistent grad requirement + weight_requires_grad = requires_grad and self.weight0.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) + grad_output_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + + def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: + super().reset_recipe_state(recipe=recipe) + + for group_idx in range(self.num_groups): + # Input/grad output quantizers use internal tensors + input_quantizer = self.get_quantizer("forward", 2 * group_idx) + grad_output_quantizer = self.get_quantizer("backward", group_idx) + if input_quantizer is not None: + input_quantizer.internal = True + if grad_output_quantizer is not None: + grad_output_quantizer.internal = True + + # Handle weight quantizer + # 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 + getattr(self, f"weight{group_idx}").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) + ) + + # 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 + input_quantizer.amax_epsilon_scales = recipe.fp8_quant_fwd_inp.amax_epsilon + weight_quantizer.force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale + weight_quantizer.amax_epsilon_scales = recipe.fp8_quant_fwd_weight.amax_epsilon + grad_output_quantizer.force_pow_2_scales = ( + recipe.fp8_quant_bwd_grad.power_2_scale + ) + grad_output_quantizer.amax_epsilon_scales = ( + recipe.fp8_quant_bwd_grad.amax_epsilon + ) + + def op_forward(self, *args, **kwargs): + 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): + 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], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + num_groups = self.num_groups + has_bias = self.has_bias + device = self.weight0.device + + # Check which grads are required + ctx = basic_op_ctxs[0] + input_requires_grad = ctx.requires_grad + weight_requires_grad = ctx.requires_grad and self.weight0.requires_grad + + # 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(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = self.weight0.dtype + + # Extract split sizes from extra input + 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)}.") + + # Extract params + weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + bs = None + if has_bias: + 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) + + # Split input tensor and convert dtypes if needed + x = maybe_dequantize(input_, dtype) + xs = None + if with_quantized_compute: + for quantizer in input_quantizers: + quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + xs = tex.split_quantize(x, split_sizes_int, input_quantizers) + else: + xs = torch.split(x, split_sizes_int) + + # Allocate output tensor + in_shape = list(input_.size()) + out_shape = in_shape[:-1] + [self.out_features] + out = torch.empty(out_shape, dtype=dtype, device=device) + + # Perform GEMMs + general_grouped_gemm( + ws, + xs, + [out], + [None] * num_groups, # quantization_params + dtype, + m_splits=split_sizes_int, + bias=bs, + use_bias=has_bias, + use_split_accumulator=_2X_ACC_FPROP, + single_output=True, + ) + + # Prepare weight tensors for backward pass + if not input_requires_grad: + ws = [None] * num_groups + elif with_quantized_compute: + for w, weight_param in zip(ws, weights): + if w is not weight_param: + w.update_usage(rowwise_usage=False, columnwise_usage=True) + + # Prepare input tensor for backward pass + if not weight_requires_grad: + xs = [None] * num_groups + elif with_quantized_compute: + for x in xs: + x.update_usage(rowwise_usage=False, columnwise_usage=True) + + # Save state for backward pass + if ctx.requires_grad: + ctx.save_for_backward(split_sizes, *xs, *ws) + 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 + + return out, [()] + + 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]]], + ]: + num_groups = self.num_groups + has_bias = self.has_bias + device = self.weight0.device + + # Saved tensors from forward pass + ctx = basic_op_ctxs[0] + saved_tensors = ctx.saved_tensors + split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] + xs, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] + ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] + + # Split grad output tensor and convert dtypes if needed + split_sizes_int = [int(s) for s in split_sizes.tolist()] + dy = maybe_dequantize(grad_output, ctx.dtype) + dys = None + grad_biases = [None] * num_groups + if ctx.with_quantized_compute: + for quantizer in ctx.grad_output_quantizers: + quantizer.set_usage( + rowwise=ctx.input_requires_grad, + columnwise=ctx.weight_requires_grad, + ) + dys = tex.split_quantize(dy, split_sizes_int, ctx.grad_output_quantizers) + if has_bias: + grad_biases = [ + dy.reshape(-1, dy.size(-1)).sum(dim=0) + for dy in torch.split(grad_output, split_sizes_int) + ] + else: + dys = torch.split(dy, split_sizes_int) + if has_bias: + grad_biases = [dy.reshape(-1, dy.size(-1)).sum(dim=0) for dy in dys] + + # Initialize grad weight buffers + accumulate_into_main_grad = self._accumulate_into_main_grad + grad_weights = [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. + 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) + else: + weight_shape = ws[0].size() + 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 + + # Perform dgrad GEMMs + grad_input = None + if ctx.input_requires_grad: + out_shape = list(grad_output.size()) + in_shape = out_shape[:-1] + [self.in_features] + grad_input = torch.empty( + in_shape, + dtype=ctx.dtype, + device=device, + ) + general_grouped_gemm( + ws, + dys, + [grad_input], + [None] * num_groups, # quantization_params + ctx.dtype, + layout="NN", + m_splits=split_sizes_int, + use_split_accumulator=_2X_ACC_DGRAD, + single_output=True, + ) + + # Perform wgrad GEMMs + if ctx.weight_requires_grad: + general_grouped_gemm( + xs, + dys, + grad_weights, + [None] * num_groups, # quantization_params + ctx.dtype, + layout="NT", + m_splits=split_sizes_int, + use_split_accumulator=_2X_ACC_WGRAD, + accumulate=accumulate_into_main_grad, + ) + + # Clear input tensors if possible + 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 + 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), + ) + + grad_params = grad_weights + grad_biases if has_bias else grad_weights + return grad_input, [grad_params], [(None,)] diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py new file mode 100644 index 0000000000..eaffbeee02 --- /dev/null +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -0,0 +1,498 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fusible operation for SwiGLU and variants.""" + +from __future__ import annotations +from collections.abc import Iterable +from typing import Any, Optional + +import torch + +import transformer_engine_torch as tex +from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload +from ...tensor import Float8CurrentScalingQuantizer, Quantizer +from ...utils import clear_tensor_data +from ..op import BasicOperation, OperationContext +from .._common import maybe_dequantize + +__all__ = ["SwiGLU", "ClampedSwiGLU", "ScaledSwiGLU"] + + +class SwiGLU(BasicOperation): + r"""Swish gated linear unit + + The input tensor is split into chunks :math:``a`` and :math:``b`` + along the last dimension and the following is computed: + + .. math:: + + \text{SwiGLU}(a,b) = \text{SiLU}(a) * b + + where + + .. math:: + + \text{SiLU}(x) = x \sigma(x) = \frac{x}{1+\exp(-x)} + + .. warning:: + + Transformer Engine's gated activations and PyTorch's GLU + activation follow opposite conventions for :math:``a`` and + :math:``b``. Transformer Engine applies the gating function to + the first half of the input tensor, while PyTorch applies it to + the second half. + + The Sigmoid Linear Unit (SiLU) gating function is also known as + the swish function. See + ``GLU Variants Improve Transformer``__. + + Parameters + ---------- + cache_quantized_input : bool, default = False + Quantize input tensor when caching for use in the backward + pass. This will typically reduce memory usage but require + extra compute and increase numerical error. This feature is + highly experimental. + glu_interleave_size : int, optional + When set, the GLU activations will use a block interleaved + format. Instead of interpreting the input tensor as a + concatenation of gates and linear units (e.g. + :math:``[a_1, a_2, a_3, a_4, b_1, b_2, b_3, b_4]`` + in the above notation), it will be interpreted + as alternating blocks of gates and linear units (e.g. + :math:``[a_1, a_2, b_1, b_2, a_3, a_4, b_3, b_4]`` + when the interleave size is 2). This data format is highly + experiental and is primarily intended to support some advanced + fused kernels. + + """ + + def __init__( + self, + *, + cache_quantized_input: bool = False, + glu_interleave_size: Optional[int] = None, + ): + super().__init__() + self.cache_quantized_input: bool = cache_quantized_input + self.glu_interleave_size: Optional[int] = glu_interleave_size + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + + # Compute dtype + dtype: torch.dtype + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = input_.dtype + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise RuntimeError(f"Unsupported dtype ({dtype})") + + # Check input tensor + input_ = maybe_dequantize(input_.contiguous(), dtype) + + # Remove interleaving if needed + swiglu_in = input_ + if self.glu_interleave_size is not None: + shape = swiglu_in.size() + swiglu_in = swiglu_in.reshape( + -1, + shape[-1] // (2 * self.glu_interleave_size), + 2, + self.glu_interleave_size, + ) + swiglu_in = swiglu_in.transpose(1, 2).contiguous() + swiglu_in = swiglu_in.view(shape) + + # Launch kernel + out = tex.swiglu(swiglu_in, next_op_input_quantizer) + + # Quantize input to FP8 before caching if needed + if self.cache_quantized_input: + input_quantizer = Float8CurrentScalingQuantizer( + tex.DType.kFloat8E4M3, + input_.device, + ) + input_quantizer.set_usage(rowwise=True, columnwise=False) + input_ = input_quantizer(input_) + + # Save state for backward pass + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(input_) + ctx.save_for_backward(input_) + ctx.dtype = dtype + ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer + + return out + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + + # Saved tensors from forward pass + (input_,) = ctx.saved_tensors + + # Make sure tensors have correct dtypes + x = maybe_dequantize(input_.contiguous(), ctx.dtype) + dy = maybe_dequantize(grad_output.contiguous(), ctx.dtype) + + # Remove interleaving if needed + swiglu_in = x + if self.glu_interleave_size is not None: + shape = swiglu_in.size() + swiglu_in = swiglu_in.reshape( + -1, + shape[-1] // (2 * self.glu_interleave_size), + 2, + self.glu_interleave_size, + ) + swiglu_in = swiglu_in.transpose(1, 2).contiguous() + swiglu_in = swiglu_in.view(shape) + + # Quantizer for grad input + quantizer = ctx.prev_op_grad_output_quantizer + if self.glu_interleave_size is not None: + quantizer = None + + # Launch kernel + grad_swiglu_in = tex.dswiglu(dy, swiglu_in, quantizer) + + # Apply interleaving if needed + dx = grad_swiglu_in + if self.glu_interleave_size is not None: + shape = dx.size() + dx = dx.reshape( + -1, + 2, + shape[-1] // (2 * self.glu_interleave_size), + self.glu_interleave_size, + ) + dx = dx.transpose(1, 2).contiguous() + dx = dx.view(shape) + + # Clear input tensor if possible + clear_tensor_data(input_) + + return dx, () + + +class ClampedSwiGLU(BasicOperation): + r"""GPT-OSS + Implementation based on ``GPT-OSS``__. + + 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. + + .. warning:: The input tensor is chunked along the last dimension to get gates/pre-activations which is different + from GPT OSS implementation where the gates/pre-activations are assumed to be interleaved in the input tensor. + + Parameters + ---------- + limit : float + The clamp limit. + alpha : float + The scaling factor for the sigmoid function used in the activation. + cache_quantized_input : bool, default = ``False`` + Quantize input tensor when caching for use in the backward pass. + glu_interleave_size : int, optional + When set, the GLU activations will use an experimental block + interleaved format. See the corresponding option in the SwiGLU + operation for more details. + + """ + + def __init__( + self, + *, + limit: float = 7.0, + alpha: float = 1.702, + cache_quantized_input: bool = False, + glu_interleave_size: Optional[int] = None, + ): + super().__init__() + self.limit: float = limit + self.alpha: float = alpha + self.cache_quantized_input: bool = cache_quantized_input + self.glu_interleave_size: Optional[int] = glu_interleave_size + + def op_forward( + self, + ctx: OperationContext, + input_: torch.Tensor, + prev_op_grad_output_quantizer: Optional[Quantizer], + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + + # Compute dtype + dtype: torch.dtype + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = input_.dtype + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise RuntimeError(f"Unsupported dtype ({dtype})") + + # Check input tensor + x = maybe_dequantize(input_.contiguous(), dtype) + + # Remove interleaving if needed + swiglu_in = input_ + if self.glu_interleave_size is not None: + shape = swiglu_in.size() + swiglu_in = swiglu_in.reshape( + -1, + shape[-1] // (2 * self.glu_interleave_size), + 2, + self.glu_interleave_size, + ) + swiglu_in = swiglu_in.transpose(1, 2).contiguous() + swiglu_in = swiglu_in.view(shape) + + # Launch kernel + out = tex.clamped_swiglu( + swiglu_in, + next_op_input_quantizer, + limit=self.limit, + alpha=self.alpha, + ) + + # Quantize input to FP8 before caching if needed + if self.cache_quantized_input: + input_quantizer = Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, x.device) + input_quantizer.set_usage(rowwise=True, columnwise=False) + x = input_quantizer(x) + + # Save state for backward pass + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(x) + ctx.save_for_backward(x) + ctx.dtype = dtype + ctx.prev_op_grad_output_quantizer = prev_op_grad_output_quantizer + + return out + + def op_backward( + self, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, tuple[()]]: + + # Saved tensors from forward pass + (input_,) = ctx.saved_tensors + + # Make sure tensors have correct dtypes + x = maybe_dequantize(input_.contiguous(), ctx.dtype) + dy = maybe_dequantize(grad_output.contiguous(), ctx.dtype) + + # Remove interleaving if needed + swiglu_in = x + if self.glu_interleave_size is not None: + shape = swiglu_in.size() + swiglu_in = swiglu_in.reshape( + -1, + shape[-1] // (2 * self.glu_interleave_size), + 2, + self.glu_interleave_size, + ) + swiglu_in = swiglu_in.transpose(1, 2).contiguous() + swiglu_in = swiglu_in.view(shape) + + # Quantizer for grad input + quantizer = ctx.prev_op_grad_output_quantizer + if self.glu_interleave_size is not None: + quantizer = None + + # Launch kernel + grad_swiglu_in = tex.clamped_dswiglu( + dy, + swiglu_in, + quantizer, + limit=self.limit, + alpha=self.alpha, + ) + + # Apply interleaving if needed + dx = grad_swiglu_in + if self.glu_interleave_size is not None: + shape = dx.size() + dx = dx.reshape( + -1, + 2, + shape[-1] // (2 * self.glu_interleave_size), + self.glu_interleave_size, + ) + dx = dx.transpose(1, 2).contiguous() + dx = dx.view(shape) + + # Clear input tensor if possible + clear_tensor_data(input_) + + return dx, () + + +class ScaledSwiGLU(BasicOperation): + r"""SwiGLU with post-scaling. + + If the SwiGLU output has shape ``(d_1, ..., d_n)``, it is + multiplied with an extra input tensor of shape + ``(d_1, ..., d_{n-1})``. + + Parameters + ---------- + glu_interleave_size : int, optional + When set, the GLU activations will use an experimental block + interleaved format. See the corresponding option in the SwiGLU + operation for more details. + + """ + + # Operation expects scales + num_extra_inputs: int = 1 + + def __init__(self, glu_interleave_size: Optional[int] = None): + super().__init__() + self.glu_interleave_size: Optional[int] = glu_interleave_size + + 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], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + extra_input = basic_op_extra_inputs[0][0] + + # Determine compute dtype + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + elif isinstance(input_, torch.Tensor): + dtype = input_.dtype + else: + dtype = extra_input.dtype + + # Make sure inputs are in correct dtype + input_ = maybe_dequantize(input_, dtype) + scales = maybe_dequantize(extra_input, dtype) + + # Remove gate interleaving if needed + swiglu_in = input_ + if self.glu_interleave_size is not None: + shape = swiglu_in.size() + swiglu_in = swiglu_in.reshape( + -1, + shape[-1] // (2 * self.glu_interleave_size), + 2, + self.glu_interleave_size, + ) + swiglu_in = swiglu_in.transpose(1, 2).contiguous() + swiglu_in = swiglu_in.view(shape) + + # Compute scaled SwiGLU + swiglu_out = tex.swiglu(swiglu_in, None) + out = swiglu_out * scales.unsqueeze(-1) + + # Save state for backward pass + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(input_) + ctx.input_requires_grad = True + ctx.extra_input_requires_grad = extra_input.requires_grad + ctx.dtype = dtype + ctx.save_for_backward( + input_, + scales if ctx.input_requires_grad else None, + ) + + return out, [()] + + 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]]], + ]: + ctx = basic_op_ctxs[0] + input_, scales = ctx.saved_tensors + input_ = maybe_dequantize(input_, ctx.dtype) + if scales is not None: + scales = maybe_dequantize(scales, ctx.dtype) + grad_output = maybe_dequantize(grad_output, ctx.dtype) + + # Remove gate interleaving if needed + swiglu_in = input_ + if self.glu_interleave_size is not None: + shape = swiglu_in.size() + swiglu_in = swiglu_in.reshape( + -1, + shape[-1] // (2 * self.glu_interleave_size), + 2, + self.glu_interleave_size, + ) + swiglu_in = swiglu_in.transpose(1, 2).contiguous() + swiglu_in = swiglu_in.view(shape) + + # Compute input grad + grad_input = None + if ctx.input_requires_grad: + grad_swiglu_out = grad_output * scales.unsqueeze(-1) + grad_swiglu_in = tex.dswiglu(grad_swiglu_out, swiglu_in, None) + grad_input = grad_swiglu_in + if self.glu_interleave_size is not None: + shape = grad_input.size() + grad_input = grad_input.reshape( + -1, + 2, + shape[-1] // (2 * self.glu_interleave_size), + self.glu_interleave_size, + ) + grad_input = grad_input.transpose(1, 2).contiguous() + grad_input = grad_input.view(shape) + + # Compute scales grad by recomputing SwiGLU + grad_extra_input = None + if ctx.extra_input_requires_grad: + swiglu_out = tex.swiglu(swiglu_in, None) + grad_extra_input = torch.linalg.vecdot(swiglu_out, grad_output) + + # Clear input tensor if possible + clear_tensor_data(ctx.saved_tensors[0]) # input_ + + return grad_input, [()], [(grad_extra_input,)] From 33ca6150c5006d872c6798e9d1a1c1745a3021f4 Mon Sep 17 00:00:00 2001 From: "Kim, Jin (Jay@SKT)" Date: Fri, 13 Feb 2026 04:18:59 +0900 Subject: [PATCH 208/521] Add sigmoid GLU (#2656) * Add sigmoid GLU Signed-off-by: Kim, Jin * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Kim, Jin * Add test for GLU op 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 reshape Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Apply suggestion from @timmoon10 Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Add omitted tests for GLU op Signed-off-by: Kim, Jin * Add GLU activation type support in JAX extension Signed-off-by: Kim, Jin * [PyTorch] Add Sigmoid activation for GLU support in numerics test (#2656) Signed-off-by: Kim, Jin --------- Signed-off-by: Kim, Jin Signed-off-by: Tim Moon 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 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- tests/pytorch/test_fusible_ops.py | 24 ++++++++++++-- tests/pytorch/test_numerics.py | 2 ++ tests/pytorch/test_sanity.py | 1 + transformer_engine/common/CMakeLists.txt | 2 ++ transformer_engine/common/activation/glu.cu | 24 ++++++++++++++ .../include/transformer_engine/activation.h | 27 +++++++++++++++ .../jax/cpp_extensions/activation.py | 1 + .../jax/csrc/extensions/activation.cpp | 6 ++++ .../jax/csrc/extensions/pybind.cpp | 1 + transformer_engine/pytorch/csrc/extensions.h | 5 +++ .../pytorch/csrc/extensions/activation.cpp | 8 +++++ .../pytorch/csrc/extensions/pybind.cpp | 6 ++++ .../pytorch/module/layernorm_mlp.py | 16 +++++++-- .../pytorch/ops/basic/__init__.py | 1 + .../pytorch/ops/basic/activation.py | 33 +++++++++++++++++++ transformer_engine/pytorch/transformer.py | 2 +- 16 files changed, 154 insertions(+), 5 deletions(-) create mode 100644 transformer_engine/common/activation/glu.cu diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 2c1320e262..5d1a5ce61d 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -1570,7 +1570,19 @@ def test_make_extra_output( @pytest.mark.parametrize( "activation", - ("gelu", "geglu", "qgelu", "qgeglu", "relu", "reglu", "srelu", "sreglu", "silu", "swiglu"), + ( + "gelu", + "geglu", + "qgelu", + "qgeglu", + "relu", + "reglu", + "glu", + "srelu", + "sreglu", + "silu", + "swiglu", + ), ) @pytest.mark.parametrize("out_shape", ((37,), (2, 13), (32, 1, 32))) @pytest.mark.parametrize("dtype", _dtypes) @@ -1590,7 +1602,7 @@ def test_activation( # Tensor dimensions in_shape = list(out_shape) - if activation in ("geglu", "qgeglu", "reglu", "sreglu", "swiglu"): + if activation in ("geglu", "glu", "qgeglu", "reglu", "sreglu", "swiglu"): in_shape[-1] *= 2 # Skip invalid configurations @@ -1630,6 +1642,13 @@ def test_activation( elif activation == "reglu": x1, x2 = x_ref.chunk(2, dim=-1) y_ref = torch.nn.functional.relu(x1) * x2 + elif activation == "sigmoid": + y_ref = torch.nn.functional.sigmoid(x_ref) + elif activation == "glu": + x = x_ref.reshape(*in_shape[:-1], 2, in_shape[-1] // 2) + x = x.flip(-2) # PyTorch GLU swaps gate and linear unit + x = x.reshape(in_shape) + y_ref = torch.nn.functional.glu(x) elif activation == "srelu": y_ref = torch.nn.functional.relu(x_ref) ** 2 elif activation == "sreglu": @@ -1649,6 +1668,7 @@ def test_activation( make_op = dict( gelu=te_ops.GELU, geglu=te_ops.GEGLU, + glu=te_ops.GLU, qgelu=te_ops.QGELU, qgeglu=te_ops.QGEGLU, relu=te_ops.ReLU, diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index abe2806e66..8e3b0517ee 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -89,6 +89,7 @@ all_activations = [ "gelu", "geglu", + "glu", "qgelu", "qgeglu", "relu", @@ -479,6 +480,7 @@ def forward(self, inp: torch.Tensor, m_splits: List[int]) -> torch.Tensor: _supported_act = { "gelu": nn.GELU(approximate="tanh"), "geglu": nn.GELU(approximate="tanh"), + "glu": nn.Sigmoid(), "qgelu": TorchQuickGELU(), "qgeglu": TorchQuickGELU(), "relu": nn.ReLU(), diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index b94cbdcd96..033a6a7ffb 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -113,6 +113,7 @@ def nvfp4_vanilla(): all_activations = [ "gelu", "geglu", + "glu", "qgelu", "qgeglu", "relu", diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f0968c62ee..caba0bf7f1 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -168,6 +168,7 @@ list(APPEND transformer_engine_cuda_sources list(APPEND transformer_engine_cuda_arch_specific_sources activation/gelu.cu + activation/glu.cu activation/relu.cu activation/swiglu.cu cast/cast.cu @@ -354,6 +355,7 @@ 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/glu.cu activation/relu.cu activation/swiglu.cu) endif() diff --git a/transformer_engine/common/activation/glu.cu b/transformer_engine/common/activation/glu.cu new file mode 100644 index 0000000000..45a6670672 --- /dev/null +++ b/transformer_engine/common/activation/glu.cu @@ -0,0 +1,24 @@ +/************************************************************************* + * 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_glu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_glu); + using namespace transformer_engine; + Empty e = {}; + gated_act_fn>(input, output, e, stream); +} + +void nvte_dglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_dglu); + using namespace transformer_engine; + Empty e = {}; + dgated_act_fn, dsigmoid>(grad, input, output, e, + stream); +} diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 4c9eed3365..06f1c65ce2 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -31,6 +31,7 @@ extern "C" { enum class NVTE_Activation_Type { GELU, GEGLU, + GLU, SILU, SWIGLU, RELU, @@ -262,6 +263,32 @@ void nvte_dsrelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Computes the GLU (Gated Linear Unit) activation of the input. + * GLU(a,b) = sigmoid(a) * b + * See "Language Modeling with Gated Convolutional Networks" (arXiv:1612.08083) + * and "GLU Variants Improve Transformer" (arXiv:2002.05202). + * 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 sigmoid(input[N, :H]) x input[N, H:] + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_glu(const NVTETensor input, NVTETensor output, cudaStream_t stream); + +/*! \brief Computes the GLU 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. + * + * \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] stream CUDA stream used for the operation. + */ +void nvte_dglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, + cudaStream_t stream); + /*! \brief Computes the gated GeLU 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. diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index 573603ef3a..8c0edae97e 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -44,6 +44,7 @@ ActivationEnum = { ("gelu",): NVTE_Activation_Type.GELU, ("gelu", "linear"): NVTE_Activation_Type.GEGLU, + ("sigmoid", "linear"): NVTE_Activation_Type.GLU, ("silu",): NVTE_Activation_Type.SILU, ("silu", "linear"): NVTE_Activation_Type.SWIGLU, ("relu",): NVTE_Activation_Type.RELU, diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index 6c5a976344..ce5828d6f3 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -109,6 +109,9 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal case NVTE_Activation_Type::GEGLU: nvte_geglu(input_tensor.data(), output_tensor.data(), stream); break; + case NVTE_Activation_Type::GLU: + nvte_glu(input_tensor.data(), output_tensor.data(), stream); + break; case NVTE_Activation_Type::SILU: nvte_silu(input_tensor.data(), output_tensor.data(), stream); break; @@ -427,6 +430,9 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, case NVTE_Activation_Type::GEGLU: nvte_dgeglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); break; + case NVTE_Activation_Type::GLU: + nvte_dglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); + break; case NVTE_Activation_Type::SWIGLU: nvte_dswiglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); break; diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index a5986404c9..bd4b8fe2c2 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -150,6 +150,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { pybind11::enum_(m, "NVTE_Activation_Type", pybind11::module_local()) .value("GELU", NVTE_Activation_Type::GELU) .value("GEGLU", NVTE_Activation_Type::GEGLU) + .value("GLU", NVTE_Activation_Type::GLU) .value("SILU", NVTE_Activation_Type::SILU) .value("SWIGLU", NVTE_Activation_Type::SWIGLU) .value("RELU", NVTE_Activation_Type::RELU) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index e0ea3d6b78..0e91071983 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -163,6 +163,11 @@ at::Tensor swap_first_dims(at::Tensor tensor, std::optional out = st * Activations **************************************************************************************************/ +/* GLU (sigmoid gate) */ +py::object glu(const at::Tensor &input, py::handle quantizer); + +py::object dglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer); + /* GELU and variants*/ py::object gelu(const at::Tensor &input, py::handle quantizer); diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 9ea14e1af0..99b9c1fefa 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -246,6 +246,14 @@ py::object dgelu(const at::Tensor& grad, const at::Tensor& input, py::handle qua return dactivation_helper(grad, input, quantizer); } +py::object glu(const at::Tensor& input, py::handle quantizer) { + return activation_helper(input, quantizer, 2); +} + +py::object dglu(const at::Tensor& grad, const at::Tensor& input, py::handle quantizer) { + return dactivation_helper(grad, input, quantizer); +} + py::object geglu(const at::Tensor& input, py::handle quantizer) { return activation_helper(input, quantizer, 2); } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 1e907d9bc0..14f32c7b93 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -132,6 +132,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("comm_overlap") = nullptr, py::arg("comm_type") = std::nullopt, py::arg("extra_output") = std::nullopt, py::arg("bulk_overlap") = false, py::arg("alpha") = 1.0f, py::arg("beta") = std::nullopt); + /* GLU (sigmoid gate) */ + m.def("glu", transformer_engine::pytorch::glu, "GLU activation", py::arg("input"), + py::arg("quantizer")); /* GELU and variants*/ m.def("gelu", transformer_engine::pytorch::gelu, "GeLU activation", py::arg("input"), py::arg("quantizer")); @@ -158,6 +161,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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); + /* Backward of GLU */ + m.def("dglu", transformer_engine::pytorch::dglu, "Backward of GLU", py::arg("grad"), + py::arg("fwd_input"), py::arg("quantizer")); /* Backward of GELU and variants */ m.def("dgelu", transformer_engine::pytorch::dgelu, "Backward of GeLU", 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 fb88764b89..4532ea60e7 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -98,6 +98,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): return { "gelu": (tex.gelu, tex.dgelu, None), "geglu": (tex.geglu, tex.dgeglu, None), + "glu": (tex.glu, tex.dglu, None), "qgelu": (tex.qgelu, tex.dqgelu, None), "qgeglu": (tex.qgeglu, tex.dqgeglu, None), "relu": (tex.relu, tex.drelu, None), @@ -114,6 +115,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): return { "gelu": (tex.gelu, tex.dgelu, tex.dbias_dgelu), "geglu": (tex.geglu, tex.dgeglu, None), + "glu": (tex.glu, tex.dglu, None), "qgelu": (tex.qgelu, tex.dqgelu, tex.dbias_dqgelu), "qgeglu": (tex.qgeglu, tex.dqgeglu, None), "relu": (tex.relu, tex.drelu, tex.dbias_drelu), @@ -136,6 +138,7 @@ def _get_act_func_supported_list(recipe: Optional[Recipe] = None): return { "gelu": (tex.gelu, tex.dgelu, None), "geglu": (tex.geglu, tex.dgeglu, None), + "glu": (tex.glu, tex.dglu, None), "qgelu": (tex.qgelu, tex.dqgelu, None), "qgeglu": (tex.qgeglu, tex.dqgeglu, None), "relu": (tex.relu, tex.drelu, None), @@ -1665,7 +1668,7 @@ class LayerNormMLP(TransformerEngineBaseModule): type of normalization applied. activation : str, default = 'gelu' activation function used. - Options: ``'gelu'``, ``'geglu'``, ``'qgelu'``, ``'qgeglu'``, ``'relu'``, ``'reglu'``, ``'srelu'``, ``'sreglu'``, + Options: ``'gelu'``, ``'geglu'``, ``'glu'``, ``'qgelu'``, ``'qgeglu'``, ``'relu'``, ``'reglu'``, ``'srelu'``, ``'sreglu'``, ``'silu'``, ``'swiglu'``, and ``'clamped_swiglu'``. activation_params : dict, default = None Additional parameters for the activation function. @@ -1884,7 +1887,15 @@ def __init__( self.layer_norm_bias = None # FC1 init - if self.activation in ["geglu", "qgeglu", "reglu", "sreglu", "swiglu", "clamped_swiglu"]: + if self.activation in [ + "geglu", + "glu", + "qgeglu", + "reglu", + "sreglu", + "swiglu", + "clamped_swiglu", + ]: fc1_output_features = 2 * self.size_per_partition else: fc1_output_features = self.size_per_partition @@ -2308,6 +2319,7 @@ def _clamped_swiglu(x, limit, alpha): activation_map = { "gelu": lambda x: torch.nn.functional.gelu(x, approximate="tanh"), "geglu": lambda x: torch.nn.functional.gelu(x.chunk(2, -1)[0]) * x.chunk(2, -1)[1], + "glu": lambda x: torch.sigmoid(x.chunk(2, -1)[0]) * x.chunk(2, -1)[1], "qgelu": lambda x: torch.nn.functional.gelu(x, approximate="tanh"), "qgeglu": lambda x: torch.nn.functional.gelu(x.chunk(2, -1)[0], approximate="tanh") * x.chunk(2, -1)[1], diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 32da121cce..e0a3f41019 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -7,6 +7,7 @@ from .activation import ( GELU, GEGLU, + GLU, QGELU, QGEGLU, ReLU, diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 2f1debdf5e..9e23bb3fb1 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -20,6 +20,7 @@ __all__ = [ "GELU", "GEGLU", + "GLU", "QGELU", "QGEGLU", "ReLU", @@ -162,6 +163,38 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: return tex.dgelu(*args, **kwargs) +class GLU(_ActivationOperation): + r"""Gated Linear Unit + + The input tensor is split into chunks :math:`a` and :math:`b` + along the last dimension and the following is computed: + + .. math:: + + \text{GLU}(a,b) = \sigma(a) * b + + where :math:`\sigma` is the sigmoid function. + + .. warning:: + + Transformer Engine's gated activations and PyTorch's GLU + activation follow opposite conventions for :math:`a` and + :math:`b`. Transformer Engine applies the gating function to + the first half of the input tensor, while PyTorch applies it to + the second half. + + See `Language Modeling with Gated Convolutional Networks`__ + and `GLU Variants Improve Transformer`__. + + """ + + def _activation_forward_impl(self, *args, **kwargs) -> torch.Tensor: + return tex.glu(*args, **kwargs) + + def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: + return tex.dglu(*args, **kwargs) + + class GEGLU(_ActivationOperation): r"""Gaussian Error Gated Linear Unit diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index fdb3869199..cf7ce5e1a4 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -184,7 +184,7 @@ class TransformerLayer(torch.nn.Module): if set to ``False``, the transformer layer will not learn any additive biases. activation : str, default = 'gelu' Type of activation used in MLP block. - Options are: ``'gelu'``, ``'geglu'``, ``'qgelu'``, ``'qgeglu'``, ``'relu'``, ``'reglu'``, ``'srelu'``, ``'sreglu'``, + Options are: ``'gelu'``, ``'geglu'``, ``'glu'``, ``'qgelu'``, ``'qgeglu'``, ``'relu'``, ``'reglu'``, ``'srelu'``, ``'sreglu'``, ``'silu'``, ``'swiglu'``, and ``'clamped_swiglu'``. activation_params : Optional[dict], default = None Additional parameters for the activation function. From cd098e4217696976df6e9c3c3fe0ef9fe9ca7b72 Mon Sep 17 00:00:00 2001 From: Harikrishna KP Date: Fri, 13 Feb 2026 00:52:13 +0530 Subject: [PATCH 209/521] fix: correct FusedAdam copy-paste in FusedSGD error messages (#2675) fix: correct copy-paste error messages in FusedSGD Signed-off-by: Mr-Neutr0n <64578610+Mr-Neutr0n@users.noreply.github.com> --- transformer_engine/pytorch/optimizers/fused_sgd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/optimizers/fused_sgd.py b/transformer_engine/pytorch/optimizers/fused_sgd.py index 08e465e951..d7ab3fe9fe 100644 --- a/transformer_engine/pytorch/optimizers/fused_sgd.py +++ b/transformer_engine/pytorch/optimizers/fused_sgd.py @@ -123,7 +123,7 @@ def __init__( self.set_grad_none = set_grad_none if self.set_grad_none is not None: warnings.warn( - "set_grad_none kwarg in FusedAdam constructor is deprecated. " + "set_grad_none kwarg in FusedSGD constructor is deprecated. " "Use set_to_none kwarg in zero_grad instead.", DeprecationWarning, ) @@ -147,7 +147,7 @@ def zero_grad(self, set_to_none: Optional[bool] = None) -> None: if set_to_none is not None and set_to_none != self.set_grad_none: raise ValueError( f"Called zero_grad with set_to_none={set_to_none}, " - f"but FusedAdam was initialized with set_grad_none={self.set_grad_none}" + f"but FusedSGD was initialized with set_grad_none={self.set_grad_none}" ) set_to_none = self.set_grad_none if set_to_none is None: From 496620a950b1b8aa051daacfdbc918a507f7a054 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Thu, 12 Feb 2026 11:38:56 -0800 Subject: [PATCH 210/521] Get rid of nvshmem dependency for cuBLASMp integration (#2661) * Remove nvshmem usage Signed-off-by: Vladimir Cherepanov * Renamings Signed-off-by: Vladimir Cherepanov * NCCL dependency Signed-off-by: Vladimir Cherepanov * Check for not yet allocated workspace Signed-off-by: Vladimir Cherepanov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address greptile comments Signed-off-by: Vladimir Cherepanov * Add a comment per greptile Signed-off-by: Vladimir Cherepanov * Fix a typo Signed-off-by: Vladimir Cherepanov * Display human-readable cuBLASMp error message Signed-off-by: Vladimir Cherepanov --------- Signed-off-by: Vladimir Cherepanov Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- setup.py | 5 -- transformer_engine/common/CMakeLists.txt | 10 ++- .../common/comm_gemm/comm_gemm.cpp | 62 +++++++++++-------- .../include/transformer_engine/comm_gemm.h | 2 + transformer_engine/common/util/logging.h | 12 ++-- 5 files changed, 47 insertions(+), 44 deletions(-) diff --git a/setup.py b/setup.py index 18bb736f24..3a66e624e3 100644 --- a/setup.py +++ b/setup.py @@ -77,11 +77,6 @@ def setup_common_extension() -> CMakeExtension: f"nvidia-cublasmp-cu{cuda_version()[0]}" ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}") cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}") - nvshmem_dir = os.getenv("NVSHMEM_HOME") or metadata.distribution( - f"nvidia-nvshmem-cu{cuda_version()[0]}" - ).locate_file("nvidia/nvshmem") - cmake_flags.append(f"-DNVSHMEM_DIR={nvshmem_dir}") - print("CMAKE_FLAGS:", cmake_flags[-2:]) # Add custom CMake arguments from environment variable nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS") diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index caba0bf7f1..4579c51e9f 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -287,20 +287,18 @@ endif() option(NVTE_WITH_CUBLASMP "Use cuBLASMp for tensor parallel GEMMs" OFF) if (NVTE_WITH_CUBLASMP) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP) - target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include ${NVSHMEM_DIR}/include) + target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include) find_library(CUBLASMP_LIB NAMES cublasmp libcublasmp PATHS ${CUBLASMP_DIR} PATH_SUFFIXES lib REQUIRED) - find_library(NVSHMEM_HOST_LIB - NAMES nvshmem_host libnvshmem_host.so.3 - PATHS ${NVSHMEM_DIR} + find_library(NCCL_LIB + NAMES nccl libnccl PATH_SUFFIXES lib REQUIRED) - target_link_libraries(transformer_engine PUBLIC ${CUBLASMP_LIB} ${NVSHMEM_HOST_LIB}) + target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CUBLASMP_LIB}) message(STATUS "Using cuBLASMp at: ${CUBLASMP_DIR}") - message(STATUS "Using nvshmem at: ${NVSHMEM_DIR}") endif() # Hack to enable dynamic loading in cuDNN frontend diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 66a3da55dd..7be3d1bb4d 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -8,7 +8,6 @@ #include #include -#include #include #include @@ -236,7 +235,7 @@ void GemmArInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n ctx->grid_row_major.get(), ctx->d_desc.get())); const cublasMpMatmulEpilogue_t epilogue = CUBLASMP_MATMUL_EPILOGUE_ALLREDUCE; - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE, &epilogue, sizeof epilogue)); } @@ -273,46 +272,46 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo const cublasOperation_t trans_a = transa ? CUBLAS_OP_T : CUBLAS_OP_N; const cublasOperation_t trans_b = transb ? CUBLAS_OP_T : CUBLAS_OP_N; - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_TRANSA, &trans_a, sizeof trans_a)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_TRANSB, &trans_b, sizeof trans_b)); cublasMpMatmulAlgoType_t algo_attr = cublasmp_algo(algo); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_ALGO_TYPE, &algo_attr, sizeof algo_attr)); const cublasMpMatmulMatrixScale_t scale_mode = CUBLASMP_MATMUL_MATRIX_SCALE_SCALAR_FP32; if (is_fp8_dtype(a->dtype())) { NVTE_CHECK(a->scale_inv.dptr, "Scaling must be set for FP8 dtype"); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_MODE, &scale_mode, sizeof scale_mode)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_A_SCALE_POINTER, &a->scale_inv.dptr, sizeof(void*))); } if (is_fp8_dtype(b->dtype())) { NVTE_CHECK(b->scale_inv.dptr, "Scaling must be set for FP8 dtype"); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_MODE, &scale_mode, sizeof scale_mode)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_B_SCALE_POINTER, &b->scale_inv.dptr, sizeof(void*))); } if (is_fp8_dtype(d->dtype())) { NVTE_CHECK(d->scale.dptr, "Scaling must be set for FP8 dtype"); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_D_SCALE_MODE, &scale_mode, sizeof scale_mode)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_D_SCALE_POINTER, &d->scale.dptr, sizeof(void*))); if (d->amax.dptr) { - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_AMAX_D_POINTER, &d->amax.dptr, sizeof(void*))); } @@ -321,7 +320,7 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo // Might be set to ALLREDUCE before, need to OR with the new flags to set. cublasMpMatmulEpilogue_t epilogue{}; size_t size_read{}; - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeGet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorGetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE, &epilogue, sizeof epilogue, &size_read)); NVTE_CHECK(size_read == sizeof epilogue); @@ -339,42 +338,42 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo pre_act_out ? pre_act_out->data.dptr != nullptr : false, grad}); it != flags_to_epilogue.end()) { epilogue = static_cast(epilogue | it->second); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE, &epilogue, sizeof epilogue)); } if (bias && bias->data.dptr) { cudaDataType_t bias_type = get_cuda_dtype(bias->data.dtype); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_BIAS_DATA_TYPE, &bias_type, sizeof bias_type)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_BIAS_POINTER, &bias->data.dptr, sizeof bias->data.dptr)); } if (pre_act_out && pre_act_out->data.dptr) { cudaDataType_t aux_type = get_cuda_dtype(pre_act_out->data.dtype); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_DATA_TYPE, &aux_type, sizeof aux_type)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_POINTER, &pre_act_out->data.dptr, sizeof pre_act_out->data.dptr)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_LD, &ldd, sizeof ldd)); if (is_fp8_dtype(pre_act_out->dtype())) { NVTE_CHECK(pre_act_out->scale.dptr, "Scaling must be set for FP8 dtype"); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_SCALE_MODE, &scale_mode, sizeof scale_mode)); - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_SCALE_POINTER, &pre_act_out->scale.dptr, sizeof(void*))); if (pre_act_out->amax.dptr) { - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_EPILOGUE_AUX_AMAX_POINTER, &pre_act_out->amax.dptr, sizeof(void*))); } @@ -382,12 +381,12 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo } if (comm_sm_count) { - NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorAttributeSet( + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_COMMUNICATION_SM_COUNT, &comm_sm_count, sizeof comm_sm_count)); } - NVTE_CHECK_CUBLASMP(cublasMpStreamSet(ctx->cublas_mp.get(), main_stream)); + NVTE_CHECK_CUBLASMP(cublasMpSetStream(ctx->cublas_mp.get(), main_stream)); size_t wrksp_size_device{}; size_t wrksp_size_host{}; @@ -423,8 +422,14 @@ void cublasmp_gemm(InitMatricesFn init_matrices_fn, NVTECommGemmCtx* ctx, NVTECo std::vector workspace_host(wrksp_size_host); if (ctx->workspace_size < wrksp_size_device) { - nvshmem_free(ctx->workspace); - ctx->workspace = nvshmem_malloc(wrksp_size_device); + if (ctx->workspace) { + NVTE_CHECK_CUBLASMP(cublasMpBufferDeregister(ctx->grid_row_major.get(), ctx->workspace)); + NVTE_CHECK_CUBLASMP(cublasMpFree(ctx->grid_col_major.get(), ctx->workspace)); + } + NVTE_CHECK_CUBLASMP( + cublasMpMalloc(ctx->grid_col_major.get(), &ctx->workspace, wrksp_size_device)); + NVTE_CHECK_CUBLASMP( + cublasMpBufferRegister(ctx->grid_row_major.get(), ctx->workspace, wrksp_size_device)); ctx->workspace_size = wrksp_size_device; } @@ -473,7 +478,10 @@ NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank void nvte_comm_gemm_ctx_destroy(NVTECommGemmCtx* ctx) { NVTE_API_CALL(nvte_comm_gemm_ctx_destroy); - nvshmemx_sync_all_on_stream(ctx->stream.get()); + if (ctx->workspace) { + NVTE_CHECK_CUBLASMP(cublasMpBufferDeregister(ctx->grid_row_major.get(), ctx->workspace)); + NVTE_CHECK_CUBLASMP(cublasMpFree(ctx->grid_col_major.get(), ctx->workspace)); + } delete ctx; } diff --git a/transformer_engine/common/include/transformer_engine/comm_gemm.h b/transformer_engine/common/include/transformer_engine/comm_gemm.h index 06b56789a3..65d3aa5d9e 100644 --- a/transformer_engine/common/include/transformer_engine/comm_gemm.h +++ b/transformer_engine/common/include/transformer_engine/comm_gemm.h @@ -55,6 +55,8 @@ NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank /*! \brief Destroy a comm-gemm context. * * \param[in] ctx Context to destroy. + * + * It's the caller's responsibility to synchronize all streams involved before calling this function. */ void nvte_comm_gemm_ctx_destroy(NVTECommGemmCtx* ctx); diff --git a/transformer_engine/common/util/logging.h b/transformer_engine/common/util/logging.h index c542afa393..8031e342e2 100644 --- a/transformer_engine/common/util/logging.h +++ b/transformer_engine/common/util/logging.h @@ -96,12 +96,12 @@ #ifdef NVTE_WITH_CUBLASMP -#define NVTE_CHECK_CUBLASMP(expr) \ - do { \ - const cublasMpStatus_t status = (expr); \ - if (status != CUBLASMP_STATUS_SUCCESS) { \ - NVTE_ERROR("cuBLASMp Error: ", std::to_string(status)); \ - } \ +#define NVTE_CHECK_CUBLASMP(expr) \ + do { \ + const cublasMpStatus_t status = (expr); \ + if (status != CUBLASMP_STATUS_SUCCESS) { \ + NVTE_ERROR("cuBLASMp Error: ", cublasMpGetStatusString(status)); \ + } \ } while (false) #endif // NVTE_WITH_CUBLASMP From f8449052455c16f8db0179e8139d90344cf03790 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 13 Feb 2026 11:41:13 +0530 Subject: [PATCH 211/521] [PyTorch] Make grouped weights opt-in (#2678) * Make grouped weights opt-in Signed-off-by: Kirthi Shankar Sivamani * Change varname Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/test_sanity.py | 20 +++++++++++++++++-- .../pytorch/module/grouped_linear.py | 5 ++++- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 033a6a7ffb..d47bc553b0 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -585,10 +585,19 @@ def test_sanity_linear_with_zero_tokens(dtype, bs, model, fp8_recipe, fp8_model_ @pytest.mark.parametrize("fp8_recipe", fp8_recipes) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) +@pytest.mark.parametrize("single_param", all_boolean) @pytest.mark.parametrize("empty_split", ["first", "last", "middle"]) @pytest.mark.parametrize("num_gemms", [4]) def test_sanity_grouped_linear( - dtype, bs, model, fp8_recipe, fp8_model_params, use_bias, num_gemms, empty_split + dtype, + bs, + model, + fp8_recipe, + fp8_model_params, + use_bias, + single_param, + num_gemms, + empty_split, ): if NVTE_TEST_NVINSPECT_ENABLED and fp8_model_params: pytest.skip("FP8 model parameters are not supported in debug mode.") @@ -598,6 +607,9 @@ def test_sanity_grouped_linear( bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) + if single_param: + os.environ["NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS"] = "1" + if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") @@ -617,7 +629,8 @@ def test_sanity_grouped_linear( # Verify that weights are stored in contiguous GroupedTensor storage. weights = [getattr(te_grouped_linear, f"weight{i}") for i in range(num_gemms)] if fp8_recipe is None or not (fp8_recipe.delayed() or fp8_recipe.float8_current_scaling()): - check_grouped_tensor_pointers(weights, fp8_recipe) + if single_param: + check_grouped_tensor_pointers(weights, fp8_recipe) inp_hidden_states = torch.randn( num_tokens, config.hidden_size, dtype=dtype, requires_grad=True @@ -636,6 +649,9 @@ def test_sanity_grouped_linear( loss.backward() assert out.shape == (num_tokens, ffn_hidden_size) + if single_param: + del os.environ["NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS"] + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index b6596bc2e9..2f859e748b 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -6,6 +6,7 @@ from typing import Union, Optional, Callable, Tuple, List from itertools import chain import warnings +import os import functools import torch @@ -793,7 +794,9 @@ def make_grouped_weights(self, defer_init=False) -> None: def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) - self.make_grouped_weights(defer_init=defer_init) + # Grouped tensor weights is an opt-in feature. + if bool(int(os.getenv("NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS", "0"))): + self.make_grouped_weights(defer_init=defer_init) def set_tensor_parallel_attributes(self, defer_init=False) -> None: """Set attributes needed for TP""" From 5d112e3c1da829d4e5533235f766d32456b98ca0 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 13 Feb 2026 09:45:51 -0800 Subject: [PATCH 212/521] [JAX] TE Permutation integration to Maxtext (#2672) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * adding more stuff missing from cherry picky jeremy PR for inspecting Signed-off-by: tdophung * fix some tracing issues when intergating to maxtext Signed-off-by: tdophung * Have sort_chunks_by_index handle situations where input buffer is larger than num tokens Signed-off-by: tdophung * remove unnecessary assert and comments Signed-off-by: JAX Toolbox * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove Jeremy's PR for inspect ffi Signed-off-by: JAX Toolbox * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * untouch the amax file, also change comment on te Signed-off-by: JAX Toolbox --------- Signed-off-by: tdophung Signed-off-by: JAX Toolbox Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: JAX Toolbox --- .../common/triton/permutation.py | 12 ++++++++++ transformer_engine/jax/permutation.py | 22 +++++++++++-------- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index 4602f41cfd..147742bb05 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -563,6 +563,13 @@ def _make_chunk_sort_map_kernel( split_sizes_ptr + load_split_offset, mask=load_split_offset < num_splits, other=0 ).to(tl.int32) input_split_sizes_cumsum = tl.cumsum(input_split_sizes) + + # Compute total valid tokens and skip phantom/padding tokens. + # When the input buffer is larger than sum(split_sizes), tokens beyond + # the valid range should map to themselves (identity mapping) to avoid + # corrupting valid output positions. + total_valid_tokens = tl.sum(input_split_sizes) + input_split_sizes_mask = tl.where(input_split_sizes_cumsum <= pid, 1, 0) input_chunk_idx = tl.sum(input_split_sizes_mask) input_split_sizes_presum = tl.sum(input_split_sizes * input_split_sizes_mask) @@ -578,6 +585,11 @@ def _make_chunk_sort_map_kernel( ).to(tl.int32) output_pre_split_sizes = tl.where(load_split_offset < output_chunk_idx, output_split_sizes, 0) dst_row = tl.sum(output_pre_split_sizes) + in_chunk_offset + + # For tokens beyond the valid range (pid >= total_valid_tokens), + # use identity mapping to avoid corrupting valid data + dst_row = tl.where(pid < total_valid_tokens, dst_row, pid) + tl.store(dst_rows_ptr + pid, dst_row) diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 438511fa55..6a0a3229d9 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -581,7 +581,7 @@ def sort_chunks_by_index( return _sort_chunks_by_index(inp, split_sizes, sorted_indices) -@partial(jax.custom_vjp, nondiff_argnums=(1, 2)) +@jax.custom_vjp def _sort_chunks_by_index( inp: jnp.ndarray, split_sizes: jnp.ndarray, @@ -596,7 +596,7 @@ def _sort_chunks_by_index_fwd_rule( inp: jnp.ndarray, split_sizes: jnp.ndarray, sorted_indices: jnp.ndarray, -) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, int, int]]: +) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, int, int]]: """Forward pass rule for sort_chunks_by_index.""" # Validate input dimensions assert inp.ndim in [2, 3], f"inp must be 2D or 3D, got {inp.ndim}D" @@ -618,18 +618,17 @@ def _sort_chunks_by_index_fwd_rule( ) # Return (primals, residuals) - residuals = (row_id_map, num_tokens, hidden_size) + # Include split_sizes and sorted_indices in residuals since we removed nondiff_argnums + residuals = (row_id_map, split_sizes, sorted_indices, num_tokens, hidden_size) return (output, row_id_map), residuals def _sort_chunks_by_index_bwd_rule( - _split_sizes: jnp.ndarray, - _sorted_indices: jnp.ndarray, - residuals: Tuple[jnp.ndarray, int, int], + residuals: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, int, int], g: Tuple[jnp.ndarray, jnp.ndarray], -) -> Tuple[jnp.ndarray]: +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]: """Backward pass rule for sort_chunks_by_index.""" - row_id_map, num_tokens, hidden_size = residuals + row_id_map, split_sizes, sorted_indices, num_tokens, hidden_size = residuals output_grad, _ = g # Backward: reverse the sort @@ -642,7 +641,12 @@ def _sort_chunks_by_index_bwd_rule( is_forward=False, ) - return (inp_grad,) + # Return gradients for all inputs: (inp, split_sizes, sorted_indices) + # split_sizes and sorted_indices are integer arrays, so their gradients are zeros + split_sizes_grad = jnp.zeros_like(split_sizes, dtype=split_sizes.dtype) + sorted_indices_grad = jnp.zeros_like(sorted_indices, dtype=sorted_indices.dtype) + + return (inp_grad, split_sizes_grad, sorted_indices_grad) _sort_chunks_by_index.defvjp(_sort_chunks_by_index_fwd_rule, _sort_chunks_by_index_bwd_rule) From fa68781cba1f5673b01a9024f44c5b9b8c0ba00e Mon Sep 17 00:00:00 2001 From: Hemil Desai Date: Tue, 17 Feb 2026 08:51:35 -0800 Subject: [PATCH 213/521] Fix `build_tools` missing from sdist causing `uv` cached installs to fail (#2684) - Include `build_tools/` in the source distribution via `MANIFEST.in` so that cached builds from `uv` (and `pip`) can resolve `setup.py`'s top-level imports `setup.py` imports from `build_tools` at the top level: ```python from build_tools.build_ext import CMakeExtension, get_build_ext from build_tools.te_version import te_version from build_tools.utils import cuda_archs, cuda_version, ... ``` The `__legacy__` build backend in `pyproject.toml` adds the source root to `sys.path`, so these imports work when building directly from the source tree. However, `build_tools/` is not included in the sdist because: 1. `MANIFEST.in` did not list it 2. `build_tools/` is not discovered by `find_packages()` (it's a standalone directory at the repo root, not under `transformer_engine/`) When `uv` caches the sdist and later builds a wheel from it, the sdist is extracted to a temporary directory where `build_tools/` is absent, causing a `ModuleNotFoundError`. Passing `--no-cache` to `uv` works around this by forcing a fresh build from the full source tree. Added `build_tools` to `MANIFEST.in`: ```diff recursive-include transformer_engine/common/include *.* +recursive-include build_tools *.py *.txt ``` - [x] `python setup.py sdist` produces a tarball that contains `build_tools/` ``` $ tar tzf dist/transformer_engine-*.tar.gz | grep build_tools transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/ transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/VERSION.txt transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/__init__.py transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/build_ext.py transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/jax.py transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/pytorch.py transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/te_version.py transformer_engine-2.13.0.dev0+82f7ebeb/build_tools/utils.py ``` Signed-off-by: Hemil Desai Co-authored-by: Claude Opus 4.6 --- MANIFEST.in | 1 + 1 file changed, 1 insertion(+) diff --git a/MANIFEST.in b/MANIFEST.in index c34025772a..c2309a0370 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1 +1,2 @@ recursive-include transformer_engine/common/include *.* +recursive-include build_tools *.py *.txt From 7e48fa1bace10749e2eabe1da4bbe045bacc005d Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 17 Feb 2026 09:04:04 -0800 Subject: [PATCH 214/521] [JAX] Debugging inspect utility (#2651) * initial debug of inspect ffi Signed-off-by: Jeremy Berchtold * writing binary dumps of tensors works Signed-off-by: Jeremy Berchtold * loading works Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add tensor statistics Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint Signed-off-by: Jeremy Berchtold * Add cuda error check and tests Signed-off-by: Jeremy Berchtold * Ad __init__.py to debug folder Signed-off-by: Jeremy Berchtold * Fix lint Signed-off-by: Jeremy Berchtold * Fix lint Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address greptile comments Signed-off-by: Jeremy Berchtold * Lint Signed-off-by: Jeremy Berchtold * Gate tests behind fp8 support Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_custom_call_compute.py | 34 ++++ transformer_engine/jax/csrc/extensions.h | 3 + .../jax/csrc/extensions/amax.cpp | 2 - .../jax/csrc/extensions/inspect.cpp | 99 ++++++++++ .../jax/csrc/extensions/pybind.cpp | 3 + transformer_engine/jax/debug/__init__.py | 11 ++ .../jax/debug/experimental/__init__.py | 14 ++ .../jax/debug/experimental/inspect.py | 174 ++++++++++++++++++ 8 files changed, 338 insertions(+), 2 deletions(-) create mode 100644 transformer_engine/jax/csrc/extensions/inspect.cpp create mode 100644 transformer_engine/jax/debug/__init__.py create mode 100644 transformer_engine/jax/debug/experimental/__init__.py create mode 100644 transformer_engine/jax/debug/experimental/inspect.py diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 80fcc68843..613aefc178 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -1921,3 +1921,37 @@ def test_grouped_dense_grad_fp8(self, fwd_bwd_dtype, scaling_mode, input_shape): assert_allclose(prim_dgrad, ref_dgrad, dtype=bwd_dtype) assert_allclose(prim_wgrad, ref_wgrad, dtype=bwd_dtype) assert_allclose(prim_dbias, ref_dbias, dtype=dtype) + + +class TestDebugInspectFFI: + + @pytest_parametrize_wrapper("shape", [(256, 128)]) + @pytest_parametrize_wrapper( + "dtype", + [ + jnp.float32, + jnp.bfloat16, + jnp.float16, + # Note: fp4 currently doesn't work + # jnp.float4_e2m1fn + ] + + ([jnp.float8_e4m3fn, jnp.float8_e5m2] if is_fp8_supported else []), + ) + def test_debug_inspect_ffi(self, shape, dtype): + from transformer_engine.jax.debug.experimental import inspect_array, load_array_dump + + def f(x): + x = x + 1 + x = inspect_array(x, "my_array") + x = x + 1 + return x + + key = jax.random.PRNGKey(0) + x = jax.random.uniform(key, shape, jnp.float32) + x = x.astype(dtype) + _ = jax.jit(f)(x) + + expected = x + 1 + actual = load_array_dump("my_tensor_gpu0.bin", shape, dtype) + + assert_allclose(actual, expected, dtype=dtype) diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 3fd086e257..1c0bc52b88 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -143,6 +143,9 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationInitializeHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationHandler); +// Inspect +XLA_FFI_DECLARE_HANDLER_SYMBOL(InspectHandler); + // Cudnn helpers XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnHandleInitHandler); diff --git a/transformer_engine/jax/csrc/extensions/amax.cpp b/transformer_engine/jax/csrc/extensions/amax.cpp index 5ffccaffb4..58c89cfd32 100644 --- a/transformer_engine/jax/csrc/extensions/amax.cpp +++ b/transformer_engine/jax/csrc/extensions/amax.cpp @@ -5,8 +5,6 @@ ************************************************************************/ #include -#include - #include "../extensions.h" #include "transformer_engine/cast.h" #include "transformer_engine/hadamard_transform.h" diff --git a/transformer_engine/jax/csrc/extensions/inspect.cpp b/transformer_engine/jax/csrc/extensions/inspect.cpp new file mode 100644 index 0000000000..9012cd054c --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/inspect.cpp @@ -0,0 +1,99 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ +#include + +#include +#include + +#include "../extensions.h" +#include "xla/ffi/api/c_api.h" + +namespace transformer_engine { +namespace jax { + +Error_Type InspectFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type min_buf, + Buffer_Type max_buf, Buffer_Type mean_buf, Buffer_Type std_buf, + Result_Type output_buf) { + NVTE_CHECK(input_buf.untyped_data() != nullptr, "Input must be provided for inspect operation"); + NVTE_CHECK(output_buf->untyped_data() != nullptr, + "Output must be provided for inspect operation"); + NVTE_CHECK(input_buf.untyped_data() == output_buf->untyped_data(), + "Input and output must point to the same buffer for inspect operation"); + + std::vector input_data(input_buf.size_bytes()); + NVTE_CHECK_CUDA(cudaMemcpyAsync(input_data.data(), input_buf.untyped_data(), + input_buf.size_bytes(), cudaMemcpyDeviceToHost, stream)); + + float min_val{}, max_val{}, mean_val{}, std_val{}; + NVTE_CHECK_CUDA(cudaMemcpyAsync(&min_val, min_buf.untyped_data(), sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NVTE_CHECK_CUDA(cudaMemcpyAsync(&max_val, max_buf.untyped_data(), sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NVTE_CHECK_CUDA(cudaMemcpyAsync(&mean_val, mean_buf.untyped_data(), sizeof(float), + cudaMemcpyDeviceToHost, stream)); + NVTE_CHECK_CUDA(cudaMemcpyAsync(&std_val, std_buf.untyped_data(), sizeof(float), + cudaMemcpyDeviceToHost, stream)); + + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + int device; + NVTE_CHECK_CUDA(cudaGetDevice(&device)); + + // Write the tensor data to a file as a binary blob + std::string filename = "my_tensor_gpu" + std::to_string(device) + ".bin"; + std::ofstream file(filename, std::ios::binary); + NVTE_CHECK(file.is_open(), "Failed to create file: ", filename); + file.write(reinterpret_cast(input_data.data()), input_data.size()); + file.close(); + + // Write out a metadata file + std::string meta_filename = "my_tensor_gpu" + std::to_string(device) + "_meta.json"; + std::ofstream meta_file(meta_filename); + NVTE_CHECK(meta_file.is_open(), "Failed to create file: ", meta_filename); + meta_file << "{"; + meta_file << "\"shape\": ["; + for (size_t i = 0; i < input_buf.dimensions().size(); ++i) { + meta_file << input_buf.dimensions()[i]; + if (i < input_buf.dimensions().size() - 1) { + meta_file << ", "; + } + } + meta_file << "], "; + meta_file << "\"dtype\": " << static_cast(input_buf.element_type()); + meta_file << ", \"min\": " << min_val; + meta_file << ", \"max\": " << max_val; + meta_file << ", \"mean\": " << mean_val; + meta_file << ", \"std\": " << std_val; + meta_file << "}"; + meta_file.close(); + + // Log the tensor metadata to the console + printf("[gpu%d]: Tensor data written to %s (shape: [", device, filename.c_str()); + for (size_t i = 0; i < input_buf.dimensions().size(); ++i) { + printf("%zu", static_cast(input_buf.dimensions()[i])); + if (i < input_buf.dimensions().size() - 1) { + printf(", "); + } + } + printf("], dtype: %d", static_cast(input_buf.element_type())); + printf(", min: %f, max: %f, mean: %f, std: %f)\n", min_val, max_val, mean_val, std_val); + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(InspectHandler, InspectFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // input + .Arg() // min + .Arg() // max + .Arg() // mean + .Arg() // std + .Ret() // output +); + +} // namespace jax +} // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index bd4b8fe2c2..71de897d9b 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -81,6 +81,9 @@ pybind11::dict Registrations() { pybind11::arg("initialize") = EncapsulateFFI(RHTAmaxCalculationInitializeHandler), pybind11::arg("execute") = EncapsulateFFI(RHTAmaxCalculationHandler)); + dict["te_inspect_ffi"] = + pybind11::dict(pybind11::arg("execute") = EncapsulateFFI(InspectHandler)); + return dict; } diff --git a/transformer_engine/jax/debug/__init__.py b/transformer_engine/jax/debug/__init__.py new file mode 100644 index 0000000000..7fcf194d75 --- /dev/null +++ b/transformer_engine/jax/debug/__init__.py @@ -0,0 +1,11 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""EXPERIMENTAL debugging utilities for Transformer Engine JAX. + +This API is experimental and may change or be removed without deprecation in future releases. +""" + +__all__ = [ + "experimental", +] diff --git a/transformer_engine/jax/debug/experimental/__init__.py b/transformer_engine/jax/debug/experimental/__init__.py new file mode 100644 index 0000000000..44a4847660 --- /dev/null +++ b/transformer_engine/jax/debug/experimental/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""EXPERIMENTAL debugging utilities for Transformer Engine JAX. + +This API is experimental and may change or be removed without deprecation in future releases. +""" + +from .inspect import inspect_array, load_array_dump + +__all__ = [ + "inspect_array", + "load_array_dump", +] diff --git a/transformer_engine/jax/debug/experimental/inspect.py b/transformer_engine/jax/debug/experimental/inspect.py new file mode 100644 index 0000000000..9ce46426cf --- /dev/null +++ b/transformer_engine/jax/debug/experimental/inspect.py @@ -0,0 +1,174 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Experimental JAX array inspection utilities.""" + +from functools import partial + +import jax +import jax.numpy as jnp +from jax import ffi + +from transformer_engine.jax.cpp_extensions.base import BasePrimitive, register_primitive + +__all__ = ["inspect_array", "load_array_dump"] + + +class InspectPrimitive(BasePrimitive): + """ + No-op used for inspect array values. + """ + + name = "te_inspect_ffi" + multiple_results = False + impl_static_args = () + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + x_aval, + x_min_aval, + x_max_aval, + x_mean_aval, + x_std_aval, + ): + """ + inspect abstract + """ + assert ( + x_min_aval.shape == () and x_min_aval.dtype == jnp.float32 + ), "x_min must be a scalar with dtype float32" + assert ( + x_max_aval.shape == () and x_max_aval.dtype == jnp.float32 + ), "x_max must be a scalar with dtype float32" + assert ( + x_mean_aval.shape == () and x_mean_aval.dtype == jnp.float32 + ), "x_mean must be a scalar with dtype float32" + assert ( + x_std_aval.shape == () and x_std_aval.dtype == jnp.float32 + ), "x_std must be a scalar with dtype float32" + return x_aval + + @staticmethod + def lowering( + ctx, + x, + x_min, + x_max, + x_mean, + x_std, + ): + """ + inspect lowering rules + """ + + return ffi.ffi_lowering( + InspectPrimitive.name, + operand_output_aliases={0: 0}, # donate input buffer to output buffer + )( + ctx, + x, + x_min, + x_max, + x_mean, + x_std, + ) + + @staticmethod + def impl( + x, + x_min, + x_max, + x_mean, + x_std, + ): + """ + inspect implementation + """ + assert InspectPrimitive.inner_primitive is not None + (x) = InspectPrimitive.inner_primitive.bind( + x, + x_min, + x_max, + x_mean, + x_std, + ) + return x + + +register_primitive(InspectPrimitive) + + +def _inspect_array_inner(x: jnp.ndarray) -> jnp.ndarray: + assert InspectPrimitive.outer_primitive is not None, ( + "InspectPrimitive FFI is not registered. Please ensure the C++ extension is properly built" + " and registered." + ) + return InspectPrimitive.outer_primitive.bind( + x, + jnp.min(x).astype(jnp.float32), + jnp.max(x).astype(jnp.float32), + jnp.mean(x.astype(jnp.float32)), + jnp.std(x.astype(jnp.float32)), + ) + + +@partial(jax.custom_vjp, nondiff_argnums=()) +def _inspect( + x, +): + """ """ + output, _ = _inspect_fwd_rule( + x, + ) + return output + + +def _inspect_fwd_rule( + x, +): + """""" + ctx = () + x = _inspect_array_inner(x) + return x, ctx + + +def _inspect_bwd_rule( + ctx, + grad, +): + """""" + del ctx + return (grad,) + + +_inspect.defvjp(_inspect_fwd_rule, _inspect_bwd_rule) + + +def inspect_array(x: jnp.ndarray, name: str) -> jnp.ndarray: + """Utility function to inspect JAX arrays by printing their name, shape, dtype, and statistics. + + Args: + x (jnp.ndarray): The JAX array to inspect. + name (str): The name of the array for identification in the output. + """ + del name # Name is currently unused, but can be included in the future for more informative output + return _inspect(x) + + +def load_array_dump(filename: str, shape: tuple, dtype: jnp.dtype) -> jnp.ndarray: + """Utility function to load a JAX array from a dumped binary file. + + Args: + filename (str): The path to the binary file containing the array data. + shape (tuple): The shape of the array to be loaded. + dtype (jnp.dtype): The data type of the array to be loaded. + + Returns: + jnp.ndarray: The loaded JAX array. + """ + with open(filename, "rb") as f: + data = f.read() + array = jnp.frombuffer(data, dtype=dtype).reshape(shape) + return array From f122b07d95d184ed52441c0216522d0eb7972646 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Tue, 17 Feb 2026 21:35:19 -0800 Subject: [PATCH 215/521] Changed VERSION to 2.14.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 90cc92ea66..c7d5307735 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.13.0.dev0 +2.14.0.dev0 From 2d0d276ea2e585382a1981b957b7a7cd649b2d49 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Wed, 18 Feb 2026 15:47:28 -0800 Subject: [PATCH 216/521] [PyT] Plumbing correct bias dims from TE to cudnn, while adding support for additional bias shapes (#2537) * Plumbing correct bias dims from TE to cudnn Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make changes for cp bias code Signed-off-by: Kshitij Lakhani * Add dbias and dbias_ to run_dpa_with_cp test Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix: Use output_dBias instead of input_dBias to extract the shape Signed-off-by: Kshitij Lakhani * Add guards for bias/bias_/dbias/dbias_ being None Signed-off-by: Kshitij Lakhani * Add support for bias shape 111s in addition to the original 1hss, 11ss, b1ss and bhss Signed-off-by: Kshitij Lakhani * Add support for dbias calculation and variant packing for the dbias shapes b1ss, bhss, 11ss in addition to the already supported 1hss Signed-off-by: Kshitij Lakhani * Add support for 111s bias shape in DPA Signed-off-by: Kshitij Lakhani * Allow fused attn for dbias calculation for 11ss, b1ss, bhss. Disable fused attn if dbias calculation for 111s is required, else enable Signed-off-by: Kshitij Lakhani * Disable requires_grad for bias for shape 111s in tests Signed-off-by: Kshitij Lakhani * Disable bias grad / training flag for 111s bias in the non-CP attn tests. Add bias shape 111s to test_dpa_bias_shapes Signed-off-by: Kshitij Lakhani * Fix to correctly create the bias shape tensor instead of the hard coded shape. Fix the comparison logic shapes for bias/dbias Signed-off-by: Kshitij Lakhani * Add fused attn cp test cases for all supported bias shapes Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: switch to elif for bias grad conditional Signed-off-by: Kshitij Lakhani * Add CP support for bias/dbias shape 111s Signed-off-by: Kshitij Lakhani * Add support for is_training in CP attn tests Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Fix incorrect comment Signed-off-by: Kshitij Lakhani * nit: Fix incorrect comment and assert string Signed-off-by: Kshitij Lakhani * Create the dbias graph tensor only if it is a cuDNN supported bias shape Signed-off-by: Kshitij Lakhani * Fix the dim that is being compared for the two cp chunks in the test Signed-off-by: Kshitij Lakhani * nit: Reinstate the original test for right side swa Signed-off-by: Kshitij Lakhani --------- Signed-off-by: Kshitij Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../attention/run_attention_with_cp.py | 413 ++++++++++++------ tests/pytorch/attention/test_attention.py | 22 +- .../attention/test_attention_with_cp.py | 38 +- tests/pytorch/utils.py | 5 +- .../fused_attn_f16_arbitrary_seqlen.cu | 105 +++-- .../common/fused_attn/fused_attn_fp8.cu | 31 +- transformer_engine/common/fused_attn/utils.h | 23 +- .../dot_product_attention/context_parallel.py | 79 ++-- .../dot_product_attention.py | 7 +- .../attention/dot_product_attention/utils.py | 11 +- 10 files changed, 501 insertions(+), 233 deletions(-) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 3efb516b57..0f36a8816d 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -179,10 +179,13 @@ def run_dpa_with_cp( fp8_mha="False", scaling_mode="delayed", f16_O="False", + is_training="True", log_level=logging.WARNING, ): """Test DotProductAttention module with context parallelism""" logging.root.setLevel(log_level) + # When is_training is False, gradient outputs are None. + is_training = is_training == "True" # set up environment variables and config fp8_bwd = fp8_bwd == "True" and dtype == "fp8" @@ -257,7 +260,9 @@ def run_dpa_with_cp( softmax_type=config.softmax_type, return_max_logit=config.return_max_logit, ).cuda() - if config.softmax_type != "vanilla": + if not is_training: + core_attn.eval() + if is_training and config.softmax_type != "vanilla": core_attn.softmax_offset.requires_grad = True # generate attention inputs @@ -305,8 +310,25 @@ def run_dpa_with_cp( x.requires_grad = True if config.attn_bias_type not in ["no_bias", "alibi"]: - attn_bias_shape = (1, 1, config.max_seqlen_q, config.max_seqlen_kv) + bias_shape_map = { + "1hss": (1, config.num_heads, config.max_seqlen_q, config.max_seqlen_kv), + "11ss": (1, 1, config.max_seqlen_q, config.max_seqlen_kv), + "b1ss": (config.batch_size, 1, config.max_seqlen_q, config.max_seqlen_kv), + "bhss": ( + config.batch_size, + config.num_heads, + config.max_seqlen_q, + config.max_seqlen_kv, + ), + "111s": (1, 1, 1, config.max_seqlen_kv), + } + attn_bias_shape = bias_shape_map.get(config.bias_shape) + if attn_bias_shape is None: + assert False, f"cuDNN does not support {config.bias_shape=}" bias = torch.randn(*attn_bias_shape, dtype=dtypes[dtype]).cuda() + # cuDNN does not support dbias calculation for 111s as of cuDNN 9.18 + # TODO(KshitijLakhani): Set requires_grad to True for all shapes once 111s is supported + bias.requires_grad = True if config.bias_shape != "111s" else False else: bias = None @@ -333,15 +355,20 @@ def run_dpa_with_cp( ) if config.return_max_logit: out, max_logit = out - if fp8_bwd and fp8_mha: - dout_fp8 = dout_quantizer(dout) - out.backward(dout_fp8) - else: - out.backward(dout) - dq, dk, dv = q.grad, k.grad, v.grad - d_softmax_offset = None - if config.softmax_type != "vanilla": - d_softmax_offset = core_attn.softmax_offset.grad + if is_training: + if fp8_bwd and fp8_mha: + dout_fp8 = dout_quantizer(dout) + out.backward(dout_fp8) + else: + out.backward(dout) + if is_training: + dq, dk, dv, dbias = q.grad, k.grad, v.grad, bias.grad if bias is not None else None + d_softmax_offset = ( + core_attn.softmax_offset.grad if config.softmax_type != "vanilla" else None + ) + else: + dq, dk, dv, dbias = None, None, None, None + d_softmax_offset = None ############ run with CP ############ logging.info(f"[Rank {rank}] Run with context parallelism") @@ -387,13 +414,30 @@ def run_dpa_with_cp( 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_ = [x.requires_grad_() for x in [q_, k_, v_]] + if is_training: + q_, k_, v_ = [x.requires_grad_() for x in [q_, k_, v_]] if bias_ is not None: - bias_ = bias_.view( - *bias_.shape[:-2], 2 * world_size, bias_.shape[-2] // (2 * world_size), bias_.shape[-1] - ) - bias_ = bias_.index_select(2, seq_idx) - bias_ = bias_.view(*bias_.shape[:2], -1, bias_.shape[-1]) + ndim = bias_.ndim + seq_q_dim = ndim - 2 + if qkv_format == "thd": + bias_seq_idx = seq_idx_q + else: + bias_seq_idx = seq_idx + shape_before_seq = bias_.shape[:seq_q_dim] + seq_q_size = bias_.shape[seq_q_dim] + seq_kv_size = bias_.shape[-1] + if seq_q_size == 1: + # TODO(KshitijLakhani): Set to True always once cuDNN supports dbias for 111s + bias_.requires_grad = False + # Bias is broadcast, no need to partition along sequence dimension + pass + else: + bias_ = bias_.view( + *shape_before_seq, 2 * world_size, seq_q_size // (2 * world_size), seq_kv_size + ) + bias_ = bias_.index_select(seq_q_dim, bias_seq_idx) + bias_ = bias_.view(*shape_before_seq, -1, seq_kv_size) + bias_.requires_grad = True # set up environment core_attn.set_context_parallel_group( cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else cp_comm_group, @@ -428,90 +472,143 @@ def run_dpa_with_cp( ) if config.return_max_logit: out_, max_logit_ = out_ - if fp8_bwd and fp8_mha: - dout_fp8_ = dout_quantizer(dout_) - out_.backward(dout_fp8_) - else: - out_.backward(dout_) - dq_, dk_, dv_ = q_.grad, k_.grad, v_.grad - d_softmax_offset_ = None - if config.softmax_type != "vanilla": - d_softmax_offset_ = core_attn.softmax_offset.grad.clone() + if is_training: + if fp8_bwd and fp8_mha: + dout_fp8_ = dout_quantizer(dout_) + out_.backward(dout_fp8_) + else: + out_.backward(dout_) + if is_training: + dq_, dk_, dv_, dbias_ = ( + q_.grad, + k_.grad, + v_.grad, + bias_.grad if bias_ is not None else None, + ) + d_softmax_offset_ = ( + core_attn.softmax_offset.grad.clone() if config.softmax_type != "vanilla" else None + ) + else: + dq_, dk_, dv_, dbias_ = None, None, None, None + d_softmax_offset_ = None # get outputs - tensors = [out, dq, dk, dv, out_, dq_, dk_, dv_] + tensors = [out, dq, dk, dv, dbias, out_, dq_, dk_, dv_, dbias_] if fp8_mha: tensors_to_deq = [out, out_] if not fp8_bwd else tensors for i, tensor in enumerate(tensors_to_deq): - tensors_to_deq[i] = tensor.dequantize() + # dbias/dbias_ could be None, so skip check for it + if tensor is not None: + tensors_to_deq[i] = tensor.dequantize() if not fp8_bwd: - tensors[0], tensors[4] = tensors_to_deq + tensors[0], tensors[5] = tensors_to_deq for tensor in tensors: - assert torch.all(~torch.isnan(tensor)) - assert torch.all(~torch.isinf(tensor)) - out, dq, dk, dv, out_, dq_, dk_, dv_ = 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)) + out, dq, dk, dv, dbias, out_, dq_, dk_, dv_, dbias_ = tensors ############ compare results between CP and no-CP ############ if qkv_format == "bshd" or qkv_format == "sbhd": - dq, dk, dv, out = [ - x.view( - *x.shape[:seq_dim], + if is_training: + dq, dk, dv, out = [ + x.view( + *x.shape[:seq_dim], + 2 * world_size, + x.shape[seq_dim] // (2 * world_size), + *x.shape[(seq_dim + 1) :], + ) + for x in [dq, dk, dv, out] + ] + dq, dk, dv, out = [x.index_select(seq_dim, seq_idx) for x in [dq, dk, dv, out]] + dq_, dk_, dv_, out_ = [ + x.view(*x.shape[:seq_dim], 2, x.shape[seq_dim] // 2, *x.shape[(seq_dim + 1) :]) + for x in [dq_, dk_, dv_, out_] + ] + if dbias is not None and dbias_ is not None: + ndim = dbias.ndim + # Query seq is at dim -2 + seq_q_dim = ndim - 2 + shape_before_seq = dbias.shape[:seq_q_dim] + seq_q_size = dbias.shape[seq_q_dim] + 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 + ) + # Index select on the newly created dimension (now at position seq_q_dim) + dbias = dbias.index_select(seq_q_dim, seq_idx) + dbias_ = dbias_.view( + *shape_before_seq, 2, dbias_.shape[seq_q_dim] // 2, seq_kv_size + ) + else: + # Forward-only: reshape only out/out_ for comparison + out = out.view( + *out.shape[:seq_dim], 2 * world_size, - x.shape[seq_dim] // (2 * world_size), - *x.shape[(seq_dim + 1) :], + out.shape[seq_dim] // (2 * world_size), + *out.shape[(seq_dim + 1) :], ) - for x in [dq, dk, dv, out] - ] - dq, dk, dv, out = [x.index_select(seq_dim, seq_idx) for x in [dq, dk, dv, out]] - dq_, dk_, dv_, out_ = [ - x.view(*x.shape[:seq_dim], 2, x.shape[seq_dim] // 2, *x.shape[(seq_dim + 1) :]) - for x in [dq_, dk_, dv_, out_] - ] + out = out.index_select(seq_dim, seq_idx) + out_ = out_.view( + *out_.shape[:seq_dim], 2, out_.shape[seq_dim] // 2, *out_.shape[(seq_dim + 1) :] + ) + elif qkv_format == "thd": - 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 - ) - 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 + 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 - ) + ).item() + == 0 + ) + 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 + ) + else: + # Forward-only: reshape only out/out_ for comparison + out = out.index_select(0, seq_idx_q).contiguous() + out_ = out_ atol, rtol, rmse_tol = get_tols(config, dtype) - tensors_cp = [out_, dq_, dk_, dv_, d_softmax_offset_, max_logit_] - tensors_no_cp = [out, dq, dk, dv, d_softmax_offset, max_logit] - names = ["out", "dq", "dk", "dv", "d_softmax_offset", "max_logit"] + tensors_cp = [out_, dq_, dk_, dv_, dbias_, d_softmax_offset_, max_logit_] + tensors_no_cp = [out, dq, dk, dv, dbias, d_softmax_offset, max_logit] + names = ["out", "dq", "dk", "dv", "dbias", "d_softmax_offset", "max_logit"] names_cp = [x + "_cp" for x in names] names_no_cp = [x + "_no_cp" for x in names] is_fp8 = dtype == "fp8" @@ -519,47 +616,113 @@ def run_dpa_with_cp( if t is not None: if "softmax_offset" not in names[i] and "max_logit" not in names[i]: if qkv_format == "bshd": - compare_and_assert( - t[:, 0], - tensors_cp[i][:, 0], - names_no_cp[i], - names_cp[i], - atol, - rtol, - rmse_tol, - is_fp8, - ) - compare_and_assert( - t[:, 1], - tensors_cp[i][:, 1], - names_no_cp[i], - names_cp[i], - atol, - rtol, - rmse_tol, - is_fp8, - ) + # Compare the two sequence chunks separately + # Compare dbias + if names[i] == "dbias": + # Compare the two chunks along dimension 2 (the split sequence dimension) + seq_q_dim_bias = 2 + ndim_bias = t.ndim + slice_0 = [slice(None)] * ndim_bias + slice_0[seq_q_dim_bias] = 0 + slice_1 = [slice(None)] * ndim_bias + slice_1[seq_q_dim_bias] = 1 + compare_and_assert( + t[tuple(slice_0)], + tensors_cp[i][tuple(slice_0)], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) + compare_and_assert( + t[tuple(slice_1)], + tensors_cp[i][tuple(slice_1)], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) + # Compare Q/K/V/out + else: + # Compare the two chunks along dimension 1 (the split sequence dimension) + compare_and_assert( + t[:, 0], + tensors_cp[i][:, 0], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) + compare_and_assert( + t[:, 1], + tensors_cp[i][:, 1], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) elif qkv_format == "sbhd": - compare_and_assert( - t[0], - tensors_cp[i][0], - names_no_cp[i], - names_cp[i], - atol, - rtol, - rmse_tol, - is_fp8, - ) - compare_and_assert( - t[1], - tensors_cp[i][1], - names_no_cp[i], - names_cp[i], - atol, - rtol, - rmse_tol, - is_fp8, - ) + # Compare the two sequence chunks separately + # Compare dbias (same as BSHD) + if names[i] == "dbias": + # Same as bshd: Compare the two chunks along dimension 2 (the split sequence dimension) + seq_q_dim_bias = 2 + ndim_bias = t.ndim + slice_0 = [slice(None)] * ndim_bias + slice_0[seq_q_dim_bias] = 0 + slice_1 = [slice(None)] * ndim_bias + slice_1[seq_q_dim_bias] = 1 + compare_and_assert( + t[tuple(slice_0)], + tensors_cp[i][tuple(slice_0)], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) + compare_and_assert( + t[tuple(slice_1)], + tensors_cp[i][tuple(slice_1)], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) + # Compare Q/K/V/out + else: + # Compare the two chunks along dimension 0 (the split sequence dimension) + compare_and_assert( + t[0], + tensors_cp[i][0], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) + compare_and_assert( + t[1], + tensors_cp[i][1], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, + ) elif qkv_format == "thd": compare_and_assert( t, tensors_cp[i], names_no_cp[i], names_cp[i], atol, rtol, rmse_tol, is_fp8 diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index bd0ac41974..01b2aac453 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -162,7 +162,16 @@ def test_dot_product_attention( ) # Get backends + # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s. + # For all other shapes test fwd+bwd is_training = True + # TODO(KshitijLakhani): Set is_training to True for all cases once cuDNN supports dbias for 111s. + if config.bias_shape == "111s": + is_training = False + logging.info( + "Setting is_training to False as cuDNN does not support dbias for" + f" {config.bias_shape=} " + ) available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=dtype, @@ -636,7 +645,8 @@ def test_dpa_bias(dtype, model_configs, model): "bias_1_1": ModelConfig(2, 128, 16, 64, attn_bias_type="post_scale_bias", bias_shape="1hss"), "bias_1_2": ModelConfig(4, 2048, 24, 128, attn_bias_type="post_scale_bias", bias_shape="b1ss"), "bias_1_3": ModelConfig(2, 2048, 24, 128, attn_bias_type="post_scale_bias", bias_shape="bhss"), - "bias_1_4": ModelConfig( + "bias_1_4": ModelConfig(2, 2048, 24, 128, attn_bias_type="post_scale_bias", bias_shape="111s"), + "bias_1_5": ModelConfig( 4, 2048, 24, @@ -646,7 +656,7 @@ def test_dpa_bias(dtype, model_configs, model): bias_shape="1hss", alibi_type="custom", ), - "bias_1_5": ModelConfig( + "bias_1_6": ModelConfig( 2, 2048, 24, @@ -1143,10 +1153,16 @@ def _run_dot_product_attention( bias = None if config.attn_bias_type == "post_scale_bias": shape = "_".join(config.bias_shape) + # For 1hss, 11ss, b1ss, bhss + shape_cache = shape shape = shape.replace("_s_s", "_sq_skv") + # For 111s + if shape == shape_cache: + shape = shape.replace("_1_s", "_1_skv") tensor_shape = [dim_to_num[j] for j in shape.split("_")] bias = torch.randn(tensor_shape, dtype=dtype, device="cuda") - if config.bias_shape != "1hss": + # For 111s, dbias calculation is not supported as of cuDNN 9.18 + if config.bias_shape == "111s": bias.requires_grad = False # Create RNG diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 836598087b..ecd0090a3b 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -147,7 +147,10 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): 2, 4096, 12, 128, attn_mask_type="causal", attn_bias_type="post_scale_bias" ), # MHA "cp_1_3": ModelConfig(2, 4096, 12, 128, attn_bias_type="post_scale_bias"), # MHA - "cp_1_4": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", window_size=(512, 512)), # MHA + "cp_1_4": ModelConfig( + 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_2": ModelConfig( @@ -160,9 +163,30 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): attn_bias_type="post_scale_bias", ), # GQA "cp_2_3": ModelConfig( - 2, 4096, 12, 128, num_gqa_groups=2, attn_bias_type="post_scale_bias" + 2, + 4096, + 12, + 128, + num_gqa_groups=2, + attn_mask_type="causal", + attn_bias_type="post_scale_bias", + bias_shape="11ss", ), # GQA "cp_2_4": ModelConfig( + 2, + 4096, + 12, + 128, + num_gqa_groups=2, + attn_mask_type="causal", + attn_bias_type="post_scale_bias", + bias_shape="111s", + return_max_logit=True, + ), # GQA + "cp_2_5": ModelConfig( + 2, 4096, 12, 128, num_gqa_groups=2, attn_bias_type="post_scale_bias" + ), # GQA + "cp_2_6": ModelConfig( 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 @@ -171,6 +195,9 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): 2, 4096, 12, 128, attn_mask_type="causal", attn_bias_type="post_scale_bias", head_dim_v=64 ), # MLA "cp_3_3": ModelConfig(2, 4096, 12, 128, attn_bias_type="post_scale_bias", head_dim_v=64), # MLA + "cp_3_4": ModelConfig( + 2, 4096, 12, 128, attn_bias_type="post_scale_bias", bias_shape="b1ss", head_dim_v=64 + ), # MLA "cp_4_0": ModelConfig( 2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="vanilla" ), # GQA @@ -191,10 +218,13 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): "cp_1_0", "cp_1_1", "cp_1_4", + "cp_1_5", "cp_2_0", "cp_2_2", + "cp_2_3", "cp_2_4", "cp_3_2", + "cp_3_4", "cp_4_2", ] model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs} @@ -324,12 +354,15 @@ def test_cp_with_fused_attention( Float8CurrentScaling(fp8_dpa=True), DelayedScaling(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( config, 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, ) _, fused_attn_supported, _ = available_backends if not fused_attn_supported: @@ -348,6 +381,7 @@ def test_cp_with_fused_attention( fp8_mha=fp8_mha, scaling_mode=scaling_mode, f16_O=f16_O, + is_training=is_training, log_level=pytest_logging_level, ), check=True, diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index c54295d478..317240fb78 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -271,7 +271,6 @@ def get_available_attention_backends( os.environ["NVTE_FUSED_ATTN"] = "1" os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True - alibi_slopes_shape = None if config.attn_bias_type == "alibi" and config.alibi_type == "custom": if config.bias_shape == "1hss": @@ -289,7 +288,9 @@ def get_available_attention_backends( and config.head_dim_qk <= 128 and config.head_dim_v <= 128 ): - core_attention_bias_requires_grad = True + # TODO(KshitijLakhani): Remove this guard when cuDNN starts support dbias calculation for bias shape 111s + if core_attention_bias_shape != "111s": + core_attention_bias_requires_grad = True fused_attn_backends = [] available_backends = None 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 d13ed97de1..eb2ebcff39 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 @@ -52,13 +52,14 @@ void fused_attn_arbitrary_seqlen_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, int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t num_pages_k, int64_t num_pages_v, 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, 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, void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, + 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, + 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) { using namespace transformer_engine; @@ -121,6 +122,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( max_pages_per_seq_v, bias_b, bias_h, + bias_sq, + bias_skv, scaling_factor, is_training, dropout_probability, @@ -269,10 +272,11 @@ void fused_attn_arbitrary_seqlen_fwd_impl( sdpa_options.set_alibi_mask(is_alibi); if (is_bias) { - bias = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, s_q, s_kv}) - .set_stride({bias_h * s_q * s_kv, s_q * s_kv, s_kv, 1})); + bias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("bias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); sdpa_options.set_bias(bias); } @@ -548,16 +552,16 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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, - 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, - bool deterministic, void *devPtrQ, void *devPtrKTranspose, void *devPtrVTranspose, - void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, void *devPtrSoftmaxOffset, - void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, void *devPtrdBias, - void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, void *devPtrDropoutOffset, - void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrSeqOffsetsQ, - void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, void *workspace, - size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + 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_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, + void *devPtrSoftmaxOffset, void *devPtrdQ, void *devPtrdK, void *devPtrdV, void *devPtrdO, + void *devPtrdBias, void *devPtrdSoftmaxOffset, void *devPtrDropoutSeed, + void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, + void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, + void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -622,6 +626,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( 0, bias_b, bias_h, + bias_sq, + bias_skv, scaling_factor, true, dropout_probability, @@ -811,19 +817,20 @@ void fused_attn_arbitrary_seqlen_bwd_impl( sdpa_backward_options.set_alibi_mask(is_alibi); if (is_bias) { - bias = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("bias") - .set_dim({bias_b, bias_h, s_q, s_kv}) - .set_stride({bias_h * s_q * s_kv, s_q * s_kv, s_kv, 1})); - dBias = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("dBias") - .set_dim({bias_b, bias_h, s_q, s_kv}) - .set_stride({bias_h * s_q * s_kv, s_q * s_kv, s_kv, 1})); + bias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("bias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); sdpa_backward_options.set_bias(bias); - // shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s] - // are not supported for dbias calculation but they are - // supported for forward bias calculation - if ((bias_b == 1) && (bias_h == h)) { + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + dBias = mha_graph->tensor( + fe::graph::Tensor_attributes() + .set_name("dBias") + .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); sdpa_backward_options.set_dbias(dBias); } } @@ -974,10 +981,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( if (is_bias) { variant_pack[bias] = devPtrBias; - if ((bias_b == 1) && (bias_h == h)) { + if (dBias != nullptr) { variant_pack[dBias] = devPtrdBias; - } else { - variant_pack[dBias] = nullptr; } } @@ -1083,10 +1088,14 @@ void fused_attn_arbitrary_seqlen_fwd( void *devPtrBias = nullptr; size_t bias_b = 0; size_t bias_h = 0; + size_t bias_sq = 0; + size_t bias_skv = 0; if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { devPtrBias = input_Bias->data.dptr; bias_b = input_Bias->data.shape[0]; bias_h = input_Bias->data.shape[1]; + bias_sq = input_Bias->data.shape[2]; + bias_skv = input_Bias->data.shape[3]; } void *devPtrSoftmaxOffset = nullptr; if (softmax_type != NVTE_VANILLA_SOFTMAX) { @@ -1152,7 +1161,7 @@ void fused_attn_arbitrary_seqlen_fwd( if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { Tensor *output_bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_bias->data.dptr = nullptr; - output_bias->data.shape = {bias_b, bias_h, max_seqlen_q, max_seqlen_kv}; + output_bias->data.shape = {bias_b, bias_h, bias_sq, bias_skv}; output_bias->data.dtype = QKV_type; } @@ -1197,10 +1206,10 @@ void fused_attn_arbitrary_seqlen_fwd( fused_attn_arbitrary_seqlen_fwd_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, 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, 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, + 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); @@ -1244,11 +1253,15 @@ void fused_attn_arbitrary_seqlen_bwd( void *devPtrdBias = nullptr; size_t bias_b = 0; size_t bias_h = 0; + size_t bias_sq = 0; + size_t bias_skv = 0; if ((bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) && (bias_type != NVTE_Bias_Type::NVTE_ALIBI)) { devPtrBias = input_Bias->data.dptr; devPtrdBias = output_dBias->data.dptr; bias_b = output_dBias->data.shape[0]; bias_h = output_dBias->data.shape[1]; + bias_sq = output_dBias->data.shape[2]; + bias_skv = output_dBias->data.shape[3]; } size_t max_batch_size = 0; @@ -1291,11 +1304,11 @@ 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, 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, + 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); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index fe859b0b22..8c8a289746 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1671,6 +1671,8 @@ void fused_attn_fp8_fwd_impl_v1( bool is_dropout = (is_training && dropout_probability != 0.0f); 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 || @@ -1697,6 +1699,8 @@ void fused_attn_fp8_fwd_impl_v1( 0, bias_b, bias_h, + bias_sq, + bias_skv, scaling_factor, is_training, dropout_probability, @@ -1818,8 +1822,8 @@ void fused_attn_fp8_fwd_impl_v1( // if (is_bias) { // bias = mha_graph->tensor(fe::graph::Tensor_attributes() // .set_name("bias") - // .set_dim({bias_b, bias_h, s_q, s_kv}) - // .set_stride({bias_h * s_q * s_kv, s_q * s_kv, s_kv, 1})); + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); // sdpa_options.set_bias(bias); // } @@ -1999,6 +2003,8 @@ void fused_attn_fp8_bwd_impl_v1( bool is_dropout = (dropout_probability != 0.0f); 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 = (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || @@ -2027,6 +2033,8 @@ void fused_attn_fp8_bwd_impl_v1( 0, bias_b, bias_h, + bias_sq, + bias_skv, scaling_factor, true, dropout_probability, @@ -2194,19 +2202,18 @@ void fused_attn_fp8_bwd_impl_v1( // if (is_bias) { // bias = mha_graph->tensor(fe::graph::Tensor_attributes() // .set_name("bias") - // .set_dim({bias_b, bias_h, s_q, s_kv}) - // .set_stride({bias_h * s_q * s_kv, s_q * s_kv, s_kv, 1})); + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); // dBias = mha_graph->tensor(fe::graph::Tensor_attributes() // .set_name("dBias") - // .set_dim({bias_b, bias_h, s_q, s_kv}) - // .set_stride({bias_h * s_q * s_kv, s_q * s_kv, s_kv, 1})); + // .set_dim({bias_b, bias_h, bias_sq, bias_skv}) + // .set_stride({bias_h * bias_sq * bias_skv, bias_sq * bias_skv, bias_skv, 1})); // sdpa_backward_options.set_bias(bias); - // // shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s] - // // are not supported for dbias calculation but they are - // // supported for forward bias calculation - // if ((bias_b == 1) && (bias_h == h)) { - // sdpa_backward_options.set_dbias(dBias); - // } + // bias shapes [1, 1, s, s], [b, 1, s, s], [b, h, s, s], [1, h, s, s] are supported for dbias calculation + // bias shape [1, 1, 1, s] is not supported for dbias calculation as of cuDNN 9.18 + // if (!((bias_b == 1) && (bias_h == 1) && (bias_sq == 1))) { + // sdpa_backward_options.set_dbias(dBias); + // } // } if (is_padding) { diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index fdfc4abe82..08a56cda6b 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -101,6 +101,8 @@ struct FADescriptor_v1 { std::int64_t max_pages_per_seq_v; std::int64_t bias_b; std::int64_t bias_h; + std::int64_t bias_sq; + std::int64_t bias_skv; float attnScale; bool isTraining; float dropoutProbability; @@ -120,18 +122,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, - attnScale, isTraining, dropoutProbability, layout, 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, - generate_max_sum_exp) < + 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, + 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, generate_max_sum_exp) < 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.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.dqkv_tensor_type, rhs.generate_max_sum_exp); + 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.dqkv_tensor_type, rhs.generate_max_sum_exp); } }; 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 a5931188dc..bd6b626b64 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -840,13 +840,24 @@ def cp_p2p_fwd_fused_attn( q_part = q_part.contiguous() if attn_bias is not None: idx = (rank - step) % cp_size - attn_bias_inputs = torch.cat( - ( - attn_bias_[..., 1, :, idx, :], - attn_bias_[..., 1, :, (2 * cp_size - idx - 1), :], - ), - dim=-1, - ).contiguous() + # For bias shape 111s, only the s_kv dim is split, i.e. [b, h, sq, 2*cp, sk//(2*cp)]) + if attn_bias.shape[-3] == 1: + attn_bias_inputs = torch.cat( + ( + attn_bias_[..., :, idx, :], + attn_bias_[..., :, (2 * cp_size - idx - 1), :], + ), + dim=-1, + ).contiguous() + # For bias shapes 1hss, 11ss, bhss, b1ss, the s_kv and s_q dims are split, i.e. [b, h, 2, sq//2, 2*cp, sk//(2*cp)]) + else: + attn_bias_inputs = torch.cat( + ( + attn_bias_[..., 1, :, idx, :], + attn_bias_[..., 1, :, (2 * cp_size - idx - 1), :], + ), + dim=-1, + ).contiguous() max_seqlen_q_ = max_seqlen_q // 2 max_seqlen_kv_ = max_seqlen_kv cu_seqlens_q_ = cu_seqlens_q_per_step @@ -1442,20 +1453,33 @@ def forward( attn_bias_ = None if attn_bias is not None: assert len(attn_bias.shape) == 4, ( - "Only support bias shape of [b, h, sq, sk] for forward, " - "and [1, h, sq, sk] for backward!" - ) - assert ( - attn_bias.shape[-2] % 2 == 0 and attn_bias.shape[-1] % (2 * cp_size) == 0 - ), "Sequence length does not meet divisible requirements!" - # [b, h, sq, sk] -> [b, h, 2, sq//2, 2*cp, sk//(2*cp)] - attn_bias_ = attn_bias.view( - *attn_bias.shape[:-2], - 2, - attn_bias.shape[-2] // 2, - 2 * cp_size, - attn_bias.shape[-1] // (2 * cp_size), + "Only support bias shape of [1,1,sq,skv], [1,h,sq,skv], [b,1,sq,skv], [b,h,sq,skv]," + " [1,1,1,skv] for forward, and [1,1,sq,skv], [1,h,sq,skv], [b,1,sq,skv]," + " [b,h,sq,skv] for backward!" ) + # For all bias shapes except 111s, sq must be divisible by 2 and skv must be divisible by 2*cp_size + # For bias shape 111s, only skv must be divisible by 2*cp_size + if attn_bias.shape[-2] != 1: + assert ( + attn_bias.shape[-2] % 2 == 0 and attn_bias.shape[-1] % (2 * cp_size) == 0 + ), "Sequence length does not meet divisible requirements!" + # [b, h, sq, sk] -> [b, h, 2, sq//2, 2*cp, sk//(2*cp)] + attn_bias_ = attn_bias.view( + *attn_bias.shape[:-2], + 2, + attn_bias.shape[-2] // 2, + 2 * cp_size, + attn_bias.shape[-1] // (2 * cp_size), + ) + else: + assert ( + attn_bias.shape[-1] % (2 * cp_size) == 0 + ), "Sequence length does not meet divisible requirements!" + # [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) + ) + # [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) @@ -2076,10 +2100,13 @@ def backward(ctx, dout, *_args): attn_dbias = torch.zeros( *ctx.attn_bias_shape, dtype=attn_biases[0].dtype, device=attn_biases[0].device ) - # [b, h, sq, 2*cp, sk//(2*cp)] -> [b, h, 2, sq//2, 2*cp, sk//(2*cp)] - attn_dbias_ = attn_dbias.view( - *attn_dbias.shape[:-3], 2, attn_dbias.shape[-3] // 2, *attn_dbias.shape[-2:] - ) + # [b, h, sq, 2*cp, sk//(2*cp)] -> [b, h, 2, sq//2, 2*cp, sk//(2*cp)] only when sq > 1 (i.e. all supported bias shapes except 111s) + if attn_dbias.shape[-3] > 1: + attn_dbias_ = attn_dbias.view( + *attn_dbias.shape[:-3], 2, attn_dbias.shape[-3] // 2, *attn_dbias.shape[-2:] + ) + else: + attn_dbias_ = None else: attn_dbias = None attn_dbias_ = None @@ -2507,8 +2534,8 @@ def backward(ctx, dout, *_args): elif i >= (cp_size - rank - 1): # [b, h, sq, sk//(2*cp)] attn_dbias[..., idx, :].copy_(dbias_) - else: - # [b, h, sq//2, sk//cp] -> [b, h, sq//2, 2, sk//(2*cp)] + elif attn_dbias_ is not None: + # upper-triangle: [b, h, sq//2, sk//cp] -> [b, h, sq//2, 2, sk//(2*cp)] dbias_ = dbias_.view(*dbias_.shape[:-1], 2, dbias_.shape[-1] // 2) attn_dbias_[..., 1, :, idx, :].copy_(dbias_[..., 0, :]) attn_dbias_[..., 1, :, (2 * cp_size - idx - 1), :].copy_(dbias_[..., 1, :]) 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 5d830dca33..64db4646f6 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 @@ -1318,11 +1318,14 @@ def forward( ): core_attention_bias_shape = "b1ss" elif core_attention_bias.shape[0] == 1 and core_attention_bias.shape[1] == 1: - core_attention_bias_shape = "11ss" + if core_attention_bias.shape[2] == 1: + core_attention_bias_shape = "111s" + else: + core_attention_bias_shape = "11ss" else: assert ( False - ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss} shapes" + ), "core_attention_bias must be in one of {bhss, 1hss, b1ss, 11ss, 111s} shapes" # check if there is padding between sequences when qkv_format='thd' if pad_between_seqs is None: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 0c5a519813..3432fd832f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -966,12 +966,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt and fu_core_attention_bias_type == "post_scale_bias" and fu_core_attention_bias_shape != "1hss" ): - if fu_core_attention_bias_requires_grad: - # remove this line when cuDNN adds bwd support for - # [1, 1, s, s], [b, 1, s, s] and [b, h, s, s] - logger.debug("Disabling FusedAttention for dBias in [1, H, S, S] shape") + # dbias calculation is not supported for 111s as of cuDNN 9.18. So, use fused attention backend only if bias does not require grad. + if fu_core_attention_bias_requires_grad and fu_core_attention_bias_shape == "111s": + logger.warning( + "Disabling FusedAttention as dbias calculation is not supported for 111s" + ) use_fused_attention = False - else: + elif not fu_core_attention_bias_requires_grad: # max512 backend will only support [1, h, s, s] os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" From 63defeadaab055a389fac08e7266ab43860f7d30 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 19 Feb 2026 17:15:05 -0800 Subject: [PATCH 217/521] Update cudnn-frontend to v1.18 (#2689) update FE to 1.18 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index b372d39879..8d19d3182b 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit b372d39879d44c91a8d5b342022e74802b6a8da2 +Subproject commit 8d19d3182bfbc304046a15e9236bec9ff31511fc From e5832221aa901cd9b0e8315926728d2b1c4e942a Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 19 Feb 2026 17:39:35 -0800 Subject: [PATCH 218/521] [PyTorch] Documentation for op fuser API (#2447) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add documentation for operation fuser API Signed-off-by: Tim Moon * Include TE ops in PyTorch API docs Signed-off-by: Tim Moon * Fix error when building docs Signed-off-by: Tim Moon * Fix typo 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> * Fix swapped args to `te.ops.Linear` 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> * Update copyright year Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Reorganize TE ops guide Signed-off-by: Tim Moon * Fix typo Signed-off-by: Tim Moon * [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: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Debug test failure Signed-off-by: Tim Moon * Debug failure when autogenerating docs Signed-off-by: Tim Moon * Debug errors when building docs Signed-off-by: Tim Moon * Include fusion registration functions in docs Signed-off-by: Tim Moon * Update docs/api/pytorch.rst 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> * Debug failure when building docs Signed-off-by: Tim Moon * Poke GitHub Signed-off-by: Tim Moon --------- 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: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .gitignore | 2 +- docs/api/pytorch.rst | 80 ++++ .../op_fuser/fp8_layernorm_linear.png | Bin 0 -> 17749 bytes docs/examples/op_fuser/layernorm_mlp.png | Bin 0 -> 28980 bytes docs/examples/op_fuser/op_fuser.rst | 353 ++++++++++++++++++ .../op_fuser/residual_layernorm_mlp.png | Bin 0 -> 15620 bytes docs/index.rst | 1 + tests/pytorch/test_fusible_ops.py | 136 ++++++- .../pytorch/ops/basic/activation.py | 16 +- .../pytorch/ops/basic/add_extra_input.py | 2 +- .../pytorch/ops/basic/basic_linear.py | 54 +-- transformer_engine/pytorch/ops/basic/bias.py | 4 +- .../pytorch/ops/basic/grouped_linear.py | 2 +- .../pytorch/ops/basic/layer_norm.py | 8 +- .../pytorch/ops/basic/make_extra_output.py | 2 +- .../pytorch/ops/basic/quantize.py | 8 +- .../pytorch/ops/basic/reshape.py | 2 +- .../pytorch/ops/basic/rmsnorm.py | 6 +- .../pytorch/ops/basic/swiglu.py | 12 +- .../ops/fused/userbuffers_backward_linear.py | 8 +- .../ops/fused/userbuffers_forward_linear.py | 10 +- transformer_engine/pytorch/ops/fuser.py | 6 +- transformer_engine/pytorch/ops/linear.py | 18 +- transformer_engine/pytorch/ops/op.py | 4 +- transformer_engine/pytorch/ops/sequential.py | 6 +- 25 files changed, 645 insertions(+), 95 deletions(-) create mode 100644 docs/examples/op_fuser/fp8_layernorm_linear.png create mode 100644 docs/examples/op_fuser/layernorm_mlp.png create mode 100644 docs/examples/op_fuser/op_fuser.rst create mode 100644 docs/examples/op_fuser/residual_layernorm_mlp.png diff --git a/.gitignore b/.gitignore index 7a86041a1e..789d3b0a5f 100644 --- a/.gitignore +++ b/.gitignore @@ -41,4 +41,4 @@ compile_commands.json .nfs tensor_dumps/ artifacts/ -*.DS_Store +.DS_Store diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index d1d54c0dda..90f68653cc 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -143,6 +143,86 @@ Tensor saving and restoring functions .. autoapifunction:: transformer_engine.pytorch.restore_from_saved +Operation fuser +--------------- + +.. autoapiclass:: transformer_engine.pytorch.ops.Sequential + :members: forward + +.. autoapiclass:: transformer_engine.pytorch.ops.FusibleOperation + :members: fuser_forward, fuser_backward + +.. autoapiclass:: transformer_engine.pytorch.ops.BasicOperation + :members: op_forward, op_backward + +.. autoapiclass:: transformer_engine.pytorch.ops.FusedOperation + :members: fuser_forward, fuser_backward + +.. autoapifunction:: transformer_engine.pytorch.ops.register_forward_fusion + +.. autoapifunction:: transformer_engine.pytorch.ops.register_backward_fusion + +.. autoapiclass:: transformer_engine.pytorch.ops.Linear + +.. autoapiclass:: transformer_engine.pytorch.ops.AddExtraInput + +.. autoapiclass:: transformer_engine.pytorch.ops.AllGather + +.. autoapiclass:: transformer_engine.pytorch.ops.AllReduce + +.. autoapiclass:: transformer_engine.pytorch.ops.BasicLinear + :members: _functional_forward, _functional_backward + +.. autoapiclass:: transformer_engine.pytorch.ops.Bias + +.. autoapiclass:: transformer_engine.pytorch.ops.ClampedSwiGLU + +.. autoapiclass:: transformer_engine.pytorch.ops.ConstantScale + +.. autoapiclass:: transformer_engine.pytorch.ops.Dropout + +.. autoapiclass:: transformer_engine.pytorch.ops.GEGLU + +.. autoapiclass:: transformer_engine.pytorch.ops.GELU + +.. autoapiclass:: transformer_engine.pytorch.ops.GLU + +.. autoapiclass:: transformer_engine.pytorch.ops.GroupedLinear + +.. autoapiclass:: transformer_engine.pytorch.ops.Identity + +.. autoapiclass:: transformer_engine.pytorch.ops.L2Normalization + +.. autoapiclass:: transformer_engine.pytorch.ops.LayerNorm + +.. autoapiclass:: transformer_engine.pytorch.ops.MakeExtraOutput + +.. autoapiclass:: transformer_engine.pytorch.ops.QGELU + +.. autoapiclass:: transformer_engine.pytorch.ops.QGEGLU + +.. autoapiclass:: transformer_engine.pytorch.ops.Quantize + +.. autoapiclass:: transformer_engine.pytorch.ops.ReGLU + +.. autoapiclass:: transformer_engine.pytorch.ops.ReLU + +.. autoapiclass:: transformer_engine.pytorch.ops.ReduceScatter + +.. autoapiclass:: transformer_engine.pytorch.ops.Reshape + +.. autoapiclass:: transformer_engine.pytorch.ops.RMSNorm + +.. autoapiclass:: transformer_engine.pytorch.ops.SReGLU + +.. autoapiclass:: transformer_engine.pytorch.ops.SReLU + +.. autoapiclass:: transformer_engine.pytorch.ops.ScaledSwiGLU + +.. autoapiclass:: transformer_engine.pytorch.ops.SiLU + +.. autoapiclass:: transformer_engine.pytorch.ops.SwiGLU + Deprecated functions -------------------- diff --git a/docs/examples/op_fuser/fp8_layernorm_linear.png b/docs/examples/op_fuser/fp8_layernorm_linear.png new file mode 100644 index 0000000000000000000000000000000000000000..b5916a615281f02d9eb61182fcac05b5b9acc0dc GIT binary patch literal 17749 zcmZ^r1ymhD(x`EF4=w?MyL)hVm*DR1!3hKnuEE{i-Q6L$yKC?_?C$^XzCG_<&ONv3 z>8kFYuIaf`U)2p$kP}CQ!-WF@0YQ|M5K#gF0Zj+q5HMgsO&k`cAn*ZUCL}8a0#Y3f z|6%|M{7-Bop(G0e;zw|!BCxd`s+o!iE@&X5fOf)1- zWo1FAfiesTBq$mPI8Xuw{y;!+K|ZtrB@ii4y#FmLfqwZ{8(;{b79bG+YNG{we!S3t z_lL|sKf$v=|E)0#?0;H=re}fwZ}~$5$jfQu#|Nyvgr*Y+2w4Be8#HX;`vp)Ku~5-) z){vFqHnOv!H!!v{G@*C5vHxfe!t2fr6m3kL4T#-stZkjR-T6rW(SjQ&e^fJ&691!# zvlSnyhO7dyu$`j`F*`jYJtHYU95FF5ucNUkw~~n1zr=z6_(;v2o$a|97~I_4=-pW8 z?HtV*n7FvO7#Nuun3?H-7IaP?w$29bbhb`p|5WlndPGc|j2tcOoh|Hai9hrj7}~iw z^O2H%bo9Tke~#1H!t~!g**g7eTEGMuK8`Rj(K9mqzr;-3E&e}ZA4mQv_K$h}vpe1o zW!#!dCQf$NE+50-y|E&CP z0eMFY6W~Ptb6o%G<3H>Eqy5M1xD_niO{_IUENo0{oj&Hm!OY0Z@c%sW-;%<1)^?7{ z_69~K{7nB6`H!mqZvBsbH2>WXD zAK0i2;l%eyIYdhN=;60$Fra8N=rdjnt$!|l+o@lEDuUKt#;uLuVkVpwxI)NgSLo z#1BBk5*+9W2Ln{e5=u2g1qaY1QLKnZ3?|Z*1@y;Kbez}o(YsXSGN1wsO4KhtRY;W* zGX&72WH|=~C}#Wu^oLNwUCan*0*1)`KSL}yJ8oNh+oGONO-(H;qpiq`#V6GV0rXo` zqM@Us_xAR_&%cHJ_|cW`X(9dj*PmK57KcjR)&}2K&%0d(T2lQO2$9=f8m-dBSTga6 zbS|f@o?wZqooIpe4xi_%ok1RHf;&>QFyo6Fe#2_cFJBHPGI(m-FxRZz`?Yt4ArO*5 z?&j)Ln#ae-vFUZ(Z>PmmZV5ehV}D{Z^xTv;t}nN`G4P(8oP;71Aq^JDrcFiTv6@fi zi2k(qc)C6M#9|tO%T{7GnOWI(Ev2dXC!R`KN=k}CHnq;67iv|vN?*vq;p*?be5dd0 z{zL|R`NJ5?+?flHMRNFU=?gKp6gYRa?9=Q?ct%h znVA`dw7rO^Xm@Q+PEJJyQf7I1xy%84dlfpRoV%Bo@K3&nRj0_vNQLji!^8LY_fyJ~ z{Li%U#62%>&lf*jZVwfUxa8$01EG-W4 z?cG^zcG}-p+3B&sh)36hQ9wU{EX`365a_I{s|yK%!oy3`ZFP+#Q-1# zkmDB88XFs{;A+o%)rX!pt)a8as1L^aGN-inl_vXlFsPBV^eC{7j5DWet&M4Ec&r$; zhL?jx--&5dRaNmgY(?JT+vy@mQ^u-Cex#VNu(1WB+|#3*UTyc|)bT&=#x?cr>gwp= z@1HW>_w5oqnO};CN5e{Ph;&uK#wg9|y?&J^;Is$PI$f$Z!oe|Q4X>+?YX1K2>s-lh zw~k;FjD+hOLWaP1-=P?in}bPl zj;Wp=Fk2XZOCz#hfiO7WpjxkUu5?C1ASb^T78W!}HEN8*GBOAp{ZiPlv9Z}{n)>q6 z(&>#Zn2~6iy9_KWs17FefQDb|H8nM>`FDxI1_XgihiQ&EAr=M)6%^MOSJ6aSG0cw% z1-OSatmgNcdwH8{WM{z-DQ;z+0EOjJ`%ePTH3tJLH4cFR7w z+j6QmG-N&ta`5bIA}BLptGUqTQq6u`EzE2Ve#J$hY}cllziSV$K<7lj(Vij~91ONlt3k~f~CHJT?< zSd~H?l#Nud9W5?e9@BGKoAx(K7OF(}Dy%=p;Nvrwrmu5JR^0QhEClC=%FiU_mQu*r zR`6fD9aaqcvOC$!L#4lJu`$d@enC|r7L>1%6q1EkgTfe!fi{B(&?$$rT0SK@6o#n# zIGWTYFU@_FaX%0d?*E_WH+p0uH_+qbm;;+&suOG zJ+Ng>xq%I|%G->`w!Jn|D}G;kdwcwo6G)Y?uwRg777W#@b>fw;^ad;rI^%lgvNV}@ zc7If%y0oke^Jv3kb&Pbh0DlK_H>YcH{Wn-n76tgjcT=U0^#hHQfL)Fu*ht8u}<*L8uUv?zO*`l^MG>H$Uu}B~4t}4nQ83kd*HP6C0 zw964mkJ%D|+1i81UH%$WI>#7R;94AnrD$%_&5J9R@ijNk)mH&4UgGR9UHo%fI@wvG zpg%Pm0`Hx3PNDGAm!J0BRvH{KL*A2Jm!_)eJ)}mVm-VNT?)3;dOk6<2{)pq&XrN8x zwxxqrZpTghsDPRs6BENhah!BdDf($m7@Tvq{*@5;i_pcxz)n1wkT(oN<-DM~NB>H` z{c*8Nn1mo`;>9GYg6$U?5y2S1vS4epxQFM$13mH@#SH3}C%_3-m{9?PcsEb!xzka~95&5N%7PGHgslt$-Asutt)U zjg**~iG`z{iHWIe90b=v5m(bLemQ(vb*FbjQ47n(_KnPQC-29@QZmD;(_6=3^pdnc z59hQM6XEPaL|hr+Vw0oHS==kqUmhft>N$AmZ|RF>W}*Iynz_Q~HProkHHsz;Smy*G zl#@uXyh}G}Y&-p>0}Bk145^sU^})q-0jR@Bn!}#QkxI3U_0{Zd@#EZ-wpQ?kvEA=c zj=%dPYOC$5ncTKTYFvZND;7m43P@vN*rA|j*WsKAd6^9B!@1-ccv3YzN_xABXuW6{ z{%AwiFV2iau}WK-@C}!z`xOt#tS-Oo)N@ISb|PG+aSXj%1Wi^pd$yjU#Hs|y2}~#} z(}E00A@}=WpYo2$7%OU|-AZ38^CPtvij_Z~ACCp3ia0N0BXz`Rxo%1f@FCyOMgjbiPR42V>7BE+kRTS!iAnQ{hVTRQ+4HGxt+z2n4<04V)wH>ddmA*#5!N#ymw`-uc3Vy4moE0p#_8 zpI+eQ9y&Oi~4hUPZT)&Y8jnMmwePJOyIsz;9?N%S_)2>lt|sU-v`Id z_LP{0Mne+AOM?aWR5IE*Hn}KiI?W%7A4I>y$j96^T(FVtqa(K#nM$sW|Kj7ZYxiT> zpK5TJ!vxbAdcRh=_Y5?iG0CngoKwRe(q$*f8bv6^0YMH)(zVv~(Na~`v%3g|{yQ#^ z*7`}@>h2+Kg&L!`+(FVHag?yxrTl722;?H+vERX6;7HZsDgU^Sic>9Lv0hH+X z(R}TubF7%#K7Z{iiTGPN6J=kw^aQk0cqp)5aL$JagZMkV@0OP}SdV z)*T{GmGh)MTL^53LqYTYv>hs+X@B)0I@!6o96m+X;L6|mVm^fAT&ulWC}VQv`m}2& zD?{C*%_DlP5Pv!<#Mp&qO}GUw^Z6lLWx zX$b?r@$|B?V)&;nvbDhjl9&7>%J_xC5+TXBFH3|)TUNsDbJGKv1ny`TXd2%#CSI2Y zdNAQeM?1_~S&O5E8R}$Mz1j>X_id2u@#hxJY#p!Ulw-LrOOkRzs+}RF*jJYEeCY*Z z;M*zE8cc1Eo>UBYP^m$Mb|-)-c5 z!MZQQmBv8OU-g%3=&;3$ub1`XcKWhO^y+Vj6Qr`bNAZ};$MD48>0n9D^U zoA<@M&IZaoRRn=h$aVy^sZICE$;0a5nK2|>qJI0jd%om1EKhAGFFRSPwa&Arww8Uv z0`~~&HpRA|&GD(uud|@C+m;Yk80w%&<-uu%x(M4jy(^<$Sonu*y|tl(K8~OCcMId+ zCOGTf%VG8Eo#5Zx%SpKx?5C{W`0iZPFa~zqAsDB1bllJ@p?TmNmH!ZZ9lX0ssWo3Z zvyr^9jei+{S||{v%rb^AYR|-tqgH=qEBhmqOCb5(Y##O?h^~8-JF^dg^lV&Qhi>`84O3?ia3H4tee z*(EJ6Lk1+`+Szyy$<6EDGtM5mk2GXXdx-27$|apRGKkhmExP9L@RXjB2m-zuaGe(y zGk2-9c4F=F8f$Dy%OEmEYJX-vIpb>$Jz$k)@$!;pK=@ICY1>;Tr|a}d4OU_BFd(=k zDiy!Ep^Tq^;eCzjyFjEP8=b9)#j+Eqm$Ci~m4&W30w%%%l0t)K0fgPZsVm z+MS{8a;~wDi+v&()E$5CGmU3@`xkpzFeF|=JL~8zg5JVUdnSCVFqKH! zN%^I{nISZi0ZlS{x}#NtmKjwx(uM2Z+xl`unj!j@fMDz8Df2YeAgZ$O*_KNg#%EA0 z*2TG$k0}@z{hnfiOqxL;1XRvotzP(OltQ%3!?TM-=AENXrQcnKs$Bl$WVYe6W8u?C zSec6i5;Eo+aNnEaNk*-;kebs9>eEbKl7mcHzpu77GZ-9)bT_Nwf>2sO6gM%NkYC(i z-PoBuVj}W4cHJXdvMovWXK7ADNiGM|gblY1@yijQ{t|QXzp;@fRkn@Ls(tqK3K1Z-K+hV(2M5zoZ ze#g`e(qpML1lOTEQyy?l%sgw?ibB3n#;ZhA_?{53^N#;8*OI`)_GIOFT|F9cWq`3D z0R0()dm~QgC2iGR9MffBNQmf+ArcC?c)>OOk9KN#=4!(`{ZN=so(WC8C>0 z``_AH(P+e5-kVBJZRy&%?GW@M8|A-J{J0&IDU=Y`*(NxmAxYw9xQN}!;?k17Ep4hR zI$!X_n80h?T+SCz7{niwYL+y0PH~aMCC#bDp0k{b=LDsl+9`g%s7!=nQvJb+{7N+Z z1ixWpW^ymcz_D9fBQ!=`fAjcO znUr#TeB#AX?<}$>JU~{QSTY95jwrWsRLh3IIHcUwarN=N|FD9e#HOnRJk_`T3@ao?N zSx=Ao4%DVf?fTKL_IOT~)~j!C_plbKDSNUP1nuDY3yYIc?ei(ZH%AB>>%7x;zZ0+C z$V_uJmS6hw8xSYIzS;XP>T6e8_gf%Isc9wI>Qb|fv*Cs>_8Q?5Hvg^|?;pujSmTb+Q&)f=GKqlJb5TUAY#?@WYHxYWex z^Zqn42fKA~d+!UQy~}$_Iy(~O3SW2Kh3F<9&6MPrc|kbQi%M!I4@X;R_4r~veP3$h zrN;9NjM>v&<~K?;9A2_PA*eYgg}(bG4Qd78gW9i4y6#M5V-{b;R9k?aO*CXs{V^)x zwsHyY)**ObrA}6OZS5hPm>%3_tur7q)ld3g15OXu$W4Z{U7i%%Q7Iuny8>DJ2B3-)D zXa>Txb(s4!)kaYLzj3Nds@Om6D=%F}_<%2g_Tr`-sE|j&Y3k^B-RM6l;V=8o^3(5q zKQUPZAK-&zTWn!f7{pSxuIiEFu-5L4QsyFFYlrYZ7gHq6OgT2~u@6@Nu3d@N=Iy7X z#}@7e+?O)Wc!Yb(ly;{jn)&s_yp9OUgbbmuQs?N~{S>q8a1&%=)-QzoNg3EOjV zY2>Y2$YeKUpX_=;YlAp1QVuS>k#;Pt2}4MG3EfT$`%#6LSiLxXf-8-_wK3wrdjt)J zMm|4XmM!n2X&ACeTV@jFS#XrsB6`Z?SbEn^Qiru(C@}0yO4=| z@?%nNthZj#_eVy@ZfEzBM#i~)VXOQ+MMs9u?<*nvz9_n{ZcF_XD3tJB7(&hdoAqCSn7a2{Nt7q4E0nMzoKdmq0ee*VthH_8d3P7ISt_{M z2Xkm%2N+y6G)BZlT3a4A42S9FE(o2fo2kE>{V;yb4Xry;B#49LJwvE*f5t4h+@Pnf zFgNI6@8xl^u*EEla(q$61s`;>l=`K@yjJj8EveDCJU7~`qW>a9 zpBZlwMFlfrMOar%Fre&_TP*9c`*H3sf6{YtqcFY`F|H@pVHqFWCb3aZ*{H@4xSLq z#8$>ps9)Am#uS~+oH#sI2ARl4=n>}(K1biMg+g^qjap$%r2n-_-q8K(9jwm$Gf(nY zjr(2E9W@oRd{%s3k^4`hK-YeVxf{j)?G~0{_WcS&3v??uN^Ar&*`m(_xEP)Ap$W=M z)+`1p)mb61d}46NDK?|25Tkoi{i8?;skpe2mnASd?fOCLe6~9MGkWrorTA`C;$hQl z+lCBjBCBICu@Y&Ye{~+@Y4Q!F+S3&Dr%utV&xA$e;>z}V<6`O|Y8InrT$X^!(ud#g z&+VcWNRPvsf6KMVH2Wn8@fYxHN80sx4^u`CU~B#=JpdepM<$TJ#AZh=b7ic8KYW4u z`>d!l=IPs?4MM~KUts3^7YkOtQpS;rg%6&=?TCo>IkEthH{6o;Ls(fB-oPxE|0<^q zmmC0ZVwz_v=A`;f5D+DEUv+m7BOaTo_Q0qVg{AX_wndwR1d}h1+pvIx1Hp8&T)4vH zehJOJK~)J*YTPf0V7?FPFu|m4+?9}kFBCuDMxq@y-lR-_uC929cF82!f2+T^jWm`R zmt~7Qe}M{A__z$VoR$;R#_;d4-1Nx`Vtg>)tSpQNWm`)*q)^)`w}{_Fs{MX8_gHp# z<`>A397Fvxh~7V+*YEi#gDPj_9NP*4ls26u<9{rww2-G?y24{0_(OW$qG(SM1R4L1|GkYTD%B z+3ezqYU!C!lPH11=LO+QUKJIn+(7hW`|S;K%2kBfN3F@e2xR~Mg{;@M-Sd&ZMm=0`_4A6pqY<~wVK~cm9mc1Iwu>&L zy@KnUKM7sxG-gXUQX%TR;?cWR|U;R4?;Y9p{-nx$?}q! zAz&&I5FC#~{HCeu>t+Tr-(_En3|3Q%g#@ESC;w~|%(2HQ?G(JC>e+Js!a#(}1*ECu4yIN}9z7 zV3NAQR3%a{5kYw1g0lc+CKm!?9PU4e2%(i!7+~jg03$7c9Eyhk?0os-TJS2PA_uO8 ze;`0CVr4YodNGgolL(JPLwv3E{yqUF<7Oa>;Ox-YcjrNNA9k4jTO5h^%M7%TRmUUb`rBJNlJ?NlD zik4l4BZ1<$vS+Sh?(<^esLka1B2tK(@dEC7)2!CoLF&SO7#Zkb1Nn1)Wwp@JthkhC zVV8~pu|eiv0jtE-r{e!;Xd&$A09%QV(lW;kYNTQ*Nc+4QE4I@%j;c2i`%7>h7H7@F zk`~t$5BEghE>_qHj8=6ayIBUfqo|PRW2D)sf(m_vaxfx*8B>_x+)3mUgZ&dU6;Yx5 zSGdO%GGMg{taU2KiZ&!4tnlS4=GI=l1*)Fj6*U|-0Q{QjhgavQ3$MrgZ}S0e7!HY2 zFbUj6_6ruSNPecgz_Y{+eMndW<A0zRAsf?2Rs|PPjUDH%B6wnFS zO^#~LBuOxrufR4fSEZcOLz*>E`2t%C5|8%@Fk@r=nAtC|F=8xLmX!IX`j}GvM}(j= z=4(6o$k~#`N5XW9gpVYGxRFE^+=)_BvV)k!5xb(Mq%_L0Q}P^}i3xh};J!7@>Q$|MW(dAXedo|zj`4yAURPy?o%ofL! zU@>rf<=B*_`o}^mC%eScPn#T-z|+$uZBaYaH5f^8DXvOpMMo6pm@V~^4SH?#luRaI=gA^*u%rpFI|r(rME@!*UaN1u3SYOm;BAaxf*Z)35x%d zfE|E!n-Y>-NlczUo&l`nIQ6Hd29jC_0DU}{`&HQY(}%Rw6}(&|y5I?2r|HTOL8Aic zS0717d}cCIM}fy{ZLO0^thb0c)W2#+>$S~kTfHOG*y5H){s7_g@!99pdL*3;eEp4w zb^EJ5-^{>frh1HxD*xFV|4AlN=im%QUB?SN2Kp}_>N)q5*0d|ceN|*8TK&Rg_U}0n z0Bl~~?+)%s5uEYgbsQq&X;JHcVJ0)JjjJb8a;Ul$y9Kwf+kU*Qx$yWl7h071lz2@8 zDS$uhuFV9VO@GB=BrsViiQtSI0S?yD#;Zt=iOFy$n#TFaGQqQ1K`Ie z_sL?*7VuwDYQE;f9Z^2^q(}WPmB5+*s!fRdmsHEtNBU^K$A&}NwREVjE5JgsxSTVl z)%q()=&1?(A99VQ5pIbhFTha zV%tqx89mwiY}cu%k`|%=SDHcoOS$}=(`jfD36gXA+D#e*gS@HJUy+UWb)@zhYPkWA z@rbJ?JCgvjteedaey5YvQ8=s0g!w|2ueKN$NKOOZf~@m8V!b4O3677>6@s>w#b;nva+w*QBxK}%ISRNWc&&3iJ;&W2L{L)s z^awpvvDa+Kj#H)UE9ac&J@+mxjsC($HoOQ{==9*~HAwa2O(6Zfl-bzPv?8HiaPqY9 zkuvQs2JML6@tI%oYRIc6ueLuaih1AUyT7>`Y@`%WRPi$M+wi7X3$z$*(7|_MZnY^2 zep|NVz!qti$0xjnb6Q+uhm$q?cAO*6F!25Roe&;;A>*RrSh?&$s7*6xOA-C2B$R8} zZ6C_o#oqS4?PU20WFf6niA30qa=X1~okkHC4ttlKj9)Dn=_xGi)}J!d0Y5lkqX+I^Jytb{!pw*9Qc>F#k+e3d<&e4K*T+z>sPjwQ)QC|VHcC%*lGp3l8J1d^3n zh{%*_oV*}>T{o}pzmPiTik!hqQssA!E;tx&Au|ywQ;4=%nOgei_Fm;}(c}M}DCZ8n zuy3sXTVXy#U)%`^eIs-M9u<=4=*=Nnl5IAMJ3}W^{@%7W{=LNpeXw?Y@Dg44E~WZ6 zKsc-F?OPcs5^ZrDJ(olQg3_4uP-cg2rmr#)kE~0{dFfZIw+h4vfEmT; z%_p0g>~EUlMM-}?zTgCjIv((LF19=Js7$2}B) zrieh+cHllRV&^X4*}q%cv&GX}nPayJP#@j8ZM+iE7ePB^#=nRdgDdz`%jE711IN(G z$M(jGOM&$C$nQw~+`BiZgAU)oJnvF&&od#>HTJXCr7I!ef=P%hp<(>G&0(qKbrLRD zV(gii@it2b@uKOGSB}Pnm*$vdOMiiU^^9~&Yy)B=ygKr5R2ttleS$*$7AQrv14v|xsA}2%Kc^fen>`R z>W3P|A?~Iu^!I$5?c%W7zJa}#5wke9WnVFm;@5G3x3m)v0+*)66-j05pCFqKl+NoW zY^cJ^L9gwf(mS0kyqI0-SmX3Z)NEr=}e!r<@y zFtK&)z5W#tSfEqGUk6tvX3$w)lm9dizNdsFmf(@M+H>TRf>vJ!#j_drHOc=;d%Bpo zp88U;>iSx328A!`cN-27Q9BUQ5u9@E?<+5*)t}7MnHth2Au$kfr0SN^vx$ z_!ac|kvE_zC802+f!;cC&ft!WzsUw8km?v_(ZM~(arwEthG%(ZG`c|ocOZzqp2@9& zvG%Si@`cI*DF=cJ`4iS(w79sjktl4*l-8rY!PI`2VgTdFt3=1#VE;Kkr_WQ;lI6n0KWfMaucXi(Rg+CY}XjmPjQ|wN!nxVq>$rSca zF7#e?gU9@a=DvZ|u!}aP>%B{Ai}sQ_VY%zEB&E1n^Py6J8>q(PCK9xt88IP`V0aCV zV^=ZzQL(hOIF(KqVG+seJ6<6eKG)nZwF@A$Z=29$tz{gYP?AN|7c`7`Indrl$xw zzqq2J#{SH-J{UsVkD%|5SDl<)YVEvdE814Gm0>+bH!i!eN_m5Y`avXl)#s?$Qr2f( zR*Yg(PAbRGE||D`9_KSs0HUEsc=Wv2ynPaE!$(HvUB$8wY*nhPyZQw000vZwQ%l?h z0uT%!-vI<;l^PCb6(_7C#;Nh7&$cfVJ&faY+Sx7w|3|s6H1CeuE);)1G3jOwgDv$M zpQeJS*qs!U-N`T&0g+%d2?uW*$E2va+Q2Sbo}MS$FPOy`ge1WX+-(|KXcGu3xsM^W z<`IMj&3LCLD?t;y;$`b+&kj#N4^c#BNMw_h9b;7gL1hU5@*)@ISjd7doY#tEXaJU% zHCgnGtk2+~$D>(j6X6r}BQU;Y_>}rTJU62nk!7?c=^C5JMY}wn4JV@#k!OsRc^D)= zh(Cn#-#BTdN3Zp4@gsY?SUCsH1{m-M3beV>A6y{19dDT^&yA7v=;t&v+q^`%j@Dh? zlfdsoaPqgwBepq;DnpQC`;|(cds(i{o2SM@JUb*G%5gD@ z)_$8Yzl>R6edii2kfh-GSxOMOpn0tO)F^SB%s9myLE1Z-%Muw9GnPA_uR#uKnYMyO zz7%Z;wWM9bsnU01Xw-2~Yz{+PV&UPa%}w)XegXfLS!Z8`l}NhO*SF}Q@@6F)oCw0C>$B>Mnh4UP|4mY}3oQgH!v_!J{%* z8u!b!ZY(!0qA0znP#4)0uh0;o3~g5>dp!$s>V5ep+@;3RPTG~T1a0^Upo%p$A3JEOn+-e zP89#$d4TMCQ-0mG*76gQM7lrCJ{QA*9s1PBA-(JyQBgTbRz*-<>)cBZ zfrIcSDgrSIT87d+1?@G3R6s{p(fVnd8I7m*Q=054-B2-aXl6!PJuk!vZr$(km7qNs z>xmv)+Ea{2X~)KQr@xynJZBSRK}`tnlRp|{gBwSY&O7dy@6C8I9(1d27*ELpeV&2n zELB!hgN=c|mkCt}_4W32#qw(&JCI_8iP=OKt#}1lp|ff-k#G9iGr5`xh5CVrMt2)_ z3?U0%0Gq^T-ZDi0Oo<|nDTKR_T1|oO9jOqXuVX_SJtruyrJx?HXL*!Be)ad>7(-;tt7j9)o zDDw@MXthC+P;Eiam9J{NylTHCM^H1+H93}l`T0G@Q&N|@Omz`r_=BG31hQEFJ5WBnn}ugf~OulB|qJpL7HRrV?VQ zv{#?K@F~64g*Xb_PpS1Eyb^y`(7@nkdAJNYXN__c*0~<`-sLb7Jsb`Gdy8g`n{OxI z`n|)F(3$?;4K}sQ8_jk?!u1exqY1;-x>xPj>9SW9)A1NY$;^AKy*s*yg}9&d<xKAZoaUS_cRf7+v801w)^@yJQttd zEx{p2*1kF8yME3J)Svfl!_H)K7S>mDp><;G$%V!~;JbOrn7q<393^+gU(Z>34--%2 zQVGR`p&K>^hdN_??>iEnGFJ##T0Jq|>ohiwGNe#_i1oTmmY&abB1CK0dtR1hc;jQO zZ=4hkWl`JIp>VFBietu1`AP&l!-oZe3);;-W=jx&mW^j@;o#nqavXq|-n)_YFPDw{x^rnL%Ep$lVgLsVm%9ALURQefG;?EEbCe9u zYHFekxyKe%K1YMvPHY6`Vu3>)S+5c8jM-l6ak?Ayy{`5$dDe<%f*_!nYs|uiR!0by zZw~>|%fg&>sIS9S?0P=OhQ*K_h&CWYaCCT@Z3ms)5zN-kf9`redJ8pyWZnXz3A3j+ z2Zp)+cIf&hUoF09>$?}N6&>Qb)wQo&AAB@bthOKY~Ox8ds>(rS}N z*K$gk`||g(je)TQ&Xo}9y%bV=nn%!X+p+thYE&o`%xrU$b{MIA4>gm`d@v7<-hmjo zqxosGd-f5~NSo<8&1sEv6QUp}KC_|e(nG{v1*5zYfOtZIqc{>QPsiuZ{M@Uz&O)|E zvAJ-U^Ua}2A=&CB|HTfCd2y4K@=HAs_-XIglrAWSWK|X=9J3iHEVR38p}r<=onfcqMVL`sSkvD^lyM3qYzhbSL=y2}L)$0wxl z!RW{QxPfdHGhap6GHx<^M z{fYHCg^Zk0HhLi?4V(zbLKeZZ*F#=gxId;?WpcAst4m`nyZ78Pj8J8F7rNLSf*$_w zN9^^*lv?ntUzZO~*WQ~1Xib~)ZHamadQ0qQEq9fJPcIw9orQ7-z7eQjI2l8yGDcr9 z3jKH7@$pZmqaNt}VJHF{F}2X#9lY^~+<8{SatU@sEZdLya7zZa`?GZKc6s27`?3g> zd6RR6z99(XnVvwAiA*$QRL2J-*m+fTjz{k)1=X{QFu6Kuc=$>3)@7pk*GyNx`N2{w zys6x)h1!b7Z6XC)bStB`orTN_&apG~!Rl%H*O(jtk$!6YI#g`3m(1cU_CM6svRp4m z3UbgJ6Fn{Q6A^3w#!#S+%1_pg1T)#1DH)f0^NfP1p?%&q-KSS01juQ>>lXzzq!xE; z9L?|Tu5MxKSjAol(UGMs2XK6QQcZmq{Kx8cjHPGf} z1mFo%1Nkr6iNz5FD8%gxpcsLdr1_>O^Af=dKS7Sz>Z@S+(5Tq57g6@2)xq)>-TO_4HAeAog%PY=bCJb;MxRL_ zf2M`j?-D!6;YEAU*n|CNc)3+IB{ZS60PbEbXj3S_;ZJQ> zlS0x~!^-J~4&aCqQiVgQNUA`Ux8VLq|jYaVb1g8PdSZ!d&Z2x2livVCv zw!8kwMTK#e5X^#yB-l^j1<)ojm@a@#7syHB1OcGK{{buI>qR}t!9=p80e5+ehzMwl ztl;g)5fOd;D?gIYU*Z_mb)3#C?v6(kTM%Z+;!lF)N~jS(1Y5#wL&rF zrOJhAcO@Z$B>bk$|6tk(?EbKuu_91IqpC-$j-fx(#Mm3K~Wl)F;i7Rlm)ASBui+mw8CNt?vrE+|pczwrxd7RWJ`b*p^ zQ$g;7RT9N6W5M2RBvUJsNO;Apv4Rc>C?MKyJWc|X(&Ppak62V}Mio9CBm8!LgHd@k zG#>Td?qGP?r*78uL!UdEL_KBs{4)Am zl>e|=j3VN}&jLwxZdf5NN4!8PZyU;&a+{Q^`EmfV^yNr(`@BT6i?OVaNz zygBF3t#lhsDl<*%HQ>UqTowVDyKMn38oZLFw=3}d8#=Vjo`n@PiaEXMID#{-i3N&; z?6uD=os45E3rYfVDhrsdN|r68@<#x=>w*VORo;Q5fdv*+Q2irFA)GgVlpHin#BpyV zVSC%KA_xluBR(m~i9B-GJPyFzHZuiv?j%)C2Xe^B(dyrP>g9j)1KFuF;7D4k*k?&;xCUss3E z?Q9ODzg~7jyJf?DdWaTyPHK_lVga5;WvmDAHB^wf@lHwnE{Kwyzi=ky=V6^}O5LImA>9H#uu z2>bKA!Lh2oK0*7cZ*2MWPwn)l+Xp)q=^mVg+1VQ)GkS0gx9+nyE!}WQx&c*|oz^Rc zf{B9K9P#Gq>C1{;rvb!rsXD`L-tqBqPfrgqbjbmbQT=h6l$6lufh6%Lhwa1Z+;5&o zG^kATRNr0tdj6xEKQOaka1Doj$bru4EeY?p2&lsz27kExi5d==B z7Xu5{k~r2WjKJyf9V)=3BHpu9xgfpgsjt0;Jt~`1? zN$n(PE*`vOX%CkM=y-qvPvBaV=_a0Hz#`oWw59fWr~@!7xHN!@@f}M%H-Wa)T7u$j zJ<};rxgP^8_m8RiPI3j7`{zK*nJPsCKt}~Um;|&)GHhobu%1|mX~Edcg! z7(K}M?;lLa=Qqy3zF`WW{-Ia^{RbL)vH<2!8R|C$z#5m~Amj_lMfR;bq?7%>KTt7q zL8p+yu&t)9hpv*MfVs0Hi>ZaPnI(&lqswnBfRK*>r08hrVM^xX=-}io;3G`&7lZ($ z{JWZ!g6uC44|`z>T_sgA31>G;G9DH-7B&hIWHK@`AvX&v0d+~~ztJIo2~*g3c(@3# zvU+=avv_l|IJ;T1vh(xvv$Aopa&Rz1Aei0XIeD1+Fgv+Z{)@;TI+B*|=5Dqw9=6U- zWWVW}nmK!V2vbn}X7ul$f63`#YxNH&C-=X*1?eE`?-o{e7B<%Z3C+^S_WuU$cgw%f z{_5Ai*a`h6Ch%6>(%sp?^S4+c?0iChariHl|H$_*gn!_*oNPTrIQ~ZYZ`J=DOZT7l zf7twI<$p&|akI6A4CKG$`kTjp*8K+bs{aoC7oWHP@ZnJ!19g;#;E&zEJa{wZciM^Z=Cm^cR{)uU0T)|JvhQBU zN8DDvF@(F!MV0g)ev-6e z$WeRI^f3Px3c~RK=HL8+_h4ng=)~-n{|E&cgn)ne{eLkCxSw{#2=HTrb1Rt5Hs!Ng zIcPd#hpH?YHq6UM=}D0=lG%55;uCPTtOjrS*)Wt9Vg0EzMHpgxkdj^4_}pTw7QTSQ zX);Hskt1VE*hBP6g>v{adUGo*hq)3Zp>j~5+y}aacR#DCAY-zo3(W&(M(W zjp3kRP!QDlG{1znzT4*1naa-NR;+Nv$&t4=fyRyRB;}mSu(jUsr|+Pl5(5- z;IV+kF8FG7s!Wx+6$<8FKRK*%oF}<4Y{G*Dr^h}LP@CJFYf}a*_UD9$ReUK1AD9Bb zkRb+t&{)VNQ%%htix~AZO)HsE%n#qG@(12uUm-qcz*C zcf>DhcHQ>SQxlJnXE9@f}M_ zK~kt*M%f`;4m zi%Z#tyxlZ5FBSm(5jdTB??3WqmD32lXv?W>m!wtXCT*a%rr!ZQNIb~9(AtzDgGE5( z+#=5R^`4JGE#{4{q)W?sKVx-lEpYr?eVF+4r3od`!N^a!JN5DTxWDWnkhfxa{^jO& zF5ugvBdv%Z-+qlu(5yV)*BcN}2xNyGIWsK(Owyp4t}{W@$bNDl%ninII+qa)l~Da2 z*NLsgmkcUQZ7UXTI3lW&^kktBQK+oSsCx|9oGANf`Oj_V#!c65a2RfguWSdqP;k6N@+GG6$5!RKTiw2A+1# zZ!8c_onRyuf|P4{;LX)_^rNibGfP-Xp^6#{A~5RvDRd>}pQN`w9>u2)IbUD=a`(9e z8=o3f%vV)AG7;rgbK?Q=qaLrXq_pd~wDW2G0y&0JePLBzHj?VJjd7XeNg4Vuf5uoqf?+j-- z{UF~vYzSKIzYMZ-_8;x7;naxhTDXA)1>Fv;#Y}g=OSP8J;whdyt)kI>eETw#y-Ha3 zc3JOL-sHg|f57>`_TZDnGG$H2J@beA-67#~$M!qZ#ITX{Q(f-ejoK?6+Rl|uON5Uq zR;1na8!O7)dY)#o~CgjOSE14hN$$3?!Z`U^!m7Cb(-&+~lIQ>dXTsAk>Z7n=Lx`DF%u<9$8c z0NY;LTZH7ypz^~*QXsz{9g^E_@BB?|UXaC8+irnBHlF5}gQGZ(g5h=IfY%Mlr}{~c z?Y{HHIL;@#yFTYW!_LNr<(1XM7@rR!;GWW)cU{rRTvB$Pjh%ECb%nT^re|6n4Ef$) zj7?#QgeWH)uOdWrwfnQZj$=**2K|VbO!3bW7w9XYOQj}@cJcXD;^Yqw9wu4xI-D!< z1+yKc@c0;;nPKB13d&Sn%B21vC(d<@6q-TqVab>6Z;YJx(y|<5H8#L!uW&SVT5QO3 z-7c9$2NtT4AoeB^ayEGCGOF}}BxWibRm|N}*m{;0-bQjsin`Mc^$tEOA9cp3gTP76 z=SbCfX5XFWBdhKfE$QGNRJ#vcHd+Y1w=M;L0gr^4h@h!N%N$O34uzw_Ke0R)UL`eG zMXlDyIG&wOe5}n~N!$2o=ePC7DygTi748nEA4SUZvRNjBT*A)&aGXf7-IoS6iAR%t z2BU|JNOk8})PUs1f!1JC$I#I4tR*Lg*FCwvz*TL2!g`}nS`r=q!6Rywu=bPC_(^lw zHMN+$SS^LTxYOder;d#9;>l9BM_?iwpm!}T&*-!5Sljp$G|D9dA7V5dDw2s>dGBm*})JoQbT2aWNc+2PEm_`#Dh3a3w9{HTjn$t$uEcs}zlo@Bd z@>WZ&@wN=%JohW-JHjc}QAW|TLnU674ypA7;wR0AX#~vk$N_zq<(ACm0B^Fzlapzrwe9eci2k=W51(b`h*m>**olvypB{hw#KV)4D5cZoZjFu4otf)dY^VE1 zB!Z5bGoSzQ)@H{qXi_SCd=cyTl>Wko_MrZ-kz7A2MrN9%3nj)n_^dn-=#W@zph_hRoSY{HzA>e z7%v~tpy>Qie=kseVd`otdFS%Rv+%UQ)?t$}-WQ*aU~uvDi0(iD1Xh`$PvQ zf`nE&a#HBxg?y!XBR!!rFkX9sm(E17&inb?6=Jj}7-s1Gq(YqGq6swOnC8f7go2q( zWNM;uVr5Vj@MUwbKCSvL%J z2XE3A18@pzxJq&lJlbtPxYs;A70K~Kl!!n&{&ssi7JvvZ zOhR#Tu;@W|PsiiW#-J-;KJN|_Bw7FA;66U~f)VJ@+2*jrLZ}~Sn)D^|?CxnA7EET24<{G*TGp`_!#Hl? zCA6nT1f}NacuVl4(0%*nhFv1fTjP-Z$G`On51*deVB+0YA+t?7As$C5>eIU|y8ux` zY&pK<2pDE3MbRT_hVLgv^+aaDp&!imlC*O15=;HxZ48Y+h3a`<5#yf8tu)=gm8Rv* zdC?aJt*pG7$mwVizUMUqv)^t->Dr*-i@kEfpZj|KjQ+JI>GUg%<#yl6ExU^6#$bOh zpx?5Gx5)HOuqX-hqt>(3mn)Ec&gI~;@6peDkFKDXpLa#{aQv2Rw26VAg{$+5K035) zw28WND6qNoJR;@`$LEqo@+}WBjv(A5B3Y-ssZI=lW@Z{YD?h!Z$F&T${rM||YNe^E z(V)>D6C2STg#Z?vi^0dL3{PKA|LfBo9nXEpy3^8v-gPWeSS`#LWn~eXB&M_v=?f1= zkQ#soe}F)oc}3tF*heibT20Ggx4eqhOd3CW?bz0Ew#+`v51QQEKGoGW7#i@r@!B9G zS_~w-AC%BI+|3euzO8Kqc0w~%k6P8e29kbnMUhEcjQ4nsOAJ5xIr?RWA1;=T9{LSZ z6!&GKuHJmk(>H-zs(j+F7@4(%ydH1tZ6>e;1G96qLiKUBxdkyl4o;FzJjb?No0}xm z5(aVmO&>J+xUr{`a&_TE!rEU8C9*|x_)X1)={tiwhLhspR3_2#bEfgM6@|X(K8AM; ziweQ{B=Q9b#eaOh!)JNu2m%O7%9&DYKCTY2E-&eW=>r2EclJ?=*w=h~e09Rb-_6ws z2|rBVm%)F_R)B8^-5$>fqnmx-w0kht4z8_F?Dh!h>78Iyp~R2Rj5_HIIfs!ZE9aZG2f-?ty6RH53vUFGX-V);x%O<%(~%$1Q=VdM2&NA= znJJ)$9FkCKD}%$PsFu|Md?YM&(fkk47ztN)bcE_wCEAQmSGc0YdDQdz zMNeNXtH2cO($L)W3(d8bt}0~1{IvOaXHej5etO>Y;>EY;{JF}n%P=ri%Xgsz9hWs> zDVQ9EDVfI$w=J7`V)P)-Mfl-PTvRD*?yg29QrrM!Az*Vq7&k@Y(T?p_kbWYTwI`|jQ3 zM7#Z(5UhIk!S+@;Rk?M#lcR&xWsM;hQHYNt^xiglC&$ecaW=?V0tPY6@x;)X)(Dfs zni$fJGLs&4K#SnN7MJe94^}so##BKqC1+z)OV)Q*=tlFBoD2DeU+&*^gY>o0bT7iu zYBhVeRX>a1&L6C%Xan+(<2p5HLtqc}6`r|vA6-+n(Y0&djzb9}i3ia^Q$Sy!mnIVK zatizcZrBqVgdzjWnXpPdLue5AV;0dBi5Pdwt@L1VOEWWp(jU;3@KuSQF0Z%hesxb` zG(}qb5UPyMGW;y>Ex`wcU;~L?f1)qi$i?nyOmaAl;5uKC{&GG%!ooaU<>JvW-X6F1 z=&YR;8sK{C`)mKsI4R`0;i$w=b%bWeo?sjhv+w%9$Q588z9rA;<*0GiXdZ;7cxH*l7(5I#^zf1PbMYEE| za{g(;2i+E<2K`8`XxiRA{cwNFj85dU!G4FF4F^|MTKT>y)MPDFpLCcS4#8f1=b>Bp z@kevk)7@Us1Brh#r6%t-(c&4yRTR@18)^$@a3UW_%~_Jn zN~Al2DoRCCQm!n*Msh~lt5h{%I%oB?Xx-2kjswg>dpch7SYX3dSqC3kg{i?-BPJ^( zdE8@OQu=6`MoY=LaUNRhV+Uqwsh#yu<9-jiT4T47Vha>g3$~OZm*GLSZrv{1Al?Y? zXRk&hYjkfi0MNXSy{ugYbm?rSi{R^dS94wq{AC_$TpwIl`WJQ*iU-e4OmZe$-# ziRzM5yf@p(n|HQ={pDRf+$SKz6eEUTJyc$kf*{puW&~en{VOmY|JK!%BLp`o8^a$t z$Ayw`3+;qu05!|C?^_NYC!b12B9u7=`}y6LOprL~Qs+5CNA9Qgb`~}eDrW^k(6;R!HII+UlnBdy>IiAk(KKnJ~J6{im3EG)4z@XhJ6K+>(Lq9?dxa z4EJ!__v@hYIA}V5gy?v-ZqN$uD+OKoxxYho-_3TuSe z7(GltP+Aq)d=8 ze%v*@I$XA-ot#K>+CGbbk;XtrITDS~Mb1PAR=D0}RrMbz2z@(}Df9WjRv3eTzM~)h zEwgm9(F+kl7Z3Oxh{Z^*3KtE{Vp*LVY!rdVuN`3a0SAPj!7$0jHWHIf`AV~Z2+)6! zm&E+&bTp_>3xdBJM9X|gj1OkGo*oP$_x?`gW%ha>V#Hxw>%f77w~ez>2a28+M(^9u zwu!(ot?0mwQZgbF=K*M4^j?9SxIuOJq@+!GmI8_R-y2{sszr3M3w$;HBE%*!`MMB! zUlX2zMrtD#Ce89_$_2j%LmD4Sfc`v5exKw7m4mQ?PM|vbaqs;2tU!fg$rt1m>uG`P zYV6DUI%$kDuS@Tj|C1GdvLs^g`bZJ4**?~@@Y0>V!5pE~GkZLpF1;@*s^7cC0`5y7 z+Qq%kP`sfmD!6R?_~9@U&v%Kncj|oO{mGGC)4j*LZ_AmD(y(>+7T4m{+XIh`Coc5n zg^ZEV0lDwDIx<9hd(Q#Jf;ovQp76c_;qnQAmBcjiYBGq#ihG;u$x)xk#@U>5&W*St zX(ar0Hdq-~^Zh$%Uy(lrmHinpF%^fX>5zDBy0!(E)j)$kn#<@>$B zw>l;wlDr={Ra|07J?W6OlzN5g6?1M}K}N*3Ykz2RWiqn_G1h?rt)BCH=v^*iN_XKY ziB6(V9cgY|Uz*xDu;d;Ltaf5_L>AE4EGKoXBY?mhaz8u?TtzV#>YfdfOrH^oS0Y6w zcv`OB;NV3V0g+Ov+x(?=)6s5UV0v;~G~7W`lbd4E+{}ztx>ln`){3DLkuF_g3i2B> zu1G0s^mpZ;fN60#`8bAiV&Y6V`A|!bqiUYrApup_z))oV&!%E){i}%@8m-vqc<#%O z4S3!PbTSB5DXMa60!j*En;@tWY%<|^HNI}X5Sis%yv2rOD#OG(Dhf(Nx@@W{m zMnM82{v*U|Jg^=aG#?xuHpEpT$w1NSQb3Or^Mwf~>|fsQ=wNz?uevu~Kl4EPk8_#= z;`vRMl~85><8IF4fH;#C+dM0r{;Wt4y2Mp2J?&YsY7j z5$0s5W2tX&Ngaw8<99cCbNJKqggan`Ns3=oRV9_7Xjf~!`m^jrDzWxE9fAD5Ie?=$ zNL=XqOh~_ZTsl%6jSp(23q&`1S#Qg*bC1HiTMva?|utrc+ zFojtg28lM3$H2+7^r`T&eBH6ag^|E3nLnQ35v=@dr*tCc(C-9|3H*dv3RLvTK-!qK zGFh~5?OisR{m#_ZF6fkq(TQ@$jJ<;^EMyi_WgYr)U8{ML5^ zao)Op6a&o#llRpjDQmof(^ibPcZxUL6+{rdabOYq%ku33lwRm)m`Gvqfbi;ku?(E! zB>^q8PPNZ9%U|eU#p$ll+LSg$m}}rjlqkB^9^roUvUSdlrky7&?PAA>Yll@*Y}U2T zI^r7Y|G9|FqQO--j3xmJQIR0!MM4PJ+sB>%vK_Vo>OnumG+`QJRpEeggb~JU3eJP) zo~rSV#XuQQFJ0+}{es#p&TLXvQJ2CM*r$EuMFu8gFg@rF*O7pZOzDhN<*Rps;iP}9 zt$*L6iZ>%Gu-(^=I$A7SCgl4b)>Ay89$I0-12#t40j{T<1_Kjlx3L!#Cs67bi$hBxZbK68>tdh0^p8O;u9e+v&hO+_wo zgO(LVH5w9v={tsEX<}Wd==#iX9fbT(j6^uxQrw<60_lpwYG)a!@7v@G!?n2sn>@MI z>0*K{7YId3jjxfj3Y_FH^jBMLN#er8cB(xP<%o3=Y(B`4fBn$7`6jm3#pGK;z&T3m zqhc>qT&UzVMo%CiO^gomLo~X6!$<8uT5qBb*2yXVtOJpK8Ck>U-|>|0Ux>*dQ(s4U zZi>HyOJw|zu#ar5T)Ehvb)=XCS)rx#rJ9s~&{8QOL4%(k!jUEaUXFjrFmPtR%^6}+ zsxuofq9ibPGPN3E<1nw2#d=Trin!W;)S^9Cz}E3(ru!RG-1^VAX)ZPwG^&QdRun`Z zq`|s;SJAj5ZLmfMu2v2+Wh<8ME?+K)8XatoQwwoUfh_NuFY^Iq zPOs;DttQ;l1f4Vp@DsS{eu+kHY`8sl8we_(Wq^k?s6%2LgLH8<8hvCJ1%JFnYq?vS z8U^-57t3oWyaOWp`Hct9cz@!=MP-QfAf3}?oBOz4;s>jkv&w4wV*_;eV=<{|dBcG? zm<=7T@t*nui(vg;p=+Yi@S5jZ^~{_g4iZJN(678hj4D;*V;XaCWXgqg4kT3Ag(7!u z{)OLD>R`WbXAfrbVNq9B9mpE8;BYvOO*juT@C75z}%+T2EU!K*|~#wLhqw8&qyjy;9J58Cr!YrSSJX-xh3ph%WZ^I5IRIOzfpt z8>35Hb@N2Vx5`}ACPgviVt>I0y#nU&c3tiEYrPFx^WAc|Z|O=reoKkf2Y1gSioxhE z*p*ZhfntFu8m!hK$0yr-_aVt;vY&(Ar?Kbx#irxq8-x1~BTSqMwU2Jc{KkX8BK^Q$ z)7wY{{h22A9~LicFxLoi?!*}Tcf!Nq)tZ(oaUcr(w>>{3%bMsNA#`rUOA}*9v_?isDvoi-W@Kjb!*GMJ107b1_s%04fUUOA3Rm% zTQNoE0)-{{B{P~Lru&-L^|J+A7b79@D3|(+_VhNe*SHZC<_T^9%M*e-XH|X(tL7Z` z>gXcI=2BVFsks3~q{m)Ni}$30hQKM%bw7Zua18{BsE|K~VXjhc+{GLR22Ui{d&7r6 z)>SsY+H%lkwcy|^H=b-^<9>eALK)V(-rRE0t%^pcIkz-XLHQcG1+)sE7;kX0H90pl z?ENOEj$#E~LNIs9GiZH72ZMO*M(ewN4)Ru92z=~MfKsRxS#>xc|*a?7v8OFUF`gWm~D9R)eGK3XhUfw7$uuCPN2D$I9c}4|ykcyh{b1&)f0% z4v>HE^sI3)j3th#!rM+z4ms5gC!WJ*c@AXmu2OrVGk@&GWQI$a6USfCf(!D|ymeXv zZ%_5J6=bmnL7m>7BXuc#+lw$-}ekHtP7FKE`l%_hZd5vXh6I5TJ*`EihFer+u zUZyQ)+wi9ihVF5o$PyQ4yjFb8AI7qEPM=~GqA!@|^Ltv+r#p`fMF+=gn|nH(CP3xq znf4Q#sCfG%Vft=uj}ItC7!BnNsF^;z=H;+3fR#ft(pCSf{mo$}e8UCFSC2X>U!9r@ zRqIu)=wRRPTz8y|8o5!S3l>;st^wUnPHeUh&9LG-@TP+Fe7?ghEf@AtQ)|6Fi}4@@ zuAazFMSXmpy{fvD(9WymO{ROqyg)M?TN-``;s|fpT51gFX^ePCmQ6Bc%QYBYsDGMc z)UQ<%XDd${hBGNTO~`rKZVMSq72Q$ijNgju#89Hp{&}HiqvJ!Ymm_$Mc1K$6-Qs@J zpcP4M^TV_*j^MCS!MO)vI&{|4(nU8%5J36*JDUaUkkrK5vMzyqRpWf_eu62Ihh_D| zS$!|MSz2R4P=Q8^6VwE`c2c)3d2K*mbpSYs+8*U7^Q@2hxfCbQJk3z3*zK(%^B248 zDg2iMh|a?9rT?w5@`rTf3yC$9NCa2v$zJlDCCi}jhk{SBzdj!b=_~M3!9Dn=?|y$z zLb>dl@Ga0eC$EIi0}h4pke&G@P_*ZEbSEo0DE*_S|5^_g{YN0u$0+sy`JM`~-IG?7 zki>8G;4SBR1=7cel|ghY;I)pv4jcmGHajaI0@gbPBzlAT^i9NvT_=G)qpPPb1m_E5 ze9~Z(k-eo8n&Yckm(Ss~y zp|3=K_xbYGF2?i69ERu2^LWzw#|lJpDAs_pRF_s~7|1VH3Z7u#lWH9;@*R@dZuf<) z4I#qjy1qZx#(n3vi2f$O<>AUpQoE;;q9$G{XIoDotm~V#lzy3~#UN55=N@fBP7dT? zw&DTV=&L^USA6x~0YPssJ@tOTwPu(%fJ-hR(?0^IqCD=Sc^QIUfVOJN{fX(v1f6^n z`zWaO3v=+p*1F?9J5d~eD?={`jTg7WMdX9v!E0bm1QN1ZV5f-VNnlQ7(A>B6tU?0- z&*F!6`LSs{Aqx#uOVyO&B{Q3eEUupnv#hsTHiRD^XYaPs9P5)2pC^5Jjo$emX9xO( znN=X*P3aQ(Uat*QA{}3wfc&x)Lk0}gGqSUSVABdG8RiTF@$U_aWnRtEo#BDaZ1EB` zBEx$=;1|dAlj?HL=x$PLEdu9^>hs6|R*wEeOFhj?x*Mfe^H{7%2v=h=AaP;RliKDd z)w~D^EKBG+a%bWrB7FB>sJ%UC(UU?8bOq=a%(Wd zzQ8L*sXps|<>d5M2!1SR<^D!#X!hkJ>4?&I#cfy>bmg?fNU}Ve_9z=yRb(D!s zOpR58MX}cQW;5Aa|AjQBF~`+o3LfJcW9hK7fKVgTum;ALHacWrF@Qjr&Lc5m6j=ec zpuQjM|A6ooUM=ni-rP8n)LfWr^x^|KpKb?@oI1}j zwRg~7<)LQBaTiWXot-Xq)-=5RCZ;a-nQPRQ&Bs%C&Pn)TuJ5xuBJrGj&)T-S6>&vG z85%W{R&B2SN+yPWGZwQ^fK27=1nwmdFh1 zJblWb;V#IMLQWlmqphD!b%Rnasg+`?MY^S@>sX{UUNlB0m59YzU?XvQa14Am-1g;E zxwsnig;?goVz|u5Rb={3rJ87%t}ArAC2_7R--{WetDIwNV%Zr{{cT2!Y}rjE6Upwr z%`r@{d67s6HB4Oi+1P6*!?Alu#Sxq`LfrT*zbdEaiEdDjbVX3@;jVP{F|Tn=#=Hsp zepro2FSURu7aC7|Q^ZzQ`8UxQIeMYSdyk|rBn82Hq1e+Q6owNL1tWr-rfk%TOXK|w zYH)!gWva_${f*soYN~1erv|I!&cM$N`oqtc=FqPj-A9+Ietd9opo%a>{(|7|2b~R_ zbJ;!l(yJanF#(kQSu@pb_%kti*i!YZCsYI0D&bK#6LPqQd0M!YYLjOu0gk3b1@Hs#S!WSFDvt+`p#MeZO!snlzS&QLO zgS9vc7Pw0t$?zYRw%q1>pfq?y-@8HGdgrV-elJwq%E)kE;5S(DT^leb=ph#NWBnF6 ztioOffnOg>g5$*v>JBd3SX{6)wS2m#HNV@B$vkwRBGPQ;Lqf>VO&8|gj|!zH`fil= zn69vxHF2Jvao3W{-mA{axOaoL=K3k=I)~ogHX0`_g0wim$t`+qs;48UQhS>_JpFMc}ZkBiVkXJwL zLS>r+Sc9zCse{V$x0@q8>@PP5NF0^ux#mG+2h#7Gu^8j(ffYVjsJdR%hNLYNlwKP` zI;7@%+lAxShBtB`WAv}L*HAtUao9Y8JorW%y}nvefdjea?-UOlD(8M%dTIFl}Wa^I5Y z8~7;^4LWZKfqk2#szfb?^bOb~>7Dr& z%}-(yYyn;Y4{fV+EjrQ#neqk-+r2vY^hkO-SS%P;MgE0&ln=0TgBw-mnn6Bs{n^y> z_1pYB`n@fgx3TyLm0)B4HE=pe0Y<`lm9edX_nDZ8Bu>K`ax>m@Tz^gIw&m|~&Bo(k z7%?KYOR~@zEF7PZ5V!TFz7vNhxm%fJ;XO=;A^nW-EDd-NdL%_r(2Y69A>~ew?m*)9o6=~AMhq~7?7e-J9 zC(RKGpUzR9as|Y*%~it8tA_gXR#}lcB+YMrRW5DfbsXBSsjiVr@KzcyDVHLaIFeaH z#=+=+9tYpN5{iQ6o;s=(zH^j;`0CWd>LBeQZX(=?_O_xW7HdZkd%bR2^tpf+QXl-9 z5B=KcVoXASFWh*RAep1bNYWzZ;!*X=ovZ0FvI*|z$5_~EGYzxm z=lR*4?Zvm3A=O7*?nUn&mKlJ}gWta*ue_hPCeCT<5<9Y-R0FK9d(?m0tba}yE`%@Z zIp5OD7W`!Bw0jiU4E3@^Bg_`ToHhIn+Kv!;U*i-G{rrzoe;<8h)2};buG!}e!IL~? z1M<~Vq%zLwnrr4s@1+5`;=}o^@Vqo|3v@b-f&E8-L?;jdauil3n2pvUFkS+^KAXnt zVw@?_(Br%N@>urcuf5pgY9iFvn`HiFMX77oLy_l)^s;@lNCu}gkPb>)!d znd~yDs(m<#=Bdf&@m;Q0puVS`c_SfL)O3^-UtqUz&Ugz^CRepC8M@iB58QAjl<^t9 zO1|N!sInB>`ZE~|{I>2H#^w~QfCVnJiQ{0<)^N&eT+)q6bS4K270nmtV45yWQfFow z)lWh)sEFCzwX?gbNA*sch2skjDXUvJgZ_P1jv=v;I0Xy0I9BA_y$iRl4dnMtN9(Sy z#}~>~AjZXsEouKZM!1%8XgHkWwW-AGS!A7&9cZ`?Nd1MIKZ1!t2{hVS?!A^I^@qk zY@Z(nktpDu8%DlaF&|=daS&E!?PUm$)JmB8pv4YY>x~(41+0DlZmE0wHi_3A#jhFb zMT~p2(gcdW>lf7JH}!c6d@H%vXXAs5jeMx%OQ0BGRAq7&U^v8mH%q%$@`G6v6q5U* zXn6jyqtpYXi!Qw8)s@lvFeqQh^6&Xb8h8rn3WL$~Ugz&!0O*W`+}^y)d4Ac{D-gBh*geqw9l}^VebzPy&PI#`;dzpxwUDMKYJX_@qDL3JBFX<%-mxD2|X~my5{NTm7KB+SeZ(lh1}a9_oBvXQ^H5e z(FUa8BSlSSw>bluO&hQtx#}2SiPFjb{h-1G1c_3yfd@(|G5&s6k&Fz9QWZ~OY6AcG z(rFYSQ7TA&edX&vXcg#?hZzja;@f0@yX61hl0i2F)77ycy2rhj_k?2@oIHao&lHQ< z&rmeFQcCNK^WE>4ivHfGjTn0{MFuWNOkBGuhdE48HG~e=A@Mh6B!6>{>d=A(CaJ9@d%NXFl9*r8h5vZXSleA9G_`Rldzl;N$ELE^d_$Tf>Z z0`>KPCi$_S=L{aSH5**Gx*id|idH0>OY`Z{;Bp0zJ*I3$4-bar(cA{`L(b6X8^ebj z{56O#w$UQ$Zq@y&hYpdh)do+@2Y| zl<&E;p^GkM3{{It8Jcmn1}56WjO+~oXbxojb$h?8Pw(hM+aq00uk4rx?#F9mKD6XemwC*0a*@^Uo|oaMxGjqd`> zN9A#0I`MLQL|0Aq%wh@H2Md~+G!xowf&}>7kZBdd*it&=Z34BwpPX*HrFEh#K&BpT z7QMB<$P~xGc`=gaH8>{sU9egev8c803!DStWG>tJN;n*gg%dUVVLHl&!D0fv_aYYW zgwHvs8ETAv@r~4R`k}onDV7x>Na$qr3m27Q#=8_3)fg##6-?j~CKlD#ieUuy@Gvko zkc)Y=41bbH*}zCcr-3yvge<`^ShB#*WR+3-Q~frtSfbtMWX`7fEHt}pj%~Jl&!?x+ zNFsRy1AQHH1D}Qy97Z*^Gz#M7NNA%p3;XefCet-+bK5LvVI>whKASgk0r;4}gS%+& z`CmW7WI;Uz$Oy}8Zto1Glu~L+;(W$f9fogx(6mvRT(EfF=4p&r0Ik@g2;)welbS2f z-5cZWn1W{>&pip|{LJBScO*-+c_J>d@b#T(mZAuXRR>EuslWxaL6|^a=o#8}`XPjr za0x6rqWdp&{cw;cGmkWlm67z+uO-E-AunG7urELJn^(CbW2=S2%JRw9tHsyKZXzK_ zFump$SixjvGr|lUm5sqmjiq&k-g{H*DEl?RgN_6Vfw({o8x9%*sm_P3N;p`;I6a=X zZ$)HTqAs5V5I}2B{xvK`jTsDp=#q?Qi%Jsb_hTCSi3@9Ke$JYNreB&TrIf)=VdK*t zcN*i6=Tl%Z0W+*IF$b{8iaj6hk@DNto7&;hkEB@jCO*l-`ZQbg)(2Br#BHN*Q|!KY zdH3P@QEJkNFw;2QpYdh1eAgI;W%KV@mL14?mJJJsK|-JxjlpEffeH_cfUQ!Jno0&3 z5kigP`_WM4Ey~%;a)-}3gqf1#ms;7V(6J-^cqL%--~(1hwlV_~I? zn)HX1DmL>NM~h-L_qMQ7i7|7XK}f!(|!sK^yX>dPrZ>Ad&nE49JKz z;o9bJvp+SNk2M-r%>4;5Qb|EJK{;*&oX9^Be_KpQ#Gh+9!%pf?#9#e)#9z~0)%4#H zf12MBe+U`v{~Z~&`w6qb7WW)cCbr=bo^M~RCdNWv994!{YS_n`?PB#X)=CSM-yur` z`_{1$5z0lJ@jyp@C8Rc4!A%syt71;mm-byENn{;nt$J6PgU3Qx!?-$`3INYy;=^U4 zLmhN};(>=4vvn0TP7sVa5igspa!b6lOGdCR4};CzMzD#F{q>aQ@~Fub#aDO|oz8CkYQDHWHYE(^g>=Gq%KA0TqIR=Vcu* znJQr`uQA!2B8rs9SB6+#&Rhiamy_W&215iRid}3L9uLAWjIR9mYSVv@MUy8Y#MV>d zcHPqEQ@>(tSF4%gBATKb{YH3;`` z+vaLc>{~PHa!_GJq_n_#?&{*lYcP6w3#L(sm^z1&7(|)nlfO;L-{hNJiht^O{b~_? z(0hRMQx7w!F1o<`kSJ$8f+mtxaC*4~x>Tgh%-1Ft-bI{vKW=wL}mVgh+eWP#bA$DwSN zki>*6w?p;SKO2B-I3zIv(-x5d{*URTK?6xlpl@NBP585q{2!4CR%GG*+`EGJuPt2r zc4rWl37@1`E8U7P>x^f)ls6<+12T;;e(O0$Vsi%tmx-;`S6YUT)9o(hmI5_SX(=+Y zV8J9@htU9AR~@n{VwZ@1&!RDM*DVdcNrn$Mep5Uidwr114;ySLioXERe*@5bY*k|* zvcXZh7q%r7@HPfjaSFsP{QL_5_OE%E*Lq$9-e`&Xz%Mn$b6eoyj%6gm%der|*O$TO zYN_aplcqDG*&Rh1ypl>$|Mww-YpbWY@(b+B;B31|Z$3IkK*%^SmU!-NLXn$=cx!|C zBOL^JoLVx+rE3^J3+jQkQ=F%9lt>0j`n-*AQoIx1(dB5Q{@A8EP^`r%<*CTsi@#`R zFWw%GK|RGxfrhdHK{GKIpLa&8rqzc7>Dc6Ey zOLZKv&Z%Z}0AysuKT426&&O8ic_DD~LB#338Uu{%Bxl{d6~tJ7e^d`UQffejMH!#? zbL>^}Ouh*fQMzAVut@k#7u)|e?+pBat4skE8Od7D)))iK>wj-dQKHTQ57x5!RD*Vn zSz|ezZvJlW>V}grRoBHs;n8F3p#8sOp&795-)51RdQf@LT9=k!lcl+!H{l@GMf#xL zWFS|+&ZgG!9vow*csCY>|6j7uj47NzjIt8N?1FAt)5vmg^Fl_dZq6zzLR-h}5G$R% z87^WD^XX(j6UmQF_%U1Uw}GEYMI>pd?87~1PZm)Gi2xWZdHZed*coxyMiSvRw5m)h z+M~$G;Ut|Q&k}iO)eUE92HK5htZq)T94F)q-f&5KEB1Y25ywhno|CINqfuN)L~`g+ zbky5`(A`$?X#&|tdm^#%ZfDYF2RB1~E{8kfDo@c)mqJzaf$NLYSVT^kMgahr`F{Uj z0JjxA;}hu}S;5UA@`bqd(=+q3@}ko;)sSlqT7Gy~`6yCdvk&CnVX_KurT&FE+oHFK zr98zWR>lX;0VNk|L1d%Azthp=d?Mh+MP1kQNdo9x8)7_H`yPvK=cILQ1j?)~!tYk) ztmOq{Aci-mfD#3!krEw!IEEp(yqvzW;$fm(tq^v=C#KG!UR4&En)d5vQ>a#gv7=#v z&pxl*u;141*2o&okRHDVMfIqp7%3(Y8lq)5|AUmX|C7u$ z4)^S)?9r)T3|2?npq@fptL9gfoY%2jrguG)J;Q^Fr6Uz-b34kYk2&?XO8MN3L>`6n zKc#q&p}DZYlU&GyAn}N>HxZ&bmTz*LWxnVRT#d~88WJpS2Aq86=hJ(a_?G93*U)6F zH$tY+@;jAa4PI1x;l|yzA>x*bSJf-JHXU~`?#E2Rnvt|B8L%|XPo{4b$2Q3=PbIB| zS9bia$7HCqKHH7Dt1c|N!xwj9;6KeDV+nk*d!%EiIsjVFxu@sFLiOu%?#Cgu-`hD4bQ0bhRdzB-0|@tYtpE2OE-A*f+!YoY{(z zHHNM((#ij;owMv}>)rY_?(Pzbl;T!eytq5X9RdYfG`P07YjKKGq(E?Y_X36D9<+F% zSkY(ozxRIjIOo;*04F1FGDdXezUNBj@0yqLB#blly?*JrWQCvKV$jqqdljw zJ*UYZ#F(XVtZ-z^XzCE@c%Y*+S1k-u+<|vnS;G3S)ww(Um_2DR(Z!PPS~KGz2_X)+ z>Qbk98JsKOn}J<94?`JkX5q~j`!oZXtMdRJUD6j8s4&AAe$J)dGSu@bXFcw;>g*QF^Qj-M7+Bk1_< zRDn%!ZpjYXrDc;W^xfV3I$+P(K4?0cJ4eYN$;5+L;uf`wY-8p{$FJ$q8NTCoX2E5Y zT^&aH>Xa5sw_zj1iPeMcS%tle)k#v@*cstl6*j0!A6hN4^xn*foy7Da6CJ`r` z1Y?7}R*;3|M>P{5H#h!%BUHB?<$`#h-Fi;#co<7g;_`B2?EK(fIpc8a9MA2y|KNJf ztK7n*2JwI!A$N`{F2S_NhuEXzUt>V5E;q?hKGBlZvM+Z<%QPbMTd8F^j2^!8JvIS( zZ@LQrKRXt$)$DguMV(!bI979y+kG!HOn=2vK$@;VNt#~5y$%}ob7bGUaIi8n&2rjr z-PFTA+vQ}JhP7wcq4(w|oiZ}piO-7Fq9@XVeWu51?MKartDn!~N&zuPOLc z7V0a@KcoPhE6$;X$fXS`d}=Y+3M`~#Y=T`X(fslG{eiTgEZ@z-?@rt}dVhhwzX;vH zKL}m5DpyhlE`~SpskjmL{7X*p6GFJ3skepvx_sNMcyJ64_mAi1WeqLJrcT65emmCW zVUM>ed*TEGi%YaUaxz>4^uj{8&!6EiaDKOc_TdiMzt~Tygk}{vM_$vbPVLYGr2y;e z6VP&Ry%i1a!*dzi{5T#FA zMpb7Ct5=s|9SFMRF({br9-bj|`xq%0-$_&WUThipzvMkf#)vZX)^WCyBkwAK7kXZO z>pLV<0c0IC z(w_n1C?8&pm^sD)etBMP_mb^g`UVFln+7KK%Xq1&G1W1p;RQ%h+12;WQZ?=R$Z|Yf zclT=4og3O7P0n18wIY=krX4tluv6s(XHQ3x#+r0nLJZwq>Tu8yv|u2!XfH8F09eR= zhb_L7gN_QDlTSri(M)2hp?XQRr8t(xy^=QVhU}o>53sd1w$J-Vk)eyL zqih{v%8ciIMKwc-dfg8@Ej@M<7jmX6M>y1` zP$bfJ;gR6#DbH?3o@u_)k?{DlSNa(b4ULkqre+FiuEI6!QSE&wzW77=p&l^RBCO*? z_x9$+Hf`&dhLqnSsHB!d4K{b9JLhC4(hQ7ny&ozVVb+>l2-Ml@5aS>! z;V%SLVK21t&^}Hq=kpEvX?W*L!xcAG`0A&^*yD32ZwLv`PC^iZo|ha=3^6aC*n0`E zFtg;~*T-^SDJjEPDkcWCGK|R3O(@H;;KqIM;FoUy!t)Si{|`J*Sur@4IITHF*DcPf z!AoJMHr??3>}LK?#vP`k4q4I z>jYdFUFz*Y6U#v(Eb)YS9o&wOyX8qw&6fwQ3*jktU*7uddz^?(h_`HTqpW%DV1nE! zcli;Jh4hy5wu6W@EGCLS#+ql}=V$;Q{o_!DK;wFr_vKoZvF#TT;1KMFYGZtm>6N5y zC&QZ6^)FRdJJ5Otu+#GEia;>sMPzO&XMQmp>ChE!?0EX>On1WZdn3?y?|z~M)A#j9 zLBWMC>^9`$XF!n#erixnS(LV9SSj_ey1A`{BKvD$t8Y1l+!5S|rjRCKL#oADc7f*2 z9?Z$gfzT~)<=jj8Rc4oy-Je{qyw|K4yJcX$W9oaV@3LOQV0(&rA5J+&K?jVfZ2SAG z!O38vE3fcagtq;IcaFD_I$D+M!J=J{%r+} zMXLvGI>x)l8(A*~b4zUJWAxA2=}yg_g(wZUWF+4o+0`1UDfn&P-$+pif{SRW8*F7d zVk>ahT+@lHw$BfiTplm9z_L`EXl)mRezc%26aQzh#To5KD_JaG>{Rg-@WqAw(poeO zB?Ni57@XJ-M~G*pus|zp<|V$Yt#cWIBe)<~TfXW2%Nu-n-`H8%m=BS7HP?QdZZ2#7 zU5Q+GrGekmt*KlOOg+d5WGH0{LxJTTNH#1a0jlDXnK272N9)A{W89Itj+TUYKAW?f zk<+i3`7U0*jn$hu+fKp8YcUyJ7!K<{;xJ7}|L$6Y*{_L^zQGeV*biHI;KzM}u8SUp zmdc=n#T6vn%u#cb=~!3i{IQCJRNHBb9zBPJEyZ<>)-1qppGfhOb*%hf(HFQ*z-^X{g5{hP-dPihs(D(8_}L-LsNtoYry*05lmu`*Y}o%zF`&DKFZ ziY*BCwXJ-CGH_cpIeJFxvnaNo6&y!kfn7)0@ClD5+^2i{i#GZ8moOy%h<&GtAwN&> z>-x3uWxu{ql4W`Whu{jq`sHBEQ!pHT}sLoRU~XTp8L;86K{(rBLj ztzk32yub|I!hiwVBMa^n3gxm|8fb^PNRLU(BY?u!IOOwIJLVAWcN~G|pEZ8IfpXUg zdQ8zzItbX*^7XD#OAZkr9&69(pz$O=$A3;H{NB-9iaq|Vh-24?dctYt zH`WEW%hJ8JXCXFzkKlCl@R0GoEJuMS8~Trx>vzc5+J9iEZgzM|O)7jn*yVB|cN=P* z@qLNjFMg+)269U!!Cj9qj9&AGVj;!o2c;Q$k)?I_rZ9?~n!7uo+E;j_I;QS6Qhe@a zh-$kZw9tEWeMW`H|)Q1%j=p`xFoW)UBDzdua0LU_)DXa#)^#zB7zFa_cx@b>3Y`@K{dk~FUZha%ovX1}kXMiS> z5v%U=j)jN*#rj)y?L4^^&a%@#H;C@*d3pu{x}S_k_1C#?W&xx)YJ4SRUzX#sIpOiS zuUlwT=U8>WF;!4EKI_&<`?4>3l3p7%JDFyX@B( z;paWz(~T_}WmggPwqaa7CVCvGJgBcf=yo!Q>TY0O=L7N`7bHQxF~OO4gM>qS^wmKYs-VNV%R%KE zT>|CTd9#~8(V)Njks-%xx)efr5Q*A6Ee1wir#EwsAKl(R1P5Xg`cVeDZCF3mDmpvj ze2ioI81Tm>{gP)3qnt)8Hn@cpt?&HHtv4|jBZ_E!Lf6;%)V|mk&yR{4{Tq;q95(H} zF8@I$ndNxA+C5iuvJI%B4K&|>&}utrGje;jM6KPLRokrAsV;$4LTb`)BcecQi1fIl z$LP?0e$%XLS3^capAR$mvPeJv#6K7F1?|Y*J0RJNkiuNxznElt)j$~WD5_veE6~+o z+9=5zbUv-v-NmGs0RzOC0i@ilj#o-_lPJlmC^5Q#eddaDm<*08M zaasZYgUjo z_TKubn`$^;>z|Jc7}LhEZ_HtxCAz0n6H$my!qU4{6NW#Fm~%Jd)<2; z;kL9P7MGu7$&LnxVdTp4;}ot?)t)FGc31Ro#tTs-b!E_#B4Z+S_S5My&_{iks|-g( zPYH_s<|wMhMDpoRZ7=SR?V0Ryz_@IfNm_J4t~8n*ahnhML+UXMw^%H_mQ*BDj56@p z$~s31Yji103PF;&Pe3w)UV4?T2`}7@t2`mJNi4MJ0Dm*gs z^I++5&+1mq2a#uup!5}IkxrJL7D;9;J#UABJJI&MSN_B^XUqE2jPJ*f(ZYn0)I#kvpy?o378O;ixZCWC%oyw-IFu5c=OK1 zPb03%;Yrcp@M!?8$UTsXfxr|PTBR6!2d|}tVm;D8!zie+$qtrIpKj-?*4vZXuH!SS zU?K{ffy+|!7X>Kvkl)=DJR72P z@pa_aHROAoCIP-f)ues@v3DWC0Ug(mSm@2{d(0c!*!&zr&$a_6R=>WHm)H;E4OESG zo~CGA8kkr~zy>Ccg7^uDYcPyR%;#DC)cs$$VlPuy%aS(kEPS+aG&NvQ!)ViUv2?l@ z5j-OC-Og5jnp*DUYdjm~=q$^kYrO65)GyKjOglA#sXHThUG$zHzTT7X6|Q>3q{)s& z4DGGXKY3DFVLEd{U;N67)J(%0jA4~iUI*BtT35^*=Z#LmTiWn(TlCx^p3*2z(?OYv zeps9$Dsznk`4z`tDz*JGxH_=!Bg9WrY3hR?L3A*6KIGybP84rnBBT3+a41nM?wq1D zAurP*BG(cfMC8WmitKUeRs1`5fgQ&eUzi^>RsbH-z+;kxNqD7XDoW655#dQ)n#Pjg z)dxy_Ix3=b2$o*z1Z+gc)}rNRZ`3~ahbagSysdvb@pFyRAbPp1Wg2;e>Ri~u_Ea28$#(mky$sRz->U=)wEah;rK9+ z+hi90!rj#EVRuFNsgMbTQ8l4jswe{#1pseWCB=j54674wehDpIpTQ9Mv%$h`?-bq_ zLW1Yg7gr$YwdryD%!3?QI+IM)yK}U7_R(XeQjm=Z?;~)y!yAY$!LC)FX#dXfP|2V8 z`UgeWJ3$#3o|C0sufS0QX5Y~JV{InAj zOOa*{a;4w0Ihv2o$BC2bTrjuNQAkh?5fO@6#lzU+{Y^aolaq~QbyUvFeXAJEq1B6> z{dI&{i=^6jQ69DphSmnTKGrO%sZJ$mc%Qm=(TCn7k%9whsQdPs5AA)=c z-2=Npmz1B3cZ}D!@icaIkM7+7479qTk!1Ck5tKTA=-zX!XHWe^R2^h8jkzq=qzq5E zcVWXkh5fXX!%4s%IWaBJehUi*FZ;x=JC{F(NI@^&WIl54>Nsn@EVif#We<{U5=^OO zDmlW#Q*W$;NK-IIci*FsUTyRgBI8X`6WohqXYi@8Qe8#ep(RB@=-RIx*hG=i%OF2W zu5BwI0bOE$SOsXw-o2SVL8oBOMcKiCGUb<9_~8YX0BQDObN?!ZEb z;BaFmxA@{&8qji7FRC1bVlf$4`_VK#Dw!-pTHTzD*ARD52m?ffaD1a5ybp2Hjb1(v zPGPugdi15E^mDtz#8ztJ{r9;dBLdwrU$owQ z&mS>ep}{Ev?4WAP>#NMFH?wC=Cy$faDi1hyZ11(8i@obVOU-Ts8_O%%`!K)Z_QJ$5 z84%P3XC#9IAryoyGbo|Ruvxg!JZm~8HhQ+1!?k$3D)Y=!WZ^CfhsvmhE-60_#hb-R zv&D2Q^p+ZILP&v*bamoW;Y3^a$UB87(Qu@44d;96*F=1n=LNj_(NV16lQ-m#oGFew z#E#l%uL781qy)TA+I4r&IIRmi(){bw1Ae(FD2c-EHXQ^v@$kJKR9fhI19R!RgoQEC zE1%bdH=Sv{Dm0>nb(hfdwTHd?#smyA8othY%VvJM!#vDghu4utYog!8MTUo{-2JRk z+vxY{V!Zw4N_4L1aO~x_TD>h4L0`wvNNZqifNf@J$v;UagHD_#U9OqO&;P3jz+z;^ z!-1Ymr^H5^PO1_RXz4SIYYDok?E<`#_qW_{CKET7i?MZw7 z_sJ3e!W{@QNjSzQK-QH zM2YO%I_s}TDx-ZO(N)WnBdc|r>v2_jC5UT zrbWeCz%`X$_l|dJYPIQ%0045N6q|UgdFwRhADQOd(ZeRj6KUpfBQ|Oxs{gWGiI9z- z1S_gu?h>{UP1dR=JMf4X!X#>6-u`18v3lOl&>)MMly2JO4@gx$PYd|f0*var8$ar? z&lm!7ax4G|DfL-90}rMXDH~ftPlA`Vbg=cqL%-jGZ-*z<=v!#H1BNYQHWIPv*%HJ5 z&87DD6-kfE{L2yo%M)U!ks>m6$fw-&@@Z6&O`-g=qu92!%eYN5*O;t{f&WwZ9}v}@ zIrN+*IS|fKR|DXac zEY8kAp6u*y0`(nGA^{XEbqjv9QoJk{}wtoO(@(OmV%bv|;9BjA`>gM7LG z=cI3_E#-*cCrpj#?U`nR<<-t(?kIYlD%WC?#NDMvA`t`@Gvt}V-j`{VFS$OdTJ#fNM-xFWl>4>ChnQ7NYwk9?W+`ffB(}us`Ryi@SzHeee2V$4!Gx2Fv5tBiJv+UF)cA@R>!YwlYoQ#+d4#Hn-Qj z@Z5Y{dV(wF0QtM@ib>Z2&OO-P%d(e@zk6`jrUc*2le^<)fbj56kURV>3jLzdq}&sH zll z2nm$GJ0b4cnhSnw2bhM(09C#E<*nnEjScL#%%oGN?-L{SVbYTbyer)QZ6*1A9bnnG zxm_cqFF$u3)g2cS+>2bF1Sv>Cx0hSA(C&XMfB4fMXI5|mm;9hNu@rBMcm9Whz9o{H zT145?2~ecPFMSby^@>uQa0dj~)gEAYX5DCY7eyutKY7HFJ(>Hoqg?2bg#VUz8#j_E z6wLHj{9(}9&3%!4x&FtLt0i>lgeqXr6tKnv_UBUQw?F5rmT<+`$xJbOn%9Z(pS-roe^}?o(-hAifaMdSzF3a4Lg~qn^HO_gMzlGjPpqZ0JE!v66VO8TPB|<3LiC!B1 zho9r89poOVBs!btriN#&>T2pfEq86};u^ZfLZZ7ifYjKaJj}rGq1--ug_!$QlIW9~ zxEf5q>HKxn>+O79_C_`P)q(lNRmahzH3f$$Ps@_=L;Qx)^=}x3IJu@TVauIuDCp%4 z-DWMOTe*kZn_vNFpZTK0q5OpB%nk04My^qAdFTN53AF5+!TrzR$9%3+l{GZl_l*A% zV(Y!TIG!eM;-1$d=o0@F^YPQLx@*pY@M-xJp4!~KOtP!-`DP`YKgX#u=Ip}t-B;?3 zs=&4+L|YjH0>qENSuY=7G|J8v1xg_xty`)?zFvtO#3O7bbck?%EXv_GGP%`$Pb2;H z&BQa6etXnizC3il`(>FGpWTYJ&&M5$n-zX`{yv(p#9UgL6?$UBp`2Eu1f#}>@Ga$^3k9#T(nQLp9o)yYCBc)$9G=+wMIJmOh){;;g*An3N9} zU-AI_u+Hx5x0kXiIv3Q7NL2W3a_+yET;bkUiyJ4wQZ`)A-hshnYjfDJhN@T=e>eI8vaNS8zx zm7UB2g}d5L)UrnBWdpjyS5gaySkZz@1)mQEiVj;-WhpcI_=$6dKkfo zp>#n7WYD~7XaI=4dg+PRLuBCA?g=x)@AtT8`?YV<_vPyjCzh0L9NcSKvBvd{ z<{#tiF5Kh?oWUEF$1$tOAGx-5Bz6SEpFMZ?>|Gp*Pa=N3Ny2E3k0=bf`eA9-K=_MI zD!Y{E(=Y#YUevC0!VE$G-PT3rAC+ceHtQ)C!8Q4)h>`5%P8%s4@uCv#sm05_0=ZDW z*Sd$)kjONywK)vgZEI5rGTFI_cI386hV1}y9t*S)^`cH;`$Jd^LvE$1&in!6F#$XQ z%S-*aaX#in3RlsqRAnBLAWF%ePr>rG*Z!VprFmBVIzwt0Ndj$V2Gq~;qEB->*u@X; z?n4SGXw96AI!@-v4QL=SXaiw+lH3(?#M`dCdzHTQ?tw{fefzR111N6($_|??Bp1 zy(pf=mvQnw$d8-!7w2<*tO&6FXlie?Yq-tHut>6L`ob4<+w!y-B6j&yE=^b@h3kyW zs5YPW+&mTVUAF2s)U<}&=4oH-xqEjU$4otMpxj&QshycF#QE?+acxa<2hTeTREQn^ z0`U&>2dp}mUS0zqF*X8Dvb)gfiucibwe)!roU~ooQ;JQ__4jcnp`^k`!Q{6nh?OaG zF{%{$XV<+4CAL`j@-xAFDdV5CE&|W>JP38=<3MIddpptq1C8oIF`?-c%mH=q# z^G#ZT#^9KsetFL(s;0*0LIYF7V5L8SMNwbM203%Lsn`ypf1mpB?=P#X=e~N~v2k4A z?!A70>(kJA)^xP#jEPK0J-q1oZ0kt>S|mh<5t&`^n`B&Q%2*KU=Tz&T-M`Fb40y6D zOuvN`>$~>Y$QI%Xsso``14?Lw{Xd zeyA(hX|lIbS@yw#fjAnv)5a`Uc+BiQKcs@7eguS1$e~ zL(JzqAlMg6gLS>7?3h;2?ks6pOjP*tEZJ;&q6gR$Nr?qz`*q&+j!PP{LQbK>H`)cr_5*{Yu=Bv zDo(cZDu_iK$B3VDrc<^#7`P0^pgKn!Dg+B&#GpsEX@iu~gWEue@$yd^@nidPfr@!F zC@KXt8us>tm&d>C{L_k!IQ6?EYMJ?=40{Nf+&oOZn3oMrP+m^xZZB`(48`R+FP9xo z?4x8S%f?uK=cnun;Z35XO_5xTkcG9{vK|bt>*t((n(;R%{4M5s5k+!r8`OCakpu!Z zlQk`FZ7OP`*A+$|pskjosMs%D)`!b_q`4hzLg1#^q8%Mi%Q3%yNG`2*E8oaQqUQ9* zWTRsy1jq#7q35CyRWm`Pw^O1z<{H|W0_B;0sJDYy9-Bsg51T9|plB-xD6^IEagy2z zAhxslL^OMN+QOg=LB(Yd_@xWn+aHn|(DqK$r0b+tm;OpXfa?;v0T1ucevL zSjBEt42a@@Dvm1OvdTtW^GBVdDkA^-7G{xVcEo=E4F@lb#{}1T<^Uh@1O%Em^i7y* z6LWxH9Pj7Zw?c)E@#|Iybjq)_Rd$r3#wcQfJWzlNO(U7@+Dd|WWK>Erbg|j2__{fH zgJg~p1>S*s6mZNc=d2#}3gU{V$0QNnmr^0Wn60o8$1@wXDg1VYZuCn?4`a+X=gs>p zB~}!;JMO(6|9virUkhlRXkfnMrUSQVA&9et~NwNd2lKek6C6LPKB*LZt{~i8cjsqNja4~_Umq17p0k|m2s>xJJnFsw3 DQr_yr literal 0 HcmV?d00001 diff --git a/docs/examples/op_fuser/op_fuser.rst b/docs/examples/op_fuser/op_fuser.rst new file mode 100644 index 0000000000..9613ba74b3 --- /dev/null +++ b/docs/examples/op_fuser/op_fuser.rst @@ -0,0 +1,353 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Operation fuser API +=================== + +Motivation +---------- + +Transformer Engine relies heavily on operation fusion to achieve high +performance. A typical training workload involves many memory-bound +operations such as activation functions and normalization, so +replacing them with fused kernels can deliver a significant +performance benefit. This is especially true for low-precision +training (e.g. FP8 and FP4) because it involves extra cast operations. + +Managing these fusions can be challenging because they differ based on +operation types, communication patterns, data types, and GPU +architectures. The most straightforward solution is to provide +monolithic modules like ``Linear``, ``LayerNormLinear``, or +``TransformerLayer``. These conform to the interface of a standard +PyTorch module, but can perform arbitrary fusions internally. These +hand-tuned implementations can achieve maximum performance, but they +tend to be complicated and difficult to modify. + +As an alternative to this "top-down" design, TE exposes a "bottom-up" +operation-based API. The user constructs individual operations and +passes them into a fuser, resulting in the same fused kernels as the +monolithic modules. This approach is more flexible, making it easier +to support new model architectures or to experiment with fusions. + +Basic usage +----------- + +Sequential operations +^^^^^^^^^^^^^^^^^^^^^ + +At the most basic level, the operation fuser API involves two classes +in the ``transformer_engine.pytorch.ops`` submodule: + +- ``FusibleOperation``: An abstract base class for tensor operations. + Examples include ``Linear``, ``LayerNorm``, and ``AllReduce``. It is + a subclass of ``torch.nn.Module``, so it can hold trainable + parameters and can be called to perform the operation's forward + pass. +- ``Sequential``: A container of modules in sequential order. Its + interface is very similar to ``torch.nn.Sequential``. If it contains + any ``FusibleOperation`` s, then it may attempt to fuse them in the + forward and backward passes. + +Thus, using the operation fuser simply involves constructing +``FusibleOperation`` s and passing them into a ``Sequential``. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + # Options + hidden_size = 4096 + ffn_size = 28672 + batch_size = 16384 + + # Construct operations and fuse + mlp = te.ops.Sequential( + te.ops.LayerNorm(hidden_size), + te.ops.Linear(hidden_size, ffn_size), + te.ops.SwiGLU(), + te.ops.Linear(ffn_size // 2, hidden_size), + ) + + # Forward pass + x = torch.randn(batch_size, hidden_size, device="cuda") + y = mlp(x) + +.. figure:: ./layernorm_mlp.png + :align: center + + Operations that match ``LayerNormMLP`` module. Note that different + fusions have been applied in the forward and backward passes. + +Quantization +^^^^^^^^^^^^ + +The operation fuser respects TE's APIs for low-precision ("quantized") +data formats like FP8 and FP4. Constructing operations within a +``quantized_model_init`` context will enable quantized weights and +performing the forward pass within an ``autocast`` context will enable +quantized compute. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + # Construct layer with quantized weights + with te.quantized_model_init(): + fc1 = te.ops.Sequential( + te.ops.LayerNorm(4096), + te.ops.Linear(4096, 28672), + ) + + # Forward pass within autocast context + x = torch.randn(16384, 4096, device="cuda") + with te.autocast(): + y = fc1(x) + + # Backward pass outside of autocast context + y.sum().backward() + +Branching operations +^^^^^^^^^^^^^^^^^^^^ + +The operation fuser supports very limited branching behavior. While +the operations must be in sequential order, some operations can accept +extra inputs or produce extra outputs. For example, ``AddExtraInput`` +will add an extra input tensor to the intermediate tensor and +``MakeExtraOutput`` will return the intermediate tensor as an extra +output. When calling a ``Sequential`` that contains any of these +branching operations, the extra inputs should be passed in as +arguments and the extra outputs will be returned. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + # Construct MLP with residual connection + fc1 = te.ops.Sequential( + te.ops.LayerNorm(4096), + te.ops.MakeExtraOutput(), # Output residual + te.ops.Linear(4096, 28672), + te.ops.SwiGLU(), + ) + fc2 = te.ops.Sequential( + te.ops.Linear(14336, 4096), + te.ops.AddExtraInput(), # Add residual + ) + + # Forward pass + x = torch.randn(16384, 4096, device="cuda") + y, residual = fc1(x) + y = fc2(y, residual) + +.. figure:: ./residual_layernorm_mlp.png + :align: center + + Operations for an MLP block with a residual connection. Note that + the block has been split into two sections, each with one branching + operation. + +Developer guide +--------------- + +Infrastructure +^^^^^^^^^^^^^^ + +In addition to ``FusibleOperation`` and ``Sequential``, the fuser +infrastructure relies on the following classes: + +- ``BasicOperation``: The most basic type of ``FusibleOperation``. + Examples include ``BasicLinear``, ``Bias``, and ``ReLU``. It holds + parameters and state, and it implements both a forward and backward + pass. The ``op_forward`` and ``op_backward`` functions have an + interface reminiscent of ``torch.autograd.Function``, e.g. they + accept a context object that caches state from the forward pass to + the backward pass. +- ``FusedOperation``: A ``FusibleOperation`` that can replace one or + more ``BasicOperation`` s. Examples include + ``ForwardLinearBiasActivation`` and ``BackwardActivationBias``. Its + forward and backward passes (the ``fuser_forward`` and + ``fuser_backward`` functions) must produce equivalent results as its + corresponding ``BasicOperation`` s. This also means that the + ``FusedOperation`` is stateless since it can access parameters and + state from the ``BasicOperation`` s. Note that different fusions may + be applied in the forward and backward pass, so a ``FusedOperation`` + may be missing its forward and/or backward implementation. +- ``OperationFuser``: This is the class that manages the operation + fusions. It launches the forward and backward passes within a + ``torch.autograd.Function``. It can also replace operations with + equivalent ``FusedOperation`` s. + +The first time that a ``Sequential`` is called, it will group adjacent +``FusibleOperation`` s together into ``OperationFuser`` s. The first +time an ``OperationFuser`` is called, it will attempt to fuse +operations for the forward pass and backward pass. Subsequent calls +will reuse the same state unless it has been invalidated, e.g. by +changing the quantization recipe. + +Quantization +^^^^^^^^^^^^ + +Each operation that supports quantized compute holds one or more +``Quantizer`` s, which are builder classes for converting +high-precision tensors (e.g. in FP32 or BF16) to quantized tensors. In +order to enable fused quantization kernels, operations can access the +quantizers of neighboring operations and quantize eagerly. + +.. figure:: ./fp8_layernorm_linear.png + :align: center + + Operations that match ``LayerNormLinear`` module with FP8 + quantization. + +In some situations, like when operations are split across multiple +``Sequential`` s, it may be helpful to encourage the fuser by manually +adding ``Quantize`` operations. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + # Construct layer with quantized weights + with te.quantized_model_init(): + norm = te.ops.Sequential( + te.ops.LayerNorm(4096), + te.ops.Quantize(), + ) + fc1 = te.ops.Sequential( + te.ops.Linear(4096, 28672), + ) + + # Forward pass + x = torch.randn(16384, 4096, device="cuda") + with te.autocast(): + y = norm(x) # y is a QuantizedTensor + z = fc1(y) + +.. warning:: + + This is an expert technique. Quantizer configurations can be quite + complicated, so the ``Quantize`` operation's quantizers may be + suboptimal. + +Implementing new operations +--------------------------- + +Implementing a basic operation +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Subclasses of ``BasicOperation`` must implement ``op_forward`` and +``op_backward``, which are reminiscent of the ``forward`` and +``backward`` methods of ``torch.autograd.Function``. They have an +argument for a context object that can be used to cache state from the +forward pass for use in the backward pass. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + + class LearnableScale(te.ops.BasicOperation): + + def __init__(self) -> None: + super().__init__() + scale = torch.ones((), dtype=torch.float32, device="cuda") + self.register_parameter("scale", torch.nn.Parameter(scale)) + + def op_forward(self, ctx, input_: torch.Tensor, **unused) -> torch.Tensor: + out = self.scale * input_ + ctx.save_for_backward(self.scale, input_) + return out + + def op_backward( + self, + ctx, + grad_output: torch.Tensor, + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + scale, input_ = ctx.saved_tensors + grad_scale = torch.inner(input_.reshape(-1), grad_output.reshape(-1)).reshape(()) + grad_input = scale * grad_output + return ( + grad_input, # Input gradient + (grad_scale,), # Param gradients + ) + +Implementing a fused operation +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Subclasses of ``FusedOperation`` should declare their corresponding +``BasicOperation`` s in the constructor. They should also implement +``fuser_forward`` and ``fuser_backward``, depending on usage. These +functions are similar to ``op_forward`` and ``op_backward`` from +``BasicOperation``, but some arguments and returns are lists. For +example, instead of taking a single context object, they take a list +of context objects for all the corresponding ``BasicOperation`` s. + +.. code-block:: python + + import torch + import transformer_engine.pytorch as te + from typing import Optional + + class ForwardAxpy(te.ops.FusedOperation): + + def __init__(self, scale: te.ops.ConstantScale, add: te.ops.AddExtraInput) -> None: + super().__init__((scale, add)) # Equivalent basic ops + + def fuser_forward( + self, + basic_op_ctxs: list, + input_: torch.Tensor, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + **unused, + ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, ...]]]: + scale_op, add_op = self.basic_ops + extra_input = basic_op_extra_inputs[1][0] # Extra input to add op + out = scale_op.scale * input_ + extra_input + scale_ctx, add_ctx = basic_op_ctxs # No state needed for backward + return ( + out, # Output + [(), ()], # Extra outputs for each basic op + ) + +.. warning:: + + Remember the contract that the fused operation must produce outputs + that are interchangeable with the corresponding basic operation + outputs. + +In order to make these fused operations useful, they should be +registered with the operation fuser. To do this, first implement a +fusion function that can replace operations with the fused operation, +and then register it with the ``register_forward_fusion`` or +``register_backward_fusion`` functions. + +.. code-block:: python + + def fuse_axpy_ops( + ops: list[te.ops.FusibleOperation], + **unused, + ) -> list[te.ops.FusibleOperation]: + """Sliding window scan to perform ForwardAxpy fusion""" + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + if ( + isinstance(window[0], te.ops.ConstantScale) + and isinstance(window[1], te.ops.AddExtraInput) + ): + window = [ForwardAxpy(window[0], window[1])] + else: + out.append(window[0]) + window = window[1:] + window, ops = window + ops[:1], ops[1:] + out.extend(window + ops) + return out + + # Register fusion with operation fuser + te.ops.register_forward_fusion(fuse_axpy_ops) diff --git a/docs/examples/op_fuser/residual_layernorm_mlp.png b/docs/examples/op_fuser/residual_layernorm_mlp.png new file mode 100644 index 0000000000000000000000000000000000000000..fa95114a69b9cb7ca9eeabdf92be4800be78550c GIT binary patch literal 15620 zcmeI3c{r5+|L+kYp%ThkDTPqV5)Dd`gpWPDNOnWUHiOAt3YAdyeP^!_x*am-mlm5@tSx4Xs9rs;yOh| zMa8W8MM11Dq*0_Kt)G= zj_Sm57vPtQnv06=cbkgpDK+>1Y->~B_}d3KW~dz%?cY8|!1nm(9PkH>{<}SqLj7My zOriNtZem4kTqoTThQT5RST`%et9KDa; zY(vY&*mBU-j~}UC1Swux4u7Be=7a&V+}R0Pm^e6SpI>KBmoSU$$Bm!{jZlMyr|HFb zCM;S{!-6h8p&71RDBQc-=Y~MfOxABe8$MLh_1F4QeW?&|$xE-W9BmY3?KE=>#Cm5F ztDSth*cpp+&mslZjX*r>jocp;UncPcELM89d+l6|P;%q$uO+M!O}?p;rw%VLGR3_Y z6E?=5>CV?`}q)4=p-r1kP{609ua`iW5IEU)c`fFBX(R@7#mysq+P$w>2nV+#3aEs(1zMNRkaTmZvB&bGgE=?_s zY${hRHN|cVaf2RTU<8F?m$cq+n(6rex_9p@aqRspXs>4h>E>fr=3LW#M>Fh8BvZra z?x#~*dR{+h9cyH3!pT8sS9o9zEx1Xd`0&dBy!FDk*cLrcsgO zr&^Z(6P1WwR-CY9v@ZAm9i=GBabEiG9Q^5=%&U&^fpgWPW~ke~gU!S$f@m>*-AFG4 zD=a>5mS1z1@5UJ+f2;jv^$+Cw6a6*qcjg@D3UDr!%d7J&({6Ozx% zg@LY0a;+68XvRH<9}ar;_R}ALV!W(bmCc{pBa;T&f>mQS(Mfn9X6ZvNuC2Cjj1#c_w*OB(LRpI9r*5%P7~L29%n z^Y3gl`|{UTOoa)Ggc#P3udpT}F{s7s^FhN(ng)?P;|~n+6)$I)1+2inOktq%zYkqf<<@bwZVw-Tb^+j z;>K7iYvkggDe@nqO1GvigJko=(!7xqqBp&V4(Z)EK#om<1|c4+P8o)-yiso6Mh`vo z)Q8dX3J#dn;qdb;JbfxV2OBf;M^Q-w)g{m)PC0X^I|fYDbCieuQrMIb?Dj!Ljy_40 zRJb>Vn+g++6qUz|o@oqN%iFs=@H)@rzTnoU_nS&Z<9x#9wi7NVAyEaFKaeXHGJXz~ z2^vd@e!PI)STIvCHNFd;EyC(b3yFb*%G2!xE`RZ~e89jSnA>OaP2hI9l_SK`ZkA{v z#pUx&3A=x21b1K8d#0w$f)E!rva=em`bM7@mbG1bICF?PqC~L;4krfSz%lwO_Z?L! zVi{K<%szI;i&K)-_?!YEW)nN`z1LnphFfgSKWI2k^Pz}alp-5&z)yT7DET8G~DqVvhs6r32D|tThM()@g5$)OHn$5hc^y)R7 zZV5GrE6!Pe7qySe^LQ%ax&4HfGB+~XJZndZw(6Pb3>0zHXZOHvfG$c?Uec6@7Aj>9W^n=<%4zSc{`hA8fP#M{ zk4f%u1_*q+R73EnkDGT$P2n_uS$aC__PQJZ8?4?zk0iVO3|RQN41N2 zQw-obTYZZRKDRI3`fg>}-4v|pTe9;4^Yp_C#0SqU`AB+&{5(_{ExseyBjB`LT#Fdj zFUbKnXQY+vgn|}UIh@45ZibBUhyQpS90?vHrP8Rh1o54?qp4suk&M*-jMENDmRzVS zU^yY|KorOCZ8C}Y6|EQQsNlC%g}+-66u9_vE_>$E35R&Ey2FM*?im_MTt)33_fakc zC9WLKNH1jrXQuLlWalqWHWkbdo(H!;&}E+4t}vDY_YjD-WAlbL8fRWfSY?fL3BZ+_ zSE!dvc>Di?h6sGrnU?^* zvd7%K2RdR^AsxIk(;m6dH(Oxkh8Wa?&ClEy6~}FiwZ!x5PJRx1*hz18>p&~`916yH z8vbmaYxZ<)VUpZ-^2TkUi>j449hHoqo_jUx8+z4qIiG4is4}7CdDwX25Ja-BE+=AW z=c)7l?x$hSitw`4fUc2-XTAh=4KNt0*`}5}RysY@F+*K-D3jJ=4R!F`N01uRVU%a2 z%2FL=o_I2Mu|q8lo_B5LfY^CHlesfusa8kJN|)$)7oA5Etv#|9oF_OpegeV2#L8Pg z?22ZCnP+m>+UuROsBA6h*>hO1*mft#aF5HjiiOm%bEUinuII?jMy?h2g`-@K@ zp@(b$Ogwwzy~ON11EULA!YS0Ve7q;?XY6gB*sZO#$fW2|yRmo&FgEUmDY_6|)u>-D z>pAqp$n7B>&lgK3w8foJAs(c@68EXI%`|Q@G~J`jyhTNJR~IY!WQiT#?MpaG@L5Ze z(w;UEkreQ?KkRT8QtF>7?dIFGK!2EKuY=1FcNl9Cy6zN@1N$*c-nf0Gd&o=vnyPDdtL8T<^e53T8Sc+JLhg4~cxW!!B5rbs98~q9Es9QFJX#qIadO0O4Mdm*x%QQWb`vvWX#7C9ROu+y6b-txHS=}h`L#D@ zypbJGR^47kwm0uonZV9+MnPJC&CC6CAKX_{uFa_+P792$CJbnbC*GM!-3pOU3uf6c zoM(k?Po34g-eP@&L0P3G+F*^4B)e2Cki3>KU=aUn5mNTNi^nTgqI0z*fSEi;h%?LH z^f%1r#gkCb{QkR(uD(TDrdqKEYb#%Sw#ul{njci<;6vz8PY_s_*tGOj3bOx3VTpGB z3u6LTuUg?I<|<+||3ij+8)9F}Uon?}#9(oU9^&fj%VTKa?CDRr8^ej`U<8aNN?u~R z@L6ID-pI4s?Yanht2L|(j5wm;_#=GRqu5sWCnjiB6A{SAU{Ty*LbWyG9{SDcTDNthK!;Op=JK`6WZ^7`8Ty3A&OL`Liqqtbr}ROuQz1KNz&vwhlpj32 z50g~l&UibwW$w@@Dm~w$+uf-DeoyCCl~4IfF|a7ucfRS+*$xWjIA^K-Gv~zlO=W=$ z%M0`4;-r)uABuv4oT-7yJTvTh@pFUoabj&t9}zBdKg|=B?z&~&>97fotW+|}*2iE? z9kZkJTf~n#gD{+@Vax5xdyy@4S7cp+UJ%HWU$Kb$bEQ_$!)--OQbxsT$xOzXFF|4_ zBs6(PW{u64y&kN=p>BPdlhP*-w2snPU#(NRIItYGY!! zTSRsoO^2G}eJzly%!^@$sDv^zO^lFi4Qf(6y%gOqR2&-+Tco?6!OuPY9`kcCG&GCk zt|BN|mwztF*&i%)2!Sn1s1tKmfe2!|xg0s8V@k2@5+bFf;uDp%nVw$oSI*1#G~|6$ zIEI8{2o5&35EfXUjVw;KuYCPDkdso7>wyh;kCn>d5Z>%7d*`3w&G^nSIOLKy)#8!$!*BKuOn-N2^uO+ zB8CO9w^iKFb&Qc2>lY1F@3a{WGRxb97oD88l9fR-IDZS_@L8GOdc?d^cYRO=GVxU& z8N1k`#{*Hvk_h)_!vl=Lu5RL>qwT++NTCkB0Iz~s{60IqhYctJn(5k_myVW?= zc-{28j$8428nyt=X()k zofUa<;Z>zX+SaA#g=l4+v98*4O7&)S2~K+l*13C~ z>x<-4^Ml*_6qSYu#gM$p$9=q9BqpZ>j+-LZgpVl>>BeV)UO19CgJdg!96_r6wg z>Pno%qh0s5b*o2=Kb3XCG;&t5Q#&)Cl1{o9xpaMivi5g2Hkj6>dXvBG7FEd)!m?N0 z)(~TbT19L@Sz>^b8Gq!-f1^harzV1ykESO+YR4e0HN|RYY*hGs=rP8*0qq;U;H5LM znW1ER`+WCS`CEi9vU*nzS}eM-rgEnT&v&(%mTvzlbaq>C9;iz56x^xQ@85)0c?@@P zmdE-Xiv{`3;xI_>J(>jh$|sS-hldpZHZ!b-Nq+n@z$DUD*vw2e)YPBCZDvZLGSeLo zJPCa+gG)!Qo8R~>(Js4eVdtw{^W~VJG^CQ-Zrk$v4Rhq3Uo+XGe7%XBOaAo%RcD z&po#xc^J(2?OH!UYLHv@yp*o;*{kAKVa^IpBPd0MrPBoc0ji=rA%)A}ILicQu)D@8 zX`nAqQz*BF&%-sWPy~t6=*nwj9nn>pbkI8l==sns#G)Fx9y%{X%;Jn%Z2Y=5*$z6Fmdv(8s1{!u3wZYx| zFAwVq%yhD>)wyLAlz)*fbz`AVnikim)x`Q5JGUfuOxHHzbbIi5c^h!=T3CBcAlM&z zQ?Z+WA)r|_H@F_L7r#0f3H_?u@K|)A!JSg!Gp}}Me5rsRsjZN=fx_Z*rE9p38o0Hg z14lp9S#PY^8Jd&)d_K$>>?{n=`7-f8muVdyOezT&9Ut`1-%!C+SR0Yvm0t437~6C# zfV{RXb+8l;AG8Y=bl>oOnOj9&I5$17UlLu+*2Qm2G=TWPw8v} zWWndvxDbx*tenQKn`K$sI@8sQZr{7eCQ5n%<_cRG2-ZiplKC0ytfL3mJe>CgyRi>8 ztAG?aKz==~`@3xs#2{r0XSZFK;hm1SYo`)?SjK$SCl;fECQm;$ zjf~M-aebX+RD}1h@5GoBIlz>|(6Z+2VZr%DQ)zpGfNtokz`8?ImfYE@Xc}+sd+!Ge zImWuBx@uQrc7ByQ`&4jVH$uFdq?l6}ztvl#!b($GDz_Jj)D5lY?N!N2%^ z8E12cQf|XW0vbR6Fj3kJMYFWCSCl-qk>T>b^n-U>JuAF^T)J-g(Dd>3X9r%(_dx|^ zjxAg6g5Lt)adm~+i3!XT&Yw8i} z0H2MRSPC)A`&4Rc8$KdKssiQ0KSTOM+Gxe9^%CenBL(NZk0y;m0kZ^2iHgWrvn0p zl!uE(M?BwUn65q&rM_v8zVhILS-uLjnU^{pWHtO_RG?jV?ep;{|7k#lH!CAPN@lav z$)3Yay(>D2Zyj3r2`F|G?V}=$2#_)Iw0KPx%_q0lw^k}S+hGze+v>>!F-46?@#~MC z1%(ltX0kMGt!kq7O$b|$fBpl=(}w5FyuxWAY3~!F0=2q9Mlx5r8ZQqaGgTv=jK%_S z3n1_FH9mOPAw0slqBXP4Y?j|bW9Aje`L--dX&pqeRYja+QGflY%Z+uC2euMSyq3Tb zSL^Hud5HRwszb&NE{mdk0!Q9&XPB&LwS%itMH@#yJ}WHp7E4b{`amN&AZfSba_pcj8Kt|EylPmMt&yCkQt;fv>DIx;{AHa8x9|KuRzqP)R*L$ zYc4^zbQ6S&FJm(-<<9oZ%iaGX1^K52W9(IU$k)IGPI-gUVDSfmx!cWrZOfb%N;&`9 z7~8S5f3Ig*@Q;NS3e}574-SCAqeN3lbgF5RJQ7=$TgwHF-uXSijMF0I2SMfI;UANo zpr)n?rDr*3`sg=Ysz0tnK=41Ir|5~Bbd&})cDaZD*QPY-%qS&qRMu$aM`QbIyJ@%U;GpKYAfK+?%MshcJF}3e|lZ4;PLN? zGVz_}-XdcgkEzbo_L8f&UALL6qGmK-`Bjk9c<-f~~=rw@T3E7ng z{^fKW0H>P{>@NCyvD(Lw@AE4;1}~kIL315_%+S%pD}5enwL4_=!I9%*Bm5GHetciB zv*P6VtDmRQLVR4z_>UZCY~yJ#NcQ4hMsAYdt?CJs1F2)HO(W^6G)PL;RXf=s@leGP z1)y2m31l0$(X>2V=_nu9Gm_3h!_2f>6o>D5-=9s4lR_FOAxi3I>4fvt0OrwU>Zn(3gh z1YiGSZ;llPe+_4oX{61NrEyYM{7r7-hedvKTMzo*%TIpI3n`5W2DXn&3}r&pzIDf# zX}mT{`^du9e2}v?TN8*qa%Zvi4j3K)EKW)TeShC2Iht8a69zf!WCLeec_(kpc3Jav zy@R`L1Mpc6ni+W}x%+5noK!9e74?q=EdJ>wc(2U-({~*JO`lpfnq@}2(x~i)l>Mg~ zOr7D3OdQpe_1}6h2J<)dXX`GD{xO6%ifq916nx|L$9AY`8v!gE1b)Kw$DLm42?9b9 z#!F_(zYUPee;qK~-F3sE6Ki*dHk`J&h9G_Y-gQ-L0M#Q#Qp91X}Tfo0Z zl=}q@E6}E;3F522bi4pxdX|jrHJdIw4gP3#X)~39=KvnOks+N{~JB;IOC+a@<(Q>)> z^(z;{qTcRpWQ5T9FP2I#CI^SuZK7!tIyMHoLk_SbA@B=dcEZ(^>n67LgdJE4a9cz` zdUG5hY5izI`^|=>b8^Ijy?3|npv`tLeIJC#$3~VQDeZ^6V*kH!)$g8UrFp13Y`^s$ zQFg7|XSFA^1#(p6ce8M;(sDk(Tl~mSZ^fS@1|r8+(Ie$`f^;J-CUDZJ)zvmgY2wyY z+AZL--D=~h%dyQ2mmVC~*Vnt7xwZp9ziA48Bx{I8=c+W(VD-LZ>mN*7A6p)C)2J5a zY^`oL@f6nHwK$|r7WF8KBsS?{_2rS%!HOtC+o@ps| z&DT?x7;MLg+nMeK!DUlP(Lx)hz)Y<{u6l}i*2X2uZf+ka0D4$?`u2OQzPU7jn8?=8 z3~MNM>|Ia!yhh&0!7?I>x;TKpm7WCk+%66+|ICF~ z$Jih$wS+Pp+IsGsJh&(4{?k7Xli;22;o-vZr$jhS5Ir?cxQy-j0y8x3rl?Oy!5uwajq)7cp}q zCNedk{)n#w^Hp&1{qoUofXd3;CQEsDjZ@CS&$}ZVRfnTgy(65n?6-Hw*I$P2^ft{` z1xr>zz^;HEI40@N;hdx`K`H!T1lLR438fUb4zd`F>_$&2(xKVYKhv@Zh@BGXQKtPM z(IS2FQ1$X9(aYrC$MbeaI&4*oX)diq8%UPC=n?4@PZZe%i`e?&O?buf>pn6qOShk0 zpgIwOs;~IUqQFCjhVB+T4%S@o6%$QW@vZY^kKg&M9c;Fgkbp&RF(x}w4755z7;E8` zMGLs-Y8zkhfXK#e-k24yI)}E8QTv_Gkq-J4maEMky14BOd49(#Fi=LYgX(K^Fq$(- zJ|BA+B{077O)9IV%Ih!-_IXOeMAbCiCY!}z;0ZtrW8L;bt5UU}u{0UL*oxKb>BZnV&!nXwi^oA*q4H@iS=@SE~63n}Rh}87H!O+I!?o?~g&_K{o#J$9bu9dr!9sZt`H|JR>*+%-7`ilmyQ)=PWlb0&8wg7k5KwM%;cmumA|uZksN0mK6QsUPez#F;XLwXqpHXuJjY&&g zt*bpGRb{?@6xj}$>+DX*WBPgLFzyKQ+0J+?;jd`mbO4^a?-FVqbI z=t7}#=A5Zbb?|6w!1%88nsG|?>uI;5p4!i1yeAdDPYk^U0&~TvG zqNVAo$J@0jpYiG%DE=}i+4(KQBRAHn&QxznQ3nYn{@y;E5{@Hx2P@cOEn_*W2ScX3 z%lx`fLC?8V)E28F(3QV$X6p9Q)I9F2mNNE^CRtc2^D3m5OP#z3vv8nqAKp7f-KL?! zQr>nt6l83qeCajahrSTO3y^dugxEk&7+moM`>e1-&)-J%@2XZ}?Zf#PR+3%Q_)>?w z7T8Mt;l=3|-SJSET-_Qt%2KYW>8j1Od4;-}hkCu?D~$IHWMv7&I>JpJBCH_^C7e$p zTs#YsOjF?~`AmPma$7FWB)6Y>o~7%Bv2Ap;s9Q|rdsVJc3M{Sw@x9>C+IC*1uY_ksVC$FZ}{s^Vfn_4%zQJ7aJXhdfa{8s(Bb69=VbN^G>Vwm>Ri(hVOM5Y~9g37;S8wcJs@<9;Yw_S<6f} zYis|S){>kEseaYX&B+3Ps42rZ?9HN@#Fgxyhvzz|2i)J{?YQXbSBV2_bMb;J+8hJd zzwN9?^^bt~c`rV%yBMqV#=f}_3l6K-!Ust*l)3s(^oO|VgX-3p&{0v-s2vju^4J#ZKMq&`9xnyQf3zwr}5%<~dla_BNW1 zl$AkvY}IrGiqrd4CKpxhtTK7HGZ1ya?xbd+m<~w#==KRQiyj}G$ z%h&+tD?C-*J4of8LP`DarkheHr|P6iMuVH!joFbn$;&rsTOIDQC@m=O>irGwKZ#9L z+t>9s3HFDCAt&T4+bijkO5%O<#lTT^C&tgLPcpyW=6Vq9a;Xu~sFeOY9~eKwAn4>} zrFo?!wjm?@?Q(^*EFJh`i_quavHhdpv3*|%W}36r#yekbTq67jcZOS#(pi~+`w{1r z>VME=Fz3^YH+FC8lWOU>LZzR~{0hC|LhF4)DCA8j-G{^R5WtZ}SB+c3!RgUNM6@&$ z_#^#)We$Im9mNG8)96{r`6HB4(L4Z@(tx0}3x5)Fh9E#IV!%bj{mpI8WKj7FJh>5e z`TynMm4_(>_Fy1Ab#gS&?1(R-gasuA99P@|Ippqf4k?{MM>`Us_|S1_pvHKl>yT$8 z>W^OfWJm}|N&%^FGXxJLr9MhViYY4$kI0s<`>&T}sCa@8o0wFXLQPtsF`JHOvW|p6U4X!xez|89GdhISX6?<9P!|5n*Em zht8ej6n<)aQ6p?l>R9bmpdkY_%Yifqz44FS*Ik01sXV3~sd)1$+r3r>jOze`_u~rd z5>yx$bX89?O@ L@u=vbMZo_6*pP2# literal 0 HcmV?d00001 diff --git a/docs/index.rst b/docs/index.rst index 336cd2d47f..194e76df24 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -57,6 +57,7 @@ Transformer Engine documentation examples/te_gemma/tutorial_generation_gemma_with_te.ipynb examples/onnx/onnx_export.ipynb examples/te_jax_integration.ipynb + examples/op_fuser/op_fuser.rst .. toctree:: :hidden: diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 5d1a5ce61d..f95f065d78 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3428,8 +3428,15 @@ def test_custom_basic_op( ) -> None: """Custom basic op""" - class CustomScaleOp(te.ops.BasicOperation): - """Custom op that applies a learnable scale""" + class LearnableScale(te.ops.BasicOperation): + """Custom op that applies a learnable scale + + This class is as an example in the op fuser guide at + docs/examples/op_fuser/op_fuser.rst (see "Implementing a + basic operation"). Any changes made to this class should + also be made there. + + """ def __init__(self) -> None: super().__init__() @@ -3442,23 +3449,19 @@ def op_forward( self, ctx: OperationContext, input_: torch.Tensor, - prev_op_grad_output_quantizer: Optional[Quantizer], - next_op_input_quantizer: Optional[Quantizer], + **unused, ) -> torch.Tensor: + out = self.scale * input_ ctx.save_for_backward(self.scale, input_) - return self.scale * input_ + return out def op_backward( self, ctx: OperationContext, grad_output: torch.Tensor, - ) -> torch.Tensor: - ( - scale, - input_, - ) = ctx.saved_tensors - grad_scale = torch.inner(input_.reshape(-1), grad_output.reshape(-1)) - grad_scale = grad_scale.reshape(()) + ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]: + scale, input_ = ctx.saved_tensors + grad_scale = torch.inner(input_.reshape(-1), grad_output.reshape(-1)).reshape(()) grad_input = scale * grad_output return grad_input, (grad_scale,) @@ -3485,7 +3488,7 @@ def op_backward( y_ref.backward(dy_ref) # Implementation with fusible operation - op = CustomScaleOp() + op = LearnableScale() forward = te.ops.Sequential(te.ops.Identity(), op, te.ops.Identity()) with torch.no_grad(): op.scale.copy_(w_test) @@ -3502,7 +3505,112 @@ def op_backward( torch.testing.assert_close(dx_test, x_ref.grad, **tols) torch.testing.assert_close(dw_test, w_ref.grad, **tols) - def test_custom_forward_fused_op( + def test_custom_forward_fused_op1( + self, + *, + shape: Iterable[int] = (5, 11), + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + ): + """Custom fused op in forward pass""" + + class ForwardAxpy(te.ops.FusedOperation): + """Custom op that computes BLAS SAXPY in forward pass + + This class is as an example in the op fuser guide at + docs/examples/op_fuser/op_fuser.rst (see "Implementing a + fused operation"). Any changes made to this class should + also be made there. + + """ + + _enabled = True + + def __init__( + self, + scale: te.ops.ConstantScale, + add: te.ops.AddExtraInput, + ) -> None: + super().__init__((scale, add)) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + **unused, + ) -> tuple[torch.Tensor, list[tuple[torch.Tensor, ...]]]: + scale_op, add_op = self.basic_ops + extra_input = basic_op_extra_inputs[1][0] # Extra input to add op + out = scale_op.scale * input_ + extra_input + scale_ctx, add_ctx = basic_op_ctxs # No state needed for backward + return ( + out, # Output + [(), ()], # Extra outputs for each basic op + ) + + def fuse_axpy_ops( + ops: list[te.ops.FusibleOperation], + **unused, + ) -> list[te.ops.FusibleOperation]: + """Apply fusion the first time this function is called""" + if ForwardAxpy._enabled: + ForwardAxpy._enabled = False + else: + return ops + out = [] + window, ops = ops[:2], ops[2:] + while len(window) == 2: + if isinstance(window[0], te.ops.ConstantScale) and isinstance( + window[1], te.ops.AddExtraInput + ): + window = [ForwardAxpy(*window)] + else: + out.append(window[0]) + window = window[1:] + window, ops = window + ops[:1], ops[1:] + out.extend(window + ops) + return out + + # Random data + scale = 0.5 + x1_ref, x1_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + ) + x2_ref, x2_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + y_ref = scale * x1_ref + x2_ref + y_ref.backward(dy_ref) + + # Implementation with fusible operation + te.ops.register_forward_fusion(fuse_axpy_ops) + model = te.ops.Sequential( + te.ops.ConstantScale(scale=scale), + te.ops.AddExtraInput(), + ) + y_test = model(x1_test, x2_test) + y_test.backward(dy_test) + + # Check values + tols = dtype_tols(dtype) + assert_close(y_test, y_ref, **tols) + assert_close_grads(x1_test, x1_ref, **tols) + assert_close_grads(x2_test, x2_ref, **tols) + + def test_custom_forward_fused_op2( self, *, shape: Iterable[int] = (7, 11), diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 9e23bb3fb1..13cb519c19 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -152,7 +152,7 @@ class GELU(_ActivationOperation): \text{GELU}(x) \approx \frac{x}{2} \left( 1 + \tanh\left( 0.797x+0.036 x^3 \right) \right) - See `Gaussian Error Linear Units (GELUs)`__. + See `Gaussian Error Linear Units (GELUs) `__. """ @@ -183,8 +183,8 @@ class GLU(_ActivationOperation): the first half of the input tensor, while PyTorch applies it to the second half. - See `Language Modeling with Gated Convolutional Networks`__ - and `GLU Variants Improve Transformer`__. + See `Language Modeling with Gated Convolutional Networks `__ + and `GLU Variants Improve Transformer `__. """ @@ -219,7 +219,7 @@ class GEGLU(_ActivationOperation): the first half of the input tensor, while PyTorch applies it to the second half. - See `GLU Variants Improve Transformer`__. + See `GLU Variants Improve Transformer `__. """ @@ -233,8 +233,8 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: class QGELU(_ActivationOperation): r"""Quick Gaussian Error Linear Unit - Quick GELU from `HuggingFace`__ - and `paper`__. + Quick GELU from `HuggingFace `__ + and `paper `__. .. math:: @@ -316,7 +316,7 @@ class ReGLU(_ActivationOperation): the first half of the input tensor, while PyTorch applies it to the second half. - See `GLU Variants Improve Transformer`__. + See `GLU Variants Improve Transformer `__. """ @@ -334,7 +334,7 @@ class SReLU(_ActivationOperation): \text{SReLU}(x) = \max(x^2,0) - See `Primer: Searching for Efficient Transformers for Language Modeling`__. + See `Primer: Searching for Efficient Transformers for Language Modeling `__. """ diff --git a/transformer_engine/pytorch/ops/basic/add_extra_input.py b/transformer_engine/pytorch/ops/basic/add_extra_input.py index 47f2b6e248..fc3ca9cade 100644 --- a/transformer_engine/pytorch/ops/basic/add_extra_input.py +++ b/transformer_engine/pytorch/ops/basic/add_extra_input.py @@ -30,7 +30,7 @@ class AddExtraInput(BasicOperation): feature and most users are discouraged from it. In-place operations break some autograd assumptions and they can result in subtle, esoteric bugs. - Compare to `MakeExtraOutput`, which does a similar operation in + Compare to ``MakeExtraOutput``, which does a similar operation in the backward pass. """ diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index e640f3ffb1..48376a297f 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -48,8 +48,8 @@ def _wait_async(handle: Optional[Any]) -> None: class BasicLinear(BasicOperation): """Apply linear transformation: :math:`y = x A^T` - This is a drop-in replacement for `torch.nn.Linear` with - `bias=False`. + This is a drop-in replacement for ``torch.nn.Linear`` with + ``bias=False``. Parameters ---------- @@ -61,27 +61,27 @@ class BasicLinear(BasicOperation): Tensor device dtype : torch.dtype, default = default dtype Tensor datatype - tensor_parallel_mode : {`None`, "column", "row"}, default = `None` + tensor_parallel_mode : {None, "column", "row"}, default = None Mode for tensor parallelism tensor_parallel_group : torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel : bool, default = `False` + sequence_parallel : bool, default = False Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) rng_state_tracker_function : callable - Function that returns `CudaRNGStatesTracker`, which is used + Function that returns ``CudaRNGStatesTracker``, which is used for model-parallel weight initialization - accumulate_into_main_grad : bool, default = `False` + accumulate_into_main_grad : bool, default = False Whether to directly accumulate weight gradients into the - weight's `main_grad` attribute instead of relying on PyTorch - autograd. The weight's `main_grad` must be set externally and - there is no guarantee that `grad` will be set or be - meaningful. This is primarily intented to integrate with + weight's ``main_grad`` attribute instead of relying on PyTorch + autograd. The weight's ``main_grad`` must be set externally + and there is no guarantee that ``grad`` will be set or be + meaningful. This is primarily intended to integrate with Megatron-LM. This argument along with weight tensor having - attribute 'overwrite_main_grad' set to True will overwrite - `main_grad` instead of accumulating. + attribute ``overwrite_main_grad`` set to ``True`` will + overwrite ``main_grad`` instead of accumulating. userbuffers_options, dict, optional Options for overlapping tensor-parallel communication with compute using Userbuffers. This feature is highly @@ -184,7 +184,7 @@ def _canonicalize_tensor_parallelism( Parameters ---------- - mode: {`None`, "column", "row"} + mode: {None, "column", "row"} Mode for tensor parallelism process_group: torch.distributed.ProcessGroup Process group for tensor parallelism @@ -200,7 +200,7 @@ def _canonicalize_tensor_parallelism( Returns ------- - mode: {`None`, "column", "row"} + mode: {None, "column", "row"} Mode for tensor parallelism process_group: torch.distributed.ProcessGroup Process group for tensor parallelism @@ -446,18 +446,18 @@ def _functional_forward( Output tensor beta: float, optional Scaling factor applied to original value of out when accumulating into it - accumulate_into_out: bool, default = `False` + accumulate_into_out: bool, default = False Add result to output tensor instead of overwriting - tensor_parallel_mode: {`None`, "column", "row"}, default = `None` + tensor_parallel_mode: {None, "column", "row"}, default = None Mode for tensor parallelism tensor_parallel_group: torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel: bool, default = `False` + sequence_parallel: bool, default = False Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) - with_quantized_compute: bool, default = `False` + with_quantized_compute: bool, default = False Whether to perform compute with quantized data. input_quantizer: Quantizer, optional Builder class for quantized input tensor. @@ -465,10 +465,10 @@ def _functional_forward( Builder class for quantized weight tensor. output_quantizer: Quantizer, optional Builder class for quantized output tensor. - input_requires_grad: bool, default = `True` + input_requires_grad: bool, default = True Whether the loss gradient w.r.t. the input tensor is required in the backward pass. - weight_requires_grad: bool, default = `True` + weight_requires_grad: bool, default = True Whether the loss gradient w.r.t. the weight tensor is required in the backward pass. @@ -477,11 +477,11 @@ def _functional_forward( torch.Tensor Output tensor torch.Tensor, optional - Input tensor, ready for use in backward pass. `None` is + Input tensor, ready for use in backward pass. ``None`` is returned if loss gradient w.r.t. the weight tensor is not required. torch.Tensor, optional - Weight tensor, ready for use in backward pass. `None` is + Weight tensor, ready for use in backward pass. ``None`` is returned if loss gradient w.r.t. the input tensor is not required. @@ -682,24 +682,24 @@ def _functional_backward( Loss gradient w.r.t. weight tensor grad_weight_beta: float, optional Scaling factor applied to original value of grad_weight when accumulating into it - accumulate_into_grad_weight: bool, default = `False` + accumulate_into_grad_weight: bool, default = False Add result to weight grad instead of overwriting grad_input: torch.Tensor, optional Loss gradient w.r.t. input tensor grad_input_beta: float, optional Scaling factor applied to original value of grad_input when accumulating into it - accumulate_into_grad_input: bool, default = `False` + accumulate_into_grad_input: bool, default = False Add result to input grad instead of overwriting - tensor_parallel_mode: {`None`, "column", "row"}, default = `None` + tensor_parallel_mode: {None, "column", "row"}, default = None Mode for tensor parallelism tensor_parallel_group: torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel: bool, default = `False` + sequence_parallel: bool, default = False Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) - with_quantized_compute: bool, default = `False` + with_quantized_compute: bool, default = False Whether to perform compute with quantized data. input_quantizer: Quantizer, optional Builder class for quantized input tensor. diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index 8b60251088..d580f84866 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -18,7 +18,7 @@ class Bias(BasicOperation): """Apply additive bias - This is equivalent to the additive bias in `torch.nn.Linear`. + This is equivalent to the additive bias in ``torch.nn.Linear``. Parameters ---------- @@ -28,7 +28,7 @@ class Bias(BasicOperation): Tensor device dtype : torch.dtype, default = default dtype Tensor datatype - tensor_parallel : bool, default = `False` + tensor_parallel : bool, default = False Whether to distribute input tensor and bias tensors along inner dimension tensor_parallel_group : torch.distributed.ProcessGroup, default = world group diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index eb8a67600d..b44e77b0c6 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -65,7 +65,7 @@ class GroupedLinear(BasicOperation): weight's ``main_grad`` attribute instead of relying on PyTorch autograd. The weight's ``main_grad`` must be set externally and there is no guarantee that `grad` will be set or be - meaningful. This is primarily intented to integrate with + meaningful. This is primarily intended to integrate with Megatron-LM. This argument along with weight tensor having attribute ``overwrite_main_grad`` set to True will overwrite ``main_grad`` instead of accumulating. diff --git a/transformer_engine/pytorch/ops/basic/layer_norm.py b/transformer_engine/pytorch/ops/basic/layer_norm.py index 631f0fafc9..3fda5145c6 100644 --- a/transformer_engine/pytorch/ops/basic/layer_norm.py +++ b/transformer_engine/pytorch/ops/basic/layer_norm.py @@ -31,7 +31,7 @@ class LayerNorm(BasicOperation): r"""Layer Normalization Applies Layer Normalization over a mini-batch of inputs as described in - the paper `Layer Normalization `__ + the paper `Layer Normalization `__ . .. math:: y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma + \beta @@ -51,9 +51,9 @@ class LayerNorm(BasicOperation): Tensor device dtype : torch.dtype, default = default dtype Tensor datatype - zero_centered_gamma : bool, default = 'False' - If `True`, the :math:`\gamma` parameter is initialized to zero - and the calculation changes to + zero_centered_gamma : bool, default = False + If ``True``, the :math:`\gamma` parameter is initialized to + zero and the calculation changes to .. math:: y = \frac{x - \mathrm{E}[x]}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * (1 + \gamma) + \beta diff --git a/transformer_engine/pytorch/ops/basic/make_extra_output.py b/transformer_engine/pytorch/ops/basic/make_extra_output.py index 61caaaf65d..0d9c870262 100644 --- a/transformer_engine/pytorch/ops/basic/make_extra_output.py +++ b/transformer_engine/pytorch/ops/basic/make_extra_output.py @@ -35,7 +35,7 @@ class MakeExtraOutput(BasicOperation): operations break some autograd assumptions and they can result in subtle, esoteric bugs. - Compare to `AddExtraInput`, which does a similar operation in the + Compare to ``AddExtraInput``, which does a similar operation in the backward pass. """ diff --git a/transformer_engine/pytorch/ops/basic/quantize.py b/transformer_engine/pytorch/ops/basic/quantize.py index d126b554b5..fa3efc3807 100644 --- a/transformer_engine/pytorch/ops/basic/quantize.py +++ b/transformer_engine/pytorch/ops/basic/quantize.py @@ -18,14 +18,14 @@ class Quantize(BasicOperation): """Quantize tensor data - Uses recipe from `autocast` context. When called outside - of an `autocast` context, this is an identity operation. + Uses recipe from ``autocast`` context. When called outside + of an ``autocast`` context, this is an identity operation. Parameters ---------- - forward : bool, default = `True` + forward : bool, default = True Perform quantization in forward pass - backward : bool, default = `False` + backward : bool, default = False Perform quantization in backward pass """ diff --git a/transformer_engine/pytorch/ops/basic/reshape.py b/transformer_engine/pytorch/ops/basic/reshape.py index f8ae86fecd..4a171c294b 100644 --- a/transformer_engine/pytorch/ops/basic/reshape.py +++ b/transformer_engine/pytorch/ops/basic/reshape.py @@ -20,7 +20,7 @@ class Reshape(BasicOperation): """Reshape tensor - See `torch.reshape`. + See ``torch.reshape``. Parameters ---------- diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 3179d0a447..1d8d8be971 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -32,7 +32,7 @@ class RMSNorm(BasicOperation): Applies Root Mean Square Layer Normalization over a mini-batch of inputs as described in the paper - `Root Mean Square Layer Normalization `__ + `Root Mean Square Layer Normalization `__ . .. math:: y = \frac{x}{\sqrt{\mathrm{Var}[x] + \varepsilon}} * \gamma @@ -50,8 +50,8 @@ class RMSNorm(BasicOperation): Tensor device dtype : torch.dtype, default = default dtype Tensor datatype - zero_centered_gamma : bool, default = 'False' - If `True`, the :math:`\gamma` parameter is initialized to zero + zero_centered_gamma : bool, default = False + If ``True``, the :math:`\gamma` parameter is initialized to zero and the calculation changes to .. math:: diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index eaffbeee02..b4427df41a 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -46,7 +46,7 @@ class SwiGLU(BasicOperation): The Sigmoid Linear Unit (SiLU) gating function is also known as the swish function. See - ``GLU Variants Improve Transformer``__. + `GLU Variants Improve Transformer `__. Parameters ---------- @@ -189,14 +189,18 @@ def op_backward( class ClampedSwiGLU(BasicOperation): r"""GPT-OSS - Implementation based on ``GPT-OSS``__. + Implementation based on `GPT-OSS `__. 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. - .. warning:: The input tensor is chunked along the last dimension to get gates/pre-activations which is different - from GPT OSS implementation where the gates/pre-activations are assumed to be interleaved in the input tensor. + .. warning:: + + The input tensor is chunked along the last dimension to get + gates/pre-activations which is different from GPT OSS + implementation where the gates/pre-activations are assumed to + be interleaved in the input tensor. Parameters ---------- diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index 90ade030c8..fbaf69d75d 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -125,18 +125,18 @@ def _functional_backward( Tensor datatype grad_weight: torch.Tensor, optional Loss gradient w.r.t. weight tensor - accumulate_into_grad_weight: bool, default = `False` + accumulate_into_grad_weight: bool, default = False Add result to weight grad instead of overwriting - tensor_parallel_mode: {`None`, "column", "row"}, default = `None` + tensor_parallel_mode: {None, "column", "row"}, default = None Mode for tensor parallelism tensor_parallel_group: torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel: bool, default = `False` + sequence_parallel: bool, default = False Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) - with_quantized_compute: bool, default = `False` + with_quantized_compute: bool, default = False Whether to perform compute with quantized data. input_quantizer: Quantizer, optional Builder class for quantized input tensor. diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 6ef9bf083b..0d3e1d0416 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -115,16 +115,16 @@ def _functional_forward( Tensor device dtype: torch.dtype Tensor datatype - tensor_parallel_mode: {`None`, "column", "row"}, default = `None` + tensor_parallel_mode: {None, "column", "row"}, default = None Mode for tensor parallelism tensor_parallel_group: torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel: bool, default = `False` + sequence_parallel: bool, default = False Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing along inner dimension (embedding dim) - with_quantized_compute: bool, default = `False` + with_quantized_compute: bool, default = False Whether to perform compute with quantized data. input_quantizer: Quantizer, optional Builder class for quantized input tensor. @@ -132,10 +132,10 @@ def _functional_forward( Builder class for quantized weight tensor. output_quantizer: Quantizer, optional Builder class for quantized output tensor. - input_requires_grad: bool, default = `True` + input_requires_grad: bool, default = True Whether the loss gradient w.r.t. the input tensor is required in the backward pass. - weight_requires_grad: bool, default = `True` + weight_requires_grad: bool, default = True Whether the loss gradient w.r.t. the weight tensor is required in the backward pass. ub_comm_name: str diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 7fe6ea37ed..bd3bc94b60 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -31,7 +31,7 @@ def _split_tuple(t: tuple, idx: int) -> tuple[tuple, tuple]: def _is_graph_capturing() -> bool: - """Whether function is called within `make_graphed_callables` + """Whether function is called within ``make_graphed_callables`` Avoid circular import with lazy import. @@ -519,6 +519,8 @@ def register_forward_fusion( The fusion function should have the following signature: + .. code-block:: python + func(ops, *, recipe) -> updated ops Parameters @@ -545,6 +547,8 @@ def register_backward_fusion( The fusion function should have the following signature: + .. code-block:: python + func(ops, *, recipe) -> updated ops Parameters diff --git a/transformer_engine/pytorch/ops/linear.py b/transformer_engine/pytorch/ops/linear.py index d5829b0c50..c6ca4786b8 100644 --- a/transformer_engine/pytorch/ops/linear.py +++ b/transformer_engine/pytorch/ops/linear.py @@ -23,7 +23,7 @@ class Linear(FusedOperation): """Apply linear transformation: :math:`y = x A^T + b` - This is a drop-in replacement for `torch.nn.Linear`. + This is a drop-in replacement for ``torch.nn.Linear``. Parameters ---------- @@ -31,17 +31,17 @@ class Linear(FusedOperation): Inner dimension of input tensor out_features : int Inner dimension of output tensor - bias : bool, default = `True` + bias : bool, default = True Apply additive bias device : torch.device, default = default CUDA device Tensor device dtype : torch.dtype, default = default dtype Tensor datatype - tensor_parallel_mode : {`None`, "column", "row"}, default = `None` + tensor_parallel_mode : {None, "column", "row"}, default = None Mode for tensor parallelism tensor_parallel_group : torch.distributed.ProcessGroup, default = world group Process group for tensor parallelism - sequence_parallel : bool, default = `False` + sequence_parallel : bool, default = False Whether to apply sequence parallelism together with tensor parallelism, i.e. distributing input or output tensors along outer dimension (sequence or batch dim) when not distributing @@ -49,12 +49,12 @@ class Linear(FusedOperation): rng_state_tracker_function : callable Function that returns CudaRNGStatesTracker, which is used for model-parallel weight initialization - accumulate_into_main_grad : bool, default = `False` + accumulate_into_main_grad : bool, default = False Whether to directly accumulate weight gradients into the - weight's `main_grad` attribute instead of relying on PyTorch - autograd. The weight's `main_grad` must be set externally and - there is no guarantee that `grad` will be set or be - meaningful. This is primarily intented to integrate with + weight's ``main_grad`` attribute instead of relying on PyTorch + autograd. The weight's ``main_grad`` must be set externally and + there is no guarantee that ``grad`` will be set or be + meaningful. This is primarily intended to integrate with Megatron-LM. """ diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 47286dfced..54b3f00117 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -94,7 +94,7 @@ def fuser_forward( several of this function's arguments are lists of arguments to forward functions of corresponding basic ops. - Called by `OperationFuser`. + Called by ``OperationFuser``. Parameters ---------- @@ -141,7 +141,7 @@ def fuser_backward( several of this function's arguments are lists of arguments to backward functions of corresponding basic ops. - Called by `OperationFuser`. + Called by ``OperationFuser``. Parameters ---------- diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py index a0db3cd2d0..592ddae23a 100644 --- a/transformer_engine/pytorch/ops/sequential.py +++ b/transformer_engine/pytorch/ops/sequential.py @@ -15,10 +15,10 @@ class Sequential(torch.nn.Module): - """Sequential container for fusible operations + """Sequential container for fusible operations. - This is a drop-in replacement for `torch.nn.Sequential`, with - support for fusing `FusibleOperation`s. + This is a drop-in replacement for ``torch.nn.Sequential`` with + support for fusing ``FusibleOperation`` s. Parameters ---------- From 57b5b6076d568a7a189a40e592ceea40cdf34cec Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Sat, 21 Feb 2026 06:23:46 +0530 Subject: [PATCH 219/521] Fix race condition in RHT amax kernels (#2695) Fix race condition in HadamardAmaxTmaKernel Signed-off-by: Kirthi Shankar Sivamani --- .../graph_safe_group_hadamard_transform.cu | 5 +++-- .../common/hadamard_transform/group_hadamard_transform.cu | 5 +++-- .../common/hadamard_transform/hadamard_transform.cu | 5 +++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 986229aabf..58b0640249 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -335,8 +335,6 @@ __global__ void GraphSafeGroupHadamardAmaxTmaKernel( is_master_thread); } - ptx::fence_proxy_async_shared_cta(); - // Wait for the data to have arrived ptx::mbarrier_wait_parity(&mbar[stage], 0); @@ -368,6 +366,9 @@ __global__ void GraphSafeGroupHadamardAmaxTmaKernel( // memory. __syncthreads(); } + + // Ensure generic shared-memory accesses are visible before the next TMA write. + ptx::fence_proxy_async_shared_cta(); } } diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index 5d45996dc8..07813be059 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -323,8 +323,6 @@ __global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap t is_master_thread); } - ptx::fence_proxy_async_shared_cta(); - // Wait for the data to have arrived ptx::mbarrier_wait_parity(&mbar[stage], 0); @@ -356,6 +354,9 @@ __global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap t // memory. __syncthreads(); } + + // Ensure generic shared-memory accesses are visible before the next TMA write. + ptx::fence_proxy_async_shared_cta(); } } diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform.cu b/transformer_engine/common/hadamard_transform/hadamard_transform.cu index de930aa2cb..4adc836886 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform.cu @@ -266,8 +266,6 @@ __global__ void HadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor is_master_thread); } - ptx::fence_proxy_async_shared_cta(); - // Wait for the data to have arrived ptx::mbarrier_wait_parity(&mbar[stage], 0); @@ -299,6 +297,9 @@ __global__ void HadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor // memory. __syncthreads(); } + + // Ensure generic shared-memory accesses are visible before the next TMA write. + ptx::fence_proxy_async_shared_cta(); } } From e8f7c5a263c1eb030db34ace6c95e43236c63a6a Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 24 Feb 2026 09:21:41 -0800 Subject: [PATCH 220/521] Add and verify support for `deterministic` fp8 dpa/mha on SM100 (#2621) * add fp8 determinism support Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update cudnn fe to 1.18 Signed-off-by: Sudhakar Singh * enable determinism for sm90 Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/pytorch/attention/dot_product_attention/utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestion from @greptile-apps[bot] Actually switch off fused-attention backend Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Sudhakar Singh * remove extraneous `deterministic` test input arg 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> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 11 +++-- .../common/fused_attn/fused_attn.cpp | 22 +++++----- .../common/fused_attn/fused_attn_fp8.cu | 41 +++++++++++-------- .../common/fused_attn/fused_attn_fp8.h | 10 ++--- .../attention/dot_product_attention/utils.py | 11 ++++- 5 files changed, 56 insertions(+), 39 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 01b2aac453..243fcac882 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1834,10 +1834,16 @@ def get_model(dtype, config): @pytest.mark.parametrize("is_training", [True, False]) @pytest.mark.parametrize("scaling_mode", ["delayed", "current"]) def test_mha_fp8_vs_f16( - dtype, model, qkv_format, input_layernorm, fp8_dpa_bwd, RoPE, is_training, scaling_mode + dtype, + model, + qkv_format, + input_layernorm, + fp8_dpa_bwd, + RoPE, + is_training, + scaling_mode, ): """Test MultiHeadAttention module in FP8""" - os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "1" os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" config = model_configs_fp8_vs_f16[model] @@ -2094,7 +2100,6 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal # config.dropout_p = 0.1 os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_dpa_bwd else "0" - os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "1" os.environ["NVTE_UnfusedDPA_Emulate_FP8"] = "1" # Test backend availability diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 4f8367aac7..b5679280c6 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -770,10 +770,10 @@ void nvte_fused_attn_bwd_qkvpacked( Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); fused_attn_fp8_bwd(b, h, h, max_seqlen, max_seqlen, d, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, &Q_view, &K_view, &V_view, input_O, input_dO, - input_M, input_ZInv, input_S, input_output_dP, &dQ_view, &dK_view, &dV_view, - input_cu_seqlens, input_cu_seqlens, input_rng_state, wkspace, stream, - handle); + bias_type, attn_mask_type, deterministic, &Q_view, &K_view, &V_view, input_O, + input_dO, input_M, input_ZInv, input_S, input_output_dP, &dQ_view, &dK_view, + &dV_view, input_cu_seqlens, input_cu_seqlens, input_rng_state, wkspace, + stream, handle); #else NVTE_ERROR("cuDNN 8.9.0 is required for FP8 fused attention. \n"); #endif @@ -1087,10 +1087,10 @@ void nvte_fused_attn_bwd_kvpacked( Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, input_Q, &K_view, &V_view, input_O, - input_dO, input_M, input_ZInv, input_S, input_output_dP, output_dQ, &dK_view, - &dV_view, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, - stream, handle); + qkv_layout, bias_type, attn_mask_type, deterministic, input_Q, &K_view, + &V_view, input_O, input_dO, input_M, input_ZInv, input_S, input_output_dP, + output_dQ, &dK_view, &dV_view, 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 @@ -1323,9 +1323,9 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso 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, 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, + 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"); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 8c8a289746..80e64370f9 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -1982,13 +1982,13 @@ void fused_attn_fp8_fwd_impl_v1( 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, 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, + 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, 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, @@ -2003,6 +2003,7 @@ void fused_attn_fp8_bwd_impl_v1( bool is_dropout = (dropout_probability != 0.0f); 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!"); @@ -2045,7 +2046,7 @@ void fused_attn_fp8_bwd_impl_v1( 0, 0, true, - false, + deterministic, qkv_tensor_type, o_tensor_type, do_tensor_type, @@ -2216,6 +2217,10 @@ void fused_attn_fp8_bwd_impl_v1( // } // } + if (cudnn_runtime_version >= 91900) { + sdpa_backward_options.set_deterministic_algorithm(deterministic); + } + if (is_padding) { seq_q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("seq_q") @@ -2519,11 +2524,11 @@ void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou 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, 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, + 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) { @@ -2581,11 +2586,11 @@ void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD)) { 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, 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, + 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, 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); diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index a1a932fdf5..225e700eff 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -28,11 +28,11 @@ void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou 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, 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, + 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); diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 3432fd832f..567fd17c34 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1067,8 +1067,15 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_fused_attention = False fused_attention_backend = None - if fused_attention_backend == FusedAttnBackend["FP8"] and is_training: - logger.debug("Disabling FusedAttention for determinism reasons with FP8") + if ( + fused_attention_backend == FusedAttnBackend["FP8"] + and is_training + and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0)) + ): + logger.debug( + "Disabling FusedAttention for determinism reasons with FP8 on arch < sm90 or cuDNN" + " < 9.19.0" + ) use_fused_attention = False fused_attention_backend = None if ( From 39b6dd9e749c8103dbff496986e0ea3615a1efce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 24 Feb 2026 18:38:32 +0100 Subject: [PATCH 221/521] [PyTorch Debug] Custom feature tutorial. (#2216) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add custom feature tutorial for NVInspect debug tools Signed-off-by: Pawel Gadzinski * Fix absolute value in threshold comparison and sync notebook copyright year - Use torch.abs(tensor) for threshold comparison to count values exceeding threshold in both directions, matching tutorial description - Update notebook cell output copyright year from 2025 to 2026 to match source file Signed-off-by: Pawel Gadzinski * Fix stale notebook outputs: abs threshold and copyright year in all output formats - Replace tensor > threshold with torch.abs(tensor) > threshold in text/plain output - Fix copyright year 2025 -> 2026 in text/latex output (LaTeX-encoded PYZhy format) Signed-off-by: Pawel Gadzinski * Remove copyright line from notebook outputs, fix stale count in LaTeX output Copyright line omitted from all rendered outputs (text/plain, text/html, text/latex) to avoid year staleness on future re-runs. Also corrected tensor > threshold to torch.abs(tensor) > threshold in LaTeX output. Signed-off-by: Pawel Gadzinski * Fix import placement and YAML syntax highlighting in tutorial - Move 'import re' from inside plot_stats() to module level (PEP 8) - Use language='yaml' instead of language='python' for YAML config display Signed-off-by: Pawel Gadzinski * Update docs/debug/5_custom_feature_tutorial.ipynb Co-authored-by: Przemyslaw Tredak Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak --- docs/debug.rst | 3 +- docs/debug/5_custom_feature_tutorial.ipynb | 609 ++++++++++++++++++ .../custom_feature_example_config.yaml | 15 + .../percentage_greater_than_threshold.py | 78 +++ docs/debug/custom_feature_dir/utils.py | 48 ++ 5 files changed, 752 insertions(+), 1 deletion(-) create mode 100644 docs/debug/5_custom_feature_tutorial.ipynb create mode 100644 docs/debug/custom_feature_dir/custom_feature_example_config.yaml create mode 100644 docs/debug/custom_feature_dir/percentage_greater_than_threshold.py create mode 100644 docs/debug/custom_feature_dir/utils.py diff --git a/docs/debug.rst b/docs/debug.rst index 7c3735ee12..d4ed106897 100644 --- a/docs/debug.rst +++ b/docs/debug.rst @@ -12,4 +12,5 @@ Precision debug tools debug/1_getting_started.rst debug/2_config_file_structure.rst debug/api - debug/4_distributed.rst \ No newline at end of file + debug/4_distributed.rst + debug/5_custom_feature_tutorial.ipynb \ No newline at end of file diff --git a/docs/debug/5_custom_feature_tutorial.ipynb b/docs/debug/5_custom_feature_tutorial.ipynb new file mode 100644 index 0000000000..cd7dfddc46 --- /dev/null +++ b/docs/debug/5_custom_feature_tutorial.ipynb @@ -0,0 +1,609 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "85a94734", + "metadata": {}, + "source": "# Adding custom feature to precision debug tools\n \nTE comes with several built-in features, such as `LogFp8TensorStats`, which can log statistics for each tensor involved in matrix multiplication (GEMM) operations.\nIn this tutorial, we'll demonstrate how to extend TE by adding a custom feature. This custom feature will log the percentage of elements in a tensor whose absolute values exceed a configurable threshold `t`, as specified in the config file.\n\nCustom features can be used for example for:\n\n1. Logging custom statistics.\n2. Dumping intermediate tensors.\n3. Experiments with modifying intermediate tensors.\n\nHow to add custom feature:\n\n1. Add Python with feature class definition which inherits from `transformer_engine.debug.features.api.TEConfigAPIMapper`.\n2. Wrap the class with `@Registry.register_feature(namespace=\"transformer_engine\")`.\n3. Implement some of API calls to nvidia-dl-framework-inspect described [here](3_api_te_calls.rst).\n\nLet's define a new file at `.../custom_feature_dir/percentage_greater_than_threshold.py` containing the following code:\n" + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "b4e7562d", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
stats:\n",
+       "  enabled: True\n",
+       "  layers:\n",
+       "    layer_name_regex_pattern: .*\n",
+       "  transformer_engine:\n",
+       "    PercentageGreaterThanThreshold:\n",
+       "      enabled: True\n",
+       "      tensors: [activation]\n",
+       "      threshold: 0.1\n",
+       "      freq: 5\n",
+       "    LogTensorStats:\n",
+       "      enabled: True\n",
+       "      tensors: [activation]\n",
+       "      stats: [min]\n",
+       "      freq: 5\n",
+       "
\n" + ], + "text/latex": [ + "\\begin{Verbatim}[commandchars=\\\\\\{\\}]\n", + "\\PY{n}{stats}\\PY{p}{:}\n", + " \\PY{n}{enabled}\\PY{p}{:} \\PY{k+kc}{True}\n", + " \\PY{n}{layers}\\PY{p}{:}\n", + " \\PY{n}{layer\\PYZus{}name\\PYZus{}regex\\PYZus{}pattern}\\PY{p}{:} \\PY{o}{.}\\PY{o}{*}\n", + " \\PY{n}{transformer\\PYZus{}engine}\\PY{p}{:}\n", + " \\PY{n}{PercentageGreaterThanThreshold}\\PY{p}{:}\n", + " \\PY{n}{enabled}\\PY{p}{:} \\PY{k+kc}{True}\n", + " \\PY{n}{tensors}\\PY{p}{:} \\PY{p}{[}\\PY{n}{activation}\\PY{p}{]}\n", + " \\PY{n}{threshold}\\PY{p}{:} \\PY{l+m+mf}{0.1}\n", + " \\PY{n}{freq}\\PY{p}{:} \\PY{l+m+mi}{5}\n", + " \\PY{n}{LogTensorStats}\\PY{p}{:}\n", + " \\PY{n}{enabled}\\PY{p}{:} \\PY{k+kc}{True}\n", + " \\PY{n}{tensors}\\PY{p}{:} \\PY{p}{[}\\PY{n}{activation}\\PY{p}{]}\n", + " \\PY{n}{stats}\\PY{p}{:} \\PY{p}{[}\\PY{n+nb}{min}\\PY{p}{]}\n", + " \\PY{n}{freq}\\PY{p}{:} \\PY{l+m+mi}{5}\n", + "\\end{Verbatim}\n" + ], + "text/plain": [ + "stats:\n", + " enabled: True\n", + " layers:\n", + " layer_name_regex_pattern: .*\n", + " transformer_engine:\n", + " PercentageGreaterThanThreshold:\n", + " enabled: True\n", + " tensors: [activation]\n", + " threshold: 0.1\n", + " freq: 5\n", + " LogTensorStats:\n", + " enabled: True\n", + " tensors: [activation]\n", + " stats: [min]\n", + " freq: 5" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from IPython.display import Code\n", + "Code(filename='./custom_feature_dir/custom_feature_example_config.yaml', language='yaml')" + ] + }, + { + "cell_type": "markdown", + "id": "3929f293-7ac1-48b0-8a4d-23bb6976aa0b", + "metadata": {}, + "source": [ + "To use this feature one needs to add `.../custom_feature_dir` to `debug_api.initialize(feature_dirs=...`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "d82f1c82", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "NVDLFW INSPECT - 2025-10-17 14:16:42,204 - WARNING - Reduction group initialized for tensor reduction before logging statistics. If per-rank statistics are required, pass `skip_reduction=True` when invoking the API. To pass another reduction group, use `reduction_group` kwarg when invoking the API.\n" + ] + } + ], + "source": [ + "import os, time\n", + "import torch\n", + "import transformer_engine.pytorch as te\n", + "import nvdlfw_inspect.api as debug_api\n", + "\n", + "te_dir = os.environ[\"TE_PATH\"] # setup TE dir as environment variable to run this script\n", + "log_dir = os.environ.get(\"LOG_PATH\", \"./log\")\n", + "\n", + "debug_api.initialize(\n", + " config_file=te_dir + \"/docs/debug/custom_feature_dir/custom_feature_example_config.yaml\",\n", + " feature_dirs=[\n", + " te_dir + \"/transformer_engine/debug/features\", \n", + " te_dir + \"/docs/debug/custom_feature_dir\" # One needs to add path to the custom feature dir here\n", + " ],\n", + " log_dir=log_dir,\n", + " default_logging_enabled=True)\n", + "\n", + "debug_api.set_tensor_reduction_group(None) # For distributed training one needs to set the reduction group\n", + "\n", + "module = te.Linear(128, 128, name=\"linear_1\")\n", + "inp = torch.randn(128, 128).cuda()\n", + "\n", + "# Simple training loop with measuring the time\n", + "times = []\n", + "for _ in range(100):\n", + " time_start = time.time()\n", + " inp.normal_()\n", + " out = module(inp)\n", + " out.sum().backward()\n", + " torch.cuda.synchronize()\n", + " time_end = time.time()\n", + " times.append(time_end - time_start)\n", + "\n", + " debug_api.step()" + ] + }, + { + "cell_type": "markdown", + "id": "e4f129a9", + "metadata": {}, + "source": [ + "Now, let's plot the gathered stats." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "b68a21ea", + "metadata": {}, + "outputs": [ + { + "data": { + "image/png": "iVBORw0KGgoAAAANSUhEUgAABKUAAAHqCAYAAADVi/1VAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjMsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvZiW1igAAAAlwSFlzAAAPYQAAD2EBqD+naQABAABJREFUeJzs3XdYk+f6B/Dvm4QkbJAtIggooOKeuCfuXav2uLXW1tpqbaunQ+342Z7a2lattrZVa4daR6sd7tG6J25QEGQIhL1HSJ7fH8n7SmRDFnB/rsvrnGY+CZC8uXPf34djjDEQQgghhBBCCCGEEGJEIlMvgBBCCCGEEEIIIYQ0PlSUIoQQQgghhBBCCCFGR0UpQgghhBBCCCGEEGJ0VJQihBBCCCGEEEIIIUZHRSlCCCGEEEIIIYQQYnRUlCKEEEIIIYQQQgghRkdFKUIIIYQQQgghhBBidFSUIoQQQgghhBBCCCFGR0UpQgghhBBCCCGEEGJ0VJQihJgcx3FYtWqVqZdRoVmzZsHHx8fUyyCEEEII0SuO47Bo0SJTL6NKMTEx4DgOa9euNfVSABhmPdu2bQPHcYiJianysj4+Ppg1a5be7rs+qMnzQ+oXKkoRQvSCf6PgOA5nzpwpcz5jDF5eXuA4DqNGjTLBCgkhhBDDK/1+yHEc5HI5WrVqhUWLFiE5OdnUy6uzu3fvYtWqVfTBsJZM8fydO3cOq1atQmZmptHus7b++usvs/6isj46cOAAOnXqBLlcjubNm2PlypUoKSmp1nU//PBDjBkzBm5ubmb/JTKpv6goRQjRK7lcjp9//rnM6adPn0Z8fDxkMlmZ8woKCvD2228bY3mEEEKIUbz33nvYsWMHNmzYgJCQEGzatAk9e/ZEfn6+qZdWJ3fv3sXq1aupKFVLpnj+zp07h9WrV9ebotTq1atNvYwG4++//8a4cePg4OCA9evXY9y4cfjggw/w8ssvV+v6b7/9Ni5fvoyOHTsaeKWkMZOYegGEkIZlxIgR+PXXX/Hll19CInnyEvPzzz+jc+fOSE1NLXMduVxuzCUSQgghBjd8+HB06dIFADBv3jw4OTnhs88+w++//46pU6fW6bbz8/NhZWWlj2WSChQWFkIqlUIkqh/f4efl5cHa2trUyzBbjfX5WbZsGdq1a4cjR44Ix+V2dnb4v//7P7zyyisIDAys9PrR0dHw8fFBamoqXFxcjLFk0gjVj1dZQki9MXXqVKSlpeHo0aPCacXFxdizZw+mTZtW7nWebgdetWoVOI5DZGQkZs2aBQcHB9jb22P27NlVfsO8aNEi2NjYlHu5qVOnwt3dHSqVCgDw+++/Y+TIkWjatClkMhn8/Pzw/vvvC+dX5NSpU+A4DqdOndI5nc8X2LZtm87p4eHhmDRpEpo0aQK5XI4uXbrgwIEDOpdRKpVYvXo1WrZsCblcDicnJ/Tu3VvneSSEEFJ/DRw4EIDmQx7vxx9/ROfOnWFpaYkmTZpgypQpiIuL07le//790bZtW1y9ehV9+/aFlZUV/vvf/wLQFE5WrVqFVq1aQS6Xw8PDAxMmTEBUVJRwfbVajc8//xxt2rSBXC6Hm5sbFixYgIyMDJ378fHxwahRo3DmzBl069YNcrkcvr6++OGHH4TLbNu2Dc888wwAYMCAAcKIIv9+WJP31Y0bN8LX1xeWlpbo1q0b/v33X/Tv3x/9+/fXuVxRURFWrlwJf39/yGQyeHl54Y033kBRUVG1nvfq3A//vr5z5068/fbb8PT0hJWVFbKzswEAFy9exLBhw2Bvbw8rKyv069cPZ8+e1bmfR48e4cUXX0RAQAAsLS3h5OSEZ555RqcjqqrnD9B0tvTp0wfW1tawtbXFyJEjcefOHZ37mjVrFmxsbBAVFYURI0bA1tYWzz33XLmPf9WqVXj99dcBAC1atBDu8+lOrd9++w1t27aFTCZDmzZtcOjQoRo/Pv4xchyHs2fPYunSpXBxcYG1tTXGjx+PlJSUctdY+nFt3LgRAHRGYJ/2zTffwM/PDzKZDF27dsXly5er/fxU9+/hypUrCA0NhbOzMywtLdGiRQvMmTOn3HVXtR4AOHHihPBzdXBwwNixY3Hv3r1Knw9AE3/xwQcfoFmzZrCyssKAAQPK/D5U5O7du7h79y6ef/55nS+KX3zxRTDGsGfPnipvo7aZqnv27AHHcTh9+nSZ877++mtwHIfbt28DAG7evIlZs2bB19cXcrkc7u7umDNnDtLS0qq8n4pGCsvL3MrMzMSrr74KLy8vyGQy+Pv74+OPP4Zarda53M6dO9G5c2fY2trCzs4OwcHB+OKLL6r/4EmNUacUIUSvfHx80LNnT/zyyy8YPnw4AM0BVlZWFqZMmYIvv/yy2rc1efJktGjRAmvWrMG1a9fw7bffwtXVFR9//HGF13n22WexceNG/Pnnn8KBH6D5VvngwYOYNWsWxGIxAM2Bk42NDZYuXQobGxucOHEC7777LrKzs/HJJ5/U8hnQdefOHfTq1Quenp5Yvnw5rK2tsXv3bowbNw579+7F+PHjAWgOGtesWYN58+ahW7duyM7OxpUrV3Dt2jUMGTJEL2shhBBiOnyhyMnJCYAmq+Wdd97B5MmTMW/ePKSkpGD9+vXo27cvrl+/DgcHB+G6aWlpGD58OKZMmYL//Oc/cHNzg0qlwqhRo3D8+HFMmTIFr7zyCnJycnD06FHcvn0bfn5+AIAFCxZg27ZtmD17NhYvXozo6Ghs2LAB169fx9mzZ2FhYSHcT2RkJCZNmoS5c+di5syZ+P777zFr1ix07twZbdq0Qd++fbF48WJ8+eWX+O9//4ugoCAAEP63uu+rmzZtwqJFi9CnTx8sWbIEMTExGDduHBwdHdGsWTPhcmq1GmPGjMGZM2fw/PPPIygoCLdu3cK6detw//59/Pbbb5U+59W9H977778PqVSKZcuWoaioCFKpFCdOnMDw4cPRuXNnrFy5EiKRCFu3bsXAgQPx77//olu3bgCAy5cv49y5c5gyZQqaNWuGmJgYbNq0Cf3798fdu3dhZWVV5fO3Y8cOzJw5E6Ghofj444+Rn5+PTZs2oXfv3rh+/bpOgaCkpAShoaHo3bs31q5dW2Hn3IQJE3D//n388ssvWLduHZydnQFAp+vlzJkz2LdvH1588UXY2triyy+/xMSJExEbGyv8vlbn8ZX28ssvw9HREStXrkRMTAw+//xzLFq0CLt27arw57VgwQI8fvwYR48exY4dO8q9zM8//4ycnBwsWLAAHMfhf//7HyZMmICHDx/q/C5X9PxU5+9BoVBg6NChcHFxwfLly+Hg4ICYmBjs27evVus5duwYhg8fDl9fX6xatQoFBQVYv349evXqhWvXrlVa+Hn33XfxwQcfYMSIERgxYgSuXbuGoUOHori4uMLr8K5fvw4AQscmr2nTpmjWrJlwviGMHDkSNjY22L17N/r166dz3q5du9CmTRu0bdsWAHD06FE8fPgQs2fPhru7O+7cuYNvvvkGd+7cwYULF8otTNZUfn4++vXrh4SEBCxYsADNmzfHuXPnsGLFCiQmJuLzzz8X1jJ16lQMGjRI+Lxx7949nD17Fq+88kqd10EqwAghRA+2bt3KALDLly+zDRs2MFtbW5afn88YY+yZZ55hAwYMYIwx5u3tzUaOHKlzXQBs5cqVwn+vXLmSAWBz5szRudz48eOZk5NTpetQq9XM09OTTZw4Uef03bt3MwDsn3/+EU7j11faggULmJWVFSssLBROmzlzJvP29hb+++TJkwwAO3nypM51o6OjGQC2detW4bRBgwax4OBgndtTq9UsJCSEtWzZUjitffv2ZZ4XQggh9Q//fnjs2DGWkpLC4uLi2M6dO5mTkxOztLRk8fHxLCYmhonFYvbhhx/qXPfWrVtMIpHonN6vXz8GgG3evFnnst9//z0DwD777LMya1Cr1Ywxxv79918GgP3000865x86dKjM6d7e3mXeJxUKBZPJZOy1114TTvv111/LfQ9krHrvq0VFRczJyYl17dqVKZVK4XLbtm1jAFi/fv2E03bs2MFEIhH7999/dW5z8+bNDAA7e/Zsmfvj1eR++Pd1X19fncegVqtZy5YtWWhoqPCc8o+zRYsWbMiQIZU+9vPnzzMA7IcffhBOq+j5y8nJYQ4ODmz+/Pk6pyclJTF7e3ud02fOnMkAsOXLl1f4+Ev75JNPGAAWHR1d5jwATCqVssjISOG0GzduMABs/fr1NX58/O//4MGDdZ6zJUuWMLFYzDIzMytd60svvcTK+4jKH2M5OTmx9PR04fTff/+dAWAHDx4UTqvo+anu38P+/fuFY9qK1GQ9HTp0YK6uriwtLU047caNG0wkErEZM2YIp/HPHf9zUigUTCqVspEjR+o8l//9738ZADZz5swK18fYk597bGxsmfO6du3KevToUen1S0tJSSlzvF6VqVOnMldXV1ZSUiKclpiYyEQiEXvvvfeE08r73frll1/KvB49/fwwVvYzBM/b21vn+Xn//feZtbU1u3//vs7lli9fzsRisfAcvfLKK8zOzk5nzcTwaHyPEKJ3kydPRkFBAf744w/k5OTgjz/+qHB0rzIvvPCCzn/36dMHaWlpQjt9eTiOwzPPPIO//voLubm5wum7du2Cp6cnevfuLZxmaWkp/P+cnBykpqaiT58+yM/PR3h4eI3X+7T09HScOHECkydPFm4/NTUVaWlpCA0NxYMHD5CQkAAAcHBwwJ07d/DgwYM63y8hhBDTGzx4MFxcXODl5YUpU6bAxsYG+/fvh6enJ/bt2we1Wo3JkycL7w2pqalwd3dHy5YtcfLkSZ3bkslkmD17ts5pe/fuhbOzc7mBxXxnwa+//gp7e3sMGTJE5346d+4MGxubMvfTunVr9OnTR/hvFxcXBAQE4OHDh9V6zNV5X71y5QrS0tIwf/58nZGi5557Do6Ojjq39+uvvyIoKAiBgYE66+dHIZ9ef2k1uR/ezJkzdR5DWFgYHjx4gGnTpiEtLU24/7y8PAwaNAj//POPMPpT+npKpRJpaWnw9/eHg4MDrl27VuVzd/ToUWRmZmLq1Kk6j1UsFqN79+7lPtaFCxdWebvVMXjwYKGzDgDatWsHOzs7nZ97TR/f888/r9Ph0qdPH6hUKjx69KhOa3322Wd1fn7872t5v6NPPz/V/XvguxT/+OMPKJXKOq0nMTERYWFhmDVrFpo0aSJcrl27dhgyZAj++uuvCm/72LFjKC4uxssvv6zzXL766quVrolXUFAAAOVuMiSXy4XzDeXZZ5+FQqHQGU/ds2cP1Go1nn32WeG00r9bhYWFSE1NRY8ePQCgWn871fHrr7+iT58+cHR01PnZDx48GCqVCv/88w8Azc8+Ly+P4jOMjMb3CCF65+LigsGDB+Pnn39Gfn4+VCoVJk2aVOPbad68uc5/82/6GRkZsLOzq/B6zz77LD7//HMcOHAA06ZNQ25uLv766y+htZp3584dvP322zhx4kSZQldWVlaN1/u0yMhIMMbwzjvv4J133in3MgqFAp6ennjvvfcwduxYtGrVCm3btsWwYcMwffp0tGvXrs7rIIQQYnwbN25Eq1atIJFI4ObmhoCAACE0+8GDB2CMoWXLluVet/QYEgB4enpCKpXqnBYVFYWAgACdgsvTHjx4gKysLLi6upZ7vkKh0Pnvp993Ac1779N5OxWpzvsqX5Tw9/fXOV8ikZQZY3rw4AHu3btXYcDy0+svrSb3w2vRokWZ+wc0xaqKZGVlwdHREQUFBVizZg22bt2KhIQEMMZ0LlMV/r74gtvTnj7ukUgk5Y4g1kZ1fu41fXyVHcPpc60V3W55z091/x769euHiRMnYvXq1Vi3bh369++PcePGYdq0aWUKPFWth/89DAgIKHN/QUFBOHz4cIUh7Px1n36dcHFxqbCwWhpf7Ckvf62wsFCnGGQIfA7brl27MGjQIACaL4k7dOiAVq1aCZdLT0/H6tWrsXPnzjJ/0/o4Hgc0P/ubN29W+Vry4osvYvfu3Rg+fDg8PT0xdOhQTJ48GcOGDdPLOkj5qChFCDGIadOmYf78+UhKSsLw4cN1sjGqi89+elrpA6Hy9OjRAz4+Pti9ezemTZuGgwcPoqCgQOdbmczMTPTr1w92dnZ477334OfnB7lcjmvXruHNN98sE3pYWkWz7U8HufK3sWzZMoSGhpZ7Hf5guW/fvoiKisLvv/+OI0eO4Ntvv8W6deuwefNmzJs3r9LHSwghxPx069atTJYLT61Wg+M4/P333+W+19nY2Oj8d20/PKrVari6uuKnn34q9/ynP6DV9n0XqNv7amXrDw4OxmeffVbu+V5eXjW+zco8/Tzza/7kk0/QoUOHcq/D/6xefvllbN26Fa+++ip69uwJe3t7cByHKVOmVOux85fZsWMH3N3dy5z/dPFRJpPpbWfA6vzca/r46vK7VNe1AuU/P9X9e+A4Dnv27MGFCxdw8OBBHD58GHPmzMGnn36KCxcu6Px9Gupx6oOHhwcATbfW038riYmJQh6aochkMowbNw779+/HV199heTkZJw9exb/93//p3O5yZMn49y5c3j99dfRoUMH2NjYQK1WY9iwYbV63QDKPyYfMmQI3njjjXIvzxfJXF1dERYWhsOHD+Pvv//G33//ja1bt2LGjBnYvn17rdZCqkZFKUKIQYwfPx4LFizAhQsXKg21NJTJkyfjiy++QHZ2Nnbt2gUfHx+hFRjQ7LSTlpaGffv2oW/fvsLppXdFqgj/7VRmZqbO6U+3pPv6+gLQfOM9ePDgKm+3SZMmmD17NmbPno3c3Fz07dsXq1atoqIUIYQ0MH5+fmCMoUWLFjodAzW9jYsXL0KpVJbprCp9mWPHjqFXr15664qo6IuZ6r6vent7A9B0Ew8YMEA4vaSkBDExMTodwn5+frhx4wYGDRpU47DjmtxPRfiRNjs7uyrfx/fs2YOZM2fi008/FU4rLCwsc6xQ0ePg78vV1bVaxww1oY+g6Oo+vrrSx1orUtO/hx49eqBHjx748MMP8fPPP+O5557Dzp07a3Rcxv8eRkRElDkvPDwczs7O5XZJlb7ugwcPhGNKAEhJSalWxxlfSL1y5YpOAerx48eIj4/H888/X+3HUVvPPvsstm/fjuPHj+PevXtgjOl8SZyRkYHjx49j9erVePfdd4XTqxtn4ejoWOZ3sLi4GImJiTqn+fn5ITc3t1p/W1KpFKNHj8bo0aOhVqvx4osv4uuvv8Y777xTpvOS6AdlShFCDMLGxgabNm3CqlWrMHr0aKPf/7PPPouioiJs374dhw4dwuTJk3XO57/ZKv1NVnFxMb766qsqb9vb2xtisViYP+c9fV1XV1f0798fX3/9dZk3RwA6WyM/ve2tjY0N/P39q73lNSGEkPpjwoQJEIvFWL16dZmOCsZYtbZCnzhxIlJTU7Fhw4Yy5/G3OXnyZKhUKrz//vtlLlNSUlKrggL/Afrp61b3fbVLly5wcnLCli1bUFJSIpz+008/lfmgPXnyZCQkJGDLli1l1lFQUIC8vLwK11mT+6lI586d4efnh7Vr1+rkVPJKv4+LxeIyP8v169eX6dio6PkLDQ2FnZ0d/u///q/cHKPS91VTFd1nTVT38dWVPtZaker+PWRkZJR5rHyBp6bHZR4eHujQoQO2b9+u85hu376NI0eOYMSIERVed/DgwbCwsMD69et11sPvFFeVNm3aIDAwEN98843Oz2nTpk3gOE4nWiMrKwvh4eF6G5fjDR48GE2aNMGuXbuwa9cudOvWTWdMtrzXDaD6j9HPz6/M8fjTjxfQ/OzPnz+Pw4cPl7mNzMxM4TXi6ddekUgkFLDpmNxwqFOKEGIwlWUwGFqnTp3g7++Pt956C0VFRTrfygBASEgIHB0dMXPmTCxevBgcx2HHjh3Vare2t7fHM888g/Xr14PjOPj5+eGPP/4oN9ti48aN6N27N4KDgzF//nz4+voiOTkZ58+fR3x8PG7cuAFAEy7bv39/dO7cGU2aNMGVK1ewZ88eLFq0SD9PCCGEELPh5+eHDz74ACtWrEBMTAzGjRsHW1tbREdHY//+/Xj++eexbNmySm9jxowZ+OGHH7B06VJcunQJffr0QV5eHo4dO4YXX3wRY8eORb9+/bBgwQKsWbMGYWFhGDp0KCwsLPDgwQP8+uuv+OKLL2qc+dihQweIxWJ8/PHHyMrKgkwmw8CBA6v9viqVSrFq1Sq8/PLLGDhwICZPnoyYmBhs27YNfn5+Op0y06dPx+7du/HCCy/g5MmT6NWrF1QqFcLDw7F7924cPny4whHJmtxPRUQiEb799lsMHz4cbdq0wezZs+Hp6YmEhAScPHkSdnZ2OHjwIABg1KhR2LFjB+zt7dG6dWucP38ex44dg5OTU7WeP1dXV2zatAnTp09Hp06dMGXKFLi4uCA2NhZ//vknevXqVW4Bsjo6d+4MAHjrrbcwZcoUWFhYYPTo0RV26JSnuo+vrvi1Ll68GKGhoRCLxZgyZYpebru6fw/bt2/HV199hfHjx8PPzw85OTnYsmUL7OzsKi0iVeSTTz7B8OHD0bNnT8ydOxcFBQVYv3497O3tsWrVqgqv5+LigmXLlmHNmjUYNWoURowYgevXr+Pvv/+Gs7Nzte97zJgxGDp0KKZMmYLbt29jw4YNmDdvHoKCgoTL7d+/H7Nnz8bWrVsxa9Ys4fQdO3bg0aNHyM/PBwD8888/+OCDDwBo/j75bq6KWFhYYMKECdi5cyfy8vKwdu1anfPt7OzQt29f/O9//4NSqYSnpyeOHDlSrckFAJg3bx5eeOEFTJw4EUOGDMGNGzdw+PDhMs/P66+/jgMHDmDUqFGYNWsWOnfujLy8PNy6dQt79uxBTEwMnJ2dMW/ePKSnp2PgwIFo1qwZHj16hPXr16NDhw46zxfRM+Nt9EcIacj4bVor2z6XMc0WrSNHjtQ5DU9t57py5UoGgKWkpJR7H+VtaVyet956iwFg/v7+5Z5/9uxZ1qNHD2ZpacmaNm3K3njjDXb48OEyWzXPnDmTeXt761w3JSWFTZw4kVlZWTFHR0e2YMECdvv2bQaAbd26VeeyUVFRbMaMGczd3Z1ZWFgwT09PNmrUKLZnzx7hMh988AHr1q0bc3BwYJaWliwwMJB9+OGHrLi4uFqPlRBCiHmo7vshY4zt3buX9e7dm1lbWzNra2sWGBjIXnrpJRYRESFcpl+/fqxNmzblXj8/P5+99dZbrEWLFszCwoK5u7uzSZMmsaioKJ3LffPNN6xz587M0tKS2drasuDgYPbGG2+wx48fC5cp7/2Zv/9+/frpnLZlyxbm6+vLxGKxzntmdd9XGWPsyy+/ZN7e3kwmk7Fu3bqxs2fPss6dO7Nhw4bpXK64uJh9/PHHrE2bNkwmkzFHR0fWuXNntnr1apaVlVXVU1yt+zl58iQDwH799ddyb+P69etswoQJzMnJiclkMubt7c0mT57Mjh8/LlwmIyODzZ49mzk7OzMbGxsWGhrKwsPDy2xNX9nzx68lNDSU2dvbM7lczvz8/NisWbPYlStXhMvMnDmTWVtbV/nYS3v//feZp6cnE4lEOsdSANhLL71U5vJPr7u6j6+i33/+OX769+BpJSUl7OWXX2YuLi6M4zjGf1yNjo5mANgnn3xS5jpPH0dW9fxU9fdw7do1NnXqVNa8eXMmk8mYq6srGzVqlM7PoCbrYYyxY8eOsV69ejFLS0tmZ2fHRo8eze7evatzmfKOc1UqFVu9ejXz8PBglpaWrH///uz27dvl/l5VZP/+/axDhw5MJpOxZs2asbfffrvM8SV/308fw/br148BKPdfVT9L3tGjRxkAxnEci4uLK3N+fHw8Gz9+PHNwcGD29vbsmWeeYY8fPy7zPFb0/Lz55pvM2dmZWVlZsdDQUBYZGVnu85OTk8NWrFjB/P39mVQqZc7OziwkJIStXbtWeD727NnDhg4dylxdXZlUKmXNmzdnCxYsYImJidV6rKR2OMbMIIWNEEIIIYQQ0mip1Wq4uLhgwoQJ5Y7r1bf7IYQQUj2UKUUIIYQQQggxmsLCwjJjfT/88APS09PRv3//enc/hBBCao86pQghhBBCCCFGc+rUKSxZsgTPPPMMnJyccO3aNXz33XcICgrC1atXIZVK69X9EEIIqT0KOieEEEIIIYQYjY+PD7y8vPDll18iPT0dTZo0wYwZM/DRRx/ptVBkrPshhBBSe9QpRQghhBBCCCGEEEKMjjKlCCGEEEIIIYQQQojRUVGKEEIIIYQQQgghhBhdg8uUUqvVePz4MWxtbcFxnKmXQwghhJB6iDGGnJwcNG3aFCJR4/sOj46nCCGEEFIX1T2WanBFqcePH8PLy8vUyyCEEEJIAxAXF4dmzZqZehlGR8dThBBCCNGHqo6lGlxRytbWFoDmgdvZ2Zl4NYQQQgipj7Kzs+Hl5SUcVzQ2dDxFCCGEkLqo7rFUgytK8S3mdnZ2dBBFCCGEkDpprKNrdDxFCCGEEH2o6liq8YUkEEIIIYQQQgghhBCTo6IUIYQQQgghhBBCCDE6KkoRQgghhBBCCCGEEKNrcJlShBBCSGOnUqmgVCpNvQyzZmFhAbFYbOplEELqGbVajeLiYlMvgxBCTE5fx1JUlCKEEEIaCMYYkpKSkJmZaeql1AsODg5wd3dvtGHmhJCaKS4uRnR0NNRqtamXQgghZkEfx1JUlCKEEEIaCL4g5erqCisrKyq2VIAxhvz8fCgUCgCAh4eHiVdECDF3jDEkJiZCLBbDy8sLIhGloBBCGi99HktRUYoQQghpAFQqlVCQcnJyMvVyzJ6lpSUAQKFQwNXV1WCjfBs3bsQnn3yCpKQktG/fHuvXr0e3bt2qvN7OnTsxdepUjB07Fr/99pvOeffu3cObb76J06dPo6SkBK1bt8bevXvRvHlz4TLnz5/HW2+9hYsXL0IsFqNDhw44fPiw8LgJITVTUlKC/Px8NG3aFFZWVqZeDiGEmJy+jqWoxE8IIYQ0AHyGFH1Yqj7+uTJU/tauXbuwdOlSrFy5EteuXUP79u0RGhoqfKtYkZiYGCxbtgx9+vQpc15UVBR69+6NwMBAnDp1Cjdv3sQ777wDuVwuXOb8+fMYNmwYhg4dikuXLuHy5ctYtGgRdXYQUgcqlQoAIJVKTbwSQggxH/o4lqJOKUIIIaQBoZG96jP0c/XZZ59h/vz5mD17NgBg8+bN+PPPP/H9999j+fLl5V5HpVLhueeew+rVq/Hvv/+WyQd76623MGLECPzvf/8TTvPz89O5zJIlS7B48WKd+wgICNDToyKkcaPXWEIIeUIfr4n0lRkhhBBCiJ4VFxfj6tWrGDx4sHCaSCTC4MGDcf78+Qqv995778HV1RVz584tc55arcaff/6JVq1aITQ0FK6urujevbvOeJ9CocDFixfh6uqKkJAQuLm5oV+/fjhz5kyl6y0qKkJ2drbOP0IIIYQQQ6OiFCGEEELqlVOnToHjOLPeZTA1NRUqlQpubm46p7u5uSEpKanc65w5cwbfffcdtmzZUu75CoUCubm5+OijjzBs2DAcOXIE48ePx4QJE3D69GkAwMOHDwEAq1atwvz583Ho0CF06tQJgwYNwoMHDypc75o1a2Bvby/88/Lyqs3DJoQQ0sD1798fr776qqmXUUZMTAw4jkNYWJhR71dfxyQcx5XJkCzNVI/PGKgoRQghhJB6JSQkBImJibC3tzf1UvQmJycH06dPx5YtW+Ds7FzuZfht6MeOHYslS5agQ4cOWL58OUaNGoXNmzfrXGbBggWYPXs2OnbsiHXr1iEgIADff/99hfe/YsUKZGVlCf/i4uL0/AgJIaRmqvqQ3lgZ63kx5y+AZs2ahXHjxpl6GWYtNjYWI0eOhJWVFVxdXfH666+jpKSk0ut8+OGHCAkJgZWVFRwcHIyzUFCmFCGEEELqGalUCnd3d1Mvo1LOzs4Qi8VITk7WOT05ObnctUdFRSEmJgajR48WTuMLTBKJBBEREfDy8oJEIkHr1q11rhsUFCSM5/FbMpd3mdjY2ArXK5PJIJPJavAICSGNkUqlAsdxtHFCDRUXF5tdSL5SqYSFhYWpl2FSDfX3WaVSYeTIkXB3d8e5c+eQmJiIGTNmwMLCAv/3f/9X4fWKi4vxzDPPoGfPnvjuu++Mtt6G9ewTQgghpN7p378/Xn75Zbz66qtwdHSEm5sbtmzZgry8PMyePRu2trbw9/fH33//DaDst7fbtm2Dg4MDDh8+jKCgINjY2GDYsGFITEw02WOSSqXo3Lkzjh8/LpymVqtx/Phx9OzZs8zlAwMDcevWLYSFhQn/xowZgwEDBiAsLAxeXl6QSqXo2rUrIiIidK57//59eHt7AwB8fHzQtGnTSi9DCGk8+vfvj0WLFmHRokWwt7eHs7Mz3nnnHTDGAGjy5JYtWwZPT09YW1uje/fuOHXqlHB9/vX1wIEDaN26NWQyGWJjY1FUVIQ333wTXl5ekMlk8Pf31/kQe/v2bQwfPhw2NjZwc3PD9OnTkZqaqrOuxYsX44033kCTJk3g7u6OVatWCef7+PgAAMaPHw+O44T/joqKwtixY+Hm5gYbGxt07doVx44d03nMiYmJGDlyJCwtLdGiRQv8/PPP8PHxweeffy5cJjMzE/PmzYOLiwvs7OwwcOBA3Lhxo9rP6wcffABXV1fY2tpi3rx5WL58OTp06CCcz3fyfPjhh2jatKmw2URcXBwmT54MBwcHNGnSBGPHjkVMTIxwvcuXL2PIkCFwdnaGvb09+vXrh2vXrlX5vADA77//jk6dOkEul8PX1xerV6/W6YzhOA6bNm3CmDFjYG1tjQ8//LDCxxcTE4MBAwYAABwdHcFxHGbNmiWcr1arK/zZAZqNPoKDg2FtbQ0vLy+8+OKLyM3NFc6vy/v2qlWrsH37dvz+++/gOA4cx+n8zj58+BADBgyAlZUV2rdvr5PjWNnvc2V/B48ePcLo0aPh6OgIa2trtGnTBn/99ZfOuq5evYouXbrAysoKISEhZd6HN23aBD8/P0ilUgQEBGDHjh2VPs5Lly6hY8eOkMvl6NKlC65fv17lc8M7cuQI7t69ix9//BEdOnTA8OHD8f7772Pjxo0oLi6u8HqrV6/GkiVLEBwcXO370gvWwGRlZTEALCsry9RLIYQQQoymoKCA3b17lxUUFAinqdVqllekNMk/tVpd7bX369eP2drasvfff5/dv3+fvf/++0wsFrPhw4ezb775ht2/f58tXLiQOTk5sby8PHby5EkGgGVkZDDGGNu6dSuzsLBggwcPZpcvX2ZXr15lQUFBbNq0aTV+znj6OJ7YuXMnk8lkbNu2bezu3bvs+eefZw4ODiwpKYkxxtj06dPZ8uXLK7z+zJkz2dixY3VO27dvH7OwsGDffPMNe/DgAVu/fj0Ti8Xs33//FS6zbt06Zmdnx3799Vf24MED9vbbbzO5XM4iIyOrvXY6niJE19OvF/Xp9dXGxoa98sorLDw8nP3444/MysqKffPNN4wxxubNm8dCQkLYP//8wyIjI9knn3zCZDIZu3//PmPsyetrSEgIO3v2LAsPD2d5eXls8uTJzMvLi+3bt49FRUWxY8eOsZ07dzLGGMvIyGAuLi5sxYoV7N69e+zatWtsyJAhbMCAATrrsrOzY6tWrWL3799n27dvZxzHsSNHjjDGGFMoFAwA27p1K0tMTGQKhYIxxlhYWBjbvHkzu3XrFrt//77w+vbo0SPhtgcPHsw6dOjALly4wK5evcr69evHLC0t2bp163QuM3r0aHb58mV2//599tprrzEnJyeWlpZW5XP6448/Mrlczr7//nsWERHBVq9ezezs7Fj79u2Fy8ycOZPZ2Niw6dOns9u3b7Pbt2+z4uJiFhQUxObMmcNu3rzJ7t69y6ZNm8YCAgJYUVERY4yx48ePsx07drB79+6xu3fvsrlz5zI3NzeWnZ1d6fPyzz//MDs7O7Zt2zYWFRXFjhw5wnx8fNiqVauENQFgrq6u7Pvvv2dRUVE6z9nTSkpK2N69exkAFhERwRITE1lmZma1fnaMad6HTpw4waKjo9nx48dZQEAAW7hwoXB+bd+3GWMsJyeHTZ48mQ0bNowlJiayxMREVlRUxKKjoxkAFhgYyP744w8WERHBJk2axLy9vZlSqdS536d/n6v6Oxg5ciQbMmQIu3nzJouKimIHDx5kp0+fZowx4Zike/fu7NSpU+zOnTusT58+LCQkRFgz/969ceNGFhERwT799FMmFovZiRMndH4++/fvFx6ji4sLmzZtGrt9+zY7ePAg8/X1ZQDY9evXq3yO3nnnHZ3fR8YYe/jwIQPArl27VuX1t27dyuzt7au8HGP6OZai8T1CCCHExCIVOcjIV6KrTxO93m6BUoXW7x7W621W1933QmElrf5hRvv27fH2228D0OQbffTRR3B2dsb8+fMBAO+++y42bdqEmzdvlnt9pVKJzZs3w8/PDwCwaNEivPfee3V8FHXz7LPPIiUlBe+++y6SkpLQoUMHHDp0SAg/j42NrfHIwPjx47F582asWbMGixcvRkBAAPbu3YvevXsLl3n11VdRWFiIJUuWID09He3bt8fRo0eF54YQUnf16fXVy8sL69atA8dxCAgIwK1bt7Bu3TqEhoZi69atePToEWwcXWAplWDZsmU4dOgQtm7dKoz5KJVKfPXVV2jfvj0ATefl7t27cfToUWGHUV9fX+H+NmzYgI4dO+qMCX3//ffw8vLC/fv30apVKwBAu3btsHLlSgBAy5YtsWHDBhw/fhxDhgyBi4sLAMDBwUFn5Ll9+/bCOgDg/fffx/79+3HgwAEsWrQI4eHhOHbsGC5fvowuXboAAL799lu0bNlSuM6ZM2dw6dIlKBQKYWx57dq1+O2337Bnzx48//zzlT6f69evx9y5czF79mwAmvenI0eO6HQCAYC1tTW+/fZbYWzvxx9/hFqtxrfffguO4wAAW7duhYODA06dOoWhQ4di4MCBOrfxzTffwMHBAadPn8aoUaMqfF5Wr16N5cuXY+bMmcLP4/3338cbb7whPMcAMG3aNGHdlRGLxWjSRHNM4urqWiZfqLKfHQCdIHQfHx988MEHeOGFF/DVV18Jp9f2fdvGxgaWlpYoKioqdxx+2bJlGDlyJADN89KmTRtERkYiMDBQuN/Sv8+xsbHYunUrYmNj0bRpU+E2Sv8dxMbGYuLEiUIHUenfd96HH36Ifv36AQCWL1+OkSNHorCwEHK5HGvXrsWsWbPw4osvAgCWLl2KCxcuYO3atUJHWmk///wz1Go1vvvuO8jlcrRp0wbx8fFYuHBhlc8PACQlJZW70Qp/nrmhohQhhBBiQowxTP/uEhQ5RTj5Wn80d7Iy9ZJMol27dsL/F4vFcHJy0mkf5w+mFAoF7OzsylzfyspKp+ji4eEBhUJhwBVXDz82U57SowHl2bZtW7mnz5kzB3PmzKn0usuXL8fy5curs0RCSAPXo0cPoQgCAD179sSnn36KW7duQaVSISAgAGoGcAA4TjPS5+TkJFxeKpXqvEaHhYVBLBYLH8CfduPGDZw8eRI2NjZlzouKitIpSpVWndft3NxcrFq1Cn/++ScSExNRUlKCgoICITMvIiICEokEnTp1Eq7j7+8PR0dHnfXl5ubqPEYAKCgoQFRUVKX3z98HX1zgdevWDSdOnNA5LTg4WCdH6saNG4iMjIStra3O5QoLC4X7TU5Oxttvv41Tp05BoVBApVIhPz+/0kxA/rbPnj2rM5KnUqlQWFiI/Px8WFlpji34Ql1dVfWzO3bsGNasWYPw8HBkZ2ejpKSkzFoM9b5dem18zqJCoRCKUk//PvN/B/zvJa/038HixYuxcOFCHDlyBIMHD8bEiRPLPAcV3W/z5s1x7969MsXOXr164Ysvvij3Mdy7dw/t2rWDXC4XTitv9L+hoKIUIYQQYkKP0vKRmFUIALgWm6HXopSlhRh33wvV2+3V9L5r4umwVY7jdE7jP1Dx4d/VuT7TZqYQQoi+1afX14rk5uZCLBbj2D/nkZqnhIWEg5+LpmBSuqBkaWmpU9SytLSs8nZHjx6Njz/+uMx5/Id1oPzX7Ype43nLli3D0aNHsXbtWvj7+8PS0hKTJk2qNCenvPV5eHiU+8WAPnccs7a2LnO/nTt3xk8//VTmsnwH1MyZM5GWloYvvvgC3t7ekMlk6NmzZ5WPLzc3F6tXr8aECRPKnFe6sPH0mmqrsp9dTEwMRo0ahYULF+LDDz9EkyZNcObMGcydOxfFxcVCUcpQ79tVHTs8/fvM/x1cvXoVYrHu3xb/dzBv3jyEhobizz//xJEjR7BmzRp8+umnePnll6t9v8bk7u6OS5cu6ZzGb7xijhvFGLQoNWbMGISFhUGhUMDR0RGDBw/Gxx9/LLTFVYYxhhEjRuDQoUPYv38/bflICCGkQboRnyn8/1sJWRjX0VNvt81xXI1GPAghhFRPfXp9vXjxos5/X7hwAS1btkTHjh2hUqmQlJwM3+Au4DgOfk3tdD6wlyc4OBhqtRqnT58WxvdK69SpE/bu3QsfHx9IJLV/jiwsLKBSqXROO3v2LGbNmoXx48cD0BQUSgeFBwQEoKSkBNevX0fnzp0BAJGRkcjIyNBZX1JSEiQSiU5QeHUFBATg8uXLmDFjhnDa5cuXq7xep06dsGvXLri6upbb8cs/vq+++gojRowAoAlGLx0QD5T/vHTq1AkRERHw9/ev6cOpEN/l9fR9VeXq1atQq9X49NNPhRH13bt3621d/Npquq6K8H8HCoUCffr0qfByXl5eeOGFF/DCCy9gxYoV2LJli05RqjJBQUE4e/asMF4JaH7WT++UW/ryO3bsEMb/AM3fbXX17NkTH374IRQKBVxdXQEAR48ehZ2dXYX3aUoG3X1vwIAB2L17NyIiIrB3715ERUVh0qRJ1bru559/XuULIiGEEFLf3YzPEv7/rYSsSi5JCCGE1FxsbCyWLl2KiIgI/PLLL1i/fj1eeeUVtGrVCs899xwWvzAPx/4+iLhHMTh/4SLWrFmDP//8s8Lb8/HxwcyZMzFnzhz89ttviI6OxqlTp4TCw0svvYT09HRMnToVly9fRlRUFA4fPozZs2fXqJDg4+OD48ePIykpSSgqtWzZEvv27UNYWBhu3LiBadOm6XSjBAYGYvDgwXj++edx6dIlXL9+Hc8//7xOd8zgwYPRs2dPjBs3DkeOHEFMTAzOnTuHt956C1euXKlyXS+//DK+++47bN++HQ8ePMAHH3yAmzdvVvnZ9bnnnoOzszPGjh2Lf//9V3jeFi9ejPj4eOHx7dixA/fu3cPFixfx3HPPlelMK+95effdd/HDDz9g9erVuHPnDu7du4edO3cKWY214e3tDY7j8McffyAlJaVMZlZF/P39oVQqsX79ejx8+BA7duzA5s2ba72O8vj4+ODmzZuIiIhAamoqlEplrW+L/zuYMWMG9u3bh+joaFy6dEnn7+DVV1/F4cOHER0djWvXruHkyZMICgqq9n28/vrr2LZtGzZt2oQHDx7gs88+w759+7Bs2bJyLz9t2jRwHIf58+fj7t27+Ouvv7B27dpq39/QoUPRunVrTJ8+HTdu3MDhw4fx9ttv46WXXhJy1C5duoTAwEAkJCQI14uNjUVYWBhiY2OhUqmE3YCr+7OvLYMWpZYsWYIePXrA29sbISEhWL58OS5cuFDlL01YWBg+/fRTfP/994ZcHiGEEGJyN0t1St1JyIJaTSNnhBBC9GfGjBkoKChAt27d8NJLL+GVV14R8m22bt2Kcc9Mxafvv42x/bti4sQJuHz5Mpo3b17pbW7atAmTJk3Ciy++iMDAQMyfPx95eXkAgKZNm+Ls2bNQqVQYOnQogoOD8eqrr8LBwaFGmzt8+umnOHr0KLy8vNCxY0cAwGeffQZHR0eEhIRg9OjRCA0N1cmPAoAffvgBbm5u6Nu3L8aPH4/58+fD1tZW6DjhOA5//fUX+vbti9mzZ6NVq1aYMmUKHj16VCYcujzPPfccVqxYgWXLlqFTp06Ijo7GrFmzdMbkymNlZYV//vkHzZs3x4QJExAUFIS5c+eisLBQ6Jz67rvvkJGRgU6dOmH69OlYvHix0OlS2fMSGhqKP/74A0eOHEHXrl3Ro0cPrFu3Dt7e3tV7ssvh6ekpBKi7ublVmI/4tPbt2+Ozzz7Dxx9/jLZt2+Knn37CmjVrar2O8syfPx8BAQHo0qULXFxccPbs2Trd3tatWzFjxgy89tprCAgIwLhx43T+DlQqFV566SUEBQVh2LBhaNWqlU5oe1XGjRuHL774AmvXrkWbNm3w9ddfY+vWrejfv3+5l7exscHBgwdx69YtdOzYEW+99Va547AVEYvF+OOPPyAWi9GzZ0/85z//wYwZM3SC5PPz8xEREaFTm3n33XfRsWNHrFy5Erm5uejYsSM6duxYrWJtXXDMSIEL6enpWLhwIRISEnDmzJkKL5efn48uXbpgzZo1GDt2LDiOq3R8r6ioCEVFRcJ/Z2dnw8vLC1lZWRW2RRJCCCHmQKVmaLvyMAqUT745Pra0H/xdy4bDVqWwsBDR0dFo0aJFlQfGRKOy5yw7Oxv29vaN9niisT9+Qp5WX19j+/fvjw4dOuDzzz+v8DL3k3NQqH0f8nGyhp2lRYWXrY/i4+Ph5eWFY8eOYdCgQQa5jyFDhsDd3R07duwwyO0TYq70cSxl0E4pAHjzzTdhbW0NJycnxMbG4vfff6/08kuWLEFISAjGjh1brdtfs2YN7O3thX9eXl76WDYhhBBicJGKXBQoVbCSitHBywEAcJtG+AghhBhRierJ+JtSZZpgZn06ceIEDhw4gOjoaJw7dw5TpkyBj48P+vbtq5fbz8/Px2effYY7d+4gPDwcK1euxLFjx3Tygggh1VfjotTy5cvBcVyl/8LDw4XLv/7667h+/TqOHDkCsViMGTNmVJiqf+DAAZw4caLSSv7TVqxYgaysLOFfXFxcTR8SIYQQYhJ8yHlbT3u0b2YPgHKlCCGEGI+aMZSUGhtXqur/CLlSqcR///tftGnTBuPHj4eLiwtOnTpVZre3irRp0wY2Njbl/vvpp590xv86d+6MgwcPYu/eveWGvpuzF154ocLH+cILL5h0bRWty8bGBv/++69J12YOzPlnVxs13g7htddew6xZsyq9jK+vr/D/nZ2d4ezsjFatWiEoKAheXl64cOECevbsWeZ6J06cQFRUVJmtOCdOnIg+ffqUu22nTCYTwroIIYSQ+oTPk2rfzB6t3DTbcFNRihBCiL6U9/mptJKnOqMaQqdUaGgoQkNDa339v/76q8IMZDc3N1haWuLYsWO1vn1z8d5771UYtG3qse2wsLAKz/P01N8uxfWVOf/saqPGRSkXFxe4uLjU6s74nRFKZ0CVtnz5csybN0/ntODgYKxbtw6jR4+u1X0SQggh5orfea9dMwehKHX3cTbUagaRiHagJYQQYlhPd0Y1hKJUXdUlHLw+cXV1LROibi78/f1NvQSzZs4/u9qocVGqui5evIjLly+jd+/ecHR0RFRUFN555x34+fkJXVIJCQkYNGgQfvjhB3Tr1g3u7u5wd3cvc1vNmzdHixYtDLVUQgghxOiKSlS4l5gNAGjfzAFNHeSQW4iQW1SC6LQ8+LnUPOycEEIIqYkSbdMAB4ChYYzvEULqF4MFnVtZWWHfvn0YNGgQAgICMHfuXLRr1w6nT58Wxu2USiUiIiKQn59vqGUQQgghZik8MQdKFYOjlQW8mlhCIhahtYem5bouYed8VzKpGj1XhJCaMtLG5UbDF6FkFmLtf6sb3GMkhBiOPl4vDNYpFRwcjBMnTlR6GR8fnyofBL0oEkIIaYj4PKngZg7gOM2oXrCnPa7FZuJWfBbGdqhZZoJUKoVIJMLjx4/h4uICqVQq3C7RxRhDcXExUlJSIBKJIJVKTb0kQoiZE4s1RZvi4mJYWlqaeDX6w2dKWVmIUahUQc0YVIxBQu8fhJBq4BuMqruRQHkMVpQihBBCSMVuaPOk+F33AM0ufABwsxadUiKRCC1atEBiYiIeP36sn0U2cFZWVmjevDlEIoM1jhNCGgiJRAIrKyukpKTAwsKiwbxuFBQUgpUowak5cGol1GqGvLwCoXOKEELKwxhDfn4+FAoFHBwchMJ9bVBRihBCCDEBvlOqXTMH4bRgbYGqtmHnUqkUzZs3R0lJCVQqlb6W2iCJxWJIJBLqJiOEVAvHcfDw8EB0dDQePXpk6uXoTWpuEQqVaiitLJBbVAKlikGdLYWcilKEkGpwcHAoNxe8JqgoRQghhBhZXlEJIhW5AHQ7pfxdbOocds5xHCwsLOrURk0IIaQsqVSKli1bori42NRL0ZsPt19GdGoe1kxoh9/vxuNSdDpeG9oKI1o1NfXSCCFmzsLCok4dUjwqShFCCCFGdjshC2oGuNvJ4WonF06XiEUI8rDD9dhM3E7Ioh34CCHEzIhEIsjl8qovWE/cVRQiPU8FV0cbyOVyJOSoEJtV0qAeIyHEvDWMYWhCCCGkHrmpzZNqV6pLiheszZW6FV/7HfgIIYSQqhSXqJGep+n6crWVw91OE+CelFVoymURQhoZKkoRQgghRnZDmyfV3suhzHlCUaoWYeeEEEJIdaXkFgEALMQcHK0s4GGv6Y5KpKIUIcSIqChFCCGEGFmlnVLa0+5ow84JIYQQQ0jO1hSfXG3l4DgO7tqiFHVKEUKMiYpShBBCiBFl5BUjNj0fANDO06HM+aXDzmPS8oy8OkIIIY2FIlvTKeVqJwOAUp1SBSZbEyGk8aGiFCGEEGJE/Fiet5MV7K3K7pDHh52XviwhhBCib4ocTUeUm62mGMV3SmUXliCvqMRk6yKENC5UlCKEEEKM6KY2T6pdM4cKL0Nh54QQQgzt6U4pW7kFbGSazdmTsmmEjxBiHFSUIoQQQozohrbQ1L6cPCleWwo7J4QQYmB8ppSbnVw4zU1boKJcKUKIsVBRihBCCDGimnRKUdg5IYQQQ1HkaDqlXGxlwmke9pYAaAc+QojxUFGKEEIIMZLk7EIkZxdBxAFtPe0qvFxLVxvIJBR2TgghxHDK65Tic6WSaXyPEGIkVJQihBBCjORGXCYAoKWrLaykkgovJxGL0LophZ0TQggxnBRtp5SrTqcU7cBHCDEuKkoRQgghRnJTmyfVrpI8KR4/wnebilKEEEL0rLhEjbS8YgDld0pRphQhxFioKEUIIYQYyQ0+T8rLocrLUtg5IYQQQ0nN1XRJWYg5OFpZCKc/6ZSiohQhxDioKEUIIYQYAWNMKDBVtvMeTwg7T6Cwc0IIIfrFZ0a52srBcZxwurudJuicOqUIIcZCRSlCCCHECGLT85GZr4RULEKge8Uh5zw+7DyHws4JIYToWXJ22Z33gCedUml5xShUqoy+LkJI40NFKUIIIcQIbmjzpII8bCGVVP32KxGLEORBYeeEEEL0LyWH33lPtyjlYGUBmfY9SqEtXBFCiCFRUYoQQggxgpvanffaNXOo9nUo7JwQQogh8J1SrrZyndM5jqMd+AghRkVFKUIIIcQIarLzHi+Yws4JIYQYgKKCTimg1A582fUvV4oxhkhFLlSUxUhIvUFFKUIIIcTAVGqG24+1IefV2HmP15bCzgkhhBhARZ1SAOBuV3934Pv1SjwGf3Yam09HmXophJBqoqIUIYQQYmCRilzkF6tgJRXDz8Wm2tdr6fYk7PxRer4BV0gIIaQxUeRoi1LldkrV3x34rmtH5S/HpJt2IYSQaqOiFCGEEGJgN+IzAWg6n8QirvILl2JBYeeEEEIMQKEdzSuvU4rPlKqPRan4DM0XOFEpuSZeCSGkuqgoRQghhBjYLW2eVPsa5EnxKOycEEKIPilVaqTlFQOoPFMqsR5mSsVpu4rjMwpQqFSZeDWEkOqgohQhhBBiYDe1nVLBNdh5jyeEncdTUYoQQkjdpWhH9yQiDo5W0jLnP+mUql+776nVDAmZmjUzBkSn5pl4RYSQ6qCiFCGEEGJAxSVq3EvMAVC7Tqm2pTqlKOycEEJIXQl5UrYyiMoZKec7pRQ5RVCq1EZdW10k5xRCqXryPhmpoBE+QuoDKkoRQgghBhSelI1ilRoOVhZo3sSqxtdv6WYDKYWdE0II0ZNkPk/KrmyeFAA4W8sgEXFg7ElXVX0Ql67b2UW5UoTUD1SUIoQQQgzohnbsLtjTHhxX/ZBzHoWdE0II0afSnVLlEYk4uGkLVon1KOw87qkvbqJSaHyPkPqAilKEEEKIAd3Ubk/dvhZ5UrxgT01RisLOCSGE1BW/855bBZ1SQP3cgS8+Q9Mp5WyjKbZF0fgeIfUCFaUIIYQQA7qp7ZRqV4s8KV47TwcAFHZOCCGk7hTZlXdKAaV24KtHYedxGZpOqX6tXAAAD1NzKYuRkHqAilKEEEKIgeQXl+CBQhty7uVQ69sRws4fZ4ExOsAmhBBSe8k5DbNTih/fC/FzglQsQqFSLezGRwgxX1SUIoQQQgzkdkI21Axws5NVevBfFSHsvLAEj9Io7JwQQkjt8Z1SLnYVd0oJmVLZ9acoxY/v+ThbwcdZs7EIhZ0TYv6oKEUIIYQYyM34TABAuzrkSQEUdk4IIUR/FHynlG1lnVKWAIDketIppVSphVFDL0cr+LnYAKCwc0LqAypKEUIIIQbC77zXvg55UjwKOyeEEFJXSpUaqbnFAADXSjqlnmRK1Y+iVGJmIdQMkEpEcLaRwd9VU5SKpLBzQsweFaUIIYQQA9FXpxQABGtzpW5S2DkhhJBaSs3VjO5JRByaWEkrvByfKZWcXVgvwsL5kPNmjpYQibhSnVJUlCLE3FFRqhErVKoo/I8QQgwkM79YyH+qy857PAo7J4QQUlfJfJ6UrQwiEVfh5VxsZRBxQImaITWvyFjLq7V4bVHKy1GTJcUXpR5SUYoQs0dFqUbs/T/uovfHJ3A+Ks3USyGEkAaH72jydrKCQyXfRldXKzdbCjsnhBBSJwptcLlrFZtvWIhFcLHVjPfVhx344tI1X7Q3c9RkYfm6WAMAUnOLkZlfbLJ1EUKqRkWpRuyfBylgDPjteoKpl0IIIQ2OPkf3AG3YubstAAo7J4QQUjvJOZquJ1fbivOkeO7asPP6kCvFj+95NdF0SlnLJGiqHUGkET5CzBsVpRqp3KIS4RuFkxGKejErTggh9Yk+Q855wdrborDz+mPjxo3w8fGBXC5H9+7dcenSpWpdb+fOneA4DuPGjStz3r179zBmzBjY29vD2toaXbt2RWxsbJnLMcYwfPhwcByH3377rY6PhBDSEKRoO6XcKgk553lou6nqQ6dUfMaTnfd4ftqw8ygF7cBHiDmjolQNXXyYhnVH7+Piw/o98haRlCP8f0VOEe48zjbhagghpOG5pS1K6atTCngSdk6dUvXDrl27sHTpUqxcuRLXrl1D+/btERoaCoVCUen1YmJisGzZMvTp06fMeVFRUejduzcCAwNx6tQp3Lx5E++88w7k8rKjOJ9//jk4ruLMGEJI48NnSrnaVj6+B9SvHfji0vlOKUvhNAo7J6R+oKJUDR28+RhfHH+AE+GVH1Cau/Ak3SJUfX88hBBiThTZhUjKLoSIA9o0tdPb7Qph5wkUdl4ffPbZZ5g/fz5mz56N1q1bY/PmzbCyssL3339f4XVUKhWee+45rF69Gr6+vmXOf+uttzBixAj873//Q8eOHeHn54cxY8bA1dVV53JhYWH49NNPK70vQkjjo8ipQaeUPd8pZd4bIxUqVVBoxxKble6U0uZKRSqoKEWIOTNYUWrMmDFo3rw55HI5PDw8MH36dDx+/LjK650/fx4DBw6EtbU17Ozs0LdvXxQUmM8LYZCH5sPF3cT63VkUnqjplOLnyU9EUFGKEEL0hR/d83e1gbVMorfb5cPOswtLEJtOYefmrLi4GFevXsXgwYOF00QiEQYPHozz589XeL333nsPrq6umDt3bpnz1Go1/vzzT7Rq1QqhoaFwdXVF9+7dy4zm5efnY9q0adi4cSPc3d319pgIIfVfQ+yU4kf3rKViOFpZCKcL43vUKUWIWTNYUWrAgAHYvXs3IiIisHfvXkRFRWHSpEmVXuf8+fMYNmwYhg4dikuXLuHy5ctYtGgRRCLzaejii1Lhpcbf6iN+fG9u7xYAgBtxmUjJMf/tXgkhpD7Qd8g5j8LO64/U1FSoVCq4ubnpnO7m5oakpKRyr3PmzBl899132LJlS7nnKxQK5Obm4qOPPsKwYcNw5MgRjB8/HhMmTMDp06eFyy1ZsgQhISEYO3ZstddbVFSE7OxsnX+EkIaH7yhyrUanlDufKZVt7kWpJyHnpUeW/bXje7Hp+SgqUZlkbYSQqunv69unLFmyRPj/3t7eWL58OcaNGwelUgkLC4sKr7N48WIsX75cOC0gIMBQS6yVQHdbcByQklOE1NwiONtU/YJubhhjuKcd3+vT0gV/3EzErYQsnIpQ4JkuXiZeHSGE1H+GCDnntfW0x434LNyKz8Kodk31fvvENHJycjB9+nRs2bIFzs7O5V5GrVYDAMaOHSscZ3Xo0AHnzp3D5s2b0a9fPxw4cAAnTpzA9evXa3T/a9aswerVq+v2IAghZq1EpUZaXvU7pTy0u+8lZRWCMWa2GXVx2k6pZo6WOqe72MpgK5Mgp6gEj9Ly0crN1hTLI4RUwSgtSOnp6fjpp58QEhJSYUFKoVDg4sWLcHV1RUhICNzc3NCvXz+cOXPGGEusNiupBD5Omvnke/V0hC8xqxA5hSWQiDj4uVpjYKAmh+IkjfARQkidMcYM1ikFUNh5feHs7AyxWIzk5GSd05OTk8sdqYuKikJMTAxGjx4NiUQCiUSCH374AQcOHIBEIkFUVBScnZ0hkUjQunVrnesGBQUJu++dOHECUVFRcHBwEG4HACZOnIj+/ftXuN4VK1YgKytL+BcXF1fHZ4AQYm5Sc4vBGCAWcXCyllZ5eb6bqqhEjcx8paGXV2vx2nH20nlSAMBxHHy1I3yUK0WI+TJoUerNN9+EtbU1nJycEBsbi99//73Cyz58+BAAsGrVKsyfPx+HDh1Cp06dMGjQIDx48KDC65mi3TxQOzpRX4tSfMi5r4s1ZBKxUJT6534qikvUplwaIYTUe3HpBcjMV8JCzCHQQ//fylLYef0glUrRuXNnHD9+XDhNrVbj+PHj6NmzZ5nLBwYG4tatWwgLCxP+jRkzBgMGDEBYWBi8vLwglUrRtWtXRERE6Fz3/v378Pb2BgAsX74cN2/e1LkdAFi3bh22bt1a4XplMhns7Ox0/hFCGpZk7Rieq60MIlHVXU9yC7FQvDLnXCk+U8qriVWZ8/gRvigqShFitmpUlFq+fDk4jqv0X3h4uHD5119/HdevX8eRI0cgFosxY8aMCg+g+Zb0BQsWYPbs2ejYsSPWrVuHgICASneOWbNmDezt7YV/Xl6GHz/jc6XuJdbPXCk+DyvQXfM4gj3t4WwjQ25RCa7EpJtyaYQQUu/d0HZJBXnYQSYR6/32W7nZQiqmsPP6YOnSpdiyZQu2b9+Oe/fuYeHChcjLy8Ps2bMBADNmzMCKFSsAAHK5HG3bttX55+DgAFtbW7Rt2xZSqeaD4euvv45du3Zhy5YtiIyMxIYNG3Dw4EG8+OKLAAB3d/cytwMAzZs3R4sWLUzwLBBCzIWQJ2Vb/fgRPuw8Kdt8Np56WlwG3yllWeY8P1fNhAuFnRNivmqUKfXaa69h1qxZlV6m9PbFzs7OcHZ2RqtWrRAUFAQvLy9cuHCh3G8IPTw8AKDSlvTyrFixAkuXLhX+Ozs72+CFqSdFqXraKaUtpgVoO75EIg4DAlzw69V4nAhXIMS//CwLQgghVXsyuqf/PCkAkEpECPKw1eRKJWTBWztSTszPs88+i5SUFLz77rtISkpChw4dcOjQISH8PDY2tsabuYwfPx6bN2/GmjVrsHjxYgQEBGDv3r3o3bu3IR4CIaQBETql7KrOk+J52Mtx53G2WXdKxWm/oPFyLNsp5cd3SqXkGXVNhJDqq1FRysXFBS4uLrW6I74Tqqio/B3efHx80LRp03Jb0ocPH17h7cpkMshkxg0bD9KOY0Sl5KK4RA2pxHx2B6wOfnwvqNRYycBAV6Eo9fao1hVdlRBCSBX4kHND5EnxhLDzBAo7N3eLFi3CokWLyj3v1KlTlV5327Zt5Z4+Z84czJkzp9proDFPQghQx04pMy1K5RaVIEObd+XVpJxOKaEolQu1mlVrbJEQYlwGqaZcvHgRGzZsQFhYGB49eoQTJ05g6tSp8PPzE7qkEhISEBgYiEuXLgHQBNG9/vrr+PLLL7Fnzx5ERkbinXfeQXh4OObOnWuIZdaap4Ml7OQSKFWs3oXmFZWo8FD7TQE/vgcAvVs6w0LM4WFqHqJT6ZsEQgipDZWa4XYCv/Oeg8HuJ7hUrhQhhBBSHQptp5RbjTqlNIUec+2UiteO7jlYWcBWXnZDLW8nK0hEHPKLVUjKNs/HQEhjZ5CilJWVFfbt24dBgwYhICAAc+fORbt27XD69Gmhq0mpVCIiIgL5+U/yMF599VWsWLECS5YsQfv27XH8+HEcPXoUfn5+hlhmrXEch8B6OsIXpchDiZrBVi6Bh/2TNyRbuQW6tWgCADgRTrvwEUJIbUSl5CK/WAUrqRj+2h1/DIEPO78VT2HnhBBCqqdWnVJ25t0pFZeuyboqL08KACzEIng7acb6KFeKEPNUo/G96goODsaJEycqvYyPj0+5B9LLly/H8uXLDbEsvQpyt8Wl6PR6V5SKSNaO7rnbgeN021cHBLjibGQaToYrMLc3haESYmwHbzxGZn4x/tPDu8zfJ6kfbsRlAgDaNrWH2IAjAk+HnVOuFCGEkKok16pTSnPZxCzzDDqvLE+K5+dig6iUPEQpctGnZe2iaAghhlO/wpDMiBB2nlS/ilJPh5yXNjDQFQBwMToNuUUlRl0XIY1dcnYhXtl5He/8fgcfH4qo+grELN0U8qQME3LOk0pECNTmAt6iET5CCCHVkJyt6ZRyaUCZUvEZmmKZV5NKilLazuVI6pQixCxRUaqWnuzAl1OvRifuJWmKUoEeZYtSvi42aOFsDaWK4cyDFGMvjZBG7bfrCVBrX0o2n47CxpORpl0QqRVh5z0vB4PflzDCR0UpQgghVShRqZGWpylK1aRTii9K5RWrkFOoNMja6iJOmylV0fgeAPjzYecKys0lxBxRUaqWAtxtIeKA9LxipOSUv6OgOYrQdnaVDjkvbUCApluKcqUIMR7GGPZdSwAAdNdmu31yOAI7zseYcFWkpopL1Lin7UZtb+BOKYDCzgkhhFRfam4xGAPEIg5O1tJqX89KKoGdXJP4Yo7dUtUa33N9sgMfIcT8UFGqluQWYrRw1mR43K0nuVIZecVC225543vAkxG+kxEpUKvrTwcYIfXZ3cRsRCTnQCoR4ZsZXbB4oD8A4J3f72D/9XgTr45UV0RSDopVajhYWaB5JWME+vKkKJVdrzp2CSGEGJ8iR1NQcrGRQVTDzENz3YGPMYYEYXyv4k4pXxfNZzZFThGyzbDbi5DGjopSdVB6hK8+CNeO7nk1sYSNrPyM+24tmsBaKkZKThFuP6Zv3wkxBr5LakiQG+wtLbBkSCvMCvEBACz79SaO3Eky4epIdd3Qju4Fe9obJaieDzvPKlAKuw8RQggh5eG/mHa1q36eFM9cc6WyCpTI0ebgNqukU8pObgE37eOOUlC3FCHmhopSdfCkKFU/OqXCtaN7AW7lj+4BmvBcflcKGuEjxPBKVGr8HvYYADC+oycAgOM4vDuqNSZ2agaVmmHRz9dxNjLVlMsk1SDkSRlhdA+gsHNCCCHVx3dKudpWP0+K92QHPvMqSvFfyDjbyCC3EFd6WT8+VyqFcqUIMTdUlKqDIO2HgXpTlNJ2dAWVE3Je2sAg7QgfFaUIMbh/I1ORmluEJtZS9At4sk2xSMTh44nBCG3jhmKVGvN/uILrsRkmXCmpypOd9xyMdp8Udk4IIaQ69NIplW1eXbl8yHllo3u8J0Up6pQixNxQUaoO+E6ph6l5KFSqTLyaqoUna3feqyDknNdf+8H4RnyW8K0KIcQw9mtH98a0bwoLse5LskQswpdTO6JPS2fkF6swa+tloeORmJf84hLcT+ZDzh2Mdr8Udk4IIaQ6UrTH9G4NqFMqPqPqkHOenzZXisb3CDE/VJSqA3c7ORysLKBSM0Sa+QucWs1wX5spVVHIOc/VVi6Mn5yKSDH42ghprHIKlTiszYua0Mmz3MvIJGJ8Pb0zOjV3QFaBEv/59hJiUqn13NzceZwNNQNcbWXCN8rGEFyqU4rCzgkhhFSkbp1Smk4kc8uU4sf3mjlW3Snl76r5/BNJnVKEmB0qStUBx3EI0nYdmfsOfLHp+ShQqiCTiODjVPW3CcIufDTCR4jB/H07CUUlavi72gjFhfJYSSXYOrsbgjzskJpbhOe+vYjELPNqoW/sbsRlAjDu6B5AYeeEEEKqh59+cKtFUcpcO6WejO9Vo1PKVdMpFZuWD6VKbdB1EUJqhopSdVRfws75kZ+WbjaQiKv+sfNFqX8fpKK4hF64CTGEfdfiAWgCzqvarc3e0gI/zOmGFs7WSMgswH++vYi03CJjLJNUA58n1d5IIec8qUQkdL9SrhQhhJCKCJ1StRjf4zuAswqUyC8u0eu66iI+Q/NlTHXG99zt5LCSilGiZniUlm/opRFCaoCKUnUUWE/Czu8lVi9Pite2qT2cbWTILSrB5Zh0Qy6NkEYpPiMfFx6mg+OAcR3LH917moutDD/O646m9nJEpeRhxveXkF2oNPBKSXUIO+95ORj9vinsnBBCSGVKVGrhi6zajO/ZyiSwlmp2tzOXET7G2JNMqWoEnXMcR2HnhJgpKkrVUWuhUyrHrPM8IpL4olTleVI8kYjDwEBN4PnxezTCR4i+/R72GADQo4UTPB2qPpjieTpY4sd53eFkLcWdx9mYu+0yCorNf6OFhiwrX4kY7beu7SoZwzQUPgOQws4JIYSUJy2vGGoGiDjAybrmRSmO4+Am7MBnHkWplNwiFCrV4DjAw756x1H+rpqilLlnARPS2FBRqo78XW0gFnHIKlCazYt0efjxvep2SgGlcqUiqChFiD4xxrBXO7pXUcB5ZXxdbPDD3G6wlUtwOSYDL/x4lcZsTehmQiYAoHkTKzhaS41+/xR2TgghpDIK7eiei60MYlHlcQEV4XOlzKVTis9R9LCTQyqp3kdaYQc+6pQixKxQUaqO5BZi4QXOXEf48otL8Chd8y0+P25YHb1busBCzCE6NQ8P6cWbEL25GZ+Fhyl5kFuIMDzYo1a30aapPbbO6gpLCzFO30/Bkl1hUKmpIGEKfJ5UOyPnSfFKh53z+RqEEEIILzmbDzmv/e6w7naabiRzCTvnR/eaVSPknPdkfI92MSbEnFBRSg+CSo3wmaP7yblgDHC2kcLZpvotuzYyCbq3cAIAnKBd+AjRGz7gPLSNO2xkklrfThefJvh6emdIxSL8eSsR/913izplTIDfea+9kXfe41HYOSGEkMoocviQ85qP7vHMrVOK/xKmmWP1IxD8tON7DxW5dLxEiBmhopQe8EWpu2baKRWeWPPRPd4AGuEjRK+KS9Q4eDMRADChU7M6317fVi74cmoHiDhg15U4fPDnPTrQMjJTd0oBFHZOCCGkYnynlGtdOqW0RSlz6ZSK006BVGfnPZ63kxXEIg45RSVCoY4QYnpUlNIDPjzcXMf3wmsYcl4anyt1KTodObTLFyF1dvp+CtLziuFiK0MvPye93Oawth74eGI7AMB3Z6Kx/kSkXm6XVE2RXYik7EKIuCeFIVMQcqXiqShFCCFEl147pbLNY0yc75TyqsH4nkwiRnPt5aMo7JwQs0FFKT3gd+CLSc0zy12w+JDzgFoUpVo4W8PX2RpKFcOZB6n6XhohjQ4/ujeuQ1NIxPp7CX6mixdWjm4NAPjs6H18fyZab7dNKnZDWwTyd7WBdR1GMeuKws4JIYRURKGPTCkzG9+L4zOlajC+B1DYOSHmiIpSeuBiK4OTtRRqBtxPNq9cKcYYIrSdUvyYYU3xI3yUK0VI3WTlK3H8nubvSB+je0+b3asFlg5pBQB474+7+PVKnN7vg+i6FZ8JAGhnojwpXit3G1iIOQo7J4QQUkZyjnZ8r06dUpriT2puMYpKTPslvErN8Diz5p1SAIWdE2KOqCilBxzHlQo7N68RPkVOETLylRBxmm/ya2OgkCuVAjXt7kVIrf1x6zGKVWoEutvWukhclZcH+mNe7xYAgDf33sSh24kGuR+iwXdKtTdhnhSgGUngcwMpV4oQQkhpimzN+F5dOqUcrSwglYh0bs9UkrMLoVQxWIg5uNfwMfFh55E0vkeI2aCilJ4EeZhnrhSfJ9XC2RpyC3GtbqOrTxPYyCRIzS2iDzuE1MH+awkAgIkG6JLicRyHt0YG4dkuXlAz4OVfruOf+ykGu7/GjDGGm9pOqWATd0oBFHZOCCGkLJWaITW37plSHMcJuVKmDjvnQ86bOlhCLOJqdN0nnVJUlCLEXFBRSk+edEqZ1/iesPNeHboypBIR+rR0BkAjfITU1qO0PFx5lAERB4zt0NSg98VxHP5vQjBGBntAqWJYsOMqrsSkG/Q+G6P4jAJk5CthIeaELyZMic+Vuk1FKUIIIVppuUVQM0DEAU42tS9KARC6khKzTDsmHqcdU69pnhTwJFMqMasQuUUlel0XIaR2qCilJ/zYxL2kbLMKmRV23nOr2wemgZQrRUid7NN2SfVu6VKnLZmrSyzisO7ZDujXygUFShVmb7uMO4+pWKFPN7RdUoHudpBJateJqk8Udk4IIeRpydpRO2cbWY27ip7Gh50nZ5tHp5SXY83ypADAwUoKZxspACCacqUIMQtUlNITf1dNyGxOYYlZhcwKRak65tf0D9AUpW4lZAk7eBBCqocxhv3XNUWpCR09jXa/UokIm//TGd18miCnsAQzvruEh9Surjc3tXlS7UycJ8Xjw84z8ynsnBBCiIYip+477/HczWR8j3+Pq2nIOY8f4YtMMa8JF0IaKypK6YlUIhJe4PhCkKkpVWpEKrRFKfe6dUq52MqEIN9TEZRPQ0hNXH2Ugdj0fFhLxRjaxs2o920pFePbWV3Q1tMOaXnF+M+3F5GQSQULfbgRlwkAaG8GeVKAJuw8QPtaTyN8hBBCgCedUnXJk+J5aAtbSabOlMrQdErVZnwPeBJ2HqWgTilCzAEVpfSotZntwBedmgelisFGJoGnQ+1etEsbGKj5MH08PLnOt0VIY7JP2yU1PNgDVlKJ0e/fTm6B7bO7wc/FGo+zCvGfby8K35yS2lGpmVD4aedlHp1SwJMRvptUlCKE1IIiu5B2Wm5g+Pd7fUQHuNtrPk+YvFMqnS9K1a1TisLOCTEPVJTSoyAzK0rx62jlZgNRHWfIgSe5UmcepKKoRFXn2yOkMShUqvDHjccAjDu69zQnGxl+nNcdng6WiE7Nw6RN5xGdSt8Q1tbDlFzkFatgaSGGv/bg1hy0pbBzQkgtnY1MRbf/O47/7r9l6qUQPdJrp5S96TulikvUSNJGiXg1qWWnlDbsnIpShJgHKkrpkbkVpfSVJ8Vr09QOLrYy5BWrcDk6Qy+3SUhDdzJcgezCEjS1l6OHr5NJ1+Jhb4lf5veAt5MVYtPzMXHTOYRpR9BIzdzQ5km19bSDRGw+b6UUdk4Iqa2D2i9Qdl6Ow614Kmw3FCl6zJTii1KKnEKUqNR1vr3aSMwqgJoBMokILrXcTdBfO74XnZpnssdBCHnCfI6kG4BA7Zbgj9LzkWcGW4xGaItSQXXMk+KJRBwGagPPaYSPkOrZq911b2xHT710LNZVcycr7F0YgmBPe6TnFWPqNxdwMoJ21aypm9qd99qZSZ4UL8DdlsLOCSG1cjYqVfj/a/6+R4XtBkKfnVJONjJIRBzUDEjJLarz7dVGXLrmva2ZoyU4rnbHVU3tLSG3EEGpYoij90pCTI6KUnrkbCODi60MjAERyaYPOw/XdmwFuOunUwoABmhH+E6G04dYQqqSlluEU9qCjylH957mbCPDzud7oG8rFxQoVZi3/Qp+vRJn6mXVKzfMbOc9HoWdE0JqIzYtH3HpBZCIOEjFIpyLSsPp+7SxTUOgz933xCJOuB1T5UrxIee13XkP0HzR7uvMh53TCB8hpkZFKT0zlxG+rAIlHmvfLAL01CkFAL1bOsNCzCEmLZ+2liekCn/cTESJmqFdM3u0dNPf36E+WMsk+G5mF0zo5AmVmuH1PTex8WQkfTNeDcUlatx7rHmNN5ed90orPcJHCCHVwXdJdWruiBk9vQEAH/0dDhWFntdrKjVDSo62U8qu7p1SAOBu4lypeL4oVcuQc56wAx99niHE5KgopWdB2hE+Uxel+NG9pvZy2Fta6O12bWQSIRfnBHVLEVKpfdfiAQDjzahLqjQLsQifPtMeC/v7AQA+ORyBd3+/Qx9CqhCRlINilRr2lhbwdqrbQbEhtKWiFCGkhs5EaopSvfyd8dIAf9jKJQhPysF+7e6xpH5KyyuCmgEiDnCylurlNvmilMk6pUqN79UFv0lJJHVKEWJyVJTSs9ZCp5Rpx/fCkzRFMX2FnJc2QJsrRUUpQioWqcjFjfgsSEQcRrdvaurlVIjjOLw5LBCrRrcGxwE7LjzCSz9dQ6GSdtisyA0hT8q+1nkWhhRcagc+6nwjhFRFrWY4H5UGAOjl7wRHayleGuAPAPjsSAS9H9RjCm2elLONTG+bcnhox/eSs+vv+B4A+LnSDnyEmAsqSulZoDa/KTwxG2oTdhsIO+/pcXSPN1CbK3UpOh05hUq93z4hDcH+65ouqX6tXOBcy91hjGlWrxbYMLUTpGIRDt1JwozvLiErn/6+y3OzVFHKHPFh5xn5SiRkUoArIaRy95KykZ5XDGupGO29HAAAs0J84GEvx+OsQmw7F2PS9ZHa4/Ok9DW6B5i+U4rfxKPO43su/PheHn2BQ4iJUVFKz3xdrCEVi5BXrBIq+abwJORc/0UpH2dr+LpYo0TN8O+D1KqvQEgjo1Yz/HZds7X2hE7NTLya6hvZzgPb53SDrVyCSzHpeObrc0jMoqLG024KIecOpl1IBWQSMVppM8xoW3dCSFXORWq6pLr7OsFC200jtxDjtaEBAICNJyORkVdssvWR2uN33nOzrXvIOe9JppTxjw8KlSohI6uu43stnK3BcZoc3jT6/SbEpKgopWcWYhFaumkq76Ya4VOrGe4na1pRgwwwvgcAA2mEj5AKXYxOR0JmAWzlEgwKcjX1cmqkp58Tfn2hJ9zsZLifnIsJX53DfTPYTdRcFBSr8ECbP2GOIec8CjsnhFRX6Typ0sZ39ESguy1yCkuw8WSkKZZG6ogfsdNnp5SHCTul+JBzG5kEDlZ1y8yVW4iFbivKlSLEtKgoZQCm3oEvIbMAuUUlsBBzaOFsbZD74Ef4TkUoTDqmSIg54gPOR7XzgNxCbOLV1Fygux32vdgL/q42SMwqxKRN53ApOt3UyzILdx5nQaVmcLGVwU2PB/n6FtyMilKEkKoVl6iF1/de/k4654lFHJYPDwQA/HD+EeLSTTcBQGpHwe+8p9dOKU2HUnJ2odE/A5QOOddHpqOfC+VKEWIOqChlAKYuSvH36+9qK7Rh61sXnyawlUmQmluMm/ShhxBBQbEKf99OAlC/Rvee5ulgiT0v9ERnb0dkF5bgP99dxKHbiaZelsnd0I7DtTfTkHMehZ0TQqrjemwGCpQqONtIEeBWNvKhXysX9PJ3QrFKjU+PRJhghaQuFAbolHK1lYHjAKWKGX3sLV5PIec8IVdKkaeX2yOE1A4VpQwgyEPzpn4vyTRFqQhtyHmQAfKkeFKJCH1aadq8aYSPkCeO3E1CblEJvJpYoou3o6mXUycOVlL8NK87hrZ2Q3GJGgt/uoYd52NMvSyTehJy7mDSdVSFws4JIdVxVrvrXoifc7mFdo7jsHxYEADgt7DHuE1fRNYrfKeUPjOlLMQiuGg3cEky8ghfXMaTTil98HPVFKUiqVOKEJOiopQBBGl34ItLLzDJ7nT8znuGCDkvbWCgGwDgRHiyQe+HkPpk//UEAMD4js3MupOmuuQWYmz6T2dM694cjAHv/H4HnxwOb7TdN09Czs1z5z1e6bBz+hBJCKnIWW2eVO+n8qRKC25mj7EdmgIA1vx9r9G+/tdHhsiUAkrnShn3Sw9+hLSuO+/x/F35TikqShFiSlSUMgBHaync7TQv1nyByJjCtR1agQYKOef1D3ABxwG3E7KF9mBCGjNFTiH+uZ8CQBMQ21CIRRw+HNcWS4e0AgBsPBmFN/bchFKlNvHKjCurQInoVE2Lv7l3SgEUdk4IqVxOoRJhcZkAgJCn8qSetmxoAKRiEc5GpuEf2nm5XlCpGVJzNeN1bnb665QCSu3AZ+Tj/3htp5S+x/cSMgtQUKzSy20SQmqOilIGwo/whRs5V6pQqRI+NAUauFPK2UYmfDA7GUEjfIQcCHsMNQM6NXcw2CYDpsJxHBYPaomPJgRDLOLw69V4PP/DFeQXl5h6aUZzS9sl5dXEEk2spSZeTdXaCkUp04ySE0LM26XodKjUDD5OVmhWReeJVxMrTO/pDQD46O9wqGiTG7OXllcElZqB4wAnPb9neWjDzo29A1+cNlNKX+N7TaylcNTu4vcwlbqlCDEVgxWlxowZg+bNm0Mul8PDwwPTp0/H48ePK71OUlISpk+fDnd3d1hbW6NTp07Yu3evoZZoUHzY+d1E43ZKPUjOhZoBjlYWcLU1/M5Qg7S78B2/R0UpYhq3E7Lw88VYlJhB186+a9rRvXoccF6VKd2a45vpnSG3EOFkRAqmbrmItNwiUy/LKG7UkzwpntApFZ9J4zaEkDLOaEf3QioZ3Stt0QB/2MoluJeYjd+0o+rEfCmyNe/NzjYySPS88ZHQKWXEolROoRKZ+ZpYFH11SgFPuqUiaYSPEJMxWFFqwIAB2L17NyIiIrB3715ERUVh0qRJlV5nxowZiIiIwIEDB3Dr1i1MmDABkydPxvXr1w21TIMx1Q58wuieu51R8mwGaotSZyJTUVRCba/EuFRqhgU7ruK/+2/h1V1hJi1MhSdl425iNizEHEa38zDZOoxhUJAbfp7fA45WFrgRl4lJm88jNq3hbxXOh5y3N/M8KV6Auy0kIgo7J4SU71ykJuS8sjyp0hytpXixvz8A4NMjEShU0nGfOVPkaPOkDPAltYcJilJx6Zr3MUcrC9jIJHq7XSFXKoV24CPEVAxWlFqyZAl69OgBb29vhISEYPny5bhw4QKUyoqDv8+dO4eXX34Z3bp1g6+vL95++204ODjg6tWrhlqmwfBFqYikHKO2OBsr5JzXpqkdXG1lyC9W4VJ0ulHukxDe+ag04cP2HzcTTVqY2q/tkhoU6AYHK/Mf7aqrTs0dsWdhCDwdLBGdmocJm841+EDtJyHnDqZdSDXJLSjsnBBSPkVOISKSc8BxQE/fyvOkSpvdywce9nI8zirE9nMxhlsgqTO+U0rfeVKlb9OYmVLx2tE9fXZJAU86paJoBz5CTMYomVLp6en46aefEBISAgsLiwovFxISgl27diE9PR1qtRo7d+5EYWEh+vfvb4xl6pWPkxVkEhEKlCo8SjNe5T1CW5TiM60MjeM4oVuKRviIse25GgcA6ODlAAsxZ7LClErN8FsYP7rXcALOq+LnYoP9L4YgyMMOqblFePbr8/j3QYqpl2UQipxCJGYVguOeZDXVB/wugRR2Tggp7XyUpkuqTVM7ONYgb0huIS616UUkMvOLDbI+UnfJ2qKUITulErMKjDYeHqcNOddXnhTPz1WTAUo78BFiOgYtSr355puwtraGk5MTYmNj8fvvv1d6+d27d0OpVMLJyQkymQwLFizA/v374e/vX+F1ioqKkJ2drfPPHEjEIqFb6Z4Rc6X48b0Ad8PuvFfaAG1R6mSEgnJLiNHkFCpx6E4SAGDVmDbYOK2TyQpT56JSkZxdBAcrCwwIcDXa/ZoDVzs5di/ogRA/J+QVqzB76+UGmTVyM05T1PF3sdHr2IChUdg5IaQ8Z7Q76PXyq97oXmkTOjVDoLstsgtLsPFkpL6XRvREGN8zYKdUoVKNrIKKp2D0KS5d2ylVRSh/TfGdUg9T8yjAnxATqVFRavny5eA4rtJ/4eHhwuVff/11XL9+HUeOHIFYLMaMGTMqLVq88847yMzMxLFjx3DlyhUsXboUkydPxq1btyq8zpo1a2Bvby/88/LyqslDMqggbWGILxQZWkpOEVJzi8FxQCs3G6PcJ6DJIpCKRXiUlo+HqTSPTYzjr1uJKFSq4edijfbN7DG0jbvJClN8wPnodk0hlTS+TU1t5RbYOrsrRrdvihI1w6u7wvDNP1ENqkh9s56FnPP4sPPbCVkN6udRn2zcuBE+Pj6Qy+Xo3r07Ll26VK3r7dy5ExzHYdy4cWXOu3fvHsaMGQN7e3tYW1uja9euiI2NBaDpTn/55ZcREBAAS0tLNG/eHIsXL0ZWFnXLEQ3GGM5qQ857VTNPqjSxiMObwwMBANvPPRKKBcS8GLJTSm4hFnahNdYOfPF8p5Sex/eaOVpBKhGhuESNhAzKXyTEFGr06em1117DvXv3Kv3n6+srXN7Z2RmtWrXCkCFDsHPnTvz111+4cOFCubcdFRWFDRs24Pvvv8egQYPQvn17rFy5El26dMHGjRsrXNOKFSuQlZUl/IuLi6vJQzIofoTOWGHn/Oiej5M1rKTG+ybfWiZBd98mAIATNMJHjGTP1XgAwKTOXkKo/9OFqSW7bxi8MJVXVIJDtzUdWxMa0eje02QSMb54tgPm9W4BAPi/v8Lx0aHwKq5Vf9zQ5km196o/o3vAk7Dz9LxiPDby1t0E2LVrF5YuXYqVK1fi2rVraN++PUJDQ6FQVP5eGRMTg2XLlqFPnz5lzouKikLv3r0RGBiIU6dO4ebNm3jnnXcgl2s6Fx4/fozHjx9j7dq1uH37NrZt24ZDhw5h7ty5BnmMpP6JScvH46xCSMUidPVpUqvb6N/KBSF+TihWqfHZ0ft6XiHRhxRtp5QhMqUAwN3OuGHnfKaUvsf3xCIOvs7aET7KlSLEJGpUlHJxcUFgYGCl/6TS8ufS1WrNB8OiovK3Ds/P17zQiES6SxKLxcJ1yyOTyWBnZ6fzz1w82YHPOON7wuiem3HypErjc6VOhFNRihheTGoeLsdkQMQB4zvqFoJKF6YO3nhs8MLUodtJKFCq4OtsjQ5eDga7n/pAJOLw9qjWeGtEEADg69MPseN8jGkXpQeMsXrbKVU67PxWPHXKGNtnn32G+fPnY/bs2WjdujU2b94MKysrfP/99xVeR6VS4bnnnsPq1at1vujjvfXWWxgxYgT+97//oWPHjvDz88OYMWPg6qp5H27bti327t2L0aNHw8/PDwMHDsSHH36IgwcPoqSkxGCPldQffJdUJ28HWErFtboNjuOwYrjmtX7/9QTaTMEMGbJTCiidK2X4ohRjzGDjewCFnRNiagaZM7l48SI2bNiAsLAwPHr0CCdOnMDUqVPh5+eHnj17AgASEhIQGBgotLEHBgbC398fCxYswKVLlxAVFYVPP/0UR48eLbd1vT4I1BalEjILkJVv+Hlrfue9QCOFnJfGF6Uux6Qju9A4s+Wk8dp7TdMl1aelC9zty34DaMzC1L7rmrWM7+gpdGw1dvP7+uKNYQEAgFUH7wrZJfVVfEYBMvKVsBBzRttEQp9Kj/AR4ykuLsbVq1cxePBg4TSRSITBgwfj/PnzFV7vvffeg6ura7mdTWq1Gn/++SdatWqF0NBQuLq6onv37vjtt98qXUtWVhbs7OwgkdSfPDRiOHxRqnctRvdKC25mjzHtmwIAPm5AnbENgVrNkJJruN33AAjHX0lZhh95y8xXIq9YBUD/nVIA4Oei6ZSKpLBzQkzCIEUpKysr7Nu3D4MGDUJAQADmzp2Ldu3a4fTp05DJNNV6pVKJiIgIoUPKwsICf/31F1xcXDB69Gi0a9cOP/zwA7Zv344RI0YYYpkGZ29pAU8HzQvnPSPkSvGdUoHuxv/Q5O1kDT8Xa5SoGf69X78/gBLzplYzIcNpUudmFV7OGIWpxKwCnNPuYDSuY+Md3SvPwn5+mNDREyo1w4s/Xa3X3z7e1HYYBbrbQSapXVeBKbXV7sB3k4pSRpWamgqVSgU3Nzed093c3JCUlFTudc6cOYPvvvsOW7ZsKfd8hUKB3NxcfPTRRxg2bBiOHDmC8ePHY8KECTh9+nSF63j//ffx/PPPV7pec904huiXSs1w/qHmfSukjkUpAHg9NAAWYg7/PkjFP/cb5u6r9VFaXjFUagaOA5xtqr+7Yk0Ys1MqTju652org9xC/+/Dfq7UKUWIKRnkK7Pg4GCcOHGi0sv4+PiUCV1t2bIl9u7da4glmUyQhy0SMgtwLzEbPXydDHY/JSo1HiRrXkgDjbjzXmkDA10RlRKNE+EKjGznYZI1kIbvwsM0JGQWwFYuwZDWbpVeli9MvfjTNRy88RgAsG5ye0jE+qnH/3b9MRgDurVoAi89B2/WdxzH4f8mBCMmLQ/XYjMxb/sV7H8xBA5Whjk4NqQno3v1K0+K93TYOXX0maecnBxMnz4dW7ZsgbNz+cUCPs5g7NixWLJkCQCgQ4cOOHfuHDZv3ox+/frpXD47OxsjR45E69atsWrVqkrvf82aNVi9enXdHwgxa3cfZyMzXwlbmQTtPOv+mubVxArTe/jg+7PRWPN3OHr7O0MkotcYU+N33nOyluntmOdp7vaaL96Tso1QlErXhpwboEsKKD2+Rxs2EWIKjW+bKCPjc6XCDZwrFZOWj6ISNSwtxGhuog/HA7QjfKciFFDTlqrEQPiA89Htm1br27Khbdzx1XOdIBFpOqaW6qljijGGfdoxwomNOOC8MnILMb6e3gWeDpaITs3DSz9fg9JIOyLq0416XpQKpLBzk3B2doZYLEZycrLO6cnJyXB3dy9z+aioKMTExGD06NGQSCSQSCT44YcfcODAAUgkEkRFRcHZ2RkSiQStW7fWuW5QUJCw+x4vJycHw4YNg62tLfbv3w8LC4tK12vOG8cQ/Tkbpelm7+7rpLdixcsD/WErl+BeYjZ+C0vQy22SulFk86N7hsmTAp50Shkj6JzvlDLUF4C+2vG99LxipOcVG+Q+CCEVo6KUgQlh5wYe3+NH91q525rsG6quPk1gK5MgLa9Y+BBHiD7lFCrx1+1EAJWP7j2tdGHqgJ4KU3ceZ+OBIhcyiQjDg6kzsCIutjJ8O7MLrKRinI1Mw+qDd0y9pBpRqxluJ2heX+tbyDmPws5NQyqVonPnzjh+/LhwmlqtxvHjx4V8zdICAwNx69YthIWFCf/GjBmDAQMGICwsDF5eXpBKpejatSsiIiJ0rnv//n14e3sL/52dnY2hQ4dCKpXiwIEDws58lTHnjWOI/jzJk9Jf976jtRQL+/sBAD49ch+FSpXebpvUTrK2e8lQIedA6Uwpwxel+J33DBFyDgBWUokQuUIjfIQYHxWlDIwvSkUk5Rh0B7AIbch5kAnypHgWYhH6tnIBAJykXfiIAfx9KwmFSjV8XazRsYY73em7MMWHrQ9p7QY7eeUdCI1dkIcdvpjSERwH/HghFj/Uox35HqbmIreoBHILEVpqMyfqIwo7N42lS5diy5Yt2L59O+7du4eFCxciLy8Ps2fPBgDMmDEDK1asAADI5XK0bdtW55+DgwNsbW3Rtm1bYXfj119/Hbt27cKWLVsQGRmJDRs24ODBg3jxxRcBPClI5eXl4bvvvkN2djaSkpKQlJQElYqKBY1ZoVKFyzHpAIBeesiTKm1OrxbwsJcjIbOgXr3GN1SKHMOGnAOAu/a2c4pKkGPgTY4MPb4HlMqVorBzQoyOilIG5t3ECpYWYhSVqBGTZrg55Xva8cAAExalgCe78B2nohQxAH50b1LnZrXKxdFXYUqpUgsZVRM7Vb9jqzEb0toNb4QGAgBWH7yLfx+YfyBuel4xXvv1JgCgfTMHg+VyGAMfdn6LilJG9eyzz2Lt2rV499130aFDB4SFheHQoUNC+HlsbCwSExNrdJvjx4/H5s2b8b///Q/BwcH49ttvsXfvXvTu3RsAcO3aNVy8eBG3bt2Cv78/PDw8hH80kte4XYvNQKFSDVdbGfz1XGSXW4ixZEgrAMCGE5HIzKcRKFMyRqeUtUwCW7lE5/4MxdDje8CTHfioU4oQ46u/R9j1hEjECYWiuwbMlYpI5nfeM227ff8AF3CcZrTJ0G9QpHF5lJaHSzHpEHHAhI61LwTpozD174MUpOYWw9lGij4t9fttc0P2Qj9fTOjE78h3zay3Xn6cWYBnNp/DjbhMOFhZ4O2Rrau+khl7OuycGM+iRYvw6NEjFBUV4eLFi+jevbtw3qlTp7Bt27YKr7tt2zb89ttvZU6fM2cOHjx4gIKCAoSFhWHs2LHCef379wdjrNx/Pj4+enxkpL45F6nZda+Xv7NBNjyY2KkZAtxskV1Ygq9ORen99kn18Z1SrgbslAKMswMfYwwJGZpOKUON7wEUdk6IKVFRygiEXKlEw+RK5RQqhbbWQBN3SjnZyNBem7tCI3xEn/Ze04Sn9vJ3FnIMamtoG3dsLFWYeu3XmhWm9mnXMqa9Z73unjE2juOwZkIwung7IqewBPO2XzbLb9MjFTmYuOkcolLy4GEvx54XeiK4noac8/iw87S8YqNs300IMT9ntHlS+h7d44lFHJaP0HTEbjsbI+QAEeNTGKFTCniyA58h31dScopQVKKGiAM8HAxXZOOLUub8hRkhDRV9mjKC1h6aQlG4gYpS95M1HVhudjI4Wpt+u/VBNMJH9EytZthbanRPH0JLFaZ+D6t+YSqrQIkjdzW7aU2gXfdqTCYRY/P0zvB0sERMWj5e/Mm8duS7HpuBSZvPIzGrEH4u1tizMAT+rqYt9uuD3EKMlnzYOY3wkUaEOgM1sgqUuKndhKaXHkPOn9a/lQt6+jqhWKXGZ0fuG+x+SOWMkSkFAB52hg8750f3POwtYWHALwL5kda4jHwK6yfEyKgoZQRPOqUMM74Xrg05N/XoHm+Atih1NjIVRSX0ok7q7kJ0GhIyC2ArkyC0Tdmt1GurNoWpv28lorhEjVZuNmjT1Dz+5uobZxvNjnzWUjHORaVh1YE7ZvHB8fT9FEzbchGZ+Uq093LAry+ECLvxNATBnprf18vR6SZeCSGGxxjD2sMR6PLBMVx9lGHq5ZjcxYdpUDPA18UaHvaGe13jOA4rtN1S+8MScOcxFcGNTa1mSBHG9wzdKWX48T1jhJwDgLONFHZyCRiDQXOACSFlUVHKCAK1Ramk7EJk5Ol/VCU8kS9Kmce3+W2a2sHNTob8YhUuPqQPP6Tu+IDzUe2bQm4h1utt17Qwte+6ZnRvQqfaha0TjdI78v10MRY/nH9k0vX8HpaAedsvo0CpQp+Wzvh5Xnc0MYPOU30aHKQJ1/7lUixSc4tMvBpCDIcxhvf+uIsNJyORlleMQ7drFibfEJ3lR/f8DJ+D2K6ZA0a3bwrGgI/+Djf4/RFd6fnFKFEzcJzmSyBD4jOlkrIKDHYf8UYIOQc0BdUnO/BRUYoQY6KilBHYyCRorn0hNUSuVATfKeVhHkUpjuOEXfhO0AgfqaO8ohIcup0EQH+je0+rbmEqLj0fl6LTwXHA2A5NDbKWxmRwazcsH8bvyHcH/9w3zY5828/F4NVdYVCqGEa188B3M7vCWiYxyVoMaUhrN7RvZo+8YhU2now09XIIMQi1muGd329j69kY4TRDdarXJ2ejnoScG8PrQwNgIebw74PUerHbakPCbzTkZC016Lgb0LA6pQDKlSLEVKgoZSSBwg58+i1KMcZwL8k8dt4rbUAAnyuVbBZjOaT++utWIvKLVWjhbI1OzR0Mdj/lFaZUat3f3d+0XVK9/JwNOv7QmDzf1xcTOzWDmgEv/WzcHfkYY/js6H2sPHAHjAEze3rjyykdIZU0zLdGjuPweqimCPjThVgKISYNjlrN8N/9t/DjhVhwHDCjpzcAIDzJMJme9UVSViEiFbkQcUBPX8PlSZXW3MkK/+mhef7X/BUOtZqOBY1F2HnP1rB5UgCEY6EkA+64zWdKGXLnPR6fKxWVQkUpQoypYR55myE+V4rPf9KXx1mFyCksgUTECdV9c9DL3xlSsQhx6QX0wk7qZE+pgHNDj8uFtnHHhmlPClNLd4cJhSnGmDC6N74jBZzrC8dx+L8JbYUd+eZuv2yQMeenqdQMb/92G18efwAAWDqkFVaNaQORqGGPZPZu6YwQP00I8RfHHph6OYTojUrN8Pqem9h5OQ4iDvj0mfZYMTwIHAek5hZDkdN4d508F6UZ3Qv2tIe9lYXR7vflgS1hK5PgbmI2fr+RYLT7beyEnfcMnCcFPOmUysxXGiwcPD5D0yll6PE94EmnFH12IcS4qChlJE/CzvX7bV2E9ts/Pxcbs/p231omQQ8/zbdxNMJHais2LR8XteNyxioEDWtbfmEqLC4T0al5sLQQY1hb/YWtE90d+R6l5WPhT1cNuiNfUYkKL/9yDT9d1HRTfDCuLRYPatloMsJeDw0AAOy9Fo9IBY01kfqvRKXG0t1h2HstHmIRh3XPdsCETs1gKRWjhZM1gCf5m43RGW2eVIiRRvd4TayleKG/HwBg7eH7tKOZkSiytTvvGaFTyk4ugZVUk/VpiB34VGqGx5nGHN/TvF48TMmj7j5CjMh8qhgNXGttUepBcq5eP2zxOQkBZhJyXtrAABcAVJQitbf3mqZLqre/M5oacRe08gpTv2o7toa1dW+QeUOm5mwjw3ezNDvyXXiYjnd/N8yOfLlFJZi99TL+upUEqViEjdM6CSMmjUXH5o4Y2toNagZ8Slu2k3pOqVLjlV1h+D3sMSQiDuundsTYDk++xDDUl4L1BWMM5yI1eVK9jVyUAoA5vVrA3U6OhMwC7DDxhhaNRXKO8TqlOI6Du53hcqUSswpQomawEHNwszN8ka15EytYiDkUKFV4bMDwdkKILipKGUkzR0vYyCQoVqnxMEV/OzqYW8h5aQMDNTs9XY7JQFaB0sSrIfWNWs2EopShAs4r83Rh6ueLsQCACZ1odM9QAt3t8OVUzY58v1yKxbZzMXq9/dTcIkz95gLORaXBWirG1tldMSLYQ6/3UV8sCw0AxwF/307CzfhMUy+HkFopLlFj0c/X8OfNRFiIOWx8rlOZv2k+01Pf8Qn1RVRKHpKyCyGViNDZ29Ho928pFWPp0FYAgA0nI5GVT8eDhsZ3SrkaoYgDPBnhS8rWfxGHDzn3dLCE2Ajj9RKxCD7a7sooPX5eI4RUjopSRiISccKBkT6/rePDO4PMKOSc19zJCv6uNlCpGe28QmrsUkw64jMKYCOTYGhr04zLlS5MAYCbnQwhRthOuzEbFOSGFcM1Ydzv/3EXpyL002kZl56PZzafx62ELDSxluKX53sYbRcqc9TKzVYYif3kcISJV0NIzRWVqPDiT1dx+E4ypGIRvp7eGaFtyr5XNPZOKT5PqquPI+QWYpOsYWKnZghws0VWgRJfnaKdPw0tWQg6N3ynFGDYHfj4DTmMkSfFE3KlaAc+QoyGilJGxHcz6evAqKhEJVTxzXF8DwAGBmp24fvkcATW/H0PR+4kITW3yMSrIvUBH3A+qp0HLKWmOZAGNIWpjc91gqutDC8N8DfKN3WN3fw+vnims2ZHvpd/vl7n3KOIpBxM2nwO0al58HSwxJ4XeqJdMwf9LLYeWzK4lbBl+zlt5gwh9UGhUoXnf7iKY/cUkElE2DKzi9Cd/bSgppqiVKQiF8UlhsuqM1dnHmjzpEz4hYpYxGG59suGredikJBJY1GGlKINOjfGuBsAePCdUgYoSsVlGC9PiufnyndKUVGKEGOhopQRCd/W6amFPEqRB5WawU4uEd4QzM2odh4QccCjtHx8ffohnt9xFV0+OIZ+n5zE0l1h+PHCI9x9nC3scEYIAOQVleCvW4kATDO697TQNu649NZgzOjpY+qlNAocx+GD8W3R1ccROUUlmLv9Sq135LsSk45nNp9DcnYRWrnZYO/CEPia0U6lpuTVxArTujUHAHx8OMIgGV6E6FtBsQrztl/B6fspkFuI8P2srujXyqXCyze1l8NOLkGJmiGykXU+qNQM5x+aLk+qtP4BLujh2wTFJWp8eoS6Mw1FrWZQGL1TSlMwMkinVLqmU6qZo/E7pRrb6wUhpkRFKSPSdws5P7oX6G5ntrtGtWvmgNOvD8Ank9phajcvtHLTvNA/SsvHvusJePu32xjx5b9ov/oI/vPtRXx29D5ORSgog6qR+/t2EvKLVfBxsjJJBgYxPZlEjM3/6Yxmjk925Ktpl8OJ8GT857uLyC4sQWdvR+xe0FMYMyAaiwa2hJVUjBtxmThyN9nUyyGkUnlFJZi97RLORKbCSirGttndqhzD5TgOgdrjL/64qbG4lZCFnMIS2MklaOtpb9K1cByHFcODAAD7ryfg7uPG9bMwloz8YpRov+h1MVJRysPOcJ1S8dpOKWOO7/m7asf3KFOKEKOhLaSMKNDdFhwHpOQUITW3CM42dXuzMOeQ89K8mljBq4kVnuniBQDIKlAiLC4TVx9l4NqjDFyPzUBuUQnORKYK2xZzHNDS1QadvR3RsbkjOns7wtfZ2myLb0S/9lyNA6DpkqKfeePlZCPDdzO7YsJXZ3HhYTpWHriN/xsfXK3fib1X4/HG3ptQqRkGBrpi47ROJh0DNVcutjLM6dUCG05GYu3hCAwOcqMRVWKWNDtnXsLlmAzYyCTYNrsruvg0qdZ1g9xtcSk6vdHlSp3VHlP19HMyi7/r9l4OGNXOA3/cTMSav+9hx9zupl5Sg5OsDTl3tpHCQmyc3gNDZkrFZfCdUsYb3+O7qVNzi5CVr4S9lYXR7puQxoqKUkZkJZXAx8ka0al5uJeYjT4tK243rw5+DDDQDEPOK2NvaYF+rVyEdnuVmuF+co5QpLoam4FHafm4n5yL+8m5+OWSpkDhaGWBTs0d0cnbEZ2aO6K9lz2spPQr3NDEpefjwsN0cBwwvpPpR/eIaQW422L9tI6Yu/0KfrkUh5autpjTu0Wl1/n234f44M97AIAJHT3x8aR2Rjs4r4/m9/XFjguP8ECRi9+uJ2CiGYzMElJadqESM7+/hOuxmbCVS7B9Tjd0al79LtogoVOqce3AxxelTD26V9rroQH461Yi/n2QiuTsQqPlHjUWyTmawpCLrfGeVz5CJDW3CMUlakgl+nm/LSpRIUmbj+VlxPE9G5kE7nZyJGUXIio1t0avNYSQ2qFP9EYW5GGrt6JUuPYbP3MNOa8usYhDkIcdgjzs8J8e3gA0b2x8geraowzcjM9CRr4Sx8MVOB6uKHU9W3TWFqp6+joZbftbYjj7riUAAEL8nODpYLxvxoj5Ghjohv8OD8KHf93DB3/eha+LNfoHuJa5HGMMHx+KwObTUQCAub1b4K0RQRCZQYeAObO3tMDC/n746O9wrDt2H6PbN9XbhwpC6iorX4kZ31/Ejfgs2FtaYMfcbjXeqCCwEe7AV6hU4cqjDABAiBkVpbydrOFuJ8fjrEI8ziygopSepWg7pdzsjDO6BwBNrKWQikUoVqmRnF2ot1G7xMxCMAZYWojhbCPVy21Wl5+rNZKyCxGpoKIUIcZARSkjC3S3w1+3knAvsW7f1qXnFQtBhvW9KFUeZxsZhrZxx1Dt9s7FJWrcTcx+0k31KANJ2YW4nZCN2wnZ2H7+ESzEHF4e2BIL+/tRV0Q9pVYz7Ln2ZHSPEN68Pi3wQJGD3Vfi8fLP17HvxRC0dHvy2leiUuO/+29h9xXNro1vDgvEC/18afyzmmb29MH3Z6IRn1GAXy7FYmaIj6mXRAgy8orxn+8u4s7jbDhaWeDHed3RpmnNs5EC3DTxCam5xUjJKTJa1o4pXYnJQHGJGh72cvg6W5t6OTrc7DVFKX7UjOhPsrazyFgh54AmL8zdXo7Y9Hy9FqVKj+4Z+73c38UGZyPTaAc+QoyEPrkbmb7CzvmwTq8mlrCRNfzaolQiQgcvB8zt3QIbn+uEC/8dhHPLB2L91I6YFeKD1h52UKoYPjt6H2M3nMXthCxTL5nUwuWYdMSlF8BGJkGotiBJCKDdkW9cMLr5NBF25EvX7shXqFRh4U/XsPtKPEQc8PHEYCzs70cFqRqwlIqxeFBLAMD6E5HIKyox8YpIY5eWW4SpWy7gzuNsOFlL8cvzPWpVkAI0v98tnDSFmcbSLXU2SjO6F+LnbHavhW7a0TK+gEL0h//C2tgdaIbIlYpL14ScGzNPiufHh50rKOycEGOgopSRBWlDyaNScmu8k1RpEfU0T0qfmjpYYnT7plg1pg3+XNwbX0zpAAcrC9xNzMa4jWfx6ZEIFJWoTL1MUgN7rmq6XEYGe1BeGClDKhFh0386wauJJWLT87Hwx6tIyy3CjO8v4ejdZO35nfFs1+amXmq99GxXL3g7WSE1twhbz0abejmkEVPkFGLKNxcQnpQDZxsZdj7fo87HO0GNbAc+IU+qpZOJV1IWX8CgopT+maJTCgDcDbADH98pZcyd93h+2rDzh9QpRYhRUFHKyDwdLGEnl0CpYohU1P6FLlw7/hfUAEf3aoPjOIzt4ImjS/phRLA7StQM609EYvT6M7gRl2nq5ZFqyC8uwV+3EgEAk7rQ6B4pH78jn41MgovR6ej7v5O4FJ0OW5kEP8zpRh12dWAhFmHpkFYAgK//eYjM/GITr4g0RsnZmoLUA0Uu3Oxk2LWgh86obm0Fao+X6hqfUB9k5StxS9sxHuJnPnlSPFdt3lESFaX0ju+UMnbGqocBOqXiMzSdUsYMOefxRalH6fl1aiIghFQPFaWMjOM4vQRuhidrDqoCGnGnVHlcbGX46rnO+Oq5TnCyluJ+ci7Gf3UWa/6+h0IldU2Zs0O3k5BXrIK3kxW6eFOoJKlYKzdbrJ/aESIOyCtWaTopFvRAD1/z6wiob0a3a4pAd1vkFJZgkzYwnhBjeZxZgGe/Po+HKXloai/Hrud7Ch8O60pf8Qn1wfmHqWAMaOlqY5ZB4nxXjYIypfROYapOKW1RKim7QG+3GZf+JFPK2NzsZLCRSaBSMzxKoxE+QgyNilImECR8W1e7AyOVmuE+P77nQZ1S5RkR7IGjS/thbIemUDPg69MPMeLLf3H1Ubqpl0YqwI/uTezUzOzyL4j5GRDoinXPdsDIYA/sXdiz1lkzRJdIxOGNYQEAgG1nY2i8hhhNXHo+nv3mPGLS8tHM0RK7FvSEjx4DugP1FJ9QH5zRju71MqNd90rjC2XUKaVfajVDSq5pMqUM0ylluvE9juPg56J5/aGwc0IMj4pSJiB8W1fLXIPY9HwUKFWQSUTwcTKvHVXMSRNrKb6Y0hFbZnSBq60MD1PyMGnzebx38C7yiynE15zEZ+TjXFQaAGBCJ08Tr4bUF2M7eGLjc53gTa+DejUgwBVdvB1RVKLGl8cfmHo5pBGITcvHlG8uIC69AM2bWGHXgp56/yDq6WAJW218QkP/kHkuUvN+au5FKSp661dGfjGUKgZAs4u1Mbnba7qZ9JUpVVCsQmquZoTcFON7wJMRvqgU6pQixNCoKGUCT1rIc8AYq/H1I7TFrFZuthCLqKOkKkNau+Hokn6Y1LkZGAO+PxuN4V/8i/PaIggxvX3XEgAAIX5OaGaigw9CiAbHcXhjWCAAYNflOBpdIAYVnZqHyV+fR0JmAXydrbF7QU94Ouh/XIfjOAS5N/wRvseZBXiYmgcRB3T3bWLq5ZTLTZsplVNYQl8S6hGfJ+VkLYVUYtyPeHynlCKnCCWqunci8l1StjIJ7CxNs/ENvwNfXTKACSHVQ0UpEwhwt4WIA9LzipGSU/N5ej6kM5BCzqvN3soCa59pj22zu8LDXo5HafmYuuUC3vntNnJp63OTYoxh7zXN6N6kzhRwTog56NaiCfoHuKBEzfDZ0fumXg5poCIVuXj26/NIyi6Ev6sNdj7fQ8imMQR+B+SGXJTid91r7+UAO7mFiVdTPlu5BaylYgBAMuVK6Q3feeZi5DwpQNOZJRZxUKmZ0OFUF/zOe82aWJks0uFJpxQVpQgxNCpKmYDcQowW2pyEu7U4MIpI4kPOqShVU/0DXHFkSV9M7abZMn7HhUcIXfcP/n2QYuKVNV5XHmXgUVo+rKViDGtLO6cRYi6WDdVkSx248Rh3HzfcD/HENO4n52DKN+ehyClCgJstfpnfw+A7hvEbzYQnNdwd+PiiVC8z3HWvNCFXSo8ZRI0d3yllinB7sYiDm7YYlphV97DzuHR+5z3jh5zz/F21mVKK3FpNthBCqo+KUiZSeoSvpsK143v8bZCasZVbYM2EYPw0rzuaOVoiIbMA07+7hOV7byK7UGnq5TU6e65ouqRGBHvASmqaFm1CSFltPe0xqp0HGAPWHokw9XJIA6IpSF1Aam4xWnvY4Zfnexilu6Oh78DHGMPZKPPOk+LxhRNFDhWl9MVUO+/x+C5HfWSFmTLknNe8iTXEIg55xSrq6CPEwKgoZSK1PTDKLy7BI+0WqdQpVTe9/J1x+NW+mBXiAwDYeTkOQz/7ByfCk027sEYkv7gEf95KBECje4SYo9eGBkAs4nAiXIErMbR7KdEPV1sZ3OzkCPa0x8/zu6OJtdQo99vKzQYcB6Tm1i4+wdw9UOQiJacIcgsROnk7mHo5leJzpahTSn9M2SkFAB7asHN97MDHd0o1M2GnlFQigreTpihGuVKEGBYVpUyktrkG95NzwZhmdtvYO2s0RNYyCVaNaYPdC3rCx8kKSdmFmLPtCpbuCkNmft1n4knlDt9JQm5RCZo3sUJXH/MMZCWkMWvhbI3JXTQF4/8diqARBqIXDlZS/Di3G36c1x0OVsYpSAGAlVSCFtrdOhtitxQ/utfVpwlkErGJV1M5N6GrpuEVB02F71BytTNtp5Q+Co18ppSpdt7jUa4UIcZBRSkT4TulHqbmoVCpqvb1whP50T3qktKnbi2a4O9X+mJe7xbgOGDf9QQMWfcPDt9JMvXSGrQ9VzWjexM7NYOIdpIkxCwtHtQSUokIl2LSceo+5e8R/XCykcHe0vhB3IHa4yc+CqEhEfKkzHx0DwDcbPU36kU0+E4pV1vTdEq5azu09NEpFZ+hzZQy4fgeQEUpQoyFilIm4m4nh4OVBVRqVqOWUD6cM8CNilL6ZikV4+1RrbF3YQj8XKyRklOEBTuuYtHP15CWS9/k6VtCZgHOabMvJnTyNPFqCCEV8bC3FMacPzkUAbWauqVI/RXkXvtMT3NWolLjwkPNiG3velCU0mf+ENFQaLvO6nunVHahElkFmoxXU47vAYCfizbsnIpShBgUFaVMhOM44cCoJjvw8d/sBVLIucF0au6IPxf3wcL+fhCLOPxxMxFD1v2Dgzce0+iKHu2/Fg/GgB6+TUz+TRghpHIL+/nBVibB3cRsIQeOkPoosIGGnd+Iz0JuUQkcrCzQuh4cIwqZUlSU0gvGmBAab7pMKW2nVHbddt+L02bnNrGWwlpm2g1w/F01nVKUKUWIYVFRyoRqGnbOGBM6pQIp5Nyg5BZivDksEPtfDEGguy3S84rx8i/X8cKPV4Vvb0jtMcaE0b1Jnb1MvBpCSFUcraWY39cXAPDZ0ftQqtQmXhEhtcPHH0Sl5KK4pOH8Hp/Tju6F+DnVi3F4Yfe97CL6wk8PMvKVUKo0z6OLiTJnhe63rKI6ddTyIedeJu6SAgBf7fhecnYRcmiHbkIMhopSJhRYw7BzRU4RMvOVEHFPKvfEsNo1c8CBRb3xyqCWkIg4HL6TjK9ORpp6WfXe1UcZiEnLh5VUjOFt3U29HEJINczp3QJO1lJEp+bh1yvxpl4OIbXi6WCJ/2/vv8Pjqs+88f99pqtLVht1WXKRjRtY2DHd4GCbhLLwbMguX8CEZQMPpGASFj+7CSHJXk42LGEhbMiPhWCy318geUiyCQnVhVAMBhs3sOUqWb1LM2pTz/ePOZ8jyVaZkWbmnDPzfl2XrgtLU8548OjMPff9vjMcFvgCckKN5LxroDwpYDT3yBsIoneIb/ZnS4xBzkmzwWbR5u1dQYYDkhR6TntmsSyoSQk5L9VBF31WihX5GaEi36nOQY2PhihxsSilocVFo7kG4XxKJIpXVfnpcFj1vVUlkdgsJtz/+QX4zhcXAwDq2hMrh0ILokvqmqVFmrdmE1F40u0W3Lt2HgDgP7Yfi2hJB5FejI1PSJQRviGvH5+c6QMAXFxtjKKUzWJCblpo8yJzpWZvNORcu83cNotJ3Qw+m1wpEXKudZ6UIHKlOMJHFDssSmloXkE6zCYJ/cO+sGbq60TIOUf3NFGRG/rEprWPJ0+zMewN4JWDoUya/7WyVOOjIaJI3PK5cpRkp6Dd5cELu+u1PhyiGRndwJcYHzJ9VN8LbyCIkuwU9VzFCAqUET7mSs2eKOwVaJQnJai5UrMoSolMqbIcffy/LKZTEqmzkkhvWJTSkMNqVqvv4XxaJ06eFrEopYni7NAnNq39swtwTHZvfNaGAY8fpTkpWFU5R+vDIaII2C1mfGPdfADAf+46CRczNsiAIs301Lv31dG9XEiS/vOkBKcSdt7BotSsdSqdUoUadkoBoe3iANA2i3Nl0SmllyU41fksShHFWlyKUh6PBytWrIAkSdi/f/+Ulx0ZGcG9996L3NxcpKen46abbkJ7e3s8DlMTi4rCX00sTp4WOvW/VSURiU9/XCN+DHr8Gh+NcYnRvZsuKDVEGCsRjXfj+SWozk9D35AP//XXU1ofDlHExLKYcM69jMBoeVJCoVrA8Gh8JMY32imlbVFKnCvPtPtNlmU0ikwp3YzviaIUM6WIYiUuRakHH3wQxcXFYV32/vvvx5/+9Cf89re/xdtvv42WlhbceOONMT5C7Yii1GfTfFrnCwTVCj0372kjw2FFupJ/xG6pmWnpG1ZPnm+6gKN7REZkMZvwrasXAgD+693T6BrgG0oyloXODEgS0DXgUTtMjKpn0KueQ15kkDwpQRSl2t3slJqtDpfSKaXx+J4zS0wVzOw57Rn0YsgbyissydZJUUoZ36vvGuTmWUooQ14/rvmPd/DPvz8Ej1/bnNCYF6VeffVVvPHGG3j00UenvWx/fz+effZZPPbYY7jyyiuxcuVK/PKXv8T777+PDz74INaHqonRT+umLkqd6hyELyAj3W7RzScHyUh8AtTCXKkZ+f0nzZBlYPXcOSg3UO4FEY23YYkTy0qzMOQN4CluJCWDSbVZUJkbik842mbsEb7dJ7shy6HzyXyNR7cipRalZpE/RCGisKdl0DkwplNqhs9pozK6V5hp181Sp6JMB1JtZviDMs4oeVdEiWB/Yx8+a3Vhx9EO2MzapjrF9N7b29tx11134Ve/+hVSU6d/A7p37174fD6sW7dO/V5NTQ3Ky8uxe/fuWB6qZsQGvvquQQx7J69QipOm0Kd7HHnSSpHyqc1stookK1mW1dE9BpwTGZskSXhwfQ0A4P/94Iy6wpvIKBYVhfehoN6J7mOjdUkBgDMrVEBhp9TsiU4prYPOR0cyZ/acNqmje/r54NJkklClZACf5AY+SiAf1/cCAGor52heX4hZUUqWZWzatAl33303amtrw7pOW1sbbDYbsrOzx32/sLAQbW1tE17H4/HA5XKN+zKS/Aw7ctNsCMrAsfbJsw1EyDlH97RVLDqlOL4XsX1nenG6axCpNjOuWVqk9eEQ0SxdMj8PF1XnwhsI4vG3jmt9OEQRqVHyOY8aPFfq/ZOhotQl83M1PpLIFWQwUyoaZFlWx1D10inV2j8CWZYjvn5jjxJyrrOpEOZKUSL6qL4HAHBhZY7GRzKDotRDDz0ESZKm/Dp69CiefPJJuN1ubNmyJRbHrdq6dSuysrLUr7KyspjeX7RJkhTWFpijys9YlNKWU/yy5fhexP7v3mYAobGfNCWbi4iM7dvrQ9lSv9vXhONTfLBCpDfhZnrqWWPPEBq6h2AxSVg113hFKXFO1T3oYVbPLPQN+eBV/v60HuEUz+mwLwDXcORLgUTIuV427wmiKHWCnVKUIPyBIPY1hDqlVlYYsCj1wAMP4MiRI1N+VVVVYceOHdi9ezfsdjssFgvmzZsHAKitrcXtt98+4W07nU54vV709fWN+357ezucTueE19myZQv6+/vVr8bGxkgfkubCaSGvE51SRdy8p6ViEeDI9cURGfEF8MqBFgAc3SNKJOeX5+DqxYUIysC/v3FM68MhCps49zrZOQCv35gFEdEltaIsW13EYiRzUm2wmiXIMgwfOK8lMf6Yk2qF3aJtDpPDakZOqhUA0OqKfKqgqVd0SumrKDWvQHRKsShFieFomxuD3gDS7Ra1c1hLEf8Gy8/PR35+/rSXe+KJJ/DDH/5Q/XNLSwvWr1+Pl156CatXr57wOitXroTVasX27dtx0003AQDq6upw5swZrFmzZsLr2O122O3GCnY822in1MSfMvcP+dCizGYvZKeUpoqyRacUx/ci8fqnbXB7/CjJTsHnDPhpLhFN7lvrF+LNI+147dM2HGjsw/KybK0PiWhaJdkpyHBY4B7x42TngHouZiTvnugGAFw0z3h5UkAoq6cgw4HmvmG0u0ZQrJNta0ajl817gjMrBb1DPrT2j0T8ZrepR2RK6ev/hdHxvQHIsqx5/g5NzOsP4n/2N+PyBfma56vp3V6lS+qCihyYTdr//xyzTKny8nIsWbJE/VqwYAEAoLq6GqWloU6J5uZm1NTUYM+ePQCArKws3Hnnndi8eTN27tyJvXv34o477sCaNWvwuc99LlaHqjnxgn2kzTXh/HWdMhJRkp2CTIc1rsdG4xXNctVtshIB5zetLIVJBy98RBQ9Cwoz8DfnlwAAfvJ6ncZHQxQeSZKwSORKGXADXzAo430l5PwSgxalAKAgUwk7Zwf6jIm/O61H94SZbuALBuXRTimdje9V5KbCJAHuET86B9jVp1ePvXkM3/6/B/HoGzwXmY6aJ6WD0T0gxtv3puPz+VBXV4ehodGtPT/96U/xxS9+ETfddBMuu+wyOJ1O/O53v9PwKGNvXkE6rGYJ7hE/mifowBEnS8yT0p74RTvg8cM14tP4aIyhtX9Y3Q500wUlGh8NEcXC/esWwGqW8O6JLvWNMpHe1ajxCcbLQ6trd6N70IsUqxkrDNyd6FS6GdpdfKM/Ux1uvXVKjYadR6JzwANvIAizSVLPt/XCYTWrhTLmSulTp9uDbe/XA+BzNB1ZltWiVG3lHI2PJiRuRanKykrIsowVK1ac870rrrhC/Z7D4cBTTz2Fnp4eDA4O4ne/+92keVKJwmYxqW2hE50Yie9xdE97aXYLMh2hqdeZrrtNNr//pBmyDKyqnIOK3DStD4eIYqBsTir+flU5AODHr9fNaOsSUbyFs2hGr95Tir+rq+bAZtH0M+ZZEYWUNnZKzViH8nen9eY9oUg8pxFuqm5URvecmQ5YzPr7f3oeN/Dp2tNvn8SwLwAAEzZ50Kim3mG0uzywmCTdfKihv3/xSWrxFCdGdaJTyoB5B4lIZB608AVvWrIsq6N7DDgnSmz3XTkfKVYzDjT24fVP27U+HKJpiQ50I3ZKiaLUxdXGHd0DRotSHN+bOb12SrVF2P02unlPX3lSQrUIO2cXju60u0bw3x80qH/ucHsMu8AiHj5uCHVJnVeShRSbtssRBBaldGKyT+uCQVndvLeInVK6MNO25GT0SWMfTnUOIsVqxjXLirQ+HCKKofwMO+68ZC4A4PG3uImP9G+hMwOSBHQNeAy1/c0XCOLD06E3FRcbOE8KAAqZKTVr7XrrlFLyVyPvlNLn5j2hOj/U7c8NfPrz1M4T8PiDWFmRA7vFBFnma8pUPq4PhZzrJU8KYFFKNyYrSjX3DWPQG4DNbEJlHkef9EANO2en1LREl9TGJU5DrqsmosiIotTRNjdz9xRPPfUUKisr4XA4sHr1anW5y3RefPFFSJKEG2644ZyfHTlyBNdddx2ysrKQlpaGCy+8EGfOnFF/PjIygnvvvRe5ublIT0/HTTfdhPZ2dq+dLdVmQaUyVm6ksPP9jX0Y8gYwJ81m+LxRZkrNnvi708u2sZl+eNukdkrptSjFTik9au4bxot7GgEAD1y9ACXKRIsIzadziaKUXvKkABaldEOEbTb0DGHQ41e/L4pUoTB0Pl16UMxOqbCM+AL404EWABzdI0oWOWk25KbZAABnuoemuXTie+mll7B582Y8/PDD2LdvH5YvX47169ejo6NjyuvV19fjW9/6Fi699NJzfnby5ElccsklqKmpwa5du3Dw4EF85zvfgcMx+ob0/vvvx5/+9Cf89re/xdtvv42WlhbceOONUX98iWCRGnZunKKUGN27qDrX8BttRSGlnedUMyLLstrlp5dOKVGUco/4MTDmPc10RKdUaY5Ox/eUolRL/8i492qkrZ/tOA5vIIjPVc3BRdV5jFmZRv+QD3XtoSms2kp2StFZ8tLtyM+wQ5ah/o8CQB3dM/onYYmkSHmxY1Fqan891gn3iB/FWQ58ripX68MhojgRn3KL0Npk9thjj+Guu+7CHXfcgcWLF+Ppp59GamoqnnvuuUmvEwgEcMstt+CRRx5BVVXVOT//53/+Z1xzzTX4t3/7N5x//vmorq7Gddddh4KCAgBAf38/nn32WTz22GO48sorsXLlSvzyl7/E+++/jw8++CBmj9WoapyhTvWjBsqVUvOkDD66B4wpYHj8fKM/A31DPngDoeycgkx9FKXS7RZk2CNfCtSo806psR+6nO5i2LkenOkewm8/Dk1lPHD1QgBAcXboNYVFqYntPRMa/Z6bl4a8dH28ZgAsSunKRCN8R0VRqohFKb0Qa2pbIpyVTzY760KdAJ9fXGj4T3KJKHzlyhuKM0lelPJ6vdi7dy/WrVunfs9kMmHdunXYvXv3pNf7/ve/j4KCAtx5553n/CwYDOLPf/4zFixYgPXr16OgoACrV6/GH/7wB/Uye/fuhc/nG3e/NTU1KC8vn/J+k5U49/rMIJ1Sgx4/PjnTBwC4JAGKUul2C9KUoF1mwEROhJznpFpht+gjsBgYE3YeZlHKHwiqH/bqNVMKGDPCx1wpXXhix3H4gzIunZ+HC5VRtJLs0P8/3MA3sY/E6J6O8qQAFqV0ZaIW8iNi856Tm/f0QhSlWvtGuPZ8ErIsY+fRTgDA2poCjY+GiOKpIpdFKQDo6upCIBBAYWHhuO8XFhaira1twuu8++67ePbZZ/HMM89M+POOjg4MDAzgRz/6ETZs2IA33ngDf/M3f4Mbb7wRb7/9NgCgra0NNpsN2dnZYd8vAHg8HrhcrnFfyUB0op/sHDDEtqY9p3vgD8oon5Oq246SSBVmMVdqpkZDzvWRJyWM5kqFVxho7R9BICjDZjbpZgxxItUFoQy6E8yV0typzgH8bt/4LilgtFOKRamJfVwf6pS6UEd5UgCLUrqyWO2UCnVHjfgCqFfaQzm+px8i6HzYF4BrmK3mE/ms1YU21whSrGaO7hElmTJ2Ss2I2+3GrbfeimeeeQZ5eRN3wASDoaLJ9ddfj/vvvx8rVqzAQw89hC9+8Yt4+umnZ3X/W7duRVZWlvpVVlY2q9szitKcFGTYLfAFZEN0P4yO7iXO79bCDFGUYqdUpESnlF5G94SiCDulxOheSU6Krrvr2SmlH/+x/TiCMnBVTQFWlGWr3y9hptSkPP4ADjT1A9BXnhTAopSuiBbyo60uBIMyjrcPICgDc9JsyNfxpwbJJsVmRk6qFQBH+Caz82hodO/ieblwWPXTTk5EscfxvZC8vDyYzeZztt61t7fD6XSec/mTJ0+ivr4e1157LSwWCywWC1544QX88Y9/hMViwcmTJ5GXlweLxYLFixePu+6iRYvU7XtOpxNerxd9fX1h3a+wZcsW9Pf3q1+NjY0zfOTGIkmSGpFghA187yZQnpTgzGJRaqb02yml5K+G+Zw26TzkXKguEBv4mCmlpWPtbvxRWaZ0/+cXjPtZifL/UHPfMCdaznK4uR9efxC5aTbMzUvT+nDGYVFKR+bmpcFmNmHQG0Bj75B6crSwMAOSpN9PDZKR+suWRakJ7azj6B5RshJFqebeYfgD+h+HihWbzYaVK1di+/bt6veCwSC2b9+ONWvWnHP5mpoaHDp0CPv371e/rrvuOqxduxb79+9HWVkZbDYbLrzwQtTV1Y277rFjx1BRUQEAWLlyJaxW67j7raurw5kzZya8X8FutyMzM3PcV7JYdFanul51DXjUrNE1CdSFLLp82liUipjYvFdo8E6pJp2HnAvzlE6p012DCARjV/CQZRnvHu/Czb/YjaXfex2Hm/tjdl9G9PhbxyDLwIbznFhSkjXuZ6LIPeILonfIp8Xh6ZaaJ1WZo7vagkXrA6BRVrMJ8wvT8WmLC0da3Qw517HiLAeOtLq4gW8CvYNefHIm9KJ3xUIWpYiSTWGmAzazCV4luFbvbzJiafPmzbj99ttRW1uLVatW4fHHH8fg4CDuuOMOAMBtt92GkpISbN26FQ6HA0uWLBl3fZELNfb73/72t3HzzTfjsssuw9q1a/Haa6/hT3/6E3bt2gUAyMrKwp133onNmzdjzpw5yMzMxNe+9jWsWbMGn/vc5+LyuI1G5HYe0XnY+fsnuwGE4h5ydbQ1abacmaE3kR3MlIrYaKeUvv5/iDTovLHXGJ1SxdkpsFtM8PiDaOwZQmWUu01kWcauY514YvtxdaEBAPzxQMs5xZdk9VmLC3851AZJAr75+fnn/NxuMSM/w45OtwctfcOYo2xMpNE8qdoKfeVJASxK6c6iokylKOVSO6UWMeRcd4qyR8POaby3j3UiKIdy0MRcNxElD7NJQumcFJzqHMSZnqGkLkrdfPPN6OzsxHe/+120tbVhxYoVeO2119Tw8zNnzsBkiqxp/W/+5m/w9NNPY+vWrfj617+OhQsX4uWXX8Yll1yiXuanP/0pTCYTbrrpJng8Hqxfvx7/+Z//GdXHlkhGF83ou1Pq/QTMkwJChWyAnVIz0aF2SulrfE/tlArzOW1Uxr31vHkPCP1+q8pPx5FWF052DkStKCXLMt78rB0/23kCB5XMH7vFhKUlWfi4oRcHm/qicj+J4KdvHQMAfGFp0aSLwEqyU9Dp9qCpd5jFPEUwKOPjhtFOKb1hUUpnRlvIXahTOqUWMuRcd0TYOTOlzrVDyZPi6B5R8iqfk6oWpS7W+mA0dt999+G+++6b8Geiu2kyzz///ITf/8pXvoKvfOUrk17P4XDgqaeewlNPPRXuYSa1hc4MSFJoPK7T7dFljqcsy3jneOLlSQGjBRVmSkVO7ZTS2/heZug8uWfQixFfYNp80SalU8oIH2JU56epRamrFhVOf4UpBIMyXj3chid3HFcnZFKsZty6pgL/cOlcdA94sfE/3sHh5lDesJ5D4OPhYFMf3vysHSYJ+Oa6BZNeriQ7Bfsb+xh2PsaprgH0DfngsJpwXrH+CnUsSumM+LTuw9M96B/2QZKABYUsSumN+ASInVLj+QNBvH0slCd1JYtSREmLYedkJKk2Cypz03C6axBH21zIz8jX+pDOcaZnCM19w7CaJayaq7/Ri9kQeUgdLg9kWdZd1oleybI8un1PZ0HnmSkWpFjNGPYF0O4aQUXu5B1FHn8A7e7Q+bTex/eAMRv4ZhF2HgjKeOVgC3624wSOd4Q2+aXbLbhtTQXuvGSuOp47J9UGh9WEAY8fp7oGMU8JWk9Wj70Z6pK6YUXJlH8XxcpEC4tSo0Se1IqybNgs+osVZ1FKZ8SoXv9wKJitMjcNKTZuL9Mb0SnFVvPx9jf2oX/Yh6wUK84fs56ViJILi1JkNDXOjFBRqtWNS+frryj13olQntT55TlItSXW6bsoqHgDoWBiZsCEp3/YB68/tExCb919kiTBmeXA6a5BtPZPXZRq7h2GLIc6hHIN8NyLDXwnOgcivq4vEMQfPmnGf+46idNdoaJWpsOCOy6eizsurkR26vjHbzGHulr2NvTiUHNfUhel9jb0YlddJ8wmCV+/6twsqbFEfEgzi1Kqj5Q8qQsr9fmhRmL9VksAOWk2ODMdarGjhqN7ujS2As9P9UaJ0b3LFuTDYtZfFZ6I4kMtSnWzKEXGsKgoE68ebtNt2Pl7Ik+qOrFG9wDAZjEhN82G7kEv2vpHWJQKk+iSyk61TjsepwVnZqgoNV3YeaM6updiiPNpsYHvRMdA2O8BvP4gXt7XhP/cdQKNPaHHm51qxT9cMhe3XVSJTId10usuLQkVpQ429eNvzi+NzoMwoJ8qXVL/64LSabO8ipWiFDulRn2sdEqtrNBfnhTAopQuLSrKUItSzJPSJ5F/4PHzU72xRFHqyhr9fcpMRPFTnstOKTIW8SHgkTb9hZ17/AH8VRmNv2R+4hWlgNB5VfegF+3uESwGF/yEQ6+b9wQ16mKaolRTb+j3RKnOQ86FuXlpkKRQp1rPoHfKTZgjvgB+83Ejnt51Ei3K30Neug13XVqF/+dzFUizT/9WfHlZKP9HBKAnow9PdePdE12wmiXcd+W8aS9frHZKcaIFADpcIzjTMwRJAi5gUYrCtagoEzvrQicfk20VIG05lBbj7kEv140qWvqGcbTNDUkCLl/APCmiZCY2KPUP+9A/5ENW6uSfAhPpgVg0c6LDDa8/qKvMjfdPdMPt8aMw056wo/GFmXZ81gq0T1PAoFEdLn1u3hOcYgPfNEuBROdQmQHypAAgxWZGSXYKmnqHcbJzcMKi1LA3gP//njP4xdsnx+R+2XH35dX4u1XlEUWzLC3JBgB82tIPfyCYdJMIsizj35UuqS/VloUVhi+yyboGPGEF7Sc6sXWvxpk5ZVeelliU0iFxYhT6b3ZK6VVRtkNtNee6UWBnXahL6vyybBbpiJJcmt2CvHQbuga8aOwdQlYqXyNJ30pzUpBht8Dt8eNU14CuPhT8y6FWAMCG85wJu31LFDDalUILTU+Eg+stT0oIt1OqUemUMsLmPaE6Px1NvcM40TEwbvHAoMePX33QgP965xS6BrwAgOIsB+65ohp/W1s2o+JIVV4a0u0WDHj8ONGpr9emeHjvRDf2nO6BzWIKq0sKALJSrEi1mTHkDaC1fwRzpxn3S3SjeVL67JICgOQqtRrEstIsSFJo1rjMIK2syUiEnbdO8wlQsth5lFv3iGgUw87JSCRJQo3yQaCecqV8gSDe+KwdALBxaZHGRxM7IuycC2TCp/9OqfCWAjUpmVJGGd8DoAaOn1TCzl0jPvxsx3Fc/OMd+NGrR9E14EXZnBRsvXEpdn17LW5dUznjbh2TScKSklAh6mBjco3whbqk6gAAf7+qXH3vNR1JkpgrNYbIk6rVacg5wE4pXarITcMv/p+VyE23JewnYolAfALUwlZzjPgCagjrWhaliAihotS+M31oYNg5GUSNMxMf1ffiaKsbOF/rownZfbIb/cM+5KXbdLs1KRpEp1QHi1Jh63AnSKZUj8iUMsb4HhDqlAKAQ039+Ombx/DL907DNeIHEMqcunftPFy/ohjWKI3aLSvNxgenenCwuQ9furAsKrdpBLvqOvHJmT44rCb877XVEV23JDsFJzoG0Nyb3EWpAY8fn7aEipm1Os2TAliU0q2rz3NqfQg0DbVTihV4fHi6B8O+AAoz7VhclFxtxUQ0MXZKkdGI+ITPdNQp9erhNgDA+vOcMCfwB5WFmaHCCjulwteu+06p0HF1DXjgCwQnLNAMevzoHgyNuRlrfC80Dranvgd7lNGo+QXpuO/KefjisuKo/1tdqsSEHEqisHNZlvGYkiV125pKtZsyXKNh58n9Pm3/mT4E5VCRTvyd6BGLUkQzVJwd3idAyWCnsnVv7cICQ6zzJaLYE28wGlmUIoMQ43tHdbKBzx8I4o1PQ0WpjUsSd3QPGC2sMFMqfHrvlJqTaoPNbII3EESH24OSCd4Qi4JBpsOCrBR9BjBPZEFhBiwmCf6gjEVFmfjalfNimvm2vDQbAHCkVX+LGGLlzc/acai5H6k2M756WVXE1y9R3qcl+/jexw2hommtjvOkABaliGZsNFMquYtSsixjhyhKcXSPiBQVuaFPktkpRUaxsDADkgR0uj3oGvAgb4pV7/Gwp74H3YNe5KRasboqcUf3gNGiVPfg5F01NEqWZd13SplMEgqz7GjsGUZb//CERalGdXTPOF1SAJCTZsPzd6xCQJZx2fy8mH8gWzYnBVkpVvQP+1DX5sbS0sReHhIMjnZJbbqocsINh9MpyWGnFGCMPCmAQedEM1akrrodQTAoa3w02jnZOYgzPUOwmU24ZF6e1odDRDohxvea+4bhCwQ1Phqi6aXZLahQ/r892qp9t9Srh0JdUlcvdiZ8kWZOqg1WswRZDhUFaWquYT+8/tDrql637wFAUebUH+CKolTZHP2OFU3mkvl5uHxBflwmBCRJwjKlEHWwuS/m96e1Vw+34WibGxl2C/5xBl1SAFCcxaBzfyCIfWdCRSk9b94DWJQimrHCTAckCfAGguo8fDLaVRfqklpdNQdpdjZfElFIQYYdNosJgaCM1r7k7igl4xC5Ulpv4AsGZbwmRveWJn7OqMkkcQNfBNqV0b2sFOuMt7rFQ+GYD3An0qiEUHPb+PTUolSCb+ALBGX89K1Ql9RXLpmL7FTbjG5H3b6XxM0DR1rdGPIGkOGwYEFBhtaHMyUWpYhmyGYxqa39k/2yTQY7xuRJEREJJpOEMqV9niN8ZBQ1TqUo1aZtUWrvmV50uj3IcFhwUXVydCEXKGHn3MA3vQ51dE+/XVLA9Bv4mnpFpxSLUtNZWpINADjYnNhFqT8daMGJjgFkpVhx56VzZ3w7ziwHTBLg9QfRNZic3ZcfKSH8KytyYpZ3Fi0sShHNQrHyy7alPzlbQ90jPuw5HXrBY54UEZ2NG/jIaBYpYedHNB7f+8uhVgDA5xcXJkWoMQA4M6fuqqFR7S4Rcq7PPClhuue0sSd0/lyaY7zxvXgTnVLH2t0Y8QU0PprY8AeC+I/txwEA/3hZFTIdMw+/t5pNat5aS5J2a4uQ8wt1nicFsChFNCti3W1rks4rv3u8C/6gjLl5aZibl6b14RCRzoiw84aeQY2PhCg8YnzvRIdbsyy0YFDGa4dDo3vXJPjWvbHUDXzMlJpWh/J3VGCYTqmJz5Mb2SkVtqIsB/LS7QgEZXzaom0nZ6z8/pNmnO4axJw0G26/qHLWt6eO8CXh+zRZlkdDziv0nScFsChFNCvqBr4kbTXn6B4RTUW80WhkpxQZRGlOCjLsFvgCMk52DmhyDAea+tDaP4I0mxmXzE+O0T1gTFGKnVLTMkyn1BSZUv3DPrhH/ADYKRWOsWHnh5r6tD2YGPAFgnhiR6hL6quXVSE9Cjm1YuNjc2/yFaUae4bR4fbAapawvCxb68OZFotSRLNQnC06pZLvBCoYlLHrWCcA4EqO7hHRBDi+R0YjSRJq1BE+bboRXlW6pK5aVKjrEOtoE/lIIsTbCBq6B9E1EP/OLrGhUP+ZUqGiQLvbg8BZYdPiw4rcNBtSbVyUE46lJWIDX+LlSv324yY09gwjL92O29ZURuU2RadUcxJ2Sok8qaUlWYb4PcKiFNEsqJ1SSZgp9WmLC51uD9JsZlw4V/9toUQUf2pRqptFKTIOEXZ+VINcKVmW1Typa5Jg695YRsuU6nR7sP7xv+L6n72HYW98M36M0imVn2GH2SQhEJTPKd6JkPNSju6FbXmZUpRqSqyilMcfwM+ULqn/fUU1UmzRKaKUZItMqeR7nybypGoNkCcFsChFNCtiVj4ZA/TE6N7F8/Jgt+i/Ak9E8SeKUq4RP/qGvBofDVF4RK7UZxp0Sn3a4kJT7zBSrGZcviC5upALlKKU2Cynd0daXRjxBdHcN4z/98OGuN53h0E6pcwmCQUZE2+qFiHnZRzdC9sSpVPqZOcABjx+jY8mel76qBEt/SNwZjrw96vLo3a7JTnJ3CllnDwpgEUpolkpUtpC210jCJ7VlpzodtSFilIc3SOiyaTYzMhX3pBwhI+MQozvHW2Lf6eU6JJaW5MftW4BoxD5Q26PH4MGeMPd0D26wOHnu07G7ZhlWTZMpxQwZinQWUWpJoacR6wgw4GiLAdkGfg0QUb4RnwB/GzHCQDAvVfOi+qoWbIGnfcOenGiI5SJuJJFKaLEV5hhh0kC/BO0JSeyrgEPDiohi2tZlCKiKTBXioxmYWEGJCk0nhXP3+1jR/c2JtHWPSHdbkGaUohrN8ACmfoxY8ndg15s210fl/t1jfjh8Yc2Q+p9+x4wOlXQdlbURaMSPs2Q88iIsPNEGeH77w8a0OH2oCQ7BTfXlkX1tkVRqnfIhyGv/gvd0bK3IdQlVZ2fhtx0/b9GACxKEc2KxWxSP6VqMUgGQjS8XdcJWQbOK85Ut+UQEU2ERSkymjS7BRXK/7fxzJU62uZGffcQ7BZT0n7gUygKGEYoSnWFOqXEeMz/76+n4B7xxfx+O5S/m0yHxRABxs7MiTdVi6Dzshx2SkViWWk2gMQIOx/y+vH02ycBAF+7ch5sluiWJjIdVmQ4QiH6ydQt9ZGSJ3WhQfKkABaliGZNbUtOohc7ju4RUbhEUaqRRSkyEJErFc8NfGLr3uUL8qOyDt2ICjOMkytVr4zv3XflPFTnp6FvyIfn3q2P+f2O5kkZ40NBZ9a5mVKyLKNJ6ZTi+F5kxAa+Q8rEgpFte78BXQNelM9JxU0rS2NyHyXqBj79F7qj5WMlT8ooo3sAi1JEs1acPfGsfKLyBYL467FOABzdI6LpiaJUAzfwkYGIDXxH2uJYlBKje0m2dW8sp0E6pQJBWQ3qrs5PxzfXLQAA/Ne7p9A/FNtuKTVPygCjewDgVDdVjz6n3YNeDPsCkKTR82gKjxjfq+8eivn/a7HkHvHhF38NdUl946r5sJpjU5ZItlypEV8Ah5TRTnZKESWRIvWXbXK82O1t6IV7xI85aTYsV1qIiYgmU57L8T0ynkVK2PmROI3vnehw43jHAKxmCVctKozLfeqRKLToPVOqtX8Y3kAQVrOE4uwUfGFpEWqcGXCP+PFf756K6X23K11khQYIOQfGZkqNPqeic7Yww8ENzhHKTrWpH/YcMvAI3/Pv1aNvyIeq/DTccH5JzO5H7ZTqTY73aYea++ENBJGXbkdFrnG6EFmUIpol8cs2WTKldh4Nje5dviAfZpOk8dEQkd6Jk+eWvmH4AkGNj4YoPGJ870SHOy7/3756KDS6d+n8fGQ6rDG/P71yKiNpei9Kic7PsjmpMJskmEyS2i313Lun0TPojdl9d7hDfzf5RumUyhwtSslyaFN1ozq6x5DzmVgqws6b+7Q9kBnqH/bhmXdCxdtvrlsQ0/cTydYp9VG9yJPKgSQZ530ai1JEs6R2SiXJi90OpSjF0T0iCkdBhh12iwlBOXlOCsn4SnNSkG63wBeQcbJzIOb39xclT2rDkuQd3QNGc5LadZ4pJfKkKnPT1O+tP68QS0oyMegNqGNJsdBhsE4p8Zx6A0G1WNfUGyrqlTLkfEaWKblSBxuN2Sn17Dun4BrxY2FhBr64NLabRsV4aHOSnH+IPKlaA43uASxKEc1aUfa5bcmJqrFnCMc7BmA2Sbh8fr7Wh0NEBiBJEnOlyHAkSUKNMzTCF+sNfKe7BnGk1QWLScLVi5N3dA8YLWDo/ZxKvJaNHY+RJAmbPx/qlnrh/QZ0umNTWBOdUkbJlLJZTMhT1tKLXCmRx1WWw06pmRAb+Iw4vtc76MVz79UDAO7//HyYYjx1UZojgs4TvygVDMr4WOmUqjVQyDnAohTRrBUrnVLtbg8CQVnjo4mtXcrWvZXlOchKTd7xAiKKjChKMVeKjCReG/hePRwKOF9TnYvsVFtM70vvCpVCS4d7dNRLj+q7zu2UAoC1Cwuwoiwbw74Afr4rNt1SaqaUQbbvAaNRF2IsU+2U4ua9GVlSEnptau4bRteAvrsKz/aLv57CgMePxUWZuHpx7DtDxfheW/9Iwr9PO94xANeIHylWMxYXZ2p9OBGJS1HK4/FgxYoVkCQJ+/fvn/RyPT09+NrXvoaFCxciJSUF5eXl+PrXv47+fuNVgSl55GfYYTZJCARl9dOrRMXRPSKaCbHyu5FFKTKQGhF23hbbTqnXlNG9a2I8xmIEBcpImi8gxzSXabYm6pQCxndL/feHDVHPxpLl0XNNo4zvAaNbFUWnVJPIlOL43oxkOKyoyg8VRMWmNSPodHuw7f16AMDmzy+IeZcUEHpNsZgk+INyzLoX9eLjhlCX1Pnl2THbZhgrcTnaBx98EMXFxdNerqWlBS0tLXj00Udx+PBhPP/883jttddw5513xuEoiWbGbJJQmDG+LTkRDXsDeP9kNwBgbQ1H94gofOyUIiOKR6dUY88QDjb1wyQh6Uf3gNCoV25aqFtMr7lSwaCMhp6JO6UA4NL5ebiwMgdefxBP7TwR1ft2jfgx4gsF7xtlfA8Yv4EvGJTVTWilHN+bMbEB+6CBilK/ePskhn0BLC/LxlWL4vMBt9kkqUXR5r7EPgcxap4UEIei1Kuvvoo33ngDjz766LSXXbJkCV5++WVce+21qK6uxpVXXol//dd/xZ/+9Cf4/f5YHyrRjBVli7DzxC1K7T7VBY8/iOIsBxYWZmh9OERkIKKbgJlSZCQLCzMgSaFP92M1IiO6pFbPzUVuunGKDLFUqPMNfB1uD0Z8QZhNEkomKKqEuqUWAgBe3NMY1SybTqVLKtNhgcNqjtrtxtrYTql29wi8gdDfnyhWUeSWKmHnhwyyga/dNYJffdAAINQlFc/NcGKErzmB36cB4zfvGU1Mi1Lt7e2466678Ktf/QqpqTNrz+zv70dmZiYsFsuEP/d4PHC5XOO+iOKtSP1lm7ghejuPdgIIje4ZacUoEWmvfMz4np5zYojGSrNbUKH8vxursPO/KHlS1yxN7q17Y4lcKb0WpcTmvdKclElHZNZU5+Ki6lx4A0H8bMfxqN236B4rMFCeFDCmU8o1rIacF2c7YDHYiJGeLCtVNvAZpFPq57tOwuMPorYiB5fNz4vrfZcoRalE3gDc2j+Mpt5hmCTg/HIWpVSyLGPTpk24++67UVtbO6Pb6Orqwg9+8AP84z/+46SX2bp1K7KystSvsrKymR4y0YyJX7YtCVqBl2VZzZO6knlSRBQhsfbb7fGjb8in8dEQha/GGRrhO9oW/Q89W/uH8cmZPkgSsP48FqUEp1rA0Oc5lQg5r5hgdG8skS3124+bcCZKXaJqnpSBRveA0e631v6R0ZDzbOZJzcZ5xVkwSaHOPb1vq/QFgnh5XxMA4Bvr5sf9w21RlBJjo4lIjO4tKspEun3iZh49i7go9dBDD0GSpCm/jh49iieffBJutxtbtmyZ0YG5XC584QtfwOLFi/G9731v0stt2bIF/f396ldjY+OM7o9oNoqUDXxtrsR8sTveMYDmvmHYLSZcVB3fTzeIyPhSbGYUKNl7zJUiIxG5Up/FIFdKjO7VVuQYrvMllkTYuV4zpeqVAlNl7tRFldrKObhsQT78QRn/sT063VJqp5SBQs6BMefJ/SNqp1TZHOZJzUaKzYwFSpzGwaY+bQ9mGnsbeuEe8WNOmk2T9xHFSdAptbchVJS60IB5UgAQcRntgQcewKZNm6a8TFVVFXbs2IHdu3fDbh9fya+trcUtt9yCbdu2TXp9t9uNDRs2ICMjA7///e9htU6+et5ut59zH0TxVpyd2J1SoktqTXUuUmzGyTAgIv2oyE1Fh9uDMz1DWF6WrfXhEIVFbOCLxfjeq4dCRamNS7h1byzRKaXX8b2G7slDzs+2+fML8Ndjnfj9J024d201qvLTZ3XfHer4nrHe+ziVouuQN4BPW0LjZty8N3tLS7JwtM2NQ839uFrH3ZY7lfcRly/IhzkOG/fOJt6nRTPfTW9EnlStAfOkgBkUpfLz85GfP/3mrSeeeAI//OEP1T+3tLRg/fr1eOmll7B69epJr+dyubB+/XrY7Xb88Y9/hMNhrE8CKDmJT4ASNVNKFKXWLuToHhHNTNmcVHxU38tOKTKUxUqn1ImOAfgCwait2e5wj+AjZX33hiX6fTOpBf1nSimdUnnTF1VWlGVj3aICvHWkA/+x/Tj+48vnz+q+25XxPaN1SqXYzMhOtaJvyIePlY6OsjksSs3WsrJs/HZvEw7oPFdqZ53yPkKjCBCx5TFRi1LuEZ+6Jba2wpidUjHLlCovL8eSJUvUrwULQnPV1dXVKC0tBQA0NzejpqYGe/bsARAqSF199dUYHBzEs88+C5fLhba2NrS1tSEQCMTqUIlmTWRKdbg98AWCGh9NdPUP+9SWUOZJEdFMibDzaGWrEMVDSXYK0u0WeANBnOocjNrtvv5pO2Q5VLQQoyUUoufte7Isq51S02VKCfcr2VJ/PNCCY+2z67jrVDqljJYpBYx2S/UMegGMFgpo5paJDXxNfbpdItLUO4Rj7QMwmyRcPn/6xpZYEM0D7hE/XCOJl2v5yZk+BOXQSKzToBstNV154PP5UFdXh6Gh0Anqvn378OGHH+LQoUOYN28eioqK1C9mRZGe5aXbYTVLkOVQYSqRvHO8E4GgjHkF6fxUi4hmTC1KsVOKDMRkklDjDI3wHYlirtSrh7h1bzKiKNU14NXdB32dAx4MeQMwSeEXVc4rzsLGJU7IMvD4W8dmdf9G7ZQCRj/AFXhOOXs1RRmwmiX0DvnQpNMQbzG6t7I8B1mpk0fyxFKa3YJs5b5bEzBq5WMxumfQLikgjkWpyspKyLKMFStWnPO9K664AgBwxRVXQJblCb8qKyvjdahEETOZpNHNIgnWGsqte0QUDRW5LEqRMYlcqSNR2sDXPeDBB6e6ATBPaiJzUm2wmkO5M3r7oK9B6fQszk6B3RJ+xuY31y2AJAF/OdSmZipFSpZlNVPKkJ1SWaNFPJvFhPx04z0GvbFbzOqG0IM6HeET7yOuqNGmS0pQN/D1Jd45yEfK5j2j5kkBGndKESWSYuWXbYvO17JGIhiU8XZdJwDmSRHR7IhPxVv7h+H166v7gWgqYgPfkSiFnb/5WTuCciikmN0i5zKZpDEb+PR1TlXfFX7I+VgLnRn44rJiAMBP35zZJj63x49hXyjOxOidUqXZKTBpEHidiJaWhkb4Djb3aXsgExjxBfD+yVABXusPt4vVopS+XlNmyxcIYn9jHwDjbt4DWJQiihoxw5tInVIHmvrQPehFht1i6Oo7EWkvP90Oh9WEoJy4YaOUmEQnwtEoje/95XBo6x4Dzienhp3r7IM+0SklOj8j8c1182GSgLeOtONgU1/E1xddUhkOiyE3IY/NuillMTZqRnOl9NcptftkNzz+IIqzHFhYmKHpsYhOqZYEO//4rMWFYV8AWSlWzJvldk8tsShFFCVFyrrRVp2dQM2GmAO/dEFe1DYOEVFykiSJuVJkSCJTqsPtQffA7MbJ+oa8eP9EFwBgI4tSk9Jr2Hl998w6pQCgOj8dN5xfAgB47M3Is6U6XCJPyphjb2M7pcoYch41y0qzAYSKUsGgvsLO1e3dNQWQJG0749TxPZ1mb83UR2qeVI6huw/5LpMoSsT4Xmt/4rzY7eToHhFFUfmc0Bs5FqXISNLsFrUzZrYjfG9+1g5/UEaNMwNVBv5UO9ZEUarNpc9MqZl0SgHAN66aD7NJwq66TnWzcbhEyLn4uzGaceN7OeyUipb5hemwW0xwe/xq0VQPZFkeLUrp4H1EcYJ2Sn2s5EmtNPhEC4tSRFGiju8lSKdUh2sEh5pDrcBX6OCXCREZn9oppaMTZ6JwLBIjfLMMO39NGd1jwPnUROGlQ0edUrIsj2ZK5UXeKQUAFblp+F8XlAIAHnuzLqLrivE9o3ZKjS2mlc1hp1S0WM0mLC4OvT6J83Y9ONExgOa+YdgsJlw0L1frw0FJTuIVpWRZxscNoU4pI+dJASxKEUWNGnSeIAF6u5QuqeWlWcg36AkQEelLufJGhJ1SZDRiA99ns8iVco348M7x0OjeNUs5ujcVZ1bovKNNR0WpnkEv3B4/gNEC+0x87ap5sJolvHeiW93CGI52dfOeMTulMhxWZDgsAGb390fnWq6M8B1o1E9RSnRJranKRarNovHRAMXZovtyBL5AYixbaegeQteAFzazCUuVbDGjYlGKKEpEplTXgCchNkuNnQMnIoqG8lyRKZU4n1RSchAb+I7OYnxvx5EOeANBzCtIx3yNQ3/1rlCH2/fqldG94iwHHNaZB42X5qTi5gvLAISypWQ5vBygDmV8z8gfFD60sQa3fq4CS4qN/QZab0RB4pCONvCJ9xFab90T8tLssJlDy1b09LoyGyJPallp1qxek/SARSmiKMlNs8GmhIEb/cXO6w/iXSWIVQ9z4ESUGMSn4409Q2G/ESPSg8VKUepEx8CMP2V/9XArAOAaBpxPqzBLFKX0kynVoIwdV8wg5Pxs962dD5vFhD2ne/DeifC6pToM3ikFALesrsAPblhi6EBmPVpWGipKHW52IaCDsPP+YR8+VjLT9PI+wmSS1G6pRAk7F3lStQYf3QNYlCKKGkmS1Fwpo88rf1zfgwGPH3npdsO3gxKRfohw2wGPHz2DXo2Phih8JdkpSLdb4A0Ecaoz8ky0QY9fHYvfwDypaYnCy4DHjwFlZE5rolOqMm/2o2fOLAduWV0OAPj3N+vCKtKLTimjZkpR7FTlpyPNZsawL4ATHQNaHw7ePd6FQFBGdX6a2iGtB2rYeYIspfqoYXTzntGxKEUURWKziJ4yEGZCtNxesTCfn2YRUdQ4rGY4lTebzJUiIzGZJNQ4QyN3R2aQK7WzrgMefxCVualYVMTRvemk2y1It4dyaPTSfR7NTikAuOeKajisJnxypk8tWE5GlmXDZ0pR7JhNEs5TPkQ+2NSn7cFAf6N7wugGPn28psxG94BH/YBkJYtSRDRWorzY7ajT5y8TIjI+dQMfi1JkMCLs/MgMNvC9ekjZure0CJLED3vCUZAZ6gjSS1FK7ZSKUudHQYYDt62pBDB9ttSAx49hXyB0vUx2StG5lqm5UtqGnQeDMt4+ps9c2hLlfVpTAozv7VXGI+cXpCMnzabx0cwei1JEUSQ6pVoN3Bba0D2IU52DsJgkXDI/T+vDIaIEUzYmVyoZPPXUU6isrITD4cDq1auxZ8+esK734osvQpIk3HDDDeO+v2nTJkiSNO5rw4YN4y5z7NgxXH/99cjLy0NmZiYuueQS7Ny5M1oPKWmJsPMjEYadD3sD2Kl82HMNR/fCJroqO3SSKxXtTikA+OplVUi1mXGouR9vfNY+6eVEl1SG3aKLTWakP8vKsgEAB5q0LUodbO5H14AX6XYLLtRZ1lGJ2jxg3PdpgsjsSoQ8KYBFKaKoKlIzpfTxqd5MiJbbCyvnINNh1fhoiCjRVOQmT6fUSy+9hM2bN+Phhx/Gvn37sHz5cqxfvx4dHR1TXq++vh7f+ta3cOmll0748w0bNqC1tVX9+vWvfz3u51/84hfh9/uxY8cO7N27F8uXL8cXv/hFtLW1Re2xJaMap9jAF1mn1NvHOjHkDaA0JwVLSjJjcWgJSYyp6SESoW/Ii74hH4DR17BoyE23446LKwEAP33zGIKThFSreVLskqJJiE6pI60uTbeAi/cRl87Pg9Wsr1JDcQIVpcTmvQsrjT+6B7AoRRRVRVmhF7s2l3Ff7HYquQZra/I1PhIiSkRifK+hO/GLUo899hjuuusu3HHHHVi8eDGefvpppKam4rnnnpv0OoFAALfccgseeeQRVFVVTXgZu90Op9OpfuXkjJ6UdnV14fjx43jooYewbNkyzJ8/Hz/60Y8wNDSEw4cPR/0xJhORKdXh9qB7IPzuHbF1b+MSJ0f3IiCKUnoY3xOvVwUZ9qh3Kt11aRUy7BYcbXPj1cMTF45Ft1hBBvOkaGIVuanIdFjg9QdxrD2ybs5o2lWnz9E9ACjJCb1Pa+4bNvQG4GFvAIeVMU29daPNFItSRFFUpKwabTVop9SQ148PToVWEzNPiohiIVnG97xeL/bu3Yt169ap3zOZTFi3bh1279496fW+//3vo6CgAHfeeeekl9m1axcKCgqwcOFC3HPPPejuHl0pn5ubi4ULF+KFF17A4OAg/H4/fvGLX6CgoAArV66c9DY9Hg9cLte4LxovzW5Ru2SOtoX3ps/jD2D7kdCbtI1LOboXiUIdZUrVK6N7lVEc3ROyU22489K5AICfvnUMgQm6pUSnVCE7pWgSkiRhWWk2AOCgRiN8He4R9b6vWKi/D7fFRMuQN4D+YZ/GRzNzB5r64AvIKMiwo1QptBkdi1JEUVSsdEp1D3oxogRSGsl7J7rh9QdRNicF1fnpWh8OESUg0SnV6hqBx2+818lwdXV1IRAIoLCwcNz3CwsLJx2je/fdd/Hss8/imWeemfR2N2zYgBdeeAHbt2/Hj3/8Y7z99tvYuHEjAoHQ36UkSXjrrbfwySefICMjAw6HA4899hhee+21cR1VZ9u6dSuysrLUr7Kyshk86sS3yClypcIr2r17vAsDHj+cmQ6sUN4wUnicaqeU9plS9V2hIno0R/fG+solc5GVYsWJjgH86UDLOT8XfwcF3LxHU1haKsLO+zS5f7FFcllpli67+hxWM/LSQ6HgzQYe4RMh5xdWzkmY7lsWpYiiKDvVCrsl9M+qrV/7T/Yipa5wXViQMC9yRKQveek2pNrMkGWgOQE24ESL2+3GrbfeimeeeQZ5eZMvmfjyl7+M6667DkuXLsUNN9yAV155BR999BF27doFILQ6/t5770VBQQHeeecd7NmzBzfccAOuvfZatLa2Tnq7W7ZsQX9/v/rV2NgY7YeYENQNfGGGnf9F2bq3YYkTJhN/r0ZCFGD0cD4lQs4r86LfKQUAmQ4r/vGy0Ljuf2w/Dn9gfCZQh1uM77FTiiYncqW06pTaqbyPuGKhfqctRNi5kc8/RJ5UbYLkSQEsShFFlSRJaoheqw5OoiIhy7Ku58CJKDFIkjSaK5XAI3x5eXkwm81obx+/Uau9vR1Op/Ocy588eRL19fW49tprYbFYYLFY8MILL+CPf/wjLBYLTp48OeH9VFVVIS8vDydOnAAA7NixA6+88gpefPFFXHzxxbjgggvwn//5n0hJScG2bdsmPV673Y7MzMxxX3Su0Q1803dKef1BvPlZqCh1DUf3IuZURm063COa57/Uq5v3YtMpBQCbLqrEnDQbTncN4nefNI/7mRhhZKcUTUVs4Ktrc8d9YsMXCOKd410A9B0BYvSw80BQHtcplShYlCKKMjGv3NpvrBe7I61utPaPwGE14XNVuVofDhElsGTIlbLZbFi5ciW2b9+ufi8YDGL79u1Ys2bNOZevqanBoUOHsH//fvXruuuuw9q1a7F///5Jx+mamprQ3d2NoqJQ0WNoKPR3ajKNP8UzmUwIBrXbyJQoxPjeiY4B+AJT/32+f7ILrhE/8jPsWFmROJ9ox0t+eqgryBeQ0TPo1fRYRNB5LDKlhDS7BXdfHuqWemL78XH/f3UqnVKF7JSiKRRnOZCbZoM/KIc9YhwtH9X3YMDjR166Te3Y0iO1KGWw5gHhWLsb7hE/0mxmdflGImBRiijKxAY+o3VK7VS6pC6uzoPDatb4aIgokYlOqTMJvoFv8+bNeOaZZ7Bt2zYcOXIE99xzDwYHB3HHHXcAAG677TZs2bIFAOBwOLBkyZJxX9nZ2cjIyMCSJUtgs9kwMDCAb3/72/jggw9QX1+P7du34/rrr8e8efOwfv16AMCaNWuQk5OD22+/HQcOHMCxY8fw7W9/G6dPn8YXvvAFzf4uEkVpTgrS7RZ4A0Gc6hyc8rKvKZvU1p9XCDNH9yJms5jU/Bctc6VcIz50K0WxWHZKAcCtn6tEfoYdTb3D+O3HTer32SlF4ZAkaUyuVHxH+MTo3uULCnQ9qmz08b2PldG988tzYDEnTikncR4JkU6ITimjtYWKXyYc3SOiWBNv7M4kcKcUANx888149NFH8d3vfhcrVqzA/v378dprr6nh52fOnJky5+lsZrMZBw8exHXXXYcFCxbgzjvvxMqVK/HOO+/Abg91UOTl5eG1117DwMAArrzyStTW1uLdd9/F//zP/2D58uUxeZzJxGSSsFD5dPpo2+SdCP5AEK9/qozuLeHo3kyJsGQtN/CJ4nleug0ZDmtM7yvFZsb/vqIaAPCzHcfh8Qcw4PFjyBsaxWKmFE1HbOA70BjfotQO9X2E/rbujSU6pYwadP6xMrqXSHlSAGDR+gCIEk1Rtn6COcPVO+jFvjOhFzkWpYgo1sT4XqIXpQDgvvvuw3333Tfhz0Q4+WSef/75cX9OSUnB66+/Pu191tbWhnU5mplFRRnY29CLz1pduH5FyYSX+fB0D3qHfJiTZsOquYmT+xFvziwHPmt1aVqUGs2Tit3o3lh/t6ocv3j7FFr6R/DinkZcMj+0+CDdbkGanW/daGpidC6eG/jOdA/hZOcgzCYJl87Xd1GqxOCZUh/XJ16eFMBOKaKoK84y3qzyX493IigDNc4M9cWaiChWyscUpbQOMCaKVI2SK3V0ig18rx4OdcCtP68woUYs4q0wM9QZ1KZhUUrkScV6dE9wWM2498p5AICndp5QO7UKMtklRdNbpozvnegYwKDHH5f7FBEgtRU5yEqJbTfhbJXkhN7ndLg98PjjGwY/W819w2juG4bZJGGFEmqfKPhbkijKnAYMOt/B0T0iiqOS7BRIEjDkDahZLURGMd0GvkBQxmuHQ1sXN3B0b1YKM8X4nnaZUvVdoU6pWIacn+3m2jKUZKegw+3B49uPA+DoHoWnINMBZ6YDQRn4tCU+YefifYSet+4JOalWOKyhEoiRplqA0Typ84ozE65rkkUpoigTnVJ9Qz4Me/VfgQ8EZbx9rBMAsHah/n+ZEJHxOaxmOJU3m8kwwkeJRWRKdbg96B44t1jycX0PugY8yEqx4qJqbrOdjdGiVPJ0SgGhkPevXxXqljrQ2Adg9O+CaDoi7PxgU1/M72vI68fuU90AjPHhtiRJhs2VEqN7ibjNlUUpoijLTLEg1RbaXmeEbqn9jb3oG/IhK8WKC8qztT4cIkoSybKBjxJPut2iFiiOtp07wveqsnXv84sLYeXo3qw4dVCUEplS8eyUAoAbLygdVwhjpxSFazRXKvZh57tPdsPrD6IkOwXzC9Jjfn/RYNQNfB8pnVKJlicFsChFFHWSJKkb+FoN0BYqWm4vW5DP3AsiipvyJAo7p8RTo3RLnT3CFwzKeE0pSm1c4oz7cSUakaOkVVFqyOtHhzvUDRfvopTVbMI3rpqv/pmdUhSuZUre0MGm2Belxo7uSZIU8/uLhtGwc/2/TxNcIz7UtYc+BKllpxQRhaMoyzibHXYcDY3uXanzFa5ElFhYlCIjG82VGt8p9UljH9pcI0i3W9StaTRzolOqa8ALXyAY9/uv7wq9PmWnWpGVGv8A5+tXlKA6P1QMi9f2PzK+pUqn1OmuQfQP+2J2P7IsY6eB8qSEYgNu4NvX0AtZDo0RFyRggZpFKaIYEJ1Seg/Qa+0fxpFWFyQJuHyBcX6ZEJHxleeyKEXGpW7gaxvfKfXqodDWvXWLCmC3mON+XIkmJ9UGqznUfSE6luKpQRnd06ogZDZJ2PaVVfjRjUtxlYHe9JO25qTZUKpsmfs0hiN8de1utPSPwG4x4XNVxsnPKzFgppTIk6qtSLzRPYBFKaKYKBIVeJ0XpXbVhbqkzi/Lxpw0m8ZHQ0TJRHRKNbIoRQa0WOmUOt4+oHbwyLKs5kltXMqte9FgMkkoyNAuV6peybybG8eQ87OV5qTiy6vKYTIZYzSK9GF5aTYA4EAMR/h2KtMWF1XnIsVmnCK8ETulRvOkEm90D2BRiigmitVMKX2/2Ik5cG7dI6J4E0WpNtcIRnz631RKNFZpTgrSbGZ4A0Gc6gx10xxq7kdz3zBSbWZcvoAj8dFSKHKlNPigT+tOKaKZEhv4DjX3xew+jDi6B4zvlJJlWeOjmZ7XH8R+ZQtnLYtSRBQupyhK6ThAz+MP4L0TXQCMscKViBLLnDQb0mxmyDLQZLANOEQmk4SaovEjfH85FOqSWltTAIfVOF0DeifOqbTplFI27+Vp1ylFNBNiA1+sws77h3zYeyY0UnaFwT7cdmY5IEmAxx9E96BX68OZ1qct/fD4g8hJtaI63xgbDiPFohRRDIi2UD13Sn14qgdD3gAKM+04rzhT68MhoiQjSRLKOMJHBiY28H3W6lJG90J5Utcs4eheNInxvTaXFplSodcmdkqR0SxROqWaeofRPRD9fztvH+9EIChjfkG6+rvcKGwWEwoyQh2YRhjhE3lSKyvmGGbDYaRYlCKKARF07hrxY9Dj1/hoJjZ2dC9RX+CISN8qGHZOBiY28B1tdeNIqxsN3UNwWE24YiFH96JJdEp1xLlTasQXQKsyMljJohQZTKbDiqq80P+3h2IQdr7LoKN7gpFypRI9TwpgUYooJjIcVmTYLQD02S0lyzJ21ilFKYP+MiEi4xO5UqIbgchIFhWFOqWOtLrULqnLF+QjTfn9T9EhMqXa4lyUEsXyDIcFOanWuN43UTSouVJRHuELBGXsOhYKOTfq+wiRK6X3+ABZlvFxg7J5rzIxN+8BLEoRxYz4ZK9Fh7lSp7sG0dA9BKtZwsXz8rQ+HCJKUqIoxU4pMqKFzlCnVIfbg99+3AQAuIZb96KuMFObTKn6LiVPKjeNHeVkSMtitIHvQFMfega9yHBYsLLCmN07JWqnlP7ep411qmsQPYNe2CwmLClJ3LgVFqWIYqRIebFr02BbzHTE6N7qublI5ye6RKQRZkqRkaXbLeO2SNrMJsOOsujZaFEqvplSo3lSxsrLIRKWxWgDn9i6d9n8fFjNxiwnGGV8b6+SJ7WiNBt2S+Iu0DDm/0VEBlAsOqV0OL73/sluAGDuBRFpSoQHn+kZMsRaZqKziRE+ALh0fh4yHBzzijZRlBrw+DEQx5xOdfMe86TIoM4rzoRJChV0o9lpmAgRIKJTqlnnRSmRJ1WbwHlSAItSRDFTlKVs4NNZW6gsyzjQ2AcAhm25JaLEUJKdAkkChn0BdMZgOxBRrNU4R8cpNnJ0LybS7Ra1qzueI3yiKMVOKTKqVJsF8wtChfODURrh63CN4HCzC5Jk7A+3jdIpJfKkLkzgPCmARSmimCnSaadUU+8wuge9sJoldXMQEZEWbBYTipUCPkf4yIjE71GLScLnFxVqfDSJS4Sdx7Uo1RV6TarMY6cUGddo2HlfVG5PdEktK81GXro9KrepBdEp1T3oxYgvoPHRTKzT7cHprkFIEnBBeWI3ErAoRRQjRdmhopTeMqUOKL+UFhVlwmFN3NlkIjKGsjmhE0OGnZMRXTI/D6sq5+Duy6uRxQ1tMRPvsHOPP6B+qMjxPTIykSt1sDk6nVIil3atgbukACAzZbQDU68jfHsbQqN7CwoyEv73C4tSRDGiju/prSiljO4tVzZyEBFpSd3A163Pk0KiqaTbLfjN3WvwrfULtT6UhOaMc9h5Y88wZBlIs5mRl26Ly30SxYLYwHewqX/W2Y1efxDvHu8CAMMvdZAkCcXZYlO6Ps8/PlZCzhM9TwpgUYooZsT43oDHD9eIT+OjGSXWwi4vy9b2QIiIMD7snIhoIgWZ8e0+b1DzpNIgSVJc7pMoFmqcGbCYJPQMemfdEfRRfQ8GvQHkpduxpDgrSkeoHb3nSn2UJHlSQJyKUh6PBytWrIAkSdi/f39Y15FlGRs3boQkSfjDH/4Q0+MjioU0uwWZjlBbqF5G+PyBIA4pRakVZcb/ZUJExlcmOqV6BjU+EiLSK6eSKdXhjs/5VH23yJNiyDkZm8NqxkJnKOz80CzDzseO7plMxi/Wqhv4evVXlBry+vGpMnLJTqkoefDBB1FcXBzRdR5//HF+MkGGp7cK/InOAQz7Aki3W1CVl6714RARjY7vsVOKiCZRqGGnFJHRiRG+A7MsSu1UilJGH90TxPu0Zp1tSgeA/Y198AdlFGU51OJZIot5UerVV1/FG2+8gUcffTTs6+zfvx///u//jueeey6GR0YUe2KETy+5UiJPallpVkJ8wkFExieKUu0uj2434BCRtgqz4psppXZK5bJTioxPhJ0fau6b8W3Udw3iVNcgLCYJF8/Pi9KRaatEZ80DY4k8qZUVOUnRqGOJ5Y23t7fjrrvuwh/+8Aekpob3oj40NIS///u/x1NPPQWn0znt5T0eDzye0V9QLpdrxsdLFG1OEXaukxe7/Y3MkyIifclJtSLDboHb40dT7xDmFWRofUhEpDOiU6rDPYJgUI75B2vslKJEsrRE2cCnhJ3PpMixsy7UJXVh5RxkOhJjE1xJjuiU0sf7tLH2KnlStRWJP7oHxLBTSpZlbNq0CXfffTdqa2vDvt7999+Piy66CNdff31Yl9+6dSuysrLUr7KyspkeMlHUFeu0U4qb94hILyRJUnOlGro5wkdE5yrICGVK+QIyeoe8Mb0vXyCIJiVjppJFKUoAC50ZsFlMcI/41S7ASO1IsNE9YHR8r7V/GMHg7DYTRlMwKGPfGbF5L/FDzoEZFKUeeughSJI05dfRo0fx5JNPwu12Y8uWLWHf9h//+Efs2LEDjz/+eNjX2bJlC/r7+9WvxsbGSB8SUcwUqS922helhr0B1LW7AQAr2ClFRDrCXCkimorVbEJeug0A0OaK7TlVc+8wAkEZDqtJLYYRGZnVbMLiokwAwMGmvoivP+jx48NTPQCAtQlUlCrMsMNskuALyOgaiM9ocDiOdwzAPeJHqs2MGmdydI9HPL73wAMPYNOmTVNepqqqCjt27MDu3btht49/Ma+trcUtt9yCbdu2nXO9HTt24OTJk8jOzh73/ZtuugmXXnopdu3adc517Hb7OfdBpBeiU6qlX/u20E9b+hEIyijMtMOpHBcRkR6U57IoRURTK8x0oGvAiw6XB+dFtj8pIqfF6N6cNOZvUsJYVpqF/Y19ONTUj+tXlER03fdOdMEbCKJsTgqq8xOne9BiNsGZ6UBz3zCa+oZRkKmP90didG9FWTYs5rjspdNcxEWp/Px85OfnT3u5J554Aj/84Q/VP7e0tGD9+vV46aWXsHr16gmv89BDD+Ef/uEfxn1v6dKl+OlPf4prr7020kMl0pwo/rT2jcx4hjta9qsh59maHQMR0UREp1Qji1JENInCTAc+bXHFvFOqoUvkSTHknBJH6Py/AQdnsIFvZ10nAODKhQUJF7pdnB0qSrX0DeOCcn3kN4mi1MokyZMCYhh0Xl5ePu7P6emh9fPV1dUoLS0FADQ3N+Oqq67CCy+8gFWrVsHpdE4Ybl5eXo65c+fG6lCJYqZICTof9gXgGvYjK1W7YECxBpaje0SkN+XMlCKiaYiw8/YYF6VE5s7cvMTpCCESG/gOK5MT5jC7AGVZxi4l5DyRRveEUK5Ur6428Ik8qQuSqCilaT+Yz+dDXV0dhoZ4EkqJKcVmRo5SiNJ6hI8h50SkV2MzpWRZP2GjRKQfhZmhuI5YF6W4eY8SUXV+OlJtZgx5AzjVORD29Y60utHaPwKH1YTPVeXG8Ai1UaLk/zb36qMo1T3gwWmlW/OCsuQpSsWsU+pslZWV55xoTvS9s/HklIyuKCsFvUM+tPYPY5ESMhhvPYNeNatlqfJJCRGRXhRnp8AkAR5/EJ1uj25yHYhIP5xqp1RsA4lFx2Ylx/cogZhNEpYUZ2FPfQ8ONPVjfmF4Ado7lS6pi6vz4LCaY3mImhAb+Jr7tF9KBQD7zvQBAOYXpGs6YRNvyZGcRaShIhF2ruGL3QFl00ZVfhqyUpLnBY6IjMFmMaknhgw7J6KJiPG9thhuNPYHgmjsDb0GVXB8jxKM+GD6UAQb+HYeTdzRPWC0U0ov43sfN4S2HCZTnhTAohRRzBVlx/4kajpidG8FR/eISKeYK0VEUxFFqQ537M6nWvtH4AvIsFlMKGLHJiUYkSt1sDm8sPPeQa+ab5SwRakc0Smlj6LUvobky5MCWJQiijkRdq5lppSaJ8WQcyLSqbG5UkREZxOZUl0DXnj9wZjcR72SJ1U+JxWmMIOgiYxCbOD+rMUFX2D6f0N/Pd6JoAzUODPUjqJEIyZa+od9GPD4NT0Wrz+oLqaqZVGKiKKpWOmUatVofE+WZXX9K4tSRKRXZUpRqpFFKSKawJw0G6zmUKGocyA2uVL1zJOiBFYxJxUZDgs8/iCOtbunvfwOZXTvioWJ2SUFABkOKzIdoZjtVo27pT5t6YfXH0ROqjXptn+yKEUUY87M0CcLbTHeFjOZpt5hdA96YTVLWFQUXqghEVG8sVOKiKYiSRIKMmIbidDQxc17lLhMJglLS0Su1NQjfIGgjLePdQIArkzQ0T2hJCd0/tGkcVFqrzK6t7IiB5KUXJ2aLEoRxZjolGrpG9Zkm6QIOV9clAm7JfG2ZhBRYqjIZVGKiKbmVEZtOmL0QR87pSjRiRG+6XKl9jf2om/Ih6wUKy4oz479gWmoZMx7NS2J/K5ky5MCWJQiijkRzOnxB9E75Iv7/TNPioiMQHRKdbg9GPYGND4aItIjkSsVq+7zhm52SlFiU8POp9nAJ0b3LluQD4s5sUsGxTrYwCfL8minVDmLUkQUZQ6rGblpNgDavNgdaFTypLh5j4h0LCvFigwl10GsZCciGkt80Nfuin6mVDAoo6FHdEqxKEWJSYzv1bW5MeKb/AOgHUdDo3trF+bH5bi0JELcm3u1K0o19Q6j3eWBxSSp3WzJhEUpojgoyo5tBsJk/IEgDjUz5JyI9E+SpNFcqW4WpYjoXKNFqeifT7W6RuD1B2ExSWr0AlGiKc1JwZw0G3wBGXVtE4edt/WP4EirC5IEXL4g8YtSo51S2uT/AqOje+cVZyLFlnxxKyxKEcVBUVboxa61P74V+OMdAxj2BZBht6AqybY4EJHxMFeKiKbijGFRSoScl89JTfhxJUpekjQadj7ZCN/OutDo3oqybOSm2+N1aJoRRalmDcf3RkPO52h2DFriKy5RHBQrwZwtce6UEnlSy8qyYDIl1xYHIjKeMm7gI6IpFMQwU0qEnFcw5JwS3Giu1MRh5yJP6sqFib11TyjNGd2U7g8ENTmGsZv3khGLUkRx4BSdUnGuwIvNe8k4m0xExlPOohQRTUGM73XEIFOKIeeULMT7gkMTbODz+AN470QXAGBtTXIUpfLT7bCaJQSCMjrc0X9tmc6gx48jrS4AwAUV2XG/fz1gUYooDkQ2QWucO6X2M+SciAyERSkimoooSg14/Bjw+KN62/VKUaqSnVKU4ESn1LF2N4a84/8d7TndgyFvAAUZdpxXnKnF4cWdySTBKaZaNFlK1YegHApcF5EvyYZFKaI4GM2Uil9Rasjrx7H2UIDhCoacE5EBVMwJdSg09gwhGJQ1Phoi0pt0uwXp9tCWzmjnSjWI8T1mcFKCK8x0oCDDjqAMfNbiGvczMbq3dmEBJCl5oj9KNMyVEqN7FyTp6B7AohRRXBRljW7fi9cbrU9bXAgEZRRm2tXqPxGRnhVlO2A2SfD4g5q00BOR/hUquVLtUfygT5blMZ1SLEpR4hMjfGfnSu0URakkGd0TtAw7/1jkSZVnx/2+9YJFKaI4KMx0QJIAbyCI7kFvXO5ThJxzdI+IjMJqNqnjzhzhI6KJiBG+dnf0ilIdbg9GfEGYTZLaMUGUyEbDzvvU753qHEB99xCsZgmXzM/T6Mi0If7dx3t8LxiUse9Mcm/eA1iUIooLm8WEPGWlalucRvj2i6IUR/eIyECYK0VEU3Fmiu7z6HVT1neFuqRKslNgs/DtESW+paIoNSbsfGddJwBg1dw56phsslDH93rjW5Q60TkA94gfKVYzFhVlxPW+9YSvukRxUiwC9Prj82InNu8xT4qIjIRFKSKaSoHolIpippSaJ8WQc0oSy0pCRalTnYNwj/gAjBndW5hco3vA6PheS198l1KJPKkVZdmwmJO3NJO8j5woztSw8zi0hXYPeNDYE7of8UkIEZERlCth52eUfBciorGcIlMqikWp08yToiSTm25Xu4MONfdjwOPHh6e7AQBXJlmeFACU5GgzvieKUiuTOOQcYFGKKG5E2Hg8NvCJVtzq/DRkOqwxvz8iomhJtE6pp556CpWVlXA4HFi9ejX27NkT1vVefPFFSJKEG264Ydz3N23aBEmSxn1t2LDhnOv/+c9/xurVq5GSkoKcnJxzbofIqApj0ikVKkqxU4qSiciVOtTUj3ePd8EXkFGZm4qq/HSNjyz+ipXmAbfHj/5hX9zudx+LUgBYlCKKGxHeG4+i1AHmSRGRQY0WpeK/ASfaXnrpJWzevBkPP/ww9u3bh+XLl2P9+vXo6OiY8nr19fX41re+hUsvvXTCn2/YsAGtra3q169//etxP3/55Zdx66234o477sCBAwfw3nvv4e///u+j9riItFSYJYpS0cyUChXB5+axU4qSh7qBr7kfu+pCv5euSMLRPQBIsZkxJ80GIH7dUj2DXpxS8uzOT+LNewCLUkRxo47vxSFTShSlmCdFREYjilJdAx4Mef0aH83sPPbYY7jrrrtwxx13YPHixXj66aeRmpqK5557btLrBAIB3HLLLXjkkUdQVVU14WXsdjucTqf6lZMz+gmr3+/HN77xDfzkJz/B3XffjQULFmDx4sX40pe+FPXHR6QF0SnV4R5BMCjP+vZkWR7TKcWiFCUP0Sl1oLEPO5WiVDKO7gnx3sAnRvfmFaQjO9UWl/vUKxaliOJEdErFOkBPlmUcaAqN7y1XPgEhIjKKrFQrslJCY8eNBu6W8nq92Lt3L9atW6d+z2QyYd26ddi9e/ek1/v+97+PgoIC3HnnnZNeZteuXSgoKMDChQtxzz33oLu7W/3Zvn370NzcDJPJhPPPPx9FRUXYuHEjDh8+POXxejweuFyucV9EelSQEcqU8gVk9Ax5Z317XQNeDHoDkCSgbE7KrG+PyCiWKGHnTb3DaHd5kGozY3XVHI2PSjvivVpznItStUk+ugewKEUUN06lU6rdFZ1P9ibT1DuMnkEvbGYTapJ4tSgRGZfolmowcNh5V1cXAoEACgsLx32/sLAQbW1tE17n3XffxbPPPotnnnlm0tvdsGEDXnjhBWzfvh0//vGP8fbbb2Pjxo0IBAIAgFOnTgEAvve97+Ff/uVf8MorryAnJwdXXHEFenp6Jr3drVu3IisrS/0qKyuL9CETxYXVbEJeeqirIBq5UuJ1pjgrBXaLeda3R2QUWSnWcSOrF8/LS+p/A2IDX7yKUiJP6gIWpViUIoqXwgw7TBLgD8roGoheDsLZ9iuje4uKM5P6FwsRGVeihZ2Hw+1249Zbb8UzzzyDvLy8SS/35S9/Gddddx2WLl2KG264Aa+88go++ugj7Nq1CwAQDAYBAP/8z/+Mm266CStXrsQvf/lLSJKE3/72t5Pe7pYtW9Df369+NTY2RvXxEUVTNMPO67tDrzOVeQw5p+SztGR0S/faJM2TEkbH92Kf/+v1B3GgqQ8AQ84BwKL1ARAlC4vZhIIMB9pcI2jpH0GBckIVbWqeVGnW1BckItKpMqUo1WjgolReXh7MZjPa29vHfb+9vR1Op/Ocy588eRL19fW49tpr1e+JApPFYkFdXR2qq6vPuV5VVRXy8vJw4sQJXHXVVSgqKgIALF68WL2M3W5HVVUVzpw5M+nx2u122O32yB4kkUYKMx34tMUVlbBz5klRMltWmoU/HmgBAKytydf4aLQlilLNvbE/9/is1QWPP4jsVCuquGCBnVJE8eRUNsa0xrAtVFTduXmPiIxKrGU3cqeUzWbDypUrsX37dvV7wWAQ27dvx5o1a865fE1NDQ4dOoT9+/erX9dddx3Wrl2L/fv3TzpO19TUhO7ubrUYtXLlStjtdtTV1amX8fl8qK+vR0VFRZQfJZE2RKdUWxQ2GqudUrnslKLks6Y6F5IUyjUSS5mSVXEcO6VEntTK8hxIkhTz+9M7dkoRxVFxtgP7G4HWKJxETcQfCOJQcyjkfBlDzonIoNRMKQMXpQBg8+bNuP3221FbW4tVq1bh8ccfx+DgIO644w4AwG233YaSkhJs3boVDocDS5YsGXf97OxsAFC/PzAwgEceeQQ33XQTnE4nTp48iQcffBDz5s3D+vXrAQCZmZm4++678fDDD6OsrAwVFRX4yU9+AgD427/92zg9cqLYKswMdfV1uKOXKcVOKUpG5xVn4ZWvXQJnjCY4jKQkR8n/dY/AFwjCao5d/w7zpMZjUYoojsQnEK39semUOtY+gBFfEBl2C1tBiciwRFGqqWcYwaAMk8mYnyLefPPN6OzsxHe/+120tbVhxYoVeO2119Tw8zNnzsBkCv+k12w24+DBg9i2bRv6+vpQXFyMq6++Gj/4wQ/Gjd795Cc/gcViwa233orh4WGsXr0aO3bsQE4OT34pMTij1CklyzJOd4WKUpUsSlGSOq+YkR8AkJtmg81igtcfRFv/iBolEG2yLOPjhtDiEeZJhbAoRRRHRcr4XkuMOqXE6N6ysizDvokjIirKcsBikuANBNHuHjH0SMF9992H++67b8KfiXDyyTz//PPj/pySkoLXX3992vu0Wq149NFH8eijj4Z7mESGMhp0PrtMqd4hH9wjfgCjxXAiSk6SJKEkOwWnuwbR3Dccs6JUc98w2l0eWEwSlnOyBQAzpYjiSu2UilGmlAg55wscERmZxWxS2+jPdBt7hI+Ioi9a2/fqldE9Z6YDKTZuLCZKdqMb+GKX/yvypM4rzuTrjoJFKaI4KsqOXjDnRPaLohRDzonI4BIlV4qIok9kSnUPeuH1B2d8OyJPqjKPXVJEFMr/BYDm3tgVpZgndS4WpYjiqDhLBOh5EAjKUb3tIa8fx9rdAIAVLEoRkcGJtvlGFqWI6Cxz0mywmkMxBbMJO6/vEpv3mCdFRGM28MUo/xcA9p5RNu+xKKViUYoojvIz7DCbJASCclQ2xox1uNmFoBxqQS/kBg0iMjjRKXWGRSkiOoskSSjImH2uFDfvEdFYYnyvuS82Uy2DHj+OtIaaCFiUGsWiFFEcmU0SCjNCLectUX6xO6iEnC8v4wYNIjI+FqWIaCrOrNnnStV3i04pju8R0ZiiVG9szj0ONPUhEJRRnOUw9BKXaGNRiijOipQXu2jnSjFPiogSSTnH94hoCiJXajZFKXZKEdFY6vhe3whkObpRKwCwt555UhNhUYoozoqUT/ZaozyrfEDplFrBzXtElADKlc6FrgEvBjx+jY+GiPRGRBW0zbAo1T/kQ++QDwBQwU4pIsLoUqphXwB9yutDNIk8qVoWpcZhUYoozsZW4KOle8CDxp5hSBKwpJTje0RkfJkOK7JTrQDYLUVE5xJFqY4ZZko19IS6pPIz7EizW6J2XERkXHaLGflK1EpzX3QbCIJBWd28t7JiTlRv2+hYlCKKM6f6yV70XugONvUDAKrz05HpsEbtdomItMRcKSKajHo+NcM4BOZJEdFEitWw8+gWpU52DsA14keK1Yyaooyo3rbRsShFFGfFSltoNDul1Dwpju4RUQJhrhQRTaZAZErNcJtxQxfzpIjoXKXqVEt0i1J7lS6p5WVZsJpZhhmLfxtEcSY2LUQzU0rNk+LmPSJKIKIo1dDNohQRjSc6pdpn2Cl1Wgk5Z6cUEY0lGgiae2NTlFrJPKlzsChFFGciQK/D7YEvEJz17cmyjAPcvEdECYjje0Q0GZEpNegNzGgZgih2s1OKiMZS83+jvJRKhJyzKHUuFqWI4iwvzQ6rWYIshwpTs9XYM4zeIR9sZhNqnJlROEIiIn3g+B4RTSbNbkGGElA+k1ypBrVTikUpIhpVIjKlotgp1TPoxanO0GvOBeUsSp0t5kUpj8eDFStWQJIk7N+/f9rL7969G1deeSXS0tKQmZmJyy67DMPD0a1SEmnJZJLUT/daozCrvF8Z3VtUnAmbhXVmIkoc5cpYTVPvMAJBWeOjISK9EblSHa7IilLuER+6BrwAgIo8ju8R0ajRoPPo5f+KrXvzCtKRnWqL2u0mipi/g33wwQdRXFwc1mV3796NDRs24Oqrr8aePXvw0Ucf4b777oPJxDfalFiKs0Rb6Oxf7MTo3opS5kkRUWIpykqBxSTBGwiiLcI3nUSU+JxZYqNxZK8PYnQvN83GrcVENI7olOoa8GDEF4jKbaqje+ySmpAlljf+6quv4o033sDLL7+MV199ddrL33///fj617+Ohx56SP3ewoULY3mIRJoQJ1HR6JRinhQRJSqzSUJpTgrqu4dwpntIPVEkIgKAwgwl7NwVWRzCaJ4Uu6SIaLzsVCtSbWYMeQNo7R/B3LzZj/gy5HxqMWtBam9vx1133YVf/epXSE2d/gW/o6MDH374IQoKCnDRRRehsLAQl19+Od59990pr+fxeOByucZ9EemdCDtvnWWnlC8QxOGWfgAsShFRYipjrhQRTaIwSxSlIjufqmeeFBFNQpKk0bDzKDQQ+AJBtYngAhalJhSTopQsy9i0aRPuvvtu1NbWhnWdU6dOAQC+973v4a677sJrr72GCy64AFdddRWOHz8+6fW2bt2KrKws9ausrCwqj4EolsT4Xusstzoca3djxBdEhsOCuTyxIqIEJDoZuIGPiM5WmBHKlIq0KCVCzrl5j4gmMporNfui1GctLnj8QWSnWlEVha6rRBRRUeqhhx6CJElTfh09ehRPPvkk3G43tmzZEvZtB4NBAMBXv/pV3HHHHTj//PPx05/+FAsXLsRzzz036fW2bNmC/v5+9auxsTGSh0SkiaKs6HRKHWhUuqRKs2EySbM+LiIivREb+BpYlCKis8w0U6peGd+rZMg5EU0gmhv4xOjeBeU5fL82iYgypR544AFs2rRpystUVVVhx44d2L17N+x2+7if1dbW4pZbbsG2bdvOuV5RUREAYPHixeO+v2jRIpw5c2bS+7Pb7efcD5HeFYmg81ludRjNk2LIORElJlGUYqcUEZ2tQNlm3BFxphQ7pYhociVK1Eo0xveYJzW9iIpS+fn5yM/Pn/ZyTzzxBH74wx+qf25pacH69evx0ksvYfXq1RNep7KyEsXFxairqxv3/WPHjmHjxo2RHCaR7olMqa4BD7z+IGyWmU3SHmjqAxDqlCIiSkTMlCKiyTgzRzOlgkE5rC6EIa9fDUavZNA5EU1AzZSaZdSKLMv4uKEHAItSU4nJ9r3y8vJxf05PTwcAVFdXo7S0FADQ3NyMq666Ci+88AJWrVoFSZLw7W9/Gw8//DCWL1+OFStWYNu2bTh69Cj+7//9v7E4TCLN5KbZYLOY4PUH0e4aUd90RWLI68exdjcAYAVDzokoQYlOqZ5BL9wjPmRwfTsRKfIz7JAkwB+U0TPkRV769NMTousyK8WK7FRbrA+RiAwoWuN7Lf0jaHd5YDZJbCKYQkyKUuHw+Xyoq6vD0NDoJ5/f/OY3MTIygvvvvx89PT1Yvnw53nzzTVRXV2t1mEQxIUkSirIcaOgeQkvf8IyKUoebXQjKoXwq0b5ORJRoMhxWzEmzoWfQi8aeYSwuZlGKiEKsZhNy0+zoGvCgrX8krKJUfZfYvMcuKSKa2GinVPhdmBMRo3vnFWcixWaO2vElmrgUpSorKyHL8rTfA0Jh6g899FA8DotIU87MUFFqpmHnap4Uq+5ElODK5qSiZ9CLMz2DWFycqfXhEJGOFGaGilId7hEA02dsipBz5kkR0WScWQ6YJMDrD6J70Iv8jJllWO8bE3JOk5tZkA0RzZqowM+0KLVf5ElxdI+IEhzDzoloMiJXqq0/vLBzEXJeydXsRDQJq9mEQuW1pXkWYecMOQ8Pi1JEGilS1hi3zjBAj5v3iChZlM8JFfFZlCKisxWMCTsPR31X6HWE43tENBV1hG+GRakhrx+ftboAsCg1HRaliDRSpL7QRd4p1TXgQVPvMCQJWFrCohQRJbaKOaGOhjM9s1/NTESJxRlhUUp0SnF8j4imMtui1P7GPgSCMoqzHOpt0cRYlCLSSJFoN3dF/kJ3UBndq85P5yYqIkp4YhnEGeXNJBGRUJgZynoJpyg14gugRYlNYKcUEU1FbOBrmuEGPjVPil1S02JRikgjRdnK+N4MOqX2N/YDYMg5ESWHcuXNY1PvMALBc5ekEFHyKswSnVLTZ0o1KiPAGXYL5qTZYnpcRGRsJcp7tZl2SjFPKnwsShFppDgrVH3vHvRixBeI6LoiT2oF86SIKAk4Mx2wmiX4g/KMc/iIKDEVZoQ/vqdu3stLhSTNbMU7ESUHdXxvBucdwaCMfWf6ALAoFQ4WpYg0kp1qhd0S+ifYFsEGPlmWcYCb94goiZhNEspyuIGPiM7lVDqluge98PqDU16WeVJEFK6SnFBRqnkG43unugbQP+yDw2rCoqLMaB9awmFRikgjkiSpFfjWCIpSZ3qG0Dfkg81sQo2TL3JElBxGc6VYlCKiUTmpVtjMobc0He6pz6fqlaIU86SIaDrifVrvkA9DXn9E1xWje8tLs2E1s+QyHf4NEWmoSPl0L5JxlP3K6N7i4kzYLPwnTETJoXwOO6WI6FySJKFADTufOleqQYzvsVOKiKaR6bAiw24BEPm2dOZJRYbvaIk0VJQVeafUASXkfAVH94goibAoRUSTKcwML1fqdJfolGJRioimp47wRRh2/rFSlKqtZFEqHCxKEWlIdEpFstVhNE+KIedElDzEBr5GFqWI6CzOMIpSHn9APd/i+B4RhUMNO4/gvVrPoBenOkMF8PPLWJQKB4tSRBoqUlaNhht07gsE8WlLqFNqeWl2rA6LiEh3RKdUA4tSRHQWMb7XNkVRqql3GEEZSLWZkZ9hj9ehEZGBFWdH3kDwyZlQl1R1fhpy0mwxOa5Ew6IUkYaKs8Sq0fCKUsfa3RjxBZHpsLD1nIiSigg67xvyoX/Yp/HREJGeiE6pjikypcZu3pMkKS7HRUTGVpIdOveIZAMf86Qix6IUkYZEp1S4QeciT2p5WTZMJp5QEVHySLdbkKt84sgRPiIaS2RKTdV5Xt8Vet3g6B4RhUt0SkWSKcWiVORYlCLSUFFmqFOqb8iHYW9g2ssfUDbvcXSPiJKR6JZiUYqIxlKDzt2TF6XGdkoREYWjRGRKhdlA4AsE1fxfFqXCx6IUkYYyUyxItZkBhNctNRpynh3DoyIi0qeKXG7gI6JzFSqZUu1TdUp1s1OKiCIjtu+19o0gEJSnvfyRVhdGfEFkpVhRlZce68NLGCxKEWlIkgUGO2gAABtySURBVCR1A1/rNLlSgx4/jrW7AQDLS7l5j4iSD8POiWgiolNq0BvAgMc/4WXYKUVEkSrIcMBskuAPyuh0T55ZJ4wd3WPUSvhYlCLSWLirRg839yMoA8VZDhQoJ19ERMmE43tENJE0uwUZdguAiXOlfIEgmpSg4so8dkoRUXjMJkldpBBOrtTHzJOaERaliDTmDCOcE+DoHhGR6JTi+B4Rna0wS2zgO/d8qqVvGP6gDLvFhMIMfrBHROETI3zhFKX2KUWpC8pZlIoEi1JEGitSA/SmKUopm/eWMeSciJKUyJRq7h2GPxDU+GiISE9ErlTbBEUpkSdVkZvKkRoiikhJmFMtLX3DaO0fgdkkYXkZo1YiwaIUkcaK1UypqV/o9ovNe3yRI6IkVZjhgM1sgj8oT5vDR0TJRd3A5zo396W+i3lSRDQzxdmh15bpilIiT2pxUSZSbZaYH1ciYVGKSGOiU6q1b/I3WJ1uD5r7hiFJwNISFqWIKDmZTBJK54ReMznCR0RjjRalJuqUChWluHmPiCJVkj3apT2VvcyTmjEWpYg0VhRGp9RBJU9qXn46MhzWeBwWEZEuMVeKiCbinKIo1aCM71XmsVOKiCIjOqWmy5Tad0bJk2JRKmIsShFpTBSlXCN+DE6yxviAOrqXHaejIiLSpwoWpYhoAlNnSolOKRaliCgy4WRKDXn9+LTFBQCoZVEqYixKEWksw2FV1xhP1i21vykUcs6iFBEluzJRlOpmUYqIRonxvY6zMqUCQRmNPaNB50REkShWilKuET/cI74JL3OgsR+BoIyiLId6eQofi1JEOuDMEgF65366J8uy2im1gpv3iCjJcXyPiCYyNlMqGJTV77f0DcMXkGEzm1CUxTeLRBSZNLsF2amh+JSJ3qsBHN2bLRaliHRAhJ23TbBNqqF7CP3DPtgsJix0ZsT70IiIdKU8l0UpIjpXfoYdkgT4gzJ6hrzq90WeVNmcFJhNklaHR0QGVpw19QifGnJezqLUTLAoRaQDxaJTaoLxvQNKyPl5xZmwWfhPloiSm+iU6h/2oX9o4jZ6vXnqqadQWVkJh8OB1atXY8+ePWFd78UXX4QkSbjhhhvGfX/Tpk2QJGnc14YNGya8DY/HgxUrVkCSJOzfv3+Wj4RIv6xmE3LTlFypMR/yMU+KiGarJCdUlGqaoCgVDMpqpxQ3780M3+ES6YBoJ2+doCX0QKOSJ8XRPSIipNosyEsPvfE0QrfUSy+9hM2bN+Phhx/Gvn37sHz5cqxfvx4dHR1TXq++vh7f+ta3cOmll0748w0bNqC1tVX9+vWvfz3h5R588EEUFxfP+nEQGYEzK/Ta0OEePZ9qUIpSFSxKEdEMTRV2fqprEH1DPjisJiwuzoz3oSUEFqWIdKAojE6pFQw5JyICAJTPCZ0cGqEo9dhjj+Guu+7CHXfcgcWLF+Ppp59GamoqnnvuuUmvEwgEcMstt+CRRx5BVVXVhJex2+1wOp3qV07OuZ/Ovvrqq3jjjTfw6KOPRu3xEOlZYUbofKqtfzTsvF4Z36vMY8g5Ec1McbbI/z33vdo+ZXRveWk2rGaWV2aCf2tEOlCULU6ixndK+QJBHG7m5j0iorHECF9jr76LUl6vF3v37sW6devU75lMJqxbtw67d++e9Hrf//73UVBQgDvvvHPSy+zatQsFBQVYuHAh7rnnHnR3d4/7eXt7O+666y786le/Qmrq9G/GPR4PXC7XuC8ioynMGg07F9gpRUSzVZId+j3a3HtuUUrNk+Lo3oxZtD4AIhozvndWUaquzQ2PP4hMhwWVXGNMRAQAeHBDDf7PFxYhXxnj06uuri4EAgEUFhaO+35hYSGOHj064XXeffddPPvss1PmP23YsAE33ngj5s6di5MnT+L//J//g40bN2L37t0wm82QZRmbNm3C3XffjdraWtTX1097rFu3bsUjjzwSycMj0h3RKSWKUsGgrAad8zyKiGZqqk6pjxt6ALAoNRssShHpgBjfG/D44RrxIdMRWjsqRveWl2VDkrgxhogIAIqzE3Otu9vtxq233opnnnkGeXl5k17uy1/+svrfS5cuxbJly1BdXY1du3bhqquuwpNPPgm3240tW7aEfd9btmzB5s2b1T+7XC6UlZXN7IEQaURkSomiVJtrBB5/EBaTpGbCEBFFSgSdt7lG4A8EYVHG9HoHvTjZGerGPJ+b92aMRSkiHUizW5DpsMA14kdr3wgynUpRqrEPAPOkiIiMKC8vD2azGe3t7eO+397eDqfTec7lT548ifr6elx77bXq94LBIADAYrGgrq4O1dXV51yvqqoKeXl5OHHiBK666irs2LEDu3fvht0+vpOstrYWt9xyC7Zt23bObdjt9nMuT2Q0BZlKHIIrlCklNu+V5qSobyKJiCKVl2aHzWyCNxBEm2sEpTmhzstPGkOje1X5aZiTZtPyEA2Nr85EOiE++W8dE3YuNu8t4+Y9IiLDsdlsWLlyJbZv365+LxgMYvv27VizZs05l6+pqcGhQ4ewf/9+9eu6667D2rVrsX///kk7l5qamtDd3Y2ioiIAwBNPPIEDBw6ot/GXv/wFQGgT4L/+67/G4JES6YNTKUp1KJ1S6uheHvOkiGjmTCZJzQBuGbMtXc2TYpfUrLBTikgnirIcONrmVnOlBjx+HOtwAwCWl2ZpeWhERDRDmzdvxu23347a2lqsWrUKjz/+OAYHB3HHHXcAAG677TaUlJRg69atcDgcWLJkybjrZ2dnA4D6/YGBATzyyCO46aab4HQ6cfLkSTz44IOYN28e1q9fDwAoLy8fdxvp6ekAgOrqapSWlsby4RJpqlApSnUPeuHxB9ROqUqGnBPRLJVkp6Che2hcrhRDzqODRSkinSgSnVLKC93h5n7IMlCc5VDb0YmIyFhuvvlmdHZ24rvf/S7a2tqwYsUKvPbaa2r4+ZkzZ2Ayhd+4bjabcfDgQWzbtg19fX0oLi7G1VdfjR/84Accv6Okl5NqVUdsOt0eNHSFOqUqGHJORLMkplqalfdqvkBQnWqprWRRajZYlCLSiSKl8CQ6pUSe1HLmSRERGdp9992H++67b8Kf7dq1a8rrPv/88+P+nJKSgtdffz2i+6+srIQsyxFdh8iIJElCQaYdTb3DaHeNsFOKiKLm7KLU0VY3hn0BZKVYUZWXruWhGR4zpYh0Qu2UEkWpMZv3iIiIiGh6Ileqrd+jZkqxU4qIZqtUea8mxvc+bugBAFxQng2TiVvSZ4NFKSKdKM5SwvOUoHPRDrqcIedEREREYRG5Uoea+zHsC8AkQd2URUQ0U2qnVG/ovRrzpKKHRSkinRjNlBpBh3sEzX3DkCRgKUPOiYiIiMIiilJ7TncDAEpyUmCz8C0PEc1Osbp9bxiyLGOfUpS6gEWpWeMrNJFOiHbzYV8A7xzrAgDML0hHup3Rb0REREThKMwMBf4fbAp1nDNPioiiQXRKDXoDONrmRkv/CMwmiVMtURDzopTH48GKFSsgSRL2798/5WXb2tpw6623wul0Ii0tDRdccAFefvnlWB8ikS6k2MzISbUCAF77tA0AR/eIiIiIIuFU4hD8wVC4P/OkiCgaHFYz8tJtAIA/HWgBACwqykAaGwhmLeZFqQcffBDFxcVhXfa2225DXV0d/vjHP+LQoUO48cYb8aUvfQmffPJJjI+SSB+KskIV+L8e6wTAkHMiIiKiSBRkOMb9mZ1SRBQtolvqTwdDRanaijlaHk7CiGlR6tVXX8Ubb7yBRx99NKzLv//++/ja176GVatWoaqqCv/yL/+C7Oxs7N27N5aHSaQbRcqnex5/EACwgkUpIiIiorCJTimhgkUpIoqSEqUo1dgTCjtnnlR0xKwo1d7ejrvuugu/+tWvkJoaXtvsRRddhJdeegk9PT0IBoN48cUXMTIygiuuuCJWh0mkK0XZoydSNosJC50ZGh4NERERkbGITClhbh7H94goOkSnlMDNe9ERkwFIWZaxadMm3H333aitrUV9fX1Y1/vNb36Dm2++Gbm5ubBYLEhNTcXvf/97zJs3b9LreDweeDwe9c8ul2u2h0+kGTG+BwBLijNhNXMXAREREVG4Um0WZDgscI/4IUlAaQ6LUkQUHWOLUs5MB4rP6sykmYnoHe9DDz0ESZKm/Dp69CiefPJJuN1ubNmyJaKD+c53voO+vj689dZb+Pjjj7F582Z86UtfwqFDhya9ztatW5GVlaV+lZWVRXSfRHpSPKZTinlSRERERJErVDYaF2elwGE1a3w0RJQoSsYUpVZW5ECSJA2PJnFE1Cn1wAMPYNOmTVNepqqqCjt27MDu3btht49vn62trcUtt9yCbdu2nXO9kydP4mc/+xkOHz6M8847DwCwfPlyvPPOO3jqqafw9NNPT3h/W7ZswebNm9U/u1wuFqbIsJyZoy90zJMiIiIiipwz04ETHQPcvEdEUTW2KMU8qeiJqCiVn5+P/Pz8aS/3xBNP4Ic//KH655aWFqxfvx4vvfQSVq9ePeF1hoaGAAAm0/jmLbPZjGAwOOl92e32c4pfREY1tlNqWWm2dgdCREREZFAFSq4UQ86JKJrGvldjnlT0xCRTqry8fNyf09PTAQDV1dUoLS0FADQ3N+Oqq67CCy+8gFWrVqGmpgbz5s3DV7/6VTz66KPIzc3FH/7wB7z55pt45ZVXYnGYRLpTkp2CxUWZSLGZUclP94iIiIgitnZhAV4/3Iaragq0PhQiSiBz0my4bEE+BkZ8OK84U+vDSRgxKUqFw+fzoa6uTu2Qslqt+Mtf/oKHHnoI1157LQYGBjBv3jxs27YN11xzjVaHSRRXFrMJr3ztEkgSOKNMRERENAPXLi/GF5YWwWTiuRQRRY8kSXjhK6u0PoyEE5eiVGVlJWRZnvZ78+fPx8svvxyPQyLSLZ5AEREREc0Oz6eIiIyB++aJiIiIiIiIiCjuWJQiIiIiIiIiIqK4Y1GKiIiIiIiIiIjijkUpIiIiIiIiIiKKOxaliIiIiIiIiIgo7liUIiIiIiIiIiKiuGNRioiIiIiIiIiI4o5FKSIiIiIiIiIiijsWpYiIiIiIiIiIKO5YlCIiIiIiIiIiorhjUYqIiIiIiIiIiOKORSkiIiIiIiIiIoo7FqWIiIiIiIiIiCjuWJQiIiIiIiIiIqK4Y1GKiIiIiIiIiIjizqL1AUSbLMsAAJfLpfGREBERkVGJ8whxXpFseD5FREREsxHuuVTCFaXcbjcAoKysTOMjISIiIqNzu93IysrS+jDijudTREREFA3TnUtJcoJ9BBgMBtHS0oKMjAxIkhT123e5XCgrK0NjYyMyMzOjfvt6x8efvI8/mR87kNyPP5kfO5Dcjz+ZH7ssy3C73SguLobJlHxpBzyfiq1kfvzJ/NiB5H78yfzYgeR+/Mn82IHkffzhnkslXKeUyWRCaWlpzO8nMzMzqf6HOhsff/I+/mR+7EByP/5kfuxAcj/+ZH3sydghJfB8Kj6S+fEn82MHkvvxJ/NjB5L78SfzYweS8/GHcy6VfB/9ERERERERERGR5liUIiIiIiIiIiKiuGNRKkJ2ux0PP/ww7Ha71oeiCT7+5H38yfzYgeR+/Mn82IHkfvzJ/NgptpL9/61kfvzJ/NiB5H78yfzYgeR+/Mn82AE+/ukkXNA5ERERERERERHpHzuliIiIiIiIiIgo7liUIiIiIiIiIiKiuGNRioiIiIiIiIiI4o5FKSIiIiIiIiIiijsWpSbw1FNPobKyEg6HA6tXr8aePXumvPxvf/tb1NTUwOFwYOnSpfjLX/4SpyONrq1bt+LCCy9ERkYGCgoKcMMNN6Curm7K6zz//POQJGncl8PhiNMRR9f3vve9cx5LTU3NlNdJlOe+srLynMcuSRLuvffeCS9v9Of9r3/9K6699loUFxdDkiT84Q9/GPdzWZbx3e9+F0VFRUhJScG6detw/PjxaW830tcOLUz12H0+H/7pn/4JS5cuRVpaGoqLi3HbbbehpaVlytucyb8drUz33G/atOmcx7Jhw4Zpb9fozz2ACV8DJEnCT37yk0lv00jPPcUfz6eS73wqmc+lgOQ6n0rmcykguc+nkvlcCuD5VCywKHWWl156CZs3b8bDDz+Mffv2Yfny5Vi/fj06OjomvPz777+Pv/u7v8Odd96JTz75BDfccANuuOEGHD58OM5HPntvv/027r33XnzwwQd488034fP5cPXVV2NwcHDK62VmZqK1tVX9amhoiNMRR99555037rG8++67k142kZ77jz76aNzjfvPNNwEAf/u3fzvpdYz8vA8ODmL58uV46qmnJvz5v/3bv+GJJ57A008/jQ8//BBpaWlYv349RkZGJr3NSF87tDLVYx8aGsK+ffvwne98B/v27cPvfvc71NXV4brrrpv2diP5t6Ol6Z57ANiwYcO4x/LrX/96yttMhOcewLjH3Nraiueeew6SJOGmm26a8naN8txTfPF8KnnPp5L1XApIrvOpZD6XApL7fCqZz6UAnk/FhEzjrFq1Sr733nvVPwcCAbm4uFjeunXrhJf/0pe+JH/hC18Y973Vq1fLX/3qV2N6nPHQ0dEhA5DffvvtSS/zy1/+Us7KyorfQcXQww8/LC9fvjzsyyfyc/+Nb3xDrq6uloPB4IQ/T6TnHYD8+9//Xv1zMBiUnU6n/JOf/ET9Xl9fn2y32+Vf//rXk95OpK8denD2Y5/Inj17ZAByQ0PDpJeJ9N+OXkz0+G+//Xb5+uuvj+h2EvW5v/766+Urr7xyyssY9bmn2OP51KhkOp/iudR4yXI+lcznUrKc3OdTyXwuJcs8n4oWdkqN4fV6sXfvXqxbt079nslkwrp167B79+4Jr7N79+5xlweA9evXT3p5I+nv7wcAzJkzZ8rLDQwMoKKiAmVlZbj++uvx6aefxuPwYuL48eMoLi5GVVUVbrnlFpw5c2bSyybqc+/1evHf//3f+MpXvgJJkia9XCI972OdPn0abW1t457brKwsrF69etLndiavHUbR398PSZKQnZ095eUi+bejd7t27UJBQQEWLlyIe+65B93d3ZNeNlGf+/b2dvz5z3/GnXfeOe1lE+m5p+jg+dR4yXY+xXOpkGQ+n+K51LmS7XyK51IhPJ8KD4tSY3R1dSEQCKCwsHDc9wsLC9HW1jbhddra2iK6vFEEg0F885vfxMUXX4wlS5ZMermFCxfiueeew//8z//gv//7vxEMBnHRRRehqakpjkcbHatXr8bzzz+P1157DT//+c9x+vRpXHrppXC73RNePlGf+z/84Q/o6+vDpk2bJr1MIj3vZxPPXyTP7UxeO4xgZGQE//RP/4S/+7u/Q2Zm5qSXi/Tfjp5t2LABL7zwArZv344f//jHePvtt7Fx40YEAoEJL5+oz/22bduQkZGBG2+8ccrLJdJzT9HD86lRyXY+xXOpUcl8PsVzqfGS7XyK51KjeD4VHovWB0D6dO+99+Lw4cPTzrKuWbMGa9asUf980UUXYdGiRfjFL36BH/zgB7E+zKjauHGj+t/Lli3D6tWrUVFRgd/85jdhVbcTxbPPPouNGzeiuLh40ssk0vNOE/P5fPjSl74EWZbx85//fMrLJtK/nS9/+cvqfy9duhTLli1DdXU1du3ahauuukrDI4uv5557Drfccsu0gbuJ9NwTxUKynU/xNWEUz6cISM7zKZ5LjeL5VHjYKTVGXl4ezGYz2tvbx32/vb0dTqdzwus4nc6ILm8E9913H1555RXs3LkTpaWlEV3XarXi/PPPx4kTJ2J0dPGTnZ2NBQsWTPpYEvG5b2howFtvvYV/+Id/iOh6ifS8i+cvkud2Jq8deiZOoBoaGvDmm29O+aneRKb7t2MkVVVVyMvLm/SxJNpzDwDvvPMO6urqIn4dABLruaeZ4/lUCM+nkvNcCuD5FM+lQng+FZKM51IAz6ciwaLUGDabDStXrsT27dvV7wWDQWzfvn3cpxhjrVmzZtzlAeDNN9+c9PJ6Jssy7rvvPvz+97/Hjh07MHfu3IhvIxAI4NChQygqKorBEcbXwMAATp48OeljSaTnXvjlL3+JgoICfOELX4joeon0vM+dOxdOp3Pcc+tyufDhhx9O+tzO5LVDr8QJ1PHjx/HWW28hNzc34tuY7t+OkTQ1NaG7u3vSx5JIz73w7LPPYuXKlVi+fHnE102k555mjudTPJ8SkvFcCuD5VLKfSwE8nxorGc+lAJ5PRUTbnHX9efHFF2W73S4///zz8meffSb/4z/+o5ydnS23tbXJsizLt956q/zQQw+pl3/vvfdki8UiP/roo/KRI0fkhx9+WLZarfKhQ4e0eggzds8998hZWVnyrl275NbWVvVraGhIvczZj/+RRx6RX3/9dfnkyZPy3r175S9/+cuyw+GQP/30Uy0ewqw88MAD8q5du+TTp0/L7733nrxu3To5Ly9P7ujokGU5sZ97WQ5tuSgvL5f/6Z/+6ZyfJdrz7na75U8++UT+5JNPZADyY489Jn/yySfqRpQf/ehHcnZ2tvw///M/8sGDB+Xrr79enjt3rjw8PKzexpVXXik/+eST6p+ne+3Qi6keu9frla+77jq5tLRU3r9//7jXAY/Ho97G2Y99un87ejLV43e73fK3vvUteffu3fLp06flt956S77gggvk+fPnyyMjI+ptJOJzL/T398upqanyz3/+8wlvw8jPPcUXz6eS83wq2c+lZDl5zqeS+VxKlpP7fCqZz6VkmedTscCi1ASefPJJuby8XLbZbPKqVavkDz74QP3Z5ZdfLt9+++3jLv+b3/xGXrBggWyz2eTzzjtP/vOf/xznI44OABN+/fKXv1Qvc/bj/+Y3v6n+XRUWFsrXXHONvG/fvvgffBTcfPPNclFRkWyz2eSSkhL55ptvlk+cOKH+PJGfe1mW5ddff10GINfV1Z3zs0R73nfu3Dnh/+viMQaDQfk73/mOXFhYKNvtdvmqq6465++loqJCfvjhh8d9b6rXDr2Y6rGfPn160teBnTt3qrdx9mOf7t+Onkz1+IeGhuSrr75azs/Pl61Wq1xRUSHfdddd55wQJeJzL/ziF7+QU1JS5L6+vglvw8jPPcUfz6eS73wq2c+lZDl5zqeS+VxKlpP7fCqZz6VkmedTsSDJsizPtMuKiIiIiIiIiIhoJpgpRUREREREREREcceiFBERERERERERxR2LUkREREREREREFHcsShERERERERERUdyxKEVERERERERERHHHohQREREREREREcUdi1JERERERERERBR3LEoREREREREREVHcsShFRERERERERERxx6IUERERERERERHFHYtSREREREREREQUdyxKERERERERERFR3P1/wEbI1dCmQUcAAAAASUVORK5CYII=", + "text/plain": [ + "
" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "from custom_feature_dir.utils import plot_stats\n", + "\n", + "plot_stats(log_dir)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/docs/debug/custom_feature_dir/custom_feature_example_config.yaml b/docs/debug/custom_feature_dir/custom_feature_example_config.yaml new file mode 100644 index 0000000000..ab0369866f --- /dev/null +++ b/docs/debug/custom_feature_dir/custom_feature_example_config.yaml @@ -0,0 +1,15 @@ +stats: + enabled: True + layers: + layer_name_regex_pattern: .* + transformer_engine: + PercentageGreaterThanThreshold: + enabled: True + tensors: [activation] + threshold: 0.1 + freq: 5 + LogTensorStats: + enabled: True + tensors: [activation] + stats: [min] + freq: 5 \ No newline at end of file diff --git a/docs/debug/custom_feature_dir/percentage_greater_than_threshold.py b/docs/debug/custom_feature_dir/percentage_greater_than_threshold.py new file mode 100644 index 0000000000..80311ec499 --- /dev/null +++ b/docs/debug/custom_feature_dir/percentage_greater_than_threshold.py @@ -0,0 +1,78 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PercentageGreaterThanThreshold Feature support for nvidia-dlframework-inspect""" + +from typing import Dict, Optional + +import torch + +from nvdlfw_inspect.registry import Registry, api_method +from nvdlfw_inspect.logging import MetricLogger +import nvdlfw_inspect.api as debug_api + +from transformer_engine.debug.features.api import TEConfigAPIMapper +from transformer_engine.pytorch.tensor import QuantizedTensor, Quantizer + + +# Class should inherit from TEConfigAPIMapper and be registered to the transformer_engine namespace. +@Registry.register_feature(namespace="transformer_engine") +class PercentageGreaterThanThreshold(TEConfigAPIMapper): + + @api_method + def inspect_tensor( + self, + config: Dict, + layer_name: str, + tensor_name: str, + iteration: int, + tp_group: torch.distributed.ProcessGroup, + tensor: torch.Tensor, + rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, + columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, + quantizer: Optional[Quantizer] = None, + ): + # API call inspect_tensor is used to gather the data about the tensor. + # All API calls are documented in the `Precision debug tools / API / Calls to Nvidia-DL-Framework-Inspect` + # section of the documentation. + + threshold = config["threshold"] + + # Get the reduction group from the debug tool + # one can set it using debug_api.set_tensor_reduction_group(group) + reduction_group = debug_api.get_tensor_reduction_group() + + # Compute percentage on local tensor + count = (torch.abs(tensor) > threshold).sum().float() + total = torch.tensor(tensor.numel(), dtype=torch.float32, device=tensor.device) + + # Perform reduction across the group if needed. + # Note that we perform all_reduce twice per every tensor, which is suboptimal. + # For guidance on implementing efficient statistics reduction, see the implementation in the `LogTensorStats` feature. + # In this tutorial we only showcase basic implementation of the feature. + if reduction_group is not None: + torch.distributed.all_reduce(count, group=reduction_group) + torch.distributed.all_reduce(total, group=reduction_group) + + percentage = count / total + + # MetricLogger is a class from nvidia-dlframework-inspect. + # By using it we can also use functionalities provided by nvidia-dlframework-inspect, + # like logging to TensorBoard, etc. + MetricLogger.log_scalar( + f"{layer_name}_{tensor_name}_percentage_greater_than_threshold", percentage, iteration + ) + + @api_method + def inspect_tensor_enabled( + self, config: Dict, layer_name: str, tensor_name: str, iteration: int + ): + # This call is used by TE to determine if the unfused debug layer - which is slower - needs to be run. + # It returns a tuple (bool, int), where the int indicates the next iteration when the feature will be enabled + # and bool indicates if the feature should be enabled at the current iteration. + + run_current = iteration % config["freq"] == 0 + # run in next multiple of freq + next_iter = iteration + (config["freq"] - iteration % config["freq"]) + return run_current, next_iter diff --git a/docs/debug/custom_feature_dir/utils.py b/docs/debug/custom_feature_dir/utils.py new file mode 100644 index 0000000000..cc954b12b6 --- /dev/null +++ b/docs/debug/custom_feature_dir/utils.py @@ -0,0 +1,48 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Utils for plotting stats in the tutorial""" + + +import os +import re +import matplotlib.pyplot as plt + + +def plot_stats(log_dir): + + # print and plot the stats + stat_file = os.path.join( + log_dir, "nvdlfw_inspect_statistics_logs", "nvdlfw_inspect_globalrank-0.log" + ) + + min_values = [] + custom_feature_values = [] + + with open(stat_file, "r") as f: + number_pattern = re.compile(r"[-+]?\d*\.\d+|\d+") + + for line in f: + if "min" in line: + matches = number_pattern.findall(line) + if matches: + min_values.append(float(matches[-1])) + if "percentage_greater_than_threshold" in line: + matches = number_pattern.findall(line) + if matches: + custom_feature_values.append(float(matches[-1])) + + # plot 2 figures side by side + fig, axs = plt.subplots(1, 2, figsize=(12, 5)) + + axs[0].plot(min_values, label="min") + axs[0].legend() + axs[0].set_title("Min values") + + axs[1].plot(custom_feature_values, label="percentage_greater_than_threshold_0.1") + axs[1].legend() + axs[1].set_title("Percentage greater than threshold 0.1 values") + + plt.tight_layout() + plt.show() From 7d1de3027e2506c231a4b7be4ee6c69ccb62f10c Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Tue, 24 Feb 2026 11:47:14 -0700 Subject: [PATCH 222/521] Fix vermin pre-commit hook (#2699) update .pre-commit-config Signed-off-by: Peter St. John --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5043d6ea22..76f476eb3f 100755 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -43,4 +43,4 @@ repos: rev: c75aca72f4e85c6e47252139e8695f1c8b5f9ae3 hooks: - id: vermin - args: ['-t=3.10', '--violations'] + args: ['-t=3.10-', '--violations'] From 459e7cf49d947e2590ef93605d610ab846805818 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Wed, 25 Feb 2026 03:43:10 +0800 Subject: [PATCH 223/521] [Common][PyTorch] Fuse scaling and unscaling of bf16 momentums into kernels (#2632) * fused scaling and unscaling of bf16 momentum Signed-off-by: Xin Yao * add more comments Signed-off-by: Xin Yao * enable cuda graphs for bf16 momentums Signed-off-by: Xin Yao * add tests Signed-off-by: Xin Yao * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update the check for store_param_remainders and capturable Signed-off-by: Xin Yao --------- Signed-off-by: Xin Yao Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fused_optimizer.py | 28 ++- transformer_engine/common/common.h | 15 ++ .../common/multi_tensor/adam.cu | 222 +++++++++++------- .../pytorch/optimizers/fused_adam.py | 63 +++-- 4 files changed, 217 insertions(+), 111 deletions(-) diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index f70be45918..185b9b85bc 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -407,6 +407,20 @@ def test_bf16_exp_avg_sq(self): master_atol=2e-3, ) + @pytest.mark.skipif(not is_bf16_available(), reason="bf16 if not supported") + def test_bf16_exp_avg_and_exp_avg_sq(self): + self.gen_precision_aware_test( + use_fp8_params=False, + param_dtype=torch.bfloat16, + use_master_weights=True, + master_weight_dtype=torch.float32, + grad_dtype=torch.float32, + exp_avg_dtype=torch.bfloat16, + exp_avg_sq_dtype=torch.bfloat16, + master_rtol=2e-3, + master_atol=2e-3, + ) + @pytest.mark.skipif(not is_bf16_available(), reason="bf16 if not supported") @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) def test_fp8_exp_avg_sq(self): @@ -553,7 +567,7 @@ def forward(self, x): return y -class AdamTest: +class TestAdamTest: def setup_method(self, *, seed: int = 0) -> None: torch.manual_seed(seed) @@ -569,8 +583,8 @@ def setup_method(self, *, seed: int = 0) -> None: def test_grad_scaler(self): params_ = [p for p in self.model_.parameters() if p.requires_grad] optimizer_ = te.optimizers.FusedAdam(params_, lr=self.lr, capturable=False) - scaler = torch.cuda.amp.GradScaler(enabled=True) - scaler_ = torch.cuda.amp.GradScaler(enabled=True) + scaler = torch.amp.GradScaler("cuda", enabled=True) + scaler_ = torch.amp.GradScaler("cuda", enabled=True) for i in range(100): x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last) @@ -620,8 +634,8 @@ def test_grad_scaler(self): def test_grad_scaler_capturable(self): params_ = [p for p in self.model_.parameters() if p.requires_grad] optimizer_ = te.optimizers.FusedAdam(params_, lr=self.lr, capturable=True) - scaler = torch.cuda.amp.GradScaler(enabled=True) - scaler_ = torch.cuda.amp.GradScaler(enabled=True) + scaler = torch.amp.GradScaler("cuda", enabled=True) + scaler_ = torch.amp.GradScaler("cuda", enabled=True) for i in range(100): x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last) @@ -678,8 +692,8 @@ def test_grad_scaler_capturable_master(self): optimizer_ = te.optimizers.FusedAdam( params_, lr=self.lr, capturable=True, master_weights=master_weights ) - scaler = torch.cuda.amp.GradScaler(enabled=True) - scaler_ = torch.cuda.amp.GradScaler(enabled=True) + scaler = torch.amp.GradScaler("cuda", enabled=True) + scaler_ = torch.amp.GradScaler("cuda", enabled=True) for i in range(100): x = torch.rand([32, 1, 28, 28]).cuda().to(memory_format=torch.channels_last) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 2d7f0e7e8c..0c722634f3 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -725,6 +725,21 @@ struct TypeInfo { NVTE_ERROR("Invalid type."); \ } +#define TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Invalid type, expected Float32 or BFloat16."); \ + } + // Add a pack_size argument to select the packed type for FP4 #define TRANSFORMER_ENGINE_TYPE_SWITCH_FP4x2_ONLY(dtype, pack_size, type, ...) \ switch (dtype) { \ diff --git a/transformer_engine/common/multi_tensor/adam.cu b/transformer_engine/common/multi_tensor/adam.cu index 5d89179c44..29a073be84 100644 --- a/transformer_engine/common/multi_tensor/adam.cu +++ b/transformer_engine/common/multi_tensor/adam.cu @@ -49,7 +49,7 @@ struct FP8Data { template <> struct FP8Data {}; -template +template struct AdamFunctorMaster { static constexpr bool is_fp8_type = is_fp8::value; @@ -79,10 +79,10 @@ struct AdamFunctorMaster { PARAM_T *p = reinterpret_cast(tl.addresses[1][tensor_loc]); p += chunk_idx * chunk_size; - FULL_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); + MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); m += chunk_idx * chunk_size; - FULL_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); + MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); v += chunk_idx * chunk_size; FULL_T *p_master = reinterpret_cast(tl.addresses[4][tensor_loc]); @@ -147,8 +147,8 @@ struct AdamFunctorMaster { int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) { p_master[i] = static_cast(r_p[ii]); - m[i] = static_cast(r_m[ii]); - v[i] = static_cast(r_v[ii]); + m[i] = static_cast(r_m[ii]); + v[i] = static_cast(r_v[ii]); if constexpr (is_fp8_type) { __builtin_assume(fp8_data.max >= 0); fp8_data.max = fmaxf(fabsf(r_p[ii]), fp8_data.max); @@ -175,7 +175,7 @@ struct AdamFunctorMaster { } }; -template +template struct AdamFunctorMasterParamRemainder { __device__ __forceinline__ void operator()(index_t chunk_size, volatile int *noop_gmem, TensorListMetadata<5> &tl, // NOLINT(*) @@ -194,10 +194,10 @@ struct AdamFunctorMasterParamRemainder { int16_t *p = reinterpret_cast(tl.addresses[1][tensor_loc]); p += chunk_idx * chunk_size; - FULL_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); + MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); m += chunk_idx * chunk_size; - FULL_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); + MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); v += chunk_idx * chunk_size; int16_t *p_remainder = reinterpret_cast(tl.addresses[4][tensor_loc]); @@ -283,15 +283,15 @@ struct AdamFunctorMasterParamRemainder { p_remainder[i] = local_p_rem[ii]; p[i] = local_p[ii]; - m[i] = static_cast(r_m[ii]); - v[i] = static_cast(r_v[ii]); + m[i] = static_cast(r_m[ii]); + v[i] = static_cast(r_v[ii]); } } } } }; -template +template struct AdamFunctor { __device__ __forceinline__ void operator()(index_t chunk_size, volatile int *noop_gmem, TensorListMetadata<4> &tl, // NOLINT(*) @@ -317,10 +317,10 @@ struct AdamFunctor { PARAM_T *p = reinterpret_cast(tl.addresses[1][tensor_loc]); p += chunk_idx * chunk_size; - FULL_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); + MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); m += chunk_idx * chunk_size; - FULL_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); + MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); v += chunk_idx * chunk_size; n -= chunk_idx * chunk_size; @@ -372,15 +372,15 @@ struct AdamFunctor { int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) { p[i] = static_cast(r_p[ii]); - m[i] = static_cast(r_m[ii]); - v[i] = static_cast(r_v[ii]); + m[i] = static_cast(r_m[ii]); + v[i] = static_cast(r_v[ii]); } } } } }; -template +template struct AdamCapturableFunctor { __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<4> &tl, // NOLINT(*) @@ -410,10 +410,10 @@ struct AdamCapturableFunctor { T *p = reinterpret_cast(tl.addresses[1][tensor_loc]); p += chunk_idx * chunk_size; - FULL_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); + MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); m += chunk_idx * chunk_size; - FULL_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); + MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); v += chunk_idx * chunk_size; n -= chunk_idx * chunk_size; @@ -466,15 +466,15 @@ struct AdamCapturableFunctor { int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) { p[i] = static_cast(r_p[ii]); - m[i] = static_cast(r_m[ii]); - v[i] = static_cast(r_v[ii]); + m[i] = static_cast(r_m[ii]); + v[i] = static_cast(r_v[ii]); } } } } }; -template +template struct AdamCapturableMasterFunctor { __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<5> &tl, // NOLINT(*) @@ -504,10 +504,10 @@ struct AdamCapturableMasterFunctor { T *p = reinterpret_cast(tl.addresses[1][tensor_loc]); p += chunk_idx * chunk_size; - FULL_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); + MOMENT_T *m = reinterpret_cast(tl.addresses[2][tensor_loc]); m += chunk_idx * chunk_size; - FULL_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); + MOMENT_T *v = reinterpret_cast(tl.addresses[3][tensor_loc]); v += chunk_idx * chunk_size; FULL_T *p_master = reinterpret_cast(tl.addresses[4][tensor_loc]); @@ -564,8 +564,8 @@ struct AdamCapturableMasterFunctor { if (i < n && i < chunk_size) { p[i] = static_cast(r_p[ii]); p_master[i] = static_cast(r_p[ii]); - m[i] = static_cast(r_m[ii]); - v[i] = static_cast(r_v[ii]); + m[i] = static_cast(r_m[ii]); + v[i] = static_cast(r_v[ii]); } } } @@ -606,12 +606,17 @@ void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK(tensor_lists[1][j]->dtype() == p_in_type_te, "Param tensor ", j, " has dtype=", to_string(tensor_lists[1][j]->dtype()), ", but expected dtype=", to_string(p_in_type_te)); - NVTE_CHECK(tensor_lists[2][j]->dtype() == DType::kFloat32, "First moment tensor ", j, - " has dtype=", to_string(tensor_lists[2][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); - NVTE_CHECK(tensor_lists[3][j]->dtype() == DType::kFloat32, "Second moment tensor ", j, - " has dtype=", to_string(tensor_lists[3][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); + { + const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; + const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; + const bool v_is_fp32 = tensor_lists[3][j]->dtype() == DType::kFloat32; + const bool v_is_bf16 = tensor_lists[3][j]->dtype() == DType::kBFloat16; + NVTE_CHECK((m_is_fp32 && v_is_fp32) || (m_is_bf16 && v_is_bf16), + "First and second moment tensors must both be Float32 or both be BFloat16, but " + "tensor ", + j, " has first moment dtype=", to_string(tensor_lists[2][j]->dtype()), + " and second moment dtype=", to_string(tensor_lists[3][j]->dtype())); + } if (num_tensor_lists == 5) { NVTE_CHECK(tensor_lists[4][j]->dtype() == DType::kFloat32, "Master param tensor ", j, " has dtype=", to_string(tensor_lists[4][j]->dtype()), @@ -633,6 +638,9 @@ void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, } } + // Get moment dtype (m and v have the same dtype, already validated above) + const auto moment_type_te = tensor_lists[2][0]->dtype(); + // Launch kernel if (requires_64bit_indexing) { if (num_tensor_lists == 4) { @@ -641,22 +649,26 @@ void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, p_in_type_te, p_in_type, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( g_in_type_te, g_in_type, - multi_tensor_apply<4>((int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, - tensor_lists, - AdamFunctor(), stream, - beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, - (adamMode_t)mode, weight_decay);)); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<4>( + (int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, tensor_lists, + AdamFunctor(), stream, + beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, + (adamMode_t)mode, weight_decay);))); } else { // g, p, m, v, p_master TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( p_in_type_te, p_in_type, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( g_in_type_te, g_in_type, - multi_tensor_apply<5>((int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, - tensor_lists, - AdamFunctorMaster(), - stream, beta1, beta2, bias_correction1, bias_correction2, - epsilon, lr, (adamMode_t)mode, weight_decay);)); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<5>( + (int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, tensor_lists, + AdamFunctorMaster(), + stream, beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, + (adamMode_t)mode, weight_decay);))); } } else { if (num_tensor_lists == 4) { @@ -665,20 +677,26 @@ void multi_tensor_adam_cuda(int chunk_size, Tensor noop_flag, p_in_type_te, p_in_type, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( g_in_type_te, g_in_type, - multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, - AdamFunctor(), stream, - beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, - (adamMode_t)mode, weight_decay);)); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<4>( + BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + AdamFunctor(), stream, + beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, + (adamMode_t)mode, weight_decay);))); } else { // g, p, m, v, p_master TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( p_in_type_te, p_in_type, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( g_in_type_te, g_in_type, - multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, - AdamFunctorMaster(), - stream, beta1, beta2, bias_correction1, bias_correction2, - epsilon, lr, (adamMode_t)mode, weight_decay);)); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<5>( + BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + AdamFunctorMaster(), + stream, beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, + (adamMode_t)mode, weight_decay);))); } } NVTE_CHECK_CUDA(cudaGetLastError()); @@ -716,24 +734,35 @@ void multi_tensor_adam_param_remainder_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK(tensor_lists[1][j]->dtype() == DType::kBFloat16, "Param tensor ", j, " has dtype=", to_string(tensor_lists[1][j]->dtype()), ", but expected dtype=", to_string(DType::kBFloat16)); - NVTE_CHECK(tensor_lists[2][j]->dtype() == DType::kFloat32, "First moment tensor ", j, - " has dtype=", to_string(tensor_lists[2][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); - NVTE_CHECK(tensor_lists[3][j]->dtype() == DType::kFloat32, "Second moment tensor ", j, - " has dtype=", to_string(tensor_lists[3][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); + { + const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; + const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; + const bool v_is_fp32 = tensor_lists[3][j]->dtype() == DType::kFloat32; + const bool v_is_bf16 = tensor_lists[3][j]->dtype() == DType::kBFloat16; + NVTE_CHECK((m_is_fp32 && v_is_fp32) || (m_is_bf16 && v_is_bf16), + "First and second moment tensors must both be Float32 or both be BFloat16, but " + "tensor ", + j, " has first moment dtype=", to_string(tensor_lists[2][j]->dtype()), + " and second moment dtype=", to_string(tensor_lists[3][j]->dtype())); + } NVTE_CHECK(tensor_lists[4][j]->dtype() == DType::kInt16, "Param remainder tensor ", j, " has dtype=", to_string(tensor_lists[4][j]->dtype()), ", but expected dtype=", to_string(DType::kInt16)); } + // Get moment dtype (m and v have the same dtype, already validated above) + const auto moment_type_te = tensor_lists[2][0]->dtype(); + // Launch kernel TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( g_in_type_te, g_in_type, - multi_tensor_apply<5>((int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, tensor_lists, - AdamFunctorMasterParamRemainder(), stream, - beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, - (adamMode_t)mode, weight_decay);); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<5>( + (int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, tensor_lists, + AdamFunctorMasterParamRemainder(), stream, + beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, (adamMode_t)mode, + weight_decay);)); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -812,17 +841,17 @@ void multi_tensor_adam_fp8_cuda(int chunk_size, Tensor noop_flag, g_in_type_te, g_in_type, multi_tensor_apply<5, true>( (int64_t)BLOCK_SIZE, (int64_t)chunk_size, noop_flag, tensor_lists, - AdamFunctorMaster(), stream, beta1, beta2, + AdamFunctorMaster(), stream, beta1, beta2, bias_correction1, bias_correction2, epsilon, lr, (adamMode_t)mode, weight_decay);)); } else { TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( fp8_dtype, FP8_T, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( g_in_type_te, g_in_type, - multi_tensor_apply<5, true>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, - AdamFunctorMaster(), - stream, beta1, beta2, bias_correction1, bias_correction2, - epsilon, lr, (adamMode_t)mode, weight_decay);)); + multi_tensor_apply<5, true>( + BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + AdamFunctorMaster(), stream, beta1, beta2, + bias_correction1, bias_correction2, epsilon, lr, (adamMode_t)mode, weight_decay);)); } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -852,22 +881,32 @@ void multi_tensor_adam_capturable_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK(tensor_lists[1][j]->dtype() == g_in_type_te, "Param tensor ", j, " has dtype=", to_string(tensor_lists[1][j]->dtype()), ", but expected dtype=", to_string(g_in_type_te)); - NVTE_CHECK(tensor_lists[2][j]->dtype() == DType::kFloat32, "First moment tensor ", j, - " has dtype=", to_string(tensor_lists[2][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); - NVTE_CHECK(tensor_lists[3][j]->dtype() == DType::kFloat32, "Second moment tensor ", j, - " has dtype=", to_string(tensor_lists[3][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); + { + const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; + const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; + const bool v_is_fp32 = tensor_lists[3][j]->dtype() == DType::kFloat32; + const bool v_is_bf16 = tensor_lists[3][j]->dtype() == DType::kBFloat16; + NVTE_CHECK((m_is_fp32 && v_is_fp32) || (m_is_bf16 && v_is_bf16), + "First and second moment tensors must both be Float32 or both be BFloat16, but " + "tensor ", + j, " has first moment dtype=", to_string(tensor_lists[2][j]->dtype()), + " and second moment dtype=", to_string(tensor_lists[3][j]->dtype())); + } } + // Get moment dtype (m and v have the same dtype, already validated above) + const auto moment_type_te = tensor_lists[2][0]->dtype(); + // Launch kernel TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( tensor_lists[0][0]->dtype(), dtype, - multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, - AdamCapturableFunctor(), stream, beta1, beta2, - reinterpret_cast(step.data.dptr), bias_correction, epsilon, - reinterpret_cast(lr.data.dptr), (adamMode_t)mode, weight_decay, - reinterpret_cast(inv_scale.data.dptr));) + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + AdamCapturableFunctor(), stream, beta1, + beta2, reinterpret_cast(step.data.dptr), bias_correction, + epsilon, reinterpret_cast(lr.data.dptr), (adamMode_t)mode, + weight_decay, reinterpret_cast(inv_scale.data.dptr));)) NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -897,25 +936,36 @@ void multi_tensor_adam_capturable_master_cuda(int chunk_size, Tensor noop_flag, NVTE_CHECK(tensor_lists[1][j]->dtype() == g_in_type_te, "Param tensor ", j, " has dtype=", to_string(tensor_lists[1][j]->dtype()), ", but expected dtype=", to_string(g_in_type_te)); - NVTE_CHECK(tensor_lists[2][j]->dtype() == DType::kFloat32, "First moment tensor ", j, - " has dtype=", to_string(tensor_lists[2][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); - NVTE_CHECK(tensor_lists[3][j]->dtype() == DType::kFloat32, "Second moment tensor ", j, - " has dtype=", to_string(tensor_lists[3][j]->dtype()), - ", but expected dtype=", to_string(DType::kFloat32)); + { + const bool m_is_fp32 = tensor_lists[2][j]->dtype() == DType::kFloat32; + const bool m_is_bf16 = tensor_lists[2][j]->dtype() == DType::kBFloat16; + const bool v_is_fp32 = tensor_lists[3][j]->dtype() == DType::kFloat32; + const bool v_is_bf16 = tensor_lists[3][j]->dtype() == DType::kBFloat16; + NVTE_CHECK((m_is_fp32 && v_is_fp32) || (m_is_bf16 && v_is_bf16), + "First and second moment tensors must both be Float32 or both be BFloat16, but " + "tensor ", + j, " has first moment dtype=", to_string(tensor_lists[2][j]->dtype()), + " and second moment dtype=", to_string(tensor_lists[3][j]->dtype())); + } NVTE_CHECK(tensor_lists[4][j]->dtype() == DType::kFloat32, "Master param tensor ", j, " has dtype=", to_string(tensor_lists[4][j]->dtype()), ", but expected dtype=", to_string(DType::kFloat32)); } + // Get moment dtype (m and v have the same dtype, already validated above) + const auto moment_type_te = tensor_lists[2][0]->dtype(); + // Launch kernel TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( tensor_lists[0][0]->dtype(), dtype, - multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, - AdamCapturableMasterFunctor(), stream, beta1, beta2, - reinterpret_cast(step.data.dptr), bias_correction, epsilon, - reinterpret_cast(lr.data.dptr), (adamMode_t)mode, weight_decay, - reinterpret_cast(inv_scale.data.dptr));) + TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16( + moment_type_te, moment_type, + multi_tensor_apply<5>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + AdamCapturableMasterFunctor(), stream, + beta1, beta2, reinterpret_cast(step.data.dptr), + bias_correction, epsilon, reinterpret_cast(lr.data.dptr), + (adamMode_t)mode, weight_decay, + reinterpret_cast(inv_scale.data.dptr));)) NVTE_CHECK_CUDA(cudaGetLastError()); } diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 495056d652..a87d968334 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -140,19 +140,24 @@ def __init__( if exp_avg_sq_dtype not in [torch.float32, torch.float16, torch.bfloat16, torch.uint8]: raise RuntimeError("FusedAdam only supports fp32/fp16/bf16/fp8 exp_avg_sq.") - # Currently, capturable mode only supports fp32 master weights and optimizer states. - # The reason is, if the master weights or optimizer states are not in fp32 dtype, - # they will be copied to temporary fp32 buffers first. These fp32 buffers are then - # used as inputs for the kernel. Consequently, the pointer for earch `.step()` differs, - # making CUDA Graph inapplicable in this scenario. + # Capturable mode requires fp32 master weights, and optimizer states (exp_avg/exp_avg_sq) + # must both be fp32 or both be bf16. This is because master weights in non-fp32 dtypes + # or optimizer states in non-fp32/bf16 dtypes require copying to temporary fp32 buffers + # before kernel execution, causing different pointers on each `.step()` call and making + # CUDA Graph inapplicable. if capturable and master_weights and master_weight_dtype != torch.float32: raise RuntimeError("Capturable mode only supports fp32 master weights.") - if capturable and exp_avg_dtype != torch.float32: - raise RuntimeError("Capturable mode only supports fp32 exp_avg.") - if capturable and exp_avg_sq_dtype != torch.float32: - raise RuntimeError("Capturable mode only supports fp32 exp_avg_sq") - if capturable and store_param_remainders: - raise RuntimeError("Capturable mode doesn't support storing param remainders") + if capturable: + valid_moment_dtypes = ( + exp_avg_dtype == exp_avg_sq_dtype == torch.float32 + or exp_avg_dtype == exp_avg_sq_dtype == torch.bfloat16 + ) + if not valid_moment_dtypes: + raise RuntimeError( + "Capturable mode requires exp_avg_dtype and exp_avg_sq_dtype to be " + "both torch.float32 or both torch.bfloat16, but got " + f"exp_avg_dtype={exp_avg_dtype} and exp_avg_sq_dtype={exp_avg_sq_dtype}." + ) # If the optimizer is capturable then LR should be a tensor (on GPU) lr = torch.tensor(lr, dtype=torch.float32) if capturable else lr @@ -207,6 +212,11 @@ def __init__( self.store_param_remainders = ( store_param_remainders and master_weights and master_weight_dtype == torch.float32 ) + if self.capturable and self.store_param_remainders: + raise RuntimeError("Capturable mode doesn't support storing param remainders") + # If the exp_avg and exp_avg_sq dtypes are bfloat16, we can fuse the unscaling/scaling + # operations into the fused Adam kernel. + self.fuse_unscale = self.exp_avg_dtype == self.exp_avg_sq_dtype == torch.bfloat16 # Deprecated options self.set_grad_none = set_grad_none @@ -268,10 +278,9 @@ def _apply_scale(self, state_name, unscaled_state, scaled_state, scale): dtype = self.name_to_dtype_map[state_name] if dtype == torch.uint8: assert isinstance(scaled_state, Float8Tensor) - assert len(scaled_state._quantizer.scale) == 1, ( - "Only scaling with one scaling factor per tensor is supported by the" - " FusedAdam." - ) + assert ( + len(scaled_state._quantizer.scale) == 1 + ), "Only scaling with one scaling factor per tensor is supported by the FusedAdam." else: assert scaled_state.dtype == dtype @@ -293,13 +302,22 @@ def _apply_scale(self, state_name, unscaled_state, scaled_state, scale): unscaled_state.mul_(rscale) scaled_state.copy_(unscaled_state) - def get_unscaled_state(self, param, state_name): + def get_unscaled_state( + self, param: torch.nn.Parameter, state_name: str, skip_unscale: bool = False + ) -> torch.Tensor: """Return the unscaled state corresponding to the input `param` and `state_name`. Arguments: param (torch.nn.Parameter): One of parameters in this optimizer. state_name (string): Name of optimizer states, can be one of 'exp_avg', 'exp_avg_sq', and 'master_param`. + skip_unscale (optional, bool): Whether to skip the unscaling operation. + Should only be True if 'self.fuse_unscale' is True. Default is False. + + Returns: + torch.Tensor: The unscaled state. Note that if the state is in BF16, the returned + tensor is still in BF16 because it doesn't require to be "unscaled", otherwise it + will be unscaled to FP32. """ state = self.state[param] dtype = self.name_to_dtype_map[state_name] @@ -321,7 +339,10 @@ def get_unscaled_state(self, param, state_name): unscaled = state[state_name] elif dtype == torch.bfloat16: assert state[state_name].dtype == torch.bfloat16 - unscaled = state[state_name].float() + if skip_unscale: + unscaled = state[state_name] + else: + unscaled = state[state_name].float() else: raise RuntimeError(f"Dtype of {state_name} can only be fp8/fp16/bf16/fp32.") return unscaled @@ -565,7 +586,9 @@ def step(self, closure=None, grad_scaler=None): unscaled_state[name] = self.state[p][name] assert unscaled_state[name].dtype == torch.int16 else: - unscaled = self.get_unscaled_state(p, name) + unscaled = self.get_unscaled_state( + p, name, skip_unscale=self.fuse_unscale + ) unscaled_state[name] = unscaled if self.name_to_dtype_map[name] != torch.float32: unscaled_lists[name].append(unscaled) @@ -748,6 +771,10 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N # Scaling for name in ["exp_avg", "exp_avg_sq", "master_param"]: + if self.fuse_unscale and name in ["exp_avg", "exp_avg_sq"]: + # When fused_unscale is True, the scaling is fused into the Adam kernel. + # The momentums are updated inplace, so we don't need to scale here. + continue if len(unscaled_lists[name]) > 0: for unscaled, scaled, scale in zip( unscaled_lists[name], scaled_lists[name], state_scales[name] From 9eb982e7072bf15eace5248712d81af3d189a832 Mon Sep 17 00:00:00 2001 From: Nicolas Castet <26874160+nvcastet@users.noreply.github.com> Date: Tue, 24 Feb 2026 15:18:49 -0600 Subject: [PATCH 224/521] Fix incorrect MNNVL fabric check (#2626) Signed-off-by: Nicolas Castet Co-authored-by: Przemyslaw Tredak --- .../comm_gemm_overlap/userbuffers/userbuffers-host.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 9c597be306..6ff9d63a2d 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp @@ -122,10 +122,11 @@ bool has_mnnvl_fabric(int device_id) { NVTE_CALL_CHECK_CUDA_NVML(nvmlDeviceGetHandleByIndex_v2, device_id, &local_device); nvmlGpuFabricInfoV_t fabricInfo = {}; fabricInfo.version = nvmlGpuFabricInfo_v2; - fabricInfo.clusterUuid[0] = '\0'; NVTE_CALL_CHECK_CUDA_NVML(nvmlDeviceGetGpuFabricInfoV, local_device, &fabricInfo); NVTE_CALL_CHECK_CUDA_NVML(nvmlShutdown); - if (fabricInfo.state >= NVML_GPU_FABRIC_STATE_COMPLETED && fabricInfo.clusterUuid[0] != '\0') { + const unsigned char zero_uuid[NVML_GPU_FABRIC_UUID_LEN] = {0}; + if (fabricInfo.state == NVML_GPU_FABRIC_STATE_COMPLETED && + memcmp(fabricInfo.clusterUuid, zero_uuid, NVML_GPU_FABRIC_UUID_LEN) != 0) { mnnvl_fabric_support = true; } } From f8b271fc06ac840f63907364c1dbd5833e9aef8e Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 24 Feb 2026 14:05:39 -0800 Subject: [PATCH 225/521] [JAX] Fix FSDP when FSDP+EP is active (#2649) * Fix FSDP when FSDP+EP Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/jax/cpp_extensions/gemm.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 71f133bfc4..a34cb030bf 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -936,7 +936,15 @@ def _parse_operand_output_specs( # Non-contracting dims of RHS always needs to be gathered along the FSDP axis rhs_non_cspecs = tuple( - None if spec is not None and spec == gsr.fsdp_resource else spec + ( + None + if spec is not None + and ( + spec == gsr.fsdp_resource + or (isinstance(spec, tuple) and gsr.fsdp_resource in spec) + ) + else spec + ) for spec in rhs_non_cspecs ) From 7222d8795dbfb841ef5574fb8a57833b0be87af6 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, 25 Feb 2026 02:47:59 +0100 Subject: [PATCH 226/521] [PyTorch Debug] Support precision debug tools for fp8 model parameters. (#2141) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * initial code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * Fix weight quantizer logic in debug mode Add 'or debug' / 'and not debug' conditions to weight quantizer configuration in linear.py, grouped_linear.py, layernorm_linear.py, and layernorm_mlp.py. In debug mode, quantizers are recreated every iteration, so we need to set quantizer states even when weights are already quantized. Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * Update transformer_engine/debug/pytorch/debug_quantization.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@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> --- tests/pytorch/debug/test_log.py | 52 +++++++++++++++++++ tests/pytorch/test_sanity.py | 6 --- transformer_engine/debug/features/api.py | 6 +-- .../debug/features/log_fp8_tensor_stats.py | 25 +++++++-- .../debug/features/log_tensor_stats.py | 8 ++- .../debug/features/utils/stats_buffer.py | 19 ++++--- .../debug/pytorch/debug_quantization.py | 26 +++++++++- transformer_engine/pytorch/module/base.py | 6 ++- .../pytorch/module/grouped_linear.py | 8 ++- .../pytorch/module/layernorm_linear.py | 3 +- .../pytorch/module/layernorm_mlp.py | 5 +- transformer_engine/pytorch/module/linear.py | 3 +- 12 files changed, 136 insertions(+), 31 deletions(-) diff --git a/tests/pytorch/debug/test_log.py b/tests/pytorch/debug/test_log.py index 5d6fc41ac7..b16291ff61 100644 --- a/tests/pytorch/debug/test_log.py +++ b/tests/pytorch/debug/test_log.py @@ -151,6 +151,58 @@ def test_sanity(feature_dirs): assert stat in output, f"Stat {stat} not found in output" +LOG_FP8_MODEL_PARAMETERS_CONFIG_BASE = """ +log: + layers: + layer_name_regex_pattern: .* + enabled: + True + transformer_engine: + LogTensorStats: + enabled: + True + stats: [min] + tensors: [weight, activation, gradient] + freq: 1 + LogFp8TensorStats: + enabled: + True + tensors_struct: + - tensor: activation + stats: [scale_inv_min, scale_inv_max, underflows%] + - tensor: weight + stats: [scale_inv_min, scale_inv_max] + freq: 1 +""" + + +def test_sanity_log_fp8_model_parameters(feature_dirs): + """ + Tests logging stats when model parameters are in fp8. + It tests 3 things: + - LogTensorStats for weight tensor should work without change, + - LogTensorStats and LogFp8TensorStats for non-weight tensors should work without change, + - LogFp8TensorStats should support scale_inv_min, scale_inv_max for weight tensor. + + """ + if not fp8_available: + pytest.skip(reason_for_no_fp8) + + with debug_session(LOG_FP8_MODEL_PARAMETERS_CONFIG_BASE, feature_dirs) as log_dir: + with te.fp8_model_init(recipe=recipe.DelayedScaling()): + model = te.Linear(128, 128, params_dtype=torch.bfloat16) + inp = torch.zeros(128, 128, dtype=torch.bfloat16).cuda() + for _ in range(10): + with te.fp8_autocast(fp8_recipe=recipe.DelayedScaling()): + output = model(inp) + loss = output.sum() + loss.backward() + debug_api.step() + output = read_log(log_dir) + assert output, "Output is empty" + TEDebugState._reset() + + fp8_recipes = [ recipe.MXFP8BlockScaling(), recipe.DelayedScaling(), diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index d47bc553b0..3ef8c0983f 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -551,8 +551,6 @@ def test_sanity_linear(dtype, fp8_recipe, model, skip_wgrad, skip_dgrad, microba @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) def test_sanity_linear_with_zero_tokens(dtype, bs, model, fp8_recipe, fp8_model_params, use_bias): - if NVTE_TEST_NVINSPECT_ENABLED and fp8_model_params: - pytest.skip("Quantized model parameters are not supported in debug mode.") config = model_configs[model] ffn_hidden_size = 4 * config.hidden_size num_tokens = bs * config.max_seqlen_q @@ -599,8 +597,6 @@ def test_sanity_grouped_linear( num_gemms, empty_split, ): - if NVTE_TEST_NVINSPECT_ENABLED and fp8_model_params: - pytest.skip("FP8 model parameters are not supported in debug mode.") config = model_configs[model] ffn_hidden_size = 4 * config.hidden_size # Small batch size used to catch bug from https://github.com/NVIDIA/TransformerEngine/pull/1527. @@ -1222,8 +1218,6 @@ def test_inference_mode( quantization: Optional[str], ) -> None: """Test heuristics for initializing quantized weights""" - if NVTE_TEST_NVINSPECT_ENABLED and quantization is not None: - pytest.skip("Quantized model parameters are not supported in debug mode.") # Tensor dimensions sequence_length = 32 diff --git a/transformer_engine/debug/features/api.py b/transformer_engine/debug/features/api.py index 9c30f87c3b..774fae3594 100644 --- a/transformer_engine/debug/features/api.py +++ b/transformer_engine/debug/features/api.py @@ -244,7 +244,7 @@ def inspect_tensor( config: Dict, layer_name: str, tensor_name: str, - tensor: torch.Tensor, + tensor: Optional[torch.Tensor], rowwise_quantized_tensor: Optional[torch.Tensor], columnwise_quantized_tensor: Optional[torch.Tensor], quantizer: Optional[Quantizer], @@ -262,8 +262,8 @@ def inspect_tensor( layer_name: str tensor_name: str one of [`activation`, `weight`, `gradient`, `output`, `wgrad`, `dgrad`], - tensor: torch.Tensor - tensor in high precision, + tensor: Optional[torch.Tensor] + tensor in high precision. It can be None only if fp8 model parameters are used and tensor name is `weight`. rowwise_quantized_tensor: Optional[torch.Tensor] rowwise quantized tensor, columnwise_quantized_tensor: Optional[torch.Tensor] diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index 108b33fd86..fd18d590ec 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -122,6 +122,10 @@ class LogFp8TensorStats(BaseLogTensorStats): - scale_inv_max - maximum of the inverse of the scaling factors, - mse - mean squared error of the quantized tensor and the original tensor = sum((quantized_tensor - original_tensor)**2) / num_elements, + When collecting stats for the weight tensor with FP8 model parameters enabled, + only "scale_inv_min" and "scale_inv_max" are available. + All other statistics require access to the high precision tensor. + tensors/tensors_struct: List[str] list of tensors to log - activation, @@ -159,7 +163,9 @@ class LogFp8TensorStats(BaseLogTensorStats): end_step: 80 """ - def check_if_stat_is_supported(self, stat: str, current_recipe: str): + def check_if_stat_is_supported( + self, stat: str, current_recipe: str, high_precision_tensor_provided: bool + ): """Returns True if stat is supported, raises ValueError otherwise.""" columnwise = stat.endswith("_columnwise") if columnwise: @@ -167,6 +173,17 @@ def check_if_stat_is_supported(self, stat: str, current_recipe: str): recipe_from_stat, _ = self.get_recipe_from_stat(stat, default_recipe=current_recipe) stat_without_recipe = stat.replace(recipe_from_stat + "_", "") + need_high_precision_tensor_stats = ["underflows%", "overflows%", "mse"] + if ( + stat_without_recipe in need_high_precision_tensor_stats + and not high_precision_tensor_provided + ): + raise ValueError( + f"Stat {stat} requires a high precision tensor to be provided. " + "This feature is not supported for weight tensors when using fp8 model " + "parameters." + ) + if current_recipe == "" and recipe_from_stat == "": raise ValueError( f"Stat {stat} does not contain a recipe name and the current recipe is not set." @@ -290,7 +307,7 @@ def inspect_tensor( tensor_name: str, iteration: int, tp_group: torch.distributed.ProcessGroup, - tensor: torch.Tensor, + tensor: Optional[torch.Tensor], rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, quantizer: Optional[Quantizer] = None, @@ -322,7 +339,9 @@ def inspect_tensor( recipe_name = _get_recipe_name(quantizer) for stat in config["stats"]: - self.check_if_stat_is_supported(stat, recipe_name) + self.check_if_stat_is_supported( + stat, recipe_name, high_precision_tensor_provided=tensor is not None + ) start_step = config.get("start_step", None) end_step = config.get("end_step", None) diff --git a/transformer_engine/debug/features/log_tensor_stats.py b/transformer_engine/debug/features/log_tensor_stats.py index 100fa64481..76e61fab24 100644 --- a/transformer_engine/debug/features/log_tensor_stats.py +++ b/transformer_engine/debug/features/log_tensor_stats.py @@ -180,13 +180,19 @@ def inspect_tensor( tensor_name: str, iteration: int, tp_group: torch.distributed.ProcessGroup, - tensor: torch.Tensor, + tensor: Optional[torch.Tensor], rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, quantizer: Optional[Quantizer] = None, ): # pylint: disable=unused-argument """API call used to collect the data about the tensor before process_tensor()/quantization.""" + # Tensor is None only if fp8 model parameters are used and tensor name is `weight`. + # If one wants to collect stats for this tensor, we need to dequantize it. + if tensor is None: + assert isinstance(rowwise_quantized_tensor, QuantizedTensor) + tensor = rowwise_quantized_tensor.dequantize() + assert ( type(tensor) not in [Float8Tensor, Float8TensorStorage, MXFP8Tensor, MXFP8TensorStorage] and tensor.dtype != torch.uint8 diff --git a/transformer_engine/debug/features/utils/stats_buffer.py b/transformer_engine/debug/features/utils/stats_buffer.py index 9ce56dd76d..ca7f22e2de 100644 --- a/transformer_engine/debug/features/utils/stats_buffer.py +++ b/transformer_engine/debug/features/utils/stats_buffer.py @@ -90,12 +90,19 @@ def feed(self, tensor, iteration, aux_dict=None): if self.modified[0] and not self.reduce_within_microbatch: return - if ( - tensor.numel() == 0 - if hasattr(tensor, "numel") - else all((t is None or t.numel() == 0) for t in tensor.get_data_tensors()) - ): - return + if tensor is not None: + # tensor can be None if we compute fp8 stats for weight and fp8 model parameters are used + # then high precision is not provided and quantized tensor from aux_dict is used. + + # This condition prevents computation of stats for empty tensor. + # This will not happen for weight - since it is the only situation then tensor can be None, + # we do not need to check similar condition for weight. + if ( + tensor.numel() == 0 + if hasattr(tensor, "numel") + else all((t is None or t.numel() == 0) for t in tensor.get_data_tensors()) + ): + return # save stats for tensor to tmp buffer for stat_name in self.stats_to_compute: diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 455079143b..5624970547 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -267,7 +267,7 @@ def _call_inspect_tensor_api( "rowwise_quantized_tensor": rowwise_gemm_tensor, "quantizer": self.parent_quantizer, } - if tensor is not None and self.inspect_tensor_enabled: + if self.inspect_tensor_enabled: debug_api.transformer_engine.inspect_tensor(**args) if self.output_tensor: @@ -559,6 +559,30 @@ def set_usage(self, rowwise: bool = None, columnwise: bool = None): if not self.output_tensor: self._update_parent_quantizer_usage() + def wrap_quantized_tensor(self, tensor: QuantizedTensor): + """ + Wraps the quantized tensor with the debug quantizer. + It is used for weight tensors when fp8 model parameters are enabled. + """ + + assert ( + self.rowwise_tensor_plan == STANDARD_QUANTIZE + and self.columnwise_tensor_plan == STANDARD_QUANTIZE + ), ( + "[NVTORCH INSPECT ERROR] Weight tensor with fp8 model parameters enabled cannot be" + " modified by any feature." + ) + + self._call_inspect_tensor_api(None, tensor, tensor) + + return DebugQuantizedTensor( + rowwise_gemm_tensor=tensor, + columnwise_gemm_tensor=tensor, + quantizer=self, + layer_name=self.layer_name, + tensor_name=self.tensor_name, + ) + @classmethod def multi_tensor_quantize( cls, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 09b12afa21..4858383c26 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1390,6 +1390,10 @@ def get_weight_workspace( rowwise_usage=update_rowwise_usage, columnwise_usage=update_columnwise_usage, ) + + if isinstance(quantizer, DebugQuantizer): + tensor = quantizer.wrap_quantized_tensor(tensor) + return tensor # Try getting workspace from cache @@ -1585,8 +1589,6 @@ def no_debug_features_active(self, quantizers): if not run_current: return True - if self.primary_weights_in_fp8: - raise RuntimeError("FP8 weights are not supported in debug mode.") return False def _validate_name(self): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 2f859e748b..b381073d78 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -123,8 +123,9 @@ def forward( and not in_fp8_activation_recompute_phase() ) # No need to set the quantizer states if weight is already quantized - if weight_quantizers[0] is not None and not isinstance( - weights[0], QuantizedTensorStorage + # for debug mode we create quantizer every iteration, thus we need to set the quantizer states + if weight_quantizers[0] is not None and ( + not isinstance(weights[0], QuantizedTensorStorage) or debug ): for weight_quantizer in weight_quantizers: weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) @@ -874,9 +875,6 @@ def forward( debug = False quantizers = self._get_quantizers() - if isinstance(weight_tensors, QuantizedTensorStorage): - raise RuntimeError("FP8 weights are not supported in debug mode.") - ( input_quantizers, weight_quantizers, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 702916696b..27632db15b 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -293,7 +293,8 @@ def forward( # Configure quantizer # If weight is already quantized, no need to set quantizer states - if is_weight_param_quantized: + # 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: weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 4532ea60e7..b8823e46ca 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -473,12 +473,13 @@ def _forward( # FP8 cast to workspace buffer update_workspace = is_first_microbatch is None or is_first_microbatch # No need to set the quantizer states if weights are already quantized - if isinstance(fc1_weight, QuantizedTensorStorage): + # 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) - if isinstance(fc2_weight, QuantizedTensorStorage): + 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) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 23ad8cacb0..a55429d33d 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -253,7 +253,8 @@ def forward( if fp8 or debug: # Configure quantizer # No need to set the quantizer states if weight is already quantized - if weight_quantizer is not None and not isinstance(weight, QuantizedTensor): + # 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 if not columnwise_usage: columnwise_usage = ( From df0ef6e2b9cefbfc11cc65ee6eebba66a55e76ef Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 24 Feb 2026 20:16:33 -0800 Subject: [PATCH 227/521] remove deprecated qkv/kv_packed apis (#2696) Signed-off-by: Sudhakar Singh --- .../common/fused_attn/fused_attn.cpp | 650 ------------------ .../include/transformer_engine/fused_attn.h | 284 -------- 2 files changed, 934 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index b5679280c6..abdce7fdac 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -15,74 +15,6 @@ #include "fused_attn_fp8.h" #include "utils.h" -namespace { -// Helper function to create a tensor view with modified shape and optional pointer offset -transformer_engine::Tensor make_tensor_view(const transformer_engine::Tensor *source, - const std::vector &shape, - size_t offset_bytes = 0) { - transformer_engine::Tensor view = *source; - if (offset_bytes > 0) { - view.data.dptr = static_cast(static_cast(source->data.dptr) + offset_bytes); - } - view.data.shape = shape; - view.nvte_tensor = 0; // Mark as unmanaged/local tensor view - return view; -} - -// Helper function to calculate stride in bytes for packed QKV tensor unpacking -size_t calculate_qkv_stride(NVTE_QKV_Layout_Group layout_group, transformer_engine::DType dtype, - size_t h, size_t d) { - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - stride = (transformer_engine::typeToNumBits(dtype) * h * d) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - stride = (transformer_engine::typeToNumBits(dtype) * d) / 8; - } - return stride; -} - -// Helper function to determine unpacked shape for QKV packed tensor -std::vector calculate_qkv_unpacked_shape(const transformer_engine::Tensor *qkv_tensor, - size_t h, size_t d) { - std::vector unpacked_shape; - if (qkv_tensor->data.shape.size() == 4) { - // T3HD or TH3D (4D) -> THD (3D): remove dimension "3" at position 1 - unpacked_shape = {qkv_tensor->data.shape[0], h, d}; - } else { - // BS3HD/SB3HD or BSH3D/SBH3D (5D) -> BSHD/SBHD (4D): remove dimension "3" at position 2 - unpacked_shape = {qkv_tensor->data.shape[0], qkv_tensor->data.shape[1], h, d}; - } - return unpacked_shape; -} - -// Helper function to calculate stride for packed KV tensor unpacking -size_t calculate_kv_stride(NVTE_QKV_Layout_Group layout_group, transformer_engine::DType dtype, - size_t h_kv, size_t d) { - size_t stride = 0; - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - stride = (transformer_engine::typeToNumBits(dtype) * h_kv * d) / 8; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - stride = (transformer_engine::typeToNumBits(dtype) * d) / 8; - } - return stride; -} - -// Helper function to determine unpacked shape for KV packed tensor -std::vector calculate_kv_unpacked_shape(const transformer_engine::Tensor *kv_tensor, - NVTE_QKV_Layout_Group layout_group, - NVTE_QKV_Format kv_format, size_t t_kv, size_t h_kv, - size_t d) { - std::vector unpacked_kv_shape; - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - unpacked_kv_shape = {t_kv, h_kv, d}; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD || - layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - unpacked_kv_shape = {kv_tensor->data.shape[0], kv_tensor->data.shape[1], h_kv, d}; - } - return unpacked_kv_shape; -} -} // namespace - // map NVTE_QKV_Layout to NVTE_QKV_Layout_Group NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { switch (qkv_layout) { @@ -516,588 +448,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( return backend; } -// NVTE fused attention FWD with packed QKV -// DEPRECATED: This API is deprecated. -// Please use nvte_fused_attn_fwd with separate Q, K, V tensors instead. -void nvte_fused_attn_fwd_qkvpacked( - const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, - const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, - 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_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) { - NVTE_API_CALL(nvte_flash_attn_fwd_qkvpacked); - using namespace transformer_engine; - - const Tensor *input_cu_seqlens = convertNVTETensorCheck(cu_seqlens); - const Tensor *input_cu_seqlens_padded = convertNVTETensorCheck(cu_seqlens_padded); - const Tensor *input_rng_state = convertNVTETensorCheck(rng_state); - const Tensor *input_QKV = convertNVTETensorCheck(QKV); - const Tensor *input_Bias = convertNVTETensorCheck(Bias); - const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(SoftmaxOffset); - Tensor *input_output_S = convertNVTETensorCheck(S); - Tensor *output_O = convertNVTETensorCheck(O); - Tensor *wkspace = convertNVTETensor(workspace); - - auto ndim = input_QKV->data.shape.size(); - size_t b = input_cu_seqlens->data.shape[0] - 1; - size_t h = 0; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - h = input_QKV->data.shape[ndim - 2]; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - h = input_QKV->data.shape[ndim - 3]; - } else { - NVTE_ERROR("nvte_fused_attn_fwd_qkvpacked only supports H3D and 3HD layouts!"); - } - size_t d = input_QKV->data.shape[ndim - 1]; - size_t t = 0; - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if (qkv_format == NVTE_QKV_Format::NVTE_THD) { - t = input_QKV->data.shape[0]; - } - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTEDType QKV_type = static_cast(input_QKV->data.dtype); - - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h, h, max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, return_max_logit, - cuda_graph, false); - - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { -#if (CUDNN_VERSION >= 8901) - // Unpack QKV and call the non-packed function - const auto QKV_type = input_QKV->data.dtype; - size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); - std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); - - // Create tensor views for Q, K, V - Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); - Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); - Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); - - fused_attn_max_512_fwd(b, h, max_seqlen, max_seqlen, d, is_training, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, &Q_view, &K_view, &V_view, - input_Bias, output_O, Aux_CTX_Tensors, input_cu_seqlens, - input_cu_seqlens, 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) - // Unpack QKV and call the non-packed function - const auto QKV_type = input_QKV->data.dtype; - size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); - std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); - - // Create tensor views for Q, K, V - Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); - Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); - Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); - - fused_attn_arbitrary_seqlen_fwd( - b, h, h, max_seqlen, max_seqlen, d, d, t, t, 0, 0, 0, 0, 0, 0, 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, &Q_view, &K_view, &V_view, - input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens, - input_cu_seqlens, input_cu_seqlens_padded, input_cu_seqlens_padded, nullptr, nullptr, - 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) - // Unpack QKV and call the non-packed function - const auto QKV_type = input_QKV->data.dtype; - size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); - std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); - - // Create tensor views for Q, K, V - Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); - Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); - Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); - - fused_attn_fp8_fwd(b, h, h, max_seqlen, max_seqlen, d, is_training, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, &Q_view, &K_view, &V_view, - input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens, - input_cu_seqlens, 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"); - } -} -// NVTE fused attention BWD with packed QKV -// DEPRECATED: This API is deprecated. -// Please use nvte_fused_attn_bwd with separate Q, K, V tensors instead. -void nvte_fused_attn_bwd_qkvpacked( - const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, - NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, - NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, - size_t max_seqlen, 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_API_CALL(nvte_flash_attn_bwd_qkvpacked); - using namespace transformer_engine; - - const Tensor *input_cu_seqlens = convertNVTETensorCheck(cu_seqlens); - const Tensor *input_cu_seqlens_padded = convertNVTETensorCheck(cu_seqlens_padded); - const Tensor *input_QKV = convertNVTETensorCheck(QKV); - const Tensor *input_O = convertNVTETensorCheck(O); - const Tensor *input_dO = convertNVTETensorCheck(dO); - const Tensor *input_S = convertNVTETensorCheck(S); - Tensor *input_output_dP = convertNVTETensorCheck(dP); - Tensor *output_dQKV = convertNVTETensorCheck(dQKV); - Tensor *output_dBias = convertNVTETensorCheck(dBias); - Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(dSoftmaxOffset); - Tensor *wkspace = convertNVTETensor(workspace); - - auto ndim = input_QKV->data.shape.size(); - size_t b = input_cu_seqlens->data.shape[0] - 1; - size_t h = 0; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { - h = input_QKV->data.shape[ndim - 2]; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_H3D) { - h = input_QKV->data.shape[ndim - 3]; - } else { - NVTE_ERROR("nvte_fused_attn_fwd_qkvpacked only supports H3D and 3HD layouts!"); - } - size_t d = input_QKV->data.shape[ndim - 1]; - size_t t = 0; - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if (qkv_format == NVTE_QKV_Format::NVTE_THD) { - t = input_QKV->data.shape[0]; - } - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTEDType QKV_type = static_cast(input_QKV->data.dtype); - - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - true, QKV_type, QKV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h, h, - max_seqlen, max_seqlen, d, d, window_size_left, window_size_right, false, 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]); - - // Unpack QKV and dQKV and call the non-packed function - const auto QKV_type = input_QKV->data.dtype; - size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); - std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); - - // Create tensor views for Q, K, V and dQ, dK, dV - Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); - Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); - Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); - - Tensor dQ_view = make_tensor_view(output_dQKV, unpacked_shape); - Tensor dK_view = make_tensor_view(output_dQKV, unpacked_shape, stride); - Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); - - fused_attn_max_512_bwd(b, h, max_seqlen, max_seqlen, d, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, &Q_view, &K_view, &V_view, input_dO, output_S, - &dQ_view, &dK_view, &dV_view, output_dBias, input_cu_seqlens, - input_cu_seqlens, 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++]); - Tensor *input_Bias, *input_SoftmaxOffset; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - - // Unpack QKV and dQKV and call the non-packed function - const auto QKV_type = input_QKV->data.dtype; - size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); - std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); - - // Create tensor views for Q, K, V and dQ, dK, dV - Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); - Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); - Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); - - Tensor dQ_view = make_tensor_view(output_dQKV, unpacked_shape); - Tensor dK_view = make_tensor_view(output_dQKV, unpacked_shape, stride); - Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); - - fused_attn_arbitrary_seqlen_bwd( - b, h, h, max_seqlen, max_seqlen, d, d, t, t, attn_scale, dropout, qkv_layout, bias_type, - attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, - deterministic, &Q_view, &K_view, &V_view, input_O, input_dO, input_Bias, - input_SoftmaxOffset, output_S, &dQ_view, &dK_view, &dV_view, output_dBias, - output_dSoftmaxOffset, input_cu_seqlens, input_cu_seqlens, input_cu_seqlens_padded, - input_cu_seqlens_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 - } 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]); - - // Unpack QKV and dQKV and call the non-packed function - const auto QKV_type = input_QKV->data.dtype; - size_t stride = calculate_qkv_stride(layout_group, QKV_type, h, d); - std::vector unpacked_shape = calculate_qkv_unpacked_shape(input_QKV, h, d); - - // Create tensor views for Q, K, V and dQ, dK, dV - Tensor Q_view = make_tensor_view(input_QKV, unpacked_shape); - Tensor K_view = make_tensor_view(input_QKV, unpacked_shape, stride); - Tensor V_view = make_tensor_view(input_QKV, unpacked_shape, 2 * stride); - - Tensor dQ_view = make_tensor_view(output_dQKV, unpacked_shape); - Tensor dK_view = make_tensor_view(output_dQKV, unpacked_shape, stride); - Tensor dV_view = make_tensor_view(output_dQKV, unpacked_shape, 2 * stride); - - fused_attn_fp8_bwd(b, h, h, max_seqlen, max_seqlen, d, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, deterministic, &Q_view, &K_view, &V_view, input_O, - input_dO, input_M, input_ZInv, input_S, input_output_dP, &dQ_view, &dK_view, - &dV_view, input_cu_seqlens, input_cu_seqlens, 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"); - } -} -// NVTE fused attention FWD with packed KV -// DEPRECATED: This API is deprecated. -// Please use nvte_fused_attn_fwd with separate Q, K, V tensors instead. -void nvte_fused_attn_fwd_kvpacked( - const NVTETensor Q, const NVTETensor KV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, - NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, - const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - 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_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) { - NVTE_API_CALL(nvte_flash_attn_fwd_kvpacked); - using namespace transformer_engine; - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(cu_seqlens_kv_padded); - const Tensor *input_page_table_k = convertNVTETensorCheck(page_table_k); - const Tensor *input_page_table_v = convertNVTETensorCheck(page_table_v); - const Tensor *input_rng_state = convertNVTETensorCheck(rng_state); - const Tensor *input_Q = convertNVTETensorCheck(Q); - const Tensor *input_KV = convertNVTETensorCheck(KV); - const Tensor *input_Bias = convertNVTETensorCheck(Bias); - const Tensor *input_SoftmaxOffset = convertNVTETensorCheck(SoftmaxOffset); - Tensor *input_output_S = convertNVTETensorCheck(S); - Tensor *output_O = convertNVTETensorCheck(O); - Tensor *wkspace = convertNVTETensor(workspace); - - size_t b = input_cu_seqlens_q->data.shape[0] - 1; - auto ndim = input_Q->data.shape.size(); - size_t h_q = input_Q->data.shape[ndim - 2]; - size_t d = input_Q->data.shape[ndim - 1]; - auto ndim_kv = input_KV->data.shape.size(); - size_t h_kv = 0; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - h_kv = input_KV->data.shape[ndim_kv - 2]; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - h_kv = input_KV->data.shape[ndim_kv - 3]; - } else { - NVTE_ERROR("nvte_fused_attn_fwd_kvpacked only supports HD_H2D and HD_2HD layouts!"); - } - 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); - 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_KV->data.shape[0]; - } - int64_t num_pages_k = 0; - int64_t num_pages_v = 0; - int64_t page_size_k = 0; - int64_t page_size_v = 0; - int64_t max_pages_per_seq_k = 0; - int64_t max_pages_per_seq_v = 0; - if (input_page_table_k->data.dptr != nullptr) { - max_pages_per_seq_k = input_page_table_k->data.shape[1]; - } - if (input_page_table_v->data.dptr != nullptr) { - max_pages_per_seq_v = input_page_table_v->data.shape[1]; - } - if (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD) { - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - if (kv_format == NVTE_QKV_Format::NVTE_BSHD) { - num_pages_k = input_KV->data.shape[0]; - page_size_k = input_KV->data.shape[1]; - num_pages_v = num_pages_v; - page_size_v = page_size_v; - } else if (kv_format == NVTE_QKV_Format::NVTE_SBHD) { - num_pages_k = input_KV->data.shape[1]; - page_size_k = input_KV->data.shape[0]; - num_pages_v = num_pages_v; - page_size_v = page_size_v; - } - } - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTEDType Q_type = static_cast(input_Q->data.dtype); - const NVTEDType KV_type = static_cast(input_KV->data.dtype); - - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - is_training, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, - h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, - return_max_logit, cuda_graph, false); - - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { -#if (CUDNN_VERSION >= 8901) - // Unpack KV and call the non-packed function - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = calculate_kv_stride(layout_group, input_Q->data.dtype, h_kv, d); - std::vector unpacked_kv_shape = - calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); - - Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); - Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); - - fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d, is_training, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, input_Q, &K_view, &V_view, - 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 >= 8903) - // Unpack KV and call the non-packed function - const auto Q_type = input_Q->data.dtype; - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); - std::vector unpacked_kv_shape = - calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); - - Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); - Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); - - fused_attn_arbitrary_seqlen_fwd( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, 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, &K_view, &V_view, - 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.3 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) - // Unpack KV and call the non-packed function - const auto Q_type = input_Q->data.dtype; - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); - std::vector unpacked_kv_shape = - calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); - - Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); - Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); - - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, is_training, attn_scale, - dropout, qkv_layout, bias_type, attn_mask_type, input_Q, &K_view, &V_view, - 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"); - } -} -// NVTE fused attention BWD with packed KV -// DEPRECATED: This API is deprecated. -// Please use nvte_fused_attn_bwd with separate Q, K, V tensors instead. -void nvte_fused_attn_bwd_kvpacked( - const NVTETensor Q, const NVTETensor KV, const NVTETensor O, const NVTETensor dO, - const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, - NVTETensor dKV, NVTETensor dBias, NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens_q, - const NVTETensor cu_seqlens_kv, 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_API_CALL(nvte_flash_attn_bwd_kvpacked); - using namespace transformer_engine; - const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); - const Tensor *input_cu_seqlens_kv = convertNVTETensorCheck(cu_seqlens_kv); - const Tensor *input_cu_seqlens_q_padded = convertNVTETensorCheck(cu_seqlens_q_padded); - const Tensor *input_cu_seqlens_kv_padded = convertNVTETensorCheck(cu_seqlens_kv_padded); - const Tensor *input_Q = convertNVTETensorCheck(Q); - const Tensor *input_KV = convertNVTETensorCheck(KV); - const Tensor *input_O = convertNVTETensorCheck(O); - const Tensor *input_dO = convertNVTETensorCheck(dO); - const Tensor *input_S = convertNVTETensorCheck(S); - Tensor *input_output_dP = convertNVTETensorCheck(dP); - Tensor *output_dQ = convertNVTETensorCheck(dQ); - Tensor *output_dKV = convertNVTETensorCheck(dKV); - Tensor *output_dBias = convertNVTETensorCheck(dBias); - Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(dSoftmaxOffset); - Tensor *wkspace = convertNVTETensor(workspace); - - size_t b = input_cu_seqlens_q->data.shape[0] - 1; - auto ndim = input_Q->data.shape.size(); - size_t h_q = input_Q->data.shape[ndim - 2]; - size_t d = input_Q->data.shape[ndim - 1]; - auto ndim_kv = input_KV->data.shape.size(); - size_t h_kv = 0; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_2HD) { - h_kv = input_KV->data.shape[ndim_kv - 2]; - } else if (layout_group == NVTE_QKV_Layout_Group::NVTE_HD_H2D) { - h_kv = input_KV->data.shape[ndim_kv - 3]; - } else { - NVTE_ERROR("nvte_fused_attn_fwd_kvpacked only supports HD_H2D and HD_2HD layouts!"); - } - 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); - 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_KV->data.shape[0]; - } - - auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); - const NVTEDType Q_type = static_cast(input_Q->data.dtype); - const NVTEDType KV_type = static_cast(input_KV->data.dtype); - - NVTE_Fused_Attn_Backend fused_attention_backend = nvte_get_fused_attn_backend( - true, Q_type, KV_type, qkv_layout, bias_type, attn_mask_type, softmax_type, dropout, h_q, - h_kv, max_seqlen_q, max_seqlen_kv, d, d, window_size_left, window_size_right, false, - 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]); - - // Unpack KV and dKV and call the non-packed function - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = calculate_kv_stride(layout_group, input_Q->data.dtype, h_kv, d); - std::vector unpacked_kv_shape = - calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); - - Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); - Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); - - Tensor dK_view = make_tensor_view(output_dKV, unpacked_kv_shape); - Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); - - fused_attn_max_512_bwd(b, h_q, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, qkv_layout, - bias_type, attn_mask_type, input_Q, &K_view, &V_view, input_dO, output_S, - output_dQ, &dK_view, &dV_view, 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 >= 8903) - size_t i = 0; - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - Tensor *input_Bias, *input_SoftmaxOffset; - if ((bias_type != NVTE_NO_BIAS) && (bias_type != NVTE_ALIBI)) { - input_Bias = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - if (softmax_type != NVTE_VANILLA_SOFTMAX) { - input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } - - // Unpack KV and dKV and call the non-packed function - const auto Q_type = input_Q->data.dtype; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); - std::vector unpacked_kv_shape = - calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); - - Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); - Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); - - // Create tensor views for dK, dV - Tensor dK_view = make_tensor_view(output_dKV, unpacked_kv_shape); - Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); - - fused_attn_arbitrary_seqlen_bwd( - b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, d, 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, &K_view, &V_view, input_O, input_dO, - input_Bias, input_SoftmaxOffset, output_S, output_dQ, &dK_view, &dV_view, 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.3 is required for BF16/FP16 fused attention " - "with arbitrary sequence length. \n"; - NVTE_ERROR(err_msg); -#endif - } 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]); - - // Unpack KV and dKV and call the non-packed function - const auto Q_type = input_Q->data.dtype; - NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); - size_t stride = calculate_kv_stride(layout_group, Q_type, h_kv, d); - std::vector unpacked_kv_shape = - calculate_kv_unpacked_shape(input_KV, layout_group, kv_format, t_kv, h_kv, d); - - Tensor K_view = make_tensor_view(input_KV, unpacked_kv_shape); - Tensor V_view = make_tensor_view(input_KV, unpacked_kv_shape, stride); - - Tensor dK_view = make_tensor_view(output_dKV, unpacked_kv_shape); - Tensor dV_view = make_tensor_view(output_dKV, unpacked_kv_shape, stride); - - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, deterministic, input_Q, &K_view, - &V_view, input_O, input_dO, input_M, input_ZInv, input_S, input_output_dP, - output_dQ, &dK_view, &dV_view, 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"); - } -} // NVTE fused attention FWD with separate Q, K and V void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETensor V, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index cddd3d7506..8169bf22e2 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -217,290 +217,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, int64_t window_size_left, int64_t window_size_right, bool return_max_logit, bool cuda_graph, bool deterministic); -/*! \brief Compute dot product attention with packed QKV input. - * - * \deprecated Please use `nvte_fused_attn_fwd` with separate Q, K, V tensors instead. - * - * Computes: - * - P = Q * Transpose(K) + Bias - * - S = ScaleMaskSoftmax(P) - * - D = Dropout(S) - * - O = D * Transpose(V) - * - * Support Matrix: - \verbatim - | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 0 | FP16/BF16 | BS3HD,SB3HD | 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 | - | 2 | FP8 | T3HD | NO_BIAS | PADDING_MASK | Yes | <= 512, % 64 == 0 | 64 | - \endverbatim - * - * Notes: - * - * Tensor `cu_seqlens_padded` helps identify the correct offsets of different sequences - * in tensors Q, K, V and O. - * When the QKV format (`nvte_get_qkv_format(qkv_layout)`) is `bshd` or `sbhd`, - * the offset tensor is not used in the attention calculation and can be set to empty `NVTETensor`. - * When the QKV format is `thd`, this tensor should follow the following rules. - * When there is no padding between sequences, the offset tensor should be equal to `cu_seqlens`, - * When there is padding between sequences, users are responsible to adjust the offsets as needed. - * For example, a tensor of 4 sequences `[a, PAD, b, b, c, PAD, PAD, d, d]` should have - * `cu_seqlens = [0, 1, 3, 4, 6]` and `cu_seqlens_padded= [0, 2, 4, 7, 9]`. - * - * \param[in] QKV The QKV tensor in packed format, H3D or 3HD. - * \param[in] Bias The Bias tensor. - * \param[in] SoftmaxOffset The SoftmaxOffset tensor. - * \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. - * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] cu_seqlens_padded Cumulative sequence offsets for QKV, [batch_size + 1]. - * \param[in] rng_state Seed and offset of CUDA random number generator. - * \param[in] max_seqlen Max sequence length used for computing, - * it may be >= max(seqlen_i) for i=0,...batch_size-1. - * \param[in] is_training Whether this is in training mode or inference. - * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] attn_scale Scaling factor for Q * K.T. - * \param[in] dropout Dropout probability. - * \param[in] qkv_layout QKV tensor's layout. - * \param[in] bias_type Bias type. - * \param[in] attn_mask_type Attention mask type. - * \param[in] softmax_type Attention softmax type. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. - * \param[in] workspace Workspace tensor. - * \param[in] stream CUDA stream used for this operation. - */ -[[deprecated( - "nvte_fused_attn_fwd_qkvpacked() is deprecated. Please use nvte_fused_attn_fwd() with separate " - "Q, K, V tensors instead.")]] -void nvte_fused_attn_fwd_qkvpacked( - const NVTETensor QKV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, NVTETensor S, - NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens, - const NVTETensor cu_seqlens_padded, const NVTETensor rng_state, size_t max_seqlen, - 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_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); - -/*! \brief Compute the backward of the dot product attention with packed QKV input. - * - * \deprecated Please use `nvte_fused_attn_bwd` with separate Q, K, V tensors instead. - * - * Support Matrix: - \verbatim - | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 0 | FP16/BF16 | BS3HD,SB3HD | 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 | - | 2 | FP8 | T3HD | NO_BIAS | PADDING_MASK | Yes | <= 512, % 64 == 0 | 64 | - \endverbatim - * - * Notes: - * - * Tensor `cu_seqlens_padded` helps identify the correct offsets of different sequences - * in tensors Q, K, V and O. - * When the QKV format (`nvte_get_qkv_format(qkv_layout)`) is `bshd` or `sbhd`, - * the offset tensor is not used in the attention calculation and can be set to empty `NVTETensor`. - * When the QKV format is `thd`, this tensor should follow the following rules. - * When there is no padding between sequences, the offset tensor should be equal to `cu_seqlens`, - * When there is padding between sequences, users are responsible to adjust the offsets as needed. - * For example, a tensor of 4 sequences `[a, PAD, b, b, c, PAD, PAD, d, d]` should have - * `cu_seqlens = [0, 1, 3, 4, 6]` and `cu_seqlens_padded= [0, 2, 4, 7, 9]`. - * - * \param[in] QKV The QKV tensor in packed format, H3D or 3HD. - * \param[in] O The O tensor from forward. - * \param[in] dO The gradient of the O tensor. - * \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. - * \param[out] dQKV The gradient of the QKV tensor. - * \param[out] dBias The gradient of the Bias tensor. - * \param[out] dSoftmaxOffset The gradient of the SoftmaxOffset tensor. - * \param[in] cu_seqlens Cumulative sequence lengths, [batch_size + 1]. - * \param[in] cu_seqlens_padded Cumulative sequence offsets for QKV, [batch_size + 1]. - * \param[in] max_seqlen Max sequence length used for computing, - * it may be >= max(seqlen_i) for i=0,...batch_size-1. - * \param[in] attn_scale Scaling factor for Q * K.T. - * \param[in] dropout Dropout probability. - * \param[in] qkv_layout QKV tensor's layout. - * \param[in] bias_type Bias type. - * \param[in] attn_mask_type Attention mask type. - * \param[in] softmax_type Attention softmax type. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. - * \param[in] deterministic Whether to execute with deterministic behaviours. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] workspace Workspace tensor. - * \param[in] stream CUDA stream used for this operation. - */ -[[deprecated( - "nvte_fused_attn_bwd_qkvpacked() is deprecated. Please use nvte_fused_attn_bwd() with separate " - "Q, K, V tensors instead.")]] -void nvte_fused_attn_bwd_qkvpacked( - const NVTETensor QKV, const NVTETensor O, const NVTETensor dO, const NVTETensor S, - NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQKV, NVTETensor dBias, - NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens, const NVTETensor cu_seqlens_padded, - size_t max_seqlen, 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); - -/*! \brief Compute dot product attention with packed KV input. - * - * \deprecated Please use `nvte_fused_attn_fwd` with separate Q, K, V tensors instead. - * - * Computes: - * - P = Q * Transpose(K) + Bias - * - S = ScaleMaskSoftmax(P) - * - D = Dropout(S) - * - O = D * Transpose(V) - * - * Support Matrix: - \verbatim - | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 0 | FP16/BF16 | BSHD_BS2HD,SBHD_SB2HD | NO/POST_SCALE_BIAS | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | <= 512, % 64 == 0 | 64 | - | 1 | FP16/BF16 | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | > 512, % 64 == 0 | <= 128, % 8 == 0 | - \endverbatim - * - * Notes: - * - * Tensors `cu_seqlens_q_padded` and `cu_seqlens_kv_padded` - * help identify the correct offsets of different sequences in tensors Q, K, V and O. - * When the QKV format (`nvte_get_qkv_format(qkv_layout)`) is `bshd` or `sbhd`, - * offset tensors are not used in the attention calculation and can be set to empty `NVTETensor`s. - * When the QKV format is `thd`, these tensors should follow the following rules. - * When there is no padding between sequences, the offset tensors should be equal to - * `cu_seqlens_q` and `cu_seqlens_kv` respectively. - * When there is padding between sequences, users are responsible to adjust the offsets as needed. - * For example, a tensor of 4 sequences `[a, PAD, b, b, c, PAD, PAD, d, d]` should have - * `cu_seqlens = [0, 1, 3, 4, 6]` and `cu_seqlens_padded= [0, 2, 4, 7, 9]`. - * - * \param[in] Q The Q tensor, in HD layouts. - * \param[in] KV The KV tensor, in 2HD or H2D layouts. - * \param[in] Bias The Bias tensor. - * \param[in] SoftmaxOffset The SoftmaxOffset tensor. - * \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. - * \param[in] cu_seqlens_q Cumulative sequence lengths for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv Cumulative sequence lengths for KV, [batch_size + 1]. - * \param[in] cu_seqlens_q_padded Cumulative sequence offsets for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv_padded Cumulative sequence offsets for KV, [batch_size + 1]. - * \param[in] page_table_k Page table for K cache, [batch_size, max_pages_per_seq_k]. - * \param[in] page_table_v Page table for V cache, [batch_size, max_pages_per_seq_v]. - * \param[in] rng_state Seed and offset of CUDA random number generator. - * \param[in] max_seqlen_q Max sequence length used for computing for Q. - * it may be >= max(seqlen_q_i) for i=0,...batch_size-1. - * \param[in] max_seqlen_kv Max sequence length used for computing for KV. - * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. - * \param[in] is_training Whether this is in training mode or inference. - * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] attn_scale Scaling factor for Q * K.T. - * \param[in] dropout Dropout probability. - * \param[in] qkv_layout QKV tensor's layout. - * \param[in] bias_type Bias type. - * \param[in] attn_mask_type Attention mask type. - * \param[in] softmax_type Attention softmax type. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. - * \param[in] workspace Workspace tensor. - * \param[in] stream CUDA stream used for this operation. - */ -[[deprecated( - "nvte_fused_attn_fwd_kvpacked() is deprecated. Please use nvte_fused_attn_fwd() with separate " - "Q, K, V tensors instead.")]] -void nvte_fused_attn_fwd_kvpacked( - const NVTETensor Q, const NVTETensor KV, const NVTETensor Bias, const NVTETensor SoftmaxOffset, - NVTETensor S, NVTETensor O, NVTETensorPack *Aux_CTX_Tensors, const NVTETensor cu_seqlens_q, - const NVTETensor cu_seqlens_kv, const NVTETensor cu_seqlens_q_padded, - const NVTETensor cu_seqlens_kv_padded, const NVTETensor page_table_k, - 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_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); - -/*! \brief Compute the backward of the dot product attention with packed KV input. - * - * \deprecated Please use `nvte_fused_attn_bwd` with separate Q, K, V tensors instead. - * - * Support Matrix: - \verbatim - | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 0 | FP16/BF16 | BSHD_BS2HD,SBHD_SB2HD | NO/POST_SCALE_BIAS | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | <= 512, % 64 == 0 | 64 | - | 1 | FP16/BF16 | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | > 512, % 64 == 0 | <= 128, % 8 == 0 | - \endverbatim - * - * Notes: - * - * Tensors `cu_seqlens_q_padded` and `cu_seqlens_kv_padded` - * help identify the correct offsets of different sequences in tensors Q, K, V and O. - * When the QKV format (`nvte_get_qkv_format(qkv_layout)`) is `bshd` or `sbhd`, - * offset tensors are not used in the attention calculation and can be set to empty `NVTETensor`s. - * When the QKV format is `thd`, these tensors should follow the following rules. - * When there is no padding between sequences, the offset tensors should be equal to - * `cu_seqlens_q` and `cu_seqlens_kv` respectively. - * When there is padding between sequences, users are responsible to adjust the offsets as needed. - * For example, a tensor of 4 sequences `[a, PAD, b, b, c, PAD, PAD, d, d]` should have - * `cu_seqlens = [0, 1, 3, 4, 6]` and `cu_seqlens_padded= [0, 2, 4, 7, 9]`. - * - * \param[in] Q The Q tensor, in HD layouts. - * \param[in] KV The KV tensor, in H2D or 2HD layouts. - * \param[in] O The O tensor from forward. - * \param[in] dO The gradient of the O tensor. - * \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. - * \param[out] dQ The gradient of the Q tensor. - * \param[out] dKV The gradient of the KV tensor. - * \param[out] dBias The gradient of the Bias tensor. - * \param[out] dSoftmaxOffset The gradient of the SoftmaxOffset tensor. - * \param[in] cu_seqlens_q Cumulative sequence lengths for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv Cumulative sequence lengths for KV, [batch_size + 1]. - * \param[in] cu_seqlens_q_padded Cumulative sequence offsets for Q, [batch_size + 1]. - * \param[in] cu_seqlens_kv_padded Cumulative sequence offsets for KV, [batch_size + 1]. - * \param[in] max_seqlen_q Max sequence length used for computing for Q. - * it may be >= max(seqlen_q_i) for i=0,...batch_size-1. - * \param[in] max_seqlen_kv Max sequence length used for computing for KV. - * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. - * \param[in] attn_scale Scaling factor for Q * K.T. - * \param[in] dropout Dropout probability. - * \param[in] qkv_layout QKV tensor's layout. - * \param[in] bias_type Bias type. - * \param[in] attn_mask_type Attention mask type. - * \param[in] softmax_type Attention softmax type. - * \param[in] window_size_left Sliding window size (the left half). - * \param[in] window_size_right Sliding window size (the right half). - * \param[in] bottom_right_diagonal Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. - * \param[in] deterministic Whether to execute with deterministic behaviours. - * \param[in] cuda_graph Whether cuda graph capture is enabled or not. - * \param[in] workspace Workspace tensor. - * \param[in] stream CUDA stream used for this operation. - */ -[[deprecated( - "nvte_fused_attn_bwd_kvpacked() is deprecated. Please use nvte_fused_attn_bwd() with separate " - "Q, K, V tensors instead.")]] -void nvte_fused_attn_bwd_kvpacked( - const NVTETensor Q, const NVTETensor KV, const NVTETensor O, const NVTETensor dO, - const NVTETensor S, NVTETensor dP, const NVTETensorPack *Aux_CTX_Tensors, NVTETensor dQ, - NVTETensor dKV, NVTETensor dBias, NVTETensor dSoftmaxOffset, const NVTETensor cu_seqlens_q, - const NVTETensor cu_seqlens_kv, 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); - /*! \brief Compute dot product attention with separate Q, K and V. * * Computes: From 842b770c8c58dad6495e6e0422af9cf909d68e6b Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Wed, 25 Feb 2026 19:38:42 -0600 Subject: [PATCH 228/521] [Common] Remove volatile keyword in fused router kernel utils (#2683) * remove volatile keyword in fused router kernel utils to avoid local mem spill on SM100 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/utils.h | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/fused_router/utils.h b/transformer_engine/common/fused_router/utils.h index 4ae0b467b5..669748c1ad 100644 --- a/transformer_engine/common/fused_router/utils.h +++ b/transformer_engine/common/fused_router/utils.h @@ -47,7 +47,7 @@ __device__ inline T warp_reduce_on_shmem(T *data_ptr, int data_size, ReduceFuncT // Some value is hanlded in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread - volatile double val = lane_id < data_size ? static_cast(data_ptr[lane_id]) : default_val; + double val = lane_id < data_size ? static_cast(data_ptr[lane_id]) : default_val; for (int i = lane_id + kThreadsPerWarp; i < data_size; i += kThreadsPerWarp) { val = reduce_func(val, data_ptr[i]); } @@ -85,7 +85,7 @@ __device__ inline T masked_warp_reduce_on_shmem(T *data_ptr, bool *mask, int dat // Some value is hanlded in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread - volatile double val = + double val = lane_id < data_size && mask[lane_id] ? static_cast(data_ptr[lane_id]) : default_val; for (int i = lane_id + kThreadsPerWarp; i < data_size; i += kThreadsPerWarp) { if (mask[i]) { @@ -183,16 +183,16 @@ __device__ inline void naive_topk_and_mask(T *scores, int data_size, int topk, i // After looping topk times, the topk_indices will be the topk indices for (int k = 0; k < topk; k++) { // Find the max value and its index - volatile double val = (lane_id < data_size && !is_masked(k, lane_id)) - ? static_cast(scores[lane_id]) - : -std::numeric_limits::infinity(); - volatile int index = (lane_id < data_size) ? lane_id : 0; + double val = (lane_id < data_size && !is_masked(k, lane_id)) + ? static_cast(scores[lane_id]) + : -std::numeric_limits::infinity(); + int index = (lane_id < data_size) ? lane_id : 0; // Some value is hanlded in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread for (int i = lane_id + kThreadsPerWarp; i < data_size; i += kThreadsPerWarp) { - volatile double cur_val = (is_masked(k, i)) ? -std::numeric_limits::infinity() - : static_cast(scores[i]); + double cur_val = (is_masked(k, i)) ? -std::numeric_limits::infinity() + : static_cast(scores[i]); if (cur_val > val) { val = cur_val; index = i; @@ -200,8 +200,8 @@ __device__ inline void naive_topk_and_mask(T *scores, int data_size, int topk, i } // Warp shuffle between threads for (int s = 16; s > 0; s /= 2) { - volatile auto shuffled_val = __shfl_xor_sync(0xffffffff, val, s); - volatile auto shuffled_index = __shfl_xor_sync(0xffffffff, index, s); + 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; From ad562838f3b47c0a92cdd174b6e59969d0c98b3c Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Fri, 27 Feb 2026 08:55:24 +0800 Subject: [PATCH 229/521] [CI] Cancel on concurrency (#2708) cancel previous runs when a new one is triggered Signed-off-by: Xin Yao --- .github/workflows/build.yml | 4 ++++ .github/workflows/docs.yml | 4 ++++ .github/workflows/lint.yml | 4 ++++ 3 files changed, 12 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index d80564274e..0f05dbc40a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -7,6 +7,10 @@ name: 'Build' on: pull_request: workflow_dispatch: +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 jobs: core: name: 'Core' diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 6fde0338a1..9d38d709e4 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -8,6 +8,10 @@ on: pull_request: workflow_dispatch: workflow_call: +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 jobs: build_docs: name: 'Build' diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index c5cb748c2c..1d2fb272f8 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -7,6 +7,10 @@ name: 'Lint' on: pull_request: workflow_dispatch: +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 jobs: pytorch_cpplint: name: 'PyTorch C++' From b345941b194542e8778d3b2fa3ead8513912a43a Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 27 Feb 2026 10:46:03 +0530 Subject: [PATCH 230/521] [PyTorch] `GroupedTensor` integration (#2600) * Python GroupedTensor and contiguous weights for GroupedLinear Signed-off-by: Kirthi Shankar Sivamani * Graph safe C API for grouped RHT, needs testing Signed-off-by: Zhongbo Zhu Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Zhongbo Zhu * C++ utils, untested Signed-off-by: Kirthi Shankar Sivamani * Pytorch Binding for GroupedTensor APIs (#13) * changes for pytoch extension; but everything seems to be broken probably unrelated to my changes Signed-off-by: Varun Thumbe * fix the issues Signed-off-by: Varun Thumbe * comment nvte API since Oleg's PR is not merged Signed-off-by: Varun Thumbe * test for all cases: Signed-off-by: Varun Thumbe * tensor attributes should be set later Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix make grouped tensor api Signed-off-by: Kirthi Shankar Sivamani * Fixes to tests Signed-off-by: Kirthi Shankar Sivamani * PyTorch-Python GroupedTensor Signed-off-by: Kirthi Shankar Sivamani * Fix test Signed-off-by: Kirthi Shankar Sivamani * All tests pass Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/pytorch/tensor/storage/grouped_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Remove mxfp8 gq test Signed-off-by: Kirthi Shankar Sivamani * C++ PyTorch GroupedTensor changes WIP Signed-off-by: Kirthi Shankar Sivamani * Compiles Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix runtime failure for test Signed-off-by: Kirthi Shankar Sivamani * Fix IMA in mxfp8 GQ Signed-off-by: Kirthi Shankar Sivamani * Add CG test for grouped_quantize Signed-off-by: Kirthi Shankar Sivamani * Fix recipe tests and FP8 weights Signed-off-by: Kirthi Shankar Sivamani * Fix recipe tests and FP8 weights Signed-off-by: Kirthi Shankar Sivamani * Fix device test Signed-off-by: Kirthi Shankar Sivamani * Disable grouped weights for unsupported recipes Signed-off-by: Kirthi Shankar Sivamani * Integrate NVFP4 Graph Safe Group Quantize (#14) * nvfp4 grouped quantize Signed-off-by: Zhongbo Zhu * fix for paged stashing Signed-off-by: Zhongbo Zhu * pass all edge cases Signed-off-by: Zhongbo Zhu * clean up Signed-off-by: Zhongbo Zhu * fix for other recipes Signed-off-by: Zhongbo Zhu --------- Signed-off-by: Zhongbo Zhu * improve mxfp8 unit test Signed-off-by: Zhongbo Zhu * pre-swizzle nvfp4 mxfp8 for MoE Signed-off-by: Zhongbo Zhu * avoid having nvte_get_grouped_tensor_param_v2 Signed-off-by: Zhongbo Zhu * more tests Signed-off-by: Zhongbo Zhu * fix group quantize mxfp8 kernel Signed-off-by: Zhongbo Zhu * Relaxed restriction for the last dim to be a multiple of 128 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: Kirthi Shankar Sivamani Signed-off-by: Zhongbo Zhu Signed-off-by: vthumbe1503 Signed-off-by: Varun Thumbe Signed-off-by: Oleg Goncharov Co-authored-by: Zhongbo Zhu Co-authored-by: vthumbe1503 Co-authored-by: Oleg Goncharov --- qa/L0_pytorch_unittest/test.sh | 1 + tests/cpp/operator/test_cast_mxfp8_grouped.cu | 55 +- tests/cpp/test_common.cu | 16 +- tests/pytorch/mxfp8/mxfp8_utils.py | 62 +++ .../test_mxfp8_group_quantize_graph_safe.py | 475 ++++++++++++++++++ .../test_mxfp8_quantize_swizzle_fusion.py | 134 +++++ tests/pytorch/nvfp4/nvfp4_utils.py | 159 ++++++ .../nvfp4/test_nvfp4_group_quantize.py | 128 +---- .../test_nvfp4_group_quantize_graph_safe.py | 451 +++++++++++++++++ tests/pytorch/test_grouped_tensor.py | 111 +++- transformer_engine/common/cast/cast.cu | 7 +- .../common/cast/dispatch/quantize.cuh | 8 +- .../cast/mxfp8/group_quantize_mxfp8.cuh | 70 ++- transformer_engine/common/common.h | 15 + .../graph_safe_group_hadamard_transform.cu | 10 +- ...cast_col_hadamard_transform_cast_fusion.cu | 7 +- .../transformer_engine/transformer_engine.h | 264 +++++++++- .../common/transformer_engine.cpp | 220 +++++--- transformer_engine/pytorch/csrc/common.h | 36 ++ transformer_engine/pytorch/csrc/extensions.h | 3 + .../pytorch/csrc/extensions/cast.cpp | 151 ++++++ .../pytorch/csrc/extensions/pybind.cpp | 15 +- transformer_engine/pytorch/csrc/pybind.h | 3 + transformer_engine/pytorch/csrc/quantizer.cpp | 446 ++++++++++++++++ .../pytorch/csrc/type_converters.cpp | 115 +++++ .../pytorch/tensor/storage/grouped_tensor.py | 108 ++-- 26 files changed, 2761 insertions(+), 309 deletions(-) create mode 100644 tests/pytorch/mxfp8/mxfp8_utils.py create mode 100644 tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py create mode 100644 tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py create mode 100644 tests/pytorch/nvfp4/nvfp4_utils.py create mode 100644 tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index cd2d85c91c..e0ad09200d 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -32,6 +32,7 @@ PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED 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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_nvfp4.xml $TE_PATH/tests/pytorch/nvfp4 || test_fail "test_nvfp4" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8.xml $TE_PATH/tests/pytorch/mxfp8 || test_fail "test_mxfp8" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_quantized_tensor.xml $TE_PATH/tests/pytorch/test_quantized_tensor.py || test_fail "test_quantized_tensor.py" 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" diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 8b084ca452..6557c83773 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -385,28 +385,41 @@ void performTest(const ProcessingMethod processing_method, NVTEBasicTensor grad_data_tensor = {grad_data_d, static_cast(itype), logical_shape_}; NVTEBasicTensor in_data_tensor = {in_data_d, static_cast(itype), logical_shape_}; - nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &in_data_tensor); - nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &grad_data_tensor); + nvte_set_grouped_tensor_param(in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, + &in_data_tensor, sizeof(in_data_tensor)); + nvte_set_grouped_tensor_param(grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, + &grad_data_tensor, sizeof(grad_data_tensor)); if ((shape_rep == VARYING_FIRST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { NVTEBasicTensor first_dims_tensor = {first_dims_d, kNVTEInt64, first_dims_shape_}; - nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, &first_dims_tensor); - nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, &first_dims_tensor); - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, &first_dims_tensor); + nvte_set_grouped_tensor_param(grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, + &first_dims_tensor, sizeof(first_dims_tensor)); + nvte_set_grouped_tensor_param(in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, + &first_dims_tensor, sizeof(first_dims_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedFirstDims, + &first_dims_tensor, sizeof(first_dims_tensor)); } if ((shape_rep == VARYING_LAST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { NVTEBasicTensor last_dims_tensor = {last_dims_d, kNVTEInt64, last_dims_shape_}; - nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor); - nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor); - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor); + nvte_set_grouped_tensor_param(grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, + &last_dims_tensor, sizeof(last_dims_tensor)); + nvte_set_grouped_tensor_param(in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, + &last_dims_tensor, sizeof(last_dims_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedLastDims, + &last_dims_tensor, sizeof(last_dims_tensor)); } if (shape_rep != SAME_BOTH_DIMS) { NVTEBasicTensor offsets_tensor = {offsets_d, kNVTEInt64, offsets_shape_}; - nvte_set_grouped_tensor_param(&grad_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, &offsets_tensor); - nvte_set_grouped_tensor_param(&in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, &offsets_tensor); - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, &offsets_tensor); + nvte_set_grouped_tensor_param(grad_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, + &offsets_tensor, sizeof(offsets_tensor)); + nvte_set_grouped_tensor_param(in_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, + &offsets_tensor, sizeof(offsets_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, + &offsets_tensor, sizeof(offsets_tensor)); } if (rowwise) { @@ -417,8 +430,11 @@ void performTest(const ProcessingMethod processing_method, NVTEBasicTensor out_data_rowwise_tensor = {out_data_rowwise_d, static_cast(otype), logical_shape_}; NVTEShape scales_rowwise_shape_ = nvte_make_shape(scales_rowwise_shape.data(), scales_rowwise_shape.size()); NVTEBasicTensor out_scales_rowwise_tensor = {out_scales_rowwise_d, NVTEDType::kNVTEFloat8E8M0, scales_rowwise_shape_}; - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &out_data_rowwise_tensor); - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseScaleInv, &out_scales_rowwise_tensor); + nvte_set_grouped_tensor_param(out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, + &out_data_rowwise_tensor, sizeof(out_data_rowwise_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedRowwiseScaleInv, + &out_scales_rowwise_tensor, sizeof(out_scales_rowwise_tensor)); } if (colwise) { @@ -429,8 +445,12 @@ void performTest(const ProcessingMethod processing_method, NVTEBasicTensor out_data_colwise_tensor = {out_data_colwise_d, static_cast(otype), logical_shape_}; NVTEShape scales_colwise_shape_ = nvte_make_shape(scales_colwise_shape.data(), scales_colwise_shape.size()); NVTEBasicTensor out_scales_colwise_tensor = {out_scales_colwise_d, NVTEDType::kNVTEFloat8E8M0, scales_colwise_shape_}; - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedColumnwiseData, &out_data_colwise_tensor); - nvte_set_grouped_tensor_param(&out_group_tensor, NVTEGroupedTensorParam::kNVTEGroupedColumnwiseScaleInv, &out_scales_colwise_tensor); + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedColumnwiseData, + &out_data_colwise_tensor, sizeof(out_data_colwise_tensor)); + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedColumnwiseScaleInv, + &out_scales_colwise_tensor, sizeof(out_scales_colwise_tensor)); } Tensor output_dbias("output_dbias", std::vector{ cols }, itype); @@ -695,7 +715,10 @@ TEST_P(GroupedFusedCastMXFP8TestSuite, Test) { } offsets[t+1] = offsets[t] + first_dims[t] * last_dims[t]; // Skips tests if tensor shape is not as required by the kernel - if ((first_dims[t] % 128 != 0) || (last_dims[t] % 32 != 0)) { + if (first_dims[t] % 128 != 0) { + GTEST_SKIP(); + } + if (!is_single_tensor && (last_dims[t] % 128 != 0)) { GTEST_SKIP(); } } diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index af99d9c42f..b64ae24131 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -1157,7 +1157,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, NVTEBasicTensor data_tensor{grouped.data.get(), static_cast(dtype), grouped.logical_shape}; NVTEGroupedTensor h = grouped.handle.get(); - nvte_set_grouped_tensor_param(&h, kNVTEGroupedRowwiseData, &data_tensor); + nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseData, &data_tensor, sizeof(data_tensor)); const bool include_columnwise = isFp8Type(dtype) || isFp4Type(dtype); if (include_columnwise) { @@ -1172,7 +1172,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, NVTEBasicTensor col_tensor{grouped.columnwise_data.get(), static_cast(dtype), grouped.logical_shape}; - nvte_set_grouped_tensor_param(&h, kNVTEGroupedColumnwiseData, &col_tensor); + nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseData, &col_tensor, sizeof(col_tensor)); } if (!same_first) { @@ -1181,7 +1181,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); NVTEShape fd_shape = nvte_make_shape(&num_tensors, 1); NVTEBasicTensor fd_tensor{grouped.first_dims_dev.get(), kNVTEInt64, fd_shape}; - nvte_set_grouped_tensor_param(&h, kNVTEGroupedFirstDims, &fd_tensor); + nvte_set_grouped_tensor_param(h, kNVTEGroupedFirstDims, &fd_tensor, sizeof(fd_tensor)); } if (!same_last) { @@ -1190,7 +1190,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); NVTEShape ld_shape = nvte_make_shape(&num_tensors, 1); NVTEBasicTensor ld_tensor{grouped.last_dims_dev.get(), kNVTEInt64, ld_shape}; - nvte_set_grouped_tensor_param(&h, kNVTEGroupedLastDims, &ld_tensor); + nvte_set_grouped_tensor_param(h, kNVTEGroupedLastDims, &ld_tensor, sizeof(ld_tensor)); } if (!same_first || !same_last) { @@ -1199,7 +1199,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); NVTEShape off_shape = nvte_make_shape(&num_tensors, 1); NVTEBasicTensor off_tensor{grouped.offsets_dev.get(), kNVTEInt64, off_shape}; - nvte_set_grouped_tensor_param(&h, kNVTEGroupedTensorOffsets, &off_tensor); + nvte_set_grouped_tensor_param(h, kNVTEGroupedTensorOffsets, &off_tensor, sizeof(off_tensor)); } if (isFp8Type(dtype)) { @@ -1213,8 +1213,10 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, sizeof(float) * num_tensors, cudaMemcpyHostToDevice)); NVTEShape scale_shape = nvte_make_shape(&num_tensors, 1); NVTEBasicTensor scale_tensor{grouped.scale_inv.get(), kNVTEFloat32, scale_shape}; - nvte_set_grouped_tensor_param(&h, kNVTEGroupedRowwiseScaleInv, &scale_tensor); - nvte_set_grouped_tensor_param(&h, kNVTEGroupedColumnwiseScaleInv, &scale_tensor); + nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseScaleInv, &scale_tensor, + sizeof(scale_tensor)); + nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseScaleInv, &scale_tensor, + sizeof(scale_tensor)); } return grouped; diff --git a/tests/pytorch/mxfp8/mxfp8_utils.py b/tests/pytorch/mxfp8/mxfp8_utils.py new file mode 100644 index 0000000000..99e088a201 --- /dev/null +++ b/tests/pytorch/mxfp8/mxfp8_utils.py @@ -0,0 +1,62 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import torch +import math + + +# Calculate the shape of the scaling tensor for MXFP8 1D blockwise quantization without padding +def get_mxfp8_scale_shape_no_padding(shape, columnwise): + M, K = 1, 1 + M = math.prod(shape[:-1]) + K = shape[-1] + + if columnwise: + outer = M // 32 + inner = K + return (outer, inner) + # rowwise + outer = M + inner = K // 32 + return (outer, inner) + + +def _rowwise_swizzle_mxfp8_scale(input_M, input_N, scale: torch.Tensor) -> torch.Tensor: + assert scale.dim() == 2 + assert input_M == scale.shape[0] + assert input_N // 32 == scale.shape[1] + + x = scale.view(input_M // 128, 4, 32, input_N // 128, 4) + x = x.permute(0, 3, 2, 1, 4) + x = x.contiguous() + # View back as original 2D shape + x = x.view(input_M, input_N // 32) + return x + + +def _columnwise_swizzle_mxfp8_scale(input_M, input_N, scale: torch.Tensor) -> torch.Tensor: + assert scale.dim() == 2 + assert input_M // 32 == scale.shape[0] + assert input_N == scale.shape[1] + + x = scale.view(input_M // 128, 4, input_N // 128, 4, 32) + x = x.permute(2, 0, 4, 3, 1) + x = x.contiguous() + + # alternative way: transpose the scale and do rowwise swizzle with M, N swapped + x1 = _rowwise_swizzle_mxfp8_scale(input_N, input_M, scale.transpose(0, 1).contiguous()) + torch.testing.assert_close( + x.view(-1), x1.view(-1), atol=0.0, rtol=0.0, msg="columnwise swizzle sanity check failed" + ) + + # View back as original 2D shape + x = x.view(input_M // 32, input_N) + return x + + +def swizzle_mxfp8_scale(input_M, input_N, scale: torch.Tensor, columnwise: bool) -> torch.Tensor: + if not columnwise: + return _rowwise_swizzle_mxfp8_scale(input_M, input_N, scale) + else: + return _columnwise_swizzle_mxfp8_scale(input_M, input_N, scale) diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py new file mode 100644 index 0000000000..3c197bc6f3 --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -0,0 +1,475 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch import MXFP8Quantizer + +import pytest +import torch +import random +import math + +from mxfp8_utils import swizzle_mxfp8_scale, get_mxfp8_scale_shape_no_padding + +recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) + + +def generate_random_multiples_sum(total=8192, n=4, multiple=64): + if total % multiple != 0: + raise ValueError(f"Total ({total}) must be a multiple of {multiple}") + if (total // multiple) < n: + raise ValueError("Total too small for given n and multiple.") + + # Work in units of multiples + total_units = total // multiple + + # choose n−1 random cut points in [1, total_units−1) + cuts = sorted(random.sample(range(1, total_units), n - 1)) + + # convert to segment lengths + parts = ( + [cuts[0]] + [cuts[i] - cuts[i - 1] for i in range(1, len(cuts))] + [total_units - cuts[-1]] + ) + + # convert back to multiples + return [p * multiple for p in parts] + + +def generate_split_sections(M: int, N: int, edge_cases: str) -> list[int]: + least_multiple = 128 + num_chunks = 4 + split_sections = None + + avg_split = M // num_chunks + + if M == 0 or N == 0: + # all zeros + return [0] * num_chunks + if edge_cases == "regular": + split_sections = [avg_split] * num_chunks + elif edge_cases == "zero_tokens_all": + split_sections = [0] * num_chunks + elif edge_cases == "zero_tokens_front": + split_sections = [0] + [avg_split] * (num_chunks - 2) + [avg_split * 2] + elif edge_cases == "zero_tokens_end": + split_sections = [avg_split * 2] + [avg_split] * (num_chunks - 2) + [0] + elif edge_cases == "zero_tokens_middle": + split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2] + elif edge_cases == "random_uneven_split": + split_sections = generate_random_multiples_sum(M, num_chunks, least_multiple) + else: + raise ValueError(f"Invalid edge case: {edge_cases}") + + # adds up the split_sections to make it M + assert sum(split_sections) == M, "The split_sections do not add up to M" + + # make sure every split_section is a multiple of least_multiple + for split_section in split_sections: + assert ( + split_section % least_multiple == 0 + ), "The split_sections are not multiples of least_multiple" + + return split_sections + + +def reference_group_quantize( + x: torch.Tensor, + quantizers: list[MXFP8Quantizer], + split_sections: list[int], + return_identity: bool, + return_transpose: bool, +) -> torch.Tensor: + x_chunks = torch.split(x, split_sections) + + # rowwise quantization + x_qx = [] + x_sx = [] + # columnwise quantization + x_qx_t = [] + x_sx_t = [] + + for i in range(len(x_chunks)): + x_chunk = x_chunks[i] + x_mxfp8_res = quantizers[i](x_chunk) + if return_identity: + x_qx.append(x_mxfp8_res._rowwise_data.view(dtype=torch.uint8)) + x_sx.append(x_mxfp8_res._rowwise_scale_inv) + else: + x_qx.append(None) + x_sx.append(None) + if return_transpose: + x_qx_t.append(x_mxfp8_res._columnwise_data.view(dtype=torch.uint8)) + x_sx_t.append(x_mxfp8_res._columnwise_scale_inv) + else: + x_qx_t.append(None) + x_sx_t.append(None) + + return x_qx, x_sx, x_qx_t, x_sx_t + + +def fused_grouped_quantize( + x: torch.Tensor, split_section_tensor: torch.Tensor, quantizer: MXFP8Quantizer +): + + # view x as a 2D tensor + hidden_dim = x.shape[-1] + x = x.view(-1, hidden_dim) + num_tensors = split_section_tensor.shape[0] + + grouped_output = tex.group_quantize(x, quantizer, num_tensors, split_section_tensor) + + return grouped_output + + +def assert_same_shape_and_dtype(x: torch.Tensor, y: torch.Tensor) -> None: + assert x.shape == y.shape + assert x.dtype == y.dtype + + +def check_grouped_tensor_mxfp8_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + return_identity: bool, + return_transpose: bool, + split_sections: list[int], + optimize_for_gemm: bool = False, +) -> None: + + te_dtype = tex.DType.kFloat8E4M3 + + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + # Setup device and random seed + device = "cuda" + seed = 0 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + # Input + x = torch.randn((M, N), dtype=x_dtype, device=device) + x_splits = torch.split(x, split_sections) + + # Quantize + quantizers = [ + MXFP8Quantizer( + fp8_dtype=te_dtype, + rowwise=return_identity, + columnwise=return_transpose, + ) + for _ in range(len(split_sections)) + ] + + grouped_quantizer = quantizers[0].copy() + # configure grouped quantizer with swizzle fusion + # and compare with reference without swizzle fusion + grouped_quantizer.optimize_for_gemm = optimize_for_gemm + + x_qx_ref, x_sx_ref, x_qx_t_ref, x_sx_t_ref = reference_group_quantize( + x, quantizers, split_sections, return_identity, return_transpose + ) + + group_quantized_output = fused_grouped_quantize(x, split_section_tensor, grouped_quantizer) + # get a list of MXFP8 quantized tensors for testing + split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() + + if return_identity: + x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] + x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] + + for i in range(len(x_qx)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_qx[i], x_qx_ref[i]) + assert_same_shape_and_dtype(x_sx[i], x_sx_ref[i]) + else: + torch.testing.assert_close(x_qx[i], x_qx_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, False) + assert ( + valid_scale_shape == x_sx[i].shape + ), "The scale shape is not correctly aligned" + x_sx_i = x_sx[i].clone() + x_sx_ref_i = x_sx_ref[i].clone() + if optimize_for_gemm: + x_sx_ref_i = swizzle_mxfp8_scale( + split_sections[i], N, x_sx_ref_i, columnwise=False + ) + torch.testing.assert_close(x_sx_i, x_sx_ref_i, atol=0.0, rtol=0.0) + + if return_transpose: + x_qx_t = [ + output._columnwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs + ] + x_sx_t = [output._columnwise_scale_inv for output in split_quantize_outputs] + # assert with zero tolerance + for i in range(len(x_qx_t)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_qx_t[i], x_qx_t_ref[i]) + assert_same_shape_and_dtype(x_sx_t[i], x_sx_t_ref[i]) + else: + torch.testing.assert_close(x_qx_t[i], x_qx_t_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, True) + assert ( + valid_scale_shape == x_sx_t[i].shape + ), "The scale shape is not correctly aligned" + x_sx_t_i = x_sx_t[i].clone() + x_sx_t_ref_i = x_sx_t_ref[i].clone() + if optimize_for_gemm: + x_sx_t_ref_i = swizzle_mxfp8_scale( + split_sections[i], N, x_sx_t_ref_i, columnwise=True + ) + torch.testing.assert_close(x_sx_t_i, x_sx_t_ref_i, atol=0.0, rtol=0.0) + + +def check_grouped_tensor_mxfp8_with_paged_stashing( + x_dtype: torch.dtype, + M: int, + N: int, + return_identity: bool, + return_transpose: bool, + split_sections: list[int], + valid_M: int = None, + optimize_for_gemm: bool = False, +) -> None: + + te_dtype = tex.DType.kFloat8E4M3 + + assert valid_M is not None, "valid_M must be provided when with_paged_stashing is True" + assert valid_M < M, "valid_M must be less than M when with_paged_stashing is True" + + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + # Setup device and random seed + device = "cuda" + seed = 0 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + # Input (fill the entire tensor with garbage too) + x = torch.randn((M, N), dtype=x_dtype, device=device) + valid_x = x[:valid_M, :].clone() + x_splits = torch.split(valid_x, split_sections) + + # Quantize + quantizers = [ + MXFP8Quantizer( + fp8_dtype=te_dtype, + rowwise=return_identity, + columnwise=return_transpose, + ) + for _ in range(len(split_sections)) + ] + + grouped_quantizer = quantizers[0].copy() + # configure grouped quantizer with swizzle fusion + # and compare with reference without swizzle fusion + grouped_quantizer.optimize_for_gemm = optimize_for_gemm + + x_qx_ref, x_sx_ref, x_qx_t_ref, x_sx_t_ref = reference_group_quantize( + valid_x, quantizers, split_sections, return_identity, return_transpose + ) + + # Note: for grouped quantize with paged stashing + # it's expected that we can just pass in the regular input x, not the valid_x + # the kernel is expected to porcess it correctly by becoming no-op for cuda graph + group_quantized_output = fused_grouped_quantize(x, split_section_tensor, grouped_quantizer) + + # get a list of MXFP8 quantized tensors for testing + split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() + + if return_identity: + x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] + x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] + + for i in range(len(x_qx)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_qx[i], x_qx_ref[i]) + assert_same_shape_and_dtype(x_sx[i], x_sx_ref[i]) + else: + torch.testing.assert_close(x_qx[i], x_qx_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, False) + assert ( + valid_scale_shape == x_sx[i].shape + ), "The scale shape is not correctly aligned" + x_sx_i = x_sx[i].clone() + x_sx_ref_i = x_sx_ref[i].clone() + if optimize_for_gemm: + x_sx_ref_i = swizzle_mxfp8_scale( + split_sections[i], N, x_sx_ref_i, columnwise=False + ) + torch.testing.assert_close(x_sx_i, x_sx_ref_i, atol=0.0, rtol=0.0) + + if return_transpose: + x_qx_t = [ + output._columnwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs + ] + x_sx_t = [output._columnwise_scale_inv for output in split_quantize_outputs] + # assert with zero tolerance + for i in range(len(x_qx_t)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_qx_t[i], x_qx_t_ref[i]) + assert_same_shape_and_dtype(x_sx_t[i], x_sx_t_ref[i]) + else: + torch.testing.assert_close(x_qx_t[i], x_qx_t_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, True) + assert ( + valid_scale_shape == x_sx_t[i].shape + ), "The scale shape is not correctly aligned" + x_sx_t_i = x_sx_t[i].clone() + x_sx_t_ref_i = x_sx_t_ref[i].clone() + if optimize_for_gemm: + x_sx_t_ref_i = swizzle_mxfp8_scale( + split_sections[i], N, x_sx_t_ref_i, columnwise=True + ) + torch.testing.assert_close(x_sx_t_i, x_sx_t_ref_i, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # edge case, zero tokens for all + (0, 512), + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + (16384, 16384), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + ], +) +@pytest.mark.parametrize( + "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] +) +@pytest.mark.parametrize( + "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] +) +def test_grouped_tensor_mxfp8_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, + quantize_mode: str, + optimize_for_gemm: bool, +) -> None: + + split_sections = generate_split_sections(M, N, edge_cases) + + if quantize_mode == "quantize": + return_identity = True + return_transpose = False + elif quantize_mode == "quantize_transpose": + return_identity = True + return_transpose = True + elif quantize_mode == "quantize_colwise_only": + return_identity = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + check_grouped_tensor_mxfp8_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + return_identity=return_identity, + return_transpose=return_transpose, + split_sections=split_sections, + optimize_for_gemm=optimize_for_gemm, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # M won't be empty in paged stashing + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + (16384, 16384), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + # even if buffer is not empty, but the token splits are all zero + "zero_tokens_all", + # partially zero tokens + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + ], +) +@pytest.mark.parametrize( + "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] +) +@pytest.mark.parametrize( + "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] +) +def test_grouped_tensor_mxfp8_with_paged_stashing( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, + quantize_mode: str, + optimize_for_gemm: bool, +) -> None: + + # paged stashing means that the sum of total tokens is less than + # or equal to the buffer size, you can have buffer [2048, 1024] + # and when you only receive 1024 tokens, the last half is garbage + # so input has shape [2048, 1024] + # split sections can be [256, 256, 256, 256], sums to 1024 + valid_M = 0 if edge_cases == "zero_tokens_all" else M // 2 + split_sections = generate_split_sections(valid_M, N, edge_cases) + + # sanity check + if edge_cases == "zero_tokens_all": + assert valid_M == 0, "valid_M must be 0 when edge_cases is zero_tokens_all" + else: + assert valid_M == M // 2, "valid_M must be M // 2 when edge_cases is not zero_tokens_all" + + if quantize_mode == "quantize": + return_identity = True + return_transpose = False + elif quantize_mode == "quantize_transpose": + return_identity = True + return_transpose = True + elif quantize_mode == "quantize_colwise_only": + return_identity = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + check_grouped_tensor_mxfp8_with_paged_stashing( + x_dtype=x_dtype, + M=M, + N=N, + return_identity=return_identity, + return_transpose=return_transpose, + split_sections=split_sections, + valid_M=valid_M, + optimize_for_gemm=optimize_for_gemm, + ) diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py new file mode 100644 index 0000000000..94ea699d14 --- /dev/null +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -0,0 +1,134 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch import MXFP8Quantizer +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage + +import pytest +import torch +import random +import math + +from typing import Tuple + +from mxfp8_utils import swizzle_mxfp8_scale, get_mxfp8_scale_shape_no_padding + +recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True) + + +def unpack_quantized_tensor( + quantized_tensor: MXFP8TensorStorage, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + qx, sx, qx_t, sx_t = None, None, None, None + if quantized_tensor._rowwise_data is not None: + qx = quantized_tensor._rowwise_data.view(dtype=torch.uint8) + if quantized_tensor._rowwise_scale_inv is not None: + sx = quantized_tensor._rowwise_scale_inv + if quantized_tensor._columnwise_data is not None: + qx_t = quantized_tensor._columnwise_data.view(dtype=torch.uint8) + if quantized_tensor._columnwise_scale_inv is not None: + sx_t = quantized_tensor._columnwise_scale_inv + return qx, sx, qx_t, sx_t + + +def check_mxfp8_quantize_swizzle_fusion( + x_dtype: torch.dtype, + M: int, + N: int, + return_identity: bool, + return_transpose: bool, +) -> None: + + te_dtype = tex.DType.kFloat8E4M3 + + # Setup device and random seed + device = "cuda" + seed = 0 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + # Input + x = torch.randn((M, N), dtype=x_dtype, device=device) + + # Quantize + quantizer = MXFP8Quantizer( + fp8_dtype=te_dtype, + rowwise=return_identity, + columnwise=return_transpose, + ) + + quantizer_swizzle_fusion = quantizer.copy() + quantizer_swizzle_fusion.optimize_for_gemm = True + + x_qx_swf, x_sx_swf, x_qx_t_swf, x_sx_t_swf = unpack_quantized_tensor( + quantizer_swizzle_fusion(x) + ) + x_qx_ref, x_sx_ref, x_qx_t_ref, x_sx_t_ref = unpack_quantized_tensor(quantizer(x)) + + if return_identity: + torch.testing.assert_close(x_qx_swf, x_qx_ref, atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x.shape, False) + assert valid_scale_shape == x_sx_swf.shape, ( + "The scale shape is not correctly aligned, this test assumes no padding is needed for" + " scaling factors" + ) + x_sx_ref_swizzled = swizzle_mxfp8_scale(M, N, x_sx_ref, columnwise=False) + torch.testing.assert_close(x_sx_swf, x_sx_ref_swizzled, atol=0.0, rtol=0.0) + + if return_transpose: + torch.testing.assert_close(x_qx_t_swf, x_qx_t_ref, atol=0.0, rtol=0.0) + valid_scale_shape = get_mxfp8_scale_shape_no_padding(x.shape, True) + assert valid_scale_shape == x_sx_t_swf.shape, ( + "The scale shape is not correctly aligned, this test assumes no padding is needed for" + " scaling factors" + ) + x_sx_t_ref_swizzled = swizzle_mxfp8_scale(M, N, x_sx_t_ref, columnwise=True) + torch.testing.assert_close(x_sx_t_swf, x_sx_t_ref_swizzled, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + (16384, 16384), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] +) +def test_mxfp8_quantize_swizzle_fusion( + x_dtype: torch.dtype, + M: int, + N: int, + quantize_mode: str, +) -> None: + + if quantize_mode == "quantize": + return_identity = True + return_transpose = False + elif quantize_mode == "quantize_transpose": + return_identity = True + return_transpose = True + elif quantize_mode == "quantize_colwise_only": + return_identity = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + check_mxfp8_quantize_swizzle_fusion( + x_dtype=x_dtype, + M=M, + N=N, + return_identity=return_identity, + return_transpose=return_transpose, + ) diff --git a/tests/pytorch/nvfp4/nvfp4_utils.py b/tests/pytorch/nvfp4/nvfp4_utils.py new file mode 100644 index 0000000000..5f1b5ac36c --- /dev/null +++ b/tests/pytorch/nvfp4/nvfp4_utils.py @@ -0,0 +1,159 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.pytorch import NVFP4Quantizer + +import torch +import math +import random + + +# Calculate the shape of the scaling tensor for NVFP4 1D blockwise quantization without padding +def get_nvfp4_scale_shape_no_padding(shape, columnwise): + M, K = 1, 1 + M = math.prod(shape[:-1]) + K = shape[-1] + + if columnwise: + outer = K + inner = math.ceil(M / 16) + return (outer, inner) + # rowwise + outer = M + inner = math.ceil(K / 16) + return (outer, inner) + + +def _rowwise_swizzle_nvfp4_scale(input_M, input_N, scale: torch.Tensor) -> torch.Tensor: + assert scale.dim() == 2 + assert input_M == scale.shape[0] + assert input_N // 16 == scale.shape[1] + + x = scale.view(input_M // 128, 4, 32, input_N // 64, 4) + x = x.permute(0, 3, 2, 1, 4) + x = x.contiguous() + # View back as original 2D shape + x = x.view(input_M, input_N // 16) + return x + + +# TN-only layout for NVFP4 means that there is only rowwise swizzle +# just need to switch the M, N which means transposing the input +def swizzle_nvfp4_scale(input_M, input_N, scale: torch.Tensor, columnwise: bool) -> torch.Tensor: + if not columnwise: + return _rowwise_swizzle_nvfp4_scale(input_M, input_N, scale) + else: + return _rowwise_swizzle_nvfp4_scale(input_N, input_M, scale) + + +# Helper function to generate random multiples sum +def _generate_random_multiples_sum(total=8192, n=4, multiple=64): + if total % multiple != 0: + raise ValueError(f"Total ({total}) must be a multiple of {multiple}") + if (total // multiple) < n: + raise ValueError("Total too small for given n and multiple.") + + # Work in units of multiples + total_units = total // multiple + + # choose n−1 random cut points in [1, total_units−1) + cuts = sorted(random.sample(range(1, total_units), n - 1)) + + # convert to segment lengths + parts = ( + [cuts[0]] + [cuts[i] - cuts[i - 1] for i in range(1, len(cuts))] + [total_units - cuts[-1]] + ) + + # convert back to multiples + return [p * multiple for p in parts] + + +# Generate split sections for NVFP4 1D blockwise quantization +def generate_split_sections( + M: int, N: int, edge_cases: str, least_multiple: int = 128 +) -> list[int]: + num_chunks = 4 + split_sections = None + + avg_split = M // num_chunks + + if M == 0 or N == 0: + # all zeros + return [0] * num_chunks + if edge_cases == "regular": + split_sections = [avg_split] * num_chunks + elif edge_cases == "zero_tokens_all": + split_sections = [0] * num_chunks + elif edge_cases == "zero_tokens_front": + split_sections = [0] + [avg_split] * (num_chunks - 2) + [avg_split * 2] + elif edge_cases == "zero_tokens_end": + split_sections = [avg_split * 2] + [avg_split] * (num_chunks - 2) + [0] + elif edge_cases == "zero_tokens_middle": + split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2] + elif edge_cases == "random_uneven_split": + split_sections = _generate_random_multiples_sum(M, num_chunks, least_multiple) + else: + raise ValueError(f"Invalid edge case: {edge_cases}") + + # adds up the split_sections to make it M + assert sum(split_sections) == M, "The split_sections do not add up to M" + + # make sure every split_section is a multiple of least_multiple + for split_section in split_sections: + assert ( + split_section % least_multiple == 0 + ), "The split_sections are not multiples of least_multiple" + + return split_sections + + +# Reference implementation of group quantization for NVFP4 1D blockwise quantization +def reference_group_quantize( + x: torch.Tensor, + quantizers: list[NVFP4Quantizer], + split_sections: list[int], + return_identity: bool, + return_transpose: bool, +) -> torch.Tensor: + x_view = x.reshape(-1, x.size(-1)) + x_chunks = torch.split(x, split_sections) + + # rowwise quantization + x_qx = [] + x_sx = [] + x_amax_rowwise = [] + # columnwise quantization + x_qx_t = [] + x_sx_t = [] + x_amax_colwise = [] + + for i in range(len(x_chunks)): + x_chunk = x_chunks[i] + x_nvfp4_res = quantizers[i](x_chunk) + if return_identity: + x_qx.append(x_nvfp4_res._rowwise_data.view(dtype=torch.uint8)) + x_sx.append(x_nvfp4_res._rowwise_scale_inv) + x_amax_rowwise.append(x_nvfp4_res._amax_rowwise) + else: + x_qx.append(None) + x_sx.append(None) + x_amax_rowwise.append(None) + if return_transpose: + x_qx_t.append(x_nvfp4_res._columnwise_data.view(dtype=torch.uint8)) + x_sx_t.append(x_nvfp4_res._columnwise_scale_inv) + x_amax_colwise.append(x_nvfp4_res._amax_columnwise) + else: + x_qx_t.append(None) + x_sx_t.append(None) + x_amax_colwise.append(None) + + return x_qx, x_sx, x_amax_rowwise, x_qx_t, x_sx_t, x_amax_colwise + + +# Function to assert that two tensors have the same shape and dtype +def assert_same_shape_and_dtype(x: torch.Tensor, y: torch.Tensor) -> None: + assert x.shape == y.shape + assert x.dtype == y.dtype diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 01a4a01205..5f35e9ad10 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -23,126 +23,14 @@ import random import math -recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) - - -def generate_random_multiples_sum(total=8192, n=4, multiple=64): - if total % multiple != 0: - raise ValueError(f"Total ({total}) must be a multiple of {multiple}") - if (total // multiple) < n: - raise ValueError("Total too small for given n and multiple.") - - # Work in units of multiples - total_units = total // multiple - - # choose n−1 random cut points in [1, total_units−1) - cuts = sorted(random.sample(range(1, total_units), n - 1)) - - # convert to segment lengths - parts = ( - [cuts[0]] + [cuts[i] - cuts[i - 1] for i in range(1, len(cuts))] + [total_units - cuts[-1]] - ) - - # convert back to multiples - return [p * multiple for p in parts] - - -def generate_split_sections(M: int, N: int, edge_cases: str) -> list[int]: - least_multiple = 64 - num_chunks = 4 - split_sections = None - - avg_split = M // num_chunks - - if M == 0 or N == 0: - # all zeros - return [0] * num_chunks - if edge_cases == "regular": - split_sections = [avg_split] * num_chunks - elif edge_cases == "zero_tokens_front": - split_sections = [0] + [avg_split] * (num_chunks - 2) + [avg_split * 2] - elif edge_cases == "zero_tokens_end": - split_sections = [avg_split * 2] + [avg_split] * (num_chunks - 2) + [0] - elif edge_cases == "zero_tokens_middle": - split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2] - elif edge_cases == "random_uneven_split": - split_sections = generate_random_multiples_sum(M, num_chunks, least_multiple) - else: - raise ValueError(f"Invalid edge case: {edge_cases}") - - # adds up the split_sections to make it M - assert sum(split_sections) == M, "The split_sections do not add up to M" - - # make sure every split_section is a multiple of least_multiple - for split_section in split_sections: - assert ( - split_section % least_multiple == 0 - ), "The split_sections are not multiples of least_multiple" - - return split_sections - - -# Calculate the shape of the scaling tensor for NVFP4 1D blockwise quantization without padding -def get_nvfp4_scale_shape_no_padding(shape, columnwise): - M, K = 1, 1 - M = math.prod(shape[:-1]) - K = shape[-1] - - if columnwise: - outer = K - inner = math.ceil(M / 16) - return (outer, inner) - # rowwise - outer = M - inner = math.ceil(K / 16) - return (outer, inner) - - -def reference_group_quantize( - x: torch.Tensor, - quantizers: list[NVFP4Quantizer], - split_sections: list[int], - return_identity: bool, - return_transpose: bool, -) -> torch.Tensor: - x_view = x.reshape(-1, x.size(-1)) - x_chunks = torch.split(x, split_sections) - - # rowwise quantization - x_qx = [] - x_sx = [] - x_amax_rowwise = [] - # columnwise quantization - x_qx_t = [] - x_sx_t = [] - x_amax_colwise = [] - - for i in range(len(x_chunks)): - x_chunk = x_chunks[i] - x_nvfp4_res = quantizers[i](x_chunk) - if return_identity: - x_qx.append(x_nvfp4_res._rowwise_data.view(dtype=torch.uint8)) - x_sx.append(x_nvfp4_res._rowwise_scale_inv) - x_amax_rowwise.append(x_nvfp4_res._amax_rowwise) - else: - x_qx.append(None) - x_sx.append(None) - x_amax_rowwise.append(None) - if return_transpose: - x_qx_t.append(x_nvfp4_res._columnwise_data.view(dtype=torch.uint8)) - x_sx_t.append(x_nvfp4_res._columnwise_scale_inv) - x_amax_colwise.append(x_nvfp4_res._amax_columnwise) - else: - x_qx_t.append(None) - x_sx_t.append(None) - x_amax_colwise.append(None) - - return x_qx, x_sx, x_amax_rowwise, x_qx_t, x_sx_t, x_amax_colwise - +from nvfp4_utils import ( + get_nvfp4_scale_shape_no_padding, + generate_split_sections, + assert_same_shape_and_dtype, + reference_group_quantize, +) -def assert_same_shape_and_dtype(x: torch.Tensor, y: torch.Tensor) -> None: - assert x.shape == y.shape - assert x.dtype == y.dtype +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) def check_group_quantization_nvfp4_versus_reference( @@ -279,7 +167,7 @@ def test_rht_with_quantization_block_tiling_versus_reference( with_rht: bool, ) -> None: - split_sections = generate_split_sections(M, N, edge_cases) + split_sections = generate_split_sections(M, N, edge_cases, least_multiple=64) # currently disable pre-RHT amax with_post_rht_amax = with_rht diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py new file mode 100644 index 0000000000..1e62f91eb8 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -0,0 +1,451 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + + +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 import utils +from transformer_engine.pytorch.constants import TE_DType +from transformer_engine.common.recipe import NVFP4BlockScaling +from transformer_engine.pytorch.tensor.storage.grouped_tensor import GroupedTensor + +import pytest +import torch +import random +import math + +from nvfp4_utils import ( + get_nvfp4_scale_shape_no_padding, + swizzle_nvfp4_scale, + generate_split_sections, + assert_same_shape_and_dtype, + reference_group_quantize, +) + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def fused_grouped_quantize( + x: torch.Tensor, split_section_tensor: torch.Tensor, quantizer: NVFP4Quantizer +): + + # view x as a 2D tensor + hidden_dim = x.shape[-1] + x = x.view(-1, hidden_dim) + num_tensors = split_section_tensor.shape[0] + + grouped_output = tex.group_quantize(x, quantizer, num_tensors, split_section_tensor) + + return grouped_output + + +def check_grouped_tensor_nvfp4_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + return_identity: bool, + return_transpose: bool, + split_sections: list[int], + with_rht: bool = True, + with_post_rht_amax: bool = True, + with_random_sign_mask: bool = True, + optimize_for_gemm: bool = False, +) -> None: + + te_dtype = tex.DType.kFloat4E2M1 + + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + # Setup device and random seed + device = "cuda" + seed = 0 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + # Input + x = torch.randn((M, N), dtype=x_dtype, device=device) + num_chunks = len(split_sections) + + x_splits = torch.split(x, split_sections) + + # Quantize + quantizers = [ + NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=return_identity, + columnwise=return_transpose, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=with_rht, + with_post_rht_amax=with_post_rht_amax, + with_random_sign_mask=with_random_sign_mask, + ) + for _ in range(len(split_sections)) + ] + + grouped_quantizer = quantizers[0].copy() + # configure grouped quantizer with swizzle fusion + # and compare with reference without swizzle fusion + grouped_quantizer.optimize_for_gemm = optimize_for_gemm + + x_qx_ref, x_sx_ref, x_amax_rowwise_ref, x_qx_t_ref, x_sx_t_ref, x_amax_colwise_ref = ( + reference_group_quantize(x, quantizers, split_sections, return_identity, return_transpose) + ) + + group_quantized_output = fused_grouped_quantize(x, split_section_tensor, grouped_quantizer) + # get a list of nvfp4 quantized tensors for testing + split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() + + if return_identity: + x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] + x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] + x_amax_rowwise = [output._amax_rowwise for output in split_quantize_outputs] + + for i in range(len(x_qx)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_amax_rowwise[i], x_amax_rowwise_ref[i]) + assert_same_shape_and_dtype(x_qx[i], x_qx_ref[i]) + assert_same_shape_and_dtype(x_sx[i], x_sx_ref[i]) + else: + torch.testing.assert_close( + x_amax_rowwise[i], x_amax_rowwise_ref[i], atol=0.0, rtol=0.0 + ) + torch.testing.assert_close(x_qx[i], x_qx_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, False) + assert ( + valid_scale_shape == x_sx[i].shape + ), "The scale shape is not correctly aligned" + x_sx_i = x_sx[i].clone() + x_sx_ref_i = x_sx_ref[i].clone() + if optimize_for_gemm: + x_sx_ref_i = swizzle_nvfp4_scale( + split_sections[i], N, x_sx_ref_i, columnwise=False + ) + torch.testing.assert_close(x_sx_i, x_sx_ref_i, atol=0.0, rtol=0.0) + + if return_transpose: + x_qx_t = [ + output._columnwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs + ] + x_sx_t = [output._columnwise_scale_inv for output in split_quantize_outputs] + x_amax_colwise = [output._amax_columnwise for output in split_quantize_outputs] + # assert with zero tolerance + for i in range(len(x_qx_t)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_amax_colwise[i], x_amax_colwise_ref[i]) + assert_same_shape_and_dtype(x_qx_t[i], x_qx_t_ref[i]) + assert_same_shape_and_dtype(x_sx_t[i], x_sx_t_ref[i]) + else: + torch.testing.assert_close( + x_amax_colwise[i], x_amax_colwise_ref[i], atol=0.0, rtol=0.0 + ) + torch.testing.assert_close(x_qx_t[i], x_qx_t_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, True) + assert ( + valid_scale_shape == x_sx_t[i].shape + ), "The scale shape is not correctly aligned" + x_sx_t_i = x_sx_t[i].clone() + x_sx_t_ref_i = x_sx_t_ref[i].clone() + if optimize_for_gemm: + x_sx_t_ref_i = swizzle_nvfp4_scale( + split_sections[i], N, x_sx_t_ref_i, columnwise=True + ) + torch.testing.assert_close(x_sx_t_i, x_sx_t_ref_i, atol=0.0, rtol=0.0) + + +def check_grouped_tensor_nvfp4_with_paged_stashing( + x_dtype: torch.dtype, + M: int, + N: int, + return_identity: bool, + return_transpose: bool, + split_sections: list[int], + with_rht: bool = True, + with_post_rht_amax: bool = True, + with_random_sign_mask: bool = True, + valid_M: int = None, + optimize_for_gemm: bool = False, +) -> None: + + te_dtype = tex.DType.kFloat4E2M1 + + assert valid_M is not None, "valid_M must be provided when with_paged_stashing is True" + assert valid_M < M, "valid_M must be less than M when with_paged_stashing is True" + + split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device="cuda") + + # Setup device and random seed + device = "cuda" + seed = 0 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + # Input (fill the entire tensor with garbage too) + x = torch.randn((M, N), dtype=x_dtype, device=device) + valid_x = x[:valid_M, :].clone() + num_chunks = len(split_sections) + + x_splits = torch.split(valid_x, split_sections) + + # Quantize + quantizers = [ + NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=return_identity, + columnwise=return_transpose, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=with_rht, + with_post_rht_amax=with_post_rht_amax, + with_random_sign_mask=with_random_sign_mask, + ) + for _ in range(len(split_sections)) + ] + + grouped_quantizer = quantizers[0].copy() + # configure grouped quantizer with swizzle fusion + # and compare with reference without swizzle fusion + grouped_quantizer.optimize_for_gemm = optimize_for_gemm + + x_qx_ref, x_sx_ref, x_amax_rowwise_ref, x_qx_t_ref, x_sx_t_ref, x_amax_colwise_ref = ( + reference_group_quantize( + valid_x, quantizers, split_sections, return_identity, return_transpose + ) + ) + + # Note: for grouped quantize with paged stashing + # it's expected that we can just pass in the regular input x, not the valid_x + # the kernel is expected to porcess it correctly by becoming no-op for cuda graph + group_quantized_output = fused_grouped_quantize(x, split_section_tensor, grouped_quantizer) + + # get a list of nvfp4 quantized tensors for testing + split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() + + if return_identity: + x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] + x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] + x_amax_rowwise = [output._amax_rowwise for output in split_quantize_outputs] + + for i in range(len(x_qx)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_amax_rowwise[i], x_amax_rowwise_ref[i]) + assert_same_shape_and_dtype(x_qx[i], x_qx_ref[i]) + assert_same_shape_and_dtype(x_sx[i], x_sx_ref[i]) + else: + torch.testing.assert_close( + x_amax_rowwise[i], x_amax_rowwise_ref[i], atol=0.0, rtol=0.0 + ) + torch.testing.assert_close(x_qx[i], x_qx_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, False) + assert ( + valid_scale_shape == x_sx[i].shape + ), "The scale shape is not correctly aligned" + x_sx_i = x_sx[i].clone() + x_sx_ref_i = x_sx_ref[i].clone() + if optimize_for_gemm: + x_sx_ref_i = swizzle_nvfp4_scale( + split_sections[i], N, x_sx_ref_i, columnwise=False + ) + torch.testing.assert_close(x_sx_i, x_sx_ref_i, atol=0.0, rtol=0.0) + + if return_transpose: + x_qx_t = [ + output._columnwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs + ] + x_sx_t = [output._columnwise_scale_inv for output in split_quantize_outputs] + x_amax_colwise = [output._amax_columnwise for output in split_quantize_outputs] + # assert with zero tolerance + for i in range(len(x_qx_t)): + if split_sections[i] == 0: + # then just assert the same shape and dtype because the buffer won't be zero out + assert_same_shape_and_dtype(x_amax_colwise[i], x_amax_colwise_ref[i]) + assert_same_shape_and_dtype(x_qx_t[i], x_qx_t_ref[i]) + assert_same_shape_and_dtype(x_sx_t[i], x_sx_t_ref[i]) + else: + torch.testing.assert_close( + x_amax_colwise[i], x_amax_colwise_ref[i], atol=0.0, rtol=0.0 + ) + torch.testing.assert_close(x_qx_t[i], x_qx_t_ref[i], atol=0.0, rtol=0.0) + valid_scale_shape = get_nvfp4_scale_shape_no_padding(x_splits[i].shape, True) + x_sx_t_i = x_sx_t[i].clone() + x_sx_t_ref_i = x_sx_t_ref[i].clone() + if optimize_for_gemm: + x_sx_t_ref_i = swizzle_nvfp4_scale( + split_sections[i], N, x_sx_t_ref_i, columnwise=True + ) + torch.testing.assert_close(x_sx_t_i, x_sx_t_ref_i, atol=0.0, rtol=0.0) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # edge case, zero tokens for all + (0, 512), + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + (16384, 16384), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + ], +) +@pytest.mark.parametrize( + "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] +) +@pytest.mark.parametrize( + "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] +) +@pytest.mark.parametrize("with_rht", [True], ids=["with_rht"]) +@pytest.mark.parametrize( + "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] +) +def test_grouped_tensor_nvfp4_versus_reference( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, + quantize_mode: str, + with_random_sign_mask: bool, + with_rht: bool, + optimize_for_gemm: bool, +) -> None: + + split_sections = generate_split_sections(M, N, edge_cases, least_multiple=128) + + # currently disable pre-RHT amax + with_post_rht_amax = with_rht + + if quantize_mode == "quantize": + return_identity = True + return_transpose = False + elif quantize_mode == "quantize_transpose": + return_identity = True + return_transpose = True + elif quantize_mode == "quantize_colwise_only": + return_identity = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + check_grouped_tensor_nvfp4_versus_reference( + x_dtype=x_dtype, + M=M, + N=N, + return_identity=return_identity, + return_transpose=return_transpose, + split_sections=split_sections, + with_rht=with_rht, + with_post_rht_amax=with_post_rht_amax, + with_random_sign_mask=with_random_sign_mask, + optimize_for_gemm=optimize_for_gemm, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, N", + [ + # M won't be empty in paged stashing + # full tile cases + (1024, 256), + # larger sizes + (8192, 1024), + (16384, 8192), + (16384, 16384), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) +@pytest.mark.parametrize( + "edge_cases", + [ + "regular", + # even if buffer is not empty, but the token splits are all zero + "zero_tokens_all", + # partially zero tokens + "zero_tokens_front", + "zero_tokens_end", + "zero_tokens_middle", + "random_uneven_split", + ], +) +@pytest.mark.parametrize( + "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] +) +@pytest.mark.parametrize( + "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] +) +@pytest.mark.parametrize("with_rht", [True], ids=["with_rht"]) +@pytest.mark.parametrize( + "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] +) +def test_grouped_tensor_nvfp4_with_paged_stashing( + x_dtype: torch.dtype, + M: int, + N: int, + edge_cases: str, + quantize_mode: str, + with_random_sign_mask: bool, + with_rht: bool, + optimize_for_gemm: bool, +) -> None: + + # paged stashing means that the sum of total tokens is less than + # or equal to the buffer size, you can have buffer [2048, 1024] + # and when you only receive 1024 tokens, the last half is garbage + # so input has shape [2048, 1024] + # split sections can be [256, 256, 256, 256], sums to 1024 + valid_M = 0 if edge_cases == "zero_tokens_all" else M // 2 + split_sections = generate_split_sections(valid_M, N, edge_cases, least_multiple=128) + + # sanity check + if edge_cases == "zero_tokens_all": + assert valid_M == 0, "valid_M must be 0 when edge_cases is zero_tokens_all" + else: + assert valid_M == M // 2, "valid_M must be M // 2 when edge_cases is not zero_tokens_all" + + # currently disable pre-RHT amax + with_post_rht_amax = with_rht + + if quantize_mode == "quantize": + return_identity = True + return_transpose = False + elif quantize_mode == "quantize_transpose": + return_identity = True + return_transpose = True + elif quantize_mode == "quantize_colwise_only": + return_identity = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + + check_grouped_tensor_nvfp4_with_paged_stashing( + x_dtype=x_dtype, + M=M, + N=N, + return_identity=return_identity, + return_transpose=return_transpose, + split_sections=split_sections, + with_rht=with_rht, + with_post_rht_amax=with_post_rht_amax, + with_random_sign_mask=with_random_sign_mask, + valid_M=valid_M, + optimize_for_gemm=optimize_for_gemm, + ) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 318009c669..ad08c0474d 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -55,7 +55,7 @@ def make_quantizer(quantization: str, num_tensors: int, shape: List[Tuple[int, int]]) -> Quantizer: - """Create quantizers for given quantization scheme""" + """Create quantizer for given quantization scheme""" if quantization == "fp8_delayed_scaling": quantizer = Float8Quantizer( @@ -203,12 +203,12 @@ def test_split_into_quantized_tensors_quantized(self, quantization: str) -> None """Test split_into_quantized_tensors for quantized tensors""" num_tensors = 3 shape = [(512, 512) for _ in range(num_tensors)] - quantizers = make_quantizer(quantization, num_tensors, shape) + quantizer = make_quantizer(quantization, num_tensors, shape) grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, shape=shape, - quantizer=quantizers, + quantizer=quantizer, device="cuda", ) @@ -260,12 +260,12 @@ def test_quantize_inplace(self, quantization: str) -> None: """Test that quantize is done in-place for all recipes""" num_tensors = 3 shape = [(512, 512) for _ in range(num_tensors)] - quantizers = make_quantizer(quantization, num_tensors, shape) + quantizer = make_quantizer(quantization, num_tensors, shape) grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, shape=shape, - quantizer=quantizers, + quantizer=quantizer, device="cuda", ) @@ -300,12 +300,12 @@ def test_quantize_varying_shapes(self, quantization: str) -> None: """Test quantize with varying shapes""" num_tensors = 3 shape = [(256, 512), (512, 512), (768, 512)] - quantizers = make_quantizer(quantization, num_tensors, shape) + quantizer = make_quantizer(quantization, num_tensors, shape) grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, shape=shape, - quantizer=quantizers, + quantizer=quantizer, device="cuda", ) @@ -334,7 +334,7 @@ def test_static_quantize_method(self, quantization: str) -> None: """Test the static quantize method""" num_tensors = 3 shape = [(512, 512) for _ in range(num_tensors)] - quantizers = make_quantizer(quantization, num_tensors, shape) + quantizer = make_quantizer(quantization, num_tensors, shape) # Create input tensors input_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in shape] @@ -342,7 +342,7 @@ def test_static_quantize_method(self, quantization: str) -> None: # Use static quantize method grouped_tensor = GroupedTensor.create_and_quantize( tensors=input_tensors, - quantizer=quantizers, + quantizer=quantizer, device="cuda", ) @@ -361,6 +361,99 @@ def test_static_quantize_method(self, quantization: str) -> None: expected_offset = _rowwise_offset_bytes(i * numel, quantization) assert rowwise_data.data_ptr() == original_data_ptr + expected_offset + @pytest.mark.parametrize( + "shape", + [[(256, 512), (512, 512), (768, 512)], [(512, 512), (512, 512), (512, 512)]], + ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]]) -> None: + """Test grouped quantization for MXFP8 against per-tensor quantization.""" + # Test wont pass until the grouped quantization PR from Oleg is merged. + num_tensors = 2 + shape = [(512, 1024) for _ in range(num_tensors)] + + # Create BF16 input tensors and pack into a 2D tensor + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + quantized_tensors = [ + MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(tensor) for tensor in input_tensors + ] + grouped_input = torch.cat(input_tensors, dim=0) + + # Create MXFP8 output grouped tensor (rowwise only for easier validation) + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=False) + first_dims = torch.tensor( + [shape[0][0] for _ in range(num_tensors)], + dtype=torch.int64, + device="cuda", + ) + + # Quantize using grouped API + grouped_output = tex.group_quantize( + grouped_input, + quantizer, + num_tensors, + first_dims, + ) + # Build expected output by quantizing each tensor independently + expected_data = [] + expected_scale_inv = [] + for tensor in input_tensors: + qtensor = quantizer(tensor) + expected_data.append(qtensor._rowwise_data.reshape(-1)) + expected_scale_inv.append(qtensor._rowwise_scale_inv.reshape(-1)) + + expected_data = torch.cat(expected_data) + expected_scale_inv = torch.cat(expected_scale_inv) + + assert torch.equal(grouped_output.data, expected_data) + assert torch.equal(grouped_output.scale_inv, expected_scale_inv) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_quantize_cudagraph_capturable(self) -> None: + """Ensure group_quantize is CUDA graph capturable.""" + num_tensors = 2 + shape = [(512, 1024) for _ in range(num_tensors)] + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=False) + first_dims = torch.tensor( + [shape[0][0] for _ in range(num_tensors)], + dtype=torch.int64, + device="cuda", + ) + + torch.cuda.synchronize() + static_input = grouped_input.clone() + static_first_dims = first_dims.clone() + + # Warmup to initialize kernels and allocator state + _ = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = tex.group_quantize( + static_input, + quantizer, + num_tensors, + static_first_dims, + ) + + fresh_input = torch.cat( + [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape], + dim=0, + ) + static_input.copy_(fresh_input) + graph.replay() + torch.cuda.synchronize() + + expected = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) + assert torch.equal(static_output.data, expected.data) + assert torch.equal(static_output.scale_inv, expected.scale_inv) + def test_clear(self) -> None: """Test clear method""" num_tensors = 3 diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 582172a88e..57404ae8a5 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -124,7 +124,8 @@ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, } // Group quantize assumes contiguous inputs and outputs in memory allocation -// TODO (zhongbo): find a better way to make it a more generalized API +// 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, @@ -134,6 +135,6 @@ void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *out constexpr bool IS_ACT = false; - dispatch::group_quantize_fwd_helper(input, outputs, split_sections, - num_tensors, quant_config, stream); + dispatch::group_quantize_fwd_host_aware_helper( + input, outputs, split_sections, num_tensors, quant_config, stream); } diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index b83df1dedf..98a3fb8cba 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -308,10 +308,12 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens } } +// Host-aware and not graph-safe: group quantization with split section info from the host. template -void group_quantize_fwd_helper(const NVTETensor input, NVTETensor *outputs, - const size_t *split_sections, const size_t num_tensors, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { +void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *outputs, + const size_t *split_sections, const size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { using namespace detail; const Tensor *input_tensor = convertNVTETensorCheck(input); diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index a29a09836e..6447fc4542 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -73,10 +73,10 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 __device__ __forceinline__ size_t get_current_tensor_id( const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, - const size_t first_logical_dim, const size_t last_logical_dim, + const size_t block_Y, const size_t first_logical_dim, const size_t last_logical_dim, const int64_t *const __restrict__ offsets_ptr) { if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { - const size_t current_row = current_offset / last_logical_dim; + const size_t current_row = block_Y * CHUNK_DIM_Y; const size_t rows_per_tensor = first_logical_dim / num_tensors; return current_row / rows_per_tensor; } else { @@ -261,10 +261,16 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; - const size_t block_global_offset = blockIdx.x * ELTS_PER_CHUNK; + const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); + + const size_t block_ID = blockIdx.y * gridDim.x + blockIdx.x; + const size_t block_global_offset = + is_single_tensor ? (blockIdx.y * CHUNK_DIM_Y * last_logical_dim + blockIdx.x * CHUNK_DIM_X) + : (block_ID * ELTS_PER_CHUNK); - const size_t tensor_id = get_current_tensor_id(shape_rep, num_tensors, block_global_offset, - first_logical_dim, last_logical_dim, offsets_ptr); + const size_t tensor_id = + get_current_tensor_id(shape_rep, num_tensors, block_global_offset, blockIdx.y, + first_logical_dim, last_logical_dim, offsets_ptr); const size_t rows = get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); @@ -273,10 +279,32 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t scale_stride_rowwise = DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(32)), 4); const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, 128); - const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); - // grouped tensor can be treated as continuous tensor for MXFP8 const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[tensor_id]); + // For grouped tensors represented as a single logical tensor, scale swizzle must still be + // computed per tensor (expert) and then concatenated along dim-0. + const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1) + ? static_cast(offsets_ptr[tensor_id]) + : tensor_base; + + // In graph-safe paged stashing, the logical shape can include trailing garbage. Skip CTAs that + // map outside the current tensor's valid [rows, cols] region. + if (rows == 0 || cols == 0) { + return; + } + if (shape_rep != SAME_BOTH_DIMS) { + const size_t tensor_start_offset = static_cast(offsets_ptr[tensor_id]); + const size_t tensor_end_offset = static_cast(offsets_ptr[tensor_id + 1]); + if (block_global_offset >= tensor_end_offset) { + return; + } + const size_t tensor_offset_from_start = block_global_offset - tensor_start_offset; + const size_t block_offset_Y_in_tensor = tensor_offset_from_start / cols; + const size_t block_offset_X_in_tensor = tensor_offset_from_start % cols; + if (block_offset_Y_in_tensor >= rows || block_offset_X_in_tensor >= cols) { + return; + } + } const CUtensorMap &tensor_map_input = is_single_tensor ? tensor_map_input_static : g_tensor_maps_input[tensor_id]; @@ -304,7 +332,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t blocks_X_num_in_current_tensor = DIVUP(cols, static_cast(128)); const size_t block_id_in_current_tensor = - is_single_tensor ? blockIdx.x : (blockIdx.x - tensor_base / ELTS_PER_CHUNK); + is_single_tensor ? block_ID : (block_ID - tensor_base / ELTS_PER_CHUNK); const size_t block_id_Y = block_id_in_current_tensor / blocks_X_num_in_current_tensor; const size_t block_id_X = block_id_in_current_tensor % blocks_X_num_in_current_tensor; @@ -481,7 +509,12 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel size_t scale_idx = 0; if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - scale_idx = gemm_swizzled_scale_idx(global_scales_offset_X, global_scales_offset_Y, + 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 local_scales_offset_Y = global_scales_offset_Y - tensor_scales_offset_Y_base; + scale_idx = tensor_scales_offset_colwise_base + + gemm_swizzled_scale_idx(global_scales_offset_X, local_scales_offset_Y, DIVUP(rows, static_cast(128))); } else { scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; @@ -801,12 +834,12 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations const size_t num_tensors = input->num_tensors; - size_t blocks = 0; + size_t blocks_X = 0; + size_t blocks_Y = 0; if (is_single_tensor) { - const size_t blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); - const size_t blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); - blocks = blocks_Y * blocks_X; + blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); + blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); } else { NVTE_CHECK(num_tensors < MAX_SUPPORTED_TENSOR_DESCRIPTORS, "Number of tensors in a group is larger than " @@ -814,9 +847,10 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations // Only full tiles supported NVTE_CHECK(last_logical_dim % CHUNK_DIM_X == 0, "Last dimension of a grouped tensor should be divisible by 128."); - blocks = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); + blocks_Y = 1; + blocks_X = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); } - const dim3 grid(blocks); + const dim3 grid(blocks_X, blocks_Y); const size_t block_size = THREADS_PER_CHUNK; const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; @@ -827,9 +861,9 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations "First dimension of a grouped tensor should be divisible by 128."); } - const int64_t *const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); - const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); - const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(output->last_dims.dptr); float *const workspace_ptr = IS_DBIAS ? reinterpret_cast(workspace->data.dptr) : nullptr; float *const amax_ptr = reinterpret_cast(output->amax.dptr); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 0c722634f3..1749b5734a 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -341,6 +341,21 @@ struct GroupedTensor { */ bool with_gemm_swizzled_scales = false; + /*! Map from NVTEGroupedTensorParam to parameter sizes */ + static constexpr size_t attr_sizes[] = { + sizeof(NVTEBasicTensor), // kNVTEGroupedRowwiseData + sizeof(NVTEBasicTensor), // kNVTEGroupedColumnwiseData + sizeof(NVTEBasicTensor), // kNVTEGroupedScale + sizeof(NVTEBasicTensor), // kNVTEGroupedAmax + sizeof(NVTEBasicTensor), // kNVTEGroupedRowwiseScaleInv + sizeof(NVTEBasicTensor), // kNVTEGroupedColumnwiseScaleInv + sizeof(NVTEBasicTensor), // kNVTEGroupedColumnwiseAmax + sizeof(NVTEBasicTensor), // kNVTEGroupedFirstDims + sizeof(NVTEBasicTensor), // kNVTEGroupedLastDims + sizeof(NVTEBasicTensor), // kNVTEGroupedTensorOffsets + sizeof(uint8_t) // kNVTEGroupedWithGEMMSwizzledScales + }; + GroupedTensor(NVTEScalingMode scaling_mode, size_t num_tensors) : data(), columnwise_data(), diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 58b0640249..04e965a9da 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -251,6 +251,11 @@ __global__ void GraphSafeGroupHadamardAmaxTmaKernel( // calculate the global offset to get tensor id size_t global_offset = blockIdx.y * CHUNK_DIM_Y * last_logical_dim; + // paged stashing: will have input buffer [M, N], where M is larger than sum(first_dims) + // also need to early return if this CTA is processing a region larger than the last offsets[num_tensors] + if (global_offset >= offsets_ptr[num_tensors]) { + return; + } int tensor_id = get_current_tensor_id(shape_rep, num_tensors, global_offset, first_logical_dim, last_logical_dim, offsets_ptr); output_pre_rht_amax_ptr = static_cast(amax_rowwise_ptr) + tensor_id; @@ -441,9 +446,8 @@ void group_hadamard_transform_amax_graph_safe(const GroupedTensor* input, Groupe float* const amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); float* const amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr); - const int64_t* const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); - const int64_t* const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); - // const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + const int64_t* const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t* const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); // some sanity checks if (all_return_pre_rht_amax) { 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 030dddfce4..19583b3afb 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 @@ -1428,9 +1428,8 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, float *const amax_rowwise_base_ptr = reinterpret_cast(output->amax.dptr); float *const amax_colwise_base_ptr = reinterpret_cast(output->columnwise_amax.dptr); - const int64_t *const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); - const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); - // const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(output->first_dims.dptr); const bool is_const_last_dim = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || shape_rep == ShapeRepresentation::VARYING_FIRST_DIM); @@ -1441,7 +1440,7 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, int k_tile_size = 1024; - const bool use_swizzle_sf_output = false; + const bool use_swizzle_sf_output = output->with_gemm_swizzled_scales; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, kEnableStochasticRounding, diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index ae41f238a4..e316f8be8c 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -449,6 +449,8 @@ enum NVTEGroupedTensorParam { kNVTEGroupedLastDims = 8, /*!< Last dimension sizes (device pointer to int64_t array) */ kNVTEGroupedTensorOffsets = 9, /*!< Tensor offsets for contiguous layout (device pointer to int64_t array) */ + kNVTEGroupedWithGEMMSwizzledScales = + 10, /*!< Whether scaling factors are in format expected by GEMM */ kNVTENumGroupedTensorParams }; @@ -479,25 +481,30 @@ NVTEGroupedTensor nvte_create_grouped_tensor(NVTEScalingMode scaling_mode, size_ void nvte_destroy_grouped_tensor(NVTEGroupedTensor tensor); /* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ -/*! \brief Set a parameter of the grouped tensor. +/*! \brief Set a grouped tensor parameter. * - * \param[in/out] tensor Grouped tensor. - * \param[in] param_name The parameter to be set. - * \param[in] param The value to be set (NVTEBasicTensor). + * \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. */ -void nvte_set_grouped_tensor_param(NVTEGroupedTensor *tensor, NVTEGroupedTensorParam param_name, - const NVTEBasicTensor *param); +void nvte_set_grouped_tensor_param(NVTEGroupedTensor tensor, NVTEGroupedTensorParam param, + const void *buf, size_t size_in_bytes); /* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ -/*! \brief Get a value of the parameter of the grouped tensor. - * - * \param[in] tensor Grouped tensor. - * \param[in] param_name The parameter to be queried. +/*! \brief Query a grouped tensor parameter. * - * \return NVTEBasicTensor containing the parameter data. + * \param[in] tensor Grouped tensor. + * \param[in] param Grouped tensor parameter type. + * \param[out] buf Memory address to write parameter value. + * Ignored if NULL. + * \param[in] size_in_bytes Size of buf. + * \param[out] size_written Number of bytes that have been written to + * buf. If buf is NULL, then the number of + * bytes that would have been written. */ -NVTEBasicTensor nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, - NVTEGroupedTensorParam param_name); +void nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, NVTEGroupedTensorParam param, + void *buf, size_t size_in_bytes, size_t *size_written); /* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ /*! \brief Get the number of tensors in a grouped tensor. @@ -957,8 +964,235 @@ class TensorWrapper { NVTETensor tensor_ = nullptr; }; -/*! \warning Deprecated */ -enum class Float8BlockScaleTensorFormat { GEMM_READY = 0, COMPACT = 1, INVALID }; +/*! \struct GroupedTensorWrapper + * \brief C++ wrapper for the NVTEGroupedTensor class. + */ + +class GroupedTensorWrapper { + public: + /*! \brief Constructs new GroupedTensorWrapper. + * + * Create a new TE grouped tensor with a given logical shape. + * TE grouped tensors are just wrappers on top of raw data and do not + * own memory. + * + * \param[in] num_tensors Number of tensors in the group (must be > 0). + * \param[in] logical_shape Logical 2D shape of the grouped data. + * \param[in] scaling_mode Tensor data format. + */ + GroupedTensorWrapper(const size_t num_tensors, const NVTEShape &logical_shape, + const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING) + : tensor_(nvte_create_grouped_tensor(scaling_mode, num_tensors, logical_shape)) {} + + /*! \brief Constructs new GroupedTensorWrapper. + * + * Create a new TE grouped tensor with a given logical shape. + * + * \param[in] num_tensors Number of tensors in the group (must be > 0). + * \param[in] logical_shape Logical 2D shape of the grouped data. + * \param[in] scaling_mode Tensor data format. + */ + GroupedTensorWrapper(const size_t num_tensors, const std::vector &logical_shape, + const NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING) + : GroupedTensorWrapper(num_tensors, + nvte_make_shape(logical_shape.data(), logical_shape.size()), + scaling_mode) {} + + /*! \brief GroupedTensorWrapper destructor. */ + ~GroupedTensorWrapper() { nvte_destroy_grouped_tensor(tensor_); } + + GroupedTensorWrapper &operator=(const GroupedTensorWrapper &other) = delete; + GroupedTensorWrapper(const GroupedTensorWrapper &other) = delete; + + /*! \brief Constructs new GroupedTensorWrapper from existing GroupedTensorWrapper. */ + GroupedTensorWrapper(GroupedTensorWrapper &&other) { + tensor_ = other.tensor_; + other.tensor_ = nullptr; + } + + /*! \brief Assign the data from existing GroupedTensorWrapper. */ + GroupedTensorWrapper &operator=(GroupedTensorWrapper &&other) { + if (this == &other) return *this; + nvte_destroy_grouped_tensor(tensor_); + tensor_ = other.tensor_; + other.tensor_ = nullptr; + return *this; + } + + // Parameter setters + template + GroupedTensorWrapper &set_parameter(const NVTEGroupedTensorParam param, void *dptr, DType type, + const ShapeType &shape) noexcept { + NVTEShape nvte_shape = this->convertShape(shape); + NVTEBasicTensor data = {dptr, static_cast(type), nvte_shape}; + nvte_set_grouped_tensor_param(tensor_, param, &data, sizeof(data)); + return *this; + } + + template + GroupedTensorWrapper &set_rowwise_data(void *dptr, DType type, const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedRowwiseData, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_columnwise_data(void *dptr, DType type, + const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedColumnwiseData, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_scale(void *dptr, DType type, const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedScale, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_amax(void *dptr, DType type, const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedAmax, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_rowwise_scale_inv(void *dptr, DType type, + const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedRowwiseScaleInv, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_columnwise_scale_inv(void *dptr, DType type, + const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedColumnwiseScaleInv, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_columnwise_amax(void *dptr, DType type, + const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedColumnwiseAmax, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_first_dims(void *dptr, DType type, const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedFirstDims, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_last_dims(void *dptr, DType type, const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedLastDims, dptr, type, shape); + } + + template + GroupedTensorWrapper &set_tensor_offsets(void *dptr, DType type, + const ShapeType &shape) noexcept { + return set_parameter(kNVTEGroupedTensorOffsets, dptr, type, shape); + } + + void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales) { + const auto val = static_cast(with_gemm_swizzled_scales); + nvte_set_grouped_tensor_param(tensor_, kNVTEGroupedWithGEMMSwizzledScales, &val, sizeof(val)); + } + + // Parameter getters + NVTEBasicTensor get_parameter(const NVTEGroupedTensorParam param) const noexcept { + NVTEBasicTensor ret; + nvte_get_grouped_tensor_param(tensor_, param, &ret, sizeof(ret), nullptr); + return ret; + } + + NVTEBasicTensor get_rowwise_data() const noexcept { + return get_parameter(kNVTEGroupedRowwiseData); + } + + NVTEBasicTensor get_columnwise_data() const noexcept { + return get_parameter(kNVTEGroupedColumnwiseData); + } + + NVTEBasicTensor get_scale() const noexcept { return get_parameter(kNVTEGroupedScale); } + + NVTEBasicTensor get_amax() const noexcept { return get_parameter(kNVTEGroupedAmax); } + + NVTEBasicTensor get_rowwise_scale_inv() const noexcept { + return get_parameter(kNVTEGroupedRowwiseScaleInv); + } + + NVTEBasicTensor get_columnwise_scale_inv() const noexcept { + return get_parameter(kNVTEGroupedColumnwiseScaleInv); + } + + NVTEBasicTensor get_columnwise_amax() const noexcept { + return get_parameter(kNVTEGroupedColumnwiseAmax); + } + + NVTEBasicTensor get_first_dims() const noexcept { return get_parameter(kNVTEGroupedFirstDims); } + + NVTEBasicTensor get_last_dims() const noexcept { return get_parameter(kNVTEGroupedLastDims); } + + NVTEBasicTensor get_tensor_offsets() const noexcept { + return get_parameter(kNVTEGroupedTensorOffsets); + } + + bool get_with_gemm_swizzled_scales() const { + uint8_t val = 0; + nvte_get_grouped_tensor_param(tensor_, kNVTEGroupedWithGEMMSwizzledScales, &val, sizeof(val), + nullptr); + return static_cast(val); + } + + /*! \brief Get an underlying NVTEGroupedTensor. + * + * \return NVTEGroupedTensor held by this GroupedTensorWrapper. + */ + NVTEGroupedTensor data() const noexcept { return tensor_; } + + /*! \brief Get the number of tensors in this GroupedTensorWrapper. */ + size_t num_tensors() const noexcept { + if (tensor_ == nullptr) return 0; + return nvte_grouped_tensor_num_tensors(tensor_); + } + + /*! \brief Get the data type of this GroupedTensorWrapper. */ + DType dtype() const noexcept { + if (tensor_ == nullptr) return DType::kNumTypes; + return static_cast(nvte_grouped_tensor_type(tensor_)); + } + + /*! \brief Get a scaling mode of the grouped tensor. */ + NVTEScalingMode scaling_mode() const noexcept { + if (tensor_ == nullptr) return NVTE_DELAYED_TENSOR_SCALING; + return nvte_grouped_tensor_scaling_mode(tensor_); + } + + /*! \brief Get the logical shape of this GroupedTensorWrapper. */ + const NVTEShape logical_shape() const noexcept { + if (tensor_ == nullptr) { + return emptyShape; + } + return nvte_get_grouped_tensor_logical_shape(tensor_); + } + + static constexpr size_t defaultData = 1; + static constexpr NVTEShape defaultShape = { + {defaultData, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 1}; + static constexpr NVTEShape emptyShape = {{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}, 1}; + + private: + NVTEShape convertShape(const NVTEShape &s) { return s; } + + NVTEShape convertShape(const std::vector &s) { + return nvte_make_shape(s.data(), s.size()); + } + + /*! \brief Wrapped NVTEGroupedTensor. */ + NVTEGroupedTensor tensor_ = nullptr; +}; + +/*! \enum Float8BlockScaleTensorFormat + * \brief Data format for an FP8 block-scaled tensor + */ +enum class Float8BlockScaleTensorFormat { + /*! FP8 data is transposed if needed and scales are swizzled */ + GEMM_READY = 0, + /*! FP8 data is untransposed and scales are not swizzled or padded */ + COMPACT = 1, + INVALID +}; /*! \struct QuantizationConfigWrapper * \brief C++ wrapper for NVTEQuantizationConfigWrapper. diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 06971443dd..cd02074fbd 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1145,8 +1145,8 @@ NVTEGroupedTensor nvte_create_grouped_tensor(NVTEScalingMode scaling_mode, size_ NVTEShape logical_shape) { NVTE_CHECK(num_tensors > 0, "Number of tensors must be greater than 0"); NVTE_CHECK(logical_shape.ndim == 2, "Logical shape must be 2D"); - NVTE_CHECK(logical_shape.data[0] > 0 && logical_shape.data[1] > 0, - "Logical shape must have positive dimensions"); + // NVTE_CHECK(logical_shape.data[0] > 0 && logical_shape.data[1] > 0, + // "Logical shape must have positive dimensions"); NVTEGroupedTensor ret = transformer_engine::GroupedTensorAllocator::instance().Allocate( scaling_mode, num_tensors, logical_shape); return ret; @@ -1156,88 +1156,178 @@ void nvte_destroy_grouped_tensor(NVTEGroupedTensor tensor) { transformer_engine::GroupedTensorAllocator::instance().Free(tensor); } -void nvte_set_grouped_tensor_param(NVTEGroupedTensor *tensor, NVTEGroupedTensorParam param_name, - const NVTEBasicTensor *param) { +void nvte_set_grouped_tensor_param(NVTEGroupedTensor tensor, NVTEGroupedTensorParam param, + const void *buf, size_t size_in_bytes) { + using namespace transformer_engine; + + // Check attribute and buffer + NVTE_CHECK(param < kNVTENumGroupedTensorParams, "Invalid NVTEGroupedTensorParam (got ", + static_cast(param), ")"); NVTE_CHECK(tensor != nullptr, "Grouped tensor pointer can't be NULL."); - auto *t = transformer_engine::convertNVTEGroupedTensor(*tensor); - NVTE_CHECK(t != nullptr, "Grouped tensor is not allocated."); - NVTE_CHECK(param != nullptr, "Grouped tensor param can't be NULL."); + auto &t = *convertNVTEGroupedTensorCheck(tensor); + const auto &attr_size = GroupedTensor::attr_sizes[param]; + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for grouped tensor parameter " + "(parameter ", + static_cast(param), " needs ", attr_size, " bytes, but buffer has ", + size_in_bytes, " bytes)"); + NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); - switch (param_name) { - case kNVTEGroupedRowwiseData: - t->data = *param; + // Read from buffer + switch (param) { + case kNVTEGroupedRowwiseData: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.data = *basic_tensor; break; - case kNVTEGroupedColumnwiseData: - t->columnwise_data = *param; + } + case kNVTEGroupedColumnwiseData: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.columnwise_data = *basic_tensor; break; - case kNVTEGroupedScale: - t->scale = *param; + } + case kNVTEGroupedScale: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.scale = *basic_tensor; break; - case kNVTEGroupedAmax: - t->amax = *param; + } + case kNVTEGroupedAmax: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.amax = *basic_tensor; break; - case kNVTEGroupedRowwiseScaleInv: - t->scale_inv = *param; + } + case kNVTEGroupedRowwiseScaleInv: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.scale_inv = *basic_tensor; break; - case kNVTEGroupedColumnwiseScaleInv: - t->columnwise_scale_inv = *param; + } + case kNVTEGroupedColumnwiseScaleInv: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.columnwise_scale_inv = *basic_tensor; break; - case kNVTEGroupedColumnwiseAmax: - t->columnwise_amax = *param; + } + case kNVTEGroupedColumnwiseAmax: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.columnwise_amax = *basic_tensor; break; - case kNVTEGroupedFirstDims: - t->first_dims = *param; - // Validate it's Int64 - NVTE_CHECK(t->first_dims.dtype == transformer_engine::DType::kInt64, - "first_dims must have dtype Int64"); + } + case kNVTEGroupedFirstDims: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.first_dims = *basic_tensor; + NVTE_CHECK(t.first_dims.dtype == DType::kInt64, "first_dims must have dtype Int64"); break; - case kNVTEGroupedLastDims: - t->last_dims = *param; - // Validate it's Int64 - NVTE_CHECK(t->last_dims.dtype == transformer_engine::DType::kInt64, - "last_dims must have dtype Int64"); + } + case kNVTEGroupedLastDims: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.last_dims = *basic_tensor; + NVTE_CHECK(t.last_dims.dtype == DType::kInt64, "last_dims must have dtype Int64"); break; - case kNVTEGroupedTensorOffsets: - t->tensor_offsets = *param; - // Validate it's Int64 - NVTE_CHECK(t->tensor_offsets.dtype == transformer_engine::DType::kInt64, - "tensor_offsets must have dtype Int64"); + } + case kNVTEGroupedTensorOffsets: { + const NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + t.tensor_offsets = *basic_tensor; + NVTE_CHECK(t.tensor_offsets.dtype == DType::kInt64, "tensor_offsets must have dtype Int64"); + break; + } + case kNVTEGroupedWithGEMMSwizzledScales: + t.with_gemm_swizzled_scales = static_cast(*reinterpret_cast(buf)); break; default: - NVTE_ERROR("Unknown grouped tensor parameter!"); + NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast(param), ")"); } } -NVTEBasicTensor nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, - NVTEGroupedTensorParam param_name) { - if (tensor == nullptr) { - return {nullptr, kNVTEFloat32, nvte_make_shape(nullptr, 1)}; +void nvte_get_grouped_tensor_param(const NVTEGroupedTensor tensor, NVTEGroupedTensorParam param, + void *buf, size_t size_in_bytes, size_t *size_written) { + using namespace transformer_engine; + + // Check param + NVTE_CHECK(param < kNVTENumGroupedTensorParams, "Invalid NVTEGroupedTensorParam (got ", + static_cast(param), ")"); + + // Write attribute size if provided + const auto &attr_size = GroupedTensor::attr_sizes[param]; + if (size_written != nullptr) { + *size_written = attr_size; } - const auto &t = *transformer_engine::convertNVTEGroupedTensorCheck(tensor); - switch (param_name) { - case kNVTEGroupedRowwiseData: - return t.data; - case kNVTEGroupedColumnwiseData: - return t.columnwise_data; - case kNVTEGroupedScale: - return t.scale; - case kNVTEGroupedAmax: - return t.amax; - case kNVTEGroupedRowwiseScaleInv: - return t.scale_inv; - case kNVTEGroupedColumnwiseScaleInv: - return t.columnwise_scale_inv; - case kNVTEGroupedColumnwiseAmax: - return t.columnwise_amax; - case kNVTEGroupedFirstDims: - return t.first_dims; - case kNVTEGroupedLastDims: - return t.last_dims; - case kNVTEGroupedTensorOffsets: - return t.tensor_offsets; + // Return immediately if buffer is not provided + if (buf == nullptr) { + return; + } + + // Check buffer size + NVTE_CHECK(size_in_bytes >= attr_size, + "Buffer is too small for grouped tensor parameter " + "(parameter ", + static_cast(param), " needs ", attr_size, " bytes, but buffer has ", + size_in_bytes, " bytes)"); + + // Get C++ grouped tensor + const GroupedTensor *t = convertNVTEGroupedTensor(tensor); + std::optional dummy; + if (t == nullptr) { + // Make dummy grouped tensor if provided tensor is invalid + dummy.emplace(NVTE_DELAYED_TENSOR_SCALING, 1); + t = &(*dummy); + } + + // Write to buffer + switch (param) { + case kNVTEGroupedRowwiseData: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->data); + break; + } + case kNVTEGroupedColumnwiseData: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->columnwise_data); + break; + } + case kNVTEGroupedScale: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->scale); + break; + } + case kNVTEGroupedAmax: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->amax); + break; + } + case kNVTEGroupedRowwiseScaleInv: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->scale_inv); + break; + } + case kNVTEGroupedColumnwiseScaleInv: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->columnwise_scale_inv); + break; + } + case kNVTEGroupedColumnwiseAmax: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->columnwise_amax); + break; + } + case kNVTEGroupedFirstDims: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->first_dims); + break; + } + case kNVTEGroupedLastDims: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->last_dims); + break; + } + case kNVTEGroupedTensorOffsets: { + NVTEBasicTensor *basic_tensor = reinterpret_cast(buf); + *basic_tensor = static_cast(t->tensor_offsets); + break; + } + case kNVTEGroupedWithGEMMSwizzledScales: + *reinterpret_cast(buf) = static_cast(t->with_gemm_swizzled_scales); + break; default: - NVTE_ERROR("Unknown grouped tensor parameter!"); + NVTE_ERROR("Unsupported grouped tensor parameter (", static_cast(param), ")"); } } diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index bc22e03097..6aab9938b3 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -103,6 +103,12 @@ class Quantizer { virtual std::pair create_tensor(const std::vector& shape, DType dtype) const = 0; + /*! @brief Construct a grouped tensor with uninitialized data */ + virtual std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + size_t logical_last_dim) const = 0; + /*! @brief Convert a PyTorch tensor into a Transformer Engine C++ tensor * * The PyTorch tensor's attributes are modified to match the @@ -138,6 +144,11 @@ class NoneQuantizer : public Quantizer { std::pair create_tensor(const std::vector& shape, DType dtype) const override; + std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + 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, at::Tensor data) const; @@ -164,6 +175,11 @@ class Float8Quantizer : public Quantizer { std::pair create_tensor(const std::vector& shape, DType dtype) const override; + std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + 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, @@ -196,6 +212,11 @@ class Float8CurrentScalingQuantizer : public Quantizer { std::pair create_tensor(const std::vector& shape, DType dtype) const override; + std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + 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. * * The amax is zeroed out. Most TE kernels that output amax expect @@ -253,6 +274,11 @@ class Float8BlockQuantizer : public Quantizer { std::pair create_tensor(const std::vector& shape, DType dtype) const override; + std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + size_t logical_last_dim) const override; + std::pair convert_and_update_tensor(py::object shape) const override; void quantize(const TensorWrapper& input, TensorWrapper& out, @@ -274,6 +300,11 @@ class MXFP8Quantizer : public Quantizer { std::pair create_tensor(const std::vector& shape, DType dtype) const override; + std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, + size_t logical_last_dim) const override; + std::pair convert_and_update_tensor(py::object shape) const override; void quantize(const TensorWrapper& input, TensorWrapper& out, @@ -308,6 +339,11 @@ class NVFP4Quantizer : public Quantizer { std::pair create_tensor(const std::vector& shape, DType dtype) const override; + std::pair create_grouped_tensor( + size_t num_tensors, const std::vector& logical_shape, DType dtype, + 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 NVFP4 tensor's amax pointer * * The amax is zeroed out. Most TE kernels that output amax expect diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 0e91071983..cb6d8b7c92 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -255,6 +255,9 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob py::object dequantize(const py::handle &input, DType otype); +py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, + std::optional first_dims); + std::vector multi_tensor_quantize(const std::vector &tensor_list, std::vector quantizer_list); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 5c9d0f5b07..f8f793f036 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -80,6 +80,157 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob return output_py; } +namespace { + +// helper functions for NVFP4 grouped quantization (cuda graph safe with shapes stored in device without D2H copy) +void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, + GroupedTensorWrapper &grouped_output_tensor, + NVFP4Quantizer *nvfp4_quantizer_cpp, cudaStream_t stream) { + size_t num_tensors = grouped_input_tensor.num_tensors(); + + // 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"); + + auto quant_config_cpp = QuantizationConfigWrapper(); + + // stochastic rounding + bool need_stochastic_rounding = nvfp4_quantizer_cpp->stochastic_rounding; + auto opts = at::TensorOptions().dtype(torch::kInt64).device(torch::kCUDA); + at::Tensor rng_states_tensor; // Declare tensor outside, do not allocate yet + TensorWrapper te_rng_state; + + if (need_stochastic_rounding) { + // in fused kernel, one rng state will be used by the grouped kernel to generate random + // number for different tensors in the group, so we only need to allocate one rng state + const size_t rng_elts_per_thread = 1024 * num_tensors; + rng_states_tensor = torch::empty({2}, opts); + auto gen = at::get_generator_or_default( + std::nullopt, at::cuda::detail::getDefaultCUDAGenerator()); + at::PhiloxCudaState philox_args = init_philox_state(gen, rng_elts_per_thread); + philox_unpack(philox_args, static_cast(rng_states_tensor.data_ptr())); + + te_rng_state = makeTransformerEngineTensor(rng_states_tensor); + quant_config_cpp.set_rng_state(te_rng_state.data()); + quant_config_cpp.set_stochastic_rounding(true); + } + + // fast math + const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); + if (use_fast_math) { + quant_config_cpp.set_use_fast_math(true); + } + + // so far, only the RHT path has grouped kernel support + // grouped kernels for non-RHT path will be added later + + if (nvfp4_quantizer_cpp->with_rht) { + // post-RHT amax or not + if (nvfp4_quantizer_cpp->with_post_rht_amax) { + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_amax_graph_safe( + grouped_input_tensor.data(), grouped_output_tensor.data(), 0, + nvfp4_quantizer_cpp->rht_matrix_random_sign_mask_t, stream); + }); + } else { + NVTE_ERROR("graph safe grouped quant kernel for non-RHT path is not ready yet"); + } + + // RHT cast fusion + auto tile_scheduler_workspace_torch = + at::empty({1}, at::device(at::kCUDA).dtype(torch::kInt32)); + auto nvte_tile_scheduler_workspace = + makeTransformerEngineTensor(tile_scheduler_workspace_torch); + + auto rht_matrix_nvte = makeTransformerEngineTensor(nvfp4_quantizer_cpp->rht_matrix); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_hadamard_transform_cast_fusion_graph_safe( + grouped_input_tensor.data(), grouped_output_tensor.data(), rht_matrix_nvte.data(), + quant_config_cpp, nvte_tile_scheduler_workspace.data(), stream); + }); + + } else { + NVTE_ERROR("graph safe grouped quant kernel for non-RHT path is not ready yet"); + } +} + +} // namespace + +// NOTE: Only supports varying first dim. +py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, + std::optional first_dims) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + + std::vector logical_shape; + for (const auto &d : tensor.sizes()) { + logical_shape.push_back(d); + } + const auto logical_first_dim = logical_shape[0]; + const auto logical_last_dim = logical_shape[1]; + + bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0; + + auto quantizer_cpp = convert_quantizer(quantizer); + + // Create input GroupedTensor. + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_input_tensor.set_rowwise_data( + tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + + // Create output GroupedTensor. + auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, logical_first_dim, + logical_last_dim); + + // dispatch to scaling methods + enum class GroupedQuantizationMode { + MXFP8_GROUPED_QUANTIZE, + NVFP4_GROUPED_QUANTIZE, + INVALID_FOR_GROUPED_QUANTIZE + }; + GroupedQuantizationMode grouped_quantization_mode = + GroupedQuantizationMode::INVALID_FOR_GROUPED_QUANTIZE; + if (detail::IsMXFP8Quantizers(quantizer.ptr())) { + grouped_quantization_mode = GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE; + } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) { + grouped_quantization_mode = GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE; + } + + if (empty_input_buffer) { + // early return for empty input buffer + // just return the output tensor as is + // no need to quantize + return py::reinterpret_borrow(grouped_output_py); + } + + switch (grouped_quantization_mode) { + case GroupedQuantizationMode::NVFP4_GROUPED_QUANTIZE: { + // NVFP4 grouped quantization + NVFP4Quantizer *nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); + group_quantize_nvfp4_impl(grouped_input_tensor, grouped_output_tensor_cpp, + nvfp4_quantizer_cpp, at::cuda::getCurrentCUDAStream()); + break; + } + case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: { + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), + at::cuda::getCurrentCUDAStream()); + }); + break; + } + case GroupedQuantizationMode::INVALID_FOR_GROUPED_QUANTIZE: + default: + NVTE_ERROR("group_quantize: only support NVFP4 or MXFP8 quantizer."); + break; + } + + return py::reinterpret_borrow(grouped_output_py); +} + py::object dequantize(const py::handle &input, transformer_engine::DType otype) { init_extension(); diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 14f32c7b93..5372f2f3e7 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -35,6 +35,7 @@ PyTypeObject *Float8BlockwiseQuantizerClass = nullptr; PyTypeObject *NVFP4TensorPythonClass = nullptr; PyTypeObject *NVFP4TensorStoragePythonClass = nullptr; PyTypeObject *NVFP4QuantizerClass = nullptr; +PyTypeObject *GroupedTensorStoragePythonClass = nullptr; void init_float8_extension() { if (Float8TensorPythonClass) return; @@ -104,11 +105,22 @@ void init_nvfp4_extensions() { "Internal error: could not initialize pyTorch NVFP4 extension."); } +void init_grouped_tensor_extension() { + if (GroupedTensorStoragePythonClass) return; + auto grouped_tensor_module = + py::module_::import("transformer_engine.pytorch.tensor.storage.grouped_tensor"); + GroupedTensorStoragePythonClass = reinterpret_cast( + PyObject_GetAttrString(grouped_tensor_module.ptr(), "GroupedTensor")); + NVTE_CHECK(GroupedTensorStoragePythonClass != nullptr, + "Internal error: could not initialize pyTorch grouped tensor extension."); +} + void init_extension() { init_float8_extension(); init_mxfp8_extension(); init_float8blockwise_extension(); init_nvfp4_extensions(); + init_grouped_tensor_extension(); } } // namespace transformer_engine::pytorch @@ -121,7 +133,8 @@ 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("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"), + py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims")); m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize, "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer")); m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)", diff --git a/transformer_engine/pytorch/csrc/pybind.h b/transformer_engine/pytorch/csrc/pybind.h index 25ffef0588..059eb5e3fb 100644 --- a/transformer_engine/pytorch/csrc/pybind.h +++ b/transformer_engine/pytorch/csrc/pybind.h @@ -43,6 +43,7 @@ extern PyTypeObject *Float8BlockwiseQuantizerClass; extern PyTypeObject *NVFP4TensorPythonClass; extern PyTypeObject *NVFP4TensorStoragePythonClass; extern PyTypeObject *NVFP4QuantizerClass; +extern PyTypeObject *GroupedTensorStoragePythonClass; void init_extension(); @@ -95,6 +96,8 @@ TensorWrapper NVTETensorFromFloat8BlockwiseQTensor(py::handle tensor, TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer); +GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor); + inline bool IsFloatingPointType(at::ScalarType type) { return type == at::kFloat || type == at::kHalf || type == at::kBFloat16; } diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 1c968e276d..e715d8f5ba 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -42,6 +42,35 @@ std::vector convert_shape_for_fp4(const std::vector& shape) { return ret; } +std::optional build_grouped_tensor_offsets(const size_t num_tensors, + const std::optional& first_dims, + const size_t logical_last_dim) { + if (!first_dims.has_value()) { + return std::nullopt; + } + + const auto& first_dims_tensor = first_dims.value(); + NVTE_CHECK(first_dims_tensor.scalar_type() == at::kLong, "first_dims must have dtype int64."); + NVTE_CHECK(static_cast(first_dims_tensor.numel()) == num_tensors, + "first_dims must have length ", num_tensors, "."); + + const int64_t logical_last_dim_i64 = static_cast(logical_last_dim); + auto scaled_first_dims = first_dims_tensor * logical_last_dim_i64; + + // Single kernel needed for these ops. + auto cumsum = at::cumsum(scaled_first_dims, 0); + auto zero = at::zeros({1}, cumsum.options()); + return at::cat({zero, cumsum}); +} + +at::TensorOptions grouped_tensor_data_options(const DType dtype) { + return at::TensorOptions().dtype(GetATenDType(dtype)).device(torch::kCUDA); +} + +py::object maybe_tensor_to_py(const std::optional& tensor) { + return tensor ? py::cast(*tensor) : py::none(); +} + } // namespace constexpr size_t NVFP4_BLOCK_SIZE = 16; @@ -88,6 +117,60 @@ std::pair NoneQuantizer::create_tensor(const std::vec return {std::move(out_cpp), py::cast(data)}; } +std::pair NoneQuantizer::create_grouped_tensor( + const size_t num_tensors, const std::vector& logical_shape, const DType dtype, + py::object quantizer, const std::optional& first_dims, + const size_t logical_first_dim, const size_t logical_last_dim) const { + using namespace pybind11::literals; + + const auto tensor_offsets = + build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + const int64_t total_elements = + static_cast(logical_first_dim) * static_cast(logical_last_dim); + + std::optional rowwise_data; + std::optional columnwise_data; + const bool with_rowwise_data = rowwise_usage; + const bool with_columnwise_data = columnwise_usage; + if (with_rowwise_data) { + rowwise_data = at::empty({total_elements}, grouped_tensor_data_options(dtype)); + } + if (with_columnwise_data) { + columnwise_data = at::empty({total_elements}, grouped_tensor_data_options(dtype)); + } + + GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); + if (with_rowwise_data) { + out_cpp.set_rowwise_data(rowwise_data->data_ptr(), dtype, getTensorShape(*rowwise_data)); + } + if (with_columnwise_data) { + out_cpp.set_columnwise_data(columnwise_data->data_ptr(), dtype, + getTensorShape(*columnwise_data)); + } + if (first_dims.has_value()) { + out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (tensor_offsets.has_value()) { + out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); + py::object out_py = GroupedTensorClass( + "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), + "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), + "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), "scale_inv"_a = py::none(), + "columnwise_scale_inv"_a = py::none(), "amax"_a = py::none(), + "columnwise_amax"_a = py::none(), "scale"_a = py::none(), + "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), + "last_dims"_a = py::none(), + "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), + "logical_shape"_a = std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + + return {std::move(out_cpp), std::move(out_py)}; +} + std::pair NoneQuantizer::convert_and_update_tensor( py::object tensor) const { auto tensor_pyt = tensor.cast(); @@ -184,6 +267,73 @@ std::pair Float8Quantizer::create_tensor( return {std::move(out_cpp), std::move(out_py)}; } +std::pair Float8Quantizer::create_grouped_tensor( + const size_t num_tensors, const std::vector& logical_shape, const DType dtype, + py::object quantizer, const std::optional& first_dims, + const size_t logical_first_dim, const size_t logical_last_dim) const { + using namespace pybind11::literals; + + const auto tensor_offsets = + build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + const int64_t total_elements = + static_cast(logical_first_dim) * static_cast(logical_last_dim); + + const auto uint8_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto float_opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + + std::optional rowwise_data; + std::optional columnwise_data; + std::optional rowwise_scale_inv; + std::optional columnwise_scale_inv; + at::Tensor amax = at::empty({static_cast(num_tensors)}, float_opts); + + if (rowwise_usage) { + rowwise_data = at::empty({total_elements}, uint8_opts); + rowwise_scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + } + if (columnwise_usage) { + columnwise_data = at::empty({total_elements}, uint8_opts); + columnwise_scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + } + + GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); + if (rowwise_usage) { + out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat32, + getTensorShape(*rowwise_scale_inv)); + } + if (columnwise_usage) { + out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, + getTensorShape(*columnwise_data)); + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat32, + getTensorShape(*columnwise_scale_inv)); + } + out_cpp.set_amax(amax.data_ptr(), DType::kFloat32, getTensorShape(amax)); + if (first_dims.has_value()) { + out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (tensor_offsets.has_value()) { + out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); + py::object out_py = GroupedTensorClass( + "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), + "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), + "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), + "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), + "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = amax, + "columnwise_amax"_a = py::none(), "scale"_a = py::none(), + "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), + "last_dims"_a = py::none(), + "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), + "logical_shape"_a = std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + + return {std::move(out_cpp), std::move(out_py)}; +} + std::pair Float8Quantizer::convert_and_update_tensor( py::object tensor) const { NVTE_CHECK(detail::IsFloat8Tensor(tensor.ptr()), "Float8Quantizer must output to Float8Tensor."); @@ -390,6 +540,75 @@ std::pair Float8CurrentScalingQuantizer::create_tenso return {std::move(out_cpp), std::move(out_py)}; } +std::pair Float8CurrentScalingQuantizer::create_grouped_tensor( + const size_t num_tensors, const std::vector& logical_shape, const DType dtype, + py::object quantizer, const std::optional& first_dims, + const size_t logical_first_dim, const size_t logical_last_dim) const { + using namespace pybind11::literals; + + const auto tensor_offsets = + build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + const int64_t total_elements = + static_cast(logical_first_dim) * static_cast(logical_last_dim); + + const auto uint8_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto float_opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + + std::optional rowwise_data; + std::optional columnwise_data; + std::optional rowwise_scale_inv; + std::optional columnwise_scale_inv; + at::Tensor scale = at::empty({static_cast(num_tensors)}, float_opts); + at::Tensor amax = at::empty({static_cast(num_tensors)}, float_opts); + + if (rowwise_usage) { + rowwise_data = at::empty({total_elements}, uint8_opts); + rowwise_scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + } + if (columnwise_usage) { + columnwise_data = at::empty({total_elements}, uint8_opts); + columnwise_scale_inv = at::empty({static_cast(num_tensors)}, float_opts); + } + + GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); + if (rowwise_usage) { + out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat32, + getTensorShape(*rowwise_scale_inv)); + } + if (columnwise_usage) { + out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, + getTensorShape(*columnwise_data)); + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat32, + getTensorShape(*columnwise_scale_inv)); + } + out_cpp.set_scale(scale.data_ptr(), DType::kFloat32, getTensorShape(scale)); + out_cpp.set_amax(amax.data_ptr(), DType::kFloat32, getTensorShape(amax)); + if (first_dims.has_value()) { + out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (tensor_offsets.has_value()) { + out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); + py::object out_py = GroupedTensorClass( + "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), + "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), + "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), + "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), + "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = amax, + "columnwise_amax"_a = py::none(), "scale"_a = scale, + "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), + "last_dims"_a = py::none(), + "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), + "logical_shape"_a = std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + + return {std::move(out_cpp), std::move(out_py)}; +} + std::pair Float8CurrentScalingQuantizer::create_unquantized_tensor_with_amax(const std::vector& shape, DType dtype, @@ -638,6 +857,77 @@ std::pair Float8BlockQuantizer::create_tensor( return {std::move(tensor), std::move(ret)}; } +std::pair Float8BlockQuantizer::create_grouped_tensor( + const size_t num_tensors, const std::vector& logical_shape, const DType dtype, + py::object quantizer, const std::optional& first_dims, + const size_t logical_first_dim, const size_t logical_last_dim) const { + using namespace pybind11::literals; + + const auto tensor_offsets = + build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + const int64_t total_elements = + static_cast(logical_first_dim) * static_cast(logical_last_dim); + + const auto uint8_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto float_opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + + std::optional rowwise_data; + std::optional columnwise_data; + std::optional rowwise_scale_inv; + std::optional columnwise_scale_inv; + const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; + + if (rowwise_usage) { + rowwise_data = at::empty({total_elements}, uint8_opts); + 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}, float_opts); + } + + if (columnwise_usage) { + columnwise_data = at::empty({total_elements}, uint8_opts); + const auto scale_shape = get_scale_shape(logical_shape_vec, true); + const int64_t total_scale_elements = static_cast(product(scale_shape)); + columnwise_scale_inv = at::empty({total_scale_elements}, float_opts); + } + + GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); + if (rowwise_usage) { + out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat32, + getTensorShape(*rowwise_scale_inv)); + } + if (columnwise_usage) { + out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, + getTensorShape(*columnwise_data)); + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat32, + getTensorShape(*columnwise_scale_inv)); + } + if (first_dims.has_value()) { + out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (tensor_offsets.has_value()) { + out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); + py::object out_py = GroupedTensorClass( + "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), + "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), + "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), + "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), + "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = py::none(), + "columnwise_amax"_a = py::none(), "scale"_a = py::none(), + "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), + "last_dims"_a = py::none(), + "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), + "logical_shape"_a = std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + + return {std::move(out_cpp), std::move(out_py)}; +} + std::pair Float8BlockQuantizer::convert_and_update_tensor( py::object tensor) const { const DType dtype = tensor.attr("_fp8_dtype").cast(); @@ -940,6 +1230,78 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve return {std::move(out_cpp), std::move(out_py)}; } +std::pair MXFP8Quantizer::create_grouped_tensor( + const size_t num_tensors, const std::vector& logical_shape, const DType dtype, + py::object quantizer, const std::optional& first_dims, + const size_t logical_first_dim, const size_t logical_last_dim) const { + using namespace pybind11::literals; + + const auto tensor_offsets = + build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + const int64_t total_elements = + static_cast(logical_first_dim) * static_cast(logical_last_dim); + + const auto uint8_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + + std::optional rowwise_data; + std::optional columnwise_data; + std::optional rowwise_scale_inv; + std::optional columnwise_scale_inv; + const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; + + if (rowwise_usage) { + rowwise_data = at::empty({total_elements}, uint8_opts); + 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); + } + + if (columnwise_usage) { + columnwise_data = at::empty({total_elements}, uint8_opts); + const auto scale_shape = get_scale_shape(logical_shape_vec, true); + const int64_t total_scale_elements = static_cast(product(scale_shape)); + columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); + } + + GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); + if (rowwise_usage) { + out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E8M0, + getTensorShape(*rowwise_scale_inv)); + } + if (columnwise_usage) { + out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, + getTensorShape(*columnwise_data)); + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E8M0, + getTensorShape(*columnwise_scale_inv)); + } + if (first_dims.has_value()) { + out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (tensor_offsets.has_value()) { + out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + out_cpp.set_with_gemm_swizzled_scales(this->optimize_for_gemm); + + py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); + py::object out_py = GroupedTensorClass( + "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), + "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), + "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), + "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), + "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = py::none(), + "columnwise_amax"_a = py::none(), "scale"_a = py::none(), + "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), + "last_dims"_a = py::none(), + "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), + "logical_shape"_a = std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + + return {std::move(out_cpp), std::move(out_py)}; +} + std::pair MXFP8Quantizer::convert_and_update_tensor( py::object tensor) const { NVTE_CHECK(detail::IsMXFP8Tensor(tensor.ptr()), "MXFP8Quantizer must output to MXFP8Tensor."); @@ -1240,6 +1602,90 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve return {std::move(out_cpp), std::move(out_py)}; } +std::pair NVFP4Quantizer::create_grouped_tensor( + const size_t num_tensors, const std::vector& logical_shape, const DType dtype, + py::object quantizer, const std::optional& first_dims, + const size_t logical_first_dim, const size_t logical_last_dim) const { + using namespace pybind11::literals; + + const auto tensor_offsets = + build_grouped_tensor_offsets(num_tensors, first_dims, logical_last_dim); + const int64_t total_elements = + static_cast(logical_first_dim) * static_cast(logical_last_dim); + NVTE_CHECK(total_elements % 2 == 0, "NVFP4 data size must be divisible by 2."); + + const auto uint8_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto float_opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + + std::optional rowwise_data; + std::optional columnwise_data; + std::optional rowwise_scale_inv; + std::optional columnwise_scale_inv; + std::optional rowwise_amax; + std::optional columnwise_amax; + const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; + + const int64_t total_data_elements = total_elements / 2; + + if (rowwise_usage) { + rowwise_data = at::empty({total_data_elements}, uint8_opts); + 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); + } + + if (columnwise_usage) { + columnwise_data = at::empty({total_data_elements}, uint8_opts); + const auto scale_shape = get_scale_shape(logical_shape_vec, true); + const int64_t total_scale_elements = static_cast(product(scale_shape)); + columnwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); + columnwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + } + + GroupedTensorWrapper out_cpp(num_tensors, logical_shape, this->get_scaling_mode()); + if (rowwise_usage) { + out_cpp.set_rowwise_data(rowwise_data->data_ptr(), this->dtype, getTensorShape(*rowwise_data)); + out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E4M3, + getTensorShape(*rowwise_scale_inv)); + out_cpp.set_amax(rowwise_amax->data_ptr(), DType::kFloat32, getTensorShape(*rowwise_amax)); + } + if (columnwise_usage) { + out_cpp.set_columnwise_data(columnwise_data->data_ptr(), this->dtype, + getTensorShape(*columnwise_data)); + out_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E4M3, + getTensorShape(*columnwise_scale_inv)); + out_cpp.set_columnwise_amax(columnwise_amax->data_ptr(), DType::kFloat32, + getTensorShape(*columnwise_amax)); + } + if (first_dims.has_value()) { + out_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (tensor_offsets.has_value()) { + out_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + out_cpp.set_with_gemm_swizzled_scales(this->optimize_for_gemm); + + py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); + py::object out_py = GroupedTensorClass( + "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), + "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), + "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), + "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), + "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), + "amax"_a = maybe_tensor_to_py(rowwise_amax), + "columnwise_amax"_a = maybe_tensor_to_py(columnwise_amax), "scale"_a = py::none(), + "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), + "last_dims"_a = py::none(), + "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), + "logical_shape"_a = std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + + return {std::move(out_cpp), std::move(out_py)}; +} + std::pair NVFP4Quantizer::create_unquantized_tensor_with_amax( TensorWrapper& quantized_tensor, DType dtype) { // Construct tensor diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 3f998bb66f..eda5e8fc54 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -170,6 +170,121 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) return ret; } +NVTEScalingMode ScalingModeFromQuantizer(py::handle quantizer) { + auto *quantizer_ptr = quantizer.ptr(); + if (IsMXFP8Quantizers(quantizer_ptr)) { + return NVTE_MXFP8_1D_SCALING; + } + if (IsNVFP4Quantizers(quantizer_ptr)) { + return NVTE_NVFP4_1D_SCALING; + } + if (IsFloat8BlockwiseQuantizers(quantizer_ptr)) { + const int block_scaling_dim = quantizer.attr("block_scaling_dim").cast(); + return (block_scaling_dim == 2) ? NVTE_BLOCK_SCALING_2D : NVTE_BLOCK_SCALING_1D; + } + return NVTE_DELAYED_TENSOR_SCALING; +} + +DType GetTransformerEngineDTypeForScaleInv(py::handle quantizer, at::Tensor scale_inv) { + auto *quantizer_ptr = quantizer.ptr(); + if (IsMXFP8Quantizers(quantizer_ptr)) { + return DType::kFloat8E8M0; + } + if (IsFloat8BlockwiseQuantizers(quantizer_ptr)) { + return DType::kFloat32; + } + if (IsNVFP4Quantizers(quantizer_ptr)) { + return DType::kFloat8E4M3; + } + return GetTransformerEngineDType(scale_inv.scalar_type()); +} + +GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { + // Returns a GroupedTensorWrapper from a PyTorch GroupedTensor. + const auto num_tensors = tensor.attr("num_tensors").cast(); + const auto logical_shape = tensor.attr("logical_shape").cast>(); + py::handle quantizer = py::none(); + DType quantizer_dtype = DType::kNumTypes; + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + if (!tensor.attr("quantizer").is_none()) { + quantizer = tensor.attr("quantizer"); + if (!quantizer.is_none()) { + scaling_mode = ScalingModeFromQuantizer(quantizer); + quantizer_dtype = quantizer.attr("dtype").cast(); + } + } + auto ret = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); + + // Rowwise data + if (!tensor.attr("data").is_none()) { + const auto &data = tensor.attr("data").cast(); + DType data_dtype = + quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; + ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + } + + // Columnwise data + if (!tensor.attr("columnwise_data").is_none()) { + const auto &data = tensor.attr("columnwise_data").cast(); + DType data_dtype = + quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; + ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + } + + // Scale + if (!tensor.attr("scale").is_none()) { + const auto &scale = tensor.attr("scale").cast(); + ret.set_scale(scale.data_ptr(), GetTransformerEngineDType(scale.scalar_type()), + getTensorShape(scale)); + } + + // Amax + if (!tensor.attr("amax").is_none()) { + const auto &amax = tensor.attr("amax").cast(); + ret.set_amax(amax.data_ptr(), GetTransformerEngineDType(amax.scalar_type()), + getTensorShape(amax)); + } + if (!tensor.attr("columnwise_amax").is_none()) { + const auto &amax = tensor.attr("columnwise_amax").cast(); + ret.set_columnwise_amax(amax.data_ptr(), GetTransformerEngineDType(amax.scalar_type()), + getTensorShape(amax)); + } + + // Scale inverse + if (!tensor.attr("scale_inv").is_none()) { + const auto &scale_inv = tensor.attr("scale_inv").cast(); + ret.set_rowwise_scale_inv(scale_inv.data_ptr(), + GetTransformerEngineDTypeForScaleInv(quantizer, scale_inv), + getTensorShape(scale_inv)); + } + if (!tensor.attr("columnwise_scale_inv").is_none()) { + const auto &scale_inv = tensor.attr("columnwise_scale_inv").cast(); + ret.set_columnwise_scale_inv(scale_inv.data_ptr(), + GetTransformerEngineDTypeForScaleInv(quantizer, scale_inv), + getTensorShape(scale_inv)); + } + + // Shape metadata + if (!tensor.attr("first_dims").is_none()) { + const auto &first_dims = tensor.attr("first_dims").cast(); + ret.set_first_dims(first_dims.data_ptr(), GetTransformerEngineDType(first_dims.scalar_type()), + getTensorShape(first_dims)); + } + if (!tensor.attr("last_dims").is_none()) { + const auto &last_dims = tensor.attr("last_dims").cast(); + ret.set_last_dims(last_dims.data_ptr(), GetTransformerEngineDType(last_dims.scalar_type()), + getTensorShape(last_dims)); + } + if (!tensor.attr("tensor_offsets").is_none()) { + const auto &tensor_offsets = tensor.attr("tensor_offsets").cast(); + ret.set_tensor_offsets(tensor_offsets.data_ptr(), + GetTransformerEngineDType(tensor_offsets.scalar_type()), + getTensorShape(tensor_offsets)); + } + + return ret; +} + } // namespace detail } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor.py index dad4d1d0ea..bf5792ffc9 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor.py @@ -52,7 +52,7 @@ class GroupedTensor: def __init__( self, num_tensors: int, - shape: List[Tuple[int, int]], + shape: Optional[List[Tuple[int, int]]] = None, quantizer: Optional[Quantizer] = None, dtype: Optional[torch.dtype] = None, data: Optional[torch.Tensor] = None, @@ -245,6 +245,7 @@ def clear(self) -> None: """ Reset tensor data and clear all buffers. """ + self.shape = None self.data = None self.columnwise_data = None self.scale_inv = None @@ -452,8 +453,7 @@ def make_grouped_tensor( scale_inv_shape = quantizer.get_scale_shape(s, False) scale_elements = math.prod(scale_inv_shape) total_scale_elements += scale_elements - if i < num_tensors - 1: - scale_inv_offsets.append(total_scale_elements) + scale_inv_offsets.append(total_scale_elements) scale_inv = torch.empty(total_scale_elements, dtype=torch.uint8, device=device) if columnwise_usage: @@ -466,8 +466,7 @@ def make_grouped_tensor( scale_inv_shape = quantizer.get_scale_shape(s, False) columnwise_scale_elements = math.prod(scale_inv_shape) total_columnwise_scale_elements += columnwise_scale_elements - if i < num_tensors - 1: - columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) + columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) @@ -477,16 +476,16 @@ def make_grouped_tensor( data = torch.empty(total_elements, dtype=torch.uint8, device=device) # Scale inverse - one per tensor scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) - # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 - scale_inv_offsets = list(range(num_tensors)) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors + scale_inv_offsets = list(range(num_tensors + 1)) if columnwise_usage: # Allocate columnwise data buffer (1D flattened, uint8) columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) # Columnwise scale inverse - one per tensor columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) - # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 - columnwise_scale_inv_offsets = list(range(num_tensors)) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors + columnwise_scale_inv_offsets = list(range(num_tensors + 1)) # Amax buffer for delayed scaling - one per tensor amax = torch.empty(num_tensors, dtype=torch.float32, device=device) @@ -502,8 +501,7 @@ def make_grouped_tensor( for i, s in enumerate(shape): scale_inv_shape = quantizer.get_scale_shape(s, False) total_scale_elements += math.prod(scale_inv_shape) - if i < num_tensors - 1: - scale_inv_offsets.append(total_scale_elements) + 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) @@ -519,8 +517,7 @@ def make_grouped_tensor( for i, s in enumerate(shape): columnwise_scale_inv_shape = quantizer.get_scale_shape(s, True) total_columnwise_scale_elements += math.prod(columnwise_scale_inv_shape) - if i < num_tensors - 1: - columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) + columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) @@ -537,8 +534,7 @@ def make_grouped_tensor( for i, s in enumerate(shape): scale_inv_shape = quantizer.get_scale_shape(s, False) total_scale_elements += math.prod(scale_inv_shape) - if i < num_tensors - 1: - scale_inv_offsets.append(total_scale_elements) + scale_inv_offsets.append(total_scale_elements) scale_inv = torch.empty(total_scale_elements, dtype=torch.float32, device=device) if columnwise_usage: @@ -550,8 +546,7 @@ def make_grouped_tensor( for i, s in enumerate(shape): columnwise_scale_inv_shape = quantizer.get_scale_shape(s, True) total_columnwise_scale_elements += math.prod(columnwise_scale_inv_shape) - if i < num_tensors - 1: - columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) + columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.float32, device=device ) @@ -562,16 +557,16 @@ def make_grouped_tensor( data = torch.empty(total_elements, dtype=torch.uint8, device=device) # Scale inverse - one per tensor scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) - # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 - scale_inv_offsets = list(range(num_tensors)) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors + scale_inv_offsets = list(range(num_tensors + 1)) if columnwise_usage: # Allocate columnwise data buffer (1D flattened, uint8) columnwise_data = torch.empty(total_elements, dtype=torch.uint8, device=device) # Columnwise scale inverse - one per tensor columnwise_scale_inv = torch.empty(num_tensors, dtype=torch.float32, device=device) - # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors-1 - columnwise_scale_inv_offsets = list(range(num_tensors)) + # One scale per tensor, so offsets are simply 0, 1, 2, ..., num_tensors + columnwise_scale_inv_offsets = list(range(num_tensors + 1)) # Scale and amax buffers for current scaling - one per tensor scale = torch.empty(num_tensors, dtype=torch.float32, device=device) @@ -615,6 +610,8 @@ def split_into_quantized_tensors( If quantizer.internal is True, returns QuantizedTensorStorage. Otherwise, returns QuantizedTensor. + This API is NOT graph safe, but can be used for testing & debugging. + TODO(ksivaman): Block cases where any dims are varying. This is needed only to expose the weights as separate parameters. """ @@ -623,6 +620,27 @@ def split_into_quantized_tensors( no_quantization = self.quantizer is None + # if self.shape is None, then trigger D2H copy and get the shape (not graph safe) + if self.shape is None: + first_dims_list = ( + [self.logical_shape[0]] * self.num_tensors + if self.first_dims is None + else self.first_dims.tolist() + ) + last_dims_list = ( + [self.logical_shape[1]] * self.num_tensors + if self.last_dims is None + else self.last_dims.tolist() + ) + shape_list = [] + for i in range(self.num_tensors): + shape_list.append((first_dims_list[i], last_dims_list[i])) + self.shape = shape_list + + # edge case: handle the case where tensor_offsets is given but offsets is not set + if self.offsets is None and self.tensor_offsets is not None: + self.offsets = self.tensor_offsets.tolist() + # Case 1: No quantization - return regular torch tensors if no_quantization: for i in range(self.num_tensors): @@ -667,6 +685,18 @@ def split_into_quantized_tensors( # Case 2: Quantized tensors recipe = self.quantizer._get_compatible_recipe() + # 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 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 + for i in range(self.num_tensors): # Get tensor shape tensor_shape = self.shape[i] @@ -716,10 +746,8 @@ def split_into_quantized_tensors( if self.scale_inv is not None and self.scale_inv_offsets is not None: scale_start = self.scale_inv_offsets[i] - if i < self.num_tensors - 1: - scale_end = self.scale_inv_offsets[i + 1] - else: - scale_end = self.scale_inv.numel() + # for paged stashing, scale_inv should depend on the split offsets + scale_end = self.scale_inv_offsets[i + 1] # Calculate expected scale shape for MXFP8 scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) @@ -730,10 +758,8 @@ def split_into_quantized_tensors( and self.columnwise_scale_inv_offsets is not None ): cscale_start = self.columnwise_scale_inv_offsets[i] - if i < self.num_tensors - 1: - cscale_end = self.columnwise_scale_inv_offsets[i + 1] - else: - cscale_end = self.columnwise_scale_inv.numel() + # for paged stashing, columnwise_scale_inv should depend on the split offsets + cscale_end = self.columnwise_scale_inv_offsets[i + 1] cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( @@ -788,10 +814,8 @@ def split_into_quantized_tensors( if self.scale_inv is not None and self.scale_inv_offsets is not None: scale_start = self.scale_inv_offsets[i] - if i < self.num_tensors - 1: - scale_end = self.scale_inv_offsets[i + 1] - else: - scale_end = self.scale_inv.numel() + # for paged stashing, scale_inv should depend on the split offsets + scale_end = self.scale_inv_offsets[i + 1] # Get scale shape from quantizer scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) @@ -802,10 +826,8 @@ def split_into_quantized_tensors( and self.columnwise_scale_inv_offsets is not None ): cscale_start = self.columnwise_scale_inv_offsets[i] - if i < self.num_tensors - 1: - cscale_end = self.columnwise_scale_inv_offsets[i + 1] - else: - cscale_end = self.columnwise_scale_inv.numel() + # for paged stashing, columnwise_scale_inv should depend on the split offsets + cscale_end = self.columnwise_scale_inv_offsets[i + 1] # Get columnwise scale shape from quantizer cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) @@ -844,10 +866,8 @@ def split_into_quantized_tensors( if self.scale_inv is not None and self.scale_inv_offsets is not None: scale_start = self.scale_inv_offsets[i] - if i < self.num_tensors - 1: - scale_end = self.scale_inv_offsets[i + 1] - else: - scale_end = self.scale_inv.numel() + # for paged stashing, scale_inv should depend on the split offsets + scale_end = self.scale_inv_offsets[i + 1] # Get scale shape from quantizer scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) @@ -858,10 +878,8 @@ def split_into_quantized_tensors( and self.columnwise_scale_inv_offsets is not None ): cscale_start = self.columnwise_scale_inv_offsets[i] - if i < self.num_tensors - 1: - cscale_end = self.columnwise_scale_inv_offsets[i + 1] - else: - cscale_end = self.columnwise_scale_inv.numel() + # for paged stashing, columnwise_scale_inv should depend on the split offsets + cscale_end = self.columnwise_scale_inv_offsets[i + 1] # Get columnwise scale shape from quantizer cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) From a9a9b3ab2f670f65b4e571f76f7608585bfa2c21 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Fri, 27 Feb 2026 16:26:14 +0800 Subject: [PATCH 231/521] [Common][PyTorch] Enhance the fused router and unify the precision (#2633) * add sqrtsoftplus Signed-off-by: Xin Yao * update and add tests Signed-off-by: Xin Yao * switch to fp32 math Signed-off-by: Xin Yao * add more comments Signed-off-by: Xin Yao * fix dtype Signed-off-by: Gao * use CompType instead of hard-coded float Signed-off-by: Xin Yao * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update docstring Signed-off-by: Xin Yao * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Xin Yao Signed-off-by: Xin Yao --------- Signed-off-by: Xin Yao Signed-off-by: Gao Signed-off-by: Xin Yao Co-authored-by: Gao 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_fused_router.py | 83 +++++--- .../common/fused_router/fused_moe_aux_loss.cu | 23 +- .../fused_score_for_moe_aux_loss.cu | 165 +++++++++------ .../fused_topk_with_score_function.cu | 196 ++++++++++-------- .../common/fused_router/utils.h | 130 +++++++----- .../include/transformer_engine/fused_router.h | 10 +- transformer_engine/pytorch/csrc/extensions.h | 23 +- .../pytorch/csrc/extensions/pybind.cpp | 11 +- .../pytorch/csrc/extensions/router.cpp | 62 +++--- transformer_engine/pytorch/router.py | 88 +++++--- 10 files changed, 465 insertions(+), 326 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index f559362d82..64000e109e 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -47,7 +47,7 @@ def group_limited_topk( # Pytorch-based topk softmax/sigmoid -def topk_softmax_sigmoid_pytorch( +def topk_score_function_pytorch( logits: torch.Tensor, topk: int, use_pre_softmax: bool = False, @@ -74,17 +74,20 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None): if score_function == "softmax": if use_pre_softmax: - scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + scores = torch.softmax(logits, dim=-1, dtype=torch.float32) 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) + probs = torch.softmax(scores, dim=-1, dtype=torch.float32) + elif score_function in ("sigmoid", "sqrtsoftplus"): + if score_function == "sigmoid": + scores = torch.sigmoid(logits.float()) + else: + scores = torch.nn.functional.softplus(logits.float()).sqrt() 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) + scores = torch.gather(scores, dim=1, index=top_indices) 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 @@ -94,6 +97,8 @@ def compute_topk(scores, topk, num_groups=None, group_topk=None): if scaling_factor: probs = probs * scaling_factor + probs = probs.type_as(logits) + 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() @@ -107,7 +112,10 @@ def compute_scores_for_aux_loss_pytorch( if score_function == "softmax": scores = torch.softmax(logits, dim=-1, dtype=torch.float32) elif score_function == "sigmoid": - scores = torch.sigmoid(logits) + scores = torch.sigmoid(logits.float()) + scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + elif score_function == "sqrtsoftplus": + scores = torch.nn.functional.softplus(logits.float()).sqrt() scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores else: raise ValueError(f"Invalid score_function: {score_function}") @@ -146,8 +154,9 @@ def run_comparison( enable_bias, ): # Set some parameters - if score_function == "sigmoid": - # Construct the special logits to avoid inf in the sigmoid function + if score_function in ("sigmoid", "sqrtsoftplus"): + # Construct logits with a narrow range to avoid very small activation values, + # which would cause precision loss when adding/subtracting expert bias in float32. 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 @@ -165,8 +174,8 @@ def run_comparison( ) logits = logits.view(num_tokens, num_experts) logits.requires_grad = True - if enable_bias and score_function == "sigmoid": - expert_bias = torch.arange(num_experts, device="cuda") * 0.1 + if enable_bias and score_function in ("sigmoid", "sqrtsoftplus"): + expert_bias = torch.arange(num_experts, device="cuda", dtype=dtype) * 0.1 expert_bias = torch.flip(expert_bias, dims=[0]) expert_bias.requires_grad = True else: @@ -183,7 +192,7 @@ def run_comparison( # Run the original implementation # We do not support the capacity factor case - probs, routing_map = topk_softmax_sigmoid_pytorch( + probs, routing_map = topk_score_function_pytorch( logits=logits, topk=topk, use_pre_softmax=use_pre_softmax, @@ -252,6 +261,37 @@ def test_topk_sigmoid( ) +@pytest.mark.parametrize("dtype", [torch.float32]) +@pytest.mark.parametrize("num_tokens", [2048, 7168, 8992]) +@pytest.mark.parametrize("num_experts", [128, 32]) +@pytest.mark.parametrize("topk", [4, 8]) +@pytest.mark.parametrize("group_topk", [None, 4]) +@pytest.mark.parametrize("scaling_factor", [None, 1.2]) +@pytest.mark.parametrize("enable_bias", [True, False]) +def test_topk_sqrtsoftplus( + dtype, + num_tokens, + num_experts, + topk, + group_topk, + scaling_factor, + enable_bias, +): + num_groups = 8 if group_topk else None + run_comparison( + dtype=dtype, + num_tokens=num_tokens, + num_experts=num_experts, + topk=topk, + use_pre_softmax=False, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function="sqrtsoftplus", + enable_bias=enable_bias, + ) + + @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) @pytest.mark.parametrize("num_experts", [128, 32]) @@ -287,10 +327,10 @@ def test_topk_softmax( @pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) @pytest.mark.parametrize("num_experts", [256, 128, 32]) @pytest.mark.parametrize("topk", [4, 8]) -@pytest.mark.parametrize("score_function", ["softmax", "sigmoid"]) +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function): - if score_function == "sigmoid": - # Construct the special logits to avoid inf in the sigmoid function + if score_function in ("sigmoid", "sqrtsoftplus"): + # Construct logits with a narrow range to avoid very small activation values 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 @@ -396,15 +436,6 @@ def profile_topk_softmax( test_topk_softmax( torch.float32, num_tokens, num_experts, topk, use_pre_softmax, group_topk, scaling_factor ) - - -if __name__ == "__main__": - test_topk_softmax( - dtype=torch.float32, - num_tokens=1024, - num_experts=128, - topk=4, - use_pre_softmax=False, - group_topk=None, - scaling_factor=None, + test_topk_sqrtsoftplus( + torch.float32, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias ) 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 2aa2805fed..8aff85450a 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -16,9 +16,7 @@ #include "utils.h" namespace transformer_engine { - -// Using Double to hanld all the calculations -using CompType = double; +namespace fused_router { template __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, @@ -98,7 +96,7 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, * Section: Compute the aux_loss */ float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; - aux_loss[0] = static_cast(static_cast(intermediate_result) * C_coeff); + aux_loss[0] = static_cast(intermediate_result * C_coeff); Const_buf[0] = C_coeff; } } @@ -154,7 +152,7 @@ __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, * Section: Compute the aux_loss */ float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; - aux_loss[0] = static_cast(static_cast(intermediate_result) * C_coeff); + aux_loss[0] = static_cast(intermediate_result * C_coeff); Const_buf[0] = C_coeff; } } @@ -229,8 +227,8 @@ __global__ void fused_moe_aux_loss_backward_kernel(const float* Const_buf, // Loop: for all positions in each row for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) { float C_coeff = Const_buf[0]; - double tokens_per_expert_i = static_cast(tokens_per_expert[i]); - double grad_aux_loss_value = static_cast(grad_aux_loss[0]); + CompType tokens_per_expert_i = static_cast(tokens_per_expert[i]); + CompType grad_aux_loss_value = static_cast(grad_aux_loss[0]); // Loop: for all rows for (int j = global_warp_id; j < num_rows; j += global_warp_num) { grad_probs[j * num_cols + i] = C_coeff * tokens_per_expert_i * grad_aux_loss_value; @@ -265,6 +263,7 @@ void fused_moe_aux_loss_backward(const Tensor& Const_buf, const Tensor& tokens_p reinterpret_cast(grad_probs.data.dptr), stream););); } +} // namespace fused_router } // namespace transformer_engine void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert, @@ -273,7 +272,7 @@ void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor to NVTETensor Const_buf, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_moe_aux_loss_forward); using namespace transformer_engine; - fused_moe_aux_loss_forward( + 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); @@ -285,8 +284,8 @@ void nvte_fused_moe_aux_loss_backward(const NVTETensor Const_buf, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_moe_aux_loss_backward); using namespace transformer_engine; - fused_moe_aux_loss_backward(*convertNVTETensorCheck(Const_buf), - *convertNVTETensorCheck(tokens_per_expert), num_rows, num_cols, - *convertNVTETensorCheck(grad_aux_loss), - *convertNVTETensorCheck(grad_probs), stream); + fused_router::fused_moe_aux_loss_backward(*convertNVTETensorCheck(Const_buf), + *convertNVTETensorCheck(tokens_per_expert), num_rows, + num_cols, *convertNVTETensorCheck(grad_aux_loss), + *convertNVTETensorCheck(grad_probs), stream); } 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 7540b5c41d..4f405e0a25 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 @@ -14,17 +14,16 @@ #include "utils.h" namespace transformer_engine { +namespace fused_router { template __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logits, int num_tokens, int num_experts, int topk, - int score_function, DataType *scores, + int score_function, float *scores, bool *routing_map, - DataType *intermediate_output) { + CompType *intermediate_output) { /*** * Section: Global Variables/Addresses init - * - Assume the sizeof(DataType) >= sizeof(int), - * So DataType address is assigned firstly to avoid the alignment issue * - Each warp is responsible for one token, and has own shared memory buffer. * Then __syncwarp() is used instead of __syncthreads() */ @@ -33,13 +32,13 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi int warp_id = threadIdx.x / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; extern __shared__ float shmem_scores_for_aux_loss[]; - DataType *logits_buf = reinterpret_cast(shmem_scores_for_aux_loss); - DataType *topk_logits_buf = - reinterpret_cast(logits_buf + num_experts * num_token_per_block); + CompType *logits_buf = reinterpret_cast(shmem_scores_for_aux_loss); + CompType *topk_logits_buf = + reinterpret_cast(logits_buf + num_experts * num_token_per_block); int *topk_indices_buf = reinterpret_cast(topk_logits_buf + topk * num_token_per_block); // The address of buffers on the current warp - DataType *local_logits = logits_buf + warp_id * num_experts; - DataType *topk_logits = topk_logits_buf + warp_id * topk; + CompType *local_logits = logits_buf + warp_id * num_experts; + CompType *topk_logits = topk_logits_buf + warp_id * topk; int *topk_indices = topk_indices_buf + warp_id * topk; /*** @@ -63,12 +62,12 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { routing_map[pos_offset + i] = false; if (score_function == 1) { - intermediate_output[pos_offset + i] = -std::numeric_limits::infinity(); + intermediate_output[pos_offset + i] = -std::numeric_limits::infinity(); } } // Load the logits to shmem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - local_logits[i] = logits[pos_offset + i]; + local_logits[i] = static_cast(logits[pos_offset + i]); } __threadfence_block(); __syncwarp(); @@ -78,11 +77,11 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi * Possible preprocess the scores before the topk operation * - Pre-softmax * - Sigmoid - * - Sigmoid post-processing when topk > 1 + * - Sqrtsoftplus + * - Sigmoid/Sqrtsoftplus post-processing when topk > 1 * This is in-place scores update */ - // score_function == 1 means softmax - if (score_function == 1) { + if (score_function == 1) { // score_function == 1 means softmax // Apply softmax to the logits before the topk apply_softmax_on_float(local_logits, num_experts, lane_id); __syncwarp(); @@ -90,10 +89,7 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { intermediate_output[pos_offset + i] = local_logits[i]; } - } - - // score_function == 0 means sigmoid - if (score_function == 0) { + } else if (score_function == 0) { // score_function == 0 means sigmoid // Apply sigmoid to the logits apply_sigmoid_on_float(local_logits, num_experts, lane_id); __syncwarp(); @@ -101,17 +97,25 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { intermediate_output[pos_offset + i] = local_logits[i]; } + } else if (score_function == 2) { // score_function == 2 means sqrtsoftplus + // First save the original logits for backward (needed for gradient computation) + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { + intermediate_output[pos_offset + i] = local_logits[i]; // Save original logits + } + __syncwarp(); + // Apply sqrtsoftplus to the logits + apply_sqrtsoftplus_on_float(local_logits, num_experts, lane_id); } - __syncwarp(); //Confirm the scores is written to the softmax/sigmoid output + __syncwarp(); //Confirm the scores is written to the output - if (score_function == 0) { + // Sigmoid/Sqrtsoftplus post-processing when topk > 1 + if (score_function == 0 || score_function == 2) { if (topk > 1) { auto sum_logits = warp_reduce_on_shmem(local_logits, num_experts, ReduceFuncType::SUM, lane_id); for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - local_logits[i] = static_cast(static_cast(local_logits[i]) / - (static_cast(sum_logits) + epsilon)); + local_logits[i] /= (sum_logits + epsilon); } } __syncwarp(); @@ -140,12 +144,12 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi template void fused_score_for_moe_aux_loss_forward_kernel_launcher( const DataType *logits, int num_tokens, int num_experts, int topk, int score_function, - DataType *scores, bool *routing_map, DataType *intermediate_output, cudaStream_t stream) { + float *scores, bool *routing_map, CompType *intermediate_output, cudaStream_t stream) { // Meta data for the kernel size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block; - size_t shared_memory_size = num_experts * num_token_per_block * sizeof(DataType) // logits - + topk * num_token_per_block * sizeof(DataType) // topk_logits + size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // logits + + topk * num_token_per_block * sizeof(CompType) // topk_logits + topk * num_token_per_block * sizeof(int); // topk_indices fused_score_for_moe_aux_loss_forward_kernel <<>>( @@ -162,20 +166,19 @@ void fused_score_for_moe_aux_loss_forward(const Tensor &logits, int num_tokens, logits.data.dtype, DataType, fused_score_for_moe_aux_loss_forward_kernel_launcher( reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, - score_function, reinterpret_cast(scores.data.dptr), + score_function, reinterpret_cast(scores.data.dptr), reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), stream);); + reinterpret_cast(intermediate_output.data.dptr), stream);); } template -__global__ void fused_score_for_moe_aux_loss_backward_kernel(const DataType *intermediate_output, - const DataType *grad_scores, +__global__ void fused_score_for_moe_aux_loss_backward_kernel(const CompType *intermediate_output, + const float *grad_scores, int num_tokens, int num_experts, int topk, int score_function, DataType *grad_logits) { /*** * Section: Global Variables/Addresses init - * - Assume the sizeof(DataType) >= sizeof(int), * - Each warp is responsible for one token, and has own shared memory buffer. * Then __syncwarp() is used instead of __syncthreads() */ @@ -184,16 +187,14 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const DataType *int int warp_id = threadIdx.x / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; extern __shared__ float shmem[]; - DataType *grad_scores_buf = reinterpret_cast(shmem); - // To store the output of softmax/sigmoid from the fwd - DataType *act_from_fwd_buf = - reinterpret_cast(grad_scores_buf + num_experts * num_token_per_block); - DataType *comp_buf = - reinterpret_cast(act_from_fwd_buf + num_experts * num_token_per_block); + CompType *grad_scores_buf = reinterpret_cast(shmem); + // To store the output of softmax/sigmoid from fwd, or original logits for sqrtsoftplus + CompType *act_from_fwd_buf = grad_scores_buf + num_experts * num_token_per_block; + CompType *comp_buf = act_from_fwd_buf + num_experts * num_token_per_block; // The address of buffers on the current warp - DataType *local_grad = grad_scores_buf + warp_id * num_experts; - DataType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts; - DataType *local_comp_buf = comp_buf + warp_id * num_experts; + CompType *local_grad = grad_scores_buf + warp_id * num_experts; + CompType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts; + CompType *local_comp_buf = comp_buf + warp_id * num_experts; /*** * Section: Main Loop @@ -227,31 +228,50 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const DataType *int /*** * Section: Backward of ops before the topk * - Pre-softmax bwd - * - Sigmoid Post-processing bwd when topk > 1 + * - Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 * - Sigmoid bwd + * - Sqrtsoftplus bwd * - Write the grad_logits to the global mem */ - // Sigmoid Post-processing bwd when topk > 1 - if (topk > 1 && score_function == 0) { - auto sum_fwd_input = - warp_reduce_on_shmem(local_act_from_fwd, num_experts, ReduceFuncType::SUM, lane_id); - // Put the result of output * grad to the comp_buf + // Sqrtsoftplus: First compute sqrtsoftplus output from original logits + // (needed for both post-processing bwd and activation bwd, compute once here) + // For sqrtsoftplus, intermediate_output stores original logits + if (score_function == 2) { + // Copy original logits to local_comp_buf and apply sqrtsoftplus in-place for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - local_comp_buf[i] = local_grad[i] * local_act_from_fwd[i]; + local_comp_buf[i] = local_act_from_fwd[i]; } __syncwarp(); - auto sum_Output_x_Grad = - warp_reduce_on_shmem(local_comp_buf, num_experts, ReduceFuncType::SUM, lane_id); + apply_sqrtsoftplus_on_float(local_comp_buf, num_experts, lane_id); + __syncwarp(); + } + + // Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 (normalization backward) + if (topk > 1 && (score_function == 0 || score_function == 2)) { + // Select the correct activation output buffer: + // - Sigmoid: local_act_from_fwd already contains sigmoid output + // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above + CompType *act_output = (score_function == 0) ? local_act_from_fwd : local_comp_buf; + + auto sum_fwd_input = + warp_reduce_on_shmem(act_output, num_experts, ReduceFuncType::SUM, lane_id); + // Compute sum of output * grad using registers + CompType local_sum_Output_x_Grad = 0.0; + 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; // In-place update for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - local_grad[i] = - static_cast(local_grad[i]) / (static_cast(sum_fwd_input) + epsilon) - - static_cast(sum_Output_x_Grad) / - ((static_cast(sum_fwd_input) + epsilon) * - (static_cast(sum_fwd_input) + epsilon)); + local_grad[i] = local_grad[i] / (sum_fwd_input + epsilon) - + sum_Output_x_Grad / ((sum_fwd_input + epsilon) * (sum_fwd_input + epsilon)); } + __syncwarp(); } - __syncwarp(); // Pre-softmax bwd if (score_function == 1) { @@ -264,9 +284,17 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const DataType *int apply_sigmoid_bwd_on_float(local_grad, local_act_from_fwd, num_experts, lane_id); __syncwarp(); } + // Sqrtsoftplus bwd + // For sqrtsoftplus, local_comp_buf already contains sqrtsoftplus output computed earlier + // Now compute gradient: dy/dx = sigmoid(x) / (2 * y) + if (score_function == 2) { + apply_sqrtsoftplus_bwd_on_float(local_grad, local_comp_buf, local_act_from_fwd, num_experts, + lane_id); + __syncwarp(); + } // Write the grad_logits to the global mem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - grad_logits[pos_offset + i] = local_grad[i]; + grad_logits[pos_offset + i] = static_cast(local_grad[i]); } __syncwarp(); } @@ -274,15 +302,15 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const DataType *int template void fused_score_for_moe_aux_loss_backward_kernel_launcher( - const DataType *intermediate_output, const DataType *grad_scores, int num_tokens, - int num_experts, int topk, int score_function, DataType *grad_logits, cudaStream_t stream) { + const CompType *intermediate_output, const float *grad_scores, int num_tokens, int num_experts, + int topk, int score_function, DataType *grad_logits, cudaStream_t stream) { // Meta data for the kernel size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block; - size_t shared_memory_size = num_experts * num_token_per_block * sizeof(DataType) // grad_scores + size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // grad_scores + - num_experts * num_token_per_block * sizeof(DataType) // act_from_fwd - + num_experts * num_token_per_block * sizeof(DataType); // comp_buf + num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd + + num_experts * num_token_per_block * sizeof(CompType); // comp_buf fused_score_for_moe_aux_loss_backward_kernel <<>>( intermediate_output, grad_scores, num_tokens, num_experts, topk, score_function, @@ -295,13 +323,14 @@ void fused_score_for_moe_aux_loss_backward(const Tensor &intermediate_output, int num_experts, int topk, int score_function, Tensor &grad_logits, cudaStream_t stream) { TE_ROUTER_PROBS_TYPE_SWITCH_ALL( - grad_scores.data.dtype, DataType, + grad_logits.data.dtype, DataType, fused_score_for_moe_aux_loss_backward_kernel_launcher( - reinterpret_cast(intermediate_output.data.dptr), - reinterpret_cast(grad_scores.data.dptr), num_tokens, num_experts, topk, + reinterpret_cast(intermediate_output.data.dptr), + reinterpret_cast(grad_scores.data.dptr), num_tokens, num_experts, topk, score_function, reinterpret_cast(grad_logits.data.dptr), stream);); } +} // namespace fused_router } // namespace transformer_engine void nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_tokens, @@ -311,10 +340,10 @@ void nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_ cudaStream_t stream) { NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_forward); using namespace transformer_engine; - fused_score_for_moe_aux_loss_forward(*convertNVTETensorCheck(logits), num_tokens, num_experts, - topk, score_function, *convertNVTETensorCheck(scores), - *convertNVTETensorCheck(routing_map), - *convertNVTETensorCheck(intermediate_output), stream); + fused_router::fused_score_for_moe_aux_loss_forward( + *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, score_function, + *convertNVTETensorCheck(scores), *convertNVTETensorCheck(routing_map), + *convertNVTETensorCheck(intermediate_output), stream); } void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_output, @@ -323,7 +352,7 @@ void nvte_fused_score_for_moe_aux_loss_backward(const NVTETensor intermediate_ou NVTETensor grad_logits, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_score_for_moe_aux_loss_backward); using namespace transformer_engine; - fused_score_for_moe_aux_loss_backward( + fused_router::fused_score_for_moe_aux_loss_backward( *convertNVTETensorCheck(intermediate_output), *convertNVTETensorCheck(grad_scores), num_tokens, num_experts, topk, score_function, *convertNVTETensorCheck(grad_logits), stream); } 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 2719c68c97..a9e680f06e 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 @@ -14,17 +14,16 @@ #include "utils.h" namespace transformer_engine { +namespace fused_router { template __global__ void fused_topk_with_score_function_forward_kernel( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, const BiasType *expert_bias, DataType *probs, bool *routing_map, - DataType *intermediate_output) { + CompType *intermediate_output) { /*** * Section: Global Variables/Addresses init - * - Assume the sizeof(DataType) >= sizeof(int), - * So DataType address is assigned firstly to avoid the alignment issue * - Each warp is responsible for one token, and has own shared memory buffer. * Then __syncwarp() is used instead of __syncthreads() */ @@ -33,24 +32,22 @@ __global__ void fused_topk_with_score_function_forward_kernel( int warp_id = threadIdx.x / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; extern __shared__ float shmem[]; - DataType *scores_buf = reinterpret_cast(shmem); - DataType *topk_scores_buf = - reinterpret_cast(scores_buf + num_experts * num_token_per_block); - DataType *group_scores_buf = nullptr, *masked_scores_buf = nullptr; + CompType *scores_buf = reinterpret_cast(shmem); + CompType *topk_scores_buf = scores_buf + num_experts * num_token_per_block; + CompType *group_scores_buf = nullptr, *masked_scores_buf = nullptr; int *topk_indices_buf = nullptr; if (group_topk > 0) { - masked_scores_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block); - group_scores_buf = - reinterpret_cast(masked_scores_buf + num_experts * num_token_per_block); + masked_scores_buf = topk_scores_buf + topk * num_token_per_block; + group_scores_buf = masked_scores_buf + num_experts * num_token_per_block; topk_indices_buf = reinterpret_cast(group_scores_buf + num_groups * num_token_per_block); } else { topk_indices_buf = reinterpret_cast(topk_scores_buf + topk * num_token_per_block); } // The address of buffers on the current warp - DataType *scores = scores_buf + warp_id * num_experts; - DataType *topk_scores = topk_scores_buf + warp_id * topk; - DataType *masked_scores = masked_scores_buf + warp_id * num_experts; - DataType *group_scores = group_scores_buf + warp_id * num_groups; + CompType *scores = scores_buf + warp_id * num_experts; + CompType *topk_scores = topk_scores_buf + warp_id * topk; + CompType *masked_scores = masked_scores_buf + warp_id * num_experts; + CompType *group_scores = group_scores_buf + warp_id * num_groups; int *topk_indices = topk_indices_buf + warp_id * topk; /*** @@ -72,10 +69,10 @@ __global__ void fused_topk_with_score_function_forward_kernel( int pos_offset = token_offset_cur_warp * num_experts; // Clear the probs/routing_map (num_experts) for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - probs[pos_offset + i] = 0.0f; + probs[pos_offset + i] = 0.0; routing_map[pos_offset + i] = false; if (score_function == 1) { - intermediate_output[pos_offset + i] = -std::numeric_limits::infinity(); + intermediate_output[pos_offset + i] = -std::numeric_limits::infinity(); } } // Load the logits to shmem @@ -85,7 +82,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( // If group_topk > 0, init the masked_scores to -inf if (group_topk > 0) { for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - masked_scores[i] = -std::numeric_limits::infinity(); + masked_scores[i] = -std::numeric_limits::infinity(); } } __threadfence_block(); @@ -96,11 +93,11 @@ __global__ void fused_topk_with_score_function_forward_kernel( * Possible preprocess the scores before the topk operation * - Pre-softmax * - Sigmoid + * - Sqrtsoftplus * - Expert bias * This is in-place scores update */ - // score_function == 1 means softmax - if (use_pre_softmax && score_function == 1) { + if (use_pre_softmax && score_function == 1) { // score_function == 1 means softmax // Apply softmax to the logits before the topk apply_softmax_on_float(scores, num_experts, lane_id); __syncwarp(); @@ -108,10 +105,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { intermediate_output[pos_offset + i] = scores[i]; } - } - - // score_function == 0 means sigmoid - if (score_function == 0) { + } else if (score_function == 0) { // score_function == 0 means sigmoid // Apply sigmoid to the logits apply_sigmoid_on_float(scores, num_experts, lane_id); __syncwarp(); @@ -119,18 +113,25 @@ __global__ void fused_topk_with_score_function_forward_kernel( for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { intermediate_output[pos_offset + i] = scores[i]; } + } else if (score_function == 2) { // score_function == 2 means sqrtsoftplus + // First save the original logits for backward (needed for sqrtsoftplus gradient computation) + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { + intermediate_output[pos_offset + i] = scores[i]; // Save original logits + } + __syncwarp(); + // Apply sqrtsoftplus to the logits + apply_sqrtsoftplus_on_float(scores, num_experts, lane_id); } - __syncwarp(); //Confirm the scores is written to the softmax/sigmoid output + __syncwarp(); //Confirm the scores is written to the output - // Expert bias is only used at the sigmoid case - if (expert_bias && score_function == 0) { + // Expert bias is only used at the sigmoid/sqrtsoftplus case + if (expert_bias && (score_function == 0 || score_function == 2)) { for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - scores[i] = static_cast(static_cast(scores[i]) + - static_cast(expert_bias[i])); + scores[i] += static_cast(expert_bias[i]); } + __syncwarp(); } - __syncwarp(); /*** * Section: Topk @@ -140,7 +141,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( * - topk with expert bias */ // Topk on the scores - // The bias is not empty only happens at the sigmod case + // The bias being not empty happens at the sigmoid/sqrtsoftplus case if (group_topk > 0) { int group_size = num_experts / num_groups; // Top2 @@ -155,7 +156,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( __syncwarp(); // Compute the group score if (lane_id == 0) { - DataType tmp = 0.0f; + CompType tmp = 0.0; for (int j = 0; j < topk / group_topk; j++) { tmp = tmp + topk_scores[j]; } @@ -194,17 +195,16 @@ __global__ void fused_topk_with_score_function_forward_kernel( * Possible postprocess the scores after the topk operation * - Revert Expert bias * - Softmax - * - Sigmoid post-processing when topk > 1 + * - Sigmoid/Sqrtsoftplus post-processing when topk > 1 * - Write the result with scaling_factor */ // Revert Expert bias from the topk scores - if (expert_bias && score_function == 0) { + if (expert_bias && (score_function == 0 || score_function == 2)) { for (int i = lane_id; i < topk; i += kThreadsPerWarp) { - topk_scores[i] = - static_cast(topk_scores[i]) - static_cast(expert_bias[topk_indices[i]]); + topk_scores[i] = topk_scores[i] - static_cast(expert_bias[topk_indices[i]]); } + __syncwarp(); } - __syncwarp(); // score_function == 1 means softmax if (!use_pre_softmax && score_function == 1) { @@ -215,14 +215,15 @@ __global__ void fused_topk_with_score_function_forward_kernel( for (int i = lane_id; i < topk; i += kThreadsPerWarp) { intermediate_output[pos_offset + topk_indices[i]] = topk_scores[i]; } + __syncwarp(); } - // score_function == 0 means sigmoid - if (score_function == 0) { + // Sigmoid/Sqrtsoftplus post-processing when topk > 1 + if (score_function == 0 || score_function == 2) { if (topk > 1) { - double sum_scores = warp_reduce_on_shmem(topk_scores, topk, ReduceFuncType::SUM, lane_id); + CompType sum_scores = warp_reduce_on_shmem(topk_scores, topk, ReduceFuncType::SUM, lane_id); for (int i = lane_id; i < topk; i += kThreadsPerWarp) { - topk_scores[i] = static_cast(topk_scores[i]) / (sum_scores + epsilon); + topk_scores[i] = topk_scores[i] / (sum_scores + epsilon); } } __syncwarp(); @@ -231,7 +232,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( // Write the probs/routing_map to the output tensor for (int i = lane_id; i < topk; i += kThreadsPerWarp) { routing_map[pos_offset + topk_indices[i]] = true; - probs[pos_offset + topk_indices[i]] = scaling_factor * static_cast(topk_scores[i]); + probs[pos_offset + topk_indices[i]] = scaling_factor * topk_scores[i]; } __threadfence_block(); __syncwarp(); @@ -242,16 +243,16 @@ template void fused_topk_with_score_function_forward_kernel_launcher( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, - const BiasType *expert_bias, DataType *probs, bool *routing_map, DataType *intermediate_output, + const BiasType *expert_bias, DataType *probs, bool *routing_map, CompType *intermediate_output, cudaStream_t stream) { size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block; - size_t shared_memory_size = num_experts * num_token_per_block * sizeof(DataType) // scores - + topk * num_token_per_block * sizeof(DataType) // topk_scores + size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // scores + + topk * num_token_per_block * sizeof(CompType) // topk_scores + topk * num_token_per_block * sizeof(int); // topk_indices if (group_topk > 0) { - shared_memory_size += num_groups * num_token_per_block * sizeof(DataType); // group_scores - shared_memory_size += num_experts * num_token_per_block * sizeof(DataType); // maksed_scores + shared_memory_size += num_groups * num_token_per_block * sizeof(CompType); // group_scores + shared_memory_size += num_experts * num_token_per_block * sizeof(CompType); // maksed_scores } fused_topk_with_score_function_forward_kernel <<>>( @@ -276,13 +277,13 @@ void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, reinterpret_cast(expert_bias.data.dptr), reinterpret_cast(probs.data.dptr), reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), stream););); + reinterpret_cast(intermediate_output.data.dptr), stream););); } template __global__ void fused_topk_with_score_function_backward_kernel( // Inputs tensor - const bool *routing_map, const DataType *intermediate_output, const DataType *grad_probs, + const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs, // Other parameters int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor, int score_function, @@ -290,7 +291,6 @@ __global__ void fused_topk_with_score_function_backward_kernel( DataType *grad_logits) { /*** * Section: Global Variables/Addresses init - * - Assume the sizeof(DataType) >= sizeof(int), * - Each warp is responsible for one token, and has own shared memory buffer. * Then __syncwarp() is used instead of __syncthreads() */ @@ -299,18 +299,16 @@ __global__ void fused_topk_with_score_function_backward_kernel( int warp_id = threadIdx.x / kThreadsPerWarp; int lane_id = threadIdx.x % kThreadsPerWarp; extern __shared__ float shmem[]; - DataType *grad_probs_buf = reinterpret_cast(shmem); - // To store the output of softmax/sigmoid from the fwd - DataType *act_from_fwd_buf = - reinterpret_cast(grad_probs_buf + num_experts * num_token_per_block); - DataType *comp_buf = - reinterpret_cast(act_from_fwd_buf + num_experts * num_token_per_block); + CompType *grad_probs_buf = reinterpret_cast(shmem); + // To store the output of softmax/sigmoid from fwd, or original logits for sqrtsoftplus + CompType *act_from_fwd_buf = grad_probs_buf + num_experts * num_token_per_block; + CompType *comp_buf = act_from_fwd_buf + num_experts * num_token_per_block; // To store the routing_map from the fwd bool *routing_map_buf = reinterpret_cast(comp_buf + num_experts * num_token_per_block); // The address of buffers on the current warp - DataType *local_grad = grad_probs_buf + warp_id * num_experts; - DataType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts; - DataType *local_comp_buf = comp_buf + warp_id * num_experts; + CompType *local_grad = grad_probs_buf + warp_id * num_experts; + CompType *local_act_from_fwd = act_from_fwd_buf + warp_id * num_experts; + CompType *local_comp_buf = comp_buf + warp_id * num_experts; bool *local_routing_map = routing_map_buf + warp_id * num_experts; /*** @@ -346,48 +344,68 @@ __global__ void fused_topk_with_score_function_backward_kernel( /*** * Section: Backward of ops after the topk * - Backward of the used scaling_factor - * - Sigmoid Post-processing bwd when topk > 1 + * - Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 * - Softmax bwd if use_pre_softmax is false */ // Backward of the used scaling_factor // In-place update for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { if (local_routing_map[i]) { - local_grad[i] = static_cast(local_grad[i]) * scaling_factor; + local_grad[i] = local_grad[i] * scaling_factor; } } __syncwarp(); - // Sigmoid Post-processing bwd when topk > 1 - if (topk > 1 && score_function == 0) { - double sum_fwd_input = masked_warp_reduce_on_shmem( - /*data ptr = */ local_act_from_fwd, - /*mask ptr = */ local_routing_map, - /*data size = */ num_experts, - /*reduce func = */ ReduceFuncType::SUM, lane_id); - // Put the result of output * grad to the comp_buf + + // Sqrtsoftplus: First compute sqrtsoftplus output from original logits + // (needed for both post-processing bwd and activation bwd, compute once here) + // For sqrtsoftplus, intermediate_output stores original logits + if (score_function == 2) { + // Copy original logits to local_comp_buf and apply sqrtsoftplus in-place for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - local_comp_buf[i] = (local_routing_map[i] ? static_cast(local_grad[i]) * - static_cast(local_act_from_fwd[i]) - : 0.0f); + local_comp_buf[i] = local_act_from_fwd[i]; } __syncwarp(); - double sum_Output_x_Grad = masked_warp_reduce_on_shmem( - /*data ptr = */ local_comp_buf, + apply_sqrtsoftplus_on_float(local_comp_buf, num_experts, lane_id); + __syncwarp(); + } + + // Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 (normalization backward) + if (topk > 1 && (score_function == 0 || score_function == 2)) { + // Select the correct activation output buffer: + // - Sigmoid: local_act_from_fwd already contains sigmoid output + // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above + CompType *act_output = (score_function == 0) ? local_act_from_fwd : local_comp_buf; + + CompType sum_fwd_input = masked_warp_reduce_on_shmem( + /*data ptr = */ act_output, /*mask ptr = */ local_routing_map, /*data size = */ num_experts, /*reduce func = */ ReduceFuncType::SUM, lane_id); + // Compute sum of output * grad using registers + CompType local_sum_Output_x_Grad = 0.0; + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { + if (local_routing_map[i]) { + 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; // In-place update for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { if (local_routing_map[i]) { local_grad[i] = - static_cast(local_grad[i]) / (sum_fwd_input + epsilon) - + local_grad[i] / (sum_fwd_input + epsilon) - sum_Output_x_Grad / ((sum_fwd_input + epsilon) * (sum_fwd_input + epsilon)); } else { - local_grad[i] = 0.0f; + local_grad[i] = 0.0; } } + __syncwarp(); } - __syncwarp(); + // Softmax bwd if use_pre_softmax is false if (!use_pre_softmax && score_function == 1) { apply_softmax_bwd_on_float(local_grad, local_act_from_fwd, local_comp_buf, local_routing_map, @@ -401,7 +419,7 @@ __global__ void fused_topk_with_score_function_backward_kernel( */ for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { if (!local_routing_map[i]) { - local_grad[i] = 0.0f; + local_grad[i] = 0.0; } } __syncwarp(); @@ -410,6 +428,7 @@ __global__ void fused_topk_with_score_function_backward_kernel( * Section: Backward of ops before the topk * - Pre-softmax bwd * - Sigmoid bwd + * - Sqrtsoftplus bwd * - Write the grad_logits to the global mem */ // Pre-softmax bwd @@ -423,6 +442,14 @@ __global__ void fused_topk_with_score_function_backward_kernel( apply_sigmoid_bwd_on_float(local_grad, local_act_from_fwd, num_experts, lane_id); __syncwarp(); } + // Sqrtsoftplus bwd + // For sqrtsoftplus, local_comp_buf already contains sqrtsoftplus output computed earlier + // Now compute gradient: dy/dx = sigmoid(x) / (2 * y) + if (score_function == 2) { + apply_sqrtsoftplus_bwd_on_float(local_grad, local_comp_buf, local_act_from_fwd, num_experts, + lane_id); + __syncwarp(); + } // Write the grad_logits to the global mem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { grad_logits[pos_offset + i] = local_grad[i]; @@ -433,16 +460,16 @@ __global__ void fused_topk_with_score_function_backward_kernel( template void fused_topk_with_score_function_backward_kernel_launcher( - const bool *routing_map, const DataType *intermediate_output, const DataType *grad_probs, + const bool *routing_map, const CompType *intermediate_output, const DataType *grad_probs, int num_tokens, int num_experts, int topk, bool use_pre_softmax, float scaling_factor, int score_function, DataType *grad_logits, cudaStream_t stream) { // Meta data for the kernel size_t num_token_per_block = kThreadsPerBlock / kThreadsPerWarp; size_t grid_size = (num_tokens + num_token_per_block - 1) / num_token_per_block; - size_t shared_memory_size = num_experts * num_token_per_block * sizeof(DataType) // grad_probs + size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // grad_probs + - num_experts * num_token_per_block * sizeof(DataType) // act_from_fwd - + num_experts * num_token_per_block * sizeof(DataType) // comp_buf + num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd + + num_experts * num_token_per_block * sizeof(CompType) // comp_buf + num_experts * num_token_per_block * sizeof(bool); // routing_map fused_topk_with_score_function_backward_kernel <<>>( @@ -461,12 +488,13 @@ void fused_topk_with_score_function_backward(const Tensor &routing_map, grad_logits.data.dtype, DataType, fused_topk_with_score_function_backward_kernel_launcher( reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), + reinterpret_cast(intermediate_output.data.dptr), reinterpret_cast(grad_probs.data.dptr), num_tokens, num_experts, topk, use_pre_softmax, scaling_factor, score_function, reinterpret_cast(grad_logits.data.dptr), stream);); } +} // namespace fused_router } // namespace transformer_engine void nvte_fused_topk_with_score_function_forward( @@ -476,7 +504,7 @@ void nvte_fused_topk_with_score_function_forward( NVTETensor intermediate_output, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_topk_with_score_function_forward); using namespace transformer_engine; - fused_topk_with_score_function_forward( + fused_router::fused_topk_with_score_function_forward( *convertNVTETensorCheck(logits), num_tokens, num_experts, topk, static_cast(use_pre_softmax), num_groups, group_topk, scaling_factor, score_function, *convertNVTETensorCheck(expert_bias), *convertNVTETensorCheck(probs), @@ -491,7 +519,7 @@ void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map, NVTETensor grad_logits, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_topk_with_score_function_backward); using namespace transformer_engine; - fused_topk_with_score_function_backward( + fused_router::fused_topk_with_score_function_backward( *convertNVTETensorCheck(routing_map), *convertNVTETensorCheck(intermediate_output), *convertNVTETensorCheck(grad_probs), num_tokens, num_experts, topk, static_cast(use_pre_softmax), scaling_factor, score_function, diff --git a/transformer_engine/common/fused_router/utils.h b/transformer_engine/common/fused_router/utils.h index 669748c1ad..60e731d990 100644 --- a/transformer_engine/common/fused_router/utils.h +++ b/transformer_engine/common/fused_router/utils.h @@ -10,6 +10,13 @@ #include "transformer_engine/transformer_engine.h" namespace transformer_engine { +namespace fused_router { + +// Using FP32 to handle all the calculations. +// Currently, only FP32 is supported because +// 1. The score functions (sigmoid, softmax, sqrtsoftplus) are implemented in FP32. +// 2. The intermediate buffer is initialized in FP32. +using CompType = float; constexpr size_t kThreadsPerWarp = 32; constexpr int kThreadsPerBlock = @@ -35,19 +42,19 @@ template __device__ inline T warp_reduce_on_shmem(T *data_ptr, int data_size, ReduceFuncType type, int lane_id) { T (*reduce_func)(T, T); - double default_val = 0; + CompType default_val = 0.0; if (type == ReduceFuncType::SUM) { reduce_func = sum; - default_val = 0; + default_val = 0.0; } else if (type == ReduceFuncType::MAX) { reduce_func = max; - default_val = -std::numeric_limits::infinity(); + default_val = -std::numeric_limits::infinity(); } // Some value is hanlded in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread - double val = lane_id < data_size ? static_cast(data_ptr[lane_id]) : default_val; + CompType val = lane_id < data_size ? data_ptr[lane_id] : default_val; for (int i = lane_id + kThreadsPerWarp; i < data_size; i += kThreadsPerWarp) { val = reduce_func(val, data_ptr[i]); } @@ -62,31 +69,23 @@ __device__ inline T warp_reduce_on_shmem(T *data_ptr, int data_size, ReduceFuncT return T(val); } -template -__device__ inline void apply_sigmoid_on_float(DataType *scores, int data_size, int lane_id) { - for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { - scores[i] = static_cast(1.0f / (1.0f + exp(-static_cast(scores[i])))); - } -} - template __device__ inline T masked_warp_reduce_on_shmem(T *data_ptr, bool *mask, int data_size, ReduceFuncType type, int lane_id) { T (*reduce_func)(T, T); - double default_val = 0; + CompType default_val = 0.0; if (type == ReduceFuncType::SUM) { reduce_func = sum; - default_val = 0; + default_val = 0.0; } else if (type == ReduceFuncType::MAX) { reduce_func = max; - default_val = -std::numeric_limits::infinity(); + default_val = -std::numeric_limits::infinity(); } // Some value is hanlded in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread - double val = - lane_id < data_size && mask[lane_id] ? static_cast(data_ptr[lane_id]) : default_val; + CompType val = lane_id < data_size && mask[lane_id] ? data_ptr[lane_id] : default_val; for (int i = lane_id + kThreadsPerWarp; i < data_size; i += kThreadsPerWarp) { if (mask[i]) { val = reduce_func(val, data_ptr[i]); @@ -103,28 +102,70 @@ __device__ inline T masked_warp_reduce_on_shmem(T *data_ptr, bool *mask, int dat return T(val); } -template -__device__ inline void apply_sigmoid_bwd_on_float(DataType *grad, DataType *fwd_output, - int data_size, int lane_id) { +__device__ inline void apply_sigmoid_on_float(float *scores, int data_size, int lane_id) { for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { - grad[i] = static_cast(grad[i]) * static_cast(fwd_output[i]) * - (1 - static_cast(fwd_output[i])); + scores[i] = 1.0f / (1.0f + expf(-scores[i])); } } -template -__device__ inline void apply_softmax_bwd_on_float(DataType *grad, DataType *fwd_output, - DataType *comp_buf, bool *mask, int data_size, +__device__ inline void apply_sigmoid_bwd_on_float(float *grad, float *fwd_output, int data_size, int lane_id) { + for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { + grad[i] = grad[i] * fwd_output[i] * (1.0f - fwd_output[i]); + } +} + +// sqrtsoftplus: y = sqrt(softplus(x)) = sqrt(log(1 + exp(x))) +__device__ inline void apply_sqrtsoftplus_on_float(float *scores, int data_size, int lane_id) { + for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { + float x = scores[i]; + // softplus(x) = log(1 + exp(x)), numerically stable version + // Matches PyTorch's Softplus(beta=1.0, threshold=20.0) + float softplus_val; + if (x > 20.0f) { + softplus_val = x; // for large x, softplus(x) ≈ x + } else { + softplus_val = log1pf(expf(x)); + } + scores[i] = sqrtf(softplus_val); + } +} + +// sqrtsoftplus backward: +// y = sqrt(softplus(x)) +// Matches PyTorch's Softplus(beta=1.0, threshold=20.0) +// We need the original logits (x) to compute the gradient +__device__ inline void apply_sqrtsoftplus_bwd_on_float(float *grad, float *fwd_output, + float *logits_buf, int data_size, + int lane_id) { + for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { + float x = logits_buf[i]; // original logit + float y = fwd_output[i]; // sqrtsoftplus output + float dy_dx; + if (x > 20.0f) { + // When softplus(x) = x, y = sqrt(x), dy/dx = 1/(2*y) + dy_dx = 1.0f / (2.0f * y + epsilon); + } else { + // When softplus(x) = log(1+exp(x)), dy/dx = sigmoid(x) / (2*y) + // where sigmoid(x) = 1 / (1 + exp(-x)) + float sigmoid_x = 1.0f / (1.0f + expf(-x)); + dy_dx = sigmoid_x / (2.0f * y + epsilon); + } + grad[i] = grad[i] * dy_dx; + } +} + +__device__ inline void apply_softmax_bwd_on_float(float *grad, float *fwd_output, float *comp_buf, + bool *mask, int data_size, int lane_id) { // Put the result of output * grad to the comp_buf for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { if (mask) { if (mask[i]) - comp_buf[i] = static_cast(grad[i]) * static_cast(fwd_output[i]); + comp_buf[i] = grad[i] * fwd_output[i]; else comp_buf[i] = 0.0f; } else { - comp_buf[i] = static_cast(grad[i]) * static_cast(fwd_output[i]); + comp_buf[i] = grad[i] * fwd_output[i]; } } __syncwarp(); @@ -136,40 +177,34 @@ __device__ inline void apply_softmax_bwd_on_float(DataType *grad, DataType *fwd_ for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { if (mask) { if (mask[i]) - grad[i] = - static_cast(fwd_output[i]) * (static_cast(grad[i]) - sum_Output_x_Grad); + grad[i] = fwd_output[i] * (grad[i] - sum_Output_x_Grad); else grad[i] = 0.0f; } else { - grad[i] = - static_cast(fwd_output[i]) * (static_cast(grad[i]) - sum_Output_x_Grad); + grad[i] = fwd_output[i] * (grad[i] - sum_Output_x_Grad); } } } -template -__device__ inline void apply_softmax_on_float(DataType *scores, int data_size, int lane_id) { +__device__ inline void apply_softmax_on_float(float *scores, int data_size, int lane_id) { // 1. compute the max of value - float max_val = - static_cast(warp_reduce_on_shmem(scores, data_size, ReduceFuncType::MAX, lane_id)); + float max_val = warp_reduce_on_shmem(scores, data_size, ReduceFuncType::MAX, lane_id); // 2. value -> exp_value for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { - scores[i] = static_cast(exp(static_cast(scores[i]) - max_val)); + scores[i] = expf(scores[i] - max_val); } __syncwarp(); // 3. compute the sum of exp_value - float sum_val = - static_cast(warp_reduce_on_shmem(scores, data_size, ReduceFuncType::SUM, lane_id)); + float sum_val = warp_reduce_on_shmem(scores, data_size, ReduceFuncType::SUM, lane_id); // 4. update the softmax value for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { - scores[i] = static_cast(scores[i]) / sum_val; + scores[i] = scores[i] / sum_val; } __syncwarp(); } -template -__device__ inline void naive_topk_and_mask(T *scores, int data_size, int topk, int *topk_indices, - T *topk_scores, int lane_id) { +__device__ inline void naive_topk_and_mask(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; @@ -183,16 +218,15 @@ __device__ inline void naive_topk_and_mask(T *scores, int data_size, int topk, i // After looping topk times, the topk_indices will be the topk indices for (int k = 0; k < topk; k++) { // Find the max value and its index - double val = (lane_id < data_size && !is_masked(k, lane_id)) - ? static_cast(scores[lane_id]) - : -std::numeric_limits::infinity(); + CompType val = (lane_id < data_size && !is_masked(k, lane_id)) + ? scores[lane_id] + : -std::numeric_limits::infinity(); int index = (lane_id < data_size) ? lane_id : 0; // Some value is hanlded in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread for (int i = lane_id + kThreadsPerWarp; i < data_size; i += kThreadsPerWarp) { - double cur_val = (is_masked(k, i)) ? -std::numeric_limits::infinity() - : static_cast(scores[i]); + CompType cur_val = (is_masked(k, i)) ? -std::numeric_limits::infinity() : scores[i]; if (cur_val > val) { val = cur_val; index = i; @@ -257,5 +291,7 @@ __device__ inline void naive_topk_and_mask(T *scores, int data_size, int topk, i default: \ NVTE_ERROR("Invalid type."); \ } +} // namespace fused_router } // namespace transformer_engine -#endif + +#endif // TRANSFORMER_ENGINE_FUSED_ROUTER_UTILS_H_ diff --git a/transformer_engine/common/include/transformer_engine/fused_router.h b/transformer_engine/common/include/transformer_engine/fused_router.h index 1f026a703d..794880d324 100644 --- a/transformer_engine/common/include/transformer_engine/fused_router.h +++ b/transformer_engine/common/include/transformer_engine/fused_router.h @@ -23,8 +23,8 @@ extern "C" { * \param[in] num_groups Number of groups in grouped topk. * \param[in] group_topk Grouped topk value. * \param[in] scaling_factor Scaling factor. - * \param[in] score_function Score function, 0: sigmoid, 1: softmax. - * \param[in] expert_bias Expert bias. (Only used at the sigmoid case) + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. + * \param[in] expert_bias Expert bias. (Used at the sigmoid/sqrtsoftplus cases) * \param[out] probs Output tensor for probabilities. * \param[out] routing_map Output tensor for routing map. * \param[out] intermediate_output Output tensor for intermediate output. (Softmax/sigmoid output) @@ -46,7 +46,7 @@ void nvte_fused_topk_with_score_function_forward( * \param[in] topk Topk value. * \param[in] use_pre_softmax Whether to use softmax before topk. * \param[in] scaling_factor Scaling factor. - * \param[in] score_function Score function, 0: sigmoid, 1: softmax. + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. * \param[out] grad_logits Gradient of logits. * \param[in] stream CUDA stream used for the operation. */ @@ -63,7 +63,7 @@ void nvte_fused_topk_with_score_function_backward(const NVTETensor routing_map, * \param[in] num_tokens Number of tokens. * \param[in] num_experts Number of experts. * \param[in] topk Topk value. - * \param[in] score_function Score function, 0: sigmoid, 1: softmax. + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. * \param[out] scores Output tensor for scores. * \param[in] routing_map Routing map. * \param[in] intermediate_output Intermediate output from the forward pass. (Softmax/sigmoid output) @@ -82,7 +82,7 @@ void nvte_fused_score_for_moe_aux_loss_forward(const NVTETensor logits, int num_ * \param[in] num_tokens Number of tokens. * \param[in] num_experts Number of experts. * \param[in] topk Topk value. - * \param[in] score_function Score function, 0: sigmoid, 1: softmax. + * \param[in] score_function Score function, 0: sigmoid, 1: softmax, 2: sqrtsoftplus. * \param[out] grad_logits Gradient of logits. * \param[in] stream CUDA stream used for the operation. */ diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index cb6d8b7c92..b2b0751b04 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -27,23 +27,22 @@ namespace transformer_engine::pytorch { **************************************************************************************************/ std::tuple fused_topk_with_score_function_fwd( - at::Tensor logits, int topk, bool use_pre_softmax, c10::optional num_groups, - c10::optional group_topk, c10::optional scaling_factor, std::string score_function, - c10::optional expert_bias); + at::Tensor logits, int topk, bool use_pre_softmax, std::optional num_groups, + std::optional group_topk, std::optional scaling_factor, std::string score_function, + std::optional expert_bias); -at::Tensor fused_topk_with_score_function_bwd(int num_tokens, int num_experts, - at::Tensor routing_map, - at::Tensor intermediate_output, at::Tensor grad_probs, - int topk, bool use_pre_softmax, - c10::optional scaling_factor, - std::string score_function); +void fused_topk_with_score_function_bwd(int num_tokens, int num_experts, at::Tensor routing_map, + at::Tensor intermediate_output, at::Tensor grad_probs, + at::Tensor grad_logits, int topk, bool use_pre_softmax, + std::optional scaling_factor, + std::string score_function); std::tuple fused_score_for_moe_aux_loss_fwd( at::Tensor logits, int topk, std::string score_function); -at::Tensor fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, - at::Tensor intermediate_output, at::Tensor grad_probs, - int topk, std::string score_function); +void fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, + at::Tensor intermediate_output, at::Tensor grad_probs, + at::Tensor grad_logits, int topk, std::string score_function); std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, at::Tensor tokens_per_expert, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 5372f2f3e7..e9683ca41e 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -344,19 +344,20 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &transformer_engine::pytorch::fused_topk_with_score_function_fwd, py::arg("logits"), py::arg("topk"), py::arg("use_pre_softmax"), py::arg("num_groups"), py::arg("group_topk"), py::arg("scaling_factor"), py::arg("score_function"), py::arg("expert_bias"), - "Fused topk softmax fwd"); + "Fused topk with score function fwd"); m.def("fused_topk_with_score_function_bwd", &transformer_engine::pytorch::fused_topk_with_score_function_bwd, py::arg("num_tokens"), py::arg("num_experts"), py::arg("routing_map"), py::arg("intermediate_output"), - py::arg("grad_probs"), py::arg("topk"), py::arg("use_pre_softmax"), - py::arg("scaling_factor"), py::arg("score_function"), "Fused topk softmax bwd"); + py::arg("grad_probs"), py::arg("grad_logits"), py::arg("topk"), py::arg("use_pre_softmax"), + py::arg("scaling_factor"), py::arg("score_function"), "Fused topk with score function bwd"); m.def("fused_score_for_moe_aux_loss_fwd", &transformer_engine::pytorch::fused_score_for_moe_aux_loss_fwd, py::arg("logits"), - py::arg("topk"), py::arg("score_function"), "Fused topk softmax fwd"); + py::arg("topk"), py::arg("score_function"), "Fused aux loss with score function fwd"); m.def("fused_score_for_moe_aux_loss_bwd", &transformer_engine::pytorch::fused_score_for_moe_aux_loss_bwd, py::arg("num_tokens"), py::arg("num_experts"), py::arg("intermediate_output"), py::arg("grad_scores"), - py::arg("topk"), py::arg("score_function"), "Fused topk softmax bwd"); + py::arg("grad_logits"), py::arg("topk"), py::arg("score_function"), + "Fused aux loss with score function bwd"); m.def("fused_moe_aux_loss_fwd", &transformer_engine::pytorch::fused_moe_aux_loss_fwd, py::arg("probs"), py::arg("tokens_per_expert"), py::arg("total_num_tokens"), py::arg("num_experts"), py::arg("num_rows"), py::arg("num_cols"), py::arg("topk"), diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 2ae0d648a1..94625c0f12 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -9,12 +9,13 @@ namespace transformer_engine::pytorch { -static std::map score_function_map = {{"sigmoid", 0}, {"softmax", 1}}; +static std::map score_function_map = { + {"sigmoid", 0}, {"softmax", 1}, {"sqrtsoftplus", 2}}; std::tuple fused_topk_with_score_function_fwd( - at::Tensor logits, int topk, bool use_pre_softmax, c10::optional num_groups, - c10::optional group_topk, c10::optional scaling_factor, std::string score_function, - c10::optional expert_bias) { + at::Tensor logits, int topk, bool use_pre_softmax, std::optional num_groups, + std::optional group_topk, std::optional scaling_factor, std::string score_function, + std::optional expert_bias) { int num_tokens = logits.size(0); int num_experts = logits.size(1); // Check if the input is valid @@ -22,13 +23,16 @@ std::tuple fused_topk_with_score_function_fw "num_tokens and num_experts must be greater than 0"); // Expert bias only happens at the sigmoid case if (expert_bias.has_value()) { - TORCH_CHECK(score_function == "sigmoid", - "score_function must be sigmoid when expert_bias is not None"); + TORCH_CHECK(score_function == "sigmoid" || score_function == "sqrtsoftplus", + "score_function must be sigmoid or sqrtsoftplus when expert_bias is not None"); + TORCH_CHECK(expert_bias.value().scalar_type() == at::kFloat, + "expert_bias must be a float32 tensor"); } // Check if the score function is valid - TORCH_CHECK(score_function == "softmax" || score_function == "sigmoid", - "score_function must be softmax or sigmoid for router fusion"); - if (score_function == "sigmoid") { + TORCH_CHECK(score_function == "softmax" || score_function == "sigmoid" || + score_function == "sqrtsoftplus", + "score_function must be softmax, sigmoid or sqrtsoftplus for router fusion"); + if (score_function == "sigmoid" || score_function == "sqrtsoftplus") { use_pre_softmax = false; // Pre-softmax only happens at the softmax case } @@ -44,7 +48,7 @@ std::tuple fused_topk_with_score_function_fw at::empty({num_tokens, num_experts}, at::dtype(at::kBool).device(at::kCUDA)); // Intermediate output is used to store the output of the softmax/sigmoid function at::Tensor intermediate_output = - at::empty({num_tokens, num_experts}, at::dtype(logits.scalar_type()).device(at::kCUDA)); + at::empty({num_tokens, num_experts}, at::dtype(at::kFloat).device(at::kCUDA)); auto logits_cu = makeTransformerEngineTensor(logits); auto probs_cu = makeTransformerEngineTensor(probs); @@ -64,18 +68,14 @@ std::tuple fused_topk_with_score_function_fw return std::make_tuple(probs, routing_map, intermediate_output); } -at::Tensor fused_topk_with_score_function_bwd(int num_tokens, int num_experts, - at::Tensor routing_map, - at::Tensor intermediate_output, at::Tensor grad_probs, - int topk, bool use_pre_softmax, - c10::optional scaling_factor, - std::string score_function) { +void fused_topk_with_score_function_bwd(int num_tokens, int num_experts, at::Tensor routing_map, + at::Tensor intermediate_output, at::Tensor grad_probs, + at::Tensor grad_logits, int topk, bool use_pre_softmax, + std::optional scaling_factor, + std::string score_function) { // Get the value of the parameters auto scaling_factor_value = scaling_factor.has_value() ? scaling_factor.value() : 1.0f; auto score_function_value = score_function_map[score_function]; - // Init the output tensor - at::Tensor grad_logits = at::empty( - {num_tokens, num_experts}, at::dtype(intermediate_output.scalar_type()).device(at::kCUDA)); auto routing_map_cu = makeTransformerEngineTensor(routing_map); auto intermediate_output_cu = makeTransformerEngineTensor(intermediate_output); @@ -86,8 +86,6 @@ at::Tensor fused_topk_with_score_function_bwd(int num_tokens, int num_experts, routing_map_cu.data(), intermediate_output_cu.data(), grad_probs_cu.data(), num_tokens, num_experts, topk, use_pre_softmax, scaling_factor_value, score_function_value, grad_logits_cu.data(), at::cuda::getCurrentCUDAStream()); - - return grad_logits; } std::tuple fused_score_for_moe_aux_loss_fwd( @@ -99,17 +97,17 @@ std::tuple fused_score_for_moe_aux_loss_fwd( "num_tokens and num_experts must be greater than 0"); TORCH_CHECK(topk > 0, "topk must be greater than 0"); // Check if the score function is valid - TORCH_CHECK(score_function == "softmax" || score_function == "sigmoid", - "score_function must be softmax or sigmoid for router fusion"); + TORCH_CHECK(score_function == "softmax" || score_function == "sigmoid" || + score_function == "sqrtsoftplus", + "score_function must be softmax, sigmoid or sqrtsoftplus for router fusion"); int score_function_value = score_function_map[score_function]; // Construct the output tensor - at::Tensor scores = - at::empty({num_tokens, num_experts}, at::dtype(logits.scalar_type()).device(at::kCUDA)); + at::Tensor scores = at::empty({num_tokens, num_experts}, at::dtype(at::kFloat).device(at::kCUDA)); at::Tensor routing_map = at::empty({num_tokens, num_experts}, at::dtype(at::kBool).device(at::kCUDA)); at::Tensor intermediate_output = - at::empty({num_tokens, num_experts}, at::dtype(logits.scalar_type()).device(at::kCUDA)); + at::empty({num_tokens, num_experts}, at::dtype(at::kFloat).device(at::kCUDA)); auto logits_cu = makeTransformerEngineTensor(logits); auto scores_cu = makeTransformerEngineTensor(scores); @@ -123,14 +121,12 @@ std::tuple fused_score_for_moe_aux_loss_fwd( return std::make_tuple(scores, routing_map, intermediate_output); } -at::Tensor fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, - at::Tensor intermediate_output, at::Tensor grad_scores, - int topk, std::string score_function) { +void fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, + at::Tensor intermediate_output, at::Tensor grad_scores, + at::Tensor grad_logits, int topk, + std::string score_function) { // Get the value of the parameters int score_function_value = score_function_map[score_function]; - // Init the output tensor - at::Tensor grad_logits = at::empty( - {num_tokens, num_experts}, at::dtype(intermediate_output.scalar_type()).device(at::kCUDA)); auto intermediate_output_cu = makeTransformerEngineTensor(intermediate_output); auto grad_scores_cu = makeTransformerEngineTensor(grad_scores); @@ -139,8 +135,6 @@ at::Tensor fused_score_for_moe_aux_loss_bwd(int num_tokens, int num_experts, nvte_fused_score_for_moe_aux_loss_backward( intermediate_output_cu.data(), grad_scores_cu.data(), num_tokens, num_experts, topk, score_function_value, grad_logits_cu.data(), at::cuda::getCurrentCUDAStream()); - - return grad_logits; } std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, diff --git a/transformer_engine/pytorch/router.py b/transformer_engine/pytorch/router.py index 52d1d9d6ca..b56b1cd5eb 100644 --- a/transformer_engine/pytorch/router.py +++ b/transformer_engine/pytorch/router.py @@ -3,7 +3,18 @@ # See LICENSE for license information. """ Fused functions used in the MoE router + +Precision Notes: +- FP64 is currently not supported. +- Inputs are casted into FP32 when loading from global memory. +- All the math/calculations/accumulations are in FP32 in the kernels. +- "scores" is always in FP32 (match the MCore implementation). +- "intermediate_output" is always in FP32 for better backward precision. +- Only cast to low-precision when necessary and the casting only happens in writing to + global memory. For example, the gradient is required to have the same dtype as the input. """ +from typing import Optional + import torch import transformer_engine_torch as tex @@ -11,7 +22,7 @@ class FusedTopkScoreFunction(torch.autograd.Function): """ Fused Topk with Score Function router. - Currently, only support softmax and sigmoid. + Currently, support "softmax", "sigmoid" and "sqrtsoftplus". """ @staticmethod @@ -20,11 +31,11 @@ def forward( logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], score_function: str, - expert_bias: torch.Tensor, + expert_bias: Optional[torch.Tensor], ): # pylint: disable=missing-function-docstring # Save the shape of the logits @@ -52,6 +63,7 @@ def forward( ctx.topk = topk ctx.scaling_factor = scaling_factor ctx.score_function = score_function + ctx.logits_dtype = logits.dtype return probs, routing_map @staticmethod @@ -62,12 +74,16 @@ def backward(ctx, grad_probs, _): tensor_shape = grad_probs.shape # Adjust the shape of the grad_probs to 2D shape grad_probs = grad_probs.contiguous().view(-1, tensor_shape[-1]) - grad_logits = tex.fused_topk_with_score_function_bwd( + grad_logits = torch.empty( + (ctx.num_tokens, ctx.num_experts), dtype=ctx.logits_dtype, device=grad_probs.device + ) + tex.fused_topk_with_score_function_bwd( ctx.num_tokens, ctx.num_experts, routing_map, intermediate_output, grad_probs, + grad_logits, ctx.topk, ctx.use_pre_softmax, ctx.scaling_factor, @@ -82,37 +98,37 @@ def fused_topk_with_score_function( logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], score_function: str, - expert_bias: torch.Tensor, + expert_bias: Optional[torch.Tensor], ): """ Fused topk with score function router. Parameters ---------- - logits : torch.Tensor + logits : torch.Tensor in fp32/bf16/fp16 topk : int use_pre_softmax : bool - if enabled, the computation order: softmax -> topk - num_groups : int + if enabled, the computation order: softmax -> topk. + num_groups : int, optional used in the group topk - group_topk : int + group_topk : int, optional used in the group topk - scaling_factor : float + scaling_factor : float, optional score_function : str - currently only support softmax and sigmoid - expert_bias : torch.Tensor - could be used in the sigmoid + currently support "softmax", "sigmoid" and "sqrtsoftplus". + expert_bias : torch.Tensor, optional + could be used with the sigmoid/sqrtsoftplus score functions. Returns ------- - probs : torch.Tensor - routing_map : torch.Tensor + probs : torch.Tensor in the same dtype as the "logits". + routing_map : torch.Tensor in bool. """ if logits.dtype == torch.float64: - raise ValueError("Current TE does not support float64 router type") + raise ValueError("Current TE does not support float64 router type.") return FusedTopkScoreFunction.apply( logits, topk, @@ -154,6 +170,7 @@ def forward( ctx.score_function = score_function ctx.num_tokens = num_tokens ctx.num_experts = num_experts + ctx.logits_dtype = logits.dtype return routing_map, scores @staticmethod @@ -164,11 +181,15 @@ def backward(ctx, _, grad_scores): tensor_shape = grad_scores.shape # Adjust the shape of the grad_scores to 2D shape grad_scores = grad_scores.contiguous().view(-1, tensor_shape[-1]) - grad_logits = tex.fused_score_for_moe_aux_loss_bwd( + grad_logits = torch.empty( + (ctx.num_tokens, ctx.num_experts), dtype=ctx.logits_dtype, device=grad_scores.device + ) + tex.fused_score_for_moe_aux_loss_bwd( num_tokens=ctx.num_tokens, num_experts=ctx.num_experts, intermediate_output=intermediate_output, grad_scores=grad_scores, + grad_logits=grad_logits, topk=ctx.topk, score_function=ctx.score_function, ) @@ -186,15 +207,15 @@ def fused_compute_score_for_moe_aux_loss( Fused compute scores for MoE aux loss, subset of the fused_topk_with_score_function. Parameters ---------- - logits : torch.Tensor + logits : torch.Tensor in fp32/bf16/fp16 topk : int score_function : str - currently only support softmax and sigmoid + currently support "softmax", "sigmoid" and "sqrtsoftplus". Returns ------- - routing_map : torch.Tensor - scores : torch.Tensor + routing_map : torch.Tensor in bool + scores : torch.Tensor in fp32 """ return FusedComputeScoresForMoEAuxLoss.apply(logits, topk, score_function) @@ -253,23 +274,24 @@ def fused_moe_aux_loss( num_experts: int, topk: int, coeff: float, -): +) -> torch.Tensor: """ Fused MoE aux loss. Parameters ---------- - probs : torch.Tensor - tokens_per_expert : torch.Tensor - the number of tokens per expert + probs : torch.Tensor in fp32/bf16/fp16 + tokens_per_expert : torch.Tensor in int32/int64/fp32/bf16 + the number of tokens per expert. total_num_tokens : int - the total number of tokens, involved in the aux loss calculation + the total number of tokens used in the aux loss calculation. num_experts : int topk : int coeff : float - the coefficient of the aux loss + the coefficient of the aux loss. Returns ------- - aux_loss : torch.scalar + aux_loss : torch.Tensor. + A scalar tensor in the same dtype as the "probs". """ return FusedAuxLoss.apply(probs, tokens_per_expert, total_num_tokens, num_experts, topk, coeff) From 3ecb5bf1b6659fdef149878c752e1b03ff7c96a2 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Fri, 27 Feb 2026 17:08:22 -0800 Subject: [PATCH 232/521] [PyTorch] Fix L3 FA tests (#2709) * fix L3 FA fp8 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 skip logic based on reference backend 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 | 154 +++++++++++----------- 1 file changed, 77 insertions(+), 77 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 243fcac882..31c7041897 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1865,7 +1865,7 @@ def test_mha_fp8_vs_f16( ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe - available_backends, _, fused_attn_backends = get_available_attention_backends( + available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, qkv_layout=qkv_format.replace("hd", "h3d"), @@ -1875,20 +1875,18 @@ def test_mha_fp8_vs_f16( deterministic=_deterministic, ) flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends + available_backends, _, fused_attn_backends = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout=qkv_format.replace("hd", "h3d"), + is_training=is_training, + deterministic=_deterministic, + ) + _, fused_attn_supported_f16, _ = available_backends if flash_attn_supported + fused_attn_supported_fp8 < 1: pytest.skip("No FP8 attention backend available.") - fused_attn_supported_f16 = False - if not fp8_dpa_bwd: - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, - qkv_dtype=dtype, - qkv_layout=qkv_format.replace("hd", "h3d"), - is_training=is_training, - deterministic=_deterministic, - ) - _, fused_attn_supported_f16, _ = available_backends - if not fused_attn_supported_f16: - pytest.skip("No attention backend available.") + if not fused_attn_supported_f16: + pytest.skip("No reference backend available.") if flash_attn_supported: os.environ["NVTE_FLASH_ATTN"] = "1" @@ -2118,7 +2116,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe - available_backends, _, fused_attn_backends = get_available_attention_backends( + available_backends, _, _ = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, qkv_layout=qkv_layout, @@ -2127,20 +2125,19 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal is_training=is_training, deterministic=_deterministic, ) - flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if flash_attn_supported + fused_attn_supported < 1: + flash_attn_supported, fused_attn_supported_fp8, unfused_attn_supported = available_backends + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout=qkv_layout, + is_training=is_training, + deterministic=_deterministic, + ) + _, fused_attn_supported_f16, _ = available_backends + if flash_attn_supported + fused_attn_supported_fp8 < 1: pytest.skip("No FP8 attention backend available.") - if not fp8_dpa_bwd: - available_backends, _, fused_attn_backends = get_available_attention_backends( - config, - qkv_dtype=dtype, - qkv_layout=qkv_layout, - is_training=is_training, - deterministic=_deterministic, - ) - _, fused_attn_supported, _ = available_backends - if not fused_attn_supported: - pytest.skip("No attention backend available.") + if not fused_attn_supported_f16: + pytest.skip("No reference backend available.") if config.num_heads != config.num_gqa_groups and "3" in qkv_layout: pytest.skip("qkv_layout not applicable for MQA/GQA") @@ -2164,30 +2161,32 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal dtype, config, True, qkv_layout, is_training, fp8_recipe ) - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "1" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True - logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FusedAttention)") - fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( - dtype, config, True, qkv_layout, is_training, fp8_recipe - ) - - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "1" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - if config.dropout_p == 0.0: - # test cuDNN FP8 dropout: need a FP16/BF16 reference on Blackwell - logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = False (FusedAttention)") - fused_attn_fwd_f16, fused_attn_bwd_f16 = _run_dpa_fp8_vs_f16( - dtype, config, False, qkv_layout, is_training, fp8_recipe + if fused_attn_supported_fp8: + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True + logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = True (FusedAttention)") + fused_attn_fwd_fp8, fused_attn_bwd_fp8 = _run_dpa_fp8_vs_f16( + dtype, config, True, qkv_layout, is_training, fp8_recipe ) + if fused_attn_supported_f16: + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "1" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + if config.dropout_p == 0.0: + # test cuDNN FP8 dropout: need a FP16/BF16 reference on Blackwell + logging.info("[test_dpa_fp8_vs_f16]: run with fp8_dpa = False (FusedAttention)") + fused_attn_fwd_f16, fused_attn_bwd_f16 = _run_dpa_fp8_vs_f16( + dtype, config, False, qkv_layout, is_training, fp8_recipe + ) + atol = 5e-1 rtol = 5e-2 rmse_tol = 0.11 bwd_names = ["dq", "dk", "dv"] - if flash_attn_supported: + if flash_attn_supported and fused_attn_supported_f16: logging.debug("========== {:^25s} ==========".format("flash fp8 vs fused f16:")) logging.debug("========== {:^25s} ==========".format("forward output")) compare_and_assert( @@ -2200,7 +2199,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal rmse_tol, True, ) - if unfused_attn_supported: + if unfused_attn_supported and fused_attn_supported_f16: logging.debug("========== {:^25s} ==========".format("unfused fp8 vs fused f16:")) logging.debug("========== {:^25s} ==========".format("forward output")) compare_and_assert( @@ -2226,37 +2225,38 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal rmse_tol, True, ) - if config.dropout_p != 0.0: - # test cuDNN FP8 dropout - assert torch.all( - fused_attn_fwd_fp8 == 1 - ), "fused_attn_fwd_fp8 must be all 1s when Q/K/V are all 1s." - else: - logging.debug("========== {:^25s} ==========".format("fused fp8 vs fused f16:")) - logging.debug("========== {:^25s} ==========".format("forward output")) - compare_and_assert( - fused_attn_fwd_fp8, - fused_attn_fwd_f16, - "fused_attn_fwd_fp8", - "fused_attn_fwd_f16", - atol, - rtol, - rmse_tol, - True, - ) - if is_training: - for i, _ in enumerate(fused_attn_bwd_f16): - logging.debug("========== {:^25s} ==========".format(bwd_names[i])) - compare_and_assert( - fused_attn_bwd_fp8[i], - fused_attn_bwd_f16[i], - f"fused_attn_bwd_fp8[{i}]", - f"fused_attn_bwd_f16[{i}]", - atol, - rtol, - rmse_tol, - True, - ) + if fused_attn_supported_fp8 and fused_attn_supported_f16: + if config.dropout_p != 0.0: + # test cuDNN FP8 dropout + assert torch.all( + fused_attn_fwd_fp8 == 1 + ), "fused_attn_fwd_fp8 must be all 1s when Q/K/V are all 1s." + else: + logging.debug("========== {:^25s} ==========".format("fused fp8 vs fused f16:")) + logging.debug("========== {:^25s} ==========".format("forward output")) + compare_and_assert( + fused_attn_fwd_fp8, + fused_attn_fwd_f16, + "fused_attn_fwd_fp8", + "fused_attn_fwd_f16", + atol, + rtol, + rmse_tol, + True, + ) + if is_training: + for i, _ in enumerate(fused_attn_bwd_f16): + logging.debug("========== {:^25s} ==========".format(bwd_names[i])) + compare_and_assert( + fused_attn_bwd_fp8[i], + fused_attn_bwd_f16[i], + f"fused_attn_bwd_fp8[{i}]", + f"fused_attn_bwd_f16[{i}]", + atol, + rtol, + rmse_tol, + True, + ) os.environ["NVTE_UnfusedDPA_Emulate_FP8"] = "0" From f508e662f0b038ea8c97f8ceb2304432a5be5c97 Mon Sep 17 00:00:00 2001 From: Robin Zhang Date: Mon, 2 Mar 2026 13:30:43 +0800 Subject: [PATCH 233/521] [PyTorch] Remove `is_first_microbatch` setting after cudagraph warmup (#2715) Remove is_first_microbatch setting after warmup Signed-off-by: Robin Zhang --- transformer_engine/pytorch/graph.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index f4b1fb23ae..d3320fd70f 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -507,11 +507,6 @@ def hook_fn( else: grad_inputs = None del outputs, grad_inputs - # The following code is added specifically for MCore's special requirements, - # aimed at preventing warmup from altering the control flow. - for module in func.modules(): - if hasattr(module, "is_first_microbatch"): - module.is_first_microbatch = True torch.cuda.synchronize() # All captures here share a mempool. To avoid replays corrupting each other's memory, From 537f134236d10cf50a3bb296f453301079b5d7d5 Mon Sep 17 00:00:00 2001 From: Tong Liu Date: Mon, 2 Mar 2026 16:16:10 +0800 Subject: [PATCH 234/521] [Common][PyTorch] Fix normalization for `fused_score_for_moe_aux_loss` (#2720) * fix topk=1 Signed-off-by: tongliu * add topk=1 ut Signed-off-by: tongliu --------- Signed-off-by: tongliu --- tests/pytorch/test_fused_router.py | 8 ++++---- .../fused_router/fused_score_for_moe_aux_loss.cu | 16 +++++++--------- 2 files changed, 11 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 64000e109e..36c09060ed 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -113,10 +113,10 @@ def compute_scores_for_aux_loss_pytorch( scores = torch.softmax(logits, dim=-1, dtype=torch.float32) elif score_function == "sigmoid": scores = torch.sigmoid(logits.float()) - scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) elif score_function == "sqrtsoftplus": scores = torch.nn.functional.softplus(logits.float()).sqrt() - scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) else: raise ValueError(f"Invalid score_function: {score_function}") @@ -324,9 +324,9 @@ def test_topk_softmax( @pytest.mark.parametrize("dtype", [torch.float32]) -@pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) +@pytest.mark.parametrize("num_tokens", [2048, 7168]) @pytest.mark.parametrize("num_experts", [256, 128, 32]) -@pytest.mark.parametrize("topk", [4, 8]) +@pytest.mark.parametrize("topk", [1, 4, 8]) @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function): if score_function in ("sigmoid", "sqrtsoftplus"): 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 4f405e0a25..d38fcde6bf 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 @@ -109,14 +109,12 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi __syncwarp(); //Confirm the scores is written to the output - // Sigmoid/Sqrtsoftplus post-processing when topk > 1 + // Sigmoid/Sqrtsoftplus post-processing if (score_function == 0 || score_function == 2) { - if (topk > 1) { - auto sum_logits = - warp_reduce_on_shmem(local_logits, num_experts, ReduceFuncType::SUM, lane_id); - for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - local_logits[i] /= (sum_logits + epsilon); - } + auto sum_logits = + warp_reduce_on_shmem(local_logits, num_experts, ReduceFuncType::SUM, lane_id); + for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { + local_logits[i] /= (sum_logits + epsilon); } __syncwarp(); } @@ -246,8 +244,8 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const CompType *int __syncwarp(); } - // Sigmoid/Sqrtsoftplus Post-processing bwd when topk > 1 (normalization backward) - if (topk > 1 && (score_function == 0 || score_function == 2)) { + // Sigmoid/Sqrtsoftplus Post-processing bwd (normalization backward) + if (score_function == 0 || score_function == 2) { // Select the correct activation output buffer: // - Sigmoid: local_act_from_fwd already contains sigmoid output // - Sqrtsoftplus: local_comp_buf contains sqrtsoftplus output computed above From bba7bf6a4101e150d1aaf9278608385214d09684 Mon Sep 17 00:00:00 2001 From: Hongbin Liu Date: Mon, 2 Mar 2026 16:39:41 +0800 Subject: [PATCH 235/521] [PyTorch] Support cuda graph capturing offloading module (#2435) * support cuda graph capture offloading module Signed-off-by: Hongbin Liu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove reset_hook and init_chunk_handler_hook Signed-off-by: Hongbin Liu * remove reset_hook and init_chunk_handler_hook Signed-off-by: Hongbin Liu * minor fix Signed-off-by: root * temp fix overlap-grad-reduce Signed-off-by: Hongbin Liu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reuse mark_not_offload() and do not offload scale_inv Signed-off-by: Hongbin Liu * temp fix for mxfp8 Signed-off-by: Hongbin Liu * fix bug for record_stream and from_blob Signed-off-by: Hongbin Liu * disable offloading core_attn_out and refine cpu overhead of at::empty Signed-off-by: Hongbin Liu * minor fix Signed-off-by: Hongbin Liu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * return ptr of whole buffer and offload the whole buffer Signed-off-by: Hongbin Liu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestions from code revie Signed-off-by: Hongbin Liu * remove code changes of offloading and quantizer Signed-off-by: Hongbin Liu * minor fix Signed-off-by: Hongbin Liu * minor fix Signed-off-by: Hongbin Liu * minor fix Signed-off-by: Hongbin Liu * minor fix Signed-off-by: Hongbin Liu * minor fix Signed-off-by: Hongbin Liu * add docstring Signed-off-by: Hongbin Liu --------- Signed-off-by: Hongbin Liu Signed-off-by: root Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root Co-authored-by: root --- transformer_engine/pytorch/graph.py | 64 ++++++++++++++++++++++++++--- 1 file changed, 58 insertions(+), 6 deletions(-) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index d3320fd70f..bae911b4e1 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -108,6 +108,8 @@ def _make_graphed_callables( pool: Optional[Tuple[int, ...]] = None, retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, + pre_warmup_hook: Optional[Callable] = None, + post_warmup_hook: Optional[Callable] = None, ) -> SingleOrTuple[Callable]: """ Helper method for `make_graphed_callables` @@ -440,6 +442,8 @@ def hook_fn( else: visited_te_modules[func_idx].update(modules) + if pre_warmup_hook is not None: + pre_warmup_hook() for warmup_iter in range(num_warmup_iters): hooks = [] for module in func.modules(): @@ -507,6 +511,8 @@ def hook_fn( else: grad_inputs = None del outputs, grad_inputs + if post_warmup_hook is not None: + post_warmup_hook() torch.cuda.synchronize() # All captures here share a mempool. To avoid replays corrupting each other's memory, @@ -777,14 +783,15 @@ class Graphed(torch.autograd.Function): """Autograd function for graph replay.""" @staticmethod - def forward(ctx, skip_fp8_weight_update, *inputs): + def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *inputs): # pylint: disable=missing-function-docstring # Set flag for whether to update FP8 weight updates ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module() if ctx.is_first_module and skip_fp8_weight_update is not None: FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(skip_fp8_weight_update) - + ctx.cuda_graph_stream = cuda_graph_stream + ctx.cuda_graph_event = cuda_graph_event # Copy values from new tensors into static tensors for i in range(len_user_args): if ( @@ -794,7 +801,16 @@ def forward(ctx, skip_fp8_weight_update, *inputs): static_input_surface[i].copy_(inputs[i]) # Replay forward graph - fwd_graph.replay() + if cuda_graph_stream != torch.cuda.current_stream(): + cuda_graph_stream.wait_stream(torch.cuda.current_stream()) + with cuda_graph_stream: + fwd_graph.replay() + if cuda_graph_event is not None: + torch.cuda.current_stream().wait_event(cuda_graph_event) + else: + torch.cuda.current_stream().wait_stream(cuda_graph_stream) + else: + fwd_graph.replay() assert isinstance(static_outputs, tuple) return tuple(o.detach() if o is not None else o for o in static_outputs) @@ -811,7 +827,16 @@ def backward(ctx, *grads): # incoming grad is already in the right place if g.data_ptr() != grad.data_ptr(): g.copy_(grad) - bwd_graph.replay() + if ctx.cuda_graph_stream != torch.cuda.current_stream(): + ctx.cuda_graph_stream.wait_stream(torch.cuda.current_stream()) + with ctx.cuda_graph_stream: + bwd_graph.replay() + if ctx.cuda_graph_event is not None: + torch.cuda.current_stream().wait_event(ctx.cuda_graph_event) + else: + torch.cuda.current_stream().wait_stream(ctx.cuda_graph_stream) + else: + bwd_graph.replay() # Update FP8 scale factors if needed if ctx.is_first_module: @@ -819,7 +844,7 @@ def backward(ctx, *grads): # Input args that didn't require grad expect a None gradient. assert isinstance(static_grad_inputs, tuple) - return (None,) + tuple( + return (None, None, None) + tuple( b.detach() if b is not None else b for b in static_grad_inputs ) @@ -834,6 +859,23 @@ def functionalized(*user_args, **user_kwargs): skip_fp8_weight_update = not user_kwargs["is_first_microbatch"] + # The cuda_graph_stream and cuda_graph_event are used in the TE CUDA graph replay. + # When replaying the graph in the cuda graph stream, the graph replay could overlap + # with the work on main stream. + # When cuda_graph_event is given, it should be an external event recorded + # in the cuda graph and is used to sync-back to the main stream. + # If cuda_graph_event is not given, it will be None and the graph replay will block + # the main stream until it is finished. + if "cuda_graph_stream" in user_kwargs: + cuda_graph_stream = user_kwargs["cuda_graph_stream"] + user_kwargs.pop("cuda_graph_stream") + else: + cuda_graph_stream = torch.cuda.current_stream() + if "cuda_graph_event" in user_kwargs: + cuda_graph_event = user_kwargs["cuda_graph_event"] + user_kwargs.pop("cuda_graph_event") + else: + cuda_graph_event = None # Check that required kwargs are provided for key in kwargs_keys: if key not in user_kwargs: @@ -849,7 +891,9 @@ def functionalized(*user_args, **user_kwargs): flatten_user_args, _ = _tree_flatten(user_args) flatten_user_kwargs, _ = _tree_flatten([user_kwargs[key] for key in kwargs_keys]) func_args = tuple(flatten_user_args) + tuple(flatten_user_kwargs) + module_params - out = Graphed.apply(skip_fp8_weight_update, *func_args) + out = Graphed.apply( + skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *func_args + ) return _tree_unflatten(out, output_unflatten_spec) return functionalized @@ -1035,6 +1079,8 @@ def make_graphed_callables( pool: Optional[Tuple[int, ...]] = None, retain_graph_in_backward: bool = False, _reuse_graph_input_output_buffers: bool = False, + pre_warmup_hook: Optional[Callable] = None, + post_warmup_hook: Optional[Callable] = None, ) -> Union[Callable, Tuple[Callable, ...]]: """ Make CUDA graph version of Transformer Engine modules @@ -1073,6 +1119,10 @@ def make_graphed_callables( graphs. Only supported with Mcore interleaved pipeline parallelism, i.e. when `_order` is provided. All callables in `modules` are assumed to have inputs and outputs with the same dtype and shape. + pre_warmup_hook: callable, default = None + A hook function that will be called before the warmup iterations. + post_warmup_hook: callable, default = None + A hook function that will be called after the warmup iterations. Quantization parameters ----------------------- @@ -1259,6 +1309,8 @@ def call_func(self, *args, **kwargs): pool=pool, retain_graph_in_backward=retain_graph_in_backward, _reuse_graph_input_output_buffers=_reuse_graph_input_output_buffers, + pre_warmup_hook=pre_warmup_hook, + post_warmup_hook=post_warmup_hook, ) # Ensures warmup does not affect numerics for ops such as dropout. From 3275e1a0f4da4d6a15cdc839e34c3341f9a98dd7 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 2 Mar 2026 10:13:58 -0800 Subject: [PATCH 236/521] [JAX] CGEMM with Shardy (#2714) * replace Shardy error with a warning * cleanup tests * added warnings to dense vjp and layernorm_mlp vjp --------- Signed-off-by: Phuong Nguyen --- examples/jax/collective_gemm/common.py | 4 ---- examples/jax/collective_gemm/test_gemm.py | 5 +---- .../jax/collective_gemm/test_layernorm_mlp_grad.py | 2 -- transformer_engine/jax/cpp_extensions/gemm.py | 13 ++++++++++--- transformer_engine/jax/dense.py | 7 +++++++ transformer_engine/jax/layernorm_mlp.py | 8 ++++++++ 6 files changed, 26 insertions(+), 13 deletions(-) diff --git a/examples/jax/collective_gemm/common.py b/examples/jax/collective_gemm/common.py index 0d812da057..2965896d07 100644 --- a/examples/jax/collective_gemm/common.py +++ b/examples/jax/collective_gemm/common.py @@ -131,10 +131,6 @@ def _initialize_distributed(args): ) _distributed_initialized = True - jax.clear_caches() - jax.config.update( - "jax_use_shardy_partitioner", False - ) # CollectiveGEMM does not work with Shardy yet assert jax.local_device_count() == 1, ( f"[{args.process_id}|{args.num_devices_per_process}] Expected 1 GPU per process, found" diff --git a/examples/jax/collective_gemm/test_gemm.py b/examples/jax/collective_gemm/test_gemm.py index d2994723bb..ea119713e3 100644 --- a/examples/jax/collective_gemm/test_gemm.py +++ b/examples/jax/collective_gemm/test_gemm.py @@ -88,8 +88,6 @@ def _jitted_cgemm(x, weight, bias, contracting_dims, collective_op, output_shard def run_gemm_tests(args, mesh=None): """Execute GEMM tests.""" print(args) - # Collective GEMM requires Shardy partitioner to be disabled - jax.config.update("jax_use_shardy_partitioner", False) # Initialize distributed with provided arguments _initialize_distributed(args) @@ -137,8 +135,7 @@ def run_gemm_tests(args, mesh=None): bias_sharded, contracting_dims=((2,), (0,)), collective_op=collective_op, - # CollectiveGEMM output should have a correct sharding without applying sharding constraint - output_sharding=None, + output_sharding=output_sharding, ) assert ( ref_output.sharding == output.sharding diff --git a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py index 61c960a7aa..84cb011da1 100644 --- a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py +++ b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py @@ -119,8 +119,6 @@ def _value_and_grad_layernorm_mlp( def run_layernorm_mlp_grad_tests(args, mesh=None): """Execute Dense Gradient tests.""" print(args) - # Collective GEMM requires Shardy partitioner to be disabled - jax.config.update("jax_use_shardy_partitioner", False) # Initialize distributed with provided arguments _initialize_distributed(args) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index a34cb030bf..fbaafdf6d8 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -1172,9 +1172,16 @@ def shardy_sharding_rule( del mesh, result_types, transpose_batch_sequence, sequence_dim, is_outer if not collective_op.is_none: - raise NotImplementedError( - "CollectiveGEMM with Shardy propagation is not supported yet! Please turn off" - " Shardy by exporting env var JAX_USE_SHARDY_PARTITIONER=false" + warnings.warn( + "CollectiveGEMM with Shardy propagation may produce an incorrect sharding pattern" + " for the output.\n To resolve this, apply a sharding constraint on the output" + " using one of the following options:\n" + " - TE `dense` vjp: set `output_axes`.\n" + " - TE `layernorm_mlp` vjp: set `dot_2_input_axes`.\n" + " - TE `transformer_engine.jax.cpp_extensions.gemm`: apply" + " `jax.lax.with_sharding_constraint` on the output.\n" + " - TE via MaxText: no action needed.", + UserWarning, ) prefix = "Gemm_" diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 23d91f7db0..268995281c 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -94,6 +94,13 @@ def dense( if transpose_batch_sequence: warnings.warn("transpose_batch_sequence is not well tested, use with caution!") + if collective_op_set != tex.noop_collective_op_set and not output_axes: + warnings.warn( + "Collective GEMM with Shardy propagation may produce an incorrect sharding pattern" + " for the output. Set `output_axes` to apply the correct sharding constraint.", + UserWarning, + ) + if quantizer_set == noop_quantizer_set: input_dtype = x.dtype kernel = kernel.astype(input_dtype) diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py index a8de32830b..c90d018aee 100644 --- a/transformer_engine/jax/layernorm_mlp.py +++ b/transformer_engine/jax/layernorm_mlp.py @@ -15,6 +15,7 @@ from typing import List, Tuple, Sequence, Union, Callable from functools import partial +import warnings import jax import jax.numpy as jnp @@ -275,6 +276,13 @@ def _layernorm_mlp_fwd_rule( assert not collective_op_set_1.forward.is_reduce_scatter assert not collective_op_set_2.forward.is_all_gather + if collective_op_set_1 != tex.noop_collective_op_set and not dot_2_input_axes: + warnings.warn( + "Collective GEMM with Shardy propagation may produce an incorrect sharding pattern" + " for the output. Set `dot_2_input_axes` to apply the correct sharding constraint.", + UserWarning, + ) + # x should be in shape of (batch..., hidden) # Kernel_1 should be in shape of (hidden_in, activation_len, intermediate) # Kernel_2 should be in shape of (intermediate, hidden_in) From 9dac78e76a8e6c33add4d0b1aec8b3dd2c7db8db Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 2 Mar 2026 16:58:43 -0800 Subject: [PATCH 237/521] CPU Overhead Optimizations (#2559) * add all the optimizations Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * requires_grad optimization Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * test if commenting out requires_grad works Signed-off-by: Varun Thumbe * fix minor bug Signed-off-by: Varun Thumbe * fix ci Signed-off-by: Varun Thumbe * missed a bug Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/csrc/quantizer.cpp Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * fix some bugs pointed to by copilot Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * linting error Signed-off-by: Varun Thumbe * fix the error Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix the bug Signed-off-by: Varun Thumbe * get rid of the change Signed-off-by: Varun Thumbe * fix the transpose shape bug Signed-off-by: Varun Thumbe * minor linter fix Signed-off-by: Varun Thumbe * fix lint Signed-off-by: Varun Thumbe * fix linting error Signed-off-by: Varun Thumbe * address copilot review comment regarding error check when both data and transpose are None 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 * missed a merge conflict Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * final optimizations Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ci error 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 comment from greptile 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 comment + stride optimization Signed-off-by: Varun Thumbe * address linter issue Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor lint Signed-off-by: Varun Thumbe * fix ci bug Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * another optimization to do at::native::empty_cuda directly instead of at::empty Signed-off-by: Varun Thumbe * cleanups Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * better solution for device Signed-off-by: Varun Thumbe * enum to int cache Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove unused function Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/tensor/float8_blockwise_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * index instead of device bug Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ci: Signed-off-by: Varun Thumbe * debug quantized tensor fix Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert cudnnt front end change 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> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/attention/test_attention.py | 29 +- tests/pytorch/test_custom_recipe.py | 21 +- .../common/gemm/cublaslt_gemm.cu | 12 +- transformer_engine/common/util/cuda_driver.h | 21 +- .../debug/pytorch/debug_quantization.py | 9 + transformer_engine/pytorch/constants.py | 20 ++ .../pytorch/cpp_extensions/fused_attn.py | 13 +- .../pytorch/cpp_extensions/gemm.py | 26 +- .../pytorch/csrc/extensions/pybind.cpp | 17 +- transformer_engine/pytorch/csrc/quantizer.cpp | 314 ++++++++++++++---- transformer_engine/pytorch/module/base.py | 5 +- .../pytorch/module/layernorm_linear.py | 42 +-- .../pytorch/module/layernorm_mlp.py | 62 ++-- transformer_engine/pytorch/module/linear.py | 42 +-- .../pytorch/quantized_tensor.py | 76 ++++- .../pytorch/tensor/float8_blockwise_tensor.py | 18 + .../pytorch/tensor/float8_tensor.py | 19 ++ .../pytorch/tensor/mxfp8_tensor.py | 28 ++ .../pytorch/tensor/nvfp4_tensor.py | 30 ++ .../float8_blockwise_tensor_storage.py | 9 + .../tensor/storage/float8_tensor_storage.py | 9 + .../tensor/storage/mxfp8_tensor_storage.py | 9 + .../tensor/storage/nvfp4_tensor_storage.py | 9 + 23 files changed, 622 insertions(+), 218 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 31c7041897..60ade522e3 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -44,6 +44,7 @@ scaled_init_method_normal, ) from transformer_engine.pytorch.utils import get_cudnn_version +from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx import transformer_engine_torch as tex from transformer_engine.pytorch.quantized_tensor import ( Quantizer, @@ -2581,12 +2582,12 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: _2X_ACC_DGRAD = False _2X_ACC_WGRAD = False -META_QKV = tex.FP8FwdTensors.GEMM1_OUTPUT -META_DQKV = tex.FP8BwdTensors.GRAD_OUTPUT1 -META_O = tex.FP8FwdTensors.GEMM2_INPUT -META_DO = tex.FP8BwdTensors.GRAD_INPUT2 -META_S = tex.FP8FwdTensors.GEMM3_OUTPUT -META_DP = tex.FP8BwdTensors.GRAD_INPUT3 +META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT +META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 +META_O = FP8FwdTensorIdx.GEMM2_INPUT +META_DO = FP8BwdTensorIdx.GRAD_INPUT2 +META_S = FP8FwdTensorIdx.GEMM3_OUTPUT +META_DP = FP8BwdTensorIdx.GRAD_INPUT3 class _custom_mha_fp8(torch.autograd.Function): @@ -2614,14 +2615,14 @@ def forward( d = in_features // h b = cu_seqlens.numel() - 1 - input_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] - qkv_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM2_INPUT] - qkv_weight_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] - o_quantizer = quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] - dO_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] - dQKV_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] - s_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT2] - dP_quantizer = quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT3] + input_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] + qkv_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_INPUT] + qkv_weight_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] + o_quantizer = quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] + dO_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] + dQKV_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] + s_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT2] + dP_quantizer = quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT3] inp_fp8 = input_quantizer(inp) diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 4de49115b3..536d43adc0 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -8,6 +8,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.common import recipe +from transformer_engine.pytorch.constants import FP8BwdTensorIdx, FP8FwdTensorIdx from transformer_engine.pytorch import ( autocast, Linear, @@ -169,11 +170,11 @@ def test_custom_recipe_matches_current_scaling(): with autocast(enabled=True, recipe=ref_recipe): out_ref = model_ref(inp_ref) # Assert dtypes for reference quantizers: HYBRID = E4M3 (fwd), E5M2 (bwd) - ref_fwd_in = model_ref.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] - ref_fwd_w = model_ref.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] - ref_fwd_out = model_ref.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] - ref_bwd_go = model_ref.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] - ref_bwd_gi = model_ref.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] + ref_fwd_in = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] + ref_fwd_w = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] + ref_fwd_out = model_ref.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] + ref_bwd_go = model_ref.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] + ref_bwd_gi = model_ref.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] assert ref_fwd_in.dtype == tex.DType.kFloat8E4M3 assert ref_fwd_w.dtype == tex.DType.kFloat8E4M3 assert ref_fwd_out.dtype == tex.DType.kFloat8E4M3 @@ -200,11 +201,11 @@ def quantizer_factory(role): with autocast(enabled=True, recipe=custom_recipe): out_custom = model_custom(inp_custom) # Assert dtypes for custom quantizers match reference mapping - cus_fwd_in = model_custom.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] - cus_fwd_w = model_custom.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] - cus_fwd_out = model_custom.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] - cus_bwd_go = model_custom.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] - cus_bwd_gi = model_custom.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] + cus_fwd_in = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] + cus_fwd_w = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] + cus_fwd_out = model_custom.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] + cus_bwd_go = model_custom.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] + cus_bwd_gi = model_custom.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] assert cus_fwd_in.dtype == tex.DType.kFloat8E4M3 assert cus_fwd_w.dtype == tex.DType.kFloat8E4M3 assert cus_fwd_out.dtype == tex.DType.kFloat8E4M3 diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index c58c3cb47a..144aea1a07 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -120,6 +120,10 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla // Set conditions for MXFP8 and NVFP4 gemm execution. const auto nvfp4 = is_nvfp_scaling(A.scaling_mode) && is_nvfp_scaling(B.scaling_mode); const auto mxfp8 = !nvfp4 && is_mxfp_scaling(A.scaling_mode) && is_mxfp_scaling(B.scaling_mode); + int is_nvte_non_tn_fp8_gemm_supported = 0; // needed only for per tensor scaling + if (is_tensor_scaling(A.scaling_mode) || is_tensor_scaling(B.scaling_mode)) { + is_nvte_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); + } // Configure A matrix if (is_tensor_scaling(A.scaling_mode)) { @@ -129,7 +133,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.Atype = A.data.dtype; ret.A_scale_inv = A.scale_inv.dptr; ret.lda = is_A_transposed ? k : m; - if (!nvte_is_non_tn_fp8_gemm_supported() && !is_A_transposed) { + if (!is_nvte_non_tn_fp8_gemm_supported && !is_A_transposed) { // Hopper only supports TN GEMMs for FP8. "Column-wise data" is transpose of data. if (A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype)) { ret.A = A.columnwise_data.dptr; @@ -140,7 +144,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla } else { NVTE_CHECK(!is_fp8_dtype(ret.Atype), "Input A is missing column-wise usage"); } - } else if (nvte_is_non_tn_fp8_gemm_supported() && !A.has_data()) { + } else if (is_nvte_non_tn_fp8_gemm_supported && !A.has_data()) { // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed // data with the mirrored transpose-flag if we don't have row-wise data. NVTE_CHECK(A.has_columnwise_data() && is_fp8_dtype(A.columnwise_data.dtype), @@ -220,7 +224,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla ret.Btype = B.data.dtype; ret.B_scale_inv = B.scale_inv.dptr; ret.ldb = is_B_transposed ? n : k; - if (!nvte_is_non_tn_fp8_gemm_supported() && is_B_transposed) { + if (!is_nvte_non_tn_fp8_gemm_supported && is_B_transposed) { // Hopper only supports TN GEMMs for FP8. "Column-wise data" is transpose of data. if (B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype)) { ret.B = B.columnwise_data.dptr; @@ -231,7 +235,7 @@ GemmParam CanonicalizeGemmInput(const transformer_engine::Tensor &A, const cubla } else { NVTE_CHECK(!is_fp8_dtype(ret.Btype), "Input B is missing column-wise usage"); } - } else if (nvte_is_non_tn_fp8_gemm_supported() && !B.has_data()) { + } else if (is_nvte_non_tn_fp8_gemm_supported && !B.has_data()) { // Blackwell supports any GEMM layout for FP8, so we can use column-wise/transposed // data with the mirrored transpose-flag if we don't have row-wise data. NVTE_CHECK(B.has_columnwise_data() && is_fp8_dtype(B.columnwise_data.dtype), diff --git a/transformer_engine/common/util/cuda_driver.h b/transformer_engine/common/util/cuda_driver.h index 2715d8e4e4..16242347f1 100644 --- a/transformer_engine/common/util/cuda_driver.h +++ b/transformer_engine/common/util/cuda_driver.h @@ -9,7 +9,9 @@ #include +#include #include +#include #include "../common.h" #include "../util/string.h" @@ -29,13 +31,30 @@ void *get_symbol(const char *symbol, int cuda_version = 12010); * without GPUs. Indirect function calls into a lazily-initialized * library ensures we are accessing the correct version. * + * Symbol pointers are cached to avoid repeated lookups. + * * \param[in] symbol Function name * \param[in] args Function arguments */ template inline CUresult call(const char *symbol, ArgTs... args) { using FuncT = CUresult(ArgTs...); - FuncT *func = reinterpret_cast(get_symbol(symbol)); + + static std::unordered_map symbol_cache; + static std::mutex cache_mutex; + FuncT *func; + + { + std::lock_guard lock(cache_mutex); + auto it = symbol_cache.find(symbol); + if (it == symbol_cache.end()) { + void *ptr = get_symbol(symbol); + symbol_cache[symbol] = ptr; + func = reinterpret_cast(ptr); + } else { + func = reinterpret_cast(it->second); + } + } return (*func)(args...); } diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 5624970547..57a5967079 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -697,3 +697,12 @@ def update_usage(self, rowwise_usage: bool = None, columnwise_usage: bool = None raise RuntimeError( "Cannot recreate columnwise tensor from rowwise tensor is debug mode." ) + + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self.rowwise_gemm_tensor is not None: + return self.rowwise_gemm_tensor.device + if self.columnwise_gemm_tensor is not None: + return self.columnwise_gemm_tensor.device + raise RuntimeError("DebugQuantizedTensor has no data!") diff --git a/transformer_engine/pytorch/constants.py b/transformer_engine/pytorch/constants.py index 3cce4600d9..2aff4fd8e8 100644 --- a/transformer_engine/pytorch/constants.py +++ b/transformer_engine/pytorch/constants.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """Enums for e2e transformer""" +from types import SimpleNamespace import torch import torch.distributed import transformer_engine_torch as tex @@ -40,6 +41,25 @@ tex.DType.kBFloat16: torch.bfloat16, } +# Cache enum -> int conversions to avoid repeated PyObject lookups. +FP8FwdTensorIdx = SimpleNamespace( + GEMM1_INPUT=int(tex.FP8FwdTensors.GEMM1_INPUT), + GEMM1_WEIGHT=int(tex.FP8FwdTensors.GEMM1_WEIGHT), + GEMM1_OUTPUT=int(tex.FP8FwdTensors.GEMM1_OUTPUT), + GEMM2_INPUT=int(tex.FP8FwdTensors.GEMM2_INPUT), + GEMM2_WEIGHT=int(tex.FP8FwdTensors.GEMM2_WEIGHT), + GEMM2_OUTPUT=int(tex.FP8FwdTensors.GEMM2_OUTPUT), + GEMM3_OUTPUT=int(tex.FP8FwdTensors.GEMM3_OUTPUT), +) +FP8BwdTensorIdx = SimpleNamespace( + GRAD_INPUT1=int(tex.FP8BwdTensors.GRAD_INPUT1), + GRAD_INPUT2=int(tex.FP8BwdTensors.GRAD_INPUT2), + GRAD_INPUT3=int(tex.FP8BwdTensors.GRAD_INPUT3), + GRAD_OUTPUT1=int(tex.FP8BwdTensors.GRAD_OUTPUT1), + GRAD_OUTPUT2=int(tex.FP8BwdTensors.GRAD_OUTPUT2), + GRAD_OUTPUT3=int(tex.FP8BwdTensors.GRAD_OUTPUT3), +) + AttnMaskTypes = ( "no_mask", "padding", diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 101e5b2525..e9f64bb693 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -16,6 +16,7 @@ NVTE_Fused_Attn_Backend, ) from ..quantized_tensor import Quantizer +from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx __all__ = [ @@ -103,12 +104,12 @@ BACKEND_F16m512_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 -META_QKV = tex.FP8FwdTensors.GEMM1_OUTPUT -META_DQKV = tex.FP8BwdTensors.GRAD_OUTPUT1 -META_O = tex.FP8FwdTensors.GEMM2_INPUT -META_DO = tex.FP8BwdTensors.GRAD_INPUT2 -META_S = tex.FP8FwdTensors.GEMM3_OUTPUT -META_DP = tex.FP8BwdTensors.GRAD_INPUT3 +META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT +META_DQKV = FP8BwdTensorIdx.GRAD_OUTPUT1 +META_O = FP8FwdTensorIdx.GEMM2_INPUT +META_DO = FP8BwdTensorIdx.GRAD_INPUT2 +META_S = FP8FwdTensorIdx.GEMM3_OUTPUT +META_DP = FP8BwdTensorIdx.GRAD_INPUT3 def fused_attn_fwd( diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 406e7075f7..a37f1c2d4d 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -67,28 +67,6 @@ def validate_gemm_scale(scale: Optional[float], required: bool) -> float: return 0.0 -def get_tensor_device(tensor: torch.Tensor) -> int: - """ - Returns tensor device as an integer. - - This method is used because checking instances of - QuantizedTensor or Storage incurs more CPU overhead. - The order of attributes checked is important to also - minimize overhead. - """ - if hasattr(tensor, "device"): - return tensor.device.index - if hasattr(tensor, "_rowwise_data") and tensor._rowwise_data is not None: - return tensor._rowwise_data.device.index - if hasattr(tensor, "_columnwise_data") and tensor._columnwise_data is not None: - return tensor._columnwise_data.device.index - if hasattr(tensor, "_data") and tensor._data is not None: - return tensor._data.device.index - if hasattr(tensor, "_transpose") and tensor._transpose is not None: - return tensor._transpose.device.index - return torch.cuda.current_device() - - def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -117,7 +95,7 @@ def general_gemm( alpha = validate_gemm_scale(alpha, True) beta = validate_gemm_scale(beta, accumulate) - workspace = get_cublas_workspace(get_tensor_device(A), ub is not None, False) + workspace = get_cublas_workspace(A.device.index, ub is not None, False) if ub_type is not None: assert ub is not None, ( @@ -235,7 +213,7 @@ def general_grouped_gemm( out_dtype = TE_DType[out[0].dtype] if D_dtype is None else D_dtype sm_count = get_sm_count() - workspaces = get_cublas_workspace(get_tensor_device(A[0]), False, True) + workspaces = get_cublas_workspace(A[0].device.index, False, True) if grad and use_bias: grad_bias = [ diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index e9683ca41e..b9fc65363d 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -35,10 +35,10 @@ PyTypeObject *Float8BlockwiseQuantizerClass = nullptr; PyTypeObject *NVFP4TensorPythonClass = nullptr; PyTypeObject *NVFP4TensorStoragePythonClass = nullptr; PyTypeObject *NVFP4QuantizerClass = nullptr; +std::once_flag extension_init_flag; PyTypeObject *GroupedTensorStoragePythonClass = nullptr; void init_float8_extension() { - if (Float8TensorPythonClass) return; auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.float8_tensor"); Float8QuantizerClass = reinterpret_cast(PyObject_GetAttrString(fp8_module.ptr(), "Float8Quantizer")); @@ -55,7 +55,6 @@ void init_float8_extension() { } void init_mxfp8_extension() { - if (MXFP8TensorPythonClass) return; auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.mxfp8_tensor"); MXFP8QuantizerClass = reinterpret_cast(PyObject_GetAttrString(fp8_module.ptr(), "MXFP8Quantizer")); @@ -70,7 +69,6 @@ void init_mxfp8_extension() { } void init_float8blockwise_extension() { - if (Float8BlockwiseQTensorStoragePythonClass) return; auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.float8_blockwise_tensor"); auto fp8_base_module = py::module_::import( @@ -91,7 +89,6 @@ void init_float8blockwise_extension() { } void init_nvfp4_extensions() { - if (NVFP4TensorPythonClass) return; auto nvfp4_module = py::module_::import("transformer_engine.pytorch.tensor.nvfp4_tensor"); NVFP4QuantizerClass = reinterpret_cast( PyObject_GetAttrString(nvfp4_module.ptr(), "NVFP4Quantizer")); @@ -116,11 +113,13 @@ void init_grouped_tensor_extension() { } void init_extension() { - init_float8_extension(); - init_mxfp8_extension(); - init_float8blockwise_extension(); - init_nvfp4_extensions(); - init_grouped_tensor_extension(); + std::call_once(extension_init_flag, []() { + init_float8_extension(); + init_mxfp8_extension(); + init_float8blockwise_extension(); + init_nvfp4_extensions(); + init_grouped_tensor_extension(); + }); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index e715d8f5ba..0da5f69197 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -31,6 +31,23 @@ std::vector make_transpose_shape(const std::vector& shape) { return ret; } +/*! @brief Calculate stride from shape for contiguous tensors */ +template +std::vector stride_from_shape(const std::vector& shape) { + std::vector stride; + if (shape.empty()) { + return stride; + } + std::vector rstride; + rstride.reserve(shape.size()); + rstride.push_back(static_cast(1)); + for (size_t i = shape.size(); i > 1; --i) { + rstride.push_back(rstride.back() * shape[i - 1]); + } + stride.assign(rstride.rbegin(), rstride.rend()); + return stride; +} + /*! @brief Convert shape for FP4 data by dividing the last dimension by 2 */ template std::vector convert_shape_for_fp4(const std::vector& shape) { @@ -206,9 +223,9 @@ std::pair Float8Quantizer::create_tensor( const std::vector& shape, DType dtype, std::optional data, std::optional transpose, std::optional scale_inv) const { 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 || nvte_is_non_tn_fp8_gemm_supported(); + 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); @@ -219,7 +236,7 @@ std::pair Float8Quantizer::create_tensor( py::object data_py = with_data ? py::cast(*data) : py::none(); // Initialize transpose tensor - const bool with_transpose = columnwise_usage && !nvte_is_non_tn_fp8_gemm_supported(); + 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); @@ -228,26 +245,58 @@ std::pair Float8Quantizer::create_tensor( transpose.reset(); } py::object transpose_py = with_transpose ? py::cast(*transpose) : py::none(); - // Initialize scale-inverse tensor if (!scale_inv) { 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) { - py::handle Float8TensorClass(reinterpret_cast(Float8TensorStoragePythonClass)); - out_py = Float8TensorClass("data"_a = data_py, "fp8_scale_inv"_a = *scale_inv, - "fp8_dtype"_a = this->dtype, "data_transpose"_a = transpose_py, - "quantizer"_a = this->quantizer); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + py::tuple args(0); + kwargs["data"] = data_py; + kwargs["fp8_scale_inv"] = scale_inv_py; + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["data_transpose"] = transpose_py; + kwargs["quantizer"] = this->quantizer; + + PyObject* result = PyObject_Call(reinterpret_cast(Float8TensorStoragePythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create Float8TensorStorage instance"); + out_py = py::reinterpret_steal(result); } else { - py::handle Float8TensorClass(reinterpret_cast(Float8TensorPythonClass)); const std::vector shape_int64(shape.begin(), shape.end()); - out_py = Float8TensorClass("shape"_a = shape_int64, "dtype"_a = GetATenDType(dtype), - "data"_a = data_py, "fp8_scale_inv"_a = *scale_inv, - "fp8_dtype"_a = this->dtype, "data_transpose"_a = transpose_py, - "quantizer"_a = this->quantizer); + const auto stride_int64 = stride_from_shape(shape_int64); + + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(shape_int64); + kwargs["stride"] = py::cast(stride_int64); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["data"] = data_py; + kwargs["fp8_scale_inv"] = scale_inv_py; + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["data_transpose"] = transpose_py; + kwargs["quantizer"] = this->quantizer; + kwargs["device"] = py::cast(device); + PyObject* result = PyObject_Call(reinterpret_cast(Float8TensorPythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create Float8Tensor instance"); + out_py = py::reinterpret_steal(result); } // Construct C++ FP8 tensor @@ -337,10 +386,10 @@ std::pair Float8Quantizer::create_grouped_tens std::pair Float8Quantizer::convert_and_update_tensor( py::object tensor) const { NVTE_CHECK(detail::IsFloat8Tensor(tensor.ptr()), "Float8Quantizer must output to Float8Tensor."); - + int is_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); // Expected buffers - const bool need_data = rowwise_usage || nvte_is_non_tn_fp8_gemm_supported(); - const bool need_transpose = columnwise_usage && !nvte_is_non_tn_fp8_gemm_supported(); + const bool need_data = rowwise_usage || is_non_tn_fp8_gemm_supported; + const bool need_transpose = columnwise_usage && !is_non_tn_fp8_gemm_supported; NVTE_CHECK(need_data || need_transpose, "Invalid usages for Float8Quantizer."); // Extract buffers from Python tensor @@ -480,7 +529,8 @@ std::pair Float8CurrentScalingQuantizer::create_tenso // Initialize data tensor at::Tensor data_tensor; - const bool with_data = rowwise_usage || nvte_is_non_tn_fp8_gemm_supported(); + int is_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); + 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); @@ -489,13 +539,12 @@ std::pair Float8CurrentScalingQuantizer::create_tenso // Initialize transpose tensor at::Tensor transpose_tensor; - const bool with_transpose = columnwise_usage && !nvte_is_non_tn_fp8_gemm_supported(); + 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); transpose_tensor = at::empty(transpose_shape, opts); } - // Initialize scale-inverse tensor at::Tensor scale_inv_tensor; { @@ -503,23 +552,55 @@ std::pair Float8CurrentScalingQuantizer::create_tenso const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); 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); py::object data_py = with_data ? py::cast(data_tensor) : py::none(); py::object transpose_py = with_transpose ? py::cast(transpose_tensor) : py::none(); if (internal) { - py::handle Float8TensorClass(reinterpret_cast(Float8TensorStoragePythonClass)); - out_py = Float8TensorClass("data"_a = data_py, "fp8_scale_inv"_a = scale_inv_tensor, - "fp8_dtype"_a = this->dtype, "data_transpose"_a = transpose_py, - "quantizer"_a = this->quantizer); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + kwargs["data"] = data_py; + kwargs["fp8_scale_inv"] = scale_inv_py; + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["data_transpose"] = transpose_py; + kwargs["quantizer"] = this->quantizer; + + py::tuple args(0); + PyObject* result = PyObject_Call(reinterpret_cast(Float8TensorStoragePythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create Float8TensorStorage instance"); + out_py = py::reinterpret_steal(result); } else { - py::handle Float8TensorClass(reinterpret_cast(Float8TensorPythonClass)); const std::vector shape_int64(shape.begin(), shape.end()); - out_py = Float8TensorClass("shape"_a = shape_int64, "dtype"_a = GetATenDType(dtype), - "data"_a = data_py, "fp8_scale_inv"_a = scale_inv_tensor, - "fp8_dtype"_a = this->dtype, "data_transpose"_a = transpose_py, - "quantizer"_a = this->quantizer); + const auto stride_int64 = stride_from_shape(shape_int64); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + kwargs["shape"] = py::cast(shape_int64); + kwargs["stride"] = py::cast(stride_int64); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["data"] = data_py; + kwargs["fp8_scale_inv"] = scale_inv_py; + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["data_transpose"] = transpose_py; + kwargs["quantizer"] = this->quantizer; + kwargs["device"] = py::cast(device); + py::tuple args(0); + PyObject* result = PyObject_Call(reinterpret_cast(Float8TensorPythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create Float8Tensor instance"); + out_py = py::reinterpret_steal(result); } // Construct C++ FP8 tensor @@ -627,10 +708,10 @@ std::pair Float8CurrentScalingQuantizer::convert_and_ py::object tensor) const { NVTE_CHECK(detail::IsFloat8Tensor(tensor.ptr()), "Float8CurrentScalingQuantizer must output to Float8Tensor."); - + int is_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); // Expected buffers - const bool need_data = rowwise_usage || nvte_is_non_tn_fp8_gemm_supported(); - const bool need_transpose = columnwise_usage && !nvte_is_non_tn_fp8_gemm_supported(); + const bool need_data = rowwise_usage || is_non_tn_fp8_gemm_supported; + const bool need_transpose = columnwise_usage && !is_non_tn_fp8_gemm_supported; NVTE_CHECK(need_data || need_transpose, "Invalid quantizer usages."); // Extract buffers from Python tensor @@ -837,21 +918,49 @@ std::pair Float8BlockQuantizer::create_tensor( py::object ret; if (internal) { - py::handle Float8BlockwiseQTensorClass( - reinterpret_cast(Float8BlockwiseQTensorStoragePythonClass)); - ret = Float8BlockwiseQTensorClass( - "rowwise_data"_a = data_rowwise, "columnwise_data"_a = data_colwise, - "rowwise_scale_inv"_a = scale_inv_rowwise, "columnwise_scale_inv"_a = scale_inv_colwise, - "fp8_dtype"_a = this->dtype, "quantizer"_a = this->quantizer, - "is_2D_scaled"_a = (block_scaling_dim == 2)); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + kwargs["rowwise_data"] = py::cast(data_rowwise); + kwargs["columnwise_data"] = py::cast(data_colwise); + kwargs["rowwise_scale_inv"] = py::cast(scale_inv_rowwise); + kwargs["columnwise_scale_inv"] = py::cast(scale_inv_colwise); + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["quantizer"] = this->quantizer; + kwargs["is_2D_scaled"] = py::cast(block_scaling_dim == 2); + + py::tuple args(0); + PyObject* result = + PyObject_Call(reinterpret_cast(Float8BlockwiseQTensorStoragePythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create Float8BlockwiseQTensorStorage instance"); + ret = py::reinterpret_steal(result); } else { - py::handle Float8BlockwiseQTensorClass( - reinterpret_cast(Float8BlockwiseQTensorPythonClass)); - ret = Float8BlockwiseQTensorClass( - "shape"_a = torch_shape, "dtype"_a = GetATenDType(dtype), "rowwise_data"_a = data_rowwise, - "columnwise_data"_a = data_colwise, "rowwise_scale_inv"_a = scale_inv_rowwise, - "columnwise_scale_inv"_a = scale_inv_colwise, "fp8_dtype"_a = this->dtype, - "quantizer"_a = this->quantizer, "is_2D_scaled"_a = (block_scaling_dim == 2)); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + const auto stride_int64 = stride_from_shape(torch_shape); + kwargs["shape"] = py::cast(torch_shape); + kwargs["stride"] = py::cast(stride_int64); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["rowwise_data"] = py::cast(data_rowwise); + kwargs["columnwise_data"] = py::cast(data_colwise); + kwargs["rowwise_scale_inv"] = py::cast(scale_inv_rowwise); + kwargs["columnwise_scale_inv"] = py::cast(scale_inv_colwise); + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["quantizer"] = this->quantizer; + kwargs["is_2D_scaled"] = py::cast(block_scaling_dim == 2); + + py::tuple args(0); + PyObject* result = PyObject_Call(reinterpret_cast(Float8BlockwiseQTensorPythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create Float8BlockwiseQTensor instance"); + ret = py::reinterpret_steal(result); } return {std::move(tensor), std::move(ret)}; @@ -1198,18 +1307,49 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve // Construct Python MXFP8 tensor py::object out_py; if (internal) { - py::handle MXFP8TensorClass(reinterpret_cast(MXFP8TensorStoragePythonClass)); - out_py = MXFP8TensorClass(rowwise_data_py, rowwise_scale_inv_py, columnwise_data_py, - columnwise_scale_inv_py, this->dtype, this->quantizer, - with_gemm_swizzled_scales); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + py::tuple args(0); + kwargs["rowwise_data"] = rowwise_data_py; + kwargs["columnwise_data"] = columnwise_data_py; + kwargs["rowwise_scale_inv"] = rowwise_scale_inv_py; + kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["quantizer"] = this->quantizer; + kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + + PyObject* result = PyObject_Call(reinterpret_cast(MXFP8TensorStoragePythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create MXFP8TensorStorage instance"); + out_py = py::reinterpret_steal(result); } else { - py::handle MXFP8TensorClass(reinterpret_cast(MXFP8TensorPythonClass)); - out_py = MXFP8TensorClass( - "shape"_a = shape_int64, "dtype"_a = GetATenDType(dtype), - "rowwise_data"_a = rowwise_data_py, "columnwise_data"_a = columnwise_data_py, - "rowwise_scale_inv"_a = rowwise_scale_inv_py, - "columnwise_scale_inv"_a = columnwise_scale_inv_py, "fp8_dtype"_a = this->dtype, - "quantizer"_a = this->quantizer, "with_gemm_swizzled_scales"_a = with_gemm_swizzled_scales); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + const auto stride_int64 = stride_from_shape(shape_int64); + kwargs["shape"] = py::cast(shape_int64); + kwargs["stride"] = py::cast(stride_int64); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["rowwise_data"] = rowwise_data_py; + kwargs["columnwise_data"] = columnwise_data_py; + kwargs["rowwise_scale_inv"] = rowwise_scale_inv_py; + kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; + kwargs["fp8_dtype"] = py::cast(this->dtype); + kwargs["quantizer"] = this->quantizer; + kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + + py::tuple args(0); + PyObject* result = PyObject_Call(reinterpret_cast(MXFP8TensorPythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create MXFP8Tensor instance"); + out_py = py::reinterpret_steal(result); } // Construct C++ MXFP8 tensor @@ -1561,19 +1701,53 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve // Construct Python NVFP4 tensor py::object out_py; if (internal) { - py::handle NVFP4TensorClass(reinterpret_cast(NVFP4TensorStoragePythonClass)); - out_py = NVFP4TensorClass(rowwise_data_py, rowwise_scale_inv_py, columnwise_data_py, - columnwise_scale_inv_py, amax_rowwise_py, amax_columnwise_py, - this->dtype, this->quantizer, with_gemm_swizzled_scales); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + kwargs["rowwise_data"] = rowwise_data_py; + kwargs["columnwise_data"] = columnwise_data_py; + kwargs["rowwise_scale_inv"] = rowwise_scale_inv_py; + kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; + kwargs["amax_rowwise"] = amax_rowwise_py; + kwargs["amax_columnwise"] = amax_columnwise_py; + kwargs["fp4_dtype"] = py::cast(this->dtype); + kwargs["quantizer"] = this->quantizer; + kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + + py::tuple args(0); + + PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorStoragePythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create NVFP4TensorStorage instance"); + out_py = py::reinterpret_steal(result); } else { - py::handle NVFP4TensorClass(reinterpret_cast(NVFP4TensorPythonClass)); - out_py = NVFP4TensorClass( - "shape"_a = shape_int64, "dtype"_a = GetATenDType(dtype), - "rowwise_data"_a = rowwise_data_py, "columnwise_data"_a = columnwise_data_py, - "rowwise_scale_inv"_a = rowwise_scale_inv_py, - "columnwise_scale_inv"_a = columnwise_scale_inv_py, "amax_rowwise"_a = amax_rowwise_py, - "amax_columnwise"_a = amax_columnwise_py, "fp4_dtype"_a = this->dtype, - "quantizer"_a = this->quantizer, "with_gemm_swizzled_scales"_a = with_gemm_swizzled_scales); + // Use direct C API call bypassing pybind11 overhead + py::dict kwargs; + const auto stride_int64 = stride_from_shape(shape_int64); + kwargs["shape"] = py::cast(shape_int64); + kwargs["stride"] = py::cast(stride_int64); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["rowwise_data"] = rowwise_data_py; + kwargs["columnwise_data"] = columnwise_data_py; + kwargs["rowwise_scale_inv"] = rowwise_scale_inv_py; + kwargs["columnwise_scale_inv"] = columnwise_scale_inv_py; + kwargs["amax_rowwise"] = amax_rowwise_py; + kwargs["amax_columnwise"] = amax_columnwise_py; + kwargs["fp4_dtype"] = py::cast(this->dtype); + kwargs["quantizer"] = this->quantizer; + kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + py::tuple args(0); + PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), + args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + + NVTE_CHECK(result != nullptr, "Failed to create NVFP4Tensor instance"); + out_py = py::reinterpret_steal(result); } // Construct C++ tensor diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 4858383c26..9c21141a39 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -929,12 +929,11 @@ def set_activation_dtype(self, inp: torch.Tensor) -> None: if torch.is_autocast_enabled(): self.fast_setattr("activation_dtype", torch_get_autocast_gpu_dtype()) return - + dtype = inp.dtype # All checks after this have already been performed once, thus skip - if self.activation_dtype == inp.dtype: + if self.activation_dtype == dtype: return - dtype = inp.dtype if not self.allow_different_data_and_param_types: for name, param in self.named_parameters(): if param is not None: diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 27632db15b..ce0581024a 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -52,7 +52,7 @@ _fsdp_scatter_tensors, _fsdp_gather_tensors, ) -from ..constants import GemmParallelModes, dist_group_type +from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing from ._common import apply_normalization, noop_cat, WeightGradStore @@ -1357,7 +1357,7 @@ def __init__( torch.nn.Parameter(weight_tensor[split_start:split_end]), init_fn=init_method, get_rng_state_tracker=get_rng_state_tracker, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, + fp8_meta_index=FP8FwdTensorIdx.GEMM1_WEIGHT, ) # Construct bias parameters if needed @@ -1615,20 +1615,20 @@ def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): grad_weight_quantizer = None grad_output_quantizer = None output_quantizer = None - input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] + input_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] input_quantizer.internal = True if not (self.parallel_mode == "column" and self.sequence_parallel): input_quantizer.optimize_for_gemm = True (weight_quantizer,) = self._get_weight_quantizers() if fp8_output: - output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] + output_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] if is_grad_enabled: - grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] + grad_output_quantizer = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] grad_output_quantizer.internal = True if not (self.parallel_mode == "row" and self.sequence_parallel): grad_output_quantizer.optimize_for_gemm = True if fp8_grad: - grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] + grad_input_quantizer = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] return ( input_quantizer, @@ -1725,43 +1725,43 @@ def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe if fwd: # set configs about amax epsilon and power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon # also set weight quantizer with same amax_epsilon & power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_WEIGHT + FP8FwdTensorIdx.GEMM1_WEIGHT ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_WEIGHT + FP8FwdTensorIdx.GEMM1_WEIGHT ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon # parallel related if self.sequence_parallel and self.parallel_mode == "column": # set input_quantizer with amax reduction TP group self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].with_amax_reduction = True self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_reduction_group = self.tp_group else: # set grad_output_quantizer with amax epsilon and power_2_scale (no amax reduction here) self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon # parallel related if self.sequence_parallel and self.parallel_mode == "row": # customize grad_output_quantizer with amax reduction TP group self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].with_amax_reduction = True self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_reduction_group = self.tp_group def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None: @@ -1771,19 +1771,19 @@ def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None: if self.sequence_parallel and self.parallel_mode == "column": # set input_quantizer with amax reduction TP group self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].with_amax_reduction = True self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_reduction_group = self.tp_group else: if self.sequence_parallel and self.parallel_mode == "row": # customize grad_output_quantizer with amax reduction TP group self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].with_amax_reduction = True self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_reduction_group = self.tp_group def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]: @@ -1807,6 +1807,6 @@ def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" if not self.fp8 and not self.fp8_calibration: return [None] - weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] + weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] weight_quantizer.internal = True return [weight_quantizer] diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index b8823e46ca..16e620fd94 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -59,7 +59,7 @@ _get_cuda_rng_state, _set_cuda_rng_state, ) -from ..constants import dist_group_type +from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing from ..tensor.float8_tensor import ( @@ -1909,7 +1909,7 @@ def __init__( fc1_weight, init_fn=init_method, get_rng_state_tracker=get_rng_state_tracker, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, + fp8_meta_index=FP8FwdTensorIdx.GEMM1_WEIGHT, ) if self.use_bias: @@ -1929,7 +1929,7 @@ def __init__( fc2_weight, init_fn=output_layer_init_method, get_rng_state_tracker=get_rng_state_tracker, - fp8_meta_index=tex.FP8FwdTensors.GEMM2_WEIGHT, + fp8_meta_index=FP8FwdTensorIdx.GEMM2_WEIGHT, ) if self.use_bias: @@ -2201,11 +2201,11 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): ) = [None] * 10 fc1_weight_quantizer, fc2_weight_quantizer = self._get_weight_quantizers() if self.fp8 or self.fp8_calibration: - fc1_input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] + fc1_input_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] fc1_input_quantizer.internal = True if not self.sequence_parallel: fc1_input_quantizer.optimize_for_gemm = True - fc2_input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM2_INPUT] + fc2_input_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_INPUT] fc2_input_quantizer.set_usage( rowwise=True, columnwise=isinstance( @@ -2216,18 +2216,16 @@ def _get_quantizers(self, fp8_output, is_grad_enabled): fc2_input_quantizer.internal = True fc2_input_quantizer.optimize_for_gemm = True if fp8_output: - fc2_output_quantizer = self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM2_OUTPUT - ] + fc2_output_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_OUTPUT] if is_grad_enabled: fc2_grad_output_quantizer = self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ] fc2_grad_output_quantizer.internal = True if not self.sequence_parallel: fc2_grad_output_quantizer.optimize_for_gemm = True fc1_grad_output_quantizer = self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ] fc1_grad_output_quantizer.internal = True fc1_grad_output_quantizer.optimize_for_gemm = True @@ -2389,63 +2387,63 @@ def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe if fwd: # fc1_input_quantizer: set configs about amax epsilon and power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon # fc2_input_quantizer self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM2_INPUT + FP8FwdTensorIdx.GEMM2_INPUT ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM2_INPUT + FP8FwdTensorIdx.GEMM2_INPUT ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon # fc1_weight_quantizer: also set numerical configs about weight self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_WEIGHT + FP8FwdTensorIdx.GEMM1_WEIGHT ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_WEIGHT + FP8FwdTensorIdx.GEMM1_WEIGHT ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon # fc2_weight_quantizer self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM2_WEIGHT + FP8FwdTensorIdx.GEMM2_WEIGHT ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM2_WEIGHT + FP8FwdTensorIdx.GEMM2_WEIGHT ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon # parallel related if self.sequence_parallel and self.set_parallel_mode: # fc1_input_quantizer: customize input_quantizer with amax reduction TP group, column parallel + sequence parallel here self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].with_amax_reduction = True self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_reduction_group = self.tp_group else: # fc2_grad_output_quantizer: set configs about amax epsilon and power_2_scale for fc2_grad_output_quantizer self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon # fc1_grad_output_quantizer: also set numerical configs for fc1_grad_output_quantizer self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon if self.sequence_parallel and self.set_parallel_mode: # fc2_grad_output_quantizer: customize grad_output_quantizer with amax reduction TP group, row parallel + sequence parallel here self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ].with_amax_reduction = True self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ].amax_reduction_group = self.tp_group def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None: @@ -2455,19 +2453,19 @@ def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None: if self.sequence_parallel and self.set_parallel_mode: # fc1_input_quantizer: customize input_quantizer with amax reduction TP group, column parallel + sequence parallel here self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].with_amax_reduction = True self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_reduction_group = self.tp_group else: if self.sequence_parallel and self.set_parallel_mode: # fc2_grad_output_quantizer: customize grad_output_quantizer with amax reduction TP group, row parallel + sequence parallel here self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ].with_amax_reduction = True self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT2 + FP8BwdTensorIdx.GRAD_OUTPUT2 ].amax_reduction_group = self.tp_group def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]: @@ -2478,9 +2476,9 @@ def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" if not self.fp8 and not self.fp8_calibration: return [None, None] - fc1_weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] + fc1_weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] fc1_weight_quantizer.internal = True - fc2_weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM2_WEIGHT] + fc2_weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM2_WEIGHT] fc2_weight_quantizer.internal = True return [fc1_weight_quantizer, fc2_weight_quantizer] diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index a55429d33d..31dac4d329 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -54,7 +54,7 @@ from ..cpp_extensions import ( general_gemm, ) -from ..constants import GemmParallelModes, dist_group_type +from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing from ..quantized_tensor import ( @@ -1272,7 +1272,7 @@ def __init__( torch.nn.Parameter(weight_tensor[split_start:split_end]), init_fn=init_method, get_rng_state_tracker=get_rng_state_tracker, - fp8_meta_index=tex.FP8FwdTensors.GEMM1_WEIGHT, + fp8_meta_index=FP8FwdTensorIdx.GEMM1_WEIGHT, ) # Construct bias parameters if needed @@ -1483,20 +1483,20 @@ def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): grad_weight_quantizer = None grad_output_quantizer = None output_quantizer = None - input_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] + input_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] input_quantizer.internal = True if not (self.parallel_mode == "column" and self.sequence_parallel): input_quantizer.optimize_for_gemm = True (weight_quantizer,) = self._get_weight_quantizers() if fp8_output: - output_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_OUTPUT] + output_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_OUTPUT] if is_grad_enabled: - grad_output_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_OUTPUT1] + grad_output_quantizer = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] grad_output_quantizer.internal = True if not (self.parallel_mode == "row" and self.sequence_parallel): grad_output_quantizer.optimize_for_gemm = True if fp8_grad: - grad_input_quantizer = self.quantizers["scaling_bwd"][tex.FP8BwdTensors.GRAD_INPUT1] + grad_input_quantizer = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] return ( input_quantizer, weight_quantizer, @@ -1601,43 +1601,43 @@ def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe if fwd: # set configs about amax epsilon and power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_epsilon = recipe.fp8_quant_fwd_inp.amax_epsilon # also set weight quantizer with same amax_epsilon & power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_WEIGHT + FP8FwdTensorIdx.GEMM1_WEIGHT ].force_pow_2_scales = recipe.fp8_quant_fwd_weight.power_2_scale self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_WEIGHT + FP8FwdTensorIdx.GEMM1_WEIGHT ].amax_epsilon = recipe.fp8_quant_fwd_weight.amax_epsilon # paralle related if self.sequence_parallel and self.parallel_mode == "column": # customize input_quantizer with amax reduction TP group self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].with_amax_reduction = True self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_reduction_group = self.tp_group else: # set grad_output_quantizer with amax epsilon and power_2_scale self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].force_pow_2_scales = recipe.fp8_quant_bwd_grad.power_2_scale self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_epsilon = recipe.fp8_quant_bwd_grad.amax_epsilon # parallel related if self.sequence_parallel and self.parallel_mode == "row": # customize grad_output_quantizer with amax reduction TP group self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].with_amax_reduction = True self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_reduction_group = self.tp_group def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None: @@ -1647,25 +1647,25 @@ def _customize_quantizers_nvfp4(self, fwd: bool, recipe: Recipe) -> None: if self.sequence_parallel and self.parallel_mode == "column": # customize input_quantizer with amax reduction TP group self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].with_amax_reduction = True self.quantizers["scaling_fwd"][ - tex.FP8FwdTensors.GEMM1_INPUT + FP8FwdTensorIdx.GEMM1_INPUT ].amax_reduction_group = self.tp_group else: if self.sequence_parallel and self.parallel_mode == "row": # customize grad_output_quantizer with amax reduction TP group self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].with_amax_reduction = True self.quantizers["scaling_bwd"][ - tex.FP8BwdTensors.GRAD_OUTPUT1 + FP8BwdTensorIdx.GRAD_OUTPUT1 ].amax_reduction_group = self.tp_group def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" if not self.fp8 and not self.fp8_calibration: return [None] - weight_quantizer = self.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_WEIGHT] + weight_quantizer = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] weight_quantizer.internal = True return [weight_quantizer] diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index d78677bc83..cb697bc197 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -369,9 +369,13 @@ def __new__( *, requires_grad: bool = False, device: Optional[torch.device] = None, + stride: Optional[Iterable[int]] = None, ): - # We are assuming only contiguous tensors - stride = _stride_from_shape(shape) + # For stride, We are assuming only contiguous tensors + # Calculate stride from shape if not provided. When creating this object from + # C++ code, we provide the stride computed from shape in C++ to avoid the + # PyobjectVectorCall overhead of calling _stride_from_shape from C++ to Python. + stride = _stride_from_shape(shape) if stride is None else stride instance = torch.Tensor._make_wrapper_subclass( cls, shape, @@ -382,9 +386,75 @@ def __new__( requires_grad=requires_grad, device=torch.cuda.current_device() if device is None else device, ) - + instance._requires_grad = requires_grad + instance._dtype = dtype return instance + @property + def dtype(self) -> torch.dtype: + """ + Return the high precision data type of the tensor + Attribute access of custom tensors goes through an + expensive Pyobject lookup. Since dtype for a tensor is never + change after creation, we cache it in a member variable and return + """ + # Lazy initialization for tensors created via alternate paths + if not hasattr(self, "_dtype"): + # pylint: disable=unnecessary-dunder-call + self._dtype = torch._C.TensorBase.dtype.__get__(self, type(self)) + return self._dtype + + @dtype.setter + def dtype(self, value: torch.dtype) -> None: + """Set dtype property""" + self._dtype = value + + @property + def requires_grad(self) -> bool: + """ + Return whether or not the tensor requires gradient. + Attribute access of custom tensors goes through an + expensive Pyobject lookup. Since requires_grad is set during + initialization and may be updated, we cache it in a member variable. + """ + # Fallback to parent if not cached yet + if not hasattr(self, "_requires_grad"): + # pylint: disable=unnecessary-dunder-call + self._requires_grad = torch._C.TensorBase.requires_grad.__get__(self, type(self)) + return self._requires_grad + + @requires_grad.setter + def requires_grad(self, value: bool) -> None: + """Set requires_grad property so that autograd engine is aware of the change""" + # Update the cached value and call parent class method to ensure autograd engine is aware + self.requires_grad_(value) + + def requires_grad_(self, requires_grad: bool = True) -> QuantizedTensor: + """Cache requires_grad property and call parent class method""" + # pylint: disable=missing-function-docstring + # Update the cached value + self._requires_grad = requires_grad + # Call parent class method to ensure autograd engine is aware + super().requires_grad_(requires_grad) + return self + + def _get_data(self) -> torch.Tensor: + """Get tensor data property""" + return super().data + + def _set_data(self, tensor: torch.Tensor) -> None: + """Set tensor data property + Updates the underlying tensor data and syncs the dtype cache. + """ + # Update the parent class's data descriptor + # pylint: disable=unnecessary-dunder-call + super(QuantizedTensor, type(self)).data.__set__(self, tensor) + # Update the dtype cache + self._dtype = tensor.dtype + + # Create the data property with getter and setter + data = property(_get_data, _set_data) + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Convert quantized data to standard PyTorch tensor""" raise NotImplementedError( diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ecafb6ddfc..a3d49ea4e9 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -567,6 +567,24 @@ def _set_from_tensor(dst: Float8BlockwiseQTensor, src: Float8BlockwiseQTensor): # Cast to FP8 when setting Float8BlockwiseQTensor.data data = property(_get_data, _set_data) + @property + def shape(self): + """Return the shape of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.shape + if self._columnwise_data is not None: + return self._columnwise_data.shape + raise RuntimeError("Float8BlockwiseQTensor has no data!") + + @property + def is_cuda(self): + """Return whether the tensor is on a CUDA device.""" + if self._rowwise_data is not None: + return self._rowwise_data.is_cuda + if self._columnwise_data is not None: + return self._columnwise_data.is_cuda + raise RuntimeError("Float8BlockwiseQTensor has no data!") + class _ViewFunc(torch.autograd.Function): """View function diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 55bca49af3..f66e88740f 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -926,6 +926,25 @@ def fsdp_post_all_gather( ) return out, all_gather_outputs + @property + def shape(self): + """Return the shape of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._data is not None: + return self._data.shape + if self._transpose is not None: + transpose_shape = self._transpose.shape + return torch.Size(tuple(transpose_shape[1:]) + (transpose_shape[0],)) + raise RuntimeError("Both data and transpose are None") + + @property + def is_cuda(self): + """Return whether the tensor is on a CUDA device.""" + if self._data is not None: + return self._data.is_cuda + if self._transpose is not None: + return self._transpose.is_cuda + raise RuntimeError("Both data and transpose are None") + @classmethod def _make_in_reduce_ex( cls, diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 41d6c87f2b..96b6a67ea8 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -842,6 +842,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: ) # pylint: disable=unnecessary-dunder-call super(MXFP8Tensor, type(self)).data.__set__(self, dummy_tensor) + self._rowwise_data = tensor._rowwise_data self._columnwise_data = tensor._columnwise_data self._quantizer = tensor._quantizer.copy() @@ -861,6 +862,33 @@ def _set_data(self, tensor: torch.Tensor) -> None: # Cast to FP8 when setting MXFP8Tensor.data data = property(_get_data, _set_data) + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.device + if self._columnwise_data is not None: + return self._columnwise_data.device + raise RuntimeError("MXFP8Tensor has no data!") + + @property + def shape(self): + """Return the shape of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.shape + if self._columnwise_data is not None: + return self._columnwise_data.shape + raise RuntimeError("MXFP8Tensor has no data!") + + @property + def is_cuda(self): + """Return whether the tensor is on a CUDA device.""" + if self._rowwise_data is not None: + return self._rowwise_data.is_cuda + if self._columnwise_data is not None: + return self._columnwise_data.is_cuda + raise RuntimeError("MXFP8Tensor has no data!") + class _ViewFunc(torch.autograd.Function): """View function diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 66f986a900..a8148b5752 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -700,6 +700,7 @@ def _set_data(self, tensor: torch.Tensor) -> None: ) # pylint: disable=unnecessary-dunder-call super(NVFP4Tensor, type(self)).data.__set__(self, dummy_tensor) + self._rowwise_data = tensor._rowwise_data self._columnwise_data = tensor._columnwise_data self._quantizer = tensor._quantizer @@ -719,6 +720,35 @@ def _set_data(self, tensor: torch.Tensor) -> None: # Cast to FP8 when setting NVFP4Tensor.data data = property(_get_data, _set_data) + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.device + if self._columnwise_data is not None: + return self._columnwise_data.device + raise RuntimeError("NVFP4Tensor has no data!") + + @property + def shape(self): + """Return the shape of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + byte_shape = self._rowwise_data.shape + return torch.Size(byte_shape[:-1] + (byte_shape[-1] * 2,)) + if self._columnwise_data is not None: + byte_shape = self._columnwise_data.shape + return torch.Size(byte_shape[1:-1] + (byte_shape[-1] * 2, byte_shape[0])) + raise RuntimeError("NVFP4Tensor has no data!") + + @property + def is_cuda(self): + """Return whether the tensor is on a CUDA device.""" + if self._rowwise_data is not None: + return self._rowwise_data.is_cuda + if self._columnwise_data is not None: + return self._columnwise_data.is_cuda + raise RuntimeError("NVFP4Tensor has no data!") + class _ViewFunc(torch.autograd.Function): """View function diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 4cd6d19cd8..2a86717017 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -290,6 +290,15 @@ def size(self, *args, **kwargs): reordered.append(dims[0]) return torch.Size(reordered) + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.device + if self._columnwise_data is not None: + return self._columnwise_data.device + raise RuntimeError("Float8BlockwiseQTensorStorage has no data!") + def _create_columnwise(self): """ Update columnwise data and columnwise scale inv. Can only be used when using 2D scaling. diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 9adb86c453..a815b366b2 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -170,6 +170,15 @@ def size(self, *args, **kwargs): size = self._transpose.size(*args, **kwargs) return torch.Size([size[-1], math.prod(size[:-1])]) + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._data is not None: + return self._data.device + if self._transpose is not None: + return self._transpose.device + raise RuntimeError("Float8TensorStorage has no data!") + def view(self, shape: torch.Size): # pylint: disable=missing-function-docstring out_data = self._data.view(shape) diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 5c8510488f..12757aa58c 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -185,6 +185,15 @@ def size(self, *args, **kwargs): return self._rowwise_data.size(*args, **kwargs) return self._columnwise_data.size(*args, **kwargs) + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.device + if self._columnwise_data is not None: + return self._columnwise_data.device + raise RuntimeError("MXFP8TensorStorage has no data!") + def view(self, shape: torch.Size): # pylint: disable=missing-function-docstring diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 8be23d0c19..36bf208bcd 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -228,6 +228,15 @@ def size(self, dim: Optional[int] = None) -> Union[torch.Size, int]: return torch.Size(shape) return shape[dim] + @property + def device(self): + """Return the device of the tensor. Define this to avoid expensive PyObject lookups.""" + if self._rowwise_data is not None: + return self._rowwise_data.device + if self._columnwise_data is not None: + return self._columnwise_data.device + raise RuntimeError("NVFP4TensorStorage has no data!") + def view(self, shape: torch.Size): # pylint: disable=missing-function-docstring From c68ec3101d0dc16fe6eb40294a5fed3a9370b6a8 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 2 Mar 2026 18:16:53 -0800 Subject: [PATCH 238/521] Add fast_set_attr to modules not inheriting from base.py (#2724) fix fast_set_attr in other nn modules for fsdp Signed-off-by: Varun Thumbe --- .../pytorch/attention/dot_product_attention/backends.py | 4 ++++ .../pytorch/attention/multi_head_attention.py | 6 +++++- transformer_engine/pytorch/module/layernorm.py | 6 +++++- transformer_engine/pytorch/module/rmsnorm.py | 6 +++++- transformer_engine/pytorch/transformer.py | 6 +++++- 5 files changed, 24 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index aa6c063951..a6a8b0b26a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -293,6 +293,10 @@ def mask_func(x, y): bool(int(os.getenv("NVTE_APPLY_QK_LAYER_SCALING", "0"))) and layer_number is not None ) + def fast_setattr(self, name: str, value: Any) -> None: + """Fast attribute set for non-parameter fields.""" + self.__dict__[name] = value + def forward( self, _alibi_cache: Dict[str, Any], diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 01c4955d78..5c581849e6 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -5,7 +5,7 @@ """Multi-head Attention.""" import os import collections -from typing import Callable, List, Optional, Tuple, Union +from typing import Any, Callable, List, Optional, Tuple, Union import torch from transformer_engine.pytorch.quantization import FP8GlobalStateManager @@ -478,6 +478,10 @@ def __init__( **common_gemm_kwargs, ) + def fast_setattr(self, name: str, value: Any) -> None: + """Fast attribute set for non-parameter fields.""" + self.__dict__[name] = value + def _create_qk_norm_modules( self, qk_norm_type: Optional[str], diff --git a/transformer_engine/pytorch/module/layernorm.py b/transformer_engine/pytorch/module/layernorm.py index d4f0a78ba2..54fad8d1bc 100644 --- a/transformer_engine/pytorch/module/layernorm.py +++ b/transformer_engine/pytorch/module/layernorm.py @@ -4,7 +4,7 @@ """LayerNorm API""" import warnings -from typing import Iterable, Optional, Union +from typing import Any, Iterable, Optional, Union import torch @@ -102,6 +102,10 @@ def __init__( **kwargs, ) + def fast_setattr(self, name: str, value: Any) -> None: + """Fast attribute set for non-parameter fields.""" + self.__dict__[name] = value + def reset_layer_norm_parameters(self) -> None: """Init LN params""" warnings.warn( diff --git a/transformer_engine/pytorch/module/rmsnorm.py b/transformer_engine/pytorch/module/rmsnorm.py index ace4be31de..f8d5aade5c 100644 --- a/transformer_engine/pytorch/module/rmsnorm.py +++ b/transformer_engine/pytorch/module/rmsnorm.py @@ -4,7 +4,7 @@ """RMSNorm API""" import warnings -from typing import Iterable, Optional, Union +from typing import Any, Iterable, Optional, Union import torch @@ -106,6 +106,10 @@ def __init__( **kwargs, ) + def fast_setattr(self, name: str, value: Any) -> None: + """Fast attribute set for non-parameter fields.""" + self.__dict__[name] = value + def reset_rms_norm_parameters(self) -> None: """Deprecated""" warnings.warn( diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index cf7ce5e1a4..868cbbdac8 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -6,7 +6,7 @@ import os import warnings from contextlib import nullcontext -from typing import Callable, List, Optional, Tuple, Union +from typing import Any, Callable, List, Optional, Tuple, Union import torch @@ -545,6 +545,10 @@ def __init__( device=device, ) + def fast_setattr(self, name: str, value: Any) -> None: + """Fast attribute set for non-parameter fields.""" + self.__dict__[name] = value + def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None: """ Set the tensor parallel group for the given From 39d249baab612b3c8b5d32d616603505d2e00282 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 3 Mar 2026 08:44:57 -0800 Subject: [PATCH 239/521] [JAX] Remove GSPMD tests + adding guards and warning msg for GSPMD rules (#2702) rm gspmd tests + add deprecated warnings for gspmd rules Signed-off-by: Phuong Nguyen --- .../run_test_multiprocessing_encoder.sh | 4 - .../encoder/test_model_parallel_encoder.py | 68 --------- examples/jax/encoder/test_multigpu_encoder.py | 47 ------- .../encoder/test_multiprocessing_encoder.py | 45 +----- tests/jax/test_distributed_fused_attn.py | 130 ------------------ tests/jax/test_distributed_layernorm.py | 6 - tests/jax/test_distributed_layernorm_mlp.py | 89 ------------ tests/jax/test_distributed_permutation.py | 12 -- tests/jax/test_distributed_softmax.py | 34 ----- transformer_engine/jax/cpp_extensions/base.py | 42 +++++- 10 files changed, 38 insertions(+), 439 deletions(-) diff --git a/examples/jax/encoder/run_test_multiprocessing_encoder.sh b/examples/jax/encoder/run_test_multiprocessing_encoder.sh index f2ef33da46..3c1f2ba1fb 100644 --- a/examples/jax/encoder/run_test_multiprocessing_encoder.sh +++ b/examples/jax/encoder/run_test_multiprocessing_encoder.sh @@ -11,10 +11,6 @@ TEST_CASES=( "test_te_current_scaling_fp8" "test_te_mxfp8" "test_te_nvfp4" -"test_te_bf16_shardy" -"test_te_delayed_scaling_fp8_shardy" -"test_te_current_scaling_fp8_shardy" -"test_te_nvfp4_shardy" ) : ${TE_PATH:=/opt/transformerengine} diff --git a/examples/jax/encoder/test_model_parallel_encoder.py b/examples/jax/encoder/test_model_parallel_encoder.py index 73b93798a0..4400485f26 100644 --- a/examples/jax/encoder/test_model_parallel_encoder.py +++ b/examples/jax/encoder/test_model_parallel_encoder.py @@ -239,7 +239,6 @@ def check_fp8(state, var_collect, inputs, masks, labels): def train_and_evaluate(args): """Execute model training and evaluation loop.""" print(args) - jax.config.update("jax_use_shardy_partitioner", args.enable_shardy) train_ds, test_ds, num_embed = get_datasets(args.max_seq_len) @@ -474,9 +473,6 @@ def encoder_parser(args): parser.add_argument( "--enable-sp", action="store_true", default=False, help="Enable sequence parallelism." ) - parser.add_argument( - "--enable-shardy", action="store_true", default=False, help="Enable Shardy (experimental)." - ) return parser.parse_args(args) @@ -559,70 +555,6 @@ def test_te_nvfp4_with_sp(self): actual = train_and_evaluate(self.args) assert actual[0] < 0.40 and actual[1] > 0.82 - @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") - def test_te_bf16_shardy(self): - """Test Transformer Engine with BF16""" - self.args.enable_shardy = True - actual = train_and_evaluate(self.args) - assert actual[0] < 0.36 and actual[1] > 0.84 - - @unittest.skipIf(not is_fp8_supported, fp8_reason) - def test_te_delayed_scaling_fp8_shardy(self): - """Test Transformer Engine with DelayedScaling FP8""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "DelayedScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.362 and actual[1] > 0.84 - - @unittest.skipIf(not is_fp8_supported, fp8_reason) - def test_te_delayed_scaling_fp8_with_sp_shardy(self): - """Test Transformer Engine with DelayedScaling FP8 + SP""" - self.args.enable_shardy = True - self.args.enable_sp = True - self.args.use_fp8 = True - self.args.fp8_recipe = "DelayedScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.362 and actual[1] > 0.84 - - @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) - def test_te_mxfp8_shardy(self): - """Test Transformer Engine with MXFP8""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "MXFP8BlockScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.36 and actual[1] > 0.84 - - @unittest.skipIf(not is_nvfp4_supported, nvfp4_reason) - def test_te_nvfp4_shardy(self): - """Test Transformer Engine with NVFP4""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "NVFP4BlockScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.40 and actual[1] > 0.82 - - @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) - def test_te_mxfp8_with_sp_shardy(self): - """Test Transformer Engine with MXFP8 + SP""" - self.args.enable_shardy = True - self.args.enable_sp = True - self.args.use_fp8 = True - self.args.fp8_recipe = "MXFP8BlockScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.36 and actual[1] > 0.84 - - @unittest.skipIf(not is_nvfp4_supported, nvfp4_reason) - def test_te_nvfp4_with_sp_shardy(self): - """Test Transformer Engine with NVFP4""" - self.args.enable_shardy = True - self.args.enable_sp = True - self.args.use_fp8 = True - self.args.fp8_recipe = "NVFP4BlockScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.40 and actual[1] > 0.82 - if __name__ == "__main__": train_and_evaluate(encoder_parser(None)) diff --git a/examples/jax/encoder/test_multigpu_encoder.py b/examples/jax/encoder/test_multigpu_encoder.py index 22a89cc0a9..e2edc589b9 100644 --- a/examples/jax/encoder/test_multigpu_encoder.py +++ b/examples/jax/encoder/test_multigpu_encoder.py @@ -249,7 +249,6 @@ def replace_params(x): def train_and_evaluate(args): """Execute model training and evaluation loop.""" print(args) - jax.config.update("jax_use_shardy_partitioner", args.enable_shardy) train_ds, test_ds, num_embed = get_datasets(args.max_seq_len) num_gpu = jax.local_device_count() @@ -438,9 +437,6 @@ def encoder_parser(args): default="DelayedScaling", help="Use FP8 recipe (default: DelayedScaling)", ) - parser.add_argument( - "--enable-shardy", action="store_true", default=False, help="Enable Shardy (experimental)." - ) return parser.parse_args(args) @@ -494,49 +490,6 @@ def test_te_nvfp4(self): actual = train_and_evaluate(self.args) assert actual[0] < 0.52 and actual[1] > 0.74 - @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") - def test_te_bf16_shardy(self): - """Test Transformer Engine with BF16""" - self.args.enable_shardy = True - actual = train_and_evaluate(self.args) - assert actual[0] < 0.51 and actual[1] > 0.75 - - @unittest.skipIf(not is_fp8_supported, fp8_reason) - def test_te_delayed_scaling_fp8_shardy(self): - """Test Transformer Engine with DelayedScaling FP8""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "DelayedScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.51 and actual[1] > 0.75 - - @unittest.skipIf(not is_fp8_supported, fp8_reason) - def test_te_current_scaling_fp8_shardy(self): - """Test Transformer Engine with CurrentScaling FP8""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "Float8CurrentScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.51 and actual[1] > 0.749 - - @unittest.skipIf(not is_mxfp8_supported, mxfp8_reason) - def test_te_mxfp8_shardy(self): - """Test Transformer Engine with MXFP8""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "MXFP8BlockScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.51 and actual[1] > 0.75 - - @unittest.skipIf(not is_nvfp4_supported, nvfp4_reason) - def test_te_nvfp4_shardy(self): - """Test Transformer Engine with NVFP4""" - self.args.enable_shardy = True - self.args.use_fp8 = True - self.args.fp8_recipe = "NVFP4BlockScaling" - actual = train_and_evaluate(self.args) - assert actual[0] < 0.52 and actual[1] > 0.74 - if __name__ == "__main__": train_and_evaluate(encoder_parser(None)) diff --git a/examples/jax/encoder/test_multiprocessing_encoder.py b/examples/jax/encoder/test_multiprocessing_encoder.py index 0166b60acd..344e7d618b 100644 --- a/examples/jax/encoder/test_multiprocessing_encoder.py +++ b/examples/jax/encoder/test_multiprocessing_encoder.py @@ -359,7 +359,6 @@ def replace_params(x): def train_and_evaluate(args): """Execute model training and evaluation loop.""" print(args) - jax.config.update("jax_use_shardy_partitioner", args.enable_shardy) if args.process_id == 0: nltk.download("punkt_tab") @@ -605,9 +604,6 @@ def encoder_parser(args): default=0, help="the ID number of the current process (default: 0)", ) - parser.add_argument( - "--enable-shardy", action="store_true", default=False, help="Enable Shardy (experimental)." - ) return parser.parse_args(args) @@ -616,7 +612,7 @@ def encoder_parser(args): class TestEncoder(unittest.TestCase): """Encoder unittests""" - def exec(self, use_fp8, fp8_recipe, *, enable_shardy=False): + def exec(self, use_fp8, fp8_recipe): """Run 5 epochs for testing""" args = encoder_parser(["--epochs", "5"]) @@ -632,7 +628,6 @@ def exec(self, use_fp8, fp8_recipe, *, enable_shardy=False): args.num_process = num_gpu args.process_id = self.process_id args.fp8_recipe = fp8_recipe - args.enable_shardy = enable_shardy return train_and_evaluate(args) @@ -674,44 +669,6 @@ def test_te_nvfp4(self): result = self.exec(True, "NVFP4BlockScaling") assert result[0] < 0.451 and result[1] > 0.787 - @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") - def test_te_bf16_shardy(self): - """Test Transformer Engine with BF16""" - result = self.exec(False, None, enable_shardy=True) - assert result[0] < 0.43 and result[1] > 0.80 - - @unittest.skipIf( - not is_fp8_supported(), "Device compute capability 9.0+ is required for DelayedScaling FP8" - ) - def test_te_delayed_scaling_fp8_shardy(self): - """Test Transformer Engine with DelayedScaling FP8""" - result = self.exec(True, "DelayedScaling", enable_shardy=True) - assert result[0] < 0.43 and result[1] > 0.80 - - @unittest.skipIf( - not is_fp8_supported(), "Device compute capability 9.0+ is required for CurrentScaling FP8" - ) - def test_te_current_scaling_fp8_shardy(self): - """Test Transformer Engine with CurrentScaling FP8""" - result = self.exec(True, "Float8CurrentScaling", enable_shardy=True) - assert result[0] < 0.432 and result[1] > 0.80 - - @unittest.skipIf( - not is_mxfp8_supported(), "Device compute capability 10.0+ is required for MXFP8" - ) - def test_te_mxfp8_shardy(self): - """Test Transformer Engine with MXFP8""" - result = self.exec(True, "MXFP8BlockScaling", enable_shardy=True) - assert result[0] < 0.43 and result[1] > 0.80 - - @unittest.skipIf( - not is_nvfp4_supported(), "Device compute capability 10.0+ is required for NVFP4" - ) - def test_te_nvfp4_shardy(self): - """Test Transformer Engine with NVFP4""" - result = self.exec(True, "NVFP4BlockScaling", enable_shardy=True) - assert result[0] < 0.451 and result[1] > 0.787 - if __name__ == "__main__": train_and_evaluate(encoder_parser(None)) diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index d5ebe9f261..50c5de1db7 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -68,9 +68,7 @@ def impl_test_self_attn( attn_mask_type, dtype, softmax_type, - use_shardy, ): - jax.config.update("jax_use_shardy_partitioner", use_shardy) dropout_prob = 0.0 is_training = True batch, seqlen, num_head, hidden = data_shape @@ -178,48 +176,6 @@ def test_self_attn( attn_mask_type, dtype, softmax_type, - use_shardy=False, - ) - - @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) - @pytest.mark.parametrize( - "attn_bias_type, bias_shape", - [ - pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), - pytest.param(AttnBiasType.PRE_SCALE_BIAS, BiasShape._1HSS, id="PRE_SCALE_BIAS-1HSS"), - ], - ) - @pytest.mark.parametrize( - "softmax_type", - [ - pytest.param(AttnSoftmaxType.VANILLA_SOFTMAX, id="VANILLA_SOFTMAX"), - pytest.param(AttnSoftmaxType.OFF_BY_ONE_SOFTMAX, id="OFF_BY_ONE_SOFTMAX"), - pytest.param(AttnSoftmaxType.LEARNABLE_SOFTMAX, id="LEARNABLE_SOFTMAX"), - ], - ) - def test_self_attn_shardy( - self, - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - attn_bias_type, - bias_shape, - softmax_type, - ): - data_shape = (32, 512, 12, 64) - self.impl_test_self_attn( - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - data_shape, - attn_bias_type, - bias_shape, - AttnMaskType.PADDING_MASK, - jnp.bfloat16, - softmax_type, - use_shardy=True, ) @@ -348,7 +304,6 @@ def impl_test_context_parallel_attn( qkv_layout, load_balanced, cp_strategy, - use_shardy, use_scan_ring=False, window_size=None, stripe_size=None, @@ -366,8 +321,6 @@ def impl_test_context_parallel_attn( os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] = "1" else: os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] = "0" - - jax.config.update("jax_use_shardy_partitioner", use_shardy) attn_bias_type = AttnBiasType.NO_BIAS bias_shape = None dropout_prob = 0.0 @@ -452,45 +405,6 @@ def check_has_backend_for_mask(mask_type): runner.test_backward() del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] - @pytest_parametrize_wrapper( - "device_count,mesh_shape,mesh_axes,mesh_resource", - generate_context_parallel_configs_for_attn(), - ) - @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES) - @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) - @pytest.mark.parametrize( - "qkv_layout, attn_mask_type", - DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS, - ) - def test_context_parallel_allgather_attn_shardy( - self, - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - data_shape, - attn_mask_type, - dtype, - qkv_layout, - ): - if qkv_layout.is_thd(): - pytest.skip("Only BSHD layout is supported for CP + AG + Dual chunk attention") - kv_groups = 8 - self.impl_test_context_parallel_attn( - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - data_shape, - kv_groups, - attn_mask_type, - dtype, - qkv_layout, - load_balanced=True, - cp_strategy=CPStrategy.ALL_GATHER, - use_shardy=True, - ) - @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), @@ -551,7 +465,6 @@ def test_context_parallel_allgather_striped_attn( qkv_layout, load_balanced, CPStrategy.ALL_GATHER, - use_shardy=False, window_size=window_size, stripe_size=stripe_size, num_segments_per_seq=num_segments_per_seq, @@ -599,7 +512,6 @@ def test_context_parallel_allgather_attn( qkv_layout, load_balanced, CPStrategy.ALL_GATHER, - use_shardy=False, ) @pytest_parametrize_wrapper( @@ -664,53 +576,11 @@ def test_context_parallel_ring_attn( qkv_layout, load_balanced, CPStrategy.RING, - use_shardy=False, use_scan_ring=use_scan, window_size=window_size, stripe_size=stripe_size, ) - @pytest_parametrize_wrapper( - "device_count,mesh_shape,mesh_axes,mesh_resource", - generate_context_parallel_configs_for_attn(), - ) - @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) - @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) - @pytest.mark.parametrize( - "qkv_layout, attn_mask_type", - DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS, - ) - def test_context_parallel_ring_attn_shardy( - self, - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - data_shape, - attn_mask_type, - dtype, - qkv_layout, - ): - kv_groups = 8 - # Set the stripe size to 1 (ring attention only support stripe_size=1) - stripe_size = 1 if qkv_layout.is_thd() else None - self.impl_test_context_parallel_attn( - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - data_shape, - kv_groups, - attn_mask_type, - dtype, - qkv_layout, - load_balanced=True, - cp_strategy=CPStrategy.RING, - use_shardy=False, - use_scan_ring=True, - stripe_size=stripe_size, - ) - REORDER_CAUSAL_LOAD_BALANCING_DATA_SHAPES = { "L0": [[]], diff --git a/tests/jax/test_distributed_layernorm.py b/tests/jax/test_distributed_layernorm.py index e9a2fa49e2..bb1f38dcc8 100644 --- a/tests/jax/test_distributed_layernorm.py +++ b/tests/jax/test_distributed_layernorm.py @@ -87,7 +87,6 @@ def generate_collectives_count_ref( @pytest_parametrize_wrapper("zero_centered_gamma", [False, True]) @pytest_parametrize_wrapper("shard_weights", [False, True]) @pytest_parametrize_wrapper("fp8_recipe", SUPPORTED_RECIPES) - @pytest_parametrize_wrapper("use_shardy", [False, True]) def test_layernorm( self, device_count, @@ -99,9 +98,7 @@ def test_layernorm( zero_centered_gamma, shard_weights, fp8_recipe, - use_shardy, ): - jax.config.update("jax_use_shardy_partitioner", use_shardy) epsilon = 1e-6 ln_type = "layernorm" q_dtype = jnp.float8_e4m3fn @@ -178,7 +175,6 @@ def ref_func(x, gamma, beta): @pytest_parametrize_wrapper("dtype", DTYPES) @pytest_parametrize_wrapper("shard_weights", [False, True]) @pytest_parametrize_wrapper("fp8_recipe", SUPPORTED_RECIPES) - @pytest_parametrize_wrapper("use_shardy", [False, True]) def test_rmsnorm( self, device_count, @@ -189,9 +185,7 @@ def test_rmsnorm( dtype, shard_weights, fp8_recipe, - use_shardy, ): - jax.config.update("jax_use_shardy_partitioner", use_shardy) epsilon = 1e-6 ln_type = "rmsnorm" q_dtype = jnp.float8_e4m3fn diff --git a/tests/jax/test_distributed_layernorm_mlp.py b/tests/jax/test_distributed_layernorm_mlp.py index d214597cb3..abf579d48e 100644 --- a/tests/jax/test_distributed_layernorm_mlp.py +++ b/tests/jax/test_distributed_layernorm_mlp.py @@ -192,10 +192,8 @@ def _test_layernorm_mlp_grad( input_shape, dtype, quantization_recipe, - use_shardy, with_jax_gemm, ): - jax.config.update("jax_use_shardy_partitioner", use_shardy) device_count, mesh_shape, mesh_axes, mesh_resource = mesh_config layernorm_type = "rmsnorm" @@ -313,36 +311,6 @@ def test_layernorm_mlp_grad( dtype, quantization_recipe, with_jax_gemm, - ): - if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4(): - pytest.skip("NVFP4 GEMM + Float16 output is unsupported!") - self._test_layernorm_mlp_grad( - mesh_config, - activation_type, - use_bias, - input_shape, - dtype, - quantization_recipe, - use_shardy=False, - with_jax_gemm=with_jax_gemm, - ) - - @pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs()) - @pytest_parametrize_wrapper("input_shape", INPUT_SHAPE) - @pytest_parametrize_wrapper("activation_type", [("gelu",), ("gelu", "linear")]) - @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("use_bias", [True, False]) - @pytest_parametrize_wrapper("quantization_recipe", [None] + SUPPORTED_RECIPES) - @pytest_parametrize_wrapper("with_jax_gemm", [False, True]) - def test_layernorm_mlp_grad_shardy( - self, - mesh_config, - activation_type, - use_bias, - input_shape, - dtype, - quantization_recipe, - with_jax_gemm, ): if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4(): pytest.skip("NVFP4 GEMM + Float16 output is unsupported!") @@ -353,7 +321,6 @@ def test_layernorm_mlp_grad_shardy( input_shape, dtype, quantization_recipe=quantization_recipe, - use_shardy=True, with_jax_gemm=with_jax_gemm, ) @@ -366,10 +333,8 @@ def _test_layernorm_mlp( dtype, use_fp8, quantization_recipe, - use_shardy, with_jax_gemm, ): - jax.config.update("jax_use_shardy_partitioner", use_shardy) batch, seqlen, hidden_in = input_shape layernorm_type = "rmsnorm" @@ -481,7 +446,6 @@ def test_layernorm_mlp_layer( dtype, use_fp8=False, quantization_recipe=None, - use_shardy=False, with_jax_gemm=with_jax_gemm, ) @@ -512,58 +476,5 @@ def test_layernorm_mlp_layer_fp8( dtype, use_fp8=True, quantization_recipe=quantization_recipe, - use_shardy=False, - with_jax_gemm=with_jax_gemm, - ) - - @pytest_parametrize_wrapper("input_shape", INPUT_SHAPE) - @pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs()) - @pytest_parametrize_wrapper("activation_type", [("gelu",), ("silu", "linear")]) - @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("use_bias", [True, False]) - @pytest_parametrize_wrapper("with_jax_gemm", [False, True]) - def test_layernorm_mlp_layer_shardy( - self, mesh_config, activation_type, use_bias, input_shape, dtype, with_jax_gemm - ): - self._test_layernorm_mlp( - mesh_config, - activation_type, - use_bias, - input_shape, - dtype, - use_fp8=False, - quantization_recipe=None, - use_shardy=True, - with_jax_gemm=with_jax_gemm, - ) - - @pytest_parametrize_wrapper("mesh_config", generate_fsdp_and_tpsp_configs()) - @pytest_parametrize_wrapper("activation_type", [("gelu",), ("gelu", "linear")]) - @pytest_parametrize_wrapper("use_bias", [True, False]) - @pytest_parametrize_wrapper("input_shape", INPUT_SHAPE) - @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("quantization_recipe", SUPPORTED_RECIPES) - @pytest_parametrize_wrapper("with_jax_gemm", [False, True]) - def test_layernorm_mlp_layer_fp8_shardy( - self, - mesh_config, - activation_type, - use_bias, - input_shape, - dtype, - quantization_recipe, - with_jax_gemm, - ): - if dtype == jnp.float16 and quantization_recipe is not None and quantization_recipe.nvfp4(): - pytest.skip("NVFP4 GEMM + Float16 output is unsupported!") - self._test_layernorm_mlp( - mesh_config, - activation_type, - use_bias, - input_shape, - dtype, - use_fp8=True, - quantization_recipe=quantization_recipe, - use_shardy=True, with_jax_gemm=with_jax_gemm, ) diff --git a/tests/jax/test_distributed_permutation.py b/tests/jax/test_distributed_permutation.py index 5b6d8fec47..04ed236e81 100644 --- a/tests/jax/test_distributed_permutation.py +++ b/tests/jax/test_distributed_permutation.py @@ -135,7 +135,6 @@ def generate_routing_map( DISPATCH_COMBINE_CASES, ) @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("use_shardy", [False, True]) def test_local_token_dispatch( self, device_count, @@ -147,7 +146,6 @@ def test_local_token_dispatch( hidden_size, topk, dtype, - use_shardy, ): """ Test token_dispatch with sharded inputs. @@ -164,7 +162,6 @@ def test_local_token_dispatch( matching the sharded execution's output ordering. Tests both forward pass (output values) and backward pass (gradients). """ - jax.config.update("jax_use_shardy_partitioner", use_shardy) key = jax.random.PRNGKey(42) # Generate global inputs @@ -307,7 +304,6 @@ def ref_chunk_loss(inp_chunk, routing_chunk, probs_chunk): DISPATCH_COMBINE_CASES, ) @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("use_shardy", [False, True]) def test_local_roundtrip( self, device_count, @@ -319,7 +315,6 @@ def test_local_roundtrip( hidden_size, topk, dtype, - use_shardy, ): """ Test roundtrip: token_dispatch followed by token_combine with sharded inputs. @@ -332,7 +327,6 @@ def test_local_roundtrip( Tests both forward pass and backward pass (gradient should be 2*x). """ - jax.config.update("jax_use_shardy_partitioner", use_shardy) key = jax.random.PRNGKey(42) # Generate global inputs @@ -403,7 +397,6 @@ def roundtrip_loss(x, rm, mprobs): DISPATCH_COMBINE_PADDING_CASES, ) @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("use_shardy", [False, True]) def test_local_token_dispatch_with_padding( self, device_count, @@ -416,14 +409,12 @@ def test_local_token_dispatch_with_padding( topk, align_size, dtype, - use_shardy, ): """ Test token_dispatch with padding using sharded inputs. Tests both forward pass (output values) and backward pass (gradients). """ - jax.config.update("jax_use_shardy_partitioner", use_shardy) key = jax.random.PRNGKey(42) # Generate global inputs @@ -502,7 +493,6 @@ def loss_with_padding(x, rm, p): DISPATCH_COMBINE_PADDING_CASES, ) @pytest_parametrize_wrapper("dtype", DTYPES) - @pytest_parametrize_wrapper("use_shardy", [False, True]) def test_local_roundtrip_with_padding( self, device_count, @@ -515,7 +505,6 @@ def test_local_roundtrip_with_padding( topk, align_size, dtype, - use_shardy, ): """ Test roundtrip with padding/alignment using sharded inputs. @@ -523,7 +512,6 @@ def test_local_roundtrip_with_padding( With uniform merging probs, should recover original input. Tests both forward pass and backward pass. """ - jax.config.update("jax_use_shardy_partitioner", use_shardy) key = jax.random.PRNGKey(42) # Generate inputs diff --git a/tests/jax/test_distributed_softmax.py b/tests/jax/test_distributed_softmax.py index 0665baa4e3..ca1dcf1174 100644 --- a/tests/jax/test_distributed_softmax.py +++ b/tests/jax/test_distributed_softmax.py @@ -87,12 +87,9 @@ def impl_test_softmax( dtype, bad_sharding, broadcast_batch_mask, - use_shardy, ): if broadcast_batch_mask and softmax_fusion_type != SoftmaxFusionType.SCALED_MASKED: pytest.skip("Softmax type has no mask.") - - jax.config.update("jax_use_shardy_partitioner", use_shardy) target_func = partial( self.target_func, scale_factor=scale_factor, softmax_fusion_type=softmax_fusion_type ) @@ -181,35 +178,4 @@ def test_softmax( dtype, bad_sharding, broadcast_batch_mask, - use_shardy=True, - ) - - @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) - @pytest.mark.parametrize( - "softmax_fusion_type", [SoftmaxFusionType.SCALED, SoftmaxFusionType.SCALED_MASKED] - ) - @pytest.mark.parametrize("bad_sharding", [False, True]) - @pytest.mark.parametrize("broadcast_batch_mask", [False, True]) - def test_softmax_gspmd( - self, - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - softmax_fusion_type, - bad_sharding, - broadcast_batch_mask, - ): - self.impl_test_softmax( - device_count, - mesh_shape, - mesh_axes, - mesh_resource, - data_shape=[32, 12, 128, 128], - softmax_fusion_type=softmax_fusion_type, - scale_factor=1.0, - dtype=DTYPES[0], - bad_sharding=bad_sharding, - broadcast_batch_mask=broadcast_batch_mask, - use_shardy=False, ) diff --git a/transformer_engine/jax/cpp_extensions/base.py b/transformer_engine/jax/cpp_extensions/base.py index b26e01c0c7..ae3888cf04 100644 --- a/transformer_engine/jax/cpp_extensions/base.py +++ b/transformer_engine/jax/cpp_extensions/base.py @@ -8,15 +8,23 @@ from abc import ABCMeta, abstractmethod from functools import partial +import jax from jax.extend import core from jax.interpreters import xla, mlir from jax.experimental.custom_partitioning import custom_partitioning from jax._src.interpreters import batching from jax._src import dispatch from jax import ffi +from packaging.version import Version as PkgVersion import transformer_engine_jax +# GSPMD sharding propagation (infer_sharding_from_operands) is removed in JAX > 0.9.1. +# Only register it for older JAX versions to maintain backwards compatibility. +# For JAX > 0.9.1, infer_sharding_from_operands is also removed from def_partition's signature, +# so it must not be passed at all. +_JAX_GSPMD_SUPPORTED = PkgVersion(jax.__version__) <= PkgVersion("0.9.1") + class BasePrimitive(metaclass=ABCMeta): """ @@ -143,13 +151,15 @@ def batcher(): """ return NotImplemented - @staticmethod - @abstractmethod - def infer_sharding_from_operands(): + @classmethod + def infer_sharding_from_operands(cls, *args, **kwargs): """ to describe infer_sharding_from_operands for custom_partitioning """ - return NotImplemented + raise NotImplementedError( + f"{cls.__name__} does not support GSPMD sharding propagation." + " Please use Shardy partitioner instead." + ) @staticmethod @abstractmethod @@ -172,6 +182,22 @@ def shardy_sharding_rule(*args): # Registry to store all registered primitive classes _primitive_registry = {} +_gspmd_deprecation_warned = False + + +def _warn_gspmd_deprecation_once(): + global _gspmd_deprecation_warned + if not _gspmd_deprecation_warned: + warnings.warn( + "GSPMD sharding propagation is planned to be removed in June 2026." + " It is no longer maintained or tested. Use it at your own risk." + " Please use Shardy partitioner instead." + " In case you cannot upgrade to a JAX version that supports Shardy, please reach out!", + DeprecationWarning, + stacklevel=3, + ) + _gspmd_deprecation_warned = True + def register_primitive(cls, outer_only=False): """ @@ -208,10 +234,16 @@ def name_of_wrapper_p(): outer_p.def_abstract_eval(cls.outer_abstract) batching.primitive_batchers[outer_p] = cls.batcher outer_p_lower = custom_partitioning(cls.impl, static_argnums=cls.impl_static_args) + if _JAX_GSPMD_SUPPORTED: + if "infer_sharding_from_operands" in cls.__dict__: + _warn_gspmd_deprecation_once() + gspmd_kwargs = {"infer_sharding_from_operands": cls.infer_sharding_from_operands} + else: + gspmd_kwargs = {} outer_p_lower.def_partition( - infer_sharding_from_operands=cls.infer_sharding_from_operands, partition=cls.partition, sharding_rule=cls.shardy_sharding_rule, + **gspmd_kwargs, ) mlir.register_lowering( outer_p, mlir.lower_fun(outer_p_lower, multiple_results=cls.multiple_results) From a3bc040682672ccfab8f79f63d099b06d9c3135a Mon Sep 17 00:00:00 2001 From: Qiyu Wan <39144338+WanZzzzzz@users.noreply.github.com> Date: Tue, 3 Mar 2026 14:30:52 -0800 Subject: [PATCH 240/521] NVFP4 primary weight support (#2691) * NVFP4 primary weight support Signed-off-by: qiyuw * rename/combine APIs and add assertions Signed-off-by: qiyuw * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix minor issues Signed-off-by: qiyuw * compile nvfp4.cu arch-aware Signed-off-by: qiyuw * Fix test and make test code cleaner Signed-off-by: qiyuw * Remove dup import Signed-off-by: qiyuw * fix CI issues Signed-off-by: qiyuw * cleaner code Signed-off-by: qiyuw * add assertion in debug code Signed-off-by: qiyuw * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove pytest for debug code due to setup issue in CI Signed-off-by: qiyuw * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix nvfp4 partial cast test Signed-off-by: Kirthi Shankar Sivamani * Fix all CI Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: qiyuw Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: qiyuw Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- .../test_cast_master_weights_to_fp8.py | 532 ++++++++++- transformer_engine/common/CMakeLists.txt | 2 +- .../include/transformer_engine/recipe.h | 112 +++ .../include/transformer_engine/transpose.h | 26 + transformer_engine/common/recipe/nvfp4.cu | 859 ++++++++++++++++++ transformer_engine/pytorch/csrc/extensions.h | 46 + .../csrc/extensions/nvfp4_2d_partial_cast.cpp | 156 ++++ .../pytorch/csrc/extensions/pybind.cpp | 61 ++ .../pytorch/csrc/extensions/transpose.cpp | 276 +++++- .../tensor/storage/nvfp4_tensor_storage.py | 59 ++ transformer_engine/pytorch/tensor/utils.py | 435 ++++++++- 11 files changed, 2509 insertions(+), 55 deletions(-) create mode 100644 transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 8a434b2148..373f27be2f 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -5,6 +5,7 @@ import argparse import datetime import os +import tempfile import subprocess import sys import pathlib @@ -18,6 +19,7 @@ DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + NVFP4BlockScaling, MXFP8BlockScaling, Format, Recipe, @@ -26,13 +28,19 @@ from transformer_engine.pytorch import ( is_fp8_available, is_fp8_block_scaling_available, - is_mxfp8_available, + is_nvfp4_available, QuantizedTensor, Float8Tensor, Float8BlockwiseQTensor, + NVFP4Tensor, + is_mxfp8_available, MXFP8Tensor, ) -from transformer_engine.pytorch.tensor import cast_master_weights_to_fp8 +from transformer_engine.pytorch.tensor.utils import ( + quantize_master_weights, + cast_master_weights_to_fp8, +) +from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer from transformer_engine.pytorch.tensor.utils import post_all_gather_processing, replace_raw_data @@ -67,6 +75,12 @@ def _get_raw_data(quantized_tensor, colwise=False): quantized_tensor._rowwise_data.dtype == torch.uint8 ), "Float8BlockwiseQTensor _rowwise_data must be uint8" return quantized_tensor._rowwise_data + elif isinstance(quantized_tensor, NVFP4Tensor): + assert hasattr(quantized_tensor, "_rowwise_data"), "NVFP4Tensor missing _rowwise_data" + assert ( + quantized_tensor._rowwise_data.dtype == torch.uint8 + ), "NVFP4Tensor _rowwise_data must be uint8" + return quantized_tensor._rowwise_data elif isinstance(quantized_tensor, MXFP8Tensor): if colwise: assert hasattr( @@ -135,22 +149,45 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals self.offsets = [0] for weight in self.weights: self.offsets.append(self.offsets[-1] + weight.numel()) - # Padding to avoid global buffer cannot be divided by world size, so the offsets[-1] may # not be the end range of the last weight. if self.offsets[-1] % self.world_size != 0: self.offsets[-1] += self.world_size - self.offsets[-1] % self.world_size + self.weights_are_nvfp4 = isinstance(self.weights[0], NVFP4Tensor) + + # Storage offsets operate on the packed representation. + # For NVFP4: packed size (2 values per byte) + # For others: same as numel() + self.storage_offsets = [0] + self.storage_sizes = [] + for weight in self.weights: + if self.weights_are_nvfp4: + storage_size = _get_raw_data(weight).view(-1).numel() + else: + storage_size = weight.numel() + self.storage_sizes.append(storage_size) + self.storage_offsets.append(self.storage_offsets[-1] + storage_size) + if self.storage_offsets[-1] % self.world_size != 0: + self.storage_offsets[-1] += self.world_size - self.storage_offsets[-1] % self.world_size + self.storage_total = self.storage_offsets[-1] + self.master_weights = [] # The start offset of the master weight in the weight self.start_offsets = [] # The overlapping area of the weight and this rank's local buffer self.overlapping_areas = [] + # Storage equivalents (only populated for NVFP4 tensors). + self.storage_start_offsets = [None] * len(self.weights) + self.storage_overlapping_areas = [None] * len(self.weights) - # The start and end of this rank's local buffer in the global buffer + # The start and end of this rank's local buffer in the global buffer (logical offsets) rank_start = self.offsets[-1] // self.world_size * self.rank rank_end = rank_start + self.offsets[-1] // self.world_size + # Storage-based rank boundaries (for NVFP4: packed size, for others: same as logical) + storage_rank_start = self.storage_total // self.world_size * self.rank + storage_rank_end = storage_rank_start + self.storage_total // self.world_size for weight, offset in zip(self.weights, self.offsets[:-1]): if offset >= rank_end or (offset + weight.numel()) <= rank_start: # This weight is not in this rank's local buffer @@ -178,6 +215,20 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals self.start_offsets.append(start_offset) self.overlapping_areas.append(overlapping_area) + if self.weights_are_nvfp4: + for idx, (weight, storage_offset, storage_size) in enumerate( + zip(self.weights, self.storage_offsets[:-1], self.storage_sizes) + ): + if ( + storage_offset >= storage_rank_end + or (storage_offset + storage_size) <= storage_rank_start + ): + continue + overlap_start = max(storage_rank_start, storage_offset) + overlap_end = min(storage_rank_end, storage_offset + storage_size) + self.storage_start_offsets[idx] = overlap_start - storage_offset + self.storage_overlapping_areas[idx] = (overlap_start, overlap_end) + # Create global buffer for grads reduce-scatter self.grad_buffer = torch.empty( [self.offsets[-1]], dtype=torch.float32, device=weights[0].device @@ -190,9 +241,9 @@ def __init__(self, weights, lr, dp_group, manual_post_all_gather_processing=Fals else: weight_buffer_dtype = weights[0].dtype self.weight_buffer = torch.empty( - [self.offsets[-1]], dtype=weight_buffer_dtype, device=weights[0].device + [self.storage_total], dtype=weight_buffer_dtype, device=weights[0].device ) - self.weight_buffer_slice = self.weight_buffer[rank_start:rank_end] + self.weight_buffer_slice = self.weight_buffer[storage_rank_start:storage_rank_end] def step(self): # ----------------------------------------------------------------------------------------- @@ -231,10 +282,20 @@ def step(self): # ----------------------------------------------------------------------------------------- # Step 4: Cast master weights to BF16 or FP8, depending on the type of the weight # ----------------------------------------------------------------------------------------- - if isinstance(self.weights[0], QuantizedTensor): - # FP8 weights case - for i in range(1, len(self.weights)): - assert isinstance(self.weights[i], QuantizedTensor) + first_weight = self.weights[0] + if isinstance(first_weight, NVFP4Tensor): + for weight in self.weights: + assert isinstance(weight, NVFP4Tensor) + quantize_master_weights( + self.weights, + self.master_weights, + self.start_offsets, + self.dp_group, + manual_post_all_gather_processing=self.manual_post_all_gather_processing, + ) + elif isinstance(first_weight, (Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor)): + for weight in self.weights: + assert isinstance(weight, QuantizedTensor) cast_master_weights_to_fp8( self.weights, self.master_weights, @@ -253,20 +314,31 @@ def step(self): end = start_offset + master_weight.numel() weight.data.view(-1)[start:end].copy_(master_weight) + # ----------------------------------------------------------------------------------------- + # Step 5: Copy the updated weights (not all weights) to the weight buffer + # ----------------------------------------------------------------------------------------- colwise_list = [False] if isinstance(self.weights[0], MXFP8Tensor): colwise_list.append(True) for colwise in colwise_list: - # ------------------------------------------------------------------------------------- - # Step 5: Copy the updated weights (not all weights) to the weight buffer - # ------------------------------------------------------------------------------------- for i in range(len(self.weights)): master_weight = self.master_weights[i] if master_weight is None: continue start_offset = self.start_offsets[i] - if isinstance(self.weights[i], QuantizedTensor): + if isinstance(self.weights[i], NVFP4Tensor): + storage_start = self.storage_start_offsets[i] + storage_overlap = self.storage_overlapping_areas[i] + if storage_start is None or storage_overlap is None: + continue + weight = _get_raw_data(self.weights[i]).view(-1) + storage_len = storage_overlap[1] - storage_overlap[0] + weight_slice = weight[storage_start : storage_start + storage_len] + overlapping_start, overlapping_end = storage_overlap + self.weight_buffer[overlapping_start:overlapping_end].copy_(weight_slice) + continue + elif isinstance(self.weights[i], QuantizedTensor): weight = _get_raw_data(self.weights[i], colwise) else: weight = self.weights[i] @@ -284,12 +356,22 @@ def step(self): # ------------------------------------------------------------------------------------- # Step 7: Copy the gathered weights from weight buffer to the actual weights # ------------------------------------------------------------------------------------- - for weight, offset in zip(self.weights, self.offsets[:-1]): - start = offset - end = offset + weight.numel() - if isinstance(weight, QuantizedTensor): - weight = _get_raw_data(weight, colwise) - weight.view(-1).data.copy_(self.weight_buffer[start:end]) + if self.weights_are_nvfp4: + # NVFP4: use storage offsets (packs 2 values per byte) + for weight, storage_offset, storage_size in zip( + self.weights, self.storage_offsets[:-1], self.storage_sizes + ): + start = storage_offset + end = storage_offset + storage_size + raw_data = _get_raw_data(weight) + raw_data.view(-1).data.copy_(self.weight_buffer[start:end]) + else: + for weight, offset in zip(self.weights, self.offsets[:-1]): + start = offset + end = offset + weight.numel() + if isinstance(weight, QuantizedTensor): + weight = _get_raw_data(weight, colwise) + weight.view(-1).data.copy_(self.weight_buffer[start:end]) if self.manual_post_all_gather_processing: quantized_weights = [ @@ -464,8 +546,23 @@ def step(self): # Update the master weight using gradient descent master_weight -= grad * self.lr - # Step 3: Cast master weights to FP8 or BF16 precision - if isinstance(self.weights[0], QuantizedTensor): + # Step 3: Cast master weights to quantized or BF16 precision + first_weight = self.weights[0] + if isinstance(first_weight, NVFP4Tensor): + local_weights = [] + for local_weight in self.local_weights: + if local_weight is None: + local_weights.append(None) + continue + local_weights.append(local_weight) + quantize_master_weights( + self.weights, + self.master_weights, + [idx[0] for idx in self.weight_indices], + self.dp_group, + local_weights, + ) + elif isinstance(first_weight, QuantizedTensor): local_weights = [] for i, local_weight in enumerate(self.local_weights): if self.flatten_columnwise is not None: @@ -730,6 +827,90 @@ def _test_fsdp_cast_master_weights_to_fp8( ), f"Loss mismatch at rank {rank}, step {i} for {quantization} (FSDP)" +def _test_cast_master_weights_to_nvfp4(dp_group, manual_post_all_gather_processing): + available, reason = is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + + rank = dist.get_rank(dp_group) + world_size = dist.get_world_size(dp_group) + + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + + mock_groups = [dist.new_group(ranks=[i]) for i in range(world_size)] + mock_group = mock_groups[rank] + + linear_kwargs = {"params_dtype": torch.bfloat16, "bias": False, "fuse_wgrad_accumulation": True} + # Disable stochastic rounding for deterministic gradients + nvfp4_recipe = NVFP4BlockScaling(disable_stochastic_rounding=True) + + with te.quantized_model_init( + enabled=True, recipe=nvfp4_recipe, preserve_high_precision_init_val=True + ): + model_nvfp4 = nn.Sequential( + te.Linear(128, 256 + 64, **linear_kwargs), + te.Linear(256 + 64, 256 * 3, **linear_kwargs), + te.Linear(256 * 3, 128, **linear_kwargs), + ) + # Create model with bf16 weights + model = nn.Sequential( + te.Linear(128, 256 + 64, **linear_kwargs), + te.Linear(256 + 64, 256 * 3, **linear_kwargs), + te.Linear(256 * 3, 128, **linear_kwargs), + ) + + for w_nvfp4, w in zip(model_nvfp4.parameters(), model.parameters()): + high_precision_init_val = w_nvfp4.get_high_precision_init_val() + w.data.copy_(high_precision_init_val) + + for w_nvfp4, w in zip(model_nvfp4.parameters(), model.parameters()): + w_nvfp4.main_grad = torch.zeros_like(w_nvfp4, dtype=torch.float32, device="cuda") + w.main_grad = torch.zeros_like(w, dtype=torch.float32, device="cuda") + + optimizer_nvfp4 = MiniZero_1( + [w for w in model_nvfp4.parameters()], 10.0, dp_group, manual_post_all_gather_processing + ) + optimizer = MiniZero_1([w for w in model.parameters()], 10.0, dp_group) + + for i in range(500): + for w_nvfp4, w in zip(model_nvfp4.parameters(), model.parameters()): + w_nvfp4.main_grad.zero_() + w.main_grad.zero_() + + inputs = [ + torch.randn(2048, 128, dtype=torch.bfloat16, device="cuda") for _ in range(world_size) + ] + x = inputs[rank] + + with te.autocast( + enabled=True, + recipe=nvfp4_recipe, + amax_reduction_group=mock_group, + ): + y_nvfp4 = model_nvfp4(x) + + with te.autocast( + enabled=True, + recipe=nvfp4_recipe, + amax_reduction_group=mock_group, + ): + y = model(x) + + targets = [torch.randn_like(y) for _ in range(world_size)] + target = targets[rank] + loss_nvfp4 = nn.MSELoss()(y_nvfp4, target) + loss = nn.MSELoss()(y, target) + + loss_nvfp4.backward() + loss.backward() + + optimizer.step() + optimizer_nvfp4.step() + + torch.testing.assert_close(loss_nvfp4, loss, atol=0, rtol=0) + + def run_parallel_tests() -> None: """Run parallel tests""" @@ -762,13 +943,45 @@ def run_parallel_tests() -> None: quantizations.append("mxfp8") manual_post_all_gather_processings = [False, True] - + print("starting mini optimizer test") _test_mini_optimizer(dp_group) - + print("starting cast master weights to fp8 test") for quantization in quantizations: for post_ag_processing in manual_post_all_gather_processings: _test_cast_master_weights_to_fp8(quantization, dp_group, post_ag_processing) _test_fsdp_cast_master_weights_to_fp8(quantization, dp_group, post_ag_processing) + nvfp4_available, _ = is_nvfp4_available(return_reason=True) + if nvfp4_available: + print("starting cast master weights to nvfp4 test") + for post_ag_processing in manual_post_all_gather_processings: + _test_cast_master_weights_to_nvfp4(dp_group, post_ag_processing) + + dist.destroy_process_group() + + +def run_parallel_nvfp4_partial_cast_test() -> None: + """Run the NVFP4 partial-cast distributed worker test.""" + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + + assert WORLD_SIZE == LOCAL_SIZE # this test supports only 1 node + assert LOCAL_SIZE <= torch.cuda.device_count() + dist_init_kwargs = { + "backend": "nccl", + "rank": WORLD_RANK, + "world_size": WORLD_SIZE, + "timeout": datetime.timedelta(seconds=30), + } + dist_init_kwargs["init_method"] = "env://" + dist_init_kwargs["device_id"] = torch.device(f"cuda:{LOCAL_RANK}") + assert dist.is_nccl_available() + torch.cuda.set_device(LOCAL_RANK) + dist.init_process_group(**dist_init_kwargs) + dp_group = dist.new_group(backend="nccl") + + _test_nvfp4_partial_cast_matches_full(dp_group) dist.destroy_process_group() @@ -798,9 +1011,280 @@ def test_cast_master_weights_to_fp8(world_size: int) -> None: def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--parallel", action="store_true", help="Run parallel tests") + parser.add_argument( + "--parallel-nvfp4-partial", + action="store_true", + help="Run NVFP4 partial-cast distributed worker test", + ) args = parser.parse_args() if args.parallel: run_parallel_tests() + elif args.parallel_nvfp4_partial: + run_parallel_nvfp4_partial_cast_test() + + +# Debugging tests for NVFP4 +def test_nvfp4_transpose_kernel() -> None: + """Test that nvfp4_transpose kernel produces bitwise identical results to reference.""" + available, reason = is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + + torch.manual_seed(1234) + device = torch.device("cuda") + shape = (2048, 5120) + master_weight = torch.randn(shape, dtype=torch.float32, device=device) + + print("\n=== Testing NVFP4 transpose kernel ===") + + # Create reference with both rowwise and columnwise data + quantizer_with_colwise = NVFP4Quantizer( + rowwise=True, columnwise=True, with_2d_quantization=True + ) + reference_tensor = quantizer_with_colwise(master_weight.to(torch.bfloat16)) + assert reference_tensor._columnwise_data is not None, "Reference should have columnwise data" + assert ( + reference_tensor._columnwise_scale_inv is not None + ), "Reference should have columnwise scale_inv" + reference_columnwise_data = reference_tensor._columnwise_data.detach().clone() + reference_columnwise_scale_inv = reference_tensor._columnwise_scale_inv.detach().clone() + reference_columnwise_amax = ( + reference_tensor._amax_columnwise.detach().clone() + if reference_tensor._amax_columnwise is not None + else None + ) + + # Create tensor with only rowwise data, then call _create_columnwise() + quantizer_rowwise_only = NVFP4Quantizer( + rowwise=True, columnwise=False, with_2d_quantization=True + ) + test_tensor = quantizer_rowwise_only(master_weight.to(torch.bfloat16)) + assert test_tensor._columnwise_data is None, "Test tensor should not have columnwise data yet" + + # Now call _create_columnwise() which uses our nvfp4_transpose kernel + test_tensor.update_usage(rowwise_usage=True, columnwise_usage=True) + assert ( + test_tensor._columnwise_data is not None + ), "Test tensor should have columnwise data after _create_columnwise()" + assert ( + test_tensor._columnwise_scale_inv is not None + ), "Test tensor should have columnwise scale_inv after _create_columnwise()" + + # Compare columnwise data - should be bitwise identical + torch.testing.assert_close( + test_tensor._columnwise_data, + reference_columnwise_data, + atol=0, + rtol=0, + msg="NVFP4 transpose kernel produced different columnwise data than reference!", + ) + + torch.testing.assert_close( + test_tensor._columnwise_scale_inv, + reference_columnwise_scale_inv, + atol=0, + rtol=0, + msg="NVFP4 _create_columnwise produced different columnwise scale_inv than reference!", + ) + + torch.testing.assert_close( + test_tensor._amax_columnwise, + reference_columnwise_amax, + atol=0, + rtol=0, + msg="NVFP4 _create_columnwise produced different columnwise amax than reference!", + ) + + +def _test_nvfp4_partial_cast_matches_full(dp_group) -> None: + """Multi-GPU worker: split master weight, partial cast on each rank, gather, compare.""" + WORLD_RANK = dist.get_rank(dp_group) + WORLD_SIZE = dist.get_world_size(dp_group) + + available, reason = is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + + torch.manual_seed(1234) + device = torch.device("cuda") + # Shape must be divisible by WORLD_SIZE for even splitting + # Also ensure dimensions are multiples of 16 for NVFP4 tiles + shape = (4096, 4096) + total_elements = shape[0] * shape[1] + assert total_elements % WORLD_SIZE == 0, "Total elements must be divisible by WORLD_SIZE" + + # Full master weight (same on all ranks due to same seed) + full_master_weight = torch.randn(shape, dtype=torch.float32, device=device) + + # Create reference using full quantization + quantizer = NVFP4Quantizer(rowwise=True, columnwise=False, with_2d_quantization=True) + reference_tensor = quantizer(full_master_weight.to(torch.bfloat16)) + reference_data = reference_tensor._rowwise_data.detach().clone() + reference_scale = reference_tensor._rowwise_scale_inv.detach().clone() + reference_amax = reference_tensor._amax_rowwise.detach().clone() + + # Split master weight evenly across ranks + shard_size = total_elements // WORLD_SIZE + start_offset = WORLD_RANK * shard_size + end_offset = start_offset + shard_size + master_weight_shard = full_master_weight.view(-1)[start_offset:end_offset].clone() + + # Create empty NVFP4 tensor for this rank (full shape, but we'll only fill our shard) + nvfp4_tensor = quantizer.make_empty(shape, dtype=torch.bfloat16, device=device) + nvfp4_tensor._rowwise_data.zero_() + nvfp4_tensor._rowwise_scale_inv.zero_() + if nvfp4_tensor._amax_rowwise is not None: + nvfp4_tensor._amax_rowwise.zero_() + + # Partial cast on each rank's shard + quantize_master_weights( + [nvfp4_tensor], + [master_weight_shard], + [start_offset], + dp_group, + ) + + # All-gather the rowwise data (packed FP4 bytes) + # Each rank has the full tensor but only its shard is filled + # We need to all-gather the shards + rowwise_data_flat = nvfp4_tensor._rowwise_data.view(-1) + + # For NVFP4, 2 elements are packed per byte, so byte shard size is shard_size // 2 + byte_shard_size = shard_size // 2 + byte_start = WORLD_RANK * byte_shard_size + byte_end = byte_start + byte_shard_size + my_shard_bytes = rowwise_data_flat[byte_start:byte_end].contiguous() + + # Gather all shards + gathered_shards = [torch.empty_like(my_shard_bytes) for _ in range(WORLD_SIZE)] + dist.all_gather(gathered_shards, my_shard_bytes, group=dp_group) + + # Reconstruct the full rowwise data + gathered_data = torch.cat(gathered_shards, dim=0).view(reference_data.shape) + + # Compare with reference + torch.testing.assert_close( + gathered_data, + reference_data, + atol=0, + rtol=0, + msg=f"[Rank {WORLD_RANK}] Gathered rowwise data does not match reference!", + ) + + # Also verify scale matches (scale should be identical on all ranks after all-reduce) + torch.testing.assert_close( + nvfp4_tensor._rowwise_scale_inv, + reference_scale, + atol=0, + rtol=0, + msg=f"[Rank {WORLD_RANK}] Scale does not match reference!", + ) + + # Verify amax matches + torch.testing.assert_close( + nvfp4_tensor._amax_rowwise, + reference_amax, + atol=0, + rtol=0, + msg=f"[Rank {WORLD_RANK}] Amax does not match reference!", + ) + + +@pytest.mark.skipif( + torch.cuda.device_count() < 2, reason="NVFP4 partial-cast test needs at least 2 GPUs." +) +@pytest.mark.parametrize("world_size", [2]) +def test_nvfp4_partial_cast_matches_full(world_size: int) -> None: + """Launch a distributed job for NVFP4 partial-cast equivalence test.""" + python_exe = pathlib.Path(sys.executable).resolve() + current_file = pathlib.Path(__file__).resolve() + command = [ + python_exe, + "-m", + "torch.distributed.run", + f"--nproc_per_node={world_size}", + current_file, + "--parallel-nvfp4-partial", + ] + subprocess.run(command, check=True) + + +def test_single_gpu_partial_cast_vs_full(): + """ + Single GPU test: compare quantize_master_weights (offset=0) vs quantizer(). + This isolates whether the issue is in our manual Python scale computation or elsewhere. + """ + available, reason = is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + + torch.manual_seed(1234) + device = torch.device("cuda") + + # Test with same shape as the optimizer test + shape = (2048, 2048) + + # Create BF16 master weight + master_weight = torch.randn(shape, dtype=torch.bfloat16, device=device) + + # === Reference: Use NVFP4Quantizer directly === + quantizer = NVFP4Quantizer(rowwise=True, columnwise=False, with_2d_quantization=True) + ref = quantizer(master_weight) + ref_data = ref._rowwise_data.clone() + ref_scale = ref._rowwise_scale_inv.clone() + ref_amax = ref._amax_rowwise.clone() + + # === Test: Use quantize_master_weights with offset=0 (full tensor) === + # Create empty NVFP4 tensor + test_tensor = quantizer.make_empty(shape, dtype=torch.bfloat16, device=device) + test_tensor._rowwise_data.zero_() + test_tensor._rowwise_scale_inv.zero_() + if test_tensor._amax_rowwise is not None: + test_tensor._amax_rowwise.zero_() + + # Create a local single-rank process group when running under plain pytest. + initialized_here = False + rendezvous_file = None + if not dist.is_initialized(): + torch.cuda.set_device(0) + with tempfile.NamedTemporaryFile(delete=False) as f: + rendezvous_file = pathlib.Path(f.name) + dist.init_process_group( + backend="nccl", + init_method=rendezvous_file.resolve().as_uri(), + rank=0, + world_size=1, + ) + initialized_here = True + + if dist.get_world_size() != 1: + pytest.skip("test_single_gpu_partial_cast_vs_full requires world_size == 1") + + mock_group = dist.new_group(ranks=[0], backend="nccl") + try: + quantize_master_weights( + [test_tensor], + [master_weight.view(-1)], # Flatten as expected + [0], # offset=0 means full tensor + mock_group, + ) + finally: + if initialized_here: + dist.destroy_process_group() + if rendezvous_file is not None: + rendezvous_file.unlink(missing_ok=True) + + # Compare amax + amax_match = torch.equal(test_tensor._amax_rowwise, ref_amax) + assert amax_match, f"Amax mismatch: {test_tensor._amax_rowwise} vs {ref_amax}" + + # Compare scale + scale_match = torch.equal(test_tensor._rowwise_scale_inv, ref_scale) + assert scale_match, f"Scale mismatch: {test_tensor._rowwise_scale_inv} vs {ref_scale}" + + # Compare data + data_match = torch.equal(test_tensor._rowwise_data, ref_data) + assert data_match, f"Data mismatch" if __name__ == "__main__": diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 4579c51e9f..a105a0343f 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -163,7 +163,6 @@ list(APPEND transformer_engine_cuda_sources recipe/current_scaling.cu recipe/delayed_scaling.cu recipe/fp8_block_scaling.cu - recipe/nvfp4.cu comm_gemm_overlap/userbuffers/userbuffers.cu) list(APPEND transformer_engine_cuda_arch_specific_sources @@ -182,6 +181,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu multi_tensor/compute_scale.cu recipe/mxfp8_scaling.cu + recipe/nvfp4.cu transpose/quantize_transpose_square_blockwise.cu transpose/quantize_transpose_vector_blockwise_fp4.cu) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index c0cec8a3b9..cad27a2992 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -309,6 +309,118 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r const NVTETensor inpB, const bool use_rowwise_amax_B, float alpha_in, NVTETensor alpha_out, cudaStream_t stream); +/*! \brief Compute tile-level amax for a partial shard of a 2D tensor. + * + * For NVFP4 2D quantization with 16x16 tiles. Computes the maximum absolute + * value within each tile, but only for elements in [start_offset, start_offset + len) + * of the flattened tensor. Used in distributed settings where each rank owns a shard. + * + * \param[in] inp Input tensor (partial shard, high-precision). + * \param[out] amax Output amax buffer [tile_rows, tile_cols], float32. + * \param[in] h Number of rows in the full 2D tensor. + * \param[in] w Number of columns in the full 2D tensor. + * \param[in] amax_stride_h Stride for amax in tile-row dimension. + * \param[in] amax_stride_w Stride for amax in tile-col dimension. + * \param[in] start_offset Starting element offset in the flattened tensor. + * \param[in] block_len Tile dimension (must be 16 for NVFP4 2D). + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h, size_t w, + size_t amax_stride_h, size_t amax_stride_w, + size_t start_offset, size_t block_len, cudaStream_t stream); + +/*! \brief Cast a partial shard of a tensor to NVFP4 using 2D tile-based quantization. + * + * Quantizes elements in [start_offset, start_offset + len) of the flattened tensor + * using precomputed per-tile scales. Each 16x16 tile uses its own scale factor. + * Used in distributed settings where each rank casts its owned shard. + * + * \param[in] inp Input tensor (partial shard, high-precision). + * \param[out] out Output NVFP4 packed tensor (2 values per byte). + * \param[in] scale Per-tile scale factors [tile_rows, tile_cols], float32. + * \param[in] global_scale Global scale factor [1], float32. + * \param[in] h Number of rows in the full 2D tensor. + * \param[in] w Number of columns in the full 2D tensor. + * \param[in] scale_stride_h Stride for scale in tile-row dimension. + * \param[in] scale_stride_w Stride for scale in tile-col dimension. + * \param[in] start_offset Starting element offset in the flattened tensor. + * \param[in] block_len Tile dimension (must be 16 for NVFP4 2D). + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, + const NVTETensor global_scale, size_t h, size_t w, + size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, + size_t block_len, cudaStream_t stream); + +/*! \brief Expand tile-level scales to row-level scales and convert to FP8 E4M3, used in partial cast. + * + * Each tile row's scale is repeated block_len times in the output. + * + * \param[in] input Input tensor with tile scales [tile_rows, tile_cols], float32. + * \param[out] output Output tensor with expanded scales [rows_padded, tile_cols], uint8 (E4M3). + * \param[in] tile_rows Number of tile rows. + * \param[in] tile_cols Number of tile columns. + * \param[in] rows_padded Padded row count in output. + * \param[in] block_len Block length (typically 16 for NVFP4). + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, + size_t tile_cols, size_t rows_padded, size_t block_len, + cudaStream_t stream); + +/*! \brief Compute per-block decode scale from block amax and global amax. + * + * Computes: + * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax + * per_block_decode_scale = block_amax / fp4_max * global_scale + * + * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh. + * + * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. + * \param[out] scale Output scale tensor [tile_rows, tile_cols], float32. + * \param[in] global_amax Global amax tensor (single element), float32. Avoids D2H transfer. + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, + const NVTETensor global_amax, cudaStream_t stream); + +/*! \brief Fused kernel for NVFP4 scale computation. + * + * Fuses three operations into one kernel: + * 1. Compute per-block decode scales from block amax and global amax + * 2. Copy global amax to target tensor + * 3. Expand tile-level scales to row-level and convert to FP8 E4M3 + * + * Saves 2 kernel launches per parameter. + * + * \param[in] block_amax Input block amax tensor [tile_rows, tile_cols], float32. + * \param[in] global_amax Global amax tensor [1], float32. + * \param[out] per_block_scale Output per-block scale [tile_rows, tile_cols], float32 (for partial_cast). + * \param[out] target_scale Output scale tensor [rows_padded, tile_cols], uint8 (E4M3). + * \param[out] target_amax Output amax tensor [1], float32 (copy of global_amax). + * \param[in] tile_rows Number of tile rows. + * \param[in] tile_cols Number of tile columns. + * \param[in] rows_padded Total padded rows in output. + * \param[in] block_len Block length (16 for NVFP4). + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, + NVTETensor per_block_scale, NVTETensor target_scale, + NVTETensor target_amax, size_t tile_rows, size_t tile_cols, + size_t rows_padded, size_t block_len, cudaStream_t stream); + +/*! \brief Compute global encode scale from global amax. + * + * Computes: global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax + * If global_amax <= 0, returns 1.0. + * + * \param[in] global_amax Input global amax tensor [num_params], float32. + * \param[out] global_scale Output global scale tensor [num_params], float32. + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/include/transformer_engine/transpose.h b/transformer_engine/common/include/transformer_engine/transpose.h index 5f9a8fe149..659a48d97d 100644 --- a/transformer_engine/common/include/transformer_engine/transpose.h +++ b/transformer_engine/common/include/transformer_engine/transpose.h @@ -326,6 +326,32 @@ void nvte_dsreglu_cast_transpose(const NVTETensor input, const NVTETensor act_in */ void nvte_swap_first_dims(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Transpose NVFP4 packed data. + * + * Unlike FP8, NVFP4 packs two 4-bit values per byte. This function correctly + * handles the nibble repacking during transpose. + * + * \param[in] input Input tensor with packed FP4 data. Shape: [M, K/2] bytes. + * \param[out] output Output tensor with transposed packed data. Shape: [K, M/2] bytes. + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_data_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream); + +/*! \brief Transpose NVFP4 tile-level scales from rowwise to columnwise format. + * + * Takes rowwise_scale_inv where scales are stored at every 16th row (tile boundaries) + * and produces columnwise_scale_inv where scales are repeated 16 times per tile row. + * Scale values are stored as E4M3 (fp8) in uint8 tensors. + * + * \param[in] input Input tensor with rowwise scales [M_padded, K_tiles], uint8 (E4M3). + * \param[out] output Output tensor with columnwise scales [K_padded, M_tiles], uint8 (E4M3). + * \param[in] M_tiles Number of tiles in M dimension. + * \param[in] K_tiles Number of tiles in K dimension. + * \param[in] stream CUDA stream. + */ +void nvte_nvfp4_scale_transpose(const NVTETensor input, NVTETensor output, size_t M_tiles, + size_t K_tiles, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 682d8b53f5..36ce60eaa5 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -5,17 +5,69 @@ ************************************************************************/ #include +#include #include +#include #include "../common.h" +#include "../util/ptx.cuh" #include "../utils.cuh" namespace transformer_engine { namespace nvfp4_recipe { +/* + * --------------------------------------------------------------------------- + * NVFP4 2D PARTIAL-SHARD KERNEL DESIGN + * + * These kernels mirror the FP8 block-scaling helpers but operate on shard-local + * slices and nibble-packed FP4 rowwise buffers. One CUDA block covers a logical + * 16x16 tile (grid = ceil(W/16) x ceil(H/16), blockDim = 256 threads). + * + * 1) Partial Amax (`nvfp4_2d_compute_partial_amax_kernel`) + * - Warps sweep the tile using nested loops, accumulating local maxima only + * for elements in [start_offset, start_offset + len). + * - Shared memory reduces the 8 warp maxima; the block writes a float into + * `amax_ptr[tile_row * stride_h + tile_col * stride_w]`. + * + * Tile/warp mapping (each '#' = elements visited by that warp): + * + * +------------------+ + * |########..........| Warp 0 + * |########..........| Warp 1 + * | ... | + * |########..........| Warp 7 + * +------------------+ + * + * 2) Partial Cast (`nvfp4_2d_partial_cast_kernel`) + * - Stage the tile into shared memory (same pattern as FP8). + * - For each 4-value group, build float2 pairs and call + * `ptx::mul_cvt_fp32_to_fp4_4x`, producing packed FP4 nibbles. + * - Compute a shard-local byte index and update only the owned nibble(s) + * using read-modify-write: + * + * packed_bits = [mw3 | mw2 | mw1 | mw0] + * byte_idx = (ref_elem_idx - start_offset) >> 1 + * if elem_idx % 2 == 0: // low nibble + * byte = (byte & 0xF0) | nibble + * else: // high nibble + * byte = (byte & 0x0F) | (nibble << 4) + * + * Thread coverage inside a tile: + * + * rows: 16 columns: 16 + * Warp 0 -> rows 0-1 lanes sweep cols 0..3, 4..7, ... + * Warp 1 -> rows 2-3 (groups of 4 elements per thread) + * ... + * Warp 7 -> rows 14-15 + * --------------------------------------------------------------------------- + */ + // 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, @@ -24,9 +76,804 @@ __global__ void compute_nvfp4_per_tensor_scale_kernel(float alpha_in, const floa *alpha_out = alpha_in * (*amax_A) * (*amax_B) * factor_inv; } +template +__global__ void __launch_bounds__(kThreadsPerBlock) + nvfp4_2d_compute_partial_amax_kernel(const IType *input, float *amax_ptr, + const size_t amax_stride_h, const size_t amax_stride_w, + const size_t h, const size_t w, const size_t start_offset, + const size_t len) { + constexpr int kThreadsPerWarp = 32; + constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp; + static_assert(kTileDim * kTileDim == kThreadsPerBlock); + + const size_t tile_col = blockIdx.x; + const size_t tile_row = blockIdx.y; + const size_t end_offset = start_offset + len; + const IType *input_minus_offset = input - start_offset; + + __shared__ float smem[kNumWarps]; + float amax = 0.0f; + + size_t r = tile_row * kTileDim + threadIdx.x / kTileDim; + size_t c = tile_col * kTileDim + threadIdx.x % kTileDim; + size_t idx = r * w + c; + if (r < h && c < w && idx >= start_offset && idx < end_offset) { + amax = fabs(static_cast(input_minus_offset[idx])); + } + + for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) { + float other_amax = __shfl_down_sync(0xFFFFFFFF, amax, delta); + __builtin_assume(amax >= 0); + __builtin_assume(other_amax >= 0); + amax = fmaxf(amax, other_amax); + } + + if (threadIdx.x % kThreadsPerWarp == 0) { + smem[threadIdx.x / kThreadsPerWarp] = amax; + } + + __syncthreads(); + + if (threadIdx.x == 0) { + for (int i = 0; i < kNumWarps; ++i) { + float other_amax = smem[i]; + __builtin_assume(amax >= 0); + __builtin_assume(other_amax >= 0); + amax = fmaxf(amax, other_amax); + } + amax_ptr[tile_row * amax_stride_h + tile_col * amax_stride_w] = amax; + } +} + +template +__global__ void __launch_bounds__(kThreadsPerBlock) + nvfp4_2d_partial_cast_kernel(const IType *input, uint8_t *output, const float *decode_scale_ptr, + const size_t scale_stride_h, const size_t scale_stride_w, + const float *global_scale_ptr, const size_t h, const size_t w, + const size_t start_offset, const size_t len) { + constexpr int kNumOutputElemsPerBank = 4; + constexpr int kThreadsPerWarp = 32; + constexpr int kLoopsPerRow = (kTileDim + kThreadsPerWarp - 1) / kThreadsPerWarp; + constexpr int kNumWarps = kThreadsPerBlock / kThreadsPerWarp; + constexpr int kRowsPerWarp = (kTileDim + kNumWarps - 1) / kNumWarps; + + __shared__ float smem[kTileDim][kTileDim + kNumOutputElemsPerBank]; + + const int tile_w = blockIdx.x; + const int tile_h = blockIdx.y; + const size_t shard_end = start_offset + len; + const IType *input_minus_offset = input - start_offset; + + float global_encode_scale = global_scale_ptr[0]; + if (global_encode_scale <= 0.f) { + global_encode_scale = 1.f; + } + const float global_decode_scale = 1.0f / global_encode_scale; + + float tile_decode_scale = decode_scale_ptr[tile_h * scale_stride_h + tile_w * scale_stride_w]; + tile_decode_scale = static_cast(static_cast(tile_decode_scale)); + constexpr float kFp32Max = 3.402823466e+38F; + float tile_encode_val = + (tile_decode_scale > 0.f) ? 1.0f / (tile_decode_scale * global_decode_scale) : kFp32Max; + tile_encode_val = fminf(tile_encode_val, kFp32Max); + const float2 scale_vec = make_float2(tile_encode_val, tile_encode_val); + + bool skip_store = true; + for (int i = 0; i < kRowsPerWarp; ++i) { + for (int j = 0; j < kLoopsPerRow; ++j) { + const int h_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i; + const int w_in_smem = threadIdx.x % kThreadsPerWarp + kThreadsPerWarp * j; + if (h_in_smem >= kTileDim || w_in_smem >= kTileDim) { + continue; + } + const int h_in_input = tile_h * kTileDim + h_in_smem; + const int w_in_input = tile_w * kTileDim + w_in_smem; + const size_t idx_in_input = static_cast(h_in_input) * w + w_in_input; + if (h_in_input < h && w_in_input < w && idx_in_input >= start_offset && + idx_in_input < shard_end) { + smem[h_in_smem][w_in_smem] = static_cast(input_minus_offset[idx_in_input]); + skip_store = false; + } + } + } + + for (int delta = kThreadsPerWarp / 2; delta > 0; delta /= 2) { + bool other = __shfl_down_sync(0xFFFFFFFF, skip_store, delta); + skip_store = skip_store && other; + } + skip_store = __shfl_sync(0xFFFFFFFF, skip_store, 0); + if (skip_store) { + return; + } + + for (int i = 0; i < kRowsPerWarp; ++i) { + const int row_in_smem = threadIdx.x / kThreadsPerWarp * kRowsPerWarp + i; + const int row_in_output = tile_h * kTileDim + row_in_smem; + if (row_in_output >= h) { + continue; + } + const int col_in_smem = threadIdx.x % kThreadsPerWarp * kNumOutputElemsPerBank; + if (col_in_smem >= kTileDim) { + continue; + } + const int col_in_output = tile_w * kTileDim + col_in_smem; + + float vals[kNumOutputElemsPerBank]; + bool mask[kNumOutputElemsPerBank]; + size_t elem_idx[kNumOutputElemsPerBank]; + bool any_valid = false; + + for (int j = 0; j < kNumOutputElemsPerBank; ++j) { + const int col = col_in_output + j; + const bool in_width = col < w; + const size_t idx = static_cast(row_in_output) * w + col; + elem_idx[j] = idx; + const bool in_shard = in_width && idx >= start_offset && idx < shard_end; + mask[j] = in_shard; + const bool in_tile = (col_in_smem + j) < kTileDim; + const float tile_val = in_tile ? smem[row_in_smem][col_in_smem + j] : 0.0f; + vals[j] = in_shard ? tile_val : 0.0f; + any_valid |= in_shard; + } + + if (!any_valid) { + continue; + } + + const float2 in01 = make_float2(vals[0], vals[1]); + const float2 in23 = make_float2(vals[2], vals[3]); + const auto packed = + transformer_engine::ptx::mul_cvt_fp32_to_fp4_4x(in01, in23, scale_vec, 0); + const uint16_t packed_bits = reinterpret_cast(packed); + + for (int pair = 0; pair < 2; ++pair) { + const int first = pair * 2; + const int second = first + 1; + if (!mask[first] && !mask[second]) { + continue; + } + const size_t ref_idx = mask[first] ? elem_idx[first] : elem_idx[second]; + const size_t byte_idx = (ref_idx - start_offset) >> 1; + uint8_t byte = output[byte_idx]; + + if (mask[first]) { + const uint8_t nibble = static_cast((packed_bits >> (4 * first)) & 0xF); + if ((elem_idx[first] & 1u) == 0) { + byte = static_cast((byte & 0xF0u) | nibble); + } else { + byte = static_cast((byte & 0x0Fu) | (nibble << 4)); + } + } + + if (mask[second]) { + const uint8_t nibble = static_cast((packed_bits >> (4 * second)) & 0xF); + if ((elem_idx[second] & 1u) == 0) { + byte = static_cast((byte & 0xF0u) | nibble); + } else { + byte = static_cast((byte & 0x0Fu) | (nibble << 4)); + } + } + + output[byte_idx] = byte; + } + } +} + +void nvfp4_2d_compute_partial_amax(const Tensor inp, Tensor amax, size_t h, size_t w, + size_t amax_stride_h, size_t amax_stride_w, size_t start_offset, + size_t block_len, cudaStream_t stream) { + NVTE_CHECK(block_len == 16, "NVFP4 2D supports 16x16 tiles only (block_len = 16)."); + + size_t len = inp.numel(); + + assert(h > 0 && w > 0); + assert(start_offset < h * w); + assert(start_offset + len <= h * w); + + size_t blocks_x = (w + kTileDim - 1) / kTileDim; + size_t blocks_y = (h + kTileDim - 1) / kTileDim; + assert(blocks_x <= std::numeric_limits::max()); + assert(blocks_y <= std::numeric_limits::max()); + dim3 grid(blocks_x, blocks_y); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + inp.dtype(), inp_dtype, + nvfp4_2d_compute_partial_amax_kernel<<>>( + reinterpret_cast(inp.data.dptr), + reinterpret_cast(amax.data.dptr), amax_stride_h, amax_stride_w, h, w, + start_offset, len);) + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void nvfp4_2d_partial_cast(const Tensor inp, Tensor out, const Tensor scale, + const Tensor global_scale, size_t h, size_t w, size_t scale_stride_h, + size_t scale_stride_w, size_t start_offset, size_t block_len, + cudaStream_t stream) { + NVTE_CHECK(block_len == 16, "NVFP4 2D supports 16x16 tiles only (block_len = 16)."); + NVTE_CHECK(out.dtype() == DType::kByte, "NVFP4 rowwise data must be uint8."); + + size_t len = inp.numel(); + + assert(h > 0 && w > 0); + assert(start_offset < h * w); + assert(start_offset + len <= h * w); + + size_t blocks_x = (w + kTileDim - 1) / kTileDim; + size_t blocks_y = (h + kTileDim - 1) / kTileDim; + assert(blocks_x <= std::numeric_limits::max()); + assert(blocks_y <= std::numeric_limits::max()); + dim3 grid(blocks_x, blocks_y); + + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + inp.dtype(), inp_dtype, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + w % kTileDim == 0, kWidthAligned, + nvfp4_2d_partial_cast_kernel + <<>>( + reinterpret_cast(inp.data.dptr), + reinterpret_cast(out.data.dptr), + reinterpret_cast(scale.data.dptr), scale_stride_h, scale_stride_w, + reinterpret_cast(global_scale.data.dptr), h, w, start_offset, len);)) + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +/* + * --------------------------------------------------------------------------- + * NVFP4 TRANSPOSE KERNEL + * + * Unlike FP8, NVFP4 packs two 4-bit values into each byte. A simple byte-wise + * transpose doesn't work because the packing changes: + * - Before transpose: elements [m, 2c] and [m, 2c+1] share a byte + * - After transpose: elements [k, 2*m_packed] and [k, 2*m_packed+1] share a byte + * which were originally [2*m_packed, k] and [2*m_packed+1, k] + * --------------------------------------------------------------------------- + */ + +// Vectorized transpose kernel parameters +constexpr int TRANSPOSE_TILE_DIM = 64; // Logical FP4 elements per tile dimension +constexpr int TRANSPOSE_TILE_PACKED = 32; // TILE_DIM / 2 bytes +constexpr int TRANSPOSE_BLOCK_SIZE = 256; // threads per block + +// Shared memory: store unpacked 4-bit values as bytes for easy transpose +// Size: TILE_DIM x (TILE_DIM + 4) to avoid bank conflicts +constexpr int TRANSPOSE_SHMEM_STRIDE = TRANSPOSE_TILE_DIM + 4; + +/* + * Vectorized transpose kernel with uint2 loads/stores (256 threads) + * Tile: 64x64 logical FP4 = 64x32 packed bytes + */ +__global__ void __launch_bounds__(TRANSPOSE_BLOCK_SIZE) + nvfp4_transpose_kernel(const uint8_t *__restrict__ input, uint8_t *__restrict__ output, + const size_t M, const size_t K) { + const size_t K_packed = K / 2; + const size_t M_packed = M / 2; + + const size_t tile_m_start = blockIdx.x * TRANSPOSE_TILE_DIM; + const size_t tile_k_start = blockIdx.y * TRANSPOSE_TILE_DIM; + + __shared__ uint8_t shmem[TRANSPOSE_TILE_DIM][TRANSPOSE_SHMEM_STRIDE]; + + const int tid = threadIdx.x; + + // Phase 1: Load input tile with VECTORIZED uint2 reads + // 256 threads, each loads 8 bytes (uint2) = 2048 bytes total + // Input tile: [64 rows, 32 cols] = 2048 bytes + { + const int thread_row = tid / 4; // 64 rows, 4 threads per row + const int thread_col = (tid % 4) * 8; // 4 x 8 = 32 bytes per row + + const size_t global_m = tile_m_start + thread_row; + const size_t global_k_packed_base = tile_k_start / 2 + thread_col; + + // Load 8 bytes as uint2 + uint2 loaded = make_uint2(0, 0); + if (global_m < M && global_k_packed_base + 7 < K_packed) { + loaded = *reinterpret_cast(&input[global_m * K_packed + global_k_packed_base]); + } else if (global_m < M) { + // Boundary: scalar loads + uint8_t *bytes = reinterpret_cast(&loaded); +#pragma unroll + for (int b = 0; b < 8; ++b) { + size_t col = global_k_packed_base + b; + bytes[b] = (col < K_packed) ? input[global_m * K_packed + col] : 0; + } + } + + // Unpack 8 bytes -> 16 nibbles and store to shared memory + const uint8_t *bytes = reinterpret_cast(&loaded); +#pragma unroll + for (int b = 0; b < 8; ++b) { + const int k0 = thread_col * 2 + b * 2; + const int k1 = k0 + 1; + shmem[thread_row][k0] = bytes[b] & 0x0F; + shmem[thread_row][k1] = (bytes[b] >> 4) & 0x0F; + } + } + + __syncthreads(); + + // Phase 2: Write output with VECTORIZED uint2 stores + // Output tile: [64 rows, 32 cols] = 2048 bytes + { + const int thread_row = tid / 4; // output K dimension [0, 64) + const int thread_col_base = (tid % 4) * 8; // output M_packed [0, 32) in steps of 8 + + const size_t global_k = tile_k_start + thread_row; + const size_t global_m_packed_base = tile_m_start / 2 + thread_col_base; + + if (global_k >= K) return; + + // Build 8 output bytes in registers + uint8_t out_bytes[8]; + +#pragma unroll + for (int b = 0; b < 8; ++b) { + const int out_m_packed = thread_col_base + b; + + if (global_m_packed_base + b >= M_packed) { + out_bytes[b] = 0; + continue; + } + + // Two M positions that pack into this output byte + const int m0 = out_m_packed * 2; + const int m1 = out_m_packed * 2 + 1; + const int k = thread_row; + + // Read from shared memory (transposed access) + const uint8_t val0 = shmem[m0][k]; + const uint8_t val1 = shmem[m1][k]; + + out_bytes[b] = val0 | (val1 << 4); + } + + // Vectorized store as uint2 + if (global_m_packed_base + 7 < M_packed) { + *reinterpret_cast(&output[global_k * M_packed + global_m_packed_base]) = + *reinterpret_cast(out_bytes); + } else { + // Boundary: scalar stores + for (int b = 0; b < 8 && global_m_packed_base + b < M_packed; ++b) { + output[global_k * M_packed + global_m_packed_base + b] = out_bytes[b]; + } + } + } +} + +void nvfp4_transpose(const Tensor input, Tensor output, cudaStream_t stream) { + // Input has logical shape [M, K], stored as [M, K/2] bytes + // Output has logical shape [K, M], stored as [K, M/2] bytes + + NVTE_CHECK(input.dtype() == DType::kByte, "NVFP4 transpose input must be uint8."); + NVTE_CHECK(output.dtype() == DType::kByte, "NVFP4 transpose output must be uint8."); + + // Get dimensions from packed storage + // input.shape() = [M, K/2], so M = shape[0], K = shape[1] * 2 + const auto in_shape = input.shape(); + NVTE_CHECK(in_shape.size() == 2, "NVFP4 transpose expects 2D input (packed), got ", + in_shape.size(), "D."); + const size_t M = in_shape[0]; + const size_t K_packed = in_shape[1]; + const size_t K = K_packed * 2; + + // Output should be [K, M/2] + const size_t M_packed = M / 2; + NVTE_CHECK(M % 2 == 0, "NVFP4 transpose requires M (", M, ") to be even."); + + const auto out_shape = output.shape(); + NVTE_CHECK(out_shape.size() == 2, "NVFP4 transpose expects 2D output."); + NVTE_CHECK(out_shape[0] == K && out_shape[1] == M_packed, + "NVFP4 transpose output shape mismatch. Expected [", K, ", ", M_packed, "], got [", + out_shape[0], ", ", out_shape[1], "]."); + + if (M == 0 || K == 0) return; + + // Use vectorized kernel (faster than TMA for pure transpose) + // 128x128 tiles with 512 threads and uint4 vectorized access + dim3 block(TRANSPOSE_BLOCK_SIZE); + dim3 grid((M + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM, + (K + TRANSPOSE_TILE_DIM - 1) / TRANSPOSE_TILE_DIM); + + nvfp4_transpose_kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output.data.dptr), M, K); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +/* + * --------------------------------------------------------------------------- + * NVFP4 SCALE TRANSPOSE KERNEL + * + * Transposes tile-level scales from rowwise to columnwise format. + * Scale values are stored as E4M3 (fp8) in uint8 tensors. + * + * Input (rowwise_scale_inv): [M_padded, K_tiles] where scales are stored + * at every 16th row (i.e., row 0, 16, 32, ... contain the actual scales, + * and each row i within a tile block has the same scale as row (i // 16) * 16). + * + * Output (columnwise_scale_inv): [K_padded, M_tiles] where scales are + * repeated 16 times per tile row. + * + * Mapping: + * output[k_tile * 16 + i, m_tile] = input[m_tile * 16, k_tile] + * for i in [0, 16) and valid (k_tile, m_tile) indices. + * --------------------------------------------------------------------------- + */ +__global__ void nvfp4_scale_transpose_kernel( + const uint8_t *__restrict__ input, // [M_padded, K_tiles], E4M3 stored as uint8 + uint8_t *__restrict__ output, // [K_padded, M_tiles], E4M3 stored as uint8 + const size_t M_tiles, // Number of M tiles + const size_t K_tiles, // Number of K tiles + const size_t input_stride, // K_tiles (input row stride) + const size_t output_stride, // M_tiles (output row stride) + const size_t K_padded // Output height +) { + // Each thread handles one output element + const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y; + const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x; + + if (out_row >= K_padded || out_col >= M_tiles) return; + + // Determine which tile row this belongs to + const size_t k_tile = out_row / kTileDim; + + // Read from input: row = m_tile * 16 (first row of the tile), col = k_tile + // m_tile = out_col + if (k_tile < K_tiles) { + const size_t in_row = out_col * kTileDim; // m_tile * 16 + const uint8_t scale = input[in_row * input_stride + k_tile]; + output[out_row * output_stride + out_col] = scale; + } else { + output[out_row * output_stride + out_col] = 0; + } +} + +void nvfp4_scale_transpose(const Tensor input, Tensor output, size_t M_tiles, size_t K_tiles, + cudaStream_t stream) { + NVTE_CHECK(input.dtype() == DType::kByte, "NVFP4 scale transpose input must be uint8 (E4M3)."); + NVTE_CHECK(output.dtype() == DType::kByte, "NVFP4 scale transpose output must be uint8 (E4M3)."); + + const auto in_shape = input.shape(); + const auto out_shape = output.shape(); + NVTE_CHECK(in_shape.size() == 2, "NVFP4 scale transpose expects 2D input."); + NVTE_CHECK(out_shape.size() == 2, "NVFP4 scale transpose expects 2D output."); + + const size_t input_stride = in_shape[1]; // K_tiles + const size_t output_stride = out_shape[1]; // M_tiles + const size_t K_padded = out_shape[0]; + + if (M_tiles == 0 || K_tiles == 0 || K_padded == 0) return; + + constexpr int kBlockDim = 16; + dim3 block(kBlockDim, kBlockDim); + dim3 grid((M_tiles + kBlockDim - 1) / kBlockDim, (K_padded + kBlockDim - 1) / kBlockDim); + + nvfp4_scale_transpose_kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output.data.dptr), M_tiles, K_tiles, input_stride, output_stride, + K_padded); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +/* + * --------------------------------------------------------------------------- + * NVFP4 SCALE EXPANSION KERNEL + * + * Expands tile-level scales to row-level scales and converts to FP8 E4M3, used in partial cast. + * + * Input (per_block_decode_scale): [tile_rows, tile_cols] in float32 + * Output (target_scale): [rows_padded, tile_cols] in uint8 (E4M3) + * + * Each tile row's scale is repeated block_len times in the output. + * --------------------------------------------------------------------------- + */ +__global__ void nvfp4_expand_scale_to_fp8_kernel( + const float *__restrict__ input, // [tile_rows, tile_cols] + uint8_t *__restrict__ output, // [rows_padded, tile_cols] + const size_t tile_rows, const size_t tile_cols, const size_t rows_padded, + const size_t block_len) { + const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y; + const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x; + + if (out_row >= rows_padded || out_col >= tile_cols) return; + + // Determine which tile row this output row belongs to + const size_t tile_row = out_row / block_len; + + float scale_val = 0.0f; + if (tile_row < tile_rows) { + scale_val = input[tile_row * tile_cols + out_col]; + } + + // Convert float32 to FP8 E4M3 + // Clamp to FP8 E4M3 range and convert + fp8e4m3 fp8_val = static_cast(scale_val); + output[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val); +} + +void nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_rows, + size_t tile_cols, size_t rows_padded, size_t block_len, + cudaStream_t stream) { + NVTE_CHECK(input.dtype() == DType::kFloat32, "Scale input must be float32."); + NVTE_CHECK(output.dtype() == DType::kByte, "Scale output must be uint8 (E4M3)."); + + if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return; + + constexpr int kBlockDim = 16; + dim3 block(kBlockDim, kBlockDim); + dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim); + + nvfp4_expand_scale_to_fp8_kernel<<>>( + reinterpret_cast(input.data.dptr), + reinterpret_cast(output.data.dptr), tile_rows, tile_cols, rows_padded, block_len); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +/* + * --------------------------------------------------------------------------- + * NVFP4 COMPUTE PER-BLOCK DECODE SCALE KERNEL + * + * Computes per-block decode scale from block amax and global amax: + * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax + * per_block_decode_scale = block_amax / fp4_max * global_scale + * = block_amax * 448 / global_amax + * + * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh + * + * Input (block_amax): [tile_rows, tile_cols] in float32 + * Input (global_amax): scalar float32 (per-tensor amax after all-reduce) + * Output (scale): [tile_rows, tile_cols] in float32 + * Output (global_scale_out): scalar float32 (the computed global encode scale) + * --------------------------------------------------------------------------- + */ +__global__ void nvfp4_compute_per_block_scale_kernel( + const float *__restrict__ block_amax, // [tile_rows, tile_cols] + float *__restrict__ scale, // [tile_rows, tile_cols] + const float *__restrict__ global_amax_ptr, // Pointer to single float value (avoids D2H) + const size_t numel) { + const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= numel) return; + + constexpr float fp4_max = 6.0f; + constexpr float fp8_max = 448.0f; + constexpr float flt_max = 3.402823466e+38f; + constexpr float tiny = 1.17549435e-38f; // FLT_MIN + + // Read global_amax from device memory (avoids D2H transfer) + float global_amax = *global_amax_ptr; + + // Compute global encode scale: S_enc = (fp8_max * fp4_max) / global_amax + float safe_global_amax = fmaxf(global_amax, tiny); + float global_scale = + (global_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; + + // Compute per-block decode scale: S_dec_b = block_amax / fp4_max * S_enc + float amax_val = block_amax[idx]; + float result = fminf((amax_val / fp4_max) * global_scale, flt_max); + scale[idx] = result; +} + +// Simple kernel to compute global encode scale from global amax +__global__ void nvfp4_compute_global_scale_kernel( + const float *__restrict__ global_amax, // [num_params] + float *__restrict__ global_scale, // [num_params] + const size_t num_params) { + const size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= num_params) return; + + constexpr float fp4_max = 6.0f; + constexpr float fp8_max = 448.0f; + constexpr float flt_max = 3.402823466e+38f; + constexpr float tiny = 1.17549435e-38f; // FLT_MIN + + float amax = global_amax[idx]; + float safe_amax = fmaxf(amax, tiny); + float scale = (amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_amax, flt_max) : 1.0f; + global_scale[idx] = scale; +} + +void nvfp4_compute_per_block_scale(const Tensor block_amax, Tensor scale, const Tensor global_amax, + cudaStream_t stream) { + NVTE_CHECK(block_amax.dtype() == DType::kFloat32, "Block amax must be float32."); + NVTE_CHECK(scale.dtype() == DType::kFloat32, "Scale must be float32."); + NVTE_CHECK(global_amax.dtype() == DType::kFloat32, "Global amax must be float32."); + NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); + + size_t numel = block_amax.numel(); + if (numel == 0) return; + + constexpr int kBlockSize = 256; + int grid_size = (numel + kBlockSize - 1) / kBlockSize; + + nvfp4_compute_per_block_scale_kernel<<>>( + reinterpret_cast(block_amax.data.dptr), + reinterpret_cast(scale.data.dptr), + reinterpret_cast(global_amax.data.dptr), numel); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void nvfp4_compute_global_scale(const Tensor global_amax, Tensor global_scale, + cudaStream_t stream) { + NVTE_CHECK(global_amax.dtype() == DType::kFloat32, "Global amax must be float32."); + NVTE_CHECK(global_scale.dtype() == DType::kFloat32, "Global scale must be float32."); + + size_t num_params = global_amax.numel(); + if (num_params == 0) return; + + constexpr int kBlockSize = 256; + int grid_size = (num_params + kBlockSize - 1) / kBlockSize; + + nvfp4_compute_global_scale_kernel<<>>( + reinterpret_cast(global_amax.data.dptr), + reinterpret_cast(global_scale.data.dptr), num_params); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +/* + * --------------------------------------------------------------------------- + * FUSED NVFP4 SCALE COMPUTATION KERNEL + * + * Fuses three operations into one kernel: + * 1. nvfp4_compute_per_block_scale: compute tile-level decode scales from block amax + * 2. target_amax.copy_: copy global amax to target tensor + * 3. nvfp4_expand_scale_to_fp8: expand to row-level and convert to FP8 E4M3 + * + * Input (block_amax): [tile_rows, tile_cols] float32 + * Input (global_amax): [1] float32 + * Output (per_block_scale): [tile_rows, tile_cols] float32 (intermediate, for partial_cast) + * Output (target_scale): [rows_padded, tile_cols] uint8 (E4M3) + * Output (target_amax): [1] float32 (copy of global_amax) + * + * Saves 2 kernel launches per parameter (eliminates nvfp4_compute_per_block_scale and + * nvfp4_expand_scale_to_fp8 as separate calls, plus the amax copy). + * --------------------------------------------------------------------------- + */ +__global__ void nvfp4_fused_scale_kernel( + const float *__restrict__ block_amax, // [tile_rows, tile_cols] + const float *__restrict__ global_amax, // [1] + float *__restrict__ per_block_scale, // [tile_rows, tile_cols] - for partial_cast + uint8_t *__restrict__ target_scale, // [rows_padded, tile_cols] + float *__restrict__ target_amax, // [1] + const size_t tile_rows, const size_t tile_cols, const size_t rows_padded, + const size_t block_len) { + const size_t out_row = blockIdx.y * blockDim.y + threadIdx.y; + const size_t out_col = blockIdx.x * blockDim.x + threadIdx.x; + + // Read global amax once per thread (broadcast) + const float g_amax = *global_amax; + + // Thread (0,0) copies global_amax to target_amax + if (out_row == 0 && out_col == 0) { + *target_amax = g_amax; + } + + if (out_row >= rows_padded || out_col >= tile_cols) return; + + // Determine which tile row this output row belongs to + const size_t tile_row = out_row / block_len; + + // Compute the scale value + constexpr float fp4_max = 6.0f; + constexpr float fp8_max = 448.0f; + constexpr float flt_max = 3.402823466e+38f; + constexpr float tiny = 1.17549435e-38f; + + float scale_val = 0.0f; + if (tile_row < tile_rows) { + float safe_global_amax = fmaxf(g_amax, tiny); + float global_scale = + (g_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; + + // Read block amax and compute per-block decode scale + float amax_val = block_amax[tile_row * tile_cols + out_col]; + scale_val = fminf((amax_val / fp4_max) * global_scale, flt_max); + + // Write per-block scale (only once per tile, when out_row % block_len == 0) + if (out_row % block_len == 0) { + per_block_scale[tile_row * tile_cols + out_col] = scale_val; + } + } + + // Convert float32 to FP8 E4M3 and write expanded scale + fp8e4m3 fp8_val = static_cast(scale_val); + target_scale[out_row * tile_cols + out_col] = reinterpret_cast(fp8_val); +} + +void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor per_block_scale, + Tensor target_scale, Tensor target_amax, size_t tile_rows, size_t tile_cols, + size_t rows_padded, size_t block_len, cudaStream_t stream) { + NVTE_CHECK(block_amax.dtype() == DType::kFloat32, "Block amax must be float32."); + NVTE_CHECK(global_amax.dtype() == DType::kFloat32, "Global amax must be float32."); + NVTE_CHECK(per_block_scale.dtype() == DType::kFloat32, "Per-block scale must be float32."); + NVTE_CHECK(target_scale.dtype() == DType::kByte, "Target scale must be uint8 (E4M3)."); + NVTE_CHECK(target_amax.dtype() == DType::kFloat32, "Target amax must be float32."); + NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); + NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); + + if (tile_rows == 0 || tile_cols == 0 || rows_padded == 0) return; + + constexpr int kBlockDim = 16; + dim3 block(kBlockDim, kBlockDim); + dim3 grid((tile_cols + kBlockDim - 1) / kBlockDim, (rows_padded + kBlockDim - 1) / kBlockDim); + + nvfp4_fused_scale_kernel<<>>( + reinterpret_cast(block_amax.data.dptr), + reinterpret_cast(global_amax.data.dptr), + reinterpret_cast(per_block_scale.data.dptr), + reinterpret_cast(target_scale.data.dptr), + reinterpret_cast(target_amax.data.dptr), tile_rows, tile_cols, rows_padded, + block_len); + NVTE_CHECK_CUDA(cudaGetLastError()); +} } // namespace nvfp4_recipe } // namespace transformer_engine +void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, + size_t tile_cols, size_t rows_padded, size_t block_len, + cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_expand_scale_to_fp8(*convertNVTETensorCheck(input), + *convertNVTETensorCheck(output), tile_rows, tile_cols, + rows_padded, block_len, stream); +} + +void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, + const NVTETensor global_amax, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_compute_per_block_scale(*convertNVTETensorCheck(block_amax), + *convertNVTETensorCheck(scale), + *convertNVTETensorCheck(global_amax), stream); +} + +void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, + cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_compute_global_scale); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_compute_global_scale(*convertNVTETensorCheck(global_amax), + *convertNVTETensorCheck(global_scale), stream); +} + +void nvte_nvfp4_scale_transpose(const NVTETensor input, NVTETensor output, size_t M_tiles, + size_t K_tiles, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_scale_transpose); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_scale_transpose(*convertNVTETensorCheck(input), + *convertNVTETensorCheck(output), M_tiles, K_tiles, stream); +} + +void nvte_nvfp4_data_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_data_transpose); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), + stream); +} + +void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h, size_t w, + size_t amax_stride_h, size_t amax_stride_w, + size_t start_offset, size_t block_len, + cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_2d_compute_partial_amax); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_2d_compute_partial_amax(*convertNVTETensorCheck(inp), + *convertNVTETensorCheck(amax), h, w, amax_stride_h, + amax_stride_w, start_offset, block_len, stream); +} + +void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, + const NVTETensor global_scale, size_t h, size_t w, + size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, + size_t block_len, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_2d_partial_cast(*convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), + *convertNVTETensorCheck(scale), + *convertNVTETensorCheck(global_scale), h, w, scale_stride_h, + scale_stride_w, start_offset, block_len, stream); +} + void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_rowwise_amax_A, const NVTETensor inpB, const bool use_rowwise_amax_B, float alpha_in, NVTETensor alpha_out, @@ -52,3 +899,15 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r reinterpret_cast(amax_B_ptr), reinterpret_cast(alpha_ptr)); NVTE_CHECK_CUDA(cudaGetLastError()); } + +void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, + NVTETensor per_block_scale, NVTETensor target_scale, + NVTETensor target_amax, size_t tile_rows, size_t tile_cols, + size_t rows_padded, size_t block_len, cudaStream_t stream) { + NVTE_API_CALL(nvte_nvfp4_fused_scale); + using namespace transformer_engine; + nvfp4_recipe::nvfp4_fused_scale( + *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(global_amax), + *convertNVTETensorCheck(per_block_scale), *convertNVTETensorCheck(target_scale), + *convertNVTETensorCheck(target_amax), tile_rows, tile_cols, rows_padded, block_len, stream); +} diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index b2b0751b04..e4d4e5094c 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -156,6 +156,39 @@ std::optional> te_general_grouped_gemm( at::Tensor fp8_transpose(at::Tensor input, DType otype, std::optional output = std::nullopt); +at::Tensor nvfp4_data_transpose(at::Tensor input, std::optional output = std::nullopt); + +void nvfp4_2d_scale_transpose(at::Tensor input, at::Tensor output, int64_t M_tiles, + int64_t K_tiles); + +void nvfp4_2d_multi_tensor_transpose(std::vector rowwise_data_list, + std::vector columnwise_data_list, + std::vector rowwise_scale_inv_list, + std::vector columnwise_scale_inv_list, + std::vector M_list, std::vector K_list); + +void nvfp4_multi_tensor_compute_partial_amax( + std::vector master_weight_list, std::vector partial_amax_list, + std::vector global_amax_list, std::vector h_list, + std::vector w_list, std::vector start_offset_list, int64_t block_len); + +void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows, + int64_t tile_cols, int64_t rows_padded, int64_t block_len); + +void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, at::Tensor global_amax); + +void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, + at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows, + int64_t tile_cols, int64_t rows_padded, int64_t block_len); + +void nvfp4_multi_tensor_fused_scale( + std::vector block_amax_list, std::vector global_amax_list, + std::vector per_block_scale_list, std::vector target_scale_list, + std::vector target_amax_list, std::vector tile_rows_list, + std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len); + +void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale); + at::Tensor swap_first_dims(at::Tensor tensor, std::optional out = std::nullopt); /*************************************************************************************************** @@ -344,6 +377,19 @@ void fp8_block_scaling_partial_cast(const at::Tensor &inp, at::Tensor out, const size_t h, size_t w, size_t start_offset, size_t block_len, const DType out_dtype); +void nvfp4_2d_compute_partial_amax(const at::Tensor &tensor, at::Tensor amax, size_t h, size_t w, + size_t start_offset, size_t block_len); + +void nvfp4_2d_partial_cast(const at::Tensor &inp, py::handle out, const at::Tensor &scale, + const at::Tensor &global_scale, size_t h, size_t w, size_t start_offset, + size_t block_len); + +void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, + std::vector out_list, + std::vector scale_list, + std::vector global_scale_list, + std::vector h_list, std::vector w_list, + std::vector start_offset_list, int64_t block_len); void mxfp8_scaling_compute_partial_amax(const at::Tensor &input, at::Tensor amax_rowwise, at::Tensor amax_colwise, int rows, int cols, size_t start_offset); diff --git a/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp new file mode 100644 index 0000000000..685250d137 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/nvfp4_2d_partial_cast.cpp @@ -0,0 +1,156 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../extensions.h" + +namespace transformer_engine::pytorch { + +void nvfp4_2d_compute_partial_amax(const at::Tensor& tensor, at::Tensor amax, size_t h, size_t w, + size_t start_offset, size_t block_len) { + TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); + TORCH_CHECK(amax.dim() == 2, "amax must be a 2D tensor"); + TORCH_CHECK(amax.scalar_type() == at::ScalarType::Float, "amax must be a float tensor"); + TORCH_CHECK(tensor.scalar_type() == at::ScalarType::Float || + tensor.scalar_type() == at::ScalarType::BFloat16, + "tensor must be a float or bfloat16 tensor"); + + const TensorWrapper tensor_cu = makeTransformerEngineTensor(tensor.contiguous()); + TensorWrapper amax_cu = makeTransformerEngineTensor(amax); + + nvte_nvfp4_2d_compute_partial_amax(tensor_cu.data(), amax_cu.data(), h, w, amax.stride(0), + amax.stride(1), start_offset, block_len, + at::cuda::getCurrentCUDAStream()); +} + +void nvfp4_2d_partial_cast(const at::Tensor& inp, py::handle out, const at::Tensor& scale, + const at::Tensor& global_scale, size_t h, size_t w, size_t start_offset, + size_t block_len) { + TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); + TORCH_CHECK(scale.dim() == 2, "scale must be a 2D tensor"); + TORCH_CHECK(scale.scalar_type() == at::ScalarType::Float, "scale must be a float tensor"); + TORCH_CHECK(global_scale.numel() == 1, "global_scale must be a scalar tensor"); + TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, + "global_scale must be a float tensor"); + TORCH_CHECK( + inp.scalar_type() == at::ScalarType::Float || inp.scalar_type() == at::ScalarType::BFloat16, + "input must be a float or bfloat16 tensor"); + + const TensorWrapper inp_cu = makeTransformerEngineTensor(inp.contiguous()); + const TensorWrapper out_cu = makeTransformerEngineTensor(out, py::none()); + const TensorWrapper scale_cu = makeTransformerEngineTensor(scale); + const TensorWrapper global_scale_cu = makeTransformerEngineTensor(global_scale); + + nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), global_scale_cu.data(), + h, w, scale.stride(0), scale.stride(1), start_offset, block_len, + at::cuda::getCurrentCUDAStream()); +} + +void nvfp4_multi_tensor_2d_partial_cast(std::vector inp_list, + std::vector out_list, + std::vector scale_list, + std::vector global_scale_list, + std::vector h_list, std::vector w_list, + std::vector start_offset_list, int64_t block_len) { + TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); + + const size_t num_tensors = inp_list.size(); + TORCH_CHECK(out_list.size() == num_tensors, "out_list size mismatch"); + TORCH_CHECK(scale_list.size() == num_tensors, "scale_list size mismatch"); + TORCH_CHECK(global_scale_list.size() == num_tensors, "global_scale_list size mismatch"); + TORCH_CHECK(h_list.size() == num_tensors, "h_list size mismatch"); + TORCH_CHECK(w_list.size() == num_tensors, "w_list size mismatch"); + TORCH_CHECK(start_offset_list.size() == num_tensors, "start_offset_list size mismatch"); + + if (num_tensors == 0) { + return; + } + + auto stream = at::cuda::getCurrentCUDAStream(); + + for (size_t i = 0; i < num_tensors; ++i) { + const auto& inp = inp_list[i]; + const auto& out = out_list[i]; + const auto& scale = scale_list[i]; + const auto& global_scale = global_scale_list[i]; + const size_t h = static_cast(h_list[i]); + const size_t w = static_cast(w_list[i]); + const size_t start_offset = static_cast(start_offset_list[i]); + + TORCH_CHECK(scale.dim() == 2, "scale must be a 2D tensor"); + TORCH_CHECK(scale.scalar_type() == at::ScalarType::Float, "scale must be a float tensor"); + TORCH_CHECK(global_scale.numel() == 1, "global_scale must be a scalar tensor"); + TORCH_CHECK(global_scale.scalar_type() == at::ScalarType::Float, + "global_scale must be a float tensor"); + TORCH_CHECK( + inp.scalar_type() == at::ScalarType::Float || inp.scalar_type() == at::ScalarType::BFloat16, + "input must be a float or bfloat16 tensor"); + + const TensorWrapper inp_cu = makeTransformerEngineTensor(inp.contiguous()); + const TensorWrapper out_cu = makeTransformerEngineTensor(out); + const TensorWrapper scale_cu = makeTransformerEngineTensor(scale); + const TensorWrapper global_scale_cu = makeTransformerEngineTensor(global_scale); + + nvte_nvfp4_2d_partial_cast(inp_cu.data(), out_cu.data(), scale_cu.data(), + global_scale_cu.data(), h, w, scale.stride(0), scale.stride(1), + start_offset, static_cast(block_len), stream); + } +} + +void nvfp4_multi_tensor_compute_partial_amax( + std::vector master_weight_list, std::vector partial_amax_list, + std::vector global_amax_list, std::vector h_list, + std::vector w_list, std::vector start_offset_list, int64_t block_len) { + TORCH_CHECK(block_len == 16, "Currently only block_len = 16 is supported for NVFP4 2D"); + + const size_t num_tensors = master_weight_list.size(); + TORCH_CHECK(partial_amax_list.size() == num_tensors, "partial_amax_list size mismatch"); + TORCH_CHECK(global_amax_list.size() == num_tensors, "global_amax_list size mismatch"); + TORCH_CHECK(h_list.size() == num_tensors, "h_list size mismatch"); + TORCH_CHECK(w_list.size() == num_tensors, "w_list size mismatch"); + TORCH_CHECK(start_offset_list.size() == num_tensors, "start_offset_list size mismatch"); + + if (num_tensors == 0) { + return; + } + + auto stream = at::cuda::getCurrentCUDAStream(); + + for (size_t i = 0; i < num_tensors; ++i) { + const auto& master_weight = master_weight_list[i]; + auto& partial_amax = partial_amax_list[i]; + auto& global_amax = global_amax_list[i]; + const size_t h = static_cast(h_list[i]); + const size_t w = static_cast(w_list[i]); + const size_t start_offset = static_cast(start_offset_list[i]); + + TORCH_CHECK(partial_amax.dim() == 2, "partial_amax must be a 2D tensor"); + TORCH_CHECK(partial_amax.scalar_type() == at::ScalarType::Float, + "partial_amax must be a float tensor"); + TORCH_CHECK(master_weight.scalar_type() == at::ScalarType::Float || + master_weight.scalar_type() == at::ScalarType::BFloat16, + "master_weight must be a float or bfloat16 tensor"); + TORCH_CHECK(global_amax.scalar_type() == at::ScalarType::Float, + "global_amax must be a float tensor"); + TORCH_CHECK(global_amax.numel() == 1, "global_amax must have exactly one element"); + + // Compute partial amax (per-block amax) + const TensorWrapper tensor_cu = makeTransformerEngineTensor(master_weight.contiguous()); + TensorWrapper amax_cu = makeTransformerEngineTensor(partial_amax); + + nvte_nvfp4_2d_compute_partial_amax(tensor_cu.data(), amax_cu.data(), h, w, + partial_amax.stride(0), partial_amax.stride(1), start_offset, + static_cast(block_len), stream); + + // Compute global amax + auto* global_amax_ptr = global_amax.data_ptr(); + TensorWrapper fake_te_output( + /*dptr=*/nullptr, tensor_cu.shape(), DType::kFloat32, global_amax_ptr); + + nvte_compute_amax(tensor_cu.data(), fake_te_output.data(), stream); + } +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index b9fc65363d..b1d60cc3da 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -272,6 +272,44 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fp8_transpose", &transformer_engine::pytorch::fp8_transpose, "Transpose with FP8 I/O", py::arg("input"), py::arg("dtype"), py::kw_only(), py::arg("out"), py::call_guard()); + m.def("nvfp4_data_transpose", &transformer_engine::pytorch::nvfp4_data_transpose, + "Transpose NVFP4 packed data with nibble repacking", py::arg("input"), py::kw_only(), + py::arg("out"), py::call_guard()); + m.def( + "nvfp4_2d_scale_transpose", &transformer_engine::pytorch::nvfp4_2d_scale_transpose, + "Transpose NVFP4 tile-level scales (E4M3 stored as uint8) from rowwise to columnwise format", + py::arg("input"), py::arg("output"), py::arg("M_tiles"), py::arg("K_tiles"), + py::call_guard()); + m.def("nvfp4_expand_scale_to_fp8", &transformer_engine::pytorch::nvfp4_expand_scale_to_fp8, + "Expand tile-level scales to row-level scales and convert to FP8 E4M3", py::arg("input"), + py::arg("output"), py::arg("tile_rows"), py::arg("tile_cols"), py::arg("rows_padded"), + py::arg("block_len"), py::call_guard()); + m.def("nvfp4_compute_per_block_scale", + &transformer_engine::pytorch::nvfp4_compute_per_block_scale, + "Compute per-block decode scale from block amax and global amax", py::arg("block_amax"), + py::arg("scale"), py::arg("global_amax"), py::call_guard()); + m.def("nvfp4_compute_global_scale", &transformer_engine::pytorch::nvfp4_compute_global_scale, + "Compute global encode scale from global amax", py::arg("global_amax"), + py::arg("global_scale"), py::call_guard()); + m.def("nvfp4_fused_scale", &transformer_engine::pytorch::nvfp4_fused_scale, + "Fused kernel: compute per-block decode scale, copy global amax, expand to row-level FP8", + py::arg("block_amax"), py::arg("global_amax"), py::arg("per_block_scale"), + py::arg("target_scale"), py::arg("target_amax"), py::arg("tile_rows"), py::arg("tile_cols"), + py::arg("rows_padded"), py::arg("block_len"), py::call_guard()); + m.def("nvfp4_multi_tensor_fused_scale", + &transformer_engine::pytorch::nvfp4_multi_tensor_fused_scale, + "Batched fused scale: compute per-block decode scale, copy global amax, expand to FP8 for " + "multiple tensors", + py::arg("block_amax_list"), py::arg("global_amax_list"), py::arg("per_block_scale_list"), + py::arg("target_scale_list"), py::arg("target_amax_list"), py::arg("tile_rows_list"), + py::arg("tile_cols_list"), py::arg("rows_padded_list"), py::arg("block_len"), + py::call_guard()); + m.def("nvfp4_2d_multi_tensor_transpose", + &transformer_engine::pytorch::nvfp4_2d_multi_tensor_transpose, + "Batched NVFP4 columnwise creation: transpose data and scales for multiple tensors", + py::arg("rowwise_data_list"), py::arg("columnwise_data_list"), + py::arg("rowwise_scale_inv_list"), py::arg("columnwise_scale_inv_list"), py::arg("M_list"), + py::arg("K_list"), py::call_guard()); m.def("swap_first_dims", &transformer_engine::pytorch::swap_first_dims, "Swap first two tensor dimensions", py::arg("tensor"), py::kw_only(), py::arg("out"), py::call_guard()); @@ -294,6 +332,29 @@ 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, + "Compute partial amax from master weights for NVFP4 2D", py::arg("tensor"), py::arg("amax"), + py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len") = 16, + py::call_guard()); + m.def("nvfp4_multi_tensor_compute_partial_amax", + &transformer_engine::pytorch::nvfp4_multi_tensor_compute_partial_amax, + "Batched compute partial and global amax from master weights for NVFP4 2D", + py::arg("master_weight_list"), py::arg("partial_amax_list"), py::arg("global_amax_list"), + py::arg("h_list"), py::arg("w_list"), py::arg("start_offset_list"), + py::arg("block_len") = 16, py::call_guard()); + m.def("nvfp4_2d_partial_cast", &transformer_engine::pytorch::nvfp4_2d_partial_cast, + "Partial cast from master weights for NVFP4 2D", py::arg("inp"), py::arg("out"), + py::arg("scale"), py::arg("global_scale"), py::arg("h"), py::arg("w"), + py::arg("start_offset"), py::arg("block_len") = 16, + py::call_guard()); + m.def("nvfp4_multi_tensor_2d_partial_cast", + &transformer_engine::pytorch::nvfp4_multi_tensor_2d_partial_cast, + "Batched partial cast from master weights for NVFP4 2D", py::arg("inp_list"), + py::arg("out_list"), py::arg("scale_list"), py::arg("global_scale_list"), py::arg("h_list"), + py::arg("w_list"), py::arg("start_offset_list"), py::arg("block_len") = 16, + py::call_guard()); m.def("mxfp8_scaling_compute_partial_amax", &transformer_engine::pytorch::mxfp8_scaling_compute_partial_amax, "Compute partial amax from master weights for fp8 mxfp8 scaling", py::arg("input"), diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index 477d7c87e7..aaa27a104a 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -5,6 +5,8 @@ ************************************************************************/ #include +#include +#include #include #include @@ -52,11 +54,218 @@ at::Tensor fp8_transpose(at::Tensor input, DType otype, std::optional output) { + init_extension(); + + // Input is packed FP4: logical [M, K] stored as [M, K/2] bytes + // Output is packed FP4: logical [K, M] stored as [K, M/2] bytes + const auto shape = getTensorShape(input); + NVTE_CHECK(shape.size() == 2, "NVFP4 transpose expects 2D input (packed storage)."); + + const size_t M = shape[0]; + const size_t K_packed = shape[1]; + const size_t K = K_packed * 2; // logical K + const size_t M_packed = M / 2; + + NVTE_CHECK(M % 2 == 0, "NVFP4 transpose requires M (", M, ") to be even."); + + // Output shape: [K, M/2] + std::vector output_shape = {static_cast(K), static_cast(M_packed)}; + + // Output tensor + at::Tensor out; + if (output.has_value()) { + out = *output; + NVTE_CHECK( + static_cast(out.size(0)) == K && static_cast(out.size(1)) == M_packed, + "Output shape mismatch for NVFP4 transpose."); + } else { + const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + out = at::empty(output_shape, opts); + } + + // Return immediately if tensor is empty + if (M == 0 || K == 0) { + return out; + } + + // Call the NVFP4 transpose kernel + auto input_cu = + makeTransformerEngineTensor(input.data_ptr(), std::vector{M, K_packed}, DType::kByte); + auto output_cu = + makeTransformerEngineTensor(out.data_ptr(), std::vector{K, M_packed}, DType::kByte); + nvte_nvfp4_data_transpose(input_cu.data(), output_cu.data(), at::cuda::getCurrentCUDAStream()); + + return out; +} + +void nvfp4_2d_scale_transpose(at::Tensor input, at::Tensor output, int64_t M_tiles, + int64_t K_tiles) { + init_extension(); + + // Input: rowwise_scale_inv [M_padded, K_tiles], uint8 (E4M3 stored as bytes) + // Output: columnwise_scale_inv [K_padded, M_tiles], uint8 (E4M3 stored as bytes) + const auto in_shape = getTensorShape(input); + const auto out_shape = getTensorShape(output); + NVTE_CHECK(in_shape.size() == 2, "NVFP4 scale transpose expects 2D input."); + NVTE_CHECK(out_shape.size() == 2, "NVFP4 scale transpose expects 2D output."); + NVTE_CHECK(input.scalar_type() == at::kByte, "NVFP4 scale transpose input must be uint8 (E4M3)."); + NVTE_CHECK(output.scalar_type() == at::kByte, + "NVFP4 scale transpose output must be uint8 (E4M3)."); + + auto input_cu = makeTransformerEngineTensor( + input.data_ptr(), std::vector{in_shape[0], in_shape[1]}, DType::kByte); + auto output_cu = makeTransformerEngineTensor( + output.data_ptr(), std::vector{out_shape[0], out_shape[1]}, DType::kByte); + + nvte_nvfp4_scale_transpose(input_cu.data(), output_cu.data(), static_cast(M_tiles), + static_cast(K_tiles), at::cuda::getCurrentCUDAStream()); +} + +void nvfp4_expand_scale_to_fp8(at::Tensor input, at::Tensor output, int64_t tile_rows, + int64_t tile_cols, int64_t rows_padded, int64_t block_len) { + init_extension(); + + // Input: per_block_decode_scale [tile_rows, tile_cols], float32 + // Output: target_scale [rows_padded, tile_cols], uint8 (E4M3) + const auto in_shape = getTensorShape(input); + const auto out_shape = getTensorShape(output); + NVTE_CHECK(in_shape.size() == 2, "NVFP4 expand scale expects 2D input."); + NVTE_CHECK(out_shape.size() == 2, "NVFP4 expand scale expects 2D output."); + NVTE_CHECK(input.scalar_type() == at::kFloat, "NVFP4 expand scale input must be float32."); + NVTE_CHECK(output.scalar_type() == at::kByte, "NVFP4 expand scale output must be uint8 (E4M3)."); + + auto input_cu = makeTransformerEngineTensor( + input.data_ptr(), std::vector{in_shape[0], in_shape[1]}, DType::kFloat32); + auto output_cu = makeTransformerEngineTensor( + output.data_ptr(), std::vector{out_shape[0], out_shape[1]}, DType::kByte); + + nvte_nvfp4_expand_scale_to_fp8(input_cu.data(), output_cu.data(), static_cast(tile_rows), + static_cast(tile_cols), static_cast(rows_padded), + static_cast(block_len), at::cuda::getCurrentCUDAStream()); +} + +void nvfp4_compute_per_block_scale(at::Tensor block_amax, at::Tensor scale, + at::Tensor global_amax) { + init_extension(); + + // block_amax and scale: [tile_rows, tile_cols], float32 + // global_amax: single element tensor, float32 (avoids D2H transfer) + NVTE_CHECK(block_amax.scalar_type() == at::kFloat, "Block amax must be float32."); + NVTE_CHECK(scale.scalar_type() == at::kFloat, "Scale must be float32."); + NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); + NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); + + auto block_amax_cu = makeTransformerEngineTensor(block_amax); + auto scale_cu = makeTransformerEngineTensor(scale); + auto global_amax_cu = makeTransformerEngineTensor(global_amax); + + nvte_nvfp4_compute_per_block_scale(block_amax_cu.data(), scale_cu.data(), global_amax_cu.data(), + at::cuda::getCurrentCUDAStream()); +} + +void nvfp4_fused_scale(at::Tensor block_amax, at::Tensor global_amax, at::Tensor per_block_scale, + at::Tensor target_scale, at::Tensor target_amax, int64_t tile_rows, + int64_t tile_cols, int64_t rows_padded, int64_t block_len) { + init_extension(); + + // block_amax: [tile_rows, tile_cols], float32 + // global_amax: [1], float32 + // per_block_scale: [tile_rows, tile_cols], float32 (for partial_cast) + // target_scale: [rows_padded, tile_cols], uint8 (E4M3) + // target_amax: [1], float32 + NVTE_CHECK(block_amax.scalar_type() == at::kFloat, "Block amax must be float32."); + NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); + NVTE_CHECK(per_block_scale.scalar_type() == at::kFloat, "Per-block scale must be float32."); + NVTE_CHECK(target_scale.scalar_type() == at::kByte, "Target scale must be uint8 (E4M3)."); + NVTE_CHECK(target_amax.scalar_type() == at::kFloat, "Target amax must be float32."); + NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); + NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); + + auto block_amax_cu = makeTransformerEngineTensor(block_amax); + auto global_amax_cu = makeTransformerEngineTensor(global_amax); + auto per_block_scale_cu = makeTransformerEngineTensor(per_block_scale); + auto target_scale_cu = makeTransformerEngineTensor(target_scale); + auto target_amax_cu = makeTransformerEngineTensor(target_amax); + + nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), + target_scale_cu.data(), target_amax_cu.data(), + static_cast(tile_rows), static_cast(tile_cols), + static_cast(rows_padded), static_cast(block_len), + at::cuda::getCurrentCUDAStream()); +} + +void nvfp4_multi_tensor_fused_scale( + std::vector block_amax_list, std::vector global_amax_list, + std::vector per_block_scale_list, std::vector target_scale_list, + std::vector target_amax_list, std::vector tile_rows_list, + std::vector tile_cols_list, std::vector rows_padded_list, int64_t block_len) { + init_extension(); + + const size_t num_tensors = block_amax_list.size(); + NVTE_CHECK(global_amax_list.size() == num_tensors, "global_amax_list size mismatch"); + NVTE_CHECK(per_block_scale_list.size() == num_tensors, "per_block_scale_list size mismatch"); + NVTE_CHECK(target_scale_list.size() == num_tensors, "target_scale_list size mismatch"); + NVTE_CHECK(target_amax_list.size() == num_tensors, "target_amax_list size mismatch"); + NVTE_CHECK(tile_rows_list.size() == num_tensors, "tile_rows_list size mismatch"); + NVTE_CHECK(tile_cols_list.size() == num_tensors, "tile_cols_list size mismatch"); + NVTE_CHECK(rows_padded_list.size() == num_tensors, "rows_padded_list size mismatch"); + + if (num_tensors == 0) { + return; + } + + auto stream = at::cuda::getCurrentCUDAStream(); + + for (size_t i = 0; i < num_tensors; ++i) { + const auto& block_amax = block_amax_list[i]; + const auto& global_amax = global_amax_list[i]; + auto& per_block_scale = per_block_scale_list[i]; + auto& target_scale = target_scale_list[i]; + auto& target_amax = target_amax_list[i]; + const size_t tile_rows = static_cast(tile_rows_list[i]); + const size_t tile_cols = static_cast(tile_cols_list[i]); + const size_t rows_padded = static_cast(rows_padded_list[i]); + + NVTE_CHECK(block_amax.scalar_type() == at::kFloat, "Block amax must be float32."); + NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); + NVTE_CHECK(per_block_scale.scalar_type() == at::kFloat, "Per-block scale must be float32."); + NVTE_CHECK(target_scale.scalar_type() == at::kByte, "Target scale must be uint8 (E4M3)."); + NVTE_CHECK(target_amax.scalar_type() == at::kFloat, "Target amax must be float32."); + NVTE_CHECK(global_amax.numel() == 1, "Global amax must be a single element tensor."); + NVTE_CHECK(target_amax.numel() == 1, "Target amax must be a single element tensor."); + + auto block_amax_cu = makeTransformerEngineTensor(block_amax); + auto global_amax_cu = makeTransformerEngineTensor(global_amax); + auto per_block_scale_cu = makeTransformerEngineTensor(per_block_scale); + auto target_scale_cu = makeTransformerEngineTensor(target_scale); + auto target_amax_cu = makeTransformerEngineTensor(target_amax); + + nvte_nvfp4_fused_scale(block_amax_cu.data(), global_amax_cu.data(), per_block_scale_cu.data(), + target_scale_cu.data(), target_amax_cu.data(), tile_rows, tile_cols, + rows_padded, static_cast(block_len), stream); + } +} + +void nvfp4_compute_global_scale(at::Tensor global_amax, at::Tensor global_scale) { + init_extension(); + + // global_amax and global_scale: [num_params], float32 + NVTE_CHECK(global_amax.scalar_type() == at::kFloat, "Global amax must be float32."); + NVTE_CHECK(global_scale.scalar_type() == at::kFloat, "Global scale must be float32."); + + auto global_amax_cu = makeTransformerEngineTensor(global_amax); + auto global_scale_cu = makeTransformerEngineTensor(global_scale); + + nvte_nvfp4_compute_global_scale(global_amax_cu.data(), global_scale_cu.data(), + at::cuda::getCurrentCUDAStream()); +} + at::Tensor swap_first_dims(at::Tensor tensor, std::optional out) { init_extension(); // Make sure input is contiguous - const auto &input = tensor.contiguous(); + const auto& input = tensor.contiguous(); // Allocate output tensor if needed if (!out) { @@ -77,5 +286,70 @@ at::Tensor swap_first_dims(at::Tensor tensor, std::optional out) { return std::move(*out); } +void nvfp4_2d_multi_tensor_transpose(std::vector rowwise_data_list, + std::vector columnwise_data_list, + std::vector rowwise_scale_inv_list, + std::vector columnwise_scale_inv_list, + std::vector M_list, std::vector K_list) { + init_extension(); + + const size_t num_tensors = rowwise_data_list.size(); + NVTE_CHECK(columnwise_data_list.size() == num_tensors, "Tensor list size mismatch"); + NVTE_CHECK(rowwise_scale_inv_list.size() == num_tensors, "Tensor list size mismatch"); + NVTE_CHECK(columnwise_scale_inv_list.size() == num_tensors, "Tensor list size mismatch"); + NVTE_CHECK(M_list.size() == num_tensors, "M_list size mismatch"); + NVTE_CHECK(K_list.size() == num_tensors, "K_list size mismatch"); + + if (num_tensors == 0) { + return; + } + + auto stream = at::cuda::getCurrentCUDAStream(); + + // Process each tensor - the main benefit is reduced Python overhead + // by doing the iteration in C++ rather than Python + constexpr size_t TILE_SIZE = 16; + + for (size_t i = 0; i < num_tensors; ++i) { + const auto& rowwise_data = rowwise_data_list[i]; + auto& columnwise_data = columnwise_data_list[i]; + const auto& rowwise_scale_inv = rowwise_scale_inv_list[i]; + auto& columnwise_scale_inv = columnwise_scale_inv_list[i]; + const int64_t M = M_list[i]; + const int64_t K = K_list[i]; + + // Transpose data: [M, K/2] -> [K, M/2] + const auto data_shape = getTensorShape(rowwise_data); + NVTE_CHECK(data_shape.size() == 2, "NVFP4 data must be 2D."); + const size_t M_packed = static_cast(M) / 2; + const size_t K_packed = data_shape[1]; + + auto input_cu = makeTransformerEngineTensor( + rowwise_data.data_ptr(), std::vector{static_cast(M), K_packed}, + DType::kByte); + auto output_cu = makeTransformerEngineTensor( + columnwise_data.data_ptr(), std::vector{static_cast(K), M_packed}, + DType::kByte); + nvte_nvfp4_data_transpose(input_cu.data(), output_cu.data(), stream); + + // Transpose scales + const size_t M_tiles = (static_cast(M) + TILE_SIZE - 1) / TILE_SIZE; + const size_t K_tiles = (static_cast(K) + TILE_SIZE - 1) / TILE_SIZE; + + const auto scale_in_shape = getTensorShape(rowwise_scale_inv); + const auto scale_out_shape = getTensorShape(columnwise_scale_inv); + + auto scale_input_cu = makeTransformerEngineTensor( + rowwise_scale_inv.data_ptr(), std::vector{scale_in_shape[0], scale_in_shape[1]}, + DType::kByte); + auto scale_output_cu = makeTransformerEngineTensor( + columnwise_scale_inv.data_ptr(), + std::vector{scale_out_shape[0], scale_out_shape[1]}, DType::kByte); + + nvte_nvfp4_scale_transpose(scale_input_cu.data(), scale_output_cu.data(), M_tiles, K_tiles, + stream); + } +} + } // namespace pytorch } // namespace transformer_engine diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 36bf208bcd..e7509f3994 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -325,6 +325,17 @@ def update_usage( if columnwise_usage is None: columnwise_usage = self._columnwise_data is not None + # If both rowwise and columnwise are requested, create columnwise from rowwise if needed + if rowwise_usage and columnwise_usage: + assert ( + self._rowwise_data is not None + and self._rowwise_scale_inv is not None + and self._amax_rowwise is not None + ), "Cannot update to rowwise and columnwise usage because rowwise data is None." + if self._columnwise_data is None or self._columnwise_scale_inv is None: + self._create_columnwise() + return + # Update row-scaled data if rowwise_usage: if self._rowwise_data is None: @@ -365,3 +376,51 @@ def update_usage( self._columnwise_data = None self._columnwise_scale_inv = None self._amax_columnwise = None + + def _create_columnwise(self): + """ + Update columnwise data and columnwise scale inv. Can only be used when using 2D scaling. + """ + assert ( + self._quantizer is not None and self._quantizer.with_2d_quantization + ), "Cannot create columnwise data without 2D quantization enabled." + rowwise_data = self._rowwise_data + if not rowwise_data.is_contiguous(): + rowwise_data = rowwise_data.contiguous() + # NVFP4 requires a specialized transpose that handles nibble repacking + self._columnwise_data = tex.nvfp4_data_transpose(rowwise_data, out=self._columnwise_data) + if self._columnwise_scale_inv is None: + assert self._quantizer is not None + # Use logical shape (self.size()), not packed byte shape (rowwise_data.shape) + # NVFP4 packs 2 elements per byte, so rowwise_data.shape[-1] is K/2 + logical_shape = self.size() + columnwise_scale_inv_shape = self._quantizer.get_scale_shape(logical_shape, True) + self._columnwise_scale_inv = torch.empty( + columnwise_scale_inv_shape, + dtype=self._rowwise_scale_inv.dtype, + device=self._rowwise_scale_inv.device, + ) + assert len(self._rowwise_scale_inv.shape) == 2 + assert len(self._columnwise_scale_inv.shape) == 2 + + # rowwise_scale_inv has shape [M_padded, K_tiles] where each tile's scale + # is repeated 16 times (once per row in the 16x16 tile). + # columnwise_scale_inv has shape [K_padded, M_tiles] where scales are + # repeated 16 times per tile row. + TILE_SIZE = 16 + logical_shape = self.size() + M, K = logical_shape[0], logical_shape[-1] + M_tiles = (M + TILE_SIZE - 1) // TILE_SIZE + K_tiles = (K + TILE_SIZE - 1) // TILE_SIZE + + tex.nvfp4_2d_scale_transpose( + self._rowwise_scale_inv, + self._columnwise_scale_inv, + M_tiles, + K_tiles, + ) + + # Also set columnwise amax (same as rowwise since it's just transposed data) + if self._amax_columnwise is None: + self._amax_columnwise = torch.empty_like(self._amax_rowwise) + self._amax_columnwise.copy_(self._amax_rowwise) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 05e2d22e9c..d23892af94 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Helper functions for using fp8 tensors as weights""" +"""Helper functions for using fp8/nvfp4 tensors as weights""" from typing import Optional, Union, List import torch @@ -16,10 +16,12 @@ from ..quantized_tensor import QuantizedTensor, Quantizer, QuantizedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer +from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from ..optimizers.multi_tensor_apply import multi_tensor_applier from ..utils import is_non_tn_fp8_gemm_supported +from ..constants import NVFP4_BLOCK_SCALING_SIZE def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): @@ -45,13 +47,19 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): new_raw_data.detach().copy_(old_raw_data) tensor._rowwise_data = new_raw_data del old_raw_data + elif isinstance(tensor, NVFP4Tensor): + old_rowwise = tensor._rowwise_data + assert old_rowwise.dtype == new_raw_data.dtype, "The data types of raw data don't match" + new_raw_data.detach().copy_(old_rowwise) + tensor._rowwise_data = new_raw_data + del old_rowwise elif isinstance(tensor, MXFP8Tensor): raise NotImplementedError("replace_raw_data for MXFP8Tensor is not supported yet") else: raise ValueError(f"replace_raw_data for {type(tensor)} is not supported yet") -def cast_master_weights_to_fp8( +def quantize_master_weights( model_weights, master_weights, start_offsets, @@ -59,15 +67,15 @@ def cast_master_weights_to_fp8( fsdp_shard_model_weights=None, manual_post_all_gather_processing=False, ): - r"""Helper function to cast master weights to FP8 primary weights. + r"""Helper function to cast master weights to quantized (FP8/NVFP4) primary weights. This is intended for use with ZeRO/FSDP. Each rank has a shard of the master weights (possibly empty) and a full copy of the model - weights. + weights. Supports FP8 (delayed, current, blockwise, MXFP8) and NVFP4 quantization. Parameters ---------- - model_weights : list of FP8 weights. + model_weights : list of quantized weights (FP8 or NVFP4). master_weights : list of master weights. Typically they are FP32 weights. start_offsets : list of integers, the starting index of the master weight in the model weight. master_weight may be smaller than model_weight because it could be distributed @@ -90,6 +98,7 @@ def cast_master_weights_to_fp8( current_scaling_params = [] blockwise_scaling_params = [] mxfp8_scaling_params = [] + nvfp4_params = [] if fsdp_shard_model_weights is None: use_fsdp_shard_model_weights = False @@ -97,6 +106,46 @@ def cast_master_weights_to_fp8( else: use_fsdp_shard_model_weights = True + # Batch convert master_weights to model dtype for NVFP4 (single kernel instead of N kernels) + # Check if there are any NVFP4 weights + has_nvfp4 = any( + isinstance(w._get_quantizer(), NVFP4Quantizer) + for w in model_weights + if hasattr(w, "_get_quantizer") + ) + if has_nvfp4 and len(model_weights) > 0: + # Find target dtype from first NVFP4 weight + target_dtype = None + for w in model_weights: + if hasattr(w, "_get_quantizer") and isinstance(w._get_quantizer(), NVFP4Quantizer): + target_dtype = w.dtype + break + + if target_dtype is not None: + # Collect non-None master_weights and their indices + non_none_indices = [] + non_none_weights = [] + sizes = [] + for i, mw in enumerate(master_weights): + if mw is not None: + non_none_indices.append(i) + non_none_weights.append(mw.view(-1)) + sizes.append(mw.numel()) + + if len(non_none_weights) > 0 and non_none_weights[0].dtype != target_dtype: + # Concatenate, convert once, then split + concatenated = torch.cat(non_none_weights) + converted = concatenated.to(target_dtype) + split_weights = torch.split(converted, sizes) + + # Rebuild master_weights list with converted tensors + converted_master_weights = list(master_weights) + for idx, split_w, orig_mw in zip( + non_none_indices, split_weights, [master_weights[i] for i in non_none_indices] + ): + converted_master_weights[idx] = split_w.view(orig_mw.shape) + master_weights = converted_master_weights + for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights ): @@ -115,34 +164,42 @@ def cast_master_weights_to_fp8( if hasattr(model_weight, "clear_high_precision_init_val"): model_weight.clear_high_precision_init_val() - if master_weight is not None: - # When not using fp8_primary_weights, the master_weight (fp32) is first cast to - # bf16/fp16, and then cast to fp8 during forward. Although it's not necessary when - # fp8_primary_weights is enabled, we still keep this logic to keep numerical - # consistency. So here we cast the master_weight to model_weight.dtype. - master_weight = master_weight.to(model_weight.dtype) - quantizer = model_weight._get_quantizer() - if isinstance(quantizer, Float8Quantizer): - delayed_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - elif isinstance(quantizer, Float8CurrentScalingQuantizer): - current_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - elif isinstance(quantizer, Float8BlockQuantizer): - blockwise_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - elif isinstance(quantizer, MXFP8Quantizer): - mxfp8_scaling_params.append( + + if isinstance(quantizer, NVFP4Quantizer): + # NVFP4: master_weight dtype conversion already done above + nvfp4_params.append( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) else: - raise ValueError( - f"cast_master_weights_to_fp8 for {type(quantizer)} is not supported yet" - ) + # FP8: convert master_weight to model dtype + if master_weight is not None: + # When not using fp8_primary_weights, the master_weight (fp32) is first cast to + # bf16/fp16, and then cast to fp8 during forward. Although it's not necessary when + # fp8_primary_weights is enabled, we still keep this logic to keep numerical + # consistency. So here we cast the master_weight to model_weight.dtype. + master_weight = master_weight.to(model_weight.dtype) + + if isinstance(quantizer, Float8Quantizer): + delayed_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, Float8CurrentScalingQuantizer): + current_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, Float8BlockQuantizer): + blockwise_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, MXFP8Quantizer): + mxfp8_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + else: + raise ValueError( + f"quantize_master_weights for {type(quantizer)} is not supported yet" + ) extra_args = [group, use_fsdp_shard_model_weights, manual_post_all_gather_processing] if len(delayed_scaling_params) > 0: @@ -153,6 +210,32 @@ def cast_master_weights_to_fp8( _cast_master_weights_to_fp8_blockwise_scaling(blockwise_scaling_params, *extra_args) if len(mxfp8_scaling_params) > 0: _cast_master_weights_to_fp8_mxfp8_scaling(mxfp8_scaling_params, *extra_args) + if len(nvfp4_params) > 0: + _cast_master_weights_to_nvfp4_2d(nvfp4_params, *extra_args) + + +def cast_master_weights_to_fp8( + model_weights, + master_weights, + start_offsets, + group, + fsdp_shard_model_weights=None, + manual_post_all_gather_processing=False, +): + r"""Helper function to cast master weights to FP8 primary weights. + + .. deprecated:: + Use :func:`quantize_master_weights` instead. + + """ + quantize_master_weights( + model_weights, + master_weights, + start_offsets, + group, + fsdp_shard_model_weights, + manual_post_all_gather_processing, + ) def _cast_master_weights_to_fp8_delayed_scaling( @@ -474,6 +557,216 @@ def _cast_master_weights_to_fp8_blockwise_scaling( ) +def _cast_master_weights_to_nvfp4_2d( + params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False +): + r"""Helper function to cast master weights to NVFP4 2D quantized weights. + + Parameters + ---------- + params : List of tuple, each tuple contains a model weight, a master weight, and an offset + indicating the starting index of the master weight in the model weight. + group : The distributed group to do amax reduction. Typically it's the data parallel + group. + use_fsdp_shard_model_weights : bool, if True, it means that the model weights are sharded. + """ + + device = params[0][0].device + block_len = NVFP4_BLOCK_SCALING_SIZE + + cu_amax_sizes = [0] + tile_shapes: List[tuple[int, int]] = [] + tile_widths: List[int] = [] + scale_targets: List[torch.Tensor] = [] + amax_targets: List[Optional[torch.Tensor]] = [] + for model_weight, _, _, _ in params: + quantizer = model_weight._get_quantizer() + assert isinstance(quantizer, NVFP4Quantizer) + assert quantizer.with_2d_quantization, "NVFP4 2D quantization must be enabled." + assert len(model_weight.shape) == 2 + h, w = model_weight.shape + tile_h = (h + block_len - 1) // block_len + tile_w = (w + block_len - 1) // block_len + tile_shapes.append((tile_h, tile_w)) + tile_widths.append(tile_w) + scale_targets.append(model_weight._rowwise_scale_inv) + amax_targets.append(model_weight._amax_rowwise) + num_amaxes = tile_h * tile_w + cu_amax_sizes.append(cu_amax_sizes[-1] + num_amaxes) + + packed_amaxes = torch.zeros(cu_amax_sizes[-1], dtype=torch.float32, device=device) + packed_scales = torch.zeros(cu_amax_sizes[-1], dtype=torch.float32, device=device) + + amaxes: List[torch.Tensor] = [] + scales: List[torch.Tensor] = [] + global_amaxes = torch.zeros(len(params), dtype=torch.float32, device=device) + global_amax_views: List[torch.Tensor] = [global_amaxes[i : i + 1] for i in range(len(params))] + + # Collect tensors for batched multi-tensor amax computation + master_weight_list: List[torch.Tensor] = [] + partial_amax_list: List[torch.Tensor] = [] + global_amax_list: List[torch.Tensor] = [] + h_list: List[int] = [] + w_list: List[int] = [] + start_offset_list: List[int] = [] + + for i, (model_weight, master_weight, start_offset, _) in enumerate(params): + scale_shape = tile_shapes[i] + amax = packed_amaxes[cu_amax_sizes[i] : cu_amax_sizes[i + 1]].reshape(scale_shape) + scale = packed_scales[cu_amax_sizes[i] : cu_amax_sizes[i + 1]].reshape(scale_shape) + global_amax_view = global_amax_views[i] + + assert model_weight._rowwise_scale_inv is not None + + amaxes.append(amax) + scales.append(scale) + + if master_weight is not None and master_weight.numel() > 0: + assert len(model_weight.shape) == 2 + h, w = model_weight.shape + # Collect for batched processing + master_weight_list.append(master_weight) + partial_amax_list.append(amax) + global_amax_list.append(global_amax_view) + h_list.append(h) + w_list.append(w) + start_offset_list.append(start_offset) + + # Batched multi-tensor call for partial and global amax computation + if master_weight_list: + tex.nvfp4_multi_tensor_compute_partial_amax( + master_weight_list, + partial_amax_list, + global_amax_list, + h_list, + w_list, + start_offset_list, + block_len, + ) + + if packed_amaxes.numel() > 0: + torch.distributed.all_reduce(packed_amaxes, op=torch.distributed.ReduceOp.MAX, group=group) + + if global_amaxes.numel() > 0: + torch.distributed.all_reduce(global_amaxes, op=torch.distributed.ReduceOp.MAX, group=group) + + # Use GPU kernel to compute global encode scales from global amaxes + # This replaces multiple Python tensor operations with a single kernel + global_scale_tensor = torch.empty_like(global_amaxes) + + tex.nvfp4_compute_global_scale(global_amaxes, global_scale_tensor) + global_scale_views = [global_scale_tensor[i : i + 1] for i in range(len(params))] + + # Collect tensors for batched fused scale kernel + fused_scale_block_amax_list: List[torch.Tensor] = [] + fused_scale_global_amax_list: List[torch.Tensor] = [] + fused_scale_per_block_scale_list: List[torch.Tensor] = [] + fused_scale_target_scale_list: List[torch.Tensor] = [] + fused_scale_target_amax_list: List[torch.Tensor] = [] + fused_scale_tile_rows_list: List[int] = [] + fused_scale_tile_cols_list: List[int] = [] + fused_scale_rows_padded_list: List[int] = [] + + # Collect tensors for batched partial cast kernel + partial_cast_inp_list: List[torch.Tensor] = [] + partial_cast_out_list: List[torch.Tensor] = [] + partial_cast_scale_list: List[torch.Tensor] = [] + partial_cast_global_scale_list: List[torch.Tensor] = [] + partial_cast_h_list: List[int] = [] + partial_cast_w_list: List[int] = [] + partial_cast_start_offset_list: List[int] = [] + + # First pass: collect all tensors and update usage + zipped_meta = zip( + tile_shapes, + tile_widths, + scale_targets, + amax_targets, + params, + amaxes, + scales, + global_scale_views, + ) + for idx, ( + tile_shape, + tile_col_cnt, + target_scale, + target_amax, + (model_weight, master_weight, start_offset, model_weight_fragment), + block_amax, + per_block_decode_scale, + global_scale, + ) in enumerate(zipped_meta): + + if not manual_post_all_gather_processing: + # Reset transpose cache for all model weights. + # We cannot create transpose cache here because users (like megatron) may want to + # overlap the all-gather of model weights and forward process, so the model weight is + # not updated currently. + model_weight.update_usage(rowwise_usage=True, columnwise_usage=False) + + tile_rows = tile_shape[0] + rows_padded = target_scale.shape[0] + global_amax_view = global_amaxes[idx : idx + 1] + + # Collect for fused scale kernel (only if target_amax is not None) + if target_amax is not None: + fused_scale_block_amax_list.append(block_amax) + fused_scale_global_amax_list.append(global_amax_view) + fused_scale_per_block_scale_list.append(per_block_decode_scale) + fused_scale_target_scale_list.append(target_scale) + fused_scale_target_amax_list.append(target_amax) + fused_scale_tile_rows_list.append(tile_rows) + fused_scale_tile_cols_list.append(tile_col_cnt) + fused_scale_rows_padded_list.append(rows_padded) + + # Collect for partial cast kernel (only for layers owned by this rank) + if master_weight is not None and master_weight.numel() > 0: + end_offset = start_offset + master_weight.numel() + if not use_fsdp_shard_model_weights: + rowwise_bytes = model_weight._rowwise_data.view(-1) + byte_start = start_offset // 2 + byte_end = (end_offset + 1) // 2 + model_weight_fragment = rowwise_bytes[byte_start:byte_end] + assert len(model_weight.shape) == 2 + h, w = model_weight.shape + + partial_cast_inp_list.append(master_weight) + partial_cast_out_list.append(model_weight_fragment) + partial_cast_scale_list.append(per_block_decode_scale) + partial_cast_global_scale_list.append(global_scale) + partial_cast_h_list.append(h) + partial_cast_w_list.append(w) + partial_cast_start_offset_list.append(start_offset) + + # Batched multi-tensor call for fused scale + if fused_scale_block_amax_list: + tex.nvfp4_multi_tensor_fused_scale( + fused_scale_block_amax_list, + fused_scale_global_amax_list, + fused_scale_per_block_scale_list, + fused_scale_target_scale_list, + fused_scale_target_amax_list, + fused_scale_tile_rows_list, + fused_scale_tile_cols_list, + fused_scale_rows_padded_list, + block_len, + ) + + # Batched multi-tensor call for partial cast + if partial_cast_inp_list: + tex.nvfp4_multi_tensor_2d_partial_cast( + partial_cast_inp_list, + partial_cast_out_list, + partial_cast_scale_list, + partial_cast_global_scale_list, + partial_cast_h_list, + partial_cast_w_list, + partial_cast_start_offset_list, + block_len, + ) + + def _cast_master_weights_to_fp8_mxfp8_scaling( params, group, use_fsdp_shard_model_weights=False, manual_post_all_gather_processing=False ): # pylint: disable=unused-argument @@ -605,9 +898,15 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten - Float8Tensor: may need to create a transposed view to match backend GEMM. - Float8BlockwiseQTensor: create column-wise storage. - Plain pytorch tensor: noop. + + For NVFP4 tensors, uses batched multi-tensor processing to reduce CPU overhead. """ if not isinstance(model_weights, list): model_weights = [model_weights] + + # Collect NVFP4 tensors for batched processing + nvfp4_tensors = [] + for model_weight in model_weights: if isinstance(model_weight, Float8Tensor): # Delayed scaling and per-tensor current scaling: if backend does not support @@ -617,12 +916,90 @@ def post_all_gather_processing(model_weights: Union[torch.Tensor, List[torch.Ten elif isinstance(model_weight, Float8BlockwiseQTensor): # Blockwise scaling: create column-wise storage. model_weight._create_columnwise() + elif isinstance(model_weight, NVFP4Tensor): + # Collect for batched processing + nvfp4_tensors.append(model_weight) elif isinstance(model_weight, MXFP8Tensor): # MXFP8 scaling: no need to do anything. pass elif isinstance(model_weight, QuantizedTensor): raise ValueError(f"post_processing for {type(model_weight)} is not supported") + # Batch process all NVFP4 tensors with multi-tensor approach + if nvfp4_tensors: + _nvfp4_2d_multi_tensor_transpose(nvfp4_tensors) + + +def _nvfp4_2d_multi_tensor_transpose(nvfp4_tensors: List[NVFP4Tensor]): + """ + Batched columnwise creation for multiple NVFP4 tensors. + Reduces CPU overhead by collecting all tensor metadata and dispatching to C++. + """ + # Prepare tensor lists for batched C++ call + rowwise_data_list = [] + columnwise_data_list = [] + rowwise_scale_inv_list = [] + columnwise_scale_inv_list = [] + M_list = [] + K_list = [] + + for tensor in nvfp4_tensors: + rowwise_data = tensor._rowwise_data + if not rowwise_data.is_contiguous(): + rowwise_data = rowwise_data.contiguous() + tensor._rowwise_data = rowwise_data + + logical_shape = tensor.size() + M, K = logical_shape[0], logical_shape[-1] + + # Allocate columnwise_data if needed + if tensor._columnwise_data is None: + # Output shape: [K, M/2] packed bytes + columnwise_data = torch.empty( + (K, M // 2), + dtype=torch.uint8, + device=rowwise_data.device, + ) + tensor._columnwise_data = columnwise_data + else: + columnwise_data = tensor._columnwise_data + + # Allocate columnwise_scale_inv if needed + if tensor._columnwise_scale_inv is None: + assert tensor._quantizer is not None + columnwise_scale_inv_shape = tensor._quantizer.get_scale_shape(logical_shape, True) + columnwise_scale_inv = torch.empty( + columnwise_scale_inv_shape, + dtype=tensor._rowwise_scale_inv.dtype, + device=tensor._rowwise_scale_inv.device, + ) + tensor._columnwise_scale_inv = columnwise_scale_inv + else: + columnwise_scale_inv = tensor._columnwise_scale_inv + + rowwise_data_list.append(rowwise_data) + columnwise_data_list.append(columnwise_data) + rowwise_scale_inv_list.append(tensor._rowwise_scale_inv) + columnwise_scale_inv_list.append(columnwise_scale_inv) + M_list.append(M) + K_list.append(K) + + # Copy amax if needed + if tensor._amax_columnwise is None and tensor._amax_rowwise is not None: + tensor._amax_columnwise = tensor._amax_rowwise.clone() + elif tensor._amax_rowwise is not None: + tensor._amax_columnwise.copy_(tensor._amax_rowwise) + + # Dispatch to C++ multi-tensor kernel + tex.nvfp4_2d_multi_tensor_transpose( + rowwise_data_list, + columnwise_data_list, + rowwise_scale_inv_list, + columnwise_scale_inv_list, + M_list, + K_list, + ) + def is_custom(x: Optional[Union[Quantizer, QuantizedTensorStorage]] = None) -> bool: """Check if an object is custom. From bf3201a15e3e5d722e80029ef329fa4b9b706bf8 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 3 Mar 2026 22:51:23 -0800 Subject: [PATCH 241/521] [PyTorch] Support single parameter for `GroupedLinear` (#2731) * Support single parameter for GroupedLinear Signed-off-by: Kirthi Shankar Sivamani * greptile suggestions Signed-off-by: Kirthi Shankar Sivamani * Faster python class creation in c++ Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/pytorch/module/grouped_linear.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/pytorch/tensor/grouped_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kirthi Shankar Sivamani 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> --- qa/L0_pytorch_unittest/test.sh | 1 + .../test_nvfp4_group_quantize_graph_safe.py | 2 +- tests/pytorch/test_grouped_tensor.py | 100 ++++---- tests/pytorch/test_sanity.py | 130 ++-------- .../pytorch/csrc/extensions/pybind.cpp | 15 +- transformer_engine/pytorch/csrc/pybind.h | 1 + transformer_engine/pytorch/csrc/quantizer.cpp | 231 ++++++++++++------ .../pytorch/module/grouped_linear.py | 58 +++-- transformer_engine/pytorch/tensor/__init__.py | 6 + .../pytorch/tensor/grouped_tensor.py | 205 ++++++++++++++++ .../pytorch/tensor/storage/__init__.py | 2 +- ...ed_tensor.py => grouped_tensor_storage.py} | 204 +++++++--------- 12 files changed, 564 insertions(+), 391 deletions(-) create mode 100644 transformer_engine/pytorch/tensor/grouped_tensor.py rename transformer_engine/pytorch/tensor/storage/{grouped_tensor.py => grouped_tensor_storage.py} (87%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index e0ad09200d..f2b0b07fed 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -37,6 +37,7 @@ 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" 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" 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 1e62f91eb8..8d81d578a7 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -10,7 +10,7 @@ from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling -from transformer_engine.pytorch.tensor.storage.grouped_tensor import GroupedTensor +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor import pytest import torch diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index ad08c0474d..9dd965fa94 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -8,7 +8,7 @@ import pytest import torch import transformer_engine.pytorch as te -from transformer_engine.pytorch.tensor.storage.grouped_tensor import GroupedTensor +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch import ( Quantizer, Float8Quantizer, @@ -125,7 +125,7 @@ def test_basic_construction_all_same_shape(self) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=None, device="cuda", dtype=torch.float32, @@ -147,7 +147,7 @@ def test_basic_construction_varying_first_dim(self) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=None, device="cuda", dtype=torch.float32, @@ -170,14 +170,18 @@ def test_split_into_quantized_tensors_no_quantization(self) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=None, device="cuda", dtype=torch.float32, ) - # Get the original data pointer - original_data_ptr = grouped_tensor.data.data_ptr() + # GroupedTensor is a wrapper; use backing storage buffer pointer. + storage = grouped_tensor.rowwise_data + if storage is None: + storage = grouped_tensor.columnwise_data + assert storage is not None + original_data_ptr = storage.data_ptr() # Split into tensors tensors = grouped_tensor.split_into_quantized_tensors() @@ -207,13 +211,18 @@ def test_split_into_quantized_tensors_quantized(self, quantization: str) -> None grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=quantizer, device="cuda", + dtype=torch.float32, ) - # Get the original data pointer - original_data_ptr = grouped_tensor.data.data_ptr() + # GroupedTensor is a wrapper; use backing storage buffer pointer. + storage = grouped_tensor.rowwise_data + if storage is None: + storage = grouped_tensor.columnwise_data + assert storage is not None + original_data_ptr = storage.data_ptr() # Split into tensors tensors = grouped_tensor.split_into_quantized_tensors() @@ -236,13 +245,17 @@ def test_split_varying_shapes(self) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=None, device="cuda", dtype=torch.float32, ) - original_data_ptr = grouped_tensor.data.data_ptr() + storage = grouped_tensor.rowwise_data + if storage is None: + storage = grouped_tensor.columnwise_data + assert storage is not None + original_data_ptr = storage.data_ptr() tensors = grouped_tensor.split_into_quantized_tensors() assert len(tensors) == num_tensors @@ -264,13 +277,18 @@ def test_quantize_inplace(self, quantization: str) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=quantizer, device="cuda", + dtype=torch.float32, ) # Get original data pointers before quantization - original_data_ptr = grouped_tensor.data.data_ptr() + storage = grouped_tensor.rowwise_data + if storage is None: + storage = grouped_tensor.columnwise_data + assert storage is not None + original_data_ptr = storage.data_ptr() original_scale_inv_ptr = grouped_tensor.scale_inv.data_ptr() original_scale_ptr = ( grouped_tensor.scale.data_ptr() if grouped_tensor.scale is not None else None @@ -283,7 +301,7 @@ def test_quantize_inplace(self, quantization: str) -> None: quantized_tensors = grouped_tensor.quantize(input_tensors) # Verify data pointers haven't changed (in-place operation) - assert grouped_tensor.data.data_ptr() == original_data_ptr + assert storage.data_ptr() == original_data_ptr assert grouped_tensor.scale_inv.data_ptr() == original_scale_inv_ptr if original_scale_ptr is not None: assert grouped_tensor.scale.data_ptr() == original_scale_ptr @@ -304,13 +322,18 @@ def test_quantize_varying_shapes(self, quantization: str) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=quantizer, device="cuda", + dtype=torch.float32, ) # Get original data pointers - original_data_ptr = grouped_tensor.data.data_ptr() + storage = grouped_tensor.rowwise_data + if storage is None: + storage = grouped_tensor.columnwise_data + assert storage is not None + original_data_ptr = storage.data_ptr() # Create input tensors with varying shapes input_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in shape] @@ -319,7 +342,7 @@ def test_quantize_varying_shapes(self, quantization: str) -> None: quantized_tensors = grouped_tensor.quantize(input_tensors) # Verify data pointer hasn't changed - assert grouped_tensor.data.data_ptr() == original_data_ptr + assert storage.data_ptr() == original_data_ptr # Verify each tensor points to correct location cumulative_numel = 0 @@ -329,38 +352,6 @@ def test_quantize_varying_shapes(self, quantization: str) -> None: assert rowwise_data.data_ptr() == original_data_ptr + expected_offset cumulative_numel += tensor_shape[0] * tensor_shape[1] - @pytest.mark.parametrize("quantization", _quantization_params) - def test_static_quantize_method(self, quantization: str) -> None: - """Test the static quantize method""" - num_tensors = 3 - shape = [(512, 512) for _ in range(num_tensors)] - quantizer = make_quantizer(quantization, num_tensors, shape) - - # Create input tensors - input_tensors = [torch.randn(s, dtype=torch.float32, device="cuda") for s in shape] - - # Use static quantize method - grouped_tensor = GroupedTensor.create_and_quantize( - tensors=input_tensors, - quantizer=quantizer, - device="cuda", - ) - - # Verify the grouped tensor was created correctly - assert grouped_tensor.num_tensors == num_tensors - assert grouped_tensor.has_data() - - # Verify quantized_tensors were created and point to same storage - assert grouped_tensor.quantized_tensors is not None - assert len(grouped_tensor.quantized_tensors) == num_tensors - - original_data_ptr = grouped_tensor.data.data_ptr() - for i, qtensor in enumerate(grouped_tensor.quantized_tensors): - rowwise_data = _get_rowwise_data_tensor(qtensor, quantization) - numel = shape[i][0] * shape[i][1] - expected_offset = _rowwise_offset_bytes(i * numel, quantization) - assert rowwise_data.data_ptr() == original_data_ptr + expected_offset - @pytest.mark.parametrize( "shape", [[(256, 512), (512, 512), (768, 512)], [(512, 512), (512, 512), (512, 512)]], @@ -374,9 +365,6 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]]) -> None: # Create BF16 input tensors and pack into a 2D tensor input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] - quantized_tensors = [ - MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(tensor) for tensor in input_tensors - ] grouped_input = torch.cat(input_tensors, dim=0) # Create MXFP8 output grouped tensor (rowwise only for easier validation) @@ -406,7 +394,7 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]]) -> None: expected_data = torch.cat(expected_data) expected_scale_inv = torch.cat(expected_scale_inv) - assert torch.equal(grouped_output.data, expected_data) + assert torch.equal(grouped_output.rowwise_data, expected_data) assert torch.equal(grouped_output.scale_inv, expected_scale_inv) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @@ -451,7 +439,7 @@ def test_group_quantize_cudagraph_capturable(self) -> None: torch.cuda.synchronize() expected = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) - assert torch.equal(static_output.data, expected.data) + assert torch.equal(static_output.rowwise_data, expected.rowwise_data) assert torch.equal(static_output.scale_inv, expected.scale_inv) def test_clear(self) -> None: @@ -461,7 +449,7 @@ def test_clear(self) -> None: grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_tensors, - shape=shape, + shapes=shape, quantizer=None, device="cuda", dtype=torch.float32, @@ -474,5 +462,5 @@ def test_clear(self) -> None: assert not grouped_tensor.has_data() assert grouped_tensor.num_tensors == 0 - assert grouped_tensor.data is None + assert grouped_tensor.rowwise_data is None assert grouped_tensor.logical_shape == (0, 0) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 3ef8c0983f..384b6774f6 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -138,115 +138,21 @@ def reset_global_fp8_state(): FP8GlobalStateManager.reset() -def check_grouped_tensor_pointers_helper(tensors, num_elems_in_byte=1, tensor_name="tensor"): - """ - Verify that tensors are stored in contiguous memory. - - Args: - tensors: List or iterable of tensors to check - num_elems_in_byte: Number of elements packed per byte (1 for normal, 2 for NVFP4) - tensor_name: Name to use in error messages - """ - tensor_list = list(tensors) - if len(tensor_list) < 2: - return # Nothing to check - - for i in range(1, len(tensor_list)): - prev_tensor = tensor_list[i - 1] - curr_tensor = tensor_list[i] - - # Calculate expected offset based on previous tensor size - prev_numel = prev_tensor.numel() - expected_offset = (prev_numel // num_elems_in_byte) * prev_tensor.element_size() - - # Verify current tensor's data pointer is correctly offset - expected_ptr = prev_tensor.data_ptr() + expected_offset - actual_ptr = curr_tensor.data_ptr() - - assert ( - actual_ptr == expected_ptr - ), f"{tensor_name} {i} data pointer mismatch: expected {expected_ptr}, got {actual_ptr}" - - -def check_grouped_tensor_pointers( - weights: List[torch.Tensor], fp8_recipe: Optional[recipe.Recipe] = None +def check_grouped_weight( + module: GroupedLinear, num_gemms: int, out_features: int, in_features: int ): """ - Verify that the pointers of the weights are in contiguous memory for GroupedTensor. - TODO(ksivaman): This check can be made way more efficient but for now leaving the brute force approach. + Verify GroupedLinear exposes one grouped weight parameter with shape + [num_gemms, out_features, in_features]. """ - - num_elems_in_a_data_byte = 1 if fp8_recipe is None else 2 if fp8_recipe.nvfp4() else 1 - - # Check data. - if hasattr(weights[0], "_data") and weights[0]._data is not None: - data_tensors = [w._data for w in weights] - check_grouped_tensor_pointers_helper(data_tensors, num_elems_in_byte=1, tensor_name="data") - - # Check transpose. - if hasattr(weights[0], "_transpose") and weights[0]._transpose is not None: - transpose_tensors = [w._transpose for w in weights] - check_grouped_tensor_pointers_helper( - transpose_tensors, num_elems_in_byte=1, tensor_name="transpose" - ) - - # Check scale_inv. - if hasattr(weights[0], "_scale_inv") and weights[0]._scale_inv is not None: - scale_inv_tensors = [w._scale_inv for w in weights] - check_grouped_tensor_pointers_helper( - scale_inv_tensors, num_elems_in_byte=1, tensor_name="scale_inv" - ) - - # Check rowwise scale_inv. - if hasattr(weights[0], "_rowwise_scale_inv") and weights[0]._rowwise_scale_inv is not None: - scale_inv_tensors = [w._rowwise_scale_inv for w in weights] - check_grouped_tensor_pointers_helper( - scale_inv_tensors, num_elems_in_byte=1, tensor_name="rowwise_scale_inv" - ) - - # Check columnwise scale_inv. - if ( - hasattr(weights[0], "_columnwise_scale_inv") - and weights[0]._columnwise_scale_inv is not None - ): - columnwise_scale_inv_tensors = [w._columnwise_scale_inv for w in weights] - check_grouped_tensor_pointers_helper( - columnwise_scale_inv_tensors, - num_elems_in_byte=1, - tensor_name="columnwise scale_inv", - ) - - # Check rowwise amax. - if hasattr(weights[0], "_rowwise_amax") and weights[0]._rowwise_amax is not None: - rowwise_amax_tensors = [w._rowwise_amax for w in weights] - check_grouped_tensor_pointers_helper( - rowwise_amax_tensors, num_elems_in_byte=1, tensor_name="rowwise amax" - ) - - # Check columnwise amax. - if hasattr(weights[0], "_columnwise_amax") and weights[0]._columnwise_amax is not None: - columnwise_amax_tensors = [w._columnwise_amax for w in weights] - check_grouped_tensor_pointers_helper( - columnwise_amax_tensors, num_elems_in_byte=1, tensor_name="columnwise amax" - ) - - # Check rowwise data. - if hasattr(weights[0], "_rowwise_data") and weights[0]._rowwise_data is not None: - rowwise_data_tensors = [w._rowwise_data for w in weights] - check_grouped_tensor_pointers_helper( - rowwise_data_tensors, - num_elems_in_byte=num_elems_in_a_data_byte, - tensor_name="rowwise data", - ) - - # Check columnwise data. - if hasattr(weights[0], "_columnwise_data") and weights[0]._columnwise_data is not None: - columnwise_data_tensors = [w._columnwise_data for w in weights] - check_grouped_tensor_pointers_helper( - columnwise_data_tensors, - num_elems_in_byte=num_elems_in_a_data_byte, - tensor_name="columnwise data", - ) + weight_params = [(name, p) for name, p in module.named_parameters() if "weight" in name] + assert len(weight_params) == 1, f"Expected 1 grouped weight parameter, got {len(weight_params)}" + name, weight = weight_params[0] + assert name == "weight", f"Expected grouped parameter name 'weight', got {name}" + assert tuple(weight.shape) == (num_gemms, out_features, in_features), ( + "Grouped weight has unexpected shape. " + f"Expected {(num_gemms, out_features, in_features)}, got {tuple(weight.shape)}" + ) def _test_sanity_e2e_amp(block, dtype, config, fp8_recipe, skip_wgrad): @@ -603,9 +509,6 @@ def test_sanity_grouped_linear( bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) - if single_param: - os.environ["NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS"] = "1" - if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") @@ -620,13 +523,13 @@ def test_sanity_grouped_linear( ffn_hidden_size, bias=use_bias, params_dtype=dtype, + single_grouped_parameter=single_param, ).cuda() - # Verify that weights are stored in contiguous GroupedTensor storage. - weights = [getattr(te_grouped_linear, f"weight{i}") for i in range(num_gemms)] + # Verify grouped linear exposes a single grouped weight parameter. if fp8_recipe is None or not (fp8_recipe.delayed() or fp8_recipe.float8_current_scaling()): if single_param: - check_grouped_tensor_pointers(weights, fp8_recipe) + check_grouped_weight(te_grouped_linear, num_gemms, ffn_hidden_size, config.hidden_size) inp_hidden_states = torch.randn( num_tokens, config.hidden_size, dtype=dtype, requires_grad=True @@ -645,9 +548,6 @@ def test_sanity_grouped_linear( loss.backward() assert out.shape == (num_tokens, ffn_hidden_size) - if single_param: - del os.environ["NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS"] - @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes) diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index b1d60cc3da..8302a13010 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -35,8 +35,9 @@ PyTypeObject *Float8BlockwiseQuantizerClass = nullptr; PyTypeObject *NVFP4TensorPythonClass = nullptr; PyTypeObject *NVFP4TensorStoragePythonClass = nullptr; PyTypeObject *NVFP4QuantizerClass = nullptr; -std::once_flag extension_init_flag; +PyTypeObject *GroupedTensorPythonClass = nullptr; PyTypeObject *GroupedTensorStoragePythonClass = nullptr; +std::once_flag extension_init_flag; void init_float8_extension() { auto fp8_module = py::module_::import("transformer_engine.pytorch.tensor.float8_tensor"); @@ -103,11 +104,17 @@ void init_nvfp4_extensions() { } void init_grouped_tensor_extension() { - if (GroupedTensorStoragePythonClass) return; + if (GroupedTensorPythonClass && GroupedTensorStoragePythonClass) return; auto grouped_tensor_module = - py::module_::import("transformer_engine.pytorch.tensor.storage.grouped_tensor"); - GroupedTensorStoragePythonClass = reinterpret_cast( + py::module_::import("transformer_engine.pytorch.tensor.grouped_tensor"); + GroupedTensorPythonClass = reinterpret_cast( PyObject_GetAttrString(grouped_tensor_module.ptr(), "GroupedTensor")); + auto grouped_tensor_storage_module = + py::module_::import("transformer_engine.pytorch.tensor.storage.grouped_tensor_storage"); + GroupedTensorStoragePythonClass = reinterpret_cast( + PyObject_GetAttrString(grouped_tensor_storage_module.ptr(), "GroupedTensorStorage")); + NVTE_CHECK(GroupedTensorPythonClass != nullptr, + "Internal error: could not initialize pyTorch grouped tensor extension."); NVTE_CHECK(GroupedTensorStoragePythonClass != nullptr, "Internal error: could not initialize pyTorch grouped tensor extension."); } diff --git a/transformer_engine/pytorch/csrc/pybind.h b/transformer_engine/pytorch/csrc/pybind.h index 059eb5e3fb..9e640537f9 100644 --- a/transformer_engine/pytorch/csrc/pybind.h +++ b/transformer_engine/pytorch/csrc/pybind.h @@ -43,6 +43,7 @@ extern PyTypeObject *Float8BlockwiseQuantizerClass; extern PyTypeObject *NVFP4TensorPythonClass; extern PyTypeObject *NVFP4TensorStoragePythonClass; extern PyTypeObject *NVFP4QuantizerClass; +extern PyTypeObject *GroupedTensorPythonClass; extern PyTypeObject *GroupedTensorStoragePythonClass; void init_extension(); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 0da5f69197..0135c7f01c 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -67,13 +67,13 @@ std::optional build_grouped_tensor_offsets(const size_t num_tensors, } const auto& first_dims_tensor = first_dims.value(); + NVTE_CHECK(first_dims_tensor.is_cuda(), "first_dims must be on CUDA."); NVTE_CHECK(first_dims_tensor.scalar_type() == at::kLong, "first_dims must have dtype int64."); NVTE_CHECK(static_cast(first_dims_tensor.numel()) == num_tensors, "first_dims must have length ", num_tensors, "."); const int64_t logical_last_dim_i64 = static_cast(logical_last_dim); - auto scaled_first_dims = first_dims_tensor * logical_last_dim_i64; - + auto scaled_first_dims = (first_dims_tensor * logical_last_dim_i64).contiguous(); // Single kernel needed for these ops. auto cumsum = at::cumsum(scaled_first_dims, 0); auto zero = at::zeros({1}, cumsum.options()); @@ -88,6 +88,11 @@ py::object maybe_tensor_to_py(const std::optional& tensor) { return tensor ? py::cast(*tensor) : py::none(); } +py::handle grouped_tensor_python_class(const bool internal) { + PyTypeObject* cls = internal ? GroupedTensorStoragePythonClass : GroupedTensorPythonClass; + return py::handle(reinterpret_cast(cls)); +} + } // namespace constexpr size_t NVFP4_BLOCK_SIZE = 16; @@ -172,18 +177,30 @@ std::pair NoneQuantizer::create_grouped_tensor getTensorShape(*tensor_offsets)); } - py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); - py::object out_py = GroupedTensorClass( - "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), - "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), - "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), "scale_inv"_a = py::none(), - "columnwise_scale_inv"_a = py::none(), "amax"_a = py::none(), - "columnwise_amax"_a = py::none(), "scale"_a = py::none(), - "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), - "last_dims"_a = py::none(), - "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), - "logical_shape"_a = std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["num_tensors"] = py::cast(num_tensors); + kwargs["quantizer"] = quantizer; + kwargs["data"] = maybe_tensor_to_py(rowwise_data); + kwargs["columnwise_data"] = maybe_tensor_to_py(columnwise_data); + kwargs["scale_inv"] = py::none(); + kwargs["columnwise_scale_inv"] = py::none(); + kwargs["amax"] = py::none(); + kwargs["columnwise_amax"] = py::none(); + kwargs["scale"] = py::none(); + kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); + kwargs["last_dims"] = py::none(); + kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create GroupedTensor instance"); + py::object out_py = py::reinterpret_steal(result); return {std::move(out_cpp), std::move(out_py)}; } @@ -366,19 +383,30 @@ std::pair Float8Quantizer::create_grouped_tens getTensorShape(*tensor_offsets)); } - py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); - py::object out_py = GroupedTensorClass( - "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), - "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), - "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), - "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), - "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = amax, - "columnwise_amax"_a = py::none(), "scale"_a = py::none(), - "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), - "last_dims"_a = py::none(), - "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), - "logical_shape"_a = std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["num_tensors"] = py::cast(num_tensors); + kwargs["quantizer"] = quantizer; + kwargs["data"] = maybe_tensor_to_py(rowwise_data); + kwargs["columnwise_data"] = maybe_tensor_to_py(columnwise_data); + kwargs["scale_inv"] = maybe_tensor_to_py(rowwise_scale_inv); + kwargs["columnwise_scale_inv"] = maybe_tensor_to_py(columnwise_scale_inv); + kwargs["amax"] = amax; + kwargs["columnwise_amax"] = py::none(); + kwargs["scale"] = py::none(); + kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); + kwargs["last_dims"] = py::none(); + kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create GroupedTensor instance"); + py::object out_py = py::reinterpret_steal(result); return {std::move(out_cpp), std::move(out_py)}; } @@ -673,19 +701,30 @@ std::pair Float8CurrentScalingQuantizer::creat getTensorShape(*tensor_offsets)); } - py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); - py::object out_py = GroupedTensorClass( - "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), - "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), - "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), - "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), - "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = amax, - "columnwise_amax"_a = py::none(), "scale"_a = scale, - "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), - "last_dims"_a = py::none(), - "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), - "logical_shape"_a = std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["num_tensors"] = py::cast(num_tensors); + kwargs["quantizer"] = quantizer; + kwargs["data"] = maybe_tensor_to_py(rowwise_data); + kwargs["columnwise_data"] = maybe_tensor_to_py(columnwise_data); + kwargs["scale_inv"] = maybe_tensor_to_py(rowwise_scale_inv); + kwargs["columnwise_scale_inv"] = maybe_tensor_to_py(columnwise_scale_inv); + kwargs["amax"] = amax; + kwargs["columnwise_amax"] = py::none(); + kwargs["scale"] = scale; + kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); + kwargs["last_dims"] = py::none(); + kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create GroupedTensor instance"); + py::object out_py = py::reinterpret_steal(result); return {std::move(out_cpp), std::move(out_py)}; } @@ -1020,19 +1059,30 @@ std::pair Float8BlockQuantizer::create_grouped getTensorShape(*tensor_offsets)); } - py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); - py::object out_py = GroupedTensorClass( - "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), - "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), - "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), - "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), - "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = py::none(), - "columnwise_amax"_a = py::none(), "scale"_a = py::none(), - "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), - "last_dims"_a = py::none(), - "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), - "logical_shape"_a = std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["num_tensors"] = py::cast(num_tensors); + kwargs["quantizer"] = quantizer; + kwargs["data"] = maybe_tensor_to_py(rowwise_data); + kwargs["columnwise_data"] = maybe_tensor_to_py(columnwise_data); + kwargs["scale_inv"] = maybe_tensor_to_py(rowwise_scale_inv); + kwargs["columnwise_scale_inv"] = maybe_tensor_to_py(columnwise_scale_inv); + kwargs["amax"] = py::none(); + kwargs["columnwise_amax"] = py::none(); + kwargs["scale"] = py::none(); + kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); + kwargs["last_dims"] = py::none(); + kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create GroupedTensor instance"); + py::object out_py = py::reinterpret_steal(result); return {std::move(out_cpp), std::move(out_py)}; } @@ -1425,19 +1475,30 @@ std::pair MXFP8Quantizer::create_grouped_tenso out_cpp.set_with_gemm_swizzled_scales(this->optimize_for_gemm); - py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); - py::object out_py = GroupedTensorClass( - "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), - "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), - "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), - "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), - "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), "amax"_a = py::none(), - "columnwise_amax"_a = py::none(), "scale"_a = py::none(), - "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), - "last_dims"_a = py::none(), - "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), - "logical_shape"_a = std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["num_tensors"] = py::cast(num_tensors); + kwargs["quantizer"] = quantizer; + kwargs["data"] = maybe_tensor_to_py(rowwise_data); + kwargs["columnwise_data"] = maybe_tensor_to_py(columnwise_data); + kwargs["scale_inv"] = maybe_tensor_to_py(rowwise_scale_inv); + kwargs["columnwise_scale_inv"] = maybe_tensor_to_py(columnwise_scale_inv); + kwargs["amax"] = py::none(); + kwargs["columnwise_amax"] = py::none(); + kwargs["scale"] = py::none(); + kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); + kwargs["last_dims"] = py::none(); + kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create GroupedTensor instance"); + py::object out_py = py::reinterpret_steal(result); return {std::move(out_cpp), std::move(out_py)}; } @@ -1842,20 +1903,30 @@ std::pair NVFP4Quantizer::create_grouped_tenso out_cpp.set_with_gemm_swizzled_scales(this->optimize_for_gemm); - py::handle GroupedTensorClass(reinterpret_cast(GroupedTensorStoragePythonClass)); - py::object out_py = GroupedTensorClass( - "num_tensors"_a = num_tensors, "quantizer"_a = std::move(quantizer), - "dtype"_a = GetATenDType(dtype), "data"_a = maybe_tensor_to_py(rowwise_data), - "columnwise_data"_a = maybe_tensor_to_py(columnwise_data), - "scale_inv"_a = maybe_tensor_to_py(rowwise_scale_inv), - "columnwise_scale_inv"_a = maybe_tensor_to_py(columnwise_scale_inv), - "amax"_a = maybe_tensor_to_py(rowwise_amax), - "columnwise_amax"_a = maybe_tensor_to_py(columnwise_amax), "scale"_a = py::none(), - "first_dims"_a = first_dims.has_value() ? py::cast(*first_dims) : py::none(), - "last_dims"_a = py::none(), - "tensor_offsets"_a = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(), - "logical_shape"_a = std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); + py::dict kwargs; + py::tuple args(0); + kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), + static_cast(logical_last_dim)}); + kwargs["dtype"] = py::cast(GetATenDType(dtype)); + kwargs["num_tensors"] = py::cast(num_tensors); + kwargs["quantizer"] = quantizer; + kwargs["data"] = maybe_tensor_to_py(rowwise_data); + kwargs["columnwise_data"] = maybe_tensor_to_py(columnwise_data); + kwargs["scale_inv"] = maybe_tensor_to_py(rowwise_scale_inv); + kwargs["columnwise_scale_inv"] = maybe_tensor_to_py(columnwise_scale_inv); + kwargs["amax"] = maybe_tensor_to_py(rowwise_amax); + kwargs["columnwise_amax"] = maybe_tensor_to_py(columnwise_amax); + kwargs["scale"] = py::none(); + kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); + kwargs["last_dims"] = py::none(); + kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); + if (result == nullptr) { + PyErr_Print(); + } + NVTE_CHECK(result != nullptr, "Failed to create GroupedTensor instance"); + py::object out_py = py::reinterpret_steal(result); return {std::move(out_cpp), std::move(out_py)}; } diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index b381073d78..f3e7b57cf1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -6,7 +6,6 @@ from typing import Union, Optional, Callable, Tuple, List from itertools import chain import warnings -import os import functools import torch @@ -14,7 +13,7 @@ import transformer_engine_torch as tex from transformer_engine.common.recipe import Recipe -from transformer_engine.pytorch.tensor.storage.grouped_tensor import GroupedTensor +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from .base import ( get_dummy_wgrad, TransformerEngineBaseModule, @@ -595,6 +594,10 @@ class GroupedLinear(TransformerEngineBaseModule): cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. Cannot work with FP8 DelayedScaling recipe. + single_grouped_parameter : 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. Notes ----- @@ -625,6 +628,7 @@ def __init__( ub_name: Optional[str] = None, delay_wgrad_compute: bool = False, save_original_input: bool = False, + single_grouped_parameter: bool = False, name: Optional[str] = None, ) -> None: super().__init__(name) @@ -641,6 +645,7 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + self.single_grouped_parameter = single_grouped_parameter assert ( not ub_overlap_rs and not ub_overlap_ag ), "GroupedLinear doesn't support Userbuffer overlap." @@ -767,7 +772,7 @@ def make_grouped_weights(self, defer_init=False) -> None: # Create the weight storage. grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=self.num_gemms, - shape=[(self.out_features, self.in_features)] * self.num_gemms, + shapes=[(self.out_features, self.in_features)] * self.num_gemms, quantizer=weight_quantizers[0], dtype=self.params_dtype, device=weights[0].device, @@ -781,22 +786,27 @@ def make_grouped_weights(self, defer_init=False) -> None: else: grouped_weights.quantized_tensors[i].copy_(weights[i]) - # Re-register the grouped weights as parameters. + # Re-register as a single grouped weight parameter. + # Re-register as a single grouped weight parameter. + assert isinstance(grouped_weights, torch.Tensor) and ( + weight_quantizers[0] is None or not weight_quantizers[0].internal + ), "Found internal quantizer with `single_grouped_parameter=True`." + self.register_parameter( + "weight", + torch.nn.Parameter(grouped_weights), + init_fn=self.init_method, + get_rng_state_tracker=self.get_rng_state_tracker, + fp8_meta_index=self._offsets["weight"], + ) for i in range(self.num_gemms): - self.register_parameter( - f"weight{i}", - torch.nn.Parameter(grouped_weights.quantized_tensors[i]), - init_fn=self.init_method, - get_rng_state_tracker=self.get_rng_state_tracker, - fp8_meta_index=self._offsets["weight"] + i * self._num_fp8_tensors_per_gemm["fwd"], - ) + self.register_parameter(f"weight{i}", None) self.set_tensor_parallel_attributes(defer_init=defer_init) def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) # Grouped tensor weights is an opt-in feature. - if bool(int(os.getenv("NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS", "0"))): + if self.single_grouped_parameter: self.make_grouped_weights(defer_init=defer_init) def set_tensor_parallel_attributes(self, defer_init=False) -> None: @@ -804,13 +814,22 @@ def set_tensor_parallel_attributes(self, defer_init=False) -> None: if not defer_init: # Set parallelism attributes for linear weights - for i in range(self.num_gemms): + grouped_weight = getattr(self, "weight", None) + if grouped_weight is not None: set_tensor_model_parallel_attributes( - tensor=getattr(self, f"weight{i}"), + tensor=grouped_weight, is_parallel=True, dim=1 if self.parallel_mode == "row" else 0, stride=1, ) + else: + for i in range(self.num_gemms): + set_tensor_model_parallel_attributes( + tensor=getattr(self, f"weight{i}"), + is_parallel=True, + dim=1 if self.parallel_mode == "row" else 0, + stride=1, + ) # Set parallelism attributes for linear biases if self.use_bias: @@ -933,7 +952,7 @@ def backward_dw(self): with get_nvtx_range_context("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() wgrad_list = tensor_list[2] - weight_params = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] + weight_params = self._get_weight_tensors() bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] if not self.fuse_wgrad_accumulation: for i in range(self.num_gemms): @@ -983,7 +1002,14 @@ def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage]]: """Get the weight tensors of the module.""" - weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] + grouped_weight = getattr(self, "weight", None) + if grouped_weight is not None: + weight_tensors = grouped_weight.quantized_tensors + if weight_tensors is None: + # TODO(ksivaman): Remove this after GEMM integration. + weight_tensors = grouped_weight.split_into_quantized_tensors() + else: + weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)] if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors): warnings.warn( "You are using quantized weights without quantized compute. " diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index cb199d24b5..5668056700 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -17,10 +17,12 @@ from .storage.mxfp8_tensor_storage import MXFP8TensorStorage from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage from .storage.nvfp4_tensor_storage import NVFP4TensorStorage +from .storage.grouped_tensor_storage import GroupedTensorStorage from .float8_tensor import Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer from .mxfp8_tensor import MXFP8Tensor, MXFP8Quantizer from .float8_blockwise_tensor import Float8BlockwiseQTensor, Float8BlockQuantizer from .nvfp4_tensor import NVFP4Tensor, NVFP4Quantizer +from .grouped_tensor import GroupedTensor from .utils import cast_master_weights_to_fp8, replace_raw_data __all__ = [ @@ -35,11 +37,13 @@ "MXFP8TensorStorage", "Float8BlockwiseQTensorStorage", "NVFP4TensorStorage", + "GroupedTensorStorage", "QuantizedTensor", "Float8Tensor", "MXFP8Tensor", "Float8BlockwiseQTensor", "NVFP4Tensor", + "GroupedTensor", "prepare_for_saving", "restore_from_saved", ] @@ -89,5 +93,7 @@ def get_all_tensor_types(): Float8BlockwiseQTensorStorage, NVFP4Tensor, NVFP4TensorStorage, + GroupedTensor, + GroupedTensorStorage, ] return all_tensor_types diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py new file mode 100644 index 0000000000..767b0ccb35 --- /dev/null +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -0,0 +1,205 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Grouped tensor class for handling collections of tensors with different shapes""" +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch +from torch.utils._pytree import tree_map + +from ..quantized_tensor import QuantizedTensorStorage, Quantizer +from .storage.grouped_tensor_storage import GroupedTensorStorage + + +# For now, conservatively ban all shape manipulating ops. +BANNED_SHAPE_OPS = { + torch.ops.aten.view.default, + torch.ops.aten._unsafe_view.default, + torch.ops.aten.reshape.default, + torch.ops.aten._reshape_alias.default, + torch.ops.aten.flatten.using_ints, + torch.ops.aten.unflatten.int, + torch.ops.aten.squeeze.dim, + torch.ops.aten.squeeze.dims, + torch.ops.aten.unsqueeze.default, + torch.ops.aten.transpose.int, + torch.ops.aten.permute.default, + torch.ops.aten.movedim.int, + torch.ops.aten.t.default, + torch.ops.aten.slice.Tensor, + torch.ops.aten.narrow.default, + torch.ops.aten.select.int, + torch.ops.aten.split.Tensor, + torch.ops.aten.chunk.default, + torch.ops.aten.expand.default, + torch.ops.aten.expand_as.default, + torch.ops.aten.cat.default, + torch.ops.aten.stack.default, +} + + +class GroupedTensor(GroupedTensorStorage, torch.Tensor): + """Tensor wrapper class for grouped tensor storage.""" + + def __new__( + cls, + shape: Tuple[int, int], + dtype: torch.dtype, + num_tensors: int, + shapes: Optional[List[Tuple[int, int]]] = None, + quantizer: Optional[Quantizer] = None, + data: Optional[torch.Tensor] = None, + columnwise_data: Optional[torch.Tensor] = None, + scale_inv: Optional[torch.Tensor] = None, + columnwise_scale_inv: Optional[torch.Tensor] = None, + amax: Optional[torch.Tensor] = None, + columnwise_amax: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + first_dims: Optional[torch.Tensor] = None, + last_dims: Optional[torch.Tensor] = None, + tensor_offsets: Optional[torch.Tensor] = None, + offsets: Optional[List[int]] = None, + scale_inv_offsets: Optional[List[int]] = None, + columnwise_scale_inv_offsets: Optional[List[int]] = None, + ): + del quantizer + del offsets + del scale_inv_offsets + del columnwise_scale_inv_offsets + + if ( + shapes is not None + and len(shapes) == num_tensors + and num_tensors > 0 + and all(shapes[0] == s for s in shapes) + ): + wrapper_shape = (num_tensors, shapes[0][0], shapes[0][1]) + else: + wrapper_shape = shape + + device = None + for maybe_tensor in ( + data, + columnwise_data, + scale_inv, + columnwise_scale_inv, + amax, + columnwise_amax, + scale, + first_dims, + last_dims, + tensor_offsets, + ): + if maybe_tensor is not None: + device = maybe_tensor.device + break + if device is None: + device = torch.device("cuda") + + strides = [1] * len(wrapper_shape) + for i in range(len(wrapper_shape) - 2, -1, -1): + strides[i] = strides[i + 1] * wrapper_shape[i + 1] + return torch.Tensor._make_wrapper_subclass( + cls, + wrapper_shape, + strides=tuple(strides), + storage_offset=0, + dtype=dtype, + layout=torch.strided, + requires_grad=False, + device=device, + ) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs=None): + """Dispatch by dequantizing grouped members, then requantizing writes.""" + if kwargs is None: + kwargs = {} + + # Parameter construction calls detach()/alias-like paths. + if func in (torch.ops.aten.detach.default, torch.ops.aten.alias.default): + return args[0] + + # Don't allow reshape/view etc. + if func in BANNED_SHAPE_OPS: + raise RuntimeError(f"{cls.__name__} forbids shape-manipulation op: {func} ") + + def grouped_to_stacked_tensor(grouped: GroupedTensor) -> torch.Tensor: + if not grouped.all_same_shape(): + raise NotImplementedError( + "GroupedTensor __torch_dispatch__ currently supports only uniform member shapes" + ) + grouped_members = grouped.quantized_tensors + if grouped_members is None: + grouped_members = grouped.split_into_quantized_tensors() + dequantized_members = [ + ( + member.dequantize(dtype=grouped.get_dtype()) + if isinstance(member, QuantizedTensorStorage) + else member + ) + for member in grouped_members + ] + return torch.stack(dequantized_members, dim=0) + + def maybe_unwrap(arg): + if isinstance(arg, GroupedTensor): + return grouped_to_stacked_tensor(arg) + return arg + + def update_grouped_tensor_inplace(grouped: GroupedTensor, updated: torch.Tensor): + if not grouped.all_same_shape(): + raise NotImplementedError( + "GroupedTensor __torch_dispatch__ currently supports only uniform member shapes" + ) + updated_members = list(updated.unbind(dim=0)) + if grouped.quantizer is None: + grouped_members = grouped.quantized_tensors + if grouped_members is None: + grouped_members = grouped.split_into_quantized_tensors() + for dst, src in zip(grouped_members, updated_members): + dst.copy_(src) + else: + grouped.quantize(updated_members) + + def maybe_update_inplace(arg, new_arg, schema_arg): + if ( + isinstance(arg, GroupedTensor) + and isinstance(new_arg, torch.Tensor) + and hasattr(schema_arg, "alias_info") + and hasattr(schema_arg.alias_info, "is_write") + and schema_arg.alias_info.is_write + ): + update_grouped_tensor_inplace(arg, new_arg) + elif isinstance(arg, list) and isinstance(new_arg, list): + for a, na in zip(arg, new_arg): + maybe_update_inplace(a, na, schema_arg) + + # In-place op: dequantize members, perform op, write back into grouped storage. + if func._schema.is_mutable: + new_args = tree_map(maybe_unwrap, args) + new_kwargs = tree_map(maybe_unwrap, kwargs) + schema_args = func._schema.arguments + args_len = len(args) + super().__torch_dispatch__(func, types, new_args, new_kwargs) + for arg, new_arg, schema_arg in zip(args, new_args, schema_args): + maybe_update_inplace(arg, new_arg, schema_arg) + for kwarg, new_kwarg, schema_arg in zip(kwargs, new_kwargs, schema_args[args_len:]): + assert kwarg == new_kwarg == schema_arg.name, "name of kwarg should match schema" + maybe_update_inplace(kwargs[kwarg], new_kwargs[new_kwarg], schema_arg) + return None + + # Default op: operate on dequantized stacked tensors. + new_args = tree_map(maybe_unwrap, args) + new_kwargs = tree_map(maybe_unwrap, kwargs) + return super().__torch_dispatch__(func, types, new_args, new_kwargs) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + # Do not force GroupedTensor on outputs. + return torch._C._disabled_torch_function_impl(func, types, args, kwargs) diff --git a/transformer_engine/pytorch/tensor/storage/__init__.py b/transformer_engine/pytorch/tensor/storage/__init__.py index 7c8a014c1d..44a77d975f 100644 --- a/transformer_engine/pytorch/tensor/storage/__init__.py +++ b/transformer_engine/pytorch/tensor/storage/__init__.py @@ -7,4 +7,4 @@ from .mxfp8_tensor_storage import MXFP8TensorStorage # noqa: F401 from .float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage # noqa: F401 from .nvfp4_tensor_storage import NVFP4TensorStorage # noqa: F401 -from .grouped_tensor import GroupedTensor # noqa: F401 +from .grouped_tensor_storage import GroupedTensorStorage # noqa: F401 diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py similarity index 87% rename from transformer_engine/pytorch/tensor/storage/grouped_tensor.py rename to transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index bf5792ffc9..92006ba45b 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -2,13 +2,12 @@ # # See LICENSE for license information. -"""Grouped tensor class for handling collections of tensors with different shapes""" +"""Grouped tensor storage class for handling collections of tensors with different shapes""" from __future__ import annotations from typing import Optional, Tuple, List, Union import math import torch - from ...quantized_tensor import QuantizedTensorStorage, Quantizer from ..mxfp8_tensor import MXFP8Tensor @@ -21,7 +20,7 @@ from .nvfp4_tensor_storage import NVFP4TensorStorage -class GroupedTensor: +class GroupedTensorStorage: """ EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. @@ -51,10 +50,11 @@ class GroupedTensor: def __init__( self, + shape: Tuple[int, int], + dtype: torch.dtype, num_tensors: int, - shape: Optional[List[Tuple[int, int]]] = None, + shapes: Optional[List[Tuple[int, int]]] = None, quantizer: Optional[Quantizer] = None, - dtype: Optional[torch.dtype] = None, data: Optional[torch.Tensor] = None, columnwise_data: Optional[torch.Tensor] = None, scale_inv: Optional[torch.Tensor] = None, @@ -68,15 +68,16 @@ def __init__( offsets: Optional[List[int]] = None, scale_inv_offsets: Optional[List[int]] = None, columnwise_scale_inv_offsets: Optional[List[int]] = None, - logical_shape: Optional[Tuple[int, int]] = None, ) -> None: """ Initialize a GroupedTensor. Args: + shape: 2D tuple representing conceptual shape + dtype: Data type of the grouped tensor num_tensors: Number of tensors in the group - shape: 2D shape of each tensor (len num_tensors) - quantizer: Quantizer for the grouped tensor + shapes: 2D shape of each tensor (len num_tensors) + quantizer: Quantizer used for all tensors in the group data: Row-wise data buffer (1D flattened) columnwise_data: Column-wise data buffer (1D flattened) scale_inv: Row-wise scale inverse buffer @@ -88,17 +89,14 @@ def __init__( last_dims: Device tensor of int64 array of length num_tensors (or None if uniform) tensor_offsets: Device tensor of int64 array of length num_tensors (or None if uniform) offsets: Vector of integer offsets for each tensor. - logical_shape: 2D tuple representing conceptual shape """ self.num_tensors = num_tensors self.quantizer = quantizer - self.shape = shape - self.dtype = ( - dtype if dtype is not None else torch.float32 - ) # Default to float32 if not provided + self.tensor_shapes = shapes + self.fake_dtype = dtype # Data buffers - self.data = data + self.rowwise_data = data self.columnwise_data = columnwise_data self.scale_inv = scale_inv self.columnwise_scale_inv = columnwise_scale_inv @@ -132,7 +130,7 @@ def __init__( # Logical shape: conceptual 2D shape of the grouped data (REQUIRED) # Represents how the 1D flattened data should be interpreted as 2D # Always 2D with positive dimensions - self.logical_shape = logical_shape if logical_shape is not None else (0, 0) + self.logical_shape = shape # Hold a reference to the quantized tensors that occupy same storage as the GroupedTensor. # Used as a convenience. @@ -145,7 +143,7 @@ def has_data(self) -> bool: Returns: True if data buffer is initialized, False otherwise """ - return self.data is not None + return self.rowwise_data is not None def has_columnwise_data(self) -> bool: """ @@ -239,14 +237,13 @@ def get_dtype(self) -> torch.dtype: The high precision dtype of the data buffer """ - return self.dtype + return self.fake_dtype def clear(self) -> None: """ Reset tensor data and clear all buffers. """ - self.shape = None - self.data = None + self.rowwise_data = None self.columnwise_data = None self.scale_inv = None self.columnwise_scale_inv = None @@ -263,49 +260,34 @@ def clear(self) -> None: self.offsets = None self.scale_inv_offsets = None self.columnwise_scale_inv_offsets = None + self.tensor_shapes = [] + self.fake_dtype = torch.float32 def __repr__(self) -> str: - """String representation of the GroupedTensor.""" + """String representation of the GroupedTensorStorage.""" return ( - f"GroupedTensor(num_tensors={self.num_tensors}, " - f"shape={self.shape}, " + f"GroupedTensorStorage(num_tensors={self.num_tensors}, " + f"shapes={self.tensor_shapes}, " f"logical_shape={self.logical_shape}, " + f"quantizer={self.quantizer}, " f"dtype={self.get_dtype()})" ) - def __str__(self) -> str: - """User-friendly string representation.""" - shape_info = [] - if self.all_same_shape(): - shape_info.append("uniform shape") - else: - if not self.all_same_first_dim(): - shape_info.append("varying first dim") - if not self.all_same_last_dim(): - shape_info.append("varying last dim") - - return ( - f"GroupedTensor with {self.num_tensors} tensors " - f"({', '.join(shape_info) if shape_info else 'uniform'}), " - f"logical_shape={self.logical_shape}, " - f"dtype={self.get_dtype()}" - ) - @staticmethod def make_grouped_tensor_with_shapes( num_tensors: int, - shape: List[Tuple[int, int]], + shapes: List[Tuple[int, int]], quantizer: Optional[Quantizer] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, - ) -> GroupedTensor: + ) -> GroupedTensorStorage: """ Create a GroupedTensor for storing multiple weight tensors of the same shape. Args: num_tensors: Number of tensors - shape: 2D shape of each tensor (len num_tensors) - quantizer: Quantizer for each tensor + shapes: 2D shape of each tensor (len num_tensors) + quantizer: Quantizer used for all tensors device: Device to allocate tensors on, defaults to current cuda device dtype: Data type of the tensor (for high precision case) @@ -314,20 +296,20 @@ def make_grouped_tensor_with_shapes( """ # First dim - first_dim_list = [s[0] for s in shape] + first_dim_list = [s[0] for s in shapes] uniform_first_dim = all(first_dim_list[0] == x for x in first_dim_list) logical_first_dim = sum(first_dim_list) if uniform_first_dim: first_dims = None else: - first_dims = torch.tensor([s[0] for s in shape], dtype=torch.int64, device=device) + first_dims = torch.tensor([s[0] for s in shapes], dtype=torch.int64, device=device) # Last dim - last_dim_list = [s[1] for s in shape] + last_dim_list = [s[1] for s in shapes] logical_last_dim = last_dim_list[0] assert all(logical_last_dim == x for x in last_dim_list), "Last dims should be uniform" - return GroupedTensor.make_grouped_tensor( + return GroupedTensorStorage.make_grouped_tensor( num_tensors=num_tensors, first_dims=first_dims, last_dims=None, @@ -348,7 +330,7 @@ def make_grouped_tensor( quantizer: Optional[Quantizer] = None, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, - ) -> GroupedTensor: + ) -> GroupedTensorStorage: """ Create a GroupedTensor for storing multiple weight tensors of the same shape. @@ -358,8 +340,8 @@ def make_grouped_tensor( last_dims: Device tensor of int64 array of length num_tensors (or None if uniform) logical_first_dim: Logical first dimension logical_last_dim: Logical last dimension - quantizer: Quantizer for each tensor - Used to figure out the recipe and what to allocate. + quantizer: Quantizer used for all tensors. Used to figure out recipe + and what to allocate. device: Device to allocate tensors on, defaults to current cuda device dtype: Data type of the tensor (for high precision case) @@ -574,10 +556,22 @@ def make_grouped_tensor( else: raise ValueError(f"Unsupported quantizer for GroupedTensor: {quantizer}") - grouped_tensor = GroupedTensor( + # Construct wrapper vs storage based on quantizer.internal. + # If quantizer is None (high precision path), default to wrapper class. + # TODO(ksivaman): Properly handle high precision path. + internal = False if quantizer is None else quantizer.internal + if internal: + grouped_tensor_class = GroupedTensorStorage + else: + from ..grouped_tensor import GroupedTensor + + grouped_tensor_class = GroupedTensor + + grouped_tensor = grouped_tensor_class( + logical_shape, + dtype, num_tensors=num_tensors, - shape=shape, - dtype=dtype, + shapes=shape, quantizer=quantizer, data=data, columnwise_data=columnwise_data, @@ -592,7 +586,6 @@ def make_grouped_tensor( offsets=offsets, scale_inv_offsets=scale_inv_offsets, columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, - logical_shape=logical_shape, ) grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() @@ -620,8 +613,8 @@ def split_into_quantized_tensors( no_quantization = self.quantizer is None - # if self.shape is None, then trigger D2H copy and get the shape (not graph safe) - if self.shape is None: + # if self.tensor_shapes is None, then trigger D2H copy and get the shape (not graph safe) + if self.tensor_shapes is None: first_dims_list = ( [self.logical_shape[0]] * self.num_tensors if self.first_dims is None @@ -635,7 +628,7 @@ def split_into_quantized_tensors( shape_list = [] for i in range(self.num_tensors): shape_list.append((first_dims_list[i], last_dims_list[i])) - self.shape = shape_list + self.tensor_shapes = shape_list # edge case: handle the case where tensor_offsets is given but offsets is not set if self.offsets is None and self.tensor_offsets is not None: @@ -645,7 +638,7 @@ def split_into_quantized_tensors( if no_quantization: for i in range(self.num_tensors): # Get tensor shape - tensor_shape = self.shape[i] + tensor_shape = self.tensor_shapes[i] # Get tensor data slice if self.offsets is not None: @@ -654,7 +647,7 @@ def split_into_quantized_tensors( end_offset = start_offset + numel if self.has_data(): - tensor_data = self.data[start_offset:end_offset].view(tensor_shape) + tensor_data = self.rowwise_data[start_offset:end_offset].view(tensor_shape) result.append(tensor_data) elif self.has_columnwise_data(): tensor_data = self.columnwise_data[start_offset:end_offset].view( @@ -670,7 +663,7 @@ def split_into_quantized_tensors( end_offset = start_offset + numel if self.has_data(): - tensor_data = self.data[start_offset:end_offset].view(tensor_shape) + tensor_data = self.rowwise_data[start_offset:end_offset].view(tensor_shape) result.append(tensor_data) elif self.has_columnwise_data(): tensor_data = self.columnwise_data[start_offset:end_offset].view( @@ -698,8 +691,9 @@ def split_into_quantized_tensors( self.columnwise_scale_inv_offsets = self.tensor_offsets // 32 for i in range(self.num_tensors): + quantizer = self.quantizer # Get tensor shape - tensor_shape = self.shape[i] + tensor_shape = self.tensor_shapes[i] numel = tensor_shape[0] * tensor_shape[1] # Get data offsets @@ -712,7 +706,7 @@ def split_into_quantized_tensors( data_end = data_start + numel # Special shape handling for NVFP4. - nvfp4 = self.quantizer._get_compatible_recipe().nvfp4() + nvfp4 = quantizer._get_compatible_recipe().nvfp4() if nvfp4: data_start = data_start // 2 data_end = data_end // 2 @@ -723,15 +717,15 @@ def split_into_quantized_tensors( if self.has_data(): if nvfp4: - rowwise_tensor_shape = self.quantizer.convert_shape_for_fp4(tensor_shape) + rowwise_tensor_shape = quantizer.convert_shape_for_fp4(tensor_shape) else: rowwise_tensor_shape = tensor_shape - rowwise_data = self.data[data_start:data_end].view(rowwise_tensor_shape) + rowwise_data = self.rowwise_data[data_start:data_end].view(rowwise_tensor_shape) if self.has_columnwise_data(): - columnwise_tensor_shape = self.quantizer.get_columnwise_shape(tensor_shape) + columnwise_tensor_shape = quantizer.get_columnwise_shape(tensor_shape) if nvfp4: - columnwise_tensor_shape = self.quantizer.convert_shape_for_fp4( + columnwise_tensor_shape = quantizer.convert_shape_for_fp4( columnwise_tensor_shape ) columnwise_data = self.columnwise_data[data_start:data_end].view( @@ -750,7 +744,7 @@ def split_into_quantized_tensors( scale_end = self.scale_inv_offsets[i + 1] # Calculate expected scale shape for MXFP8 - scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + scale_shape = quantizer.get_scale_shape(tensor_shape, False) rowwise_scale_inv = self.scale_inv[scale_start:scale_end].view(scale_shape) if ( @@ -761,25 +755,25 @@ def split_into_quantized_tensors( # for paged stashing, columnwise_scale_inv should depend on the split offsets cscale_end = self.columnwise_scale_inv_offsets[i + 1] - cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + cscale_shape = quantizer.get_scale_shape(tensor_shape, True) columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( cscale_shape ) - if self.quantizer.internal: + if quantizer.internal: mxfp8_tensor_class = MXFP8TensorStorage else: mxfp8_tensor_class = MXFP8Tensor tensor = mxfp8_tensor_class( shape=tensor_shape, - dtype=self.dtype, + dtype=self.fake_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, columnwise_data=columnwise_data, columnwise_scale_inv=columnwise_scale_inv, - fp8_dtype=self.quantizer.dtype, - quantizer=self.quantizer, - with_gemm_swizzled_scales=self.quantizer.optimize_for_gemm, + fp8_dtype=quantizer.dtype, + quantizer=quantizer, + with_gemm_swizzled_scales=quantizer.optimize_for_gemm, ) result.append(tensor) @@ -790,18 +784,18 @@ def split_into_quantized_tensors( if self.scale_inv is not None: scale_inv = self.scale_inv[i : i + 1] - if self.quantizer.internal: + if quantizer.internal: float8_tensor_class = Float8TensorStorage else: float8_tensor_class = Float8Tensor tensor = float8_tensor_class( shape=tensor_shape, - dtype=self.dtype, + dtype=self.fake_dtype, data=rowwise_data, fp8_scale_inv=scale_inv, - fp8_dtype=self.quantizer.dtype, - quantizer=self.quantizer, + fp8_dtype=quantizer.dtype, + quantizer=quantizer, data_transpose=columnwise_data, ) result.append(tensor) @@ -818,7 +812,7 @@ def split_into_quantized_tensors( scale_end = self.scale_inv_offsets[i + 1] # Get scale shape from quantizer - scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + scale_shape = quantizer.get_scale_shape(tensor_shape, False) rowwise_scale_inv = self.scale_inv[scale_start:scale_end].view(scale_shape) if ( @@ -830,28 +824,28 @@ def split_into_quantized_tensors( cscale_end = self.columnwise_scale_inv_offsets[i + 1] # Get columnwise scale shape from quantizer - cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + cscale_shape = quantizer.get_scale_shape(tensor_shape, True) columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( cscale_shape ) # Compute is_2D_scaled and data_format from quantizer attributes - is_2D_scaled = self.quantizer.block_scaling_dim == 2 + is_2D_scaled = quantizer.block_scaling_dim == 2 - if self.quantizer.internal: + if quantizer.internal: float8_blockwise_q_tensor_class = Float8BlockwiseQTensorStorage else: float8_blockwise_q_tensor_class = Float8BlockwiseQTensor tensor = float8_blockwise_q_tensor_class( shape=tensor_shape, - dtype=self.dtype, + dtype=self.fake_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, columnwise_data=columnwise_data, columnwise_scale_inv=columnwise_scale_inv, - fp8_dtype=self.quantizer.dtype, - quantizer=self.quantizer, + fp8_dtype=quantizer.dtype, + quantizer=quantizer, is_2D_scaled=is_2D_scaled, ) result.append(tensor) @@ -870,7 +864,7 @@ def split_into_quantized_tensors( scale_end = self.scale_inv_offsets[i + 1] # Get scale shape from quantizer - scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + scale_shape = quantizer.get_scale_shape(tensor_shape, False) rowwise_scale_inv = self.scale_inv[scale_start:scale_end].view(scale_shape) if ( @@ -882,7 +876,7 @@ def split_into_quantized_tensors( cscale_end = self.columnwise_scale_inv_offsets[i + 1] # Get columnwise scale shape from quantizer - cscale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + cscale_shape = quantizer.get_scale_shape(tensor_shape, True) columnwise_scale_inv = self.columnwise_scale_inv[cscale_start:cscale_end].view( cscale_shape ) @@ -894,23 +888,23 @@ def split_into_quantized_tensors( if self.columnwise_amax is not None: amax_columnwise = self.columnwise_amax[i : i + 1] - if self.quantizer.internal: + if quantizer.internal: nvfp4_tensor_class = NVFP4TensorStorage else: nvfp4_tensor_class = NVFP4Tensor tensor = nvfp4_tensor_class( shape=tensor_shape, - dtype=self.dtype, + dtype=self.fake_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, - fp4_dtype=self.quantizer.dtype, - quantizer=self.quantizer, - with_gemm_swizzled_scales=self.quantizer.optimize_for_gemm, + fp4_dtype=quantizer.dtype, + quantizer=quantizer, + with_gemm_swizzled_scales=quantizer.optimize_for_gemm, ) result.append(tensor) @@ -919,32 +913,6 @@ def split_into_quantized_tensors( return result - @staticmethod - def create_and_quantize( - tensors: int, - quantizer: None | Quantizer, - *, - device: Optional[torch.device] = None, - dtype: Optional[torch.dtype] = None, - noop_flag: Optional[torch.Tensor] = None, - ) -> Tuple[QuantizedTensorStorage, ...]: - """ - Quantize given tensors into quantized tensors with underlying - storage allocated in a GroupedTensor. - """ - - grouped_tensor = GroupedTensor.make_grouped_tensor_with_shapes( - num_tensors=len(tensors), - shape=[t.shape for t in tensors], - quantizer=quantizer, - device=device, - dtype=dtype, - ) - - grouped_tensor.quantize(tensors, noop_flag=noop_flag) - - return grouped_tensor - def quantize( self, tensors: List[torch.Tensor], From 00ba0b493c27f32e2f210b0022132c50da78dac7 Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Tue, 3 Mar 2026 23:51:57 -0700 Subject: [PATCH 242/521] pass params_dtype to qk_norm creation (#2718) Signed-off-by: Peter St. John Co-authored-by: Xin Yao Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_qk_norm.py | 12 ++++++------ .../pytorch/attention/multi_head_attention.py | 9 ++++++++- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/tests/pytorch/test_qk_norm.py b/tests/pytorch/test_qk_norm.py index 873bd91863..b182d175e7 100644 --- a/tests/pytorch/test_qk_norm.py +++ b/tests/pytorch/test_qk_norm.py @@ -11,7 +11,8 @@ @pytest.mark.parametrize("qk_norm_type", [None, "L2Normalization", "RMSNorm", "LayerNorm"]) @pytest.mark.parametrize("attention_type", ["self", "cross"]) @pytest.mark.parametrize("qk_norm_eps", [1e-6, 1e-5]) -def test_qk_norm_functionality(qk_norm_type, attention_type, qk_norm_eps) -> None: +@pytest.mark.parametrize("params_dtype", [torch.float32, torch.bfloat16]) +def test_qk_norm_functionality(qk_norm_type, attention_type, qk_norm_eps, params_dtype) -> None: """Test QK normalization functionality, module structure, and numerical behavior.""" hidden_size = 256 num_attention_heads = 8 @@ -26,6 +27,7 @@ def test_qk_norm_functionality(qk_norm_type, attention_type, qk_norm_eps) -> Non qk_norm_eps=qk_norm_eps, bias=False, device="cuda", + params_dtype=params_dtype, ).cuda() # Check module structure based on qk_norm_type parameter @@ -78,13 +80,11 @@ def test_qk_norm_functionality(qk_norm_type, attention_type, qk_norm_eps) -> Non # Create input tensors batch_size = 2 # Use a fixed batch size for testing - hidden_states = torch.randn( - seq_len, batch_size, hidden_size, device="cuda", dtype=torch.float32 - ) + hidden_states = torch.randn(seq_len, batch_size, hidden_size, device="cuda", dtype=params_dtype) if attention_type == "cross": encoder_output = torch.randn( - seq_len, batch_size, hidden_size, device="cuda", dtype=torch.float32 + seq_len, batch_size, hidden_size, device="cuda", dtype=params_dtype ) else: encoder_output = None @@ -109,7 +109,7 @@ def test_qk_norm_functionality(qk_norm_type, attention_type, qk_norm_eps) -> Non if attention_type == "self": head_dim = hidden_size // num_attention_heads rotary_dim = head_dim // 2 - rotary_pos_emb = torch.randn(seq_len, 1, 1, rotary_dim, device="cuda", dtype=torch.float32) + rotary_pos_emb = torch.randn(seq_len, 1, 1, rotary_dim, device="cuda", dtype=params_dtype) with torch.no_grad(): output_with_rope = mha(hidden_states, rotary_pos_emb=rotary_pos_emb) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 5c581849e6..d95d327c78 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -355,7 +355,7 @@ def __init__( } self.q_norm, self.k_norm = self._create_qk_norm_modules( - qk_norm_type, qk_norm_eps, device, seq_length, micro_batch_size + qk_norm_type, qk_norm_eps, device, seq_length, micro_batch_size, params_dtype ) qkv_parallel_mode = "column" if set_parallel_mode else None @@ -489,6 +489,7 @@ def _create_qk_norm_modules( device: Union[torch.device, str], seq_length: Optional[int] = None, micro_batch_size: Optional[int] = None, + params_dtype: Optional[torch.dtype] = None, ) -> Tuple[Optional[torch.nn.Module], Optional[torch.nn.Module]]: """ Create query and key normalization modules based on the specified normalization type. @@ -505,6 +506,8 @@ def _create_qk_norm_modules( Sequence length for L2Normalization optimization micro_batch_size : Optional[int], default = None Micro batch size for L2Normalization optimization + params_dtype : Optional[torch.dtype], default = None + Data type for the normalization modules Returns ------- @@ -528,11 +531,13 @@ def _create_qk_norm_modules( normalized_shape=self.hidden_size_per_attention_head, eps=qk_norm_eps, device=device, + params_dtype=params_dtype, ) k_norm = RMSNorm( normalized_shape=self.hidden_size_per_attention_head, eps=qk_norm_eps, device=device, + params_dtype=params_dtype, ) return q_norm, k_norm @@ -541,11 +546,13 @@ def _create_qk_norm_modules( normalized_shape=self.hidden_size_per_attention_head, eps=qk_norm_eps, device=device, + params_dtype=params_dtype, ) k_norm = LayerNorm( normalized_shape=self.hidden_size_per_attention_head, eps=qk_norm_eps, device=device, + params_dtype=params_dtype, ) return q_norm, k_norm From 505b89699c639fdc10883248aae06a125549ed35 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 4 Mar 2026 08:09:06 -0800 Subject: [PATCH 243/521] [JAX] GSPMD Deprecation Warning - Only trigger when the primitive is invoked (#2729) * gspmd warning - only trigger when the primitive is invoked --------- Signed-off-by: Phuong Nguyen --- transformer_engine/jax/cpp_extensions/base.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/base.py b/transformer_engine/jax/cpp_extensions/base.py index ae3888cf04..6eb588c849 100644 --- a/transformer_engine/jax/cpp_extensions/base.py +++ b/transformer_engine/jax/cpp_extensions/base.py @@ -189,12 +189,12 @@ def _warn_gspmd_deprecation_once(): global _gspmd_deprecation_warned if not _gspmd_deprecation_warned: warnings.warn( - "GSPMD sharding propagation is planned to be removed in June 2026." - " It is no longer maintained or tested. Use it at your own risk." - " Please use Shardy partitioner instead." + "GSPMD sharding propagation rules in TE-JAX are planned to be removed in June 2026." + " They are no longer maintained or tested. Use them at your own risk." + " Please use Shardy propagation instead." " In case you cannot upgrade to a JAX version that supports Shardy, please reach out!", DeprecationWarning, - stacklevel=3, + stacklevel=2, ) _gspmd_deprecation_warned = True @@ -234,12 +234,24 @@ def name_of_wrapper_p(): outer_p.def_abstract_eval(cls.outer_abstract) batching.primitive_batchers[outer_p] = cls.batcher outer_p_lower = custom_partitioning(cls.impl, static_argnums=cls.impl_static_args) + if _JAX_GSPMD_SUPPORTED: - if "infer_sharding_from_operands" in cls.__dict__: - _warn_gspmd_deprecation_once() - gspmd_kwargs = {"infer_sharding_from_operands": cls.infer_sharding_from_operands} + fn = cls.__dict__.get("infer_sharding_from_operands") + if fn is not None: + actual_fn = ( + cls.infer_sharding_from_operands + ) # Use descriptor protocol to unwrap staticmethod + + def _gspmd_wrapper(*args, **kwargs): + _warn_gspmd_deprecation_once() + return actual_fn(*args, **kwargs) + + gspmd_kwargs = {"infer_sharding_from_operands": _gspmd_wrapper} + else: + gspmd_kwargs = {"infer_sharding_from_operands": cls.infer_sharding_from_operands} else: gspmd_kwargs = {} + outer_p_lower.def_partition( partition=cls.partition, sharding_rule=cls.shardy_sharding_rule, From 139c863f92420271bbae2cbce49d9b170b7d03f9 Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Wed, 4 Mar 2026 12:37:53 -0700 Subject: [PATCH 244/521] Add fused_adam, quantized_model_init, and fsdp2 example (#2698) Expand fsdp2 test suite, add example for FusedAdam w/ and w/o fully_shard Signed-off-by: Peter St. John Co-authored-by: vthumbe1503 --- .../quantized_model_init/fully_shard.py | 266 +++++++ examples/pytorch/quantized_model_init/main.py | 151 ++++ .../distributed/run_fsdp2_fused_adam.py | 653 ++++++++++++++++++ tests/pytorch/distributed/run_fsdp2_model.py | 63 +- tests/pytorch/distributed/test_torch_fsdp2.py | 210 +++++- .../quantization_current_scaling.py | 6 + .../pytorch/optimizers/fused_adam.py | 17 +- .../pytorch/tensor/float8_blockwise_tensor.py | 33 +- .../pytorch/tensor/float8_tensor.py | 6 + .../pytorch/tensor/nvfp4_tensor.py | 23 +- 10 files changed, 1366 insertions(+), 62 deletions(-) create mode 100644 examples/pytorch/quantized_model_init/fully_shard.py create mode 100644 examples/pytorch/quantized_model_init/main.py create mode 100644 tests/pytorch/distributed/run_fsdp2_fused_adam.py diff --git a/examples/pytorch/quantized_model_init/fully_shard.py b/examples/pytorch/quantized_model_init/fully_shard.py new file mode 100644 index 0000000000..6131712001 --- /dev/null +++ b/examples/pytorch/quantized_model_init/fully_shard.py @@ -0,0 +1,266 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FSDP2 distributed training with quantized model initialization. + +Extends the single-GPU ``main.py`` example to multi-GPU training using +PyTorch-native FSDP2 (``fully_shard``). The script demonstrates: + +1. **Meta-device initialization** -- Model parameters are created on the + ``meta`` device (zero memory), then FSDP2 sharding is applied, and + finally ``reset_parameters()`` materializes and quantizes only the + local shards on each rank's GPU. +2. ``quantized_model_init`` -- Flags the model for FP8 weight initialization + (actual quantization happens in ``reset_parameters`` after sharding). +3. ``fully_shard`` -- PyTorch FSDP2 sharding of each TransformerLayer. +4. ``FusedAdam`` with FP32 master weights for full-precision training updates. + +.. note:: + ``fuse_wgrad_accumulation`` is **not** used here. That feature writes + weight gradients directly into ``main_grad`` buffers, bypassing the + autograd gradient flow. FSDP2 requires gradients to go through its + reduce-scatter, so ``fuse_wgrad_accumulation`` needs Megatron-Core's + FSDP integration (which provides ``get_main_grad()``). + +Usage:: + + torchrun --nproc-per-node 2 fully_shard.py +""" + +import os + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed._composable.fsdp import fully_shard +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import QuantizedTensor +from transformer_engine.pytorch.module.base import TransformerEngineBaseModule + +# ── Configuration (matches main.py) ────────────────────────────────── +HIDDEN_SIZE = 256 +FFN_HIDDEN_SIZE = 1024 +NUM_ATTENTION_HEADS = 8 +NUM_LAYERS = 3 +SEQ_LEN = 32 +BATCH_PER_RANK = 2 +NUM_STEPS = 5 +DTYPE = torch.bfloat16 + + +def dist_print(msg): + """Print only on rank 0.""" + if int(os.environ.get("RANK", "0")) == 0: + print(msg) + + +def main(): + # ── 1. Distributed setup ───────────────────────────────────────── + assert "TORCHELASTIC_RUN_ID" in os.environ, ( + "This script must be launched with torchrun, e.g.:\n" + " torchrun --nproc-per-node 2 fully_shard.py" + ) + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="nccl") + device = torch.device(f"cuda:{local_rank}") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + # ── 2. Create model on meta device (zero memory) ──────────────── + # quantized_model_init sets the flag for FP8 weight initialization, + # but with device="meta" no actual memory is allocated yet. + with te.quantized_model_init(enabled=True): + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + fuse_qkv_params=True, + params_dtype=DTYPE, + hidden_dropout=0.0, + attention_dropout=0.0, + device="meta", + ) + for _ in range(NUM_LAYERS) + ] + ) + + # Verify all parameters are on meta device (no GPU memory used). + for name, param in model.named_parameters(): + assert param.device == torch.device("meta"), f"{name} is not on meta device" + dist_print("Model created on meta device (zero GPU memory).") + + # ── 3. FSDP2 sharding ──────────────────────────────────────────── + # Apply sharding to the meta-device model. FSDP2 wraps parameters + # as DTensors but no GPU memory is allocated yet. + mesh = DeviceMesh("cuda", list(range(world_size))) + for child in model.children(): + fully_shard(child, mesh=mesh) + fully_shard(model, mesh=mesh) + dist_print("FSDP2 sharding applied to meta-device model.") + + # ── 4. Materialize parameters on GPU ────────────────────────────── + # reset_parameters() on each TE module materializes the local shard + # on CUDA, applies weight initialization, and quantizes to FP8. + for module in model.modules(): + if isinstance(module, TransformerEngineBaseModule): + module.reset_parameters() + + # Post-materialization verification. + for name, param in model.named_parameters(): + assert isinstance(param, DTensor), f"{name} is not a DTensor after sharding" + qt_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) + ) + assert qt_count > 0, "No QuantizedTensor local tensors after materialization" + dist_print( + f"Parameters materialized: {qt_count} FP8 (QuantizedTensor) weight params " + "wrapped in DTensors." + ) + + # ── 5. Optimizer ───────────────────────────────────────────────── + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + dist_print("Using FusedAdam with master_weights=True.") + + # ── 6. Training loop ───────────────────────────────────────────── + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=DTYPE, device=device) + target = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=DTYPE, device=device) + + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + + with te.autocast(enabled=True): + output = model(x) + + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + dist_print(f" Step {step}: loss = {loss.item():.6f}") + + # ── 7. Post-training assertions ────────────────────────────────── + dist_print("\nVerifying invariants ...") + + qt_after = 0 + for name, param in model.named_parameters(): + assert isinstance(param, DTensor), f"{name} lost DTensor wrapping" + if isinstance(param._local_tensor, QuantizedTensor): + qt_after += 1 + assert qt_after > 0, "No QuantizedTensor local tensors after training" + dist_print(f" {qt_after} params still have QuantizedTensor local tensors.") + + # Optimizer states: master weights and moments should be float32. + for param in model.parameters(): + state = optimizer.state[param] + if "master_param" in state: + assert ( + state["master_param"].dtype == torch.float32 + ), f"Master weight dtype {state['master_param'].dtype}, expected float32" + assert state["exp_avg"].dtype == torch.float32, "exp_avg should be float32" + assert state["exp_avg_sq"].dtype == torch.float32, "exp_avg_sq should be float32" + + dist_print("All assertions passed!") + dist_print(" - Linear weight parameters: QuantizedTensor (FP8) wrapped in DTensor") + dist_print(" - Optimizer master weights: float32") + dist_print(" - Optimizer states (exp_avg, exp_avg_sq): float32") + + # ── 8. Distributed checkpoint: save and load ───────────────────── + # torch.distributed.checkpoint (DCP) saves sharded state — each rank + # writes only its local shard. This preserves FP8 compute weights + # and the full optimizer state (master weights, moments, step count). + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + ) + + # Use a fixed path so all ranks agree on the checkpoint location. + checkpoint_dir = "/tmp/te_fsdp2_example_checkpoint" + dist_print(f"\nSaving distributed checkpoint to {checkpoint_dir} ...") + + # Save sharded checkpoint. DCP handles DTensor shards natively — + # each rank writes only its local shard to the filesystem. + dcp.save( + {"model": model.state_dict(), "optimizer": optimizer.state_dict()}, + checkpoint_id=checkpoint_dir, + ) + dist_print(" Checkpoint saved (FP8 weights + optimizer state).") + + # Load checkpoint back. Provide empty state dict containers with the + # same structure; DCP fills them from the saved files. + state_to_load = {"model": model.state_dict(), "optimizer": optimizer.state_dict()} + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model.load_state_dict(state_to_load["model"]) + optimizer.load_state_dict(state_to_load["optimizer"]) + dist_print(" Checkpoint loaded — FP8 weights and optimizer state restored.") + + # Verify training continues after checkpoint load. + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + dist_print(f" Post-checkpoint training step: loss = {loss.item():.6f}") + + # ── 9. Save full-precision (FP32) model to safetensors ─────────── + # For inference or fine-tuning you typically want FP32 weights, not + # FP8 compute weights. The optimizer's master weight copies are the + # authoritative FP32 values (more precise than dequantizing FP8). + # All ranks must participate in gathering; only rank 0 saves. + from safetensors.torch import save_file + + full_opts = StateDictOptions(full_state_dict=True, cpu_offload=True) + + full_model_state = get_model_state_dict(model, options=full_opts) + full_opt_state = get_optimizer_state_dict(model, optimizer, options=full_opts) + + rank = int(os.environ.get("RANK", "0")) + if rank == 0: + fp32_state = {} + opt_param_states = full_opt_state.get("state", {}) + + for key, value in full_model_state.items(): + if key in opt_param_states and "master_param" in opt_param_states[key]: + # Prefer optimizer's FP32 master weight (maintained throughout training). + fp32_state[key] = opt_param_states[key]["master_param"].float() + elif isinstance(value, QuantizedTensor): + # Fallback: dequantize FP8 → FP32 (e.g. if master_weights was off). + fp32_state[key] = value.dequantize().float() + else: + # Non-FP8 params (e.g. LayerNorm weights): cast to FP32. + fp32_state[key] = value.float() + + save_path = "/tmp/te_fsdp2_example_model_fp32.safetensors" + save_file(fp32_state, save_path) + dist_print(f"\nSaved FP32 model ({len(fp32_state)} params) to {save_path}") + + # Quick verification: all saved tensors are float32. + from safetensors.torch import load_file + + loaded = load_file(save_path) + for k, v in loaded.items(): + assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + dist_print(f" Verified: all {len(loaded)} tensors are float32.") + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/examples/pytorch/quantized_model_init/main.py b/examples/pytorch/quantized_model_init/main.py new file mode 100644 index 0000000000..a9d3480cad --- /dev/null +++ b/examples/pytorch/quantized_model_init/main.py @@ -0,0 +1,151 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Quantized model initialization with FusedAdam and gradient accumulation fusion. + +Demonstrates three Transformer Engine features working together: + +1. ``quantized_model_init`` -- Initialize a model with low-precision (FP8) + parameters, avoiding the memory cost of storing both high-precision and + quantized copies of every weight. + +2. ``FusedAdam`` with master weights -- Maintain FP32 master copies of the + weights inside the optimizer so that the training update retains full + precision despite the model parameters being FP8. + +3. Gradient accumulation fusion -- Use ``fuse_wgrad_accumulation=True`` + together with per-parameter ``main_grad`` buffers so that weight + gradients are accumulated directly in FP32 via Tensor Cores, avoiding a + separate FP8-to-FP32 cast kernel. + +Usage:: + + python main.py +""" + +import torch +import transformer_engine.pytorch as te +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor + +# ── Configuration ────────────────────────────────────────────────────── +HIDDEN_SIZE = 256 +FFN_HIDDEN_SIZE = 1024 +NUM_ATTENTION_HEADS = 8 +SEQ_LEN = 32 +BATCH_SIZE = 2 +NUM_STEPS = 5 +DTYPE = torch.bfloat16 + + +def main(): + # ── 1. Create model with quantized parameters ───────────────────── + # + # Inside quantized_model_init, TransformerEngine modules store only the + # FP8 quantized copy of each parameter (a Float8Tensor), eliminating the + # memory overhead of a high-precision shadow copy. + with te.quantized_model_init(enabled=True): + model = te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + fuse_wgrad_accumulation=True, + fuse_qkv_params=True, # required for fuse_wgrad_accumulation + params_dtype=DTYPE, + hidden_dropout=0.0, # disable dropout for this synthetic example + attention_dropout=0.0, + ) + + # Verify that linear-layer weight parameters are quantized. + # Biases and LayerNorm parameters are *not* quantized. + quantized_count = 0 + for name, param in model.named_parameters(): + if isinstance(param, QuantizedTensor): + quantized_count += 1 + assert quantized_count > 0, "No QuantizedTensor parameters found" + print(f"Found {quantized_count} QuantizedTensor (FP8) weight parameters.") + + # ── 2. Allocate main_grad buffers (FP32) ────────────────────────── + # + # fuse_wgrad_accumulation causes weight-gradient GEMMs to write directly + # into ``param.main_grad`` in FP32 (via Tensor Core accumulation). + # Non-weight parameters (e.g. LayerNorm) still receive gradients through + # the normal ``param.grad`` path. + for param in model.parameters(): + param.main_grad = torch.zeros(param.shape, dtype=torch.float32, device=param.device) + + # ── 3. Optimizer with FP32 master weights ───────────────────────── + # + # use_decoupled_grad=True tells FusedAdam to read gradients from + # ``param.decoupled_grad`` instead of ``param.grad``. This avoids + # the dtype-mismatch error that would occur when assigning FP32 + # gradients to bfloat16 parameters via ``.grad``. + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + ) + + # ── 4. Training loop ────────────────────────────────────────────── + # + # Use a fixed synthetic dataset so that loss decreases over steps. + x = torch.randn(SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE, dtype=DTYPE, device="cuda") + target = torch.randn(SEQ_LEN, BATCH_SIZE, HIDDEN_SIZE, dtype=DTYPE, device="cuda") + + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + for param in model.parameters(): + param.main_grad.zero_() + + # Forward pass inside autocast to enable FP8 compute. + with te.autocast(enabled=True): + output = model(x) + + loss = torch.nn.functional.mse_loss(output, target) + loss.backward() + + # Consolidate gradients into main_grad. + # * Weight params with fuse_wgrad_accumulation: backward already + # accumulated the gradient directly into main_grad (FP32). + # * Other params (e.g. LayerNorm): autograd set param.grad. + for param in model.parameters(): + if param.grad is not None: + param.main_grad.copy_(param.grad) + param.grad = None + + # Expose main_grad as decoupled_grad so FusedAdam can read it. + for param in model.parameters(): + param.decoupled_grad = param.main_grad + + optimizer.step() + print(f" Step {step}: loss = {loss.item():.6f}") + + # ── 5. Post-training assertions ─────────────────────────────────── + print("\nVerifying invariants ...") + + # Optimizer states. + for param in model.parameters(): + state = optimizer.state[param] + if "master_param" in state: + master = state["master_param"] + assert ( + master.dtype == torch.float32 + ), f"Master weight dtype {master.dtype}, expected float32" + assert state["exp_avg"].dtype == torch.float32, "exp_avg should be float32" + assert state["exp_avg_sq"].dtype == torch.float32, "exp_avg_sq should be float32" + + # main_grad buffers. + for param in model.parameters(): + assert param.main_grad.dtype == torch.float32, "main_grad should be float32" + + print("All assertions passed!") + print(" - Linear weight parameters: QuantizedTensor (FP8)") + print(" - Optimizer master weights: float32") + print(" - Optimizer states (exp_avg, exp_avg_sq): float32") + print(" - Gradient accumulation buffers (main_grad): float32") + + +if __name__ == "__main__": + main() diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/run_fsdp2_fused_adam.py new file mode 100644 index 0000000000..0439bf1b5a --- /dev/null +++ b/tests/pytorch/distributed/run_fsdp2_fused_adam.py @@ -0,0 +1,653 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FSDP2 + FusedAdam compatibility tests. + +Launched via torchrun from test_fused_optimizer.py. +""" + +import argparse +import functools +import os + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed._composable.fsdp import fully_shard +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.tensor import DTensor + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import QuantizedTensor +import transformer_engine.common.recipe + + +def get_recipe_from_string(recipe): + return getattr(transformer_engine.common.recipe, recipe)() + + +HIDDEN_SIZE = 256 +FFN_HIDDEN_SIZE = 1024 +NUM_ATTENTION_HEADS = 8 +NUM_LAYERS = 2 +SEQ_LEN = 32 +BATCH_PER_RANK = 2 +NUM_STEPS = 3 + + +def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + if isinstance(param, QuantizedTensor): + ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] + else: + ignore_keys = [] + attrs = vars(param) + custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} + return custom_attrs + + +def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) + + +def _setup(): + """Common distributed setup. Returns (world_size, local_rank, device).""" + world_size = int(os.environ["WORLD_SIZE"]) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + # CPU backend required for async save + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + device = torch.device(f"cuda:{local_rank}") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + return world_size, local_rank, device + + +def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None): + """Build a Sequential of TransformerLayers, optionally with FP8 init.""" + if fp8_init: + ctx = te.quantized_model_init(enabled=True, recipe=recipe) + else: + from contextlib import nullcontext + + ctx = nullcontext() + with ctx: + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + for _ in range(NUM_LAYERS) + ] + ) + return model + + +def _shard_model(model, world_size): + """Apply FSDP2 sharding with save/restore custom attrs.""" + custom_attrs = save_custom_attrs(model) + mesh = DeviceMesh("cuda", list(range(world_size))) + for child in model.children(): + fully_shard(child, mesh=mesh) + fully_shard(model, mesh=mesh) + restore_custom_attrs(model, custom_attrs) + return model + + +def test_fused_adam_fp8_master_weights(recipe=None): + """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). + + Verifies: + - Optimizer states are created with correct dtype (float32) + - Training loop completes without error + - DTensor wrapping and QuantizedTensor local tensors are preserved + """ + world_size, _, device = _setup() + + model = _build_model(fp8_init=True, recipe=recipe) + + # Verify FP8 params created + qt_count = sum(1 for _, p in model.named_parameters() if isinstance(p, QuantizedTensor)) + assert qt_count > 0, "No QuantizedTensor local tensors before training" + + model = _shard_model(model, world_size) + + # Verify params are DTensors + for name, param in model.named_parameters(): + assert isinstance(param, DTensor), f"{name} is not DTensor" + + # Verify FP8 params after sharding + qt_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) + ) + assert qt_count > 0, "No QuantizedTensor local tensors after sharding" + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Verify optimizer states + for param in model.parameters(): + state = optimizer.state[param] + assert ( + state["exp_avg"].dtype == torch.float32 + ), f"exp_avg dtype {state['exp_avg'].dtype}, expected float32" + assert ( + state["exp_avg_sq"].dtype == torch.float32 + ), f"exp_avg_sq dtype {state['exp_avg_sq'].dtype}, expected float32" + if "master_param" in state: + assert ( + state["master_param"].dtype == torch.float32 + ), f"master_param dtype {state['master_param'].dtype}, expected float32" + + # Verify FP8 params preserved + qt_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) + ) + assert qt_count > 0, "No QuantizedTensor local tensors after training" + + dist.destroy_process_group() + + +def test_fused_adam_bf16(recipe=None): + """FusedAdam with master_weights + FSDP2 + bf16 params (no FP8). + + Verifies the non-FP8 DTensor param path in step() works correctly. + """ + world_size, _, device = _setup() + + model = _build_model(fp8_init=False) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + # Verify optimizer states are float32 + for param in model.parameters(): + state = optimizer.state[param] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + + # Verify loss decreased (basic sanity) + assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + + dist.destroy_process_group() + + +def test_fused_adam_fp8_no_master(recipe=None): + """FusedAdam without master_weights + FSDP2 + FP8 params. + + Verifies FusedAdam works with FSDP2 even without master weights enabled. + """ + world_size, _, device = _setup() + + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=False, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Verify DTensors preserved + for name, param in model.named_parameters(): + assert isinstance(param, DTensor), f"{name} lost DTensor wrapping" + + dist.destroy_process_group() + + +def test_fused_adam_bf16_store_param_remainders(recipe=None): + """FusedAdam with master_weights + store_param_remainders + FSDP2 + bf16 params. + + store_param_remainders stores only the trailing 16 remainder bits (int16) + instead of full FP32 master params. The FP32 master can be reconstructed + from BF16 params + int16 remainders. Only works with bf16 params + fp32 + master weights. + + Verifies: + - Training loop completes without error + - Optimizer master_param states are int16 (remainder bits) + - exp_avg and exp_avg_sq are float32 + - Loss decreases (basic sanity) + """ + world_size, _, device = _setup() + + model = _build_model(fp8_init=False) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + store_param_remainders=True, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + losses = [] + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + optimizer.step() + + # Verify model params are bf16 (required for store_param_remainders) + for name, param in model.named_parameters(): + assert ( + param.dtype == torch.bfloat16 + ), f"{name}: param dtype {param.dtype}, expected bfloat16" + + # Verify optimizer states + for name, param in model.named_parameters(): + state = optimizer.state[param] + assert ( + state["exp_avg"].dtype == torch.float32 + ), f"{name}: exp_avg dtype {state['exp_avg'].dtype}, expected float32" + assert ( + state["exp_avg_sq"].dtype == torch.float32 + ), f"{name}: exp_avg_sq dtype {state['exp_avg_sq'].dtype}, expected float32" + # store_param_remainders stores master_param as int16 remainder bits + if "master_param" in state: + assert ( + state["master_param"].dtype == torch.int16 + ), f"{name}: master_param dtype {state['master_param'].dtype}, expected int16" + + # Verify loss decreased (basic sanity) + assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + + dist.destroy_process_group() + + +def test_fuse_wgrad_accumulation(recipe=None): + """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail. + + With vanilla FSDP2, PyTorch's autograd Function.apply unwraps DTensor + inputs to local tensors. The local Float8Tensor inside the autograd + function does not have the `main_grad` attribute (which is set on the + DTensor parameter). This causes an AttributeError during backward. + + Additionally, even if main_grad were accessible, fuse_wgrad_accumulation + writes the gradient directly into main_grad and returns None to autograd, + bypassing FSDP2's reduce-scatter. + """ + world_size, _, device = _setup() + + model = _build_model(fp8_init=True, fuse_wgrad_accumulation=True, recipe=recipe) + + # Allocate main_grad buffers on the DTensor params + for param in model.parameters(): + param.main_grad = torch.zeros(param.shape, dtype=torch.float32, device=param.device) + + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # This is currently failing during backward because the local Float8Tensor + # inside the autograd function doesn't have main_grad. + optimizer.zero_grad(set_to_none=True) + for param in model.parameters(): + param.main_grad.zero_() + + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + + loss = F.mse_loss(output, target) + loss.backward() # Expected to raise AttributeError + + dist.destroy_process_group() + + +def test_safetensors_fp32_export(recipe=None): + """Export full-precision (FP32) model to safetensors from optimizer master weights. + + Verifies: + - get_model_state_dict with full_state_dict gathers all params + - get_optimizer_state_dict with full_state_dict gathers optimizer state + - FP32 state dict is built from optimizer master weights + - All saved tensors are float32 + - Saved tensor shapes match expected (unsharded) shapes + """ + from safetensors.torch import load_file, save_file + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + ) + + world_size, _, device = _setup() + + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # Train a few steps. + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Gather full state dicts (all ranks participate). + full_opts = StateDictOptions(full_state_dict=True, cpu_offload=True) + full_model_state = get_model_state_dict(model, options=full_opts) + full_opt_state = get_optimizer_state_dict(model, optimizer, options=full_opts) + + rank = int(os.environ.get("RANK", "0")) + save_path = "/tmp/te_test_fsdp2_model_fp32.safetensors" + + if rank == 0: + # Build FP32 state dict from optimizer master weights. + fp32_state = {} + opt_param_states = full_opt_state.get("state", {}) + + for key, value in full_model_state.items(): + if key in opt_param_states and "master_param" in opt_param_states[key]: + fp32_state[key] = opt_param_states[key]["master_param"].float() + else: + fp32_state[key] = value.float() + + assert len(fp32_state) > 0, "FP32 state dict is empty" + + # Save and verify. + save_file(fp32_state, save_path) + loaded = load_file(save_path) + + assert len(loaded) == len( + fp32_state + ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" + for k, v in loaded.items(): + assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + + # Clean up. + os.remove(save_path) + + dist.destroy_process_group() + + +def test_dcp_output_parity(recipe=None, async_save=False): + """DCP save/load round-trip produces bitwise-identical model outputs. + + 1. Builds and trains a model for NUM_STEPS + 2. Runs a forward pass and records the output + 3. Saves model + optimizer state via DCP + 4. Builds a *fresh* model + optimizer (same architecture) + 5. Loads the DCP checkpoint into the fresh model + 6. Runs the same forward pass and asserts outputs are identical + 7. Runs one more training step on both models and asserts outputs still match + """ + import torch.distributed.checkpoint as dcp + + world_size, local_rank, device = _setup() + + # ── Build and train the original model ─────────────────────────── + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Record reference output from the trained model. + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + ref_output = model(x).clone() + + # ── Save checkpoint ────────────────────────────────────────────── + checkpoint_dir = "/tmp/te_test_fsdp2_dcp_parity" + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # We need to remove the _extra_state keys from the model state dict for DelayedScaling, + # since otherwise we'll run into an error that the tensor sizes are different. The + # alternative is a LoadPlanner that dynamically re-sizes the input tensors, see + # NVIDIA/TransformerEngine#1860 for more details. + model_state = { + k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") + } + else: + model_state = model.state_dict() + + if not async_save: + dcp.save( + {"model": model_state, "optimizer": optimizer.state_dict()}, + checkpoint_id=checkpoint_dir, + ) + future = None + else: + future = dcp.async_save( + {"model": model_state, "optimizer": optimizer.state_dict()}, + checkpoint_id=checkpoint_dir, + ) + + # ── Build a fresh model and load the checkpoint ────────────────── + model2 = _build_model(fp8_init=True, recipe=recipe) + model2 = _shard_model(model2, world_size) + + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + # Populate optimizer state so load_state_dict has matching structure. + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + model2_state = { + k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") + } + else: + model2_state = model2.state_dict() + + state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + + if async_save: + future.result() # Block on async save completion + + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict( + state_to_load["model"], + strict=( + False if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) else True + ), + ) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + # ── Verify identical forward-pass output ───────────────────────── + with torch.no_grad(): + 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. + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0.05, + atol=0.1, + msg="Fresh model loaded from DCP checkpoint produces different output", + ) + else: + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0, + atol=0, + msg="Fresh model loaded from DCP checkpoint produces different output", + ) + + # ── Verify one more training step produces identical results ───── + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out1 = model(x) + loss1 = F.mse_loss(out1, target) + loss1.backward() + optimizer.step() + + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out2 = model2(x) + loss2 = F.mse_loss(out2, target) + loss2.backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + torch.testing.assert_close( + out2, + out1, + rtol=0.05, + atol=0.1, + msg="Training step after DCP load produces different output", + ) + else: + torch.testing.assert_close( + out2, out1, msg="Training step after DCP load produces different output" + ) + + # ── Cleanup ────────────────────────────────────────────────────── + import shutil + + if int(os.environ.get("RANK", "0")) == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + + dist.destroy_process_group() + + +TESTS = { + "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, + "fused_adam_bf16": test_fused_adam_bf16, + "fused_adam_fp8_no_master": test_fused_adam_fp8_no_master, + "fused_adam_bf16_store_param_remainders": test_fused_adam_bf16_store_param_remainders, + "fuse_wgrad_accumulation": test_fuse_wgrad_accumulation, + "dcp_output_parity": functools.partial(test_dcp_output_parity, async_save=False), + "dcp_output_parity_async": functools.partial(test_dcp_output_parity, async_save=True), + "safetensors_fp32_export": test_safetensors_fp32_export, +} + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--test", required=True, choices=list(TESTS.keys())) + parser.add_argument( + "--recipe", + type=str, + default="MXFP8BlockScaling", + help="Quantizer type.", + choices=[ + "DelayedScaling", + "Float8CurrentScaling", + "Float8BlockScaling", + "MXFP8BlockScaling", + "NVFP4BlockScaling", + ], + ) + args = parser.parse_args() + recipe = get_recipe_from_string(args.recipe) + TESTS[args.test](recipe) diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/run_fsdp2_model.py index 5df3468861..60d7cd2023 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/run_fsdp2_model.py @@ -9,12 +9,7 @@ import argparse import transformer_engine.pytorch as te -from transformer_engine.common.recipe import ( - Format, - DelayedScaling, - Float8CurrentScaling, - MXFP8BlockScaling, -) +import transformer_engine.common.recipe import torch import torch.distributed as dist @@ -43,14 +38,23 @@ def _parse_args(argv=None, namespace=None): parser.add_argument("--seq-length", type=int, default=128, help="Sequence length of input") parser.add_argument("--params-dtype", type=str, default="float32", help="Parameter dtype.") parser.add_argument( - "--fp8-init", action="store_true", default=False, help="Initialize primary weights in FP8." + "--fp8-init", + action="store_true", + default=False, + help="Initialize primary weights in FP8.", ) parser.add_argument( "--recipe", type=str, - default="mx_fp8_block_scaling", + default="MXFP8BlockScaling", help="Quantizer type.", - choices=["delayed_scaling", "current_scaling", "mx_fp8_block_scaling"], + choices=[ + "DelayedScaling", + "Float8CurrentScaling", + "Float8BlockScaling", + "MXFP8BlockScaling", + "NVFP4BlockScaling", + ], ) parser.add_argument( "--layer-type", @@ -110,15 +114,8 @@ def get_te_layer_from_string(layer_name): return te_layer_map[layer_name.lower()] -def get_recipe_from_string(recipe, fp8_format=Format.HYBRID): - if recipe == "delayed_scaling": - return DelayedScaling(fp8_format=fp8_format, amax_history_len=16, amax_compute_algo="max") - elif recipe == "current_scaling": - return Float8CurrentScaling(fp8_format=fp8_format) - elif recipe == "mx_fp8_block_scaling": - return MXFP8BlockScaling(fp8_format=fp8_format) - else: - raise ValueError(f"Unknown quantizer type: {recipe}") +def get_recipe_from_string(recipe): + return getattr(transformer_engine.common.recipe, recipe)() def init_te_model(config): @@ -244,7 +241,7 @@ def test_fp8_fsdp2_allgather(model): module.unshard() # Make sure allgathered parameters match exactly for name, param in model.named_parameters(): - assert torch.allclose(param.dequantize(), fp32_allgathered_params[name]) + torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name]) # Revert model to original sharded state for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. @@ -278,8 +275,7 @@ def _train(args): device = torch.device(f"cuda:{LOCAL_RANK}") # FP8 Configuration - fp8_format = Format.HYBRID - fp8_recipe = get_recipe_from_string(args.recipe, fp8_format) + fp8_recipe = get_recipe_from_string(args.recipe) build_model_context_args = {} if not args.fp8_init: @@ -292,13 +288,13 @@ def _train(args): build_model_context_args["enabled"] = True build_model_context_args["recipe"] = fp8_recipe - dist_print(f"Memory before model init: {torch.cuda.memory_allocated(device)/1e6} MB") + dist_print(f"Memory before model init: {torch.cuda.memory_allocated(device) / 1e6} MB") # Create the model on the meta/cuda device as per args with build_model_context(**build_model_context_args): model, inp_shape, out_shape = init_te_model(args) dist_print( f"Memory after model init on device {args.device}:" - f" {torch.cuda.memory_allocated(device)/1e6} MB" + f" {torch.cuda.memory_allocated(device) / 1e6} MB" ) # Creating a DeviceMesh for fully_shard @@ -319,7 +315,7 @@ def _train(args): dist_print(f" Sharded parameters materialized and initialized on cuda device.") dist_print( - f"FSDP2 model in cuda, memory allocated: {torch.cuda.memory_allocated(device)/1e6} MB" + f"FSDP2 model in cuda, memory allocated: {torch.cuda.memory_allocated(device) / 1e6} MB" ) optimizer = optim.Adam(model.parameters(), lr=1e-3) @@ -327,11 +323,20 @@ def _train(args): for iteration in range(args.iter): # Zero the parameter gradients optimizer.zero_grad() - input_data = torch.randn(inp_shape).to(device) - with te.autocast(enabled=True, recipe=fp8_recipe): - output = model(input_data) - target = torch.randn(out_shape).to(device) - loss = F.mse_loss(output, target) + + input_data = torch.randn(inp_shape, device=device) + target = torch.randn(out_shape, device=device) + + # NVFP4BlockScaling requires bfloat16 inputs in both the forward and backward passes. + with ( + torch.autocast(device_type="cuda", dtype=torch.bfloat16) + if args.recipe == "NVFP4BlockScaling" + else nullcontext() + ): + with te.autocast(enabled=True, recipe=fp8_recipe): + output = model(input_data) + loss = F.mse_loss(output, target) + loss.backward() optimizer.step() dist_print(f"Iteration {iteration} completed with loss {loss.item()}") diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index e328e57758..b10f31ea07 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -3,19 +3,64 @@ # See LICENSE for license information. import os -import pytest import subprocess from pathlib import Path -import transformer_engine.pytorch as te +import pytest import torch +import transformer_engine.pytorch as te +from transformer_engine.pytorch import fp8 -fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) -mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) NUM_PROCS: int = torch.cuda.device_count() +def check_nvfp4_support(): + supported, reason = fp8.check_nvfp4_support() + if supported and torch.cuda.get_device_capability()[0] == 12: + return ( + False, + ( + "NVFP4BlockScaling is failing on SM120 with " + "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " + "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" + ), + ) + + return supported, reason + + +# Each entry: (recipe_class_name, check_fn) +_FP8_RECIPE_CONFIGS = [ + ("DelayedScaling", fp8.check_fp8_support), + ("Float8CurrentScaling", fp8.check_fp8_support), + ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), + ("MXFP8BlockScaling", fp8.check_mxfp8_support), + ("NVFP4BlockScaling", check_nvfp4_support), +] + + +def _parametrize_fp8_recipes(): + """Generate pytest.param objects with xfail marks for unsupported FP8 recipes.""" + params = [] + for name, check_fn in _FP8_RECIPE_CONFIGS: + supported, reason = check_fn() + params.append( + pytest.param( + name, + id=name, + marks=pytest.mark.xfail(condition=not supported, reason=reason), + ) + ) + return params + + +@pytest.fixture(params=_parametrize_fp8_recipes()) +def fp_recipe(request): + """Parametrized fixture providing FP8 recipe Hydra overrides for each supported TE recipe.""" + return request.param + + def _run_test(fp_init, sharding_dims, recipe, layer_type): test_path = Path(__file__).parent.resolve() / "run_fsdp2_model.py" test_cmd = ["torchrun", f"--nproc_per_node={NUM_PROCS}", str(test_path)] @@ -32,28 +77,155 @@ def _run_test(fp_init, sharding_dims, recipe, layer_type): test_cmd += ["--recipe", recipe] test_cmd += ["--layer-type", layer_type] - result = subprocess.run(test_cmd, env=os.environ, check=True) + subprocess.run(test_cmd, env=os.environ, check=True) -@pytest.mark.skipif(NUM_PROCS < 4, reason="Requires 4+ GPUs") @pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") @pytest.mark.parametrize("sharding_dims", ([NUM_PROCS], [2, NUM_PROCS // 2])) @pytest.mark.parametrize("fp8_init", (False, True)) -@pytest.mark.parametrize("recipe", ("delayed_scaling", "current_scaling", "mx_fp8_block_scaling")) @pytest.mark.parametrize("layer_type", ("LayerNormLinear", "TransformerLayer")) -def test_distributed(fp8_init, sharding_dims, recipe, layer_type): - - # Skip invalid configurations - if torch.cuda.device_count() < 4: - pytest.skip("FSDP2 test requires at least 4 GPUs") - - if recipe == "mx_fp8_block_scaling" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) - elif not fp8_available: - pytest.skip(reason_for_no_fp8) - - _run_test(fp8_init, sharding_dims, recipe, layer_type) +def test_distributed(fp8_init, sharding_dims, fp_recipe, layer_type): + + if fp_recipe in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: + pytest.xfail(f"{fp_recipe} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") + + _run_test(fp8_init, sharding_dims, fp_recipe, layer_type) + + +## ── FusedAdam + FSDP2 tests ───────────────────────────────────────── + + +def _run_fused_adam_test(test_name, recipe="delayed_scaling"): + """Launch an FSDP2 + FusedAdam test via torchrun.""" + test_path = Path(__file__).parent.resolve() / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) # These tests only need 2 GPUs + test_cmd = [ + "torchrun", + f"--nproc_per_node={nproc}", + str(test_path), + "--test", + test_name, + "--recipe", + recipe, + ] + + subprocess.run(test_cmd, env=os.environ, check=True) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_fused_adam_fp8_master_weights(fp_recipe): + """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init.""" + if fp_recipe in ("Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{fp_recipe}: quantized_model_init and FSDP2 is not currently supported, since the " + "block tensor is dequantized before we flatten it for FSDP2." + ) + _run_fused_adam_test("fused_adam_fp8_master_weights", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_fused_adam_bf16(fp_recipe): + """FusedAdam(master_weights=True) + FSDP2 + bf16 params (no FP8).""" + _run_fused_adam_test("fused_adam_bf16", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_fused_adam_fp8_no_master(fp_recipe): + """FusedAdam(master_weights=False) + FSDP2 + FP8 params.""" + if fp_recipe == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access" + ) + _run_fused_adam_test("fused_adam_fp8_no_master", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_fused_adam_bf16_store_param_remainders(fp_recipe): + """FusedAdam(master_weights=True, store_param_remainders=True) + FSDP2 + bf16.""" + _run_fused_adam_test("fused_adam_bf16_store_param_remainders", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_dcp_output_parity(fp_recipe): + """DCP save/load round-trip into a fresh model produces identical outputs.""" + if fp_recipe == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access" + ) + + if fp_recipe == "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 " + "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " + "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " + "requires using power of two scaling factors." + ) + + _run_fused_adam_test("dcp_output_parity", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_dcp_output_parity_async(fp_recipe): + """DCP save/load round-trip into a fresh model produces identical outputs.""" + if fp_recipe in ("DelayedScaling", "Float8CurrentScaling"): + pytest.xfail( + f"async DCP save/load with {fp_recipe} uses StateDictStager._offload_tensor() which " + "tries to deep-copy the tensor's underlying storage. Float8Tensor is a wrapper subclass" + "(_make_wrapper_subclass) with data_ptr() == 0 (empty storage). The staging code at " + "line 215 skips the storage copy for wrapper subclasses, creating a plain tensor with " + "uninitialized garbage data. The actual FP8 data (in _data, _scale_inv attributes) is " + "deep-copied but ignored by DCP when writing." + ) + + if fp_recipe == "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" + ) + + if fp_recipe == "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 " + "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " + "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " + "requires using power of two scaling factors." + ) + + _run_fused_adam_test("dcp_output_parity_async", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_safetensors_fp32_export(fp_recipe): + """Export FP32 model from optimizer master weights to safetensors.""" + if fp_recipe == "MXFP8BlockScaling": + pytest.xfail( + "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " + "MXFP8 quantized tensors, causing illegal memory access" + ) + _run_fused_adam_test("safetensors_fp32_export", fp_recipe) + + +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.xfail( + reason=( + "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " + "autograd Function.apply unwraps DTensors to local tensors, so " + "main_grad (set on the DTensor) is inaccessible during backward. " + "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." + ), + raises=subprocess.CalledProcessError, + strict=True, +) +def test_fsdp2_fuse_wgrad_accumulation(fp_recipe): + """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail.""" + _run_fused_adam_test("fuse_wgrad_accumulation", fp_recipe) def test_dummy() -> None: diff --git a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py b/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py index 5bdc537e4b..8580cf4a33 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py @@ -218,6 +218,12 @@ def __init__( self.with_amax_reduction = False self.amax_reduction_group = None + def __getstate__(self): + """Exclude unpicklable process group from serialized state.""" + state = self.__dict__.copy() + state["amax_reduction_group"] = None + return state + @property def custom(self) -> bool: """Flag to indicate this quantizer is custom.""" diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index a87d968334..46b038d922 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -394,8 +394,14 @@ def _initialize_state( store_param_remainders (bool): Store only trailing remainder bits. """ dtype = self.name_to_dtype_map[state_name] + # Extract local tensor from DTensor (e.g. from FSDP2) to avoid + # QuantizedTensor.__torch_dispatch__ ignoring the dtype kwarg in + # torch.empty_like, and to ensure optimizer states are plain tensors. + local_param = param._local_tensor if isinstance(param, DTensor) else param # Handle QuantizedTensor by dequantizing first - param_for_empty = param.dequantize() if isinstance(param, QuantizedTensor) else param + param_for_empty = ( + local_param.dequantize() if isinstance(local_param, QuantizedTensor) else local_param + ) if store_param_remainders: data = torch.zeros_like(param_for_empty, dtype=torch.int16) else: @@ -440,7 +446,14 @@ def initialize_state(self, param, store_param_remainders): store_param_remainders=store_param_remainders, ) if not store_param_remainders: - self.set_scaled_state(param, "master_param", param.clone().detach().float()) + # Extract local tensor from DTensor and dequantize QuantizedTensor + # to get a plain float32 copy for the master weight. + local_param = param._local_tensor if isinstance(param, DTensor) else param + if isinstance(local_param, QuantizedTensor): + master = local_param.dequantize(dtype=torch.float32).clone().detach() + else: + master = local_param.clone().detach().float() + self.set_scaled_state(param, "master_param", master) def state_dict(self): """Override the state_dict() of pytorch. Before returning the state_dict, cast all diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index a3d49ea4e9..e65730c015 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Iterable import math +import warnings from typing import Any, Optional, Tuple, Union import torch @@ -622,19 +623,27 @@ def forward( if tensor._is_2D_scaled: # For the case of 2D scaled tensor, the last 2 dimensions should not change if shape[-1] != ctx.shape[-1] or shape[-2] != ctx.shape[-2]: - raise RuntimeError( + warnings.warn( "2D scaled Float8BlockwiseQTensor does not support view " "the last 2 dimensions " - f"(attempted to view dims={tuple(tensor.shape)} to {tuple(shape)})" + f"(attempted to view dims={tuple(tensor.shape)} to {tuple(shape)}). " + "If you are using this for FSDP2 without compiled_autograd_enabled, " + "then ignore this warning since this view is not going to be used anywhere.", + stacklevel=2, ) + return tensor.dequantize().view(*shape) else: # For the case of 1D scaled tensor, the last dimension should not change if shape[-1] != ctx.shape[-1]: - raise RuntimeError( + warnings.warn( "1D scaled Float8BlockwiseQTensor does not support view " "the last dimension " - f"(attempted to view dims={tuple(tensor.shape)} to {tuple(shape)})" + f"(attempted to view dims={tuple(tensor.shape)} to {tuple(shape)}). " + "If you are using this for FSDP2 without compiled_autograd_enabled, " + "then ignore this warning since this view is not going to be used anywhere.", + stacklevel=2, ) + return tensor.dequantize().view(*shape) if list(shape) == list(tensor.shape): return tensor @@ -729,19 +738,27 @@ def forward( if tensor._is_2D_scaled: # For the case of 2D scaled tensor, the last 2 dimensions should not change if shape[-1] != ctx.shape[-1] or shape[-2] != ctx.shape[-2]: - raise RuntimeError( + warnings.warn( "2D scaled Float8BlockwiseQTensor does not support reshaping " "the last 2 dimensions " - f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)})" + f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)}). " + "If you are using this for FSDP2 without compiled_autograd_enabled, " + "then ignore this warning since this view is not going to be used anywhere.", + stacklevel=2, ) + return tensor.dequantize().reshape(*shape) else: # For the case of 1D scaled tensor, the last dimension should not change if shape[-1] != ctx.shape[-1]: - raise RuntimeError( + warnings.warn( "1D scaled Float8BlockwiseQTensor does not support reshaping " "the last dimension " - f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)})" + f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)}). " + "If you are using this for FSDP2 without compiled_autograd_enabled, " + "then ignore this warning since this view is not going to be used anywhere.", + stacklevel=2, ) + return tensor.dequantize().reshape(*shape) if list(shape) == list(tensor.shape): return tensor diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index f66e88740f..c60bb2308d 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -284,6 +284,12 @@ def __init__( self.force_pow_2_scales = force_pow_2_scales self.amax_epsilon = amax_epsilon + def __getstate__(self): + """Exclude unpicklable process group from serialized state.""" + state = self.__dict__.copy() + state["amax_reduction_group"] = None + return state + def copy(self) -> Float8CurrentScalingQuantizer: """Create shallow copy""" diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index a8148b5752..4314fd248c 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -6,6 +6,7 @@ from __future__ import annotations from collections.abc import Iterable import math +import warnings from typing import Dict, Optional, Tuple, Union import functools @@ -157,6 +158,12 @@ def __init__( ) self.rht_matrix = get_rht_matrix(with_random_sign_mask, torch.cuda.current_device()) + def __getstate__(self): + """Exclude unpicklable process group from serialized state.""" + state = self.__dict__.copy() + state["amax_reduction_group"] = None + return state + def update_quantized( self, src: torch.Tensor, @@ -785,10 +792,14 @@ def forward( shape[i] = d_inferred break if shape[-1] != cur_shape[-1]: - raise RuntimeError( + warnings.warn( "NVFP4Tensor does not support reshaping inner dimension " - f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)})" + f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)}). " + "If you are using this for FSDP2 without compiled_autograd_enabled, " + "then ignore this warning since this view is not going to be used anywhere.", + stacklevel=2, ) + return tensor.dequantize().view(*shape) # Reshape data new_rowwise_data = None @@ -907,10 +918,14 @@ def forward( shape[i] = d_inferred break if shape[-1] != cur_shape[-1]: - raise RuntimeError( + warnings.warn( "NVFP4Tensor does not support reshaping inner dimension " - f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)})" + f"(attempted to reshape dims={tuple(tensor.shape)} to {tuple(shape)}). " + "If you are using this for FSDP2 without compiled_autograd_enabled, " + "then ignore this warning since this view is not going to be used anywhere.", + stacklevel=2, ) + return tensor.dequantize().reshape(*shape) # Reshape data new_rowwise_data = None From 56c2fa64b2c7c59f6a5679fb65aa25f54696d2ee Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 4 Mar 2026 13:58:50 -0800 Subject: [PATCH 245/521] [JAX] Support calling MOE router kernels from JAX side (#2711) * initial draft Signed-off-by: tdophung * tests should pass Signed-off-by: tdophung * fix lint Signed-off-by: tdophung * Address comments, minus the returning scalar request to reduce squeeze op Signed-off-by: tdophung * add notImplemented infer_sharding_from_operands and partition back in to make basePrimitive class happy Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * properly merge the jax top level APIs of score_for_moe_aux_loss with topk_and_score Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint + import issues Signed-off-by: tdophung * address Phuong's comments Signed-off-by: tdophung * remove infer_sharding_from_operand Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address jeremy's comments Signed-off-by: tdophung * address greptile comments Signed-off-by: tdophung * rename num_row num_cols Signed-off-by: tdophung * address greptile comments Signed-off-by: tdophung * rename remaining num_rows, num_cols to meaning names in jax Signed-off-by: tdophung --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_distributed_router.py | 454 ++++++++++++ tests/jax/test_fused_router.py | 539 ++++++++++++++ .../jax/cpp_extensions/__init__.py | 1 + .../jax/cpp_extensions/router.py | 678 ++++++++++++++++++ transformer_engine/jax/csrc/extensions.h | 7 + transformer_engine/jax/csrc/extensions/misc.h | 5 + .../jax/csrc/extensions/pybind.cpp | 13 + .../jax/csrc/extensions/router.cpp | 237 ++++++ transformer_engine/jax/router.py | 318 ++++++++ 9 files changed, 2252 insertions(+) create mode 100644 tests/jax/test_distributed_router.py create mode 100644 tests/jax/test_fused_router.py create mode 100644 transformer_engine/jax/cpp_extensions/router.py create mode 100644 transformer_engine/jax/csrc/extensions/router.cpp create mode 100644 transformer_engine/jax/router.py diff --git a/tests/jax/test_distributed_router.py b/tests/jax/test_distributed_router.py new file mode 100644 index 0000000000..1b3fe14e75 --- /dev/null +++ b/tests/jax/test_distributed_router.py @@ -0,0 +1,454 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for distributed/sharded execution of fused MoE router primitives. + +Testing Strategy: +================= +Router operations process each token independently (1 warp per token), so +sharded execution on the token dimension should produce identical results +to processing each shard independently with the reference implementation. + +For fused_topk_with_score_function (including compute_aux_scores mode): +- Input logits [num_tokens, num_experts] are sharded on num_tokens (DP axis) +- Expert dimension is replicated +- Each GPU processes its local tokens independently +- We verify sharded output matches per-shard reference, concatenated + +For fused_moe_aux_loss: +- This is a global reduction to a scalar +- All inputs and outputs are replicated (partition function forces this) +- We verify the op works correctly under a mesh context + +These tests exercise: batcher and shardy_sharding_rule from the router primitives. +""" + +import pytest + +import jax +import jax.numpy as jnp +import numpy as np +from jax.sharding import Mesh, NamedSharding, PartitionSpec + +from distributed_test_base import generate_configs +from utils import assert_allclose, pytest_parametrize_wrapper + +from transformer_engine.jax.router import ( + fused_topk_with_score_function, + fused_moe_aux_loss, +) + +jax.config.update("jax_use_shardy_partitioner", True) + +from test_fused_router import ( + reference_topk_softmax_sigmoid, + reference_compute_scores_for_aux_loss, + reference_aux_loss, + make_logits, +) + +# (num_tokens, num_experts, topk) +ALL_TOPK_CASES = [ + (128, 32, 4), + (2048, 128, 8), +] +TOPK_CASES = { + "L0": ALL_TOPK_CASES[0:1], + "L2": ALL_TOPK_CASES, +} + +ALL_AUX_LOSS_CASES = [ + (128, 32, 4), + (2048, 128, 4), +] +AUX_LOSS_CASES = { + "L0": ALL_AUX_LOSS_CASES[0:1], + "L2": ALL_AUX_LOSS_CASES, +} + + +class TestDistributedFusedTopk: + """Test distributed execution of fused_topk_with_score_function. + + Shards logits on the token dimension. Each GPU independently runs the + fused kernel on its local tokens. We compare against the reference + implementation run per-shard and concatenated. + """ + + def _impl_test( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + score_function, + ): + logits = make_logits(num_tokens, num_experts, score_function) + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + dp_axis = mesh_resource.dp_resource + sharded_pspec = PartitionSpec(dp_axis, None) + num_dp_devices = mesh.shape[dp_axis] if dp_axis else 1 + local_num_tokens = num_tokens // num_dp_devices + + with mesh: + logits_sharding = NamedSharding(mesh, sharded_pspec) + logits_sharded = jax.device_put(logits, logits_sharding) + + # === Forward === + @jax.jit + def target_fwd(x): + return fused_topk_with_score_function( + x, + topk=topk, + score_function=score_function, + ) + + target_probs, target_routing_map = target_fwd(logits_sharded) + + logits_shards = jnp.reshape(logits, (num_dp_devices, local_num_tokens, num_experts)) + ref_fwd_fn = jax.jit( + lambda x: reference_topk_softmax_sigmoid( + x, + topk=topk, + score_function=score_function, + ) + ) + ref_probs_list = [] + ref_routing_list = [] + for i in range(num_dp_devices): + p, rm = ref_fwd_fn(logits_shards[i]) + ref_probs_list.append(p) + ref_routing_list.append(rm) + + ref_probs = jnp.concatenate(ref_probs_list, axis=0) + ref_routing = jnp.concatenate(ref_routing_list, axis=0) + + assert_allclose( + jax.device_get(target_probs), + ref_probs, + dtype=jnp.float32, + ) + assert jnp.array_equal( + jax.device_get(target_routing_map), + ref_routing, + ), "Routing map mismatch in distributed fused_topk" + + # === Backward === + def target_loss(x): + p, _ = fused_topk_with_score_function( + x, + topk=topk, + score_function=score_function, + ) + return jnp.sum(p) + + def ref_chunk_loss(x_chunk): + p, _ = reference_topk_softmax_sigmoid( + x_chunk, + topk=topk, + score_function=score_function, + ) + return jnp.sum(p) + + target_grad = jax.jit(jax.grad(target_loss))(logits_sharded) + + ref_grads = [] + ref_chunk_grad_fn = jax.jit(jax.grad(ref_chunk_loss)) + for i in range(num_dp_devices): + ref_grads.append(ref_chunk_grad_fn(logits_shards[i])) + ref_grad = jnp.concatenate(ref_grads, axis=0) + + assert_allclose( + jax.device_get(target_grad), + ref_grad, + dtype=jnp.float32, + ) + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + TOPK_CASES, + ) + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid"]) + def test_distributed_topk( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + score_function, + ): + self._impl_test( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + score_function, + ) + + +class TestDistributedScoreForAuxLoss: + """Test distributed execution of fused_topk_with_score_function with compute_aux_scores=True. + + Same sharding strategy as fused_topk: shard on token dim, replicate experts. + Each GPU independently computes scores and routing map for its local tokens. + """ + + def _impl_test( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + score_function, + ): + logits = make_logits(num_tokens, num_experts, score_function) + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + dp_axis = mesh_resource.dp_resource + sharded_pspec = PartitionSpec(dp_axis, None) + num_dp_devices = mesh.shape[dp_axis] if dp_axis else 1 + local_num_tokens = num_tokens // num_dp_devices + + with mesh: + logits_sharding = NamedSharding(mesh, sharded_pspec) + logits_sharded = jax.device_put(logits, logits_sharding) + + # === Forward === + @jax.jit + def target_fwd(x): + return fused_topk_with_score_function( + x, + topk=topk, + score_function=score_function, + compute_aux_scores=True, + ) + + target_scores, target_routing_map = target_fwd(logits_sharded) + + logits_shards = jnp.reshape(logits, (num_dp_devices, local_num_tokens, num_experts)) + ref_fwd_fn = jax.jit( + lambda x: reference_compute_scores_for_aux_loss( + x, + topk=topk, + score_function=score_function, + ) + ) + ref_routing_list = [] + ref_scores_list = [] + for i in range(num_dp_devices): + rm, s = ref_fwd_fn(logits_shards[i]) + ref_routing_list.append(rm) + ref_scores_list.append(s) + + ref_routing = jnp.concatenate(ref_routing_list, axis=0) + ref_scores = jnp.concatenate(ref_scores_list, axis=0) + + assert_allclose( + jax.device_get(target_scores), + ref_scores, + dtype=jnp.float32, + ) + assert jnp.array_equal( + jax.device_get(target_routing_map), + ref_routing, + ), "Routing map mismatch in distributed score_for_aux_loss" + + # === Backward === + def target_loss(x): + s, _ = fused_topk_with_score_function( + x, + topk=topk, + score_function=score_function, + compute_aux_scores=True, + ) + return jnp.sum(s) + + def ref_chunk_loss(x_chunk): + _, s = reference_compute_scores_for_aux_loss( + x_chunk, + topk=topk, + score_function=score_function, + ) + return jnp.sum(s) + + target_grad = jax.jit(jax.grad(target_loss))(logits_sharded) + + ref_grads = [] + ref_chunk_grad_fn = jax.jit(jax.grad(ref_chunk_loss)) + for i in range(num_dp_devices): + ref_grads.append(ref_chunk_grad_fn(logits_shards[i])) + ref_grad = jnp.concatenate(ref_grads, axis=0) + + assert_allclose( + jax.device_get(target_grad), + ref_grad, + dtype=jnp.float32, + ) + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + TOPK_CASES, + ) + @pytest.mark.parametrize("score_function", ["softmax", "sigmoid"]) + def test_distributed_score_for_aux_loss( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + score_function, + ): + self._impl_test( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + score_function, + ) + + +class TestDistributedMoEAuxLoss: + """Test distributed execution of fused_moe_aux_loss. + + Aux loss is a global reduction to a scalar. The partition function forces + all inputs to be replicated. We verify the op produces correct results + under a mesh context with replicated sharding, testing both forward + (scalar loss) and backward (gradient w.r.t. probs). + """ + + def _impl_test( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + ): + key = jax.random.PRNGKey(42) + _, subkey1, _ = jax.random.split(key, 3) + + offset = jnp.arange(-num_tokens // 2, num_tokens // 2, dtype=jnp.float32) * 1e-4 + probs = jnp.arange(-num_experts // 2, num_experts // 2, dtype=jnp.float32) * 1e-2 + probs = probs[None, :].repeat(num_tokens, axis=0) + offset[:, None] + + tokens_per_expert = jax.random.randint(subkey1, (num_experts,), 1, 1000).astype(jnp.int32) + coeff = 0.01 + + devices = np.asarray(jax.devices()[:device_count]).reshape(*mesh_shape) + mesh = Mesh(devices, mesh_axes) + + replicated_2d_pspec = PartitionSpec(None, None) + replicated_1d_pspec = PartitionSpec(None) + + with mesh: + probs_sharding = NamedSharding(mesh, replicated_2d_pspec) + tpe_sharding = NamedSharding(mesh, replicated_1d_pspec) + + probs_dev = jax.device_put(probs, probs_sharding) + tpe_dev = jax.device_put(tokens_per_expert, tpe_sharding) + + # === Forward === + @jax.jit + def target_fwd(p, tpe): + return fused_moe_aux_loss(p, tpe, topk=topk, coeff=coeff) + + target_loss = target_fwd(probs_dev, tpe_dev) + + ref_fwd_fn = jax.jit( + lambda p: reference_aux_loss( + p, + tokens_per_expert, + num_tokens, + topk, + num_experts, + coeff, + ) + ) + ref_loss = ref_fwd_fn(probs) + + assert_allclose( + jax.device_get(target_loss), + ref_loss, + dtype=jnp.float32, + ) + + # === Backward === + def target_loss_fn(p): + return fused_moe_aux_loss( + p, + tokens_per_expert, + topk=topk, + coeff=coeff, + ) + + def ref_loss_fn(p): + return reference_aux_loss( + p, + tokens_per_expert, + num_tokens, + topk, + num_experts, + coeff, + ) + + target_grad = jax.jit(jax.grad(target_loss_fn))(probs_dev) + ref_grad = jax.jit(jax.grad(ref_loss_fn))(probs) + + assert_allclose( + jax.device_get(target_grad), + ref_grad, + dtype=jnp.float32, + ) + + @pytest.mark.parametrize("device_count,mesh_shape,mesh_axes,mesh_resource", generate_configs()) + @pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + AUX_LOSS_CASES, + ) + def test_distributed_aux_loss( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + ): + self._impl_test( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + num_tokens, + num_experts, + topk, + ) diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py new file mode 100644 index 0000000000..77e89457c8 --- /dev/null +++ b/tests/jax/test_fused_router.py @@ -0,0 +1,539 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for fused MoE router CUDA kernels (JAX wrappers).""" + +from functools import partial +from typing import Optional + +import jax +import jax.numpy as jnp +import pytest + +from utils import pytest_parametrize_wrapper + +from transformer_engine.jax.router import ( + fused_topk_with_score_function, + fused_moe_aux_loss, +) + +# ============================================================================= +# Test case definitions (L0 = fast smoke, L2 = comprehensive) +# ============================================================================= + +# (num_tokens, num_experts, topk) +ALL_TOPK_CASES = [ + (128, 32, 4), + (2048, 32, 4), + (2048, 128, 8), + (7168, 128, 4), + (7168, 32, 8), +] +TOPK_CASES = { + "L0": ALL_TOPK_CASES[0:2], + "L2": ALL_TOPK_CASES, +} + +ALL_GROUP_TOPK_OPTIONS = [None, 4] +GROUP_TOPK_OPTIONS = { + "L0": [None], + "L2": ALL_GROUP_TOPK_OPTIONS, +} + +ALL_SCALING_FACTOR_OPTIONS = [None, 1.2] +SCALING_FACTOR_OPTIONS = { + "L0": [None], + "L2": ALL_SCALING_FACTOR_OPTIONS, +} + +ALL_ENABLE_BIAS_OPTIONS = [True, False] +ENABLE_BIAS_OPTIONS = { + "L0": [False], + "L2": ALL_ENABLE_BIAS_OPTIONS, +} + +ALL_USE_PRE_SOFTMAX_OPTIONS = [True, False] +USE_PRE_SOFTMAX_OPTIONS = { + "L0": [False], + "L2": ALL_USE_PRE_SOFTMAX_OPTIONS, +} + +# (num_tokens, num_experts, topk) +ALL_SCORE_AUX_LOSS_CASES = [ + (128, 32, 4), + (2048, 128, 4), + (2048, 256, 8), + (7168, 128, 8), + (7168, 32, 4), +] +SCORE_AUX_LOSS_CASES = { + "L0": ALL_SCORE_AUX_LOSS_CASES[0:2], + "L2": ALL_SCORE_AUX_LOSS_CASES, +} + +ALL_SCORE_FUNCTIONS = ["softmax", "sigmoid"] +SCORE_FUNCTIONS = { + "L0": ["softmax"], + "L2": ALL_SCORE_FUNCTIONS, +} + +# (num_tokens, num_experts, topk) +ALL_AUX_LOSS_CASES = [ + (128, 32, 4), + (2048, 128, 4), + (2048, 256, 4), + (7168, 128, 4), + (7168, 32, 4), +] +AUX_LOSS_CASES = { + "L0": ALL_AUX_LOSS_CASES[0:2], + "L2": ALL_AUX_LOSS_CASES, +} + +ALL_DTYPES = [jnp.float32] +DTYPES = { + "L0": [jnp.float32], + "L2": ALL_DTYPES, +} + +SEED = 42 + + +# ============================================================================= +# Reference Implementations +# ============================================================================= + + +def reference_group_limited_topk( + scores: jnp.ndarray, + topk: int, + num_tokens: int, + num_experts: int, + num_groups: int, + group_topk: int, +): + """Reference implementation for grouped top-k. + + Only valid when num_groups and group_topk are both positive integers. + For plain top-k without grouping, use jax.lax.top_k directly. + """ + assert num_groups is not None and num_groups > 0, ( + "reference_group_limited_topk requires valid num_groups > 0. " + "For plain top-k, use jax.lax.top_k directly." + ) + assert ( + group_topk is not None and group_topk > 0 + ), "reference_group_limited_topk requires valid group_topk > 0." + assert ( + num_experts % num_groups == 0 + ), f"num_experts ({num_experts}) must be divisible by num_groups ({num_groups})" + group_size = num_experts // num_groups + experts_per_group = topk // group_topk + + group_scores = ( + scores.reshape(num_tokens, num_groups, group_size) + .sort(axis=-1)[..., -experts_per_group:] + .sum(axis=-1) + ) + group_idx = jax.lax.top_k(group_scores, k=group_topk)[1] + group_mask = jnp.zeros_like(group_scores).at[jnp.arange(num_tokens)[:, None], group_idx].set(1) + + score_mask = (group_mask[:, :, None] * jnp.ones((num_tokens, num_groups, group_size))).reshape( + num_tokens, -1 + ) + + masked_scores = jnp.where(score_mask.astype(bool), scores, -jnp.inf) + probs, top_indices = jax.lax.top_k(masked_scores, k=topk) + return probs, top_indices + + +def reference_topk_softmax_sigmoid( + logits: jnp.ndarray, + 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[jnp.ndarray] = None, +): + """Reference implementation for topk + softmax/sigmoid.""" + num_tokens, num_experts = logits.shape + + def compute_topk(scores, topk, num_groups=None, group_topk=None): + if group_topk: + return reference_group_limited_topk( + scores=scores, + topk=topk, + num_tokens=num_tokens, + num_experts=num_experts, + num_groups=num_groups, + group_topk=group_topk, + ) + else: + return jax.lax.top_k(scores, k=topk) + + if score_function == "softmax": + if use_pre_softmax: + scores = jax.nn.softmax(logits.astype(jnp.float32), axis=-1).astype(logits.dtype) + probs, top_indices = compute_topk(scores, topk, num_groups, group_topk) + else: + scores, top_indices = compute_topk(logits, topk, num_groups, group_topk) + probs = jax.nn.softmax(scores.astype(jnp.float32), axis=-1).astype(logits.dtype) + elif score_function == "sigmoid": + scores = jax.nn.sigmoid(logits.astype(jnp.float32)).astype(logits.dtype) + 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 = jnp.take_along_axis(scores, top_indices, axis=1).astype(logits.dtype) + else: + scores, top_indices = compute_topk(scores, topk, num_groups, group_topk) + probs = scores / (scores.sum(axis=-1, keepdims=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 = ( + jnp.zeros_like(logits).at[jnp.arange(num_tokens)[:, None], top_indices].set(probs) + ) + topk_map = ( + jnp.zeros_like(logits, dtype=jnp.bool_) + .at[jnp.arange(num_tokens)[:, None], top_indices] + .set(True) + ) + + return topk_masked_gates, topk_map + + +def reference_compute_scores_for_aux_loss(logits: jnp.ndarray, topk: int, score_function: str): + """Reference implementation for computing routing scores for aux loss.""" + if score_function == "softmax": + scores = jax.nn.softmax(logits.astype(jnp.float32), axis=-1) + elif score_function == "sigmoid": + scores = jax.nn.sigmoid(logits.astype(jnp.float32)) + scores = scores / (scores.sum(axis=-1, keepdims=True) + 1e-20) if topk > 1 else scores + else: + raise ValueError(f"Invalid score_function: {score_function}") + + _, top_indices = jax.lax.top_k(scores, k=topk) + num_tokens = logits.shape[0] + routing_map = ( + jnp.zeros_like(logits, dtype=jnp.bool_) + .at[jnp.arange(num_tokens)[:, None], top_indices] + .set(True) + ) + return routing_map, scores + + +def reference_aux_loss( + probs: jnp.ndarray, + tokens_per_expert: jnp.ndarray, + total_num_tokens: int, + topk: int, + num_experts: int, + moe_aux_loss_coeff: float, +): + """Reference implementation for MoE auxiliary loss.""" + aggregated_probs_per_expert = probs.sum(axis=0) + aux_loss = jnp.sum(aggregated_probs_per_expert * tokens_per_expert) * ( + num_experts * moe_aux_loss_coeff / (topk * total_num_tokens * total_num_tokens) + ) + return aux_loss + + +# ============================================================================= +# Helper: logits generation +# ============================================================================= + + +def make_logits(num_tokens, num_experts, score_function, dtype=jnp.float32): + """Create deterministic logits for testing.""" + if score_function == "sigmoid": + offset = jnp.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype) * 1e-4 + logits = jnp.arange(-num_experts // 2, num_experts // 2, dtype=dtype) * 1e-2 + logits = logits[None, :].repeat(num_tokens, axis=0) + offset[:, None] + else: + logits = ( + jnp.arange( + -num_tokens * num_experts // 2, + num_tokens * num_experts // 2, + dtype=dtype, + ) + * 1e-4 + ) + logits = logits.reshape(num_tokens, num_experts) + return logits + + +# ============================================================================= +# Test: Fused Top-K with Score Function +# ============================================================================= + + +def run_topk_comparison( + dtype, + num_tokens, + num_experts, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + enable_bias, +): + """Compare fused vs reference top-k implementation, both jitted.""" + logits = make_logits(num_tokens, num_experts, score_function, dtype) + + if enable_bias and score_function == "sigmoid": + expert_bias = jnp.arange(num_experts, dtype=jnp.float32) * 0.1 + expert_bias = jnp.flip(expert_bias) + else: + expert_bias = None + + # Forward: reference (jitted) + ref_fwd_fn = jax.jit( + partial( + reference_topk_softmax_sigmoid, + 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, + ) + ) + probs_ref, routing_map_ref = ref_fwd_fn(logits) + + # Forward: fused (jitted) + fused_fwd_fn = jax.jit( + partial( + fused_topk_with_score_function, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups if num_groups else -1, + group_topk=group_topk if group_topk else -1, + scaling_factor=scaling_factor if scaling_factor else 1.0, + score_function=score_function, + expert_bias=expert_bias, + ) + ) + probs_fused, routing_map_fused = fused_fwd_fn(logits) + + assert jnp.allclose( + probs_ref, probs_fused, atol=1e-5, rtol=1e-5 + ), f"Probs mismatch: max diff = {jnp.abs(probs_ref - probs_fused).max()}" + assert jnp.array_equal(routing_map_ref, routing_map_fused), "Routing map mismatch" + + # Backward: reference (jitted) + def loss_ref(logits_): + p, _ = reference_topk_softmax_sigmoid( + logits_, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + ) + return p.sum() + + def loss_fused(logits_): + p, _ = fused_topk_with_score_function( + logits_, + topk, + use_pre_softmax, + num_groups if num_groups else -1, + group_topk if group_topk else -1, + scaling_factor if scaling_factor else 1.0, + score_function, + expert_bias, + ) + return p.sum() + + grad_ref = jax.jit(jax.grad(loss_ref))(logits) + grad_fused = jax.jit(jax.grad(loss_fused))(logits) + assert jnp.allclose( + grad_ref, grad_fused, atol=1e-5, rtol=1e-5 + ), f"Grad mismatch: max diff = {jnp.abs(grad_ref - grad_fused).max()}" + + +@pytest_parametrize_wrapper("dtype", DTYPES) +@pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + TOPK_CASES, +) +@pytest_parametrize_wrapper("group_topk", GROUP_TOPK_OPTIONS) +@pytest_parametrize_wrapper("scaling_factor", SCALING_FACTOR_OPTIONS) +@pytest_parametrize_wrapper("enable_bias", ENABLE_BIAS_OPTIONS) +def test_topk_sigmoid( + dtype, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias +): + num_groups = 8 if group_topk else None + run_topk_comparison( + dtype=dtype, + num_tokens=num_tokens, + num_experts=num_experts, + topk=topk, + use_pre_softmax=False, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function="sigmoid", + enable_bias=enable_bias, + ) + + +@pytest_parametrize_wrapper("dtype", DTYPES) +@pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + TOPK_CASES, +) +@pytest_parametrize_wrapper("use_pre_softmax", USE_PRE_SOFTMAX_OPTIONS) +@pytest_parametrize_wrapper("group_topk", GROUP_TOPK_OPTIONS) +@pytest_parametrize_wrapper("scaling_factor", SCALING_FACTOR_OPTIONS) +def test_topk_softmax( + dtype, num_tokens, num_experts, topk, use_pre_softmax, group_topk, scaling_factor +): + num_groups = 8 if group_topk else None + run_topk_comparison( + dtype=dtype, + num_tokens=num_tokens, + num_experts=num_experts, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function="softmax", + enable_bias=False, + ) + + +# ============================================================================= +# Test: Fused Score for MoE Aux Loss +# ============================================================================= + + +@pytest_parametrize_wrapper("dtype", DTYPES) +@pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + SCORE_AUX_LOSS_CASES, +) +@pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS) +def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function): + logits = make_logits(num_tokens, num_experts, score_function, dtype) + + # Forward: reference (jitted) + ref_fwd_fn = jax.jit( + partial( + reference_compute_scores_for_aux_loss, + topk=topk, + score_function=score_function, + ) + ) + routing_map_ref, scores_ref = ref_fwd_fn(logits) + + # Forward: fused (jitted) + fused_fwd_fn = jax.jit( + partial( + fused_topk_with_score_function, + topk=topk, + score_function=score_function, + compute_aux_scores=True, + ) + ) + scores_fused, routing_map_fused = fused_fwd_fn(logits) + + assert jnp.allclose( + scores_ref, scores_fused, atol=1e-5, rtol=1e-5 + ), f"Scores mismatch: max diff = {jnp.abs(scores_ref - scores_fused).max()}" + assert jnp.array_equal(routing_map_ref, routing_map_fused), "Routing map mismatch" + + # Backward (jitted) + def loss_ref(logits_): + _, s = reference_compute_scores_for_aux_loss(logits_, topk, score_function) + return s.sum() + + def loss_fused(logits_): + s, _ = fused_topk_with_score_function( + logits_, + topk, + score_function=score_function, + compute_aux_scores=True, + ) + return s.sum() + + grad_ref = jax.jit(jax.grad(loss_ref))(logits) + grad_fused = jax.jit(jax.grad(loss_fused))(logits) + assert jnp.allclose( + grad_ref, grad_fused, atol=1e-5, rtol=1e-5 + ), f"Grad mismatch: max diff = {jnp.abs(grad_ref - grad_fused).max()}" + + +# ============================================================================= +# Test: Fused MoE Aux Loss +# ============================================================================= + + +@pytest_parametrize_wrapper("dtype", DTYPES) +@pytest_parametrize_wrapper( + "num_tokens,num_experts,topk", + AUX_LOSS_CASES, +) +def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk): + key = jax.random.PRNGKey(SEED) + + offset = jnp.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype) * 1e-4 + probs = jnp.arange(-num_experts // 2, num_experts // 2, dtype=dtype) * 1e-2 + probs = probs[None, :].repeat(num_tokens, axis=0) + offset[:, None] + probs = probs.reshape(num_tokens, num_experts) + + tokens_per_expert = jax.random.randint(key, (num_experts,), 1, 1000).astype(jnp.int32) + coeff = 0.01 + + # Forward: reference (jitted) + ref_fwd_fn = jax.jit( + partial( + reference_aux_loss, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + topk=topk, + num_experts=num_experts, + moe_aux_loss_coeff=coeff, + ) + ) + aux_loss_ref = ref_fwd_fn(probs) + + # Forward: fused (jitted) + fused_fwd_fn = jax.jit( + partial( + fused_moe_aux_loss, + tokens_per_expert=tokens_per_expert, + topk=topk, + coeff=coeff, + ) + ) + aux_loss_fused = fused_fwd_fn(probs) + + assert jnp.allclose( + aux_loss_ref, aux_loss_fused, atol=1e-5, rtol=1e-5 + ), f"Aux loss mismatch: ref={aux_loss_ref}, fused={aux_loss_fused}" + + # Backward (jitted) + def loss_ref_fn(probs_): + return reference_aux_loss(probs_, tokens_per_expert, num_tokens, topk, num_experts, coeff) + + def loss_fused_fn(probs_): + return fused_moe_aux_loss(probs_, tokens_per_expert, topk, coeff) + + grad_ref = jax.jit(jax.grad(loss_ref_fn))(probs) + grad_fused = jax.jit(jax.grad(loss_fused_fn))(probs) + assert jnp.allclose( + grad_ref, grad_fused, atol=1e-5, rtol=1e-5 + ), f"Grad mismatch: max diff = {jnp.abs(grad_ref - grad_fused).max()}" diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py index 6a2f9b7378..d203fcea9d 100644 --- a/transformer_engine/jax/cpp_extensions/__init__.py +++ b/transformer_engine/jax/cpp_extensions/__init__.py @@ -9,3 +9,4 @@ from .quantization import * from .softmax import * from .gemm import * +from .router import * diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py new file mode 100644 index 0000000000..1fce6d2fd7 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -0,0 +1,678 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""JAX/TE custom ops for fused MoE router""" +from enum import IntEnum + +import jax.numpy as jnp +from jax import dtypes, ffi +from jax.sharding import NamedSharding, PartitionSpec +from transformer_engine_jax import JAXX_Score_Function + +from .base import BasePrimitive, register_primitive +from .misc import get_padded_spec + +__all__ = [ + "ScoreFunction", + "fused_topk_with_score_function_fwd", + "fused_topk_with_score_function_bwd", + "fused_moe_aux_loss_fwd", + "fused_moe_aux_loss_bwd", +] + + +class ScoreFunction(IntEnum): + """Score function enum for fused MoE router kernels, synced with C++ JAXX_Score_Function.""" + + SIGMOID = int(JAXX_Score_Function.SIGMOID) + SOFTMAX = int(JAXX_Score_Function.SOFTMAX) + + +# =========================================== ================================== +# Fused Top-K with Score Function - Forward +# ============================================================================= + + +class FusedTopkWithScoreFunctionFwdPrimitive(BasePrimitive): + """ + Fused Top-K with Score Function Forward Primitive. + Computes score_function(logits) -> top-k -> probs, routing_map. + When compute_aux_scores=1, instead computes clean scores for aux loss. + """ + + name = "te_fused_topk_with_score_function_forward_ffi" + multiple_results = True + impl_static_args = ( + 2, + 3, + 4, + 5, + 6, + 7, + 8, + ) # topk, use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, compute_aux_scores + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + logits_aval, + expert_bias_aval, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ): + """Abstract evaluation: describe output shapes and dtypes.""" + del expert_bias_aval, topk, use_pre_softmax, num_groups, group_topk + del scaling_factor, score_function, compute_aux_scores + i_dtype = dtypes.canonicalize_dtype(logits_aval.dtype) + i_shape = logits_aval.shape + probs_aval = logits_aval.update(shape=i_shape, dtype=i_dtype) + routing_map_aval = logits_aval.update(shape=i_shape, dtype=jnp.bool_) + intermediate_aval = logits_aval.update(shape=i_shape, dtype=i_dtype) + return probs_aval, routing_map_aval, intermediate_aval + + @staticmethod + def lowering( + ctx, + logits, + expert_bias, + *, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ): + return ffi.ffi_lowering(FusedTopkWithScoreFunctionFwdPrimitive.name)( + ctx, + logits, + expert_bias, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=compute_aux_scores, + ) + + @staticmethod + def impl( + logits, + expert_bias, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ): + assert FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive is not None + return FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive.bind( + logits, + expert_bias, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=compute_aux_scores, + ) + + @staticmethod + def batcher( + batched_args, + batch_dims, + *, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ): + assert FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive is not None + logits, expert_bias = batched_args + logits_bdim, _ = batch_dims + return ( + FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive.bind( + logits, + expert_bias, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=compute_aux_scores, + ), + (logits_bdim, logits_bdim, logits_bdim), + ) + + @staticmethod + def partition( + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + mesh, + arg_infos, + result_infos, + ): + del result_infos + logits_spec = get_padded_spec(arg_infos[0]) + out_sharding = NamedSharding(mesh, PartitionSpec(*logits_spec)) + routing_sharding = NamedSharding(mesh, PartitionSpec(*logits_spec)) + intermediate_sharding = NamedSharding(mesh, PartitionSpec(*logits_spec)) + out_shardings = [out_sharding, routing_sharding, intermediate_sharding] + arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding) + + def sharded_impl(logits, expert_bias): + return FusedTopkWithScoreFunctionFwdPrimitive.impl( + logits, + expert_bias, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule(*args): + del args + return ( + "num_tokens num_experts, bias_dim -> num_tokens num_experts, num_tokens num_experts," + " num_tokens num_experts" + ) + + +register_primitive(FusedTopkWithScoreFunctionFwdPrimitive) + + +# ============================================================================= +# Fused Top-K with Score Function - Backward +# ============================================================================= + + +class FusedTopkWithScoreFunctionBwdPrimitive(BasePrimitive): + """ + Fused Top-K with Score Function Backward Primitive. + When compute_aux_scores=1, runs the score-for-aux-loss backward instead. + """ + + name = "te_fused_topk_with_score_function_backward_ffi" + multiple_results = False + impl_static_args = ( + 3, + 4, + 5, + 6, + 7, + ) # topk, use_pre_softmax, scaling_factor, score_function, compute_aux_scores + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + routing_map_aval, + intermediate_aval, + grad_probs_aval, + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + ): + del topk, use_pre_softmax, scaling_factor, score_function + del compute_aux_scores, routing_map_aval + return intermediate_aval.update( + shape=intermediate_aval.shape, + dtype=dtypes.canonicalize_dtype(grad_probs_aval.dtype), + ) + + @staticmethod + def lowering( + ctx, + routing_map, + intermediate, + grad_probs, + *, + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + ): + return ffi.ffi_lowering(FusedTopkWithScoreFunctionBwdPrimitive.name)( + ctx, + routing_map, + intermediate, + grad_probs, + topk=topk, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=compute_aux_scores, + ) + + @staticmethod + def impl( + routing_map, + intermediate, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + ): + assert FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive is not None + return FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive.bind( + routing_map, + intermediate, + grad_probs, + topk=topk, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=compute_aux_scores, + ) + + @staticmethod + def batcher( + batched_args, + batch_dims, + *, + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + ): + assert FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive is not None + routing_map, intermediate, grad_probs = batched_args + _, _, grad_probs_bdim = batch_dims + return ( + FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive.bind( + routing_map, + intermediate, + grad_probs, + topk=topk, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=compute_aux_scores, + ), + grad_probs_bdim, + ) + + @staticmethod + def partition( + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + mesh, + arg_infos, + result_infos, + ): + del result_infos + grad_spec = get_padded_spec(arg_infos[2]) + out_sharding = NamedSharding(mesh, PartitionSpec(*grad_spec)) + arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding, arg_infos[2].sharding) + + def sharded_impl(routing_map, intermediate, grad_probs): + return FusedTopkWithScoreFunctionBwdPrimitive.impl( + routing_map, + intermediate, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + ) + + return mesh, sharded_impl, out_sharding, arg_shardings + + @staticmethod + def shardy_sharding_rule(*args): + del args + return ( + "num_tokens num_experts, num_tokens num_experts, num_tokens num_experts -> num_tokens" + " num_experts" + ) + + +register_primitive(FusedTopkWithScoreFunctionBwdPrimitive) + + +# ============================================================================= +# Fused MoE Aux Loss - Forward +# ============================================================================= + + +class FusedMoEAuxLossFwdPrimitive(BasePrimitive): + """ + Fused MoE Aux Loss Forward Primitive. + """ + + name = "te_fused_moe_aux_loss_forward_ffi" + multiple_results = True + impl_static_args = (2, 3) # topk, coeff + inner_primitive = None + outer_primitive = None + + @staticmethod + 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) + return aux_loss_aval, const_buf_aval + + @staticmethod + def lowering(ctx, probs, tokens_per_expert, *, topk, coeff): + return ffi.ffi_lowering(FusedMoEAuxLossFwdPrimitive.name)( + ctx, + probs, + tokens_per_expert, + topk=topk, + coeff=coeff, + ) + + @staticmethod + def impl(probs, tokens_per_expert, topk, coeff): + assert FusedMoEAuxLossFwdPrimitive.inner_primitive is not None + return FusedMoEAuxLossFwdPrimitive.inner_primitive.bind( + probs, + tokens_per_expert, + topk=topk, + coeff=coeff, + ) + + @staticmethod + def batcher(batched_args, batch_dims, *, topk, coeff): + assert FusedMoEAuxLossFwdPrimitive.outer_primitive is not None + probs, tokens_per_expert = batched_args + probs_bdim, _ = batch_dims + return ( + FusedMoEAuxLossFwdPrimitive.outer_primitive.bind( + probs, + tokens_per_expert, + topk=topk, + coeff=coeff, + ), + (probs_bdim, probs_bdim), + ) + + @staticmethod + def partition(topk, coeff, mesh, arg_infos, result_infos): + del result_infos + aux_loss_sharding = NamedSharding(mesh, PartitionSpec()) + const_buf_sharding = NamedSharding(mesh, PartitionSpec(None)) + out_shardings = [aux_loss_sharding, const_buf_sharding] + arg_shardings = (arg_infos[0].sharding, arg_infos[1].sharding) + + def sharded_impl(probs, tokens_per_expert): + return FusedMoEAuxLossFwdPrimitive.impl( + probs, + tokens_per_expert, + topk, + coeff, + ) + + return mesh, sharded_impl, out_shardings, arg_shardings + + @staticmethod + def shardy_sharding_rule(*args): + del args + return "num_tokens num_experts, num_experts -> , const_buf_one" + + +register_primitive(FusedMoEAuxLossFwdPrimitive) + + +# ============================================================================= +# Fused MoE Aux Loss - Backward +# ============================================================================= + + +class FusedMoEAuxLossBwdPrimitive(BasePrimitive): + """ + Fused MoE Aux Loss Backward Primitive. + """ + + name = "te_fused_moe_aux_loss_backward_ffi" + multiple_results = False + impl_static_args = (3,) # num_tokens + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract(const_buf_aval, tokens_per_expert_aval, grad_aux_loss_aval, num_tokens): + del const_buf_aval + num_experts = tokens_per_expert_aval.shape[0] + out_dtype = dtypes.canonicalize_dtype(grad_aux_loss_aval.dtype) + return grad_aux_loss_aval.update( + shape=(num_tokens, num_experts), + dtype=out_dtype, + ) + + @staticmethod + def lowering(ctx, const_buf, tokens_per_expert, grad_aux_loss, *, num_tokens): + del num_tokens + return ffi.ffi_lowering(FusedMoEAuxLossBwdPrimitive.name)( + ctx, + const_buf, + tokens_per_expert, + grad_aux_loss, + ) + + @staticmethod + def impl(const_buf, tokens_per_expert, grad_aux_loss, num_tokens): + assert FusedMoEAuxLossBwdPrimitive.inner_primitive is not None + return FusedMoEAuxLossBwdPrimitive.inner_primitive.bind( + const_buf, + tokens_per_expert, + grad_aux_loss, + num_tokens=num_tokens, + ) + + @staticmethod + def batcher(batched_args, batch_dims, *, num_tokens): + assert FusedMoEAuxLossBwdPrimitive.outer_primitive is not None + const_buf, tokens_per_expert, grad_aux_loss = batched_args + _, _, grad_bdim = batch_dims + return ( + FusedMoEAuxLossBwdPrimitive.outer_primitive.bind( + const_buf, + tokens_per_expert, + grad_aux_loss, + num_tokens=num_tokens, + ), + grad_bdim, + ) + + @staticmethod + def partition( + num_tokens, + mesh, + arg_infos, + result_infos, + ): + del result_infos + out_sharding = NamedSharding(mesh, PartitionSpec(None, None)) + arg_shardings = ( + arg_infos[0].sharding, + arg_infos[1].sharding, + arg_infos[2].sharding, + ) + + def sharded_impl(const_buf, tokens_per_expert, grad_aux_loss): + return FusedMoEAuxLossBwdPrimitive.impl( + const_buf, + tokens_per_expert, + grad_aux_loss, + num_tokens, + ) + + return mesh, sharded_impl, out_sharding, arg_shardings + + @staticmethod + def shardy_sharding_rule(*args): + del args + # num_tokens only appears in the output (not in any input) because the + # backward reconstructs the full [num_tokens, num_experts] grad_probs from + # scalar inputs. Shardy will leave num_tokens unsharded, which matches the + # replicated PartitionSpec(None, None) in partition(). + return "const_buf_one, num_experts, grad_one -> i num_experts" + + +register_primitive(FusedMoEAuxLossBwdPrimitive) + + +# ============================================================================= +# Public API functions +# ============================================================================= + + +def fused_topk_with_score_function_fwd( + logits: jnp.ndarray, + topk: int, + use_pre_softmax: bool, + num_groups: int, + group_topk: int, + scaling_factor: float, + score_function, + expert_bias: jnp.ndarray, + compute_aux_scores: bool = False, +): + """ + Fused top-k with score function forward pass. + + When compute_aux_scores=True, runs the clean score-for-aux-loss kernel + instead of the full top-k kernel (expert_bias, use_pre_softmax, num_groups, + group_topk, and scaling_factor are ignored). + + Parameters + ---------- + logits : jnp.ndarray + [num_tokens, num_experts] logits from gating GEMM. + topk : int + Number of top experts to select. + use_pre_softmax : bool + If True, apply softmax before top-k. + num_groups : int + Number of groups for grouped top-k (1 to disable). + group_topk : int + Top-k at group level (1 to disable). + scaling_factor : float + Scaling factor for output probs. + score_function : ScoreFunction + ScoreFunction.SOFTMAX or ScoreFunction.SIGMOID. + expert_bias : jnp.ndarray + Expert bias (only used with sigmoid). Pass empty array if unused. + compute_aux_scores : bool + If True, compute clean scores for aux loss instead of full top-k. + + Returns + ------- + probs_or_scores, routing_map, saved_scores + """ + return FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive.bind( + logits, + expert_bias, + topk=int(topk), + use_pre_softmax=int(use_pre_softmax), + num_groups=int(num_groups), + group_topk=int(group_topk), + scaling_factor=float(scaling_factor), + score_function=int(score_function), + compute_aux_scores=int(compute_aux_scores), + ) + + +def fused_topk_with_score_function_bwd( + routing_map: jnp.ndarray, + saved_scores: jnp.ndarray, + grad_probs: jnp.ndarray, + topk: int, + use_pre_softmax: bool, + scaling_factor: float, + score_function, + compute_aux_scores: bool = False, +): + """ + Fused top-k with score function backward pass. + + When compute_aux_scores=True, routing_map is ignored and the + score-for-aux-loss backward kernel is used instead. + """ + return FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive.bind( + routing_map, + saved_scores, + grad_probs, + topk=int(topk), + use_pre_softmax=int(use_pre_softmax), + scaling_factor=float(scaling_factor), + score_function=int(score_function), + compute_aux_scores=int(compute_aux_scores), + ) + + +def fused_moe_aux_loss_fwd( + probs: jnp.ndarray, + tokens_per_expert: jnp.ndarray, + topk: int, + coeff: float, +): + """ + Fused MoE aux loss forward pass. + + Returns + ------- + aux_loss, const_buf + """ + return FusedMoEAuxLossFwdPrimitive.outer_primitive.bind( + probs, + tokens_per_expert, + topk=int(topk), + coeff=float(coeff), + ) + + +def fused_moe_aux_loss_bwd( + const_buf: jnp.ndarray, + tokens_per_expert: jnp.ndarray, + grad_aux_loss: jnp.ndarray, + num_tokens: int, +): + """ + Fused MoE aux loss backward pass. + """ + return FusedMoEAuxLossBwdPrimitive.outer_primitive.bind( + const_buf, + tokens_per_expert, + grad_aux_loss, + num_tokens=int(num_tokens), + ) diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 1c0bc52b88..e844100c5a 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -152,6 +152,12 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(CudnnHandleInitHandler); // CuBLAS helpers XLA_FFI_DECLARE_HANDLER_SYMBOL(CublasHandleInitHandler); +// Router +XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossForwardHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler); + } // namespace jax } // namespace transformer_engine @@ -165,6 +171,7 @@ XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( // ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); +XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Score_Function); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Collective_Op); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Quantize_Layout); diff --git a/transformer_engine/jax/csrc/extensions/misc.h b/transformer_engine/jax/csrc/extensions/misc.h index eb7be0a66a..c6f6f87cb4 100644 --- a/transformer_engine/jax/csrc/extensions/misc.h +++ b/transformer_engine/jax/csrc/extensions/misc.h @@ -122,6 +122,11 @@ void hash_combine(int64_t &seed, const T &v, Rest... rest) { (hash_combine(seed, rest), ...); } +enum class JAXX_Score_Function : int64_t { + SIGMOID = 0, + SOFTMAX = 1, +}; + enum class JAXX_Collective_Op : int64_t { NONE = 0, ALL_GATHER = 1, diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 71de897d9b..deea64caa1 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -84,6 +84,14 @@ pybind11::dict Registrations() { dict["te_inspect_ffi"] = pybind11::dict(pybind11::arg("execute") = EncapsulateFFI(InspectHandler)); + // Router + dict["te_fused_topk_with_score_function_forward_ffi"] = + EncapsulateFFI(FusedTopkWithScoreFunctionForwardHandler); + dict["te_fused_topk_with_score_function_backward_ffi"] = + EncapsulateFFI(FusedTopkWithScoreFunctionBackwardHandler); + dict["te_fused_moe_aux_loss_forward_ffi"] = EncapsulateFFI(FusedMoEAuxLossForwardHandler); + dict["te_fused_moe_aux_loss_backward_ffi"] = EncapsulateFFI(FusedMoEAuxLossBackwardHandler); + return dict; } @@ -191,6 +199,11 @@ PYBIND11_MODULE(transformer_engine_jax, m) { .value("ROWWISE_COLWISE", JAXX_Quantize_Layout::ROWWISE_COLWISE) .export_values(); + pybind11::enum_(m, "JAXX_Score_Function", pybind11::module_local()) + .value("SIGMOID", JAXX_Score_Function::SIGMOID) + .value("SOFTMAX", JAXX_Score_Function::SOFTMAX) + .export_values(); + pybind11::enum_(m, "JAXX_Collective_Op", pybind11::module_local()) .value("NONE", JAXX_Collective_Op::NONE) .value("ALL_GATHER", JAXX_Collective_Op::ALL_GATHER) diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp new file mode 100644 index 0000000000..0190d3fd75 --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -0,0 +1,237 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../extensions.h" +#include "xla/ffi/api/c_api.h" + +namespace transformer_engine { +namespace jax { + +// ============================================================================ +// Fused Top-K with Score Function - Forward +// ============================================================================ + +Error_Type FusedTopkWithScoreFunctionForwardFFI( + cudaStream_t stream, + Buffer_Type logits_buf, // [num_tokens, num_experts] + Buffer_Type expert_bias_buf, // [num_experts] or empty + Result_Type probs_buf, // [num_tokens, num_experts] (or scores when compute_aux_scores) + Result_Type routing_map_buf, // [num_tokens, num_experts] + Result_Type intermediate_buf, // [num_tokens, num_experts] + int64_t topk, int64_t use_pre_softmax, int64_t num_groups, int64_t group_topk, + double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores) { + auto dtype = convert_ffi_datatype_to_te_dtype(logits_buf.element_type()); + auto dims = logits_buf.dimensions(); + auto num_tokens = static_cast(product(dims, 0, dims.size() - 1)); + auto num_experts = static_cast(dims[dims.size() - 1]); + + auto *logits = logits_buf.untyped_data(); + auto *expert_bias = expert_bias_buf.untyped_data(); + auto *probs = probs_buf->untyped_data(); + auto *routing_map = routing_map_buf->untyped_data(); + auto *intermediate = intermediate_buf->untyped_data(); + + auto flat_shape = + std::vector{static_cast(num_tokens), static_cast(num_experts)}; + auto logits_tensor = TensorWrapper(logits, flat_shape, dtype); + auto probs_tensor = TensorWrapper(probs, flat_shape, dtype); + auto routing_map_tensor = TensorWrapper(routing_map, flat_shape, DType::kByte); + auto intermediate_tensor = TensorWrapper(intermediate, flat_shape, dtype); + + if (compute_aux_scores) { + nvte_fused_score_for_moe_aux_loss_forward( + logits_tensor.data(), num_tokens, num_experts, static_cast(topk), + static_cast(score_function), probs_tensor.data(), routing_map_tensor.data(), + intermediate_tensor.data(), stream); + } else { + auto bias_dims = expert_bias_buf.dimensions(); + auto expert_bias_tensor = + (bias_dims.size() > 0 && bias_dims[0] > 0) + ? TensorWrapper(expert_bias, std::vector{static_cast(bias_dims[0])}, + convert_ffi_datatype_to_te_dtype(expert_bias_buf.element_type())) + : TensorWrapper(); + + nvte_fused_topk_with_score_function_forward( + logits_tensor.data(), num_tokens, num_experts, static_cast(topk), + static_cast(use_pre_softmax), static_cast(num_groups), + static_cast(group_topk), static_cast(scaling_factor), + static_cast(score_function), expert_bias_tensor.data(), probs_tensor.data(), + routing_map_tensor.data(), intermediate_tensor.data(), stream); + } + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionForwardHandler, + FusedTopkWithScoreFunctionForwardFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // logits + .Arg() // expert_bias + .Ret() // probs (or scores) + .Ret() // routing_map + .Ret() // intermediate_output + .Attr("topk") + .Attr("use_pre_softmax") + .Attr("num_groups") + .Attr("group_topk") + .Attr("scaling_factor") + .Attr("score_function") + .Attr("compute_aux_scores"), + FFI_CudaGraph_Traits); + +// ============================================================================ +// Fused Top-K with Score Function - Backward +// ============================================================================ + +Error_Type FusedTopkWithScoreFunctionBackwardFFI( + cudaStream_t stream, + Buffer_Type routing_map_buf, // [num_tokens, num_experts] (unused when compute_aux_scores) + Buffer_Type intermediate_buf, // [num_tokens, num_experts] + Buffer_Type grad_probs_buf, // [num_tokens, num_experts] (grad_scores when compute_aux_scores) + Result_Type grad_logits_buf, // [num_tokens, num_experts] + int64_t topk, int64_t use_pre_softmax, double scaling_factor, + JAXX_Score_Function score_function, int64_t compute_aux_scores) { + auto dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf.element_type()); + auto dims = intermediate_buf.dimensions(); + auto num_tokens = static_cast(product(dims, 0, dims.size() - 1)); + auto num_experts = static_cast(dims[dims.size() - 1]); + + auto flat_shape = + std::vector{static_cast(num_tokens), static_cast(num_experts)}; + + auto intermediate_tensor = TensorWrapper(intermediate_buf.untyped_data(), flat_shape, dtype); + auto grad_probs_tensor = TensorWrapper(grad_probs_buf.untyped_data(), flat_shape, dtype); + auto grad_logits_tensor = TensorWrapper(grad_logits_buf->untyped_data(), flat_shape, dtype); + + if (compute_aux_scores) { + nvte_fused_score_for_moe_aux_loss_backward(intermediate_tensor.data(), grad_probs_tensor.data(), + num_tokens, num_experts, static_cast(topk), + static_cast(score_function), + grad_logits_tensor.data(), stream); + } else { + auto routing_map_tensor = + TensorWrapper(routing_map_buf.untyped_data(), flat_shape, DType::kByte); + + nvte_fused_topk_with_score_function_backward( + routing_map_tensor.data(), intermediate_tensor.data(), grad_probs_tensor.data(), num_tokens, + num_experts, static_cast(topk), static_cast(use_pre_softmax), + static_cast(scaling_factor), static_cast(score_function), + grad_logits_tensor.data(), stream); + } + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler, + FusedTopkWithScoreFunctionBackwardFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // routing_map + .Arg() // intermediate_output + .Arg() // grad_probs + .Ret() // grad_logits + .Attr("topk") + .Attr("use_pre_softmax") + .Attr("scaling_factor") + .Attr("score_function") + .Attr("compute_aux_scores"), + FFI_CudaGraph_Traits); + +// ============================================================================ +// Fused MoE Aux Loss - Forward +// ============================================================================ + +Error_Type FusedMoEAuxLossForwardFFI(cudaStream_t stream, + Buffer_Type probs_buf, // [num_tokens, num_experts] + Buffer_Type tokens_per_expert_buf, // [num_experts] + Result_Type aux_loss_buf, // scalar + Result_Type const_buf, // scalar + int64_t topk, double coeff) { + auto dtype = convert_ffi_datatype_to_te_dtype(probs_buf.element_type()); + auto probs_dims = probs_buf.dimensions(); + auto num_tokens = static_cast(probs_dims[0]); + auto num_experts = static_cast(probs_dims[1]); + + auto probs_shape = + 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); + + nvte_fused_moe_aux_loss_forward(probs_tensor.data(), tpe_tensor.data(), num_tokens, num_experts, + num_tokens, num_experts, static_cast(topk), + static_cast(coeff), aux_loss_tensor.data(), + const_buf_tensor.data(), stream); + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedMoEAuxLossForwardHandler, FusedMoEAuxLossForwardFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // probs + .Arg() // tokens_per_expert + .Ret() // aux_loss + .Ret() // const_buf + .Attr("topk") + .Attr("coeff"), + FFI_CudaGraph_Traits); + +// ============================================================================ +// Fused MoE Aux Loss - Backward +// ============================================================================ + +Error_Type FusedMoEAuxLossBackwardFFI(cudaStream_t stream, + Buffer_Type const_buf_in, // scalar float32 + Buffer_Type tokens_per_expert_buf, // [num_experts] + Buffer_Type grad_aux_loss_buf, // scalar + Result_Type grad_probs_buf) { // [num_tokens, num_experts] + auto grad_dtype = convert_ffi_datatype_to_te_dtype(grad_aux_loss_buf.element_type()); + auto tpe_dtype = convert_ffi_datatype_to_te_dtype(tokens_per_expert_buf.element_type()); + + auto grad_probs_dims = grad_probs_buf->dimensions(); + 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 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); + auto grad_probs_tensor = + TensorWrapper(grad_probs_buf->untyped_data(), grad_probs_shape, grad_dtype); + + nvte_fused_moe_aux_loss_backward(const_buf_tensor.data(), tpe_tensor.data(), num_tokens, + num_experts, grad_aux_loss_tensor.data(), + grad_probs_tensor.data(), stream); + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler, FusedMoEAuxLossBackwardFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // const_buf + .Arg() // tokens_per_expert + .Arg() // grad_aux_loss + .Ret(), // grad_probs + FFI_CudaGraph_Traits); + +} // namespace jax +} // namespace transformer_engine diff --git a/transformer_engine/jax/router.py b/transformer_engine/jax/router.py new file mode 100644 index 0000000000..65f2e8a7ff --- /dev/null +++ b/transformer_engine/jax/router.py @@ -0,0 +1,318 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused MoE Router API for JAX. + +This module provides high-level fused router operations for Mixture of Experts (MoE) +models with proper automatic differentiation support. These wrap the CUDA kernels in +transformer_engine/common/fused_router/. + +Functions: + fused_topk_with_score_function: + Fused score_function + top-k selection. Supports softmax/sigmoid, + grouped top-k, expert bias, and scaling factor. When compute_aux_scores=True, + switches to the clean score-for-aux-loss kernel (no bias/groups/scaling, + dense output). + + fused_moe_aux_loss: + Compute the MoE auxiliary load-balancing loss scalar. +""" + +from functools import partial +from typing import Optional, Tuple, Union + +import jax +import jax.numpy as jnp + +from transformer_engine.jax.cpp_extensions.router import ( + ScoreFunction, + fused_topk_with_score_function_fwd, + fused_topk_with_score_function_bwd, + fused_moe_aux_loss_fwd, + fused_moe_aux_loss_bwd, +) + +__all__ = [ + "ScoreFunction", + "fused_topk_with_score_function", + "fused_moe_aux_loss", +] + + +def _validate_score_function(score_function: Union[str, ScoreFunction]) -> ScoreFunction: + """Validate and convert score_function to a ScoreFunction enum.""" + if isinstance(score_function, ScoreFunction): + return score_function + try: + return ScoreFunction[score_function.upper()] + except (KeyError, AttributeError): + raise ValueError( + "score_function must be 'softmax', 'sigmoid', or a ScoreFunction enum, " + f"got {score_function!r}" + ) from None + + +# ============================================================================= +# Fused Top-K with Score Function +# ============================================================================= + + +def fused_topk_with_score_function( + logits: jnp.ndarray, + topk: int, + use_pre_softmax: bool = False, + num_groups: int = -1, + group_topk: int = -1, + scaling_factor: float = 1.0, + score_function: Union[str, ScoreFunction] = ScoreFunction.SOFTMAX, + expert_bias: Optional[jnp.ndarray] = None, + compute_aux_scores: bool = False, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """ + Fused top-k with score function router. + + When compute_aux_scores=False (default), runs the main routing kernel: + score_function(logits) -> [optional bias] -> top-k -> [optional post-softmax] -> scale. + Returns sparse probs (only top-k positions nonzero) and routing_map. + + When compute_aux_scores=True, runs the score-for-aux-loss kernel instead: + score_function(logits) -> top-k (clean, no bias/groups/scaling). + Returns dense scores (all expert positions) and routing_map. + The expert_bias, use_pre_softmax, num_groups, group_topk, and scaling_factor + parameters are ignored in this mode. + + Parameters + ---------- + logits : jnp.ndarray + Logits from the gating GEMM, shape [num_tokens, num_experts]. + topk : int + Number of top experts to select per token. + use_pre_softmax : bool + If True, apply softmax before top-k (only for softmax score function). Else, apply post top-k. + Ignored when compute_aux_scores=True. + num_groups : int + Number of groups for grouped top-k. <= 0 disables grouping (default). + Ignored when compute_aux_scores=True. + group_topk : int + Top-k at group level. <= 0 disables group-level selection (default). + Ignored when compute_aux_scores=True. + scaling_factor : float + Scaling factor applied to output probs. + Ignored when compute_aux_scores=True. + score_function : Union[str, ScoreFunction] + Score function: "softmax" / "sigmoid" or ScoreFunction.SOFTMAX / ScoreFunction.SIGMOID. + expert_bias : Optional[jnp.ndarray] + Expert bias, shape [num_experts]. Only used with sigmoid. + Ignored when compute_aux_scores=True. + compute_aux_scores : bool + If True, use the clean score-for-aux-loss kernel. Returns dense scores + over all experts instead of sparse probs. + + Returns + ------- + probs_or_scores : jnp.ndarray + When compute_aux_scores=False: Sparse probability tensor, shape [num_tokens, num_experts]. + Non-zero only at selected expert positions. + When compute_aux_scores=True: Dense score tensor, shape [num_tokens, num_experts]. + All expert positions contain scores. + routing_map : jnp.ndarray + Boolean mask, shape [num_tokens, num_experts]. + True at selected expert positions. + """ + if not isinstance(scaling_factor, (int, float)): + raise TypeError( + f"scaling_factor must be a Python float or int, not {type(scaling_factor).__name__}. " + "If you used jnp.sqrt() or similar, use math.sqrt() instead." + ) + + score_function = _validate_score_function(score_function) + + if compute_aux_scores: + expert_bias = jnp.empty((0,), dtype=logits.dtype) + use_pre_softmax = False + num_groups = -1 + group_topk = -1 + scaling_factor = 1.0 + else: + if expert_bias is not None and score_function != ScoreFunction.SIGMOID: + raise ValueError( + "expert_bias is only supported with score_function='sigmoid'. " + f"Got score_function='{score_function.name}'." + ) + if expert_bias is None: + expert_bias = jnp.empty((0,), dtype=logits.dtype) + + probs_or_scores, routing_map = _fused_topk_with_score_function( + logits, + expert_bias, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ) + + return probs_or_scores, routing_map + + +@partial(jax.custom_vjp, nondiff_argnums=(2, 3, 4, 5, 6, 7, 8)) +def _fused_topk_with_score_function( + logits: jnp.ndarray, + expert_bias: jnp.ndarray, + topk: int, + use_pre_softmax: bool, + num_groups: int, + group_topk: int, + scaling_factor: float, + score_function: ScoreFunction, + compute_aux_scores: bool, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + (probs, routing_map), _ = _fused_topk_with_score_function_fwd( + logits, + expert_bias, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + ) + return probs, routing_map + + +def _fused_topk_with_score_function_fwd( + logits, + expert_bias, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, +): + probs, routing_map, saved_scores = fused_topk_with_score_function_fwd( + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, + compute_aux_scores, + ) + residuals = (routing_map, saved_scores) + return (probs, routing_map), residuals + + +def _fused_topk_with_score_function_bwd( + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + compute_aux_scores, + residuals, + g, +): + del num_groups, group_topk + routing_map, saved_scores = residuals + grad_probs, _ = g + + grad_logits = fused_topk_with_score_function_bwd( + routing_map, + saved_scores, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, + compute_aux_scores, + ) + return grad_logits, None + + +_fused_topk_with_score_function.defvjp( + _fused_topk_with_score_function_fwd, + _fused_topk_with_score_function_bwd, +) + + +# ============================================================================= +# Fused MoE Aux Loss +# ============================================================================= + + +def fused_moe_aux_loss( + probs: jnp.ndarray, + tokens_per_expert: jnp.ndarray, + topk: int, + coeff: float, +) -> jnp.ndarray: + """ + Compute the MoE auxiliary load-balancing loss. + + loss = (E * coeff / (k * T^2)) * sum_i(sum_t(probs[t,i]) * tokens_per_expert[i]) + + where T = probs.shape[0] (num_tokens) and E = probs.shape[1] (num_experts). + + Parameters + ---------- + probs : jnp.ndarray + Probability/score tensor, shape [num_tokens, num_experts]. + tokens_per_expert : jnp.ndarray + Token counts per expert, shape [num_experts]. Integer tensor. + topk : int + Top-k value. + coeff : float + Loss coefficient. + + Returns + ------- + aux_loss : jnp.ndarray + Scalar loss value. + """ + return _fused_moe_aux_loss(probs, tokens_per_expert, topk, coeff) + + +@partial(jax.custom_vjp, nondiff_argnums=(2, 3)) +def _fused_moe_aux_loss( + probs: jnp.ndarray, + tokens_per_expert: jnp.ndarray, + topk: int, + coeff: float, +) -> jnp.ndarray: + aux_loss, _ = _fused_moe_aux_loss_fwd(probs, tokens_per_expert, topk, coeff) + return aux_loss + + +def _fused_moe_aux_loss_fwd(probs, tokens_per_expert, topk, coeff): + aux_loss, const_buf = fused_moe_aux_loss_fwd(probs, tokens_per_expert, topk, coeff) + residuals = (const_buf, tokens_per_expert, probs.shape[0]) + return aux_loss, residuals + + +def _fused_moe_aux_loss_bwd(topk, coeff, residuals, g): + del topk, coeff + const_buf, tokens_per_expert, num_tokens = residuals + grad_aux_loss = g.reshape(1) + + grad_probs = fused_moe_aux_loss_bwd( + const_buf, + tokens_per_expert, + grad_aux_loss, + num_tokens, + ) + return grad_probs, None + + +_fused_moe_aux_loss.defvjp( + _fused_moe_aux_loss_fwd, + _fused_moe_aux_loss_bwd, +) From d2e4755d3d87ba5e45f38a2a59a8bb3384e4ad1b Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 4 Mar 2026 19:36:27 -0800 Subject: [PATCH 246/521] [PyTorch] Skip `test_nvfp4_partial_cast_matches_full` test when NVFP4 is not available. (#2735) Skip test_nvfp4_partial_cast_matches_full test on hopper Signed-off-by: Kirthi Shankar Sivamani --- .../distributed/test_cast_master_weights_to_fp8.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 373f27be2f..1606641b78 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -1101,10 +1101,6 @@ def _test_nvfp4_partial_cast_matches_full(dp_group) -> None: WORLD_RANK = dist.get_rank(dp_group) WORLD_SIZE = dist.get_world_size(dp_group) - available, reason = is_nvfp4_available(return_reason=True) - if not available: - pytest.skip(reason) - torch.manual_seed(1234) device = torch.device("cuda") # Shape must be divisible by WORLD_SIZE for even splitting @@ -1196,6 +1192,11 @@ def _test_nvfp4_partial_cast_matches_full(dp_group) -> None: @pytest.mark.parametrize("world_size", [2]) def test_nvfp4_partial_cast_matches_full(world_size: int) -> None: """Launch a distributed job for NVFP4 partial-cast equivalence test.""" + + available, reason = is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + python_exe = pathlib.Path(sys.executable).resolve() current_file = pathlib.Path(__file__).resolve() command = [ From 145e88c3a05dc389012f2db65175572a1126c47a Mon Sep 17 00:00:00 2001 From: aagallo Date: Thu, 5 Mar 2026 01:18:06 -0500 Subject: [PATCH 247/521] Add multi-precision training support to FSDP script (#2662) * Add precision parameter support for multiple training formats Enable configurable precision training with support for FP32, FP16, FP8, MXFP8, and NVFP4 formats. Added precision argument parser and match statement to configure appropriate dtype and recipe based on selected precision. - Add precision() type validator function - Implement precision-based configuration in train() - Support FP32, FP16, FP8, MXFP8, and NVFP4 formats - Configure format-specific recipes (DelayedScaling, MXFP8BlockScaling, NVFP4BlockScaling) - Set appropriate no_fp8 flags based on precision selection Signed-off-by: aagallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix FP16 dtype mapping and implement CLI flag precedence Correct FP16 precision to use torch.float16 instead of torch.bfloat16, and add precedence logic where --dtype and --no-fp8 flags override --precision when explicitly set, with warnings issued for conflicts. - Fix case fp16 to use torch.float16 instead of torch.bfloat16 - Add flag precedence detection by comparing against default values - Implement warning messages when --dtype or --no-fp8 override --precision - Update argument parser help text to document precedence behavior - Ensure --dtype and --no-fp8 take precedence over --precision presets Signed-off-by: Andrea Gallo * Add logging and documentation for precision configuration Add informative log messages and enhanced help text to clarify precision configuration behavior and flag precedence for better user transparency. - Add log message showing which precision preset is being used - Add warning logs when --dtype or --no-fp8 override --precision - Add final training configuration log (dtype, FP8 status, recipe) - Enhance argument parser help text with precedence examples - Add inline code comments explaining precedence logic Signed-off-by: Andrea Gallo * Initialize recipe variable in all precision cases Add recipe initialization for fp32 and fp16 precision cases to prevent undefined variable errors, even though recipe is not used when no_fp8 is set to True. - Add DelayedScaling recipe setup for fp32 case with no_fp8=True - Add DelayedScaling recipe setup for fp16 case with no_fp8=True - Add inline comments explaining recipe is set up but not used by autocast - Ensure recipe variable is defined in all precision branches for consistency Signed-off-by: Andrea Gallo * Fix dtype flag detection to support explicit override behavior Update flag precedence detection to use sys.argv for checking if --dtype was explicitly set, ensuring dtype always overrides precision regardless of whether it matches the default value. - Add sys import for command-line argument detection - Change dtype_explicitly_set check to use '--dtype' in sys.argv - Change no_fp8_explicitly_set check to use '--no-fp8' in sys.argv - Ensure --dtype bf16 correctly overrides --precision even when matching default - Maintain warning messages when explicit flags override precision presets Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Replace sys.argv parsing with custom action and fix default case Replace fragile sys.argv parsing with robust custom argparse action class to track explicitly set arguments, and fix default precision case to explicitly set no_fp8 to False for consistent FP8-enabled behavior. - Add StoreExplicitAction custom action class for tracking explicit arguments - Update --dtype argument to use StoreExplicitAction - Replace sys.argv check with getattr for dtype_explicitly_set attribute - Remove sys import from train() function - Fix default case to set no_fp8 = False instead of opts.no_fp8 - Ensure recipe variable is properly initialized in all code paths - Support all argument passing methods including config files and = syntax Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix params_dtype to use computed dtype from precision logic Remove params_dtype initialization from get_layer_args() and update FSDP MixedPrecision to use computed dtype variable instead of raw opts.dtype, ensuring precision presets are properly applied throughout the model. - Remove params_dtype from get_layer_args() layer_kwargs initialization - Update FSDP MixedPrecision param_dtype to use computed dtype variable - Ensure precision preset logic is respected in both layer initialization and FSDP - Maintain backward compatibility with original FP8-enabled default behavior Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix type conversion in StoreExplicitAction for --dtype argument Add type converter application in StoreExplicitAction custom action to ensure --dtype values are properly converted from strings to torch dtype objects, preventing runtime errors in torch operations. - Store type converter in StoreExplicitAction.__init__ - Apply type conversion in __call__ before setting attribute value - Add error handling for invalid type conversions - Ensure opts.dtype contains torch dtype object, not raw string - Fix runtime errors in torch.rand() and MixedPrecision() calls Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix precision preset recipe selection and add incompatibility validation Address critical bugs where FP8 recipes were incorrectly selected when explicit flags were set, and add validation to prevent incompatible flag combinations that would silently disable FP8 training. - Remove default value from --precision parameter (set to None for backward compatibility) - Add get_precision_preset() and get_recipe_for_precision() helper functions - Implement two-path configuration logic: backward compatibility mode vs. precision preset mode - Add incompatibility validation: raise ValueError when --no-fp8 used with fp8/mxfp8/nvfp4 presets - Preserve FP8 recipe selection when --dtype explicitly overrides precision preset dtype - Fix fp16 case to correctly map to torch.float16 instead of torch.bfloat16 - Update parameter help text with precedence rules and usage examples - Ensure backward compatibility: scripts without --precision work identically to original version Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix unreachable default case and redundant recipe recreation Remove dead code in get_precision_preset() default case and eliminate redundant recipe recreation when dtype is explicitly overridden, ensuring cleaner logic flow and preventing duplicate recipe instantiation. - Remove unreachable case _: branch from get_precision_preset() function - Delete redundant recipe recreation when dtype_explicitly_set is true - Preserve existing recipe from preset when dtype override occurs - Ensure dtype override only affects parameter storage, not FP8 recipe selection Signed-off-by: Andrea Gallo * Add explicit error handling for invalid precision presets Prevent silent failures when precision validation is bypassed or new presets are added without updating get_precision_preset() function by adding explicit ValueError for unhandled cases. - Add case _: branch to get_precision_preset() that raises ValueError - Ensure invalid precision values fail loudly with clear error message - Prevent TypeError on tuple unpacking if function returns None - Improve maintainability when adding new precision presets Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: address argparse robustness and cleanup issues in fsdp.py Resolve three code review issues in examples/pytorch/fsdp/fsdp.py: dead commented-out code, unhelpful TypeError in precision(), and rigid __init__ signature in StoreTrueExplicitAction. - Remove commented-out layer_kwargs["params_dtype"] = dtype at line 106; dead code after params_dtype was moved to train() - Replace bare raise TypeError in precision() with argparse.ArgumentTypeError and explicit list of supported values (fp32, fp16, fp8, mxfp8, nvfp4) for a meaningful error message - Add **kwargs to StoreTrueExplicitAction.__init__ and forward to super().__init__(); aligns with StoreExplicitAction for robustness Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: suppress spurious dtype override warning when value matches preset Guard the --dtype override warning and dtype reassignment behind an actual value change check to avoid a false positive when the user explicitly passes --dtype with the same value the precision preset would have selected. - Add new_dtype != preset_dtype guard inside the dtype_explicitly_set branch so warning and dtype reassignment only trigger on a real override - Suppress redundant recipe re-creation when dtype matches preset default; recipe is already correctly set from preset_recipe above - No behavioral change when --dtype differs from preset default Signed-off-by: Andrea Gallo * fix: add type conversion with error handling in StoreExplicitAction.__call__ Apply type_converter inside __call__ with proper exception handling to ensure --dtype values are converted and validated at parse time rather than silently passing raw strings through. - Wrap type_converter call in try/except catching ValueError, TypeError, and argparse.ArgumentTypeError to surface conversion failures via parser.error() with a descriptive message - Guard conversion behind if self.type_converter is not None check for cases where no converter is registered - Ensures --dtype argument is correctly converted and validated consistently with standard argparse type= behavior Signed-off-by: Andrea Gallo * fix: remove redundant condition, deduplicate recipe logic, guard re-instantiation Address three code review issues in examples/pytorch/fsdp/fsdp.py: redundant opts.no_fp8 check, duplicated recipe construction, and unnecessary recipe re-instantiation when dtype matches preset. - Remove redundant 'and opts.no_fp8' from no_fp8_explicitly_set guard at line 373; StoreTrueExplicitAction always sets opts.no_fp8 to True when it fires, making the extra check always True - Refactor get_recipe_for_precision() to delegate to get_precision_preset() and extract the recipe, eliminating duplicated recipe construction logic and silent drift hazard when recipe parameters are tuned in one place but not the other - Guard recipe re-creation inside new_dtype != preset_dtype branch to avoid unnecessary re-instantiation when dtype_explicitly_set but the value matches the preset default Signed-off-by: Andrea Gallo * fix: validate flags before dist.init_process_group and remove redundant arg Move incompatible-flags check before dist.init_process_group() to avoid leaving the NCCL process group partially initialized, and remove redundant fp8_format=Format.E4M3 from MXFP8BlockScaling(). - Move no_fp8_explicitly_set + precision conflict check to the top of train() before dist.init_process_group() to prevent deadlocks or 'Address already in use' errors on other ranks that are still waiting inside init_process_group when rank 0 raises ValueError - Remove explicit fp8_format=Format.E4M3 from MXFP8BlockScaling() call; Format.E4M3 is already the dataclass default and passing it explicitly adds noise without adding clarity Signed-off-by: Andrea Gallo * fix: simplify StoreExplicitAction and improve training config log Delegate type conversion to argparse in StoreExplicitAction and include active FP8 recipe type in the training configuration log. - Remove self.type_converter field and manual try/except block from StoreExplicitAction.__call__; forward type= kwarg to super().__init__() so argparse handles conversion natively before __call__ is invoked, restoring standard error messages and %(type)s help interpolation - Simplify StoreExplicitAction.__init__ to use **kwargs passthrough, removing the now-unnecessary type= interception logic - Include active recipe type in training configuration log output using type(recipe).__name__ so log emits messages like 'FP8=enabled (MXFP8BlockScaling)' or 'FP8=disabled', making it easier to verify the intended quantization scheme is in use Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: forward kwargs in StoreTrueExplicitAction, improve dtype log, document rank assumption Address three code review issues in examples/pytorch/fsdp/fsdp.py: silent kwargs drop in StoreTrueExplicitAction, missing confirmation log when dtype matches preset, and undocumented torchrun assumption. - Forward **kwargs to super().__init__() in StoreTrueExplicitAction to prevent silent discard of unexpected keyword arguments (e.g. metavar, choices) if argument registration is ever extended - Add info log when dtype_explicitly_set but new_dtype == preset_dtype so user receives confirmation their --dtype flag was acknowledged even when it matches the preset default and no override is needed - Add comment above no_fp8_explicitly_set validation documenting that raising ValueError before dist.init_process_group is safe because torchrun guarantees all ranks receive identical CLI arguments Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixing typo in code documentation Signed-off-by: Andrea Gallo * fix: format dtype in log messages and document recipe=None intent Strip 'torch.' prefix from dtype in user-facing log messages and add a comment documenting the intentional recipe=None behavior when FP8 is disabled. - Replace raw dtype formatting with str(dtype).replace('torch.', '') in both Warning and Info log messages so users see 'float32' or 'bfloat16' instead of 'torch.float32' or 'torch.bfloat16' - Add inline comment on recipe=None assignment explaining that te.autocast safely substitutes get_default_fp8_recipe() internally when recipe is None, and skips check_recipe_support when enabled=False, so the assignment is intentional and safe despite populating global FP8 state with a default recipe Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor: remove redundant __init__ override in StoreExplicitAction Remove the __init__ override from StoreExplicitAction since it only calls super().__init__() with the same arguments, which Python does automatically. The class now consists solely of __call__, eliminating dead code without any behavioral change. Signed-off-by: Andrea Gallo * fix: use 'quantization' label in log and remove redundant recipe re-instantiation Replace misleading 'FP8' label with 'quantization' in training configuration log and remove redundant recipe re-instantiation in the dtype_explicitly_set path. - Replace 'FP8=enabled/disabled' with 'quantization=enabled/disabled' in dist_print configuration log to accurately cover all TE precision modes including NVFP4 which is 4-bit, not FP8 - Remove get_recipe_for_precision() call inside dtype_explicitly_set block; recipe is already correctly assigned from preset_recipe above and re-instantiating it is wasteful and creates a second object discarding the first - Add inline comment clarifying that recipe requires no update in the dtype_explicitly_set path since it is determined by opts.precision, not dtype Signed-off-by: Andrea Gallo * fix: remove redundant recipe re-instantiation in equal-dtype path Remove unnecessary get_recipe_for_precision() call in the else branch of the dtype_explicitly_set block where new_dtype == preset_dtype. - recipe is already correctly assigned from preset_recipe before the dtype_explicitly_set block; no re-instantiation is needed in either branch since recipe is determined by opts.precision, not dtype - Previous else branch was re-creating the recipe (wasteful) while the if branch was not, inverting the logic implied by the comment - Replace with a comment clarifying that recipe requires no update in the dtype_explicitly_set path Signed-off-by: Andrea Gallo * fix: define dtype_name unconditionally and guard dtype override warning Fix two bugs in train() precision configuration block. - Define dtype_name unconditionally before the dtype_explicitly_set block to prevent NameError in the config log when dtype_explicitly_set is False (the common case when --dtype is not explicitly passed) - Guard dtype override warning behind 'dtype_explicitly_set and opts.precision is not None' to prevent spurious warning when user passes --dtype without --precision (original behavior path) Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: recompute dtype_name after override and restore DelayedScaling default Fix two bugs in train() precision configuration block: stale dtype_name in log messages after dtype override, and behavioral regression where recipe=None was passed to te.autocast when FP8 was enabled in backward-compatible mode. - Recompute dtype_name immediately after dtype = new_dtype in the dtype override branch so warning and config log reflect the effective dtype rather than the stale preset dtype - Restore original default behavior in opts.precision is None path: when no_fp8 is False (FP8 enabled), supply DelayedScaling recipe to preserve the original te.autocast behavior instead of passing recipe=None which changed the implicit fallback behavior Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * refactor: simplify opts.dtype None check and remove redundant no_fp8 assignment Remove two redundant lines in train() precision configuration block. - Remove 'if opts.dtype is not None' guard in opts.precision is None branch; --dtype has default=torch.bfloat16 so opts.dtype is never None and the condition is always True - Remove redundant 'no_fp8 = preset_no_fp8' assignment in the else branch; no_fp8 is already assigned from preset_no_fp8 at the tuple unpack above and reassigning it in the else branch adds noise without changing behavior Signed-off-by: Andrea Gallo * fix: guard recipe=None, remove dead case None, shorten help text Address three code review issues in examples/pytorch/fsdp/fsdp.py: recipe=None passed to te.autocast, dead case None in get_precision_preset, and excessively verbose help strings for --no-fp8 and --dtype. - Use 'recipe or DelayedScaling()' fallback at te.autocast call site to preserve original defensive pattern of always passing a concrete recipe instance, even when enabled=False - Remove case None from get_precision_preset() and guard call site in train() with 'if opts.precision is not None' to eliminate dead-code path whose return values were immediately overridden by the caller - Replace multi-paragraph help strings for --no-fp8 and --dtype with concise one-liner synopses; move detailed precedence rules to module-level docstring or README Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: resolve recipe=None fallback before training loop, not per-iteration Move 'recipe or DelayedScaling()' fallback to a one-time assignment before the training loop to avoid allocating a new DelayedScaling() object on every iteration when FP8 is disabled. - Add 'if recipe is None: recipe = DelayedScaling()' after the configuration block and before the training loop so the fallback object is created once and reused across all iterations - Restore clean 'recipe=recipe' in te.autocast call, matching the original code pattern - Add comment explaining why recipe is always set to a concrete object even when FP8 is disabled Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: move recipe=None fallback before training loop with consistent parameters Move 'if recipe is None: recipe = DelayedScaling(...)' guard to just before the training loop instead of inside it to avoid redundant is-None checks on every iteration and variable mutation inside the loop. - Use consistent DelayedScaling parameters (fp8_format=Format.HYBRID, amax_history_len=32, amax_compute_algo='max') matching the rest of the file, rather than plain DelayedScaling() with default args - Guard runs once before the loop; recipe is stable for all iterations - Restores clean 'recipe=recipe' in te.autocast call with no inline fallback expression Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: pass amax_reduction_group only for DelayedScaling and shorten help text Pass amax_reduction_group conditionally based on recipe type and replace verbose multi-paragraph help strings with concise one-liners. - Compute amax_group = all_gpus if isinstance(recipe, DelayedScaling) else None and pass amax_group to te.autocast; amax_reduction_group is a DelayedScaling-specific parameter for per-tensor amax aggregation and is not accepted by MXFP8BlockScaling or NVFP4BlockScaling which use block-level scaling - Replace multi-paragraph help strings for --no-fp8, --dtype, and --precision (with PRECEDENCE/BEHAVIOR/RATIONALE/EXAMPLES sections) with concise one-liner synopses suitable for terminal --help output - Move detailed precedence rules to module-level docstring Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: guard amax_group with not no_fp8 to prevent spurious distributed comms Add 'not no_fp8' condition to amax_group assignment to prevent amax_reduction_group=all_gpus being passed to te.autocast when FP8 is disabled. - When no_fp8=True and recipe was None, the DelayedScaling fallback causes isinstance(recipe, DelayedScaling) to return True, which incorrectly set amax_group=all_gpus even though enabled=False - Add 'not no_fp8' guard so amax_group is only set to all_gpus when FP8 is active AND the recipe is DelayedScaling (per-tensor amax); all other cases (FP8 disabled, block-scaling recipes) use None Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: warn on redundant --no-fp8 with fp32/fp16 and document amax_group=None Emit a warning when --no-fp8 is combined with a non-FP8 precision preset and add an inline comment explaining why amax_reduction_group is None for block-scaling recipes. - Add warning when opts.precision in ['fp32', 'fp16'] and opts.no_fp8 is set; FP8 is already disabled by these presets so the flag is redundant and silently ignored without this feedback - Add inline comment on amax_group assignment explaining that MXFP8BlockScaling and NVFP4BlockScaling use local block scaling and do not require a distributed amax reduction group, and that None is also correct when FP8 is disabled to avoid unnecessary distributed communication Signed-off-by: Andrea Gallo * fix: initialize preset_dtype and preset_recipe before conditional block Initialize preset_dtype and preset_recipe with fallback values before the 'if opts.precision is not None' block to prevent static analyzer warnings about potentially unbound variables. - Assign preset_dtype = opts.dtype and preset_recipe = None as sensible fallbacks before the if-else block; these are overwritten by get_precision_preset() when opts.precision is not None and are never accessed in the else branch - Satisfies mypy, pylint, and pyflakes 'possibly undefined' / 'unbound' warnings that would otherwise trigger CI lint failures in projects treating unbound-variable warnings as errors - No behavioral change; the if-else logic is unchanged Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: compute amax_group before recipe=None fallback to avoid isinstance race Move amax_group computation before the 'if recipe is None' fallback assignment so isinstance(recipe, DelayedScaling) reflects the actual user-selected recipe rather than the defensive fallback object. - When recipe is None (non-FP8 presets or --no-fp8), isinstance correctly returns False and amax_group is set to None before the fallback substitutes a DelayedScaling instance - Prevents the fragile ordering dependency where not no_fp8 was the sole guard against passing all_gpus to a recipe that doesn't need it - Add inline comment explaining why amax_group must be computed before the recipe fallback to preserve the invariant Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix: warn on potentially incompatible --dtype float16 with FP8-family presets Add explicit warning when --dtype float16 is combined with --precision fp8, mxfp8, or nvfp4, which expect bfloat16 accumulation. - Emit compatibility warning before applying the dtype override when opts.precision is in ['fp8', 'mxfp8', 'nvfp4'] and new_dtype is torch.float16; these presets are designed for bfloat16 accumulation and pairing with float16 may produce incorrect or undefined results - Warning is emitted in addition to the existing dtype override warning so users see both the compatibility concern and the override confirmation - Override is still applied (not blocked) to preserve user control; users who know their TE version supports float16 accumulation can proceed with awareness of the risk Signed-off-by: Andrea Gallo * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: aagallo Signed-off-by: Andrea Gallo Co-authored-by: aagallo Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 --- examples/pytorch/fsdp/fsdp.py | 180 +++++++++++++++++++++++++++++++--- 1 file changed, 168 insertions(+), 12 deletions(-) diff --git a/examples/pytorch/fsdp/fsdp.py b/examples/pytorch/fsdp/fsdp.py index b469ef56b7..ac7a2fac7b 100644 --- a/examples/pytorch/fsdp/fsdp.py +++ b/examples/pytorch/fsdp/fsdp.py @@ -18,7 +18,12 @@ ) import transformer_engine.pytorch as te -from transformer_engine.common.recipe import Format, DelayedScaling +from transformer_engine.common.recipe import ( + Format, + DelayedScaling, + MXFP8BlockScaling, + NVFP4BlockScaling, +) from transformer_engine.pytorch.distributed import prepare_te_modules_for_fsdp LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) @@ -64,10 +69,21 @@ def torch_dtype(d): "bfloat16": torch.bfloat16, } if lowercase(d) not in typemap.keys(): - raise TypeError + raise argparse.ArgumentTypeError( + f"invalid dtype '{d}'. Supported values: fp32/float32, fp16/float16, bf16/bfloat16" + ) return typemap[lowercase(d)] +def precision(d): + typemap = ["fp32", "fp16", "fp8", "mxfp8", "nvfp4"] + if lowercase(d) not in typemap: + raise argparse.ArgumentTypeError( + f"invalid precision '{d}'. Supported values: {', '.join(typemap)}" + ) + return lowercase(d) + + te_layer_map = { "linear": te.Linear, "layernorm": te.LayerNorm, @@ -91,7 +107,6 @@ def get_layer_args(opts): hidden_size = opts.num_heads * opts.head_dim layer_args = (hidden_size,) layer_kwargs = { - "params_dtype": opts.dtype, "device": "cuda" if opts.no_defer_init else "meta", "get_rng_state_tracker": get_cuda_rng_tracker, } @@ -112,6 +127,15 @@ def get_layer_args(opts): return layer_args, layer_kwargs +class StoreExplicitAction(argparse.Action): + """Custom action that tracks whether an argument was explicitly set.""" + + def __call__(self, parser, namespace, values, option_string=None): + # values already converted by argparse via action.type + setattr(namespace, self.dest, values) + setattr(namespace, f"{self.dest}_explicitly_set", True) + + def parse_fsdp_args(): parser = argparse.ArgumentParser( description="Run Transformer Engine modules with the " @@ -173,7 +197,10 @@ def parse_fsdp_args(): "--no-fp8", action="store_true", default=False, - help="Disables the te.autocast() context.", + help=( + "Disable te.autocast() FP8 context. Incompatible with --precision fp8/mxfp8/nvfp4." + " Default: False." + ), ) parser.add_argument( "--no-defer-init", @@ -189,7 +216,21 @@ def parse_fsdp_args(): "--dtype", type=torch_dtype, default=torch.bfloat16, - help="Data type for input tensor and Transformer Engine module parameters.", + action=StoreExplicitAction, + help=( + "Parameter dtype: fp32/float32, fp16/float16, bf16/bfloat16. Overrides --precision" + " dtype when explicitly set. Default: bfloat16." + ), + ) + parser.add_argument( + "--precision", + type=precision, + default=None, + help=( + "Precision preset: fp32, fp16, fp8, mxfp8, nvfp4. Configures dtype and FP8 recipe" + " automatically. Overridden by explicit --dtype. Default: None (use --dtype and" + " --no-fp8 directly)." + ), ) return parser.parse_args() @@ -200,15 +241,118 @@ def dist_print(text, all_ranks=False, no_new_line=False): print(f"[GPU-{LOCAL_RANK}] " + text, end=end) +def get_precision_preset(precision_value): + """Get dtype, no_fp8, and recipe based on precision preset. + + Returns: + tuple: (dtype, no_fp8, recipe) + """ + match precision_value: + case "fp32": + return torch.float32, True, None + case "fp16": + return torch.float16, True, None + case "fp8": + recipe = DelayedScaling( + fp8_format=Format.HYBRID, amax_history_len=32, amax_compute_algo="max" + ) + return torch.bfloat16, False, recipe + case "mxfp8": + recipe = MXFP8BlockScaling() + return torch.bfloat16, False, recipe + case "nvfp4": + recipe = NVFP4BlockScaling() + return torch.bfloat16, False, recipe + case _: + raise ValueError( + f"Invalid precision preset: {precision_value}. " + "Supported values: fp32, fp16, fp8, mxfp8, nvfp4" + ) + + def train(opts): + # Check which flags were explicitly set + dtype_explicitly_set = getattr(opts, "dtype_explicitly_set", False) + + # Validate flag combinations before touching distributed state. + # Error if user requests FP8-based precision but also sets --no-fp8 + # Safe to raise here because torchrun guarantees all ranks receive + # identical CLI arguments; all ranks will raise simultaneously. + if opts.precision in ["fp8", "mxfp8", "nvfp4"] and opts.no_fp8: + raise ValueError( + f"Cannot use --no-fp8 with --precision {opts.precision}. " + "These flags are incompatible. " + f"Either remove --no-fp8 to use {opts.precision} training, " + "or use --precision fp32/fp16 for non-FP8 training." + ) + if opts.precision in ["fp32", "fp16"] and opts.no_fp8: + dist_print( + f"Warning: --no-fp8 is redundant when using --precision {opts.precision} " + "(FP8 is already disabled by this preset). The flag will be ignored." + ) + # Initialize torch.distributed global process group dist.init_process_group(backend="nccl") torch.cuda.set_device(LOCAL_RANK) dist_print(f"WORLD_SIZE = {WORLD_SIZE}") torch.manual_seed(opts.seed) + preset_dtype: torch.dtype = opts.dtype # sensible fallback + preset_recipe = None + + if opts.precision is not None: + preset_dtype, preset_no_fp8, preset_recipe = get_precision_preset(opts.precision) + dtype, no_fp8, recipe = preset_dtype, preset_no_fp8, preset_recipe + dist_print(f"Using precision preset: {opts.precision}") + else: + # Original behavior: --dtype and --no-fp8 control training directly + dtype = opts.dtype + no_fp8 = opts.no_fp8 + recipe = ( + DelayedScaling(fp8_format=Format.HYBRID, amax_history_len=32, amax_compute_algo="max") + if not no_fp8 + else None + ) + + dtype_name = str(dtype).replace("torch.", "") + + # Apply explicit dtype override with warning + if dtype_explicitly_set and opts.precision is not None: + new_dtype = opts.dtype + if new_dtype != preset_dtype: + if opts.precision in ["fp8", "mxfp8", "nvfp4"] and new_dtype == torch.float16: + dist_print( + "Warning: --dtype float16 may be incompatible with --precision" + f" {opts.precision}, which expects bfloat16 accumulation." + ) + + dtype = new_dtype + dtype_name = str(dtype).replace("torch.", "") + + dist_print( + f"Warning: --dtype {dtype_name} overrides --precision {opts.precision} dtype" + " setting" + ) + else: + new_dtype_name = str(new_dtype).replace("torch.", "") + dist_print( + f"Info: --dtype {new_dtype_name} matches --precision {opts.precision} preset" + " default, no override needed" + ) + + # recipe is already set correctly from preset_recipe above; + # dtype only affects parameter storage, not the quantization recipe + + # Always log the final configuration being used + dist_print( + f"Training configuration: dtype={dtype_name}, " + f"quantization={'disabled' if no_fp8 else f'enabled ({type(recipe).__name__})'}" + ) + # Construct a simple homogeneous model (only one layer type) with NO PARALLELISM layer_args, layer_kwargs = get_layer_args(opts) + layer_kwargs["params_dtype"] = dtype + if opts.num_layers > 1: te_layer_list = [] for i in range(opts.num_layers): @@ -239,7 +383,7 @@ def train(opts): process_group=all_gpus, use_orig_params=True, mixed_precision=MixedPrecision( - param_dtype=opts.dtype, + param_dtype=dtype, reduce_dtype=torch.float32, ), auto_wrap_policy=fsdp_wrap_policy, @@ -258,10 +402,6 @@ def train(opts): dist_print(f"Post-FSDP memory use = {post_mem_use}MiB") dist_print(f"FSDP-Wrapped + Checkpointed TE Model:\n{te_model}") - # Fp8 setup for TE - fp8_format = Format.HYBRID - fp8_recipe = DelayedScaling(fp8_format=fp8_format, amax_history_len=32, amax_compute_algo="max") - # Optimizer must be created after the model is wrapped in FSDP and the parameters are sharded optim = torch.optim.Adam(te_model.parameters(), lr=0.0001) @@ -275,17 +415,33 @@ def train(opts): torch.cuda.synchronize() start.record() + # MXFP8 and NVFP4 use local block scaling — no distributed amax reduction group needed. + # amax_reduction_group is only required for DelayedScaling (global AMAX allreduce). + # Also skip when FP8 is disabled to avoid unnecessary distributed communication. + # Compute amax_group BEFORE the recipe fallback so isinstance() reflects the actual + # recipe, not the defensive DelayedScaling() substituted for None. + amax_group = all_gpus if (not no_fp8 and isinstance(recipe, DelayedScaling)) else None + + # Ensure recipe is always a concrete object before passing to te.autocast. + # When FP8 is disabled, te.autocast ignores the recipe, but some TE versions + # perform attribute access on it regardless of the enabled flag. + if recipe is None: + recipe = DelayedScaling( + fp8_format=Format.HYBRID, amax_history_len=32, amax_compute_algo="max" + ) + for i in range(opts.num_iters): # Generate a random input batch x = torch.rand( opts.seq_length, opts.batch_size, opts.num_heads * opts.head_dim, - dtype=opts.dtype, + dtype=dtype, device="cuda", ) + # autocast needs to be given the FSDP process group for amax reductions - with te.autocast(enabled=not opts.no_fp8, recipe=fp8_recipe, amax_reduction_group=all_gpus): + with te.autocast(enabled=not no_fp8, recipe=recipe, amax_reduction_group=amax_group): y = te_model(x) loss = y.sum() # calculate gradient and take training step outside the autocast context From d9152b0f2f595d0594a323f76df239269428a401 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 4 Mar 2026 23:57:05 -0800 Subject: [PATCH 248/521] [PyTorch] Support `GroupedTensor` torch ops for DDP and distributed optimizer (#2736) * Fix e2e execution of GroupedTensor in distributed settings Signed-off-by: Kirthi Shankar Sivamani * Minor fixes Signed-off-by: Kirthi Shankar Sivamani * fix Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * fix greptile commit Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- transformer_engine/pytorch/csrc/quantizer.cpp | 42 +++-- .../pytorch/tensor/grouped_tensor.py | 170 ++++++++++++++++-- .../tensor/storage/grouped_tensor_storage.py | 101 ++++++++--- 3 files changed, 263 insertions(+), 50 deletions(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 0135c7f01c..0214f7ff71 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -180,8 +180,11 @@ std::pair NoneQuantizer::create_grouped_tensor py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; py::tuple args(0); - kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + const std::vector grouped_shape = {static_cast(logical_first_dim), + static_cast(logical_last_dim)}; + const std::vector grouped_stride = stride_from_shape(grouped_shape); + kwargs["shape"] = py::cast(grouped_shape); + kwargs["stride"] = py::cast(grouped_stride); kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["num_tensors"] = py::cast(num_tensors); kwargs["quantizer"] = quantizer; @@ -386,8 +389,11 @@ std::pair Float8Quantizer::create_grouped_tens py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; py::tuple args(0); - kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + const std::vector grouped_shape = {static_cast(logical_first_dim), + static_cast(logical_last_dim)}; + const std::vector grouped_stride = stride_from_shape(grouped_shape); + kwargs["shape"] = py::cast(grouped_shape); + kwargs["stride"] = py::cast(grouped_stride); kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["num_tensors"] = py::cast(num_tensors); kwargs["quantizer"] = quantizer; @@ -704,8 +710,11 @@ std::pair Float8CurrentScalingQuantizer::creat py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; py::tuple args(0); - kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + const std::vector grouped_shape = {static_cast(logical_first_dim), + static_cast(logical_last_dim)}; + const std::vector grouped_stride = stride_from_shape(grouped_shape); + kwargs["shape"] = py::cast(grouped_shape); + kwargs["stride"] = py::cast(grouped_stride); kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["num_tensors"] = py::cast(num_tensors); kwargs["quantizer"] = quantizer; @@ -1062,8 +1071,11 @@ std::pair Float8BlockQuantizer::create_grouped py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; py::tuple args(0); - kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + const std::vector grouped_shape = {static_cast(logical_first_dim), + static_cast(logical_last_dim)}; + const std::vector grouped_stride = stride_from_shape(grouped_shape); + kwargs["shape"] = py::cast(grouped_shape); + kwargs["stride"] = py::cast(grouped_stride); kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["num_tensors"] = py::cast(num_tensors); kwargs["quantizer"] = quantizer; @@ -1478,8 +1490,11 @@ std::pair MXFP8Quantizer::create_grouped_tenso py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; py::tuple args(0); - kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + const std::vector grouped_shape = {static_cast(logical_first_dim), + static_cast(logical_last_dim)}; + const std::vector grouped_stride = stride_from_shape(grouped_shape); + kwargs["shape"] = py::cast(grouped_shape); + kwargs["stride"] = py::cast(grouped_stride); kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["num_tensors"] = py::cast(num_tensors); kwargs["quantizer"] = quantizer; @@ -1906,8 +1921,11 @@ std::pair NVFP4Quantizer::create_grouped_tenso py::handle GroupedTensorClass = grouped_tensor_python_class(this->internal); py::dict kwargs; py::tuple args(0); - kwargs["shape"] = py::cast(std::vector{static_cast(logical_first_dim), - static_cast(logical_last_dim)}); + const std::vector grouped_shape = {static_cast(logical_first_dim), + static_cast(logical_last_dim)}; + const std::vector grouped_stride = stride_from_shape(grouped_shape); + kwargs["shape"] = py::cast(grouped_shape); + kwargs["stride"] = py::cast(grouped_stride); kwargs["dtype"] = py::cast(GetATenDType(dtype)); kwargs["num_tensors"] = py::cast(num_tensors); kwargs["quantizer"] = quantizer; diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 767b0ccb35..685b2c5548 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -14,10 +14,36 @@ from .storage.grouped_tensor_storage import GroupedTensorStorage -# For now, conservatively ban all shape manipulating ops. +def _stride_from_shape(shape: Tuple[int, ...]) -> Tuple[int, ...]: + """Calculate contiguous stride from shape.""" + if len(shape) == 0: + return () + stride = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + stride[i] = stride[i + 1] * shape[i + 1] + return tuple(stride) + + +class _GroupedIdentityFunc(torch.autograd.Function): + """Identity autograd function used to create a dummy grad_fn node.""" + + @staticmethod + def forward(ctx, tensor: "GroupedTensor") -> "GroupedTensor": + # pylint: disable=missing-function-docstring + ctx.input_dtype = tensor.dtype + return tensor.detach() + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + # pylint: disable=missing-function-docstring + grad_input = grad_output + if grad_input.dtype != ctx.input_dtype: + grad_input = grad_input.to(ctx.input_dtype) + return grad_input + + +# For now, conservatively ban 'most' shape manipulating ops. BANNED_SHAPE_OPS = { - torch.ops.aten.view.default, - torch.ops.aten._unsafe_view.default, torch.ops.aten.reshape.default, torch.ops.aten._reshape_alias.default, torch.ops.aten.flatten.using_ints, @@ -34,8 +60,6 @@ torch.ops.aten.select.int, torch.ops.aten.split.Tensor, torch.ops.aten.chunk.default, - torch.ops.aten.expand.default, - torch.ops.aten.expand_as.default, torch.ops.aten.cat.default, torch.ops.aten.stack.default, } @@ -48,6 +72,7 @@ def __new__( cls, shape: Tuple[int, int], dtype: torch.dtype, + *, num_tensors: int, shapes: Optional[List[Tuple[int, int]]] = None, quantizer: Optional[Quantizer] = None, @@ -64,12 +89,9 @@ def __new__( offsets: Optional[List[int]] = None, scale_inv_offsets: Optional[List[int]] = None, columnwise_scale_inv_offsets: Optional[List[int]] = None, + requires_grad: bool = False, + stride: Optional[List[int]] = None, ): - del quantizer - del offsets - del scale_inv_offsets - del columnwise_scale_inv_offsets - if ( shapes is not None and len(shapes) == num_tensors @@ -99,19 +121,41 @@ def __new__( if device is None: device = torch.device("cuda") - strides = [1] * len(wrapper_shape) - for i in range(len(wrapper_shape) - 2, -1, -1): - strides[i] = strides[i + 1] * wrapper_shape[i + 1] - return torch.Tensor._make_wrapper_subclass( + # Match QuantizedTensor __new__: accept externally-computed stride to + # avoid Python-side stride computation overhead for C++ construction. + strides = _stride_from_shape(tuple(wrapper_shape)) if stride is None else tuple(stride) + instance = torch.Tensor._make_wrapper_subclass( cls, wrapper_shape, - strides=tuple(strides), + strides=strides, storage_offset=0, dtype=dtype, layout=torch.strided, - requires_grad=False, + requires_grad=requires_grad, device=device, ) + GroupedTensorStorage._initialize_storage_fields( + instance=instance, + shape=shape, + dtype=dtype, + num_tensors=num_tensors, + shapes=shapes, + quantizer=quantizer, + data=data, + columnwise_data=columnwise_data, + scale_inv=scale_inv, + columnwise_scale_inv=columnwise_scale_inv, + amax=amax, + columnwise_amax=columnwise_amax, + scale=scale, + first_dims=first_dims, + last_dims=last_dims, + tensor_offsets=tensor_offsets, + offsets=offsets, + scale_inv_offsets=scale_inv_offsets, + columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, + ) + return instance @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): @@ -119,9 +163,94 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if kwargs is None: kwargs = {} + def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> None: + """Shallow-copy grouped-storage metadata onto wrapper outputs.""" + dst.num_tensors = src.num_tensors + dst.quantizer = src.quantizer + dst.tensor_shapes = src.tensor_shapes + dst.fake_dtype = src.fake_dtype + dst.rowwise_data = src.rowwise_data + dst.columnwise_data = src.columnwise_data + dst.scale_inv = src.scale_inv + dst.columnwise_scale_inv = src.columnwise_scale_inv + dst.amax = src.amax + dst.columnwise_amax = src.columnwise_amax + dst.scale = src.scale + dst.first_dims = src.first_dims + dst.last_dims = src.last_dims + dst.tensor_offsets = src.tensor_offsets + dst.offsets = src.offsets + dst.scale_inv_offsets = src.scale_inv_offsets + dst.columnwise_scale_inv_offsets = src.columnwise_scale_inv_offsets + dst.logical_shape = src.logical_shape + dst.quantized_tensors = src.quantized_tensors + + def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: + """Create a wrapper of the same type and tensor metadata as src.""" + out = torch.Tensor._make_wrapper_subclass( + type(src), + tuple(src.shape), + strides=tuple(src.stride()), + storage_offset=src.storage_offset(), + dtype=src.dtype, + layout=src.layout, + requires_grad=requires_grad, + device=src.device, + ) + copy_grouped_storage_metadata(out, src) + return out + # Parameter construction calls detach()/alias-like paths. if func in (torch.ops.aten.detach.default, torch.ops.aten.alias.default): - return args[0] + src = args[0] + assert isinstance(src, GroupedTensor) + if func == torch.ops.aten.detach.default: + return make_wrapper_like(src, requires_grad=False) + return make_wrapper_like(src, requires_grad=src.requires_grad) + + # Parameter construction may invoke aten.expand on tensor subclasses. + # Handle this explicitly so grouped parameters can be created safely. + if func == torch.ops.aten.expand.default: + src = args[0] + assert isinstance(src, GroupedTensor) + expanded_shape = tuple(args[1]) + src_shape = tuple(src.shape) + if len(expanded_shape) == len(src_shape): + normalized_shape = tuple( + src_shape[i] if dim == -1 else dim for i, dim in enumerate(expanded_shape) + ) + if normalized_shape == src_shape: + return make_wrapper_like(src, requires_grad=src.requires_grad) + return super().__torch_dispatch__(func, types, args, kwargs) + + # DDP and mcore use expand_as(self) to build a dummy autograd node and + # access gradient accumulators during parameter hook registration. + if func == torch.ops.aten.expand_as.default: + src = args[0] + other = args[1] + assert isinstance(src, GroupedTensor) + if other is src: + return _GroupedIdentityFunc.apply(src) + if tuple(other.shape) == tuple(src.shape): + return make_wrapper_like(src, requires_grad=src.requires_grad) + return super().__torch_dispatch__(func, types, args, kwargs) + + # Distributed optimizer flattens detached parameters via + # model_param.detach().view(-1). Support this path explicitly by + # returning a flat view of grouped backing storage. + if func in (torch.ops.aten.view.default, torch.ops.aten._unsafe_view.default): + src = args[0] + assert isinstance(src, GroupedTensor) + target_shape = tuple(args[1]) + if target_shape in ((-1,), (src.numel(),)): + if src.rowwise_data is not None: + return src.rowwise_data.view(-1) + raise RuntimeError( + f"{cls.__name__} view(-1) requires rowwise_data to be initialized" + ) + raise RuntimeError( + f"{cls.__name__} only supports view(-1) for distributed optimizer flattening" + ) # Don't allow reshape/view etc. if func in BANNED_SHAPE_OPS: @@ -203,3 +332,10 @@ def __torch_function__(cls, func, types, args=(), kwargs=None): kwargs = {} # Do not force GroupedTensor on outputs. return torch._C._disabled_torch_function_impl(func, types, args, kwargs) + + def expand_as(self, other: torch.Tensor) -> torch.Tensor: + # pylint: disable=missing-function-docstring + # Needed during parameter creation/hook registration paths. + if other is self: + return _GroupedIdentityFunc.apply(self) + return super().expand_as(other) diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 92006ba45b..3b7b9bc169 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -48,8 +48,9 @@ class GroupedTensorStorage: Note: This structure is used only for combined storage of multiple tensors with the same dtype and scaling mode. """ - def __init__( - self, + @staticmethod + def _initialize_storage_fields( + instance: "GroupedTensorStorage", shape: Tuple[int, int], dtype: torch.dtype, num_tensors: int, @@ -68,6 +69,8 @@ def __init__( offsets: Optional[List[int]] = None, scale_inv_offsets: Optional[List[int]] = None, columnwise_scale_inv_offsets: Optional[List[int]] = None, + requires_grad: bool = False, + stride: Optional[List[int]] = None, ) -> None: """ Initialize a GroupedTensor. @@ -90,31 +93,37 @@ def __init__( tensor_offsets: Device tensor of int64 array of length num_tensors (or None if uniform) offsets: Vector of integer offsets for each tensor. """ - self.num_tensors = num_tensors - self.quantizer = quantizer - self.tensor_shapes = shapes - self.fake_dtype = dtype + # `requires_grad` and `stride` are accepted for API symmetry with + # GroupedTensor.__new__ but are not relevant for storage-only + # initialization; they are intentionally ignored here. + del requires_grad + del stride + + instance.num_tensors = num_tensors + instance.quantizer = quantizer + instance.tensor_shapes = shapes + instance.fake_dtype = dtype # Data buffers - self.rowwise_data = data - self.columnwise_data = columnwise_data - self.scale_inv = scale_inv - self.columnwise_scale_inv = columnwise_scale_inv - self.amax = amax - self.columnwise_amax = columnwise_amax - self.scale = scale + instance.rowwise_data = data + instance.columnwise_data = columnwise_data + instance.scale_inv = scale_inv + instance.columnwise_scale_inv = columnwise_scale_inv + instance.amax = amax + instance.columnwise_amax = columnwise_amax + instance.scale = scale # For convenient indexing for python GroupedTensor API. - self.scale_inv_offsets = scale_inv_offsets - self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets + instance.scale_inv_offsets = scale_inv_offsets + instance.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets # Shape information (OPTIONAL - None if dimension is uniform across all tensors) # first_dims[i] = first dimension of tensor i (None if all tensors have same first dim) # last_dims[i] = last dimension of tensor i (None if all tensors have same last dim) - self.first_dims = ( + instance.first_dims = ( first_dims # Device pointer to int64_t array of length num_tensors (or None) ) - self.last_dims = ( + instance.last_dims = ( last_dims # Device pointer to int64_t array of length num_tensors (or None) ) @@ -122,19 +131,69 @@ def __init__( # tensor_offsets[i] = element offset to start of tensor i (cumulative sum of numel for tensors 0..i-1) # Usage: tensor_i_ptr = data.data_ptr() + tensor_offsets[i] * element_size # If None and all_same_shape(): offset[i] = i * M * N (where M, N are common dimensions) - self.tensor_offsets = ( + instance.tensor_offsets = ( tensor_offsets # Device pointer to int64_t array of length num_tensors (or None) ) - self.offsets = offsets # Vector of integer offsets for each tensor. + instance.offsets = offsets # Vector of integer offsets for each tensor. # Logical shape: conceptual 2D shape of the grouped data (REQUIRED) # Represents how the 1D flattened data should be interpreted as 2D # Always 2D with positive dimensions - self.logical_shape = shape + instance.logical_shape = shape # Hold a reference to the quantized tensors that occupy same storage as the GroupedTensor. # Used as a convenience. - self.quantized_tensors = None + instance.quantized_tensors = None + + def __new__( + cls, + shape: Tuple[int, int], + dtype: torch.dtype, + *, + num_tensors: int, + shapes: Optional[List[Tuple[int, int]]] = None, + quantizer: Optional[Quantizer] = None, + data: Optional[torch.Tensor] = None, + columnwise_data: Optional[torch.Tensor] = None, + scale_inv: Optional[torch.Tensor] = None, + columnwise_scale_inv: Optional[torch.Tensor] = None, + amax: Optional[torch.Tensor] = None, + columnwise_amax: Optional[torch.Tensor] = None, + scale: Optional[torch.Tensor] = None, + first_dims: Optional[torch.Tensor] = None, + last_dims: Optional[torch.Tensor] = None, + tensor_offsets: Optional[torch.Tensor] = None, + offsets: Optional[List[int]] = None, + scale_inv_offsets: Optional[List[int]] = None, + columnwise_scale_inv_offsets: Optional[List[int]] = None, + requires_grad: bool = False, + stride: Optional[List[int]] = None, + ): + instance = object.__new__(cls) + cls._initialize_storage_fields( + instance=instance, + shape=shape, + dtype=dtype, + num_tensors=num_tensors, + shapes=shapes, + quantizer=quantizer, + data=data, + columnwise_data=columnwise_data, + scale_inv=scale_inv, + columnwise_scale_inv=columnwise_scale_inv, + amax=amax, + columnwise_amax=columnwise_amax, + scale=scale, + first_dims=first_dims, + last_dims=last_dims, + tensor_offsets=tensor_offsets, + offsets=offsets, + scale_inv_offsets=scale_inv_offsets, + columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, + requires_grad=requires_grad, + stride=stride, + ) + return instance def has_data(self) -> bool: """ From d226ce288b8dafa8ec7b51ed0c2962226ae8af82 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 5 Mar 2026 09:05:21 -0800 Subject: [PATCH 249/521] [JAX] Integrate BF16 Grouped GEMM with on-device group sizes (#2680) * Grouped GEMM Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Jeremy Berchtold * disable cuda-graph for GMM Signed-off-by: Jeremy Berchtold * proper workspace size Signed-off-by: Jeremy Berchtold * remove duplicate workspace size logic in Python gemm.py Signed-off-by: Jeremy Berchtold * use group_sizes as int32 and handle int64 and offsets inside FFI to avoid enabling JAX x64 globally Signed-off-by: Jeremy Berchtold * restore previous non-cuda-graphable grouped GEMM FFI and move new version to a different suffix Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup and lint fixes Signed-off-by: Jeremy Berchtold * re-add cublas alignment checks Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix symbol export when building with older cublas Signed-off-by: Jeremy Berchtold * Fix backend selection depending on whether TE was compiled with the right cuBLAS version Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Consolidate grouped GEMM primitives Signed-off-by: Jeremy Berchtold * fixes Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixes Signed-off-by: Jeremy Berchtold * Update C++ grouped GEMM tests to address row-major/col-major bugfix Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/jax/cpp_extensions/gemm.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review comments Signed-off-by: Jeremy Berchtold * Fix GMM runtime errors with nvte_set_grouped_tensor_param in gemm.cpp Signed-off-by: Jeremy Berchtold * Rename CudaGraphable to V2 and remove unnecessary FFI attributes Signed-off-by: Jeremy Berchtold * Lint Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- 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> --- tests/cpp/operator/test_grouped_gemm.cu | 8 +- .../common/gemm/cublaslt_grouped_gemm.cu | 73 +++- .../common/include/transformer_engine/gemm.h | 25 ++ transformer_engine/jax/cpp_extensions/gemm.py | 185 +++++++-- .../jax/cpp_extensions/quantization.py | 6 +- transformer_engine/jax/csrc/extensions.h | 1 + .../jax/csrc/extensions/gemm.cpp | 381 ++++++++++++++++++ .../jax/csrc/extensions/pybind.cpp | 5 + transformer_engine/jax/flax/__init__.py | 7 +- transformer_engine/jax/flax/module.py | 30 +- 10 files changed, 672 insertions(+), 49 deletions(-) diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu index a694052b15..a7aabbbcb6 100644 --- a/tests/cpp/operator/test_grouped_gemm.cu +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -123,10 +123,10 @@ void run_grouped_gemm_case(const TestParams& params) { 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{M, K} - : std::vector{K, M}; - const std::vector b_shape = params.transb ? std::vector{K, N} - : std::vector{N, K}; + 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)); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index b3e216dc4f..dc4757ab90 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -440,6 +440,29 @@ inline cublasLtMatmulAlgo_t select_grouped_gemm_algo(cublasLtHandle_t handle, return heuristicResult.algo; } +// Device helper: compute the element offset for tensor `idx` given shape metadata. +// Three cases: +// 1. Explicit per-tensor offset array provided → use it directly. +// 2. Per-tensor first/last dims provided but no offsets → cumulative sum of (first*last) products. +// 3. Fully uniform shapes → idx * uniform_first * uniform_last. +__forceinline__ __device__ int64_t compute_grouped_tensor_offset(const TensorShapeInfo &meta, + size_t idx) { + if (meta.offsets) { + return meta.offsets[idx]; + } else if (meta.first_dims != nullptr || meta.last_dims != nullptr) { + // offset[i] = sum_{j < i} (first_dims[j] * last_dims[j]) + int64_t cumsum = 0; + for (size_t i = 0; i < idx; i++) { + int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first; + int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last; + cumsum += f * l; + } + return cumsum; + } else { + return static_cast(idx) * meta.uniform_first * meta.uniform_last; + } +} + // Single kernel that sets up all GEMM parameters. // Rationale: cuBLASLt grouped matmul API needs flat arrays of pointers and per-matrix dimensions, // but NVTEGroupedTensor stores a single contiguous buffer + optional per-tensor offsets/shapes. @@ -464,15 +487,11 @@ __global__ void setup_grouped_gemm_kernel( int64_t d_first = D_meta.first_dims ? D_meta.first_dims[idx] : D_meta.uniform_first; int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last; - // Compute offsets (from array or compute from uniform dims) - int64_t a_offset = - A_meta.offsets ? A_meta.offsets[idx] : (idx * A_meta.uniform_first * A_meta.uniform_last); - int64_t b_offset = - B_meta.offsets ? B_meta.offsets[idx] : (idx * B_meta.uniform_first * B_meta.uniform_last); - int64_t c_offset = - C_meta.offsets ? C_meta.offsets[idx] : (idx * C_meta.uniform_first * C_meta.uniform_last); - int64_t d_offset = - D_meta.offsets ? D_meta.offsets[idx] : (idx * D_meta.uniform_first * D_meta.uniform_last); + // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform) + int64_t a_offset = compute_grouped_tensor_offset(A_meta, idx); + int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx); + int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx); + int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx); // Compute data pointers A_ptrs[idx] = a_base + a_offset * a_elem_size; @@ -487,9 +506,8 @@ __global__ void setup_grouped_gemm_kernel( a_cols[idx] = static_cast(a_first); b_rows[idx] = static_cast(b_last); b_cols[idx] = static_cast(b_first); - // For OUTPUTS (D, C): cuBLAS writes in column-major, so rows=first (M), cols=last (N). - d_rows[idx] = static_cast(d_first); - d_cols[idx] = static_cast(d_last); + d_rows[idx] = static_cast(d_last); + d_cols[idx] = static_cast(d_first); // Fill alpha/beta pointers (per-matrix) alpha_ptrs[idx] = alpha_ptr + idx; @@ -535,6 +553,11 @@ inline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) { } // namespace +size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) { + NVTE_API_CALL(nvte_get_grouped_gemm_setup_workspace_size); + return grouped_gemm_setup_workspace_size(num_tensors); +} + 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, @@ -642,4 +665,30 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 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.2+, but compile-time cuBLAS " + "version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); + return 0; +} + #endif // CUBLAS_VERSION >= 130200 + +namespace { + +__global__ void convert_int32_to_int64_kernel(const int32_t *src, int64_t *dst, size_t n) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) dst[idx] = static_cast(src[idx]); +} + +} // namespace + +void nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cudaStream_t stream) { + NVTE_API_CALL(nvte_convert_int32_to_int64); + if (n == 0) return; + const int threads = 256; + const int blocks = static_cast((n + threads - 1) / threads); + convert_int32_to_int64_kernel<<>>(src, dst, n); + NVTE_CHECK_CUDA(cudaGetLastError()); +} diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 7403448722..0f3b0ebd6b 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -329,6 +329,31 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor * - 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 + * by the grouped GEMM kernel. Its size depends only on the number of tensors (GEMMs) + * in the group and is independent of matrix dimensions. + * + * Pass the result as the size of the workspace_setup tensor in nvte_grouped_gemm. + * + * \param[in] num_tensors Number of tensors (GEMMs) in the group. + * \return Required size in bytes for workspace_setup. + */ +size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors); + +/*! \brief Convert a device array of int32 values to int64 values. + * + * Useful for preparing group_sizes for nvte_grouped_gemm when the caller + * holds int32 sizes and needs int64 values on the device. + * + * \param[in] src Device pointer to source int32 array. + * \param[out] dst Device pointer to destination int64 array. + * \param[in] n Number of elements. + * \param[in] stream CUDA stream. + */ +void nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cudaStream_t stream); + 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, diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index fbaafdf6d8..ab2be7f799 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -24,6 +24,7 @@ get_device_compute_capability, initialize_cgemm_communicator, get_cgemm_num_max_streams, + get_grouped_gemm_setup_workspace_size, ) from .base import BasePrimitive, register_primitive @@ -71,6 +72,18 @@ num_cublas_streams = get_num_compute_streams() +# Cache whether the CUDA-graphable grouped GEMM implementation is available at import time. +# Calling get_grouped_gemm_setup_workspace_size raises a RuntimeError mentioning "cublas" when +# compiled against cuBLAS < 13.2, in which case the cuda-graphable path is unavailable. +try: + get_grouped_gemm_setup_workspace_size(1) + _v2_grouped_gemm_available = True +except RuntimeError as e: + if "cublas" in str(e).lower(): + _v2_grouped_gemm_available = False + else: + raise + def get_cublas_workspace_size_bytes() -> None: """Return 32 MiB if using hopper, 4 MiB for all other architectures.""" @@ -591,7 +604,7 @@ def lowering( assert_cublas_requirements( scaling_mode, lhs_contracting_size, - "LHS", + f"LHS {lhs_aval.shape} with contracting dims {lhs_cdims}", ) rhs_axis_boundary = get_rhs_axis_boundary(rhs_cdims, rhs_transposed) rhs_contracting_size = ( @@ -602,7 +615,7 @@ def lowering( assert_cublas_requirements( scaling_mode, rhs_contracting_size, - "RHS", + f"RHS {rhs_aval.shape} with contracting dims {rhs_cdims}", ) args = (lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alpha, beta) @@ -1430,12 +1443,15 @@ def impl( class GroupedGemmPrimitive(BasePrimitive): """ - Primitive for grouped GEMM + Primitive for grouped GEMM using nvte_multi_tensor_gemm (supports all scaling modes) or nvte_grouped_gemm (supporting BF16). """ + # args = lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, group_sizes, group_offset, unused_placeholder name = "te_grouped_gemm_ffi" + # args = lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, group_sizes, alpha, beta + name_graph_safe = "te_grouped_gemm_v2_ffi" multiple_results = True - impl_static_args = (7, 8, 9, 10, 11, 12, 13, 14, 15, 16) + impl_static_args = (8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18) inner_primitive = None outer_primitive = None @@ -1447,8 +1463,7 @@ def abstract( rhs_scale_inv_aval, bias_aval, group_sizes_aval, - group_offset_aval, - *, + *additional_args, # group_offset_aval, unused_placeholder OR alpha_aval, beta_aval M, N, K, @@ -1459,6 +1474,7 @@ def abstract( has_bias, is_grouped_dense_wgrad, use_async_d2h_group_sizes, + use_v2_ffi, ): """ Grouped GEMM operation. @@ -1470,7 +1486,11 @@ def abstract( rhs_scale_inv: Right-hand side input scale_inv matrix, 1D flattened array bias: Bias matrix of shape (G, N) group_sizes: 1D array containing the sizes of each group - group_offset: 1D array containing offsets for each group (not yet implemented) + additional_args: Either + * group_offsets: 1D array containing offsets for each group (not yet implemented) + OR + * alpha: 1D array of shape (G,) containing alpha values for each group + * beta: 1D array of shape (G,) containing beta values for each group M: Number of rows in the output matrix N: Number of columns in the output matrix K: Number of columns in the left-hand side matrix @@ -1485,10 +1505,69 @@ def abstract( Returns: A jnp.ndarray containing the result of the grouped GEMM operation """ - del lhs_data_aval, rhs_data_aval, bias_aval, group_offset_aval + del lhs_data_aval, rhs_data_aval, bias_aval del K, lhs_is_trans, rhs_is_trans, has_bias, use_async_d2h_group_sizes + + num_groups = group_sizes_aval.size + + cublas_workspace_aval = jax.core.ShapedArray( + shape=( + GroupedGemmPrimitive._compute_cublas_workspace_size( + scaling_mode, lhs_scale_inv_aval, rhs_scale_inv_aval, use_v2_ffi + ), + ), + dtype=jnp.uint8, + ) + + out_shape = (M, N) + if is_grouped_dense_wgrad: + out_shape = (num_groups, M, N) + out_aval = jax.core.ShapedArray(shape=out_shape, dtype=out_dtype) + + if use_v2_ffi: + setup_workspace_aval = jax.core.ShapedArray( + shape=(get_grouped_gemm_setup_workspace_size(num_groups),), dtype=jnp.uint8 + ) + # Temporary buffer for int32 -> int64 conversion of group_sizes on device. + int64_workspace_size = num_groups * jnp.dtype(jnp.int64).itemsize + int64_workspace_aval = jax.core.ShapedArray( + shape=(int64_workspace_size,), dtype=jnp.uint8 + ) + + assert len(additional_args) == 2, ( + "Expected additional_args to contain alpha, beta for the graph-safe grouped GEMM" + f" primitive, but got {len(additional_args)} arguments." + ) + alpha_aval, beta_aval = additional_args + assert alpha_aval.shape == ( + num_groups, + ), f"Expected alpha shape {(num_groups,)}, got {alpha_aval.shape}" + assert ( + alpha_aval.dtype == jnp.float32 + ), f"Expected alpha dtype float32, got {alpha_aval.dtype}" + assert beta_aval.shape == ( + num_groups, + ), f"Expected beta shape {(num_groups,)}, got {beta_aval.shape}" + assert ( + beta_aval.dtype == jnp.float32 + ), f"Expected beta dtype float32, got {beta_aval.dtype}" + + return (out_aval, cublas_workspace_aval, setup_workspace_aval, int64_workspace_aval) + + return (out_aval, cublas_workspace_aval) + + @staticmethod + def _compute_cublas_workspace_size( + scaling_mode: ScalingMode, + lhs_scale_inv_aval, + rhs_scale_inv_aval, + use_v2_ffi: bool, + ): + """Compute the required cuBLAS workspace size based on the scaling mode and alignment requirements.""" + stream_count = 1 if use_v2_ffi else num_cublas_streams + # TODO(Phuong): move some shape checks from Cpp to here - workspace_size = get_cublas_workspace_size_bytes() * num_cublas_streams + workspace_size = get_cublas_workspace_size_bytes() * stream_count workspace_alignment_padding = 256 tensor_scaling_sinv_aligment = 16 mxfp8_scaling_sinv_alignment_padding = 256 @@ -1507,18 +1586,12 @@ def abstract( # We also pad scale_inv swizzle buffers size for 256 bytes alignment. workspace_size += lhs_scale_inv_aval.size + mxfp8_scaling_sinv_alignment_padding workspace_size += rhs_scale_inv_aval.size + mxfp8_scaling_sinv_alignment_padding - workspace_aval = jax.core.ShapedArray(shape=(workspace_size,), dtype=jnp.uint8) - - out_shape = (M, N) - if is_grouped_dense_wgrad: - out_shape = (group_sizes_aval.size, M, N) - out_aval = jax.core.ShapedArray(shape=out_shape, dtype=out_dtype) - return (out_aval, workspace_aval) + return workspace_size @staticmethod def outer_abstract(*args, **kwargs): - (out_aval, _) = GroupedGemmPrimitive.abstract(*args, **kwargs) - return (out_aval,) + (out, *_) = GroupedGemmPrimitive.abstract(*args, **kwargs) + return (out,) @staticmethod def lowering( @@ -1534,9 +1607,24 @@ def lowering( has_bias, is_grouped_dense_wgrad, use_async_d2h_group_sizes, + use_v2_ffi, ): del out_dtype - return jax.ffi.ffi_lowering(GroupedGemmPrimitive.name)( + if use_v2_ffi: + ffi_name = GroupedGemmPrimitive.name_graph_safe + return jax.ffi.ffi_lowering(ffi_name)( + ctx, + *args, + M=M, + N=N, + K=K, + lhs_is_trans=lhs_is_trans, + rhs_is_trans=rhs_is_trans, + scaling_mode=scaling_mode.value, + is_grouped_dense_wgrad=is_grouped_dense_wgrad, + ) + ffi_name = GroupedGemmPrimitive.name + return jax.ffi.ffi_lowering(ffi_name)( ctx, *args, M=M, @@ -1558,7 +1646,8 @@ def impl( rhs_scale_inv, bias, group_sizes, - group_offset, + additional_arg_0, # group_offset (non-graph-safe) OR alpha (graph-safe) + additional_arg_1, # unused placeholder (non-graph-safe) OR beta (graph-safe) M, N, K, @@ -1569,16 +1658,21 @@ def impl( has_bias, is_grouped_dense_wgrad, use_async_d2h_group_sizes, + use_v2_ffi, ): assert GroupedGemmPrimitive.inner_primitive is not None - (out, _) = GroupedGemmPrimitive.inner_primitive.bind( + if use_v2_ffi: + additional_args = (additional_arg_0, additional_arg_1) + else: + additional_args = (additional_arg_0,) + (out, *_) = GroupedGemmPrimitive.inner_primitive.bind( lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, group_sizes, - group_offset, + *additional_args, M=M, N=N, K=K, @@ -1589,6 +1683,7 @@ def impl( has_bias=has_bias, is_grouped_dense_wgrad=is_grouped_dense_wgrad, use_async_d2h_group_sizes=use_async_d2h_group_sizes, + use_v2_ffi=use_v2_ffi, ) return (out,) @@ -1910,6 +2005,23 @@ def grouped_gemm_copy_group_sizes( return out +def _can_use_v2_grouped_gemm( + scaling_mode: ScalingMode, + dtype: jnp.dtype, + has_bias: bool, +) -> bool: + """Determine whether the cuda-graphable grouped GEMM implementation can be used based on the input parameters.""" + # Use the cuda-graphable path for plain BF16 non-quantized inputs; fall back to the legacy + # nvte_multi_tensor_gemm path for all other cases (FP8, MXFP8, etc.) to stay + # feature-compatible with the main branch. + # Bias can be supported in a kernel or in pure-JAX in the future. + + if not _v2_grouped_gemm_available: + return False + + return scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16 and not has_bias + + def grouped_gemm( lhs: Union[jnp.ndarray, GroupedScaledTensor1x], rhs: Union[jnp.ndarray, GroupedScaledTensor1x], @@ -1944,8 +2056,6 @@ def grouped_gemm( lhs: [M, K] or [K, N] rhs: [G, N, K] or [G, K, N] or [G * K, N] or [N, G * K] """ - # TODO(Phuong): implement the group_offset - group_offset = group_offset or jnp.zeros((1,), jnp.int32) # TODO(Phuong): implement the precision del precision @@ -2081,12 +2191,29 @@ def grouped_gemm( else: assert group_sizes.size == rhs_shape[0] - assert group_offset.size == 1 - has_bias = bias is not None - assert not has_bias or bias.shape == (group_sizes.size, N) + if has_bias: + assert bias.shape == ( + group_sizes.size, + N, + ), f"bias shape {bias.shape} does not match expected shape {(group_sizes.size, N)}" bias = jnp.empty((), jnp.float32) if bias is None else bias + assert group_offset is None, ( + "group_offset is not supported yet and is instead computed" + " internally assuming contiguous grouping. Any padding is included in the group_sizes" + " and padded with zeros to not affect the result of the MoE block." + ) + + use_v2_ffi = _can_use_v2_grouped_gemm(scaling_mode, lhs_data.dtype, has_bias) + if use_v2_ffi: + num_gemms = group_sizes.shape[0] + additional_arg_0 = jnp.ones((num_gemms,), jnp.float32) # alpha + additional_arg_1 = jnp.zeros((num_gemms,), jnp.float32) # beta + else: + additional_arg_0 = jnp.zeros((1,), jnp.int32) # group_offset + additional_arg_1 = jnp.zeros((0,), jnp.int32) # unused placeholder + (out,) = GroupedGemmPrimitive.outer_primitive.bind( lhs_data, lhs_scale_inv, @@ -2094,7 +2221,8 @@ def grouped_gemm( rhs_scale_inv, bias, group_sizes, - group_offset, + additional_arg_0, + additional_arg_1, M=M, N=N, K=K_lhs, @@ -2105,5 +2233,6 @@ def grouped_gemm( has_bias=has_bias, is_grouped_dense_wgrad=is_grouped_dense_wgrad, use_async_d2h_group_sizes=use_async_d2h_group_sizes, + use_v2_ffi=use_v2_ffi, ) return out diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index 1fcecb0e96..bf4e833c89 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -97,7 +97,9 @@ def abstract( dtype = dtypes.canonicalize_dtype(x_aval.dtype) assert dtype in [jnp.float32, jnp.float16, jnp.bfloat16] out_shape = x_aval.shape - assert scale_aval is None or scale_aval.dtype == jnp.float32 + assert ( + scale_aval is None or scale_aval.dtype == jnp.float32 + ), f"scale must be float32 but received {scale_aval}" if stochastic_rounding: assert ScalingMode( scaling_mode @@ -1213,7 +1215,7 @@ def grouped_quantize( assert n_groups == len( quantizer.quantizers ), f"n_groups={n_groups} != n_quantizers = {len(quantizer.quantizers)}" - scale = jnp.empty((n_groups,), jnp.float32) + scale = jnp.ones((n_groups,), jnp.float32) if quantizer.scaling_mode == ScalingMode.DELAYED_TENSOR_SCALING: for i, quantizer_i in enumerate(quantizer.quantizers): diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index e844100c5a..93c85aaacc 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -138,6 +138,7 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(CollectiveGemmInitHandler); // Grouped GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmD2HGroupSizesHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmV2Handler); // Amax XLA_FFI_DECLARE_HANDLER_SYMBOL(RHTAmaxCalculationInitializeHandler); diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 4303682bfb..4cbec405a4 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -409,6 +409,387 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmD2HGroupSizesHandler, GroupedGemmD2HGro .Ret() // dummy_output .Attr("num_gemms")); +class JAXX_GroupedTensorWrapper { + public: + JAXX_GroupedTensorWrapper() = delete; + JAXX_GroupedTensorWrapper(JAXX_Scaling_Mode scaling_mode, size_t num_tensors, + NVTEShape const &dataShape); + JAXX_GroupedTensorWrapper(JAXX_GroupedTensorWrapper const &) = delete; + JAXX_GroupedTensorWrapper &operator=(JAXX_GroupedTensorWrapper const &) = delete; + JAXX_GroupedTensorWrapper(JAXX_GroupedTensorWrapper &&other) noexcept + : m_data_shape(other.m_data_shape), + m_grouped_tensor(other.m_grouped_tensor), + m_data_tensor(other.m_data_tensor), + m_scale_inv_tensor(other.m_scale_inv_tensor), + m_sizes_tensor(other.m_sizes_tensor), + m_offsets_tensor(other.m_offsets_tensor) { + other.m_grouped_tensor = nullptr; + } + JAXX_GroupedTensorWrapper &operator=(JAXX_GroupedTensorWrapper &&) = delete; + ~JAXX_GroupedTensorWrapper(); + + void set_rowwise(Buffer_Type const &data, std::optional const &scale_inv); + void set_group_info(Buffer_Type const &group_sizes, Buffer_Type const &group_offsets, + NVTEGroupedTensorParam group_sizes_param_name); + // Set only group sizes (no offsets); the setup kernel will compute offsets from sizes. + void set_group_sizes_only(const int64_t *sizes_ptr, size_t num_tensors, + NVTEGroupedTensorParam group_sizes_param_name); + + operator NVTEGroupedTensor() const { return m_grouped_tensor; } + NVTEGroupedTensor const &get_grouped_tensor() const; + + private: + NVTEShape m_data_shape{}; + NVTEGroupedTensor m_grouped_tensor{}; + + // Internal tensors. These need to be kept alive as long as the grouped tensor is alive. + NVTEBasicTensor m_data_tensor{}; + NVTEBasicTensor m_scale_inv_tensor{}; + + NVTEBasicTensor m_sizes_tensor{}; + NVTEBasicTensor m_offsets_tensor{}; +}; + +JAXX_GroupedTensorWrapper::JAXX_GroupedTensorWrapper(JAXX_Scaling_Mode scaling_mode, + size_t num_tensors, + NVTEShape const &dataShape) { + m_data_shape = dataShape; + m_grouped_tensor = + nvte_create_grouped_tensor(get_nvte_scaling_mode(scaling_mode), num_tensors, dataShape); +} + +JAXX_GroupedTensorWrapper::~JAXX_GroupedTensorWrapper() { + if (m_grouped_tensor != nullptr) { + nvte_destroy_grouped_tensor(m_grouped_tensor); + } +} + +void JAXX_GroupedTensorWrapper::set_rowwise(Buffer_Type const &data, + std::optional const &scale_inv) { + NVTEDType data_dtype = + static_cast(convert_ffi_datatype_to_te_dtype(data.element_type())); + m_data_tensor = + NVTEBasicTensor{reinterpret_cast(data.untyped_data()), data_dtype, m_data_shape}; + + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedRowwiseData, &m_data_tensor, + sizeof(m_data_tensor)); + + if (scale_inv.has_value()) { + NVTEDType scale_inv_dtype = + static_cast(convert_ffi_datatype_to_te_dtype(scale_inv->element_type())); + NVTEShape logical_scale_shape{}; + if (scale_inv->dimensions().size() == 1) { + logical_scale_shape.ndim = 1; + logical_scale_shape.data[0] = scale_inv->dimensions()[0]; + } else if (scale_inv->dimensions().size() == 2) { + logical_scale_shape.ndim = 2; + logical_scale_shape.data[0] = scale_inv->dimensions()[0]; + logical_scale_shape.data[1] = scale_inv->dimensions()[1]; + } else { + NVTE_CHECK(false, "Expected 1D or 2D tensor for GEMM scale_inv but received ndim=", + scale_inv->dimensions().size()); + } + m_scale_inv_tensor = NVTEBasicTensor{reinterpret_cast(scale_inv->untyped_data()), + scale_inv_dtype, logical_scale_shape}; + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedRowwiseScaleInv, + &m_scale_inv_tensor, sizeof(m_scale_inv_tensor)); + } +} + +void JAXX_GroupedTensorWrapper::set_group_info(Buffer_Type const &group_sizes, + Buffer_Type const &group_offsets, + NVTEGroupedTensorParam group_sizes_param_name) { + NVTEDType sizes_dtype = + static_cast(convert_ffi_datatype_to_te_dtype(group_sizes.element_type())); + NVTEDType offsets_dtype = + static_cast(convert_ffi_datatype_to_te_dtype(group_offsets.element_type())); + + NVTE_CHECK(sizes_dtype == NVTEDType::kNVTEInt64, "group_sizes must be of type int64."); + NVTE_CHECK(offsets_dtype == NVTEDType::kNVTEInt64, "group_offsets must be of type int64."); + + size_t num_tensors = group_sizes.dimensions()[0]; + NVTE_CHECK(group_sizes.dimensions().size() == 1, + "group_sizes must be a 1D tensor with length equal to the number of tensors."); + NVTE_CHECK(group_offsets.dimensions().size() == 1, + "group_offsets must be a 1D tensor with length equal to the number of tensors."); + NVTE_CHECK(group_offsets.dimensions()[0] == num_tensors, + "group_sizes and group_offsets must have the same number of elements."); + + NVTEShape shape{}; + shape.ndim = 1; + shape.data[0] = num_tensors; + + m_sizes_tensor = NVTEBasicTensor{reinterpret_cast(group_sizes.untyped_data()), + NVTEDType::kNVTEInt64, shape}; + m_offsets_tensor = NVTEBasicTensor{reinterpret_cast(group_offsets.untyped_data()), + NVTEDType::kNVTEInt64, shape}; + + nvte_set_grouped_tensor_param(m_grouped_tensor, group_sizes_param_name, &m_sizes_tensor, + sizeof(m_sizes_tensor)); + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedTensorOffsets, &m_offsets_tensor, + sizeof(m_offsets_tensor)); +} + +void JAXX_GroupedTensorWrapper::set_group_sizes_only( + const int64_t *sizes_ptr, size_t num_tensors, NVTEGroupedTensorParam group_sizes_param_name) { + NVTEShape shape{}; + shape.ndim = 1; + shape.data[0] = num_tensors; + m_sizes_tensor = NVTEBasicTensor{reinterpret_cast(const_cast(sizes_ptr)), + NVTEDType::kNVTEInt64, shape}; + nvte_set_grouped_tensor_param(m_grouped_tensor, group_sizes_param_name, &m_sizes_tensor, + sizeof(m_sizes_tensor)); + // Intentionally no offset tensor: offsets will be computed by the setup kernel. +} + +NVTEGroupedTensor const &JAXX_GroupedTensorWrapper::get_grouped_tensor() const { + return m_grouped_tensor; +} + +JAXX_GroupedTensorWrapper make_grouped_tensor(Buffer_Type const &data, + std::optional scale_inv, + JAXX_Scaling_Mode scaling_mode, size_t num_tensors, + NVTEShape const &dataShape) { + JAXX_GroupedTensorWrapper grouped_tensor_wrapper(scaling_mode, num_tensors, dataShape); + if (scaling_mode == JAXX_Scaling_Mode::NO_SCALING) { + scale_inv = std::nullopt; + } + grouped_tensor_wrapper.set_rowwise(data, scale_inv); + + return std::move(grouped_tensor_wrapper); +} + +// This FFI is EXPERIMENTAL and subject to change without deprecation, intended for use in JAX's internal implementation of grouped GEMM. +Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, + Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, + Buffer_Type group_sizes, Buffer_Type alpha, Buffer_Type beta, + Result_Type output, Result_Type cublas_workspace, + Result_Type setup_workspace, Result_Type int64_workspace, size_t m, + size_t n, size_t k, bool lhs_is_trans, bool rhs_is_trans, + JAXX_Scaling_Mode scaling_mode, bool is_grouped_dense_wgrad) { + // Notes on matrix layouts and transpose: + // Jax uses row-major data_layout, on entering this function, each input matrix pair: + // A: row-major [m, k] for N - [k, m] for T + // B: row-major [k, n] for N - [n, k] for T + // on exiting this function, JAX expect: + // C: row-major with size [m, n]. + // cuBLAS uses column-major data_layout, in this view, each input matrix pair: + // A: column-major with size [k, m] for T - [m, k] for N + // B: column-major with size [n, k] for T - [k, n] for N + // + // If we call cuBLAS GEMM for A * B, the output will be: + // C: column-major with size [m, n] --> row-major with size [n, m]. + // To make the output compatible with JAX, we need to swap A and B in cuBLAS GEMM call. + + // Inputs + auto lhs_ptr = reinterpret_cast(lhs_data.untyped_data()); + auto rhs_ptr = reinterpret_cast(rhs_data.untyped_data()); + auto lhs_sinv_ptr = reinterpret_cast(lhs_sinv.untyped_data()); + auto rhs_sinv_ptr = reinterpret_cast(rhs_sinv.untyped_data()); + auto lhs_dtype = convert_ffi_datatype_to_te_dtype(lhs_data.element_type()); + auto rhs_dtype = convert_ffi_datatype_to_te_dtype(rhs_data.element_type()); + auto lhs_sinv_dtype = convert_ffi_datatype_to_te_dtype(lhs_sinv.element_type()); + auto rhs_sinv_dtype = convert_ffi_datatype_to_te_dtype(rhs_sinv.element_type()); + bool has_bias = product(bias.dimensions()) > 0; + auto bias_ptr = has_bias ? reinterpret_cast(bias.untyped_data()) : nullptr; + auto bias_dtype = convert_ffi_datatype_to_te_dtype(bias.element_type()); + + NVTE_CHECK(group_sizes.dimensions().size() == 1); + size_t num_gemms = group_sizes.dimensions()[0]; + + // Convert int32 group_sizes to int64 into the dedicated output buffer. + NVTE_CHECK(group_sizes.element_type() == xla::ffi::DataType::S32, "group_sizes must be int32."); + auto *int64_sizes_ptr = reinterpret_cast(int64_workspace->untyped_data()); + nvte_convert_int32_to_int64(reinterpret_cast(group_sizes.untyped_data()), + int64_sizes_ptr, num_gemms, stream); + + NVTE_CHECK(scaling_mode == JAXX_Scaling_Mode::NO_SCALING, + "Only non-quantized grouped GEMM is supported in current implementation."); + + // It is weird that TE/Common GEMM only use colwise for MXFP8 + const bool is_fp8_gemm = is_fp8_dtype(lhs_dtype); + const bool is_tensor_scaling = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING || + scaling_mode == JAXX_Scaling_Mode::CURRENT_TENSOR_SCALING; + const bool is_mxfp8_scaling = scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING; + const bool rhs_use_colwise = is_mxfp8_scaling && !rhs_is_trans; + const bool lhs_use_colwise = is_mxfp8_scaling && lhs_is_trans; + + // Outputs + auto out_ptr = reinterpret_cast(output->untyped_data()); + auto out_dtype = convert_ffi_datatype_to_te_dtype(output->element_type()); + auto setup_workspace_ptr = reinterpret_cast(setup_workspace->untyped_data()); + // Here we clear the lower 8 bits of the buffer address to ensure the buffer is 256-aligned + auto cublas_workspace_ptr = reinterpret_cast(cublas_workspace->untyped_data()); + cublas_workspace_ptr = move_ptr_to_next_256B_aligned(cublas_workspace_ptr); + auto workspace_total_size = product(cublas_workspace->dimensions()); + + auto lhs_sinv_size = product(lhs_sinv.dimensions()); + auto rhs_sinv_size = product(rhs_sinv.dimensions()); + const size_t workspace_alignment_padding = 256; + const size_t tensor_scaling_sinv_aligment = 16; + const size_t mxfp8_scaling_sinv_alignment_padding = 256; + auto workspace_size = workspace_total_size - workspace_alignment_padding; + if (is_mxfp8_scaling) { + // For MXFP8 swizzled scale_inv buffers, only the first pointer needs to be with 256B alignment padding. Later pointers are guaranteed to be 256-aligned as the scale_inv shapes are padded by 128x4. + workspace_size -= (lhs_sinv_size + rhs_sinv_size + 2 * mxfp8_scaling_sinv_alignment_padding); + } else if (is_tensor_scaling) { + // For tensor scaling, each matrix has a single scale value, and all scales need to be aligned + // by 16 bytes to meet the requirement of CUDA 12.9.1 and later. + workspace_size -= tensor_scaling_sinv_aligment * (lhs_sinv_size + rhs_sinv_size); + } + auto swizzled_lhs_sinv_ptr = cublas_workspace_ptr + workspace_size; + swizzled_lhs_sinv_ptr = move_ptr_to_next_256B_aligned(swizzled_lhs_sinv_ptr); + auto swizzled_rhs_sinv_ptr = swizzled_lhs_sinv_ptr + lhs_sinv_size; + swizzled_rhs_sinv_ptr = move_ptr_to_next_256B_aligned(swizzled_rhs_sinv_ptr); + auto lhs_scatter_aligned_ptr = swizzled_lhs_sinv_ptr; // Already 256B aligned + auto rhs_scatter_aligned_ptr = lhs_scatter_aligned_ptr + num_gemms * tensor_scaling_sinv_aligment; + + size_t lhs_dtype_bytes = te_dtype_bytes(lhs_dtype); + size_t rhs_dtype_bytes = te_dtype_bytes(rhs_dtype); + size_t lhs_sinv_dtype_bytes = te_dtype_bytes(lhs_sinv_dtype); + size_t rhs_sinv_dtype_bytes = te_dtype_bytes(rhs_sinv_dtype); + size_t bias_dtype_bytes = te_dtype_bytes(bias_dtype); + size_t out_dtype_bytes = te_dtype_bytes(out_dtype); + + NVTE_CHECK(lhs_dtype_bytes == rhs_dtype_bytes, "sizeof(lhs_dtype) != sizeof(rhs_dtype)"); + NVTE_CHECK(lhs_sinv_dtype_bytes == rhs_sinv_dtype_bytes, + "sizeof(lhs_sinv_dtype) != sizeof(rhs_sinv_dtype)"); + + size_t expected_lhs_size = m * k; + size_t expected_rhs_size = is_grouped_dense_wgrad ? (k * n) : (num_gemms * k * n); + size_t expected_out_size = is_grouped_dense_wgrad ? (num_gemms * m * n) : (m * n); + size_t actual_lhs_size = product(lhs_data.dimensions()); + size_t actual_rhs_size = product(rhs_data.dimensions()); + size_t actual_out_size = product(output->dimensions()); + NVTE_CHECK(expected_lhs_size == actual_lhs_size, "Unexpected lhs size! Expect ", + expected_lhs_size, ", got ", actual_lhs_size); + if (!is_grouped_dense_wgrad) { + NVTE_CHECK(expected_rhs_size == actual_rhs_size, + "Unexpected rhs size! Expect num_gemms * n * k = ", num_gemms, " * ", n, " * ", k, + " = ", expected_rhs_size, ", got ", actual_rhs_size); + NVTE_CHECK(expected_out_size == actual_out_size, "Unexpected output size! Expect m * n = ", m, + " * ", n, " = ", expected_out_size, ", got ", actual_out_size); + } else { + NVTE_CHECK(expected_rhs_size == actual_rhs_size, "Unexpected rhs size! Expect k * n = ", k, + " * ", n, " = ", expected_rhs_size, ", got ", actual_rhs_size); + NVTE_CHECK(expected_out_size == actual_out_size, + "Unexpected output size! Expect num_gemms * m * n = ", num_gemms, " * ", m, " * ", n, + " = ", expected_out_size, ", got ", actual_out_size); + } + + auto num_math_sm = cuda::sm_count() - getenv("NVTE_EXT_MARGIN_SM", 0); + bool grad = false; + bool accumulate = false; + bool use_split_accumulator = false; + auto bias_shape = std::vector{has_bias ? n : 0}; + const int arch = cuda::sm_arch(); + + if (arch < 100 && is_fp8_gemm) { + NVTE_CHECK(!lhs_is_trans && rhs_is_trans, + "For SM90 or older archs and FP8 input, only NT (row-major) GEMM is supported, ", + "got lhs_is_trans=", lhs_is_trans, ", rhs_is_trans=", rhs_is_trans); + } + + TensorWrapper workspace_setup(setup_workspace_ptr, + std::vector{product(setup_workspace->dimensions())}, + DType::kByte); + TensorWrapper workspace_cublas(cublas_workspace_ptr, std::vector{workspace_size}, + DType::kByte); + + TensorWrapper alpha_tensor(static_cast(alpha.untyped_data()), + std::vector{num_gemms}, + convert_ffi_datatype_to_te_dtype(alpha.element_type())); + TensorWrapper beta_tensor(static_cast(beta.untyped_data()), + std::vector{num_gemms}, + convert_ffi_datatype_to_te_dtype(beta.element_type())); + + if (is_grouped_dense_wgrad) { + NVTE_CHECK(lhs_is_trans && !rhs_is_trans, + "For grouped dense wgrad, only TN GEMM is supported in TE/JAX currently."); + + //// RHS + NVTEShape rhsShape{.data = {k, n}, .ndim = 2}; + auto rhs_tensor = make_grouped_tensor(rhs_data, rhs_sinv, scaling_mode, num_gemms, rhsShape); + rhs_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, kNVTEGroupedFirstDims); + + //// LHS + NVTEShape lhsShape{.data = {k, m}, .ndim = 2}; + lhs_is_trans = true; + auto lhs_tensor = make_grouped_tensor(lhs_data, lhs_sinv, scaling_mode, num_gemms, lhsShape); + lhs_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, kNVTEGroupedFirstDims); + + //// OUTPUT + NVTEShape outShape{.data = {num_gemms * m, n}, .ndim = 2}; + auto out_tensor = make_grouped_tensor(*output, std::nullopt, JAXX_Scaling_Mode::NO_SCALING, + num_gemms, outShape); + + nvte_grouped_gemm(rhs_tensor, rhs_is_trans, lhs_tensor, lhs_is_trans, nullptr, out_tensor, + alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), + workspace_cublas.data(), + nullptr, // config (use defaults) + stream); + + return ffi_with_cuda_error_check(); + } + + // Nominal case for FWD or DGRAD + + //// RHS + NVTEShape rhsShape{.data = {num_gemms * k, n}, .ndim = 2}; + if (rhs_is_trans) { + rhsShape.data[0] = num_gemms * n; + rhsShape.data[1] = k; + } + auto rhs_tensor = make_grouped_tensor(rhs_data, rhs_sinv, scaling_mode, num_gemms, rhsShape); + + //// LHS + NVTEShape lhsShape{.data = {m, k}, .ndim = 2}; + if (lhs_is_trans) { + std::swap(lhsShape.data[0], lhsShape.data[1]); + } + auto lhs_tensor = make_grouped_tensor(lhs_data, lhs_sinv, scaling_mode, num_gemms, lhsShape); + lhs_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, + lhs_is_trans ? kNVTEGroupedLastDims : kNVTEGroupedFirstDims); + + //// OUTPUT + NVTEShape outShape{.data = {m, n}, .ndim = 2}; + auto out_tensor = make_grouped_tensor(*output, std::nullopt, JAXX_Scaling_Mode::NO_SCALING, + num_gemms, outShape); + out_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, kNVTEGroupedFirstDims); + + nvte_grouped_gemm(rhs_tensor, rhs_is_trans, lhs_tensor, lhs_is_trans, nullptr, out_tensor, + alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), + workspace_cublas.data(), + nullptr, // config (use defaults) + stream); + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, + FFI::Bind() + .Ctx() // stream + .Arg() // lhs_data + .Arg() // lhs_sinv + .Arg() // rhs_data + .Arg() // rhs_sinv + .Arg() // bias + .Arg() // group_sizes (int32) + .Arg() // alpha + .Arg() // beta + .Ret() // output + .Ret() // cublas_workspace + .Ret() // setup_workspace + .Ret() // int64_workspace + .Attr("M") + .Attr("N") + .Attr("K") + .Attr("lhs_is_trans") + .Attr("rhs_is_trans") + .Attr("scaling_mode") + .Attr("is_grouped_dense_wgrad"), + FFI_CudaGraph_Traits); + Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, Buffer_Type group_sizes, Buffer_Type group_offset, Result_Type output, diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index deea64caa1..837dd55f9c 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -7,6 +7,7 @@ #include "../extensions.h" #include "cgemm_helper.h" #include "common/util/cuda_runtime.h" +#include "transformer_engine/gemm.h" namespace transformer_engine { namespace jax { @@ -75,6 +76,9 @@ pybind11::dict Registrations() { dict["te_grouped_gemm_ffi"] = pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CublasHandleInitHandler), pybind11::arg("execute") = EncapsulateFFI(GroupedGemmHandler)); + dict["te_grouped_gemm_v2_ffi"] = + pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CublasHandleInitHandler), + pybind11::arg("execute") = EncapsulateFFI(GroupedGemmV2Handler)); // Amax dict["te_rht_amax_ffi"] = pybind11::dict( @@ -113,6 +117,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); m.def("initialize_cgemm_communicator", &InitializeCgemmCommunicator); m.def("get_cgemm_num_max_streams", &GetCgemmNumMaxStreams); + m.def("get_grouped_gemm_setup_workspace_size", &nvte_get_grouped_gemm_setup_workspace_size); pybind11::enum_(m, "DType", pybind11::module_local()) .value("kByte", DType::kByte) diff --git a/transformer_engine/jax/flax/__init__.py b/transformer_engine/jax/flax/__init__.py index dd7d2a47ba..92a968f061 100644 --- a/transformer_engine/jax/flax/__init__.py +++ b/transformer_engine/jax/flax/__init__.py @@ -4,7 +4,11 @@ """Transformer Engine bindings for JAX""" from .module import DenseGeneral, LayerNorm from .module import LayerNormDenseGeneral, LayerNormMLP -from .module import wrap_function_in_te_state_module, make_dot_general_cls +from .module import ( + wrap_function_in_te_state_module, + make_dot_general_cls, + make_grouped_dense_cls, +) from .transformer import extend_logical_axis_rules from .transformer import DotProductAttention, MultiHeadAttention, RelativePositionBiases from .transformer import TransformerLayer, TransformerLayerType @@ -16,6 +20,7 @@ "LayerNormMLP", "wrap_function_in_te_state_module", "make_dot_general_cls", + "make_grouped_dense_cls", "extend_logical_axis_rules", "DotProductAttention", "MultiHeadAttention", diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 3d82d8f0b4..7decfca6c6 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -17,7 +17,7 @@ from jax.ad_checkpoint import checkpoint_name -from ..dense import dense +from ..dense import dense, grouped_dense from ..layernorm import canonicalize_norm_type from ..layernorm import layernorm @@ -377,6 +377,7 @@ def generate_quantizer_set( variable_collection: str = None, quantization_checkpoint_name: Optional[str] = None, fp8_recipe=None, + n_groups: int = None, ): """ Generate a set of FP8 meta for a GEMM. @@ -409,6 +410,7 @@ def generate_quantizer_set( fp8_recipe=fp8_recipe, quantize_meta_set=quantize_meta_set, checkpoint_name=quantization_checkpoint_name, + n_groups=n_groups, ) return quantizer_set @@ -1379,12 +1381,13 @@ def wrap_function_in_te_state_module(f, quantization_recipe, name: Optional[str] class TEWrapper(te.flax.module.TransformerEngineBase): """Wrapper Flax module for TransformerEngine quantization support.""" - def generate_quantizer_set(self, postfix: str = ""): + def generate_quantizer_set(self, postfix: str = "", n_groups: int = None): OVERWRITE_WITH_GRADIENT = "_overwrite_with_gradient" return super().generate_quantizer_set( postfix=postfix, variable_collection=OVERWRITE_WITH_GRADIENT, fp8_recipe=quantization_recipe, + n_groups=n_groups, ) @nn.compact @@ -1438,3 +1441,26 @@ def te_dot_general(generate_quantizer_set, x, kernel, dims, **kwargs): ) return wrap_function_in_te_state_module(te_dot_general, quantization_recipe, "dot_general") + + +def make_grouped_dense_cls(quantization_recipe): + """Creates a grouped dense (grouped GEMM) instance for use with TE state module.""" + assert quantization_recipe is None, "Ragged dot grouped GEMM does not support quantization yet" + + def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwargs): + del kwargs # Unused + num_groups = group_sizes.shape[0] + quantizer_set = generate_quantizer_set(n_groups=num_groups) + + out = grouped_dense( + x, + kernel, + group_sizes=group_sizes, + contracting_dims=((1,), (1,)), + quantizer_set=quantizer_set, + ) + return out + + return wrap_function_in_te_state_module( + te_grouped_dot_general, quantization_recipe, "ragged_dot" + )() From d40b9de312c3cacec73b10a71b8d8419a11824c8 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 5 Mar 2026 14:08:55 -0800 Subject: [PATCH 250/521] WAR sort_chunks_by_index intermittent failures in L0 JAX unitttest part 1 (#2730) * Add WAR to alias input (output_grads) to output (input_grads) of the sort_chunk_by_map bwd function Signed-off-by: tdophung * adapt pytorch call site to extra dummy pointer Signed-off-by: tdophung --------- Signed-off-by: tdophung --- .../common/triton/permutation.py | 4 +++ .../jax/triton_extensions/permutation.py | 28 +++++++++++++++++-- .../pytorch/triton/permutation.py | 1 + 3 files changed, 31 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index 147742bb05..75bb85f5ec 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -599,6 +599,10 @@ def _sort_chunks_by_map_kernel( input_ptr, row_id_map_ptr, probs_ptr, + # Pre-allocated output buffer for JAX input_output_aliases. + # Aliased to output_ptr in JAX so they point to the same memory. + # In PyTorch, pass the same tensor as output_ptr. + output_buf_ptr, # pylint: disable=unused-argument # strides stride_input_token, stride_input_hidden, diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index 0c80f9f18c..98c54e52bb 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -1666,10 +1666,19 @@ class SortChunksByMapPrimitive(BasePrimitive): @staticmethod def abstract( - inp_aval, row_id_map_aval, probs_aval, *, num_tokens, hidden_size, is_forward, with_probs + inp_aval, + row_id_map_aval, + probs_aval, + output_buf_aval=None, # Pre-allocated output buffer (inner primitive only) + *, + num_tokens, + hidden_size, + is_forward, + with_probs, ): """Shape/dtype inference.""" del row_id_map_aval, is_forward + del output_buf_aval # Used for input_output_aliases only output_aval = jax.core.ShapedArray((num_tokens, hidden_size), inp_aval.dtype) @@ -1684,10 +1693,14 @@ def abstract( def impl(inp, row_id_map, probs, num_tokens, hidden_size, is_forward, with_probs): """Forward to inner primitive.""" assert SortChunksByMapPrimitive.inner_primitive is not None + + output_buf = jnp.empty((num_tokens, hidden_size), dtype=inp.dtype) + return SortChunksByMapPrimitive.inner_primitive.bind( inp, row_id_map, probs, + output_buf, num_tokens=num_tokens, hidden_size=hidden_size, is_forward=is_forward, @@ -1695,7 +1708,9 @@ def impl(inp, row_id_map, probs, num_tokens, hidden_size, is_forward, with_probs ) @staticmethod - def lowering(ctx, inp, row_id_map, probs, *, num_tokens, hidden_size, is_forward, with_probs): + def lowering( + ctx, inp, row_id_map, probs, output_buf, *, num_tokens, hidden_size, is_forward, with_probs + ): """MLIR lowering using triton_call_lowering.""" # Compute strides inp_stride_token = hidden_size @@ -1709,13 +1724,22 @@ def lowering(ctx, inp, row_id_map, probs, *, num_tokens, hidden_size, is_forward 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 + # other input (like output_grad in backward) to the output buffer. + # Input indices: 0=inp, 1=row_id_map, 2=probs, 3=output_buf + # Output indices: 0=output, 1=permuted_probs + input_output_aliases = {3: 0} + return triton_call_lowering( ctx, _sort_chunks_by_map_kernel, inp, row_id_map, probs, + output_buf, grid=grid, + input_output_aliases=input_output_aliases, constexprs={ "stride_input_token": inp_stride_token, "stride_input_hidden": inp_stride_hidden, diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 6b5de9ab0f..4902bc686c 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -427,6 +427,7 @@ def sort_chunks_by_map( inp, row_id_map, probs, + output, # no use in Pytorch side, serves as WAR for JAX side inp.stride(0), inp.stride(1), output.stride(0), From 5fd5c35780e5f7220024ffa5285aa292ae1d25de Mon Sep 17 00:00:00 2001 From: Chen Cui Date: Sun, 8 Mar 2026 12:30:33 -0700 Subject: [PATCH 251/521] Fix FP8 block scaling with sequence parallel (#2637) * fix subchannel fp8 + sp Signed-off-by: Chen Cui * Support sequence-parallel all-gather with small inputs Perform all-gather in high-precision if the input tensor is too small to quantize. Signed-off-by: Tim Moon * Fix lint error Signed-off-by: Przemek Tredak * Keep the previous behavior with dtype Signed-off-by: Przemek Tredak --------- Signed-off-by: Chen Cui Signed-off-by: Tim Moon Signed-off-by: Przemek Tredak Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon Co-authored-by: Przemek Tredak Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/distributed.py | 15 ++++++++++++--- .../pytorch/module/layernorm_linear.py | 2 -- .../pytorch/module/layernorm_mlp.py | 2 -- transformer_engine/pytorch/module/linear.py | 2 -- transformer_engine/pytorch/utils.py | 11 ----------- 5 files changed, 12 insertions(+), 20 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index f269e21b8c..2a65fa272b 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1100,6 +1100,9 @@ def _start_all_gather_fp8_blockwise( # Fall back to high-precision all-gather if FP8 is not supported if not quantizer.is_quantizable(inp) or quantizer.block_scaling_dim != 1: + warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.") + if isinstance(inp, QuantizedTensorStorage): + inp = inp.dequantize(dtype=dtype) # Dequantize if needed out = torch.empty(out_shape, dtype=dtype, device=device) torch.distributed.all_gather_into_tensor(out, inp, group=process_group, async_op=False) out = quantizer(out) @@ -1115,7 +1118,7 @@ def _start_all_gather_fp8_blockwise( "Input and quantizer do not have matching usages. " "Dequantizing and requantizing to Float8BlockwiseQTensor." ) - inp = quantizer(inp.dequantize()) + inp = quantizer(inp.dequantize(dtype=dtype)) # Construct Float8BlockwiseQTensor output tensor out = quantizer.make_empty(out_shape, dtype=dtype, device=device) @@ -1338,6 +1341,9 @@ def _all_gather_nvfp4( and quantizer is not None and not quantizer.is_quantizable(inp) ): + warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.") + if isinstance(inp, QuantizedTensorStorage): + inp = inp.dequantize(dtype=dtype) # Dequantize if needed out = torch.empty( out_shape, dtype=dtype, @@ -1358,7 +1364,7 @@ def _all_gather_nvfp4( "Input and quantizer do not have matching usages. " "Dequantizing and requantizing to NVFP4." ) - inp = quantizer(inp.dequantize()) + inp = quantizer(inp.dequantize(dtype=dtype)) # Construct NVFP4 output tensor out = quantizer.make_empty(out_shape, dtype=dtype, device=device) @@ -1505,6 +1511,9 @@ def _all_gather_mxfp8( and quantizer is not None and not quantizer.is_quantizable(inp) ): + warnings.warn("Cannot quantize input tensor. Performing all-gather in high precision.") + if isinstance(inp, QuantizedTensorStorage): + inp = inp.dequantize(dtype=dtype) # Dequantize if needed out = torch.empty( out_shape, dtype=dtype, @@ -1525,7 +1534,7 @@ def _all_gather_mxfp8( "Input and quantizer do not have matching usages. " "Dequantizing and requantizing to MXFP8." ) - inp = quantizer(inp.dequantize()) + inp = quantizer(inp.dequantize(dtype=dtype)) # Construct MXFP8 output tensor out = quantizer.make_empty(out_shape, dtype=dtype, device=device) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index ce0581024a..9a716207c1 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -29,7 +29,6 @@ from ..quantization import FP8GlobalStateManager from ..utils import ( assert_dim_for_fp8_exec, - assert_dim_for_all_gather, cast_if_needed, clear_tensor_data, divide, @@ -158,7 +157,6 @@ def forward( inputmat = inp if fp8: assert_dim_for_fp8_exec(inputmat, weight) - assert_dim_for_all_gather(inputmat, with_input_all_gather, input_quantizer) # Cast for native AMP nvtx_range_push(f"{nvtx_label}.norm_input_cast") diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 16e620fd94..ba92fb32ed 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -40,7 +40,6 @@ init_method_constant, cast_if_needed, assert_dim_for_fp8_exec, - assert_dim_for_all_gather, clear_tensor_data, requires_grad, needs_quantized_gemm, @@ -334,7 +333,6 @@ def _forward( inputmat = inp.view((-1, in_features)) if fp8: assert_dim_for_fp8_exec(inputmat, fc1_weight, fc2_weight) - assert_dim_for_all_gather(inputmat, sequence_parallel, fc1_input_quantizer) activation_func = _act_func( activation, FP8GlobalStateManager.get_fp8_recipe() if fp8 else None diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 31dac4d329..de2a421b53 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -34,7 +34,6 @@ requires_grad, needs_quantized_gemm, assert_dim_for_fp8_exec, - assert_dim_for_all_gather, nvtx_range_pop, nvtx_range_push, get_nvtx_range_context, @@ -175,7 +174,6 @@ def forward( own_quantized_input = False if fp8: assert_dim_for_fp8_exec(inputmat, weight) - assert_dim_for_all_gather(inputmat, with_input_all_gather_nccl, input_quantizer) if save_original_input: assert not isinstance( input_quantizer, Float8Quantizer diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 47af9fabe1..b1cc3be19d 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -12,7 +12,6 @@ import numpy as np import torch -from .quantized_tensor import Quantizer from .torch_version import torch_version from ..debug.pytorch.debug_quantization import DebugQuantizedTensor @@ -447,16 +446,6 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None: ) -def assert_dim_for_all_gather( - tensor: torch.Tensor, with_all_gather: bool, quantizer: Quantizer -) -> None: - """Assert that tensor dimensions are supported for all-gather""" - if with_all_gather: - assert quantizer.is_quantizable(tensor), ( - "All-gather requires quantizable tensor for quantizer " + quantizer.__class__.__name__ - ) - - def is_bf16_compatible() -> bool: """Replaces torch.cuda.is_bf16_compatible() with an explicit check on device compute capability to enforce sm_80 or higher. From ab9d60e0a6188c554456f255cf5eff6c70cbf397 Mon Sep 17 00:00:00 2001 From: Fabian Joswig Date: Sun, 8 Mar 2026 20:30:46 +0100 Subject: [PATCH 252/521] [PyTorch] Zero-initialize learnable softmax_offset in DotProductAttention (#2694) DotProductAttention used torch.empty() for the learnable softmax_offset parameter. Unlike all other TransformerEngineBaseModule subclasses, DotProductAttention does not call reset_parameters() in __init__, so the deferred initialization that would normally overwrite the empty tensor is never invoked, leaving the parameter with uninitialized memory. The JAX implementation explicitly uses nn.initializers.zeros for this parameter. This aligns the PyTorch behavior by using torch.zeros(). Signed-off-by: Fabian Joswig Co-authored-by: Kirthi Shankar Sivamani --- .../attention/dot_product_attention/dot_product_attention.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 64db4646f6..2dc42be18a 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 @@ -439,7 +439,7 @@ def __init__( if self.softmax_type == "learnable": self.register_parameter( "softmax_offset", - Parameter(torch.empty(self.num_attention_heads // self.tp_size, device="cuda")), + Parameter(torch.zeros(self.num_attention_heads // self.tp_size, device="cuda")), get_rng_state_tracker=get_rng_state_tracker, ) From e9ea352912cec5d1ea26bbdaa4f68fd2880c72e8 Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Sun, 8 Mar 2026 16:41:22 -0500 Subject: [PATCH 253/521] docs: update cuDNN sliding window attention support (#2624) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * docs: update cuDNN sliding window attention support Update documentation to reflect that cuDNN now supports causal sliding window attention (SWA) starting from version 9.2+. Changes: - Updated backend support matrix table to show cuDNN supports SWA (cuDNN 9.2+, causal masks only) - Added SWA comparison between flash-attention and cuDNN in section 1.3 - Added clarifying note in cp_ag_thd_dpa_jax_deep_dive.ipynb that cuDNN supports SWA but not all striping patterns for context parallelism Technical details: - cuDNN 9.2+: Supports causal SWA with window_size=(left, 0) - cuDNN 9.6+: Enhanced support for asymmetric windows (left, right) - Constraints: Requires dropout=0.0 and bias_type="no_bias" - Only works with causal mask types Signed-off-by: Santosh Bhavani * docs: update SWA notebook notes Signed-off-by: Santosh Bhavani --------- Signed-off-by: Santosh Bhavani Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> --- docs/examples/attention/attention.ipynb | 3 ++- docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/examples/attention/attention.ipynb b/docs/examples/attention/attention.ipynb index 4b2ed80497..e7253415d2 100644 --- a/docs/examples/attention/attention.ipynb +++ b/docs/examples/attention/attention.ipynb @@ -151,6 +151,7 @@ "- flash-attention does not support `post_scale_bias`, and cuDNN attention does.\n", "- flash-attention supports KV-caching and paged attention, and cuDNN attention does not.\n", "- flash-attention uses bottom right diagonal for `causal` mask in cross attention (see [change log](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#21-change-behavior-of-causal-flag)), and cuDNN attention supports both top left and bottom right.\n", + "- **Sliding window attention (SWA):** flash-attention has SWA(left, right) support for all mask types except top-left causal masks, with or without dropout, and without bias. cuDNN attention supports SWA(left, 0) starting from 9.2 and SWA(left, right) starting from 9.6, without dropout, and with `bias_type=\"no_bias\"`.\n", "- flash-attention outperforms cuDNN attention on Ampere architectures, and cuDNN attention has 20-50% advantages on Hopper architectures, based on our benchmarks for a number of commonly-used model configurations.\n", "\n", "To compare cuDNN attention and flash-attention, users can modify the `model_configs` dictionary in [benchmarks/attention/benchmark_attention.py](https://github.com/NVIDIA/TransformerEngine/blob/main/benchmarks/attention/benchmark_attention.py) to collect performance numbers. The script runs each entry in `model_configs` for `num_iters` times, each time with one forward pass and one backward pass. Both backends are tried, and if one backend does not have support for the specific user input, the runtimes and speedups in the final table would be 0." @@ -389,7 +390,7 @@ "\n", "| Attention Backend | Precision | Architecture | Sliding Window Attention | MQA/GQA | Multi-Latent Attention | Context Parallelism | Determinism Possible |\n", "| :---------------- | :-------- | :----------- | :----------------------- | :------ | :--------------------- | :------------------ | :------------ |\n", - "| cuDNN attention (all frameworks) | BF16, FP16, FP8 (PyTorch only) | sm80+ | No | Yes | Yes | Yes (`bshd`,`sbhd`, `thd`) | Yes |\n", + "| cuDNN attention (all frameworks) | BF16, FP16, FP8 (PyTorch only) | sm80+ | Yes (cuDNN 9.2+) | Yes | Yes | Yes (`bshd`,`sbhd`, `thd`) | Yes |\n", "| flash-attention (PyTorch) | BF16, FP16 | sm80+ | Yes | Yes | Yes | Yes (`bshd`,`thd`) | Yes |\n", "| Framework-native attention | BF16, FP16, FP32 | Any | No, unless used as a mask | Yes | Yes (PyTorch only) | No | Yes |\n", "\n", diff --git a/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb b/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb index 56bc3b13cf..338ce7fdd2 100644 --- a/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb +++ b/docs/examples/attention/cp_ag_thd_dpa_jax_deep_dive.ipynb @@ -28,6 +28,7 @@ "source": [ "### Question 1: Why choose Striped>1 ?\n", "\n", + "\n", "Prior to the addition of this feature, Transformer Engine JAX attention already supported load balancing via a striping pattern, i.e., `stripe_size=1` for `CP + THD + P2P(Ring) + Striped + SWA`. However, this reordering technique does not lend itself well to an all-gathered (post-AG) pattern. The following example illustrates this distinction. For this example, `cp_size=4`, `num_segments=4`, `window_size=(8,0)`, and the pattern is for a single rank after striped reordering has been performed: \n", "\n", "#### I. Striped (`stripe_size=1`)\n", From 6638fefb8b27e7126b247f1a7d0ece99562b07c0 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Sun, 8 Mar 2026 16:58:23 -0700 Subject: [PATCH 254/521] [JAX] GEMM tex and FFI cleanup (#2739) * cleanup Signed-off-by: Phuong Nguyen * rename FFI v2 API to match the style of GroupedGEMM Signed-off-by: Phuong Nguyen * format + lint Signed-off-by: Phuong Nguyen * fix: forward use_split_accumulator to _jax_gemm Signed-off-by: Phuong Nguyen * cpp warnings with call_once Signed-off-by: Phuong Nguyen * rm 2X_ACC_XGRAD Signed-off-by: Phuong Nguyen * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fix Signed-off-by: Phuong Nguyen * minor updates Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- tests/jax/test_distributed_dense.py | 23 +- transformer_engine/jax/cpp_extensions/gemm.py | 372 ++++-------------- transformer_engine/jax/csrc/extensions.h | 22 ++ .../jax/csrc/extensions/cgemm_helper.cpp | 4 +- .../jax/csrc/extensions/gemm.cpp | 230 +++++++---- .../jax/csrc/extensions/pybind.cpp | 4 + transformer_engine/jax/dense.py | 15 +- transformer_engine/jax/layernorm_dense.py | 15 +- transformer_engine/jax/layernorm_mlp.py | 32 +- transformer_engine/jax/quantize/helper.py | 9 - 10 files changed, 290 insertions(+), 436 deletions(-) diff --git a/tests/jax/test_distributed_dense.py b/tests/jax/test_distributed_dense.py index b8caf188d4..0c2ac8b24b 100644 --- a/tests/jax/test_distributed_dense.py +++ b/tests/jax/test_distributed_dense.py @@ -161,16 +161,21 @@ def test_distributed_gemm( # Compare results assert_allclose(gathered_te, gathered_jax, dtype=dtype) - def _te_sum_dense(self, x, weight, bias, contracting_dims): + def _te_sum_dense(self, x, weight, bias, contracting_dims, output_sharding): """TE GEMM function for gradient testing""" - return jnp.sum(dense(x, weight, bias=bias, contracting_dims=contracting_dims)) + output = dense(x, weight, bias=bias, contracting_dims=contracting_dims) + if output_sharding is not None: + output = jax.lax.with_sharding_constraint(output, output_sharding) + return jnp.sum(output) - def _jax_sum_dense(self, x, weight, bias, contracting_dims): + def _jax_sum_dense(self, x, weight, bias, contracting_dims, output_sharding): """JAX dot function for gradient testing""" - result = ( + output = ( jax.lax.dot_general(x, weight, dimension_numbers=(contracting_dims, ((), ()))) + bias ) - return jnp.sum(result) + if output_sharding is not None: + output = jax.lax.with_sharding_constraint(output, output_sharding) + return jnp.sum(output) @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", @@ -213,18 +218,18 @@ def test_te_distributed_dense_grad( # Test gradients w.r.t. all inputs te_grad_func = jax.jit( jax.value_and_grad(self._te_sum_dense, argnums=(0, 1, 2)), - static_argnames=("contracting_dims",), + static_argnames=("contracting_dims", "output_sharding"), ) jax_grad_func = jax.jit( jax.value_and_grad(self._jax_sum_dense, argnums=(0, 1, 2)), - static_argnames=("contracting_dims",), + static_argnames=("contracting_dims", "output_sharding"), ) te_val, te_grads = te_grad_func( - x_sharded, weight_sharded, bias_sharded, contracting_dims + x_sharded, weight_sharded, bias_sharded, contracting_dims, output_sharding ) jax_val, jax_grads = jax_grad_func( - x_sharded, weight_sharded, bias_sharded, contracting_dims + x_sharded, weight_sharded, bias_sharded, contracting_dims, output_sharding ) # Compare forward pass diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index ab2be7f799..70557f29c7 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -5,9 +5,10 @@ import math import operator +import os from collections.abc import Iterable from dataclasses import dataclass -from functools import partial, reduce +from functools import partial, reduce, cache from typing import Tuple, Sequence, Union from enum import Enum import warnings @@ -43,8 +44,6 @@ noop_quantizer_set, is_fp8_gemm_with_all_layouts_supported, apply_padding_to_scale_inv, - get_quantize_config_with_recipe, - get_global_quantize_recipe, QuantizeLayout, ) from .misc import get_padded_spec, is_all_reduce_in_float32 @@ -63,7 +62,6 @@ "gemm", "grouped_gemm_copy_group_sizes", "grouped_gemm", - "gemm_uses_jax_dot", "sanitize_dims", "get_non_contracting_dims", "transpose_dims", @@ -380,6 +378,12 @@ def get_rhs_axis_boundary(rhs_cdims, is_transposed): return min(rhs_cdims) if is_transposed else max(rhs_cdims) + 1 +@cache +def _get_high_precision_accumulation_from_env() -> bool: + """Read NVTE_FP8_GEMM_HIGH_PRECISION_ACCUMULATION once per process (cached).""" + return os.getenv("NVTE_FP8_GEMM_HIGH_PRECISION_ACCUMULATION", "0") == "1" + + def assert_cublas_requirements(scaling_mode, contracting_size, tensor_name): """Assert that the given tensor shape and layout meet the requirements for cuBLAS GEMM.""" if scaling_mode != ScalingMode.NO_SCALING: @@ -397,9 +401,9 @@ class GemmPrimitive(BasePrimitive): Primitive for cuBLAS GEMM """ - name = "te_gemm_ffi" + name = "te_gemm_v2_ffi" multiple_results = True - impl_static_args = (8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18) + impl_static_args = (7, 8, 9, 10, 11, 12, 13, 14) inner_primitive = None outer_primitive = None @@ -410,15 +414,11 @@ def abstract( rhs, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -457,6 +457,8 @@ def _dims_are_consecutive(dims): "cuBLAS GEMM operands have incompatible contracting dimensions: " f"{lhs.shape} @ idx {lhs_contracting_dims} X {rhs.shape} @ idx {rhs_contracting_dims}." ) + assert_cublas_requirements(scaling_mode, lhs_contracting_size, "LHS") + assert_cublas_requirements(scaling_mode, rhs_contracting_size, "RHS") lhs_is_transposed, rhs_is_transposed = _get_gemm_layout(operand_ndims, contracting_dims) if scaling_mode != ScalingMode.NO_SCALING: @@ -509,8 +511,8 @@ def _dims_are_consecutive(dims): assert out_dtype == jnp.bfloat16, f"Unsupported out_dtype={out_dtype}" output = jax.core.ShapedArray(shape=overlap_out_shape, dtype=out_dtype) - # Validate bias - if fuse_bias: + # Validate bias when present (bias.size > 0 means fuse bias) + if bias.size > 0: assert bias.shape == tuple(rhs_non_contracting_shape), ( "cuBLAS GEMM bias tensor has incorrect shape, " f"expected ({tuple(rhs_non_contracting_shape)}, ) but found {bias.shape}." @@ -519,29 +521,7 @@ def _dims_are_consecutive(dims): "cuBLAS GEMM bias tensor has incorrect data type, " f"expected {out_dtype} but found {bias.dtype}." ) - # WAR: allocate dbias regardless of fuse_bias so that the sharding propagation works as we - # change the fuse_bias value in the sharded_impl - dbias_shape = bias.shape if grad else (0,) - bias_grad = jax.core.ShapedArray(shape=dbias_shape, dtype=bias.dtype) - - # Validate pre-GeLU - pre_gelu_shape = (0,) - pre_gelu_dtype = out_dtype - if fuse_gelu: - pre_gelu_shape = out_shape - if grad: - pre_gelu_ndim = len(pre_gelu_shape) - assert gelu_input.ndim == pre_gelu_shape and all( - gelu_input.shape[i] == pre_gelu_shape[i] for i in range(pre_gelu_ndim) - ), ( - "cuBLAS GEMM pre-GeLU tensor has incorrect shape, " - f"expected {pre_gelu_shape} but found {gelu_input.shape}." - ) - assert gelu_input.dtype == out_dtype, ( - "cuBLAS GEMM pre-GeLU tensor has incorrect data type, " - f"expected {pre_gelu_dtype} but found {gelu_input.dtype}." - ) - pre_gelu_out = jax.core.ShapedArray(shape=pre_gelu_shape, dtype=pre_gelu_dtype) + assert alpha.size == 1 and alpha.dtype == jnp.float32 assert beta.size == 1 and beta.dtype == jnp.float32 @@ -557,12 +537,12 @@ def _dims_are_consecutive(dims): workspace_size += 256 workspace = jax.core.ShapedArray(shape=(workspace_size,), dtype=jnp.uint8) - return output, bias_grad, pre_gelu_out, workspace + return output, workspace @staticmethod def outer_abstract(*args, **kwargs): - outputs = GemmPrimitive.abstract(*args, **kwargs) - return outputs[:-1] # discard workspace array + output, _ = GemmPrimitive.abstract(*args, **kwargs) + return (output,) @staticmethod def lowering( @@ -572,15 +552,11 @@ def lowering( rhs, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -595,53 +571,18 @@ def lowering( (lhs_aval.ndim, rhs_aval.ndim), (lhs_cdims, rhs_cdims) ) - lhs_axis_boundary = get_lhs_axis_boundary(lhs_cdims, lhs_transposed) - lhs_contracting_size = ( - reduce(operator.mul, lhs_aval.shape[lhs_axis_boundary:]) - if lhs_transposed - else reduce(operator.mul, lhs_aval.shape[:lhs_axis_boundary]) - ) - assert_cublas_requirements( - scaling_mode, - lhs_contracting_size, - f"LHS {lhs_aval.shape} with contracting dims {lhs_cdims}", - ) - rhs_axis_boundary = get_rhs_axis_boundary(rhs_cdims, rhs_transposed) - rhs_contracting_size = ( - reduce(operator.mul, rhs_aval.shape[:rhs_axis_boundary]) - if rhs_transposed - else reduce(operator.mul, rhs_aval.shape[rhs_axis_boundary:]) - ) - assert_cublas_requirements( - scaling_mode, - rhs_contracting_size, - f"RHS {rhs_aval.shape} with contracting dims {rhs_cdims}", - ) - - args = (lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alpha, beta) + args = (lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, alpha, beta) kwargs = { "scaling_mode": int(scaling_mode.value), + "collective_op": int(collective_op.value), "lhs_axis_boundary": get_lhs_axis_boundary(lhs_cdims, lhs_transposed), "rhs_axis_boundary": get_rhs_axis_boundary(rhs_cdims, rhs_transposed), "lhs_transposed": lhs_transposed, "rhs_transposed": rhs_transposed, - "fuse_bias": fuse_bias, - "fuse_gelu": fuse_gelu, - "grad": grad, "use_split_accumulator": use_split_accumulator, - "collective_op": int(collective_op.value), } - operand_output_aliases = {} - if grad: - operand_output_aliases.update({4: 1}) # bias <-> bias_grad - if fuse_gelu and grad: - operand_output_aliases.update({5: 2}) # gelu_input <-> pre_gelu_out - - return jax.ffi.ffi_lowering( - GemmPrimitive.name, - operand_output_aliases=operand_output_aliases, - )(ctx, *args, **kwargs) + return jax.ffi.ffi_lowering(GemmPrimitive.name)(ctx, *args, config=kwargs) @staticmethod def impl( @@ -650,15 +591,11 @@ def impl( rhs, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -679,6 +616,7 @@ def impl( rhs_scale_inv = apply_padding_to_scale_inv( rhs_scale_inv, scaling_mode, rhs.shape, not rhs_transposed, rhs_flatten_axis ) + # Only perform JAX-based swizzle for MXFP8, NVFP4 swizzle will go though nvte kernel if scaling_mode.is_mxfp8_scaling: lhs_scale_inv = swizzled_scale(lhs_scale_inv, lhs_flatten_axis, lhs_transposed) @@ -711,26 +649,22 @@ def impl( reordered = reshaped.transpose(2, 0, 1, 3, *range(4, reshaped.ndim)) lhs = reordered.reshape(original_shape) - (output, bias_grad, pre_gelu_out, _) = GemmPrimitive.inner_primitive.bind( + (output, _) = GemmPrimitive.inner_primitive.bind( lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype=out_dtype, contracting_dims=contracting_dims, scaling_mode=scaling_mode, - fuse_bias=fuse_bias, - fuse_gelu=fuse_gelu, - grad=grad, use_split_accumulator=use_split_accumulator, - collective_op=collective_op, transpose_batch_sequence=transpose_batch_sequence, sequence_dim=sequence_dim, is_outer=is_outer, + collective_op=collective_op, ) # Alter output blocks for CGEMM AG if ( @@ -759,7 +693,7 @@ def impl( reordered = reshaped.transpose(1, 2, 0, 3, *range(4, reshaped.ndim)) output = reordered.reshape(original_shape) - return [output, bias_grad, pre_gelu_out] + return (output,) @staticmethod def outer_impl( @@ -768,15 +702,11 @@ def outer_impl( rhs, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -789,15 +719,11 @@ def outer_impl( rhs, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -812,9 +738,6 @@ def batcher( out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, collective_op, transpose_batch_sequence, @@ -831,30 +754,19 @@ def batcher( ), f"(Batching is not supported, got lhs_bdims={lhs_bdims}, rhs_bdims={rhs_bdims})" out_bdims = (None,) - # Bias gradient is never batched - bias_bdims = (None,) - - # Pre-GeLU output, if exists, is batched like GEMM output - pre_gelu_bdims = (None,) - if fuse_gelu and not grad: - pre_gelu_bdims = out_bdims - return ( GemmPrimitive.outer_primitive.bind( *batched_args, out_dtype=out_dtype, contracting_dims=contracting_dims, scaling_mode=scaling_mode, - fuse_bias=fuse_bias, - fuse_gelu=fuse_gelu, - grad=grad, use_split_accumulator=use_split_accumulator, collective_op=collective_op, transpose_batch_sequence=transpose_batch_sequence, sequence_dim=sequence_dim, is_outer=is_outer, ), - (out_bdims, bias_bdims, pre_gelu_bdims), + (out_bdims,), ) @staticmethod @@ -996,16 +908,15 @@ def _parse_operand_output_specs( (lhs_non_cspecs, rhs_non_cspecs), ) - # Bias and Pre-GeLU sharding is based on GEMM output before any scatter - bias_specs = tuple(list(rhs_non_cspecs).copy()) - gelu_specs = tuple(list(out_specs).copy()) + # Bias sharding is based on GEMM output before any scatter + bias_specs = rhs_non_cspecs if arg_infos[4].size > 0 else (None,) # bias is operand index 4 if not collective_op.is_none: assert sequence_dim >= 0, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" return ( - (lhs_specs, rhs_specs, bias_specs, gelu_specs), - (out_specs, bias_specs, gelu_specs), + (lhs_specs, rhs_specs, bias_specs), + out_specs, reduce_spec, sequence_dim, ) @@ -1015,9 +926,6 @@ def infer_sharding_from_operands( out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -1036,33 +944,18 @@ def infer_sharding_from_operands( sequence_dim, ) - (_, (out_specs, dbias_specs, pre_gelu_specs), *_) = ( - GemmPrimitive._parse_operand_output_specs( - arg_infos, contracting_dims, transpose_batch_sequence, collective_op - ) + (_, out_specs, *_) = GemmPrimitive._parse_operand_output_specs( + arg_infos, contracting_dims, transpose_batch_sequence, collective_op ) out_sharding = NamedSharding(mesh, PartitionSpec(*out_specs)) - # Discard dbias gradient spec if there is no bias and grad fusion - if not (fuse_bias and grad): - dbias_specs = (None,) - dbias_sharding = NamedSharding(mesh, PartitionSpec(*dbias_specs)) - - # Discard pre-GeLU output spec if there is no GeLU fusion - if not fuse_gelu: - pre_gelu_specs = (None,) - pre_gelu_sharding = NamedSharding(mesh, PartitionSpec(*pre_gelu_specs)) - - return [out_sharding, dbias_sharding, pre_gelu_sharding] + return (out_sharding,) @staticmethod def partition( out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -1075,8 +968,8 @@ def partition( del result_infos, is_outer, sequence_dim ( - (lhs_specs, rhs_specs, bias_input_specs, gelu_input_specs), - (out_specs, dbias_specs, pre_gelu_specs), + (lhs_specs, rhs_specs, bias_input_specs), + out_specs, reduce_spec, inferred_sequence_dim, ) = GemmPrimitive._parse_operand_output_specs( @@ -1097,50 +990,31 @@ def partition( rhs_sharding if scaling_mode.is_1d_block_scaling() else none_sharding, ) - # Discard bias input spec if there is no bias fusion - if not fuse_bias: - bias_input_specs = (None,) + # Bias arg_shardings += (NamedSharding(mesh, PartitionSpec(*bias_input_specs)),) - # Discard pre-GeLU input spec if there is no GeLU fusion - if not fuse_gelu: - gelu_input_specs = (None,) - arg_shardings += (NamedSharding(mesh, PartitionSpec(*gelu_input_specs)),) - # Alpha, beta arg_shardings += (none_sharding, none_sharding) # Assemble output shardings - out_shardings = [NamedSharding(mesh, PartitionSpec(*out_specs))] - - # Discard bias gradient spec if there is no bias and grad fusion - if not (fuse_bias and grad): - dbias_specs = (None,) - out_shardings.append(NamedSharding(mesh, PartitionSpec(*dbias_specs))) + out_sharding = (NamedSharding(mesh, PartitionSpec(*out_specs)),) - # Discard pre-GeLU output spec if there is no GeLU fusion - if not fuse_gelu: - pre_gelu_specs = (None,) - out_shardings.append(NamedSharding(mesh, PartitionSpec(*pre_gelu_specs))) - - def _sharded_impl(lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alpha, beta): + def _sharded_impl(lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, alpha, beta): # We should not fuse bias in the output reduction case - sharded_fuse_bias = fuse_bias and reduce_spec is None - outputs = GemmPrimitive.impl( + has_bias = bias.size > 0 + fuse_bias = has_bias and reduce_spec is None + bias_for_impl = bias if fuse_bias else jnp.empty(0, dtype=bias.dtype) + (output,) = GemmPrimitive.impl( lhs, lhs_scale_inv, rhs, rhs_scale_inv, - bias, - gelu_input, + bias_for_impl, alpha, beta, out_dtype=out_dtype, contracting_dims=contracting_dims, scaling_mode=scaling_mode, - fuse_bias=sharded_fuse_bias, - fuse_gelu=fuse_gelu, - grad=grad, use_split_accumulator=use_split_accumulator, transpose_batch_sequence=transpose_batch_sequence, sequence_dim=inferred_sequence_dim, @@ -1151,27 +1025,24 @@ def _sharded_impl(lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, gelu_input, alph if reduce_spec is not None: if not collective_op.is_reduce_scatter: if is_all_reduce_in_float32(): # For unittest only - outputs[0] = jax.lax.psum( - outputs[0].astype(jnp.float32), reduce_spec - ).astype(out_dtype) + output = jax.lax.psum(output.astype(jnp.float32), reduce_spec).astype( + out_dtype + ) else: - outputs[0] = jax.lax.psum(outputs[0], reduce_spec) + output = jax.lax.psum(output, reduce_spec) - if fuse_bias: # TODO(Phuong): rename fuse_bias to has_bias - outputs[0] += bias + if has_bias: + output += bias - return outputs + return (output,) - return mesh, _sharded_impl, out_shardings, arg_shardings + return mesh, _sharded_impl, out_sharding, arg_shardings @staticmethod def shardy_sharding_rule( out_dtype, contracting_dims, scaling_mode, - fuse_bias, - fuse_gelu, - grad, use_split_accumulator, transpose_batch_sequence, sequence_dim, @@ -1231,11 +1102,10 @@ def _generate_operand_rules(name, ndim, cdims): lhs_non_cspec = tuple(lhs_specs[i] for i in range(operand_ndims[0]) if i not in lhs_cdims) rhs_non_cspec = tuple(rhs_specs[i] for i in range(operand_ndims[1]) if i not in rhs_cdims) out_spec = (*lhs_non_cspec, *rhs_non_cspec) - bias_spec = rhs_non_cspec if fuse_bias else ("…4",) - gelu_spec = out_spec if fuse_gelu else ("…5",) - alpha_spec = ("_6",) - beta_spec = ("_7",) - dbias_spec = bias_spec if grad else ("…8") + bias_aval = operand_types[4] + bias_spec = rhs_non_cspec if math.prod(bias_aval.shape) > 0 else ("…4",) + alpha_spec = ("_5",) + beta_spec = ("_6",) return SdyShardingRule( operand_mappings=( @@ -1244,56 +1114,30 @@ def _generate_operand_rules(name, ndim, cdims): rhs_specs, rhs_scale_specs, bias_spec, - gelu_spec, alpha_spec, beta_spec, ), - result_mappings=( - out_spec, - dbias_spec, - gelu_spec, - ), + result_mappings=(out_spec,), ) register_primitive(GemmPrimitive) -def gemm_uses_jax_dot() -> bool: - """Check if the GEMM call directs to the TE custom cuBLAS call or native JAX dot.""" - return not GemmPrimitive.enabled() - - +# TODO(Phuong): move this function down after GroupedGemmPrimitive after initial review. Keep it +# here for now to minimize line changes. def _te_gemm( lhs: Union[jax.Array, ScaledTensor], rhs: Union[jax.Array, ScaledTensor], bias: jax.Array = None, - gelu_input: jax.Array = None, lhs_quantizer: Quantizer = None, rhs_quantizer: Quantizer = None, contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((-1,), (0,)), - fuse_bias: bool = False, - fuse_gelu: bool = False, - grad: bool = False, - use_split_accumulator: bool = None, + use_split_accumulator: bool = False, transpose_batch_sequence: bool = False, collective_op: CollectiveOp = CollectiveOp.NONE, ) -> Tuple[jax.Array, ...]: - if grad or fuse_gelu: - warnings.warn( - "GEMM + fused grad or fused gelu is not well tested and will be deprecated in the" - " future", - DeprecationWarning, - ) - - if use_split_accumulator is None: - # TODO(jberchtold): Rework GEMM API to provide the context here instead of relying on global state and also - # use context of the GEMM type so we can decide between fprop, dgrad, and wgrad - use_split_accumulator = get_quantize_config_with_recipe( - get_global_quantize_recipe() - ).FP8_2X_ACC_FPROP - # Prepare non-quantized GEMM operands lhs_data = lhs rhs_data = rhs @@ -1354,34 +1198,28 @@ def _te_gemm( rhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(rhs_amax) alpha = lhs_tensor_scale_inv * rhs_tensor_scale_inv - # Dummy empties for bias and gelu out_dtype = lhs_q.dq_dtype if isinstance(lhs_q, ScaledTensor) else lhs_data.dtype - if bias is None or not (fuse_bias and not grad): + if bias is None: bias = jnp.empty(0, dtype=out_dtype) - if gelu_input is None or not (fuse_gelu and grad): - gelu_input = jnp.empty(0, dtype=out_dtype) - return GemmPrimitive.outer_primitive.bind( + (output,) = GemmPrimitive.outer_primitive.bind( lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, - gelu_input, alpha, beta, out_dtype=out_dtype, contracting_dims=(lhs_cdims, rhs_cdims), scaling_mode=scaling_mode, - fuse_bias=fuse_bias, - fuse_gelu=fuse_gelu, - grad=grad, use_split_accumulator=use_split_accumulator, transpose_batch_sequence=transpose_batch_sequence, sequence_dim=-1, # Dummy value and will be set in the primitive is_outer=True, collective_op=collective_op, ) + return output class GroupedGemmCopySizesPrimitive(BasePrimitive): @@ -1827,6 +1665,7 @@ def _jax_gemm( contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((1,), (0,)), lhs_quantizer: Quantizer = None, rhs_quantizer: Quantizer = None, + use_split_accumulator: bool = False, ) -> jnp.ndarray: """ FP8 GEMM via JAX @@ -1839,12 +1678,6 @@ def _jax_gemm_impl(lhs, rhs): rhs.scaling_mode == lhs.scaling_mode ), f"rhs.scaling_mode={rhs.scaling_mode} != lhs.scaling_mode={lhs.scaling_mode}" - # TODO(jberchtold): Rework GEMM API to provide the context here instead of relying on global state and also - # use context of the GEMM type so we can decide between fprop, dgrad, and wgrad - use_split_accumulator = get_quantize_config_with_recipe( - get_global_quantize_recipe() - ).FP8_2X_ACC_FPROP - precision = ( jax.lax.Precision.HIGHEST if use_split_accumulator else jax.lax.Precision.DEFAULT ) @@ -1874,6 +1707,7 @@ def _jax_gemm_impl(lhs, rhs): def gemm( lhs: Union[jnp.ndarray, AbstractBaseTensor], rhs: Union[jnp.ndarray, AbstractBaseTensor], + bias: jnp.ndarray = None, contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((-1,), (0,)), lhs_quantizer: Quantizer = None, rhs_quantizer: Quantizer = None, @@ -1889,30 +1723,15 @@ def gemm( Left-hand side operand in the matrix multiplication. rhs: Union[jax.Array, ScaledTensor] Right-hand side operand in the matrix multiplication. + bias: jax.Array, default = None + Optional additive bias term. When provided (non-empty), bias is added to the result of the Matrix Multiplication operation. + This bias addition is fused when using the TE's custom call to cuBLAS GEMM. + contracting_dims: Tuple[Sequence[int], Sequence[int]], default = ((-1, ), (0, )) + Tuple of sequences representing the contracting dimensions of the operands. lhs_quantizer: Quantizer, default = None Object for down-casting the LHS operand for quantized GEMM. rhs_quantizer: Quantizer, default = None Object for down-casting the RHS operand for quantized GEMM. - contracting_dims: Tuple[Sequence[int], Sequence[int]], default = ((-1, ), (0, )) - Tuple of sequences representing the contracting dimensions of the operands. - bias: jax.Array, default = None - Optional additive bias term, required for forward GEMM with bias fusion. Only supported - with TE's custom call to cuBLAS GEMM. - gelu_input: jax.Array, default = None - Pre-GeLU output from forward GEMM, required for backward/grad GEMM with dGeLU fusion. Only - supported with TE's custom call to cuBLAS GEMM. - fuse_bias: bool, default = False - Enable bias addition in forward GEMM or bias gradient in backward GEMM. Only supported with - TE's custom call to cuBLAS GEMM. - fuse_gelu: bool, default = False - Enable GeLU activation in forward GEMM or GeLU gradient in backward GEMM. Only supported - with TE's custom call to cuBLAS GEMM. - grad: bool, default = False - Flag for switching bias and GeLU fusions from forward to backward mode. Only supported with - TE's custom call to cuBLAS GEMM. - use_split_accumulator: bool, default = True - Enable promoting some intermediate sums to higher precision when accumulating the result in - the cuBLAS GEMM kernel. Disabling this trades off numerical accuracy for speed. transpose_batch_sequence: bool, default = False Transpose the batch and sequence dimensions of the input tensor. collective_op: CollectiveOp, default = CollectiveOp.NONE @@ -1921,18 +1740,7 @@ def gemm( Returns ------- jax.Array: - Result of the operation. For TE's custom call to cuBLAS GEMM, this result can include the - GeLU application when `fuse_gelu=True` and `grad=False`, the GeLU gradient contribution - when `fuse_gelu=True` and `grad=True`, and the additive bias when `fuse_bias=True` and - `grad=False`. - Optional[jax.Array]: - Bias gradient when `fuse_bias=True` and `grad=True`. Only supported with TE's custom call - to cuBLAS GEMM. - Optional[jax.Array]: - Pre-GeLU GEMM output when `fuse_gelu=True` and `grad=False`. This is required as an input - to `_te_gemm()` with `fuse_gelu=True` and `grad=True` in the backward pass in order to - compute the GeLU contribution to the gradient. Only supported with TE's custom call to - cuBLAS GEMM. + Result of the operation lhs * rhs + bias. """ if isinstance(lhs, NoScaleTensor): lhs = lhs.data @@ -1946,45 +1754,33 @@ def gemm( lhs_quantizer = quantizer_set.x rhs_quantizer = quantizer_set.kernel + # This option enable promoting some intermediate sums to higher precision when accumulating the result in + # the cuBLAS GEMM kernel. Disabling this trades off numerical accuracy for speed. + use_split_accumulator = _get_high_precision_accumulation_from_env() + # Fall back on a native JAX implementation when the custom call to cuBLAS GEMM is disabled - # TODO(Phuong): fuse_bias -> has_bias and has_bias = bias is not None - fuse_bias = kwargs.get("fuse_bias", False) - fuse_gelu = kwargs.get("fuse_gelu", False) if not GemmPrimitive.enabled(): - assert kwargs.get("bias", None) is None and not fuse_gelu, ( - "TE GEMM was invoked with bias fusion options that are not supported by the " - "`jax.lax.dot_general` and `jax.nn.scaled_matmul` backends used when the custom cuBLAS " - "GEMM primitive is disabled." - ) - assert kwargs.get("gelu_input", None) is None and not fuse_bias, ( - "TE GEMM was invoked with GeLU fusion options that are not supported by the " - "`jax.lax.dot_general` and `jax.nn.scaled_matmul` backends used when the custom cuBLAS " - "GEMM primitive is disabled." - ) assert collective_op.is_none, "JAX GEMM does not support collective GEMM" - return _jax_gemm(lhs, rhs, contracting_dims, lhs_quantizer, rhs_quantizer) + output = _jax_gemm( + lhs, rhs, contracting_dims, lhs_quantizer, rhs_quantizer, use_split_accumulator + ) + if bias is not None: + output += bias # Unfused + return output - outputs = _te_gemm( + output = _te_gemm( lhs, rhs, + bias, lhs_quantizer=lhs_quantizer, rhs_quantizer=rhs_quantizer, contracting_dims=contracting_dims, + use_split_accumulator=use_split_accumulator, transpose_batch_sequence=transpose_batch_sequence, collective_op=collective_op, - **kwargs, ) - # Discard empty outputs - grad = kwargs.get("grad", False) - clean_outputs = outputs[0] # first output is the final result and is never empty - if (fuse_bias and grad) or (fuse_gelu and not grad): - clean_outputs = (outputs[0],) - if fuse_bias and grad: # only return bias gradient if it exists - clean_outputs += (outputs[1],) - if fuse_gelu and not grad: # only return pre-GeLU output if it exists - clean_outputs += (outputs[2],) - return clean_outputs + return output def grouped_gemm_copy_group_sizes( diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 93c85aaacc..0fe4e99239 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -45,6 +45,16 @@ struct ActivationConfig { ClampedSwigluConfig clamped_swiglu; }; +struct GemmConfig { + JAXX_Scaling_Mode scaling_mode; + JAXX_Collective_Op collective_op; + int64_t lhs_axis_boundary; + int64_t rhs_axis_boundary; + bool lhs_transposed; + bool rhs_transposed; + bool use_split_accumulator; +}; + inline bool use_fp8(DType type) { return type == DType::kFloat8E4M3 || type == DType::kFloat8E5M2; } // Activation @@ -133,7 +143,9 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( // GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmV2Handler); XLA_FFI_DECLARE_HANDLER_SYMBOL(CollectiveGemmInitHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(GemmInitV2Handler); // Grouped GEMM XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedGemmD2HGroupSizesHandler); @@ -170,6 +182,16 @@ XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( transformer_engine::jax::ActivationConfig, ::xla::ffi::StructMember("clamped_swiglu")); +XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( + transformer_engine::jax::GemmConfig, + ::xla::ffi::StructMember("scaling_mode"), + ::xla::ffi::StructMember("collective_op"), + ::xla::ffi::StructMember("lhs_axis_boundary"), + ::xla::ffi::StructMember("rhs_axis_boundary"), + ::xla::ffi::StructMember("lhs_transposed"), + ::xla::ffi::StructMember("rhs_transposed"), + ::xla::ffi::StructMember("use_split_accumulator")); + // ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Score_Function); diff --git a/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp b/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp index 87a889621e..36a4a068a4 100644 --- a/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp +++ b/transformer_engine/jax/csrc/extensions/cgemm_helper.cpp @@ -138,8 +138,8 @@ void CommunicatorHandler::init(int num_total_devices, int num_devices_per_proces // Bootstrap UB via creating a dummy CommOverlapP2PBase object std::vector buffer_shape{1, 1}; - auto _ = CollectiveGemmPlanRegistry::getInstance().get_executor(buffer_shape, DType::kFloat32, - JAXX_Collective_Op::ALL_GATHER); + [[maybe_unused]] auto _ = CollectiveGemmPlanRegistry::getInstance().get_executor( + buffer_shape, DType::kFloat32, JAXX_Collective_Op::ALL_GATHER); } void InitializeCgemmCommunicator(int num_total_devices, int num_devices_per_process, int process_id, diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 4cbec405a4..737dd65622 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -98,46 +98,75 @@ std::tuple> xla_buffer_to_nvte_gemm_operand( return std::make_tuple(std::move(input), input_shape); } -Error_Type CollectiveGemmInitFFI(Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, - Buffer_Type rhs_scale_inv, Buffer_Type bias, - Buffer_Type gelu_input, Buffer_Type alpha, Buffer_Type beta, - Result_Type output, Result_Type bias_grad, - Result_Type pre_gelu_out, Result_Type workspace, - JAXX_Scaling_Mode scaling_mode, int64_t lhs_axis_boundary, - int64_t rhs_axis_boundary, bool lhs_transposed, - bool rhs_transposed, bool fuse_bias, bool fuse_gelu, bool grad, - bool use_split_accumulator, JAXX_Collective_Op collective_op) { +Error_Type GemmInitV2FFI(Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, + Buffer_Type rhs_scale_inv, Buffer_Type bias, Buffer_Type alpha, + Buffer_Type beta, Result_Type output, Result_Type workspace, + GemmConfig config) { nvte_cublas_handle_init(); // Init UB buffer - if (collective_op != JAXX_Collective_Op::NONE) { + if (config.collective_op != JAXX_Collective_Op::NONE) { auto &comm_handler = CommunicatorHandler::get(); std::vector lhs_shape = { - product(lhs.dimensions(), 0, lhs_axis_boundary), - product(lhs.dimensions(), lhs_axis_boundary, lhs.dimensions().size())}; + product(lhs.dimensions(), 0, config.lhs_axis_boundary), + product(lhs.dimensions(), config.lhs_axis_boundary, lhs.dimensions().size())}; std::vector rhs_shape = { - product(rhs.dimensions(), 0, rhs_axis_boundary), - product(rhs.dimensions(), rhs_axis_boundary, rhs.dimensions().size())}; + product(rhs.dimensions(), 0, config.rhs_axis_boundary), + product(rhs.dimensions(), config.rhs_axis_boundary, rhs.dimensions().size())}; - std::vector out_shape = {(lhs_transposed) ? lhs_shape[1] : lhs_shape[0], - (rhs_transposed) ? rhs_shape[0] : rhs_shape[1]}; + std::vector out_shape = {(config.lhs_transposed) ? lhs_shape[1] : lhs_shape[0], + (config.rhs_transposed) ? rhs_shape[0] : rhs_shape[1]}; std::vector buffer_shape{0, 0}; DType buffer_dtype = convert_ffi_datatype_to_te_dtype(output->element_type()); - if (collective_op == JAXX_Collective_Op::ALL_GATHER) { + if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) { buffer_shape[0] = lhs_shape[0] * comm_handler.tp_size; buffer_shape[1] = lhs_shape[1]; buffer_dtype = convert_ffi_datatype_to_te_dtype(lhs.element_type()); - } else if (collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { + } else if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { buffer_shape[0] = out_shape[0]; buffer_shape[1] = out_shape[1]; } - auto _ = CollectiveGemmPlanRegistry::getInstance().get_executor(buffer_shape, buffer_dtype, - collective_op); + [[maybe_unused]] auto _ = CollectiveGemmPlanRegistry::getInstance().get_executor( + buffer_shape, buffer_dtype, config.collective_op); } return ffi_with_cuda_error_check(); } +XLA_FFI_DEFINE_HANDLER_SYMBOL(GemmInitV2Handler, GemmInitV2FFI, + FFI::Bind() + .Arg() // lhs + .Arg() // lhs_scale_inv + .Arg() // rhs + .Arg() // rhs_scale_inv + .Arg() // bias + .Arg() // alpha + .Arg() // beta + .Ret() // output + .Ret() // workspace + .Attr("config"), + FFI_CudaGraph_Traits); + +Error_Type CollectiveGemmInitFFI(Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, + Buffer_Type rhs_scale_inv, Buffer_Type bias, + Buffer_Type gelu_input, Buffer_Type alpha, Buffer_Type beta, + Result_Type output, Result_Type bias_grad, + Result_Type pre_gelu_out, Result_Type workspace, + JAXX_Scaling_Mode scaling_mode, int64_t lhs_axis_boundary, + int64_t rhs_axis_boundary, bool lhs_transposed, + bool rhs_transposed, bool fuse_bias, bool fuse_gelu, bool grad, + bool use_split_accumulator, JAXX_Collective_Op collective_op) { + static std::once_flag gemm_init_warned; + std::call_once(gemm_init_warned, []() { + std::cerr << "[CollectiveGemmInitFFI] Deprecation: This API is deprecated and will be removed " + "in September 2026. Use GemmInitV2FFI instead." + << std::endl; + }); + return GemmInitV2FFI(lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, alpha, beta, output, workspace, + GemmConfig{scaling_mode, collective_op, lhs_axis_boundary, rhs_axis_boundary, + lhs_transposed, rhs_transposed, use_split_accumulator}); +} + XLA_FFI_DEFINE_HANDLER_SYMBOL(CollectiveGemmInitHandler, CollectiveGemmInitFFI, FFI::Bind() .Arg() // lhs @@ -161,21 +190,19 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(CollectiveGemmInitHandler, CollectiveGemmInitFFI, .Attr("fuse_gelu") .Attr("grad") .Attr("use_split_accumulator") - .Attr("collective_op")); + .Attr("collective_op"), + FFI_CudaGraph_Traits); -Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, - Buffer_Type rhs_scale_inv, Buffer_Type bias, Buffer_Type gelu_input, - Buffer_Type alpha, Buffer_Type beta, Result_Type output, Result_Type bias_grad, - Result_Type pre_gelu_out, Result_Type workspace, JAXX_Scaling_Mode scaling_mode, - int64_t lhs_axis_boundary, int64_t rhs_axis_boundary, bool lhs_transposed, - bool rhs_transposed, bool fuse_bias, bool fuse_gelu, bool grad, - bool use_split_accumulator, JAXX_Collective_Op collective_op) { +Error_Type GemmV2FFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_inv, + Buffer_Type rhs, Buffer_Type rhs_scale_inv, Buffer_Type bias, + Buffer_Type alpha, Buffer_Type beta, Result_Type output, Result_Type workspace, + GemmConfig config) { // cuBLAS workspace + 256 alignment enforcement (+ swizzle scales) uint8_t *lhs_swizzle_scale_ptr = nullptr, *rhs_swizzle_scale_ptr = nullptr; auto workspace_ptr = reinterpret_cast(workspace->untyped_data()); workspace_ptr = move_ptr_to_next_256B_aligned(workspace_ptr); size_t workspace_size = static_cast(workspace->element_count()) - 256; - if (is_nvfp4_scaling(scaling_mode)) { + if (is_nvfp4_scaling(config.scaling_mode)) { auto lhs_scale_size = product(lhs_scale_inv.dimensions()); auto rhs_scale_size = product(rhs_scale_inv.dimensions()); workspace_size = workspace_size - lhs_scale_size - rhs_scale_size; @@ -187,60 +214,42 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i // NOTE: TensorWrapper operands are always rowwise for full-precision GEMM, or FP8 GEMM when // device supports non-TN layouts (compute capability >= 10.0, excluding 12.x) - bool always_rowwise = (scaling_mode == JAXX_Scaling_Mode::NO_SCALING || - (is_tensor_scaling(scaling_mode) && nvte_is_non_tn_fp8_gemm_supported())); - bool make_lhs_rowwise = (always_rowwise) ? true : !lhs_transposed; - bool make_rhs_rowwise = (always_rowwise) ? true : rhs_transposed; - - auto [lhs_, lhs_shape] = - xla_buffer_to_nvte_gemm_operand(stream, lhs, lhs_scale_inv, lhs_swizzle_scale_ptr, - scaling_mode, lhs_axis_boundary, make_lhs_rowwise); - auto [rhs_, rhs_shape] = - xla_buffer_to_nvte_gemm_operand(stream, rhs, rhs_scale_inv, rhs_swizzle_scale_ptr, - scaling_mode, rhs_axis_boundary, make_rhs_rowwise); - - std::vector out_shape = {(lhs_transposed) ? lhs_shape[1] : lhs_shape[0], - (rhs_transposed) ? rhs_shape[0] : rhs_shape[1]}; + bool always_rowwise = + (config.scaling_mode == JAXX_Scaling_Mode::NO_SCALING || + (is_tensor_scaling(config.scaling_mode) && nvte_is_non_tn_fp8_gemm_supported())); + bool make_lhs_rowwise = (always_rowwise) ? true : !config.lhs_transposed; + bool make_rhs_rowwise = (always_rowwise) ? true : config.rhs_transposed; + + auto [lhs_, lhs_shape] = xla_buffer_to_nvte_gemm_operand( + stream, lhs, lhs_scale_inv, lhs_swizzle_scale_ptr, config.scaling_mode, + config.lhs_axis_boundary, make_lhs_rowwise); + auto [rhs_, rhs_shape] = xla_buffer_to_nvte_gemm_operand( + stream, rhs, rhs_scale_inv, rhs_swizzle_scale_ptr, config.scaling_mode, + config.rhs_axis_boundary, make_rhs_rowwise); + + std::vector out_shape = {(config.lhs_transposed) ? lhs_shape[1] : lhs_shape[0], + (config.rhs_transposed) ? rhs_shape[0] : rhs_shape[1]}; auto out_dtype = convert_ffi_datatype_to_te_dtype(output->element_type()); // Bias input to forward pass or bias gradient output from backward pass void *bias_ptr = nullptr; size_t bias_size = 0; DType bias_dtype = out_dtype; + auto fuse_bias = bias.element_count() > 0; if (fuse_bias) { - if (grad) { - NVTE_CHECK(bias_grad->untyped_data() == bias.untyped_data(), - "Missing operand-output aliasing in GemmPrimitive: bias <-> bias_grad"); - } bias_ptr = bias.untyped_data(); bias_size = product(bias.dimensions()); bias_dtype = convert_ffi_datatype_to_te_dtype(bias.element_type()); } auto bias_ = TensorWrapper(bias_ptr, std::vector{bias_size}, bias_dtype); - // Pre-GeLU output from forward pass or input to backward pass - void *pre_gelu_ptr = nullptr; - std::vector pre_gelu_shape = {0}; - DType pre_gelu_dtype = out_dtype; - if (gelu_input.element_count() > 0) { - if (grad) { - NVTE_CHECK(pre_gelu_out->untyped_data() == gelu_input.untyped_data(), - "Missing operand-output aliasing in GemmPrimitive: gelu_input <-> pre_gelu_out"); - } - pre_gelu_ptr = pre_gelu_out->untyped_data(); - pre_gelu_shape = {product(pre_gelu_out->dimensions(), 0, pre_gelu_out->dimensions().size() - 1), - static_cast(pre_gelu_out->dimensions().back())}; - pre_gelu_dtype = convert_ffi_datatype_to_te_dtype(pre_gelu_out->element_type()); - } - auto pre_gelu_ = TensorWrapper(pre_gelu_ptr, pre_gelu_shape, pre_gelu_dtype); - auto num_math_sm = cuda::sm_count() - getenv("NVTE_EXT_MARGIN_SM", 0); float one = 1.; float zero = 0.; // alpha, beta float *alpha_ptr = &one, *beta_ptr = &zero; - if (is_nvfp4_scaling(scaling_mode)) { + if (is_nvfp4_scaling(config.scaling_mode)) { NVTE_CHECK(alpha.element_count() == 1 && convert_ffi_datatype_to_te_dtype(alpha.element_type()) == DType::kFloat32); alpha_ptr = reinterpret_cast(alpha.untyped_data()); @@ -250,16 +259,12 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i } // Construct GEMM config - transformer_engine::MatmulConfigWrapper config; - config.set_use_split_accumulator(use_split_accumulator); - config.set_sm_count(num_math_sm); - if (fuse_bias) config.set_bias_tensor(bias_.data()); - if (fuse_gelu) { - config.set_with_gelu_epilogue(true); - config.set_epilogue_aux_tensor(pre_gelu_.data()); - } + transformer_engine::MatmulConfigWrapper matmul_config; + matmul_config.set_use_split_accumulator(config.use_split_accumulator); + matmul_config.set_sm_count(num_math_sm); + if (fuse_bias) matmul_config.set_bias_tensor(bias_.data()); - if (collective_op == JAXX_Collective_Op::NONE) { + if (config.collective_op == JAXX_Collective_Op::NONE) { auto out_ = TensorWrapper(output->untyped_data(), out_shape, out_dtype); NVTE_CHECK(out_.numel() == output->element_count(), "cuBLAS GEMM output buffer size is incorrect, expected ", out_.numel(), " elements ", @@ -269,19 +274,20 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i ", out_shape[1]=", out_shape[1]); // Launch TE/common kernel with swapped LHS/RHS for cuBLAS column-major order - nvte_cublas_gemm_v2(rhs_transposed /*transa*/, lhs_transposed /*transb*/, alpha_ptr, - rhs_.data() /*A*/, lhs_.data() /*B*/, beta_ptr, out_.data() /*C*/, - out_.data() /*D*/, workspace_.data(), config, stream); + nvte_cublas_gemm_v2(config.rhs_transposed /*transa*/, config.lhs_transposed /*transb*/, + alpha_ptr, rhs_.data() /*A*/, lhs_.data() /*B*/, beta_ptr, + out_.data() /*C*/, out_.data() /*D*/, workspace_.data(), matmul_config, + stream); } else { std::vector buffer_shape{0, 0}; DType buffer_dtype = out_dtype; auto &comm_handler = CommunicatorHandler::get(); - if (collective_op == JAXX_Collective_Op::ALL_GATHER) { + if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) { buffer_shape[0] = lhs_shape[0] * comm_handler.tp_size; buffer_shape[1] = lhs_shape[1]; out_shape[0] = out_shape[0] * comm_handler.tp_size; buffer_dtype = convert_ffi_datatype_to_te_dtype(lhs.element_type()); - } else if (collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { + } else if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { buffer_shape[0] = out_shape[0]; buffer_shape[1] = out_shape[1]; out_shape[0] = out_shape[0] / comm_handler.tp_size; @@ -289,8 +295,9 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i NVTE_CHECK(!fuse_bias || bias_size == out_shape[1], "bias_size=", bias_size, ", out_shape[1]=", out_shape[1]); auto executor = CollectiveGemmPlanRegistry::getInstance().get_executor( - buffer_shape, buffer_dtype, collective_op); - if (collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { + buffer_shape, buffer_dtype, config.collective_op); + auto pre_gelu_ = TensorWrapper(nullptr, std::vector{0}, DType::kByte); + if (config.collective_op == JAXX_Collective_Op::REDUCE_SCATTER) { auto ubuf_out_ = TensorWrapper(executor->get_ubuf_dptr(), buffer_shape, out_dtype); // Prepare the auxiliary buffer for the reduce-scattered GEMM output auto out_ = TensorWrapper(output->untyped_data(), out_shape, out_dtype); @@ -300,11 +307,11 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i " elements ", to_string_like(output->dimensions())); // Launch GEMM+RS - executor->split_overlap_rs(rhs_, rhs_transposed, lhs_, lhs_transposed, ubuf_out_, bias_, - pre_gelu_, workspace_, grad, false, use_split_accumulator, out_, - stream); + executor->split_overlap_rs(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, + ubuf_out_, bias_, pre_gelu_, workspace_, false /*grad*/, + false /*accumulate*/, config.use_split_accumulator, out_, stream); - } else if (collective_op == JAXX_Collective_Op::ALL_GATHER) { + } else if (config.collective_op == JAXX_Collective_Op::ALL_GATHER) { auto aux_out_ = TensorWrapper(nullptr, std::vector{0}, out_dtype); // Empty auto out_ = TensorWrapper(output->untyped_data(), out_shape, out_dtype); @@ -315,14 +322,65 @@ Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_i // Copy the distributed LHS operand into the local chunk of the communication buffer executor->copy_into_buffer(stream, lhs_, true, make_lhs_rowwise); // Launch AG+GEMM - executor->split_overlap_ag(rhs_, rhs_transposed, lhs_, lhs_transposed, out_, bias_, pre_gelu_, - workspace_, grad, false, use_split_accumulator, aux_out_, stream); + executor->split_overlap_ag(rhs_, config.rhs_transposed, lhs_, config.lhs_transposed, out_, + bias_, pre_gelu_, workspace_, false /*grad*/, false /*accumulate*/, + config.use_split_accumulator, aux_out_, stream); } } return ffi_with_cuda_error_check(); } +XLA_FFI_DEFINE_HANDLER_SYMBOL(GemmV2Handler, GemmV2FFI, + FFI::Bind() + .Ctx() // stream + .Arg() // lhs + .Arg() // lhs_scale_inv + .Arg() // rhs + .Arg() // rhs_scale_inv + .Arg() // bias + .Arg() // alpha + .Arg() // beta + .Ret() // output + .Ret() // workspace + .Attr("config"), + FFI_CudaGraph_Traits); + +Error_Type GemmFFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale_inv, Buffer_Type rhs, + Buffer_Type rhs_scale_inv, Buffer_Type bias, Buffer_Type gelu_input, + Buffer_Type alpha, Buffer_Type beta, Result_Type output, Result_Type bias_grad, + Result_Type pre_gelu_out, Result_Type workspace, JAXX_Scaling_Mode scaling_mode, + int64_t lhs_axis_boundary, int64_t rhs_axis_boundary, bool lhs_transposed, + bool rhs_transposed, bool fuse_bias, bool fuse_gelu, bool grad, + bool use_split_accumulator, JAXX_Collective_Op collective_op) { + static std::once_flag once_fuse_bias; + static std::once_flag once_fuse_gelu_grad; + static std::once_flag once_api; + if (fuse_bias) { + std::call_once(once_fuse_bias, [] { + std::cerr << "[GemmFFI] Deprecation: fuse_bias is deprecated; bias fusion is inferred from " + "non-empty bias. This parameter will be removed in future release." + << std::endl; + }); + } + if (fuse_gelu || grad) { + std::call_once(once_fuse_gelu_grad, [] { + std::cerr << "[GemmFFI] Deprecation: fuse_gelu and grad are deprecated. These options are " + "ignored as there is no support for them in the current implementation. " + << std::endl; + }); + } + std::call_once(once_api, [] { + std::cerr << "[GemmFFI] Deprecation: This API is deprecated in Sep 2026. Use GemmV2FFI instead." + << std::endl; + }); + + return GemmV2FFI(stream, lhs, lhs_scale_inv, rhs, rhs_scale_inv, bias, alpha, beta, output, + workspace, + GemmConfig{scaling_mode, collective_op, lhs_axis_boundary, rhs_axis_boundary, + lhs_transposed, rhs_transposed, use_split_accumulator}); +} + XLA_FFI_DEFINE_HANDLER_SYMBOL(GemmHandler, GemmFFI, FFI::Bind() .Ctx() // stream diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 837dd55f9c..28cb39b5d1 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -69,6 +69,10 @@ pybind11::dict Registrations() { pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CollectiveGemmInitHandler), pybind11::arg("execute") = EncapsulateFFI(GemmHandler)); + dict["te_gemm_v2_ffi"] = + pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(GemmInitV2Handler), + pybind11::arg("execute") = EncapsulateFFI(GemmV2Handler)); + // Grouped GEMM dict["te_grouped_gemm_d2h_group_sizes_ffi"] = pybind11::dict(pybind11::arg("prepare") = EncapsulateFFI(CublasHandleInitHandler), diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index 268995281c..fe02e61fc0 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -217,30 +217,25 @@ def _dense_fwd_rule( casted_kernel = with_sharding_constraint_by_logical_axes(casted_kernel, kernel_axes) # GEMM NN - use_bias = bias is not None output = tex.gemm( casted_x.get_tensor(usage=TensorUsage.LHS), casted_kernel.get_tensor(usage=TensorUsage.RHS), + bias=bias, contracting_dims=(x_contracting_dims, k_contracting_dims), transpose_batch_sequence=transpose_batch_sequence, - bias=bias if not tex.gemm_uses_jax_dot() else None, - fuse_bias=use_bias if not tex.gemm_uses_jax_dot() else False, collective_op=collective_op_set.forward, ) output = with_sharding_constraint_by_logical_axes(output, output_axes) - if use_bias and tex.gemm_uses_jax_dot(): - bias_new_shape = (1,) * (output.ndim - bias.ndim) + bias.shape - output += jnp.reshape(bias, bias_new_shape) - + has_bias = bias is not None ctx = ( casted_x.get_tensor(usage=TensorUsage.LHS_TRANS).checkpoint(quantizer_set.x), casted_kernel.get_tensor(usage=TensorUsage.RHS_TRANS).checkpoint(quantizer_set.kernel), x.shape, kernel.shape, - use_bias, quantizer_set, flatten_axis_k, + has_bias, ) return output, ctx @@ -265,9 +260,9 @@ def _dense_bwd_rule( casted_kernel_rhs, x_shape, kernel_shape, - use_bias, quantizer_set, flatten_axis_k, + has_bias, ) = ctx grad = with_sharding_constraint_by_logical_axes(grad, output_axes) @@ -277,7 +272,7 @@ def _dense_bwd_rule( casted_grad, dbias = tex.quantize_dbias( grad, - is_dbias=use_bias, + is_dbias=has_bias, flatten_axis=flatten_axis_k, quantizer=quantizer_set.dgrad, amax_scope=AmaxScope.TPSP, diff --git a/transformer_engine/jax/layernorm_dense.py b/transformer_engine/jax/layernorm_dense.py index 8c21496ffe..63e6daf9d5 100644 --- a/transformer_engine/jax/layernorm_dense.py +++ b/transformer_engine/jax/layernorm_dense.py @@ -220,20 +220,15 @@ def _layernorm_dense_fwd_rule( # NN GEMM # (batch..., hidden_in) x (hidden_in, hidden_out...) - use_bias = bias is not None output = tex.gemm( casted_ln_out.get_tensor(TensorUsage.LHS), casted_kernel.get_tensor(TensorUsage.RHS), contracting_dims=(x_contracting_dims, k_contracting_dims), transpose_batch_sequence=transpose_batch_sequence, - bias=bias if not tex.gemm_uses_jax_dot() else None, - fuse_bias=use_bias if not tex.gemm_uses_jax_dot() else False, + bias=bias, ) - if use_bias and tex.gemm_uses_jax_dot(): - bias_new_shape = (1,) * (output.ndim - bias.ndim) + bias.shape - output += jnp.reshape(bias, bias_new_shape) - + has_bias = bias is not None ctx = ( casted_ln_out.get_tensor(TensorUsage.LHS_TRANS).checkpoint(quantizer_set.x), casted_kernel.get_tensor(TensorUsage.RHS_TRANS).checkpoint(quantizer_set.kernel), @@ -246,7 +241,7 @@ def _layernorm_dense_fwd_rule( beta, x_contracting_dims, k_contracting_dims, - use_bias, + has_bias, quantizer_set, flatten_axis, ) @@ -289,14 +284,14 @@ def _layernorm_dense_bwd_rule( beta, x_contracting_dims_in_fwd, k_contracting_dims_in_fwd, - use_bias, + has_bias, quantizer_set, flatten_axis, ) = ctx casted_grad, dbias = tex.quantize_dbias( grad, - is_dbias=use_bias, + is_dbias=has_bias, flatten_axis=flatten_axis, quantizer=quantizer_set.dgrad, amax_scope=AmaxScope.TPSP, diff --git a/transformer_engine/jax/layernorm_mlp.py b/transformer_engine/jax/layernorm_mlp.py index c90d018aee..4c324c208e 100644 --- a/transformer_engine/jax/layernorm_mlp.py +++ b/transformer_engine/jax/layernorm_mlp.py @@ -295,8 +295,8 @@ def _layernorm_mlp_fwd_rule( assert x.shape[x_contracting_dims[0]] == kernel_1.shape[k_contracting_dims[0]] - use_bias_1 = bias_1 is not None - use_bias_2 = bias_1 is not None + has_bias_1 = bias_1 is not None + has_bias_2 = bias_2 is not None x = with_sharding_constraint_by_logical_axes(x, norm_input_axes) @@ -328,16 +328,10 @@ def _layernorm_mlp_fwd_rule( casted_kernel_1.get_tensor(TensorUsage.RHS), contracting_dims=(x_contracting_dims, k_contracting_dims), transpose_batch_sequence=transpose_batch_sequence, - bias=bias_1 if not tex.gemm_uses_jax_dot() else None, - fuse_bias=use_bias_1 if not tex.gemm_uses_jax_dot() else False, + bias=bias_1, collective_op=collective_op_set_1.forward, ) - if use_bias_1 and tex.gemm_uses_jax_dot(): - bias_1_shape = bias_1.shape - bias_1_new_shape = (1,) * (dot_1_output.ndim - bias_1.ndim) + bias_1_shape - dot_1_output += jnp.reshape(bias_1, bias_1_new_shape) - # This sharding constraint is needed to correct the Shardy sharding propagation if dot_2_input_axes is not None: dot_1_output_axes = ( @@ -377,16 +371,10 @@ def _layernorm_mlp_fwd_rule( casted_kernel_2.get_tensor(TensorUsage.RHS), contracting_dims=(x_contracting_dims, k_contracting_dims), transpose_batch_sequence=transpose_batch_sequence, - bias=bias_2 if not tex.gemm_uses_jax_dot() else None, - fuse_bias=use_bias_2 if not tex.gemm_uses_jax_dot() else False, + bias=bias_2, collective_op=collective_op_set_2.forward, ) - if use_bias_2 and tex.gemm_uses_jax_dot(): - bias_2_shape = bias_2.shape - bias_2_new_shape = (1,) * (dot_2_output.ndim - bias_2.ndim) + bias_2_shape - dot_2_output += jnp.reshape(bias_2, bias_2_new_shape) - # sharding of outputs should be the same as dot_1's input dot_2_output = with_sharding_constraint_by_logical_axes(dot_2_output, dot_1_input_axes) dot_2_output = checkpoint_name(dot_2_output, ffn2_ckpt_name) @@ -406,8 +394,8 @@ def _layernorm_mlp_fwd_rule( k_contracting_dims, kernel_1.shape, kernel_2.shape, - use_bias_1, - use_bias_2, + has_bias_1, + has_bias_2, quantizer_sets, ) @@ -461,8 +449,8 @@ def _layernorm_mlp_bwd_rule( k_contracting_dims_in_fwd, kernel_1_shape, kernel_2_shape, - use_bias_1, - use_bias_2, + has_bias_1, + has_bias_2, quantizer_sets, ) = ctx @@ -477,7 +465,7 @@ def _layernorm_mlp_bwd_rule( casted_grad, dbias_2 = tex.quantize_dbias( grad, - is_dbias=use_bias_2, + is_dbias=has_bias_2, quantizer=ffn1_quantizer_set.dgrad, amax_scope=AmaxScope.TPSP, transpose_batch_sequence=transpose_batch_sequence, @@ -522,7 +510,7 @@ def _layernorm_mlp_bwd_rule( dgrad_2, dot_1_output, activation_type=activation_type, - is_dbias=use_bias_1, + is_dbias=has_bias_1, quantizer=ffn2_quantizer_set.dgrad, act_params=( tex.activation.ActivationParams.create(activation_type, **activation_params) diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index 7d81d71bc8..c5256aef5c 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -274,9 +274,6 @@ class BaseQuantizeConfig(ABC): COLLECTION_NAME: Name of the collection for quantization metadata FWD_DTYPE: Forward pass data type BWD_DTYPE: Backward pass data type - FP8_2X_ACC_FPROP: Whether to use 2x accumulation for forward pass - FP8_2X_ACC_DGRAD: Whether to use 2x accumulation for data gradients - FP8_2X_ACC_WGRAD: Whether to use 2x accumulation for weight gradients INFERENCE_MODE: Whether to enable optimization for inference AMAX_HISTORY_LEN: Length of AMAX history for delayed scaling AMAX_COMPUTE_ALGO: Algorithm for AMAX computation @@ -287,9 +284,6 @@ class BaseQuantizeConfig(ABC): COLLECTION_NAME: str = NVTE_FP8_COLLECTION_NAME FWD_DTYPE: DType = None BWD_DTYPE: DType = None - FP8_2X_ACC_FPROP: bool = False - FP8_2X_ACC_DGRAD: bool = False - FP8_2X_ACC_WGRAD: bool = False INFERENCE_MODE: bool = False # DelayedScaling @@ -435,9 +429,6 @@ def initialize_from_recipe(self, fp8_recipe: Recipe) -> None: } self.AMAX_COMPUTE_ALGO = string_to_amax_compute_algo[fp8_recipe.amax_compute_algo] - self.FP8_2X_ACC_DGRAD = True - self.FP8_2X_ACC_WGRAD = True - def get_scaling_mode(self, tensor_source: TensorSource) -> ScalingMode: """Gets the scaling mode for a specific tensor's usage type.""" return ScalingMode.DELAYED_TENSOR_SCALING From 34a6c0a5a3b8ad9630a4b199fb1d0fe32805ded1 Mon Sep 17 00:00:00 2001 From: Chaoyang Mei <1192554423@qq.com> Date: Tue, 10 Mar 2026 02:30:11 +0800 Subject: [PATCH 255/521] Fix Flash Attention 3 API compatibility for window size parameters (#2704) * Fix Flash Attention 3 API compatibility for window size parameters Replace single window_size parameter with window_size_left and window_size_right in flash_attn_fwd function to align with flash-attn v2.7.0+ API changes. - Update function signature in flash_attn_interface - Maintain backward compatibility where possible - Ensure consistency with Flash Attention v2 implementation Signed-off-by: Chaoyang Mei <1192554423@qq.com> Signed-off-by: meichaoyang001 * Fix Flash Attention 3 backward API parameter naming Rename causal parameter to is_causal in flash_attn_bwd function to align with flash-attn v2.7.0+ API changes. This ensures consistency with the updated flash-attn library interface for backward pass operations. Signed-off-by: meichaoyang001 * Fix Flash Attention 3 backward API parameter naming Rename causal parameter to is_causal in flash_attn_bwd function to align with flash-attn v2.7.0+ API changes. This ensures consistency with the updated flash-attn library interface for backward pass operations. Signed-off-by: meichaoyang001 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor Flash Attention 3 to use positional args instead of kwargs Replace keyword arguments with positional arguments in flash_attn_fwd and flash_attn_bwd to abstract away parameter naming differences (causal vs is_causal) between flash-attn versions. This provides a more robust interface that is resilient to future API changes in the flash-attn library. - Convert window_size_left, window_size_right, and causal parameters to positional args in both forward and backward functions - Eliminate version-specific parameter naming dependencies - Simplify compatibility handling across flash-attn v2.7.0+ variants Signed-off-by: meichaoyang001 * Fix Flash Attention 3 backward API parameter naming Rename causal parameter to is_causal in flash_attn_bwd function to align with flash-attn v3 API changes. This ensures consistency with the updated flash-attn library interface for backward pass operations. Signed-off-by: meichaoyang001 --------- Signed-off-by: Chaoyang Mei <1192554423@qq.com> Signed-off-by: meichaoyang001 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh --- .../dot_product_attention/context_parallel.py | 46 +++++++++++-------- 1 file changed, 28 insertions(+), 18 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 bd6b626b64..10ba99595b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -937,9 +937,9 @@ def cp_p2p_fwd_flash_attn( elif section == "upper-triangle": max_seqlen_q_ = max_seqlen_q // 2 if section in ["lower-triangle", "upper-triangle"]: - if use_flash_attn_3 or (fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus): + if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: fa_forward_kwargs["window_size"] = (-1, -1) - elif fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or fa_utils.v2_7_0_plus: fa_forward_kwargs["window_size_left"] = -1 fa_forward_kwargs["window_size_right"] = -1 @@ -1189,9 +1189,9 @@ def cp_p2p_bwd_flash_attn( ): """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 use_flash_attn_3 or (fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus): + if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size"] = (-1, -1) - elif fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size_left"] = -1 fa_backward_kwargs["window_size_right"] = -1 if not use_flash_attn_3: @@ -1201,9 +1201,9 @@ def cp_p2p_bwd_flash_attn( softmax_lse__ = softmax_lse causal_ = False if section == "diagonal": - if use_flash_attn_3 or (fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus): + if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size"] = (-1, 0) - elif fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size_left"] = -1 fa_backward_kwargs["window_size_right"] = 0 causal_ = True @@ -1225,6 +1225,10 @@ def cp_p2p_bwd_flash_attn( dk=dk, dv=dv, ) + if use_flash_attn_3: + fa_backward_kwargs["is_causal"] = causal_ + else: + fa_backward_kwargs["causal"] = causal_ flash_attn_bwd( dout_part, q_part, @@ -1233,7 +1237,6 @@ def cp_p2p_bwd_flash_attn( out_part, softmax_lse__, *fa_backward_args_thd, - causal=causal_, **fa_backward_kwargs, ) @@ -1508,7 +1511,8 @@ def forward( flash_attn_fwd = ( _flash_attn_fwd_v3 # pylint: disable=possibly-used-before-assignment ) - fa_forward_kwargs["window_size"] = (-1, 0) if causal else (-1, -1) + fa_forward_kwargs["window_size_left"] = -1 + fa_forward_kwargs["window_size_right"] = 0 if causal else -1 else: if qkv_format == "thd": from transformer_engine.pytorch.attention.dot_product_attention.backends import ( @@ -2985,9 +2989,9 @@ def forward( max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv_, ) - if use_flash_attn_3 or (fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus): + if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: fa_forward_kwargs["window_size"] = window_size_per_step[i] - elif fa_utils.v2_7_0_plus: + elif use_flash_attn_3 or fa_utils.v2_7_0_plus: 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( @@ -3206,13 +3210,15 @@ def backward(ctx, dout, *_args): ) if not ctx.use_flash_attn_3: fa_backward_kwargs["rng_state"] = rng_states[i] - if ctx.use_flash_attn_3 or ( - fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus - ): + if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size"] = window_size_per_step[i] - elif fa_utils.v2_7_0_plus: + elif ctx.use_flash_attn_3 or fa_utils.v2_7_0_plus: 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 + else: + fa_backward_kwargs["causal"] = "causal" in ctx.attn_mask_type flash_attn_bwd( dout_, q_, @@ -3221,7 +3227,6 @@ def backward(ctx, dout, *_args): out_, softmax_lse_per_step[i], *fa_backward_args_thd, - causal="causal" in ctx.attn_mask_type, **fa_backward_kwargs, ) @@ -3361,7 +3366,8 @@ def forward( ) flash_attn_fwd = _flash_attn_fwd_v3 - fa_forward_kwargs["window_size"] = window_size + fa_forward_kwargs["window_size_left"] = window_size[0] + fa_forward_kwargs["window_size_right"] = window_size[1] else: if qkv_format == "thd": from transformer_engine.pytorch.attention.dot_product_attention.backends import ( @@ -3738,7 +3744,8 @@ def backward(ctx, dout, *_args): flash_attn_bwd = ( _flash_attn_bwd_v3 # pylint: disable=possibly-used-before-assignment ) - fa_backward_kwargs["window_size"] = ctx.window_size + fa_backward_kwargs["window_size_left"] = ctx.window_size[0] + fa_backward_kwargs["window_size_right"] = ctx.window_size[1] fa_backward_kwargs["deterministic"] = ctx.deterministic else: if qkv_format == "thd": @@ -3821,6 +3828,10 @@ def backward(ctx, dout, *_args): ) if not ctx.use_flash_attn_3: fa_backward_kwargs["rng_state"] = rng_state + fa_backward_kwargs["causal"] = causal + else: + fa_backward_kwargs["is_causal"] = causal + flash_attn_bwd( dout, q, @@ -3829,7 +3840,6 @@ def backward(ctx, dout, *_args): out, softmax_lse, *fa_backward_args_thd, - causal=causal, **fa_backward_kwargs, ) From 6e0085ad5d535820f39699e8bdfbae4b0708dd56 Mon Sep 17 00:00:00 2001 From: Sung Hyun Cho Date: Tue, 10 Mar 2026 03:37:43 +0900 Subject: [PATCH 256/521] [Common] Remove redundant grad_logits zero-initialization in fused router backward kernels (#2745) Remove redundant grad_logits zero-initialization in fused router backward kernels Signed-off-by: Sung Hyun Cho Co-authored-by: Xin Yao --- .../common/fused_router/fused_score_for_moe_aux_loss.cu | 4 ---- .../common/fused_router/fused_topk_with_score_function.cu | 4 ---- 2 files changed, 8 deletions(-) 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 d38fcde6bf..ebdcb293e0 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 @@ -211,10 +211,6 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const CompType *int * - Load the dgrad/output_from_fwd to shmem */ int pos_offset = token_offset_cur_warp * num_experts; - // Clear the logits_grad in global mem - for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - grad_logits[pos_offset + i] = 0.0f; - } // Load the dgrad/output_from_fwd to shmem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { local_grad[i] = grad_scores[pos_offset + i]; 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 a9e680f06e..1bed871de8 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 @@ -328,10 +328,6 @@ __global__ void fused_topk_with_score_function_backward_kernel( * - Load the dgrad/output_from_fwd to shmem */ int pos_offset = token_offset_cur_warp * num_experts; - // Clear the logits_grad in global mem - for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { - grad_logits[pos_offset + i] = 0.0f; - } // Load the dgrad/output_from_fwd to shmem for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { local_grad[i] = grad_probs[pos_offset + i]; From f64941a0ab1832775396d2e3efc5def346cc012d Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Mon, 9 Mar 2026 17:01:25 -0700 Subject: [PATCH 257/521] Enable dequantization from MXFP8 tensor with only columnwise data (#2712) Enable dequantization from just columnwise data Signed-off-by: Przemek Tredak --- tests/pytorch/test_quantized_tensor.py | 83 +++++++++++++++++++ .../tensor/storage/mxfp8_tensor_storage.py | 4 +- 2 files changed, 85 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index b2e8fca7cb..978ec09b40 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -656,3 +656,86 @@ def test_chunk( tols = dict(rtol=0, atol=0) # Chunking is exact y_test = y_test.to(dtype=torch.float64, device="cpu") torch.testing.assert_close(y_test, y_ref, **tols) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +class TestMXFP8Tensor: + + @staticmethod + def setup_class(cls) -> None: + # Configure RNG + seed = 1234 + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + @pytest.mark.parametrize("fp8_dtype", _fp8_dtypes) + @pytest.mark.parametrize("dtype", _dtypes) + @pytest.mark.parametrize("dims", [[128, 128], [256, 256], [128, 256]]) + def test_mxfp8_dequantize_columnwise_only( + self, + fp8_dtype: tex.DType, + dtype: torch.dtype, + dims: DimsType, + ) -> None: + """Check dequantization of MXFP8 tensor with only columnwise data""" + + # Initialize random data + x_ref = 2 * torch.rand(_to_list(dims), dtype=dtype, device="cuda") - 1 + + # Quantize with both rowwise and columnwise + quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=True, columnwise=True) + x_mxfp8 = quantizer(x_ref) + + # Dequantize from rowwise (default path) + x_deq_rowwise = x_mxfp8.dequantize(dtype=dtype) + + # Rowwise dequantization should be close to the original + torch.testing.assert_close(x_deq_rowwise, x_ref, **_tols[fp8_dtype]) + + # Strip rowwise data, keeping only columnwise + x_mxfp8.update_usage(rowwise_usage=False, columnwise_usage=True) + assert x_mxfp8._rowwise_data is None + assert x_mxfp8._columnwise_data is not None + + # Dequantize from columnwise only + x_deq_columnwise = x_mxfp8.dequantize(dtype=dtype) + + # Columnwise dequantization should be close to the original + torch.testing.assert_close(x_deq_columnwise, x_ref, **_tols[fp8_dtype]) + + # Rowwise and columnwise dequantizations should match each other + torch.testing.assert_close(x_deq_columnwise, x_deq_rowwise, **_tols[fp8_dtype]) + + # Make sure we are not trivially passing the test + with pytest.raises(AssertionError): + torch.testing.assert_close(x_deq_columnwise, -x_ref, **_tols[fp8_dtype]) + + @pytest.mark.parametrize("fp8_dtype", _fp8_dtypes) + @pytest.mark.parametrize("dims", [[128, 128], [256, 256]]) + def test_mxfp8_dequantize_columnwise_only_quantized_separately( + self, + fp8_dtype: tex.DType, + dims: DimsType, + ) -> None: + """Check dequantization of MXFP8 tensor quantized with columnwise only""" + + dtype = torch.bfloat16 + + # Initialize random data + x_ref = 2 * torch.rand(_to_list(dims), dtype=dtype, device="cuda") - 1 + + # Quantize with columnwise only (no rowwise) + quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype, rowwise=False, columnwise=True) + x_mxfp8 = quantizer(x_ref) + assert x_mxfp8._rowwise_data is None + assert x_mxfp8._columnwise_data is not None + + # Dequantize from columnwise only + x_deq = x_mxfp8.dequantize(dtype=dtype) + + # Should be close to the original + torch.testing.assert_close(x_deq, x_ref, **_tols[fp8_dtype]) + + # Make sure we are not trivially passing the test + with pytest.raises(AssertionError): + torch.testing.assert_close(x_deq, -x_ref, **_tols[fp8_dtype]) diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 12757aa58c..64344b78a1 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -33,9 +33,9 @@ def forward( dtype = torch_to_transformer_engine_dtype[dtype] # Make sure FP8 data is in expected format - if tensor._rowwise_data is not None: + if tensor._rowwise_data is not None or tensor._columnwise_data is not None: return tex.dequantize(tensor, dtype) - raise NotImplementedError("Casting back from the transpose not implemented yet!") + raise ValueError("Cannot dequantize MXFP8 tensor with no data") @staticmethod def backward( From e6d97ffc7369515fe319d2b700f6b583e0519735 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EB=85=B8=EB=9E=80=ED=86=A0=EB=81=BC?= <83907395+Bias92@users.noreply.github.com> Date: Tue, 10 Mar 2026 11:30:08 +0900 Subject: [PATCH 258/521] [PyTorch] Fix cross_entropy_forward stride guard for non-contiguous input (#2746) * Fix cross_entropy_forward stride guard for non-contiguous input Signed-off-by: Bias92 * Add regression test for non-contiguous transposed input Signed-off-by: Bias92 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Bias92 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_parallel_cross_entropy.py | 21 +++++++++++++++++++ .../pytorch/triton/cross_entropy.py | 2 +- 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_parallel_cross_entropy.py b/tests/pytorch/test_parallel_cross_entropy.py index 7b92672af7..b4ea193f06 100644 --- a/tests/pytorch/test_parallel_cross_entropy.py +++ b/tests/pytorch/test_parallel_cross_entropy.py @@ -167,3 +167,24 @@ def test_ignore_idx_reduced_loss(self): reduce_loss=True, ignore_idx=True, ) + + +def test_non_contiguous_transposed_input(): + """Regression test: stride(-2) != shape[-1] should not produce wrong results.""" + s, b, v = 4, 2, 8 + torch.manual_seed(42) + logits = torch.randn(s, b, v, device="cuda") + target = torch.randint(0, v, (b, s), device="cuda") + + logits_transposed = logits.transpose(0, 1) # stride(-2) != shape[-1] + logits_contiguous = logits_transposed.contiguous() + + assert logits_transposed.stride(-1) == 1 + assert logits_transposed.stride(-2) != logits_transposed.shape[-1] + + loss_t = parallel_cross_entropy(logits_transposed, target, 0.0, False, None) + loss_c = parallel_cross_entropy(logits_contiguous, target, 0.0, False, None) + + assert torch.allclose( + loss_t, loss_c + ), f"Non-contiguous transposed input gave wrong results: {loss_t} vs {loss_c}" diff --git a/transformer_engine/pytorch/triton/cross_entropy.py b/transformer_engine/pytorch/triton/cross_entropy.py index b574d69e0f..1401383c8f 100644 --- a/transformer_engine/pytorch/triton/cross_entropy.py +++ b/transformer_engine/pytorch/triton/cross_entropy.py @@ -49,7 +49,7 @@ def cross_entropy_forward( n_non_ignore = torch.zeros(1, dtype=torch.int64, device=_input.device) # ensure _input and target are contiguous in the last dimension - if _input.stride(-1) != 1: + if _input.stride(-1) != 1 or _input.stride(-2) != _input.shape[-1]: _input = _input.contiguous() if target.stride(-1) != 1: target = target.contiguous() From 7c2aa2cd3f8a9cdc32b13cfaf9299391d71fc1b3 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Tue, 10 Mar 2026 19:46:39 +0100 Subject: [PATCH 259/521] [Common] MOE Split dBias (#2674) * Implemented the kernel with split dbias 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 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Oleg Goncharov * Relaxed constraints on the last dimension Signed-off-by: Oleg Goncharov * Added notes on group tensor restrictions into documentation Signed-off-by: Oleg Goncharov * Fixes per the review Signed-off-by: Oleg Goncharov * Fixed pointer Signed-off-by: Oleg Goncharov * More fixes Signed-off-by: Oleg Goncharov * Fixed kernel grid size 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> --- tests/cpp/operator/test_cast_mxfp8_grouped.cu | 150 +++++++++--------- transformer_engine/common/activation/gelu.cu | 8 +- transformer_engine/common/activation/relu.cu | 8 +- .../common/activation/swiglu.cu | 4 +- transformer_engine/common/cast/cast.cu | 2 +- .../common/cast/core/common.cuh | 84 ++++++++++ .../common/cast/dispatch/quantize.cuh | 11 +- .../cast/mxfp8/group_quantize_mxfp8.cuh | 41 +++-- .../include/transformer_engine/activation.h | 10 ++ .../common/include/transformer_engine/cast.h | 42 +++-- 10 files changed, 236 insertions(+), 124 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 6557c83773..e469ad0845 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -58,8 +58,7 @@ void compute_ref(const ProcessingMethod processing_method, const size_t rows, const size_t cols, const size_t scales_stride_rowwise, - const size_t scales_stride_colwise, - const bool is_single_tensor) + const size_t scales_stride_colwise) { const size_t tile_size_Y = 32; const size_t tile_size_X = 32; @@ -169,10 +168,8 @@ void compute_ref(const ProcessingMethod processing_method, } } - if (is_single_tensor) { - for (size_t j = 0; j < cols; ++j) { - output_dbias[j] = static_cast(output_dbias_fp32[j]); - } + for (size_t j = 0; j < cols; ++j) { + output_dbias[j] = static_cast(output_dbias_fp32[j]); } } @@ -250,12 +247,16 @@ void performTest(const ProcessingMethod processing_method, DType itype = TypeInfo::dtype; DType otype = TypeInfo::dtype; + const bool compute_dbias = (processing_method == ProcessingMethod::CAST_DBIAS + || processing_method == ProcessingMethod::CAST_DBIAS_DACT); + const size_t rows = logical_shape_vec[0]; const size_t cols = logical_shape_vec[1]; size_t elts_num = 0; size_t rowwise_sfs_num = 0; size_t colwise_sfs_num = 0; + size_t sum_of_last_dims = 0; std::vector rowwise_scales_first_dim(num_tensors, 0); std::vector rowwise_scales_last_dim(num_tensors, 0); @@ -263,6 +264,7 @@ void performTest(const ProcessingMethod processing_method, std::vector colwise_scales_first_dim(num_tensors, 0); std::vector colwise_scales_last_dim(num_tensors, 0); std::vector colwise_scales_offset(num_tensors + 1, 0); + std::vector dbias_offsets(num_tensors + 1, 0); for (size_t t = 0; t < num_tensors; ++t) { const size_t M = first_dims_h[t]; @@ -285,13 +287,13 @@ void performTest(const ProcessingMethod processing_method, rowwise_sfs_num += rowwise_sfs; colwise_sfs_num += colwise_sfs; + sum_of_last_dims += K; rowwise_scales_offset[t+1] = rowwise_sfs_num; colwise_scales_offset[t+1] = colwise_sfs_num; + dbias_offsets[t+1] = sum_of_last_dims; } - const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS) || (shape_rep == VARYING_FIRST_DIM); - std::vector scales_rowwise_shape = {rowwise_sfs_num}; std::vector scales_colwise_shape = {colwise_sfs_num}; @@ -311,7 +313,7 @@ void performTest(const ProcessingMethod processing_method, std::vector out_scales_rowwise_ref(rowwise ? rowwise_sfs_num : 0); std::vector out_scales_colwise_ref(colwise ? colwise_sfs_num : 0); - std::vector ref_output_dbias(is_single_tensor ? cols : 0); + std::vector ref_output_dbias(sum_of_last_dims, static_cast(0.0f)); for (size_t i = 0; i < elts_num; ++i) { const float val = dis(gen); @@ -336,6 +338,7 @@ void performTest(const ProcessingMethod processing_method, const size_t in_data_size = elts_num * sizeof(InputType); const size_t out_data_size = elts_num * sizeof(OutputType); + const size_t dbias_data_size = sum_of_last_dims * sizeof(InputType); const size_t rowwise_scales_size = rowwise_sfs_num * sizeof(fp8e8m0); const size_t colwise_scales_size = colwise_sfs_num * sizeof(fp8e8m0); @@ -343,15 +346,16 @@ void performTest(const ProcessingMethod processing_method, const size_t last_dims_size = num_tensors * sizeof(size_t); const size_t offsets_size = (num_tensors + 1) * sizeof(size_t); - InputType* grad_data_d; - InputType* in_data_d; - OutputType* out_data_rowwise_d; - OutputType* out_data_colwise_d; - fp8e8m0* out_scales_rowwise_d; - fp8e8m0* out_scales_colwise_d; - size_t* first_dims_d; - size_t* last_dims_d; - size_t* offsets_d; + InputType* grad_data_d = nullptr; + InputType* in_data_d = nullptr; + InputType* dbias_out_data_d = nullptr; + OutputType* out_data_rowwise_d = nullptr; + OutputType* out_data_colwise_d = nullptr; + fp8e8m0* out_scales_rowwise_d = nullptr; + fp8e8m0* out_scales_colwise_d = nullptr; + size_t* first_dims_d = nullptr; + size_t* last_dims_d = nullptr; + size_t* offsets_d = nullptr; cudaMalloc((void**)&grad_data_d, in_data_size); cudaMalloc((void**)&in_data_d, in_data_size); @@ -367,6 +371,10 @@ void performTest(const ProcessingMethod processing_method, NVTEShape logical_shape_ = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); + std::vector dbias_logical_shape_vec= {num_tensors, cols}; + NVTEShape dbias_logical_shape_ = nvte_make_shape(dbias_logical_shape_vec.data(), + dbias_logical_shape_vec.size()); + NVTEShape first_dims_shape_; NVTEShape last_dims_shape_; NVTEShape offsets_shape_; @@ -382,6 +390,7 @@ void performTest(const ProcessingMethod processing_method, NVTEGroupedTensor grad_group_tensor = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape_); NVTEGroupedTensor in_group_tensor = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape_); NVTEGroupedTensor out_group_tensor = nvte_create_grouped_tensor(NVTE_MXFP8_1D_SCALING, num_tensors, logical_shape_); + NVTEGroupedTensor output_dbias_tensor = nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, dbias_logical_shape_); NVTEBasicTensor grad_data_tensor = {grad_data_d, static_cast(itype), logical_shape_}; NVTEBasicTensor in_data_tensor = {in_data_d, static_cast(itype), logical_shape_}; @@ -453,52 +462,41 @@ void performTest(const ProcessingMethod processing_method, &out_scales_colwise_tensor, sizeof(out_scales_colwise_tensor)); } - Tensor output_dbias("output_dbias", std::vector{ cols }, itype); + if (compute_dbias) { + cudaMalloc((void**)&dbias_out_data_d, dbias_data_size); + cudaMemset(dbias_out_data_d, 0, dbias_data_size); + NVTEBasicTensor output_dbias_data_tensor = {dbias_out_data_d, static_cast(itype), dbias_logical_shape_}; + nvte_set_grouped_tensor_param(output_dbias_tensor, NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, + &output_dbias_data_tensor, sizeof(output_dbias_data_tensor)); + } // Reference (CPU) - if (is_single_tensor) { - - const size_t unpadded_rowwise_blocks_X = divide_round_up(cols, 32); - const size_t unpadded_colwise_blocks_X = cols; + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; - const size_t scales_stride_rowwise = round_up_to_nearest_multiple(unpadded_rowwise_blocks_X, 4); - const size_t scales_stride_colwise = round_up_to_nearest_multiple(unpadded_colwise_blocks_X, 128); + const size_t scales_stride_rowwise = rowwise_scales_last_dim[t]; + const size_t scales_stride_colwise = colwise_scales_last_dim[t]; + const size_t data_offset = offsets_h[t]; + const size_t rowwise_sfs_offset = rowwise_scales_offset[t]; + const size_t colwise_sfs_offset = colwise_scales_offset[t]; + const size_t dbias_offset = dbias_offsets[t]; + + const InputType* const grad_ptr = grad_data.data() + data_offset; + const InputType* const in_ptr = in_data.data() + data_offset; + OutputType* const out_data_rowwise_ptr = out_data_rowwise_ref.data() + data_offset; + OutputType* const out_data_colwise_ptr = out_data_colwise_ref.data() + data_offset; + fp8e8m0* const out_scales_rowwise_ptr = out_scales_rowwise_ref.data() + rowwise_sfs_offset; + fp8e8m0* const out_scales_colwise_ptr = out_scales_colwise_ref.data() + colwise_sfs_offset; + InputType* const ref_output_dbias_ptr = ref_output_dbias.data() + dbias_offset; compute_ref( - processing_method, OP, rowwise, colwise, in_data.data(), grad_data.data(), - out_data_rowwise_ref.data(), out_data_colwise_ref.data(), - out_scales_rowwise_ref.data(), out_scales_colwise_ref.data(), - ref_output_dbias.data(), rows, cols, + processing_method, OP, rowwise, colwise, in_ptr, grad_ptr, + out_data_rowwise_ptr, out_data_colwise_ptr, + out_scales_rowwise_ptr, out_scales_colwise_ptr, + ref_output_dbias_ptr, M, K, scales_stride_rowwise, - scales_stride_colwise, - is_single_tensor); - } else { - for (size_t t = 0; t < num_tensors; ++t) { - const size_t M = first_dims_h[t]; - const size_t K = last_dims_h[t]; - - const size_t scales_stride_rowwise = rowwise_scales_last_dim[t]; - const size_t scales_stride_colwise = colwise_scales_last_dim[t]; - const size_t data_offset = offsets_h[t]; - const size_t rowwise_sfs_offset = rowwise_scales_offset[t]; - const size_t colwise_sfs_offset = colwise_scales_offset[t]; - - const InputType* const grad_ptr = grad_data.data() + data_offset; - const InputType* const in_ptr = in_data.data() + data_offset; - OutputType* const out_data_rowwise_ptr = out_data_rowwise_ref.data() + data_offset; - OutputType* const out_data_colwise_ptr = out_data_colwise_ref.data() + data_offset; - fp8e8m0* const out_scales_rowwise_ptr = out_scales_rowwise_ref.data() + rowwise_sfs_offset; - fp8e8m0* const out_scales_colwise_ptr = out_scales_colwise_ref.data() + colwise_sfs_offset; - - compute_ref( - processing_method, OP, rowwise, colwise, in_ptr, grad_ptr, - out_data_rowwise_ptr, out_data_colwise_ptr, - out_scales_rowwise_ptr, out_scales_colwise_ptr, - ref_output_dbias.data(), M, K, - scales_stride_rowwise, - scales_stride_colwise, - is_single_tensor); - } + scales_stride_colwise); } // GPU @@ -509,9 +507,9 @@ void performTest(const ProcessingMethod processing_method, break; } case ProcessingMethod::CAST_DBIAS: { - nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias.data(), workspace.data(), 0); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0); workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); - nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias.data(), workspace.data(), 0); + nvte_group_quantize_dbias(grad_group_tensor, out_group_tensor, output_dbias_tensor, workspace.data(), 0); break; } case ProcessingMethod::CAST_DBIAS_DACT: { @@ -522,10 +520,10 @@ void performTest(const ProcessingMethod processing_method, else if (OP == &dsrelu) { nvte_group_quantize_dbias_dact = &nvte_group_quantize_dbias_dsrelu; } nvte_group_quantize_dbias_dact(grad_group_tensor, in_group_tensor, out_group_tensor, - output_dbias.data(), workspace.data(), 0); + output_dbias_tensor, workspace.data(), 0); workspace = Tensor("workspace", workspace.rowwise_shape(), workspace.dtype()); nvte_group_quantize_dbias_dact(grad_group_tensor, in_group_tensor, out_group_tensor, - output_dbias.data(), workspace.data(), 0); + output_dbias_tensor, workspace.data(), 0); break; } case ProcessingMethod::CAST_ACT: { @@ -586,9 +584,10 @@ void performTest(const ProcessingMethod processing_method, out_data_colwise_h.data(), rows, cols, false, mismatches_elts); } - if (processing_method == ProcessingMethod::CAST_DBIAS - || processing_method == ProcessingMethod::CAST_DBIAS_DACT) - { + if (compute_dbias) { + Tensor output_dbias("output_dbias", std::vector{ sum_of_last_dims }, itype); + cudaMemcpy(output_dbias.rowwise_dptr(), dbias_out_data_d, dbias_data_size, cudaMemcpyDeviceToDevice); + auto [atol_dbias, rtol_dbias] = getTolerances(itype); if (itype == DType::kFloat32) { atol_dbias = 1e-4; @@ -601,6 +600,7 @@ void performTest(const ProcessingMethod processing_method, cudaFree(grad_data_d); cudaFree(in_data_d); + cudaFree(dbias_out_data_d); cudaFree(first_dims_d); cudaFree(last_dims_d); cudaFree(offsets_d); @@ -648,7 +648,8 @@ std::vector> input_config = { {SAME_BOTH_DIMS, 1, 128,128}, {SAME_BOTH_DIMS, 2, 256,128}, {VARYING_FIRST_DIM, 2, 512,128, 128,384}, - {VARYING_FIRST_DIM, 2, 384,160, 128,256}, + {VARYING_FIRST_DIM, 3, 1024,144, 128,384,512}, + {VARYING_FIRST_DIM, 4, 1536,160, 128,384,512,512}, {VARYING_FIRST_DIM, 5, 4096,512, 128,256,384,1024,2304}, {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, {VARYING_BOTH_DIMS, 2, 1,(128*128)+(256*256), 128,256, 128,256}, @@ -714,26 +715,31 @@ TEST_P(GroupedFusedCastMXFP8TestSuite, Test) { } } offsets[t+1] = offsets[t] + first_dims[t] * last_dims[t]; - // Skips tests if tensor shape is not as required by the kernel - if (first_dims[t] % 128 != 0) { + // Skip tests when the tensor shape is incompatible with the kernel. + // The TMA engine requires strides to be 16-byte aligned. + if ((first_dims[t] % 128 != 0) || (last_dims[t] % 16 != 0)) { GTEST_SKIP(); } - if (!is_single_tensor && (last_dims[t] % 128 != 0)) { + // If a grouped tensor has a varying last dimension, it must be a multiple of 128. + // Otherwise, computing the grid size adds runtime overhead in the non-persistent kernel, + // since the relevant tensor metadata resides in device memory. + constexpr size_t CHUNK_DIM_X = 128; + if (!is_single_tensor && (last_dims[t] % CHUNK_DIM_X != 0)) { GTEST_SKIP(); } } - // Skips DBias tests if last dimension of tensors variates + // Skip dBias tests when tensors in the group have different last dimensions. if ((processing_method == ProcessingMethod::CAST_DBIAS || processing_method == ProcessingMethod::CAST_DBIAS_DACT) && !is_single_tensor) { GTEST_SKIP(); } - // Skips non Act tests if the Activation type is not an identity + // Skip non-activation tests when the activation type is not Identity. if ((processing_method == ProcessingMethod::CAST_ONLY || processing_method == ProcessingMethod::CAST_DBIAS) && activation != ActivationKind::Identity) { GTEST_SKIP(); } - // Skips Act tests if the Activation is an identity + // Skip activation tests when the activation type is Identity. if ((processing_method == ProcessingMethod::CAST_DBIAS_DACT || processing_method == ProcessingMethod::CAST_DACT || processing_method == ProcessingMethod::CAST_ACT) && (activation == ActivationKind::Identity)) { diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index d209ea8d47..ea864813bf 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -32,7 +32,7 @@ void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inpu NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_dgelu); using namespace transformer_engine; - NVTETensor dbias = nullptr; + NVTEGroupedTensor dbias = nullptr; NVTETensor workspace = nullptr; constexpr bool IS_DBIAS = false; @@ -57,7 +57,7 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor activati void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTETensor dbias, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias_dgelu); using namespace transformer_engine; @@ -110,7 +110,7 @@ void nvte_group_dqgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inp NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_dqgelu); using namespace transformer_engine; - NVTETensor dbias = nullptr; + NVTEGroupedTensor dbias = nullptr; NVTETensor workspace = nullptr; constexpr bool IS_DBIAS = false; @@ -135,7 +135,7 @@ void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor activat void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTETensor dbias, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias_dqgelu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/relu.cu b/transformer_engine/common/activation/relu.cu index b6f758caf6..fc9122b7ec 100644 --- a/transformer_engine/common/activation/relu.cu +++ b/transformer_engine/common/activation/relu.cu @@ -32,7 +32,7 @@ void nvte_group_drelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inpu NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_drelu); using namespace transformer_engine; - NVTETensor dbias = nullptr; + NVTEGroupedTensor dbias = nullptr; NVTETensor workspace = nullptr; constexpr bool IS_DBIAS = false; @@ -57,7 +57,7 @@ void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor activati void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTETensor dbias, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias_drelu); using namespace transformer_engine; @@ -110,7 +110,7 @@ void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inp NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_dsrelu); using namespace transformer_engine; - NVTETensor dbias = nullptr; + NVTEGroupedTensor dbias = nullptr; NVTETensor workspace = nullptr; constexpr bool IS_DBIAS = false; @@ -135,7 +135,7 @@ void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor activat void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTETensor dbias, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias_dsrelu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index 77d5b6867f..12478af4cf 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -32,7 +32,7 @@ void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor inpu NVTEGroupedTensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_group_dsilu); using namespace transformer_engine; - NVTETensor dbias = nullptr; + NVTEGroupedTensor dbias = nullptr; NVTETensor workspace = nullptr; constexpr bool IS_DBIAS = false; @@ -57,7 +57,7 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor activati void nvte_group_quantize_dbias_dsilu(const NVTEGroupedTensor input, const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTETensor dbias, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias_dsilu); using namespace transformer_engine; diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 57404ae8a5..4f9ddb4fc5 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -70,7 +70,7 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d } void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream) { + NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_group_quantize_dbias); using namespace transformer_engine; diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh index 0997b01f7e..a4e033939b 100644 --- a/transformer_engine/common/cast/core/common.cuh +++ b/transformer_engine/common/cast/core/common.cuh @@ -22,6 +22,14 @@ namespace transformer_engine { namespace dispatch { namespace common { + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_block) { const size_t N = product(t->data.shape); const bool isFullTile = (N % elems_per_block == 0); @@ -78,6 +86,56 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) } stg_vec.store_to(thread_out_base); } + +template +__global__ void __launch_bounds__(THREADS_PER_BLOCK) + group_reduce_dbias_kernel(const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const offsets_ptr, const int64_t *const first_dims_ptr, + const int64_t *const last_dims_ptr, OType *const dbias_output, + const float *dbias_partial, const size_t chunk_dim_Y) { + using ComputeVec = Vec; + using OutputVec = Vec; + + const size_t tensor_id = blockIdx.y; + const size_t tensor_rows = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) + ? (first_logical_dim / num_tensors) + : first_dims_ptr[tensor_id]; + + const size_t rows = tensor_rows / chunk_dim_Y; + const size_t cols = last_logical_dim; + + const size_t dbias_in_offset_Y = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) + ? (tensor_id * (tensor_rows / chunk_dim_Y)) + : (offsets_ptr[tensor_id] / cols / chunk_dim_Y); + + const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; + + if (thread_id * nvec >= cols) { + return; + } + + const float *const thread_in_base = dbias_partial + dbias_in_offset_Y * cols + thread_id * nvec; + OType *const thread_out_base = dbias_output + tensor_id * cols + thread_id * nvec; + + ComputeVec ldg_vec; + ComputeVec acc_vec; + acc_vec.clear(); + for (int i = 0; i < rows; ++i) { + ldg_vec.load_from(thread_in_base + i * cols); +#pragma unroll + for (int e = 0; e < nvec; ++e) { + acc_vec.data.elt[e] += ldg_vec.data.elt[e]; + } + } + + OutputVec stg_vec; +#pragma unroll + for (int e = 0; e < nvec; ++e) { + stg_vec.data.elt[e] = static_cast(acc_vec.data.elt[e]); + } + stg_vec.store_to(thread_out_base); +} } // namespace kernel template @@ -96,6 +154,32 @@ void reduce_dbias(const float *workspace_ptr, Tensor *dbias, const size_t rows, NVTE_CHECK_CUDA(cudaGetLastError()); } +template +void grouped_reduce_dbias(const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const data_tensor_offsets_ptr, + const int64_t *const data_tensor_first_dims_ptr, + const int64_t *const data_tensor_last_dims_ptr, GroupedTensor *dbias, + const float *workspace_ptr, const size_t chunk_dim_Y, + cudaStream_t stream) { + using namespace kernel; + constexpr size_t reduce_dbias_store_bytes = 8; // stg.64 + constexpr size_t reduce_dbias_nvec = reduce_dbias_store_bytes / sizeof(IType); + + NVTE_CHECK(last_logical_dim % reduce_dbias_nvec == 0, "Unsupported shape."); + + const size_t blocks_X = DIVUP(last_logical_dim, THREADS_PER_BLOCK * reduce_dbias_nvec); + const size_t blocks_Y = num_tensors; + const dim3 grid(blocks_X, blocks_Y); + + group_reduce_dbias_kernel<<>>( + shape_rep, num_tensors, first_logical_dim, last_logical_dim, data_tensor_offsets_ptr, + data_tensor_first_dims_ptr, data_tensor_last_dims_ptr, + reinterpret_cast(dbias->data.dptr), workspace_ptr, chunk_dim_Y); + + NVTE_CHECK_CUDA(cudaGetLastError()); +} + } // namespace common } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 98a3fb8cba..f7823b4c58 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -382,13 +382,13 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor NVTEScalingMode scaling_mode = nvte_grouped_tensor_scaling_mode(output); const NVTEGroupedTensor activation = nullptr; - NVTETensor dbias = nullptr; + NVTEGroupedTensor dbias = nullptr; NVTETensor workspace = nullptr; const GroupedTensor *input_tensor = convertNVTEGroupedTensorCheck(input); GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); const GroupedTensor *activations_tensor = convertNVTEGroupedTensor(activation); - Tensor *dbias_tensor = convertNVTETensor(dbias); + GroupedTensor *dbias_tensor = convertNVTEGroupedTensor(dbias); Tensor *workspace_tensor = convertNVTETensor(workspace); // Quantization config @@ -419,8 +419,9 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor template void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, NVTETensor dbias, NVTETensor workspace, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { using namespace detail; NVTEScalingMode scaling_mode = nvte_grouped_tensor_scaling_mode(output); @@ -428,7 +429,7 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe const GroupedTensor *grad_tensor = convertNVTEGroupedTensorCheck(grad); const GroupedTensor *input_tensor = convertNVTEGroupedTensor(input); GroupedTensor *output_tensor = convertNVTEGroupedTensorCheck(output); - Tensor *dbias_tensor = convertNVTETensor(dbias); + GroupedTensor *dbias_tensor = convertNVTEGroupedTensor(dbias); Tensor *workspace_tensor = convertNVTETensor(workspace); // Quantization config diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 6447fc4542..129d6724ac 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -28,19 +28,14 @@ namespace dispatch { namespace mxfp8 { namespace group_quantize_kernel { +using namespace dispatch::common; + constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; __device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; __device__ alignas(128) CUtensorMap g_tensor_maps_act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; __device__ alignas(128) CUtensorMap g_tensor_maps_output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; __device__ alignas(128) CUtensorMap g_tensor_maps_output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; -enum ShapeRepresentation { - SAME_BOTH_DIMS = 0, - VARYING_FIRST_DIM = 1, - VARYING_LAST_DIM = 2, - VARYING_BOTH_DIMS = 3 -}; - constexpr size_t SCALE_DIM_Y = 32; constexpr size_t SCALE_DIM_X = 32; @@ -111,6 +106,9 @@ __device__ __forceinline__ size_t get_tensor_rows_num( rows_num = static_cast(first_dims_ptr[tensor_id]); break; } + if (rows_num % 128 != 0) { + NVTE_DEVICE_ERROR("First dimension of each tensor in a group must be divisible by 128."); + } return rows_num; } @@ -144,7 +142,7 @@ __device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_te if constexpr (is_blackwell) { const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Shape not supported, as data stride must be 16B aligned."); + NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); } if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); @@ -782,8 +780,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel template void group_quantize(const GroupedTensor *input, const GroupedTensor *activations, - const Tensor *noop, GroupedTensor *output, Tensor *dbias, Tensor *workspace, - cudaStream_t stream) { + const Tensor *noop, GroupedTensor *output, GroupedTensor *dbias, + Tensor *workspace, cudaStream_t stream) { using namespace group_quantize_kernel; checkCuDriverContext(stream); @@ -841,12 +839,9 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); } else { - NVTE_CHECK(num_tensors < MAX_SUPPORTED_TENSOR_DESCRIPTORS, + NVTE_CHECK(num_tensors <= MAX_SUPPORTED_TENSOR_DESCRIPTORS, "Number of tensors in a group is larger than " "the MAX number of supported descriptors (64)."); - // Only full tiles supported - NVTE_CHECK(last_logical_dim % CHUNK_DIM_X == 0, - "Last dimension of a grouped tensor should be divisible by 128."); blocks_Y = 1; blocks_X = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); } @@ -858,7 +853,7 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations // Logical shape of a tensor with varying all dims is [1, M*K] if (shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS) { NVTE_CHECK(first_logical_dim % 128 == 0, - "First dimension of a grouped tensor should be divisible by 128."); + "First logical dimension of a grouped tensor must be divisible by 128."); } const int64_t *const offsets_ptr = reinterpret_cast(output->tensor_offsets.dptr); @@ -879,18 +874,20 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations NVTE_CHECK(scales_colwise_ptr != nullptr, "Columnwise scaling tensor must be allocated"); } - const size_t dbias_rows = DIVUP(first_logical_dim, CHUNK_DIM_Y); - const size_t dbias_cols = last_logical_dim; if constexpr (IS_DBIAS) { NVTE_CHECK(is_single_tensor, "DBias is only supported for tensors with the const last dimension."); NVTE_CHECK(dbias->data.dtype == input->dtype(), "DBias must have the same type as input_tensor."); - NVTE_CHECK(dbias->data.shape == std::vector{last_logical_dim}, "Wrong shape of DBias."); - NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor."); + std::vector 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."); + const size_t dbias_workspace_rows = DIVUP(first_logical_dim, CHUNK_DIM_Y); + const size_t dbias_workspace_cols = last_logical_dim; if (workspace->data.dptr == nullptr) { - workspace->data.shape = {dbias_rows, dbias_cols}; + workspace->data.shape = {dbias_workspace_rows, dbias_workspace_cols}; workspace->data.dtype = DType::kFloat32; return; } @@ -1007,7 +1004,9 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr); if constexpr (IS_DBIAS) { - common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); + common::grouped_reduce_dbias( + shape_rep, num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, + first_dims_ptr, last_dims_ptr, dbias, workspace_ptr, CHUNK_DIM_Y, stream); } NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 06f1c65ce2..854f52c203 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -56,6 +56,7 @@ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Computes the GeLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor for activation. * \param[in,out] output Output grouped tensor. @@ -76,6 +77,7 @@ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Computes the SiLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor for activation. * \param[in,out] output Output grouped tensor. @@ -96,6 +98,7 @@ void nvte_relu(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Computes the ReLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor for activation. * \param[in,out] output Output grouped tensor. @@ -116,6 +119,7 @@ void nvte_qgelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Computes the Quick GeLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor for activation. * \param[in,out] output Output grouped tensor. @@ -136,6 +140,7 @@ void nvte_srelu(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Computes the Squared ReLU activation of the grouped input. * If the scaling mode of the grouped output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor for activation. * \param[in,out] output Output grouped tensor. @@ -158,6 +163,7 @@ void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output /*! \brief Computes the GeLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] grad Incoming grouped gradient. * \param[in] input Input grouped tensor for activation. @@ -182,6 +188,7 @@ void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output /*! \brief Computes the SiLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] grad Incoming grouped gradient. * \param[in] input Input grouped tensor for activation. @@ -206,6 +213,7 @@ void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output /*! \brief Computes the ReLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] grad Incoming grouped gradient. * \param[in] input Input grouped tensor for activation. @@ -230,6 +238,7 @@ void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu /*! \brief Computes the Quick GeLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] grad Incoming grouped gradient. * \param[in] input Input grouped tensor for activation. @@ -254,6 +263,7 @@ void nvte_dsrelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu /*! \brief Computes the Squared ReLU activation gradient of the grouped input. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] grad Incoming grouped gradient. * \param[in] input Input grouped tensor for activation. diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 04712d3003..755052d6dd 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -92,6 +92,7 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea /*! \brief Casts input grouped tensor to MXFP8. * The type of quantized tensor in the output depends on the scaling mode of the output * tensor. See file level comments. + * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor to be cast. * \param[in,out] output Output grouped MXFP8 tensor. @@ -146,6 +147,7 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d /*! \brief Casts input grouped tensor to MXFP8. Additionally, reduces the input along columns. * 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. + * Grouped dbias is not yet supported for grouped tensors with a varying last dimension. * * This function produces 2 results: * - `output` is equal to `cast(dact(input))` @@ -161,7 +163,7 @@ void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor d * \param[in] stream CUDA stream used for the operation. */ void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream); /*! \brief Computes backward of GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the GeLU backward along columns. @@ -190,6 +192,7 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor act_inpu * Additionally, reduces the result of the GeLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * Grouped dbias is not yet supported for grouped tensors with a varying last dimension. * * This function produces 2 results: * - `output` is equal to `cast(dact(input))` @@ -207,7 +210,8 @@ void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor act_inpu */ void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, const NVTEGroupedTensor act_input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + cudaStream_t stream); /*! \brief Computes backward of SiLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the SiLU backward along columns. @@ -236,6 +240,7 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor act_inpu * Additionally, reduces the result of the SiLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * Grouped dbias is not yet supported for grouped tensors with a varying last dimension. * * This function produces 2 results: * - `output` is equal to `cast(dact(input))` @@ -253,7 +258,8 @@ void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor act_inpu */ void nvte_group_quantize_dbias_dsilu(const NVTEGroupedTensor input, const NVTEGroupedTensor act_input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + cudaStream_t stream); /*! \brief Computes backward of ReLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the ReLU backward along columns. @@ -282,6 +288,7 @@ void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor act_inpu * Additionally, reduces the result of the ReLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * Grouped dbias is not yet supported for grouped tensors with a varying last dimension. * * This function produces 2 results: * - `output` is equal to `cast(dact(input))` @@ -299,7 +306,8 @@ void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor act_inpu */ void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, const NVTEGroupedTensor act_input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + cudaStream_t stream); /*! \brief Computes backward of Quick GeLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the Quick GeLU backward along columns. @@ -328,6 +336,7 @@ void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor act_inp * Additionally, reduces the result of the Quick GeLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * Grouped dbias is not yet supported for grouped tensors with a varying last dimension. * * This function produces 2 results: * - `output` is equal to `cast(dact(input))` @@ -345,7 +354,8 @@ void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor act_inp */ void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, const NVTEGroupedTensor act_input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + cudaStream_t stream); /*! \brief Computes backward of Squared ReLU operation on the input, then casts to FP8/MXFP8. * Additionally, reduces the result of the Squared ReLU backward along columns. @@ -374,6 +384,7 @@ void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor act_inp * Additionally, reduces the result of the Squared ReLU backward along columns. * If the scaling mode of the output grouped tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. + * Grouped dbias is not yet supported for grouped tensors with a varying last dimension. * * This function produces 2 results: * - `output` is equal to `cast(dact(input))` @@ -391,7 +402,8 @@ void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor act_inp */ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, const NVTEGroupedTensor act_input, NVTEGroupedTensor output, - NVTETensor dbias, NVTETensor workspace, cudaStream_t stream); + NVTEGroupedTensor dbias, NVTETensor workspace, + cudaStream_t stream); /*! \brief Casts input tensor from reduced to higher precision. * If the scaling mode of the input tensor is set to NVTE_MXFP8_1D_SCALING, @@ -407,11 +419,11 @@ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t str /*! \brief Casts multiple input tensors to quantized output tensors. * - * \param[in] inputs List of input tensors to be cast. - * \param[in,out] outputs List of output quantized tensors. - * \param[in] quant_config (Optional) Quantization configurations. - * \param[in] num_tensors Number of input and output tensors. - * \param[in] stream CUDA stream used for the operation. + * \param[in] inputs List of input tensors to be cast. + * \param[in,out] outputs List of output quantized tensors. + * \param[in] quant_config (Optional) Quantization configurations. + * \param[in] num_tensors Number of input and output tensors. + * \param[in] stream CUDA stream used for the operation. */ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, const NVTEQuantizationConfig quant_config, const size_t num_tensors, @@ -420,11 +432,11 @@ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, /*! \brief Casts grouped input tensor to quantized output tensors. * * \param[in] input Input tensor to be cast. - * \param[in,out] outputs Output quantized tensors. - * \param[in] split_sections Split sections of the input tensor. - * \param[in] num_tensors Number of output tensors. + * \param[in,out] outputs Output quantized tensors. + * \param[in] split_sections Split sections of the input tensor. + * \param[in] num_tensors Number of output tensors. * \param[in] quant_config (Optional) Quantization configurations. - * \param[in] stream CUDA stream used for the operation. + * \param[in] stream CUDA stream used for the operation. */ void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *outputs, const size_t *split_sections, size_t num_tensors, From 3846bf77f0a2c59a4abf50f705488deecdfebc41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Tue, 10 Mar 2026 23:26:24 +0100 Subject: [PATCH 260/521] Fix deploy nightly docs issue (#2636) * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * ccache size limit Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * ccache Signed-off-by: Pawel Gadzinski * added blackwell Signed-off-by: Pawel Gadzinski * added blackwell Signed-off-by: Pawel Gadzinski * Revert build.yml changes unrelated to deploy nightly docs Restore .github/workflows/build.yml to upstream/main state. Only deploy_nightly_docs.yml changes are relevant to this PR. Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski --- .github/workflows/deploy_nightly_docs.yml | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/workflows/deploy_nightly_docs.yml b/.github/workflows/deploy_nightly_docs.yml index b4e015d2da..a8e5ee5ba2 100644 --- a/.github/workflows/deploy_nightly_docs.yml +++ b/.github/workflows/deploy_nightly_docs.yml @@ -7,6 +7,7 @@ name: Deploy nightly docs on: push: branches: [ "main" ] + workflow_dispatch: jobs: build: uses: ./.github/workflows/docs.yml @@ -21,9 +22,8 @@ jobs: name: "te_docs" path: "html" - name: Prepare for pages - uses: actions/upload-pages-artifact@v1.0.7 + uses: actions/upload-pages-artifact@v3 with: - name: github-pages path: "html" deploy: needs: prepare @@ -36,4 +36,5 @@ jobs: runs-on: ubuntu-latest steps: - name: Deploy - uses: actions/deploy-pages@v2.0.0 + id: deployment + uses: actions/deploy-pages@v4 From d32f9e422ba7961265a159507ba7336cf173881f Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Tue, 10 Mar 2026 19:37:48 -0700 Subject: [PATCH 261/521] [JAX] Fix get_seqlens_and_offsets() to accept vmapped seg ids and non vmapped seg offsets (#2692) * Fix batcher for when segment ids received are batched/vmapped whereas the TE constructed segment pos are not thereby causing mismatches in impl() Signed-off-by: Kshitij Lakhani * nit: Fix the shape check for assert Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix the batcher logic to check for q and kv seg ids separately Signed-off-by: Kshitij Lakhani * Remove batcher logic to expand segment pos. Keep the shape check asserts. Signed-off-by: Kshitij Lakhani * Add support for vmapped seg id and non vmapped seg pos when computing the seqlens and offsets for fused attn Signed-off-by: Kshitij Lakhani * Undo batcher check logic for seg pos and seg ids as it is already moved to get_seqlens_and_offsets() Signed-off-by: Kshitij Lakhani * nit: Remove unnecessary assert check Signed-off-by: Kshitij Lakhani * 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 * nit: Use partial instead of a single use function Signed-off-by: Kshitij Lakhani --------- Signed-off-by: Kshitij Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/jax/attention.py | 83 ++++++++++++++++--- .../jax/cpp_extensions/attention.py | 6 +- 2 files changed, 75 insertions(+), 14 deletions(-) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 21db296c34..765cf2872f 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -569,6 +569,8 @@ def _segment_ids_pos_to_seqlens_offsets( # 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() @@ -693,26 +695,83 @@ def get_seqlens_and_offsets( self, attn_mask_type, qkv_layout, window_size, max_segments_per_seq ): """ - Acquire the seqlens/offsets for cuDNN backend + Acquire the seqlens/offsets for cuDNN backend. """ q_segment_ids, kv_segment_ids = self.segment_ids q_segment_pos, kv_segment_pos = self.segment_pos - assert q_segment_ids.shape == q_segment_pos.shape - assert kv_segment_ids.shape == kv_segment_pos.shape # No segment_ids/segment_pos if q_segment_ids.size + kv_segment_ids.size == 0: return self.seqlens, self.seq_offsets - if qkv_layout.is_thd(): - q_seqlens, kv_seqlens, q_offsets, kv_offsets = _segment_ids_pos_to_seqlens_offsets( - q_segment_ids, - kv_segment_ids, - q_segment_pos, - kv_segment_pos, - attn_mask_type, - window_size, - max_segments_per_seq, + # Allow segment_pos to have fewer leading dims than segment_ids if vmapped segment_ids and non-vmapped segment_pos + # e.g. when using from_segment_ids_and_pos() for segment_pos generation from segment_ids it is acceptable to have + # something like : segment_ids (B, batch, seq), segment_pos (batch, seq)). + if q_segment_ids.ndim < q_segment_pos.ndim or kv_segment_ids.ndim < kv_segment_pos.ndim: + raise AssertionError( + "segment_ids must not have fewer dims than segment_pos; got" + f" q_segment_ids.ndim={q_segment_ids.ndim}," + f" q_segment_pos.ndim={q_segment_pos.ndim}," + f" kv_segment_ids.ndim={kv_segment_ids.ndim}," + f" kv_segment_pos.ndim={kv_segment_pos.ndim}" + ) + if not ( + q_segment_ids.shape[-q_segment_pos.ndim :] == q_segment_pos.shape + and kv_segment_ids.shape[-kv_segment_pos.ndim :] == kv_segment_pos.shape + ): + raise AssertionError( + "segment_pos trailing shape must match segment_ids; got" + f" q_segment_ids.shape={q_segment_ids.shape}," + f" q_segment_pos.shape={q_segment_pos.shape}," + f" kv_segment_ids.shape={kv_segment_ids.shape}," + f" kv_segment_pos.shape={kv_segment_pos.shape}" ) + # THD: compute seqlens/offsets. + if qkv_layout.is_thd(): + # If there are more leading dims on segment_ids, e.g. vmap + if q_segment_ids.ndim > q_segment_pos.ndim or kv_segment_ids.ndim > kv_segment_pos.ndim: + # Flatten leading batch dims so that segment_ids and segment_pos have the same number of leading dims, + # vmap seqlens/offsets computation with segment_pos broadcast, + # reshape back to the original leading batch dims. + n_extra_batch_dims_q = q_segment_ids.ndim - q_segment_pos.ndim + n_extra_batch_dims_kv = kv_segment_ids.ndim - kv_segment_pos.ndim + extra_batch_shape_q = q_segment_ids.shape[:n_extra_batch_dims_q] + extra_batch_shape_kv = kv_segment_ids.shape[:n_extra_batch_dims_kv] + extra_flat_batch_size_q = jnp.prod(extra_batch_shape_q) + extra_flat_batch_size_kv = jnp.prod(extra_batch_shape_kv) + # vmap below requires same batch size on axis 0 for q_flat and kv_flat; JAX will raise if they differ. + q_flat = q_segment_ids.reshape( + extra_flat_batch_size_q, *q_segment_ids.shape[n_extra_batch_dims_q:] + ) + kv_flat = kv_segment_ids.reshape( + extra_flat_batch_size_kv, *kv_segment_ids.shape[n_extra_batch_dims_kv:] + ) + + single_extra_batch = partial( + _segment_ids_pos_to_seqlens_offsets, + attn_mask_type=attn_mask_type, + window_size=window_size, + max_segments_per_seq=max_segments_per_seq, + ) + + q_sl, kv_sl, q_off, kv_off = jax.vmap( + single_extra_batch, in_axes=(0, 0, None, None) + )(q_flat, kv_flat, q_segment_pos, kv_segment_pos) + + q_seqlens = q_sl.reshape(*extra_batch_shape_q, *q_sl.shape[1:]) + kv_seqlens = kv_sl.reshape(*extra_batch_shape_kv, *kv_sl.shape[1:]) + q_offsets = q_off.reshape(*extra_batch_shape_q, *q_off.shape[1:]) + kv_offsets = kv_off.reshape(*extra_batch_shape_kv, *kv_off.shape[1:]) + else: + q_seqlens, kv_seqlens, q_offsets, kv_offsets = _segment_ids_pos_to_seqlens_offsets( + q_segment_ids, + kv_segment_ids, + q_segment_pos, + kv_segment_pos, + attn_mask_type, + window_size, + max_segments_per_seq, + ) + # BSHD: compute seqlens/offsets. else: q_seqlens, kv_seqlens = _segment_ids_to_seqlens( q_segment_ids, diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index e5d75e1501..f4d914062d 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -626,10 +626,12 @@ def convert_to_2d(offsets, batch, max_seqlen): @staticmethod def batcher(batched_args, batch_dims, *, config): + # batch_dims: each element is the batch axis (0, ...) or None. Only 0 or None allowed. check_valid_batch_dims(batch_dims) assert FusedAttnFwdPrimitive.outer_primitive is not None q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims - + # Pass through; segment_ids/segment_pos may have different batch dims (e.g. vmapped ids, + # replicated pos). get_seqlens_and_offsets() in attention.py handles conversion without expanding. out_bdims = q_bdim, q_bdim, seed_bdim return ( FusedAttnFwdPrimitive.outer_primitive.bind(*batched_args, config=config), @@ -1084,7 +1086,7 @@ def batcher(batched_args, batch_dims, *, config): check_valid_batch_dims(batch_dims) assert FusedAttnBwdPrimitive.outer_primitive is not None q_bdim, k_bdim, v_bdim, bias_bdim, softmax_offset_bdim, *_ = batch_dims - + # Pass through; segment_ids/segment_pos may have different batch dims. Conversion is in attention.py. out_bdims = q_bdim, k_bdim, v_bdim, bias_bdim, softmax_offset_bdim return ( FusedAttnBwdPrimitive.outer_primitive.bind(*batched_args, config=config), From 61d5865397a45deb9c6eebfcb7972142c1e3beaa Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Tue, 10 Mar 2026 21:45:22 -0700 Subject: [PATCH 262/521] [NVFP4][MOE] Add unfused quantization fallback when input shape is not aligned (#2747) * fallback Signed-off-by: Zhongbo Zhu * warn once Signed-off-by: Zhongbo Zhu --------- Signed-off-by: Zhongbo Zhu --- tests/pytorch/nvfp4/test_nvfp4_group_quantize.py | 2 ++ transformer_engine/pytorch/csrc/extensions/cast.cpp | 13 ++++++++++++- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 5f35e9ad10..d4bf1fd3a1 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -130,6 +130,8 @@ def check_group_quantization_nvfp4_versus_reference( [ # edge case, zero tokens for all (0, 512), + # edge case, not 128 multiple hidden dimension + (1024, 320), # full tile cases (256, 1024), (1024, 256), diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index f8f793f036..89cd90f347 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -1355,9 +1356,19 @@ std::vector split_quantize(const at::Tensor &tensor, for (auto &quantizer : quantizer_cpp_list) { nvfp4_quantizers.push_back(static_cast(quantizer.get())); } - bool contiguous_data_and_scale; + 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) { + static std::once_flag once_unfused_nvfp4_fallback_warning; + std::call_once(once_unfused_nvfp4_fallback_warning, []() { + NVTE_WARN( + "Unfused NVFP4 quantization fallback is triggered because the input tensor inner " + "dimension is not a multiple of 128, disabling NVFP4 grouped kernel fusion. " + "NVFP4 might bring performance regressions for this input tensor shape."); + }); + quantization_method = QuantizationMethod::UNFUSED; + } if (!contiguous_data_and_scale) { // Avoid fused quantize kernel if data is not contiguous quantization_method = QuantizationMethod::UNFUSED; From 7545d8c143fa274b87c690894e1b909cdd2151b1 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, 11 Mar 2026 05:47:50 +0100 Subject: [PATCH 263/521] [PyTorch debug] Fix issue with tp_group=None (#2733) * code drop Signed-off-by: Pawel Gadzinski * [Debug] Pass tp_size to DebugQuantizer and use it in get_reduction_params Use tp_size to determine whether tensor parallelism is active instead of checking tp_group is None (which is ambiguous since None means world group in torch.distributed). Also add tp_size to the backward-compat kwargs filtering in call_feature so custom features without tp_size in their inspect_tensor signature continue to work. Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/debug/features/api.py | 11 ++++++++++- .../debug/features/log_fp8_tensor_stats.py | 3 ++- .../debug/features/log_nvfp4_tensor_stats.py | 3 ++- .../debug/features/log_tensor_stats.py | 3 ++- transformer_engine/debug/features/utils/__init__.py | 12 +++++++++--- .../debug/pytorch/debug_quantization.py | 3 +++ transformer_engine/pytorch/module/grouped_linear.py | 2 +- .../pytorch/module/layernorm_linear.py | 2 +- transformer_engine/pytorch/module/layernorm_mlp.py | 1 + transformer_engine/pytorch/module/linear.py | 2 +- 10 files changed, 32 insertions(+), 10 deletions(-) diff --git a/transformer_engine/debug/features/api.py b/transformer_engine/debug/features/api.py index 774fae3594..a1cf80dd25 100644 --- a/transformer_engine/debug/features/api.py +++ b/transformer_engine/debug/features/api.py @@ -479,7 +479,12 @@ def call_feature(self, call, feat_config, layer_name, **kwargs): """ if call.__name__ == "inspect_tensor": kwargs_copy = kwargs.copy() - for k in ["quantizer", "columnwise_quantized_tensor", "rowwise_quantized_tensor"]: + for k in [ + "quantizer", + "columnwise_quantized_tensor", + "rowwise_quantized_tensor", + "tp_size", + ]: if k not in call.__code__.co_varnames: kwargs_copy.pop(k) else: @@ -490,6 +495,10 @@ def call_feature(self, call, feat_config, layer_name, **kwargs): "inspect_tensor_postquantize is deprecated, use inspect_tensor instead.", DeprecationWarning, ) + kwargs_copy = kwargs.copy() + for k in ["tp_size"]: + if k not in call.__code__.co_varnames: + kwargs_copy.pop(k, None) return call(feat_config, layer_name, **kwargs_copy) diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index fd18d590ec..cf11964e25 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -311,6 +311,7 @@ def inspect_tensor( rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, quantizer: Optional[Quantizer] = None, + tp_size: int = 1, ): """ API call used to collect the data about the tensor after process_tensor()/quantization. @@ -357,7 +358,7 @@ def inspect_tensor( ) skip_reduction, reduction_group, reduce_within_microbatch = get_reduction_params( - tensor_name, tp_group + tensor_name, tp_group, tp_size ) STATS_BUFFERS.try_add_buffer( diff --git a/transformer_engine/debug/features/log_nvfp4_tensor_stats.py b/transformer_engine/debug/features/log_nvfp4_tensor_stats.py index 18ac8619f3..8a76f4edcf 100644 --- a/transformer_engine/debug/features/log_nvfp4_tensor_stats.py +++ b/transformer_engine/debug/features/log_nvfp4_tensor_stats.py @@ -148,6 +148,7 @@ def inspect_tensor( rowwise_quantized_tensor: Optional[QuantizedTensor] = None, columnwise_quantized_tensor: Optional[QuantizedTensor] = None, quantizer: Optional[Quantizer] = None, + tp_size: int = 1, ): """ API call used to collect the data about the tensor after process_tensor()/quantization. @@ -199,7 +200,7 @@ def inspect_tensor( ) skip_reduction, reduction_group, reduce_within_microbatch = get_reduction_params( - tensor_name, tp_group + tensor_name, tp_group, tp_size ) # Add nvfp4_ prefix to all stats for internal use diff --git a/transformer_engine/debug/features/log_tensor_stats.py b/transformer_engine/debug/features/log_tensor_stats.py index 76e61fab24..5e6ce137bd 100644 --- a/transformer_engine/debug/features/log_tensor_stats.py +++ b/transformer_engine/debug/features/log_tensor_stats.py @@ -184,6 +184,7 @@ def inspect_tensor( rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, quantizer: Optional[Quantizer] = None, + tp_size: int = 1, ): # pylint: disable=unused-argument """API call used to collect the data about the tensor before process_tensor()/quantization.""" @@ -214,7 +215,7 @@ def inspect_tensor( ) skip_reduction, reduction_group, reduce_within_microbatch = get_reduction_params( - tensor_name, tp_group + tensor_name, tp_group, tp_size ) for stat in config["stats"]: diff --git a/transformer_engine/debug/features/utils/__init__.py b/transformer_engine/debug/features/utils/__init__.py index d691c1828c..813fb2addc 100644 --- a/transformer_engine/debug/features/utils/__init__.py +++ b/transformer_engine/debug/features/utils/__init__.py @@ -12,7 +12,7 @@ from transformer_engine.debug.pytorch.debug_state import TEDebugState -def get_reduction_params(tensor_name: str, tp_group: torch.distributed.ProcessGroup): +def get_reduction_params(tensor_name: str, tp_group: torch.distributed.ProcessGroup, tp_size: int): """ Returns the statistics reduction parameters for the tensor. """ @@ -20,8 +20,14 @@ def get_reduction_params(tensor_name: str, tp_group: torch.distributed.ProcessGr reduction_group = debug_api.get_tensor_reduction_group() reduce_within_microbatch = tensor_name != "weight" if tensor_name == "weight": - if TEDebugState.weight_tensor_tp_group_reduce: - reduction_group = tp_group + if TEDebugState.weight_tensor_tp_group_reduce and tp_size > 1: + # Do not overwrite with `None`: in torch.distributed collectives + # group=None means the default/world process group. + if tp_group is not None: + reduction_group = tp_group + else: + # "Reduce in TP group" requested, but TP group is missing. + skip_reduction = True else: skip_reduction = True return skip_reduction, reduction_group, reduce_within_microbatch diff --git a/transformer_engine/debug/pytorch/debug_quantization.py b/transformer_engine/debug/pytorch/debug_quantization.py index 57a5967079..ed5fdd4660 100644 --- a/transformer_engine/debug/pytorch/debug_quantization.py +++ b/transformer_engine/debug/pytorch/debug_quantization.py @@ -53,6 +53,7 @@ def __init__( tensor_name: str, parent_quantizer: Optional[Quantizer], tp_group: torch.distributed.ProcessGroup, + tp_size: int, ): super().__init__(rowwise=True, columnwise=True) @@ -60,6 +61,7 @@ def __init__( self.tensor_name = tensor_name self.parent_quantizer = parent_quantizer self.tp_group = tp_group # used in inspect_tensor calls + self.tp_size = tp_size self.iteration = TEDebugState.get_iteration() # Configure parent quantizer @@ -263,6 +265,7 @@ def _call_inspect_tensor_api( "tensor_name": self.tensor_name, "iteration": TEDebugState.get_iteration(), "tp_group": self.tp_group, + "tp_size": self.tp_size, "columnwise_quantized_tensor": columnwise_gemm_tensor, "rowwise_quantized_tensor": rowwise_gemm_tensor, "quantizer": self.parent_quantizer, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f3e7b57cf1..02607e45e5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1082,7 +1082,7 @@ def _get_debug_quantizers(self): names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"] return tuple( [ - DebugQuantizer(self.name + f".gemm_{q_id}", name, q, self.tp_group) + DebugQuantizer(self.name + f".gemm_{q_id}", name, q, self.tp_group, self.tp_size) for q_id, q in enumerate(qs) ] for name, qs in zip(names, original_quantizers) diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 9a716207c1..a90105477c 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1644,7 +1644,7 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"] return tuple( - DebugQuantizer(self.name, name, q, self.tp_group) + DebugQuantizer(self.name, name, q, self.tp_group, self.tp_size) for name, q in zip(names, original_quantizers) ) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ba92fb32ed..037fb6c858 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -2371,6 +2371,7 @@ def make_debug(prefix, offset): label, None if label in ("dgrad", "wgrad") else base_quantizers[i + offset], self.tp_group, + self.tp_size, ) for i, label in enumerate(labels) ] diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index de2a421b53..1e3eadc405 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1511,7 +1511,7 @@ def _get_debug_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"] return tuple( - DebugQuantizer(self.name, name, q, self.tp_group) + DebugQuantizer(self.name, name, q, self.tp_group, self.tp_size) for name, q in zip(names, original_quantizers) ) From 107f5585f7e61c0f5e3d9dd8828a227e0ba95b4e 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, 11 Mar 2026 14:59:16 +0100 Subject: [PATCH 264/521] Documentation for cpu offloading (#2520) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add CPU offloading documentation Documents the get_cpu_offload_context() API with examples for basic usage, manual synchronization, and CUDA graphs integration. Adds new other_optimizations section to the documentation structure. Signed-off-by: Pawel Gadzinski * Fix copyright year to 2026 in cpu_offloading.rst Signed-off-by: Pawel Gadzinski * Add missing copyright headers to CPU offloading example files Signed-off-by: Pawel Gadzinski * Improve cpu_offloading docs: legacy params note, manual sync clarifications - Add comment about legacy parameters in function signature - Clarify that num_layers is ignored (not forbidden) when manual_synchronization=True - Document num_layers constraint: must be <= model_layers-2 for overlap - Add note that offload_stream.synchronize() must precede release_activation_forward_gpu_memory() Signed-off-by: Pawel Gadzinski * Remove incorrect note about offload_stream.synchronize() release_activation_forward_gpu_memory() internally waits for offload completion via CUDA events - explicit synchronize() is not required. Signed-off-by: Pawel Gadzinski * Update docs/features/other_optimizations/cpu_offloading/cpu_offloading.rst Co-authored-by: Przemyslaw Tredak Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * update Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * docs: review fixes for cpu offloading documentation - Export mark_not_offload and ManualOffloadSynchronizer from transformer_engine.pytorch and add them to API reference - Fix condition 3 description (xi is needed as input to backward, not computed after it completes) - Fix offload_weights docstring default value (False, not True) - Fix legacy params comment to link to API reference - Fix RST spacing around inline code in Fig. 3/4 captions - Add note explaining retain_pinned_cpu_buffers history and pytorch#167507 fix landing in PyTorch 2.11 - Fix typo: seqeuences -> sequences - Add trailing newline to other_optimizations/index.rst Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak Co-authored-by: Claude Sonnet 4.6 --- docs/api/pytorch.rst | 4 + .../cpu_offloading/cpu_offloading.rst | 290 ++++++++++++++++++ .../cpu_offloading/img/layer_sequence.svg | 66 ++++ .../cpu_offloading/img/pcie_vs_nvlink.svg | 132 ++++++++ .../cpu_offloading/img/scheduling.svg | 110 +++++++ .../cpu_offloading/img/scheduling_stall.svg | 143 +++++++++ .../pytorch_basic_offload_example.py | 36 +++ .../pytorch_cuda_graphs_example.py | 46 +++ .../pytorch_manual_offload_example.py | 40 +++ docs/features/other_optimizations/index.rst | 12 + docs/index.rst | 1 + transformer_engine/pytorch/__init__.py | 6 +- transformer_engine/pytorch/cpu_offload.py | 8 +- 13 files changed, 889 insertions(+), 5 deletions(-) create mode 100644 docs/features/other_optimizations/cpu_offloading/cpu_offloading.rst create mode 100644 docs/features/other_optimizations/cpu_offloading/img/layer_sequence.svg create mode 100644 docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg create mode 100644 docs/features/other_optimizations/cpu_offloading/img/scheduling.svg create mode 100644 docs/features/other_optimizations/cpu_offloading/img/scheduling_stall.svg create mode 100644 docs/features/other_optimizations/cpu_offloading/pytorch_basic_offload_example.py create mode 100644 docs/features/other_optimizations/cpu_offloading/pytorch_cuda_graphs_example.py create mode 100644 docs/features/other_optimizations/cpu_offloading/pytorch_manual_offload_example.py create mode 100644 docs/features/other_optimizations/index.rst diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 90f68653cc..1fe4f19990 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -49,6 +49,10 @@ PyTorch .. autoapifunction:: transformer_engine.pytorch.get_cpu_offload_context +.. autoapifunction:: transformer_engine.pytorch.mark_not_offload + +.. autoapiclass:: transformer_engine.pytorch.ManualOffloadSynchronizer + .. autoapifunction:: transformer_engine.pytorch.parallel_cross_entropy Recipe availability diff --git a/docs/features/other_optimizations/cpu_offloading/cpu_offloading.rst b/docs/features/other_optimizations/cpu_offloading/cpu_offloading.rst new file mode 100644 index 0000000000..47ea35a834 --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/cpu_offloading.rst @@ -0,0 +1,290 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +CPU Offloading +=================================== + +.. note:: + + CPU Offloading in Transformer Engine is currently available only for **PyTorch**. + It supports all PyTorch modules, not just TE layers. + +CPU offloading moves activation tensors from GPU to CPU memory during the +forward pass and reloads them during backward. Transfers are **asynchronous**, +enabling significant GPU memory savings with minimal overhead. + +Unlike activation checkpointing, offloading avoids recomputation — activations +are stored on CPU instead of being recalculated, making it faster when +CPU-GPU bandwidth is sufficient. + + +Hardware Support +---------------- + +CPU offloading benefits greatly from fast CPU-GPU interconnects. +The faster the link, the more effectively transfer time can be hidden +behind computation. + +.. raw:: html + :file: img/pcie_vs_nvlink.svg + +*Figure 1. Traditional PCIe system vs GB200 Superchip with NVLink-C2C.* + +Traditional **PCIe Gen5 x16** systems offer **128 GB/s** bidirectional bandwidth +between CPU and GPU, which limits offloading benefits. + +With **NVLink-C2C** (GB200), bandwidth jumps to **900 GB/s** bidirectional per link, +making offloading increasingly attractive on modern NVIDIA superchips. +The GB200 pairs a Grace CPU with 480 GB LPDDR5X memory and two Blackwell GPUs, +each with 192 GB HBM3e (384 GB total), providing ample CPU memory for offloading +activations. + +Offloading/reloading consumes HBM bandwidth, which may compete with +other GPU operations — even when transfers are asynchronous. +This is unlikely to affect compute-bound operations like GEMMs, but the impact on +memory-bound operations like quantization may be noticeable. + + +CPU Offloading in Transformer Engine +------------------------------------ + +Transformer Engine supports CPU offloading of activations for **sequential models**. +A model is considered sequential if it satisfies the following conditions: + +1. The model is a sequence of layers: ``x₁ = Layer₁(x₀)``, ``x₂ = Layer₂(x₁)``, ..., ``xₙ = Layerₙ(xₙ₋₁)``. + **The layers may be any PyTorch modules**, not just TE layers. +2. Each intermediate tensor ``xᵢ`` is used only as input to the next layer (not elsewhere in the model). +3. ``xᵢ`` is only needed as input to ``Layerᵢ₊₁``'s backward pass and can be freed once that pass completes. + +Most LLM architectures (stacked Transformer blocks) satisfy these conditions. + +.. raw:: html + :file: img/layer_sequence.svg + +*Figure 2. Sequential model: xᵢ₊₁ = Layerᵢ₊₁(xᵢ). Each layer consumes only the output of the previous one.* + +The example below shows how to offload activations for a sequence of ``torch.nn.Linear`` layers using the default scheduling algorithm: + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: pytorch_basic_offload_example.py + :language: python + :start-after: # START_BASIC_EXAMPLE + :end-before: # END_BASIC_EXAMPLE + + + +Let's take a look at the API in detail: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + def get_cpu_offload_context( + enabled: bool = False, + num_layers: Optional[int] = 1, + model_layers: int = 1, + manual_synchronization: bool = False, + offload_stream: Optional[torch.cuda.Stream] = None, + # ... (legacy parameters omitted, see :func:`get_cpu_offload_context`) + ) -> Union[Tuple[ContextManager, Callable], Tuple[ContextManager, Callable, ManualOffloadSynchronizer]]: + ... + +The ``model_layers`` parameter must always be set to the total number of layers in the model. +There are two modes of operation: + +1. **Default scheduling** — set ``num_layers`` to the number of layers to offload. + The algorithm automatically schedules offload/reload operations to overlap with computation. + +2. **Manual synchronization** — set ``manual_synchronization=True`` (``num_layers`` is ignored in this mode). + This mode provides explicit control over when to start offload/reload using the returned ``ManualOffloadSynchronizer``. + +The :func:`transformer_engine.pytorch.get_cpu_offload_context` function returns: + +- **context manager** — wraps each layer's forward pass to intercept tensors saved for backward. +- **sync function** — registers a backward hook on the output tensor to trigger activation reload. +- **ManualOffloadSynchronizer** *(only in manual mode)* — provides explicit control over offload/reload. + +The usage pattern for default scheduling is: + +.. tabs:: + + .. tab:: PyTorch + + .. code-block:: python + + cpu_offload_context, sync_function = get_cpu_offload_context(...) + + for layer in layers: + with cpu_offload_context: + x = layer(x) + x = sync_function(x) + + +Default Offloading Scheduling +----------------------------- + +Default scheduling is enabled when ``manual_synchronization=False`` (the default). +The ``num_layers`` parameter must be specified to set the number of layers to offload. +The algorithm then automatically determines when to offload and reload activations +to maximize overlap with computation. + +For ``num_layers`` layers offloaded of ``model_layers`` layers: + +- First ``num_layers`` layers are offloaded to CPU. +- Offloading starts as soon as tensors are saved for backward — it does not wait + for the layer's forward pass to complete. +- At most ``(model_layers - num_layers)`` sets of activations are on GPU at any time; + both compute and reload may be stalled to enforce this limit. +- Reloading must complete by the time the tensor is needed for the layer's backward pass. +- ``num_layers`` must be at most ``model_layers - 1`` (setting it to ``model_layers`` + raises an assertion error). However, ``model_layers - 1`` leaves only 1 activation set + on GPU at a time — compute and transfers cannot overlap, and a warning is raised. + For full overlap, use ``model_layers - 2`` or less. + +Specifying a low enough ``num_layers`` enables full overlap of computation +and offload/reload. The following two scenarios illustrate this — one with full overlap, and one with stalls. + +.. raw:: html + :file: img/scheduling.svg + +*Figure 3. With* ``num_layers=2``\ *and* ``model_layers=5``\ *, at most 3 sets of activations are on GPU. Layer 1 offloading starts during its forward pass (when the first tensor is saved for backward). Offloading fully overlaps with forward, reloading fully overlaps with backward.* + +When ``num_layers`` is too high, the GPU memory limit forces stalls: + +.. raw:: html + :file: img/scheduling_stall.svg + +*Figure 4. With* ``num_layers=3``\ *and* ``model_layers=5``\ *, at most 2 sets of activations can be on GPU (5-3=2), which causes stalls. In forward, Layer 4 cannot start until Layer 2 is offloaded, otherwise there would be 3 sets of activations on GPU (Layers 2, 3, 4). In backward, Layer 3 cannot start immediately — its activations are still on CPU and must be reloaded first. Some tensors may finish reloading earlier, allowing parts of the layer (e.g., a sublayer) to run while the rest waits. The same applies to Layers 2 and 1.* + + +Manual Synchronization +---------------------- + +For custom scheduling, set ``manual_synchronization=True``. +Optionally, pass a custom ``offload_stream`` for fine-grained synchronization. +This mode returns a ``ManualOffloadSynchronizer`` with explicit control over transfers. + +This mode is useful when training does not follow the standard "all forwards then all backwards" +pattern — for example, in pipeline parallelism. Providing a custom ``offload_stream`` enables +additional synchronization logic (e.g., waiting, recording events) tailored to the specific workload. + +The ``ManualOffloadSynchronizer`` object provides the following methods: + +- ``start_offload_layer(layer_id)`` — queue async GPU→CPU copies on the offload stream. + Before each copy, the offload stream waits for an event recorded when that tensor + was saved for backward. +- ``release_activation_forward_gpu_memory(layer_id)`` — make the current stream wait for + this layer's offload to complete, then release GPU memory. +- ``start_reload_layer(layer_id)`` — queue async CPU→GPU copies on the offload stream. + When tensors are accessed in backward, compute stream waits for each tensor's reload + to complete. + +To skip offloading for a specific layer, simply do not call any of these methods for that layer. + +.. tabs:: + + .. tab:: PyTorch + + The example demonstrates: + + 1. **Forward pass**: After each layer, call ``start_offload_layer(i)`` to begin + async copy of layer ``i``'s activations to CPU. + 2. **Release GPU memory**: Call ``release_activation_forward_gpu_memory(i)`` to free + the GPU tensors. Each call waits internally for that layer's offload to complete. + 3. **Before backward**: Call ``start_reload_layer(i)`` to begin async reload. + The compute stream will automatically wait for each tensor to be reloaded + before it's accessed in backward. + + .. literalinclude:: pytorch_manual_offload_example.py + :language: python + :start-after: # START_MANUAL_EXAMPLE + :end-before: # END_MANUAL_EXAMPLE + + +CPU Offloading and CUDA Graphs +------------------------------ + +CPU offloading works with CUDA graphs — async copies and stream synchronization +are GPU operations that can be captured and replayed, even when accessing +pinned CPU memory (via PCIe DMA, without CPU involvement). + +.. note:: + + We recommend capturing the entire forward and backward pass in a single graph. + Async copy operations (offload/reload) must complete within the same graph where + they started. If the graph ends before copies finish, PyTorch will block waiting + for them, defeating the purpose of graph capture. + +.. tabs:: + + .. tab:: PyTorch + + .. literalinclude:: pytorch_cuda_graphs_example.py + :language: python + :start-after: # START_CUDA_GRAPHS_EXAMPLE + :end-before: # END_CUDA_GRAPHS_EXAMPLE + +.. note:: + + In PyTorch versions prior to 2.11, CPU offloading with CUDA graphs required passing + ``retain_pinned_cpu_buffers=True`` to :func:`get_cpu_offload_context`. The root cause + was that ``torch.empty`` with pinned CPU memory was not supported inside CUDA graph + capture — buffers had to be pre-allocated and reused across iterations to avoid + invalidating DMA addresses captured in the graph. This was fixed in + `pytorch#167507 `_ (merged December 2025, + shipping in PyTorch 2.11). On PyTorch 2.11+, ``retain_pinned_cpu_buffers`` is no longer needed. + +Caveats +------- + +.. warning:: + + **Heuristic activation detection**: + + CPU Offloading is implemented using + `PyTorch saved tensors hooks `_. + PyTorch saves various tensors for backward — not just activations, but also weights and other data. + + Activation detection is heuristic. A CUDA tensor is offloaded if it: + + - has at least 256×1024 elements (~1 MB for float32), + - is not a ``torch.nn.Parameter``, + - is not marked with ``mark_not_offload()``. + + Additionally, non-contiguous tensors are skipped to avoid memory layout changes (see below). + For TE layers, tensors that should not be offloaded are manually excluded. + For non-TE layers, no such exclusions exist, so some tensors may remain pinned in GPU memory + even after being copied to CPU (e.g., if the layer stores references in ``ctx``), + resulting in wasted bandwidth with no memory savings. + + To exclude specific tensors from offloading, use :func:`mark_not_offload`: + + .. code-block:: python + + from transformer_engine.pytorch import mark_not_offload + mark_not_offload(tensor) + +.. warning:: + + **Memory layout changes**: + + Offloading/reloading can change tensor memory layout and relations: + + 1. Views of the same storage may be restored as separate allocations. + 2. Adjacent tensors may not be adjacent after reload. + + CUDA kernels that rely on specific memory layout may produce unexpected results. + To mitigate (1), non-trivial views are excluded from offloading by default. + TE attention kernels are an exception — they use internal handling that is tested and supported. + Issue (2) is not mitigated — custom kernels that assume adjacent tensors share + contiguous memory may still fail. + + If you encounter layout-related issues, use :func:`mark_not_offload` to exclude + problematic tensors from offloading. diff --git a/docs/features/other_optimizations/cpu_offloading/img/layer_sequence.svg b/docs/features/other_optimizations/cpu_offloading/img/layer_sequence.svg new file mode 100644 index 0000000000..cdb8814a97 --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/img/layer_sequence.svg @@ -0,0 +1,66 @@ + + + + + + + + + + + x₀ + + + + + + Layer 1 + + + + + + x₁ + + + + + + Layer 2 + + + + + + x₂ + + + + + + Layer 3 + + + + + ··· + + + + + + Layer N + + + + + + xₙ + + diff --git a/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg b/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg new file mode 100644 index 0000000000..0b8ec3912a --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/img/pcie_vs_nvlink.svg @@ -0,0 +1,132 @@ + + + + + + + + + + Traditional PCIe System + + + + + + + CPU + + + + RAM + + + + + + + + + + GPU + + + + HBM + + + + + + + + PCIe + + 128 GB/s + + + + GB200 Superchip + NVIDIA Grace Blackwell + + + + + + + + + + Blackwell + GPU 1 + + + + HBM + + + + + + + NVLink + C2C + + + + + + + Grace CPU + + + + RAM + + + + + + + NVLink + C2C + + + + + + + Blackwell + GPU 2 + + + + HBM + + + + 900 GB/s per NVLink-C2C link + + diff --git a/docs/features/other_optimizations/cpu_offloading/img/scheduling.svg b/docs/features/other_optimizations/cpu_offloading/img/scheduling.svg new file mode 100644 index 0000000000..19255c3474 --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/img/scheduling.svg @@ -0,0 +1,110 @@ + + + + + + + Model (model_layers = 5) + + + + Layer 1 + + + Layer 2 + + + Layer 3 + + + Layer 4 + + + Layer 5 + + + + num_layers = 2 (offloaded) + + + + + + Forward Pass + + + compute stream + offload stream + + + + Layer 1 fwd + + + Layer 2 fwd + + + Layer 3 fwd + + + Layer 4 fwd + + + Layer 5 fwd + + + + Layer 1 offload + + + Layer 2 offload + + + + + + Backward Pass + + + compute stream + reload stream + + + + Layer 5 bwd + + + Layer 4 bwd + + + Layer 3 bwd + + + Layer 2 bwd + + + Layer 1 bwd + + + + Layer 2 reload + + + Layer 1 reload + + diff --git a/docs/features/other_optimizations/cpu_offloading/img/scheduling_stall.svg b/docs/features/other_optimizations/cpu_offloading/img/scheduling_stall.svg new file mode 100644 index 0000000000..cd2d1a660c --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/img/scheduling_stall.svg @@ -0,0 +1,143 @@ + + + + + + + Model (model_layers = 5) + + + + Layer 1 + + + Layer 2 + + + Layer 3 + + + Layer 4 + + + Layer 5 + + + + num_layers = 3 (offloaded) + + + + + + Forward Pass + + + compute stream + offload stream + + + Layer 1 fwd + + + Layer 2 fwd + + + Layer 3 fwd + + + + wait + + + Layer 4 fwd + + + + wait + + + Layer 5 fwd + + + + Layer 1 offload + + + Layer 2 offload + + + Layer 3 offload + + + + + + Backward Pass + + + compute stream + reload stream + + + Layer 5 bwd + + + Layer 4 bwd + + + + Layer 3 bwd + + + wait + + + + + Layer 2 bwd + + + wait + + + + + Layer 1 bwd + + wait + + + wait + + + + + Layer 3 reload + + + Layer 2 reload + + + Layer 1 reload + + diff --git a/docs/features/other_optimizations/cpu_offloading/pytorch_basic_offload_example.py b/docs/features/other_optimizations/cpu_offloading/pytorch_basic_offload_example.py new file mode 100644 index 0000000000..b453e824a5 --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/pytorch_basic_offload_example.py @@ -0,0 +1,36 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_BASIC_EXAMPLE +import torch +from transformer_engine.pytorch import get_cpu_offload_context + +# Setup +num_layers = 12 +offloaded_layers = 3 +layers = [torch.nn.Linear(1024, 1024).cuda() for _ in range(num_layers)] +x = torch.randn(16, 1024, 1024, device="cuda") + +# Get offloading context and sync function +cpu_offload_context, sync_function = get_cpu_offload_context( + enabled=True, + model_layers=num_layers, + num_layers=offloaded_layers, +) + +# Forward pass +for i in range(num_layers): + # Context manager captures tensors saved for backward. + # These tensors will be offloaded to CPU asynchronously. + with cpu_offload_context: + x = layers[i](x) + + # sync_function must be called after each layer's forward pass. + # This cannot be done inside the context manager because + # it needs the output tensor after the layer has finished. + x = sync_function(x) + +loss = x.sum() +loss.backward() +# END_BASIC_EXAMPLE diff --git a/docs/features/other_optimizations/cpu_offloading/pytorch_cuda_graphs_example.py b/docs/features/other_optimizations/cpu_offloading/pytorch_cuda_graphs_example.py new file mode 100644 index 0000000000..a42bd89089 --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/pytorch_cuda_graphs_example.py @@ -0,0 +1,46 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_CUDA_GRAPHS_EXAMPLE +import torch +from transformer_engine.pytorch import get_cpu_offload_context, make_graphed_callables + +# Setup +num_layers = 12 +offloaded_layers = 3 +layers = [torch.nn.Linear(1024, 1024).cuda() for _ in range(num_layers)] + +# Enable offloading for CUDA graphs +cpu_offload_context, sync_function = get_cpu_offload_context( + enabled=True, + model_layers=num_layers, + num_layers=offloaded_layers, +) + + +# Wrap layers in a module that uses offloading +class OffloadedModel(torch.nn.Module): + def __init__(self, layers): + super().__init__() + self.layers = torch.nn.ModuleList(layers) + + def forward(self, x): + for layer in self.layers: + with cpu_offload_context: + x = layer(x) + x = sync_function(x) + return x + + +model = OffloadedModel(layers) +sample_input = (torch.randn(16, 1024, 1024, device="cuda"),) + +# Create graphed callable (warmup is handled internally) +graphed_model = make_graphed_callables(model, sample_input) + +# Use the graphed model +x = torch.randn(16, 1024, 1024, device="cuda") +out = graphed_model(x) +out.sum().backward() +# END_CUDA_GRAPHS_EXAMPLE diff --git a/docs/features/other_optimizations/cpu_offloading/pytorch_manual_offload_example.py b/docs/features/other_optimizations/cpu_offloading/pytorch_manual_offload_example.py new file mode 100644 index 0000000000..92e0768c80 --- /dev/null +++ b/docs/features/other_optimizations/cpu_offloading/pytorch_manual_offload_example.py @@ -0,0 +1,40 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# START_MANUAL_EXAMPLE +import torch +from transformer_engine.pytorch import get_cpu_offload_context + +# Setup +num_layers = 12 +layers = [torch.nn.Linear(1024, 1024).cuda() for _ in range(num_layers)] +x = torch.randn(16, 1024, 1024, device="cuda") + +offload_stream = torch.cuda.Stream() +cpu_offload_context, sync_function, manual_controller = get_cpu_offload_context( + enabled=True, + model_layers=num_layers, + manual_synchronization=True, + offload_stream=offload_stream, +) + +# Forward pass - manually trigger offload after each layer +for i in range(num_layers): + with cpu_offload_context: + x = layers[i](x) + x = sync_function(x) + manual_controller.start_offload_layer(i) + +# Release GPU memory (each call waits for that layer's offload to complete) +for i in range(num_layers): + manual_controller.release_activation_forward_gpu_memory(i) + +# Start reloading before backward +for i in range(num_layers - 1, -1, -1): + manual_controller.start_reload_layer(i) + +# Backward pass +loss = x.sum() +loss.backward() +# END_MANUAL_EXAMPLE diff --git a/docs/features/other_optimizations/index.rst b/docs/features/other_optimizations/index.rst new file mode 100644 index 0000000000..05e89c4b05 --- /dev/null +++ b/docs/features/other_optimizations/index.rst @@ -0,0 +1,12 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +Other optimizations +=================================== + +.. toctree:: + + cpu_offloading/cpu_offloading.rst + diff --git a/docs/index.rst b/docs/index.rst index 194e76df24..7389553679 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -45,6 +45,7 @@ Transformer Engine documentation :caption: Features features/low_precision_training/index.rst + features/other_optimizations/index.rst .. toctree:: diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 5e1eb6954b..cd18ca75ad 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -54,7 +54,11 @@ from transformer_engine.pytorch.graph import make_graphed_callables from transformer_engine.pytorch.distributed import checkpoint from transformer_engine.pytorch.distributed import CudaRNGStatesTracker -from transformer_engine.pytorch.cpu_offload import get_cpu_offload_context +from transformer_engine.pytorch.cpu_offload import ( + get_cpu_offload_context, + mark_not_offload, + ManualOffloadSynchronizer, +) from transformer_engine.pytorch import ops from transformer_engine.pytorch import optimizers from transformer_engine.pytorch.export import onnx_export diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index 05219b7b18..e7390de365 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -685,7 +685,7 @@ def get_cpu_offload_context( offload_stream: Optional[torch.cuda.Stream] = None, ): """ - CPU Offloading feature for seqeuences of layers. Can be used for arbitrary layers, not necessarily + CPU Offloading feature for sequences of layers. Can be used for arbitrary layers, not necessarily for these provided by the TE. Usage: @@ -710,7 +710,7 @@ def get_cpu_offload_context( Number of layers in the model that will be used under this context. offload_activations : bool, default = True Deprecated. - offload_weights : bool, default = True + offload_weights : bool, default = False Deprecated. double_buffering : bool, default = False Deprecated. @@ -769,14 +769,14 @@ def get_cpu_offload_context( out[i] = sync_function(out[i]) manual_controller.start_offload_layer(i) - offload_stream.synchronize() + # Release GPU memory - each call inserts a GPU-side wait_event on the compute stream for i in range(num_layers): manual_controller.release_activation_forward_gpu_memory(i) + # Start reloading - backward will wait for each tensor's reload via wait_event for i in range(num_layers - 1, -1, -1): manual_controller.start_reload_layer(i) - offload_stream.synchronize() for i in range(num_layers): out[i].sum().backward() From d5ce4166327d60f4822531c4f1d74bb8804dc885 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 11 Mar 2026 09:47:41 -0700 Subject: [PATCH 265/521] Add guard at lowest JAX version that still supports triton kernel calling (#2741) add guard at bisected jax version where lower is segfault Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pylint: remove unused lru_cache import and fix import order in helper.py Signed-off-by: tdophung * Guard Triton tests against JAX < 0.8.0 using release version check - Add version_utils.py with is_triton_extension_supported() checking JAX >= 0.8.0 (release version, not dev snapshot) and TRITON_EXTENSION_MIN_JAX_VERSION constant - Add pytest.mark.triton marker and conftest hook to skip marked tests on old JAX - Add require_triton() for module-level skipping in test files - Rewrite triton_extensions to use is_triton_extension_supported() instead of direct jaxlib dev-version comparison Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address review: allow_module_level, drop is_triton_extension_supported re-export, revert test.sh - require_triton(): add allow_module_level=True to pytest.skip() so module-level calls on old JAX produce a proper skip instead of a collection failure - Remove is_triton_extension_supported from triton_extensions/utils.py __all__: importing triton_extensions on JAX < 0.8.0 raises immediately, so re-exporting the check from there defeats its purpose; callers should import directly from transformer_engine.jax.version_utils - Revert qa/L0_jax_lint/test.sh TE_PATH to /opt/transformerengine (local dev path was accidentally committed; pass TE_PATH= at invocation time instead) Signed-off-by: tdophung * Address review: move version guard before gpu_triton import, fix __all__ and hardcoded version - Move is_triton_extension_supported() guard before the gpu_triton import block with a comment clarifying the segfault is at dispatch time, not import time - Remove _jax_version_meet_requirement from version_utils __all__ (private helper, not a public API; callers import it explicitly as needed) - Use TRITON_EXTENSION_MIN_JAX_VERSION constant in conftest marker description instead of hardcoded '0.8.0' Signed-off-by: tdophung * address more comments Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- qa/L0_jax_lint/test.sh | 0 tests/jax/conftest.py | 27 ++++++++++++ tests/jax/test_distributed_permutation.py | 28 +++++++++--- tests/jax/test_distributed_router.py | 29 +++++++++++-- tests/jax/test_fused_router.py | 30 +++++++++++-- tests/jax/test_permutation.py | 34 ++++++++++++--- tests/jax/test_triton_custom_calls.py | 5 ++- tests/jax/utils.py | 15 +++++++ transformer_engine/jax/quantize/helper.py | 19 ++------ .../jax/triton_extensions/__init__.py | 3 ++ .../jax/triton_extensions/utils.py | 20 +++++++++ transformer_engine/jax/version_utils.py | 43 +++++++++++++++++++ 12 files changed, 218 insertions(+), 35 deletions(-) mode change 100644 => 100755 qa/L0_jax_lint/test.sh create mode 100644 transformer_engine/jax/version_utils.py diff --git a/qa/L0_jax_lint/test.sh b/qa/L0_jax_lint/test.sh old mode 100644 new mode 100755 diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py index 6b7520d147..db30f0ed39 100644 --- a/tests/jax/conftest.py +++ b/tests/jax/conftest.py @@ -11,6 +11,10 @@ import transformer_engine.jax from transformer_engine_jax import get_device_compute_capability +from transformer_engine.jax.version_utils import ( + TRITON_EXTENSION_MIN_JAX_VERSION, + is_triton_extension_supported, +) @pytest.fixture(autouse=True, scope="function") @@ -83,5 +87,28 @@ def pytest_sessionfinish(self, session, exitstatus): def pytest_configure(config): + config.addinivalue_line( + "markers", + "triton: mark test (or test class) as requiring JAX Triton kernel support" + f" (JAX >= {TRITON_EXTENSION_MIN_JAX_VERSION})." + " Apply per test/class with @pytest.mark.triton so non-Triton tests in the same file run on" + " old JAX.", + ) if os.getenv("NVTE_JAX_TEST_TIMING", "0") == "1": config.pluginmanager.register(TestTimingPlugin(), "test_timing") + + +def pytest_collection_modifyitems(config, items): + """Skip tests marked 'triton' when JAX is too old for Triton kernel dispatch.""" + if is_triton_extension_supported(): + return + skip_triton = pytest.mark.skip( + reason=( + f"JAX >= {TRITON_EXTENSION_MIN_JAX_VERSION} required for Triton kernel support. " + "Triton kernel dispatch segfaults with older jaxlib. " + "Upgrade with: pip install --upgrade jax jaxlib" + ) + ) + for item in items: + if item.get_closest_marker("triton"): + item.add_marker(skip_triton) diff --git a/tests/jax/test_distributed_permutation.py b/tests/jax/test_distributed_permutation.py index 04ed236e81..ee7a56a7ec 100644 --- a/tests/jax/test_distributed_permutation.py +++ b/tests/jax/test_distributed_permutation.py @@ -34,11 +34,28 @@ from distributed_test_base import generate_configs from utils import assert_allclose, pytest_parametrize_wrapper -# High-level API with VJP support -from transformer_engine.jax.permutation import ( - token_dispatch, - token_combine, -) + +@pytest.fixture(autouse=True, scope="function") +def _inject_permutation(request): + """Lazy-load permutation API only for tests marked 'triton'. Other tests run without importing. + + We inject into sys.modules[__name__] so test code in this module can use + token_dispatch, token_combine as module-level names (fixture locals are not + visible to test methods). + """ + if not request.node.get_closest_marker("triton"): + yield + return + import sys + from transformer_engine.jax.permutation import token_dispatch, token_combine + + mod = sys.modules[__name__] + mod.token_dispatch = token_dispatch + mod.token_combine = token_combine + yield + + +# High-level API with VJP support (injected by _inject_permutation) # Reference implementations from test_permutation.py from test_permutation import ( @@ -80,6 +97,7 @@ } +@pytest.mark.triton class TestDistributedPermutation: """Test distributed/sharded execution of MoE permutation primitives. diff --git a/tests/jax/test_distributed_router.py b/tests/jax/test_distributed_router.py index 1b3fe14e75..35f59c897d 100644 --- a/tests/jax/test_distributed_router.py +++ b/tests/jax/test_distributed_router.py @@ -34,10 +34,28 @@ from distributed_test_base import generate_configs from utils import assert_allclose, pytest_parametrize_wrapper -from transformer_engine.jax.router import ( - fused_topk_with_score_function, - fused_moe_aux_loss, -) + +@pytest.fixture(autouse=True, scope="function") +def _inject_router(request): + """Lazy-load router API only for tests marked 'triton'. Other tests run without importing. + + We inject into sys.modules[__name__] so test code can use fused_topk_with_score_function, + fused_moe_aux_loss as module-level names (fixture locals are not visible to tests). + """ + if not request.node.get_closest_marker("triton"): + yield + return + import sys + from transformer_engine.jax.router import ( + fused_topk_with_score_function, + fused_moe_aux_loss, + ) + + mod = sys.modules[__name__] + mod.fused_topk_with_score_function = fused_topk_with_score_function + mod.fused_moe_aux_loss = fused_moe_aux_loss + yield + jax.config.update("jax_use_shardy_partitioner", True) @@ -68,6 +86,7 @@ } +@pytest.mark.triton class TestDistributedFusedTopk: """Test distributed execution of fused_topk_with_score_function. @@ -200,6 +219,7 @@ def test_distributed_topk( ) +@pytest.mark.triton class TestDistributedScoreForAuxLoss: """Test distributed execution of fused_topk_with_score_function with compute_aux_scores=True. @@ -333,6 +353,7 @@ def test_distributed_score_for_aux_loss( ) +@pytest.mark.triton class TestDistributedMoEAuxLoss: """Test distributed execution of fused_moe_aux_loss. diff --git a/tests/jax/test_fused_router.py b/tests/jax/test_fused_router.py index 77e89457c8..89a32f1ce2 100644 --- a/tests/jax/test_fused_router.py +++ b/tests/jax/test_fused_router.py @@ -4,6 +4,7 @@ """Tests for fused MoE router CUDA kernels (JAX wrappers).""" +import sys from functools import partial from typing import Optional @@ -13,10 +14,27 @@ from utils import pytest_parametrize_wrapper -from transformer_engine.jax.router import ( - fused_topk_with_score_function, - fused_moe_aux_loss, -) + +@pytest.fixture(autouse=True, scope="function") +def _inject_router(request): + """Lazy-load router API only for tests marked 'triton'. Other tests run without importing. + + We inject into sys.modules[__name__] so test code can use fused_topk_with_score_function, + fused_moe_aux_loss as module-level names (fixture locals are not visible to tests). + """ + if not request.node.get_closest_marker("triton"): + yield + return + from transformer_engine.jax.router import ( + fused_topk_with_score_function, + fused_moe_aux_loss, + ) + + mod = sys.modules[__name__] + mod.fused_topk_with_score_function = fused_topk_with_score_function + mod.fused_moe_aux_loss = fused_moe_aux_loss + yield + # ============================================================================= # Test case definitions (L0 = fast smoke, L2 = comprehensive) @@ -371,6 +389,7 @@ def loss_fused(logits_): @pytest_parametrize_wrapper("group_topk", GROUP_TOPK_OPTIONS) @pytest_parametrize_wrapper("scaling_factor", SCALING_FACTOR_OPTIONS) @pytest_parametrize_wrapper("enable_bias", ENABLE_BIAS_OPTIONS) +@pytest.mark.triton def test_topk_sigmoid( dtype, num_tokens, num_experts, topk, group_topk, scaling_factor, enable_bias ): @@ -397,6 +416,7 @@ def test_topk_sigmoid( @pytest_parametrize_wrapper("use_pre_softmax", USE_PRE_SOFTMAX_OPTIONS) @pytest_parametrize_wrapper("group_topk", GROUP_TOPK_OPTIONS) @pytest_parametrize_wrapper("scaling_factor", SCALING_FACTOR_OPTIONS) +@pytest.mark.triton def test_topk_softmax( dtype, num_tokens, num_experts, topk, use_pre_softmax, group_topk, scaling_factor ): @@ -426,6 +446,7 @@ def test_topk_softmax( SCORE_AUX_LOSS_CASES, ) @pytest_parametrize_wrapper("score_function", SCORE_FUNCTIONS) +@pytest.mark.triton def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function): logits = make_logits(num_tokens, num_experts, score_function, dtype) @@ -486,6 +507,7 @@ def loss_fused(logits_): "num_tokens,num_experts,topk", AUX_LOSS_CASES, ) +@pytest.mark.triton def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk): key = jax.random.PRNGKey(SEED) diff --git a/tests/jax/test_permutation.py b/tests/jax/test_permutation.py index 138a817240..38fbee18e3 100644 --- a/tests/jax/test_permutation.py +++ b/tests/jax/test_permutation.py @@ -5,20 +5,41 @@ """Tests for permutation Triton kernels and high-level APIs""" import functools +import sys import jax import jax.numpy as jnp import pytest -# High-level API with VJP support -from transformer_engine.jax.permutation import ( - token_dispatch, - token_combine, - sort_chunks_by_index, -) from utils import assert_allclose, pytest_parametrize_wrapper +@pytest.fixture(autouse=True, scope="function") +def _inject_permutation(request): + """Lazy-load permutation API only for tests marked 'triton'. Other tests run without importing. + + We inject into sys.modules[__name__] so that test code in this module can use + token_dispatch, token_combine, etc. as module-level names. A plain import inside + this fixture would only bind those names in the fixture's local scope; the test + methods (e.g. in TestHighLevelPermutationAPI) reference them as globals, so they + must exist on the module's namespace. + """ + if not request.node.get_closest_marker("triton"): + yield + return + from transformer_engine.jax.permutation import ( + token_dispatch, + token_combine, + sort_chunks_by_index, + ) + + mod = sys.modules[__name__] + mod.token_dispatch = token_dispatch + mod.token_combine = token_combine + mod.sort_chunks_by_index = sort_chunks_by_index + yield + + ALL_DISPATCH_COMBINE_CASES = [ (128, 5, 128, 3), (1024, 8, 128, 8), @@ -449,6 +470,7 @@ def reference_sort_chunks_by_map( return output, permuted_probs +@pytest.mark.triton class TestHighLevelPermutationAPI: """Test high-level permutation APIs (token_dispatch, token_combine, etc.) diff --git a/tests/jax/test_triton_custom_calls.py b/tests/jax/test_triton_custom_calls.py index 6d969de0d3..846d26a417 100644 --- a/tests/jax/test_triton_custom_calls.py +++ b/tests/jax/test_triton_custom_calls.py @@ -7,7 +7,9 @@ import jax.numpy as jnp import pytest -from utils import assert_allclose, pytest_parametrize_wrapper +from utils import assert_allclose, pytest_parametrize_wrapper, require_triton_or_skip_test_file + +require_triton_or_skip_test_file() import triton import triton.language as tl @@ -23,6 +25,7 @@ def init(): yield +@pytest.mark.triton class TestTritonBinding: """Test Triton binding primitive.""" diff --git a/tests/jax/utils.py b/tests/jax/utils.py index c22b0a6063..c5e564dbc7 100644 --- a/tests/jax/utils.py +++ b/tests/jax/utils.py @@ -26,6 +26,10 @@ make_swa_mask, ) from transformer_engine.jax.quantize.helper import DType as TEDType +from transformer_engine.jax.version_utils import ( + TRITON_EXTENSION_MIN_JAX_VERSION, + is_triton_extension_supported, +) PRNGKey = Any Shape = Tuple[int, ...] @@ -40,6 +44,17 @@ NVTE_DEBUG_NUMERICS = bool(int(os.getenv("NVTE_DEBUG_NUMERICS", 0))) +def require_triton_or_skip_test_file(): + """Skip the current test file if JAX is too old for Triton kernel support (calls pytest.skip).""" + if not is_triton_extension_supported(): + pytest.skip( + f"JAX >= {TRITON_EXTENSION_MIN_JAX_VERSION} required for Triton kernel support. " + "Triton kernel dispatch segfaults with older jaxlib. " + "Upgrade with: pip install --upgrade jax jaxlib", + allow_module_level=True, + ) + + def is_devices_enough(required): """ Check if the available GPUs is enough diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index c5256aef5c..c491bb8638 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -14,11 +14,9 @@ from enum import Enum import hashlib from typing import Optional, Tuple, Dict, Union, Sequence, Type, List -from functools import reduce, lru_cache +from functools import reduce import operator -from importlib.metadata import version as get_pkg_version import warnings -from packaging.version import Version as PkgVersion import jax import jax.numpy as jnp @@ -40,6 +38,7 @@ get_all_mesh_axes, with_sharding_constraint, ) +from transformer_engine.jax.version_utils import jax_version_meet_requirement from .metadata import QuantizeMeta from .scaling_modes import ScalingMode @@ -68,16 +67,6 @@ NVTE_FP8_COLLECTION_NAME = "fp8_metas" -@lru_cache(maxsize=None) -def _jax_version_meet_requirement(version: str): - """ - Helper function checking if required JAX version is available - """ - jax_version = PkgVersion(get_pkg_version("jax")) - jax_version_required = PkgVersion(version) - return jax_version >= jax_version_required - - def _check_delayed_scaling_fp8_support(gpu_arch) -> Tuple[bool, str]: """Check if delayed scaling FP8 is supported on the given GPU architecture. @@ -111,7 +100,7 @@ def _check_block_scaling_fp8_support(gpu_arch) -> Tuple[bool, str]: return False, "CublasLt version 12.8.0 or higher required for MXFP8 execution." if get_cuda_version() < 12080: return False, "Cuda version 12.8 or higher required for MXFP8 execution." - if not _jax_version_meet_requirement("0.5.3"): + if not jax_version_meet_requirement("0.5.3"): return False, "Jax version 0.5.3 or higher required for MXFP8 execution." return True, "" @@ -124,7 +113,7 @@ def _check_fp4_support(gpu_arch) -> Tuple[bool, str]: return False, "CublasLt version 12.8.0 or higher required for NVFP4 execution." if get_cuda_version() < 12080: return False, "Cuda version 12.8 or higher required for NVFP4 execution." - if not _jax_version_meet_requirement("0.5.3"): + if not jax_version_meet_requirement("0.5.3"): return False, "Jax version 0.5.3 or higher required for NVFP4 execution." return True, "" diff --git a/transformer_engine/jax/triton_extensions/__init__.py b/transformer_engine/jax/triton_extensions/__init__.py index d9708fde9f..150a5fbf12 100644 --- a/transformer_engine/jax/triton_extensions/__init__.py +++ b/transformer_engine/jax/triton_extensions/__init__.py @@ -54,6 +54,9 @@ def lowering(ctx, x, **kwargs): from transformer_engine.jax.triton_extensions import get_triton_info info = get_triton_info() print(f"Using Triton {info['version']} from {info['source']}") + + # Check if JAX version supports Triton (without importing triton_extensions) + from transformer_engine.jax.version_utils import is_triton_extension_supported """ from .utils import * diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 2627a08929..28e3f08e18 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -42,6 +42,11 @@ import jax import jax.numpy as jnp +from ..version_utils import ( + TRITON_EXTENSION_MIN_JAX_VERSION, + is_triton_extension_supported, +) + # Placeholder package version on PyPI that should never be used _PYTORCH_TRITON_PLACEHOLDER_VERSION = "0.0.1" @@ -150,6 +155,21 @@ def _check_triton_compatibility(): # Perform compatibility check and get triton info _TRITON_VERSION, _IS_PYTORCH_TRITON = _check_triton_compatibility() +# Enforce minimum JAX version before importing gpu_triton. The segfault on old +# jaxlib occurs at Triton kernel dispatch time, not at import time, so gpu_triton +# itself is safe to import on older jaxlib. The guard is placed here (before the +# import) as a belt-and-suspenders measure so that if the import behaviour ever +# changes, we still fail fast with a clear error rather than a cryptic crash. +if not is_triton_extension_supported(): + raise RuntimeError( + f"JAX >= {TRITON_EXTENSION_MIN_JAX_VERSION} required for " + "transformer_engine.jax.triton_extensions. " + "Triton kernel dispatch segfaults with older jaxlib. " + f"Current jax version: {jax.__version__}. " + "Please upgrade: pip install --upgrade jax jaxlib. " + "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead." + ) + try: from jax._src.lib import gpu_triton from triton.compiler import compiler as tc diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py new file mode 100644 index 0000000000..04b7ff879a --- /dev/null +++ b/transformer_engine/jax/version_utils.py @@ -0,0 +1,43 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +""" +JAX version helpers. + +Provides version checks for JAX that can be used across TE JAX (quantize, triton +extensions, etc.) without pulling in feature-specific code. +""" + +from functools import lru_cache +from importlib.metadata import version as get_pkg_version + +from packaging.version import Version as PkgVersion + + +@lru_cache(maxsize=None) +def jax_version_meet_requirement(version: str): + """Return True if the installed JAX version is >= the required version.""" + jax_version = PkgVersion(get_pkg_version("jax")) + jax_version_required = PkgVersion(version) + return jax_version >= jax_version_required + + +# Minimum JAX version required for Triton kernel dispatch (jaxlib < 0.8.0 segfaults). +TRITON_EXTENSION_MIN_JAX_VERSION = "0.8.0" + + +def is_triton_extension_supported() -> bool: + """Return True if the current JAX version supports Triton kernel dispatch. + + JAX/jaxlib >= 0.8.0 is required. Older versions segfault when dispatching + Triton kernels. Use this to skip tests or gate features without importing + triton_extensions (which would raise immediately on old jax). + """ + return jax_version_meet_requirement(TRITON_EXTENSION_MIN_JAX_VERSION) + + +__all__ = [ + "jax_version_meet_requirement", + "is_triton_extension_supported", + "TRITON_EXTENSION_MIN_JAX_VERSION", +] From f6001c492270742d70bce88fd886d3e5c1a08cd7 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 11 Mar 2026 10:09:15 -0700 Subject: [PATCH 266/521] Support configurable number of philox rounds for stochastic rounding (#2751) * Support configurable number of philox rounds for SR during build Signed-off-by: Kirthi Shankar Sivamani * format and lint Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- docs/envvars.rst | 6 ++++++ transformer_engine/common/CMakeLists.txt | 16 ++++++++++++++++ .../common/cast/nvfp4/core_nvfp4.cuh | 9 +++++---- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 2 +- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 4 ++-- .../quantize_transpose_nvfp4_tuned_1D.cuh | 3 ++- transformer_engine/common/common.h | 6 ++++++ ...ow_cast_col_hadamard_transform_cast_fusion.cu | 8 ++++++-- .../group_hadamard_transform_cast_fusion.cu | 3 ++- ...ow_cast_col_hadamard_transform_cast_fusion.cu | 8 ++++++-- .../hadamard_transform_cast_fusion.cu | 3 ++- .../quantize_transpose_vector_blockwise_fp4.cu | 10 +++++----- 12 files changed, 59 insertions(+), 19 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 86b313b133..85445430f8 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -72,6 +72,12 @@ Build Configuration :Default: Not set :Description: Internal flag set to ``1`` during the build process to indicate that the project is being built. Not intended for external use. +.. envvar:: NVTE_BUILD_NUM_PHILOX_ROUNDS + + :Type: ``int`` (positive integer) + :Default: ``10`` + :Description: Number of Philox4x32 rounds used by stochastic rounding kernels. Must be a positive integer. + Optional Dependencies ^^^^^^^^^^^^^^^^^^^^^ diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a105a0343f..b3d48f68bd 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -301,6 +301,22 @@ if (NVTE_WITH_CUBLASMP) message(STATUS "Using cuBLASMp at: ${CUBLASMP_DIR}") endif() +# Number of philox4x32 rounds for stochastic rounding (build-time constant). +set(NVTE_BUILD_NUM_PHILOX_ROUNDS_STR $ENV{NVTE_BUILD_NUM_PHILOX_ROUNDS}) +if (NOT NVTE_BUILD_NUM_PHILOX_ROUNDS_STR) + set(NVTE_BUILD_NUM_PHILOX_ROUNDS_STR "10") +endif() +if (NOT NVTE_BUILD_NUM_PHILOX_ROUNDS_STR MATCHES "^[1-9][0-9]*$") + message(FATAL_ERROR + "Environment variable NVTE_BUILD_NUM_PHILOX_ROUNDS must be a positive integer, " + "but got '${NVTE_BUILD_NUM_PHILOX_ROUNDS_STR}'.") +endif() +set(NVTE_BUILD_NUM_PHILOX_ROUNDS ${NVTE_BUILD_NUM_PHILOX_ROUNDS_STR}) + +target_compile_definitions(transformer_engine + PUBLIC NVTE_BUILD_NUM_PHILOX_ROUNDS=${NVTE_BUILD_NUM_PHILOX_ROUNDS}) +message(STATUS "Philox rounds for stochastic rounding: ${NVTE_BUILD_NUM_PHILOX_ROUNDS}") + # Hack to enable dynamic loading in cuDNN frontend target_compile_definitions(transformer_engine PUBLIC NV_CUDNN_FRONTEND_USE_DYNAMIC_LOADING) diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index bdbe5cddc3..8d2d806559 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -88,10 +88,11 @@ __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const return global_encode_scale; } -__device__ __forceinline__ uint32_t -get_rbits(transformer_engine::curanddx::detail::philox4x32_native_state<10> &rng, - // philox4x32_native_state<10>: 10 rounds of philox4_32 - uint4 &random_uint4, int &rnd_idx) { +__device__ __forceinline__ uint32_t get_rbits( + transformer_engine::curanddx::detail::philox4x32_native_state + &rng, + // philox4x32_native_state: compile-time configurable rounds + uint4 &random_uint4, int &rnd_idx) { if (rnd_idx == 4) { rnd_idx = 0; random_uint4 = rng.generate4(); 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 1ceb08a9d0..a2f3dac15a 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -191,7 +191,7 @@ __global__ void __launch_bounds__(THREADS_NUM) threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state rng; rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 99776db281..f164636e38 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -135,7 +135,7 @@ __global__ void __launch_bounds__(THREADS_NUM) threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state rng; rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; // Index of the random number. It increments each time when used and resets to 0 if reaches 4x @@ -647,7 +647,7 @@ __global__ void __launch_bounds__(THREADS_NUM) threadIdx.x + blockIdx.x * THREADS_NUM + blockIdx.y * gridDim.x * THREADS_NUM; const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state rng; rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = USE_STOCHASTIC_ROUNDING ? rng.generate4() : uint4{0, 0, 0, 0}; int rnd_idx = 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 061a88fd6d..fc337f6078 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 @@ -142,7 +142,8 @@ using OType2x3D = fp4e2m1x2[BUFFS_NUM_OUT][BUFF_OUT_DIM_Y][BUFF_OUT_DIM_X]; using OType2xt3D = fp4e2m1x2[BUFFS_NUM_OUT_TR][BUFF_OUT_TR_DIM_Y][BUFF_OUT_TR_DIM_X]; using ScalesType2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_Y][SCALES_PER_CHUNK_X]; using ScalesTypeTr2D = nvfp4_scale_t[TunableConfig::CHUNK_DIM_X][SCALES_PER_CHUNK_Y]; -using RNG_t = typename transformer_engine::curanddx::detail::philox4x32_native_state<10>; +using RNG_t = typename transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS>; template struct SCALING_COEFFICIENT_TYPE {}; diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 1749b5734a..452784bac8 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -10,6 +10,12 @@ #include #define FP4_TYPE_SUPPORTED (CUDA_VERSION >= 12080) +#ifndef NVTE_BUILD_NUM_PHILOX_ROUNDS +#define NVTE_BUILD_NUM_PHILOX_ROUNDS 10 +#endif +static_assert(NVTE_BUILD_NUM_PHILOX_ROUNDS > 0, + "NVTE_BUILD_NUM_PHILOX_ROUNDS must be a positive integer."); + #include #include #include 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 19583b3afb..6f3cf90d90 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 @@ -913,7 +913,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g // Prepare stochastic rounding random state if enabled uint4 random_uint4 = uint4{0, 0, 0, 0}; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; // "Prefetch" a stochastic rounding state for the first tile if constexpr (kEnableStochasticRounding) { const size_t rng_sequence = global_thread_idx + k_tile * 512 + @@ -1072,7 +1074,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g Tensor amax = make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); Tensor pvscales = make_tensor_like(amax); - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; if constexpr (kEnableStochasticRounding) { const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + 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 85bb98f0f1..1e40fd4a58 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 @@ -660,7 +660,8 @@ __global__ static void group_rht_gemm_device( // Initialize RNG for tile const size_t rng_sequence = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state + rng; rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = uint4{0, 0, 0, 0}; 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 1ef1f81e82..4013fdf119 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 @@ -891,7 +891,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( // Prepare stochastic rounding random state if enabled uint4 random_uint4 = uint4{0, 0, 0, 0}; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; // "Prefetch" a stochastic rounding state for the first tile if constexpr (kEnableStochasticRounding) { const size_t rng_sequence = global_thread_idx + k_tile * 512 + @@ -1048,7 +1050,9 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( Tensor amax = make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); Tensor pvscales = make_tensor_like(amax); - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; if constexpr (kEnableStochasticRounding) { const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + 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 0696deaaa7..1a2462e6fa 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -516,7 +516,8 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, const size_t rng_sequence = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state + rng; rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = uint4{0, 0, 0, 0}; 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 798c712fda..e25cc607e5 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -201,10 +201,10 @@ __device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_ return global_encode_scale; } -__device__ __forceinline__ uint32_t -get_rbits(transformer_engine::curanddx::detail::philox4x32_native_state<10>& - rng, // philox4x32_native_state<10>: 10 rounds of philox4_32 - uint4& random_uint4, int& rnd_idx) { +__device__ __forceinline__ uint32_t get_rbits( + transformer_engine::curanddx::detail::philox4x32_native_state& + rng, // NVTE_BUILD_NUM_PHILOX_ROUNDS rounds of philox4x32 + uint4& random_uint4, int& rnd_idx) { if (rnd_idx == 4) { rnd_idx = 0; random_uint4 = rng.generate4(); @@ -344,7 +344,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - transformer_engine::curanddx::detail::philox4x32_native_state<10> rng; + transformer_engine::curanddx::detail::philox4x32_native_state rng; rng.init(rng_seed, rng_sequence, rng_offset); uint4 random_uint4 = kApplyStochasticRounding ? rng.generate4() : uint4{0, 0, 0, 0}; From 61f95942178b8d0cd3bdad886586c7a65aa1839a Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 11 Mar 2026 10:17:37 -0700 Subject: [PATCH 267/521] [All] Added better error messages (#2705) * Added better error messages Signed-off-by: Przemek Tredak * Update transformer_engine/pytorch/distributed.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixes Signed-off-by: Przemek Tredak --------- Signed-off-by: Przemek Tredak Signed-off-by: Przemyslaw Tredak 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> --- transformer_engine/common/common.h | 405 ++++++++-------- .../common/fused_attn/fused_attn.cpp | 94 +++- .../common/fused_router/utils.h | 79 ++-- .../common/gemm/cutlass_grouped_gemm.cuh | 6 +- .../common/normalization/common.cpp | 8 +- .../common/transformer_engine.cpp | 4 +- .../quantize_transpose_square_blockwise.cu | 3 +- .../jax/cpp_extensions/attention.py | 95 +++- transformer_engine/jax/cpp_extensions/gemm.py | 442 +++++++++++------- .../jax/cpp_extensions/normalization.py | 95 +++- .../jax/cpp_extensions/router.py | 40 +- transformer_engine/jax/flax/module.py | 3 +- transformer_engine/jax/flax/transformer.py | 75 ++- transformer_engine/jax/layernorm.py | 6 +- .../pytorch/cpp_extensions/fused_attn.py | 93 ++-- transformer_engine/pytorch/cpu_offload.py | 85 ++-- transformer_engine/pytorch/csrc/common.cpp | 3 +- transformer_engine/pytorch/csrc/quantizer.cpp | 7 +- .../pytorch/custom_recipes/gemm.py | 42 +- .../custom_recipes/quantization_nvfp4.py | 120 +++-- transformer_engine/pytorch/distributed.py | 120 +++-- transformer_engine/pytorch/graph.py | 253 ++++++---- transformer_engine/pytorch/module/base.py | 170 ++++--- .../pytorch/module/grouped_linear.py | 44 +- transformer_engine/pytorch/permutation.py | 156 +++++-- transformer_engine/pytorch/quantization.py | 3 +- .../pytorch/tensor/grouped_tensor.py | 18 +- .../tensor/storage/nvfp4_tensor_storage.py | 35 +- transformer_engine/pytorch/tensor/utils.py | 147 ++++-- transformer_engine/pytorch/transformer.py | 89 ++-- transformer_engine/pytorch/utils.py | 65 +-- 31 files changed, 1856 insertions(+), 949 deletions(-) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 452784bac8..b1543d55a2 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -44,6 +44,8 @@ namespace transformer_engine { std::string to_string(const DType type); std::string to_string(const NVTEScalingMode &mode); +inline std::string to_string_like(const DType &val) { return to_string(val); } + inline bool is_tensor_scaling(const NVTEScalingMode &mode) { return mode == NVTE_DELAYED_TENSOR_SCALING; } @@ -625,140 +627,149 @@ struct TypeInfo { #define SWITCH_FP4_TYPE_HANDLE(type, ...) // do nothing #endif -#define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kByte: { \ - using type = unsigned char; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kInt16: { \ - using type = int16_t; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kInt32: { \ - using type = int32_t; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kInt64: { \ - using type = int64_t; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat16: { \ - using type = fp16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E4M3: { \ - using type = fp8e4m3; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E5M2: { \ - using type = fp8e5m2; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E8M0: { \ - using type = byte; \ - { __VA_ARGS__ } \ - } break; \ - SWITCH_FP4_TYPE_HANDLE(type, __VA_ARGS__) \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kByte: { \ + using type = unsigned char; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kInt16: { \ + using type = int16_t; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kInt32: { \ + using type = int32_t; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kInt64: { \ + using type = int64_t; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat16: { \ + using type = fp16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E4M3: { \ + using type = fp8e4m3; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E5M2: { \ + using type = fp8e5m2; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E8M0: { \ + using type = byte; \ + { __VA_ARGS__ } \ + } break; \ + SWITCH_FP4_TYPE_HANDLE(type, __VA_ARGS__) \ + default: \ + NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Byte, Int16, Int32, Int64, Float32, " \ + "Float16, BFloat16, Float8E4M3, Float8E5M2, " \ + "Float8E8M0, Float4E2M1."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat16: { \ - using type = fp16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E4M3: { \ - using type = fp8e4m3; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E5M2: { \ - using type = fp8e5m2; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_FLOAT(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat16: { \ + using type = fp16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E4M3: { \ + using type = fp8e4m3; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E5M2: { \ + using type = fp8e5m2; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float32, Float16, BFloat16, " \ + "Float8E4M3, Float8E5M2."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat16: { \ - using type = fp16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E5M2: { \ - using type = fp8e5m2; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E4M3: { \ - using type = fp8e4m3; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_OUTPUT(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat16: { \ + using type = fp16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E5M2: { \ + using type = fp8e5m2; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E4M3: { \ + using type = fp8e4m3; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported output dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float32, Float16, BFloat16, " \ + "Float8E5M2, Float8E4M3."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat16: { \ - using type = fp16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat16: { \ + using type = fp16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float32, Float16, BFloat16."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type, expected Float32 or BFloat16."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_FP32_BF16(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float32, BFloat16."); \ } // Add a pack_size argument to select the packed type for FP4 @@ -770,80 +781,90 @@ struct TypeInfo { { __VA_ARGS__ } \ } break; \ default: \ - NVTE_ERROR("Invalid type."); \ + NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ + ". Expected: Float4E2M1."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat8E5M2: { \ - using type = fp8e5m2; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E4M3: { \ - using type = fp8e4m3; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat8E5M2: { \ + using type = fp8e5m2; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E4M3: { \ + using type = fp8e4m3; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float8E5M2, Float8E4M3."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat16: { \ - using type = fp16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat8E5M2: \ - case DType::kFloat8E4M3: { \ - NVTE_ERROR("FP8 type not instantiated for input."); \ - } break; \ - case DType::kFloat4E2M1: { \ - NVTE_ERROR("FP4 type not instantiated for input."); \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat16: { \ + using type = fp16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat8E5M2: \ + case DType::kFloat8E4M3: { \ + NVTE_ERROR("FP8 dtype ", to_string(static_cast(dtype)), \ + " is not instantiated for input. " \ + "Expected one of: Float32, Float16, BFloat16."); \ + } break; \ + case DType::kFloat4E2M1: { \ + NVTE_ERROR( \ + "FP4 dtype Float4E2M1 is not instantiated " \ + "for input. Expected one of: Float32, Float16, " \ + "BFloat16."); \ + } break; \ + default: \ + NVTE_ERROR("Unsupported input dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float32, Float16, BFloat16."); \ } -#define TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat16: { \ - using type = fp16; \ - __VA_ARGS__; \ - break; \ - } \ - case DType::kBFloat16: { \ - using type = bf16; \ - __VA_ARGS__; \ - break; \ - } \ - default: \ - NVTE_ERROR("Invalid type for 16 bit."); \ +#define TRANSFORMER_ENGINE_TYPE_SWITCH_16BIT(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat16: { \ + using type = fp16; \ + __VA_ARGS__; \ + break; \ + } \ + case DType::kBFloat16: { \ + using type = bf16; \ + __VA_ARGS__; \ + break; \ + } \ + default: \ + NVTE_ERROR("Unsupported 16-bit dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float16, BFloat16."); \ } -#define TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH(SCALE_DIM, DIM, ...) \ - switch (SCALE_DIM) { \ - case 1: { \ - constexpr size_t DIM = 1; \ - { __VA_ARGS__ } \ - } break; \ - case 32: { \ - constexpr size_t DIM = 32; \ - { __VA_ARGS__ } \ - } break; \ - default: { \ - NVTE_ERROR("Invalid size of the MX scaling factor."); \ - } \ +#define TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH(SCALE_DIM, DIM, ...) \ + switch (SCALE_DIM) { \ + case 1: { \ + constexpr size_t DIM = 1; \ + { __VA_ARGS__ } \ + } break; \ + case 32: { \ + constexpr size_t DIM = 32; \ + { __VA_ARGS__ } \ + } break; \ + default: { \ + NVTE_ERROR("Unsupported MX scaling factor dimension ", SCALE_DIM, \ + ". Expected one of: 1, 32."); \ + } \ } #define TRANSFORMER_ENGINE_SWITCH_CONDITION(CONDITION, FLAG, ...) \ diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index abdce7fdac..6a136c67e4 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -15,6 +15,88 @@ #include "fused_attn_fp8.h" #include "utils.h" +namespace transformer_engine { + +std::string to_string(NVTE_QKV_Layout layout) { + switch (layout) { + case NVTE_SB3HD: + return "NVTE_SB3HD"; + case NVTE_SBH3D: + return "NVTE_SBH3D"; + case NVTE_SBHD_SB2HD: + return "NVTE_SBHD_SB2HD"; + case NVTE_SBHD_SBH2D: + return "NVTE_SBHD_SBH2D"; + case NVTE_SBHD_SBHD_SBHD: + return "NVTE_SBHD_SBHD_SBHD"; + case NVTE_BS3HD: + return "NVTE_BS3HD"; + case NVTE_BSH3D: + return "NVTE_BSH3D"; + case NVTE_BSHD_BS2HD: + return "NVTE_BSHD_BS2HD"; + case NVTE_BSHD_BSH2D: + return "NVTE_BSHD_BSH2D"; + case NVTE_BSHD_BSHD_BSHD: + return "NVTE_BSHD_BSHD_BSHD"; + case NVTE_T3HD: + return "NVTE_T3HD"; + case NVTE_TH3D: + return "NVTE_TH3D"; + case NVTE_THD_T2HD: + return "NVTE_THD_T2HD"; + case NVTE_THD_TH2D: + return "NVTE_THD_TH2D"; + case NVTE_THD_THD_THD: + return "NVTE_THD_THD_THD"; + case NVTE_SBHD_BSHD_BSHD: + return "NVTE_SBHD_BSHD_BSHD"; + case NVTE_BSHD_SBHD_SBHD: + return "NVTE_BSHD_SBHD_SBHD"; + case NVTE_THD_BSHD_BSHD: + return "NVTE_THD_BSHD_BSHD"; + case NVTE_THD_SBHD_SBHD: + return "NVTE_THD_SBHD_SBHD"; + case NVTE_Paged_KV_BSHD_BSHD_BSHD: + return "NVTE_Paged_KV_BSHD_BSHD_BSHD"; + case NVTE_Paged_KV_BSHD_SBHD_SBHD: + return "NVTE_Paged_KV_BSHD_SBHD_SBHD"; + case NVTE_Paged_KV_SBHD_BSHD_BSHD: + return "NVTE_Paged_KV_SBHD_BSHD_BSHD"; + case NVTE_Paged_KV_SBHD_SBHD_SBHD: + return "NVTE_Paged_KV_SBHD_SBHD_SBHD"; + case NVTE_Paged_KV_THD_BSHD_BSHD: + return "NVTE_Paged_KV_THD_BSHD_BSHD"; + case NVTE_Paged_KV_THD_SBHD_SBHD: + return "NVTE_Paged_KV_THD_SBHD_SBHD"; + default: + return "UNKNOWN_QKV_LAYOUT(" + std::to_string(static_cast(layout)) + ")"; + } +} + +std::string to_string(NVTE_QKV_Format format) { + switch (format) { + case NVTE_SBHD: + return "NVTE_SBHD"; + case NVTE_BSHD: + return "NVTE_BSHD"; + case NVTE_THD: + return "NVTE_THD"; + case NVTE_BSHD_2SBHD: + return "NVTE_BSHD_2SBHD"; + case NVTE_SBHD_2BSHD: + return "NVTE_SBHD_2BSHD"; + case NVTE_THD_2BSHD: + return "NVTE_THD_2BSHD"; + case NVTE_THD_2SBHD: + return "NVTE_THD_2SBHD"; + default: + return "UNKNOWN_QKV_FORMAT(" + std::to_string(static_cast(format)) + ")"; + } +} + +} // namespace transformer_engine + // map NVTE_QKV_Layout to NVTE_QKV_Layout_Group NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { switch (qkv_layout) { @@ -50,7 +132,8 @@ NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: return NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD; default: - NVTE_ERROR("qkv_layout not supported!"); + NVTE_ERROR("Unsupported qkv_layout ", transformer_engine::to_string(qkv_layout), + " in nvte_get_qkv_layout_group."); } } @@ -90,7 +173,8 @@ NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: return NVTE_QKV_Format::NVTE_THD_2SBHD; default: - NVTE_ERROR("qkv_layout not supported!"); + NVTE_ERROR("Unsupported qkv_layout ", transformer_engine::to_string(qkv_layout), + " in nvte_get_qkv_format."); } } @@ -109,7 +193,8 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Format::NVTE_THD_2SBHD: return NVTE_QKV_Format::NVTE_THD; default: - NVTE_ERROR("qkv_layout not supported!"); + NVTE_ERROR("Unsupported qkv_format ", transformer_engine::to_string(qkv_format), + " in nvte_get_q_format."); } } @@ -128,7 +213,8 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Format::NVTE_THD: return NVTE_QKV_Format::NVTE_THD; default: - NVTE_ERROR("qkv_layout not supported!"); + NVTE_ERROR("Unsupported qkv_format ", transformer_engine::to_string(qkv_format), + " in nvte_get_kv_format."); } } diff --git a/transformer_engine/common/fused_router/utils.h b/transformer_engine/common/fused_router/utils.h index 60e731d990..372efdc490 100644 --- a/transformer_engine/common/fused_router/utils.h +++ b/transformer_engine/common/fused_router/utils.h @@ -250,46 +250,49 @@ __device__ inline void naive_topk_and_mask(CompType *scores, int data_size, int } // Current TE only support float32/bf16/fp16, float64 probs should be considered in the future -#define TE_ROUTER_PROBS_TYPE_SWITCH_ALL(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat16: { \ - using type = fp16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TE_ROUTER_PROBS_TYPE_SWITCH_ALL(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat16: { \ + using type = fp16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported router probs dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Float32, Float16, BFloat16."); \ } -#define TE_ROUTER_INDEX_TYPE_SWITCH_ALL(dtype, type, ...) \ - switch (dtype) { \ - using namespace transformer_engine; \ - case DType::kInt32: { \ - using type = int32_t; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kInt64: { \ - using type = int64_t; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kBFloat16: { \ - using type = bf16; \ - { __VA_ARGS__ } \ - } break; \ - case DType::kFloat32: { \ - using type = float; \ - { __VA_ARGS__ } \ - } break; \ - default: \ - NVTE_ERROR("Invalid type."); \ +#define TE_ROUTER_INDEX_TYPE_SWITCH_ALL(dtype, type, ...) \ + switch (dtype) { \ + using namespace transformer_engine; \ + case DType::kInt32: { \ + using type = int32_t; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kInt64: { \ + using type = int64_t; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kBFloat16: { \ + using type = bf16; \ + { __VA_ARGS__ } \ + } break; \ + case DType::kFloat32: { \ + using type = float; \ + { __VA_ARGS__ } \ + } break; \ + default: \ + NVTE_ERROR("Unsupported router index dtype ", to_string(static_cast(dtype)), \ + ". Expected one of: Int32, Int64, BFloat16, " \ + "Float32."); \ } } // namespace fused_router } // namespace transformer_engine diff --git a/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh b/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh index eb99edc4d3..aa2bde4203 100644 --- a/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh +++ b/transformer_engine/common/gemm/cutlass_grouped_gemm.cuh @@ -326,17 +326,17 @@ void CutlassGroupedGemm(const NVTETensor* A, const NVTETensor* B, NVTETensor* D, // Check can implement the kernel. if (gemm.can_implement(arguments) != cutlass::Status::kSuccess) { - NVTE_CHECK(false, "Failed to implement CUTLASS Grouped GEMM"); + NVTE_ERROR("Failed to implement CUTLASS Grouped GEMM with ", num_gemms, " GEMMs"); } // Initialize the kernel. if (gemm.initialize(arguments, kernel_workspace_ptr) != cutlass::Status::kSuccess) { - NVTE_CHECK(false, "Failed to initialize CUTLASS Grouped GEMM"); + NVTE_ERROR("Failed to initialize CUTLASS Grouped GEMM with ", num_gemms, " GEMMs"); } // Execute the kernel in the current stream. if (gemm.run(stream) != cutlass::Status::kSuccess) { - NVTE_CHECK(false, "Failed to run CUTLASS Grouped GEMM"); + NVTE_ERROR("Failed to run CUTLASS Grouped GEMM with ", num_gemms, " GEMMs"); } } diff --git a/transformer_engine/common/normalization/common.cpp b/transformer_engine/common/normalization/common.cpp index 852b418b39..11f12775c5 100644 --- a/transformer_engine/common/normalization/common.cpp +++ b/transformer_engine/common/normalization/common.cpp @@ -116,7 +116,9 @@ void TeNormalizationPlan::execute(Tensor* z, void* x_dptr, void* beta_dptr, void* mean_dptr, void* eps_dptr, void* rsigma_dptr, void* workspace_dptr, cudaStream_t stream) { - NVTE_ERROR("Backward normalization should not call the forward execute function!"); + NVTE_ERROR( + "Backward normalization should not call the forward execute function. " + "Use the backward-specific execute overload instead."); } template @@ -165,7 +167,9 @@ void TeNormalizationPlan::execute(void* x_dptr, void* gamma void* dx_dptr, void* dz_dptr, void* add_dptr, void* dbeta_dptr, void* dgamma_dptr, void* workspace_dptr, cudaStream_t stream) { - NVTE_ERROR("Forward normalization should not call the backward execute function!"); + NVTE_ERROR( + "Forward normalization should not call the backward execute function. " + "Use the forward-specific execute overload instead."); } template <> diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index cd02074fbd..1875f4f690 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -650,7 +650,7 @@ NVTEShape nvte_make_shape(const size_t *data, size_t ndim) { NVTEShape nvte_tensor_shape(const NVTETensor tensor) { auto *t = transformer_engine::convertNVTETensor(tensor); if (t == nullptr) { - NVTE_ERROR("Invalid tensor"); + NVTE_ERROR("Invalid tensor: received null pointer in nvte_tensor_shape"); } // Determine tensor shape depending on tensor format @@ -662,7 +662,7 @@ NVTEShape nvte_tensor_shape(const NVTETensor tensor) { NVTEShape nvte_tensor_columnwise_shape(const NVTETensor tensor) { auto *t = transformer_engine::convertNVTETensor(tensor); if (t == nullptr) { - NVTE_ERROR("Invalid tensor"); + 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()); diff --git a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu index 0e286009a5..3a8536587c 100644 --- a/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu +++ b/transformer_engine/common/transpose/quantize_transpose_square_blockwise.cu @@ -463,7 +463,8 @@ CUtensorMap get_tensor_map(const SimpleTensor& tensor, size_t global_dim_x, size std::is_same_v) { dataType = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_UINT8; } else { - NVTE_CHECK(false, "Invalid Output type (must be FP8)."); + NVTE_ERROR( + "Invalid output type for blockwise transpose (must be FP8: Float8E4M3 or Float8E5M2)."); } CUtensorMap tensor_map_output_trans{}; diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index f4d914062d..40d02f40e1 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -165,13 +165,25 @@ def parse_qkv_aval(q_aval, k_aval, v_aval, qkv_layout): kv_max_seqlen = q_max_seqlen num_gqa_groups = attn_heads v_head_dim = q_head_dim - assert nqkv == 3 + assert nqkv == 3, ( + f"Expected nqkv == 3 for qkvpacked layout, but got nqkv={nqkv} from" + f" q_aval.shape={q_aval.shape}" + ) elif qkv_layout.is_kvpacked(): *q_batch_shape, q_max_seqlen, attn_heads, q_head_dim = q_aval.shape *kv_batch_shape, kv_max_seqlen, nkv, num_gqa_groups, v_head_dim = k_aval.shape - assert q_batch_shape == kv_batch_shape - assert q_head_dim == v_head_dim - assert nkv == 2 + assert q_batch_shape == kv_batch_shape, ( + f"Mismatched batch shapes for kvpacked layout: q_batch_shape={q_batch_shape}," + f" kv_batch_shape={kv_batch_shape}" + ) + assert q_head_dim == v_head_dim, ( + f"Mismatched head dims for kvpacked layout: q_head_dim={q_head_dim}," + f" v_head_dim={v_head_dim}" + ) + assert nkv == 2, ( + f"Expected nkv == 2 for kvpacked layout, but got nkv={nkv} from" + f" k_aval.shape={k_aval.shape}" + ) elif qkv_layout.is_separate(): *q_batch_shape, q_max_seqlen, attn_heads, q_head_dim = q_aval.shape *k_batch_shape, k_max_seqlen, k_num_gqa_groups, k_head_dim = k_aval.shape @@ -244,9 +256,13 @@ def check_seed(self, seed, dropout_probability, is_training): ) seed = seed.astype(self.rng_state_dtype) - assert seed.dtype == self.rng_state_dtype + assert ( + seed.dtype == self.rng_state_dtype + ), f"Expected seed.dtype={self.rng_state_dtype}, but got seed.dtype={seed.dtype}" # Backend takes an int64_t seed, so only the first two u32 elements are taken - assert seed.size >= self.seed_size + assert ( + seed.size >= self.seed_size + ), f"Expected seed.size >= {self.seed_size}, but got seed.size={seed.size}" return seed @@ -363,7 +379,9 @@ def abstract( # 32-bit unsigned int to get the buffer size we need in the C++ kernel checker = _FusedAttnRNGStateChecker() seed_dtype = dtypes.canonicalize_dtype(seed_aval.dtype) - assert seed_dtype == checker.rng_state_dtype + assert ( + seed_dtype == checker.rng_state_dtype + ), f"Expected seed_dtype={checker.rng_state_dtype}, but got seed_dtype={seed_dtype}" rng_state_shape = (seed_aval.shape[0], checker.rng_state_size) rng_state_aval = seed_aval.update(shape=rng_state_shape, dtype=checker.rng_state_dtype) @@ -408,11 +426,19 @@ def abstract( shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) ) - assert softmax_offset_aval.dtype == jnp.float32 + assert ( + softmax_offset_aval.dtype == jnp.float32 + ), f"Expected softmax_offset_aval.dtype=float32, but got {softmax_offset_aval.dtype}" if config.softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: - assert softmax_offset_aval.shape == (1, attn_heads, 1, 1) + assert softmax_offset_aval.shape == (1, attn_heads, 1, 1), ( + f"Expected softmax_offset_aval.shape=(1, {attn_heads}, 1, 1) for" + f" {config.softmax_type}, but got {softmax_offset_aval.shape}" + ) else: - assert softmax_offset_aval.shape == (0,) + assert softmax_offset_aval.shape == (0,), ( + "Expected softmax_offset_aval.shape=(0,) for VANILLA_SOFTMAX, but got" + f" {softmax_offset_aval.shape}" + ) return out_aval, softmax_aux_aval, rng_state_aval, wkspace_aval @@ -533,7 +559,9 @@ def impl( _kv_segment_pos, config: _FusedAttnConfig, ): - assert FusedAttnFwdPrimitive.inner_primitive is not None + assert ( + FusedAttnFwdPrimitive.inner_primitive is not None + ), "FusedAttnFwdPrimitive.inner_primitive has not been registered" sequence_descriptor = SequenceDescriptor( seqlens=(q_seqlen, kv_seqlen), @@ -628,7 +656,9 @@ def convert_to_2d(offsets, batch, max_seqlen): def batcher(batched_args, batch_dims, *, config): # batch_dims: each element is the batch axis (0, ...) or None. Only 0 or None allowed. check_valid_batch_dims(batch_dims) - assert FusedAttnFwdPrimitive.outer_primitive is not None + assert ( + FusedAttnFwdPrimitive.outer_primitive is not None + ), "FusedAttnFwdPrimitive.outer_primitive has not been registered" q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims # Pass through; segment_ids/segment_pos may have different batch dims (e.g. vmapped ids, # replicated pos). get_seqlens_and_offsets() in attention.py handles conversion without expanding. @@ -780,8 +810,15 @@ def abstract( v_dtype = dtypes.canonicalize_dtype(v_aval.dtype) bias_dtype = dtypes.canonicalize_dtype(bias_aval.dtype) doutput_dtype = dtypes.canonicalize_dtype(doutput_aval.dtype) - assert q_dtype == k_dtype == v_dtype == bias_dtype == doutput_dtype - assert q_seqlen_or_cu_seqlen_aval.dtype == kv_seqlen_or_cu_seqlen_aval.dtype + assert q_dtype == k_dtype == v_dtype == bias_dtype == doutput_dtype, ( + f"Mismatched dtypes: q_dtype={q_dtype}, k_dtype={k_dtype}, v_dtype={v_dtype}," + f" bias_dtype={bias_dtype}, doutput_dtype={doutput_dtype}" + ) + assert q_seqlen_or_cu_seqlen_aval.dtype == kv_seqlen_or_cu_seqlen_aval.dtype, ( + "Mismatched seqlen dtypes:" + f" q_seqlen_or_cu_seqlen_aval.dtype={q_seqlen_or_cu_seqlen_aval.dtype}," + f" kv_seqlen_or_cu_seqlen_aval.dtype={kv_seqlen_or_cu_seqlen_aval.dtype}" + ) ( batch_shape, @@ -985,7 +1022,9 @@ def impl( _kv_segment_pos, config, ): - assert FusedAttnBwdPrimitive.inner_primitive is not None + assert ( + FusedAttnBwdPrimitive.inner_primitive is not None + ), "FusedAttnBwdPrimitive.inner_primitive has not been registered" sequence_descriptor = SequenceDescriptor( seqlens=(q_seqlen, kv_seqlen), @@ -1025,7 +1064,9 @@ def convert_to_2d(offsets, batch, max_seqlen): batch, q_max_seqlen, kv_max_seqlen, *_ = FusedAttnHelper.parse_qkv_aval( q, k, v, config.qkv_layout ) - assert len(batch) == 1 + assert ( + len(batch) == 1 + ), f"Expected len(batch) == 1, but got len(batch)={len(batch)}, batch={batch}" kv_batch = q_batch = batch[0] # Gather valid q_seqlen, which is greater than 0 @@ -1084,7 +1125,9 @@ def convert_to_2d(offsets, batch, max_seqlen): @staticmethod def batcher(batched_args, batch_dims, *, config): check_valid_batch_dims(batch_dims) - assert FusedAttnBwdPrimitive.outer_primitive is not None + assert ( + FusedAttnBwdPrimitive.outer_primitive is not None + ), "FusedAttnBwdPrimitive.outer_primitive has not been registered" q_bdim, k_bdim, v_bdim, bias_bdim, softmax_offset_bdim, *_ = batch_dims # Pass through; segment_ids/segment_pos may have different batch dims. Conversion is in attention.py. out_bdims = q_bdim, k_bdim, v_bdim, bias_bdim, softmax_offset_bdim @@ -3398,7 +3441,9 @@ def fused_attn_fwd( raise ValueError(f"Unknown {qkv_layout=}") if attn_bias_type == AttnBiasType.NO_BIAS: - assert bias is None + assert ( + bias is None + ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias={bias}" bias = jnp.zeros(0, dtype=qkv[0].dtype) if softmax_offset is None: @@ -3416,10 +3461,16 @@ def fused_attn_fwd( softmax_offset, (None, HEAD_AXES, None, None) ) else: - assert softmax_type == AttnSoftmaxType.VANILLA_SOFTMAX + assert softmax_type == AttnSoftmaxType.VANILLA_SOFTMAX, ( + "Expected VANILLA_SOFTMAX when softmax_offset is None and not OFF_BY_ONE_SOFTMAX," + f" but got softmax_type={softmax_type}" + ) softmax_offset = jnp.zeros(0, dtype=jnp.float32) else: - assert softmax_offset.dtype == jnp.float32 + assert softmax_offset.dtype == jnp.float32, ( + "Expected softmax_offset.dtype=float32, but got" + f" softmax_offset.dtype={softmax_offset.dtype}" + ) # Shard by heads dimension if not VANILLA_SOFTMAX if softmax_type != AttnSoftmaxType.VANILLA_SOFTMAX: softmax_offset = with_sharding_constraint_by_logical_axes( @@ -3558,7 +3609,9 @@ def fused_attn_bwd( raise ValueError(f"Unknown {qkv_layout=}") if attn_bias_type == AttnBiasType.NO_BIAS: - assert bias is None + assert ( + bias is None + ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias with type={type(bias)}" bias = jnp.zeros(0, dtype=qkv[0].dtype) if softmax_offset is None: diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 70557f29c7..4506adf33b 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -177,17 +177,26 @@ def _quantize_gemm_operands(lhs, rhs, lhs_quantizer, rhs_quantizer, contracting_ flatten_axis=flatten_axis, ) - assert not isinstance(lhs_q, ScaledTensor2x) - assert not isinstance(rhs_q, ScaledTensor2x) + if isinstance(lhs_q, ScaledTensor2x): + raise TypeError( + "Expected lhs_q to not be ScaledTensor2x after quantization, but got" + f" type={type(lhs_q)}" + ) + if isinstance(rhs_q, ScaledTensor2x): + raise TypeError( + "Expected rhs_q to not be ScaledTensor2x after quantization, but got" + f" type={type(rhs_q)}" + ) def has_rht_applied(q: AbstractBaseTensor) -> bool: return isinstance(q, ScaledTensor1x) and q.has_rht_applied - assert has_rht_applied(lhs_q) == has_rht_applied(rhs_q), ( - "With NVFP4_1D_SCALING, if one operand is quantized with RHT, the other must be quantized" - " with RHT as well. This is to ensure the RHT is applied to both and will cancel out in the" - " GEMM." - ) + if has_rht_applied(lhs_q) != has_rht_applied(rhs_q): + raise ValueError( + "With NVFP4_1D_SCALING, if one operand is quantized with RHT, the other must be" + " quantized with RHT as well. This is to ensure the RHT is applied to both and will" + " cancel out in the GEMM." + ) return lhs_q, rhs_q @@ -284,14 +293,15 @@ def collective_gemm_bootstrap( this function with its own unique process_id. """ - assert ( - num_devices_per_process == 1 and jax.local_device_count() == 1 - ), "Only single device per process is supported at the moment!" - assert num_total_devices % num_devices_per_process == 0, ( - f"Invalid num_total_devices={num_total_devices}," - f" num_devices_per_process={num_devices_per_process}" - ) - assert 0 <= process_id < num_total_devices, f"Invalid process_id={process_id}" + if not (num_devices_per_process == 1 and jax.local_device_count() == 1): + raise RuntimeError("Only single device per process is supported at the moment!") + if num_total_devices % num_devices_per_process != 0: + raise ValueError( + f"Invalid num_total_devices={num_total_devices}," + f" num_devices_per_process={num_devices_per_process}" + ) + if not 0 <= process_id < num_total_devices: + raise ValueError(f"Invalid process_id={process_id}") initialize_cgemm_communicator( num_total_devices, num_devices_per_process, @@ -390,10 +400,11 @@ def assert_cublas_requirements(scaling_mode, contracting_size, tensor_name): # Requirements from https://docs.nvidia.com/cuda/cublas/#tensor-core-usage alignment = 32 if scaling_mode.is_nvfp4_scaling else 16 - assert contracting_size % alignment == 0, ( - f"cuBLAS GEMM {tensor_name} tensor's contracting dimension must be a multiple of" - f" {alignment} when using quantized inputs. Got contracting_size={contracting_size}" - ) + if contracting_size % alignment != 0: + raise ValueError( + f"cuBLAS GEMM {tensor_name} tensor's contracting dimension must be a multiple of" + f" {alignment} when using quantized inputs. Got contracting_size={contracting_size}" + ) class GemmPrimitive(BasePrimitive): @@ -439,57 +450,63 @@ def _dims_are_consecutive(dims): lhs_contracting_dims, rhs_contracting_dims, ) = map(sanitize_dims, operand_ndims, contracting_dims) - assert _dims_are_consecutive(lhs_contracting_dims), ( - "cuBLAS GEMM expected consecutive contracting dimensions for LHS operand, but got " - f"{lhs_contracting_dims}." - ) - assert _dims_are_consecutive(rhs_contracting_dims), ( - "cuBLAS GEMM expected consecutive contracting dimensions for RHS operand, but got " - f"{rhs_contracting_dims}." - ) + if not _dims_are_consecutive(lhs_contracting_dims): + raise ValueError( + "cuBLAS GEMM expected consecutive contracting dimensions for LHS operand, but got " + f"{lhs_contracting_dims}." + ) + if not _dims_are_consecutive(rhs_contracting_dims): + raise ValueError( + "cuBLAS GEMM expected consecutive contracting dimensions for RHS operand, but got " + f"{rhs_contracting_dims}." + ) lhs_contracting_size, rhs_contracting_size = map( lambda shape, dims: reduce(operator.mul, [shape[dim] for dim in dims]), (lhs.shape, rhs.shape), (lhs_contracting_dims, rhs_contracting_dims), ) - assert lhs_contracting_size == rhs_contracting_size, ( - "cuBLAS GEMM operands have incompatible contracting dimensions: " - f"{lhs.shape} @ idx {lhs_contracting_dims} X {rhs.shape} @ idx {rhs_contracting_dims}." - ) + if lhs_contracting_size != rhs_contracting_size: + raise ValueError( + f"cuBLAS GEMM operands have incompatible contracting dimensions: {lhs.shape} @ idx" + f" {lhs_contracting_dims} X {rhs.shape} @ idx {rhs_contracting_dims}." + ) assert_cublas_requirements(scaling_mode, lhs_contracting_size, "LHS") assert_cublas_requirements(scaling_mode, rhs_contracting_size, "RHS") lhs_is_transposed, rhs_is_transposed = _get_gemm_layout(operand_ndims, contracting_dims) if scaling_mode != ScalingMode.NO_SCALING: - assert scaling_mode.is_nvfp4_scaling or _compatible_fp8_gemm_dtypes( - lhs.dtype, rhs.dtype - ), ( - "cuBLAS GEMM quantized operands have incompatible data types: " - f"{lhs.dtype} x {rhs.dtype}." - ) - assert ( - lhs_scale_inv.size > 0 and rhs_scale_inv.size > 0 - ), "Quantized cuBLAS GEMM requires inverse scaling factors for both operands." + if not ( + scaling_mode.is_nvfp4_scaling or _compatible_fp8_gemm_dtypes(lhs.dtype, rhs.dtype) + ): + raise ValueError( + "cuBLAS GEMM quantized operands have incompatible data types: " + f"{lhs.dtype} x {rhs.dtype}." + ) + if not (lhs_scale_inv.size > 0 and rhs_scale_inv.size > 0): + raise ValueError( + "Quantized cuBLAS GEMM requires inverse scaling factors for both operands." + ) if ( scaling_mode != ScalingMode.MXFP8_1D_SCALING and not is_fp8_gemm_with_all_layouts_supported() ): - assert not lhs_is_transposed and rhs_is_transposed, ( - "cuBLAS FP8 GEMM on devices with compute capability < 10.0 (Hopper) " - "require non-transposed LHS and transposed RHS operands " - "(`contracting_dims=((-1, ), (-1, ))`)." - ) + if lhs_is_transposed or not rhs_is_transposed: + raise ValueError( + "cuBLAS FP8 GEMM on devices with compute capability < 10.0 (Hopper) " + "require non-transposed LHS and transposed RHS operands " + "(`contracting_dims=((-1, ), (-1, ))`)." + ) else: - assert lhs.dtype == rhs.dtype, ( - "For TE cuBLAS GEMM for non-quantized inputs, the operand dtypes must be equal." - f" LHS dtype != RHS dtype, lhs.dtype={lhs.dtype}, rhs.dtype={rhs.dtype}" - ) + if lhs.dtype != rhs.dtype: + raise ValueError( + "For TE cuBLAS GEMM for non-quantized inputs, the operand dtypes must be equal." + f" LHS dtype != RHS dtype, lhs.dtype={lhs.dtype}, rhs.dtype={rhs.dtype}" + ) # Determine output shape and dtype - assert ( - dtypes.canonicalize_dtype(out_dtype).itemsize > 1 - ), "cuBLAS GEMM custom op does not support 8-bit quantized output types." + if not dtypes.canonicalize_dtype(out_dtype).itemsize > 1: + raise ValueError("cuBLAS GEMM custom op does not support 8-bit quantized output types.") lhs_non_contracting_shape, rhs_non_contracting_shape = map( lambda shape, dims: [shape[dim] for dim in range(len(shape)) if dim not in dims], (lhs.shape, rhs.shape), @@ -500,7 +517,8 @@ def _dims_are_consecutive(dims): # Adjust output shape for comm+GEMM overlap if not collective_op.is_none and not is_outer: # Inner abstract - assert sequence_dim == 1, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + if sequence_dim != 1: + raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") overlap_out_shape = list(out_shape).copy() if collective_op.is_all_gather: overlap_out_shape[1] *= tpsp_axis_size() @@ -508,23 +526,34 @@ def _dims_are_consecutive(dims): overlap_out_shape[sequence_dim] = ( overlap_out_shape[sequence_dim] // tpsp_axis_size() ) - assert out_dtype == jnp.bfloat16, f"Unsupported out_dtype={out_dtype}" + if out_dtype != jnp.bfloat16: + raise ValueError(f"Unsupported out_dtype={out_dtype}") output = jax.core.ShapedArray(shape=overlap_out_shape, dtype=out_dtype) # Validate bias when present (bias.size > 0 means fuse bias) if bias.size > 0: - assert bias.shape == tuple(rhs_non_contracting_shape), ( - "cuBLAS GEMM bias tensor has incorrect shape, " - f"expected ({tuple(rhs_non_contracting_shape)}, ) but found {bias.shape}." + if bias.shape != tuple(rhs_non_contracting_shape): + raise ValueError( + "cuBLAS GEMM bias tensor has incorrect shape, " + f"expected ({tuple(rhs_non_contracting_shape)}, ) but found {bias.shape}." + ) + if bias.dtype != out_dtype: + raise ValueError( + "cuBLAS GEMM bias tensor has incorrect data type, " + f"expected {out_dtype} but found {bias.dtype}." + ) + + if alpha.size != 1 or alpha.dtype != jnp.float32: + raise ValueError( + f"Expected alpha to be a single float32 scalar, but got alpha.size={alpha.size}," + f" alpha.dtype={alpha.dtype}" ) - assert bias.dtype == out_dtype, ( - "cuBLAS GEMM bias tensor has incorrect data type, " - f"expected {out_dtype} but found {bias.dtype}." + if beta.size != 1 or beta.dtype != jnp.float32: + raise ValueError( + f"Expected beta to be a single float32 scalar, but got beta.size={beta.size}," + f" beta.dtype={beta.dtype}" ) - assert alpha.size == 1 and alpha.dtype == jnp.float32 - assert beta.size == 1 and beta.dtype == jnp.float32 - # Declare cuBLAS workspace workspace_size = get_cublas_workspace_size_bytes() # NVFP4 swizzling happen in via nvte kernel instead of JAX transposes @@ -629,16 +658,19 @@ def impl( and not is_outer and not lhs.shape[0] == 1 ): - assert sequence_dim == 1, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + if sequence_dim != 1: + raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") original_shape = lhs.shape - assert original_shape[0] % dp_or_fsdp_axis_size() == 0 or original_shape[0] == 1, ( - f"Original_shape[0]={original_shape[0]} is not divisible by" - f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" - ) - assert original_shape[1] % tpsp_axis_size() == 0 or original_shape[1] == 1, ( - f"Original_shape[1]={original_shape[1]} is not divisible by" - f" tpsp_axis_size()={tpsp_axis_size()}" - ) + if original_shape[0] % dp_or_fsdp_axis_size() != 0 and original_shape[0] != 1: + raise ValueError( + f"Original_shape[0]={original_shape[0]} is not divisible by" + f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" + ) + if original_shape[1] % tpsp_axis_size() != 0 and original_shape[1] != 1: + raise ValueError( + f"Original_shape[1]={original_shape[1]} is not divisible by" + f" tpsp_axis_size()={tpsp_axis_size()}" + ) reshaped = lhs.reshape( dp_or_fsdp_axis_size(), int(original_shape[0] / dp_or_fsdp_axis_size()), @@ -673,16 +705,19 @@ def impl( and not is_outer and not output.shape[0] == 1 ): - assert sequence_dim == 1, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + if sequence_dim != 1: + raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") original_shape = output.shape - assert original_shape[0] % dp_or_fsdp_axis_size() == 0 or original_shape[0] == 1, ( - f"Original_shape[0]={original_shape[0]} is not divisible by" - f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" - ) - assert original_shape[1] % tpsp_axis_size() == 0 or original_shape[1] == 1, ( - f"Original_shape[1]={original_shape[1]} is not divisible by" - f" tpsp_axis_size()={tpsp_axis_size()}" - ) + if original_shape[0] % dp_or_fsdp_axis_size() != 0 and original_shape[0] != 1: + raise ValueError( + f"Original_shape[0]={original_shape[0]} is not divisible by" + f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" + ) + if original_shape[1] % tpsp_axis_size() != 0 and original_shape[1] != 1: + raise ValueError( + f"Original_shape[1]={original_shape[1]} is not divisible by" + f" tpsp_axis_size()={tpsp_axis_size()}" + ) reshaped = output.reshape( tpsp_axis_size(), dp_or_fsdp_axis_size(), @@ -745,13 +780,15 @@ def batcher( is_outer, ): del transpose_batch_sequence, sequence_dim, is_outer - assert GemmPrimitive.outer_primitive is not None + if GemmPrimitive.outer_primitive is None: + raise RuntimeError("GemmPrimitive.outer_primitive has not been registered") lhs_bdims, _, rhs_bdims, *_ = batch_dims # Batched GEMM is not supported - assert ( - lhs_bdims is None and rhs_bdims is None - ), f"(Batching is not supported, got lhs_bdims={lhs_bdims}, rhs_bdims={rhs_bdims})" + if not (lhs_bdims is None and rhs_bdims is None): + raise RuntimeError( + f"Batching is not supported, got lhs_bdims={lhs_bdims}, rhs_bdims={rhs_bdims}" + ) out_bdims = (None,) return ( @@ -806,7 +843,8 @@ def _parse_operand_output_specs( for l in lhs_cspecs: for r in rhs_cspecs: if l is not None and l == r: - assert reduce_spec is None, "Multiple reduce dimension is detected!" + if reduce_spec is not None: + raise RuntimeError("Multiple reduce dimension is detected!") reduce_spec = l sequence_dim = None @@ -822,18 +860,20 @@ def _parse_operand_output_specs( " Please check your sharding configuration." ) from exc sequence_dim = tpsp_idx - assert (sequence_dim == 1) ^ transpose_batch_sequence, ( - "CollectiveGEMM supports only (sequence_dim=1 and transpose_batch_sequence=False)" - " or (sequence_dim=0 and transpose_batch_sequence=True). Received:" - f" sequence_dim={sequence_dim}," - f" transpose_batch_sequence={transpose_batch_sequence}." - ) + if not (sequence_dim == 1) ^ transpose_batch_sequence: + raise ValueError( + "CollectiveGEMM supports only (sequence_dim=1 and" + " transpose_batch_sequence=False) or (sequence_dim=0 and" + f" transpose_batch_sequence=True). Received: sequence_dim={sequence_dim}," + f" transpose_batch_sequence={transpose_batch_sequence}." + ) elif collective_op.is_reduce_scatter: - assert reduce_spec == gsr.tpsp_resource, ( - "Only CollectiveGemm RS with the Reduction over the TPSP axis is supported! Got" - f" reduce_spec={reduce_spec}, tpsp_resource={gsr.tpsp_resource}" - ) + if reduce_spec != gsr.tpsp_resource: + raise ValueError( + "Only CollectiveGemm RS with the Reduction over the TPSP axis is supported! Got" + f" reduce_spec={reduce_spec}, tpsp_resource={gsr.tpsp_resource}" + ) sequence_dim = int(not transpose_batch_sequence) if reduce_spec is not None: @@ -886,14 +926,18 @@ def _parse_operand_output_specs( # Only do AG Sequence dim if not Overlap RS if collective_op.is_all_gather: - assert sequence_dim <= len( - lhs_non_cspecs - ), f"Sequence dim {sequence_dim} is out of bounds for lhs_non_cspecs: {lhs_non_cspecs}" + if sequence_dim > len(lhs_non_cspecs): + raise ValueError( + f"Sequence dim {sequence_dim} is out of bounds for lhs_non_cspecs:" + f" {lhs_non_cspecs}" + ) out_specs = out_specs[:sequence_dim] + (None,) + out_specs[sequence_dim + 1 :] elif collective_op.is_reduce_scatter: - assert sequence_dim <= len( - lhs_non_cspecs - ), f"Sequence dim {sequence_dim} is out of bounds for lhs_non_cspecs: {lhs_non_cspecs}" + if sequence_dim > len(lhs_non_cspecs): + raise ValueError( + f"Sequence dim {sequence_dim} is out of bounds for lhs_non_cspecs:" + f" {lhs_non_cspecs}" + ) out_specs = ( out_specs[:sequence_dim] + (gsr.tpsp_resource,) + out_specs[sequence_dim + 1 :] ) @@ -912,7 +956,8 @@ def _parse_operand_output_specs( bias_specs = rhs_non_cspecs if arg_infos[4].size > 0 else (None,) # bias is operand index 4 if not collective_op.is_none: - assert sequence_dim >= 0, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + if sequence_dim < 0: + raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") return ( (lhs_specs, rhs_specs, bias_specs), @@ -1154,10 +1199,11 @@ def _te_gemm( lhs_amax = rhs_amax = None # Extract GEMM custom op inputs from quantized operands if isinstance(lhs_q, ScaledTensor): - assert isinstance(rhs_q, ScaledTensor) or rhs_quantizer is not None, ( - "cuBLAS GEMM with quantized LHS and non-quantized RHS operands requires a valid " - "`Quantizer` object to quantize the RHS operand." - ) + if not isinstance(rhs_q, ScaledTensor) and rhs_quantizer is None: + raise ValueError( + "cuBLAS GEMM with quantized LHS and non-quantized RHS operands requires a valid " + "`Quantizer` object to quantize the RHS operand." + ) if isinstance(lhs_q, ScaledTensor2x): # Choose the quantization of the contracting dimension(s) lhs_q = lhs_q.get_colwise_tensor() if lhs_is_transposed else lhs_q.get_rowwise_tensor() @@ -1169,21 +1215,23 @@ def _te_gemm( lhs_amax = lhs_q.amax if isinstance(rhs_q, ScaledTensor): - assert isinstance(lhs_q, ScaledTensor) or lhs_quantizer is not None, ( - "cuBLAS GEMM with non-quantized LHS and quantized RHS operands requires a valid " - "`Quantizer` object to quantize the LHS operand." - ) + if not isinstance(lhs_q, ScaledTensor) and lhs_quantizer is None: + raise ValueError( + "cuBLAS GEMM with non-quantized LHS and quantized RHS operands requires a valid " + "`Quantizer` object to quantize the LHS operand." + ) if isinstance(rhs_q, ScaledTensor2x): # Choose the quantization of the contracting dimension(s) rhs_q = rhs_q.get_rowwise_tensor() if rhs_is_transposed else rhs_q.get_colwise_tensor() - assert ( + if not ( rhs_q.scaling_mode == lhs_q.scaling_mode or rhs_q.scaling_mode.is_nvfp4_scaling and lhs_q.scaling_mode.is_nvfp4_scaling - ), ( - "cuBLAS GEMM quantized operands have mismatched scaling types, " - f"LHS:{lhs_q.scaling_mode} x RHS:{rhs_q.scaling_mode}." - ) + ): + raise ValueError( + "cuBLAS GEMM quantized operands have mismatched scaling types, " + f"LHS:{lhs_q.scaling_mode} x RHS:{rhs_q.scaling_mode}." + ) rhs_data = rhs_q.data rhs_scale_inv = rhs_q.scale_inv if rhs_q.data_layout == "T": @@ -1193,7 +1241,8 @@ def _te_gemm( alpha = jnp.ones((1,), jnp.float32) beta = jnp.zeros((1,), jnp.float32) if scaling_mode.is_nvfp4_scaling: - assert lhs_amax is not None and rhs_amax is not None + if lhs_amax is None or rhs_amax is None: + raise ValueError("NVFP4 scaling requires non-None amax for both LHS and RHS operands") lhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(lhs_amax) rhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(rhs_amax) alpha = lhs_tensor_scale_inv * rhs_tensor_scale_inv @@ -1268,7 +1317,10 @@ def impl( group_sizes, num_gemms, ): - assert GroupedGemmCopySizesPrimitive.inner_primitive is not None + if GroupedGemmCopySizesPrimitive.inner_primitive is None: + raise RuntimeError( + "GroupedGemmCopySizesPrimitive.inner_primitive has not been registered" + ) out = GroupedGemmCopySizesPrimitive.inner_primitive.bind( group_sizes, num_gemms=num_gemms, @@ -1372,23 +1424,20 @@ def abstract( shape=(int64_workspace_size,), dtype=jnp.uint8 ) - assert len(additional_args) == 2, ( - "Expected additional_args to contain alpha, beta for the graph-safe grouped GEMM" - f" primitive, but got {len(additional_args)} arguments." - ) + if len(additional_args) != 2: + raise ValueError( + "Expected additional_args to contain alpha, beta for the graph-safe grouped" + f" GEMM primitive, but got {len(additional_args)} arguments." + ) alpha_aval, beta_aval = additional_args - assert alpha_aval.shape == ( - num_groups, - ), f"Expected alpha shape {(num_groups,)}, got {alpha_aval.shape}" - assert ( - alpha_aval.dtype == jnp.float32 - ), f"Expected alpha dtype float32, got {alpha_aval.dtype}" - assert beta_aval.shape == ( - num_groups, - ), f"Expected beta shape {(num_groups,)}, got {beta_aval.shape}" - assert ( - beta_aval.dtype == jnp.float32 - ), f"Expected beta dtype float32, got {beta_aval.dtype}" + if alpha_aval.shape != (num_groups,): + raise ValueError(f"Expected alpha shape {(num_groups,)}, got {alpha_aval.shape}") + if alpha_aval.dtype != jnp.float32: + raise ValueError(f"Expected alpha dtype float32, got {alpha_aval.dtype}") + if beta_aval.shape != (num_groups,): + raise ValueError(f"Expected beta shape {(num_groups,)}, got {beta_aval.shape}") + if beta_aval.dtype != jnp.float32: + raise ValueError(f"Expected beta dtype float32, got {beta_aval.dtype}") return (out_aval, cublas_workspace_aval, setup_workspace_aval, int64_workspace_aval) @@ -1498,7 +1547,8 @@ def impl( use_async_d2h_group_sizes, use_v2_ffi, ): - assert GroupedGemmPrimitive.inner_primitive is not None + if GroupedGemmPrimitive.inner_primitive is None: + raise RuntimeError("GroupedGemmPrimitive.inner_primitive has not been registered") if use_v2_ffi: additional_args = (additional_arg_0, additional_arg_1) else: @@ -1586,30 +1636,37 @@ def _jax_scaled_matmul( """ JAX GEMM for MXFP8 via scaled_matmul """ - assert rhs.scaling_mode in ( + if rhs.scaling_mode not in ( ScalingMode.MXFP8_1D_SCALING, ScalingMode.NVFP4_1D_SCALING, ScalingMode.NVFP4_2D_SCALING, - ), f"rhs does not have MXFP8 or NVFP4 scaling mode, got rhs.scaling_mode={rhs.scaling_mode}" + ): + raise ValueError( + "rhs does not have MXFP8 or NVFP4 scaling mode, got" + f" rhs.scaling_mode={rhs.scaling_mode}" + ) (lhs_contract, rhs_contract), (lhs_batch, rhs_batch) = dim_nums expected_lhs_is_colwise = lhs_contract[-1] != lhs.data.ndim - 1 expected_rhs_is_colwise = rhs_contract[-1] != rhs.data.ndim - 1 - assert lhs.is_colwise is expected_lhs_is_colwise, ( - f"LHS with unexpected quantize dimension.\nExpect is_colwise={expected_lhs_is_colwise}, got" - f" {lhs.is_colwise}" - ) - assert rhs.is_colwise is expected_rhs_is_colwise, ( - f"RHS with unexpected quantize dimension.\nExpect is_colwise={expected_rhs_is_colwise}, got" - f" {rhs.is_colwise}" - ) + if lhs.is_colwise is not expected_lhs_is_colwise: + raise ValueError( + f"LHS with unexpected quantize dimension.\nExpect is_colwise={expected_lhs_is_colwise}," + f" got {lhs.is_colwise}" + ) + if rhs.is_colwise is not expected_rhs_is_colwise: + raise ValueError( + f"RHS with unexpected quantize dimension.\nExpect is_colwise={expected_rhs_is_colwise}," + f" got {rhs.is_colwise}" + ) if lhs.scaling_mode == ScalingMode.MXFP8_1D_SCALING: out_dtype = lhs.dq_dtype - assert ( - lhs.data_layout == "N" and rhs.data_layout == "N" - ), f"Got lhs.data_layout={lhs.data_layout}, rhs.data_layout={rhs.data_layout}" + if not (lhs.data_layout == "N" and rhs.data_layout == "N"): + raise ValueError( + f"Got lhs.data_layout={lhs.data_layout}, rhs.data_layout={rhs.data_layout}" + ) else: if lhs.data_layout == "T": lhs_contract = transpose_dims( @@ -1641,7 +1698,8 @@ def _jax_scaled_matmul( lhs_3d, rhs_3d, lhs_scale_3d, rhs_scale_3d, preferred_element_type=out_dtype ) if lhs.scaling_mode.is_nvfp4_scaling: - assert lhs.amax is not None and rhs.amax is not None + if lhs.amax is None or rhs.amax is None: + raise ValueError("NVFP4 scaling requires non-None amax for both LHS and RHS operands") lhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(lhs.amax) rhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(rhs.amax) alpha = lhs_tensor_scale_inv * rhs_tensor_scale_inv @@ -1674,9 +1732,10 @@ def _jax_gemm( def _jax_gemm_impl(lhs, rhs): if lhs.scaling_mode.is_tensor_scaling(): - assert ( - rhs.scaling_mode == lhs.scaling_mode - ), f"rhs.scaling_mode={rhs.scaling_mode} != lhs.scaling_mode={lhs.scaling_mode}" + if rhs.scaling_mode != lhs.scaling_mode: + raise ValueError( + f"rhs.scaling_mode={rhs.scaling_mode} != lhs.scaling_mode={lhs.scaling_mode}" + ) precision = ( jax.lax.Precision.HIGHEST if use_split_accumulator else jax.lax.Precision.DEFAULT @@ -1760,7 +1819,8 @@ def gemm( # Fall back on a native JAX implementation when the custom call to cuBLAS GEMM is disabled if not GemmPrimitive.enabled(): - assert collective_op.is_none, "JAX GEMM does not support collective GEMM" + if not collective_op.is_none: + raise RuntimeError("JAX GEMM does not support collective GEMM") output = _jax_gemm( lhs, rhs, contracting_dims, lhs_quantizer, rhs_quantizer, use_split_accumulator ) @@ -1857,7 +1917,10 @@ def grouped_gemm( del precision if isinstance(lhs, jnp.ndarray): - assert isinstance(rhs, jnp.ndarray) + if not isinstance(rhs, jnp.ndarray): + raise TypeError( + f"Expected rhs to be jnp.ndarray when lhs is jnp.ndarray, but got type={type(rhs)}" + ) out_dtype = lhs.dtype lhs_shape = lhs.shape rhs_shape = rhs.shape @@ -1866,7 +1929,11 @@ def grouped_gemm( lhs_scale_inv = rhs_scale_inv = jnp.empty((0,), jnp.float32) scaling_mode = ScalingMode.NO_SCALING elif isinstance(lhs, GroupedScaledTensor1x): - assert isinstance(rhs, GroupedScaledTensor1x) + if not isinstance(rhs, GroupedScaledTensor1x): + raise TypeError( + "Expected rhs to be GroupedScaledTensor1x when lhs is GroupedScaledTensor1x, but" + f" got type={type(rhs)}" + ) out_dtype = lhs.dq_dtype lhs_shape = lhs.original_shape rhs_shape = rhs.original_shape @@ -1874,7 +1941,11 @@ def grouped_gemm( rhs_data = rhs.data lhs_scale_inv = lhs.scale_inv rhs_scale_inv = rhs.scale_inv - assert lhs.scaling_mode == rhs.scaling_mode + if lhs.scaling_mode != rhs.scaling_mode: + raise ValueError( + f"Mismatched scaling modes: lhs.scaling_mode={lhs.scaling_mode}," + f" rhs.scaling_mode={rhs.scaling_mode}" + ) scaling_mode = lhs.scaling_mode else: raise TypeError("Unsupported lhs type object!") @@ -1911,8 +1982,16 @@ def grouped_gemm( and not isinstance(rhs, ScaledTensor) and quantizer_set != noop_quantizer_set ): - assert isinstance(quantizer_set.x, GroupedQuantizer) - assert type(quantizer_set.x) is type(quantizer_set.kernel) + if not isinstance(quantizer_set.x, GroupedQuantizer): + raise TypeError( + "Expected quantizer_set.x to be GroupedQuantizer, but got" + f" type={type(quantizer_set.x)}" + ) + if type(quantizer_set.x) is not type(quantizer_set.kernel): + raise TypeError( + "Expected quantizer_set.x and quantizer_set.kernel to have the same type, but got" + f" {type(quantizer_set.x)} and {type(quantizer_set.kernel)}" + ) scaling_mode = quantizer_set.x.scaling_mode if ( quantizer_set.x.scaling_mode.is_tensor_scaling() @@ -1939,9 +2018,8 @@ def grouped_gemm( lhs_shape = lhs_q.original_shape rhs_shape = rhs_q.original_shape - assert not ( - lhs_data.dtype == jnp.float8_e5m2 and rhs_data.dtype == jnp.float8_e5m2 - ), "FP8 GEMM does not support E5M2 * E5M2" + if lhs_data.dtype == jnp.float8_e5m2 and rhs_data.dtype == jnp.float8_e5m2: + raise ValueError("FP8 GEMM does not support E5M2 * E5M2") # Only support FP8 GEMM with NT layout on Hopper and other earlier GPUs # thus additional transpose is required @@ -1954,12 +2032,10 @@ def grouped_gemm( rhs_layout_is_T = rhs_q.data_layout == "T" # we can't apply _shape_normalization on the grouped input # thus we need to ensure that lhs is in N and rhs is in T - assert ( - lhs_is_trans == lhs_layout_is_T - ), "lhs input must be transposed before calling grouped_gemm" - assert ( - not rhs_is_trans == rhs_layout_is_T - ), "rhs input must be transposed before calling grouped_gemm" + if lhs_is_trans != lhs_layout_is_T: + raise RuntimeError("lhs input must be transposed before calling grouped_gemm") + if (not rhs_is_trans) != rhs_layout_is_T: + raise RuntimeError("rhs input must be transposed before calling grouped_gemm") lhs_is_trans = False rhs_is_trans = True lhs_ndim = len(lhs_shape) @@ -1978,28 +2054,36 @@ def grouped_gemm( # Calling GroupedGEMM Custom Call K_lhs = math.prod(lhs_shape[i] for i in lhs_contract_dim) K_rhs = math.prod(rhs_shape[i] for i in rhs_contract_dim) - assert K_lhs == K_rhs + if K_lhs != K_rhs: + raise ValueError( + f"Mismatched contracting dimensions: K_lhs={K_lhs}, K_rhs={K_rhs} (from" + f" lhs_shape={lhs_shape}, rhs_shape={rhs_shape})" + ) M = math.prod(_calculate_remaining_shape(lhs_shape, lhs_contract_dim)) N = math.prod(_calculate_remaining_shape(rhs_shape, rhs_contract_dim)[1:]) # Exclude G if is_grouped_dense_wgrad: N = math.prod(_calculate_remaining_shape(rhs_shape, rhs_contract_dim)) else: - assert group_sizes.size == rhs_shape[0] + if group_sizes.size != rhs_shape[0]: + raise ValueError( + "Expected group_sizes.size == rhs_shape[0], but got" + f" group_sizes.size={group_sizes.size}, rhs_shape[0]={rhs_shape[0]}" + ) has_bias = bias is not None - if has_bias: - assert bias.shape == ( - group_sizes.size, - N, - ), f"bias shape {bias.shape} does not match expected shape {(group_sizes.size, N)}" + if has_bias and bias.shape != (group_sizes.size, N): + raise ValueError( + f"Expected bias.shape=({group_sizes.size}, {N}), but got bias.shape={bias.shape}" + ) bias = jnp.empty((), jnp.float32) if bias is None else bias - assert group_offset is None, ( - "group_offset is not supported yet and is instead computed" - " internally assuming contiguous grouping. Any padding is included in the group_sizes" - " and padded with zeros to not affect the result of the MoE block." - ) + if group_offset is not None: + raise RuntimeError( + "group_offset is not supported yet and is instead computed" + " internally assuming contiguous grouping. Any padding is included in the group_sizes" + " and padded with zeros to not affect the result of the MoE block." + ) use_v2_ffi = _can_use_v2_grouped_gemm(scaling_mode, lhs_data.dtype, has_bias) if use_v2_ffi: diff --git a/transformer_engine/jax/cpp_extensions/normalization.py b/transformer_engine/jax/cpp_extensions/normalization.py index 70fdf4c474..29292f946b 100644 --- a/transformer_engine/jax/cpp_extensions/normalization.py +++ b/transformer_engine/jax/cpp_extensions/normalization.py @@ -132,9 +132,17 @@ def abstract( ) x_dtype = dtypes.canonicalize_dtype(x_aval.dtype) - assert x_dtype in [jnp.float32, jnp.float16, jnp.bfloat16] - assert scale_aval is None or scale_aval.dtype == jnp.float32 - assert amax_aval is None or amax_aval.dtype == jnp.float32 + assert x_dtype in [ + jnp.float32, + jnp.float16, + jnp.bfloat16, + ], f"Unsupported x_dtype={x_dtype}, expected one of [float32, float16, bfloat16]" + assert ( + scale_aval is None or scale_aval.dtype == jnp.float32 + ), f"Expected scale_aval.dtype=float32, but got scale_aval.dtype={scale_aval.dtype}" + assert ( + amax_aval is None or amax_aval.dtype == jnp.float32 + ), f"Expected amax_aval.dtype=float32, but got amax_aval.dtype={amax_aval.dtype}" assert ( scaling_mode != ScalingMode.MXFP8_1D_SCALING.value @@ -159,7 +167,10 @@ def abstract( mu_rsigama_dtype = jnp.float32 if norm_type == NVTE_Norm_Type.LayerNorm: - assert gamma_aval.size == beta_aval.size + assert gamma_aval.size == beta_aval.size, ( + "Expected gamma_aval.size == beta_aval.size, but got" + f" gamma_aval.size={gamma_aval.size}, beta_aval.size={beta_aval.size}" + ) assert gamma_aval.dtype == beta_aval.dtype, ( f"gamma and beta should have the same dtype, but got {gamma_aval.dtype} and " f"{beta_aval.dtype}" @@ -265,18 +276,35 @@ def lowering( del out_dtype, scale_dtype, is_outer, amax_scope, transpose_batch_sequence x_aval, scale_aval, amax_aval, gamma_aval, beta_aval = ctx.avals_in - assert x_aval.dtype in [jnp.float32, jnp.float16, jnp.bfloat16] - assert scale_aval is None or scale_aval.dtype == jnp.float32 - assert amax_aval is None or amax_aval.dtype == jnp.float32 + assert x_aval.dtype in [ + jnp.float32, + jnp.float16, + jnp.bfloat16, + ], f"Unsupported x_aval.dtype={x_aval.dtype}, expected one of [float32, float16, bfloat16]" + assert ( + scale_aval is None or scale_aval.dtype == jnp.float32 + ), f"Expected scale_aval.dtype=float32, but got scale_aval.dtype={scale_aval.dtype}" + assert ( + amax_aval is None or amax_aval.dtype == jnp.float32 + ), f"Expected amax_aval.dtype=float32, but got amax_aval.dtype={amax_aval.dtype}" g_type = ir.RankedTensorType(gamma.type) g_shape = g_type.shape if norm_type == NVTE_Norm_Type.LayerNorm: - assert gamma_aval.dtype == beta_aval.dtype + assert gamma_aval.dtype == beta_aval.dtype, ( + "Expected gamma and beta to have the same dtype, but got" + f" gamma_aval.dtype={gamma_aval.dtype}, beta_aval.dtype={beta_aval.dtype}" + ) b_type = ir.RankedTensorType(beta.type) b_shape = b_type.shape - assert g_type == b_type - assert g_shape == b_shape + assert g_type == b_type, ( + f"Expected gamma and beta to have the same IR type, but got gamma_type={g_type}," + f" beta_type={b_type}" + ) + assert g_shape == b_shape, ( + f"Expected gamma and beta to have the same shape, but got gamma_shape={g_shape}," + f" beta_shape={b_shape}" + ) sm_margin = get_forward_sm_margin() return ffi.ffi_lowering( @@ -321,7 +349,9 @@ def impl( to describe implementation """ del is_outer - assert NormFwdPrimitive.inner_primitive is not None + assert ( + NormFwdPrimitive.inner_primitive is not None + ), "NormFwdPrimitive.inner_primitive has not been registered" ( out, colwise_out, @@ -391,7 +421,9 @@ def batcher( to describe batch rules for vmap """ check_valid_batch_dims(batch_dims) - assert NormFwdPrimitive.outer_primitive is not None + assert ( + NormFwdPrimitive.outer_primitive is not None + ), "NormFwdPrimitive.outer_primitive has not been registered" x, scale, amax, gamma, beta = batched_args x_bdim, scale_bdim, _, _, _ = batch_dims @@ -706,13 +738,26 @@ def abstract(dz_aval, x_aval, mu_aval, rsigma_aval, gamma_aval, norm_type, zero_ w_dtype = dtypes.canonicalize_dtype(gamma_aval.dtype) rsigma_dtype = dtypes.canonicalize_dtype(rsigma_aval.dtype) - assert dtypes.canonicalize_dtype(dz_aval.dtype) == w_dtype - assert dz_aval.shape == x_aval.shape + assert dtypes.canonicalize_dtype(dz_aval.dtype) == w_dtype, ( + f"Expected dz_aval.dtype={w_dtype} (matching gamma dtype), but got" + f" dz_aval.dtype={dtypes.canonicalize_dtype(dz_aval.dtype)}" + ) + assert dz_aval.shape == x_aval.shape, ( + f"Expected dz_aval.shape == x_aval.shape, but got dz_aval.shape={dz_aval.shape}," + f" x_aval.shape={x_aval.shape}" + ) if norm_type == NVTE_Norm_Type.LayerNorm: mu_dtype = dtypes.canonicalize_dtype(mu_aval.dtype) - assert mu_aval.shape == rsigma_aval.shape == x_aval.shape[:-1] - assert mu_dtype == rsigma_dtype == jnp.float32 + assert mu_aval.shape == rsigma_aval.shape == x_aval.shape[:-1], ( + "Expected mu_aval.shape == rsigma_aval.shape == x_aval.shape[:-1], but got" + f" mu_aval.shape={mu_aval.shape}, rsigma_aval.shape={rsigma_aval.shape}," + f" x_aval.shape[:-1]={x_aval.shape[:-1]}" + ) + assert mu_dtype == rsigma_dtype == jnp.float32, ( + f"Expected mu_dtype == rsigma_dtype == float32, but got mu_dtype={mu_dtype}," + f" rsigma_dtype={rsigma_dtype}" + ) dx_aval = dz_aval dgamma_aval = dbeta_aval = gamma_aval @@ -756,8 +801,14 @@ def lowering(ctx, dz, x, mu, rsigma, gamma, *, norm_type, zero_centered_gamma): g_shape = g_type.shape b_type = ir.RankedTensorType(gamma.type) b_shape = b_type.shape - assert g_type == b_type - assert g_shape == b_shape + assert g_type == b_type, ( + f"Expected gamma and beta to have the same IR type, but got gamma_type={g_type}," + f" beta_type={b_type}" + ) + assert g_shape == b_shape, ( + f"Expected gamma and beta to have the same shape, but got gamma_shape={g_shape}," + f" beta_shape={b_shape}" + ) sm_margin = get_backward_sm_margin() return ffi.ffi_lowering(NormBwdPrimitive.name)( @@ -774,7 +825,9 @@ def lowering(ctx, dz, x, mu, rsigma, gamma, *, norm_type, zero_centered_gamma): @staticmethod def impl(dz, x, mu, rsigma, gamma, norm_type, zero_centered_gamma): - assert NormBwdPrimitive.inner_primitive is not None + assert ( + NormBwdPrimitive.inner_primitive is not None + ), "NormBwdPrimitive.inner_primitive has not been registered" dx, dgamma, dbeta, _ = NormBwdPrimitive.inner_primitive.bind( dz, x, mu, rsigma, gamma, norm_type=norm_type, zero_centered_gamma=zero_centered_gamma ) @@ -783,7 +836,9 @@ def impl(dz, x, mu, rsigma, gamma, norm_type, zero_centered_gamma): @staticmethod def batcher(batched_args, batch_dims, *, norm_type, zero_centered_gamma): check_valid_batch_dims(batch_dims) - assert NormBwdPrimitive.outer_primitive is not None + assert ( + NormBwdPrimitive.outer_primitive is not None + ), "NormBwdPrimitive.outer_primitive has not been registered" dz, x, mu, rsigma, gamma = batched_args _, x_bdim, _, _, gamma_bdim = batch_dims diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 1fce6d2fd7..031ab483a0 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -115,7 +115,10 @@ def impl( score_function, compute_aux_scores, ): - assert FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive is not None + if FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive is None: + raise RuntimeError( + "FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive has not been registered" + ) return FusedTopkWithScoreFunctionFwdPrimitive.inner_primitive.bind( logits, expert_bias, @@ -141,7 +144,10 @@ def batcher( score_function, compute_aux_scores, ): - assert FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive is not None + if FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive is None: + raise RuntimeError( + "FusedTopkWithScoreFunctionFwdPrimitive.outer_primitive has not been registered" + ) logits, expert_bias = batched_args logits_bdim, _ = batch_dims return ( @@ -284,7 +290,10 @@ def impl( score_function, compute_aux_scores, ): - assert FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive is not None + if FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive is None: + raise RuntimeError( + "FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive has not been registered" + ) return FusedTopkWithScoreFunctionBwdPrimitive.inner_primitive.bind( routing_map, intermediate, @@ -307,7 +316,10 @@ def batcher( score_function, compute_aux_scores, ): - assert FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive is not None + if FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive is None: + raise RuntimeError( + "FusedTopkWithScoreFunctionBwdPrimitive.outer_primitive has not been registered" + ) routing_map, intermediate, grad_probs = batched_args _, _, grad_probs_bdim = batch_dims return ( @@ -402,7 +414,10 @@ def lowering(ctx, probs, tokens_per_expert, *, topk, coeff): @staticmethod def impl(probs, tokens_per_expert, topk, coeff): - assert FusedMoEAuxLossFwdPrimitive.inner_primitive is not None + if FusedMoEAuxLossFwdPrimitive.inner_primitive is None: + raise RuntimeError( + "FusedMoEAuxLossFwdPrimitive.inner_primitive has not been registered" + ) return FusedMoEAuxLossFwdPrimitive.inner_primitive.bind( probs, tokens_per_expert, @@ -412,7 +427,10 @@ def impl(probs, tokens_per_expert, topk, coeff): @staticmethod def batcher(batched_args, batch_dims, *, topk, coeff): - assert FusedMoEAuxLossFwdPrimitive.outer_primitive is not None + if FusedMoEAuxLossFwdPrimitive.outer_primitive is None: + raise RuntimeError( + "FusedMoEAuxLossFwdPrimitive.outer_primitive has not been registered" + ) probs, tokens_per_expert = batched_args probs_bdim, _ = batch_dims return ( @@ -490,7 +508,10 @@ def lowering(ctx, const_buf, tokens_per_expert, grad_aux_loss, *, num_tokens): @staticmethod def impl(const_buf, tokens_per_expert, grad_aux_loss, num_tokens): - assert FusedMoEAuxLossBwdPrimitive.inner_primitive is not None + if FusedMoEAuxLossBwdPrimitive.inner_primitive is None: + raise RuntimeError( + "FusedMoEAuxLossBwdPrimitive.inner_primitive has not been registered" + ) return FusedMoEAuxLossBwdPrimitive.inner_primitive.bind( const_buf, tokens_per_expert, @@ -500,7 +521,10 @@ def impl(const_buf, tokens_per_expert, grad_aux_loss, num_tokens): @staticmethod def batcher(batched_args, batch_dims, *, num_tokens): - assert FusedMoEAuxLossBwdPrimitive.outer_primitive is not None + if FusedMoEAuxLossBwdPrimitive.outer_primitive is None: + raise RuntimeError( + "FusedMoEAuxLossBwdPrimitive.outer_primitive has not been registered" + ) const_buf, tokens_per_expert, grad_aux_loss = batched_args _, _, grad_bdim = batch_dims return ( diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 7decfca6c6..31ce6e72e9 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -1445,7 +1445,8 @@ def te_dot_general(generate_quantizer_set, x, kernel, dims, **kwargs): def make_grouped_dense_cls(quantization_recipe): """Creates a grouped dense (grouped GEMM) instance for use with TE state module.""" - assert quantization_recipe is None, "Ragged dot grouped GEMM does not support quantization yet" + if quantization_recipe is not None: + raise ValueError("Ragged dot grouped GEMM does not support quantization yet") def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwargs): del kwargs # Unused diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index ad5a60e4c2..513677e4a1 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -182,7 +182,9 @@ def __call__( is_gqa = h_q != h_kv if is_gqa: - assert (h_q % h_kv == 0) and (h_q >= h_kv) + assert (h_q % h_kv == 0) and ( + h_q >= h_kv + ), f"num_query_heads ({h_q}) must be divisible by and >= num_kv_heads ({h_kv})" group_size = h_q // h_kv grouped_query = query.reshape((*query.shape[:2], h_kv, group_size, query.shape[-1])) @@ -428,7 +430,9 @@ def __call__( if self.transpose_batch_sequence: x = x.transpose([1, 0, 2, 3]) - assert x.dtype == query.dtype + assert ( + x.dtype == query.dtype + ), f"output dtype {x.dtype} does not match query dtype {query.dtype}" return x @@ -713,9 +717,13 @@ def __call__( del self.attn_bias_type, self.attn_mask_type, self.qkv_layout if attn_bias_type == AttnBiasType.NO_BIAS: - assert bias is None + assert ( + bias is None + ), f"bias must be None when attn_bias_type is NO_BIAS, but got bias={bias}" else: - assert bias is not None + assert ( + bias is not None + ), f"bias must not be None when attn_bias_type is {attn_bias_type}" bias = bias.astype(input_dtype) self._assert_dtypes(query, key, value, qkv_layout) @@ -823,11 +831,13 @@ def __call__( key, value = jnp.split(key, [1], axis=-3) key, value = map(functools.partial(jnp.squeeze, axis=-3), [key, value]) else: - assert qkv_layout.is_separate() + assert ( + qkv_layout.is_separate() + ), f"Expected separate qkv_layout, but got {qkv_layout}" assert sequence_descriptor is None or isinstance( sequence_descriptor, (jnp.ndarray, np.ndarray) - ) + ), f"sequence_descriptor must be None or ndarray, but got {type(sequence_descriptor)}" x = _UnfusedDotProductAttention( attention_dropout=self.attention_dropout, @@ -994,7 +1004,7 @@ def _canonicalize_lora_scope(scope): SCOPE_EX_QKV_PROJ, SCOPE_EX_OUTPUT_PROJ, SCOPE_EX_MLP, - ] + ], f"Unsupported LoRA scope: {scope}" lora_scope = LoRAScope() @@ -1307,8 +1317,10 @@ def query_init(*args): return self.kernel_init(*args) / (depth_scaling if self.scaled_query_init else 1.0) def qkv_init(key, shape, dtype): - assert len(shape) == 3 - assert shape[-2] == 3 + assert ( + len(shape) == 3 + ), f"qkv_init expects 3D shape, but got {len(shape)}D shape {shape}" + assert shape[-2] == 3, f"qkv_init expects shape[-2] == 3, but got shape={shape}" q_key, k_key, v_key = jax_random.split(key, num=3) @@ -1323,8 +1335,8 @@ def qkv_init(key, shape, dtype): return jnp.stack([q_kernel, k_kernel, v_kernel], axis=-2, dtype=dtype) def kv_init(key, shape, dtype): - assert len(shape) == 3 - assert shape[-2] == 2 + assert len(shape) == 3, f"kv_init expects 3D shape, but got {len(shape)}D shape {shape}" + assert shape[-2] == 2, f"kv_init expects shape[-2] == 2, but got shape={shape}" k_key, v_key = jax_random.split(key) @@ -1415,7 +1427,7 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): )(inputs_q) if is_self_attn: - assert ln_out is not None + assert ln_out is not None, "ln_out must not be None for self-attention" inputs_kv = ln_out kv_proj = DenseGeneral( @@ -1475,7 +1487,7 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): )(inputs_q) if is_self_attn: - assert ln_out is not None + assert ln_out is not None, "ln_out must not be None for self-attention" inputs_kv = ln_out query = query.astype(input_dtype) @@ -1494,7 +1506,9 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): elif qkv_layout == QKVLayout.BSHD_BS2HD: key, value = jnp.split(kv_proj, [1], axis=-2) else: - assert qkv_layout == QKVLayout.BSHD_BSHD_BSHD + assert ( + qkv_layout == QKVLayout.BSHD_BSHD_BSHD + ), f"Expected QKVLayout.BSHD_BSHD_BSHD, but got {qkv_layout}" # No changes to memory layout, should trigger bitcast only (Ideally no Perf impact) query = query.reshape((*query.shape[:2], self.num_attention_heads, self.head_dim)) @@ -1520,7 +1534,9 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): value = value.reshape((*value.shape[:2], self.num_gqa_groups, self.head_dim)) if decode: - assert qkv_layout == QKVLayout.BSHD_BSHD_BSHD + assert ( + qkv_layout == QKVLayout.BSHD_BSHD_BSHD + ), f"decode mode requires QKVLayout.BSHD_BSHD_BSHD, but got {qkv_layout}" is_initialized = self.has_variable("cache", "cached_key") cached_key = self.variable("cache", "cached_key", jnp.zeros, key.shape, key.dtype) @@ -1588,7 +1604,9 @@ def generate_batch_seqlen_logical_axes(is_sharded_seq): kv_proj = with_sharding_constraint_by_logical_axes(kv_proj, kv_sharding_constraint) dpa_args = [query, kv_proj, None] else: - assert qkv_layout == QKVLayout.BSHD_BSHD_BSHD + assert ( + qkv_layout == QKVLayout.BSHD_BSHD_BSHD + ), f"Expected QKVLayout.BSHD_BSHD_BSHD, but got {qkv_layout}" query = query.reshape((*query.shape[:2], self.num_attention_heads, self.head_dim)) key = key.reshape((*key.shape[:2], self.num_gqa_groups, self.head_dim)) value = value.reshape((*value.shape[:2], self.num_gqa_groups, self.head_dim)) @@ -2101,7 +2119,9 @@ def generate_batch_seqlen_logical_axes(is_shared_seq=None): l = inputs.shape[sequence_dim] attn_bias = rel_emb(l, l, False) - assert inputs.ndim == 3 + assert ( + inputs.ndim == 3 + ), f"inputs must be 3D (batch, sequence, hidden), but got {inputs.ndim}D" # Make name be the exactly same as T5X, since names would affect # RNGKey during init and apply. Myabe no need in the feature. @@ -2151,10 +2171,15 @@ def generate_batch_seqlen_logical_axes(is_shared_seq=None): )(inputs, inputs, attention_mask, attn_bias, deterministic=deterministic, decode=decode) def hidden_dropout(x, deterministic): - assert isinstance(self.hidden_dropout_dims, Sequence) + assert isinstance( + self.hidden_dropout_dims, Sequence + ), f"hidden_dropout_dims must be a Sequence, but got {type(self.hidden_dropout_dims)}" x_shape_len = len(x.shape) for dims in self.hidden_dropout_dims: - assert -x_shape_len <= dims < x_shape_len + assert -x_shape_len <= dims < x_shape_len, ( + f"hidden_dropout_dims value {dims} is out of range " + f"[{-x_shape_len}, {x_shape_len}) for input with {x_shape_len} dimensions" + ) return nn.Dropout( rate=self.hidden_dropout, @@ -2179,7 +2204,9 @@ def hidden_dropout(x, deterministic): )(x, deterministic=deterministic) if self.apply_residual_connection_post_layernorm: - assert ln_out is not None + assert ( + ln_out is not None + ), "ln_out must not be None when apply_residual_connection_post_layernorm is True" residual = ln_out x = x + residual @@ -2239,7 +2266,9 @@ def hidden_dropout(x, deterministic): y = hidden_dropout(y, deterministic) if self.apply_residual_connection_post_layernorm: - assert ln_out is not None + assert ( + ln_out is not None + ), "ln_out must not be None when apply_residual_connection_post_layernorm is True" residual = ln_out mlp_input = y + residual @@ -2284,7 +2313,9 @@ def hidden_dropout(x, deterministic): )(mlp_input, deterministic=deterministic) if self.apply_residual_connection_post_layernorm: - assert ln_out is not None + assert ( + ln_out is not None + ), "ln_out must not be None when apply_residual_connection_post_layernorm is True" residual = ln_out z = with_sharding_constraint_by_logical_axes( diff --git a/transformer_engine/jax/layernorm.py b/transformer_engine/jax/layernorm.py index 3f3f3802db..0f173a89e3 100644 --- a/transformer_engine/jax/layernorm.py +++ b/transformer_engine/jax/layernorm.py @@ -31,7 +31,11 @@ def canonicalize_norm_type(x): Canonicalized normalization type string """ canonicalized = x.lower().strip().replace("-", "").replace("_", "") - assert canonicalized in ["layernorm", "rmsnorm"] + if canonicalized not in ["layernorm", "rmsnorm"]: + raise ValueError( + f"Unsupported normalization type '{x}' (canonicalized: '{canonicalized}'). " + "Valid options are: 'layernorm', 'rmsnorm'." + ) return canonicalized diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index e9f64bb693..2de4576e05 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -271,14 +271,23 @@ def fused_attn_fwd( attn_scale = 1.0 / math.sqrt(d) if attn_bias_type not in ["no_bias", "alibi"]: - assert ( - attn_bias is not None - ), "attn_bias tensor cannot be None when attn_bias_type is not no_bias or alibi." - assert attn_bias.dtype == q.dtype, "attn_bias tensor must be in the same dtype as q and kv." - - assert ( - fused_attention_backend != FusedAttnBackend["No_Backend"] - ), "Fused attention does not support this input combination." + if attn_bias is None: + raise ValueError( + f"attn_bias tensor cannot be None when attn_bias_type={attn_bias_type!r}." + ) + if attn_bias.dtype != q.dtype: + raise ValueError( + "attn_bias tensor must have the same dtype as q and kv: " + f"attn_bias.dtype={attn_bias.dtype} but q.dtype={q.dtype}." + ) + + if fused_attention_backend == FusedAttnBackend["No_Backend"]: + raise ValueError( + "Fused attention does not support this input combination:" + f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," + f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," + 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"]: @@ -294,12 +303,16 @@ def fused_attn_fwd( max_seqlen_q * max_seqlen_q + BACKEND_F16m512_FP8_THREADS_PER_CTA - 1 ) // BACKEND_F16m512_FP8_THREADS_PER_CTA - assert ( - s_quantizer is not None - ), "s_quantizer is required as an input for FP8 fused attention." - assert ( - o_quantizer is not None - ), "o_quantizer is required as an input for FP8 fused attention." + 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}") @@ -488,28 +501,44 @@ def fused_attn_bwd( d = q.size(-1) attn_scale = 1.0 / math.sqrt(d) - assert ( - fused_attention_backend != FusedAttnBackend["No_Backend"] - ), "Fused attention does not support this input combination." + if fused_attention_backend == FusedAttnBackend["No_Backend"]: + raise ValueError( + "Fused attention backward does not support this input combination:" + f" qkv_layout={qkv_layout!r}, attn_bias_type={attn_bias_type!r}," + f" attn_mask_type={attn_mask_type!r}, q.shape={list(q.shape)}," + f" q.dtype={q.dtype}, backend={fused_attention_backend}." + ) if fused_attention_backend != FusedAttnBackend["F16_max512_seqlen"]: - assert ( - len(aux_ctx_tensors) >= 1 - ), "aux_ctx_tensors must contain rng_state as its last element." + 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 fused_attention_backend == FusedAttnBackend["FP8"]: - assert ( - s_quantizer is not None - ), "s_quantizer is required as an input for FP8 fused attention backward." - assert ( - dp_quantizer is not None - ), "dp_quantizer is required as an input for FP8 fused attention backward." - assert ( - dqkv_dtype is not None - ), "dqkv_dtype is required as an input for FP8 fused attention backward." - assert ( - len(aux_ctx_tensors) == 3 - ), "aux_ctx_tensors is required to be [M, ZInv, rng_state] for FP8 fused attention." + 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, diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index e7390de365..d0b314a64f 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -124,7 +124,11 @@ def tensor_group_process_after_reload(tensor_group: TensorGroup): """ Call for a tensor group, just after reload logic. """ - assert tensor_group.aux is not None + if tensor_group.aux is None: + raise RuntimeError( + "TensorGroup.aux must be set before post-reload processing, " + f"but got aux=None for tensor_group with {len(tensor_group.tensor_list)} tensors" + ) tensor_group = TensorGroupProcessor._restore_tensor_duplicates(tensor_group) tensor_group = TensorGroupProcessor._switch_to_views(tensor_group) return tensor_group @@ -158,9 +162,8 @@ def _check_if_offload_base_tensor(tensor: torch.Tensor) -> bool: if _check_if_offload_base_tensor(tensor): aux["views"].append((tensor.shape, tensor.stride(), tensor.storage_offset())) tensor = tensor._base - assert ( - tensor is not None - ), "Cannot offload base tensor, if the tensor is not a view." + if tensor is None: + raise RuntimeError("Cannot offload base tensor, if the tensor is not a view.") tensor_group.tensor_list[tensor_id] = tensor else: aux["views"].append(None) @@ -247,9 +250,10 @@ def __init__( self.state = "not_offloaded" def _validate_state(self, func_name: str, allowed_states: list[str]): - assert ( - self.state in allowed_states - ), f"Invalid state: {self.state} for {func_name}, must be one of {allowed_states}" + if self.state not in allowed_states: + raise RuntimeError( + f"Invalid state: {self.state} for {func_name}, must be one of {allowed_states}" + ) def start_offload(self): """ @@ -271,7 +275,12 @@ def start_offload(self): ) for tensor_id, tensor in enumerate(self.fwd_gpu_tensor_group.tensor_list): - assert tensor.is_contiguous() + if not tensor.is_contiguous(): + raise ValueError( + f"Tensor at index {tensor_id} must be contiguous for CPU offloading, " + f"but got non-contiguous tensor with shape={tensor.shape}, " + f"stride={tensor.stride()}, dtype={tensor.dtype}" + ) # Wait for the moment the tensor is ready to be offloaded. self.offload_stream.wait_event(self.fwd_gpu_tensor_group.events[tensor_id]) # type: ignore[arg-type] @@ -284,12 +293,13 @@ def start_offload(self): self.cpu_tensor_group.tensor_list.append(offloaded_tensor) else: offloaded_tensor = self.cpu_tensor_group.tensor_list[tensor_id] - assert offloaded_tensor.shape == tensor.shape, ( - "CPU buffer shape does not match the offloaded tensor shape:" - f" {offloaded_tensor.shape} != {tensor.shape} " - "Make sure that tensor shapes do not change between" - " iterations if retain_pinned_cpu_buffers is True." - ) + if offloaded_tensor.shape != tensor.shape: + raise ValueError( + "CPU buffer shape does not match the offloaded tensor shape:" + f" {offloaded_tensor.shape} != {tensor.shape} " + "Make sure that tensor shapes do not change between" + " iterations if retain_pinned_cpu_buffers is True." + ) offloaded_tensor.copy_(tensor, non_blocking=True) # aux is a dictionary that contains auxiliary data like information which tensors were deduplicated, @@ -420,7 +430,11 @@ def pop_tensor( return self.fwd_gpu_tensor_group.tensor_list[tensor_or_tensor_id] # 4. the layer was offloaded - assert self.state == "reload_started" + if self.state != "reload_started": + raise RuntimeError( + "Expected state='reload_started' when popping an offloaded tensor, " + f"but got state='{self.state}' for tensor={tensor_or_tensor_id}" + ) # wait for the tensor to be reloaded torch.cuda.current_stream().wait_event( self.bwd_gpu_tensor_group.events[tensor_or_tensor_id] @@ -824,18 +838,19 @@ def get_cpu_offload_context( raise RuntimeError("CPU offload is not supported in debug mode.") if not manual_synchronization: - assert ( - num_layers <= model_layers - 1 - ), "Cannot offload all layers without manual synchronization - last layer is not offloaded." + if num_layers > model_layers - 1: + raise ValueError( + "Cannot offload all layers without manual synchronization - last layer is not" + f" offloaded. Got num_layers={num_layers}, model_layers={model_layers}." + ) if num_layers == model_layers - 1: warnings.warn( "Offloading num_layers == model_layers - 1 is not recommended, it prevents" " overlapping of computation and offload/reload." ) - assert ( - offload_stream is None or manual_synchronization - ), "offload_stream can be provided only if manual_synchronization is True" + if offload_stream is not None and not manual_synchronization: + raise ValueError("offload_stream can be provided only if manual_synchronization is True") if manual_synchronization: offload_synchronizer = ManualOffloadSynchronizer( @@ -858,9 +873,10 @@ def __init__(self): self.inside_context = False def __enter__(self): - assert ( - self.inside_context is False - ), "Offloading context was entered without synchronization function being called." + if self.inside_context: + raise RuntimeError( + "Offloading context was entered without synchronization function being called." + ) self.inside_context = True self._hooks_ctx = saved_tensors_hooks( offload_synchronizer.push_tensor, offload_synchronizer.pop_tensor @@ -882,12 +898,23 @@ def synchronization_function(self, tensor): """ This function is used to catch the backward pass of the model. """ - assert tensor.requires_grad is True - assert self.current_layer is not None + if not tensor.requires_grad: + raise ValueError( + "Tensor passed to synchronization_function must require grad to " + "register backward hooks, but got requires_grad=False for tensor " + f"with shape={tensor.shape}, dtype={tensor.dtype}" + ) + if self.current_layer is None: + raise RuntimeError( + "synchronization_function called but no layer has been set via __enter__. " + f"inside_context={self.inside_context}, " + f"offload_synchronizer num_layers={self.offload_synchronizer.num_layers}" + ) cur_layer = self.current_layer - assert ( - self.inside_context is False - ), "Synchronization function was called without offloading context being entered." + if self.inside_context: + raise RuntimeError( + "Synchronization function was called without offloading context being entered." + ) def hook(_): # offload_synchronizer.finish_part_of_bwd needs diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index 645dbb48d2..b06f6f5619 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -272,7 +272,8 @@ at::Tensor allocateSpace(const NVTEShape& shape, const transformer_engine::DType } else if (size == 1) { return at::empty({static_cast(shape.data[0])}, at::CUDA(GetATenDType(type))); } - NVTE_CHECK(false, "Should never reach here! func: allocateSpace"); + NVTE_ERROR("Unsupported tensor allocation: ndim=", size, ", init_to_zeros=", init_to_zeros, + ". Only 1D and 2D tensors are supported."); } at::Tensor allocateTorchTensor(int M, int N, transformer_engine::DType dtype) { diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 0214f7ff71..7e13cc105f 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2184,7 +2184,8 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // Compute amax. if (this->with_rht) { if (input.dtype() != DType::kBFloat16) { - NVTE_CHECK(false, "RHT is only supported for bfloat16 input"); + NVTE_ERROR("RHT is only supported for bfloat16 input, got dtype enum value ", + static_cast(input.dtype())); } if (this->with_post_rht_amax) { // We need: @@ -2196,7 +2197,9 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou }); } else { // raise error since it's not supported yet - NVTE_CHECK(false, "Pre-RHT amax is not supported yet"); + NVTE_ERROR( + "Pre-RHT amax is not supported yet. " + "Use with_post_rht_amax=true instead."); } } else { // Without RHT if (compute_amax) { diff --git a/transformer_engine/pytorch/custom_recipes/gemm.py b/transformer_engine/pytorch/custom_recipes/gemm.py index 8f853ff093..3d1e1cc43e 100644 --- a/transformer_engine/pytorch/custom_recipes/gemm.py +++ b/transformer_engine/pytorch/custom_recipes/gemm.py @@ -32,7 +32,8 @@ def custom_gemm( grad: bool = False, ) -> Iterable[Optional[torch.Tensor]]: """Dispatch GEMM to quantizer's qgemm method.""" - assert is_custom(A) and is_custom(B), "A and B must be custom tensors" + if not (is_custom(A) and is_custom(B)): + raise TypeError("A and B must be custom tensors") A, B = B, A @@ -68,11 +69,16 @@ def custom_gemm( if gemm_type == GEMMType.FPROP: qx, sx = A.data, A.scale qw, sw = B.data, B.scale - assert qx is not None - assert sx is not None - assert qw is not None - assert sw is not None - assert A.original_shape is not None + if qx is None: + raise ValueError("FPROP GEMM: quantized activation data (A.data) is None") + if sx is None: + raise ValueError("FPROP GEMM: activation scale (A.scale) is None") + if qw is None: + raise ValueError("FPROP GEMM: quantized weight data (B.data) is None") + if sw is None: + raise ValueError("FPROP GEMM: weight scale (B.scale) is None") + if A.original_shape is None: + raise ValueError("FPROP GEMM: A.original_shape is None, cannot determine output shape") # Call quantizer's qgemm method result = quantizer.qgemm( @@ -95,10 +101,14 @@ def custom_gemm( elif gemm_type == GEMMType.DGRAD: qdy, sdy = A.data, A.scale qw_t, sw_t = B.data_t, B.scale_t - assert qdy is not None - assert sdy is not None - assert qw_t is not None - assert sw_t is not None + if qdy is None: + raise ValueError("DGRAD GEMM: quantized gradient data (A.data) is None") + if sdy is None: + raise ValueError("DGRAD GEMM: gradient scale (A.scale) is None") + if qw_t is None: + raise ValueError("DGRAD GEMM: transposed quantized weight data (B.data_t) is None") + if sw_t is None: + raise ValueError("DGRAD GEMM: transposed weight scale (B.scale_t) is None") result = quantizer.qgemm( qdy, @@ -115,10 +125,14 @@ def custom_gemm( elif gemm_type == GEMMType.WGRAD: qdy_t, sdy_t = A.data_t, A.scale_t qx_t, sx_t = B.data_t, B.scale_t - assert qdy_t is not None - assert sdy_t is not None - assert qx_t is not None - assert sx_t is not None + if qdy_t is None: + raise ValueError("WGRAD GEMM: transposed quantized gradient data (A.data_t) is None") + if sdy_t is None: + raise ValueError("WGRAD GEMM: transposed gradient scale (A.scale_t) is None") + if qx_t is None: + raise ValueError("WGRAD GEMM: transposed quantized activation data (B.data_t) is None") + if sx_t is None: + raise ValueError("WGRAD GEMM: transposed activation scale (B.scale_t) is None") result = quantizer.qgemm( qdy_t, diff --git a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py index d00d0c8b94..f42183ec09 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py @@ -169,7 +169,8 @@ def high_precision_gemm_ref( y_shape = (mat1.size(0), mat2.size(1)) if bias is not None: - assert not accumulate, "Bias is not supported with accumulation" + if accumulate: + raise ValueError("Bias is not supported with accumulation") bias = bias.to(out_dtype) # With bias case if out_dtype == torch.float32: @@ -325,7 +326,8 @@ def size(self, *args, **kwargs): # pylint: disable=unused-argument the second dimension by half. This method returns the logical shape that users expect, not the internal packed storage shape. """ - assert self.original_shape is not None + if self.original_shape is None: + raise RuntimeError("NVFP4TensorRef.size() called but original_shape has not been set") return torch.Size(self.original_shape) @@ -374,7 +376,8 @@ def _build_hadamard_matrix( Uses Sylvester construction to avoid SciPy dependency. """ - assert (size & (size - 1)) == 0, "Hadamard size must be a power of two" + if (size & (size - 1)) != 0: + raise ValueError(f"Hadamard size must be a power of two, got {size}") h = torch.ones((1, 1), device=device, dtype=torch.float32) while h.shape[0] < size: h = torch.cat( @@ -402,9 +405,10 @@ def _apply_rht(self, x: torch.Tensor) -> torch.Tensor: # RHT dimension equals the quantization tile length (NVFP4 uses 16) rht_dim = self.quant_tile_shape[1] - assert ( - x.shape[-1] % rht_dim == 0 - ), f"Inner dimension {x.shape[-1]} must be divisible by hadamard dimension {rht_dim}" + if x.shape[-1] % rht_dim != 0: + raise ValueError( + f"Inner dimension {x.shape[-1]} must be divisible by hadamard dimension {rht_dim}" + ) # Build H and scale H = self._build_hadamard_matrix(rht_dim, x.device, x.dtype, self.with_random_sign_mask) @@ -446,7 +450,11 @@ def _quantize_blockwise_reference( eps: float, # pylint: disable=unused-argument ) -> Tuple[torch.Tensor, torch.Tensor]: - assert x.ndim == 2 + if x.ndim != 2: + raise ValueError( + f"_quantize_blockwise_reference expects a 2D tensor, got {x.ndim}D with shape" + f" {x.shape}" + ) using_2d_quantization = tile_len_x == 16 and tile_len_y == 16 m, n = x.shape # Compute vec_max based on the original x (before reshape) @@ -525,7 +533,11 @@ def _pad_tensor( tensor: torch.Tensor, row_divisor: Optional[int], col_divisor: Optional[int] ) -> torch.Tensor: - assert tensor.dim() == 2, "only supports 2D tensors" + if tensor.dim() != 2: + raise ValueError( + f"_pad_tensor only supports 2D tensors, got {tensor.dim()}D tensor with shape" + f" {tensor.shape}" + ) M, N = tensor.shape padding_needed_rows = 0 padding_needed_cols = 0 @@ -553,7 +565,11 @@ def _pad_tensor( @staticmethod def _rm_pad_tensor(tensor: torch.Tensor, original_size: tuple[int, ...]) -> torch.Tensor: - assert tensor.dim() == 2, "only supports 2D tensors" + if tensor.dim() != 2: + raise ValueError( + f"_rm_pad_tensor only supports 2D tensors, got {tensor.dim()}D tensor with shape" + f" {tensor.shape}" + ) M, N = original_size out = tensor[:M, :N].contiguous() return out @@ -584,19 +600,20 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ - sx_t: scale tensor for qx_t (if columnwise_usage), None otherwise - global_amax_row, global_amax_col: global amax tensors """ + global_amax_col = None if self.pow_2_scales: - assert self.quant_tile_shape == ( - 1, - 32, - ), "MXFP4 only supports 1x32 tile shape." + if self.quant_tile_shape != (1, 32): + raise ValueError( + f"MXFP4 only supports 1x32 tile shape, got {self.quant_tile_shape}" + ) # 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) else: - assert self.quant_tile_shape in ( - (1, 16), - (16, 16), - ), "NVFP4 only supports 1x16 or 16x16 tile shape." + if self.quant_tile_shape not in ((1, 16), (16, 16)): + raise ValueError( + f"NVFP4 only supports 1x16 or 16x16 tile shape, got {self.quant_tile_shape}" + ) # Prepare inputs once so we can reuse for both amax and quantization # Row-input will always be the original input. row_input = tensor @@ -670,7 +687,11 @@ def quantize( **kwargs, # pylint: disable=unused-argument ) -> NVFP4TensorRef: # sanity checks - assert tensor.dtype in utils.HIGH_PRECISION_FLOAT_DTYPES, "Unsupported input dtype." + if tensor.dtype not in utils.HIGH_PRECISION_FLOAT_DTYPES: + raise TypeError( + f"Unsupported input dtype {tensor.dtype}, expected one of" + f" {utils.HIGH_PRECISION_FLOAT_DTYPES}" + ) # Make it work with 3D tensors original_shape = tensor.shape @@ -766,7 +787,10 @@ def is_data_t_transposed_in_memory(self) -> bool: TODO(etsykunov): Confirm docstring is correct. """ - raise NotImplementedError("Not implemented yet") + raise NotImplementedError( + "NVFP4QuantizerRef.is_data_t_transposed_in_memory is not implemented for FP4" + " quantization" + ) def qgemm( self, @@ -784,7 +808,8 @@ def qgemm( qresult_w: QuantizedTensorStorage | None = None, ) -> torch.Tensor: """Python implementation of microblock FP4 GEMM.""" - assert bias is None, "Bias is implemented for FP4 GEMM." + if bias is not None: + raise ValueError("Bias is not supported in NVFP4QuantizerRef.qgemm") high_precision_x = cast_from_fp4x2(qx, out_dtype) high_precision_w = cast_from_fp4x2(qw, out_dtype) @@ -814,11 +839,22 @@ def qgemm( else: - assert qresult_x is not None - assert qresult_w is not None - - assert qresult_x.global_amax_row is not None - assert qresult_w.global_amax_col is not None + if qresult_x is None: + raise ValueError( + "qresult_x is required for non-pow_2_scales NVFP4 GEMM (needed for global_amax)" + ) + if qresult_w is None: + raise ValueError( + "qresult_w is required for non-pow_2_scales NVFP4 GEMM (needed for global_amax)" + ) + if qresult_x.global_amax_row is None: + raise ValueError( + "qresult_x.global_amax_row must be set for non-pow_2_scales NVFP4 GEMM" + ) + if qresult_w.global_amax_col is None: + raise ValueError( + "qresult_w.global_amax_col must be set for non-pow_2_scales NVFP4 GEMM" + ) sx = sx.to(torch.float32) sw = sw.to(torch.float32) @@ -833,23 +869,27 @@ def qgemm( M, K = high_precision_x.shape N, K_w = high_precision_w.shape - assert K == K_w, "K dimension mismatch between qx and qw" - - assert K % 32 == 0, "K dimension must be divisible by 32" - assert N % 8 == 0, "N dimension must be divisible by 8" + if K != K_w: + raise ValueError( + f"K dimension mismatch between qx and qw: qx has K={K}, qw has K={K_w}" + ) + if K % 32 != 0: + raise ValueError(f"K dimension must be divisible by 32, got K={K}") + if N % 8 != 0: + raise ValueError(f"N dimension must be divisible by 8, got N={N}") block_length = 32 if self.pow_2_scales else 16 grid_k = K // block_length - assert sx.shape == ( - M, - K // block_length, - ), f"sx shape mismatch: expected ({M}, {K//block_length}), got {sx.shape}" - assert sw.shape == ( - N, - K // block_length, - ), f"sw shape mismatch: expected ({N}, {K//block_length}), got {sw.shape}" + if sx.shape != (M, K // block_length): + raise ValueError( + f"sx shape mismatch: expected ({M}, {K // block_length}), got {sx.shape}" + ) + if sw.shape != (N, K // block_length): + raise ValueError( + f"sw shape mismatch: expected ({N}, {K // block_length}), got {sw.shape}" + ) y = torch.zeros(M, N, dtype=torch.float32, device=qx.device) @@ -878,10 +918,12 @@ def qgemm( # accumulation happens at epilogue in float32 if accumulate: - assert out is not None, "Output tensor must be provided for accumulation." + if out is None: + raise ValueError("Output tensor must be provided for accumulation.") y += out.to(torch.float32) else: - assert out is None, "Output tensor should be None when accumulate is False." + if out is not None: + raise ValueError("Output tensor should be None when accumulate is False.") y = y.to(out_dtype) return y diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 2a65fa272b..03ac9c1595 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -152,7 +152,11 @@ def set_tensor_model_parallel_attributes( ) -> None: """set attributes needed for TP""" for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS: - assert not hasattr(tensor, attribute) + if hasattr(tensor, attribute): + raise RuntimeError( + f"Tensor already has attribute '{attribute}' set. Cannot set " + "tensor model parallel attributes on a tensor that already has them." + ) # Set the attributes. setattr(tensor, "tensor_model_parallel", is_parallel) setattr(tensor, "partition_dim", dim) @@ -170,7 +174,11 @@ def get_distributed_world_size(group: Optional[dist_group_type] = None) -> int: @lru_cache def get_distributed_rank(group: Optional[dist_group_type] = None) -> int: """Return my rank for the distributed group.""" - assert torch.distributed.is_initialized(), "torch.distributed is not initialized." + if not torch.distributed.is_initialized(): + raise RuntimeError( + "torch.distributed is not initialized. Call torch.distributed.init_process_group() " + "before calling get_distributed_rank()." + ) return torch.distributed.get_rank(group=group) @@ -743,7 +751,12 @@ def checkpoint( # If saved activations need to be distributed but there is no process group, # default to the world group. if distribute_saved_activations: - assert torch.distributed.is_initialized(), "torch.distributed is not initialized." + if not torch.distributed.is_initialized(): + raise RuntimeError( + "torch.distributed is not initialized. Call " + "torch.distributed.init_process_group() before using " + "distribute_saved_activations=True." + ) tp_group = torch.distributed.GroupMember.WORLD if tp_group is None else tp_group return _CheckpointFunction.apply( @@ -917,9 +930,12 @@ def reduce_scatter_along_first_dim( return inp, None dim_size = list(inp.size()) - assert ( - dim_size[0] % world_size == 0 - ), "First dimension of the tensor should be divisible by tensor parallel size" + if dim_size[0] % world_size != 0: + raise ValueError( + "First dimension of the tensor should be divisible by tensor parallel size, " + f"but got dim_size[0]={dim_size[0]} and world_size={world_size} " + f"(remainder={dim_size[0] % world_size})." + ) dim_size[0] = dim_size[0] // world_size @@ -984,7 +1000,11 @@ def _all_gather_fp8( # Note: We cannot directly all-gather the transposed FP8 tensor, # so temporarily modify quantizer to avoid creating FP8 transpose. if not isinstance(inp, Float8TensorStorage): - assert isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)) + if not isinstance(quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)): + raise TypeError( + "Expected quantizer to be Float8Quantizer or Float8CurrentScalingQuantizer " + f"when input is not Float8TensorStorage, but got {type(quantizer).__name__}." + ) # we cannot directly gather the transposed fp8 tensor # so we need to disable columnwise usage for the quantizer # and then set it back to the original value after quantizing @@ -1234,10 +1254,18 @@ def _swap_first_dims(tensor: torch.Tensor, world_size: int): """ shape = tensor.shape - assert len(shape) >= 2, "Wrong number of dimensions for fixing interleave." + if len(shape) < 2: + raise ValueError( + f"Wrong number of dimensions for fixing interleave: got {len(shape)}, " + f"expected at least 2 (shape={shape})." + ) first_dim = shape[0] flattened_trailing = math.prod(shape[1:]) - assert first_dim % world_size == 0, "Wrong dimensions for fixing interleave." + if first_dim % world_size != 0: + raise ValueError( + f"Wrong dimensions for fixing interleave: first_dim={first_dim} is not divisible " + f"by world_size={world_size} (remainder={first_dim % world_size})." + ) tensor = tensor.reshape(world_size, first_dim // world_size, flattened_trailing) tensor = tex.swap_first_dims(tensor, out=None) return tensor.reshape(first_dim // world_size, flattened_trailing * world_size) @@ -1327,7 +1355,11 @@ def _all_gather_nvfp4( f"found {inp.__class__.__name__})" ) - assert in_shape is not None or in_shape_t is not None, "No data found." + if in_shape is None and in_shape_t is None: + raise ValueError( + "No data found: both in_shape and in_shape_t are None. " + "Input tensor must have rowwise or columnwise data." + ) world_size = get_distributed_world_size(process_group) @@ -1380,7 +1412,11 @@ def _all_gather_nvfp4( if quantizer.rowwise_usage: # Remove padding from NVFP4 scale-inverses - assert in_shape is not None, "Shape not found." + if in_shape is None: + raise RuntimeError( + "Shape not found: in_shape is None but rowwise_usage is True. " + "Input tensor must have rowwise data for NVFP4 rowwise gathering." + ) in_scale_inv = inp._rowwise_scale_inv out_scale_inv = out._rowwise_scale_inv flattened_in_shape0 = math.prod(in_shape[:-1]) @@ -1681,7 +1717,10 @@ def gather_along_first_dim( # MXFP8 case if isinstance(inp, MXFP8TensorStorage) or isinstance(quantizer, MXFP8Quantizer): - assert isinstance(quantizer, MXFP8Quantizer) + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError( + f"Expected MXFP8Quantizer for MXFP8 all-gather, but got {type(quantizer).__name__}." + ) return _all_gather_mxfp8( inp, process_group, @@ -1692,7 +1731,10 @@ def gather_along_first_dim( # NVFP4 case if isinstance(inp, NVFP4TensorStorage) or isinstance(quantizer, NVFP4Quantizer): - assert isinstance(quantizer, NVFP4Quantizer) + if not isinstance(quantizer, NVFP4Quantizer): + raise TypeError( + f"Expected NVFP4Quantizer for NVFP4 all-gather, but got {type(quantizer).__name__}." + ) return _all_gather_nvfp4( inp, process_group, @@ -1835,8 +1877,15 @@ def symmetric_all_reduce( - The second element is the async work handle if async_op=True, otherwise None. """ - assert async_op is False, "Async symmetric ops no supported yet" - assert HAS_TORCH_SYMMETRIC, "Could not import symetric memory from torch" + if async_op: + raise RuntimeError( + f"Async symmetric ops are not supported yet, but async_op={async_op!r} was passed." + ) + if not HAS_TORCH_SYMMETRIC: + raise RuntimeError( + "Could not import symmetric memory from torch. " + "Please ensure torch.distributed._symmetric_memory is available." + ) if get_distributed_world_size(tp_group) == 1: return inp, None @@ -1969,10 +2018,19 @@ def _fsdp_gather_tensors( *tensors: torch.Tensor, ): if fsdp_group is not None: - assert len(shapes) == len(tensors), "Number of tensors and tensor shapes must be equal." + if len(shapes) != len(tensors): + raise ValueError( + "Number of tensors and tensor shapes must be equal, " + f"but got {len(shapes)} shapes and {len(tensors)} tensors." + ) for s, t in zip(shapes, tensors): if isinstance(t, torch.Tensor): - assert s is not None, "Internal TE error." + if s is None: + raise RuntimeError( + "Internal TE error: shape is None for a non-None tensor in " + "post_optimizer_step_fwd_amax_reduction. " + f"Tensor type: {type(t).__name__}, tensor shape: {t.shape}." + ) targets = t.get_data_tensors() if isinstance(t, QuantizedTensor) else [t] for target in targets: safely_set_viewless_tensor_data( @@ -2020,17 +2078,23 @@ def prepare_te_modules_for_fsdp(fsdp_root: torch.nn.Module) -> None: fsdp_root : torch.nn.Module FSDP-wrapped root module that may contain FSDP-wrapped TE modules. """ - assert isinstance(fsdp_root, FSDP), "Root module must be FSDP-wrapped." + if not isinstance(fsdp_root, FSDP): + raise TypeError(f"Root module must be FSDP-wrapped, but got {type(fsdp_root).__name__}.") # If the root module is a TE module, inject FSDP information into it if _is_te_module(fsdp_root.module): if hasattr(fsdp_root, "primary_weights_in_fp8"): - assert not fsdp_root.primary_weights_in_fp8, ( - "TE modules with primary weights in FP8 cannot be FSDP-wrapped. " - "Please initialize your model without the te.quantized_model_init(...) context." - ) + if fsdp_root.primary_weights_in_fp8: + raise RuntimeError( + "TE modules with primary weights in FP8 cannot be FSDP-wrapped. " + "Please initialize your model without the te.quantized_model_init(...) context." + ) root_state = _get_module_fsdp_state(fsdp_root) - assert root_state is not None, "Root module does not have a valid _FSDPState." + if root_state is None: + raise RuntimeError( + f"Root module ({type(fsdp_root.module).__name__}) does not have a valid " + "_FSDPState. Ensure the module is properly wrapped with FSDP." + ) fsdp_root.module.fast_setattr("fsdp_group", root_state.process_group) # Iterate through all FSDP-wrapped submodules and inject FSDP information into TE modules @@ -2038,10 +2102,12 @@ def prepare_te_modules_for_fsdp(fsdp_root: torch.nn.Module) -> None: for state, fsdp_module in zip(fsdp_states, fsdp_modules): if _is_te_module(fsdp_module.module): if hasattr(fsdp_module.module, "primary_weights_in_fp8"): - assert not fsdp_module.module.primary_weights_in_fp8, ( - "TE modules with primary weights in FP8 cannot be FSDP-wrapped. " - "Please initialize your model without the te.quantized_model_init(...) context." - ) + if fsdp_module.module.primary_weights_in_fp8: + raise RuntimeError( + f"TE module '{type(fsdp_module.module).__name__}' with primary weights " + "in FP8 cannot be FSDP-wrapped. Please initialize your model without " + "the te.quantized_model_init(...) context." + ) fsdp_module.module.fast_setattr("fsdp_group", state.process_group) diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index bae911b4e1..86b8a4acf4 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -139,7 +139,7 @@ def _make_graphed_callables( # Check training/inference is_training = all(c.training for c in callables) if not is_training and any(c.training for c in callables): - assert False, ( + raise RuntimeError( "make_graphed_callables only supports when modules are all in training or all in" " inference mode." ) @@ -148,8 +148,16 @@ def _make_graphed_callables( _order_without_wgrad = None delay_wgrad_compute = False if _order is None: - assert len(sample_args) == len(callables) - assert len(sample_kwargs) == len(callables) + if len(sample_args) != len(callables): + raise ValueError( + "Expected sample_args to have the same length as callables, " + f"but got {len(sample_args)} sample_args for {len(callables)} callables" + ) + if len(sample_kwargs) != len(callables): + raise ValueError( + "Expected sample_kwargs to have the same length as callables, " + f"but got {len(sample_kwargs)} sample_kwargs for {len(callables)} callables" + ) else: # Custom logic for interleaved pipeline parallelism # Note: This is tightly coupled with the Megatron-core @@ -173,48 +181,62 @@ def _make_graphed_callables( _order_without_wgrad.append(c_id) num_model_chunks = max(_order_without_wgrad) num_microbatches = len(_order_without_wgrad) // num_model_chunks // 2 - assert num_model_chunks * num_microbatches * 2 == len(_order_without_wgrad) + if num_model_chunks * num_microbatches * 2 != len(_order_without_wgrad): + raise ValueError( + f"Pipeline-parallel order dimension mismatch: num_model_chunks ({num_model_chunks})" + f" * num_microbatches ({num_microbatches}) * 2 =" + f" {num_model_chunks * num_microbatches * 2}, but len(_order_without_wgrad) =" + f" {len(_order_without_wgrad)}" + ) # When delay_wgrad_compute is enabled, each layer is treated as a model chunk, which # allows for fine-grained graph capture order. if delay_wgrad_compute: - assert ( - _num_layers_per_chunk is not None - ), "'_num_layers_per_chunk' must be provided when delay_wgrad_compute is True." + if _num_layers_per_chunk is None: + raise ValueError( + "'_num_layers_per_chunk' must be provided when delay_wgrad_compute is True." + ) for num_layers in _num_layers_per_chunk: - assert ( - num_layers == 1 - ), "Each model chunk must have only one layer when delay_wgrad_compute is True." + if num_layers != 1: + raise ValueError( + "Each model chunk must have only one layer when delay_wgrad_compute is" + f" True, but got {num_layers} layers." + ) # Determine number of layers in each model chunk. if _num_layers_per_chunk is None: - assert len(sample_args) * 2 >= len(_order_without_wgrad) and ( - len(sample_args) * 2 % len(_order_without_wgrad) == 0 - ), ( - f"{len(sample_args)} * 2 >= {len(_order_without_wgrad)} and {len(sample_args)} * 2" - f" % {len(_order_without_wgrad)} == 0" - ) + if not ( + len(sample_args) * 2 >= len(_order_without_wgrad) + and (len(sample_args) * 2 % len(_order_without_wgrad) == 0) + ): + raise ValueError( + f"{len(sample_args)} * 2 >= {len(_order_without_wgrad)} and" + f" {len(sample_args)} * 2 % {len(_order_without_wgrad)} == 0" + ) num_layers = len(sample_args) // num_model_chunks // num_microbatches _num_layers_per_chunk = [num_layers] * num_model_chunks else: - assert ( + if not ( isinstance(_num_layers_per_chunk, int) or len(_num_layers_per_chunk) == num_model_chunks - ), ( - "If _num_layers_per_chunk is provided, it must be an integer or a list of" - f" {num_model_chunks} integers, but got {_num_layers_per_chunk}." - ) + ): + raise ValueError( + "If _num_layers_per_chunk is provided, it must be an integer or a list of" + f" {num_model_chunks} integers, but got {_num_layers_per_chunk}." + ) if isinstance(_num_layers_per_chunk, int): _num_layers_per_chunk = [_num_layers_per_chunk] * num_model_chunks total_num_layers = sum(_num_layers_per_chunk) - assert len(callables) == total_num_layers, ( - f"Callables should have ({total_num_layers}) " - + f"entries when order input is provided but got {len(callables)}." - ) - assert len(sample_args) == total_num_layers * num_microbatches, ( - f"Expected {total_num_layers * num_microbatches} " - + f"args tuple, but got {len(sample_args)}." - ) + if len(callables) != total_num_layers: + raise ValueError( + f"Callables should have ({total_num_layers}) " + + f"entries when order input is provided but got {len(callables)}." + ) + if len(sample_args) != total_num_layers * num_microbatches: + raise ValueError( + f"Expected {total_num_layers * num_microbatches} " + + f"args tuple, but got {len(sample_args)}." + ) # Calculate the starting index of each chunk in callables for future use. _prefix_num_layers = [0] @@ -222,19 +244,26 @@ def _make_graphed_callables( num_layers = _num_layers_per_chunk[m_chunk] _prefix_num_layers.append(_prefix_num_layers[-1] + num_layers) - assert len(sample_kwargs) == len(sample_args) + if len(sample_kwargs) != len(sample_args): + raise ValueError( + "Pipeline-parallel schedule requires sample_kwargs and sample_args to have " + f"the same length, but got {len(sample_kwargs)} sample_kwargs " + f"for {len(sample_args)} sample_args" + ) # Check reuse graph conditions and reorganize sample_args and sample_kwargs. # Note: When capturing a graph, we hold onto the args and kwargs so we have static buffers # when the graph is replayed. If two model chunk microbatches have no overlap between their # forward and backward, then we can reduce memory usage by reusing the same static buffers. if _reuse_graph_input_output_buffers: - assert ( - _order is not None - ), "`_order` must be provided when `_reuse_graph_input_output_buffers` is True." - assert ( - is_training - ), "`_reuse_graph_input_output_buffers` is only available in training mode." + if _order is None: + raise ValueError( + "`_order` must be provided when `_reuse_graph_input_output_buffers` is True." + ) + if not is_training: + raise RuntimeError( + "`_reuse_graph_input_output_buffers` is only available in training mode." + ) if isinstance(sample_args, tuple): sample_args = list(sample_args) if isinstance(sample_kwargs, tuple): @@ -300,20 +329,22 @@ def _make_graphed_callables( # Check callables for c in callables: if isinstance(c, torch.nn.Module): - assert ( + if not ( len(c._backward_hooks) == 0 and len(c._forward_hooks) == 0 and len(c._forward_pre_hooks) == 0 - ), ( - "Modules must not have hooks registered at the time they are passed. " - + "However, registering hooks on modules after passing them " - + "through make_graphed_callables is allowed." - ) - assert all(b.requires_grad is False for b in c.buffers()), ( - "In any :class:`~torch.nn.Module` passed to " - + ":func:`~make_graphed_callables`, only parameters may be trainable. " - + "All buffers must have ``requires_grad=False``." - ) + ): + raise RuntimeError( + "Modules must not have hooks registered at the time they are passed. " + + "However, registering hooks on modules after passing them " + + "through make_graphed_callables is allowed." + ) + if not all(b.requires_grad is False for b in c.buffers()): + raise RuntimeError( + "In any :class:`~torch.nn.Module` passed to " + + ":func:`~make_graphed_callables`, only parameters may be trainable. " + + "All buffers must have ``requires_grad=False``." + ) # Flatten callable arguments per_callable_kwargs_keys = [list(kwargs.keys()) for kwargs in sample_kwargs] @@ -322,10 +353,11 @@ def _make_graphed_callables( flatten_arg, _ = _tree_flatten(args) flatten_kwarg, _ = _tree_flatten([kwargs[key] for key in kwargs_keys]) flatten_sample_args.append(tuple(flatten_arg + flatten_kwarg)) - assert all(isinstance(arg, torch.Tensor) for arg in flatten_arg), ( - "In the beta API, sample_args " - + "for each callable must contain only Tensors. Other types are not allowed." - ) + if not all(isinstance(arg, torch.Tensor) for arg in flatten_arg): + raise TypeError( + "In the beta API, sample_args " + + "for each callable must contain only Tensors. Other types are not allowed." + ) # If a callable is an nn.Module, its graph's full input surface is the args the user explicitly # passes to forward (ie, its sample_args) AND the module's parameter attributes. @@ -354,7 +386,12 @@ def _make_graphed_callables( ) else () ) - assert len(per_callable_module_params) == len(flatten_sample_args) + if len(per_callable_module_params) != len(flatten_sample_args): + raise ValueError( + "Pipeline-parallel dimension mismatch: " + f"per_callable_module_params has {len(per_callable_module_params)} entries, " + f"but flatten_sample_args has {len(flatten_sample_args)} entries" + ) per_callable_static_input_surfaces = [ flatten_sample_args[i] + per_callable_module_params[i] for i in range(len(flatten_sample_args)) @@ -400,12 +437,12 @@ def _make_graphed_callables( warmup_func_idx.append(func_idx) warmup_func.append(func) fwd_idx[m_chunk] += 1 - assert len(warmup_func) == len( - sample_args - ), f"Warmup runs {len(warmup_func)} don't match args {len(sample_args)}." - assert len(warmup_func_idx) == len( - set(warmup_func_idx) - ), f"Warmup runs {len(warmup_func)} but only {len(set(warmup_func_idx))} are unique." + if len(warmup_func) != len(sample_args): + raise ValueError(f"Warmup runs {len(warmup_func)} don't match args {len(sample_args)}.") + if len(warmup_func_idx) != len(set(warmup_func_idx)): + raise RuntimeError( + f"Warmup runs {len(warmup_func)} but only {len(set(warmup_func_idx))} are unique." + ) # Filter the TE modules that cudagraph can access. visited_te_modules = {} @@ -429,9 +466,10 @@ def hook_fn( modules.add(module) # If forward is called on a te.ops.Sequential it is not called on its constituent ops elif isinstance(module, Sequential): - assert ( - module._module_groups is not None - ), "Should have been initialized by warmup" + if module._module_groups is None: + raise RuntimeError( + "module._module_groups should have been initialized by warmup" + ) for module_group in module._module_groups: if isinstance(module_group, OperationFuser): for basic_op in module_group._basic_ops: @@ -480,20 +518,22 @@ def hook_fn( grad_inputs[grad_inputs_idx] is None and grad_inputs_idx < num_required_grad_sample_args ): - assert allow_unused_input, ( - "The input tensor requires grad, but the grad is None after" - " backward pass." - ) + if not allow_unused_input: + raise RuntimeError( + "The input tensor requires grad, but the grad is None after" + " backward pass." + ) elif ( grad_inputs[grad_inputs_idx] is not None and grad_inputs_idx >= num_required_grad_sample_args ): module_params_with_grad.append(static_input_surface[inputs_idx]) if len(module_params_with_grad) != len(per_callable_module_params[func_idx]): - assert warmup_iter == 0, ( - "no-grad params should only be used as inputs in the first warmup" - " iteration" - ) + if warmup_iter != 0: + raise RuntimeError( + "no-grad params should only be used as inputs in the first warmup" + f" iteration, but found in iteration {warmup_iter}" + ) per_callable_module_params[func_idx] = tuple(module_params_with_grad) static_input_surface = flatten_sample_args[func_idx] + tuple( module_params_with_grad @@ -531,7 +571,10 @@ def hook_fn( previous_chunk_last_callable_bwd_idx = None for i, c_id in enumerate(_order): if c_id > 0: - assert isinstance(c_id, int), "Forward order value must be an integer." + if not isinstance(c_id, int): + raise TypeError( + f"Forward order value must be an integer, but got {type(c_id).__name__}." + ) # Capture forward graph for model chunk c_id, microbatch fwd_idx[c_id-1] m_chunk = c_id - 1 for l_no in range(_num_layers_per_chunk[m_chunk]): @@ -583,23 +626,27 @@ def hook_fn( break if wgrad_validation_list[i] is None: wgrad_validation_list[i] = False - assert wgrad_validation_list[i], ( - f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number " - f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}." - ) + if not wgrad_validation_list[i]: + raise RuntimeError( + f"Number of wgrad graph({num_wgrad_c_id}) doesn't match number " + f"of dgrad graphs ({len(same_bwd_c_id_list)}) for chunk {c_id}." + ) elif ceil(c_id) != c_id: per_callable_bwd_idx -= _num_layers_per_chunk[m_chunk] - assert is_training, "Only training mode supports backward_dw." + if not is_training: + raise RuntimeError("Only training mode supports backward_dw.") # If no one module needs the backward_dw, the bwd_dw_graph will be empty. # So skip capturing it. For backward_dw, the order value is c_id - 0.5 to indicate # the specific order of backward_dw. - assert ceil(c_id) - c_id == 0.5, ( - "The order diff of wgrad and dgrad must be 0.5, " - f"get {ceil(c_id) - c_id}." - ) - assert need_bwd_dw_graph[ - per_callable_bwd_idx - ], "No module needs wgrad computation but get float in order" + if ceil(c_id) - c_id != 0.5: + raise ValueError( + "The order diff of wgrad and dgrad must be 0.5, " + f"get {ceil(c_id) - c_id}." + ) + if not need_bwd_dw_graph[per_callable_bwd_idx]: + raise RuntimeError( + "No module needs wgrad computation but get float in order" + ) bwd_dw_graph = bwd_dw_graphs[per_callable_bwd_idx] with _graph_context_wrapper(bwd_dw_graph, pool=mempool): for module in visited_te_modules[per_callable_bwd_idx]: @@ -811,7 +858,11 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i torch.cuda.current_stream().wait_stream(cuda_graph_stream) else: fwd_graph.replay() - assert isinstance(static_outputs, tuple) + if not isinstance(static_outputs, tuple): + raise TypeError( + "Expected static_outputs to be a tuple, but got" + f" {type(static_outputs).__name__}" + ) return tuple(o.detach() if o is not None else o for o in static_outputs) @staticmethod @@ -820,7 +871,12 @@ def backward(ctx, *grads): # pylint: disable=missing-function-docstring # Replay backward graph - assert len(grads) == len(static_grad_outputs) + if len(grads) != len(static_grad_outputs): + raise ValueError( + "Backward graph grad dimension mismatch: " + f"received {len(grads)} grads, " + f"but expected {len(static_grad_outputs)} static_grad_outputs" + ) for g, grad in zip(static_grad_outputs, grads): if g is not None: # don't copy if autograd gods have been kind and the @@ -843,7 +899,11 @@ def backward(ctx, *grads): FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) # Input args that didn't require grad expect a None gradient. - assert isinstance(static_grad_inputs, tuple) + if not isinstance(static_grad_inputs, tuple): + raise TypeError( + "Expected static_grad_inputs to be a tuple, but got" + f" {type(static_grad_inputs).__name__}" + ) return (None, None, None) + tuple( b.detach() if b is not None else b for b in static_grad_inputs ) @@ -853,9 +913,13 @@ def functionalized(*user_args, **user_kwargs): # Decide whether to update FP8 weights skip_fp8_weight_update = None if cache_quantized_params: - assert "is_first_microbatch" in user_kwargs and isinstance( + if "is_first_microbatch" not in user_kwargs or not isinstance( user_kwargs["is_first_microbatch"], bool - ), "`is_first_microbatch` boolean kwarg must be provided for FP8 weight caching." + ): + raise ValueError( + "`is_first_microbatch` boolean kwarg must be provided for FP8 weight" + " caching." + ) skip_fp8_weight_update = not user_kwargs["is_first_microbatch"] @@ -1237,12 +1301,16 @@ def make_graphed_callables( modules = (modules,) if not isinstance(enabled, tuple): - assert isinstance(enabled, bool), "enabled must be a bool or a tuple of bools" + if not isinstance(enabled, bool): + raise TypeError( + f"enabled must be a bool or a tuple of bools, but got {type(enabled).__name__}" + ) enabled = (enabled,) * len(modules) else: - assert len(enabled) == len( - modules - ), f"enabled length ({len(enabled)}) must match modules length ({len(modules)})" + if len(enabled) != len(modules): + raise ValueError( + f"enabled length ({len(enabled)}) must match modules length ({len(modules)})" + ) if any(enabled) and recipe is None: recipe = get_default_fp8_recipe() elif not any(enabled): @@ -1278,7 +1346,8 @@ def call_func(self, *args, **kwargs): forward_funcs = [] for module in modules: - assert isinstance(module, torch.nn.Module), f"Graphing for {type(module)} is not supported." + if not isinstance(module, torch.nn.Module): + raise TypeError(f"Graphing for {type(module)} is not supported.") wrap_autocast(module) forward_funcs.append(module) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 9c21141a39..2d4583e936 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -80,7 +80,8 @@ class UserBufferQuantizationMode(Enum): def get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor: """Returns a dummy tensor of given shape.""" - assert len(shape) == 2 + if len(shape) != 2: + raise ValueError(f"Expected 2D shape, got {len(shape)}D: {shape}") global _dummy_wgrads if (shape[0], shape[1], dtype) not in _dummy_wgrads: _dummy_wgrads[(shape[0], shape[1], dtype)] = torch.empty( @@ -156,10 +157,11 @@ def initialize_ub( which also requires ``MPI_HOME=/path/to/mpi/root`` to be set at compile time. """ if not tex.device_supports_multicast(): - assert bool(int(os.getenv("UB_SKIPMC", "0"))), ( - "CUDA device, driver and/or toolkit version does not support comm+GEMM overlap with " - + "CUDA Multicast. Launch app with UB_SKIPMC=1 to try CUDA IPC instead." - ) + if not bool(int(os.getenv("UB_SKIPMC", "0"))): + raise RuntimeError( + "CUDA device, driver and/or toolkit version does not support comm+GEMM overlap " + "with CUDA Multicast. Launch app with UB_SKIPMC=1 to try CUDA IPC instead." + ) if not quantization_modes: warnings.warn( @@ -171,34 +173,48 @@ def initialize_ub( UserBufferQuantizationMode.FP8 if use_fp8 else UserBufferQuantizationMode.NONE ] else: - assert isinstance(quantization_modes, list), "quantization_modes must be a list" - assert all( - isinstance(mode, UserBufferQuantizationMode) for mode in quantization_modes - ), "quantization_modes must be a list of UserBufferQuantizationMode" + if not isinstance(quantization_modes, list): + raise TypeError( + f"quantization_modes must be a list, got {type(quantization_modes).__name__}" + ) + invalid_modes = [ + mode for mode in quantization_modes if not isinstance(mode, UserBufferQuantizationMode) + ] + if invalid_modes: + raise TypeError( + "quantization_modes must be a list of UserBufferQuantizationMode, " + f"got invalid entries: {invalid_modes}" + ) if isinstance(ub_cfgs, dict) or ub_cfgs is None: ub_cfgs = [ub_cfgs] * len(quantization_modes) else: - assert len(ub_cfgs) == len( - quantization_modes - ), "Number of ub_cfgs settings must match number of quantization configurations" + if len(ub_cfgs) != len(quantization_modes): + raise ValueError( + f"Number of ub_cfgs settings ({len(ub_cfgs)}) must match number of " + f"quantization configurations ({len(quantization_modes)})" + ) global _ub_communicators - assert _ub_communicators is None, "UB communicators are already initialized." + if _ub_communicators is not None: + raise RuntimeError("UB communicators are already initialized.") _ub_communicators = {} if tex.ubuf_built_with_mpi(): # We're bootstrapping with direct calls to MPI in Userbuffers code so we need to force # an MPI_Init() here by creating a new MPI process group... - assert torch.distributed.is_mpi_available() + if not torch.distributed.is_mpi_available(): + raise RuntimeError( + "MPI backend is not available in torch.distributed but is required " + "when Userbuffers is built with MPI support" + ) _ = torch.distributed.new_group(backend="mpi") helper = tex.CommOverlapHelper() else: # Bootstrapping with torch.distributed API, so check backend and construct # intra/inter-node process groups... - assert ( - torch.distributed.is_initialized() - ), "torch.distributed must be initialized before Userbuffers" + if not torch.distributed.is_initialized(): + raise RuntimeError("torch.distributed must be initialized before using Userbuffers") if bootstrap_backend is None: bootstrap_backend = "nccl" if torch.distributed.is_mpi_available(): @@ -206,15 +222,16 @@ def initialize_ub( elif torch.distributed.is_gloo_available(): bootstrap_backend = "gloo" else: - assert bootstrap_backend in [ - "gloo", - "mpi", - "nccl", - ], "Invalid torch.distributed backend for bootstrapping Userbuffers!" - assert torch.distributed.is_backend_available(bootstrap_backend), ( - f"PyTorch must be compiled with '{bootstrap_backend}' support in order to " - f"bootstrap Userbuffers with '{bootstrap_backend}' collectives." - ) + if bootstrap_backend not in ["gloo", "mpi", "nccl"]: + raise ValueError( + f"Invalid torch.distributed backend '{bootstrap_backend}' for bootstrapping " + "Userbuffers. Must be one of: 'gloo', 'mpi', 'nccl'" + ) + if not torch.distributed.is_backend_available(bootstrap_backend): + raise RuntimeError( + f"PyTorch must be compiled with '{bootstrap_backend}' support in order to " + f"bootstrap Userbuffers with '{bootstrap_backend}' collectives." + ) world_group = torch.distributed.new_group(backend=bootstrap_backend) world_rank = torch.distributed.get_rank(world_group) @@ -333,9 +350,11 @@ def add_ub( warnings.warn( "Atomic GEMM uses a beta API from cublas and is not tested for all use cases." ) - assert ( - quantization_mode == UserBufferQuantizationMode.FP8 - ), "Atomic GEMM overlap supported only for FP8 GEMM." + if quantization_mode != UserBufferQuantizationMode.FP8: + raise ValueError( + "Atomic GEMM overlap supported only for FP8 GEMM, " + f"got quantization_mode={quantization_mode}" + ) if method in ("bulk", "external"): warnings.warn( f"At {name}, atoimic GEMM not is supported for a bulk overlap." @@ -360,20 +379,24 @@ def add_ub( "for functionality." ) if name in layers_atomic_ring_exchange: - assert atomic_gemm and method == "ring_exchange", assert_message + if not (atomic_gemm and method == "ring_exchange"): + raise ValueError(assert_message) else: if atomic_gemm and method == "ring_exchange": - assert rs_ag_pairs[name] in layers_atomic_ring_exchange, assert_message + if rs_ag_pairs[name] not in layers_atomic_ring_exchange: + raise ValueError(assert_message) if name in external_gemm_to_overlap: - assert method == "external", ( - f"At {name}, `external` overlap method is specified, but the selected method is" - f" {method}" - ) - assert external_gemm_to_overlap[name] in methods["ring_exchange"], ( - f"At {name}, `external` overlap method is specified, but the external gemm" - f" {external_gemm_to_overlap[name]} is not using `ring_exchange` overlap method" - ) + if method != "external": + raise ValueError( + f"At {name}, `external` overlap method is specified, but the selected method " + f"is {method}" + ) + if external_gemm_to_overlap[name] not in methods["ring_exchange"]: + raise ValueError( + f"At {name}, `external` overlap method is specified, but the external gemm " + f"{external_gemm_to_overlap[name]} is not using `ring_exchange` overlap method" + ) buffer_dtype = ( torch.uint8 @@ -424,7 +447,12 @@ def add_ub( and user_ub_cfg[name]["method"] != "bulk" ): wgrad_name = name.replace("dgrad", "wgrad") - assert wgrad_name not in user_ub_cfg + if wgrad_name in user_ub_cfg: + raise ValueError( + f"Cannot specify user UB config for '{wgrad_name}' when its " + f"corresponding dgrad '{name}' uses a non-bulk overlap method " + f"('{user_ub_cfg[name]['method']}')" + ) layers_reduce_scatter_overlap.remove(wgrad_name) layers_all_gather_overlap.remove(name) layers_reduce_scatter_overlap.append(name) @@ -451,8 +479,10 @@ def get_ub(name: str, use_fp8: bool): # So favour simplicity until the correct design becomes clear. # This is mainly an internal API so we don't need to worry about future changes key = (name, UserBufferQuantizationMode.FP8 if use_fp8 else UserBufferQuantizationMode.NONE) - assert _ub_communicators is not None, "UB manager is not initialized." - assert key in _ub_communicators, f"UB for {name} with use_fp8={use_fp8} is not registered." + if _ub_communicators is None: + raise RuntimeError("UB manager is not initialized.") + if key not in _ub_communicators: + raise KeyError(f"UB for {name} with use_fp8={use_fp8} is not registered.") return _ub_communicators[key] @@ -608,7 +638,8 @@ class TransformerEngineBaseModule(torch.nn.Module, ABC): def __init__(self, name: Optional[str] = None) -> None: super().__init__() - assert torch.cuda.is_available(), "TransformerEngine needs CUDA." + if not torch.cuda.is_available(): + raise RuntimeError("TransformerEngine needs CUDA.") self.name = name self.next_iter_when_debug_should_be_run = 0 self.fp8_initialized = False @@ -694,9 +725,12 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> ] for pos, buffer_key in zip((fwd_pos, bwd_pos), (fwd_key, bwd_key)): if buffer_key in FP8GlobalStateManager.global_amax_buffer: - assert ( - buffer_key in FP8GlobalStateManager.global_amax_history_buffer - ), "TE internal error during amax history change." + if buffer_key not in FP8GlobalStateManager.global_amax_history_buffer: + raise RuntimeError( + "TE internal error during amax history change: " + f"buffer_key '{buffer_key}' found in global_amax_buffer " + "but missing from global_amax_history_buffer" + ) FP8GlobalStateManager.global_amax_buffer[buffer_key][pos] = self.fp8_meta[ meta_key ].amax_history[0] @@ -745,10 +779,11 @@ def _update_weight_quantizers(self) -> None: """Update the quantizers for the weight tensors.""" weight_tensors = self._get_weight_tensors() weight_quantizers = self._get_weight_quantizers() - assert len(weight_tensors) == len(weight_quantizers), ( - f"Number of weight tensors ({len(weight_tensors)}) and quantizers " - f"({len(weight_quantizers)}) must match" - ) + if len(weight_tensors) != len(weight_quantizers): + raise ValueError( + f"Number of weight tensors ({len(weight_tensors)}) and quantizers " + f"({len(weight_quantizers)}) must match" + ) for weight, quantizer in zip(weight_tensors, weight_quantizers): if quantizer is not None and isinstance(weight, QuantizedTensorStorage): weight.update_quantizer(quantizer) @@ -796,7 +831,11 @@ def reset(key): torch.zeros_like(self.fp8_meta[key].amax_history) ) else: - assert key in fp8_meta_tensors, "Cannot reset fp8 tensors." + if key not in fp8_meta_tensors: + raise KeyError( + f"Cannot reset fp8 tensors: key '{key}' not found in fp8_meta_tensors. " + f"Available keys: {list(fp8_meta_tensors.keys())}" + ) self.fp8_meta[key].scale.copy_(fp8_meta_tensors[key][0]) self.fp8_meta[key].amax_history.copy_(fp8_meta_tensors[key][1]) @@ -937,10 +976,11 @@ def set_activation_dtype(self, inp: torch.Tensor) -> None: if not self.allow_different_data_and_param_types: for name, param in self.named_parameters(): if param is not None: - assert dtype == param.dtype, ( - "Data types for parameters must match when outside of autocasted region. " - f" Found input dtype: {dtype} and {name!r} dtype: {param.dtype}" - ) + if dtype != param.dtype: + raise TypeError( + "Data types for parameters must match when outside of autocasted " + f"region. Found input dtype: {dtype} and {name!r} dtype: {param.dtype}" + ) self.fast_setattr("activation_dtype", dtype) def set_tensor_parallel_group(self, tp_group: Union[dist_group_type, None]) -> None: @@ -1045,10 +1085,17 @@ def prepare_forward( delayed_scaling_recipe = self.fp8_meta["recipe"].delayed() FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) else: - assert inp.is_cuda, "TransformerEngine needs CUDA." + if not inp.is_cuda: + raise RuntimeError( + f"TransformerEngine needs CUDA. Got input on device: {inp.device}" + ) if self.tp_size > 1: - assert self.tp_group_initialized, "TP group not initialized." + if not self.tp_group_initialized: + raise RuntimeError( + "Tensor parallel group not initialized. Call " + "set_tensor_parallel_group() before forward pass when tp_size > 1." + ) self.set_activation_dtype(inp) self.init_fp8_metadata(num_gemms=num_gemms) @@ -1057,10 +1104,11 @@ def prepare_forward( delayed_scaling_recipe = self.fp8 and self.fp8_meta["recipe"].delayed() if delayed_scaling_recipe: if self.sequence_parallel: - assert self.fp8_meta["recipe"].reduce_amax, ( - "Amax reduction across tensor parallel group is " - "necessary when using sequence parallelism with FP8." - ) + if not self.fp8_meta["recipe"].reduce_amax: + raise ValueError( + "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) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 02607e45e5..fade2957d5 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -646,9 +646,8 @@ def __init__( self.ub_name = ub_name self.save_original_input = save_original_input self.single_grouped_parameter = single_grouped_parameter - assert ( - not ub_overlap_rs and not ub_overlap_ag - ), "GroupedLinear doesn't support Userbuffer overlap." + if ub_overlap_rs or ub_overlap_ag: + raise ValueError("GroupedLinear doesn't support Userbuffer overlap.") self.init_method = init_method self.get_rng_state_tracker = get_rng_state_tracker self.rng_tracker_name = rng_tracker_name @@ -683,9 +682,11 @@ def __init__( ) self.parallel_mode = parallel_mode - assert ( - self.parallel_mode in GemmParallelModes - ), f"parallel_mode {parallel_mode} not supported" + if self.parallel_mode not in GemmParallelModes: + raise ValueError( + f"parallel_mode {parallel_mode!r} not supported." + f" Supported modes: {GemmParallelModes}" + ) if self.parallel_mode == "column": self.out_features = divide(self.out_features, self.tp_size) @@ -788,9 +789,11 @@ def make_grouped_weights(self, defer_init=False) -> None: # Re-register as a single grouped weight parameter. # Re-register as a single grouped weight parameter. - assert isinstance(grouped_weights, torch.Tensor) and ( - weight_quantizers[0] is None or not weight_quantizers[0].internal - ), "Found internal quantizer with `single_grouped_parameter=True`." + if not ( + isinstance(grouped_weights, torch.Tensor) + and (weight_quantizers[0] is None or not weight_quantizers[0].internal) + ): + raise RuntimeError("Found internal quantizer with `single_grouped_parameter=True`.") self.register_parameter( "weight", torch.nn.Parameter(grouped_weights), @@ -875,10 +878,13 @@ def forward( """ debug = self.is_debug_iter() - assert not isinstance( - inp, QuantizedTensorStorage - ), "GroupedLinear doesn't support input tensor in FP8." - assert len(m_splits) == self.num_gemms, "Number of splits should match number of GEMMs." + if isinstance(inp, QuantizedTensorStorage): + raise TypeError("GroupedLinear doesn't support input tensor in FP8.") + if len(m_splits) != self.num_gemms: + raise ValueError( + f"Number of splits ({len(m_splits)}) should match number of" + f" GEMMs ({self.num_gemms})." + ) is_grad_enabled = torch.is_grad_enabled() @@ -969,10 +975,11 @@ def backward_dw(self): def _customize_quantizers_float8_current_scaling(self, fwd: bool, recipe: Recipe) -> None: """Customize quantizers based on current scaling recipe + linear.""" - assert not self.tp_size > 1, ( - "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " - "Because the TP communication is handled outside of this module." - ) + if self.tp_size > 1: + raise ValueError( + "GroupedLinear doesn't support TP > 1 with Float8 current scaling. " + "Because the TP communication is handled outside of this module." + ) if fwd: for i in range(self.num_gemms): @@ -1077,7 +1084,8 @@ def _get_quantizers(self): def _get_debug_quantizers(self): original_quantizers = self._get_quantizers() - assert TEDebugState.debug_enabled + if not TEDebugState.debug_enabled: + raise RuntimeError("TEDebugState.debug_enabled must be True to get debug quantizers") names = ["activation", "weight", "output", "dgrad", "wgrad", "gradient"] return tuple( diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index 5beeed1262..ca59a0ebf8 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -42,10 +42,16 @@ def forward( return inp, torch.tensor([], device=inp.device) # Device check - assert inp.is_cuda, "TransformerEngine needs CUDA." - assert index.is_cuda, "TransformerEngine needs CUDA." + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not index.is_cuda: + raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") # Shape check - assert inp.size(0) == index.size(0), "Permute not possible" + if inp.size(0) != index.size(0): + raise ValueError( + f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " + f"index.size(0) ({index.size(0)})." + ) # Data type check dtype = TE_DType[inp.dtype] @@ -119,7 +125,8 @@ def forward( # None probs check if probs is not None: - assert probs.is_cuda, "TransformerEngine needs CUDA." + if not probs.is_cuda: + raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") if probs.dtype != torch.float32: warnings.warn( @@ -136,8 +143,12 @@ def forward( probs = torch.empty(0) # Device check - assert inp.is_cuda, "TransformerEngine needs CUDA." - assert row_id_map.is_cuda, "TransformerEngine needs CUDA." + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not row_id_map.is_cuda: + raise ValueError( + f"row_id_map must be a CUDA tensor, but got tensor on {row_id_map.device}." + ) # Data type check dtype = TE_DType[inp.dtype] @@ -198,19 +209,30 @@ def forward( ctx.probs = probs return inp, torch.tensor([], device=inp.device), torch.tensor([], device=inp.device) - assert inp.is_cuda, "TransformerEngine needs CUDA." - assert routing_map.is_cuda, "TransformerEngine needs CUDA." + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not routing_map.is_cuda: + raise ValueError( + f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + ) if probs is not None: - assert probs.is_cuda, "TransformerEngine needs CUDA." + if not probs.is_cuda: + raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") if pad_offsets is not None: - assert pad_offsets.is_cuda, "TransformerEngine needs CUDA." + if not pad_offsets.is_cuda: + raise ValueError( + f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + ) - assert inp.size(0) == routing_map.size(0), "Permute not possible" + if inp.size(0) != routing_map.size(0): + raise ValueError( + f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " + f"routing_map.size(0) ({routing_map.size(0)})." + ) num_tokens, hidden_size = inp.size() num_experts = routing_map.size(1) - assert ( - num_out_tokens is not None - ), "num_out_tokens must be provided to the fused permute function." + if num_out_tokens is None: + raise ValueError("num_out_tokens must be provided to the fused permute function.") row_id_map = triton_permutation.make_row_id_map(routing_map, num_tokens, num_experts) @@ -226,13 +248,25 @@ def forward( if blockwise_recipe: fp8_scale = inp._rowwise_scale_inv.T.contiguous() scale_hidden_dim = fp8_scale.shape[1] - assert num_tokens == fp8_scale.shape[0], "scale and input shape mismatch" + if num_tokens != fp8_scale.shape[0]: + raise ValueError( + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Input shape: ({num_tokens}, {hidden_size}), " + f"scale shape: {tuple(fp8_scale.shape)}." + ) inp = inp._rowwise_data # mxfp8 scaling elif mxfp8_recipe: fp8_scale = inp._rowwise_scale_inv.contiguous() scale_hidden_dim = fp8_scale.shape[1] - assert num_tokens == fp8_scale.shape[0], "scale and input shape mismatch" + if num_tokens != fp8_scale.shape[0]: + raise ValueError( + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Input shape: ({num_tokens}, {hidden_size}), " + f"scale shape: {tuple(fp8_scale.shape)}." + ) inp = inp._rowwise_data # per-tensor scaling elif per_tensor_recipe: @@ -318,9 +352,11 @@ def backward( probs_grad = None if ctx.needs_input_grad[0]: row_id_map, pad_offsets = ctx.saved_tensors - assert not isinstance( - permuted_act_grad, QuantizedTensor - ), "The backward of moe_permute does not support FP8." + if isinstance(permuted_act_grad, QuantizedTensor): + raise TypeError( + "The backward of moe_permute does not support FP8, but got " + f"QuantizedTensor of type {type(permuted_act_grad).__name__}." + ) act_grad, probs_grad = triton_permutation.unpermute_with_mask_map( permuted_act_grad, row_id_map, @@ -360,17 +396,30 @@ def forward( with_probs = merging_probs is not None if with_probs: - assert merging_probs.is_cuda, "TransformerEngine needs CUDA." + if not merging_probs.is_cuda: + raise ValueError( + "merging_probs must be a CUDA tensor, but got tensor on " + f"{merging_probs.device}." + ) # Device check - assert inp.is_cuda, "TransformerEngine needs CUDA." - assert row_id_map.is_cuda, "TransformerEngine needs CUDA." + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not row_id_map.is_cuda: + raise ValueError( + f"row_id_map must be a CUDA tensor, but got tensor on {row_id_map.device}." + ) if pad_offsets is not None: - assert pad_offsets.is_cuda, "TransformerEngine needs CUDA." + if not pad_offsets.is_cuda: + raise ValueError( + f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + ) - assert not isinstance( - inp, QuantizedTensor - ), "The forward of moe_unpermute does not support FP8." + if isinstance(inp, QuantizedTensor): + raise TypeError( + "The forward of moe_unpermute does not support FP8, but got " + f"QuantizedTensor of type {type(inp).__name__}." + ) unpermuted_output, _ = triton_permutation.unpermute_with_mask_map( inp, row_id_map, @@ -427,13 +476,23 @@ def backward(ctx, unpermuted_act_grad): fp8_scale = unpermuted_act_grad._rowwise_scale_inv.T.contiguous() unpermuted_act_grad = unpermuted_act_grad._rowwise_data scale_hidden_dim = fp8_scale.shape[1] - assert ctx.num_tokens == fp8_scale.shape[0], "scale and input shape mismatch" + if ctx.num_tokens != fp8_scale.shape[0]: + raise ValueError( + f"Scale and input shape mismatch: num_tokens ({ctx.num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Scale shape: {tuple(fp8_scale.shape)}." + ) # mxfp8 scaling elif mxfp8_recipe: fp8_scale = unpermuted_act_grad._rowwise_scale_inv.contiguous() unpermuted_act_grad = unpermuted_act_grad._rowwise_data scale_hidden_dim = fp8_scale.shape[1] - assert ctx.num_tokens == fp8_scale.shape[0], "scale and input shape mismatch" + if ctx.num_tokens != fp8_scale.shape[0]: + raise ValueError( + f"Scale and input shape mismatch: num_tokens ({ctx.num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Scale shape: {tuple(fp8_scale.shape)}." + ) else: raise ValueError("Unsupported FP8 recipe") else: @@ -441,10 +500,13 @@ def backward(ctx, unpermuted_act_grad): fp8_dtype = None fp8_scale = None + permuted_scale = None if ctx.with_probs: - assert ( - not fp8 - ), "The backward of moe_unpermute with merging probs does not support FP8." + if fp8: + raise TypeError( + "The backward of moe_unpermute with merging probs does not support FP8, " + f"but got FP8 gradient with dtype {fp8_dtype}." + ) act_grad, probs_grad = ( triton_permutation.unpermute_with_mask_map_bwd_with_merging_probs( unpermuted_act_grad, @@ -619,10 +681,12 @@ def moe_permute_and_pad_with_probs( align_size : int the alignment size for the input tensor. """ - assert ( - tokens_per_expert is not None - ), "tokens_per_expert must be provided to the fused permute padding function." - assert align_size > 0, f"align_size must be positive, got {align_size}" + if tokens_per_expert is None: + raise ValueError( + "tokens_per_expert must be provided to the fused permute padding function." + ) + if align_size <= 0: + raise ValueError(f"align_size must be positive, got {align_size}.") # Ensure tokens_per_expert is on the same device as input to avoid device transfers if tokens_per_expert.device != inp.device: @@ -713,15 +777,27 @@ def forward( if not inp.numel(): return inp, probs - assert inp.is_cuda, "TransformerEngine needs CUDA." - assert split_sizes.is_cuda, "TransformerEngine needs CUDA." - assert sorted_idxs.is_cuda, "TransformerEngine needs CUDA." + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not split_sizes.is_cuda: + raise ValueError( + f"split_sizes must be a CUDA tensor, but got tensor on {split_sizes.device}." + ) + if not sorted_idxs.is_cuda: + raise ValueError( + f"sorted_idxs must be a CUDA tensor, but got tensor on {sorted_idxs.device}." + ) if probs is not None: - assert probs.is_cuda, "TransformerEngine needs CUDA." + if not probs.is_cuda: + raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") num_tokens, hidden_size = inp.shape num_splits = split_sizes.size(0) - assert num_splits == sorted_idxs.size(0) + if num_splits != sorted_idxs.size(0): + raise ValueError( + f"split_sizes.size(0) ({num_splits}) must match " + f"sorted_idxs.size(0) ({sorted_idxs.size(0)})." + ) fp8 = isinstance(inp, Float8Tensor) if fp8: diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index eba547afb0..47e6d5c8dc 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -97,7 +97,8 @@ def check_recipe_support(recipe: Recipe) -> None: recipe_supported, unsupported_reason = check_fp8_block_scaling_support() elif isinstance(recipe, MXFP8BlockScaling): recipe_supported, unsupported_reason = check_mxfp8_support() - assert recipe_supported, unsupported_reason + if not recipe_supported: + raise RuntimeError(unsupported_reason) def get_default_fp8_recipe() -> Recipe: diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 685b2c5548..22a6a41eb1 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -203,7 +203,8 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: # Parameter construction calls detach()/alias-like paths. if func in (torch.ops.aten.detach.default, torch.ops.aten.alias.default): src = args[0] - assert isinstance(src, GroupedTensor) + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") if func == torch.ops.aten.detach.default: return make_wrapper_like(src, requires_grad=False) return make_wrapper_like(src, requires_grad=src.requires_grad) @@ -212,7 +213,8 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: # Handle this explicitly so grouped parameters can be created safely. if func == torch.ops.aten.expand.default: src = args[0] - assert isinstance(src, GroupedTensor) + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") expanded_shape = tuple(args[1]) src_shape = tuple(src.shape) if len(expanded_shape) == len(src_shape): @@ -228,7 +230,8 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: if func == torch.ops.aten.expand_as.default: src = args[0] other = args[1] - assert isinstance(src, GroupedTensor) + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") if other is src: return _GroupedIdentityFunc.apply(src) if tuple(other.shape) == tuple(src.shape): @@ -240,7 +243,8 @@ def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: # returning a flat view of grouped backing storage. if func in (torch.ops.aten.view.default, torch.ops.aten._unsafe_view.default): src = args[0] - assert isinstance(src, GroupedTensor) + if not isinstance(src, GroupedTensor): + raise TypeError(f"Expected GroupedTensor, got {type(src).__name__}") target_shape = tuple(args[1]) if target_shape in ((-1,), (src.numel(),)): if src.rowwise_data is not None: @@ -317,7 +321,11 @@ def maybe_update_inplace(arg, new_arg, schema_arg): for arg, new_arg, schema_arg in zip(args, new_args, schema_args): maybe_update_inplace(arg, new_arg, schema_arg) for kwarg, new_kwarg, schema_arg in zip(kwargs, new_kwargs, schema_args[args_len:]): - assert kwarg == new_kwarg == schema_arg.name, "name of kwarg should match schema" + if kwarg != new_kwarg or kwarg != schema_arg.name: + raise RuntimeError( + f"Name of kwarg should match schema, got kwarg={kwarg!r}," + f" new_kwarg={new_kwarg!r}, schema_arg.name={schema_arg.name!r}" + ) maybe_update_inplace(kwargs[kwarg], new_kwargs[new_kwarg], schema_arg) return None diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index e7509f3994..3e8bf3f2f3 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -327,11 +327,14 @@ def update_usage( # If both rowwise and columnwise are requested, create columnwise from rowwise if needed if rowwise_usage and columnwise_usage: - assert ( - self._rowwise_data is not None - and self._rowwise_scale_inv is not None - and self._amax_rowwise is not None - ), "Cannot update to rowwise and columnwise usage because rowwise data is None." + if ( + self._rowwise_data is None + or self._rowwise_scale_inv is None + or self._amax_rowwise is None + ): + raise RuntimeError( + "Cannot update to rowwise and columnwise usage because rowwise data is None." + ) if self._columnwise_data is None or self._columnwise_scale_inv is None: self._create_columnwise() return @@ -381,16 +384,16 @@ def _create_columnwise(self): """ Update columnwise data and columnwise scale inv. Can only be used when using 2D scaling. """ - assert ( - self._quantizer is not None and self._quantizer.with_2d_quantization - ), "Cannot create columnwise data without 2D quantization enabled." + if self._quantizer is None or not self._quantizer.with_2d_quantization: + raise RuntimeError("Cannot create columnwise data without 2D quantization enabled.") rowwise_data = self._rowwise_data if not rowwise_data.is_contiguous(): rowwise_data = rowwise_data.contiguous() # NVFP4 requires a specialized transpose that handles nibble repacking self._columnwise_data = tex.nvfp4_data_transpose(rowwise_data, out=self._columnwise_data) if self._columnwise_scale_inv is None: - assert self._quantizer is not None + if self._quantizer is None: + raise RuntimeError("Cannot create columnwise scale inverse: quantizer is None.") # Use logical shape (self.size()), not packed byte shape (rowwise_data.shape) # NVFP4 packs 2 elements per byte, so rowwise_data.shape[-1] is K/2 logical_shape = self.size() @@ -400,8 +403,18 @@ def _create_columnwise(self): dtype=self._rowwise_scale_inv.dtype, device=self._rowwise_scale_inv.device, ) - assert len(self._rowwise_scale_inv.shape) == 2 - assert len(self._columnwise_scale_inv.shape) == 2 + if len(self._rowwise_scale_inv.shape) != 2: + raise ValueError( + "Expected rowwise_scale_inv to be 2D, but got" + f" {len(self._rowwise_scale_inv.shape)}D with shape" + f" {self._rowwise_scale_inv.shape}." + ) + if len(self._columnwise_scale_inv.shape) != 2: + raise ValueError( + "Expected columnwise_scale_inv to be 2D, but got" + f" {len(self._columnwise_scale_inv.shape)}D with shape" + f" {self._columnwise_scale_inv.shape}." + ) # rowwise_scale_inv has shape [M_padded, K_tiles] where each tile's scale # is repeated 16 times (once per row in the 16x16 tile). diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index d23892af94..c80bc8aaa4 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -37,19 +37,31 @@ def replace_raw_data(tensor: QuantizedTensor, new_raw_data: torch.Tensor): """ if isinstance(tensor, Float8Tensor): old_raw_data = tensor._data - assert old_raw_data.dtype == new_raw_data.dtype, "The data types of raw data don't match" + if old_raw_data.dtype != new_raw_data.dtype: + raise ValueError( + "The data types of raw data don't match: " + f"old dtype={old_raw_data.dtype}, new dtype={new_raw_data.dtype}" + ) new_raw_data.detach().copy_(old_raw_data) tensor._data = new_raw_data del old_raw_data elif isinstance(tensor, Float8BlockwiseQTensor): old_raw_data = tensor._rowwise_data - assert old_raw_data.dtype == new_raw_data.dtype, "The data types of raw data don't match" + if old_raw_data.dtype != new_raw_data.dtype: + raise ValueError( + "The data types of raw data don't match: " + f"old dtype={old_raw_data.dtype}, new dtype={new_raw_data.dtype}" + ) new_raw_data.detach().copy_(old_raw_data) tensor._rowwise_data = new_raw_data del old_raw_data elif isinstance(tensor, NVFP4Tensor): old_rowwise = tensor._rowwise_data - assert old_rowwise.dtype == new_raw_data.dtype, "The data types of raw data don't match" + if old_rowwise.dtype != new_raw_data.dtype: + raise ValueError( + f"The data types of raw data don't match: {old_rowwise.dtype} vs" + f" {new_raw_data.dtype}" + ) new_raw_data.detach().copy_(old_rowwise) tensor._rowwise_data = new_raw_data del old_rowwise @@ -276,10 +288,16 @@ def _cast_master_weights_to_fp8_delayed_scaling( continue # If master weight is not None, start_offset must be a valid value. - assert start_offset is not None - assert start_offset >= 0 + if start_offset is None: + raise ValueError("start_offset must not be None when master_weight is provided") + if start_offset < 0: + raise ValueError(f"start_offset must be non-negative, got {start_offset}") end_offset = start_offset + master_weight.numel() - assert end_offset <= model_weight.numel() + if end_offset > model_weight.numel(): + raise ValueError( + f"end_offset ({end_offset}) exceeds model_weight numel ({model_weight.numel()}), " + f"start_offset={start_offset}, master_weight numel={master_weight.numel()}" + ) # master_weight may be smaller than model_weight because it could be distributed across # multiple ranks. So we need to create a dummy weight using the raw data from model_weight. @@ -363,9 +381,21 @@ def _cast_master_weights_to_fp8_current_scaling( # Make sure all the model weights have the same numerical options. quantizer = model_weight._get_quantizer() - assert quantizer.dtype == fp8_dtype - assert quantizer.force_pow_2_scales == force_pow_2_scales - assert quantizer.amax_epsilon == amax_epsilon + if quantizer.dtype != fp8_dtype: + raise ValueError( + "All model weights must have the same fp8 dtype, " + f"expected {fp8_dtype} but got {quantizer.dtype}" + ) + if quantizer.force_pow_2_scales != force_pow_2_scales: + raise ValueError( + "All model weights must have the same force_pow_2_scales, " + f"expected {force_pow_2_scales} but got {quantizer.force_pow_2_scales}" + ) + if quantizer.amax_epsilon != amax_epsilon: + raise ValueError( + "All model weights must have the same amax_epsilon, " + f"expected {amax_epsilon} but got {quantizer.amax_epsilon}" + ) scales.append(quantizer.scale.view(1)) scale_invs.append(model_weight._scale_inv.view(1)) @@ -479,19 +509,47 @@ def _cast_master_weights_to_fp8_blockwise_scaling( # Make sure all the model weights have the same numerical options. quantizer = model_weight._get_quantizer() - assert block_len == quantizer.block_len - assert fp8_dtype == quantizer.dtype - assert force_pow_2_scales == quantizer.force_pow_2_scales - assert amax_epsilon == quantizer.amax_epsilon + if block_len != quantizer.block_len: + raise ValueError( + "All model weights must have the same block_len, " + f"expected {block_len} but got {quantizer.block_len}" + ) + if fp8_dtype != quantizer.dtype: + raise ValueError( + "All model weights must have the same fp8 dtype, " + f"expected {fp8_dtype} but got {quantizer.dtype}" + ) + if force_pow_2_scales != quantizer.force_pow_2_scales: + raise ValueError( + "All model weights must have the same force_pow_2_scales, " + f"expected {force_pow_2_scales} but got {quantizer.force_pow_2_scales}" + ) + if amax_epsilon != quantizer.amax_epsilon: + raise ValueError( + "All model weights must have the same amax_epsilon, " + f"expected {amax_epsilon} but got {quantizer.amax_epsilon}" + ) scale_shape = quantizer.get_scale_shape(model_weight.shape, False) amax = packed_amaxes[cu_amax_sizes[i] : cu_amax_sizes[i + 1]].reshape(scale_shape) scale = torch.empty(scale_shape, dtype=torch.float32, device=device) scale_inv = model_weight._rowwise_scale_inv - assert len(scale_shape) == 2 - assert len(scale_inv.shape) == 2 - assert scale_inv.shape[0] == scale_shape[0] - assert scale_inv.shape[1] == scale_shape[1] + if len(scale_shape) != 2: + raise ValueError(f"scale_shape must be 2D, got {len(scale_shape)}D shape {scale_shape}") + if len(scale_inv.shape) != 2: + raise ValueError( + f"scale_inv must be 2D, got {len(scale_inv.shape)}D shape {scale_inv.shape}" + ) + if scale_inv.shape[0] != scale_shape[0]: + raise ValueError( + f"scale_inv dim 0 mismatch: scale_inv.shape={scale_inv.shape}," + f" scale_shape={scale_shape}" + ) + if scale_inv.shape[1] != scale_shape[1]: + raise ValueError( + f"scale_inv dim 1 mismatch: scale_inv.shape={scale_inv.shape}," + f" scale_shape={scale_shape}" + ) amaxes.append(amax) scales.append(scale) @@ -499,7 +557,11 @@ def _cast_master_weights_to_fp8_blockwise_scaling( # Compute amax of the master weight and store it in packed_amaxes. if master_weight is not None: - assert len(model_weight.shape) == 2 + if len(model_weight.shape) != 2: + raise ValueError( + "model_weight must be 2D for blockwise scaling, " + f"got {len(model_weight.shape)}D shape {model_weight.shape}" + ) h, w = model_weight.shape tex.fp8_block_scaling_compute_partial_amax( master_weight, amax, h, w, start_offset, block_len @@ -550,7 +612,11 @@ def _cast_master_weights_to_fp8_blockwise_scaling( end_offset = start_offset + master_weight.numel() if not use_fsdp_shard_model_weights: model_weight_fragment = model_weight._rowwise_data.reshape(-1)[start_offset:end_offset] - assert len(model_weight.shape) == 2 + if len(model_weight.shape) != 2: + raise ValueError( + "model_weight must be 2D for blockwise scaling partial cast, " + f"got {len(model_weight.shape)}D shape {model_weight.shape}" + ) h, w = model_weight.shape tex.fp8_block_scaling_partial_cast( master_weight, model_weight_fragment, scale, h, w, start_offset, block_len, fp8_dtype @@ -581,9 +647,12 @@ def _cast_master_weights_to_nvfp4_2d( amax_targets: List[Optional[torch.Tensor]] = [] for model_weight, _, _, _ in params: quantizer = model_weight._get_quantizer() - assert isinstance(quantizer, NVFP4Quantizer) - assert quantizer.with_2d_quantization, "NVFP4 2D quantization must be enabled." - assert len(model_weight.shape) == 2 + if not isinstance(quantizer, NVFP4Quantizer): + raise TypeError(f"Expected NVFP4Quantizer, got {type(quantizer).__name__}") + if not quantizer.with_2d_quantization: + raise ValueError("NVFP4 2D quantization must be enabled.") + if len(model_weight.shape) != 2: + raise ValueError(f"Expected 2D model weight, got {len(model_weight.shape)}D") h, w = model_weight.shape tile_h = (h + block_len - 1) // block_len tile_w = (w + block_len - 1) // block_len @@ -616,13 +685,15 @@ def _cast_master_weights_to_nvfp4_2d( scale = packed_scales[cu_amax_sizes[i] : cu_amax_sizes[i + 1]].reshape(scale_shape) global_amax_view = global_amax_views[i] - assert model_weight._rowwise_scale_inv is not None + if model_weight._rowwise_scale_inv is None: + raise RuntimeError("model_weight._rowwise_scale_inv must not be None") amaxes.append(amax) scales.append(scale) if master_weight is not None and master_weight.numel() > 0: - assert len(model_weight.shape) == 2 + if len(model_weight.shape) != 2: + raise ValueError(f"Expected 2D model weight, got {len(model_weight.shape)}D") h, w = model_weight.shape # Collect for batched processing master_weight_list.append(master_weight) @@ -728,7 +799,8 @@ def _cast_master_weights_to_nvfp4_2d( byte_start = start_offset // 2 byte_end = (end_offset + 1) // 2 model_weight_fragment = rowwise_bytes[byte_start:byte_end] - assert len(model_weight.shape) == 2 + if len(model_weight.shape) != 2: + raise ValueError(f"Expected 2D model weight, got {len(model_weight.shape)}D") h, w = model_weight.shape partial_cast_inp_list.append(master_weight) @@ -793,9 +865,15 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( cu_colwise_amax_sizes = [0] for model_weight, _, _, _ in params: rowwise_shape = model_weight._rowwise_scale_inv.shape - assert len(rowwise_shape) == 2 + if len(rowwise_shape) != 2: + raise ValueError( + f"rowwise_scale_inv must be 2D, got {len(rowwise_shape)}D shape {rowwise_shape}" + ) colwise_shape = model_weight._columnwise_scale_inv.shape - assert len(colwise_shape) == 2 + if len(colwise_shape) != 2: + raise ValueError( + f"columnwise_scale_inv must be 2D, got {len(colwise_shape)}D shape {colwise_shape}" + ) cu_rowwise_amax_sizes.append( cu_rowwise_amax_sizes[-1] + rowwise_shape[0] * rowwise_shape[1] ) @@ -834,7 +912,11 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( # Compute amax of the master weight and store it in packed_amaxes. if master_weight is not None: - assert len(model_weight.shape) == 2 + if len(model_weight.shape) != 2: + raise ValueError( + "model_weight must be 2D for MXFP8 scaling, " + f"got {len(model_weight.shape)}D shape {model_weight.shape}" + ) h, w = model_weight.shape tex.mxfp8_scaling_compute_partial_amax( master_weight, amax_rowwise, amax_colwise, h, w, start_offset @@ -878,7 +960,11 @@ def _cast_master_weights_to_fp8_mxfp8_scaling( else: rowwise_fragment = model_weight._rowwise_data.reshape(-1)[start_offset:end_offset] colwise_fragment = model_weight._columnwise_data.reshape(-1)[start_offset:end_offset] - assert len(model_weight.shape) == 2 + if len(model_weight.shape) != 2: + raise ValueError( + "model_weight must be 2D for MXFP8 scaling partial cast, " + f"got {len(model_weight.shape)}D shape {model_weight.shape}" + ) h, w = model_weight.shape tex.mxfp8_scaling_partial_cast( master_weight, @@ -966,7 +1052,8 @@ def _nvfp4_2d_multi_tensor_transpose(nvfp4_tensors: List[NVFP4Tensor]): # Allocate columnwise_scale_inv if needed if tensor._columnwise_scale_inv is None: - assert tensor._quantizer is not None + if tensor._quantizer is None: + raise RuntimeError("tensor._quantizer must not be None") columnwise_scale_inv_shape = tensor._quantizer.get_scale_shape(logical_shape, True) columnwise_scale_inv = torch.empty( columnwise_scale_inv_shape, diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 868cbbdac8..4b96ccf739 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -373,23 +373,35 @@ def __init__( self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm if parallel_attention_mlp: - assert self.layer_type == "encoder", "parallel_attention requires layer_type='encoder'" - assert not self.apply_residual_connection_post_layernorm, ( - "parallel_attention and apply_residual_connection_post_layernorm " - "not supported simultaneously." - ) - assert ( - not self.output_layernorm - ), "parallel_attention and output_layernorm not supported simultaneously" + if self.layer_type != "encoder": + raise ValueError( + "parallel_attention requires layer_type='encoder', " + f"but got layer_type={self.layer_type!r}" + ) + if self.apply_residual_connection_post_layernorm: + raise ValueError( + "parallel_attention and apply_residual_connection_post_layernorm " + "are not supported simultaneously." + ) + if self.output_layernorm: + raise ValueError( + "parallel_attention and output_layernorm are not supported simultaneously." + ) self.parallel_attention_mlp = parallel_attention_mlp - assert layer_type in LayerTypes, f"layer_type {layer_type} not supported" + if layer_type not in LayerTypes: + raise ValueError( + f"layer_type {layer_type!r} is not supported. " + f"Supported types are: {', '.join(repr(t) for t in LayerTypes)}" + ) if not fuse_qkv_params: - assert ( - not fuse_wgrad_accumulation - ), "Gradient accumulation fusion requires single QKV parameter." + if fuse_wgrad_accumulation: + raise ValueError( + "Gradient accumulation fusion (fuse_wgrad_accumulation=True) " + "requires fuse_qkv_params=True, but fuse_qkv_params is False." + ) if not fuse_qkv_params: qkv_weight_interleaved = False @@ -796,32 +808,57 @@ def forward( }: enc_dec_bottom_right_diagonal = True - assert ( - self_attn_mask_type in AttnMaskTypes - ), f"self_attn_mask_type {self_attn_mask_type} not supported" - assert ( - enc_dec_attn_mask_type in AttnMaskTypes - ), f"enc_dec_attn_mask_type {enc_dec_attn_mask_type} not supported" + if self_attn_mask_type not in AttnMaskTypes: + raise ValueError( + f"self_attn_mask_type {self_attn_mask_type!r} is not supported. " + f"Supported types are: {', '.join(repr(t) for t in AttnMaskTypes)}" + ) + if enc_dec_attn_mask_type not in AttnMaskTypes: + raise ValueError( + f"enc_dec_attn_mask_type {enc_dec_attn_mask_type!r} is not supported. " + f"Supported types are: {', '.join(repr(t) for t in AttnMaskTypes)}" + ) hidden_states = hidden_states.contiguous() if self.sequence_parallel and self.seq_length is not None: - assert ( - hidden_states.shape[0] == self.seq_length // self.tp_size - ), "Sequence dimension must be split across TP group when using sequence parallel." + if hidden_states.shape[0] != self.seq_length // self.tp_size: + raise ValueError( + "Sequence dimension must be split across TP group when using " + "sequence parallel. Expected hidden_states.shape[0] to be " + f"{self.seq_length // self.tp_size} " + f"(seq_length={self.seq_length} // tp_size={self.tp_size}), " + f"but got {hidden_states.shape[0]}." + ) if ( "padding" in self_attn_mask_type or self_attn_mask_type == "arbitrary" ) and attention_mask is not None: - assert all( - attention_mask[i].dtype == torch.bool for i in range(len(attention_mask)) - ), "Attention mask must be a boolean tensor or a list/tuple of two boolean tensors" + if not all(attention_mask[i].dtype == torch.bool for i in range(len(attention_mask))): + non_bool_dtypes = [ + (i, attention_mask[i].dtype) + for i in range(len(attention_mask)) + if attention_mask[i].dtype != torch.bool + ] + raise TypeError( + "Attention mask must be a boolean tensor or a list/tuple of boolean " + f"tensors, but found non-bool dtypes at indices: {non_bool_dtypes}" + ) if ( "padding" in enc_dec_attn_mask_type or enc_dec_attn_mask_type == "arbitrary" ) and enc_dec_attn_mask is not None: - assert all( + if not all( enc_dec_attn_mask[i].dtype == torch.bool for i in range(len(enc_dec_attn_mask)) - ), "Encoder-decoder attention mask must be boolean tensor(s)" + ): + non_bool_dtypes = [ + (i, enc_dec_attn_mask[i].dtype) + for i in range(len(enc_dec_attn_mask)) + if enc_dec_attn_mask[i].dtype != torch.bool + ] + raise TypeError( + "Encoder-decoder attention mask must be boolean tensor(s), " + f"but found non-bool dtypes at indices: {non_bool_dtypes}" + ) # For AMP if torch.is_autocast_enabled(): diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index b1cc3be19d..a23e822f91 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -146,7 +146,8 @@ def compare_tensors(a: torch.Tensor, b: torch.Tensor) -> None: def ensure_divisibility(numerator: int, denominator: int) -> None: """Ensure that numerator is divisible by the denominator.""" - assert numerator % denominator == 0, f"{numerator} is not divisible by {denominator}" + if numerator % denominator != 0: + raise ValueError(f"{numerator} is not divisible by {denominator}") def divide(numerator: int, denominator: int) -> int: @@ -270,13 +271,16 @@ def forward( @staticmethod def backward(ctx, *grad_outputs): # pylint: disable=missing-function-docstring - assert len(grad_outputs) > 0, "No gradients received for backprop!" + if len(grad_outputs) == 0: + raise RuntimeError("No gradients received for backprop!") if isinstance(ctx.split_size_or_sections, (list, tuple)): split_sizes = ctx.split_size_or_sections - assert len(grad_outputs) == len( - split_sizes - ), "Unequal number of gradients vs split sections for backprop!" + if len(grad_outputs) != len(split_sizes): + raise RuntimeError( + f"Unequal number of gradients ({len(grad_outputs)}) vs " + f"split sections ({len(split_sizes)}) for backprop!" + ) if isinstance(ctx.split_size_or_sections, int): split_sizes = [ctx.split_size_or_sections] * len(grad_outputs) dims = len(grad_outputs[0].shape) @@ -370,7 +374,8 @@ def validate_rng_states_func(get_rng_tracker: Callable) -> None: """Checks if passed in param function has everything required for tensor/model and sequence parallel. """ - assert callable(get_rng_tracker), "get_rng_tracker is not a valid function" + if not callable(get_rng_tracker): + raise TypeError(f"get_rng_tracker must be callable, got {type(get_rng_tracker).__name__}") rng_tracker = None try: @@ -378,15 +383,13 @@ def validate_rng_states_func(get_rng_tracker: Callable) -> None: except Exception as e: raise RuntimeError("Cannot call get_rng_tracker function") from e - assert hasattr(rng_tracker, "get_states") and callable( - rng_tracker.get_states - ), "rng_tracker object does not have valid method get_states" - assert hasattr(rng_tracker, "set_states") and callable( - rng_tracker.set_states - ), "rng_tracker object does not have valid method set_states" - assert hasattr(rng_tracker, "fork") and callable( - rng_tracker.fork - ), "rng_tracker object does not have valid method fork" + for method_name in ("get_states", "set_states", "fork"): + if not hasattr(rng_tracker, method_name) or not callable(getattr(rng_tracker, method_name)): + raise TypeError( + f"rng_tracker object ({type(rng_tracker).__name__}) does not have " + f"a valid callable method '{method_name}'. " + "Required methods: get_states, set_states, fork." + ) validate_ctx_manager(rng_tracker.fork) @@ -397,11 +400,12 @@ def assert_viewless_tensor(tensor: torch.Tensor, extra_msg: Optional[str] = None return [assert_viewless_tensor(t) for t in tensor] if not isinstance(tensor, torch.Tensor): return tensor - assert tensor._base is None, ( - "Ensure tensor._base is None before setting tensor.data or storing " - "tensor to memory buffer. Otherwise, a memory leak will occur (and " - f"likely accumulate over iterations). {extra_msg}" - ) + if tensor._base is not None: + raise ValueError( + "Ensure tensor._base is None before setting tensor.data or storing " + "tensor to memory buffer. Otherwise, a memory leak will occur (and " + f"likely accumulate over iterations). {extra_msg}" + ) return tensor @@ -439,11 +443,13 @@ def assert_dim_for_fp8_exec(*tensors: List[torch.Tensor]) -> None: """Assert that tensor or tensors dimensions are supported for FP8 TN GEMM.""" for tensor in tensors: - assert math.prod(tensor.shape[:-1]) % 8 == 0 and tensor.shape[-1] % 16 == 0, ( - "FP8 execution requires the product of all dimensions except the last to be divisible" - " by 8 and the last dimension to be divisible by 16, but got tensor with" - f" dims={list(tensor.size())}" - ) + if math.prod(tensor.shape[:-1]) % 8 != 0 or tensor.shape[-1] % 16 != 0: + raise ValueError( + "FP8 execution requires the product of all dimensions except the last to be" + " divisible by 8 and the last dimension to be divisible by 16, but got tensor" + f" with dims={list(tensor.size())} (product of leading dims =" + f" {math.prod(tensor.shape[:-1])}, last dim = {tensor.shape[-1]})" + ) def is_bf16_compatible() -> bool: @@ -741,7 +747,9 @@ def __cuda_array_interface__(self): def torch_dtype_to_np_typestr(self): """Convert PyTorch dtype to numpy typestr.""" ret = _torch_dtype_to_np_typestr_dict.get(self.dtype) - assert ret is not None, f"Unsupported dtype: {self.dtype}" + if ret is None: + supported = ", ".join(str(d) for d in _torch_dtype_to_np_typestr_dict) + raise TypeError(f"Unsupported dtype: {self.dtype}. Supported dtypes: {supported}") return ret @@ -780,4 +788,7 @@ def convert_to_torch_tensor(tensor: Union[_WeakRefTensor, torch.Tensor]) -> torc return x if x is None: return None - raise TypeError(f"Invalid type {type(x)} to make weak ref") + raise TypeError( + f"Invalid type {type(x).__name__} to make weak ref. " + "Valid types are: torch.Tensor, tuple, list, dict, int, float, bool, and None." + ) From c021e7e353cf79e7f3dfd3dd1e071f395fe89763 Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Wed, 11 Mar 2026 10:18:58 -0700 Subject: [PATCH 268/521] [PyTorch] Fix fuser so it releases tensors properly (#2750) Signed-off-by: Kaining Zhong --- transformer_engine/pytorch/ops/fuser.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index bd3bc94b60..80386db2d9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -213,6 +213,7 @@ def backward( # Restore saved tensors saved_tensors = restore_from_saved(func_ctx.tensor_objects, func_ctx.saved_tensors) + func_ctx.tensor_objects = None # Unflatten list of saved tensors for ctx in basic_op_ctxs: From 7fb10d33621c75f4071487232d824fc159a24da7 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 11 Mar 2026 10:26:44 -0700 Subject: [PATCH 269/521] [PyTorch] Add dtype information to QuantizedTensorStorage class (#2676) * First pass Signed-off-by: Przemek Tredak * Cleaning the dtype usage in dequantize and distributed Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add fake_dtype to get_metadata Signed-off-by: Przemek Tredak * Fix to make_like Signed-off-by: Przemek Tredak --------- 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 | 5 +++++ transformer_engine/pytorch/distributed.py | 6 +++--- transformer_engine/pytorch/module/base.py | 2 ++ transformer_engine/pytorch/quantized_tensor.py | 9 +++++++-- .../pytorch/tensor/float8_blockwise_tensor.py | 2 +- transformer_engine/pytorch/tensor/float8_tensor.py | 4 +++- transformer_engine/pytorch/tensor/mxfp8_tensor.py | 2 +- transformer_engine/pytorch/tensor/nvfp4_tensor.py | 2 +- .../storage/float8_blockwise_tensor_storage.py | 13 ++++++++++--- .../tensor/storage/float8_tensor_storage.py | 10 ++++++++-- .../pytorch/tensor/storage/mxfp8_tensor_storage.py | 10 ++++++++-- .../pytorch/tensor/storage/nvfp4_tensor_storage.py | 14 +++++++++++--- transformer_engine/pytorch/utils.py | 1 + 13 files changed, 61 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 7e13cc105f..27dc87697f 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -285,6 +285,7 @@ std::pair Float8Quantizer::create_tensor( kwargs["fp8_dtype"] = py::cast(this->dtype); kwargs["data_transpose"] = transpose_py; kwargs["quantizer"] = this->quantizer; + kwargs["fake_dtype"] = GetATenDType(dtype); PyObject* result = PyObject_Call(reinterpret_cast(Float8TensorStoragePythonClass), args.ptr(), kwargs.ptr()); @@ -603,6 +604,7 @@ std::pair Float8CurrentScalingQuantizer::create_tenso kwargs["fp8_dtype"] = py::cast(this->dtype); kwargs["data_transpose"] = transpose_py; kwargs["quantizer"] = this->quantizer; + kwargs["fake_dtype"] = GetATenDType(dtype); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(Float8TensorStoragePythonClass), @@ -975,6 +977,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["fake_dtype"] = GetATenDType(dtype); py::tuple args(0); PyObject* result = @@ -1379,6 +1382,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["fake_dtype"] = GetATenDType(dtype); PyObject* result = PyObject_Call(reinterpret_cast(MXFP8TensorStoragePythonClass), args.ptr(), kwargs.ptr()); @@ -1788,6 +1792,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["fake_dtype"] = GetATenDType(dtype); py::tuple args(0); diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index 03ac9c1595..b80e58fe20 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -1099,7 +1099,7 @@ def _start_all_gather_fp8_blockwise( device = inp._columnwise_data.device else: raise ValueError("Got Float8BlockwiseQTensorStorage input tensor without any data") - dtype = torch.bfloat16 # Only has fp8 dtype. Guess BF16 for dequant. + dtype = inp._dtype else: raise ValueError( "Invalid type for input tensor (expected torch.Tensor or" @@ -1348,7 +1348,7 @@ def _all_gather_nvfp4( if inp._columnwise_data is not None: in_shape_t = inp._columnwise_data.size() device = inp._columnwise_data.device - dtype = torch.bfloat16 + dtype = inp._dtype else: raise ValueError( "Invalid type for input tensor (expected torch.Tensor or NVFP4TensorStorage, " @@ -1528,7 +1528,7 @@ def _all_gather_mxfp8( device = inp._columnwise_data.device else: raise ValueError("Got MXFP8 input tensor without any data") - dtype = torch.bfloat16 # Guess high-precision dtype. + dtype = inp._dtype else: raise ValueError( "Invalid type for input tensor (expected torch.Tensor or MXFP8TensorStorage, " diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 2d4583e936..28da4873f0 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -552,6 +552,7 @@ def fill_userbuffers_buffer_for_all_gather( data=global_tensor_data, fp8_scale_inv=local_tensor._scale_inv, fp8_dtype=local_tensor._fp8_dtype, + fake_dtype=local_tensor._dtype, quantizer=quantizer, ) return global_tensor, local_tensor @@ -626,6 +627,7 @@ def fill_userbuffers_buffer_for_all_gather( fp8_dtype=local_tensor._fp8_dtype, quantizer=quantizer, with_gemm_swizzled_scales=False, + fake_dtype=local_tensor._dtype, ) return global_tensor, local_tensor diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index cb697bc197..07171914f5 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -37,6 +37,7 @@ class QuantizedTensorStorage: XTensor should only implement the functionality needed to behave like regular torch.Tensor (like __torch_dispatch__).""" + _dtype: torch.dtype _quantizer: Optional[Quantizer] def update_usage( @@ -367,10 +368,13 @@ def __new__( shape: Iterable[int], dtype: torch.dtype, *, + fake_dtype: Optional[torch.dtype] = None, requires_grad: bool = False, device: Optional[torch.device] = None, stride: Optional[Iterable[int]] = None, ): + if fake_dtype is not None and fake_dtype != dtype: + raise ValueError(f"fake_dtype ({fake_dtype}) does not match dtype ({dtype})") # For stride, We are assuming only contiguous tensors # Calculate stride from shape if not provided. When creating this object from # C++ code, we provide the stride computed from shape in C++ to avoid the @@ -485,7 +489,7 @@ def clear(self): ) def __repr__(self, *, tensor_contents=None) -> str: - return f"{self.__class__.__name__}(data={self.dequantize(dtype=self.dtype)})" + return f"{self.__class__.__name__}(data={self.dequantize()})" def float(self) -> torch.Tensor: # pylint: disable=missing-function-docstring @@ -588,7 +592,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): def maybe_unwrap(arg): if isinstance(arg, QuantizedTensor): - return arg.dequantize(dtype=arg.dtype) + return arg.dequantize() return arg def maybe_update_inplace(arg, new_arg, schema_arg): @@ -671,6 +675,7 @@ def make_like( 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) def to_dtype(self, dtype: torch.dtype) -> QuantizedTensor: diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index e65730c015..0fae40f786 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -326,7 +326,7 @@ def __repr__(self, *, tensor_contents=None): return ( f"Float8BlockwiseQTensor(fp8_dtype={self._fp8_dtype}," f" is_2D_scaled={self._is_2D_scaled}," - f" data={self.dequantize(dtype=self.dtype)})" + f" data={self.dequantize()})" ) def quantize_( diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index c60bb2308d..9cc00855cd 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -182,6 +182,7 @@ def create_tensor_from_data( data=data, fp8_scale_inv=1 / self.scale, fp8_dtype=self.dtype, + fake_dtype=fake_dtype, requires_grad=requires_grad, data_transpose=None, quantizer=self, @@ -407,6 +408,7 @@ def create_tensor_from_data( data=data, fp8_scale_inv=torch.empty(1, dtype=torch.float32, device=data.device), fp8_dtype=self.dtype, + fake_dtype=fake_dtype, requires_grad=requires_grad, data_transpose=None, quantizer=self, @@ -498,7 +500,7 @@ def __repr__(self, *, tensor_contents=None): "Float8Tensor(" f"fp8_dtype={self._fp8_dtype}, " f"scale_inv={self._scale_inv.item()}, " - f"data={self.dequantize(dtype=self.dtype)}" + f"data={self.dequantize()}" ")" ) diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 96b6a67ea8..baff9cc2aa 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -296,7 +296,7 @@ def __new__( ) def __repr__(self, *, tensor_contents=None): - return f"MXFP8Tensor(fp8_dtype={self._fp8_dtype}, data={self.dequantize(dtype=self.dtype)})" + return f"MXFP8Tensor(fp8_dtype={self._fp8_dtype}, data={self.dequantize()})" def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 4314fd248c..8ed1b4682c 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -450,7 +450,7 @@ def __new__( return instance def __repr__(self, *, tensor_contents=None): - return f"NVFP4Tensor, data={self.dequantize(dtype=self.dtype)})" + return f"NVFP4Tensor, data={self.dequantize()})" def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 2a86717017..52e292125e 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -46,12 +46,14 @@ def __new__( quantizer: Quantizer, is_2D_scaled: bool, *args, + fake_dtype: Optional[torch.dtype] = None, **kwargs, ): if cls is Float8BlockwiseQTensorStorage: instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 else: - instance = super().__new__(cls, *args, **kwargs) + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data instance._quantizer = quantizer.copy() if quantizer is not None else None @@ -101,6 +103,7 @@ def get_metadata(self) -> Dict[str, Any]: "fp8_dtype": self._fp8_dtype, "quantizer": self._quantizer, "is_2D_scaled": self._is_2D_scaled, + "fake_dtype": self._dtype, } def prepare_for_saving( @@ -149,7 +152,9 @@ def _transpose_dq_columnwise_output(self, columnwise_dq: torch.Tensor) -> torch. permute_dims.append(0) return torch.permute(columnwise_dq, tuple(permute_dims)).contiguous() - def _dequantize_vectorwise(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: + def _dequantize_vectorwise(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if dtype is None: + dtype = self._dtype block_len = 128 q_M, q_K = 1, 1 @@ -211,10 +216,12 @@ def _dequantize_vectorwise(self, *, dtype: torch.dtype = torch.float32) -> torch return self._transpose_dq_columnwise_output(result) return result - def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ Construct plain PyTorch tensor from Float8BlockwiseQTensor """ + if dtype is None: + dtype = self._dtype block_len = 128 if not self._is_2D_scaled: return self._dequantize_vectorwise(dtype=dtype) diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index a815b366b2..0fb7966c2f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -75,14 +75,16 @@ def __new__( data: Optional[torch.Tensor], fp8_scale_inv: torch.Tensor, fp8_dtype: TE_DType, + fake_dtype: Optional[torch.dtype] = None, data_transpose: Optional[torch.Tensor] = None, quantizer: Optional[Quantizer] = None, **kwargs, ): if cls is Float8TensorStorage: instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 else: - instance = super().__new__(cls, *args, **kwargs) + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) instance._data = data instance._quantizer = quantizer.copy() if quantizer is not None else None instance._fp8_dtype = fp8_dtype @@ -130,6 +132,7 @@ def get_metadata(self) -> Dict[str, Any]: "fp8_dtype": self._fp8_dtype, "data_transpose": self._transpose, "quantizer": self._quantizer, + "fake_dtype": self._dtype, } def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], QuantizedTensorStorage]: @@ -159,8 +162,10 @@ def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = Tr return self._transpose raise ValueError("No data to get, both rowwise_data and columnwise_data are False") - def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Dequantize to a higher precision.""" + if dtype is None: + dtype = self._dtype return _FromFloat8Func.forward(None, self, dtype) def size(self, *args, **kwargs): @@ -192,6 +197,7 @@ def view(self, shape: torch.Size): data=out_data, fp8_scale_inv=self._scale_inv, fp8_dtype=self._fp8_dtype, + fake_dtype=self._dtype, data_transpose=out_transpose, quantizer=self._quantizer, ) diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 64344b78a1..7bbe809c9d 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -84,12 +84,14 @@ def __new__( quantizer: Optional[Quantizer], with_gemm_swizzled_scales: bool, *args, + fake_dtype: Optional[torch.dtype] = None, **kwargs, ): if cls is MXFP8TensorStorage: instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 else: - instance = super().__new__(cls, *args, **kwargs) + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data instance._rowwise_scale_inv = rowwise_scale_inv @@ -139,6 +141,7 @@ def get_metadata(self) -> Dict[str, Any]: "fp8_dtype": self._fp8_dtype, "quantizer": self._quantizer, "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, + "fake_dtype": self._dtype, } def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], MXFP8TensorStorage]: @@ -175,8 +178,10 @@ def get_data_tensors(self, rowwise_data: bool = True, columnwise_data: bool = Tr return self._columnwise_data raise ValueError("No data to get, both rowwise_data and columnwise_data are False") - def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Dequantize to a higher precision.""" + if dtype is None: + dtype = self._dtype return _FromMXFP8Func.forward(None, self, dtype) def size(self, *args, **kwargs): @@ -238,6 +243,7 @@ def view(self, shape: torch.Size): fp8_dtype=self._fp8_dtype, quantizer=self._quantizer, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, + fake_dtype=self._dtype, ) def __repr__(self): diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 3e8bf3f2f3..fb163c9032 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -106,10 +106,14 @@ def __new__( quantizer: Optional[Quantizer], with_gemm_swizzled_scales: bool, *args, + fake_dtype: Optional[torch.dtype] = None, **kwargs, ): - - instance = super().__new__(cls, *args, **kwargs) + if cls is NVFP4TensorStorage: + instance = object.__new__(cls) + instance._dtype = fake_dtype if fake_dtype is not None else torch.float32 + else: + instance = super().__new__(cls, *args, fake_dtype=fake_dtype, **kwargs) instance._rowwise_data = rowwise_data instance._columnwise_data = columnwise_data @@ -168,6 +172,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, + "fake_dtype": self._dtype, } def prepare_for_saving(self) -> Tuple[list[Optional[torch.Tensor]], NVFP4TensorStorage]: @@ -204,8 +209,10 @@ def get_data_tensors(self): """Get this Tensor's data.""" return self._rowwise_data, self._columnwise_data - def dequantize(self, *, dtype: torch.dtype = torch.float32) -> torch.Tensor: + def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Dequantize to a higher precision.""" + if dtype is None: + dtype = self._dtype return _FromNVFP4Func.forward(None, self, dtype) def size(self, dim: Optional[int] = None) -> Union[torch.Size, int]: @@ -295,6 +302,7 @@ def view(self, shape: torch.Size): quantizer=self._quantizer, fp4_dtype=self._fp4_dtype, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, + fake_dtype=self._dtype, ) def __repr__(self): diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index a23e822f91..db2f28aa47 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -242,6 +242,7 @@ def forward( fp8_dtype=mixed_x_layer._fp8_dtype, data=x.squeeze(split_dim) if squeeze else x, shape=x.squeeze(split_dim).shape if squeeze else x.shape, + fake_dtype=mixed_x_layer._dtype, quantizer=mixed_x_layer._quantizer, ) for x in torch.split( From 4c5b1a2de5cb420b3a7b4ea74b78ecdc138055ac Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 11 Mar 2026 16:01:08 -0700 Subject: [PATCH 270/521] [JAX] Change dtype of intermediate result aval of fused_topk_and_score_function_fwd to fp32 (#2752) * Fixed aval for intermediate results (softmaxed/sigmoided logits) to pass as residuals to CompType, which is currently fp32. This prevents incorrect reading of this buffer when logits dtype used are not fp32 Signed-off-by: tdophung * address comments of inconsitency in style and NVTE CHECK for fp32 type Signed-off-by: tdophung * revert the remaining Comptype checking, address greptile suggestion Signed-off-by: tdophung * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../jax/cpp_extensions/router.py | 4 ++- .../jax/csrc/extensions/router.cpp | 25 +++++++++++++++---- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index 031ab483a0..f2affacdaa 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -73,7 +73,9 @@ def abstract( i_shape = logits_aval.shape probs_aval = logits_aval.update(shape=i_shape, dtype=i_dtype) routing_map_aval = logits_aval.update(shape=i_shape, dtype=jnp.bool_) - intermediate_aval = logits_aval.update(shape=i_shape, dtype=i_dtype) + # The CUDA kernel always uses float32 (CompType) for intermediate + # computations (softmax/sigmoid values saved for backward). + intermediate_aval = logits_aval.update(shape=i_shape, dtype=jnp.float32) return probs_aval, routing_map_aval, intermediate_aval @staticmethod diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp index 0190d3fd75..c81671f104 100644 --- a/transformer_engine/jax/csrc/extensions/router.cpp +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -41,7 +41,14 @@ Error_Type FusedTopkWithScoreFunctionForwardFFI( auto logits_tensor = TensorWrapper(logits, flat_shape, dtype); auto probs_tensor = TensorWrapper(probs, flat_shape, dtype); auto routing_map_tensor = TensorWrapper(routing_map, flat_shape, DType::kByte); - auto intermediate_tensor = TensorWrapper(intermediate, flat_shape, dtype); + // intermediate is always float32 (CompType) regardless of logits dtype. + auto intermediate_dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf->element_type()); + NVTE_CHECK( + intermediate_dtype == DType::kFloat32, + "intermediate_output must be float32 (CompType); got dtype ", + static_cast(intermediate_dtype), + ". Check FusedTopkWithScoreFunctionFwdPrimitive.abstract in cpp_extensions/router.py."); + auto intermediate_tensor = TensorWrapper(intermediate, flat_shape, DType::kFloat32); if (compute_aux_scores) { nvte_fused_score_for_moe_aux_loss_forward( @@ -97,7 +104,14 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( Result_Type grad_logits_buf, // [num_tokens, num_experts] int64_t topk, int64_t use_pre_softmax, double scaling_factor, JAXX_Score_Function score_function, int64_t compute_aux_scores) { - auto dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf.element_type()); + // intermediate is always float32 (CompType) regardless of logits dtype. + auto intermediate_dtype = convert_ffi_datatype_to_te_dtype(intermediate_buf.element_type()); + NVTE_CHECK( + intermediate_dtype == DType::kFloat32, + "intermediate_output must be float32 (CompType); got dtype ", + static_cast(intermediate_dtype), + ". Check FusedTopkWithScoreFunctionFwdPrimitive.abstract in cpp_extensions/router.py."); + auto grad_dtype = convert_ffi_datatype_to_te_dtype(grad_probs_buf.element_type()); auto dims = intermediate_buf.dimensions(); auto num_tokens = static_cast(product(dims, 0, dims.size() - 1)); auto num_experts = static_cast(dims[dims.size() - 1]); @@ -105,9 +119,10 @@ Error_Type FusedTopkWithScoreFunctionBackwardFFI( auto flat_shape = std::vector{static_cast(num_tokens), static_cast(num_experts)}; - auto intermediate_tensor = TensorWrapper(intermediate_buf.untyped_data(), flat_shape, dtype); - auto grad_probs_tensor = TensorWrapper(grad_probs_buf.untyped_data(), flat_shape, dtype); - auto grad_logits_tensor = TensorWrapper(grad_logits_buf->untyped_data(), flat_shape, dtype); + auto intermediate_tensor = + TensorWrapper(intermediate_buf.untyped_data(), flat_shape, DType::kFloat32); + auto grad_probs_tensor = TensorWrapper(grad_probs_buf.untyped_data(), flat_shape, grad_dtype); + auto grad_logits_tensor = TensorWrapper(grad_logits_buf->untyped_data(), flat_shape, grad_dtype); if (compute_aux_scores) { nvte_fused_score_for_moe_aux_loss_backward(intermediate_tensor.data(), grad_probs_tensor.data(), From 06a23e364009a09c34bb51bdd1e0c24247f7f4d3 Mon Sep 17 00:00:00 2001 From: vasunvidia <108759426+vasunvidia@users.noreply.github.com> Date: Wed, 11 Mar 2026 22:40:39 -0700 Subject: [PATCH 271/521] Initial commit to pass scale as Tensor for multi_tensor_scale op (#2594) * Initial commit to pass scale as Tensor for multi_tensor_scale op 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 * Enable capturable mode for optimizer if store_param_remainders is passed but not actually enabled Signed-off-by: Vasudevan Rengasamy * Revert "Enable capturable mode for optimizer if store_param_remainders is passed but not actually enabled" This reverts commit 74a9bccf0fadd4159f70d28da49a533ea7c76108. Signed-off-by: Vasudevan Rengasamy * Apply suggestion from @greptile-apps[bot] 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> Signed-off-by: Vasudevan Rengasamy * Change noop_flag to is_infinite 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 * Update transformer_engine/pytorch/csrc/extensions.h Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Remove duplication Signed-off-by: Kirthi Shankar Sivamani * Add test for scale tensor cuda Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Vasudevan Rengasamy Signed-off-by: Kirthi Shankar Sivamani 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: Kirthi Shankar Sivamani --- tests/pytorch/test_multi_tensor.py | 111 ++++++++++++ .../include/transformer_engine/multi_tensor.h | 21 ++- .../common/multi_tensor/scale.cu | 170 +++++++++++------- transformer_engine/pytorch/csrc/extensions.h | 4 + .../csrc/extensions/multi_tensor/scale.cpp | 20 ++- .../pytorch/csrc/extensions/pybind.cpp | 3 + .../pytorch/optimizers/__init__.py | 1 + 7 files changed, 261 insertions(+), 69 deletions(-) diff --git a/tests/pytorch/test_multi_tensor.py b/tests/pytorch/test_multi_tensor.py index b7caa094ae..6f1b6948ab 100644 --- a/tests/pytorch/test_multi_tensor.py +++ b/tests/pytorch/test_multi_tensor.py @@ -137,6 +137,117 @@ def find_inf( ) +@pytest.mark.parametrize("input_size_pair", input_size_pairs) +@pytest.mark.parametrize("applier", appliers) +@pytest.mark.parametrize("repeat", [1, 55]) +@pytest.mark.parametrize("in_type", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("out_type", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("inplace", [False, True]) +def test_multi_tensor_scale_tensor(input_size_pair, applier, repeat, in_type, out_type, inplace): + if inplace is True and (out_type is not in_type): + pytest.skip("inplace=True and out_type != in_type is not supported.") + elif (in_type == torch.float16 and out_type == torch.bfloat16) or ( + in_type == torch.bfloat16 and out_type == torch.float16 + ): + pytest.skip("float16 to bfloat16 is not necessary and vice versa.") + + device = torch.device("cuda") + scale = 4.0 + inv_scale_cuda = torch.tensor([1.0 / scale], dtype=torch.float32, device=device) + overflow_buf = torch.zeros(1, dtype=torch.int32, device=device) + ref = torch.tensor([1.0], dtype=torch.float32, device=device) + sizea, sizeb = input_size_pair + + def downscale(sizea, sizeb, applier, repeat, in_type, out_type, inplace=False): + overflow_buf.zero_() + a = torch.full([sizea], scale, dtype=torch.float32, device=device) + b = torch.full([sizeb], scale, dtype=torch.float32, device=device) + + out_list = [] + for _ in range(repeat): + out_list += [a.clone().to(out_type), b.clone().to(out_type)] + + if inplace: + in_list = out_list + else: + in_list = [out.clone().to(in_type) for out in out_list] + + applier(tex.multi_tensor_scale_tensor, overflow_buf, [in_list, out_list], inv_scale_cuda) + + assert all([torch.allclose(out, ref.to(out_type)) for out in out_list]) + assert overflow_buf.item() == 0 + + def find_inf( + sizea, + sizeb, + applier, + repeat, + in_type, + out_type, + t, + ind, + val, + inplace=False, + ): + overflow_buf.zero_() + a = torch.full([sizea], scale, dtype=torch.float32, device=device) + b = torch.full([sizeb], scale, dtype=torch.float32, device=device) + + out_list = [] + for _ in range(repeat): + out_list += [a.clone().to(out_type), b.clone().to(out_type)] + + if inplace: + in_list = out_list + else: + in_list = [out.clone().to(in_type) for out in out_list] + + applier(tex.multi_tensor_scale_tensor, overflow_buf, [in_list, out_list], inv_scale_cuda) + + overflow_buf.zero_() + in_list[t][ind] = val + applier(tex.multi_tensor_scale_tensor, overflow_buf, [in_list, out_list], inv_scale_cuda) + assert overflow_buf.item() > 0 + + downscale(sizea, sizeb, applier, repeat, in_type, out_type, inplace=inplace) + find_inf( + sizea, + sizeb, + applier, + repeat, + in_type, + out_type, + 0, + 0, + float("nan"), + inplace=inplace, + ) + find_inf( + sizea, + sizeb, + applier, + repeat, + in_type, + out_type, + 2 * repeat - 1, + sizeb - 1, + float("inf"), + inplace=inplace, + ) + find_inf( + sizea, + sizeb, + applier, + repeat, + in_type, + out_type, + 2 * (repeat // 2), + sizea // 2, + float("inf"), + inplace=inplace, + ) + + @pytest.mark.parametrize("input_size_pair", input_size_pairs) @pytest.mark.parametrize("applier", appliers) @pytest.mark.parametrize("repeat", [1, 55]) diff --git a/transformer_engine/common/include/transformer_engine/multi_tensor.h b/transformer_engine/common/include/transformer_engine/multi_tensor.h index b5eadcf678..09ab260f15 100644 --- a/transformer_engine/common/include/transformer_engine/multi_tensor.h +++ b/transformer_engine/common/include/transformer_engine/multi_tensor.h @@ -233,17 +233,34 @@ void nvte_multi_tensor_sgd_cuda(int chunk_size, NVTETensor noop_flag, NVTETensor * \warning This API is **experimental** and subject to change. * * \param[in] chunk_size Number of tensor elements processed by a CUDA block. - * \param[in] noop_flag If this single element tensor has non-zero value, kernel will exit immediately. + * \param[out] is_infinite Whether the kernel detected a non-finite input value. * \param[in,out] tensor_lists 2D array of input tensors. * \param[in] num_tensor_lists Size (dim0) of tensor_lists. * \param[in] num_tensors_per_list Size (dim1) of tensor_lists. * \param[in] scale Scalar for the scaling operation. * \param[in] stream CUDA stream used for this operation. */ -void nvte_multi_tensor_scale_cuda(int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, +void nvte_multi_tensor_scale_cuda(int chunk_size, NVTETensor is_infinite, NVTETensor **tensor_lists, const size_t num_tensor_lists, const size_t num_tensors_per_list, float scale, cudaStream_t stream); +/*! \brief Check overflow and scale a list of tensors. scale is tensor input. + * + * \warning This API is **experimental** and subject to change. + * + * \param[in] chunk_size Number of tensor elements processed by a CUDA block. + * \param[out] is_infinite Whether the kernel detected a non-finite input value. + * \param[in,out] tensor_lists 2D array of input tensors. + * \param[in] num_tensor_lists Size (dim0) of tensor_lists. + * \param[in] num_tensors_per_list Size (dim1) of tensor_lists. + * \param[in] scale Tensor for the scaling operation. + * \param[in] stream CUDA stream used for this operation. + */ +void nvte_multi_tensor_scale_tensor_cuda(int chunk_size, NVTETensor is_infinite, + NVTETensor **tensor_lists, const size_t num_tensor_lists, + const size_t num_tensors_per_list, NVTETensor scale, + cudaStream_t stream); + /*! \brief Check overflow and scale a list of tensors. * * \warning This API is **experimental** and subject to change. diff --git a/transformer_engine/common/multi_tensor/scale.cu b/transformer_engine/common/multi_tensor/scale.cu index b3266200c4..6b9b66faa8 100644 --- a/transformer_engine/common/multi_tensor/scale.cu +++ b/transformer_engine/common/multi_tensor/scale.cu @@ -33,97 +33,141 @@ __device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset, int s ((LT *)dst)[dst_offset] = ((LT *)src)[src_offset]; // NOLINT(*) } -template -struct ScaleFunctor { - __device__ __forceinline__ void operator()(int chunk_size, volatile int *noop_gmem, - TensorListMetadata<2> &tl, // NOLINT(*) - float scale) { - // I'd like this kernel to propagate infs/nans. - // if(*noop_gmem == 1) - // return; - - int tensor_loc = tl.block_to_tensor[blockIdx.x]; - int chunk_idx = tl.block_to_chunk[blockIdx.x]; - int n = tl.sizes[tensor_loc]; - - in_t *in = reinterpret_cast(tl.addresses[0][tensor_loc]); - in += chunk_idx * chunk_size; - - out_t *out = reinterpret_cast(tl.addresses[1][tensor_loc]); - out += chunk_idx * chunk_size; - - n -= chunk_idx * chunk_size; - - bool finite = true; - in_t r_in[ILP]; - out_t r_out[ILP]; - - // to make things simple, we put aligned case in a different code path - if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(in) && is_aligned(out)) { - for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; - i_start += blockDim.x) { - // load - load_store(r_in, in, 0, i_start); +__device__ __forceinline__ float get_scale_value(float scale) { return scale; } + +__device__ __forceinline__ float get_scale_value(const float *scale_ptr) { return *scale_ptr; } + +template +__device__ __forceinline__ void scale_chunk(int chunk_size, volatile int *is_infinite_gmem, + TensorListMetadata<2> &tl, scale_t scale_arg) { + // I'd like this kernel to propagate infs/nans. + // if(*noop_gmem == 1) + // return; + const float scale = get_scale_value(scale_arg); + int tensor_loc = tl.block_to_tensor[blockIdx.x]; + int chunk_idx = tl.block_to_chunk[blockIdx.x]; + int n = tl.sizes[tensor_loc]; + + in_t *in = reinterpret_cast(tl.addresses[0][tensor_loc]); + in += chunk_idx * chunk_size; + + out_t *out = reinterpret_cast(tl.addresses[1][tensor_loc]); + out += chunk_idx * chunk_size; + + n -= chunk_idx * chunk_size; + + bool finite = true; + in_t r_in[ILP]; + out_t r_out[ILP]; + + // to make things simple, we put aligned case in a different code path + if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(in) && is_aligned(out)) { + for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; + i_start += blockDim.x) { + // load + load_store(r_in, in, 0, i_start); #pragma unroll - for (int ii = 0; ii < ILP; ii++) { - r_out[ii] = static_cast(r_in[ii]) * scale; - finite = finite && isfinite(static_cast(r_in[ii])); - } - // store - load_store(out, r_out, i_start, 0); + for (int ii = 0; ii < ILP; ii++) { + r_out[ii] = static_cast(r_in[ii]) * scale; + finite = finite && isfinite(static_cast(r_in[ii])); } - } else { - // Non-divergent exit condition for __syncthreads, not necessary here - for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) { + // store + load_store(out, r_out, i_start, 0); + } + } else { + // Non-divergent exit condition for __syncthreads, not necessary here + for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * ILP) { #pragma unroll - for (int ii = 0; ii < ILP; ii++) { - r_in[ii] = 0.f; - int i = i_start + threadIdx.x + ii * blockDim.x; - if (i < n && i < chunk_size) r_in[ii] = in[i]; - } - // note for clarification to future michael: - // From a pure memory dependency perspective, there's likely no point unrolling - // the write loop, since writes just fire off once their LDGs arrive. - // Put another way, the STGs are dependent on the LDGs, but not on each other. - // There is still compute ILP benefit from unrolling the loop though. + for (int ii = 0; ii < ILP; ii++) { + r_in[ii] = 0.f; + int i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) r_in[ii] = in[i]; + } + // From a pure memory dependency perspective, there's likely no point unrolling + // the write loop, since writes just fire off once their LDGs arrive. + // Put another way, the STGs are dependent on the LDGs, but not on each other. + // There is still compute ILP benefit from unrolling the loop though. #pragma unroll - for (int ii = 0; ii < ILP; ii++) { - r_out[ii] = static_cast(r_in[ii]) * scale; - finite = finite && isfinite(static_cast(r_in[ii])); - } + for (int ii = 0; ii < ILP; ii++) { + r_out[ii] = static_cast(r_in[ii]) * scale; + finite = finite && isfinite(static_cast(r_in[ii])); + } #pragma unroll - for (int ii = 0; ii < ILP; ii++) { - int i = i_start + threadIdx.x + ii * blockDim.x; - if (i < n && i < chunk_size) out[i] = r_out[ii]; - } + for (int ii = 0; ii < ILP; ii++) { + int i = i_start + threadIdx.x + ii * blockDim.x; + if (i < n && i < chunk_size) out[i] = r_out[ii]; } } - if (!finite) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok. + } + if (!finite) *is_infinite_gmem = 1; // Blindly fire off a write. These will race but that's ok. +} + +template +struct ScaleFunctor { + __device__ __forceinline__ void operator()(int chunk_size, volatile int *is_infinite_gmem, + TensorListMetadata<2> &tl, // NOLINT(*) + float scale) { + scale_chunk(chunk_size, is_infinite_gmem, tl, scale); } }; -void multi_tensor_scale_cuda(int chunk_size, Tensor noop_flag, +template +struct ScalePtrFunctor { + __device__ __forceinline__ void operator()(int chunk_size, volatile int *is_infinite_gmem, + TensorListMetadata<2> &tl, // NOLINT(*) + float *scale_ptr) { + scale_chunk(chunk_size, is_infinite_gmem, tl, scale_ptr); + } +}; + +void multi_tensor_scale_cuda(int chunk_size, Tensor is_infinite, std::vector> tensor_lists, float scale, cudaStream_t stream) { TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( tensor_lists[0][0]->dtype(), p_in_type, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( tensor_lists[1][0]->dtype(), g_in_type, - multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, + multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, is_infinite, tensor_lists, ScaleFunctor(), stream, scale);)) NVTE_CHECK_CUDA(cudaGetLastError()); } +void multi_tensor_scale_tensor_cuda(int chunk_size, Tensor is_infinite, + std::vector> tensor_lists, float *scale, + cudaStream_t stream) { + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + tensor_lists[0][0]->dtype(), p_in_type, + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( + tensor_lists[1][0]->dtype(), g_in_type, + multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, is_infinite, tensor_lists, + ScalePtrFunctor(), stream, scale);)) + NVTE_CHECK_CUDA(cudaGetLastError()); +} + } // namespace multi_tensor_scale } // namespace transformer_engine -void nvte_multi_tensor_scale_cuda(int chunk_size, NVTETensor noop_flag, NVTETensor **tensor_lists, +void nvte_multi_tensor_scale_cuda(int chunk_size, NVTETensor is_infinite, NVTETensor **tensor_lists, const size_t num_tensor_lists, const size_t num_tensors_per_list, float scale, cudaStream_t stream) { NVTE_API_CALL(nvte_multi_tensor_scale_cuda); using namespace transformer_engine; multi_tensor_scale::multi_tensor_scale_cuda( - chunk_size, *convertNVTETensorCheck(noop_flag), + chunk_size, *convertNVTETensorCheck(is_infinite), convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list), scale, stream); } + +void nvte_multi_tensor_scale_tensor_cuda(int chunk_size, NVTETensor is_infinite, + NVTETensor **tensor_lists, const size_t num_tensor_lists, + const size_t num_tensors_per_list, NVTETensor scale, + cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_scale_tensor_cuda); + using namespace transformer_engine; + + Tensor *scale_tensor = convertNVTETensorCheck(scale); + multi_tensor_scale::multi_tensor_scale_tensor_cuda( + chunk_size, *convertNVTETensorCheck(is_infinite), + convert_tensor_array(tensor_lists, num_tensor_lists, num_tensors_per_list), + reinterpret_cast(scale_tensor->data.dptr), stream); +} diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index e4d4e5094c..d8f00becb5 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -466,6 +466,10 @@ at::Tensor thd_get_partitioned_indices(const at::Tensor &cu_seqlens, int total_t void multi_tensor_scale_cuda(int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, float scale); +void multi_tensor_scale_tensor_cuda(int chunk_size, at::Tensor is_infinite, + std::vector> tensor_lists, + at::Tensor scale); + std::tuple multi_tensor_l2norm_cuda( int chunk_size, at::Tensor noop_flag, std::vector> tensor_lists, at::optional per_tensor_python); diff --git a/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp b/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp index 4bb83bfeed..687eb34f32 100644 --- a/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp +++ b/transformer_engine/pytorch/csrc/extensions/multi_tensor/scale.cpp @@ -8,14 +8,26 @@ namespace transformer_engine::pytorch { -void multi_tensor_scale_cuda(int chunk_size, at::Tensor noop_flag, +void multi_tensor_scale_cuda(int chunk_size, at::Tensor is_infinite, std::vector> tensor_lists, float scale) { - auto noop_flag_cu = makeTransformerEngineTensor(noop_flag); + auto is_infinite_cu = makeTransformerEngineTensor(is_infinite); auto [_, __, tensor_lists_ptr, num_lists, num_tensors] = makeTransformerEngineTensorList(tensor_lists); - nvte_multi_tensor_scale_cuda(chunk_size, noop_flag_cu.data(), tensor_lists_ptr.data(), num_lists, - num_tensors, scale, at::cuda::getCurrentCUDAStream()); + nvte_multi_tensor_scale_cuda(chunk_size, is_infinite_cu.data(), tensor_lists_ptr.data(), + num_lists, num_tensors, scale, at::cuda::getCurrentCUDAStream()); +} + +void multi_tensor_scale_tensor_cuda(int chunk_size, at::Tensor is_infinite, + std::vector> tensor_lists, + at::Tensor scale) { + auto is_infinite_cu = makeTransformerEngineTensor(is_infinite); + auto scale_cu = makeTransformerEngineTensor(scale); + auto [_, __, tensor_lists_ptr, num_lists, num_tensors] = + makeTransformerEngineTensorList(tensor_lists); + nvte_multi_tensor_scale_tensor_cuda(chunk_size, is_infinite_cu.data(), tensor_lists_ptr.data(), + num_lists, num_tensors, scale_cu.data(), + at::cuda::getCurrentCUDAStream()); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 8302a13010..817a481808 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -491,6 +491,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("multi_tensor_scale", &transformer_engine::pytorch::multi_tensor_scale_cuda, "Fused overflow check + scale for a list of contiguous tensors", py::call_guard()); + m.def("multi_tensor_scale_tensor", &transformer_engine::pytorch::multi_tensor_scale_tensor_cuda, + "Fused overflow check + scale for a list of contiguous tensors with scale passed as tensor", + py::call_guard()); m.def("multi_tensor_l2norm", &transformer_engine::pytorch::multi_tensor_l2norm_cuda, "Computes L2 norm for a list of contiguous tensors", py::call_guard()); diff --git a/transformer_engine/pytorch/optimizers/__init__.py b/transformer_engine/pytorch/optimizers/__init__.py index 792eab094a..7220f1924a 100644 --- a/transformer_engine/pytorch/optimizers/__init__.py +++ b/transformer_engine/pytorch/optimizers/__init__.py @@ -5,6 +5,7 @@ """Fused optimizers and multi-tensor kernels.""" from transformer_engine_torch import ( multi_tensor_scale, + multi_tensor_scale_tensor, multi_tensor_l2norm, multi_tensor_unscale_l2norm, multi_tensor_adam, From ef703e551e154eb7c77799a708aab459f8ebe7ba Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 12 Mar 2026 10:44:57 -0700 Subject: [PATCH 272/521] [Core] MXFP8 grouped GEMM + tensor-scaled FP8 fixes (#2748) * MXFP8 grouped GEMM + tensor-scaled FP8 fixes Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change version to 13.3 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Random padding condition shouldnt be done for mxfp8 Signed-off-by: vthumbe1503 * Remove incorrect comment Signed-off-by: vthumbe1503 * CUBLAS > 13.2 is enough Signed-off-by: vthumbe1503 * CUBLAS version needed for MXFP8 indeed seems to be 13.3 Signed-off-by: vthumbe1503 * Accidental line removal added back. Plus need changes ci t trigger Add documentation for scaling factors in common.h Signed-off-by: vthumbe1503 * Update cuBLAS version requirement for MXFP8 support Signed-off-by: vthumbe1503 * grouped gemm: address code review comments - Replace nvte_set/get_grouped_tensor_swizzled_scales with nvte_set_grouped_tensor_param - Add host-side validation: A and B must use same scaling mode (both MXFP8 or both tensor scaling) - Add host-side validation: A and B must both be FP8 or both non-FP8; restrict inputs to FP8/BF16 - Restrict output (C/D) to BF16/FP32; remove FP16 from supported types - Refactor workspace allocation: replace manual offset arithmetic with moving pointer pattern - Use void* + NVTEScalingMode in setup kernel instead of separate float*/char* scale params - Extract use_columnwise(swap_dims) helper to eliminate duplicated MXFP8 columnwise blocks - Split set_fp8_scale_pointers into set_fp8_scale_pointers / set_mxfp8_scale_pointers - Remove scale_inv_ptrs from GroupedOperandSelection; pass workspace pointers directly - Move swizzled-scales validation into validate_grouped_gemm_inputs for fail-fast behavior - Add use_split_accumulator to GroupedMatmulConfig (Hopper only, default false) - Add FP8 test case with per-tensor scales; add BF16/MXFP8 shape-varying test cases Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Signed-off-by: vthumbe1503 Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Pawel Gadzinski --- tests/cpp/operator/test_grouped_gemm.cu | 111 +++++++- tests/cpp/test_common.cu | 116 ++++++-- tests/cpp/test_common.h | 1 + transformer_engine/common/common.h | 3 +- transformer_engine/common/gemm/config.h | 9 +- .../common/gemm/cublaslt_grouped_gemm.cu | 252 ++++++++++++++---- .../common/include/transformer_engine/gemm.h | 9 + 7 files changed, 413 insertions(+), 88 deletions(-) diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu index a7aabbbcb6..34bb729b25 100644 --- a/tests/cpp/operator/test_grouped_gemm.cu +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -20,6 +20,7 @@ #include #include #include +#include #include #include "../test_common.h" @@ -32,6 +33,7 @@ namespace { enum class InputCase { kFP8Current, kBF16, + kMXFP8, }; enum class ShapeCase { @@ -44,8 +46,8 @@ enum class ShapeCase { 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: 6 pointer arrays (A, B, C, D, alpha, beta) + 6 int arrays (a_rows, a_cols, b_rows, b_cols, d_rows, d_cols) - size_t size = 6 * ptr_bytes + 6 * int_bytes; + // 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; @@ -53,7 +55,20 @@ size_t grouped_setup_workspace_size(const size_t num_tensors) { Tensor make_fp8_operand(const std::string& name, const std::vector& shape) { Tensor input_fp32(name + "_fp32", shape, DType::kFloat32); - fillUniform(&input_fp32); + + const size_t numel = shape[0] * shape[1]; + std::vector data(numel); + std::mt19937 gen(std::hash{}(name)); + // Random mean and stddev -> different amax per tensor -> different scales + std::uniform_real_distribution param_dis(0.1f, 10.0f); + float mean = param_dis(gen); + float stddev = param_dis(gen); + std::normal_distribution dis(mean, stddev); + for (size_t i = 0; i < numel; ++i) { + data[i] = dis(gen); + } + NVTE_CHECK_CUDA(cudaMemcpy(input_fp32.rowwise_dptr(), data.data(), + numel * sizeof(float), cudaMemcpyHostToDevice)); Tensor fp8(name, shape, TypeInfo::dtype, true, true, NVTE_DELAYED_TENSOR_SCALING); @@ -73,6 +88,64 @@ 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 +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 + 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, + 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); + 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)); + } + + // Swizzle scales for GEMM + nvte_swizzle_scaling_factors(mxfp8.data(), mxfp8_swizzled.data(), 0); + + // Sync to ensure operations are complete + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + + return mxfp8_swizzled; +} + struct TestParams { InputCase input_case; bool transa; @@ -88,16 +161,16 @@ struct TestParams { std::vector> make_shapes(ShapeCase scase) { switch (scase) { case ShapeCase::kAllSame: - return {{64, 64, 32}, {64, 64, 32}, {64, 64, 32}}; + return {{128, 256, 384}, {128, 256, 384}, {128, 256, 384}}; case ShapeCase::kSameFirst: // Same M (first dim), varying N and K - return {{64, 80, 32}, {64, 96, 48}, {64, 112, 64}}; + return {{128, 256, 384}, {128, 384, 512}, {128, 512, 640}}; case ShapeCase::kSameLast: // Same N (last dim), varying M and K - return {{64, 80, 32}, {80, 80, 48}, {96, 80, 64}}; + return {{128, 256, 384}, {256, 256, 512}, {384, 256, 640}}; case ShapeCase::kAllDifferent: default: - return {{64, 96, 32}, {80, 112, 48}, {96, 128, 64}}; + return {{128, 256, 384}, {256, 384, 512}, {384, 512, 640}}; } } @@ -138,6 +211,13 @@ void run_grouped_gemm_case(const TestParams& params) { 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}, @@ -246,7 +326,9 @@ void run_grouped_gemm_case(const TestParams& params) { cublas_ws.data(), nullptr, // config (use defaults) 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])), @@ -277,7 +359,7 @@ TEST_P(GroupedGemmTest, CompareWithMultiTensorGemm) { } std::string MakeGroupedGemmTestName(const testing::TestParamInfo& info) { - constexpr const char* kInputNames[] = {"FP8Current", "BF16"}; + constexpr const char* kInputNames[] = {"FP8Current", "BF16", "MXFP8"}; constexpr const char* kShapeNames[] = {"AllSame", "SameM", "SameN", "AllDiff"}; const std::string layout = std::string("ta") + (info.param.transa ? "T" : "N") + "tb" + (info.param.transb ? "T" : "N"); @@ -288,16 +370,27 @@ std::string MakeGroupedGemmTestName(const testing::TestParamInfo kTestParams = { - // Basic tests + // 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}, + // 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}, // Test NULL C (valid when beta=0) {InputCase::kBF16, false, false, ShapeCase::kAllSame, 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}, + // MXFP8 with NULL C + {InputCase::kMXFP8, true, false, ShapeCase::kAllSame, true}, }; INSTANTIATE_TEST_SUITE_P(OperatorTest, diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index b64ae24131..ff26d1b6c5 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -1061,7 +1061,14 @@ std::array get_scale_tensor_dims(const size_t rows, GroupedBuffers build_grouped_tensor(const std::vector& tensors, const NVTEScalingMode scaling_mode) { NVTE_CHECK(!tensors.empty(), "No tensors provided for grouped tensor build."); - const NVTEShape shape = tensors[0]->rowwise_shape(); + + // Check which data layouts are available (all tensors must have the same) + const bool has_rowwise = tensors[0]->rowwise(); + 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; @@ -1076,7 +1083,8 @@ 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 = tensors[i]->rowwise_shape(); + const auto s = has_rowwise ? tensors[i]->rowwise_shape() + : tensors[i]->columnwise_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]); @@ -1105,10 +1113,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) { - offsets[i] = offsets[i - 1] + numel(i - 1) + random_padding(); + offsets[i] = offsets[i - 1] + numel(i - 1) + (use_random_padding ? random_padding() : 0); } } else { for (size_t i = 0; i < num_tensors; ++i) { @@ -1146,21 +1155,24 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, : (logical_first * logical_last); const size_t total_bytes = static_cast(total_elems) * elem_size; - 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, - tensors[i]->rowwise_dptr(), - grouped.tensor_bytes[i], - cudaMemcpyDeviceToDevice)); - } - - NVTEBasicTensor data_tensor{grouped.data.get(), static_cast(dtype), grouped.logical_shape}; NVTEGroupedTensor h = grouped.handle.get(); - nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseData, &data_tensor, sizeof(data_tensor)); - const bool include_columnwise = isFp8Type(dtype) || isFp4Type(dtype); - if (include_columnwise) { + // Copy rowwise data if available + 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, + tensors[i]->rowwise_dptr(), + grouped.tensor_bytes[i], + cudaMemcpyDeviceToDevice)); + } + NVTEBasicTensor data_tensor{grouped.data.get(), static_cast(dtype), grouped.logical_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseData, &data_tensor, sizeof(data_tensor)); + } + + // Copy columnwise data if available + 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; @@ -1202,11 +1214,17 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, nvte_set_grouped_tensor_param(h, kNVTEGroupedTensorOffsets, &off_tensor, sizeof(off_tensor)); } - if (isFp8Type(dtype)) { + 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 std::vector scale_inv_cpu(num_tensors, 1.f); for (size_t i = 0; i < num_tensors; ++i) { tensors[i]->to_cpu(); - scale_inv_cpu[i] = tensors[i]->rowwise_cpu_scale_inv_ptr()[0]; + if (has_rowwise) { + scale_inv_cpu[i] = tensors[i]->rowwise_cpu_scale_inv_ptr()[0]; + } else { + scale_inv_cpu[i] = tensors[i]->columnwise_cpu_scale_inv_ptr()[0]; + } } grouped.scale_inv = cuda_alloc(sizeof(float) * num_tensors); NVTE_CHECK_CUDA(cudaMemcpy(grouped.scale_inv.get(), scale_inv_cpu.data(), @@ -1217,6 +1235,68 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, sizeof(scale_tensor)); 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 + if (has_rowwise) { + auto [row_buffer, row_total] = gather_scales( + [](Tensor* t) { return t->rowwise_scale_inv_shape(); }, + [](Tensor* t) { 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( + [](Tensor* t) { return t->columnwise_scale_inv_shape(); }, + [](Tensor* t) { 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)); } return grouped; diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 5bb6400629..927407f478 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -535,6 +535,7 @@ struct GroupedBuffers { GroupedTensorHandle handle; CudaPtr<> data; CudaPtr<> scale_inv; + CudaPtr<> columnwise_scale_inv; CudaPtr first_dims_dev; CudaPtr last_dims_dev; CudaPtr offsets_dev; diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index b1543d55a2..41a8fd1112 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -378,7 +378,8 @@ struct GroupedTensor { last_dims(nullptr, std::vector{0}, DType::kInt64), tensor_offsets(nullptr, std::vector{0}, DType::kInt64), logical_shape(nvte_make_shape(nullptr, 1)), - nvte_tensor(0) {} + nvte_tensor(0), + with_gemm_swizzled_scales(false) {} explicit operator NVTEGroupedTensor() const noexcept { return nvte_tensor; } diff --git a/transformer_engine/common/gemm/config.h b/transformer_engine/common/gemm/config.h index ad38e88334..eed47e23d9 100644 --- a/transformer_engine/common/gemm/config.h +++ b/transformer_engine/common/gemm/config.h @@ -44,10 +44,13 @@ struct GroupedMatmulConfig { // Number of streaming multiprocessors to use in GEMM kernel int sm_count = 0; + // Split accumulator mode. Only taken into account on Hopper. + bool use_split_accumulator = false; + // Note: API transfers the value type, not std::optional - static constexpr size_t attr_sizes[] = {sizeof(decltype(avg_m)::value_type), - sizeof(decltype(avg_n)::value_type), - sizeof(decltype(avg_k)::value_type), sizeof(sm_count)}; + static constexpr size_t attr_sizes[] = { + sizeof(decltype(avg_m)::value_type), sizeof(decltype(avg_n)::value_type), + sizeof(decltype(avg_k)::value_type), sizeof(sm_count), sizeof(uint8_t)}; }; } // namespace transformer_engine diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index dc4757ab90..7069debc56 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -11,6 +11,7 @@ #include #include +#include #include "../common.h" #include "../util/cuda_runtime.h" @@ -26,6 +27,9 @@ inline void CreateCublasHandle(cublasLtHandle_t *handle) { } // namespace +// MXFP8 support for grouped GEMM requires cuBLAS 13.2+ +#define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130200 + #if CUBLAS_VERSION >= 130200 namespace { @@ -109,6 +113,10 @@ struct GroupedGemmSetupWorkspace { 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) // Storage dimensions for cuBLAS matrix layouts int *a_rows; int *a_cols; @@ -118,28 +126,47 @@ struct GroupedGemmSetupWorkspace { int *d_cols; // N (last dim) - also used for C // Initialize from workspace buffer - // Layout: all pointer arrays first (8-byte aligned), then int arrays (4-byte aligned) + // 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) { GroupedGemmSetupWorkspace ws; size_t offset = 0; 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 + + // Helper to align offset to kPtrAlignment + auto align_offset = [&]() { + offset = (offset + kPtrAlignment - 1) / kPtrAlignment * kPtrAlignment; + }; - // Pointer arrays first (all 8-byte aligned) + // 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) + // 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); @@ -159,8 +186,12 @@ struct GroupedGemmSetupWorkspace { 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); - // Layout: 6 ptr arrays, then 6 int arrays - size_t size = 6 * ptr_size + 6 * int_size; + 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; } @@ -169,12 +200,17 @@ struct GroupedGemmSetupWorkspace { // ----------------------------------------------------------------------------- // Helper routines to keep nvte_grouped_gemm readable // ----------------------------------------------------------------------------- -inline void validate_grouped_gemm_inputs(const transformer_engine::GroupedTensor *inputA, - const transformer_engine::GroupedTensor *inputB, - const transformer_engine::GroupedTensor *inputC, - const transformer_engine::GroupedTensor *outputD, - const transformer_engine::Tensor *alpha_tensor, - const transformer_engine::Tensor *beta_tensor) { +struct GroupedGemmInputProperties { + bool is_fp8; + bool is_mxfp8; +}; + +inline GroupedGemmInputProperties validate_grouped_gemm_inputs( + const transformer_engine::GroupedTensor *inputA, + const transformer_engine::GroupedTensor *inputB, + const transformer_engine::GroupedTensor *inputC, + const transformer_engine::GroupedTensor *outputD, + const transformer_engine::Tensor *alpha_tensor, const transformer_engine::Tensor *beta_tensor) { const size_t num_tensors = inputA->num_tensors; NVTE_CHECK(num_tensors >= 1, "Grouped GEMM: number of tensors must be at least 1"); NVTE_CHECK(inputB->num_tensors == num_tensors, @@ -195,28 +231,41 @@ inline void validate_grouped_gemm_inputs(const transformer_engine::GroupedTensor NVTE_CHECK(beta_numel == num_tensors, "Grouped GEMM: beta must have num_tensors (", num_tensors, ") elements, got ", beta_numel); - auto is_fp8_or_16bit = [](transformer_engine::DType dtype) { + 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::kBFloat16; }; 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_fp8_or_16bit(inputA->dtype()) && is_fp8_or_16bit(inputB->dtype()), - "Grouped GEMM inputs must be FP8, BF16, or FP16."); + NVTE_CHECK(is_supported_input_dtype(inputA->dtype()) && is_supported_input_dtype(inputB->dtype()), + "Grouped GEMM inputs must be FP8 or BF16."); + NVTE_CHECK(is_fp8_dtype(inputA->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(inputA->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, " + "mixed configurations are not supported."); + const bool is_fp8 = is_fp8_dtype(inputA->dtype()); + const bool is_mxfp8 = transformer_engine::is_mxfp_scaling(inputA->scaling_mode); + if (is_mxfp8) { + NVTE_CHECK(inputA->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"); + } // Only check C dtype if C is provided if (inputC != nullptr) { - NVTE_CHECK(is_output_dtype(inputC->dtype()), "Grouped GEMM: C must be BF16, FP16, or FP32."); + NVTE_CHECK(is_output_dtype(inputC->dtype()), "Grouped GEMM: C must be BF16 or FP32."); } - NVTE_CHECK(is_output_dtype(outputD->dtype()), "Grouped GEMM: D must be BF16, FP16, or FP32."); + NVTE_CHECK(is_output_dtype(outputD->dtype()), "Grouped GEMM: D must be BF16 or FP32."); NVTE_CHECK(inputA->has_data() || inputA->has_columnwise_data(), "Grouped GEMM: A tensor is missing both row-wise and column-wise data"); NVTE_CHECK(inputB->has_data() || inputB->has_columnwise_data(), "Grouped GEMM: B tensor is missing both row-wise and column-wise data"); + return {is_fp8, is_mxfp8}; } // Select row-wise vs column-wise storage and adjust transpose flag for grouped GEMM. @@ -226,8 +275,10 @@ inline void validate_grouped_gemm_inputs(const transformer_engine::GroupedTensor struct GroupedOperandSelection { TensorShapeInfo shape; // Shape info with dims already swapped for columnwise if needed char *dptr = nullptr; - void *scale_inv = nullptr; + void *scale_inv = nullptr; // Contiguous array of scales (input) transformer_engine::DType dtype = transformer_engine::DType::kNumTypes; + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + bool with_gemm_swizzled_scales = false; bool trans = false; }; @@ -266,26 +317,33 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: NVTE_CHECK(has_row || has_col, "Grouped GEMM operand is missing both row-wise and column-wise data"); - // Currently only unquantized data and tensor-scaled FP8 are supported. const auto sm = t->scaling_mode; - NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING, - "Grouped GEMM is only supported with unquantized data and tensor-scaled FP8 data"); + const bool mxfp8 = is_mxfp_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"); const DType row_dtype = t->data.dtype; const DType col_dtype = t->columnwise_data.dtype; GroupedOperandSelection sel; sel.trans = trans; + sel.scaling_mode = sm; + sel.with_gemm_swizzled_scales = t->with_gemm_swizzled_scales; 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(); - // Helper to select columnwise storage (swaps dims in shape) - auto use_columnwise = [&]() { + // 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) { sel.dptr = static_cast(t->columnwise_data.dptr); sel.scale_inv = t->columnwise_scale_inv.dptr; sel.dtype = col_dtype; - sel.shape = create_shape_info(t, /*swap_dims=*/true); + sel.shape = create_shape_info(t, swap_dims); }; // Helper to select row-wise storage @@ -296,6 +354,28 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: sel.shape = create_shape_info(t, /*swap_dims=*/false); }; + // MXFP8: Row-wise and column-wise data are scaled along different dimensions. + if (mxfp8) { + if (is_A) { + if (trans) { + NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 transposed A is missing row-wise data"); + use_rowwise(); + } else { + NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 non-transposed A is missing column-wise data"); + use_columnwise(/*swap_dims=*/false); + } + } else { // B + if (trans) { + NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 transposed B is missing column-wise data"); + use_columnwise(/*swap_dims=*/false); + } else { + NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 non-transposed B is missing row-wise data"); + use_rowwise(); + } + } + return sel; + } + // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed. if (is_fp8 && !non_tn_fp8_ok) { if (is_A) { @@ -364,7 +444,7 @@ inline void init_matrix_layouts(cublasLtMatrixLayoutOpaque_t &descA, } inline void init_matmul_desc(cublasLtMatmulDescOpaque_t &matmulDesc, cublasOperation_t op_A, - cublasOperation_t op_B) { + cublasOperation_t op_B, bool use_fp8, bool use_split_accumulator) { NVTE_CHECK_CUBLAS(cublasLtMatmulDescInit(&matmulDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &op_A, @@ -383,27 +463,53 @@ inline void init_matmul_desc(cublasLtMatmulDescOpaque_t &matmulDesc, cublasOpera 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); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_FAST_ACCUM, + &fastAccuMode, sizeof(fastAccuMode))); } -inline void set_fp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, - const GroupedOperandSelection &A_sel, - const GroupedOperandSelection &B_sel) { - const bool is_fp8_a = is_fp8_dtype(A_sel.dtype); - const bool is_fp8_b = is_fp8_dtype(B_sel.dtype); - if (!is_fp8_a && !is_fp8_b) return; - - if (is_fp8_a) { - void *a_scale_inv = A_sel.scale_inv; - NVTE_CHECK(a_scale_inv != nullptr, "FP8 grouped GEMM: A scale_inv is required"); - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - &matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_scale_inv, sizeof(a_scale_inv))); - } - if (is_fp8_b) { - void *b_scale_inv = B_sel.scale_inv; - NVTE_CHECK(b_scale_inv != nullptr, "FP8 grouped GEMM: B scale_inv is required"); - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - &matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale_inv, sizeof(b_scale_inv))); - } +// Configures cuBLAS for MXFP8 grouped GEMM: sets VEC32_UE8M0 scale mode and scale pointers +// for both A and B. +inline void set_mxfp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, + void **a_scale_inv_ptrs, void **b_scale_inv_ptrs) { +#if CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION, + "MXFP8 grouped GEMM requires cuBLAS ", CUBLAS_MXFP8_GROUPED_GEMM_VERSION, + "+, but run-time cuBLAS version is ", transformer_engine::cuda::cublas_version()); + const cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0; + 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, "MXFP8 grouped GEMM requires cuBLAS ", CUBLAS_MXFP8_GROUPED_GEMM_VERSION, + "+, but compile-time cuBLAS version is ", CUBLAS_VERSION); +#endif // CUBLAS_VERSION >= CUBLAS_MXFP8_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, + void **b_scale_inv_ptrs) { + const cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_PER_BATCH_SCALAR_32F; + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_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_MODE, + &scale_mode, sizeof(scale_mode))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs))); } // Constants for grouped GEMM workspace (declared early for use in heuristics) @@ -471,11 +577,13 @@ __global__ void setup_grouped_gemm_kernel( // Output arrays void **A_ptrs, void **B_ptrs, void **C_ptrs, void **D_ptrs, int *a_rows, int *a_cols, int *b_rows, int *b_cols, int *d_rows, int *d_cols, float **alpha_ptrs, float **beta_ptrs, + 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, - size_t num_tensors) { + // Scale inputs: contiguous scale buffers and the shared scaling recipe for A and B + void *a_scale_base, void *b_scale_base, NVTEScalingMode scaling_mode, size_t num_tensors) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_tensors) return; @@ -512,6 +620,25 @@ __global__ void setup_grouped_gemm_kernel( // Fill alpha/beta pointers (per-matrix) alpha_ptrs[idx] = alpha_ptr + idx; beta_ptrs[idx] = beta_ptr + idx; + + // 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) { + a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + a_offset / 32; + } else { + a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + idx; + } + } + if (b_scale_base) { + if (scaling_mode == NVTE_MXFP8_1D_SCALING) { + b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + b_offset / 32; + } else { + b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + idx; + } + } } // Launch the setup kernel to populate workspace arrays @@ -537,12 +664,15 @@ inline void launch_grouped_gemm_setup( 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. 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, 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), - num_tensors); + 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), A_sel.scale_inv, B_sel.scale_inv, + A_sel.scaling_mode, num_tensors); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -590,8 +720,9 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT config_ = *reinterpret_cast(config); } - // Validate inputs and num_tensors - validate_grouped_gemm_inputs(inputA, inputB, inputC_raw, outputD, alpha_tensor, beta_tensor); + // Validate inputs and num_tensors; returns dtype properties shared by A and B. + const auto [is_fp8, is_mxfp8] = + validate_grouped_gemm_inputs(inputA, inputB, inputC_raw, outputD, alpha_tensor, beta_tensor); // 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; @@ -599,8 +730,8 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT // Select operand storage (row-wise vs column-wise) and adjust transpose flags to // mirror the non-grouped GEMM logic for FP8 layout constraints. - const auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true); - const auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false); + 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); // Workspaces: setup (pointer arrays) and cuBLAS const size_t setup_workspace_size = grouped_gemm_setup_workspace_size(num_tensors); @@ -613,6 +744,7 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers( static_cast(setup_workspace_ptr), num_tensors); + launch_grouped_gemm_setup(setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, beta_tensor, num_tensors, stream); @@ -631,8 +763,14 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT // Create matmul descriptor cublasLtMatmulDescOpaque_t matmulDesc; - init_matmul_desc(matmulDesc, op_A, op_B); - set_fp8_scale_pointers(matmulDesc, A_sel, B_sel); + init_matmul_desc(matmulDesc, op_A, op_B, is_fp8, config_.use_split_accumulator); + if (is_mxfp8) { + set_mxfp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, + setup_workspace.b_scale_inv_ptrs); + } else if (is_fp8) { + set_fp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, + setup_workspace.b_scale_inv_ptrs); + } // Compute average dimensions for heuristics // K dimension: if transa, K is A's first dim; if not, K is A's last dim diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 0f3b0ebd6b..35d327f085 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -84,6 +84,8 @@ enum NVTEGroupedMatmulConfigAttribute { kNVTEGroupedMatmulConfigAvgK = 2, /*! Number of streaming multiprocessors to use in GEMM kernel. */ kNVTEGroupedMatmulConfigSMCount = 3, + /*! Split accumulator mode. Only taken into account on Hopper. Default: true. */ + kNVTEGroupedMatmulConfigUseSplitAccumulator = 4, kNVTEGroupedMatmulConfigNumAttributes }; @@ -522,6 +524,13 @@ class GroupedMatmulConfigWrapper { sizeof(int)); } + /*! \brief Set split accumulator mode. Only taken into account on Hopper. */ + void set_use_split_accumulator(bool use_split_accumulator) { + const auto val = static_cast(use_split_accumulator); + nvte_set_grouped_matmul_config_attribute(config_, kNVTEGroupedMatmulConfigUseSplitAccumulator, + &val, sizeof(val)); + } + private: /*! \brief Wrapped NVTEGroupedMatmulConfig. */ NVTEGroupedMatmulConfig config_ = nullptr; From 67898a7c6c7897ee9c41b9dbe220f26ff72b1ac6 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Thu, 12 Mar 2026 13:16:10 -0700 Subject: [PATCH 273/521] Cherry pick "Adds dst.dtype information in copy_ method of quantized tensors. #2120" (#2673) * Adds dst.dtype information in copy_ method of quantized tensors. Signed-off-by: Zhiyi Su * Update transformer_engine/pytorch/tensor/quantized_tensor.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: ZhiyiDanielSu <35579247+zobeideThePlayer@users.noreply.github.com> * Update transformer_engine/pytorch/quantized_tensor.py Signed-off-by: Kirthi Shankar Sivamani * Fix reference tensor copy Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Zhiyi Su Signed-off-by: ZhiyiDanielSu <35579247+zobeideThePlayer@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Zhiyi Su Co-authored-by: ZhiyiDanielSu <35579247+zobeideThePlayer@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fusible_ops.py | 4 ++-- tests/pytorch/test_quantized_tensor.py | 2 +- transformer_engine/pytorch/quantized_tensor.py | 5 ++++- 3 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index f95f065d78..b97afbc191 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -178,7 +178,7 @@ def make_reference_and_test_tensors( test = test.dequantize() # Make sure reference and test tensors match each other - ref.copy_(test) + ref.copy_(test.to(dtype=ref.dtype)) ref.requires_grad_(requires_grad) test.requires_grad_(requires_grad) @@ -1956,7 +1956,7 @@ def test_dropout( ) with torch.no_grad(): x_test += 1 - x_ref.copy_(x_test) + x_ref.copy_(x_test.to(dtype=x_ref.dtype)) dy_ref, dy_test = make_reference_and_test_tensors( shape, test_dtype=dtype, diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 978ec09b40..620fc834dd 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -173,7 +173,7 @@ def make_reference_and_test_tensors( raise ValueError(f"Unsupported quantization scheme ({quantization})") # Make sure reference and test tensors match each other - ref.copy_(test) + ref.copy_(test.to(dtype=ref.dtype)) ref.requires_grad_(requires_grad) test.requires_grad_(requires_grad) diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 07171914f5..2129fd486f 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -552,7 +552,10 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dst.quantize_(src) else: if isinstance(src, QuantizedTensor): - src = src.dequantize() + dtype = dst.dtype + if dtype not in (torch.float32, torch.float16, torch.bfloat16): + dtype = torch.float32 + src = src.dequantize(dtype=dtype) dst.copy_(src) return None From 134304efaefe2fa4ba32ba2592dc7f57b46dc57a Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 12 Mar 2026 13:16:27 -0700 Subject: [PATCH 274/521] Fused kernel for calculating offsets from first dim splits (#2755) * Fuse scale + 0 + cumulative sum for splits to offsets calc Signed-off-by: Kirthi Shankar Sivamani * Add unit test and fix bug in kernel for >256 size Signed-off-by: Kirthi Shankar Sivamani * fix race Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * check for logical_last_dim > 0 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * suggestions Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_splits_to_offsets.cu | 80 +++++++++++++++++++ transformer_engine/common/common.cu | 55 +++++++++++++ .../transformer_engine/transformer_engine.h | 16 ++++ transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/misc.cpp | 18 +++++ .../pytorch/csrc/extensions/pybind.cpp | 3 + transformer_engine/pytorch/csrc/quantizer.cpp | 14 ++-- 8 files changed, 184 insertions(+), 5 deletions(-) create mode 100644 tests/cpp/operator/test_splits_to_offsets.cu diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 56880a428d..5e73675f4f 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -25,6 +25,7 @@ add_executable(test_operator test_normalization.cu test_normalization_mxfp8.cu test_memset.cu + test_splits_to_offsets.cu test_multi_cast_transpose.cu test_multi_padding.cu test_multi_unpadding.cu diff --git a/tests/cpp/operator/test_splits_to_offsets.cu b/tests/cpp/operator/test_splits_to_offsets.cu new file mode 100644 index 0000000000..faac4b7b6f --- /dev/null +++ b/tests/cpp/operator/test_splits_to_offsets.cu @@ -0,0 +1,80 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include +#include + +#include +#include "../test_common.h" + +class SplitsToOffsetsTestSuite : public ::testing::TestWithParam> {}; + +TEST_P(SplitsToOffsetsTestSuite, TestSplitsToOffsets) { + const size_t num_tensors = std::get<0>(GetParam()); + const int64_t logical_last_dim = std::get<1>(GetParam()); + + std::vector h_first_dims(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + h_first_dims[i] = static_cast((i % 17) + 1); + } + + std::vector h_expected(num_tensors + 1, 0); + for (size_t i = 0; i < num_tensors; ++i) { + h_expected[i + 1] = h_expected[i] + h_first_dims[i] * logical_last_dim; + } + + std::vector h_output(num_tensors + 1, -1); + + int64_t *d_first_dims = nullptr; + int64_t *d_output = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&d_first_dims, sizeof(int64_t) * num_tensors)); + NVTE_CHECK_CUDA(cudaMalloc(&d_output, sizeof(int64_t) * (num_tensors + 1))); + NVTE_CHECK_CUDA(cudaMemcpy(d_first_dims, h_first_dims.data(), sizeof(int64_t) * num_tensors, + cudaMemcpyHostToDevice)); + + nvte_splits_to_offsets(d_first_dims, d_output, num_tensors, logical_last_dim, 0 /* stream */); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + + NVTE_CHECK_CUDA(cudaMemcpy(h_output.data(), d_output, sizeof(int64_t) * (num_tensors + 1), + cudaMemcpyDeviceToHost)); + + NVTE_CHECK_CUDA(cudaFree(d_first_dims)); + NVTE_CHECK_CUDA(cudaFree(d_output)); + + for (size_t i = 0; i < h_output.size(); ++i) { + EXPECT_EQ(h_output[i], h_expected[i]) + << "Mismatch at index " << i << ": expected " << h_expected[i] << ", got " << h_output[i]; + } +} + +namespace { + +std::vector splits_to_offsets_num_tensors = { + 1, + 4, + 255, + 256, + 257, + 1024, +}; + +} // namespace + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, SplitsToOffsetsTestSuite, + ::testing::Combine(::testing::ValuesIn(splits_to_offsets_num_tensors), + ::testing::Values(static_cast(1), static_cast(7), + static_cast(128))), + [](const testing::TestParamInfo &info) { + std::string name = std::to_string(std::get<0>(info.param)) + "X" + + std::to_string(std::get<1>(info.param)); + return name; + }); diff --git a/transformer_engine/common/common.cu b/transformer_engine/common/common.cu index 0ec40dc01c..1bdd80a369 100644 --- a/transformer_engine/common/common.cu +++ b/transformer_engine/common/common.cu @@ -87,6 +87,48 @@ __global__ void __launch_bounds__(kThreadsPerBlock) reinterpret_cast(ptr)[idx] = data.value; } +__global__ void __launch_bounds__(kThreadsPerBlock) + splits_to_offsets_kernel(const int64_t *__restrict__ first_dims, int64_t *__restrict__ output, + size_t num_tensors, int64_t logical_last_dim) { + __shared__ int64_t block_scan[kThreadsPerBlock]; + __shared__ int64_t chunk_prefix; + + const size_t tid = threadIdx.x; + if (tid == 0) { + output[0] = 0; + chunk_prefix = 0; + } + __syncthreads(); + + for (size_t chunk_start = 0; chunk_start < num_tensors; chunk_start += kThreadsPerBlock) { + const size_t idx = chunk_start + tid; + int64_t value = 0; + if (idx < num_tensors) { + value = first_dims[idx] * logical_last_dim; + } + block_scan[tid] = value; + __syncthreads(); + + // Inclusive scan in shared memory. + for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) { + const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0; + __syncthreads(); + block_scan[tid] += addend; + __syncthreads(); + } + + if (idx < num_tensors) { + output[idx + 1] = chunk_prefix + block_scan[tid]; + } + __syncthreads(); + + if (tid == kThreadsPerBlock - 1) { + chunk_prefix += block_scan[tid]; + } + __syncthreads(); + } +} + } // namespace #define MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, vectorizedType, stream) \ @@ -116,6 +158,19 @@ void nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, float, stream); MEMSET_VECTORIZED_KERNEL_DISPATCH(ptr, size_in_bytes, value, uint8_t, stream); } + +void nvte_splits_to_offsets(const int64_t *first_dims, int64_t *output, size_t num_tensors, + int64_t logical_last_dim, cudaStream_t stream) { + NVTE_API_CALL(nvte_splits_to_offsets); + NVTE_CHECK(output != nullptr, "Output pointer must be allocated."); + NVTE_CHECK(num_tensors > 0, "num_tensors must be greater than 0."); + NVTE_CHECK(first_dims != nullptr, "first_dims pointer must be allocated."); + NVTE_CHECK(logical_last_dim > 0, "logical_last_dim must be greater than 0."); + + splits_to_offsets_kernel<<<1, kThreadsPerBlock, 0, stream>>>(first_dims, output, num_tensors, + logical_last_dim); + NVTE_CHECK_CUDA(cudaGetLastError()); +} } // extern "C" void checkCuDriverContext(CUstream stream) { diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index e316f8be8c..b7461a85d1 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -427,6 +427,22 @@ int nvte_is_non_tn_fp8_gemm_supported(); */ void nvte_memset(void *ptr, int value, size_t size_in_bytes, cudaStream_t stream); +/*! \brief Compute scaled prefix-sum offsets for grouped tensors. + * + * Computes: + * output[0] = 0 + * output[i + 1] = sum_{j=0..i}(first_dims[j] * logical_last_dim) + * for i in [0, num_tensors - 1]. + * + * \param[in] first_dims Pointer to device int64 array of size num_tensors. + * \param[out] output Pointer to device int64 array of size num_tensors + 1. + * \param[in] num_tensors Number of entries in first_dims. + * \param[in] logical_last_dim Scale factor applied to each first_dims entry. + * \param[in] stream CUDA stream to use for the operation. + */ +void nvte_splits_to_offsets(const int64_t *first_dims, int64_t *output, size_t num_tensors, + int64_t logical_last_dim, cudaStream_t stream); + /*! \brief TE Grouped Tensor type * * NVTEGroupedTensor is a collection of tensors with potentially different shapes diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index d8f00becb5..66c01aaf7a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -435,6 +435,8 @@ size_t get_cublasLt_version(); size_t get_cudnn_version(); +at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_dim); + /*************************************************************************************************** * Support THD format for Context Parallel **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/misc.cpp b/transformer_engine/pytorch/csrc/extensions/misc.cpp index d667a61d44..c5707fa53c 100644 --- a/transformer_engine/pytorch/csrc/extensions/misc.cpp +++ b/transformer_engine/pytorch/csrc/extensions/misc.cpp @@ -12,4 +12,22 @@ size_t get_cublasLt_version() { return cublasLtGetVersion(); } size_t get_cudnn_version() { return cudnnGetVersion(); } +at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_dim) { + NVTE_CHECK(first_dims.is_cuda(), "first_dims must be on CUDA."); + NVTE_CHECK(first_dims.scalar_type() == at::kLong, "first_dims must have dtype int64."); + NVTE_CHECK(first_dims.dim() == 1, "first_dims must be a 1D tensor."); + NVTE_CHECK(logical_last_dim > 0, "logical_last_dim must be greater than 0."); + + auto first_dims_contiguous = first_dims.contiguous(); + const auto num_tensors = static_cast(first_dims_contiguous.numel()); + auto output = at::empty({static_cast(num_tensors) + 1}, + first_dims_contiguous.options().dtype(at::kLong)); + + nvte_splits_to_offsets(static_cast(first_dims_contiguous.data_ptr()), + static_cast(output.data_ptr()), num_tensors, logical_last_dim, + at::cuda::getCurrentCUDAStream()); + + return output; +} + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 817a481808..7721671a36 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -445,6 +445,9 @@ 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("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()); m.def("get_num_cublas_streams", &nvte_get_num_compute_streams, "Get number of compute streams", py::call_guard()); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 27dc87697f..c904057e97 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -73,11 +73,15 @@ std::optional build_grouped_tensor_offsets(const size_t num_tensors, "first_dims must have length ", num_tensors, "."); const int64_t logical_last_dim_i64 = static_cast(logical_last_dim); - auto scaled_first_dims = (first_dims_tensor * logical_last_dim_i64).contiguous(); - // Single kernel needed for these ops. - auto cumsum = at::cumsum(scaled_first_dims, 0); - auto zero = at::zeros({1}, cumsum.options()); - return at::cat({zero, cumsum}); + const auto first_dims_contiguous = first_dims_tensor.contiguous(); + auto tensor_offsets = + at::empty({static_cast(num_tensors) + 1}, first_dims_contiguous.options()); + NVTE_SCOPED_GIL_RELEASE({ + nvte_splits_to_offsets(static_cast(first_dims_contiguous.data_ptr()), + static_cast(tensor_offsets.data_ptr()), num_tensors, + logical_last_dim_i64, at::cuda::getCurrentCUDAStream()); + }); + return tensor_offsets; } at::TensorOptions grouped_tensor_data_options(const DType dtype) { From a5d746425948c03497d680c13e52c652d3529438 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Thu, 12 Mar 2026 14:44:01 -0700 Subject: [PATCH 275/521] Added new users to CI (#2756) * Added new people to CI Signed-off-by: Przemek Tredak * Removing duplicate Signed-off-by: Przemek Tredak --------- Signed-off-by: Przemek Tredak --- .github/workflows/trigger-ci.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/trigger-ci.yml b/.github/workflows/trigger-ci.yml index c56601ae98..de26531a98 100644 --- a/.github/workflows/trigger-ci.yml +++ b/.github/workflows/trigger-ci.yml @@ -58,6 +58,8 @@ jobs: || github.actor == 'vthumbe1503' || github.actor == 'shengfangd' || github.actor == 'kainzhong' + || github.actor == 'cspades' + || github.actor == 'jomitchellnv' ) steps: - name: Check if comment is issued by authorized person From 6a68c7336f5ac50b27eeab852965620df5fcee09 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 12 Mar 2026 15:15:06 -0700 Subject: [PATCH 276/521] [PyTorch] Error out if constructing `LayerNormLinear` with row tensor parallelism (#2688) * Error out if constructing LayerNormLinear with row tensor parallelism Signed-off-by: Tim Moon * Disable Userbuffers test for row-TP LayerNormLinear Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/distributed/test_comm_gemm_overlap.py | 4 ---- transformer_engine/pytorch/module/layernorm_linear.py | 4 ++++ 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py index 95bf5aa05a..7a81f93bd6 100644 --- a/tests/pytorch/distributed/test_comm_gemm_overlap.py +++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py @@ -210,7 +210,6 @@ def test_bulk_overlaps(comm_type, quantization, connections): (te.Linear.__name__, "row", False), (te.Linear.__name__, "column", False), (te.Linear.__name__, "column", True), - (te.LayerNormLinear.__name__, "row", False), (te.LayerNormLinear.__name__, "column", False), (te.LayerNormLinear.__name__, "column", True), ] @@ -225,7 +224,6 @@ def test_bulk_overlaps(comm_type, quantization, connections): f" {te.Linear.__name__} - ROW-PARALLEL ", f" {te.Linear.__name__} - COL-PARALLEL - BULK DGRAD/WGRAD ", f" {te.Linear.__name__} - COL-PARLALEL - DGRAD+RS ", - f" {te.LayerNormLinear.__name__} - ROW-PARALLEL ", f" {te.LayerNormLinear.__name__} - COL-PARALLEL - BULK DGRAD/WGRAD ", f" {te.LayerNormLinear.__name__} - COL-PARALLEL - DGRAD+RS ", ] @@ -254,7 +252,6 @@ def test_layers_with_overlap_bf16(layer_type, linear_parallel_mode, overlap_rs_d (te.Linear.__name__, "row", False), (te.Linear.__name__, "column", False), (te.Linear.__name__, "column", True), - (te.LayerNormLinear.__name__, "row", False), (te.LayerNormLinear.__name__, "column", False), (te.LayerNormLinear.__name__, "column", True), ] @@ -269,7 +266,6 @@ def test_layers_with_overlap_bf16(layer_type, linear_parallel_mode, overlap_rs_d f"{te.Linear.__name__}-row_tensor_parallel", f"{te.Linear.__name__}-col_tensor_parallel-BULK DGRAD/WGRAD", f"{te.Linear.__name__}-col_tensor_parallel-DGRAD+RS", - f"{te.LayerNormLinear.__name__}-row_tensor_parallel", f"{te.LayerNormLinear.__name__}-col_tensor_parallel-BULK DGRAD/WGRAD", f"{te.LayerNormLinear.__name__}-col_tensor_parallel-DGRAD+RS", ] diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index a90105477c..d775dc3e8e 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1192,6 +1192,10 @@ def __init__( assert ( self.parallel_mode in GemmParallelModes ), f"parallel_mode {parallel_mode} not supported" + if self.parallel_mode == "row": + raise NotImplementedError( + "Normalization does not support tensor-parallel distribution." + ) if self.parallel_mode == "column": self.out_features = divide(self.out_features, self.tp_size) From 14c29da8ade97370610ae71c6fe494f351316d5f Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 13 Mar 2026 08:27:40 -0700 Subject: [PATCH 277/521] [JAX] Collective GEMM with FP8 and MXFP8 support (#2740) * Enable cgemm + FP8 tests * Implement CGEMM + MXFP8 * added size check for mxfp8 * added tols for assertions * update tests with recipes * enable tests + is_quantize_recipe_supported Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- examples/jax/collective_gemm/common.py | 81 +++----- .../jax/collective_gemm/run_test_cgemm.sh | 72 +++++-- .../jax/collective_gemm/test_dense_grad.py | 116 ++++++++++- examples/jax/collective_gemm/test_gemm.py | 127 +++++++++--- .../test_layernorm_mlp_grad.py | 78 ++++++- transformer_engine/jax/cpp_extensions/gemm.py | 193 ++++++++++++------ .../jax/csrc/extensions/gemm.cpp | 2 + transformer_engine/jax/quantize/helper.py | 51 +++++ 8 files changed, 540 insertions(+), 180 deletions(-) diff --git a/examples/jax/collective_gemm/common.py b/examples/jax/collective_gemm/common.py index 2965896d07..6815932395 100644 --- a/examples/jax/collective_gemm/common.py +++ b/examples/jax/collective_gemm/common.py @@ -1,47 +1,52 @@ # Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # See LICENSE for license information. -"""Shared functions for the comm_overlap tests""" +"""Shared functions for the collective GEMM tests""" +import argparse + +import jax import jax.numpy as jnp import numpy as np +from jax.experimental import mesh_utils + +from transformer_engine.jax.cpp_extensions.gemm import collective_gemm_bootstrap -# Add this after your existing imports def dtype_tols(dtype, rtol=None, atol=None): """Expected numerical tolerance for a data type.""" - # Return immediately if tolerances are fully specified if rtol is not None and atol is not None: return {"rtol": rtol, "atol": atol} - # Default tolerances for common dtypes if dtype in [jnp.float32, "float32"]: return {"rtol": 1e-5, "atol": 1e-8} elif dtype in [jnp.float16, "float16"]: return {"rtol": 1e-3, "atol": 1e-6} elif dtype in [jnp.bfloat16, "bfloat16"]: return {"rtol": 1e-2, "atol": 1e-5} + elif dtype in [jnp.float8_e4m3fn, "float8_e4m3fn", jnp.float8_e5m2, "float8_e5m2"]: + # FP8 quantization introduces ~1% error; match C++ getTolerances for fp8 types + return {"rtol": 1e-2, "atol": 1e-2} else: return {"rtol": 1e-5, "atol": 1e-8} -def assert_allclose( - actual, - desired, - rtol=None, - atol=None, - dtype=None, - **kwargs, -): +def get_tolerance_dtype(quantizer_set): + """Return the dtype used to select numerical tolerances based on the active quantizer. + + Reads q_dtype from quantizer_set.x; falls back to bfloat16 when no quantizer is + active (NO_SCALING / noop path, where quantizer_set.x is None). + """ + if quantizer_set.x is not None: + return quantizer_set.x.q_dtype + return jnp.bfloat16 + + +def assert_allclose(actual, desired, rtol=None, atol=None, dtype=None, **kwargs): """Check if two tensors are close.""" - # Infer data type if needed if dtype is None: - if isinstance(actual, float): - dtype = "float32" - else: - dtype = actual.dtype + dtype = "float32" if isinstance(actual, float) else actual.dtype - # Determine tolerances tols = {} if rtol is None or atol is None: tols = dtype_tols(dtype) @@ -50,49 +55,26 @@ def assert_allclose( if atol is not None: tols["atol"] = atol - # Cast tensors to fp32 if not isinstance(actual, float): actual = actual.astype(jnp.float32) if not isinstance(desired, float): desired = desired.astype(jnp.float32) - # Check if tensors are close np.testing.assert_allclose(actual, desired, **tols, **kwargs) -def assert_allclose_print_index(ref_output, gathered_output, rtol=1e-5, atol=1e-8): - if not jnp.allclose(ref_output, gathered_output, rtol=rtol, atol=atol): - diff = jnp.abs(ref_output - gathered_output) - mask = diff > (atol + rtol * jnp.abs(gathered_output)) - print(mask.astype(int)) - print(jnp.where(mask, diff, 0)) - - -# Shared constants for all tests +# Shared constants DP_AXIS = "data" TPSP_AXIS = "tensor_sequence" -PARAMS_KEY = "params" - -# Shared functions for distributed testing -import argparse -import jax -from jax.experimental import mesh_utils -from transformer_engine.jax.cpp_extensions.gemm import collective_gemm_bootstrap # Global flag to track if distributed has been initialized _distributed_initialized = False -def _is_distributed_initialized(): - """Check if JAX distributed has been initialized.""" - return _distributed_initialized - - def _initialize_distributed(args): """Initialize JAX distributed with custom arguments.""" global _distributed_initialized - # Check if already initialized if _distributed_initialized: return @@ -105,14 +87,10 @@ def _initialize_distributed(args): assert ( args.num_devices_per_process is not None ), "Either local_device_ids or num_devices_per_process must be provided" - # Calculate device range for this process - # Single process single device: each process gets one unique device - # Single process multiple devices: each process gets a unique range of devices start_device = args.process_id * args.num_devices_per_process device_range = range(start_device, start_device + args.num_devices_per_process) global_device_ids_for_this_process = ",".join(map(str, device_range)) else: - # Use explicitly provided global device IDs global_device_ids_for_this_process = args.local_device_ids args.num_devices_per_process = len(args.local_device_ids.split(",")) @@ -229,7 +207,16 @@ def cgemm_parser(description="Collective GEMM test on multi-GPU with tensor para help="Type of collective operation", ) parser.add_argument( - "--fp8-recipe", type=str, default="DelayedScaling", help="FP8 recipe to use" + "--quantize-recipe", + type=str, + default=None, + choices=[ + "DelayedScaling", + "Float8CurrentScaling", + "MXFP8BlockScaling", + "NVFP4BlockScaling", + ], + help="Quantization recipe to use. Omit for BF16 (no quantization).", ) parser.add_argument( "--enable-data-parallel", action="store_true", help="Enable data parallelism" diff --git a/examples/jax/collective_gemm/run_test_cgemm.sh b/examples/jax/collective_gemm/run_test_cgemm.sh index 388c878376..8340d2010f 100644 --- a/examples/jax/collective_gemm/run_test_cgemm.sh +++ b/examples/jax/collective_gemm/run_test_cgemm.sh @@ -23,11 +23,36 @@ else echo "NVLINK support detected" fi -# Define the test files to run -TEST_FILES=( -"test_gemm.py" -"test_dense_grad.py" -"test_layernorm_mlp_grad.py" +# Define individual test cases to run (file::class::method) +# DelayedScalingFP8 and CurrentScalingFP8 use the same GEMM so we don't need to test both cases all +# the time. +TEST_CASES=( +# test_gemm.py cases +"test_gemm.py::TestCollectiveGemmWithDP::test_te_bf16_all_gather_with_dp" +"test_gemm.py::TestCollectiveGemmWithDP::test_te_bf16_reduce_scatter_with_dp" +"test_gemm.py::TestCollectiveGemmWithDP::test_te_delayed_scaling_fp8_all_gather_with_dp" +"test_gemm.py::TestCollectiveGemmWithDP::test_te_delayed_scaling_fp8_reduce_scatter_with_dp" +"test_gemm.py::TestCollectiveGemmWithDP::test_te_mxfp8_all_gather_with_dp" +"test_gemm.py::TestCollectiveGemmWithDP::test_te_mxfp8_reduce_scatter_with_dp" +# # "test_gemm.py::TestCollectiveGemmWithDP::test_te_nvfp4_all_gather_with_dp" +# # "test_gemm.py::TestCollectiveGemmWithDP::test_te_nvfp4_reduce_scatter_with_dp" +# +# # test_dense_grad.py cases +"test_dense_grad.py::TestCollectiveDenseGradient::test_te_bf16_all_gather" +"test_dense_grad.py::TestCollectiveDenseGradient::test_te_bf16_reduce_scatter" +"test_dense_grad.py::TestCollectiveDenseGradient::test_te_current_scaling_fp8_all_gather" +"test_dense_grad.py::TestCollectiveDenseGradient::test_te_current_scaling_fp8_reduce_scatter" +"test_dense_grad.py::TestCollectiveDenseGradient::test_te_mxfp8_all_gather" +"test_dense_grad.py::TestCollectiveDenseGradient::test_te_mxfp8_reduce_scatter" +# "test_dense_grad.py::TestCollectiveDenseGradient::test_te_nvfp4_all_gather" +# "test_dense_grad.py::TestCollectiveDenseGradient::test_te_nvfp4_reduce_scatter" + +# test_layernorm_mlp_grad.py cases +"test_layernorm_mlp_grad.py::TestCollectiveLayerNormMLPGradient::test_te_bf16_layernorm_mlp_grad" +"test_layernorm_mlp_grad.py::TestCollectiveLayerNormMLPGradient::test_te_delayed_scaling_fp8_layernorm_mlp_grad" +"test_layernorm_mlp_grad.py::TestCollectiveLayerNormMLPGradient::test_te_current_scaling_fp8_layernorm_mlp_grad" +"test_layernorm_mlp_grad.py::TestCollectiveLayerNormMLPGradient::test_te_mxfp8_layernorm_mlp_grad" +# "test_layernorm_mlp_grad.py::TestCollectiveLayerNormMLPGradient::test_te_nvfp4_layernorm_mlp_grad" ) echo @@ -57,24 +82,27 @@ cleanup() { # Set up signal handlers to cleanup on exit trap cleanup EXIT INT TERM -# Run each test file across all GPUs -for TEST_FILE in "${TEST_FILES[@]}"; do +# Run each test case across all GPUs +for TEST_CASE in "${TEST_CASES[@]}"; do echo - echo "=== Starting test file: $TEST_FILE ..." + echo "=== Starting test: $TEST_CASE ..." + + # Extract just the test method name for log/xml file naming + TEST_NAME=$(echo "$TEST_CASE" | awk -F'::' '{print $NF}') - # Clear PIDs array for this test file + # Clear PIDs array for this test case PIDS=() for i in $(seq 0 $(($NUM_GPUS - 1))); do # Define output file for logs - LOG_FILE="${TEST_FILE}_gpu_${i}.log" + LOG_FILE="${TEST_NAME}_gpu_${i}.log" if [ $i -eq 0 ]; then # For process 0: show live output AND save to log file using tee echo "=== Live output from process 0 ===" pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \ - -vs --junitxml=$XML_LOG_DIR/collective_gemm_${TEST_FILE}.xml \ - "$TE_PATH/examples/jax/collective_gemm/$TEST_FILE" \ + -vs --junitxml=$XML_LOG_DIR/collective_gemm_${TEST_NAME}.xml \ + "$TE_PATH/examples/jax/collective_gemm/$TEST_CASE" \ --num-processes=$NUM_GPUS \ --process-id=$i 2>&1 | tee "$LOG_FILE" & PID=$! @@ -82,7 +110,7 @@ for TEST_FILE in "${TEST_FILES[@]}"; do else # For other processes: redirect to log files only pytest -s -c "$TE_PATH/tests/jax/pytest.ini" \ - -vs "$TE_PATH/examples/jax/collective_gemm/$TEST_FILE" \ + -vs "$TE_PATH/examples/jax/collective_gemm/$TEST_CASE" \ --num-processes=$NUM_GPUS \ --process-id=$i > "$LOG_FILE" 2>&1 & PID=$! @@ -93,22 +121,22 @@ for TEST_FILE in "${TEST_FILES[@]}"; do # Wait for all processes to finish wait - # Check and print the log content from process 0 (now has log file thanks to tee) - if grep -q "SKIPPED" "${TEST_FILE}_gpu_0.log"; then - echo "... $TEST_FILE SKIPPED" - elif grep -q "FAILED" "${TEST_FILE}_gpu_0.log"; then - echo "... $TEST_FILE FAILED" + # Check and print the log content from process 0 + if grep -q "SKIPPED" "${TEST_NAME}_gpu_0.log"; then + echo "... $TEST_CASE SKIPPED" + elif grep -q "FAILED" "${TEST_NAME}_gpu_0.log"; then + echo "... $TEST_CASE FAILED" HAS_FAILURE=1 - elif grep -q "PASSED" "${TEST_FILE}_gpu_0.log"; then - echo "... $TEST_FILE PASSED" + elif grep -q "PASSED" "${TEST_NAME}_gpu_0.log"; then + echo "... $TEST_CASE PASSED" else - echo "... $TEST_FILE INVALID" + echo "... $TEST_CASE INVALID" HAS_FAILURE=1 fi # Remove the log files after processing them wait - rm ${TEST_FILE}_gpu_*.log + rm ${TEST_NAME}_gpu_*.log done wait diff --git a/examples/jax/collective_gemm/test_dense_grad.py b/examples/jax/collective_gemm/test_dense_grad.py index 94c7dc5b66..1d300f8e90 100644 --- a/examples/jax/collective_gemm/test_dense_grad.py +++ b/examples/jax/collective_gemm/test_dense_grad.py @@ -2,7 +2,6 @@ # # See LICENSE for license information. """Collective Dense Gradient test on multi-GPU with tensor parallelism""" -import argparse import unittest import os @@ -13,18 +12,24 @@ from common import ( assert_allclose, + get_tolerance_dtype, _initialize_distributed, _get_dp_and_tp_sizes, _create_mesh, DP_AXIS, TPSP_AXIS, - PARAMS_KEY, cgemm_parser, ) from transformer_engine.jax.dense import dense -from transformer_engine.jax.quantize import autocast +from transformer_engine.jax.quantize import ( + autocast, + is_quantize_recipe_supported, + get_quantization_recipe, + QuantizerFactory, + noop_quantizer_set, +) from transformer_engine.jax.cpp_extensions.gemm import ( CollectiveOp, CollectiveOpSet, @@ -56,7 +61,9 @@ def _get_operand_sharding(mesh, collective_op): return x_sharding, weight_sharding, bias_sharding -def _mean_dense(x, weight, bias, input_axes, weight_axes, output_axes, collective_op_set): +def _mean_dense( + x, weight, bias, input_axes, weight_axes, output_axes, collective_op_set, quantizer_set +): output = dense( x, weight, @@ -66,13 +73,16 @@ def _mean_dense(x, weight, bias, input_axes, weight_axes, output_axes, collectiv kernel_axes=weight_axes, output_axes=output_axes, collective_op_set=collective_op_set, + quantizer_set=quantizer_set, ) return jnp.mean(output.astype(jnp.float32)) -def _value_and_grad_dense(x, weight, bias, input_axes, weight_axes, output_axes, collective_op_set): +def _value_and_grad_dense( + x, weight, bias, input_axes, weight_axes, output_axes, collective_op_set, quantizer_set +): return jax.jit(jax.value_and_grad(_mean_dense, (0, 1, 2)), static_argnums=(3, 4, 5, 6))( - x, weight, bias, input_axes, weight_axes, output_axes, collective_op_set + x, weight, bias, input_axes, weight_axes, output_axes, collective_op_set, quantizer_set ) @@ -98,11 +108,16 @@ def run_dense_grad_tests(args, mesh=None): ) collective_op_set = CollectiveOpSet.create(forward_collective_op=collective_op) + use_quantization = args.quantize_recipe is not None + recipe = get_quantization_recipe(args.quantize_recipe) if use_quantization else None with mesh, autocast( - enabled=False, - recipe=None, + enabled=use_quantization, + recipe=recipe, mesh_resource=MeshResource(dp_resource=DP_AXIS, tpsp_resource=TPSP_AXIS), ): + # Build quantizer_set inside autocast so create_set() reads the global recipe + # for correct fwd/bwd dtypes. + quantizer_set = QuantizerFactory.create_set() if use_quantization else noop_quantizer_set # Get the base axis rules and extend them with TE's rules. This must be done inside autocast axis_rules = flax.linen.get_logical_axis_rules() axis_rules += ((TPSP_AXIS, TPSP_AXIS), (DP_AXIS, DP_AXIS)) @@ -123,6 +138,7 @@ def run_dense_grad_tests(args, mesh=None): weight_axes, output_axes, noop_collective_op_set, + quantizer_set, ) output, sharded_grads = _value_and_grad_dense( x_sharded, @@ -132,6 +148,7 @@ def run_dense_grad_tests(args, mesh=None): weight_axes, output_axes, collective_op_set, + quantizer_set, ) jax.block_until_ready(ref_output) jax.block_until_ready(output) @@ -148,9 +165,10 @@ def run_dense_grad_tests(args, mesh=None): jax.block_until_ready(gathered_ref_grads) if args.enable_result_check and args.process_id == 0: - assert_allclose(ref_output, output, dtype=jnp.bfloat16) + tol_dtype = get_tolerance_dtype(quantizer_set) + assert_allclose(ref_output, output, dtype=tol_dtype) for ref_grad, gathered_grad in zip(gathered_ref_grads, gathered_grads): - assert_allclose(ref_grad, gathered_grad, dtype=jnp.bfloat16) + assert_allclose(ref_grad, gathered_grad, dtype=tol_dtype) class TestCollectiveDenseGradient(unittest.TestCase): @@ -187,6 +205,82 @@ def test_te_bf16_reduce_scatter(self): self.args.collective_type = "reduce_scatter" run_dense_grad_tests(self.args, self.mesh) + def test_te_delayed_scaling_fp8_all_gather(self): + """Test Collective Dense Gradient with FP8 DelayedScaling + AllGather""" + self.args.quantize_recipe = "DelayedScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "all_gather" + run_dense_grad_tests(self.args, self.mesh) + + def test_te_delayed_scaling_fp8_reduce_scatter(self): + """Test Collective Dense Gradient with FP8 DelayedScaling + ReduceScatter""" + self.args.quantize_recipe = "DelayedScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "reduce_scatter" + run_dense_grad_tests(self.args, self.mesh) + + def test_te_current_scaling_fp8_all_gather(self): + """Test Collective Dense Gradient with FP8 Float8CurrentScaling + AllGather""" + self.args.quantize_recipe = "Float8CurrentScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "all_gather" + run_dense_grad_tests(self.args, self.mesh) + + def test_te_current_scaling_fp8_reduce_scatter(self): + """Test Collective Dense Gradient with FP8 Float8CurrentScaling + ReduceScatter""" + self.args.quantize_recipe = "Float8CurrentScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "reduce_scatter" + run_dense_grad_tests(self.args, self.mesh) + + def test_te_mxfp8_all_gather(self): + """Test Collective Dense Gradient with MXFP8BlockScaling + AllGather""" + self.args.quantize_recipe = "MXFP8BlockScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + self.args.collective_type = "all_gather" + run_dense_grad_tests(self.args, self.mesh) + + def test_te_mxfp8_reduce_scatter(self): + """Test Collective Dense Gradient with MXFP8BlockScaling + ReduceScatter""" + self.args.quantize_recipe = "MXFP8BlockScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + self.args.collective_type = "reduce_scatter" + run_dense_grad_tests(self.args, self.mesh) + + # def test_te_nvfp4_all_gather(self): + # """Test Collective Dense Gradient with NVFP4BlockScaling + AllGather""" + # self.args.quantize_recipe = "NVFP4BlockScaling" + # is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + # if not is_supported: + # self.skipTest(reason) + # self.args.collective_type = "all_gather" + # run_dense_grad_tests(self.args, self.mesh) + + # def test_te_nvfp4_reduce_scatter(self): + # """Test Collective Dense Gradient with NVFP4BlockScaling + ReduceScatter""" + # self.args.quantize_recipe = "NVFP4BlockScaling" + # is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + # if not is_supported: + # self.skipTest(reason) + # self.args.collective_type = "reduce_scatter" + # run_dense_grad_tests(self.args, self.mesh) + if __name__ == "__main__": import sys @@ -209,6 +303,6 @@ def test_te_bf16_reduce_scatter(self): args = cgemm_parser( "Collective Dense Gradient test on multi-GPU with tensor parallelism" - ).parse_args([]) + ).parse_args() _initialize_distributed(args) run_dense_grad_tests(args, mesh=None) diff --git a/examples/jax/collective_gemm/test_gemm.py b/examples/jax/collective_gemm/test_gemm.py index ea119713e3..c2db8fc44a 100644 --- a/examples/jax/collective_gemm/test_gemm.py +++ b/examples/jax/collective_gemm/test_gemm.py @@ -22,17 +22,23 @@ from common import ( assert_allclose, + get_tolerance_dtype, _initialize_distributed, _get_dp_and_tp_sizes, _create_mesh, DP_AXIS, TPSP_AXIS, - PARAMS_KEY, cgemm_parser, ) import transformer_engine.jax.cpp_extensions as tex -from transformer_engine.jax.quantize import autocast +from transformer_engine.jax.quantize import ( + autocast, + is_quantize_recipe_supported, + get_quantization_recipe, + QuantizerFactory, + noop_quantizer_set, +) from transformer_engine.jax.cpp_extensions.gemm import CollectiveOp from transformer_engine.jax.sharding import MeshResource @@ -54,31 +60,15 @@ def _get_operand_sharding(mesh, collective_op, is_with_dp): return x_sharding, weight_sharding, bias_sharding, output_sharding -def _get_dp_and_tp_sizes(args): - num_gpu = args.num_processes * args.num_devices_per_process - if args.tensor_parallel_size is None: - num_gpu_dp = 2 if args.enable_data_parallel else 1 - assert ( - num_gpu > 1 and num_gpu % num_gpu_dp == 0 - ), "Number of GPUs must be greater than 1 and divisible by number of data parallel GPUs" - num_gpu_tp = num_gpu // num_gpu_dp - else: - num_gpu_tp = args.tensor_parallel_size - assert ( - num_gpu > 1 and num_gpu % num_gpu_tp == 0 - ), "Number of GPUs must be greater than 1 and divisible by number of data parallel GPUs" - num_gpu_dp = num_gpu // num_gpu_tp - return num_gpu_dp, num_gpu_tp - - @partial(jax.jit, static_argnames=("contracting_dims", "collective_op", "output_sharding")) -def _jitted_cgemm(x, weight, bias, contracting_dims, collective_op, output_sharding): +def _jitted_cgemm(x, weight, bias, quantizer_set, contracting_dims, collective_op, output_sharding): output = tex.gemm( x, weight, bias=bias, contracting_dims=contracting_dims, collective_op=collective_op, + quantizer_set=quantizer_set, ) if output_sharding is not None: output = jax.lax.with_sharding_constraint(output, output_sharding) @@ -107,11 +97,20 @@ def run_gemm_tests(args, mesh=None): else CollectiveOp.REDUCE_SCATTER ) + use_quantization = args.quantize_recipe is not None + recipe = get_quantization_recipe(args.quantize_recipe) if use_quantization else None + + # autocast sets the global recipe (fwd/bwd dtypes) AND the global MeshResource + # (via global_shard_guard) required for collective GEMM sharding axis resolution. with mesh, autocast( - enabled=False, - recipe=None, + enabled=use_quantization, + recipe=recipe, mesh_resource=MeshResource(dp_resource=DP_AXIS, tpsp_resource=TPSP_AXIS), ): + # Build quantizer_set inside autocast so create_set() can read the global recipe + # for correct fwd/bwd dtypes. autocast does not inject quantizers into raw + # tex.gemm() calls, so we must pass quantizer_set explicitly. + quantizer_set = QuantizerFactory.create_set() if use_quantization else noop_quantizer_set print(f"Device mesh: {mesh}") x_sharding, weight_sharding, bias_sharding, output_sharding = _get_operand_sharding( @@ -125,6 +124,7 @@ def run_gemm_tests(args, mesh=None): x_sharded, weight_sharded, bias_sharded, + quantizer_set, contracting_dims=((2,), (0,)), collective_op=CollectiveOp.NONE, output_sharding=output_sharding, @@ -133,6 +133,7 @@ def run_gemm_tests(args, mesh=None): x_sharded, weight_sharded, bias_sharded, + quantizer_set, contracting_dims=((2,), (0,)), collective_op=collective_op, output_sharding=output_sharding, @@ -150,7 +151,9 @@ def run_gemm_tests(args, mesh=None): jax.block_until_ready(gathered_output) if args.enable_result_check and args.process_id == 0: - assert_allclose(gathered_ref_output, gathered_output) + assert_allclose( + gathered_ref_output, gathered_output, dtype=get_tolerance_dtype(quantizer_set) + ) class TestCollectiveGemmWithDP(unittest.TestCase): @@ -186,6 +189,84 @@ def test_te_bf16_reduce_scatter_with_dp(self): self.args.collective_type = "reduce_scatter" run_gemm_tests(self.args, self.mesh) + def test_te_delayed_scaling_fp8_all_gather_with_dp(self): + """Test Collective GEMM with FP8 DelayedScaling + AllGather""" + self.args.quantize_recipe = "DelayedScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "all_gather" + run_gemm_tests(self.args, self.mesh) + + def test_te_delayed_scaling_fp8_reduce_scatter_with_dp(self): + """Test Collective GEMM with FP8 DelayedScaling + ReduceScatter""" + self.args.quantize_recipe = "DelayedScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "reduce_scatter" + run_gemm_tests(self.args, self.mesh) + + def test_te_current_scaling_fp8_all_gather_with_dp(self): + """Test Collective GEMM with FP8 Float8CurrentScaling + AllGather""" + self.args.quantize_recipe = "Float8CurrentScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "all_gather" + run_gemm_tests(self.args, self.mesh) + + def test_te_current_scaling_fp8_reduce_scatter_with_dp(self): + """Test Collective GEMM with FP8 Float8CurrentScaling + ReduceScatter""" + self.args.quantize_recipe = "Float8CurrentScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "reduce_scatter" + run_gemm_tests(self.args, self.mesh) + + def test_te_mxfp8_all_gather_with_dp(self): + """Test Collective GEMM with MXFP8BlockScaling + AllGather""" + self.args.quantize_recipe = "MXFP8BlockScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "all_gather" + run_gemm_tests(self.args, self.mesh) + + def test_te_mxfp8_reduce_scatter_with_dp(self): + """Test Collective GEMM with MXFP8BlockScaling + ReduceScatter""" + self.args.quantize_recipe = "MXFP8BlockScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + self.args.collective_type = "reduce_scatter" + run_gemm_tests(self.args, self.mesh) + + # def test_te_nvfp4_all_gather_with_dp(self): + # """Test Collective GEMM with NVFP4BlockScaling + AllGather""" + # self.args.quantize_recipe = "NVFP4BlockScaling" + # is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + # if not is_supported: + # self.skipTest(reason) + # self.args.collective_type = "all_gather" + # run_gemm_tests(self.args, self.mesh) + + # def test_te_nvfp4_reduce_scatter_with_dp(self): + # """Test Collective GEMM with NVFP4BlockScaling + ReduceScatter""" + # self.args.quantize_recipe = "NVFP4BlockScaling" + # is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + # if not is_supported: + # self.skipTest(reason) + # self.args.collective_type = "reduce_scatter" + # run_gemm_tests(self.args, self.mesh) + if __name__ == "__main__": import sys diff --git a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py index 84cb011da1..be94c68d37 100644 --- a/examples/jax/collective_gemm/test_layernorm_mlp_grad.py +++ b/examples/jax/collective_gemm/test_layernorm_mlp_grad.py @@ -2,7 +2,6 @@ # # See LICENSE for license information. """Collective Dense Gradient test on multi-GPU with tensor parallelism""" -import argparse import unittest import os @@ -13,18 +12,24 @@ from common import ( assert_allclose, + get_tolerance_dtype, _initialize_distributed, _get_dp_and_tp_sizes, _create_mesh, DP_AXIS, TPSP_AXIS, - PARAMS_KEY, cgemm_parser, ) from transformer_engine.jax.layernorm_mlp import layernorm_mlp -from transformer_engine.jax.quantize import autocast +from transformer_engine.jax.quantize import ( + autocast, + is_quantize_recipe_supported, + get_quantization_recipe, + QuantizerFactory, + noop_quantizer_set, +) from transformer_engine.jax.cpp_extensions.gemm import ( CollectiveOpSet, CollectiveOp, @@ -68,6 +73,7 @@ def _mean_layernorm_mlp( weight_1_axes, weight_2_axes, collective_op_sets, + quantizer_sets, ): output = layernorm_mlp( x, @@ -82,6 +88,7 @@ def _mean_layernorm_mlp( kernel_2_axes=weight_2_axes, activation_type=("gelu",), collective_op_sets=collective_op_sets, + quantizer_sets=quantizer_sets, ) return jnp.mean(output) @@ -98,6 +105,7 @@ def _value_and_grad_layernorm_mlp( weight_1_axes, weight_2_axes, collective_op_sets, + quantizer_sets, ): return jax.jit( jax.value_and_grad(_mean_layernorm_mlp, (0, 1, 2, 3, 4, 5)), static_argnums=(6, 7, 8, 9, 10) @@ -113,11 +121,12 @@ def _value_and_grad_layernorm_mlp( weight_1_axes, weight_2_axes, collective_op_sets, + quantizer_sets, ) def run_layernorm_mlp_grad_tests(args, mesh=None): - """Execute Dense Gradient tests.""" + """Execute LayerNorm MLP Gradient tests.""" print(args) # Initialize distributed with provided arguments @@ -149,11 +158,21 @@ def run_layernorm_mlp_grad_tests(args, mesh=None): collective_op_sets = (collective_op_set_1, collective_op_set_2) noop_collective_op_sets = (noop_collective_op_set, noop_collective_op_set) + use_quantization = args.quantize_recipe is not None + recipe = get_quantization_recipe(args.quantize_recipe) if use_quantization else None with mesh, autocast( - enabled=False, - recipe=None, + enabled=use_quantization, + recipe=recipe, mesh_resource=MeshResource(dp_resource=DP_AXIS, tpsp_resource=TPSP_AXIS), ): + # Build quantizer_sets inside autocast so create_set() reads the global recipe + # for correct fwd/bwd dtypes. One set per dense layer (GEMM1=AG, GEMM2=RS). + quantizer_sets = ( + QuantizerFactory.create_set(n_quantizer_sets=2) + if use_quantization + else (noop_quantizer_set, noop_quantizer_set) + ) + # Get the base axis rules and extend them with TE's rules. This must be done inside autocast axis_rules = flax.linen.get_logical_axis_rules() axis_rules += ((TPSP_AXIS, TPSP_AXIS), (DP_AXIS, DP_AXIS)) @@ -181,6 +200,7 @@ def run_layernorm_mlp_grad_tests(args, mesh=None): weight_1_axes, weight_2_axes, noop_collective_op_sets, + quantizer_sets, ) output, sharded_grads = _value_and_grad_layernorm_mlp( x_sharded, @@ -194,6 +214,7 @@ def run_layernorm_mlp_grad_tests(args, mesh=None): weight_1_axes, weight_2_axes, collective_op_sets, + quantizer_sets, ) jax.block_until_ready(ref_output) jax.block_until_ready(output) @@ -210,13 +231,14 @@ def run_layernorm_mlp_grad_tests(args, mesh=None): jax.block_until_ready(gathered_ref_grads) if args.enable_result_check and args.process_id == 0: - assert_allclose(ref_output, output, dtype=jnp.bfloat16) + tol_dtype = get_tolerance_dtype(quantizer_sets[0]) + assert_allclose(ref_output, output, dtype=tol_dtype) for ref_grad, gathered_grad in zip(gathered_ref_grads, gathered_grads): - assert_allclose(ref_grad, gathered_grad, dtype=jnp.bfloat16) + assert_allclose(ref_grad, gathered_grad, dtype=tol_dtype) class TestCollectiveLayerNormMLPGradient(unittest.TestCase): - """Collective Dense Gradient unittests""" + """Collective LayerNorm MLP Gradient unittests""" def setUp(self): self.args = cgemm_parser( @@ -240,9 +262,43 @@ def tearDown(self): os.environ.pop("NVTE_JAX_ALL_REDUCE_IN_FP32", None) def test_te_bf16_layernorm_mlp_grad(self): - """Test Collective Dense Gradient with AllGather""" + """Test Collective LayerNorm MLP Gradient with BF16""" + run_layernorm_mlp_grad_tests(self.args, self.mesh) + + def test_te_delayed_scaling_fp8_layernorm_mlp_grad(self): + """Test Collective LayerNorm MLP Gradient with FP8 DelayedScaling""" + self.args.quantize_recipe = "DelayedScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + run_layernorm_mlp_grad_tests(self.args, self.mesh) + def test_te_current_scaling_fp8_layernorm_mlp_grad(self): + """Test Collective LayerNorm MLP Gradient with FP8 Float8CurrentScaling""" + self.args.quantize_recipe = "Float8CurrentScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + + run_layernorm_mlp_grad_tests(self.args, self.mesh) + + def test_te_mxfp8_layernorm_mlp_grad(self): + """Test Collective LayerNorm MLP Gradient with MXFP8BlockScaling""" + self.args.quantize_recipe = "MXFP8BlockScaling" + is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + if not is_supported: + self.skipTest(reason) + run_layernorm_mlp_grad_tests(self.args, self.mesh) + + # def test_te_nvfp4_layernorm_mlp_grad(self): + # """Test Collective LayerNorm MLP Gradient with NVFP4BlockScaling""" + # self.args.quantize_recipe = "NVFP4BlockScaling" + # is_supported, reason = is_quantize_recipe_supported(self.args.quantize_recipe) + # if not is_supported: + # self.skipTest(reason) + # run_layernorm_mlp_grad_tests(self.args, self.mesh) + if __name__ == "__main__": import sys @@ -265,6 +321,6 @@ def test_te_bf16_layernorm_mlp_grad(self): args = cgemm_parser( "Collective LayerNorm MLP Gradient test on multi-GPU with tensor parallelism" - ).parse_args([]) + ).parse_args() _initialize_distributed(args) run_layernorm_mlp_grad_tests(args, mesh=None) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 4506adf33b..515f02af6e 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -407,6 +407,48 @@ def assert_cublas_requirements(scaling_mode, contracting_size, tensor_name): ) +def _reorder_tpsp_leading(tensor, original_shape): + """Reorder tensor so the tpsp axis is leading: reshape (dp, n, tpsp, m, ...), transpose (2, 0, 1, 3, ...).""" + assert original_shape[0] % dp_or_fsdp_axis_size() == 0 or original_shape[0] == 1, ( + f"Original_shape[0]={original_shape[0]} is not divisible by" + f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" + ) + assert original_shape[1] % tpsp_axis_size() == 0 or original_shape[1] == 1, ( + f"Original_shape[1]={original_shape[1]} is not divisible by" + f" tpsp_axis_size()={tpsp_axis_size()}" + ) + reshaped = tensor.reshape( + dp_or_fsdp_axis_size(), + int(original_shape[0] / dp_or_fsdp_axis_size()), + tpsp_axis_size(), + int(original_shape[1] / tpsp_axis_size()), + *original_shape[2:], + ) + reordered = reshaped.transpose(2, 0, 1, 3, *range(4, reshaped.ndim)) + return reordered.reshape(original_shape) + + +def _reorder_dp_leading(tensor, original_shape): + """Reorder tensor so the dp axis is leading: reshape (tpsp, dp, n, m, ...), transpose (1, 2, 0, 3, ...).""" + assert original_shape[0] % dp_or_fsdp_axis_size() == 0 or original_shape[0] == 1, ( + f"Original_shape[0]={original_shape[0]} is not divisible by" + f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" + ) + assert original_shape[1] % tpsp_axis_size() == 0 or original_shape[1] == 1, ( + f"Original_shape[1]={original_shape[1]} is not divisible by" + f" tpsp_axis_size()={tpsp_axis_size()}" + ) + reshaped = tensor.reshape( + tpsp_axis_size(), + dp_or_fsdp_axis_size(), + int(original_shape[0] / dp_or_fsdp_axis_size()), + int(original_shape[1] / tpsp_axis_size()), + *original_shape[2:], + ) + reordered = reshaped.transpose(1, 2, 0, 3, *range(4, reshaped.ndim)) + return reordered.reshape(original_shape) + + class GemmPrimitive(BasePrimitive): """ Primitive for cuBLAS GEMM @@ -639,47 +681,64 @@ def impl( lhs_flatten_axis = max(lhs_cdims) + 1 if lhs_transposed else min(lhs_cdims) rhs_flatten_axis = min(rhs_cdims) if rhs_transposed else max(rhs_cdims) + 1 - lhs_scale_inv = apply_padding_to_scale_inv( - lhs_scale_inv, scaling_mode, lhs.shape, lhs_transposed, lhs_flatten_axis - ) - rhs_scale_inv = apply_padding_to_scale_inv( - rhs_scale_inv, scaling_mode, rhs.shape, not rhs_transposed, rhs_flatten_axis - ) + if not collective_op.is_none and not is_outer: + # MXFP8 + Collective AG/RS: both sides of flatten_axis must be multiples of 128. + # No padding is needed in this case + lhs_first, lhs_last = math.prod(lhs.shape[:lhs_flatten_axis]), math.prod( + lhs.shape[lhs_flatten_axis:] + ) + assert lhs_first % 128 == 0 and lhs_last % 128 == 0, ( + "MXFP8 + Collective AG/RS requires LHS dimensions before and after the flatten" + f" axis to be multiples of 128. Got lhs.shape={lhs.shape}," + f" lhs_flatten_axis={lhs_flatten_axis}" + ) + rhs_first, rhs_last = math.prod(rhs.shape[:rhs_flatten_axis]), math.prod( + rhs.shape[rhs_flatten_axis:] + ) + assert rhs_first % 128 == 0 and rhs_last % 128 == 0, ( + "MXFP8 + Collective AG/RS requires LHS dimensions before and after the flatten" + f" axis to be multiples of 128. Got rhs.shape={rhs.shape}," + f" rhs_flatten_axis={rhs_flatten_axis}" + ) + # The scale needs to be in good shape for reordering + assert lhs_scale_inv.shape[sequence_dim] % tpsp_axis_size() == 0, ( + "MXFP8 + Collective AG/RS requires RHS scale inv sequence dimension to be" + f" multiples of tpsp_axis_size. Got lhs_scale_inv.shape={lhs_scale_inv.shape}," + f" tpsp_axis_size={tpsp_axis_size()}, sequence_dim={sequence_dim}" + ) + else: + lhs_scale_inv = apply_padding_to_scale_inv( + lhs_scale_inv, + scaling_mode, + lhs.shape, + lhs_transposed, + lhs_flatten_axis, + ) + rhs_scale_inv = apply_padding_to_scale_inv( + rhs_scale_inv, scaling_mode, rhs.shape, not rhs_transposed, rhs_flatten_axis + ) # Only perform JAX-based swizzle for MXFP8, NVFP4 swizzle will go though nvte kernel if scaling_mode.is_mxfp8_scaling: lhs_scale_inv = swizzled_scale(lhs_scale_inv, lhs_flatten_axis, lhs_transposed) rhs_scale_inv = swizzled_scale(rhs_scale_inv, rhs_flatten_axis, not rhs_transposed) + # Determine if we need to reorder the tensor so that the input/output are in the correct layout for the collective operation + need_reorder = not transpose_batch_sequence and not is_outer and not collective_op.is_none + # Alter lhs blocks so that CGEMM RS outputs correctly + if need_reorder and collective_op.is_reduce_scatter and lhs.shape[0] != 1: + assert sequence_dim == 1, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + lhs = _reorder_tpsp_leading(lhs, lhs.shape) + if ( - collective_op.is_reduce_scatter - and not transpose_batch_sequence - and not is_outer - and not lhs.shape[0] == 1 + need_reorder + and (collective_op.is_reduce_scatter or collective_op.is_all_gather) + and lhs_scale_inv.shape[0] != 1 + and scaling_mode.is_1d_block_scaling() ): - if sequence_dim != 1: - raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") - original_shape = lhs.shape - if original_shape[0] % dp_or_fsdp_axis_size() != 0 and original_shape[0] != 1: - raise ValueError( - f"Original_shape[0]={original_shape[0]} is not divisible by" - f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" - ) - if original_shape[1] % tpsp_axis_size() != 0 and original_shape[1] != 1: - raise ValueError( - f"Original_shape[1]={original_shape[1]} is not divisible by" - f" tpsp_axis_size()={tpsp_axis_size()}" - ) - reshaped = lhs.reshape( - dp_or_fsdp_axis_size(), - int(original_shape[0] / dp_or_fsdp_axis_size()), - tpsp_axis_size(), - int(original_shape[1] / tpsp_axis_size()), - *original_shape[2:], - ) - reordered = reshaped.transpose(2, 0, 1, 3, *range(4, reshaped.ndim)) - lhs = reordered.reshape(original_shape) + assert sequence_dim == 1, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + lhs_scale_inv = _reorder_tpsp_leading(lhs_scale_inv, lhs_scale_inv.shape) (output, _) = GemmPrimitive.inner_primitive.bind( lhs, @@ -699,34 +758,9 @@ def impl( collective_op=collective_op, ) # Alter output blocks for CGEMM AG - if ( - collective_op.is_all_gather - and not transpose_batch_sequence - and not is_outer - and not output.shape[0] == 1 - ): - if sequence_dim != 1: - raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") - original_shape = output.shape - if original_shape[0] % dp_or_fsdp_axis_size() != 0 and original_shape[0] != 1: - raise ValueError( - f"Original_shape[0]={original_shape[0]} is not divisible by" - f" dp_or_fsdp_axis_size()={dp_or_fsdp_axis_size()}" - ) - if original_shape[1] % tpsp_axis_size() != 0 and original_shape[1] != 1: - raise ValueError( - f"Original_shape[1]={original_shape[1]} is not divisible by" - f" tpsp_axis_size()={tpsp_axis_size()}" - ) - reshaped = output.reshape( - tpsp_axis_size(), - dp_or_fsdp_axis_size(), - int(original_shape[0] / dp_or_fsdp_axis_size()), - int(original_shape[1] / tpsp_axis_size()), - *original_shape[2:], - ) - reordered = reshaped.transpose(1, 2, 0, 3, *range(4, reshaped.ndim)) - output = reordered.reshape(original_shape) + if need_reorder and collective_op.is_all_gather and output.shape[0] != 1: + assert sequence_dim == 1, f"Invalid sequence_dim. Got sequence_dim={sequence_dim}" + output = _reorder_dp_leading(output, output.shape) return (output,) @@ -812,6 +846,7 @@ def _parse_operand_output_specs( contracting_dims, transpose_batch_sequence, collective_op, + scaling_mode, ): lhs_specs, _, rhs_specs, *_ = map(get_padded_spec, arg_infos) @@ -955,12 +990,25 @@ def _parse_operand_output_specs( # Bias sharding is based on GEMM output before any scatter bias_specs = rhs_non_cspecs if arg_infos[4].size > 0 else (None,) # bias is operand index 4 + # Scale shardings are based on the scaling_mode and collective_op + lhs_scale_specs = rhs_scale_specs = (None,) + if scaling_mode.is_1d_block_scaling(): + rhs_scale_specs = rhs_specs + # Set the seq spec to None to trigger AG the scales as TE/Common CGEMM does not handle + # scale collecting yet + if collective_op.is_all_gather: + lhs_scale_specs = tuple( + None if i == sequence_dim else s for i, s in enumerate(lhs_specs) + ) + else: + lhs_scale_specs = lhs_specs + if not collective_op.is_none: if sequence_dim < 0: raise ValueError(f"Invalid sequence_dim. Got sequence_dim={sequence_dim}") return ( - (lhs_specs, rhs_specs, bias_specs), + (lhs_specs, lhs_scale_specs, rhs_specs, rhs_scale_specs, bias_specs), out_specs, reduce_spec, sequence_dim, @@ -982,7 +1030,6 @@ def infer_sharding_from_operands( ): del ( out_dtype, - scaling_mode, use_split_accumulator, result_infos, is_outer, @@ -990,7 +1037,11 @@ def infer_sharding_from_operands( ) (_, out_specs, *_) = GemmPrimitive._parse_operand_output_specs( - arg_infos, contracting_dims, transpose_batch_sequence, collective_op + arg_infos, + contracting_dims, + transpose_batch_sequence, + collective_op, + scaling_mode, ) out_sharding = NamedSharding(mesh, PartitionSpec(*out_specs)) @@ -1013,7 +1064,7 @@ def partition( del result_infos, is_outer, sequence_dim ( - (lhs_specs, rhs_specs, bias_input_specs), + (lhs_specs, lhs_scale_specs, rhs_specs, rhs_scale_specs, bias_input_specs), out_specs, reduce_spec, inferred_sequence_dim, @@ -1022,17 +1073,21 @@ def partition( contracting_dims, transpose_batch_sequence, collective_op, + scaling_mode, ) # Block scale inverses match their operands, but tensor scale inverses are unsharded. none_sharding = NamedSharding(mesh, PartitionSpec(None)) lhs_sharding = NamedSharding(mesh, PartitionSpec(*lhs_specs)) + lhs_scale_sharding = NamedSharding(mesh, PartitionSpec(*lhs_scale_specs)) rhs_sharding = NamedSharding(mesh, PartitionSpec(*rhs_specs)) + rhs_scale_sharding = NamedSharding(mesh, PartitionSpec(*rhs_scale_specs)) + arg_shardings = ( lhs_sharding, - lhs_sharding if scaling_mode.is_1d_block_scaling() else none_sharding, + lhs_scale_sharding, rhs_sharding, - rhs_sharding if scaling_mode.is_1d_block_scaling() else none_sharding, + rhs_scale_sharding, ) # Bias @@ -1247,6 +1302,12 @@ def _te_gemm( rhs_tensor_scale_inv = _get_nvfp4_tensor_scale_inv(rhs_amax) alpha = lhs_tensor_scale_inv * rhs_tensor_scale_inv + if not collective_op.is_none: + assert not scaling_mode.is_nvfp4_scaling, ( + f"Collective GEMM is not yet supported with {scaling_mode} quantization. Only" + " DELAYED_TENSOR_SCALING, CURRENT_TENSOR_SCALING, and MXFP8_1D_SCALING are supported." + ) + out_dtype = lhs_q.dq_dtype if isinstance(lhs_q, ScaledTensor) else lhs_data.dtype if bias is None: bias = jnp.empty(0, dtype=out_dtype) diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 737dd65622..2acefa2d30 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -58,6 +58,7 @@ std::tuple> xla_buffer_to_nvte_gemm_operand( std::vector scale_shape = {1}; auto is_nvfp4 = is_nvfp4_scaling(scaling_mode); auto scale_dtype = convert_ffi_datatype_to_te_dtype(scale_inv.element_type()); + if (scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING || is_nvfp4) { // Block scaling also needs to be collapsed to match 2D data scale_shape = {product(scale_inv.dimensions(), 0, axis_boundary), @@ -202,6 +203,7 @@ Error_Type GemmV2FFI(cudaStream_t stream, Buffer_Type lhs, Buffer_Type lhs_scale auto workspace_ptr = reinterpret_cast(workspace->untyped_data()); workspace_ptr = move_ptr_to_next_256B_aligned(workspace_ptr); size_t workspace_size = static_cast(workspace->element_count()) - 256; + if (is_nvfp4_scaling(config.scaling_mode)) { auto lhs_scale_size = product(lhs_scale_inv.dimensions()); auto rhs_scale_size = product(rhs_scale_inv.dimensions()); diff --git a/transformer_engine/jax/quantize/helper.py b/transformer_engine/jax/quantize/helper.py index c491bb8638..3a93af4a68 100644 --- a/transformer_engine/jax/quantize/helper.py +++ b/transformer_engine/jax/quantize/helper.py @@ -51,6 +51,8 @@ "fp8_autocast", "is_fp8_available", "is_scaling_mode_supported", + "is_quantize_recipe_supported", + "get_quantization_recipe", "get_supported_scaling_modes", "get_supported_quantization_recipes", "update_collections", @@ -162,6 +164,54 @@ def is_scaling_mode_supported( return _is_scaling_mode_supported[scaling_mode], _reason_for_no_scaling_mode[scaling_mode] +_RECIPE_NAME_TO_RECIPE = { + "DelayedScaling": DelayedScaling, + "Float8CurrentScaling": Float8CurrentScaling, + "MXFP8BlockScaling": MXFP8BlockScaling, + "NVFP4BlockScaling": NVFP4BlockScaling, +} + + +def get_quantization_recipe(name: str) -> Recipe: + """Return a recipe object from a recipe name string. + + Args: + name: Recipe name. One of "DelayedScaling", "Float8CurrentScaling", + "MXFP8BlockScaling", or "NVFP4BlockScaling". + + Returns: + A new instance of the corresponding recipe class. + + Raises: + ValueError: If ``name`` does not match any known recipe. + """ + recipe_cls = _RECIPE_NAME_TO_RECIPE.get(name) + if recipe_cls is None: + valid = list(_RECIPE_NAME_TO_RECIPE) + raise ValueError(f"Invalid quantization recipe '{name}'. Valid options: {valid}") + return recipe_cls() + + +def is_quantize_recipe_supported(recipe_name: str) -> Tuple[bool, str]: + """Check if the given quantization recipe (by name) is supported on the current GPU. + + Args: + recipe_name: Name of the recipe, e.g. "DelayedScaling", "Float8CurrentScaling", + "MXFP8BlockScaling", "NVFP4BlockScaling". + + Returns: + A tuple of (supported: bool, reason: str). + """ + recipe = get_quantization_recipe(recipe_name) + config = get_quantize_config_with_recipe(recipe) + for tensor_source in TensorSource: + scaling_mode = config.get_scaling_mode(tensor_source) + is_supported, reason = is_scaling_mode_supported(scaling_mode) + if not is_supported: + return is_supported, reason + return True, None + + def is_fp8_available( scaling_mode=ScalingMode.DELAYED_TENSOR_SCALING, gpu_id=None, @@ -916,6 +966,7 @@ def apply_padding_to_scale_inv( unpadded_scale_shape = scaling_mode.get_scale_shape( data_shape, is_colwise=is_colwise, is_padded=False, flatten_axis=flatten_axis ) + assert scale_inv.shape == unpadded_scale_shape, ( f"Unpadded inverse scale factor has wrong shape, expected {unpadded_scale_shape} but got " f"{scale_inv.shape}." From fcceeb961950bfc1e814db3fd438ed2f5e15e47e Mon Sep 17 00:00:00 2001 From: jomitchellnv <148147880+jomitchellnv@users.noreply.github.com> Date: Fri, 13 Mar 2026 12:29:34 -0700 Subject: [PATCH 278/521] [Pytorch] Add QuantizedTensor support in FusedAdam.step for MXFP8BlockScaling and Float8BlockScaling quantized model init. (#2753) * Updates FusedAdam with FSDP2 and MXFP8 Signed-off-by: Jonathan Mitchell * removes xfailing unit test for MXFPr MXFP8 Signed-off-by: Jonathan Mitchell * addresses comments related to reset parameters and guard against self.capturable Signed-off-by: Jonathan Mitchell * adds e2e unit test Signed-off-by: Jonathan Mitchell * adds test to non meta device init Signed-off-by: Jonathan Mitchell * attempts to add float8block scaling fsdp hooks Signed-off-by: Jonathan Mitchell * adds e2e test for Float8BlockScaling Signed-off-by: Jonathan Mitchell * addresses review comments and code cleanup Signed-off-by: Jonathan Mitchell * more review comments addressed Signed-off-by: Jonathan Mitchell * removes unused block_len param Signed-off-by: Jonathan Mitchell * fixes failing unit test because we still need to xfail nvfp4 dcp Signed-off-by: Jonathan Mitchell * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint - replacing todo with note Signed-off-by: Jonathan Mitchell --------- Signed-off-by: Jonathan Mitchell Signed-off-by: Jonathan Mitchell Signed-off-by: Jonathan Mitchell Co-authored-by: Jonathan Mitchell Co-authored-by: Jonathan Mitchell Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 Co-authored-by: Jonathan Mitchell --- .../distributed/run_fsdp2_fused_adam.py | 97 +++++-- tests/pytorch/distributed/test_torch_fsdp2.py | 56 +++- tests/pytorch/test_fused_optimizer.py | 266 +++++++++++++++++- .../pytorch/optimizers/fused_adam.py | 36 +++ .../pytorch/tensor/float8_blockwise_tensor.py | 164 +++++++++++ 5 files changed, 589 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/run_fsdp2_fused_adam.py index 0439bf1b5a..f97f2bb22e 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/run_fsdp2_fused_adam.py @@ -70,14 +70,36 @@ def _setup(): return world_size, local_rank, device -def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None): - """Build a Sequential of TransformerLayers, optionally with FP8 init.""" +def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_device=True): + """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(). + + 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. + """ if fp8_init: ctx = te.quantized_model_init(enabled=True, recipe=recipe) else: from contextlib import nullcontext ctx = nullcontext() + kwargs = dict( + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if fp8_init and use_meta_device: + kwargs["device"] = "meta" with ctx: model = torch.nn.Sequential( *[ @@ -85,11 +107,7 @@ def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None): HIDDEN_SIZE, FFN_HIDDEN_SIZE, NUM_ATTENTION_HEADS, - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - fuse_qkv_params=True, - params_dtype=torch.bfloat16, - hidden_dropout=0.0, - attention_dropout=0.0, + **kwargs, ) for _ in range(NUM_LAYERS) ] @@ -98,12 +116,29 @@ def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None): def _shard_model(model, world_size): - """Apply FSDP2 sharding with save/restore custom attrs.""" + """Apply FSDP2 sharding with save/restore custom attrs. + + If the model was created on the meta device (e.g. for FP8 init), + parameters are materialized after sharding via reset_parameters(). + + restore_custom_attrs is called last so it applies to the final parameter + objects. For meta-device models, reset_parameters() replaces params via + module_setattr (base.py:1336-1339), so attrs must be restored afterward. + """ + has_meta_params = any(p.is_meta for p in model.parameters()) custom_attrs = save_custom_attrs(model) mesh = DeviceMesh("cuda", list(range(world_size))) for child in model.children(): fully_shard(child, mesh=mesh) fully_shard(model, mesh=mesh) + if has_meta_params: + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + # Restore after reset_parameters so attrs land on the final param objects. + # save_custom_attrs skips private attrs (_*) on QuantizedTensor params; + # reset_parameters fully reinitializes quantizer state from + # self.param_init_meta, so no private attrs need restoring. restore_custom_attrs(model, custom_attrs) return model @@ -119,18 +154,11 @@ def test_fused_adam_fp8_master_weights(recipe=None): world_size, _, device = _setup() model = _build_model(fp8_init=True, recipe=recipe) - - # Verify FP8 params created - qt_count = sum(1 for _, p in model.named_parameters() if isinstance(p, QuantizedTensor)) - assert qt_count > 0, "No QuantizedTensor local tensors before training" - model = _shard_model(model, world_size) - # Verify params are DTensors + # Verify params are DTensors with QuantizedTensor local shards for name, param in model.named_parameters(): assert isinstance(param, DTensor), f"{name} is not DTensor" - - # Verify FP8 params after sharding qt_count = sum( 1 for _, p in model.named_parameters() @@ -181,6 +209,42 @@ def test_fused_adam_fp8_master_weights(recipe=None): dist.destroy_process_group() +def test_fused_adam_fp8_master_weights_no_meta(recipe=None): + """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. + + For per-tensor FP8 (DelayedScaling, Float8CurrentScaling) this works + because Float8Tensor's storage is accessible via data_ptr(). + """ + world_size, _, device = _setup() + + model = _build_model(fp8_init=True, recipe=recipe, use_meta_device=False) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + dist.destroy_process_group() + + def test_fused_adam_bf16(recipe=None): """FusedAdam with master_weights + FSDP2 + bf16 params (no FP8). @@ -622,6 +686,7 @@ def test_dcp_output_parity(recipe=None, async_save=False): TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, + "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, "fused_adam_bf16": test_fused_adam_bf16, "fused_adam_fp8_no_master": test_fused_adam_fp8_no_master, "fused_adam_bf16_store_param_remainders": test_fused_adam_bf16_store_param_remainders, diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index b10f31ea07..042028a949 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -41,7 +41,7 @@ def check_nvfp4_support(): def _parametrize_fp8_recipes(): - """Generate pytest.param objects with xfail marks for unsupported FP8 recipes.""" + """Generate pytest.param objects with skip marks for unsupported FP8 recipes.""" params = [] for name, check_fn in _FP8_RECIPE_CONFIGS: supported, reason = check_fn() @@ -49,7 +49,7 @@ def _parametrize_fp8_recipes(): pytest.param( name, id=name, - marks=pytest.mark.xfail(condition=not supported, reason=reason), + marks=pytest.mark.skipif(not supported, reason=reason), ) ) return params @@ -115,8 +115,8 @@ def _run_fused_adam_test(test_name, recipe="delayed_scaling"): @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") def test_fsdp2_fused_adam_fp8_master_weights(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init.""" - if fp_recipe in ("Float8BlockScaling", "MXFP8BlockScaling", "NVFP4BlockScaling"): + """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (meta device init).""" + if fp_recipe in ("NVFP4BlockScaling",): pytest.xfail( f"{fp_recipe}: quantized_model_init and FSDP2 is not currently supported, since the " "block tensor is dequantized before we flatten it for FSDP2." @@ -124,6 +124,25 @@ def test_fsdp2_fused_adam_fp8_master_weights(fp_recipe): _run_fused_adam_test("fused_adam_fp8_master_weights", fp_recipe) +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +def test_fsdp2_fused_adam_fp8_master_weights_no_meta(fp_recipe): + """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (CUDA init, no meta device). + + Block-scaling QuantizedTensors (MXFP8, Float8Blockwise, NVFP4) are wrapper + subclasses with data_ptr() == 0. Without meta-device init, FSDP2's + reset_sharded_param() crashes with 'invalid python storage'. + Per-tensor FP8 (DelayedScaling, Float8CurrentScaling) works because + Float8Tensor's storage is accessible. + """ + if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{fp_recipe}: FSDP2 without meta-device init crashes on block-scaling " + "QuantizedTensor wrapper subclasses (data_ptr() == 0). " + "Use device='meta' + reset_parameters() after sharding." + ) + _run_fused_adam_test("fused_adam_fp8_master_weights_no_meta", fp_recipe) + + @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") def test_fsdp2_fused_adam_bf16(fp_recipe): """FusedAdam(master_weights=True) + FSDP2 + bf16 params (no FP8).""" @@ -133,10 +152,10 @@ def test_fsdp2_fused_adam_bf16(fp_recipe): @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") def test_fsdp2_fused_adam_fp8_no_master(fp_recipe): """FusedAdam(master_weights=False) + FSDP2 + FP8 params.""" - if fp_recipe == "MXFP8BlockScaling": + if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" + f"{fp_recipe}: FusedAdam without master_weights does not support " + "block-scaling quantized tensors. Use master_weights=True." ) _run_fused_adam_test("fused_adam_fp8_no_master", fp_recipe) @@ -156,6 +175,12 @@ def test_fsdp2_dcp_output_parity(fp_recipe): "MXFP8 quantized tensors, causing illegal memory access" ) + if fp_recipe == "NVFP4BlockScaling": + pytest.xfail( + "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " + "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" + ) + if fp_recipe == "Float8BlockScaling" and torch.cuda.get_device_capability()[0] == 12: pytest.xfail( "Float8BlockScaling is failing on SM120 with RuntimeError: " @@ -189,13 +214,18 @@ def test_fsdp2_dcp_output_parity_async(fp_recipe): "multi_tensor_apply: CUDA Error: an illegal memory access was encountered" ) - if fp_recipe == "Float8BlockScaling" and torch.cuda.get_device_capability()[0] == 12: + if fp_recipe == "NVFP4BlockScaling": pytest.xfail( - "Float8BlockScaling is failing on SM120 with RuntimeError: " - "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " - "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " - "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " - "requires using power of two scaling factors." + "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " + "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" + ) + + if fp_recipe == "Float8BlockScaling": + 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." ) _run_fused_adam_test("dcp_output_parity_async", fp_recipe) diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index 185b9b85bc..e72cad9db1 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -10,8 +10,9 @@ 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 +from transformer_engine.common.recipe import DelayedScaling, MXFP8BlockScaling, Float8BlockScaling from transformer_engine.pytorch import MultiheadAttention, quantized_model_init, is_bf16_available +from transformer_engine.pytorch import QuantizedTensor from transformer_engine.pytorch.utils import gpu_autocast_ctx # Check if FP8 is supported @@ -519,6 +520,269 @@ def test_fp8_model_weight_cast(self): ) +class TestFusedAdamMXFP8(TestFusedOptimizer): + """FusedAdam with MXFP8BlockScaling quantized primary weights (single GPU, no FSDP).""" + + def setup_method(self) -> None: + super().setup_method(iters=5) + mxfp8_available, self.mxfp8_reason = te.is_mxfp8_available(return_reason=True) + self.mxfp8_available = mxfp8_available + + def _build_model(self): + recipe = MXFP8BlockScaling() + with quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, recipe + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_mxfp8_linear_fused_adam_master_weights(self): + """quantized_model_init(MXFP8) + te.Linear + FusedAdam(master_weights=True). + + Verifies: + - Model params are MXFP8 QuantizedTensors after init + - FP32 master weights track a reference Adam optimizer + - Params remain QuantizedTensors after training + - Loss decreases over training steps + """ + if not self.mxfp8_available: + pytest.skip(self.mxfp8_reason) + + model, recipe = self._build_model() + + # Verify weight params are QuantizedTensors (bias stays bf16) + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, QuantizedTensor + ), f"Expected QuantizedTensor for {name}, got {type(p).__name__}" + + # Build reference: clone dequantized weights for a plain Adam + ref_params = [p.detach().clone().float() for p in model.parameters()] + + options = {"lr": 5e-4, "betas": (0.9, 0.999), "eps": 1e-8, "weight_decay": 0} + ref_optim = torch.optim.Adam(ref_params, **options) + tst_optim = te.optimizers.FusedAdam( + list(model.parameters()), + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + **options, + ) + + for _ in range(self.iters): + for p_ref, p in zip(ref_params, model.parameters()): + p_ref.grad = torch.rand_like(p_ref) + p.decoupled_grad = p_ref.grad.clone() + ref_optim.step() + tst_optim.step() + + # FP32 master weights should match reference Adam exactly + master_params = [ + tst_optim.get_unscaled_state(p, "master_param") for p in model.parameters() + ] + torch.testing.assert_close(ref_params, master_params) + + # Weight params should still be QuantizedTensors after training + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, QuantizedTensor + ), f"{name} lost QuantizedTensor type after training: {type(p).__name__}" + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_mxfp8_linear_forward_backward_step(self): + """End-to-end: quantized_model_init + autocast forward + backward + FusedAdam.step(). + + Uses te.autocast with MXFP8BlockScaling recipe for the forward pass, + verifying the full training loop works with quantized compute. + """ + if not self.mxfp8_available: + pytest.skip(self.mxfp8_reason) + + model, recipe = self._build_model() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + batch_size, seq_len, hidden = 4, 32, 256 + x = torch.randn(batch_size, seq_len, hidden, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(self.iters): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + + # Verify all params have non-None gradients after backward + for name, p in model.named_parameters(): + assert p.grad is not None, f"Step {i}: {name} has no gradient after backward" + assert ( + p.grad.shape == p.shape + ), f"Step {i}: {name} grad shape {p.grad.shape} != param shape {p.shape}" + assert torch.isfinite(p.grad).all(), f"Step {i}: {name} has non-finite gradients" + assert p.grad.any(), f"Step {i}: {name} gradient is all zeros" + + optimizer.step() + + # Verify loss decreased + assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + + # Verify weight params remain QuantizedTensors + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, QuantizedTensor + ), f"{name} lost QuantizedTensor type: {type(p).__name__}" + + # Verify optimizer states are float32 + for name, p in model.named_parameters(): + state = optimizer.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "bias" not in name: + assert state["master_param"].dtype == torch.float32 + + +class TestFusedAdamFloat8Block(TestFusedOptimizer): + """FusedAdam with Float8BlockScaling quantized primary weights (single GPU, no FSDP).""" + + def setup_method(self) -> None: + super().setup_method(iters=5) + fp8_block_available, self.fp8_block_reason = te.is_fp8_block_scaling_available( + return_reason=True + ) + self.fp8_block_available = fp8_block_available + + def _build_model(self): + recipe = Float8BlockScaling() + with quantized_model_init(enabled=True, recipe=recipe): + model = te.Linear(256, 256, params_dtype=torch.bfloat16).cuda() + return model, recipe + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_float8block_linear_fused_adam_master_weights(self): + """quantized_model_init(Float8BlockScaling) + te.Linear + FusedAdam(master_weights=True). + + Verifies: + - Model params are QuantizedTensors after init + - FP32 master weights track a reference Adam optimizer + - Params remain QuantizedTensors after training + """ + if not self.fp8_block_available: + pytest.skip(self.fp8_block_reason) + + model, recipe = self._build_model() + + # Verify weight params are QuantizedTensors (bias stays bf16) + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, QuantizedTensor + ), f"Expected QuantizedTensor for {name}, got {type(p).__name__}" + + # Build reference: clone dequantized weights for a plain Adam + ref_params = [p.detach().clone().float() for p in model.parameters()] + + options = {"lr": 5e-4, "betas": (0.9, 0.999), "eps": 1e-8, "weight_decay": 0} + ref_optim = torch.optim.Adam(ref_params, **options) + tst_optim = te.optimizers.FusedAdam( + list(model.parameters()), + master_weights=True, + master_weight_dtype=torch.float32, + use_decoupled_grad=True, + **options, + ) + + for _ in range(self.iters): + for p_ref, p in zip(ref_params, model.parameters()): + p_ref.grad = torch.rand_like(p_ref) + p.decoupled_grad = p_ref.grad.clone() + ref_optim.step() + tst_optim.step() + + # FP32 master weights should match reference Adam exactly + master_params = [ + tst_optim.get_unscaled_state(p, "master_param") for p in model.parameters() + ] + torch.testing.assert_close(ref_params, master_params) + + # Weight params should still be QuantizedTensors after training + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, QuantizedTensor + ), f"{name} lost QuantizedTensor type after training: {type(p).__name__}" + + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + def test_float8block_linear_forward_backward_step(self): + """End-to-end: quantized_model_init + autocast forward + backward + FusedAdam.step(). + + Uses te.autocast with Float8BlockScaling recipe for the forward pass, + verifying the full training loop works with quantized compute. + """ + if not self.fp8_block_available: + pytest.skip(self.fp8_block_reason) + + model, recipe = self._build_model() + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + batch_size, seq_len, hidden = 4, 32, 256 + x = torch.randn(batch_size, seq_len, hidden, dtype=torch.bfloat16, device="cuda") + target = torch.randn_like(x) + + losses = [] + for i in range(self.iters): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = torch.nn.functional.mse_loss(output, target) + losses.append(loss.item()) + loss.backward() + + # Verify all params have non-None gradients after backward + for name, p in model.named_parameters(): + assert p.grad is not None, f"Step {i}: {name} has no gradient after backward" + assert ( + p.grad.shape == p.shape + ), f"Step {i}: {name} grad shape {p.grad.shape} != param shape {p.shape}" + assert torch.isfinite(p.grad).all(), f"Step {i}: {name} has non-finite gradients" + assert p.grad.any(), f"Step {i}: {name} gradient is all zeros" + + optimizer.step() + + # Verify loss decreased + assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" + + # Verify weight params remain QuantizedTensors + for name, p in model.named_parameters(): + if "bias" not in name: + assert isinstance( + p, QuantizedTensor + ), f"{name} lost QuantizedTensor type: {type(p).__name__}" + + # Verify optimizer states are float32 + for name, p in model.named_parameters(): + state = optimizer.state[p] + assert state["exp_avg"].dtype == torch.float32 + assert state["exp_avg_sq"].dtype == torch.float32 + if "bias" not in name: + assert state["master_param"].dtype == torch.float32 + + class TestFusedSGD(TestFusedOptimizer): def setup_method(self) -> None: diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 46b038d922..bcfd2bef19 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -571,6 +571,7 @@ def step(self, closure=None, grad_scaler=None): has_fp16 = False has_bf16 = False + quantized_params_to_update = [] for p in group["params"]: state = self.state[p] @@ -622,6 +623,29 @@ 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. + 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) + g_of_f32_model.append(p_grad.data.float()) + m_of_f32_model.append(unscaled_state["exp_avg"]) + v_of_f32_model.append(unscaled_state["exp_avg_sq"]) + quantized_params_to_update.append((p, unscaled_state["master_param"])) elif p.dtype in [torch.float16, torch.bfloat16]: has_fp16 = has_fp16 or p.dtype == torch.float16 has_bf16 = has_bf16 or p.dtype == torch.bfloat16 @@ -646,6 +670,13 @@ def step(self, closure=None, grad_scaler=None): "FusedAdam does not support FP8 model weights with capturable=True." ) + if self.capturable and len(quantized_params_to_update) > 0: + raise RuntimeError( + "FusedAdam does not support block-scaling quantized weights " + "with capturable=True. The post-step quantize_() writeback " + "cannot be captured in a CUDA graph." + ) + if has_fp16 and has_bf16: if self.store_param_remainders: raise RuntimeError( @@ -782,6 +813,11 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N tensor_lists = [g_of_f32_model, p_f32_model, m_of_f32_model, v_of_f32_model] apply_multi_tensor_adam(self.multi_tensor_adam, tensor_lists) + # Write updated FP32 master weights back to quantized parameters + for qt_param, master_w in quantized_params_to_update: + local_p = qt_param._local_tensor if isinstance(qt_param, DTensor) else qt_param + local_p.quantize_(master_w.data) + # Scaling for name in ["exp_avg", "exp_avg_sq", "master_param"]: if self.fuse_unscale and name in ["exp_avg", "exp_avg_sq"]: diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 0fae40f786..ab496d5a9e 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -429,6 +429,30 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): ) return Float8BlockwiseQTensor.make_like(tensor) + # as_strided op — applied by FSDP2 on the unsharded param. + # When shape and strides match (no-op), return self to preserve the quantized type. + # If shape differs (e.g. padding needed), fall through to dequantize. + if func == aten.as_strided.default: + tensor = args[0] + shape = args[1] + strides = args[2] + if ( + len(shape) == len(strides) == 2 + and tuple(strides) == (shape[-1], 1) + and tuple(shape) == tuple(tensor.size()) + ): + return Float8BlockwiseQTensor.make_like(tensor) + + # slice op — applied by FSDP2 when shards need unpadding. + # When the slice is a no-op (covers entire dimension), return self. + if func == aten.slice.Tensor: + tensor = args[0] + dim = args[1] + start = args[2] + length = args[3] + if start == 0 and length == tensor.size(dim): + return Float8BlockwiseQTensor.make_like(tensor) + # record stream op if func == torch.ops.aten.record_stream.default: qt, stream = args @@ -586,6 +610,146 @@ def is_cuda(self): return self._columnwise_data.is_cuda raise RuntimeError("Float8BlockwiseQTensor has no data!") + def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + """Called by FSDP2 before all-gather of weights for forward and backward passes. + + Args: + mesh: DeviceMesh used by FSDP2 to shard the weights. + orig_size: Original size of the weight tensor. + contiguous_orig_stride: Original stride of the weight tensor. + module: FSDP-wrapped module containing this tensor. + mp_policy: Mixed precision policy used by FSDP2. + + Returns: + sharded_tensors: Tuple of tensors to be all-gathered. + metadata: Metadata needed for reconstructing the tensor after all-gather. + """ + # pylint: disable=unused-argument + from transformer_engine.pytorch.distributed import _get_module_fsdp_state + + if not self._is_2D_scaled: + raise NotImplementedError( + "FSDP2 is only supported for Float8BlockwiseQTensors with 2D block scaling " + "(block_scaling_dim=2). 1D block scaling is not supported because the scale " + "layout has M in dim1, which is incompatible with FSDP2 dim0 all-gather." + ) + + block_len = self._quantizer.block_len # 128 + + # Prepare rowwise tensors — for 2D scaling, M is in dim0 of both data and scale_inv, + # so they naturally align with FSDP2's dim0 all-gather. No unpadding needed. + rowwise_data = self._rowwise_data + rowwise_scale_inv = self._rowwise_scale_inv + + # Prepare columnwise tensors — columnwise data is transposed (K, M) and + # columnwise scale_inv is (ceil(K/128), round_up(ceil(M/128), 4)). + # M is in dim1 for both, so we must transpose to put M in dim0 for all-gather. + columnwise_data = self._columnwise_data + columnwise_scale_inv = self._columnwise_scale_inv + + if columnwise_data is not None: + # Transpose (K, shard_M) -> (shard_M, K) so M is in dim0 + columnwise_data = columnwise_data.t().contiguous() + + if columnwise_scale_inv is not None: + # Original shape: (ceil(K/128), round_up(ceil(shard_M/128), 4)) + # Strip padding from dim1 (the M-block dimension), transpose, then all-gather + shard_M = math.prod(self.shape[:-1]) + m_blocks = (shard_M + block_len - 1) // block_len # ceil(shard_M/128) + columnwise_scale_inv = columnwise_scale_inv[:, :m_blocks] # unpad dim1 + columnwise_scale_inv = columnwise_scale_inv.t().contiguous() # (m_blocks, k_blocks) + + # Always send both rowwise and columnwise data. + # Unlike MXFP8 (where both forms share the same shape), Float8Blockwise has + # differently-shaped rowwise (M, K) and columnwise (K, M) data. The GEMM kernel + # needs both forms available to perform forward and backward operations, so we + # cannot optimize by sending only one usage based on forward/backward pass. + rowwise_usage = True + sharded_tensors = (rowwise_data, rowwise_scale_inv) + columnwise_usage = self._quantizer.columnwise_usage + if columnwise_usage: + sharded_tensors += (columnwise_data, columnwise_scale_inv) + + metadata = (self._fp8_dtype, self._is_2D_scaled, rowwise_usage, columnwise_usage) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata: Any, + param_dtype: torch.dtype, + *, + out: Optional[Float8BlockwiseQTensor] = None, + ): + """Called by FSDP2 after all-gather of weights for forward and backward passes. + + Args: + all_gather_outputs: All-gathered tensors from fsdp_pre_all_gather. + metadata: Metadata from fsdp_pre_all_gather. + param_dtype: High-precision dtype of the tensor. + out: Existing tensor to update in-place (None on first iteration). + + Returns: + Tuple of (Float8BlockwiseQTensor, all_gather_outputs). + """ + fp8_dtype, is_2D_scaled, rowwise_usage, columnwise_usage = metadata + + # Extract rowwise tensors from all-gather outputs + rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] if rowwise_usage else (None, None) + + # Extract columnwise tensors — they were transposed in pre_all_gather, + # so we need to transpose them back. + columnwise_data, columnwise_scale_inv = ( + all_gather_outputs[-2:] if columnwise_usage else (None, None) + ) + + if columnwise_data is not None: + # All-gathered shape is (full_M, K), transpose back to (K, full_M) + columnwise_data = columnwise_data.t().contiguous() + + if columnwise_scale_inv is not None: + # All-gathered shape is (full_m_blocks, k_blocks), + # transpose back to (k_blocks, full_m_blocks) + columnwise_scale_inv = columnwise_scale_inv.t().contiguous() + # Repad dim1 (M-block dimension) to multiple of 4 for GEMM alignment + current_m_blocks = columnwise_scale_inv.shape[1] + pad_amount = (4 - current_m_blocks % 4) % 4 + if pad_amount > 0: + columnwise_scale_inv = torch.nn.functional.pad( + columnwise_scale_inv, (0, pad_amount) + ) + + # Determine the logical shape from the all-gathered data + if rowwise_data is not None: + data_shape = rowwise_data.shape + else: + # columnwise_data is (K, full_M), logical shape is (full_M, K) + data_shape = (columnwise_data.shape[1], columnwise_data.shape[0]) + + if out is not None: + # Update existing tensor in-place (subsequent iterations) + out._rowwise_data = rowwise_data + out._rowwise_scale_inv = rowwise_scale_inv + out._columnwise_data = columnwise_data + out._columnwise_scale_inv = columnwise_scale_inv + else: + # Construct new tensor (first iteration). + # Float8BlockwiseQTensor constructor copies the quantizer, + # so the sharded tensor's quantizer remains independent. + out = Float8BlockwiseQTensor( + shape=data_shape, + dtype=param_dtype, + fp8_dtype=fp8_dtype, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + quantizer=self._quantizer, + is_2D_scaled=is_2D_scaled, + ) + out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) + return out, all_gather_outputs + class _ViewFunc(torch.autograd.Function): """View function From 306e853381f0e6a0e74e035ebc62b58149ab915c Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Fri, 13 Mar 2026 14:25:47 -0600 Subject: [PATCH 279/521] add .claude to gitignore (#2762) Signed-off-by: Peter St. John --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 789d3b0a5f..8a627a7e76 100644 --- a/.gitignore +++ b/.gitignore @@ -42,3 +42,4 @@ compile_commands.json tensor_dumps/ artifacts/ .DS_Store +.claude/ From b7214fd906f36c959c1531f27acc2925d4452da1 Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Sat, 14 Mar 2026 21:43:56 -0600 Subject: [PATCH 280/521] Fix for async dcp checkpointing with Float8Tensors (#2721) * fix for async dcp checkpointing Signed-off-by: Peter St. John * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Peter St. John * Update transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address Greptile review feedback: defensive guards for edge cases - Add _quantizer None guard in new_empty dispatch - Replace self.is_cpu with explicit _data/_transpose checks in __reduce_ex__ - Make get_metadata() safe for cleared tensors (both _data and _transpose None) Signed-off-by: Peter St. John --------- Signed-off-by: Peter St. John Signed-off-by: Peter St. John Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani --- .../distributed/run_fsdp2_fused_adam.py | 21 ++++++------------ tests/pytorch/distributed/test_torch_fsdp2.py | 10 --------- .../pytorch/quantized_tensor.py | 21 ++++++++++++++++++ .../pytorch/tensor/float8_tensor.py | 22 ++++++++++++++++++- .../tensor/storage/float8_tensor_storage.py | 14 +++++++++++- 5 files changed, 62 insertions(+), 26 deletions(-) diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/run_fsdp2_fused_adam.py index f97f2bb22e..c39957cf13 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/run_fsdp2_fused_adam.py @@ -570,17 +570,13 @@ def test_dcp_output_parity(recipe=None, async_save=False): else: model_state = model.state_dict() + save_state = {"model": model_state, "optimizer": optimizer.state_dict()} + if not async_save: - dcp.save( - {"model": model_state, "optimizer": optimizer.state_dict()}, - checkpoint_id=checkpoint_dir, - ) - future = None + dcp.save(save_state, checkpoint_id=checkpoint_dir) else: - future = dcp.async_save( - {"model": model_state, "optimizer": optimizer.state_dict()}, - checkpoint_id=checkpoint_dir, - ) + future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) + future.result() # Block on async save completion # ── Build a fresh model and load the checkpoint ────────────────── model2 = _build_model(fp8_init=True, recipe=recipe) @@ -609,9 +605,6 @@ def test_dcp_output_parity(recipe=None, async_save=False): state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} - if async_save: - future.result() # Block on async save completion - dcp.load(state_to_load, checkpoint_id=checkpoint_dir) model2.load_state_dict( state_to_load["model"], @@ -636,7 +629,7 @@ def test_dcp_output_parity(recipe=None, async_save=False): ref_output, rtol=0.05, atol=0.1, - msg="Fresh model loaded from DCP checkpoint produces different output", + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", ) else: torch.testing.assert_close( @@ -644,7 +637,7 @@ def test_dcp_output_parity(recipe=None, async_save=False): ref_output, rtol=0, atol=0, - msg="Fresh model loaded from DCP checkpoint produces different output", + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", ) # ── Verify one more training step produces identical results ───── diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 042028a949..02e45d99cb 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -196,16 +196,6 @@ def test_fsdp2_dcp_output_parity(fp_recipe): @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") def test_fsdp2_dcp_output_parity_async(fp_recipe): """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe in ("DelayedScaling", "Float8CurrentScaling"): - pytest.xfail( - f"async DCP save/load with {fp_recipe} uses StateDictStager._offload_tensor() which " - "tries to deep-copy the tensor's underlying storage. Float8Tensor is a wrapper subclass" - "(_make_wrapper_subclass) with data_ptr() == 0 (empty storage). The staging code at " - "line 215 skips the storage copy for wrapper subclasses, creating a plain tensor with " - "uninitialized garbage data. The actual FP8 data (in _data, _scale_inv attributes) is " - "deep-copied but ignored by DCP when writing." - ) - if fp_recipe == "MXFP8BlockScaling": pytest.xfail( "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 2129fd486f..807671e863 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -563,6 +563,27 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): if func == torch.ops.aten.view.default: raise NotImplementedError("{cls.__name__} class does not support tensor views") + # New empty op (used by DCP async staging to create CPU copies) + if func == torch.ops.aten.new_empty.default: + tensor = args[0] + size = args[1] + dtype = kwargs.get("dtype", tensor.dtype) + device = kwargs.get("device", tensor.device) + pin_memory = kwargs.get("pin_memory", False) + if tensor._quantizer is None: + raise RuntimeError( + f"{type(tensor).__name__} does not have a quantizer; " + "cannot create new_empty QuantizedTensor" + ) + out = tensor._quantizer.make_empty( + shape=torch.Size(size), + dtype=dtype, + device=device, + requires_grad=tensor.requires_grad, + pin_memory=pin_memory, + ) + return out + # Empty like op if func == torch.ops.aten.empty_like.default: tensor = args[0] diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 9cc00855cd..5f00bc8017 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -152,6 +152,7 @@ def make_empty( requires_grad=requires_grad, data_transpose=data_transpose, quantizer=self, + device=device, ) def calibrate(self, tensor: torch.Tensor) -> None: @@ -379,6 +380,7 @@ def make_empty( requires_grad=requires_grad, data_transpose=data_transpose, quantizer=self, + device=device, ) def calibrate(self, tensor: torch.Tensor) -> None: @@ -953,6 +955,15 @@ def is_cuda(self): return self._transpose.is_cuda raise RuntimeError("Both data and transpose are None") + @property + def is_cpu(self): + """Return whether the tensor is on CPU.""" + if self._data is not None: + return self._data.is_cpu + if self._transpose is not None: + return self._transpose.is_cpu + raise RuntimeError("Both data and transpose are None") + @classmethod def _make_in_reduce_ex( cls, @@ -977,7 +988,16 @@ def _make_in_reduce_ex( ) def __reduce_ex__(self, protocol: int) -> tuple: - """Custom pickling to remove references to FP8 metadata objects""" + """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), diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index 0fb7966c2f..de7f8f58e2 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -14,7 +14,7 @@ from ...quantized_tensor import QuantizedTensorStorage, Quantizer -from ...constants import TE_DType as torch_to_transformer_engine_dtype +from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch from ...utils import is_non_tn_fp8_gemm_supported, _empty_tensor @@ -35,6 +35,13 @@ def forward( if tensor._data is not None: if tensor._data.numel() == 0: return torch.empty_like(tensor._data, dtype=dtype) + if tensor._data.is_cpu: + # CPU fallback: reinterpret uint8 as FP8, cast to target dtype, scale + fp8_torch_dtype = TE_DType_To_Torch[tensor._fp8_dtype] + return ( + tensor._data.view(fp8_torch_dtype).float() + * tensor._scale_inv.to(tensor._data.device) + ).to(dtype) # Cast from FP8 return tex.dequantize(tensor, te_dtype) @@ -132,6 +139,11 @@ 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, } From 708d7c160ad6b2bf44c9c597083d4cbb4860f068 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Sun, 15 Mar 2026 20:53:29 -0700 Subject: [PATCH 281/521] Pytorch binding for cublas grouped gemm + Grouped Bias Support + Grouped Tensor Swizzling (#2669) Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove changes not needed for bf16 Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * keep only pytorch binding for now Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * linting error Signed-off-by: Varun Thumbe * add fast accumulator support Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * MXFP8 grouped GEMM + tensor-scaled FP8 fixes Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change version to 13.3 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * fix the test Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Random padding condition shouldnt be done for mxfp8 Signed-off-by: vthumbe1503 * Remove incorrect comment Signed-off-by: vthumbe1503 * CUBLAS > 13.2 is enough Signed-off-by: vthumbe1503 * CUBLAS version needed for MXFP8 indeed seems to be 13.3 Signed-off-by: vthumbe1503 * all changes for grouped gemm Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Accidental line removal added back. Plus need changes ci t trigger Add documentation for scaling factors in common.h Signed-off-by: vthumbe1503 * Update cuBLAS version requirement for MXFP8 support Signed-off-by: vthumbe1503 * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * grouped gemm: address code review comments - Replace nvte_set/get_grouped_tensor_swizzled_scales with nvte_set_grouped_tensor_param - Add host-side validation: A and B must use same scaling mode (both MXFP8 or both tensor scaling) - Add host-side validation: A and B must both be FP8 or both non-FP8; restrict inputs to FP8/BF16 - Restrict output (C/D) to BF16/FP32; remove FP16 from supported types - Refactor workspace allocation: replace manual offset arithmetic with moving pointer pattern - Use void* + NVTEScalingMode in setup kernel instead of separate float*/char* scale params - Extract use_columnwise(swap_dims) helper to eliminate duplicated MXFP8 columnwise blocks - Split set_fp8_scale_pointers into set_fp8_scale_pointers / set_mxfp8_scale_pointers - Remove scale_inv_ptrs from GroupedOperandSelection; pass workspace pointers directly - Move swizzled-scales validation into validate_grouped_gemm_inputs for fail-fast behavior - Add use_split_accumulator to GroupedMatmulConfig (Hopper only, default false) - Add FP8 test case with per-tensor scales; add BF16/MXFP8 shape-varying test cases Signed-off-by: Pawel Gadzinski * address reviee comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * missed merged conflict handling Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor change Signed-off-by: Varun Thumbe * forgot adding a or Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * resolve merge conflicts Signed-off-by: Varun Thumbe * 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 * address minor review comments Signed-off-by: Varun Thumbe * remove unecessary code Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * one line that broke everything :( 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 * 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 * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * 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 * unecessary Signed-off-by: Varun Thumbe * revert caching changes Signed-off-by: Varun Thumbe * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * fix minor bug from greptile Signed-off-by: Varun Thumbe * revert for now 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 * address review comments Signed-off-by: Varun Thumbe * address review comments Signed-off-by: Varun Thumbe * 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 * address review comments Signed-off-by: Varun Thumbe * address review commentsgp Signed-off-by: Varun Thumbe --------- Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Signed-off-by: Jeremy Berchtold Signed-off-by: Pawel Gadzinski Co-authored-by: Kirthi Shankar Sivamani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jeremy Berchtold Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Pawel Gadzinski --- tests/cpp/test_common.cu | 6 +- tests/pytorch/test_numerics.py | 340 ++++++- transformer_engine/common/gemm/config.cpp | 18 + .../common/gemm/cublaslt_grouped_gemm.cu | 961 ++++++++++++++---- .../common/include/transformer_engine/gemm.h | 45 + .../include/transformer_engine/swizzle.h | 15 + transformer_engine/common/swizzle/swizzle.cu | 216 +++- .../pytorch/cpp_extensions/gemm.py | 112 ++ transformer_engine/pytorch/csrc/extensions.h | 19 + .../pytorch/csrc/extensions/gemm.cpp | 216 ++++ .../pytorch/csrc/extensions/pybind.cpp | 9 + .../pytorch/csrc/extensions/swizzle.cpp | 84 ++ transformer_engine/pytorch/csrc/quantizer.cpp | 6 + .../pytorch/csrc/type_converters.cpp | 10 +- transformer_engine/pytorch/csrc/util.h | 10 + .../pytorch/tensor/grouped_tensor.py | 2 + .../tensor/storage/grouped_tensor_storage.py | 7 + 17 files changed, 1863 insertions(+), 213 deletions(-) diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index ff26d1b6c5..5180a81612 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -1157,6 +1157,8 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, NVTEGroupedTensor h = grouped.handle.get(); + size_t total_elems_size = static_cast(total_elems); + NVTEShape flat_shape = nvte_make_shape(&total_elems_size, 1); // Copy rowwise data if available if (has_rowwise) { grouped.data = cuda_alloc(total_bytes); @@ -1167,7 +1169,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, grouped.tensor_bytes[i], cudaMemcpyDeviceToDevice)); } - NVTEBasicTensor data_tensor{grouped.data.get(), static_cast(dtype), grouped.logical_shape}; + NVTEBasicTensor data_tensor{grouped.data.get(), static_cast(dtype), flat_shape}; nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseData, &data_tensor, sizeof(data_tensor)); } @@ -1183,7 +1185,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, } NVTEBasicTensor col_tensor{grouped.columnwise_data.get(), static_cast(dtype), - grouped.logical_shape}; + flat_shape}; nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseData, &col_tensor, sizeof(col_tensor)); } diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 8e3b0517ee..19b94d3531 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -46,7 +46,12 @@ is_nvfp4_available, ) from transformer_engine.pytorch import checkpoint as te_checkpoint -from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm +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.common import recipe import transformer_engine_torch as tex from utils import ModelConfig, reset_rng_states @@ -2792,6 +2797,339 @@ def test_grouped_gemm(shape, dtype, layout, accumulate, 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, + ) + + +@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]) +def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate) -> None: + if tex.get_cublasLt_version() < 130200: + pytest.skip("Grouped GEMM requires cuBLAS 13.2+.") + 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.") + + 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 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 + ) + # 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, + ) + 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 = ( + [(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 + ) + 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) + + +def _make_grouped_tensor_quantized_mxfp8( + tensors: List[torch.Tensor], + *, + is_a: bool, + transposed: bool, + device: torch.device, + optimize_for_gemm: bool = True, +) -> GroupedTensor: + 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 + grouped_input = torch.cat(tensors, dim=0) + 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) + + +@pytest.mark.parametrize( + "shape", + [ + (1, 128, 128, 512), + (8, 1024, 128, 512), + (16, 4096, 128, 512), + ], +) +@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() < 130200: + pytest.skip("Grouped GEMM requires cuBLAS 13.2+.") + 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" + grouped_A = _make_grouped_tensor_quantized_mxfp8(A, is_a=True, transposed=transa, device="cuda") + grouped_B = _make_grouped_tensor_quantized_mxfp8( + B, is_a=False, transposed=transb, device="cuda" + ) + A_fp8 = grouped_A.split_into_quantized_tensors() + B_fp8 = grouped_B.split_into_quantized_tensors() + + 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( diff --git a/transformer_engine/common/gemm/config.cpp b/transformer_engine/common/gemm/config.cpp index 286fc0cc96..de533909f6 100644 --- a/transformer_engine/common/gemm/config.cpp +++ b/transformer_engine/common/gemm/config.cpp @@ -153,6 +153,12 @@ void nvte_get_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, static_cast(attr), " needs ", attr_size, " bytes, but buffer has ", size_in_bytes, " bytes)"); + // bool size is implementation-dependent, so we explicitly specify + // uint8_t in the user-facing API. + auto bool_to_uint8 = [](bool in, void *out) { + *reinterpret_cast(out) = static_cast(in); + }; + // Write to buffer NVTE_CHECK(config != nullptr, "Invalid NVTEGroupedMatmulConfig (got NULL)"); const auto &config_ = *reinterpret_cast(config); @@ -172,6 +178,9 @@ void nvte_get_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, std::memcpy(buf, &val, attr_size); break; } + case kNVTEGroupedMatmulConfigUseSplitAccumulator: + bool_to_uint8(config_.use_split_accumulator, buf); + break; case kNVTEGroupedMatmulConfigSMCount: std::memcpy(buf, &config_.sm_count, attr_size); break; @@ -194,6 +203,12 @@ void nvte_set_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, " bytes)"); NVTE_CHECK(buf != nullptr, "Invalid buffer (got NULL)"); + // bool size is implementation-dependent, so we explicitly specify + // uint8_t in the user-facing API. + auto uint8_to_bool = [](const void *in, bool &out) { + out = static_cast(*reinterpret_cast(in)); + }; + // Read from buffer NVTE_CHECK(config != nullptr, "Invalid NVTEGroupedMatmulConfig (got NULL)"); auto &config_ = *reinterpret_cast(config); @@ -216,6 +231,9 @@ void nvte_set_grouped_matmul_config_attribute(NVTEGroupedMatmulConfig config, config_.avg_k = val; break; } + case kNVTEGroupedMatmulConfigUseSplitAccumulator: + uint8_to_bool(buf, config_.use_split_accumulator); + break; case kNVTEGroupedMatmulConfigSMCount: std::memcpy(&config_.sm_count, buf, attr_size); break; diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 7069debc56..ccf1e53ba4 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1,8 +1,8 @@ /************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ +* Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +* +* See LICENSE for license information. +************************************************************************/ #include #include @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -17,6 +18,7 @@ #include "../util/cuda_runtime.h" #include "../util/handle_manager.h" #include "../util/logging.h" +#include "../util/vectorized_pointwise.h" #include "./config.h" namespace { @@ -105,6 +107,10 @@ inline int64_t compute_avg_last_dim(const transformer_engine::GroupedTensor *t) return static_cast(t->logical_shape.data[1]) / static_cast(t->num_tensors); } +// Constants for grouped GEMM workspace (declared early for use in helpers) +static constexpr size_t kGroupedGemmAlignment = 256; +static constexpr size_t kGroupedGemmCublasWorkspaceSize = 32ull * 1024 * 1024; // 32 MiB + // Workspace layout for grouped GEMM struct GroupedGemmSetupWorkspace { void **A_ptrs; @@ -197,33 +203,15 @@ struct GroupedGemmSetupWorkspace { } }; -// ----------------------------------------------------------------------------- -// Helper routines to keep nvte_grouped_gemm readable -// ----------------------------------------------------------------------------- -struct GroupedGemmInputProperties { - bool is_fp8; - bool is_mxfp8; -}; - -inline GroupedGemmInputProperties validate_grouped_gemm_inputs( - const transformer_engine::GroupedTensor *inputA, - const transformer_engine::GroupedTensor *inputB, - const transformer_engine::GroupedTensor *inputC, - const transformer_engine::GroupedTensor *outputD, +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 size_t num_tensors = inputA->num_tensors; NVTE_CHECK(num_tensors >= 1, "Grouped GEMM: number of tensors must be at least 1"); - NVTE_CHECK(inputB->num_tensors == num_tensors, - "Grouped GEMM: A and B must have the same number of tensors"); - // C can be NULL (will use D as C when beta=0) - if (inputC != nullptr) { - NVTE_CHECK(inputC->num_tensors == num_tensors, - "Grouped GEMM: A and C must have the same number of tensors"); + for (const auto *tensor : inputs) { + NVTE_CHECK(tensor->num_tensors == num_tensors, + "Grouped GEMM: inputs must have the same number of tensors"); } - NVTE_CHECK(outputD->num_tensors == num_tensors, - "Grouped GEMM: A and D must have the same number of tensors"); - // Validate alpha/beta have per-matrix values 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, @@ -234,38 +222,74 @@ inline GroupedGemmInputProperties validate_grouped_gemm_inputs( 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::kBFloat16 || + dtype == transformer_engine::DType::kFloat16; }; + bool dtype_ok = true; + for (const auto *tensor : inputs) { + dtype_ok = dtype_ok && is_supported_input_dtype(tensor->dtype()); + } + NVTE_CHECK(dtype_ok, "Grouped GEMM inputs must be FP8, BF16, or FP16."); + for (const auto *tensor : inputs) { + NVTE_CHECK(tensor->has_data() || tensor->has_columnwise_data(), + "Grouped GEMM: input tensor is missing both row-wise and column-wise data"); + } + + // Cross-operand consistency across all inputs. + const auto *ref = *inputs.begin(); + const bool ref_is_fp8 = is_fp8_dtype(ref->dtype()); + const bool ref_is_mxfp8 = transformer_engine::is_mxfp_scaling(ref->scaling_mode); + for (const auto *tensor : inputs) { + 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(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) { + NVTE_CHECK(tensor->with_gemm_swizzled_scales, + "MXFP8 grouped GEMM: scales must be swizzled for GEMM."); + } + } + return num_tensors; +} + +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; }; - NVTE_CHECK(is_supported_input_dtype(inputA->dtype()) && is_supported_input_dtype(inputB->dtype()), - "Grouped GEMM inputs must be FP8 or BF16."); - NVTE_CHECK(is_fp8_dtype(inputA->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(inputA->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, " - "mixed configurations are not supported."); - const bool is_fp8 = is_fp8_dtype(inputA->dtype()); - const bool is_mxfp8 = transformer_engine::is_mxfp_scaling(inputA->scaling_mode); - if (is_mxfp8) { - NVTE_CHECK(inputA->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"); + 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."); } - // Only check C dtype if C is provided - if (inputC != nullptr) { - NVTE_CHECK(is_output_dtype(inputC->dtype()), "Grouped GEMM: C must be BF16 or FP32."); +} + +inline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) { + return GroupedGemmSetupWorkspace::required_setup_size(num_tensors, kGroupedGemmAlignment); +} + +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() >= 130200, api_name, + " requires cuBLAS 13.2+, but run-time cuBLAS version is ", + transformer_engine::cuda::cublas_version()); +} + +inline transformer_engine::GroupedMatmulConfig parse_grouped_gemm_config( + NVTEGroupedMatmulConfig config) { + transformer_engine::GroupedMatmulConfig config_; + if (config != nullptr) { + config_ = *reinterpret_cast(config); } - NVTE_CHECK(is_output_dtype(outputD->dtype()), "Grouped GEMM: D must be BF16 or FP32."); - NVTE_CHECK(inputA->has_data() || inputA->has_columnwise_data(), - "Grouped GEMM: A tensor is missing both row-wise and column-wise data"); - NVTE_CHECK(inputB->has_data() || inputB->has_columnwise_data(), - "Grouped GEMM: B tensor is missing both row-wise and column-wise data"); - return {is_fp8, is_mxfp8}; + return config_; } // Select row-wise vs column-wise storage and adjust transpose flag for grouped GEMM. @@ -282,6 +306,219 @@ struct GroupedOperandSelection { bool trans = false; }; +constexpr int kMaxTensorsPerKernel = 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]; +}; + +// 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]; +}; +struct MultiTensorListInfo { + bool all_row = true; + bool all_col = true; + transformer_engine::DType row_dtype = transformer_engine::DType::kNumTypes; + transformer_engine::DType col_dtype = transformer_engine::DType::kNumTypes; + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; + bool with_gemm_swizzled_scales = false; +}; + +struct OperandStorageChoice { + bool use_rowwise = true; + bool swap_dims = true; + 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 has_row, bool has_col, + const char *name) { + NVTE_CHECK(has_row || has_col, "Grouped GEMM: ", name, + " is missing both row-wise and column-wise data"); + if (is_mxfp8) { + if (is_A) { + if (trans) { + NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 transposed ", name, " is missing row-wise data"); + return {true, true, trans}; + } + NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 non-transposed ", name, + " is missing column-wise data"); + return {false, false, trans}; + } + if (trans) { + NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 transposed ", name, " is missing column-wise data"); + return {false, false, trans}; + } + NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 non-transposed ", name, " is missing row-wise data"); + return {true, true, trans}; + } + + // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed. + 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"); + 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"); + return {false, true, false}; + } + } + + // If only column-wise data is available, mirror the transpose flag (pre-transposed storage). + if (!has_row && has_col) { + NVTE_CHECK(!is_fp8 || non_tn_fp8_ok, + "Grouped GEMM: FP8 on Hopper requires row-wise data for this transpose config."); + return {false, true, !trans}; + } + + NVTE_CHECK(has_row, "Grouped GEMM: ", name, " is missing row-wise data"); + return {true, true, trans}; +} + +// Build Kernel Arguments detailing out addresses and other metadata for list of C/D tensors +// passed to the grouped GEMM kernel. Use-case: C/D --> List of wgrads for experts in MOE +inline MultiTensorGroupGemmOutputArgs build_grouped_gemm_multi_out_args( + const NVTETensor *tensor_list, size_t list_size, size_t expected_num_tensors, + transformer_engine::DType expected_dtype, const char *name) { + MultiTensorGroupGemmOutputArgs args{}; + if (list_size == 0) { + NVTE_CHECK(tensor_list == nullptr, "Grouped GEMM: ", name, "_list provided with num_", name, + "_tensors=0"); + return args; + } + NVTE_CHECK(tensor_list != nullptr, "Grouped GEMM: ", name, "_list is null but num_", name, + "_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); + + for (size_t i = 0; i < list_size; ++i) { + const transformer_engine::Tensor *t = + transformer_engine::convertNVTETensorCheck(tensor_list[i]); + NVTE_CHECK(t->has_data(), "Grouped GEMM: ", name, "_list tensor ", i, " has no data"); + NVTE_CHECK(t->dtype() == expected_dtype, "Grouped GEMM: ", name, "_list tensor ", i, + " dtype mismatch. Expected ", transformer_engine::to_string(expected_dtype), " got ", + transformer_engine::to_string(t->dtype())); + const auto &shape = t->shape(); + NVTE_CHECK(shape.size() == 2, "Grouped GEMM: ", name, "_list tensor ", i, " must be 2D."); + args.data_ptrs[i] = t->data.dptr; + args.rows[i] = static_cast(shape[1]); + args.cols[i] = static_cast(shape[0]); + } + return 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, + int64_t *avg_first_dim, int64_t *avg_last_dim, const char *name) { + using namespace transformer_engine; + MultiTensorGroupGemmInputArgs args{}; + *avg_first_dim = 0; + *avg_last_dim = 0; + if (list_size == 0) { + return args; + } + for (size_t i = 0; i < list_size; ++i) { + const transformer_engine::Tensor *t = + transformer_engine::convertNVTETensorCheck(tensor_list[i]); + const transformer_engine::SimpleTensor &data = use_rowwise ? t->data : t->columnwise_data; + const transformer_engine::SimpleTensor &scale_inv = + 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]); + + if (is_fp8) { + NVTE_CHECK(scale_inv.has_data(), "Grouped GEMM: ", name, "_list tensor ", i, + " requires scale_inv for FP8."); + args.scale_inv_ptrs[i] = scale_inv.dptr; + } else { + args.scale_inv_ptrs[i] = nullptr; + } + } + *avg_first_dim /= static_cast(list_size); + *avg_last_dim /= static_cast(list_size); + return args; +} + +inline MultiTensorListInfo validate_grouped_gemm_multi_inputA_list(const NVTETensor *tensor_list, + size_t list_size, + size_t expected_num_tensors, + const char *name) { + using namespace transformer_engine; + MultiTensorListInfo info{}; + if (list_size == 0) { + NVTE_CHECK(tensor_list == nullptr, "Grouped GEMM: ", name, "_list provided with num_", name, + "_tensors=0"); + return info; + } + NVTE_CHECK(tensor_list != nullptr, "Grouped GEMM: ", name, "_list is null but num_", name, + "_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); + + const transformer_engine::Tensor *t0 = transformer_engine::convertNVTETensorCheck(tensor_list[0]); + 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."); + + for (size_t i = 0; i < list_size; ++i) { + const transformer_engine::Tensor *t = + transformer_engine::convertNVTETensorCheck(tensor_list[i]); + NVTE_CHECK(t->scaling_mode == info.scaling_mode, "Grouped GEMM: ", name, + "_list tensors must share the same scaling mode."); + NVTE_CHECK(t->with_gemm_swizzled_scales == info.with_gemm_swizzled_scales, + "Grouped GEMM: ", name, "_list tensors must share GEMM swizzled scale state."); + + if (t->has_data()) { + if (info.row_dtype == DType::kNumTypes) { + info.row_dtype = t->data.dtype; + } + // Check all tensors have the same dtype + NVTE_CHECK(t->data.dtype == info.row_dtype, "Grouped GEMM: ", name, + "_list rowwise dtypes must match."); + } else { + // All tensors must have either data or columnwise data + info.all_row = false; + } + + if (t->has_columnwise_data()) { + if (info.col_dtype == DType::kNumTypes) { + info.col_dtype = t->columnwise_data.dtype; + } + NVTE_CHECK(t->columnwise_data.dtype == info.col_dtype, "Grouped GEMM: ", name, + "_list columnwise dtypes must match."); + } else { + // All tensors must have either data or columnwise data + info.all_col = false; + } + } + + 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). @@ -326,7 +563,7 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: const DType row_dtype = t->data.dtype; const DType col_dtype = t->columnwise_data.dtype; - GroupedOperandSelection sel; + GroupedOperandSelection sel{}; sel.trans = trans; sel.scaling_mode = sm; sel.with_gemm_swizzled_scales = t->with_gemm_swizzled_scales; @@ -354,60 +591,14 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: sel.shape = create_shape_info(t, /*swap_dims=*/false); }; - // MXFP8: Row-wise and column-wise data are scaled along different dimensions. - if (mxfp8) { - if (is_A) { - if (trans) { - NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 transposed A is missing row-wise data"); - use_rowwise(); - } else { - NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 non-transposed A is missing column-wise data"); - use_columnwise(/*swap_dims=*/false); - } - } else { // B - if (trans) { - NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 transposed B is missing column-wise data"); - use_columnwise(/*swap_dims=*/false); - } else { - NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 non-transposed B is missing row-wise data"); - use_rowwise(); - } - } - return sel; - } - - // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed. - if (is_fp8 && !non_tn_fp8_ok) { - if (is_A) { - if (!sel.trans) { - NVTE_CHECK(has_col, "Grouped GEMM: A is missing column-wise data needed for FP8 TN layout"); - use_columnwise(); - sel.trans = true; // using pre-transposed storage - return sel; - } - } else { // B - if (sel.trans) { - NVTE_CHECK(has_col, "Grouped GEMM: B is missing column-wise data needed for FP8 TN layout"); - use_columnwise(); - sel.trans = false; // using pre-transposed storage - return sel; - } - } - } - - // If only column-wise data is available, mirror the transpose flag (pre-transposed storage). - if (!has_row && has_col) { - // On Hopper FP8, this would break TN requirement - should have been handled above - NVTE_CHECK( - !is_fp8 || non_tn_fp8_ok, - "Grouped GEMM: FP8 on Hopper requires row-wise data for this transpose configuration"); - use_columnwise(); - sel.trans = !trans; // flip transpose for pre-transposed storage - return sel; + 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"); + sel.trans = choice.trans; + if (choice.use_rowwise) { + use_rowwise(); + } else { + use_columnwise(choice.swap_dims); } - - // Default: use row-wise data - use_rowwise(); return sel; } @@ -420,17 +611,14 @@ inline void *validate_and_get_workspace_ptr(transformer_engine::Tensor *ws, size return ws->data.dptr; } -inline void init_matrix_layouts(cublasLtMatrixLayoutOpaque_t &descA, - cublasLtMatrixLayoutOpaque_t &descB, - cublasLtMatrixLayoutOpaque_t &descC, - cublasLtMatrixLayoutOpaque_t &descD, - const GroupedGemmSetupWorkspace &ws, - const GroupedOperandSelection &A_sel, - const GroupedOperandSelection &B_sel, - const transformer_engine::GroupedTensor *D, size_t num_tensors) { +inline void init_matrix_layouts( + cublasLtMatrixLayoutOpaque_t &descA, cublasLtMatrixLayoutOpaque_t &descB, + cublasLtMatrixLayoutOpaque_t &descC, cublasLtMatrixLayoutOpaque_t &descD, + const GroupedGemmSetupWorkspace &ws, const GroupedOperandSelection &A_sel, + const GroupedOperandSelection &B_sel, transformer_engine::DType d_dtype, size_t num_tensors) { const cudaDataType_t A_type = get_cuda_dtype(A_sel.dtype); const cudaDataType_t B_type = get_cuda_dtype(B_sel.dtype); - const cudaDataType_t D_type = get_cuda_dtype(D->dtype()); + const cudaDataType_t D_type = get_cuda_dtype(d_dtype); // Storage dimensions computed by kernel, leading dimension = rows NVTE_CHECK_CUBLAS(cublasLtGroupedMatrixLayoutInit(&descA, A_type, num_tensors, ws.a_rows, @@ -511,11 +699,6 @@ inline void set_fp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, void CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs))); } - -// Constants for grouped GEMM workspace (declared early for use in heuristics) -static constexpr size_t kGroupedGemmAlignment = 256; -static constexpr size_t kGroupedGemmCublasWorkspaceSize = 32ull * 1024 * 1024; // 32 MiB - inline cublasLtMatmulAlgo_t select_grouped_gemm_algo(cublasLtHandle_t handle, cublasLtMatmulDescOpaque_t &matmulDesc, cublasLtMatrixLayoutOpaque_t &descA, @@ -546,6 +729,64 @@ inline cublasLtMatmulAlgo_t select_grouped_gemm_algo(cublasLtHandle_t handle, return heuristicResult.algo; } +struct GroupedGemmWorkspace { + GroupedGemmSetupWorkspace setup_workspace; + void *cublas_workspace_ptr = nullptr; + size_t num_tensors = 0; +}; + +inline GroupedGemmWorkspace setup_grouped_gemm_workspace(transformer_engine::Tensor *wspace_setup, + transformer_engine::Tensor *wspace_cublas, + size_t num_tensors) { + const size_t setup_workspace_size = grouped_gemm_setup_workspace_size(num_tensors); + const size_t cublas_workspace_size = kGroupedGemmCublasWorkspaceSize; + void *setup_workspace_ptr = validate_and_get_workspace_ptr(wspace_setup, setup_workspace_size, + "Grouped GEMM setup workspace"); + void *cublas_workspace_ptr = validate_and_get_workspace_ptr(wspace_cublas, cublas_workspace_size, + "Grouped GEMM cuBLAS workspace"); + auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers( + static_cast(setup_workspace_ptr), num_tensors); + return {std::move(setup_workspace), cublas_workspace_ptr, num_tensors}; +} + +inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspace, + 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) { + using cublasHandleManager = + transformer_engine::detail::HandleManager; + cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); + + cublasOperation_t op_A = A_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N; + cublasOperation_t op_B = B_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N; + + cublasLtMatrixLayoutOpaque_t descA, descB, descC, descD; + init_matrix_layouts(descA, descB, descC, descD, setup_workspace, A_sel, B_sel, d_dtype, + num_tensors); + + cublasLtMatmulDescOpaque_t matmulDesc; + init_matmul_desc(matmulDesc, op_A, op_B, use_fp8, use_split_accumulator); + 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) { + set_fp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, + setup_workspace.b_scale_inv_ptrs); + } + + 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)); +} + // Device helper: compute the element offset for tensor `idx` given shape metadata. // Three cases: // 1. Explicit per-tensor offset array provided → use it directly. @@ -569,6 +810,46 @@ __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; + + 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; + if (m == 0 || n == 0) return; + + 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); + + auto *d_ptr = reinterpret_cast(d_base + d_offset * sizeof(T)); + const auto *bias_ptr = reinterpret_cast(bias_base + bias_offset * sizeof(T)); + + 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]; + } + storer.store(vec_id, elements); +} + // Single kernel that sets up all GEMM parameters. // Rationale: cuBLASLt grouped matmul API needs flat arrays of pointers and per-matrix dimensions, // but NVTEGroupedTensor stores a single contiguous buffer + optional per-tensor offsets/shapes. @@ -582,30 +863,47 @@ __global__ void setup_grouped_gemm_kernel( 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, - // Scale inputs: contiguous scale buffers and the shared scaling recipe for A and B - void *a_scale_base, void *b_scale_base, NVTEScalingMode scaling_mode, size_t num_tensors) { + // 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, + MultiTensorGroupGemmInputArgs a_multi_tensor_args, + MultiTensorGroupGemmOutputArgs c_multi_tensor_args, + MultiTensorGroupGemmOutputArgs d_multi_tensor_args) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_tensors) return; // Get dimensions for this tensor (from array or uniform value) - int64_t a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first; - int64_t a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last; + const bool has_a_multi_tensor = (a_base == nullptr); + const bool has_c_multi_tensor = (c_base == nullptr); + 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 { + 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; + } int64_t b_first = B_meta.first_dims ? B_meta.first_dims[idx] : B_meta.uniform_first; int64_t b_last = B_meta.last_dims ? B_meta.last_dims[idx] : B_meta.uniform_last; int64_t d_first = D_meta.first_dims ? D_meta.first_dims[idx] : D_meta.uniform_first; int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last; // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform) - int64_t a_offset = compute_grouped_tensor_offset(A_meta, idx); + int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx); int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx); int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx); int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx); // Compute data pointers - A_ptrs[idx] = a_base + a_offset * a_elem_size; + 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; - C_ptrs[idx] = c_base + c_offset * c_elem_size; - D_ptrs[idx] = d_base + d_offset * d_elem_size; + 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, @@ -614,8 +912,13 @@ __global__ void setup_grouped_gemm_kernel( a_cols[idx] = static_cast(a_first); b_rows[idx] = static_cast(b_last); b_cols[idx] = static_cast(b_first); - d_rows[idx] = static_cast(d_last); - d_cols[idx] = static_cast(d_first); + if (has_d_multi_tensor) { + d_rows[idx] = d_multi_tensor_args.rows[idx]; + d_cols[idx] = d_multi_tensor_args.cols[idx]; + } else { + d_rows[idx] = static_cast(d_last); + d_cols[idx] = static_cast(d_first); + } // Fill alpha/beta pointers (per-matrix) alpha_ptrs[idx] = alpha_ptr + idx; @@ -627,14 +930,18 @@ __global__ void setup_grouped_gemm_kernel( // otherwise : one float per tensor, indexed by tensor index if (a_scale_base) { if (scaling_mode == NVTE_MXFP8_1D_SCALING) { - a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + a_offset / 32; + a_scale_inv_ptrs[idx] = reinterpret_cast( + static_cast(static_cast(a_scale_base)) + a_offset / 32); } else { a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + idx; } + } 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) { - b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + b_offset / 32; + b_scale_inv_ptrs[idx] = reinterpret_cast( + static_cast(static_cast(b_scale_base)) + b_offset / 32); } else { b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + idx; } @@ -646,20 +953,51 @@ 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, 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; - TensorShapeInfo C_meta = TensorShapeInfo::create_shape_info_for_C(C, D); - TensorShapeInfo D_meta = TensorShapeInfo::from_tensor(D); + TensorShapeInfo C_meta{}; + TensorShapeInfo D_meta{}; + + const bool has_d_multi_tensor = (D_list != nullptr); + const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor; + MultiTensorGroupGemmOutputArgs c_multi_tensor_args{}; + MultiTensorGroupGemmOutputArgs d_multi_tensor_args{}; + if (has_d_multi_tensor) { + d_multi_tensor_args = + build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, "D"); + } + if (C_list != nullptr) { + c_multi_tensor_args = + build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, "C"); + } else if (has_d_multi_tensor) { + c_multi_tensor_args = d_multi_tensor_args; + } + + char *c_base = nullptr; + char *d_base = nullptr; - char *c_base = static_cast(C->data.dptr); - char *d_base = static_cast(D->data.dptr); + if (!has_c_multi_tensor) { + NVTE_CHECK(C != nullptr && D != nullptr, + "Grouped GEMM: C/D grouped tensors are required when no C list is provided"); + C_meta = TensorShapeInfo::create_shape_info_for_C(C, D); + c_base = static_cast(C->data.dptr); + } + if (!has_d_multi_tensor) { + NVTE_CHECK(D != nullptr, + "Grouped GEMM: D grouped tensor is required when no D list is provided"); + D_meta = TensorShapeInfo::from_tensor(D); + 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 c_elem_size = transformer_engine::typeToSize(C->dtype()); - const size_t d_elem_size = transformer_engine::typeToSize(D->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; @@ -671,16 +1009,13 @@ inline void launch_grouped_gemm_setup( 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), A_sel.scale_inv, B_sel.scale_inv, - A_sel.scaling_mode, num_tensors); + 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); NVTE_CHECK_CUDA(cudaGetLastError()); } -inline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) { - return GroupedGemmSetupWorkspace::required_setup_size(num_tensors, kGroupedGemmAlignment); -} - } // namespace size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) { @@ -697,12 +1032,7 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT using namespace transformer_engine; // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.2+ - const int current_device = transformer_engine::cuda::current_device(); - NVTE_CHECK(transformer_engine::cuda::sm_arch(current_device) >= 100, - "nvte_grouped_gemm requires Blackwell (SM100) or newer architecture."); - NVTE_CHECK(transformer_engine::cuda::cublas_version() >= 130200, - "nvte_grouped_gemm requires cuBLAS 13.2+, but run-time cuBLAS version is ", - transformer_engine::cuda::cublas_version()); + check_grouped_gemm_requirements("nvte_grouped_gemm"); // Convert to internal types const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A); @@ -715,81 +1045,279 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT Tensor *wspace_cublas = convertNVTETensor(workspace_cublas); // Parse config (if provided) - GroupedMatmulConfig config_; - if (config != nullptr) { - config_ = *reinterpret_cast(config); - } + GroupedMatmulConfig config_ = parse_grouped_gemm_config(config); - // Validate inputs and num_tensors; returns dtype properties shared by A and B. - const auto [is_fp8, is_mxfp8] = - validate_grouped_gemm_inputs(inputA, inputB, inputC_raw, outputD, alpha_tensor, beta_tensor); + // 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}); // 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 = inputA->num_tensors; - + // num_tensors validated above. // Select operand storage (row-wise vs column-wise) and adjust transpose flags to // 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); // Workspaces: setup (pointer arrays) and cuBLAS - const size_t setup_workspace_size = grouped_gemm_setup_workspace_size(num_tensors); - const size_t cublas_workspace_size = kGroupedGemmCublasWorkspaceSize; + auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); - void *setup_workspace_ptr = validate_and_get_workspace_ptr(wspace_setup, setup_workspace_size, - "Grouped GEMM setup workspace"); - void *cublas_workspace_ptr = validate_and_get_workspace_ptr(wspace_cublas, cublas_workspace_size, - "Grouped GEMM cuBLAS workspace"); + 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()); - auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers( - static_cast(setup_workspace_ptr), num_tensors); + // 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 = + 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, outputD->dtype(), num_tensors, + config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val, + workspace.cublas_workspace_ptr, stream); +} - launch_grouped_gemm_setup(setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, - beta_tensor, num_tensors, stream); +void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, + int transa, const NVTEGroupedTensor B, int transb, + const NVTEGroupedTensor C, NVTEGroupedTensor D, + const NVTETensor alpha, const NVTETensor beta, + NVTETensor workspace_setup, NVTETensor workspace_cublas, + NVTEGroupedMatmulConfig config, cudaStream_t stream) { + NVTE_API_CALL(nvte_grouped_gemm_with_discrete_inputA); + using namespace transformer_engine; - // Get cuBLAS handle - using cublasHandleManager = detail::HandleManager; - cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.2+ + check_grouped_gemm_requirements("nvte_grouped_gemm_with_discrete_inputA"); - // Setup cuBLAS operations - cublasOperation_t op_A = A_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N; - cublasOperation_t op_B = B_sel.trans ? CUBLAS_OP_T : CUBLAS_OP_N; + 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."); - // Create grouped matrix layouts - cublasLtMatrixLayoutOpaque_t descA, descB, descC, descD; - init_matrix_layouts(descA, descB, descC, descD, setup_workspace, A_sel, B_sel, outputD, - num_tensors); + const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B); + const GroupedTensor *inputC_raw = convertNVTEGroupedTensor(C); // Can be NULL + GroupedTensor *outputD = convertNVTEGroupedTensorCheck(D); + const Tensor *alpha_tensor = convertNVTETensorCheck(alpha); + const Tensor *beta_tensor = convertNVTETensorCheck(beta); + Tensor *wspace_setup = convertNVTETensor(workspace_setup); + Tensor *wspace_cublas = convertNVTETensor(workspace_cublas); - // Create matmul descriptor - cublasLtMatmulDescOpaque_t matmulDesc; - init_matmul_desc(matmulDesc, op_A, op_B, is_fp8, config_.use_split_accumulator); - if (is_mxfp8) { - set_mxfp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, - setup_workspace.b_scale_inv_ptrs); - } else if (is_fp8) { - set_fp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, - setup_workspace.b_scale_inv_ptrs); + // Parse config (if provided) + 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; + + // 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) { + return dtype == transformer_engine::DType::kFloat8E4M3 || + dtype == transformer_engine::DType::kFloat8E5M2 || + 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."); + + // 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)) { + 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."); + } + + // Select operand storage for B (row-wise vs column-wise) + auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false); + + GroupedOperandSelection A_sel{}; + 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); + + 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); + + 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"); + A_sel.trans = choice.trans; + 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"); } + // For discrete A_list, scale pointers are per-tensor; use multi-tensor args. + // Base pointer is unused when providing per-tensor pointers. + A_sel.scale_inv = nullptr; + A_sel.dptr = nullptr; + + // 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 - // 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_n_val = + 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); + 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); +} + +void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, + const NVTEGroupedTensor B, int transb, + const NVTETensor *C_list, size_t num_c_tensors, + NVTETensor *D_list, size_t num_d_tensors, + const NVTETensor alpha, const NVTETensor beta, + NVTETensor workspace_setup, NVTETensor workspace_cublas, + NVTEGroupedMatmulConfig config, cudaStream_t stream) { + NVTE_API_CALL(nvte_grouped_gemm_with_discrete_out); + using namespace transformer_engine; + + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.2+ + check_grouped_gemm_requirements("nvte_grouped_gemm_with_discrete_out"); + + 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) { + NVTE_CHECK(C_list != nullptr, "Grouped GEMM: C_list is null but num_c_tensors > 0."); + } + + const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A); + const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B); + const Tensor *alpha_tensor = convertNVTETensorCheck(alpha); + const Tensor *beta_tensor = convertNVTETensorCheck(beta); + Tensor *wspace_setup = convertNVTETensor(workspace_setup); + Tensor *wspace_cublas = convertNVTETensor(workspace_cublas); + + 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); + 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."); + + // Parse config (if provided) + GroupedMatmulConfig config_ = parse_grouped_gemm_config(config); + + // Select operand storage (row-wise vs column-wise) and adjust transpose flags to + // 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); + // 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 = + config_.avg_m.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA)); + int64_t avg_n_val = + 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(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); +} - // Heuristic selection - cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo(handle, matmulDesc, descA, descB, descC, - descD, avg_m_val, avg_n_val, avg_k_val); +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; - // Execute the grouped GEMM - 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)); + const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output); + const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); + + 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 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); + 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); + }); + + NVTE_CHECK_CUDA(cudaGetLastError()); } #else // CUBLAS_VERSION < 130200 @@ -803,6 +1331,37 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); } +void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, + int transa, const NVTEGroupedTensor B, int transb, + const NVTEGroupedTensor C, NVTEGroupedTensor D, + const NVTETensor alpha, const NVTETensor beta, + NVTETensor workspace_setup, NVTETensor workspace_cublas, + NVTEGroupedMatmulConfig config, cudaStream_t stream) { + NVTE_ERROR( + "nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.2+, but compile-time " + "cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); +} + +void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, + const NVTEGroupedTensor B, int transb, + const NVTETensor *C_list, size_t num_c_tensors, + NVTETensor *D_list, size_t num_d_tensors, + const NVTETensor alpha, const NVTETensor beta, + NVTETensor workspace_setup, NVTETensor workspace_cublas, + NVTEGroupedMatmulConfig config, cudaStream_t stream) { + NVTE_ERROR( + "nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.2+, but compile-time " + "cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); +} + +void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, + cudaStream_t stream) { + NVTE_ERROR("nvte_grouped_bias_add requires cuBLAS 13.2+, but compile-time cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 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.2+, 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 35d327f085..6999dd857f 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -362,6 +362,51 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream); +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Grouped matrix multiplication with discrete A input tensors. + * + * Identical to nvte_grouped_gemm, but A is provided as a list of tensors + * 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] num_a_tensors Number of tensors in A_list. + */ +void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, + int transa, const NVTEGroupedTensor B, int transb, + const NVTEGroupedTensor C, NVTEGroupedTensor D, + const NVTETensor alpha, const NVTETensor beta, + NVTETensor workspace_setup, NVTETensor workspace_cublas, + NVTEGroupedMatmulConfig config, cudaStream_t stream); + +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Grouped matrix multiplication with discrete output tensors. +* +* Identical to nvte_grouped_gemm, but C and D are provided as lists of tensors +* 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] 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[in] num_d_tensors Number of tensors in D_list. +* \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, + const NVTEGroupedTensor B, int transb, + const NVTETensor *C_list, size_t num_c_tensors, + NVTETensor *D_list, size_t num_d_tensors, + const NVTETensor alpha, const NVTETensor beta, + NVTETensor workspace_setup, NVTETensor workspace_cublas, + NVTEGroupedMatmulConfig config, cudaStream_t stream); + +/*! \brief Grouped bias add for grouped GEMM outputs. +* +* 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); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 5e420b2d42..904812118c 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -63,6 +63,21 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen * */ void nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Swizzling scaling factors into the required interleaved layout for GEMM (grouped tensor) + * + * \param[in] input Input grouped tensor with non-swizzled scale_inv. + * \param[in,out] output Output grouped tensor which hosts swizzled scale_inv. + * \param[in] stream CUDA stream used for the operation. + * + * Requirements(for now, more features will be added later): + * - scaling mode must be MXFP8 1D scaling. + * - scale_inv is stored in row-major per group. + * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. + * - all tensors in the grouped tensor must have the same shape. + */ +void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream); #ifdef __cplusplus } // extern "C" diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 4425c4e9f7..619987931e 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -268,6 +268,40 @@ struct MultiSwizzleArgs { int num_tensors; }; +constexpr size_t round_up_to_multiple(size_t value, size_t multiple) { + return DIVUP(value, multiple) * multiple; +} + +template +__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 int tensor_id = blockIdx.z; + 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( + input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, + gridDim.y); +} + +template +__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 int tensor_id = blockIdx.z; + 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( + input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, + gridDim.y); +} + template __global__ void multi_tensor_swizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) { // Find tensor corresponding to block @@ -569,9 +603,9 @@ void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, 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. - */ + * 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]; @@ -819,10 +853,10 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, } // namespace transformer_engine /* - * WIP (Phuong): - * - Opt for bank conflicts - * - Adding swizzle for 2d-block scaling. - */ +* WIP (Phuong): +* - Opt for bank conflicts +* - Adding swizzle for 2d-block scaling. +*/ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_swizzle_scaling_factors); using namespace transformer_engine; @@ -841,3 +875,171 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen } multi_tensor_swizzle_scaling_factors(input_list, output_list, stream); } + +namespace transformer_engine { + +void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, + cudaStream_t stream) { + // Check scaling mode + NVTE_CHECK(input->scaling_mode == NVTE_MXFP8_1D_SCALING, + "Grouped swizzle supports only MXFP8 scaling."); + + // Check tensors + CheckInputGroupedTensor(*input, "input"); + CheckOutputGroupedTensor(*output, "output", false); + NVTE_CHECK(!input->with_gemm_swizzled_scales, + "Expected input grouped tensor with scales in compact format."); + NVTE_CHECK(output->with_gemm_swizzled_scales, + "Expected output grouped tensor with scales in GEMM swizzled format."); + + // Check scaling factors availability + const bool has_rowwise_scale_inv = input->scale_inv.has_data(); + const bool has_columnwise_scale_inv = input->columnwise_scale_inv.has_data(); + if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) { + 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."); + + // 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); + const size_t scale_elems = padded_m * padded_k; + + 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."); + } 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."); + } + + 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; + + 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 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) { + 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."); + } + } 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."); + } + } + NVTE_CHECK_CUDA(cudaGetLastError()); + }; + + if (has_rowwise_scale_inv) { + launch_grouped_swizzle(true); + } + if (has_columnwise_scale_inv) { + launch_grouped_swizzle(false); + } +} + +} // namespace transformer_engine + +void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_swizzle_grouped_scaling_factors); + using namespace transformer_engine; + swizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input), + convertNVTEGroupedTensorCheck(output), stream); +} diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index a37f1c2d4d..115569ccba 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -5,6 +5,7 @@ """Python interface for GEMM extensions""" from typing import Iterable, Optional, Tuple, Union, List +import ctypes import os import functools import torch @@ -22,6 +23,7 @@ __all__ = [ "general_gemm", "general_grouped_gemm", + "general_grouped_gemm_for_grouped_tensor", ] @@ -284,3 +286,113 @@ def general_grouped_gemm( ) return out, bias, gelu_input + + +@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 + + +def general_grouped_gemm_for_grouped_tensor( + A, + B, + out, + *, + layout: str = "TN", + accumulate: bool = False, + use_split_accumulator: bool = False, + bias=None, + grad: bool = False, + alpha: Optional[torch.Tensor] = None, + beta: Optional[torch.Tensor] = None, +) -> Union[torch.Tensor, List[torch.Tensor]]: + """ + Grouped GEMM using GroupedTensor inputs. + + This uses nvte_grouped_gemm and supports different per-matrix shapes. + + The caller must ensure that GroupedTensor metadata is already compatible with the + underlying GEMM implementation (e.g., aligned offsets and output metadata layout). + """ + assert layout in ("TN", "NN", "NT"), f"GEMM layout {layout} not supported." + if grad: + raise NotImplementedError("grad is not supported for grouped_tensor GEMM yet.") + transa = layout[0] == "T" + transb = layout[1] == "T" + is_discrete_out = isinstance(out, list) + is_discrete_in = isinstance(A, list) + if is_discrete_in and is_discrete_out: + raise ValueError("Both A and out are discrete. This is not supported yet.") + + if is_discrete_out: + # wgrad case. + grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_out + elif is_discrete_in: + # Use-case: forward pass with list of weights. + grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_in + else: + # Use-case: Single Grouped Parameter for Weight/ Weight Grads. + grouped_gemm_impl = tex.te_general_grouped_gemm_for_grouped_tensor + + if is_discrete_out and bias is not None: + raise ValueError( + "Bias is not supported when out is a list (discrete_out mode) yet. " + "Apply bias manually after the GEMM." + ) + + num_tensors = B.num_tensors + rowwise = B.rowwise_data + device = rowwise.device if rowwise is not None else B.columnwise_data.device + + if alpha is None: + alpha = torch.ones(num_tensors, dtype=torch.float32, device=device) + if beta is None: + if accumulate: + beta = torch.ones(num_tensors, dtype=torch.float32, device=device) + else: + beta = torch.zeros(num_tensors, dtype=torch.float32, device=device) + + if not alpha.is_cuda or not beta.is_cuda: + raise ValueError("alpha and beta must be CUDA tensors.") + + workspace_setup = torch.empty( + get_grouped_gemm_setup_workspace_size(num_tensors), + dtype=torch.uint8, + device=device, + ) + workspace_cublas = torch.empty( + get_cublas_workspace_size_bytes(), + dtype=torch.uint8, + device=device, + ) + + sm_count = get_sm_count() + sm_count = sm_count - int(os.getenv("NVTE_EXT_MARGIN_SM", str(sm_count))) + + return grouped_gemm_impl( + A, + transa, + B, + transb, + out, + bias, + alpha, + beta, + workspace_setup, + workspace_cublas, + use_split_accumulator, + sm_count, + ) diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 66c01aaf7a..1c5116a8da 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -149,6 +149,25 @@ std::optional> te_general_grouped_gemm( std::vector pre_gelu_out, bool grad, std::vector workspace, size_t workspaceSize, bool accumulate, bool use_split_accumulator, int math_sm_count); +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); + +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, + 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_out(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); + /*************************************************************************************************** * Transpose **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index d75b0f14c7..1431ebdfb4 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -78,6 +78,46 @@ bool checkGemmShape(const std::vector& expected, const NVTEShape& actual return true; } +struct GroupedGemmConfig { + TensorWrapper te_alpha; + TensorWrapper te_beta; + TensorWrapper te_workspace_setup; + TensorWrapper te_workspace_cublas; + std::optional matmul_config; +}; + +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."); + + GroupedGemmConfig grouped_gemm_config{ + makeTransformerEngineTensor(alpha), + makeTransformerEngineTensor(beta), + makeTransformerEngineTensor(workspace_setup.data_ptr(), + std::vector{static_cast(workspace_setup.numel())}, + DType::kByte), + makeTransformerEngineTensor( + workspace_cublas.data_ptr(), + std::vector{static_cast(workspace_cublas.numel())}, DType::kByte), + std::nullopt, + }; + + if (math_sm_count > 0 || use_split_accumulator) { + grouped_gemm_config.matmul_config.emplace(); + if (math_sm_count > 0) { + grouped_gemm_config.matmul_config->set_sm_count(math_sm_count); + } + grouped_gemm_config.matmul_config->set_use_split_accumulator(use_split_accumulator); + } + + return grouped_gemm_config; +} + } // namespace detail std::pair createOutputTensor(const std::vector& shape, @@ -570,4 +610,180 @@ std::optional> te_general_grouped_gemm( return bias; } +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) { + using namespace transformer_engine::pytorch::detail; + + init_extension(); + + // Ensure that cublasLt handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(workspace_cublas.device()); + + auto grouped_A = GroupedTensorFromPyTorchGroupedTensor(A); + auto grouped_B = GroupedTensorFromPyTorchGroupedTensor(B); + auto grouped_D = GroupedTensorFromPyTorchGroupedTensor(D); + + const size_t num_tensors = grouped_A.num_tensors(); + NVTE_CHECK(num_tensors > 0, "Grouped GEMM requires non-empty inputs."); + NVTE_CHECK(grouped_B.num_tensors() == num_tensors, + "Grouped GEMM requires A and B to have the same num_tensors."); + NVTE_CHECK(grouped_D.num_tensors() == num_tensors, + "Grouped GEMM requires D to have the same num_tensors as inputs."); + + auto gemm_config = prepare_grouped_gemm_config(alpha, beta, workspace_setup, workspace_cublas, + num_tensors, math_sm_count, use_split_accumulator); + + [[maybe_unused]] auto swizzled_scales_A = maybe_swizzle_grouped_tensor_for_gemm(grouped_A); + [[maybe_unused]] auto swizzled_scales_B = maybe_swizzle_grouped_tensor_for_gemm(grouped_B); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_gemm(grouped_A.data(), transa, grouped_B.data(), transb, grouped_D.data(), + grouped_D.data(), gemm_config.te_alpha.data(), gemm_config.te_beta.data(), + gemm_config.te_workspace_setup.data(), gemm_config.te_workspace_cublas.data(), + gemm_config.matmul_config.has_value() + ? static_cast(*gemm_config.matmul_config) + : nullptr, + at::cuda::getCurrentCUDAStream()); + }); + + 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()); + }); + } + + return py::reinterpret_borrow(D); +} + +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, + 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(); + + // Ensure that cublasLt handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(workspace_cublas.device()); + + auto grouped_B = GroupedTensorFromPyTorchGroupedTensor(B); + auto grouped_D = GroupedTensorFromPyTorchGroupedTensor(D); + + const auto A_list = py::cast>(A); + const size_t num_tensors = grouped_B.num_tensors(); + NVTE_CHECK(num_tensors > 0, "Grouped GEMM requires non-empty inputs."); + NVTE_CHECK(A_list.size() == num_tensors, + "Grouped GEMM requires A_list to have num_tensors elements."); + NVTE_CHECK(grouped_D.num_tensors() == num_tensors, + "Grouped GEMM requires D to have the same num_tensors as inputs."); + + auto gemm_config = prepare_grouped_gemm_config(alpha, beta, workspace_setup, workspace_cublas, + num_tensors, math_sm_count, use_split_accumulator); + + std::vector te_A_wrappers; + std::vector te_A_vector; + te_A_wrappers.reserve(num_tensors); + te_A_vector.reserve(num_tensors); + const auto none = py::none(); + for (const auto& tensor : A_list) { + te_A_wrappers.emplace_back(makeTransformerEngineTensor(tensor, none)); + te_A_vector.emplace_back(te_A_wrappers.back().data()); + } + + std::vector> swizzled_scale_inverses_list; + swizzled_scale_inverses_list.emplace_back( + multi_tensor_swizzle_scales_for_gemm(te_A_wrappers, transa, !transa)); + + [[maybe_unused]] auto swizzled_scales_B = maybe_swizzle_grouped_tensor_for_gemm(grouped_B); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_gemm_with_discrete_inputA( + te_A_vector.data(), num_tensors, transa, grouped_B.data(), transb, grouped_D.data(), + grouped_D.data(), gemm_config.te_alpha.data(), gemm_config.te_beta.data(), + gemm_config.te_workspace_setup.data(), gemm_config.te_workspace_cublas.data(), + gemm_config.matmul_config.has_value() + ? static_cast(*gemm_config.matmul_config) + : nullptr, + at::cuda::getCurrentCUDAStream()); + }); + + 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()); + }); + } + + return py::reinterpret_borrow(D); +} + +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, + 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(); + + // Ensure that cublasLt handle is created on the correct device, + // overriding torch.cuda.set_device calls from user side. + // Assumes all tensors passed are on the same device. + at::cuda::CUDAGuard device_guard(workspace_cublas.device()); + + NVTE_CHECK(bias.is_none(), "Bias is not supported for discrete output grouped GEMM."); + + auto grouped_A = GroupedTensorFromPyTorchGroupedTensor(A); + auto grouped_B = GroupedTensorFromPyTorchGroupedTensor(B); + + const auto D_list = py::cast>(D); + const size_t num_tensors = grouped_A.num_tensors(); + NVTE_CHECK(num_tensors > 0, "Grouped GEMM requires non-empty inputs."); + NVTE_CHECK(grouped_B.num_tensors() == num_tensors, + "Grouped GEMM requires A and B to have the same num_tensors."); + NVTE_CHECK(D_list.size() == num_tensors, + "Grouped GEMM requires D_list to have num_tensors elements."); + + auto gemm_config = prepare_grouped_gemm_config(alpha, beta, workspace_setup, workspace_cublas, + num_tensors, math_sm_count, use_split_accumulator); + + std::vector te_D_wrappers; + std::vector te_D_vector; + te_D_wrappers.reserve(num_tensors); + te_D_vector.reserve(num_tensors); + const auto none = py::none(); + for (const auto& tensor : D_list) { + te_D_wrappers.emplace_back(makeTransformerEngineTensor(tensor, none)); + te_D_vector.emplace_back(te_D_wrappers.back().data()); + } + + [[maybe_unused]] auto swizzled_scales_A = maybe_swizzle_grouped_tensor_for_gemm(grouped_A); + [[maybe_unused]] auto swizzled_scales_B = maybe_swizzle_grouped_tensor_for_gemm(grouped_B); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_gemm_with_discrete_out( + grouped_A.data(), transa, grouped_B.data(), transb, te_D_vector.data(), num_tensors, + te_D_vector.data(), num_tensors, gemm_config.te_alpha.data(), gemm_config.te_beta.data(), + gemm_config.te_workspace_setup.data(), gemm_config.te_workspace_cublas.data(), + gemm_config.matmul_config.has_value() + ? static_cast(*gemm_config.matmul_config) + : nullptr, + at::cuda::getCurrentCUDAStream()); + }); + + return py::reinterpret_borrow(D); +} + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 7721671a36..c590a3c9e2 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -276,6 +276,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("quantizer_list"), py::arg("disable_bulk_allocation") = false); m.def("te_general_grouped_gemm", &transformer_engine::pytorch::te_general_grouped_gemm, "Grouped GEMM"); + m.def("te_general_grouped_gemm_for_grouped_tensor", + &transformer_engine::pytorch::te_general_grouped_gemm_for_grouped_tensor, + "Grouped GEMM for GroupedTensor"); + m.def("te_general_grouped_gemm_for_discrete_in", + &transformer_engine::pytorch::te_general_grouped_gemm_for_discrete_in, + "Grouped GEMM for discrete A input list"); + m.def("te_general_grouped_gemm_for_discrete_out", + &transformer_engine::pytorch::te_general_grouped_gemm_for_discrete_out, + "Grouped GEMM for discrete output list"); m.def("fp8_transpose", &transformer_engine::pytorch::fp8_transpose, "Transpose with FP8 I/O", py::arg("input"), py::arg("dtype"), py::kw_only(), py::arg("out"), py::call_guard()); diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index a4750d9aa0..7ff35d6b68 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -35,6 +35,13 @@ void reset_tensor_data(transformer_engine::TensorWrapper &tensor, bool rowwise, } } +bool is_empty_grouped_tensor_param(const NVTEBasicTensor &t) { + if (t.data_ptr == nullptr) { + return true; + } + return t.shape.ndim == 1 && t.shape.data[0] == 0; +} + } // namespace std::tuple, std::optional> swizzle_scales_for_gemm( @@ -331,6 +338,83 @@ at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapp return swizzled_scale_inv; } +std::optional maybe_swizzle_grouped_tensor_for_gemm( + GroupedTensorWrapper &input) { + if (input.scaling_mode() != NVTE_MXFP8_1D_SCALING) { + return std::nullopt; + } + if (input.get_with_gemm_swizzled_scales()) { + return std::nullopt; + } + + const auto row_scales = input.get_rowwise_scale_inv(); + const auto col_scales = input.get_columnwise_scale_inv(); + const bool has_rowwise_scales = !is_empty_grouped_tensor_param(row_scales); + const bool has_columnwise_scales = !is_empty_grouped_tensor_param(col_scales); + if (!has_rowwise_scales && !has_columnwise_scales) { + 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; + GroupedTensorWrapper output(input.num_tensors(), input.logical_shape(), input.scaling_mode()); + + const auto rowwise_data = input.get_rowwise_data(); + if (rowwise_data.data_ptr != nullptr) { + output.set_rowwise_data(rowwise_data.data_ptr, static_cast(rowwise_data.dtype), + rowwise_data.shape); + } + const auto columnwise_data = input.get_columnwise_data(); + if (columnwise_data.data_ptr != nullptr) { + output.set_columnwise_data(columnwise_data.data_ptr, static_cast(columnwise_data.dtype), + columnwise_data.shape); + } + const auto tensor_offsets = input.get_tensor_offsets(); + if (tensor_offsets.data_ptr != nullptr) { + output.set_tensor_offsets(tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), + tensor_offsets.shape); + } + + if (has_rowwise_scales) { + const auto scales_dtype = static_cast(row_scales.dtype); + rowwise_scales_pyt = allocateSpace(row_scales.shape, scales_dtype, false); + void *output_scales_dptr = getDataPtr(*rowwise_scales_pyt); + output.set_rowwise_scale_inv(output_scales_dptr, scales_dtype, row_scales.shape); + } + if (has_columnwise_scales) { + const auto scales_dtype = static_cast(col_scales.dtype); + columnwise_scales_pyt = allocateSpace(col_scales.shape, scales_dtype, false); + void *output_scales_dptr = getDataPtr(*columnwise_scales_pyt); + output.set_columnwise_scale_inv(output_scales_dptr, scales_dtype, col_scales.shape); + } + + output.set_with_gemm_swizzled_scales(true); + NVTE_SCOPED_GIL_RELEASE({ + nvte_swizzle_grouped_scaling_factors(input.data(), output.data(), + at::cuda::getCurrentCUDAStream()); + }); + + if (has_rowwise_scales) { + const auto scales_dtype = static_cast(row_scales.dtype); + input.set_rowwise_scale_inv(getDataPtr(*rowwise_scales_pyt), scales_dtype, row_scales.shape); + } + if (has_columnwise_scales) { + const auto scales_dtype = static_cast(col_scales.dtype); + input.set_columnwise_scale_inv(getDataPtr(*columnwise_scales_pyt), scales_dtype, + col_scales.shape); + } + input.set_with_gemm_swizzled_scales(true); + + return SwizzledGroupedScales{std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)}; +} + 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/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index c904057e97..8c5504e44b 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -202,6 +202,7 @@ std::pair NoneQuantizer::create_grouped_tensor kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); kwargs["last_dims"] = py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -412,6 +413,7 @@ std::pair Float8Quantizer::create_grouped_tens kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); kwargs["last_dims"] = py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -734,6 +736,7 @@ std::pair Float8CurrentScalingQuantizer::creat kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); kwargs["last_dims"] = py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -1096,6 +1099,7 @@ std::pair Float8BlockQuantizer::create_grouped kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); kwargs["last_dims"] = py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); + kwargs["with_gemm_swizzled_scales"] = py::cast(false); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -1516,6 +1520,7 @@ std::pair MXFP8Quantizer::create_grouped_tenso kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); 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; PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -1948,6 +1953,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso kwargs["first_dims"] = first_dims.has_value() ? py::cast(*first_dims) : py::none(); 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; PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index eda5e8fc54..e9c6ca882e 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -216,8 +216,8 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { auto ret = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode); // Rowwise data - if (!tensor.attr("data").is_none()) { - const auto &data = tensor.attr("data").cast(); + if (!tensor.attr("rowwise_data").is_none()) { + const auto &data = tensor.attr("rowwise_data").cast(); DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); @@ -282,6 +282,12 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { getTensorShape(tensor_offsets)); } + bool with_gemm_swizzled = false; + if (py::hasattr(tensor, "_with_gemm_swizzled_scales")) { + with_gemm_swizzled = tensor.attr("_with_gemm_swizzled_scales").cast(); + } + ret.set_with_gemm_swizzled_scales(with_gemm_swizzled); + return ret; } diff --git a/transformer_engine/pytorch/csrc/util.h b/transformer_engine/pytorch/csrc/util.h index 8988c18261..587ec289a4 100644 --- a/transformer_engine/pytorch/csrc/util.h +++ b/transformer_engine/pytorch/csrc/util.h @@ -33,6 +33,16 @@ std::optional multi_tensor_swizzle_scales_for_gemm(std::vector, std::optional>; + +/*! \brief Swizzle grouped tensor scales for GEMM if needed. + * Currently only works for MXFP8 1D scaling with uniform shapes. + * + * The returned swizzled scales should be kept alive during the GEMM. + */ +std::optional maybe_swizzle_grouped_tensor_for_gemm( + GroupedTensorWrapper& input); + /*! \brief Convert a block scaling tensor to an mxfp8 tensor in-place. * * If rowwise==false, the columnwise data will be reinterpreted as diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 22a6a41eb1..2fce9a38e2 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -91,6 +91,7 @@ def __new__( columnwise_scale_inv_offsets: Optional[List[int]] = None, requires_grad: bool = False, stride: Optional[List[int]] = None, + with_gemm_swizzled_scales: bool = False, ): if ( shapes is not None @@ -154,6 +155,7 @@ def __new__( offsets=offsets, scale_inv_offsets=scale_inv_offsets, columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, ) 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 3b7b9bc169..68097259c6 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -71,6 +71,7 @@ def _initialize_storage_fields( columnwise_scale_inv_offsets: Optional[List[int]] = None, requires_grad: bool = False, stride: Optional[List[int]] = None, + with_gemm_swizzled_scales: bool = False, ) -> None: """ Initialize a GroupedTensor. @@ -144,6 +145,7 @@ def _initialize_storage_fields( # Hold a reference to the quantized tensors that occupy same storage as the GroupedTensor. # Used as a convenience. instance.quantized_tensors = None + instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales def __new__( cls, @@ -168,6 +170,7 @@ def __new__( columnwise_scale_inv_offsets: Optional[List[int]] = None, requires_grad: bool = False, stride: Optional[List[int]] = None, + with_gemm_swizzled_scales: bool = False, ): instance = object.__new__(cls) cls._initialize_storage_fields( @@ -192,6 +195,7 @@ def __new__( columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, requires_grad=requires_grad, stride=stride, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, ) return instance @@ -645,6 +649,9 @@ def make_grouped_tensor( offsets=offsets, scale_inv_offsets=scale_inv_offsets, columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, + with_gemm_swizzled_scales=( + quantizer.optimize_for_gemm if quantizer is not None else False + ), ) grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() From 40588422a2032d00935d84feba12dad2b6b83c2f Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 16 Mar 2026 10:25:14 -0700 Subject: [PATCH 282/521] Changed VERSION to 2.15.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 c7d5307735..34ab1df063 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.14.0.dev0 +2.15.0.dev0 From a94584628ddf7b25859875e0bcc90b99f9c18388 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Mon, 16 Mar 2026 11:19:22 -0700 Subject: [PATCH 283/521] [Common] Fix linker error for to_string(DType) in distributed tests (#2757) * [Common] Fix linker error for to_string(DType) in distributed tests Make transformer_engine::to_string(DType) inline in common.h so that translation units outside libtransformer_engine.so can resolve it without requiring the symbol to be exported. Regression introduced by 61f95942 which added to_string(DType) calls into TRANSFORMER_ENGINE_TYPE_SWITCH_* macros, causing test object files to reference the symbol that the linker version script hides. Co-Authored-By: Claude Sonnet 4.6 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 --------- Signed-off-by: Vladimir Cherepanov 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/common/common.h | 29 ++++++++++++++++++- .../common/transformer_engine.cpp | 29 ------------------- 2 files changed, 28 insertions(+), 30 deletions(-) diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 41a8fd1112..a98668d058 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -41,7 +41,34 @@ static_assert(NVTE_BUILD_NUM_PHILOX_ROUNDS > 0, namespace transformer_engine { -std::string to_string(const DType type); +inline std::string to_string(const DType type) { + switch (type) { + case DType::kByte: + return "Byte"; + case DType::kBFloat16: + return "BFloat16"; + case DType::kFloat16: + return "Float16"; + case DType::kFloat32: + return "Float32"; + case DType::kFloat8E4M3: + return "Float8E4M3"; + case DType::kFloat8E5M2: + return "Float8E5M2"; + case DType::kFloat8E8M0: + return "Float8E8M0"; + case DType::kFloat4E2M1: + return "Float4E2M1"; + case DType::kInt16: + return "Int16"; + case DType::kInt32: + return "Int32"; + case DType::kInt64: + return "Int64"; + default: + return std::string("Invalid type ") + std::to_string(static_cast(type)); + } +} std::string to_string(const NVTEScalingMode &mode); inline std::string to_string_like(const DType &val) { return to_string(val); } diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 1875f4f690..b97504f2ae 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -33,35 +33,6 @@ size_t typeToSize(const DType type) { return typeToNumBits(type) / 8; } -std::string to_string(const DType type) { - switch (type) { - case DType::kByte: - return "Byte"; - case DType::kBFloat16: - return "BFloat16"; - case DType::kFloat16: - return "Float16"; - case DType::kFloat32: - return "Float32"; - case DType::kFloat8E4M3: - return "Float8E4M3"; - case DType::kFloat8E5M2: - return "Float8E5M2"; - case DType::kFloat8E8M0: - return "Float8E8M0"; - case DType::kFloat4E2M1: - return "Float4E2M1"; - case DType::kInt16: - return "Int16"; - case DType::kInt32: - return "Int32"; - case DType::kInt64: - return "Int64"; - default: - return concat_strings("Invalid type ", static_cast(type)); - } -} - std::string to_string(const NVTEScalingMode &mode) { switch (mode) { case NVTE_DELAYED_TENSOR_SCALING: From 523801df70c6598e6cee0e9197134b08d4b230b9 Mon Sep 17 00:00:00 2001 From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com> Date: Mon, 16 Mar 2026 11:24:49 -0700 Subject: [PATCH 284/521] [NVFP4][Dense/MoE] Integrate Cutlass NVFP4 Row-Cast-Col-RHT-Transpose-Cast Fusion Kernel (#2555) * first draft Signed-off-by: Zhongbo Zhu * pass numerical unit test Signed-off-by: Zhongbo Zhu * format Signed-off-by: Zhongbo Zhu * add benchmark script Signed-off-by: Zhongbo Zhu * lint and format Signed-off-by: Zhongbo Zhu * compile guard Signed-off-by: Zhongbo Zhu * warning fix Signed-off-by: Zhongbo Zhu * resolve greptile comment Signed-off-by: Zhongbo Zhu * minor style fixes Signed-off-by: Zhongbo Zhu * fix namespace Signed-off-by: Zhongbo Zhu * resolve some comments Signed-off-by: Zhongbo Zhu * fix comment Signed-off-by: Zhongbo Zhu * attempt to fix compile CI with guard Signed-off-by: Zhongbo Zhu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * better naming for tests Signed-off-by: Zhongbo Zhu * fix deprecate messsage Signed-off-by: Zhongbo Zhu * more compile guard Signed-off-by: Zhongbo Zhu * new API name Signed-off-by: Zhongbo Zhu * fix format all in one Signed-off-by: Zhongbo Zhu * try to fix compile CI again Signed-off-by: Zhongbo Zhu * AI code review comments Signed-off-by: Zhongbo Zhu * to pass oldest compile CI with cuda 12.1 Signed-off-by: Zhongbo Zhu * add more guards to nvfp4 Signed-off-by: Zhongbo Zhu * make multiply inverse default numerics Signed-off-by: Zhongbo Zhu * update numerics of nvfp4 partial cast as well Signed-off-by: Zhongbo Zhu * resolve comments Signed-off-by: Zhongbo Zhu * add NVTE_BUILD_NUM_PHILOX_ROUNDS after rebase Signed-off-by: Zhongbo Zhu * simplify compile guard messsages Signed-off-by: Zhongbo Zhu --------- Signed-off-by: Zhongbo Zhu Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- benchmarks/linear/benchmark_linear.py | 332 ++++ .../test_mxfp8_group_quantize_graph_safe.py | 56 +- .../test_mxfp8_quantize_swizzle_fusion.py | 24 +- tests/pytorch/nvfp4/nvfp4_utils.py | 4 +- .../nvfp4/test_nvfp4_group_quantize.py | 26 +- .../test_nvfp4_group_quantize_graph_safe.py | 52 +- .../nvfp4/test_nvfp4_quantize_exact.py | 16 +- .../nvfp4/test_nvfp4_rht_quantize_exact.py | 68 +- transformer_engine/common/CMakeLists.txt | 1 + .../common/cast/nvfp4/core_nvfp4.cuh | 8 +- ...cast_col_hadamard_transform_cast_fusion.cu | 1754 ++++++++--------- .../group_hadamard_transform_cast_fusion.cu | 999 +++++----- ...cast_col_hadamard_transform_cast_fusion.cu | 1726 ++++++++-------- .../hadamard_transform_cast_fusion.cu | 22 +- ...cast_col_hadamard_transform_cast_fusion.cu | 1370 +++++++++++++ .../transformer_engine/hadamard_transform.h | 17 +- transformer_engine/common/recipe/nvfp4.cu | 51 +- ...quantize_transpose_vector_blockwise_fp4.cu | 14 +- transformer_engine/common/util/ptx.cuh | 6 +- transformer_engine/pytorch/csrc/common.h | 5 + .../pytorch/csrc/extensions/cast.cpp | 4 + transformer_engine/pytorch/csrc/quantizer.cpp | 234 ++- .../custom_recipes/quantization_nvfp4.py | 5 +- 23 files changed, 4279 insertions(+), 2515 deletions(-) create mode 100644 benchmarks/linear/benchmark_linear.py create mode 100644 transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu diff --git a/benchmarks/linear/benchmark_linear.py b/benchmarks/linear/benchmark_linear.py new file mode 100644 index 0000000000..4230db446d --- /dev/null +++ b/benchmarks/linear/benchmark_linear.py @@ -0,0 +1,332 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import argparse +import torch +import torch.utils.benchmark as benchmark +import pandas as pd + +from transformer_engine.pytorch.module import Linear as TELinear +from transformer_engine.common.recipe import ( + Float8BlockScaling, + MXFP8BlockScaling, + NVFP4BlockScaling, +) +from transformer_engine.pytorch.quantization import autocast, FP8GlobalStateManager +from contextlib import nullcontext + +""" +# Profile BF16 recipe with Nsight Systems +nsys profile \ + --output=./benchmarks/linear/b200_linear_bf16 \ + --force-overwrite true \ + --trace=cuda,nvtx,cudnn,cublas \ + python benchmarks/linear/benchmark_linear.py --profile --recipe bf16 + +# Profile FP8 sub-channel recipe with Nsight Systems +nsys profile \ + --output=./benchmarks/linear/b200_linear_fp8_sub_channel \ + --force-overwrite true \ + --trace=cuda,nvtx,cudnn,cublas \ + python benchmarks/linear/benchmark_linear.py --profile --recipe fp8_sub_channel + +# Profile MXFP8 recipe with Nsight Systems +nsys profile \ + --output=./benchmarks/linear/b200_linear_mxfp8 \ + --force-overwrite true \ + --trace=cuda,nvtx,cudnn,cublas \ + python benchmarks/linear/benchmark_linear.py --profile --recipe mxfp8 + +# Profile NVFP4 recipe with Nsight Systems +nsys profile \ + --output=./benchmarks/linear/b200_linear_nvfp4_rht_cast_fusion \ + --force-overwrite true \ + --trace=cuda,nvtx,cudnn,cublas \ + python benchmarks/linear/benchmark_linear.py --profile --recipe nvfp4 + +# Example to look at a single kernel target with NCU, like the fused hadamard amax kernel for NVFP4 recipe +ncu -f -o ./benchmarks/linear/ncu_b200_linear_nvfp4_rht_cast_fusion \ + --set=full \ + --kernel-name "row_col_rht_gemm_device" \ + -s 5 -c 5 \ + python benchmarks/linear/benchmark_linear.py --profile --recipe nvfp4 + +""" + +RECIPES = { + "bf16": None, + "fp8_sub_channel": Float8BlockScaling(), + "mxfp8": MXFP8BlockScaling(), + "nvfp4": NVFP4BlockScaling(), +} + +mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available() +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = ( + FP8GlobalStateManager.is_fp8_block_scaling_available() +) +nvfp4_available, reason_for_no_nvfp4 = FP8GlobalStateManager.is_nvfp4_available() + + +def run_linear_multiple_steps(layer, x, mode, gradient, run_num_steps=1, recipe=None): + assert mode in ["fwd_only", "fwd_bwd"] + quantization_context = ( + autocast(enabled=True, recipe=recipe) if recipe is not None else nullcontext() + ) + + if mode == "fwd_only": + with torch.no_grad(), quantization_context: + for i in range(run_num_steps): + y_q = layer.forward( + x, + is_first_microbatch=(i == 0), + ) + return y_q + else: + # reset gradients + layer.zero_grad() + x.grad = None + + with quantization_context: + for i in range(run_num_steps): + label = f"step_{i}" + torch.cuda.nvtx.range_push(label) + y_q = layer.forward( + x, + is_first_microbatch=(i == 0), + ) + y_q.backward(gradient) + torch.cuda.nvtx.range_pop() + + grads_q = [] + grads_q.append(x.grad) + # remaining derivatives are in respect to model parameters + for p in layer.parameters(): + if p.requires_grad: + grads_q.append(p.grad) + + return y_q, grads_q + + +def benchmark_linear( + x, + w, + bias, + recipe_name, + mode, +): + params_dtype = torch.bfloat16 + recipe = RECIPES[recipe_name] + + in_features = x.shape[1] + out_features = w.shape[0] + gradient = torch.ones((x.shape[0], out_features), dtype=torch.bfloat16, device=x.device) + + layer = TELinear( + in_features, + out_features, + bias=bias is not None, + params_dtype=params_dtype, + ) + + layer = layer.to("cuda") + with torch.no_grad(): + layer.weight.copy_(w) + if bias is not None: + layer.bias.copy_(bias) + + num_microbatches = 32 + + label = f"{recipe_name}_{'linear'}" + torch.cuda.nvtx.range_push(label) + timing = benchmark.Timer( + stmt="run_linear_multiple_steps(layer, x, mode, gradient, num_microbatches, recipe)", + globals={ + "run_linear_multiple_steps": run_linear_multiple_steps, + "layer": layer, + "x": x, + "mode": mode, + "gradient": gradient, + "num_microbatches": num_microbatches, + "recipe": recipe, + }, + num_threads=1, + ).blocked_autorange(min_run_time=10) + print(f"{recipe_name}: {timing} \n") + timing_ms = timing.median * 1000 / num_microbatches + + return timing_ms + + +def run_benchmark_linear(mkns, recipe_name, use_bias, fwd_only=False): + data = [] + assert not use_bias, "Bias is not supported in this benchmark script" + + print(f"========== Benchmarking {recipe_name} ==========") + for m, k, n in mkns: + device = "cuda" + x = torch.randn((m, k), dtype=torch.bfloat16, device=device, requires_grad=True) + w = torch.randn((n, k), dtype=torch.bfloat16, device=device) + bias = None + + # Run the benchmark + print(f"fwd_m={m}, fwd_k={k}, fwd_n={n}") + print(f"fwd_only: {fwd_only}") + + linear_fwd_bwd_timing_ms = benchmark_linear( + x, + w, + bias, + recipe_name, + mode="fwd_only" if fwd_only else "fwd_bwd", + ) + + # Append the results + data.append( + [ + m, + k, + n, + recipe_name, + linear_fwd_bwd_timing_ms, + ] + ) + + timing_notation = "linear_fwd_time_ms" if fwd_only else "linear_fwd_bwd_time_ms" + + df = pd.DataFrame( + data=data, + columns=[ + "m", + "k", + "n", + "recipe", + timing_notation, + ], + ) + + print(df, "\n") + return df + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + parser.add_argument("--profile", action="store_true", help="Enable profiling mode") + parser.add_argument( + "--output-dir", + type=str, + default="benchmark_output/", + help="output path for report", + ) + # arguments for recipe, options are fp8_sub_channel, mxfp8, bf16, all + parser.add_argument( + "--recipe", + type=str, + default="bf16", + help="Recipe to use, options are fp8_sub_channel, mxfp8, bf16, or all", + ) + parser.add_argument( + "--token-dim", + type=int, + default=None, + help="Token dimension to use, calculated by SEQ_LEN * MBS / TP_SIZE", + ) + parser.add_argument( + "--hidden-dim", + type=int, + default=None, + help="Hidden dimension to use", + ) + parser.add_argument( + "--output-dim", + type=int, + default=None, + help="Output dimension to use", + ) + parser.add_argument( + "--fwd-only", + action="store_true", + default=False, + help="Run forward pass only, default is both forward and backward passes", + ) + args = parser.parse_args() + + use_bias = False + + token_dim_list = [16384] + hidden_dim_list = [4096] + output_dim_list = [4096] + + if args.token_dim is not None: + token_dim_list = [args.token_dim] + + if args.hidden_dim is not None: + hidden_dim_list = [args.hidden_dim] + + if args.output_dim is not None: + output_dim_list = [args.output_dim] + + # MKN for linear + mkns = [] + for m in token_dim_list: + for k in hidden_dim_list: + for n in output_dim_list: + mkns.append((m, k, n)) + + # default recipes to run if not specified + recipe_list = ["bf16"] + + if args.recipe == "all": + recipe_list = ["bf16", "fp8_sub_channel", "mxfp8", "nvfp4"] + else: + recipe_list = [args.recipe] + + profiler_ctx = None + if args.profile: + hidden_dim_to_profile = 4096 if args.hidden_dim is None else args.hidden_dim + output_dim_to_profile = 4096 if args.output_dim is None else args.output_dim + token_dim_to_profile = 16384 if args.token_dim is None else args.token_dim + mkns = [(token_dim_to_profile, hidden_dim_to_profile, output_dim_to_profile)] + # in profile mode, only run one recipe specified in args.recipe + assert args.recipe != "all", ( + "In profile mode, only one recipe can be specified, please specify the recipe as" + " fp8_sub_channel, mxfp8, nvfp4, or bf16" + ) + recipe_list = [args.recipe] + profiler_ctx = torch.autograd.profiler.emit_nvtx(record_shapes=True) + profiler_ctx.__enter__() + + # Initialize a dataframe to store the results + df_linears = pd.DataFrame() + + # Run the fp8 benchmarks + for recipe_name in recipe_list: + assert recipe_name in [ + "bf16", + "fp8_sub_channel", + "mxfp8", + "nvfp4", + ], "Recipe must be one of bf16, fp8_sub_channel, mxfp8, or nvfp4" + if recipe_name == "mxfp8" and not mxfp8_available: + print(f"MXFP8 is not available, skipping {recipe_name}") + continue + if recipe_name == "fp8_sub_channel" and not fp8_block_scaling_available: + print(f"FP8 block scaling is not available, skipping {recipe_name}") + continue + if recipe_name == "nvfp4" and not nvfp4_available: + print(f"NVFP4 is not available, skipping {recipe_name}") + continue + + df = run_benchmark_linear( + mkns, + recipe_name, + use_bias, + fwd_only=args.fwd_only, + ) + df_linears = pd.concat([df_linears, df]) + + print(df_linears) + + if args.profile: + profiler_ctx.__exit__(None, None, None) diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py index 3c197bc6f3..c2f8e8de12 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -79,7 +79,7 @@ def reference_group_quantize( x: torch.Tensor, quantizers: list[MXFP8Quantizer], split_sections: list[int], - return_identity: bool, + return_rowwise: bool, return_transpose: bool, ) -> torch.Tensor: x_chunks = torch.split(x, split_sections) @@ -94,7 +94,7 @@ def reference_group_quantize( for i in range(len(x_chunks)): x_chunk = x_chunks[i] x_mxfp8_res = quantizers[i](x_chunk) - if return_identity: + if return_rowwise: x_qx.append(x_mxfp8_res._rowwise_data.view(dtype=torch.uint8)) x_sx.append(x_mxfp8_res._rowwise_scale_inv) else: @@ -133,7 +133,7 @@ def check_grouped_tensor_mxfp8_versus_reference( x_dtype: torch.dtype, M: int, N: int, - return_identity: bool, + return_rowwise: bool, return_transpose: bool, split_sections: list[int], optimize_for_gemm: bool = False, @@ -157,7 +157,7 @@ def check_grouped_tensor_mxfp8_versus_reference( quantizers = [ MXFP8Quantizer( fp8_dtype=te_dtype, - rowwise=return_identity, + rowwise=return_rowwise, columnwise=return_transpose, ) for _ in range(len(split_sections)) @@ -169,14 +169,14 @@ def check_grouped_tensor_mxfp8_versus_reference( grouped_quantizer.optimize_for_gemm = optimize_for_gemm x_qx_ref, x_sx_ref, x_qx_t_ref, x_sx_t_ref = reference_group_quantize( - x, quantizers, split_sections, return_identity, return_transpose + x, quantizers, split_sections, return_rowwise, return_transpose ) group_quantized_output = fused_grouped_quantize(x, split_section_tensor, grouped_quantizer) # get a list of MXFP8 quantized tensors for testing split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() - if return_identity: + if return_rowwise: x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] @@ -229,7 +229,7 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( x_dtype: torch.dtype, M: int, N: int, - return_identity: bool, + return_rowwise: bool, return_transpose: bool, split_sections: list[int], valid_M: int = None, @@ -258,7 +258,7 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( quantizers = [ MXFP8Quantizer( fp8_dtype=te_dtype, - rowwise=return_identity, + rowwise=return_rowwise, columnwise=return_transpose, ) for _ in range(len(split_sections)) @@ -270,7 +270,7 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( grouped_quantizer.optimize_for_gemm = optimize_for_gemm x_qx_ref, x_sx_ref, x_qx_t_ref, x_sx_t_ref = reference_group_quantize( - valid_x, quantizers, split_sections, return_identity, return_transpose + valid_x, quantizers, split_sections, return_rowwise, return_transpose ) # Note: for grouped quantize with paged stashing @@ -281,7 +281,7 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( # get a list of MXFP8 quantized tensors for testing split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() - if return_identity: + if return_rowwise: x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] @@ -355,9 +355,7 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( "random_uneven_split", ], ) -@pytest.mark.parametrize( - "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] ) @@ -372,14 +370,14 @@ def test_grouped_tensor_mxfp8_versus_reference( split_sections = generate_split_sections(M, N, edge_cases) - if quantize_mode == "quantize": - return_identity = True + if quantize_mode == "rowwise_only": + return_rowwise = True return_transpose = False - elif quantize_mode == "quantize_transpose": - return_identity = True + elif quantize_mode == "both_directions": + return_rowwise = True return_transpose = True - elif quantize_mode == "quantize_colwise_only": - return_identity = False + elif quantize_mode == "columnwise_only": + return_rowwise = False return_transpose = True else: raise ValueError(f"Invalid quantize mode: {quantize_mode}") @@ -388,7 +386,7 @@ def test_grouped_tensor_mxfp8_versus_reference( x_dtype=x_dtype, M=M, N=N, - return_identity=return_identity, + return_rowwise=return_rowwise, return_transpose=return_transpose, split_sections=split_sections, optimize_for_gemm=optimize_for_gemm, @@ -422,9 +420,7 @@ def test_grouped_tensor_mxfp8_versus_reference( "random_uneven_split", ], ) -@pytest.mark.parametrize( - "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] ) @@ -451,14 +447,14 @@ def test_grouped_tensor_mxfp8_with_paged_stashing( else: assert valid_M == M // 2, "valid_M must be M // 2 when edge_cases is not zero_tokens_all" - if quantize_mode == "quantize": - return_identity = True + if quantize_mode == "rowwise_only": + return_rowwise = True return_transpose = False - elif quantize_mode == "quantize_transpose": - return_identity = True + elif quantize_mode == "both_directions": + return_rowwise = True return_transpose = True - elif quantize_mode == "quantize_colwise_only": - return_identity = False + elif quantize_mode == "columnwise_only": + return_rowwise = False return_transpose = True else: raise ValueError(f"Invalid quantize mode: {quantize_mode}") @@ -467,7 +463,7 @@ def test_grouped_tensor_mxfp8_with_paged_stashing( x_dtype=x_dtype, M=M, N=N, - return_identity=return_identity, + return_rowwise=return_rowwise, return_transpose=return_transpose, split_sections=split_sections, valid_M=valid_M, diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py index 94ea699d14..6f0700809b 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py +++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py @@ -39,7 +39,7 @@ def check_mxfp8_quantize_swizzle_fusion( x_dtype: torch.dtype, M: int, N: int, - return_identity: bool, + return_rowwise: bool, return_transpose: bool, ) -> None: @@ -57,7 +57,7 @@ def check_mxfp8_quantize_swizzle_fusion( # Quantize quantizer = MXFP8Quantizer( fp8_dtype=te_dtype, - rowwise=return_identity, + rowwise=return_rowwise, columnwise=return_transpose, ) @@ -69,7 +69,7 @@ def check_mxfp8_quantize_swizzle_fusion( ) x_qx_ref, x_sx_ref, x_qx_t_ref, x_sx_t_ref = unpack_quantized_tensor(quantizer(x)) - if return_identity: + if return_rowwise: torch.testing.assert_close(x_qx_swf, x_qx_ref, atol=0.0, rtol=0.0) valid_scale_shape = get_mxfp8_scale_shape_no_padding(x.shape, False) assert valid_scale_shape == x_sx_swf.shape, ( @@ -103,9 +103,7 @@ def check_mxfp8_quantize_swizzle_fusion( ], ) @pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) -@pytest.mark.parametrize( - "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) def test_mxfp8_quantize_swizzle_fusion( x_dtype: torch.dtype, M: int, @@ -113,14 +111,14 @@ def test_mxfp8_quantize_swizzle_fusion( quantize_mode: str, ) -> None: - if quantize_mode == "quantize": - return_identity = True + if quantize_mode == "rowwise_only": + return_rowwise = True return_transpose = False - elif quantize_mode == "quantize_transpose": - return_identity = True + elif quantize_mode == "both_directions": + return_rowwise = True return_transpose = True - elif quantize_mode == "quantize_colwise_only": - return_identity = False + elif quantize_mode == "columnwise_only": + return_rowwise = False return_transpose = True else: raise ValueError(f"Invalid quantize mode: {quantize_mode}") @@ -129,6 +127,6 @@ def test_mxfp8_quantize_swizzle_fusion( x_dtype=x_dtype, M=M, N=N, - return_identity=return_identity, + return_rowwise=return_rowwise, return_transpose=return_transpose, ) diff --git a/tests/pytorch/nvfp4/nvfp4_utils.py b/tests/pytorch/nvfp4/nvfp4_utils.py index 5f1b5ac36c..757ed249d2 100644 --- a/tests/pytorch/nvfp4/nvfp4_utils.py +++ b/tests/pytorch/nvfp4/nvfp4_utils.py @@ -115,7 +115,7 @@ def reference_group_quantize( x: torch.Tensor, quantizers: list[NVFP4Quantizer], split_sections: list[int], - return_identity: bool, + return_rowwise: bool, return_transpose: bool, ) -> torch.Tensor: x_view = x.reshape(-1, x.size(-1)) @@ -133,7 +133,7 @@ def reference_group_quantize( for i in range(len(x_chunks)): x_chunk = x_chunks[i] x_nvfp4_res = quantizers[i](x_chunk) - if return_identity: + if return_rowwise: x_qx.append(x_nvfp4_res._rowwise_data.view(dtype=torch.uint8)) x_sx.append(x_nvfp4_res._rowwise_scale_inv) x_amax_rowwise.append(x_nvfp4_res._amax_rowwise) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index d4bf1fd3a1..7bf288fff7 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -37,7 +37,7 @@ def check_group_quantization_nvfp4_versus_reference( x_dtype: torch.dtype, M: int, N: int, - return_identity: bool, + return_rowwise: bool, return_transpose: bool, split_sections: list[int], with_rht: bool = True, @@ -63,7 +63,7 @@ def check_group_quantization_nvfp4_versus_reference( quantizers = [ NVFP4Quantizer( fp4_dtype=te_dtype, - rowwise=return_identity, + rowwise=return_rowwise, columnwise=return_transpose, with_amax_reduction=False, amax_reduction_group=None, @@ -74,12 +74,12 @@ def check_group_quantization_nvfp4_versus_reference( for _ in range(len(split_sections)) ] x_qx_ref, x_sx_ref, x_amax_rowwise_ref, x_qx_t_ref, x_sx_t_ref, x_amax_colwise_ref = ( - reference_group_quantize(x, quantizers, split_sections, return_identity, return_transpose) + reference_group_quantize(x, quantizers, split_sections, return_rowwise, return_transpose) ) split_quantize_outputs = tex.split_quantize(x, split_sections, quantizers) - if return_identity: + if return_rowwise: x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] x_amax_rowwise = [output._amax_rowwise for output in split_quantize_outputs] @@ -152,9 +152,7 @@ def check_group_quantization_nvfp4_versus_reference( "random_uneven_split", ], ) -@pytest.mark.parametrize( - "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] ) @@ -174,14 +172,14 @@ def test_rht_with_quantization_block_tiling_versus_reference( # currently disable pre-RHT amax with_post_rht_amax = with_rht - if quantize_mode == "quantize": - return_identity = True + if quantize_mode == "rowwise_only": + return_rowwise = True return_transpose = False - elif quantize_mode == "quantize_transpose": - return_identity = True + elif quantize_mode == "both_directions": + return_rowwise = True return_transpose = True - elif quantize_mode == "quantize_colwise_only": - return_identity = False + elif quantize_mode == "columnwise_only": + return_rowwise = False return_transpose = True else: raise ValueError(f"Invalid quantize mode: {quantize_mode}") @@ -190,7 +188,7 @@ def test_rht_with_quantization_block_tiling_versus_reference( x_dtype=x_dtype, M=M, N=N, - return_identity=return_identity, + return_rowwise=return_rowwise, return_transpose=return_transpose, split_sections=split_sections, with_rht=with_rht, 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 8d81d578a7..cf2ae50ee9 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -46,7 +46,7 @@ def check_grouped_tensor_nvfp4_versus_reference( x_dtype: torch.dtype, M: int, N: int, - return_identity: bool, + return_rowwise: bool, return_transpose: bool, split_sections: list[int], with_rht: bool = True, @@ -75,7 +75,7 @@ def check_grouped_tensor_nvfp4_versus_reference( quantizers = [ NVFP4Quantizer( fp4_dtype=te_dtype, - rowwise=return_identity, + rowwise=return_rowwise, columnwise=return_transpose, with_amax_reduction=False, amax_reduction_group=None, @@ -92,14 +92,14 @@ def check_grouped_tensor_nvfp4_versus_reference( grouped_quantizer.optimize_for_gemm = optimize_for_gemm x_qx_ref, x_sx_ref, x_amax_rowwise_ref, x_qx_t_ref, x_sx_t_ref, x_amax_colwise_ref = ( - reference_group_quantize(x, quantizers, split_sections, return_identity, return_transpose) + reference_group_quantize(x, quantizers, split_sections, return_rowwise, return_transpose) ) group_quantized_output = fused_grouped_quantize(x, split_section_tensor, grouped_quantizer) # get a list of nvfp4 quantized tensors for testing split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() - if return_identity: + if return_rowwise: x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] x_amax_rowwise = [output._amax_rowwise for output in split_quantize_outputs] @@ -162,7 +162,7 @@ def check_grouped_tensor_nvfp4_with_paged_stashing( x_dtype: torch.dtype, M: int, N: int, - return_identity: bool, + return_rowwise: bool, return_transpose: bool, split_sections: list[int], with_rht: bool = True, @@ -196,7 +196,7 @@ def check_grouped_tensor_nvfp4_with_paged_stashing( quantizers = [ NVFP4Quantizer( fp4_dtype=te_dtype, - rowwise=return_identity, + rowwise=return_rowwise, columnwise=return_transpose, with_amax_reduction=False, amax_reduction_group=None, @@ -214,7 +214,7 @@ def check_grouped_tensor_nvfp4_with_paged_stashing( x_qx_ref, x_sx_ref, x_amax_rowwise_ref, x_qx_t_ref, x_sx_t_ref, x_amax_colwise_ref = ( reference_group_quantize( - valid_x, quantizers, split_sections, return_identity, return_transpose + valid_x, quantizers, split_sections, return_rowwise, return_transpose ) ) @@ -226,7 +226,7 @@ def check_grouped_tensor_nvfp4_with_paged_stashing( # get a list of nvfp4 quantized tensors for testing split_quantize_outputs = group_quantized_output.split_into_quantized_tensors() - if return_identity: + if return_rowwise: x_qx = [output._rowwise_data.view(dtype=torch.uint8) for output in split_quantize_outputs] x_sx = [output._rowwise_scale_inv for output in split_quantize_outputs] x_amax_rowwise = [output._amax_rowwise for output in split_quantize_outputs] @@ -307,9 +307,7 @@ def check_grouped_tensor_nvfp4_with_paged_stashing( "random_uneven_split", ], ) -@pytest.mark.parametrize( - "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] ) @@ -333,14 +331,14 @@ def test_grouped_tensor_nvfp4_versus_reference( # currently disable pre-RHT amax with_post_rht_amax = with_rht - if quantize_mode == "quantize": - return_identity = True + if quantize_mode == "rowwise_only": + return_rowwise = True return_transpose = False - elif quantize_mode == "quantize_transpose": - return_identity = True + elif quantize_mode == "both_directions": + return_rowwise = True return_transpose = True - elif quantize_mode == "quantize_colwise_only": - return_identity = False + elif quantize_mode == "columnwise_only": + return_rowwise = False return_transpose = True else: raise ValueError(f"Invalid quantize mode: {quantize_mode}") @@ -349,7 +347,7 @@ def test_grouped_tensor_nvfp4_versus_reference( x_dtype=x_dtype, M=M, N=N, - return_identity=return_identity, + return_rowwise=return_rowwise, return_transpose=return_transpose, split_sections=split_sections, with_rht=with_rht, @@ -386,9 +384,7 @@ def test_grouped_tensor_nvfp4_versus_reference( "random_uneven_split", ], ) -@pytest.mark.parametrize( - "quantize_mode", ["quantize", "quantize_transpose", "quantize_colwise_only"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "with_random_sign_mask", [True, False], ids=["with_random_sign_mask", "no_random_sign_mask"] ) @@ -424,14 +420,14 @@ def test_grouped_tensor_nvfp4_with_paged_stashing( # currently disable pre-RHT amax with_post_rht_amax = with_rht - if quantize_mode == "quantize": - return_identity = True + if quantize_mode == "rowwise_only": + return_rowwise = True return_transpose = False - elif quantize_mode == "quantize_transpose": - return_identity = True + elif quantize_mode == "both_directions": + return_rowwise = True return_transpose = True - elif quantize_mode == "quantize_colwise_only": - return_identity = False + elif quantize_mode == "columnwise_only": + return_rowwise = False return_transpose = True else: raise ValueError(f"Invalid quantize mode: {quantize_mode}") @@ -440,7 +436,7 @@ def test_grouped_tensor_nvfp4_with_paged_stashing( x_dtype=x_dtype, M=M, N=N, - return_identity=return_identity, + return_rowwise=return_rowwise, return_transpose=return_transpose, split_sections=split_sections, with_rht=with_rht, diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 80ccb2f23d..bf3f545b8b 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -147,9 +147,7 @@ def check_quantization_nvfp4_versus_reference( ], ) @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize( - "return_transpose", [True, False], ids=["quantize_transpose", "skip_transpose"] -) +@pytest.mark.parametrize("return_transpose", [True, False], ids=["both_directions", "rowwise_only"]) @pytest.mark.parametrize("swizzled_scale", [False], ids=["linear_scale"]) @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] @@ -186,9 +184,7 @@ def test_quantization_block_tiling_versus_reference( ) @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("extrema_high", [False, True], ids=["zeros", "maxes"]) -@pytest.mark.parametrize( - "return_transpose", [True, False], ids=["quantize_transpose", "skip_transpose"] -) +@pytest.mark.parametrize("return_transpose", [True, False], ids=["both_directions", "rowwise_only"]) @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @@ -286,9 +282,7 @@ def test_nvfp4_quantization_extrema_versus_reference( ], ) @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize( - "return_transpose", [True, False], ids=["quantize_transpose", "skip_transpose"] -) +@pytest.mark.parametrize("return_transpose", [True, False], ids=["both_directions", "rowwise_only"]) @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @@ -399,9 +393,7 @@ def test_nvfp4_quantization_boundary_values( ], ) @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) -@pytest.mark.parametrize( - "return_transpose", [True, False], ids=["quantize_transpose", "skip_transpose"] -) +@pytest.mark.parametrize("return_transpose", [True, False], ids=["both_directions", "rowwise_only"]) @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index 98be9a4f54..795721df04 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -35,6 +35,7 @@ def check_quantization_nvfp4_versus_reference( M: int, N: int, contiguous: bool, + return_rowwise: bool, return_transpose: bool, use_cpp_allocator: bool, swizzled_scale: bool = False, @@ -61,7 +62,7 @@ def check_quantization_nvfp4_versus_reference( # Quantize nvfp4_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, - rowwise=True, + rowwise=return_rowwise, columnwise=return_transpose, with_amax_reduction=False, amax_reduction_group=None, @@ -78,9 +79,11 @@ def check_quantization_nvfp4_versus_reference( x_nvfp4_sut = nvfp4_quantizer.update_quantized(x, x_nvfp4_sut) # Extract data from NVFP4Tensor - assert x_nvfp4_sut._rowwise_data is not None - qx: torch.Tensor = x_nvfp4_sut._rowwise_data.view(dtype=torch.uint8) - assert x_nvfp4_sut._rowwise_scale_inv is not None + qx: torch.Tensor = ( + x_nvfp4_sut._rowwise_data.view(dtype=torch.uint8) + if x_nvfp4_sut._rowwise_data is not None + else None + ) sx: torch.Tensor = x_nvfp4_sut._rowwise_scale_inv qx_t = ( x_nvfp4_sut._columnwise_data.view(dtype=torch.uint8) @@ -91,13 +94,13 @@ def check_quantization_nvfp4_versus_reference( amax_rowwise = x_nvfp4_sut._amax_rowwise amax_colwise = x_nvfp4_sut._amax_columnwise - qx = unpack_fp4(qx) + qx = unpack_fp4(qx) if qx is not None else None qx_t = unpack_fp4(qx_t) if qx_t is not None else None # Reference quantization using NVFP4QuantizerRef with built-in RHT ref_quantizer = NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, - rowwise=True, + rowwise=return_rowwise, columnwise=return_transpose, pow_2_scales=False, eps=0.0, @@ -130,13 +133,14 @@ def check_quantization_nvfp4_versus_reference( sx_t_ref = None ref_amax_colwise_t = None - torch.testing.assert_close(amax_rowwise, ref_amax_rowwise, atol=0.0, rtol=0.0) + if return_rowwise: + torch.testing.assert_close(amax_rowwise, ref_amax_rowwise, atol=0.0, rtol=0.0) - torch.testing.assert_close(qx, qx_ref, atol=0.0, rtol=0.0) - # Compare only the valid portion of scale tensors (reference may not have padding) - ref_sx_shape = sx_ref.shape - sx_valid = sx[: ref_sx_shape[0], : ref_sx_shape[1]] - torch.testing.assert_close(sx_valid, sx_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(qx, qx_ref, atol=0.0, rtol=0.0) + # Compare only the valid portion of scale tensors (reference may not have padding) + ref_sx_shape = sx_ref.shape + sx_valid = sx[: ref_sx_shape[0], : ref_sx_shape[1]] + torch.testing.assert_close(sx_valid, sx_ref, atol=0.0, rtol=0.0) if return_transpose: torch.testing.assert_close(amax_colwise, ref_amax_colwise_t, atol=0.0, rtol=0.0) @@ -184,9 +188,7 @@ def check_quantization_nvfp4_versus_reference( ], ) @pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) -@pytest.mark.parametrize( - "return_transpose", [True, False], ids=["quantize_transpose", "skip_transpose"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @@ -197,15 +199,29 @@ def test_rht_with_quantization_block_tiling_versus_reference( x_dtype: torch.dtype, M: int, N: int, - return_transpose: bool, + quantize_mode: str, use_cpp_allocator: bool, with_random_sign_mask: bool, ) -> None: + + if quantize_mode == "rowwise_only": + return_rowwise = True + return_transpose = False + elif quantize_mode == "both_directions": + return_rowwise = True + return_transpose = True + elif quantize_mode == "columnwise_only": + return_rowwise = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + check_quantization_nvfp4_versus_reference( x_dtype=x_dtype, M=M, N=N, contiguous=True, + return_rowwise=return_rowwise, return_transpose=return_transpose, use_cpp_allocator=use_cpp_allocator, with_random_sign_mask=with_random_sign_mask, @@ -220,9 +236,7 @@ def test_rht_with_quantization_block_tiling_versus_reference( ], ) @pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str) -@pytest.mark.parametrize( - "return_transpose", [True, False], ids=["quantize_transpose", "skip_transpose"] -) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @@ -233,15 +247,29 @@ def test_nvfp4_quantization_noncontiguous_inputs( x_dtype: torch.dtype, M: int, N: int, - return_transpose: bool, + quantize_mode: str, use_cpp_allocator: bool, with_random_sign_mask: bool, ): + + if quantize_mode == "rowwise_only": + return_rowwise = True + return_transpose = False + elif quantize_mode == "both_directions": + return_rowwise = True + return_transpose = True + elif quantize_mode == "columnwise_only": + return_rowwise = False + return_transpose = True + else: + raise ValueError(f"Invalid quantize mode: {quantize_mode}") + check_quantization_nvfp4_versus_reference( x_dtype=x_dtype, M=M, N=N, contiguous=False, + return_rowwise=return_rowwise, return_transpose=return_transpose, use_cpp_allocator=use_cpp_allocator, with_random_sign_mask=with_random_sign_mask, diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index b3d48f68bd..b9e2b907e0 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -176,6 +176,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources hadamard_transform/graph_safe_group_hadamard_transform.cu hadamard_transform/hadamard_transform.cu hadamard_transform/hadamard_transform_cast_fusion.cu + hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu hadamard_transform/group_hadamard_transform_cast_fusion.cu hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index 8d2d806559..792b068cbc 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -47,7 +47,8 @@ __device__ __forceinline__ nvfp4_scale_t compute_decoding_scaling_factor(const f // However, this is part of the emulation code to ensure exact match. using namespace detail; constexpr float fp4_max = TypeExtrema::max; // 6.0f; - const float S_dec_b = block_amax / fp4_max * S_enc; + constexpr float fp4_max_inv = 1.0f / fp4_max; + const float S_dec_b = block_amax * (S_enc * fp4_max_inv); return static_cast(fminf(S_dec_b, TypeExtrema::max)); } #endif // FP4_TYPE_SUPPORTED @@ -59,11 +60,12 @@ namespace quantization_SF { // Compute per-block E4M3 encoding/decoding scaling factor __device__ __forceinline__ fp8e4m3 compute_decoding_scaling_factor(const float block_amax, const float S_enc) { - constexpr float rcp_6f = 1.0f / 6.0f; + using namespace detail; + constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; // 1 / 6.0f // const float S_dec_b = block_amax * rcp_6f; // const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); // return S_dec_b_fp8; - return static_cast(block_amax * rcp_6f * S_enc); + return static_cast(block_amax * (S_enc * fp4_max_inv)); } #endif // FP4_TYPE_SUPPORTED } // namespace quantization_SF 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 6f3cf90d90..0c3a5e9299 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 @@ -193,957 +193,933 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device_g // Abort immediately if compilation is not supported constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; if constexpr (!is_blackwell_arch) { - NVTE_DEVICE_ERROR( - "group_row_col_rht_gemm_device_graph_safe is only supported on Blackwell " - "with architecture-specific compilation. " - "Try recompiling with sm_100a or similar."); + NVTE_DEVICE_ERROR("RHT fusion is only supported on Blackwell."); return; - } - static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, - "group_row_col_rht_gemm_device_graph_safe must generate row-wise " - "and/or column-wise output."); + } else { + static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, + "group_row_col_rht_gemm_device_graph_safe must generate row-wise " + "and/or column-wise output."); #if !defined(CUTLASS_ARCH_CLC_ENABLED) - CUTLASS_NOT_IMPLEMENTED(); - return; + CUTLASS_NOT_IMPLEMENTED(); + return; #endif - using X = Underscore; - // Accumulator data type for main computation - using ElementAccumulator = float; - static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); - using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; - static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( - size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); - static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; - static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; - static constexpr bool kEnableRowQuant = kEnableRowQuant_; - static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_; - static constexpr bool kUseFastMath = kUseFastMath_; - - // Constant for RHT tensor processing (tile size etc) - static int constexpr RhtTensorSize = 16; - - // Get the total number of tokens to process - // Note that here M is the hidden size, which is the last logical dimension of the input tensor x - // The kernel is designed in column major, so M is the hidden size - size_t sum_token_dims = offsets[num_tensors] / M; - - // Transaction bytes for TMA transfer on RHT tensor blocks - static int constexpr kTmaRhtTensorTransactionBytes = - cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); - static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; - static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; - - // Mainloop pipeline stage calculation, vectorization parameters for scaling factors - static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); - static int constexpr SFVecSize = 16; - // Swizzle output layout for scaling factor arrays - using SwizzledSFALayoutAtom = - cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; - using SwizzledSFDLayoutAtom = - cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; - - // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling - using MainloopPipeline = - cutlass::detail::CustomizedPipelineTmaUmmaAsync; - using MainloopPipelineState = typename MainloopPipeline::PipelineState; - using SchedPipeline = cutlass::PipelineCLCFetchAsync; - using SchedPipelineState = typename SchedPipeline::PipelineState; - using SchedThrottlePipeline = cutlass::PipelineAsync; - using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState; - - static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, "ClusterShape must be Shape<_1,_1,_1>"); - - using TmemAllocator = cute::TMEM::Allocator1Sm; - static int constexpr VectorSize = RhtTensorSize; - - // Compile-time safety: static shapes required for shared memory layouts - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - // CUTE_STATIC_ASSERT(is_static::value); - - auto cluster_size = size<0>(cluster_shape); - auto mainloop_tiler = Shape<_128, _16, _128>{}; - auto epilogue_tiler = Shape<_128, _128, _128>{}; - - static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); - - // Get the appropriate blocks for this Cluster - dim3 cluster_coord_in_grid = cluster_id_in_grid(); - - // Total number of k-tiles - int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler); - - struct TileScheduler { - uint32_t tiles_in_m = 0; - uint32_t tiles_in_n = 0; - uint32_t linear_idx = 0; - uint32_t next_linear_idx = 0; - uint32_t start_idx = 0; - uint32_t tile_m_idx = 0; - uint32_t tile_n_idx = 0; - int k_tile_max = 0; - uint32_t *atomic_tile_index_; - uint32_t *smem_tile_counter; - uint32_t atomic_offset; - cutlass::FastDivmodU64 divmod_tiles_in_m; - - CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, - uint32_t *atomic_tile_index, uint32_t *smem_tile_counter) - : tiles_in_m(tiles_m), - tiles_in_n(tiles_n), - linear_idx(blockIdx.x), - next_linear_idx(blockIdx.x), - start_idx(blockIdx.x), - k_tile_max(kmax), - atomic_tile_index_(atomic_tile_index), - smem_tile_counter(smem_tile_counter), - atomic_offset(gridDim.x), - divmod_tiles_in_m(uint64_t(tiles_m)) { - update_tile_idx(); - } - CUTLASS_DEVICE void update_tile_idx() { - uint64_t q, r; - divmod_tiles_in_m(q, r, uint64_t(linear_idx)); - tile_m_idx = static_cast(r); - tile_n_idx = static_cast(q) * uint32_t(k_tile_max); - } - CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; } - CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; } - CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; } - - CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; } + using X = Underscore; + // Accumulator data type for main computation + using ElementAccumulator = float; + static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); + static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; + static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; + static constexpr bool kEnableRowQuant = kEnableRowQuant_; + static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_; + static constexpr bool kUseFastMath = kUseFastMath_; + + // Constant for RHT tensor processing (tile size etc) + static int constexpr RhtTensorSize = 16; + + // Get the total number of tokens to process + // Note that here M is the hidden size, which is the last logical dimension of the input tensor x + // The kernel is designed in column major, so M is the hidden size + size_t sum_token_dims = offsets[num_tensors] / M; + + // Transaction bytes for TMA transfer on RHT tensor blocks + static int constexpr kTmaRhtTensorTransactionBytes = + cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; + + // Mainloop pipeline stage calculation, vectorization parameters for scaling factors + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + static int constexpr SFVecSize = 16; + // Swizzle output layout for scaling factor arrays + using SwizzledSFALayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + + // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling + using MainloopPipeline = + cutlass::detail::CustomizedPipelineTmaUmmaAsync; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + using SchedPipeline = cutlass::PipelineCLCFetchAsync; + using SchedPipelineState = typename SchedPipeline::PipelineState; + using SchedThrottlePipeline = cutlass::PipelineAsync; + using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState; + + static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, "ClusterShape must be Shape<_1,_1,_1>"); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static int constexpr VectorSize = RhtTensorSize; + + // Compile-time safety: static shapes required for shared memory layouts + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + // CUTE_STATIC_ASSERT(is_static::value); + + auto cluster_size = size<0>(cluster_shape); + auto mainloop_tiler = Shape<_128, _16, _128>{}; + auto epilogue_tiler = Shape<_128, _128, _128>{}; + + static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler); + + struct TileScheduler { + uint32_t tiles_in_m = 0; + uint32_t tiles_in_n = 0; + uint32_t linear_idx = 0; + uint32_t next_linear_idx = 0; + uint32_t start_idx = 0; + uint32_t tile_m_idx = 0; + uint32_t tile_n_idx = 0; + int k_tile_max = 0; + uint32_t *atomic_tile_index_; + uint32_t *smem_tile_counter; + uint32_t atomic_offset; + cutlass::FastDivmodU64 divmod_tiles_in_m; + + CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, + uint32_t *atomic_tile_index, uint32_t *smem_tile_counter) + : tiles_in_m(tiles_m), + tiles_in_n(tiles_n), + linear_idx(blockIdx.x), + next_linear_idx(blockIdx.x), + start_idx(blockIdx.x), + k_tile_max(kmax), + atomic_tile_index_(atomic_tile_index), + smem_tile_counter(smem_tile_counter), + atomic_offset(gridDim.x), + divmod_tiles_in_m(uint64_t(tiles_m)) { + update_tile_idx(); + } + CUTLASS_DEVICE void update_tile_idx() { + uint64_t q, r; + divmod_tiles_in_m(q, r, uint64_t(linear_idx)); + tile_m_idx = static_cast(r); + tile_n_idx = static_cast(q) * uint32_t(k_tile_max); + } + CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; } + CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; } + CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; } - CUTLASS_DEVICE bool is_valid() const { - return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()), - cute::make_coord(tiles_in_m, tiles_in_n)); - } + CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; } - CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; } + CUTLASS_DEVICE bool is_valid() const { + return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()), + cute::make_coord(tiles_in_m, tiles_in_n)); + } - CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; } + CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; } - // Fetch a new tile_id using atomics. - CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) { - uint32_t tile_id_counter = 0; - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.eq.u32 p, %2, 1;\n\t" - "@p atom.global.add.u32 %0, [%1], 1; \n\t" - "}" - : "=r"(tile_id_counter) - : "l"(atomic_tile_index_), "r"(pred)); + CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; } - return tile_id_counter; - } + // Fetch a new tile_id using atomics. + CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) { + uint32_t tile_id_counter = 0; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p atom.global.add.u32 %0, [%1], 1; \n\t" + "}" + : "=r"(tile_id_counter) + : "l"(atomic_tile_index_), "r"(pred)); - CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline, - SchedPipelineState sched_pipeline_consumer_state) { - sched_pipeline.consumer_wait(sched_pipeline_consumer_state); - next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()]; - cutlass::arch::fence_view_async_shared(); - sched_pipeline.consumer_release(sched_pipeline_consumer_state); - return; - } + return tile_id_counter; + } - CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline, - SchedPipelineState sched_pipeline_producer_state) { - uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state); - // Wait for clcID buffer to become empty with a flipped phase - sched_pipeline.producer_acquire(sched_pipeline_producer_state); - auto is_leading_thread = cute::elect_one_sync(); - uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset; - uint32_t smem_addr = - cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]); - if (is_leading_thread) { - cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0); + CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_consumer_state) { + sched_pipeline.consumer_wait(sched_pipeline_consumer_state); + next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()]; + cutlass::arch::fence_view_async_shared(); + sched_pipeline.consumer_release(sched_pipeline_consumer_state); + return; } - ++sched_pipeline_producer_state; - return sched_pipeline_producer_state; - } + CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_producer_state) { + uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state); + // Wait for clcID buffer to become empty with a flipped phase + sched_pipeline.producer_acquire(sched_pipeline_producer_state); + auto is_leading_thread = cute::elect_one_sync(); + uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset; + uint32_t smem_addr = + cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]); + if (is_leading_thread) { + cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0); + } - CUTLASS_DEVICE auto update_work_tile_info() { - linear_idx = next_linear_idx; - update_tile_idx(); - return; - } - }; - - // Allocate and alias shared memory to the kernel's shared storage type - extern __shared__ char shared_memory[]; - using SharedStorage = - SharedStorage; - SharedStorage &shared_storage = *reinterpret_cast(shared_memory); - - // Compute the number of tiles in M and N after tiling and assign scheduler - uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); - uint32_t tiles_in_n = uint32_t(size(ceil_div(sum_token_dims, size<2>(epilogue_tiler)))); - - TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace, - shared_storage.atomic_tile_counter); - - int block_rank_in_cluster = cute::block_rank_in_cluster(); - - // Shapes for accumulated tiles in mainloop and epilogue - auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{}); - auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{}); - - // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended - auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); - auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); - - // Number of threads assigned for various epilogue roles depending on quantization settings - static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; - static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; - static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0; - static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0; - static int constexpr NumSchedThreads = 32; - static int constexpr NumMainloopLoadThreads = 32; - static int constexpr NumEpilogueThreads = - NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; - - TmemAllocator tmem_allocator{}; - cutlass::arch::NamedBarrier tmem_allocation_result_barrier( - NumMmaThreadCount + NumEpilogueColQuantThreadCount, - cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); - - int warp_idx = cutlass::canonical_warp_idx_sync(); - - // warp assignment - bool is_mma_warp = (warp_idx == 0); - bool is_dma_warp = (warp_idx == 1); - bool is_sched_warp = (warp_idx == 2); - bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); - bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); - - typename MainloopPipeline::Params mainloop_pipeline_params; - if (is_dma_warp) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; - } - if (is_mma_warp) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; - } - mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; - mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; - mainloop_pipeline_params.initializing_warp = 0; - mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + ++sched_pipeline_producer_state; + return sched_pipeline_producer_state; + } - MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, - cluster_shape, cute::true_type{}, // Perform barrier init - cute::true_type{}); // Delay mask calculation + CUTLASS_DEVICE auto update_work_tile_info() { + linear_idx = next_linear_idx; + update_tile_idx(); + return; + } + }; - MainloopPipelineState mainloop_pipe_consumer_state; - MainloopPipelineState mainloop_pipe_producer_state = - cutlass::make_producer_start_state(); + // Allocate and alias shared memory to the kernel's shared storage type + extern __shared__ char shared_memory[]; + using SharedStorage = + SharedStorage; + SharedStorage &shared_storage = *reinterpret_cast(shared_memory); - using AccumulatorPipeline = - cutlass::PipelineUmmaAsync; - using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; - using AccumulatorPipelineInitBarriers = cute::bool_constant; + // Compute the number of tiles in M and N after tiling and assign scheduler + uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); + uint32_t tiles_in_n = uint32_t(size(ceil_div(sum_token_dims, size<2>(epilogue_tiler)))); - AccumulatorPipelineState accumulator_pipe_consumer_state; - AccumulatorPipelineState accumulator_pipe_producer_state = - cutlass::make_producer_start_state(); + TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace, + shared_storage.atomic_tile_counter); - typename AccumulatorPipeline::Params accumulator_pipeline_params; - if (is_mma_warp) { - accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; - } - if (is_epilogue_col_quant_warp) { - accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; - } - // Only one producer thread arrives on this barrier. - accumulator_pipeline_params.producer_arv_count = 1; - accumulator_pipeline_params.consumer_arv_count = - size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; - accumulator_pipeline_params.initializing_warp = 1; - AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, accumulator_pipeline_params, - cluster_shape, AccumulatorPipelineInitBarriers{}, - cute::true_type{}); // Delay mask calculation - typename SchedPipeline::Params sched_pipeline_params; - if (is_sched_warp) { - sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer; - } else { - sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer; - } - sched_pipeline_params.producer_blockid = 0; - sched_pipeline_params.producer_arv_count = 1; - sched_pipeline_params.consumer_arv_count = - NumSchedThreads + - cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); - sched_pipeline_params.transaction_bytes = sizeof(uint32_t); - sched_pipeline_params.initializing_warp = 3; - SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape); - SchedPipelineState sched_pipeline_consumer_state; - SchedPipelineState sched_pipeline_producer_state = - cutlass::make_producer_start_state(); - - typename SchedThrottlePipeline::Params sched_throttle_pipeline_params; - if (is_dma_warp) { - sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer; - } - if (is_sched_warp) { - sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer; - } - sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; - sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; - sched_throttle_pipeline_params.dst_blockid = 0; - sched_throttle_pipeline_params.initializing_warp = 4; - - SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle, - sched_throttle_pipeline_params); - SchedThrottlePipelineState sched_pipeline_throttle_consumer_state; - SchedThrottlePipelineState sched_pipeline_throttle_producer_state = - cutlass::make_producer_start_state(); - - if (warp_idx == 2 && elect_one_sync()) { - cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); - } - __syncthreads(); - - // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer - if (is_dma_warp) { - // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access). - cutlass::arch::warpgroup_reg_dealloc<32>(); - // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory. - Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N)); - Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); - - // Partition tensors for tiling according to the mainloop and cluster tilers. - Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); - Tensor gB_nk = - local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) - - // Shared memory tensors for pipeline - Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), - sAlayout); // (MMA,MMA_M,MMA_N,PIPE) - Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), - sBlayout); // (MMA,MMA_N,MMA_K,PIPE) - - // Determine warp/tile positioning int block_rank_in_cluster = cute::block_rank_in_cluster(); - ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx - // Partition global to local fragments for A and B - Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) - Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) - - Layout cta_layout_mnk = make_layout(cluster_shape); - Layout cta_layout_vmnk = - tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); - auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); - - auto [tAgA, tAsA] = - tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), - group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); - - auto [tBgB, tBsB] = - tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), - group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); - - uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); - uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); - if constexpr (kEnableRHTColQuant) { - if (elect_one_sync()) { - cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], - kTmaRhtTensorTransactionBytes); - copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), - tBsB(_, 0)); - } - } - do { - // is_first_wave indicates whether this scheduler wave is the first among a group. - bool is_first_wave = scheduler.is_first_wave(); - uint32_t skip_wait = is_first_wave; - auto tAgA_mk = tAgA(_, scheduler.tile_m(), _); - int k_tile = 0; - - sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state); - sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state); - ++sched_pipeline_throttle_producer_state; - CUTLASS_PRAGMA_NO_UNROLL - while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { - int k_tile_idx_n = scheduler.tile_n_base() + k_tile; - ++k_tile; - skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); - mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); - using BarrierType = typename MainloopPipeline::ProducerBarrierType; - BarrierType *tma_barrier = - mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); - int write_stage = mainloop_pipe_producer_state.index(); - ++mainloop_pipe_producer_state; - if (cute::elect_one_sync()) { - copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), - tAsA(_, write_stage)); - } - } - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - scheduler.update_work_tile_info(); - // scheduler.advance(); - } while (scheduler.is_valid()); - mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); - } else if (is_mma_warp) { - // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform. - cutlass::arch::warpgroup_reg_dealloc<32>(); - if constexpr (kEnableRHTColQuant) { - // Setup shared memory fragments for A and B tiles. + // Shapes for accumulated tiles in mainloop and epilogue + auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{}); + auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{}); + + // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended + auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); + auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); + + // Number of threads assigned for various epilogue roles depending on quantization settings + static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; + static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; + static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0; + static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0; + static int constexpr NumSchedThreads = 32; + static int constexpr NumMainloopLoadThreads = 32; + static int constexpr NumEpilogueThreads = + NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + NumMmaThreadCount + NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + // warp assignment + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_sched_warp = (warp_idx == 2); + bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); + bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + mainloop_pipeline_params.num_consumers = + NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + + MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, + cluster_shape, cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = + cutlass::make_producer_start_state(); + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + using AccumulatorPipelineInitBarriers = cute::bool_constant; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = + cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_col_quant_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = + size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; + accumulator_pipeline_params.initializing_warp = 1; + AccumulatorPipeline accumulator_pipeline( + shared_storage.accumulator, accumulator_pipeline_params, cluster_shape, + AccumulatorPipelineInitBarriers{}, cute::true_type{}); // Delay mask calculation + typename SchedPipeline::Params sched_pipeline_params; + if (is_sched_warp) { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer; + } else { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer; + } + sched_pipeline_params.producer_blockid = 0; + sched_pipeline_params.producer_arv_count = 1; + sched_pipeline_params.consumer_arv_count = + NumSchedThreads + + cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); + sched_pipeline_params.transaction_bytes = sizeof(uint32_t); + sched_pipeline_params.initializing_warp = 3; + SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape); + SchedPipelineState sched_pipeline_consumer_state; + SchedPipelineState sched_pipeline_producer_state = + cutlass::make_producer_start_state(); + + typename SchedThrottlePipeline::Params sched_throttle_pipeline_params; + if (is_dma_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer; + } + if (is_sched_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer; + } + sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; + sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; + sched_throttle_pipeline_params.dst_blockid = 0; + sched_throttle_pipeline_params.initializing_warp = 4; + + SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle, + sched_throttle_pipeline_params); + SchedThrottlePipelineState sched_pipeline_throttle_consumer_state; + SchedThrottlePipelineState sched_pipeline_throttle_producer_state = + cutlass::make_producer_start_state(); + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + + // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer + if (is_dma_warp) { + // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access). + cutlass::arch::warpgroup_reg_dealloc<32>(); + // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory. + Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N)); + Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); + + // Partition tensors for tiling according to the mainloop and cluster tilers. + Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gB_nk = + local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) + + // Shared memory tensors for pipeline Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE) Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + // Determine warp/tile positioning int block_rank_in_cluster = cute::block_rank_in_cluster(); ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx - // Allocate "fragments" -- these are actually umma smem descriptors - Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) - Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) - - mma.accumulate_ = UMMA::ScaleOut::Zero; - - tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, - &shared_storage.tmem_base_ptr); - __syncwarp(); - tmem_allocation_result_barrier.arrive(); - uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; - bulk_tmem_mma.data() = tmem_base_ptr; - // Wait until the B (Hadamard) tensor copy is complete - cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); - do { - uint32_t skip_wait = K_TILE_MAX <= 0; + // Partition global to local fragments for A and B + Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) + Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = + tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = + tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); + + auto [tBgB, tBsB] = + tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + if constexpr (kEnableRHTColQuant) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], + kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), + tBsB(_, 0)); + } + } - auto barrier_token = - mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + do { + // is_first_wave indicates whether this scheduler wave is the first among a group. + bool is_first_wave = scheduler.is_first_wave(); + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_, scheduler.tile_m(), _); + int k_tile = 0; + + sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state); + sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state); + ++sched_pipeline_throttle_producer_state; + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { + int k_tile_idx_n = scheduler.tile_n_base() + k_tile; + ++k_tile; + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType *tma_barrier = + mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + if (cute::elect_one_sync()) { + copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), + tAsA(_, write_stage)); + } + } scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); ++sched_pipeline_consumer_state; - CUTLASS_PRAGMA_NO_UNROLL - for (int k_tile = 0; - k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { - mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); - int read_stage = mainloop_pipe_consumer_state.index(); - auto tCrA_mk = tCrA(_, _, _, read_stage); - auto tCrB_nk = tCrB(_, _, 0, 0); - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) { - int accumulator_k_block = - accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; - int tCrA_k_block = k_block * EpilogueUnrollFactor; - accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + scheduler.update_work_tile_info(); + // scheduler.advance(); + } while (scheduler.is_valid()); + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform. + cutlass::arch::warpgroup_reg_dealloc<32>(); + if constexpr (kEnableRHTColQuant) { + // Setup shared memory fragments for A and B tiles. + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Allocate "fragments" -- these are actually umma smem descriptors + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + // Wait until the B (Hadamard) tensor copy is complete + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + + auto barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_, _, _, read_stage); + auto tCrB_nk = tCrB(_, _, 0, 0); CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < EpilogueUnrollFactor; i++) { - auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i); - gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators); + for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) { + int accumulator_k_block = + accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; + int tCrA_k_block = k_block * EpilogueUnrollFactor; + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < EpilogueUnrollFactor; i++) { + auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i); + gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; } - - accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); - ++accumulator_pipe_producer_state; + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); + barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); } - auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; - ++mainloop_pipe_consumer_state; - ++k_tile; - skip_wait = k_tile >= K_TILE_MAX; - mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); - barrier_token = - mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); - } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + } else if (is_sched_warp) { + // Scheduler warp manages tile assignment and pipeline progress for warps + cutlass::arch::warpgroup_reg_dealloc<32>(); + do { + sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state); + sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state); + ++sched_pipeline_throttle_consumer_state; + sched_pipeline_producer_state = + scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; scheduler.update_work_tile_info(); } while (scheduler.is_valid()); - tmem_allocator.release_allocation_lock(); - accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); - tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); - } - } else if (is_sched_warp) { - // Scheduler warp manages tile assignment and pipeline progress for warps - cutlass::arch::warpgroup_reg_dealloc<32>(); - do { - sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state); - sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state); - ++sched_pipeline_throttle_consumer_state; - sched_pipeline_producer_state = - scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state); - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - scheduler.update_work_tile_info(); - } while (scheduler.is_valid()); - } else if (is_epilogue_col_quant_warp) { - // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage, - // and writing result tensors/scales to global memory. - cutlass::arch::warpgroup_reg_alloc<192>(); - if constexpr (kEnableRHTColQuant) { - using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; - - auto acc_epilogue_pipelined_shape = - append(acc_shape_epilogue, Int{}); - auto bulk_tmem_epilogue_layout = make_layout( - acc_epilogue_pipelined_shape, - make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler))); - auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); - - // Use 256-bit fragments for aligned bulk stores - static int constexpr FragmentSize = 256 / sizeof_bits_v; - - // Wait for TMEM allocation for this pipeline to finish - tmem_allocation_result_barrier.arrive_and_wait(); - uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; - bulk_tmem_epilogue.data() = tmem_base_ptr; - int global_thread_idx = threadIdx.x; - int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; - // g2s load all global_d_amax - CUTLASS_PRAGMA_NO_UNROLL - for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) { - shared_storage.global_d_amax[g] = __ldg(reinterpret_cast(amax_colwise + g)); - } + } else if (is_epilogue_col_quant_warp) { + // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage, + // and writing result tensors/scales to global memory. + cutlass::arch::warpgroup_reg_alloc<192>(); + if constexpr (kEnableRHTColQuant) { + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + auto acc_epilogue_pipelined_shape = + append(acc_shape_epilogue, Int{}); + auto bulk_tmem_epilogue_layout = make_layout( + acc_epilogue_pipelined_shape, + make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler))); + auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); + + // Use 256-bit fragments for aligned bulk stores + static int constexpr FragmentSize = 256 / sizeof_bits_v; + + // Wait for TMEM allocation for this pipeline to finish + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; + // g2s load all global_d_amax + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueColQuantThreadCount) { + shared_storage.global_d_amax[g] = __ldg(reinterpret_cast(amax_colwise + g)); + } - size_t rng_seed = 0; - size_t rng_offset = 0; - // Setup RNG for stochastic rounding - if constexpr (kEnableStochasticRounding) { - rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; - rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; - } - // TODO(zhongbo): double check the logic here - int group_idx = get_current_tensor_id(shape_rep, num_tensors, - (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, - packed_N, M, offsets); - - // Determine quantization scale factor layouts/output splits for this group - TSFDLayout sfd_layout; - int cur_N = static_cast(first_dims[group_idx]); - if constexpr (kEnableSwizzleSFOutput) { - sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); - } else { - sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), - make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); - } - // Build output tensors for columns and their quant scales - // TODO(zhongbo): double check the logic here - Tensor mD = make_tensor(cute::subbyte_iterator(reinterpret_cast( - reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)), - make_shape(M, cur_N), DStride{}); // (M,packed_N) - Tensor gD_mn = - local_tile(mD, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{}); // (BLK_M,BLK_N) - - // for every tensor [x, y] row major, x y both a multiple of 128 - // both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3 - Tensor mSFD = make_tensor( - make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) + - offsets[group_idx] / kNVFP4BlockSize)), - sfd_layout); - Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + size_t rng_seed = 0; + size_t rng_offset = 0; + // Setup RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + // TODO(zhongbo): double check the logic here + int group_idx = get_current_tensor_id( + shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, + packed_N, M, offsets); + + // Determine quantization scale factor layouts/output splits for this group + TSFDLayout sfd_layout; + int cur_N = static_cast(first_dims[group_idx]); + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // Build output tensors for columns and their quant scales + // TODO(zhongbo): double check the logic here + Tensor mD = make_tensor(cute::subbyte_iterator(reinterpret_cast( + reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)), + make_shape(M, cur_N), DStride{}); // (M,packed_N) + Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{}); // (BLK_M,BLK_N) - Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); - - // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors - auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); - auto tiled_r2g = - make_tiled_copy_D(Copy_Atom{}, tiled_t2r); - auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); - auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); - - cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, - cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; - static constexpr float fp4_max_inv = 1.0f / fp4_max; - float c_global_amax_val = shared_storage.global_d_amax[group_idx]; - float global_encode_scale = c_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - float global_decode_scale = 1.0f / global_encode_scale; - - // Scaling factor for fast math path - float global_encode_scale_multiplier = 1.0f; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; - } + // for every tensor [x, y] row major, x y both a multiple of 128 + // both of its rowwise and colwise scaling factors will have exactly x * y / 16 elements in FP8 E4M3 + Tensor mSFD = make_tensor( + make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) + + offsets[group_idx] / kNVFP4BlockSize)), + sfd_layout); + Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + + // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); + auto tiled_r2g = + make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); + auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); + + cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float c_global_amax_val = shared_storage.global_d_amax[group_idx]; + float global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + float global_decode_scale = 1.0f / global_encode_scale; + + // Scaling factor for fast math path + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + + do { + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); + ++k_tile) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); - do { - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - CUTLASS_PRAGMA_NO_UNROLL - for (int k_tile = 0; - k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); - ++k_tile) { - int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); - - // TODO(zhongbo): double check the logic here - int cur_group_idx = get_current_tensor_id(shape_rep, num_tensors, - global_tile_n_offset * M, packed_N, M, offsets); - - if (cur_group_idx != group_idx) { - group_idx = cur_group_idx; - c_global_amax_val = shared_storage.global_d_amax[group_idx]; - // update amax - global_encode_scale = c_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - global_decode_scale = 1.0f / global_encode_scale; - if constexpr (kUseFastMath) { + // TODO(zhongbo): double check the logic here + int cur_group_idx = get_current_tensor_id( + shape_rep, num_tensors, global_tile_n_offset * M, packed_N, M, offsets); + + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + c_global_amax_val = shared_storage.global_d_amax[group_idx]; + // update amax + global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + // TODO(zhongbo): double check the logic here + cur_N = first_dims[group_idx]; + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = + tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = + make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // update tensor + mD = make_tensor(cute::subbyte_iterator(reinterpret_cast( + reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)), + make_shape(M, cur_N), DStride{}); + gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( + reinterpret_cast(SFA_COLWISE) + + offsets[group_idx] / kNVFP4BlockSize)), + sfd_layout); + gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); } - // TODO(zhongbo): double check the logic here - cur_N = first_dims[group_idx]; - if constexpr (kEnableSwizzleSFOutput) { - sfd_layout = - tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); - } else { - sfd_layout = - make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), - make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + int group_start_offset = offsets[group_idx] / M; + int local_tile_n_idx = + (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler); + Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx); + + Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx); + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); + Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + + Tensor tTR_rAcc = make_tensor( + shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDrD = make_tensor(shape(tDgD)); + Tensor tTR_rAcc_frag = + recast>(coalesce(tTR_rAcc)); + Tensor tDrD_frag = recast>(coalesce(tDrD)); + + Tensor src = thr_r2g.retile_S(tDrD); + Tensor dst = thr_r2g.retile_D(tDgD); + + Tensor tDgSFD_view = make_tensor( + tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); + Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); + Tensor tDrSFD = make_tensor(shape(tDgSFD)); + + static int constexpr NumVecs = size(tDgD) / VectorSize; + Tensor tD_rRowSFD_frg = recast>(tDrSFD); + + // Compute amax and quantization scales for this tile + cutlass::maximum_absolute_value_reduction< + cutlass::Array, true> + amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // Copy from TMEM to registers + copy(tiled_t2r, tDtAcc, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with + // unfused kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = + convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + tTR_rAcc_frag(_1{}) = + convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); } - // update tensor - mD = make_tensor(cute::subbyte_iterator(reinterpret_cast( - reinterpret_cast(QA_COLWISE) + offsets[group_idx] / 2)), - make_shape(M, cur_N), DStride{}); - gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), - Step<_1, _1, X>{}); // (BLK_M,BLK_N) - mSFD = make_tensor( - make_gmem_ptr(reinterpret_cast(reinterpret_cast(SFA_COLWISE) + - offsets[group_idx] / kNVFP4BlockSize)), - sfd_layout); - gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), - Step<_1, _1, X>{}); // (BLK_M,BLK_N) - gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); - } - int group_start_offset = offsets[group_idx] / M; - int local_tile_n_idx = - (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler); - Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx); - - Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx); - accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); - - auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); - Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - - Tensor tTR_rAcc = - make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - Tensor tDrD = make_tensor(shape(tDgD)); - Tensor tTR_rAcc_frag = - recast>(coalesce(tTR_rAcc)); - Tensor tDrD_frag = recast>(coalesce(tDrD)); - - Tensor src = thr_r2g.retile_S(tDrD); - Tensor dst = thr_r2g.retile_D(tDgD); - - Tensor tDgSFD_view = make_tensor( - tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), - make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); - Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); - Tensor tDrSFD = make_tensor(shape(tDgSFD)); - - static int constexpr NumVecs = size(tDgD) / VectorSize; - Tensor tD_rRowSFD_frg = recast>(tDrSFD); - - // Compute amax and quantization scales for this tile - cutlass::maximum_absolute_value_reduction, - true> - amax_reduction; - cutlass::Array vec_maxs; - cutlass::Array pvscales; - // Copy from TMEM to registers - copy(tiled_t2r, tDtAcc, tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); - ++accumulator_pipe_consumer_state; - - if constexpr (!kUseFastMath) { - // Downcast to BF16 for bit-wise compatibility with - // unfused kernels - auto convert_accum_to_bf16 = - cutlass::NumericArrayConverter{}; - auto convert_bf16_to_accum = - cutlass::NumericArrayConverter{}; - tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); - tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); - } - - auto compute_frgs = reinterpret_cast *>( - tTR_rAcc_frag.data()); - auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < NumVecs; v++) { - vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); - } + auto compute_frgs = reinterpret_cast *>( + tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); + } - if constexpr (kUseFastMath) { - // Fast math: multiply with precomputed reciprocal pvscales = cutlass::multiplies>{}( vec_maxs, global_encode_scale_multiplier); - } else { - // Accurate math: perform division - pvscales = - cutlass::divides>{}(vec_maxs, fp4_max); - pvscales = cutlass::multiplies>{}( - pvscales, global_encode_scale); - } - auto pvscales_cvted = - cutlass::NumericArrayConverter{}(pvscales); - - tD_rRowSFD_frg(_0{}) = pvscales_cvted; - auto qpvscale_ups = cutlass::NumericArrayConverter{}( - tD_rRowSFD_frg(_0{})); - auto qpvscale_scaled = cutlass::multiplies>{}( - qpvscale_ups, global_decode_scale); - cutlass::Array acc_scales; - if constexpr (kUseFastMath) { - // Fast math: compute approximate reciprocal - acc_scales = - cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); - } else { - // Accurate math: compute reciprocal with division - acc_scales = cutlass::divides>{}( - 1.0, qpvscale_scaled); - } - - // Prepare stochastic rounding random state if enabled - uint4 random_uint4 = uint4{0, 0, 0, 0}; - transformer_engine::curanddx::detail::philox4x32_native_state< - NVTE_BUILD_NUM_PHILOX_ROUNDS> - rng; - // "Prefetch" a stochastic rounding state for the first tile - if constexpr (kEnableStochasticRounding) { - const size_t rng_sequence = global_thread_idx + k_tile * 512 + - scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; - rng.init(rng_seed, rng_sequence, rng_offset); - } - CUTLASS_PRAGMA_UNROLL - // Apply round/quantize to each fragment, with or without stochastic rounding - for (int v = 0; v < NumVecs; v++) { - auto acc_scale = cutlass::minimum_with_nan_propagation{}( - acc_scales[v], cutlass::platform::numeric_limits::max()); - if constexpr (kEnableStochasticRounding) { - random_uint4 = rng.generate4(); - output_frgs[v] = StochasticNumericConverter( - cutlass::multiplies>{}( - compute_frgs[v], acc_scale), - *reinterpret_cast *>(&random_uint4)); - } else { - output_frgs[v] = cutlass::NumericArrayConverter{}( - cutlass::multiplies>{}( - compute_frgs[v], acc_scale)); - } - } + auto pvscales_cvted = + cutlass::NumericArrayConverter{}(pvscales); - // Write quantized FP4 tile and dequant scale to gmem - copy(tiled_r2g, src, dst); - copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD); - } - scheduler.update_work_tile_info(); - } while (scheduler.is_valid()); - } - } else if (is_epilogue_row_quant_warp) { - // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage. - cutlass::arch::warpgroup_reg_alloc<136>(); - if constexpr (kEnableRowQuant) { - using S2RVectorType = uint128_t; - - int global_thread_idx = threadIdx.x; - int local_thread_idx = global_thread_idx % 256; - size_t rng_seed = 0; - size_t rng_offset = 0; - // g2s load all global_a_amax for all groups/tensors - CUTLASS_PRAGMA_NO_UNROLL - for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) { - shared_storage.global_a_amax[g] = __ldg(reinterpret_cast(amax_rowwise + g)); - } - // RNG for stochastic rounding - if constexpr (kEnableStochasticRounding) { - rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; - rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; - } - // Input/output tensors/partitions for row quant warp - Tensor mQA = - make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA)); - Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); - Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); - - Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _), - Step<_1, X, _1>{}); // (BLK_M,BLK_N) - // Swizzled shared memory A tile, with layout - Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>( - coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), - sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) - - // Set up layouts for partitioning – tile-by-warp, with vector granularity - using S2RWarpLayout = Layout>; - using WarpGroupLayout = Layout>; - using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); - using S2RValLayout = Layout, _1>>; - using S2RAtomA = Copy_Atom; - using R2GAtomQA = Copy_Atom; - using R2GAtomSFA = Copy_Atom; - auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); - auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{}); - auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{}); - - auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); - auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); - auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx); - Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) - - // Allocate temporary register tensors for copying quantization => output - Tensor tQArA = make_tensor_like( - make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) - Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn); - Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); - - Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); - Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); - - // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8 - // in order to go over the reserved named barrier count. - constexpr int row_quant_barrier_id = 2; - cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); - - int group_idx = get_current_tensor_id(shape_rep, num_tensors, - (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, - packed_N, M, offsets); - float a_global_amax_val = shared_storage.global_a_amax[group_idx]; - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; - static constexpr float fp4_max_inv = 1.0f / fp4_max; - float global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - - float global_decode_scale = 1.0f / global_encode_scale; - float global_encode_scale_multiplier = 1.0f; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; - } - auto sfa_converter = cutlass::NumericConverter{}; - do { - CUTLASS_PRAGMA_NO_UNROLL - for (int k_tile = 0; - k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { - int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); - - int cur_group_idx = get_current_tensor_id(shape_rep, num_tensors, - global_tile_n_offset * M, packed_N, M, offsets); - if (cur_group_idx != group_idx) { - group_idx = cur_group_idx; - a_global_amax_val = shared_storage.global_a_amax[group_idx]; - // Update group quantization parameters/scaling - global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - global_decode_scale = 1.0f / global_encode_scale; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; - } - } - - auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); - auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); - auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state); - mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); - copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); - cutlass::arch::fence_view_async_shared(); - mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state); - ++mainloop_pipe_consumer_state; - ++k_tile; - - // static int constexpr NumVecs = size(tQArA) / VectorSize; - cutlass::maximum_absolute_value_reduction, - true> - amax_reduction; - auto compute_frgs = reinterpret_cast *>(tQArA.data()); - auto output_frgs = - reinterpret_cast *>(raw_pointer_cast(tQArQA.data())); - Tensor amax = - make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); - Tensor pvscales = make_tensor_like(amax); - transformer_engine::curanddx::detail::philox4x32_native_state< - NVTE_BUILD_NUM_PHILOX_ROUNDS> - rng; - if constexpr (kEnableStochasticRounding) { - const size_t rng_sequence = global_thread_idx + k_tile * 512 + - scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + - tiles_in_m * tiles_in_n * K_TILE_MAX * 512; - rng.init(rng_seed, rng_sequence, rng_offset); - } - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) { - auto amax_view = group_modes<1, rank(amax)>(amax); - auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales); - auto compute_frgs_up = - cutlass::NumericArrayConverter{}( - compute_frgs[v]); - amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up); - if constexpr (kUseFastMath) { - // Fast math: multiply with precomputed reciprocal - pvscales_view(_0{}, v) = cutlass::multiplies{}( - amax_view(_0{}, v), global_encode_scale_multiplier); - } else { - // Accurate math: perform division - pvscales_view(_0{}, v) = - cutlass::divides{}(amax_view(_0{}, v), fp4_max); - pvscales_view(_0{}, v) = cutlass::multiplies{}( - pvscales_view(_0{}, v), global_encode_scale); - } - filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v)); - auto qpvscale_ups = - cutlass::NumericConverter{}(filter(tQArSFA)(v)); + tD_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}( + tD_rRowSFD_frg(_0{})); auto qpvscale_scaled = - cutlass::multiplies{}(qpvscale_ups, global_decode_scale); - ElementAccumulator acc_scales; + cutlass::multiplies>{}( + qpvscale_ups, global_decode_scale); + cutlass::Array acc_scales; if constexpr (kUseFastMath) { // Fast math: compute approximate reciprocal acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); } else { // Accurate math: compute reciprocal with division - acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + acc_scales = cutlass::divides>{}( + 1.0, qpvscale_scaled); } - auto acc_scale = cutlass::minimum_with_nan_propagation{}( - acc_scales, cutlass::platform::numeric_limits::max()); + + // Prepare stochastic rounding random state if enabled uint4 random_uint4 = uint4{0, 0, 0, 0}; + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; + // "Prefetch" a stochastic rounding state for the first tile if constexpr (kEnableStochasticRounding) { - random_uint4 = rng.generate4(); - output_frgs[v] = StochasticNumericConverter( - cutlass::multiplies>{}( - compute_frgs_up, acc_scale), - *reinterpret_cast *>(&random_uint4)); - } else { - output_frgs[v] = - cutlass::NumericArrayConverter{}( - cutlass::multiplies>{}( - compute_frgs_up, acc_scale)); + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); } + CUTLASS_PRAGMA_UNROLL + // Apply round/quantize to each fragment, with or without stochastic rounding + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], cutlass::platform::numeric_limits::max()); + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = + cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale)); + } + } + + // Write quantized FP4 tile and dequant scale to gmem + copy(tiled_r2g, src, dst); + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD); } - copy(tiled_r2g_QA, tQArQA, tQAgQA_mn); - copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn)); + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + } else if (is_epilogue_row_quant_warp) { + // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage. + cutlass::arch::warpgroup_reg_alloc<136>(); + if constexpr (kEnableRowQuant) { + using S2RVectorType = uint128_t; + + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % 256; + size_t rng_seed = 0; + size_t rng_offset = 0; + // g2s load all global_a_amax for all groups/tensors + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < num_tensors; g += NumEpilogueRowQuantThreadCount) { + shared_storage.global_a_amax[g] = __ldg(reinterpret_cast(amax_rowwise + g)); } - // scheduler.advance(); - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - scheduler.update_work_tile_info(); - } while (scheduler.is_valid()); - } + // RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + // Input/output tensors/partitions for row quant warp + Tensor mQA = + make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA)); + Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); + + Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _), + Step<_1, X, _1>{}); // (BLK_M,BLK_N) + // Swizzled shared memory A tile, with layout + Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>( + coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) + + // Set up layouts for partitioning – tile-by-warp, with vector granularity + using S2RWarpLayout = Layout>; + using WarpGroupLayout = Layout>; + using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); + using S2RValLayout = Layout, _1>>; + using S2RAtomA = Copy_Atom; + using R2GAtomQA = Copy_Atom; + using R2GAtomSFA = Copy_Atom; + auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{}); + + auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); + auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); + auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx); + Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) + + // Allocate temporary register tensors for copying quantization => output + Tensor tQArA = make_tensor_like( + make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) + Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn); + Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); + + Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); + Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); + + // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8 + // in order to go over the reserved named barrier count. + constexpr int row_quant_barrier_id = 2; + cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); + + int group_idx = get_current_tensor_id( + shape_rep, num_tensors, (scheduler.tile_n_base() * size<1>(epilogue_tiler)) * M, + packed_N, M, offsets); + float a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + + float global_decode_scale = 1.0f / global_encode_scale; + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + auto sfa_converter = cutlass::NumericConverter{}; + do { + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + int cur_group_idx = get_current_tensor_id( + shape_rep, num_tensors, global_tile_n_offset * M, packed_N, M, offsets); + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Update group quantization parameters/scaling + global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } - } else { - cutlass::arch::warpgroup_reg_dealloc<32>(); + auto tQAgSFA_mn = + tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state); + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); + cutlass::arch::fence_view_async_shared(); + mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state); + ++mainloop_pipe_consumer_state; + ++k_tile; + + // static int constexpr NumVecs = size(tQArA) / VectorSize; + cutlass::maximum_absolute_value_reduction< + cutlass::Array, true> + amax_reduction; + auto compute_frgs = reinterpret_cast *>(tQArA.data()); + auto output_frgs = reinterpret_cast *>( + raw_pointer_cast(tQArQA.data())); + Tensor amax = + make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); + Tensor pvscales = make_tensor_like(amax); + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + + tiles_in_m * tiles_in_n * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) { + auto amax_view = group_modes<1, rank(amax)>(amax); + auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales); + auto compute_frgs_up = + cutlass::NumericArrayConverter{}( + compute_frgs[v]); + amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up); + pvscales_view(_0{}, v) = cutlass::multiplies{}( + amax_view(_0{}, v), global_encode_scale_multiplier); + filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v)); + auto qpvscale_ups = + cutlass::NumericConverter{}(filter(tQArSFA)(v)); + auto qpvscale_scaled = + cutlass::multiplies{}(qpvscale_ups, global_decode_scale); + ElementAccumulator acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = cutlass::reciprocal_approximate_ftz{}( + qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + } + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales, cutlass::platform::numeric_limits::max()); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = + cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale)); + } + } + copy(tiled_r2g_QA, tQArQA, tQAgQA_mn); + copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn)); + } + // scheduler.advance(); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + + } else { + cutlass::arch::warpgroup_reg_dealloc<32>(); + } } } // NOLINT(readability/fn_size) 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 1e40fd4a58..e6de366f52 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 @@ -171,528 +171,525 @@ __global__ static void group_rht_gemm_device( BSmemLayout sBlayout, CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, CSmemLayout, TiledMMA mma, MultiAmaxHadamardCastFusionArgs kernel_args, const size_t *rng_state) { using namespace cute; - using X = Underscore; - // static constexpr bool kApplyStochasticRounding = true; - using ElementAccumulator = float; - static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{}); - using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; - static constexpr uint32_t kTmaTransactionBytes = cutlass::bits_to_bytes( - size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); - - static constexpr int kTmaRhtTensorTransactionBytes = - cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v); - static constexpr int AccumulatorPipelineStageCount = 16; - - static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{}); - using MainloopPipeline = - cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>; - using MainloopPipelineState = typename MainloopPipeline::PipelineState; - - using TmemAllocator = cute::TMEM::Allocator1Sm; - static constexpr int VectorSize = 16; - const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; - const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; - // Preconditions - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - - // Represent the full tensors - Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, N)); - Tensor mB = tma_load_b.get_tma_tensor(make_shape(16, 16)); - - using TensorC = decltype(make_tensor(subbyte_iterator(recast_ptr(nullptr)), // engine - make_shape(int{}, int{}), // (M, N_i) - Stride2D{} // stride (dM, dN) - )); - - using TensorSFC = decltype(make_tensor( - make_gmem_ptr(recast_ptr(nullptr)), - make_layout(make_shape(int{}, // M - make_shape(make_shape(Int<16>{}, _4{}), // (16, 4) - int{}) // n_tiles = split / 64 - ), - make_stride(int{}, // dM = (split / 16) - make_stride(make_stride(_0{}, _1{}), // inner (16,4) layout - _4{}) // tiles stride - )))); - - auto cluster_shape = Shape<_1, _1, _1>{}; - - // Get the appropriate blocks for this Cluster - dim3 cluster_coord_in_grid = cluster_id_in_grid(); - - // Total number of k-tiles - const int K_TILE_MAX = min(N, K) / 64; - uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile); - uint32_t tiles_in_n = (N + 64 - 1) / 64; - uint32_t linear_tile_idx = blockIdx.x; - uint32_t tile_idx_m = linear_tile_idx % tiles_in_m; - uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; - - auto mainloop_tiler = Shape<_128, _16, _64>{}; - auto epilogue_tiler = Shape<_128, _64, _64>{}; - Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); - Tensor gB_nk = - local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) - // Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) - - using TensorGC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){}, - make_coord(_, _, _), Step<_1, _1, X>{})); - - using TensorGSFC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){}, + constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; + if constexpr (!is_blackwell_arch) { + NVTE_DEVICE_ERROR("RHT fusion is only supported on Blackwell."); + return; + } else { + using X = Underscore; + // static constexpr bool kApplyStochasticRounding = true; + using ElementAccumulator = float; + static constexpr int K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static constexpr uint32_t kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); + + static constexpr int kTmaRhtTensorTransactionBytes = + cutlass::bits_to_bytes(16 * 16 * cute::sizeof_bits_v); + static constexpr int AccumulatorPipelineStageCount = 16; + + static constexpr int MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = cutlass::PipelineTmaUmmaAsync, AtomThrShapeMNK>; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static constexpr int VectorSize = 16; + const size_t rng_seed = rng_state != nullptr ? rng_state[0] : 0; + const size_t rng_offset = rng_state != nullptr ? rng_state[1] : 0; + // Preconditions + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + + // Represent the full tensors + Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, N)); + Tensor mB = tma_load_b.get_tma_tensor(make_shape(16, 16)); + + using TensorC = decltype(make_tensor(subbyte_iterator(recast_ptr(nullptr)), // engine + make_shape(int{}, int{}), // (M, N_i) + Stride2D{} // stride (dM, dN) + )); + + using TensorSFC = decltype(make_tensor( + make_gmem_ptr(recast_ptr(nullptr)), + make_layout(make_shape(int{}, // M + make_shape(make_shape(Int<16>{}, _4{}), // (16, 4) + int{}) // n_tiles = split / 64 + ), + make_stride(int{}, // dM = (split / 16) + make_stride(make_stride(_0{}, _1{}), // inner (16,4) layout + _4{}) // tiles stride + )))); + + auto cluster_shape = Shape<_1, _1, _1>{}; + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + const int K_TILE_MAX = min(N, K) / 64; + uint32_t tiles_in_m = (M + size<0>(cluster_tile) - 1) / size<0>(cluster_tile); + uint32_t tiles_in_n = (N + 64 - 1) / 64; + uint32_t linear_tile_idx = blockIdx.x; + uint32_t tile_idx_m = linear_tile_idx % tiles_in_m; + uint32_t tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + + auto mainloop_tiler = Shape<_128, _16, _64>{}; + auto epilogue_tiler = Shape<_128, _64, _64>{}; + Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gB_nk = + local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) + // Tensor gC_mn = local_tile(mC, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + using TensorGC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){}, make_coord(_, _, _), Step<_1, _1, X>{})); - // Allocate SMEM - extern __shared__ char shared_memory[]; - using SharedStorage = SharedStorage; - SharedStorage &shared_storage = *reinterpret_cast(shared_memory); - Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), - sAlayout); // (MMA,MMA_M,MMA_N,PIPE) - Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), - sBlayout); // (MMA,MMA_N,MMA_K,PIPE) - - // - // MMA: Define C accumulators and A/B partitioning - // - - int block_rank_in_cluster = cute::block_rank_in_cluster(); - ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx - Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) - - auto mma_epilogue = make_tiled_mma( - SM100_MMA_F16BF16_SS{}, - Layout>{}); - ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster); - - using TiledMmaEpilogue = decltype(mma_epilogue); - Tensor tCgA = thr_mma.partition_A(gA_mk); - // Allocate "fragments" -- these are actually umma smem descriptors - Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) - Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) - - auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0, 2>(ClusterTileShape{})); - auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0, 2>(epilogue_tiler)); - - auto bulk_tmem_mma = - TiledMMA::make_fragment_C(append(acc_shape_mma, Int{})); - - auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C( - append(acc_shape_epilogue, Int{})); - - TmemAllocator tmem_allocator{}; - cutlass::arch::NamedBarrier tmem_allocation_result_barrier( - 32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); - - Layout cta_layout_mnk = make_layout(cluster_shape); - Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); - auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); - - auto [tAgA, tAsA] = - tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), - group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); - - auto [tBgB, tBsB] = - tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), - group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); - - uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); - uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); - - int warp_idx = cutlass::canonical_warp_idx_sync(); - - bool is_mma_warp = (warp_idx == 0); - bool is_dma_warp = (warp_idx == 1); - bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7); - - // if (is_epilogue_warp && elect_one_sync()) { - // // prefetch to make the global amax in cache - // for (size_t i = 0; i < kernel_args.num_tensors; ++i) { - // cute::prefetch(raw_pointer_cast(kernel_args.global_amax_list[i])); - // } - // } - - typename MainloopPipeline::Params mainloop_pipeline_params; - if (is_dma_warp) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; - } - if (is_mma_warp) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; - } - mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; - mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; - mainloop_pipeline_params.initializing_warp = 0; - MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, - cluster_shape, cute::true_type{}, // Perform barrier init - cute::true_type{}); // Delay mask calculation - - MainloopPipelineState mainloop_pipe_consumer_state; - MainloopPipelineState mainloop_pipe_producer_state = - cutlass::make_producer_start_state(); - - using AccumulatorPipeline = - cutlass::PipelineUmmaAsync; - using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; - - AccumulatorPipelineState accumulator_pipe_consumer_state; - AccumulatorPipelineState accumulator_pipe_producer_state = - cutlass::make_producer_start_state(); - - typename AccumulatorPipeline::Params accumulator_pipeline_params; - if (is_mma_warp) { - accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; - } - if (is_epilogue_warp) { - accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; - } - // Only one producer thread arrives on this barrier. - accumulator_pipeline_params.producer_arv_count = 1; - accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128; - accumulator_pipeline_params.initializing_warp = 1; - AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, accumulator_pipeline_params, - cluster_shape, - cute::true_type{}, // Perform barrier init - cute::true_type{}); // Delay mask calculation - - if (warp_idx == 2 && elect_one_sync()) { - cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); - } - __syncthreads(); - using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; - - if (is_dma_warp) { - if (elect_one_sync()) { - cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], - kTmaRhtTensorTransactionBytes); - copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), - tBsB(_, 0)); + using TensorGSFC = decltype(local_tile(std::declval(), decltype(epilogue_tiler){}, + make_coord(_, _, _), Step<_1, _1, X>{})); + + // Allocate SMEM + extern __shared__ char shared_memory[]; + using SharedStorage = SharedStorage; + SharedStorage &shared_storage = *reinterpret_cast(shared_memory); + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + // + // MMA: Define C accumulators and A/B partitioning + // + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + auto mma_epilogue = make_tiled_mma(SM100_MMA_F16BF16_SS{}, + Layout>{}); + ThrMMA thr_mma_epilogue = mma_epilogue.get_slice(block_rank_in_cluster); + + using TiledMmaEpilogue = decltype(mma_epilogue); + Tensor tCgA = thr_mma.partition_A(gA_mk); + // Allocate "fragments" -- these are actually umma smem descriptors + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + auto acc_shape_mma = partition_shape_C(TiledMMA{}, take<0, 2>(ClusterTileShape{})); + auto acc_shape_epilogue = partition_shape_C(TiledMmaEpilogue{}, take<0, 2>(epilogue_tiler)); + + auto bulk_tmem_mma = + TiledMMA::make_fragment_C(append(acc_shape_mma, Int{})); + + auto bulk_tmem_epilogue = TiledMmaEpilogue::make_fragment_C( + append(acc_shape_epilogue, Int{})); + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + 32 + 128, cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = + tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = + tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); + + auto [tBgB, tBsB] = + tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_epilogue_warp = (warp_idx >= 4 && warp_idx <= 7); + + // if (is_epilogue_warp && elect_one_sync()) { + // // prefetch to make the global amax in cache + // for (size_t i = 0; i < kernel_args.num_tensors; ++i) { + // cute::prefetch(raw_pointer_cast(kernel_args.global_amax_list[i])); + // } + // } + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, + cluster_shape, cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = + cutlass::make_producer_start_state(); + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = + cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * 128; + accumulator_pipeline_params.initializing_warp = 1; + AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, + accumulator_pipeline_params, cluster_shape, + cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + if (is_dma_warp) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], + kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), + tBsB(_, 0)); + } - do { - bool is_first_wave = linear_tile_idx == blockIdx.x; - uint32_t skip_wait = is_first_wave; - auto tAgA_mk = tAgA(_, tile_idx_m, _); - int k_tile = 0; - auto barrier_token = - mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait); - - CUTE_NO_UNROLL - while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) { - int k_tile_idx_n = tile_idx_n + k_tile; - ++k_tile; - skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); - mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token); - using BarrierType = typename MainloopPipeline::ProducerBarrierType; - BarrierType *tma_barrier = - mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); - int write_stage = mainloop_pipe_producer_state.index(); - ++mainloop_pipe_producer_state; - barrier_token = + do { + bool is_first_wave = linear_tile_idx == blockIdx.x; + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_, tile_idx_m, _); + int k_tile = 0; + auto barrier_token = mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait); - if (cute::elect_one_sync()) { - copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), - tAsA(_, write_stage)); + + CUTE_NO_UNROLL + while (k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n) { + int k_tile_idx_n = tile_idx_n + k_tile; + ++k_tile; + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state, barrier_token); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType *tma_barrier = + mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + barrier_token = + mainloop_pipeline.producer_try_acquire(mainloop_pipe_producer_state, skip_wait); + if (cute::elect_one_sync()) { + copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), + tAsA(_, write_stage)); + } } - } - linear_tile_idx += gridDim.x; - tile_idx_m = linear_tile_idx % tiles_in_m; - tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; - } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); - mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); - } else if (is_mma_warp) { - mma.accumulate_ = UMMA::ScaleOut::Zero; - - tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); - __syncwarp(); - tmem_allocation_result_barrier.arrive(); - uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; - bulk_tmem_mma.data() = tmem_base_ptr; - - cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); - do { - uint32_t skip_wait = K_TILE_MAX <= 0; - auto barrier_token = - mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); - CUTE_NO_UNROLL - for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n;) { - mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); - int read_stage = mainloop_pipe_consumer_state.index(); - auto tCrA_mk = tCrA(_, _, _, read_stage); - auto tCrB_nk = tCrB(_, _, 0, 0); - CUTE_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block) { - accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + linear_tile_idx += gridDim.x; + tile_idx_m = linear_tile_idx % tiles_in_m; + tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + auto barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + CUTE_NO_UNROLL + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n;) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_, _, _, read_stage); + auto tCrB_nk = tCrB(_, _, 0, 0); CUTE_UNROLL - for (int i = 0; i < 4; i++) { - auto accumulators = - bulk_tmem_mma(_, _, _, accumulator_pipe_producer_state.index() * 4 + i); - gemm(mma, tCrA_mk(_, _, k_block * 4 + i), tCrB_nk, accumulators); + for (int k_block = 0; k_block < size<2>(tCrA) / 4; ++k_block) { + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTE_UNROLL + for (int i = 0; i < 4; i++) { + auto accumulators = + bulk_tmem_mma(_, _, _, accumulator_pipe_producer_state.index() * 4 + i); + gemm(mma, tCrA_mk(_, _, k_block * 4 + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; } - - accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); - ++accumulator_pipe_producer_state; + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state); } - auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; - ++mainloop_pipe_consumer_state; - ++k_tile; - skip_wait = k_tile >= K_TILE_MAX; - barrier_token = - mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); - mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state); - } - linear_tile_idx += gridDim.x; - tile_idx_m = linear_tile_idx % tiles_in_m; - tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; - } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); - tmem_allocator.release_allocation_lock(); - accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); - tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); - } else if (is_epilogue_warp) { - static constexpr int FragmentSize = 256 / sizeof_bits_v; - - tmem_allocation_result_barrier.arrive_and_wait(); - uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; - bulk_tmem_epilogue.data() = tmem_base_ptr; - int thread_idx = threadIdx.x % 128; - - auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); - auto tiled_r2g = - make_tiled_copy_D(Copy_Atom{}, tiled_t2r); - auto thr_t2r = tiled_t2r.get_slice(thread_idx); - auto thr_r2g = tiled_r2g.get_slice(thread_idx); - - // NVFP4 non-E8 recipe constants and global scales - static constexpr float fp4_max = 6.0f; - static constexpr float fp4_max_inv = 1.0f / fp4_max; - - // get global amax pointer - int tensor_id = GetTensorId(&kernel_args, tile_idx_n * 64); - float *global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, tensor_id); - - TC *cur_output_colwise_ptr = reinterpret_cast(kernel_args.output_colwise_list[tensor_id]); - TSFC *cur_output_colwise_scale_inv_ptr = - reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); - int cur_output_colwise_n = kernel_args.split_sections[tensor_id]; - - TensorC cur_mC = - cute::make_tensor(cute::subbyte_iterator(cur_output_colwise_ptr), - cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i) - kernel_args.output_stride2d_list[tensor_id]); - - auto cur_sfc_shape = - make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64)); - - auto cur_sfc_stride = - make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{})); - - TensorSFC cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr), - make_layout(cur_sfc_shape, cur_sfc_stride)); - - TensorGC cur_gC_mn = - local_tile(cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N) - ); - - TensorGSFC cur_gSFC_mn = local_tile( - cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N-like) - ); - - Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); - - float global_amax_val = *global_amax_ptr; - float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); - - // Scaling factor for fast math path - float global_encode_scale_multiplier = 1.0f; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; - } - - float global_decode_scale = 1.0f / global_encode_scale; - - auto sfd_converter = cutlass::NumericConverter{}; - - do { - for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) { - // get the starting index of current k-tile in global tensor, to query the correct global amax - int cur_k_tile_global_elem_idx = (tile_idx_n + k_tile) * 64; - int new_tensor_id = GetTensorId(&kernel_args, cur_k_tile_global_elem_idx); - // float* new_global_amax_ptr = GetGlobalAmaxPtr(&kernel_args, cur_k_tile_global_elem_idx); - global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, new_tensor_id); - // update the scaling factors when it's no longer the same amax pointer - // TODO(zhongbo): the math operations are very expensive - // since the kernel is persistent, we can have a cache for all the possible scaling factors - if (tensor_id != new_tensor_id) { - global_amax_val = *global_amax_ptr; - global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); - if constexpr (kUseFastMath) { + linear_tile_idx += gridDim.x; + tile_idx_m = linear_tile_idx % tiles_in_m; + tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } else if (is_epilogue_warp) { + static constexpr int FragmentSize = 256 / sizeof_bits_v; + + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int thread_idx = threadIdx.x % 128; + + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); + auto tiled_r2g = + make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(thread_idx); + auto thr_r2g = tiled_r2g.get_slice(thread_idx); + + // NVFP4 non-E8 recipe constants and global scales + static constexpr float fp4_max = 6.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + + // get global amax pointer + int tensor_id = GetTensorId(&kernel_args, tile_idx_n * 64); + float *global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, tensor_id); + + TC *cur_output_colwise_ptr = + reinterpret_cast(kernel_args.output_colwise_list[tensor_id]); + TSFC *cur_output_colwise_scale_inv_ptr = + reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); + int cur_output_colwise_n = kernel_args.split_sections[tensor_id]; + + TensorC cur_mC = cute::make_tensor( + cute::subbyte_iterator(cur_output_colwise_ptr), + cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i) + kernel_args.output_stride2d_list[tensor_id]); + + auto cur_sfc_shape = + make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64)); + + auto cur_sfc_stride = + make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{})); + + TensorSFC cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr), + make_layout(cur_sfc_shape, cur_sfc_stride)); + + TensorGC cur_gC_mn = local_tile( + cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N) + ); + + TensorGSFC cur_gSFC_mn = local_tile( + cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N-like) + ); + + Tensor tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); + + float global_amax_val = *global_amax_ptr; + float global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); + + // Scaling factor for fast math path + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + + float global_decode_scale = 1.0f / global_encode_scale; + + auto sfd_converter = cutlass::NumericConverter{}; + + do { + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + tile_idx_n < tiles_in_n; ++k_tile) { + // get the starting index of current k-tile in global tensor, to query the correct global amax + int cur_k_tile_global_elem_idx = (tile_idx_n + k_tile) * 64; + int new_tensor_id = GetTensorId(&kernel_args, cur_k_tile_global_elem_idx); + // float* new_global_amax_ptr = GetGlobalAmaxPtr(&kernel_args, cur_k_tile_global_elem_idx); + global_amax_ptr = GetGlobalAmaxPtrByTensorId(&kernel_args, new_tensor_id); + // update the scaling factors when it's no longer the same amax pointer + // TODO(zhongbo): the math operations are very expensive + // since the kernel is persistent, we can have a cache for all the possible scaling factors + if (tensor_id != new_tensor_id) { + global_amax_val = *global_amax_ptr; + global_encode_scale = ComputeGlobalEncodeScaleFP4(global_amax_val); global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + global_decode_scale = 1.0f / global_encode_scale; + tensor_id = new_tensor_id; + // went through the cute operations to update the local tensors + cur_output_colwise_ptr = + reinterpret_cast(kernel_args.output_colwise_list[tensor_id]); + cur_output_colwise_scale_inv_ptr = + reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); + cur_output_colwise_n = kernel_args.split_sections[tensor_id]; + + cur_mC = cute::make_tensor( + cute::subbyte_iterator(cur_output_colwise_ptr), + cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i) + kernel_args.output_stride2d_list[tensor_id]); + + cur_sfc_shape = + make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64)); + + cur_sfc_stride = + make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{})); + + cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr), + make_layout(cur_sfc_shape, cur_sfc_stride)); + + cur_gC_mn = local_tile( + cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N) + ); + + cur_gSFC_mn = local_tile(cur_mSFC, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{} // (BLK_M, BLK_N-like) + ); + + tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); + } + // maybe udpated to the new tensor id + int tensor_start_elem = kernel_args.split_sections_range[tensor_id]; + int local_tile_idx_n = (cur_k_tile_global_elem_idx - tensor_start_elem) / 64; + + Tensor tCgC_mn = tCgC(_, _, _, tile_idx_m, local_tile_idx_n); + Tensor tCgSFC_mn = cur_gSFC_mn(_, _, tile_idx_m, local_tile_idx_n); + + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto tCtC = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); + Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + + Tensor tTR_rAcc = + make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDrC = make_tensor(shape(tDgC)); + Tensor tTR_rAcc_frag = + recast>(coalesce(tTR_rAcc)); + Tensor tDrC_frag = recast>(coalesce(tDrC)); + + Tensor src = thr_r2g.retile_S(tDrC); + Tensor dst = thr_r2g.retile_D(tDgC); + + Tensor tCgSFC = make_tensor( + tCgSFC_mn.data(), make_layout(make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{}))); + + Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC)); + Tensor tDrSFC = make_tensor(shape(tDgSFC)); + + static constexpr int NumVecs = size(tDgC) / VectorSize; + Tensor tC_rRowSFD_frg = recast>(tDrSFC); + + cutlass::maximum_absolute_value_reduction, + true> + amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // TMEM_LOAD + copy(tiled_t2r, tDtC, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with unfused + // kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); } - global_decode_scale = 1.0f / global_encode_scale; - tensor_id = new_tensor_id; - // went through the cute operations to update the local tensors - cur_output_colwise_ptr = - reinterpret_cast(kernel_args.output_colwise_list[tensor_id]); - cur_output_colwise_scale_inv_ptr = - reinterpret_cast(kernel_args.output_colwise_scale_inv_list[tensor_id]); - cur_output_colwise_n = kernel_args.split_sections[tensor_id]; - - cur_mC = cute::make_tensor( - cute::subbyte_iterator(cur_output_colwise_ptr), - cute::make_shape(static_cast(M), cur_output_colwise_n), // (M, N_i) - kernel_args.output_stride2d_list[tensor_id]); - - cur_sfc_shape = - make_shape(M, make_shape(make_shape(Int<16>{}, _4{}), cur_output_colwise_n / 64)); - - cur_sfc_stride = - make_stride(cur_output_colwise_n / 16, make_stride(make_stride(_0{}, _1{}), _4{})); - - cur_mSFC = cute::make_tensor(make_gmem_ptr(cur_output_colwise_scale_inv_ptr), - make_layout(cur_sfc_shape, cur_sfc_stride)); - - cur_gC_mn = local_tile( - cur_mC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} // (BLK_M, BLK_N) - ); - - cur_gSFC_mn = local_tile(cur_mSFC, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{} - // (BLK_M, BLK_N-like) - ); - - tCgC = thr_mma_epilogue.partition_C(cur_gC_mn); - } - // maybe udpated to the new tensor id - int tensor_start_elem = kernel_args.split_sections_range[tensor_id]; - int local_tile_idx_n = (cur_k_tile_global_elem_idx - tensor_start_elem) / 64; - - Tensor tCgC_mn = tCgC(_, _, _, tile_idx_m, local_tile_idx_n); - Tensor tCgSFC_mn = cur_gSFC_mn(_, _, tile_idx_m, local_tile_idx_n); - - accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); - - auto tCtC = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); - Tensor tDtC = thr_t2r.partition_S(tCtC); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - Tensor tDgC = thr_t2r.partition_D(tCgC_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - - Tensor tTR_rAcc = - make_tensor(shape(tDgC)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - Tensor tDrC = make_tensor(shape(tDgC)); - Tensor tTR_rAcc_frag = - recast>(coalesce(tTR_rAcc)); - Tensor tDrC_frag = recast>(coalesce(tDrC)); - - Tensor src = thr_r2g.retile_S(tDrC); - Tensor dst = thr_r2g.retile_D(tDgC); - - Tensor tCgSFC = make_tensor( - tCgSFC_mn.data(), make_layout(make_shape(shape(tCgSFC_mn), Int<1>{}, Int<1>{}), - make_stride(stride(tCgSFC_mn), Int<0>{}, Int<0>{}))); - - Tensor tDgSFC = filter(thr_t2r.partition_D(tCgSFC)); - Tensor tDrSFC = make_tensor(shape(tDgSFC)); - - static constexpr int NumVecs = size(tDgC) / VectorSize; - Tensor tC_rRowSFD_frg = recast>(tDrSFC); - - cutlass::maximum_absolute_value_reduction, - true> - amax_reduction; - cutlass::Array vec_maxs; - cutlass::Array pvscales; - // TMEM_LOAD - copy(tiled_t2r, tDtC, tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - - accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); - - ++accumulator_pipe_consumer_state; - - if constexpr (!kUseFastMath) { - // Downcast to BF16 for bit-wise compatibility with unfused - // kernels - auto convert_accum_to_bf16 = - cutlass::NumericArrayConverter{}; - auto convert_bf16_to_accum = - cutlass::NumericArrayConverter{}; - tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); - } - auto compute_frgs = reinterpret_cast *>( - tTR_rAcc_frag.data()); - auto output_frgs = reinterpret_cast *>(tDrC_frag.data()); - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < NumVecs; v++) { - vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); - } + auto compute_frgs = reinterpret_cast *>( + tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrC_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); + } - if constexpr (kUseFastMath) { - // Fast math: multiply with precomputed reciprocal pvscales = cutlass::multiplies>{}( vec_maxs, global_encode_scale_multiplier); - } else { - // Accurate math: perform division - pvscales = - cutlass::divides>{}(vec_maxs, fp4_max); - pvscales = cutlass::multiplies>{}( - pvscales, global_encode_scale); - } - auto pvscales_cvted = - cutlass::NumericArrayConverter{}(pvscales); - - tC_rRowSFD_frg(_0{}) = pvscales_cvted; - auto qpvscale_ups = cutlass::NumericArrayConverter{}( - tC_rRowSFD_frg(_0{})); - auto qpvscale_scaled = cutlass::multiplies>{}( - qpvscale_ups, global_decode_scale); - cutlass::Array acc_scales; - if constexpr (kUseFastMath) { - // Fast math: compute approximate reciprocal - acc_scales = - cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); - } else { - // Accurate math: compute reciprocal with division - acc_scales = - cutlass::divides>{}(1.0, qpvscale_scaled); - } - - // Initialize RNG for tile - const size_t rng_sequence = thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256; - - transformer_engine::curanddx::detail::philox4x32_native_state - rng; - rng.init(rng_seed, rng_sequence, rng_offset); - uint4 random_uint4 = uint4{0, 0, 0, 0}; - - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < NumVecs; v++) { - auto acc_scale = cutlass::minimum_with_nan_propagation{}( - acc_scales[v], cutlass::platform::numeric_limits::max()); - // auto acc_scale = acc_scales[v]; - if constexpr (kEnableStochasticRounding) { - random_uint4 = rng.generate4(); - output_frgs[v] = StochasticNumericConverter( - cutlass::multiplies>{}( - compute_frgs[v], acc_scale), - reinterpret_cast *>(&random_uint4)); + auto pvscales_cvted = + cutlass::NumericArrayConverter{}(pvscales); + + tC_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}( + tC_rRowSFD_frg(_0{})); + auto qpvscale_scaled = cutlass::multiplies>{}( + qpvscale_ups, global_decode_scale); + cutlass::Array acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = + cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); } else { - output_frgs[v] = cutlass::NumericArrayConverter{}( - cutlass::multiplies>{}( - compute_frgs[v], acc_scale)); + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides>{}( + 1.0, qpvscale_scaled); } - } - copy(tiled_r2g, src, dst); + // Initialize RNG for tile + const size_t rng_sequence = + thread_idx + k_tile * 256 + linear_tile_idx * K_TILE_MAX * 256; + + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; + rng.init(rng_seed, rng_sequence, rng_offset); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], cutlass::platform::numeric_limits::max()); + // auto acc_scale = acc_scales[v]; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale), + reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale)); + } + } - // copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrC, tDgC); + copy(tiled_r2g, src, dst); - copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC); - } - linear_tile_idx += gridDim.x; - tile_idx_m = linear_tile_idx % tiles_in_m; - tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; - } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + // copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrC, tDgC); + + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFC, tDgSFC); + } + linear_tile_idx += gridDim.x; + tile_idx_m = linear_tile_idx % tiles_in_m; + tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; + } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); + } } } 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 4013fdf119..1265f2711c 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 @@ -185,942 +185,918 @@ __launch_bounds__(512, 1) __global__ static void group_row_col_rht_gemm_device( // Abort immediately if compilation is not supported constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; if constexpr (!is_blackwell_arch) { - NVTE_DEVICE_ERROR( - "group_row_col_rht_gemm_device is only supported on Blackwell " - "with architecture-specific compilation. " - "Try recompiling with sm_100a or similar."); + NVTE_DEVICE_ERROR("RHT fusion is only supported on Blackwell."); return; - } - static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, - "group_row_col_rht_gemm_device must generate row-wise " - "and/or column-wise output."); + } else { + static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, + "group_row_col_rht_gemm_device must generate row-wise " + "and/or column-wise output."); #if !defined(CUTLASS_ARCH_CLC_ENABLED) - CUTLASS_NOT_IMPLEMENTED(); - return; + CUTLASS_NOT_IMPLEMENTED(); + return; #endif - using X = Underscore; - // Accumulator data type for main computation - using ElementAccumulator = float; - static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); - using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; - static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( - size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); - static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; - static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; - static constexpr bool kEnableRowQuant = kEnableRowQuant_; - static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_; - static constexpr bool kUseFastMath = kUseFastMath_; - - // Constant for RHT tensor processing (tile size etc) - static int constexpr RhtTensorSize = 16; - - // Transaction bytes for TMA transfer on RHT tensor blocks - static int constexpr kTmaRhtTensorTransactionBytes = - cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); - static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; - static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; - - // Mainloop pipeline stage calculation, vectorization parameters for scaling factors - static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); - static int constexpr SFVecSize = 16; - // Swizzle output layout for scaling factor arrays - using SwizzledSFALayoutAtom = - cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; - using SwizzledSFDLayoutAtom = - cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; - - // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling - using MainloopPipeline = - cutlass::detail::CustomizedPipelineTmaUmmaAsync; - using MainloopPipelineState = typename MainloopPipeline::PipelineState; - using SchedPipeline = cutlass::PipelineCLCFetchAsync; - using SchedPipelineState = typename SchedPipeline::PipelineState; - using SchedThrottlePipeline = cutlass::PipelineAsync; - using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState; - - static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, "ClusterShape must be Shape<_1,_1,_1>"); - - using TmemAllocator = cute::TMEM::Allocator1Sm; - static int constexpr VectorSize = RhtTensorSize; - - // Compile-time safety: static shapes required for shared memory layouts - CUTE_STATIC_ASSERT(is_static::value); - CUTE_STATIC_ASSERT(is_static::value); - // CUTE_STATIC_ASSERT(is_static::value); - - auto cluster_size = size<0>(cluster_shape); - auto mainloop_tiler = Shape<_128, _16, _128>{}; - auto epilogue_tiler = Shape<_128, _128, _128>{}; - - static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); - - // Get the appropriate blocks for this Cluster - dim3 cluster_coord_in_grid = cluster_id_in_grid(); - - // Total number of k-tiles - int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler); - - struct TileScheduler { - uint32_t tiles_in_m = 0; - uint32_t tiles_in_n = 0; - uint32_t linear_idx = 0; - uint32_t next_linear_idx = 0; - uint32_t start_idx = 0; - uint32_t tile_m_idx = 0; - uint32_t tile_n_idx = 0; - int k_tile_max = 0; - uint32_t *atomic_tile_index_; - uint32_t *smem_tile_counter; - uint32_t atomic_offset; - cutlass::FastDivmodU64 divmod_tiles_in_m; - - CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, - uint32_t *atomic_tile_index, uint32_t *smem_tile_counter) - : tiles_in_m(tiles_m), - tiles_in_n(tiles_n), - linear_idx(blockIdx.x), - next_linear_idx(blockIdx.x), - start_idx(blockIdx.x), - k_tile_max(kmax), - atomic_tile_index_(atomic_tile_index), - smem_tile_counter(smem_tile_counter), - atomic_offset(gridDim.x), - divmod_tiles_in_m(uint64_t(tiles_m)) { - update_tile_idx(); - } - CUTLASS_DEVICE void update_tile_idx() { - uint64_t q, r; - divmod_tiles_in_m(q, r, uint64_t(linear_idx)); - tile_m_idx = static_cast(r); - tile_n_idx = static_cast(q) * uint32_t(k_tile_max); - } - CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; } - CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; } - CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; } - - CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; } + using X = Underscore; + // Accumulator data type for main computation + using ElementAccumulator = float; + static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0, 3>(ASmemLayout{})) * cute::sizeof_bits_v); + static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; + static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; + static constexpr bool kEnableRowQuant = kEnableRowQuant_; + static constexpr bool kEnableSwizzleSFOutput = kEnableSwizzleSFOutput_; + static constexpr bool kUseFastMath = kUseFastMath_; + + // Constant for RHT tensor processing (tile size etc) + static int constexpr RhtTensorSize = 16; + + // Transaction bytes for TMA transfer on RHT tensor blocks + static int constexpr kTmaRhtTensorTransactionBytes = + cutlass::bits_to_bytes(RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; + + // Mainloop pipeline stage calculation, vectorization parameters for scaling factors + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + static int constexpr SFVecSize = 16; + // Swizzle output layout for scaling factor arrays + using SwizzledSFALayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = + cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + + // Mainloop pipeline types for TMA async execution and epilogue cluster scheduling + using MainloopPipeline = + cutlass::detail::CustomizedPipelineTmaUmmaAsync; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + using SchedPipeline = cutlass::PipelineCLCFetchAsync; + using SchedPipelineState = typename SchedPipeline::PipelineState; + using SchedThrottlePipeline = cutlass::PipelineAsync; + using SchedThrottlePipelineState = typename SchedThrottlePipeline::PipelineState; + + static_assert(ClusterShape{} == Shape<_1, _1, _1>{}, "ClusterShape must be Shape<_1,_1,_1>"); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static int constexpr VectorSize = RhtTensorSize; + + // Compile-time safety: static shapes required for shared memory layouts + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + // CUTE_STATIC_ASSERT(is_static::value); + + auto cluster_size = size<0>(cluster_shape); + auto mainloop_tiler = Shape<_128, _16, _128>{}; + auto epilogue_tiler = Shape<_128, _128, _128>{}; + + static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + int const K_TILE_MAX = min(packed_N, K) / size<2>(epilogue_tiler); + + struct TileScheduler { + uint32_t tiles_in_m = 0; + uint32_t tiles_in_n = 0; + uint32_t linear_idx = 0; + uint32_t next_linear_idx = 0; + uint32_t start_idx = 0; + uint32_t tile_m_idx = 0; + uint32_t tile_n_idx = 0; + int k_tile_max = 0; + uint32_t *atomic_tile_index_; + uint32_t *smem_tile_counter; + uint32_t atomic_offset; + cutlass::FastDivmodU64 divmod_tiles_in_m; + + CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, + uint32_t *atomic_tile_index, uint32_t *smem_tile_counter) + : tiles_in_m(tiles_m), + tiles_in_n(tiles_n), + linear_idx(blockIdx.x), + next_linear_idx(blockIdx.x), + start_idx(blockIdx.x), + k_tile_max(kmax), + atomic_tile_index_(atomic_tile_index), + smem_tile_counter(smem_tile_counter), + atomic_offset(gridDim.x), + divmod_tiles_in_m(uint64_t(tiles_m)) { + update_tile_idx(); + } + CUTLASS_DEVICE void update_tile_idx() { + uint64_t q, r; + divmod_tiles_in_m(q, r, uint64_t(linear_idx)); + tile_m_idx = static_cast(r); + tile_n_idx = static_cast(q) * uint32_t(k_tile_max); + } + CUTLASS_DEVICE uint32_t tile_m() const { return tile_m_idx; } + CUTLASS_DEVICE uint32_t tile_n_base() const { return tile_n_idx; } + CUTLASS_DEVICE uint32_t tiles_m() const { return tiles_in_m; } - CUTLASS_DEVICE bool is_valid() const { - return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()), - cute::make_coord(tiles_in_m, tiles_in_n)); - } + CUTLASS_DEVICE uint32_t tiles_n() const { return tiles_in_n; } - CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; } + CUTLASS_DEVICE bool is_valid() const { + return cute::elem_less(cute::make_coord(tile_m(), tile_n_base()), + cute::make_coord(tiles_in_m, tiles_in_n)); + } - CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; } + CUTLASS_DEVICE bool is_first_wave() const { return linear_idx == start_idx; } - // Fetch a new tile_id using atomics. - CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) { - uint32_t tile_id_counter = 0; - asm volatile( - "{\n\t" - ".reg .pred p;\n\t" - "setp.eq.u32 p, %2, 1;\n\t" - "@p atom.global.add.u32 %0, [%1], 1; \n\t" - "}" - : "=r"(tile_id_counter) - : "l"(atomic_tile_index_), "r"(pred)); + CUTLASS_DEVICE uint32_t get_linear_tile_idx() const { return linear_idx; } - return tile_id_counter; - } + // Fetch a new tile_id using atomics. + CUTLASS_DEVICE uint32_t fetch_tile_id_counter(int pred) { + uint32_t tile_id_counter = 0; + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.eq.u32 p, %2, 1;\n\t" + "@p atom.global.add.u32 %0, [%1], 1; \n\t" + "}" + : "=r"(tile_id_counter) + : "l"(atomic_tile_index_), "r"(pred)); - CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline, - SchedPipelineState sched_pipeline_consumer_state) { - sched_pipeline.consumer_wait(sched_pipeline_consumer_state); - next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()]; - cutlass::arch::fence_view_async_shared(); - sched_pipeline.consumer_release(sched_pipeline_consumer_state); - return; - } + return tile_id_counter; + } - CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline, - SchedPipelineState sched_pipeline_producer_state) { - uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state); - // Wait for clcID buffer to become empty with a flipped phase - sched_pipeline.producer_acquire(sched_pipeline_producer_state); - auto is_leading_thread = cute::elect_one_sync(); - uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset; - uint32_t smem_addr = - cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]); - if (is_leading_thread) { - cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0); + CUTLASS_DEVICE auto fetch_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_consumer_state) { + sched_pipeline.consumer_wait(sched_pipeline_consumer_state); + next_linear_idx = smem_tile_counter[sched_pipeline_consumer_state.index()]; + cutlass::arch::fence_view_async_shared(); + sched_pipeline.consumer_release(sched_pipeline_consumer_state); + return; } - ++sched_pipeline_producer_state; - return sched_pipeline_producer_state; - } + CUTLASS_DEVICE auto advance_to_next_work(SchedPipeline &sched_pipeline, + SchedPipelineState sched_pipeline_producer_state) { + uint32_t mbarrier_addr = sched_pipeline.producer_get_barrier(sched_pipeline_producer_state); + // Wait for clcID buffer to become empty with a flipped phase + sched_pipeline.producer_acquire(sched_pipeline_producer_state); + auto is_leading_thread = cute::elect_one_sync(); + uint32_t tile_id_counter = fetch_tile_id_counter(is_leading_thread) + atomic_offset; + uint32_t smem_addr = + cute::cast_smem_ptr_to_uint(&smem_tile_counter[sched_pipeline_producer_state.index()]); + if (is_leading_thread) { + cute::store_shared_remote(tile_id_counter, smem_addr, mbarrier_addr, 0); + } - CUTLASS_DEVICE auto update_work_tile_info() { - linear_idx = next_linear_idx; - update_tile_idx(); - return; - } - }; - - // Allocate and alias shared memory to the kernel's shared storage type - extern __shared__ char shared_memory[]; - using SharedStorage = - SharedStorage; - SharedStorage &shared_storage = *reinterpret_cast(shared_memory); - - // Compute the number of tiles in M and N after tiling and assign scheduler - uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); - uint32_t tiles_in_n = uint32_t( - size(ceil_div(args.split_sections_range[args.num_tensors], size<2>(epilogue_tiler)))); - - TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace, - shared_storage.atomic_tile_counter); - - int block_rank_in_cluster = cute::block_rank_in_cluster(); - - // Shapes for accumulated tiles in mainloop and epilogue - auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{}); - auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{}); - - // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended - auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); - auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); - - // Number of threads assigned for various epilogue roles depending on quantization settings - static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; - static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; - static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0; - static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0; - static int constexpr NumSchedThreads = 32; - static int constexpr NumMainloopLoadThreads = 32; - static int constexpr NumEpilogueThreads = - NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; - - TmemAllocator tmem_allocator{}; - cutlass::arch::NamedBarrier tmem_allocation_result_barrier( - NumMmaThreadCount + NumEpilogueColQuantThreadCount, - cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); - - int warp_idx = cutlass::canonical_warp_idx_sync(); - - // warp assignment - bool is_mma_warp = (warp_idx == 0); - bool is_dma_warp = (warp_idx == 1); - bool is_sched_warp = (warp_idx == 2); - bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); - bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); - - typename MainloopPipeline::Params mainloop_pipeline_params; - if (is_dma_warp) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; - } - if (is_mma_warp) { - mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; - } - mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; - mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; - mainloop_pipeline_params.initializing_warp = 0; - mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + ++sched_pipeline_producer_state; + return sched_pipeline_producer_state; + } - MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, - cluster_shape, cute::true_type{}, // Perform barrier init - cute::true_type{}); // Delay mask calculation + CUTLASS_DEVICE auto update_work_tile_info() { + linear_idx = next_linear_idx; + update_tile_idx(); + return; + } + }; - MainloopPipelineState mainloop_pipe_consumer_state; - MainloopPipelineState mainloop_pipe_producer_state = - cutlass::make_producer_start_state(); + // Allocate and alias shared memory to the kernel's shared storage type + extern __shared__ char shared_memory[]; + using SharedStorage = + SharedStorage; + SharedStorage &shared_storage = *reinterpret_cast(shared_memory); - using AccumulatorPipeline = - cutlass::PipelineUmmaAsync; - using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; - using AccumulatorPipelineInitBarriers = cute::bool_constant; + // Compute the number of tiles in M and N after tiling and assign scheduler + uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); + uint32_t tiles_in_n = uint32_t( + size(ceil_div(args.split_sections_range[args.num_tensors], size<2>(epilogue_tiler)))); - AccumulatorPipelineState accumulator_pipe_consumer_state; - AccumulatorPipelineState accumulator_pipe_producer_state = - cutlass::make_producer_start_state(); + TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, tile_scheduler_workspace, + shared_storage.atomic_tile_counter); - typename AccumulatorPipeline::Params accumulator_pipeline_params; - if (is_mma_warp) { - accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; - } - if (is_epilogue_col_quant_warp) { - accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; - } - // Only one producer thread arrives on this barrier. - accumulator_pipeline_params.producer_arv_count = 1; - accumulator_pipeline_params.consumer_arv_count = - size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; - accumulator_pipeline_params.initializing_warp = 1; - AccumulatorPipeline accumulator_pipeline(shared_storage.accumulator, accumulator_pipeline_params, - cluster_shape, AccumulatorPipelineInitBarriers{}, - cute::true_type{}); // Delay mask calculation - typename SchedPipeline::Params sched_pipeline_params; - if (is_sched_warp) { - sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer; - } else { - sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer; - } - sched_pipeline_params.producer_blockid = 0; - sched_pipeline_params.producer_arv_count = 1; - sched_pipeline_params.consumer_arv_count = - NumSchedThreads + - cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); - sched_pipeline_params.transaction_bytes = sizeof(uint32_t); - sched_pipeline_params.initializing_warp = 3; - SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape); - SchedPipelineState sched_pipeline_consumer_state; - SchedPipelineState sched_pipeline_producer_state = - cutlass::make_producer_start_state(); - - typename SchedThrottlePipeline::Params sched_throttle_pipeline_params; - if (is_dma_warp) { - sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer; - } - if (is_sched_warp) { - sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer; - } - sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; - sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; - sched_throttle_pipeline_params.dst_blockid = 0; - sched_throttle_pipeline_params.initializing_warp = 4; - - SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle, - sched_throttle_pipeline_params); - SchedThrottlePipelineState sched_pipeline_throttle_consumer_state; - SchedThrottlePipelineState sched_pipeline_throttle_producer_state = - cutlass::make_producer_start_state(); - - if (warp_idx == 2 && elect_one_sync()) { - cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); - } - __syncthreads(); - - // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer - if (is_dma_warp) { - // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access). - cutlass::arch::warpgroup_reg_dealloc<32>(); - // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory. - Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N)); - Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); - - // Partition tensors for tiling according to the mainloop and cluster tilers. - Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); - Tensor gB_nk = - local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) - - // Shared memory tensors for pipeline - Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), - sAlayout); // (MMA,MMA_M,MMA_N,PIPE) - Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), - sBlayout); // (MMA,MMA_N,MMA_K,PIPE) - - // Determine warp/tile positioning int block_rank_in_cluster = cute::block_rank_in_cluster(); - ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx - // Partition global to local fragments for A and B - Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) - Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) - - Layout cta_layout_mnk = make_layout(cluster_shape); - Layout cta_layout_vmnk = - tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); - auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); - - auto [tAgA, tAsA] = - tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), - group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); - - auto [tBgB, tBsB] = - tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), - group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); - - uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); - uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); - if constexpr (kEnableRHTColQuant) { - if (elect_one_sync()) { - cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], - kTmaRhtTensorTransactionBytes); - copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), - tBsB(_, 0)); - } - } - do { - // is_first_wave indicates whether this scheduler wave is the first among a group. - bool is_first_wave = scheduler.is_first_wave(); - uint32_t skip_wait = is_first_wave; - auto tAgA_mk = tAgA(_, scheduler.tile_m(), _); - int k_tile = 0; - - sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state); - sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state); - ++sched_pipeline_throttle_producer_state; - CUTLASS_PRAGMA_NO_UNROLL - while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { - int k_tile_idx_n = scheduler.tile_n_base() + k_tile; - ++k_tile; - skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); - mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); - using BarrierType = typename MainloopPipeline::ProducerBarrierType; - BarrierType *tma_barrier = - mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); - int write_stage = mainloop_pipe_producer_state.index(); - ++mainloop_pipe_producer_state; - if (cute::elect_one_sync()) { - copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), - tAsA(_, write_stage)); - } - } - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - scheduler.update_work_tile_info(); - // scheduler.advance(); - } while (scheduler.is_valid()); - mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); - } else if (is_mma_warp) { - // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform. - cutlass::arch::warpgroup_reg_dealloc<32>(); - if constexpr (kEnableRHTColQuant) { - // Setup shared memory fragments for A and B tiles. + // Shapes for accumulated tiles in mainloop and epilogue + auto acc_shape_mma = make_shape(take<0, 2>(mainloop_tiler), _1{}, _1{}); + auto acc_shape_epilogue = make_shape(take<0, 2>(epilogue_tiler), _1{}, _1{}); + + // Shape of the accumulator fragment for the main loop pipeline, with pipeline stages appended + auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); + auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); + + // Number of threads assigned for various epilogue roles depending on quantization settings + static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; + static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; + static int constexpr NumMmaThreadCount = kEnableRHTColQuant ? 32 : 0; + static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant ? 1 : 0; + static int constexpr NumSchedThreads = 32; + static int constexpr NumMainloopLoadThreads = 32; + static int constexpr NumEpilogueThreads = + NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + NumMmaThreadCount + NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + // warp assignment + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_sched_warp = (warp_idx == 2); + bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); + bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + mainloop_pipeline_params.num_consumers = + NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + + MainloopPipeline mainloop_pipeline(shared_storage.mainloop, mainloop_pipeline_params, + cluster_shape, cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = + cutlass::make_producer_start_state(); + + using AccumulatorPipeline = + cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + using AccumulatorPipelineInitBarriers = cute::bool_constant; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = + cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_col_quant_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = + size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; + accumulator_pipeline_params.initializing_warp = 1; + AccumulatorPipeline accumulator_pipeline( + shared_storage.accumulator, accumulator_pipeline_params, cluster_shape, + AccumulatorPipelineInitBarriers{}, cute::true_type{}); // Delay mask calculation + typename SchedPipeline::Params sched_pipeline_params; + if (is_sched_warp) { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::ProducerConsumer; + } else { + sched_pipeline_params.role = SchedPipeline::ThreadCategory::Consumer; + } + sched_pipeline_params.producer_blockid = 0; + sched_pipeline_params.producer_arv_count = 1; + sched_pipeline_params.consumer_arv_count = + NumSchedThreads + + cluster_size * (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); + sched_pipeline_params.transaction_bytes = sizeof(uint32_t); + sched_pipeline_params.initializing_warp = 3; + SchedPipeline sched_pipeline(shared_storage.sched, sched_pipeline_params, cluster_shape); + SchedPipelineState sched_pipeline_consumer_state; + SchedPipelineState sched_pipeline_producer_state = + cutlass::make_producer_start_state(); + + typename SchedThrottlePipeline::Params sched_throttle_pipeline_params; + if (is_dma_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Producer; + } + if (is_sched_warp) { + sched_throttle_pipeline_params.role = SchedThrottlePipeline::ThreadCategory::Consumer; + } + sched_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; + sched_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; + sched_throttle_pipeline_params.dst_blockid = 0; + sched_throttle_pipeline_params.initializing_warp = 4; + + SchedThrottlePipeline sched_throttle_pipeline(shared_storage.sched_throttle, + sched_throttle_pipeline_params); + SchedThrottlePipelineState sched_pipeline_throttle_consumer_state; + SchedThrottlePipelineState sched_pipeline_throttle_producer_state = + cutlass::make_producer_start_state(); + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + + // Warp group roles: DMA (global->shared copy), MMA (tensor core gemm), scheduler, column quantizer, row quantizer + if (is_dma_warp) { + // Warp responsible for loading input from global to shared memory using TMA (Tensor Memory Access). + cutlass::arch::warpgroup_reg_dealloc<32>(); + // Get TMA tensors for input matrix A and B (Hadamard/transform matrix) from global memory. + Tensor mA = tma_load_a.get_tma_tensor(make_shape(M, packed_N)); + Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); + + // Partition tensors for tiling according to the mainloop and cluster tilers. + Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor gB_nk = + local_tile(mB, cluster_tile, make_coord(_, _, _), Step{}); // (BLK_N,BLK_K,k) + + // Shared memory tensors for pipeline Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE) Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + // Determine warp/tile positioning int block_rank_in_cluster = cute::block_rank_in_cluster(); ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx - // Allocate "fragments" -- these are actually umma smem descriptors - Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) - Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) - - mma.accumulate_ = UMMA::ScaleOut::Zero; - - tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, - &shared_storage.tmem_base_ptr); - __syncwarp(); - tmem_allocation_result_barrier.arrive(); - uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; - bulk_tmem_mma.data() = tmem_base_ptr; - // Wait until the B (Hadamard) tensor copy is complete - cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); - do { - uint32_t skip_wait = K_TILE_MAX <= 0; + // Partition global to local fragments for A and B + Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) + Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = + tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = + tma_partition(tma_load_a, get<2>(cta_coord_vmnk), make_layout(size<2>(cta_layout_vmnk)), + group_modes<0, 3>(tCsA), group_modes<0, 3>(tCgA)); + + auto [tBgB, tBsB] = + tma_partition(tma_load_b, get<1>(cta_coord_vmnk), make_layout(size<1>(cta_layout_vmnk)), + group_modes<0, 3>(tCsB), group_modes<0, 3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + if constexpr (kEnableRHTColQuant) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], + kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_, 0, 0), + tBsB(_, 0)); + } + } - auto barrier_token = - mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; + do { + // is_first_wave indicates whether this scheduler wave is the first among a group. + bool is_first_wave = scheduler.is_first_wave(); + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_, scheduler.tile_m(), _); + int k_tile = 0; + + sched_throttle_pipeline.producer_acquire(sched_pipeline_throttle_producer_state); + sched_throttle_pipeline.producer_commit(sched_pipeline_throttle_producer_state); + ++sched_pipeline_throttle_producer_state; CUTLASS_PRAGMA_NO_UNROLL - for (int k_tile = 0; - k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { - mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); - int read_stage = mainloop_pipe_consumer_state.index(); - auto tCrA_mk = tCrA(_, _, _, read_stage); - auto tCrB_nk = tCrB(_, _, 0, 0); - CUTLASS_PRAGMA_UNROLL - for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) { - int accumulator_k_block = - accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; - int tCrA_k_block = k_block * EpilogueUnrollFactor; - accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < EpilogueUnrollFactor; i++) { - auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i); - gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators); - } - - accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); - ++accumulator_pipe_producer_state; - } - auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; - ++mainloop_pipe_consumer_state; + while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { + int k_tile_idx_n = scheduler.tile_n_base() + k_tile; ++k_tile; - skip_wait = k_tile >= K_TILE_MAX; - mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); - barrier_token = - mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType *tma_barrier = + mainloop_pipeline.producer_get_barrier(mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + if (cute::elect_one_sync()) { + copy(tma_load_a.with(*tma_barrier, tma_mcast_mask_a), tAgA_mk(_, k_tile_idx_n), + tAsA(_, write_stage)); + } } + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; scheduler.update_work_tile_info(); + // scheduler.advance(); } while (scheduler.is_valid()); - tmem_allocator.release_allocation_lock(); - accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); - tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); - } - } else if (is_sched_warp) { - // Scheduler warp manages tile assignment and pipeline progress for warps - cutlass::arch::warpgroup_reg_dealloc<32>(); - do { - sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state); - sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state); - ++sched_pipeline_throttle_consumer_state; - sched_pipeline_producer_state = - scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state); - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - scheduler.update_work_tile_info(); - } while (scheduler.is_valid()); - } else if (is_epilogue_col_quant_warp) { - // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage, - // and writing result tensors/scales to global memory. - cutlass::arch::warpgroup_reg_alloc<192>(); - if constexpr (kEnableRHTColQuant) { - using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; - - auto acc_epilogue_pipelined_shape = - append(acc_shape_epilogue, Int{}); - auto bulk_tmem_epilogue_layout = make_layout( - acc_epilogue_pipelined_shape, - make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler))); - auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); - - // Use 256-bit fragments for aligned bulk stores - static int constexpr FragmentSize = 256 / sizeof_bits_v; - - // Wait for TMEM allocation for this pipeline to finish - tmem_allocation_result_barrier.arrive_and_wait(); - uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; - bulk_tmem_epilogue.data() = tmem_base_ptr; - int global_thread_idx = threadIdx.x; - int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; - // g2s load all global_d_amax - CUTLASS_PRAGMA_NO_UNROLL - for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueColQuantThreadCount) { - shared_storage.global_d_amax[g] = - __ldg(reinterpret_cast(args.global_d_amax_list[g])); - } - - size_t rng_seed = 0; - size_t rng_offset = 0; - // Setup RNG for stochastic rounding - if constexpr (kEnableStochasticRounding) { - rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; - rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; - } - int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); - - // Determine quantization scale factor layouts/output splits for this group - TSFDLayout sfd_layout; - int cur_N = args.split_sections[group_idx]; - if constexpr (kEnableSwizzleSFOutput) { - sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); - } else { - sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), - make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); - } - // Build output tensors for columns and their quant scales - Tensor mD = make_tensor( - cute::subbyte_iterator(reinterpret_cast(args.output_colwise_list[group_idx])), - make_shape(M, cur_N), DStride{}); // (M,packed_N) - Tensor gD_mn = - local_tile(mD, epilogue_tiler, make_coord(_, _, _), Step<_1, _1, X>{}); // (BLK_M,BLK_N) - - Tensor mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( - args.output_colwise_scale_inv_list[group_idx])), - sfd_layout); - Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), - Step<_1, _1, X>{}); // (BLK_M,BLK_N) - - Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); - - // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors - auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); - auto tiled_r2g = - make_tiled_copy_D(Copy_Atom{}, tiled_t2r); - auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); - auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); - - cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, - cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; - static constexpr float fp4_max_inv = 1.0f / fp4_max; - float c_global_amax_val = shared_storage.global_d_amax[group_idx]; - float global_encode_scale = c_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - float global_decode_scale = 1.0f / global_encode_scale; - - // Scaling factor for fast math path - float global_encode_scale_multiplier = 1.0f; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + // This warp executes the main tensor core matrix-multiply-accumulate for the Hadamard transform. + cutlass::arch::warpgroup_reg_dealloc<32>(); + if constexpr (kEnableRHTColQuant) { + // Setup shared memory fragments for A and B tiles. + Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), + sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Allocate "fragments" -- these are actually umma smem descriptors + Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, + &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + // Wait until the B (Hadamard) tensor copy is complete + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + + auto barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_, _, _, read_stage); + auto tCrB_nk = tCrB(_, _, 0, 0); + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) { + int accumulator_k_block = + accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; + int tCrA_k_block = k_block * EpilogueUnrollFactor; + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < EpilogueUnrollFactor; i++) { + auto accumulators = bulk_tmem_mma(_, _, _, accumulator_k_block + i); + gemm(mma, tCrA_mk(_, _, tCrA_k_block + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; + } + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); + barrier_token = + mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state, skip_wait); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); } - + } else if (is_sched_warp) { + // Scheduler warp manages tile assignment and pipeline progress for warps + cutlass::arch::warpgroup_reg_dealloc<32>(); do { + sched_throttle_pipeline.consumer_wait(sched_pipeline_throttle_consumer_state); + sched_throttle_pipeline.consumer_release(sched_pipeline_throttle_consumer_state); + ++sched_pipeline_throttle_consumer_state; + sched_pipeline_producer_state = + scheduler.advance_to_next_work(sched_pipeline, sched_pipeline_producer_state); scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } else if (is_epilogue_col_quant_warp) { + // Warp responsible for quantizing output of Hadamard transform to FP4 for columnwise usage, + // and writing result tensors/scales to global memory. + cutlass::arch::warpgroup_reg_alloc<192>(); + if constexpr (kEnableRHTColQuant) { + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + auto acc_epilogue_pipelined_shape = + append(acc_shape_epilogue, Int{}); + auto bulk_tmem_epilogue_layout = make_layout( + acc_epilogue_pipelined_shape, + make_stride(stride<0>(bulk_tmem_mma), Int<0>{}, Int<0>{}, size<1>(epilogue_tiler))); + auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); + + // Use 256-bit fragments for aligned bulk stores + static int constexpr FragmentSize = 256 / sizeof_bits_v; + + // Wait for TMEM allocation for this pipeline to finish + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; + // g2s load all global_d_amax CUTLASS_PRAGMA_NO_UNROLL - for (int k_tile = 0; - k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); - ++k_tile) { - int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); - - int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset); - - if (cur_group_idx != group_idx) { - group_idx = cur_group_idx; - c_global_amax_val = shared_storage.global_d_amax[group_idx]; - // update amax - global_encode_scale = c_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / c_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - global_decode_scale = 1.0f / global_encode_scale; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; - } - cur_N = args.split_sections[group_idx]; - if constexpr (kEnableSwizzleSFOutput) { - sfd_layout = - tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); - } else { - sfd_layout = - make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), - make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); - } - // update tensor - mD = make_tensor(cute::subbyte_iterator( - reinterpret_cast(args.output_colwise_list[group_idx])), - make_shape(M, cur_N), DStride{}); - gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), - Step<_1, _1, X>{}); // (BLK_M,BLK_N) - mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( - args.output_colwise_scale_inv_list[group_idx])), - sfd_layout); - gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), - Step<_1, _1, X>{}); // (BLK_M,BLK_N) - - gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); - } - int group_start_offset = args.split_sections_range[group_idx]; - int local_tile_n_idx = - (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler); - Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx); - - Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx); - accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); - - auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); - Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - - Tensor tTR_rAcc = - make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) - Tensor tDrD = make_tensor(shape(tDgD)); - Tensor tTR_rAcc_frag = - recast>(coalesce(tTR_rAcc)); - Tensor tDrD_frag = recast>(coalesce(tDrD)); - - Tensor src = thr_r2g.retile_S(tDrD); - Tensor dst = thr_r2g.retile_D(tDgD); - - Tensor tDgSFD_view = make_tensor( - tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), - make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); - Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); - Tensor tDrSFD = make_tensor(shape(tDgSFD)); - - static int constexpr NumVecs = size(tDgD) / VectorSize; - Tensor tD_rRowSFD_frg = recast>(tDrSFD); - - // Compute amax and quantization scales for this tile - cutlass::maximum_absolute_value_reduction, - true> - amax_reduction; - cutlass::Array vec_maxs; - cutlass::Array pvscales; - // Copy from TMEM to registers - copy(tiled_t2r, tDtAcc, tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); - ++accumulator_pipe_consumer_state; - - if constexpr (!kUseFastMath) { - // Downcast to BF16 for bit-wise compatibility with - // unfused kernels - auto convert_accum_to_bf16 = - cutlass::NumericArrayConverter{}; - auto convert_bf16_to_accum = - cutlass::NumericArrayConverter{}; - tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); - tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); - } + for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueColQuantThreadCount) { + shared_storage.global_d_amax[g] = + __ldg(reinterpret_cast(args.global_d_amax_list[g])); + } - auto compute_frgs = reinterpret_cast *>( - tTR_rAcc_frag.data()); - auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < NumVecs; v++) { - vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); - } + size_t rng_seed = 0; + size_t rng_offset = 0; + // Setup RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); + + // Determine quantization scale factor layouts/output splits for this group + TSFDLayout sfd_layout; + int cur_N = args.split_sections[group_idx]; + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // Build output tensors for columns and their quant scales + Tensor mD = make_tensor( + cute::subbyte_iterator(reinterpret_cast(args.output_colwise_list[group_idx])), + make_shape(M, cur_N), DStride{}); // (M,packed_N) + Tensor gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) - if constexpr (kUseFastMath) { - // Fast math: multiply with precomputed reciprocal - pvscales = cutlass::multiplies>{}( - vec_maxs, global_encode_scale_multiplier); - } else { - // Accurate math: perform division - pvscales = - cutlass::divides>{}(vec_maxs, fp4_max); - pvscales = cutlass::multiplies>{}( - pvscales, global_encode_scale); - } - auto pvscales_cvted = - cutlass::NumericArrayConverter{}(pvscales); - - tD_rRowSFD_frg(_0{}) = pvscales_cvted; - auto qpvscale_ups = cutlass::NumericArrayConverter{}( - tD_rRowSFD_frg(_0{})); - auto qpvscale_scaled = cutlass::multiplies>{}( - qpvscale_ups, global_decode_scale); - cutlass::Array acc_scales; - if constexpr (kUseFastMath) { - // Fast math: compute approximate reciprocal - acc_scales = - cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); - } else { - // Accurate math: compute reciprocal with division - acc_scales = cutlass::divides>{}( - 1.0, qpvscale_scaled); - } + Tensor mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( + args.output_colwise_scale_inv_list[group_idx])), + sfd_layout); + Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + + Tensor gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + + // Setup tile-level TMEM (t2r) and global memory (r2g) copy descriptors + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_, _, _, _0{})); + auto tiled_r2g = + make_tiled_copy_D(Copy_Atom{}, tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); + auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); + + cutlass::arch::NamedBarrier::sync(NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float c_global_amax_val = shared_storage.global_d_amax[group_idx]; + float global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + float global_decode_scale = 1.0f / global_encode_scale; + + // Scaling factor for fast math path + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + + do { + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); + ++k_tile) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset); + + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + c_global_amax_val = shared_storage.global_d_amax[group_idx]; + // update amax + global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + cur_N = args.split_sections[group_idx]; + if constexpr (kEnableSwizzleSFOutput) { + sfd_layout = + tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(M, cur_N), Step<_2, _1>{}); + } else { + sfd_layout = + make_layout(make_shape(M, make_shape(Int{}, cur_N / SFVecSize)), + make_stride(cur_N / SFVecSize, make_stride(_0{}, _1{}))); + } + // update tensor + mD = make_tensor(cute::subbyte_iterator( + reinterpret_cast(args.output_colwise_list[group_idx])), + make_shape(M, cur_N), DStride{}); + gD_mn = local_tile(mD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) + mSFD = make_tensor(make_gmem_ptr(reinterpret_cast( + args.output_colwise_scale_inv_list[group_idx])), + sfd_layout); + gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_, _, _), + Step<_1, _1, X>{}); // (BLK_M,BLK_N) - // Prepare stochastic rounding random state if enabled - uint4 random_uint4 = uint4{0, 0, 0, 0}; - transformer_engine::curanddx::detail::philox4x32_native_state< - NVTE_BUILD_NUM_PHILOX_ROUNDS> - rng; - // "Prefetch" a stochastic rounding state for the first tile - if constexpr (kEnableStochasticRounding) { - const size_t rng_sequence = global_thread_idx + k_tile * 512 + - scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; - rng.init(rng_seed, rng_sequence, rng_offset); - } - CUTLASS_PRAGMA_UNROLL - // Apply round/quantize to each fragment, with or without stochastic rounding - for (int v = 0; v < NumVecs; v++) { - auto acc_scale = cutlass::minimum_with_nan_propagation{}( - acc_scales[v], cutlass::platform::numeric_limits::max()); - if constexpr (kEnableStochasticRounding) { - random_uint4 = rng.generate4(); - output_frgs[v] = StochasticNumericConverter( - cutlass::multiplies>{}( - compute_frgs[v], acc_scale), - *reinterpret_cast *>(&random_uint4)); - } else { - output_frgs[v] = cutlass::NumericArrayConverter{}( - cutlass::multiplies>{}( - compute_frgs[v], acc_scale)); + gD_mn_view = tiled_divide(gD_mn, take<0, 2>(epilogue_tiler)); + } + int group_start_offset = args.split_sections_range[group_idx]; + int local_tile_n_idx = + (global_tile_n_offset - group_start_offset) / size<1>(epilogue_tiler); + Tensor tDgD_mn = gD_mn_view(_, _, _, scheduler.tile_m(), local_tile_n_idx); + + Tensor tDgSFD_mn = gSFD_mn(_, _, scheduler.tile_m(), local_tile_n_idx); + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto Acc = bulk_tmem_epilogue(_, _, _, accumulator_pipe_consumer_state.index()); + Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + + Tensor tTR_rAcc = make_tensor( + shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + Tensor tDrD = make_tensor(shape(tDgD)); + Tensor tTR_rAcc_frag = + recast>(coalesce(tTR_rAcc)); + Tensor tDrD_frag = recast>(coalesce(tDrD)); + + Tensor src = thr_r2g.retile_S(tDrD); + Tensor dst = thr_r2g.retile_D(tDgD); + + Tensor tDgSFD_view = make_tensor( + tDgSFD_mn.data(), make_layout(make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); + Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); + Tensor tDrSFD = make_tensor(shape(tDgSFD)); + + static int constexpr NumVecs = size(tDgD) / VectorSize; + Tensor tD_rRowSFD_frg = recast>(tDrSFD); + + // Compute amax and quantization scales for this tile + cutlass::maximum_absolute_value_reduction< + cutlass::Array, true> + amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // Copy from TMEM to registers + copy(tiled_t2r, tDtAcc, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with + // unfused kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = + convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + tTR_rAcc_frag(_1{}) = + convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); } - } - // Write quantized FP4 tile and dequant scale to gmem - copy(tiled_r2g, src, dst); - copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD); - } - scheduler.update_work_tile_info(); - } while (scheduler.is_valid()); - } - } else if (is_epilogue_row_quant_warp) { - // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage. - cutlass::arch::warpgroup_reg_alloc<136>(); - if constexpr (kEnableRowQuant) { - using S2RVectorType = uint128_t; - - int global_thread_idx = threadIdx.x; - int local_thread_idx = global_thread_idx % 256; - size_t rng_seed = 0; - size_t rng_offset = 0; - // g2s load all global_a_amax for all groups/tensors - CUTLASS_PRAGMA_NO_UNROLL - for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueRowQuantThreadCount) { - shared_storage.global_a_amax[g] = - __ldg(reinterpret_cast(args.global_a_amax_list[g])); - } - // RNG for stochastic rounding - if constexpr (kEnableStochasticRounding) { - rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; - rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; - } - // Input/output tensors/partitions for row quant warp - Tensor mQA = - make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA)); - Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); - Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); - - Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _), - Step<_1, X, _1>{}); // (BLK_M,BLK_N) - // Swizzled shared memory A tile, with layout - Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>( - coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), - sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) - - // Set up layouts for partitioning – tile-by-warp, with vector granularity - using S2RWarpLayout = Layout>; - using WarpGroupLayout = Layout>; - using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); - using S2RValLayout = Layout, _1>>; - using S2RAtomA = Copy_Atom; - using R2GAtomQA = Copy_Atom; - using R2GAtomSFA = Copy_Atom; - auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); - auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{}); - auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{}); - - auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); - auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); - auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx); - Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) - - // Allocate temporary register tensors for copying quantization => output - Tensor tQArA = make_tensor_like( - make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) - Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn); - Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); - - Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); - Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); - - // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8 - // in order to go over the reserved named barrier count. - constexpr int row_quant_barrier_id = 2; - cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); - - int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); - float a_global_amax_val = shared_storage.global_a_amax[group_idx]; - // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} - static constexpr float fp4_max = 6.0f; - static constexpr float fp8_max = 448.0f; - static constexpr float fp4_max_inv = 1.0f / fp4_max; - float global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - - float global_decode_scale = 1.0f / global_encode_scale; - float global_encode_scale_multiplier = 1.0f; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; - } - auto sfa_converter = cutlass::NumericConverter{}; - do { - CUTLASS_PRAGMA_NO_UNROLL - for (int k_tile = 0; - k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { - int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); - - int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset); - if (cur_group_idx != group_idx) { - group_idx = cur_group_idx; - a_global_amax_val = shared_storage.global_a_amax[group_idx]; - // Update group quantization parameters/scaling - global_encode_scale = a_global_amax_val > 0.0f - ? cutlass::minimum_with_nan_propagation{}( - (fp8_max * fp4_max) / a_global_amax_val, - cutlass::platform::numeric_limits::max()) - : 1.0f; - global_decode_scale = 1.0f / global_encode_scale; - if constexpr (kUseFastMath) { - global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + auto compute_frgs = reinterpret_cast *>( + tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); } - } - auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); - auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); - auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state); - mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); - copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); - cutlass::arch::fence_view_async_shared(); - mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state); - ++mainloop_pipe_consumer_state; - ++k_tile; + pvscales = cutlass::multiplies>{}( + vec_maxs, global_encode_scale_multiplier); + auto pvscales_cvted = + cutlass::NumericArrayConverter{}(pvscales); - // static int constexpr NumVecs = size(tQArA) / VectorSize; - cutlass::maximum_absolute_value_reduction, - true> - amax_reduction; - auto compute_frgs = reinterpret_cast *>(tQArA.data()); - auto output_frgs = - reinterpret_cast *>(raw_pointer_cast(tQArQA.data())); - Tensor amax = - make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); - Tensor pvscales = make_tensor_like(amax); - transformer_engine::curanddx::detail::philox4x32_native_state< - NVTE_BUILD_NUM_PHILOX_ROUNDS> - rng; - if constexpr (kEnableStochasticRounding) { - const size_t rng_sequence = global_thread_idx + k_tile * 512 + - scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + - tiles_in_m * tiles_in_n * K_TILE_MAX * 512; - rng.init(rng_seed, rng_sequence, rng_offset); - } - CUTLASS_PRAGMA_UNROLL - for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) { - auto amax_view = group_modes<1, rank(amax)>(amax); - auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales); - auto compute_frgs_up = - cutlass::NumericArrayConverter{}( - compute_frgs[v]); - amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up); - if constexpr (kUseFastMath) { - // Fast math: multiply with precomputed reciprocal - pvscales_view(_0{}, v) = cutlass::multiplies{}( - amax_view(_0{}, v), global_encode_scale_multiplier); - } else { - // Accurate math: perform division - pvscales_view(_0{}, v) = - cutlass::divides{}(amax_view(_0{}, v), fp4_max); - pvscales_view(_0{}, v) = cutlass::multiplies{}( - pvscales_view(_0{}, v), global_encode_scale); - } - filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v)); - auto qpvscale_ups = - cutlass::NumericConverter{}(filter(tQArSFA)(v)); + tD_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}( + tD_rRowSFD_frg(_0{})); auto qpvscale_scaled = - cutlass::multiplies{}(qpvscale_ups, global_decode_scale); - ElementAccumulator acc_scales; + cutlass::multiplies>{}( + qpvscale_ups, global_decode_scale); + cutlass::Array acc_scales; if constexpr (kUseFastMath) { // Fast math: compute approximate reciprocal acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); } else { // Accurate math: compute reciprocal with division - acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + acc_scales = cutlass::divides>{}( + 1.0, qpvscale_scaled); } - auto acc_scale = cutlass::minimum_with_nan_propagation{}( - acc_scales, cutlass::platform::numeric_limits::max()); + + // Prepare stochastic rounding random state if enabled uint4 random_uint4 = uint4{0, 0, 0, 0}; + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; + // "Prefetch" a stochastic rounding state for the first tile if constexpr (kEnableStochasticRounding) { - random_uint4 = rng.generate4(); - output_frgs[v] = StochasticNumericConverter( - cutlass::multiplies>{}( - compute_frgs_up, acc_scale), - *reinterpret_cast *>(&random_uint4)); - } else { - output_frgs[v] = - cutlass::NumericArrayConverter{}( - cutlass::multiplies>{}( - compute_frgs_up, acc_scale)); + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + // Apply round/quantize to each fragment, with or without stochastic rounding + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], cutlass::platform::numeric_limits::max()); + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = + cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], acc_scale)); + } } + + // Write quantized FP4 tile and dequant scale to gmem + copy(tiled_r2g, src, dst); + copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, tDrSFD, tDgSFD); } - copy(tiled_r2g_QA, tQArQA, tQAgQA_mn); - copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn)); + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + } else if (is_epilogue_row_quant_warp) { + // Warp responsible for quantizing the input (before Hadamard transform) to FP4 for row-wise usage. + cutlass::arch::warpgroup_reg_alloc<136>(); + if constexpr (kEnableRowQuant) { + using S2RVectorType = uint128_t; + + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % 256; + size_t rng_seed = 0; + size_t rng_offset = 0; + // g2s load all global_a_amax for all groups/tensors + CUTLASS_PRAGMA_NO_UNROLL + for (int g = local_thread_idx; g < args.num_tensors; g += NumEpilogueRowQuantThreadCount) { + shared_storage.global_a_amax[g] = + __ldg(reinterpret_cast(args.global_a_amax_list[g])); } - // scheduler.advance(); - scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); - ++sched_pipeline_consumer_state; - scheduler.update_work_tile_info(); - } while (scheduler.is_valid()); - } + // RNG for stochastic rounding + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + // Input/output tensors/partitions for row quant warp + Tensor mQA = + make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, packed_N), dQA)); + Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_, _, _), Step<_1, X, _1>{}); + Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); + + Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_, _, _), + Step<_1, X, _1>{}); // (BLK_M,BLK_N) + // Swizzled shared memory A tile, with layout + Tensor sA = as_position_independent_swizzle_tensor(group_modes<0, 2>( + coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), + sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) + + // Set up layouts for partitioning – tile-by-warp, with vector granularity + using S2RWarpLayout = Layout>; + using WarpGroupLayout = Layout>; + using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); + using S2RValLayout = Layout, _1>>; + using S2RAtomA = Copy_Atom; + using R2GAtomQA = Copy_Atom; + using R2GAtomSFA = Copy_Atom; + auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_QA = make_tiled_copy(R2GAtomQA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_SFA = make_tiled_copy(R2GAtomSFA{}, S2RThreadLayout{}, S2RValLayout{}); + + auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); + auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); + auto thr_r2g_SFA = tiled_r2g_SFA.get_slice(local_thread_idx); + Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) + + // Allocate temporary register tensors for copying quantization => output + Tensor tQArA = make_tensor_like( + make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) + Tensor tQAgQA = thr_r2g_QA.partition_S(gQA_mn); + Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); + + Tensor tQAgSFA = thr_r2g_SFA.partition_S(gSFA_mn); + Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); + + // Will result in barrier_id=10 passed to bar.sync instr as cutlass adds 8 + // in order to go over the reserved named barrier count. + constexpr int row_quant_barrier_id = 2; + cutlass::arch::NamedBarrier::sync(NumEpilogueRowQuantThreadCount, row_quant_barrier_id); + + int group_idx = GetGroupIdx(&args, scheduler.tile_n_base() * size<1>(epilogue_tiler)); + float a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + static constexpr float fp4_max_inv = 1.0f / fp4_max; + float global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + + float global_decode_scale = 1.0f / global_encode_scale; + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + auto sfa_converter = cutlass::NumericConverter{}; + do { + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; + k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n();) { + int global_tile_n_offset = (scheduler.tile_n_base() + k_tile) * size<1>(epilogue_tiler); + + int cur_group_idx = GetGroupIdx(&args, global_tile_n_offset); + if (cur_group_idx != group_idx) { + group_idx = cur_group_idx; + a_global_amax_val = shared_storage.global_a_amax[group_idx]; + // Update group quantization parameters/scaling + global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + global_decode_scale = 1.0f / global_encode_scale; + global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + } - } else { - cutlass::arch::warpgroup_reg_dealloc<32>(); - } + auto tQAgSFA_mn = + tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto barrier_token = mainloop_pipeline.consumer_try_wait(mainloop_pipe_consumer_state); + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); + cutlass::arch::fence_view_async_shared(); + mainloop_pipeline.consumer_release(mainloop_pipe_consumer_state); + ++mainloop_pipe_consumer_state; + ++k_tile; + + // static int constexpr NumVecs = size(tQArA) / VectorSize; + cutlass::maximum_absolute_value_reduction< + cutlass::Array, true> + amax_reduction; + auto compute_frgs = reinterpret_cast *>(tQArA.data()); + auto output_frgs = reinterpret_cast *>( + raw_pointer_cast(tQArQA.data())); + Tensor amax = + make_tensor(prepend(take<1, rank(tQArA)>(tQArA.shape()), _1{})); + Tensor pvscales = make_tensor_like(amax); + transformer_engine::curanddx::detail::philox4x32_native_state< + NVTE_BUILD_NUM_PHILOX_ROUNDS> + rng; + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512 + + tiles_in_m * tiles_in_n * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < size<1>(group_modes<1, rank(tQArA)>(tQArA)); v++) { + auto amax_view = group_modes<1, rank(amax)>(amax); + auto pvscales_view = group_modes<1, rank(pvscales)>(pvscales); + auto compute_frgs_up = + cutlass::NumericArrayConverter{}( + compute_frgs[v]); + amax_view(_0{}, v) = amax_reduction(ElementAccumulator(0), compute_frgs_up); + pvscales_view(_0{}, v) = cutlass::multiplies{}( + amax_view(_0{}, v), global_encode_scale_multiplier); + filter(tQArSFA)(v) = sfa_converter(pvscales_view(_0{}, v)); + auto qpvscale_ups = + cutlass::NumericConverter{}(filter(tQArSFA)(v)); + auto qpvscale_scaled = + cutlass::multiplies{}(qpvscale_ups, global_decode_scale); + ElementAccumulator acc_scales; + if constexpr (kUseFastMath) { + // Fast math: compute approximate reciprocal + acc_scales = cutlass::reciprocal_approximate_ftz{}( + qpvscale_scaled); + } else { + // Accurate math: compute reciprocal with division + acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + } + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales, cutlass::platform::numeric_limits::max()); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale), + *reinterpret_cast *>(&random_uint4)); + } else { + output_frgs[v] = + cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs_up, acc_scale)); + } + } + copy(tiled_r2g_QA, tQArQA, tQAgQA_mn); + copy(tiled_r2g_SFA, filter(tQArSFA), filter(tQAgSFA_mn)); + } + // scheduler.advance(); + scheduler.fetch_next_work(sched_pipeline, sched_pipeline_consumer_state); + ++sched_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + + } else { + cutlass::arch::warpgroup_reg_dealloc<32>(); + } + } // sm100 compile guard end } // NOLINT(readability/fn_size) template >{}(vec_maxs, global_encode_scale_multiplier); - } else { - // Accurate math: perform division - pvscales = cutlass::divides>{}(vec_maxs, fp4_max); - pvscales = cutlass::multiplies>{}(pvscales, global_encode_scale); - } + pvscales = cutlass::multiplies>{}(vec_maxs, global_encode_scale_multiplier); auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales); tC_rRowSFD_frg(_0{}) = pvscales_cvted; @@ -548,6 +543,7 @@ rht_gemm_device(MShape M, NShape N, KShape K, ClusterTileShape cluster_tile, tile_idx_n = (linear_tile_idx / tiles_in_m) * K_TILE_MAX; } while (tile_idx_m < tiles_in_m && tile_idx_n < tiles_in_n); } + } } // this function computes RHT-GEMM for 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 new file mode 100644 index 0000000000..99060ab627 --- /dev/null +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -0,0 +1,1370 @@ +/************************************************************************* + * 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 "common/common.h" +#include "common/util/cuda_runtime.h" +#include "common/util/curanddx.hpp" +#include "common/util/ptx.cuh" +#include "common/utils.cuh" +#include "customized_pipeline.cuh" +#include "cutlass/arch/barrier.h" +#include "cutlass/arch/reg_reconfig.h" +#include "cutlass/cluster_launch.hpp" +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cutlass/fast_math.h" +#include "cutlass/float8.h" +#include "cutlass/float_subbyte.h" +#include "cutlass/gemm/collective/builders/sm100_common.inl" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_types.h" +#include "cutlass/pipeline/pipeline.hpp" +#include "cutlass/platform/platform.h" +#include "cutlass/util/GPU_Clock.hpp" +#include "cutlass/util/command_line.h" +#include "cutlass/util/print_error.hpp" + +// clang-format off + +namespace transformer_engine { +namespace detail { +namespace { + +using namespace cute; + +struct CLCResponse { uint32_t data[4] = {0}; }; + +constexpr int kFp4ConvertChunkElements = 8; +constexpr int kFp4ConvertFullElements = 16; +constexpr int kFp4RbitsPerChunk = 2; +constexpr int kFp4ChunkCount = kFp4ConvertFullElements / kFp4ConvertChunkElements; + + +CUTLASS_DEVICE +cutlass::Array StochasticNumericConverterBase( + cutlass::Array const &input, + cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + auto output_ptr = reinterpret_cast(&output); + constexpr bool has_rs = ARCH_HAS_STOCHASTIC_ROUNDING; + if constexpr (has_rs) { + asm volatile( + "{\n" + "cvt.rs.satfinite.e2m1x4.f32 %0, {%5, %4, %3, %2}, %10;\n" + "cvt.rs.satfinite.e2m1x4.f32 %1, {%9, %8, %7, %6}, %11;\n" + "}" + : "=h"(output_ptr[0]), "=h"(output_ptr[1]) + : "f"(input[0]), "f"(input[1]), "f"(input[2]), "f"(input[3]), "f"(input[4]), "f"(input[5]), + "f"(input[6]), "f"(input[7]), "r"(rbits[0]), "r"(rbits[1])); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + return output; +} + +CUTLASS_DEVICE +cutlass::Array +StochasticNumericConverter(cutlass::Array const &input, + cutlass::Array const &rbits) { + using result_type = cutlass::Array; + result_type output; + cutlass::Array *result_ptr = + reinterpret_cast *>(&output); + cutlass::Array const *source_ptr = + reinterpret_cast const *>(&input); + cutlass::Array const *rbits_ptr = + reinterpret_cast const *>(&rbits); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < kFp4ChunkCount; i++) { + result_ptr[i] = StochasticNumericConverterBase(source_ptr[i], rbits_ptr[i]); + } + return output; +} + +template < + class ElementA, + class ElementB, + class ASmemLayout, + class BSmemLayout, + class ClusterShape, + int AccumulatorPipelineStageCount_, + int EpilogueUnrollFactor_, + int SchedulerPipelineStageCount_> +struct SharedStorage { + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr EpilogueUnrollFactor = EpilogueUnrollFactor_; + using AtomThrShapeMNK = cute::Shape<_1, _1, _1>; + + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + using AccumulatorPipelineStorage = typename AccumulatorPipeline::SharedStorage; + + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = cutlass::detail::CustomizedPipelineTmaUmmaAsync< + MainloopPipelineStageCount, + Shape<_1,_1,_1>, + AtomThrShapeMNK>; + using MainloopPipelineStorage = typename MainloopPipeline::SharedStorage; + using CLCPipeline = cutlass::PipelineCLCFetchAsync; + using CLCPipelineStorage = typename CLCPipeline::SharedStorage; + using CLCThrottlePipeline = cutlass::PipelineAsync; + using CLCThrottlePipelineStorage = typename CLCThrottlePipeline::SharedStorage; + + struct TensorStorage : cute::aligned_struct<128, _1> { + // cute::array_aligned> smem_A; + cute::array_aligned> smem_A; + cute::array_aligned> smem_B; + } tensors; + + alignas(16) AccumulatorPipelineStorage accumulator; + alignas(16) MainloopPipelineStorage mainloop; + alignas(16) cute::uint64_t tma_barrier[1]; + alignas(16) CLCPipelineStorage clc; + alignas(16) CLCThrottlePipelineStorage clc_throttle; + alignas(16) CLCResponse clc_response[SchedulerPipelineStageCount_]; + uint32_t tmem_base_ptr; +}; + +template +__launch_bounds__(512, 1) +__global__ static void row_col_rht_gemm_device( + MShape M, + NShape N, + KShape K, + ClusterShape cluster_shape, + ClusterTileShape cluster_tile, + TA const* A, + AStride dA, + ASmemLayout sAlayout, + CUTE_GRID_CONSTANT TmaLoadA const tma_load_a, + TB const* B, + BStride dB, + BSmemLayout sBlayout, + CUTE_GRID_CONSTANT TmaLoadB const tma_load_b, + TD* D, + DStride dD, + DSmemLayout, + TSFD* SFD, + TSFDLayout sfd_layout, + TQA* QA, + QAStride dQA, + TSFA* SFA, + TSFALayout sfa_layout, + TiledMMA mma, + float const* a_global_amax, + float const* c_global_amax, + const size_t* rng_state) { + using namespace cute; + + // Abort immediately if compilation is not supported + constexpr bool is_blackwell_arch = ARCH_BLACKWELL_FAMILY; + if constexpr (!is_blackwell_arch) { + NVTE_DEVICE_ERROR("RHT fusion is only supported on Blackwell."); + return; + } else { + static_assert(kEnableRHTColQuant_ || kEnableRowQuant_, + "row_col_rht_gemm_device must generate row-wise " + "and/or column-wise output."); +#if !defined(CUTLASS_ARCH_CLC_ENABLED) + CUTLASS_NOT_IMPLEMENTED(); + return; +#endif + + using X = Underscore; + // static constexpr bool kApplyStochasticRounding = true; + using ElementAccumulator = float; + static int constexpr K_PIPE_MAX = size<3>(ASmemLayout{}); + using AtomThrShapeMNK = Shape(typename TiledMMA::ThrLayoutVMNK{})), _1, _1>; + static uint32_t constexpr kTmaTransactionBytes = cutlass::bits_to_bytes( + size(AtomThrShapeMNK{}) * cosize(take<0,3>(ASmemLayout{})) * cute::sizeof_bits_v); + static constexpr bool kEnableStochasticRounding = kEnableStochasticRounding_; + static constexpr bool kEnableRHTColQuant = kEnableRHTColQuant_; + static constexpr bool kEnableRowQuant = kEnableRowQuant_; + static constexpr bool kUseFastMath = kUseFastMath_; + static int constexpr RhtTensorSize = 16; + static int constexpr kTmaRhtTensorTransactionBytes = cutlass::bits_to_bytes( + RhtTensorSize * RhtTensorSize * cute::sizeof_bits_v); + static int constexpr AccumulatorPipelineStageCount = AccumulatorPipelineStageCount_; + static int constexpr SchedulerPipelineStageCount = SchedulerPipelineStageCount_; + + static int constexpr MainloopPipelineStageCount = size<3>(ASmemLayout{}); + using MainloopPipeline = cutlass::detail::CustomizedPipelineTmaUmmaAsync< + MainloopPipelineStageCount, + ClusterShape, + AtomThrShapeMNK>; + using MainloopPipelineState = typename MainloopPipeline::PipelineState; + using CLCPipeline = cutlass::PipelineCLCFetchAsync; + using CLCPipelineState = typename CLCPipeline::PipelineState; + using CLCThrottlePipeline = cutlass::PipelineAsync; + using CLCThrottlePipelineState = typename CLCThrottlePipeline::PipelineState; + + static_assert(ClusterShape{} == Shape<_1,_1,_1>{}, "ClusterShape must be Shape<_1,_1,_1>"); + + using TmemAllocator = cute::TMEM::Allocator1Sm; + static int constexpr VectorSize = RhtTensorSize; + // Preconditions + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + CUTE_STATIC_ASSERT(is_static::value); + auto cluster_size = size<0>(cluster_shape); + auto mainloop_tiler = Shape<_128,_16,_128>{}; + auto epilogue_tiler = Shape<_128,_128,_128>{}; + + static int constexpr EpilogueUnrollFactor = size<2>(epilogue_tiler) / size<2>(cluster_tile); + + // Get the appropriate blocks for this Cluster + dim3 cluster_coord_in_grid = cluster_id_in_grid(); + + // Total number of k-tiles + int const K_TILE_MAX = ceil_div(min(N, K), size<2>(epilogue_tiler)); + + struct TileScheduler { + struct WorkTileInfo { + uint32_t m_idx = 0; + uint32_t n_idx = 0; + uint32_t l_idx = 0; + bool is_valid_tile = false; + }; + uint32_t tiles_in_m = 0; + uint32_t tiles_in_n = 0; + + int k_tile_max = 0; + + int wave_cnt = 0; + WorkTileInfo work_tile_info; + WorkTileInfo next_work_tile_info; + CLCResponse* clc_response_ptr_; + CUTLASS_DEVICE TileScheduler(uint32_t tiles_m, uint32_t tiles_n, int kmax, CLCResponse* clc_response_ptr) + : tiles_in_m(tiles_m), + tiles_in_n(tiles_n), + + k_tile_max(kmax), + work_tile_info({blockIdx.x, blockIdx.y, blockIdx.z, blockIdx.x( + &clc_response_ptr[state.index()])); + asm volatile( + "{\n\t" + "clusterlaunchcontrol.try_cancel.async.shared::cta.mbarrier::complete_tx::bytes.multicast::cluster::all.b128 [%0], [%1];\n\t" + "}\n" + : + : "r"(result_addr), "r"(mbarrier_addr)); + #else + CUTLASS_NOT_IMPLEMENTED(); + #endif + } + CUTLASS_DEVICE + static WorkTileInfo + work_tile_info_from_clc_response(uint32_t result_addr) { + WorkTileInfo work_tile_info; + uint32_t valid = 0; + #if defined(CUTLASS_ARCH_CLC_ENABLED) + asm volatile( + "{\n" + ".reg .pred p1;\n\t" + ".reg .b128 clc_result;\n\t" + "ld.shared.b128 clc_result, [%4];\n\t" + "clusterlaunchcontrol.query_cancel.is_canceled.pred.b128 p1, clc_result;\n\t" + "selp.u32 %3, 1, 0, p1;\n\t" + "@p1 clusterlaunchcontrol.query_cancel.get_first_ctaid.v4.b32.b128 {%0, %1, %2, _}, clc_result;\n\t" + "}\n" + : "=r"(work_tile_info.m_idx), "=r"(work_tile_info.n_idx), "=r"(work_tile_info.l_idx), "=r"(valid) + : "r"(result_addr) + : "memory" + ); + + cutlass::arch::fence_view_async_shared(); + #else + CUTLASS_NOT_IMPLEMENTED(); + #endif + work_tile_info.is_valid_tile = (valid == 1); + return work_tile_info; + } + }; + + + + // Allocate SMEM + extern __shared__ char shared_memory[]; + using SharedStorage = SharedStorage; + SharedStorage& shared_storage = *reinterpret_cast(shared_memory); + uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile)))); + uint32_t tiles_in_n = uint32_t(size(ceil_div(N, size<2>(epilogue_tiler)))); + TileScheduler scheduler(tiles_in_m, tiles_in_n, K_TILE_MAX, shared_storage.clc_response); + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + auto acc_shape_mma = make_shape(take<0,2>(mainloop_tiler), _1{}, _1{}); + auto acc_shape_epilogue = make_shape(take<0,2>(epilogue_tiler), _1{}, _1{}); + + auto acc_mainloop_pipelined_shape = append(acc_shape_mma, Int{}); + auto bulk_tmem_mma = TiledMMA::make_fragment_C(acc_mainloop_pipelined_shape); + + static int constexpr NumEpilogueColQuantThreadCount = kEnableRHTColQuant ? 128 : 0; + static int constexpr NumEpilogueRowQuantThreadCount = kEnableRowQuant ? 256 : 0; + static int constexpr NumMmaThreadCount = kEnableRHTColQuant? 32: 0; + static int constexpr NumMmaIssueThreadCount = kEnableRHTColQuant? 1: 0; + static int constexpr NumSchedThreads = 32; + static int constexpr NumMainloopLoadThreads = 32; + static int constexpr NumEpilogueThreads = NumEpilogueColQuantThreadCount + NumEpilogueRowQuantThreadCount; + + TmemAllocator tmem_allocator{}; + cutlass::arch::NamedBarrier tmem_allocation_result_barrier( + NumMmaThreadCount + NumEpilogueColQuantThreadCount, + cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier); + + int warp_idx = cutlass::canonical_warp_idx_sync(); + + // warp assignment + bool is_mma_warp = (warp_idx == 0); + bool is_dma_warp = (warp_idx == 1); + bool is_sched_warp = (warp_idx == 2); + bool is_epilogue_col_quant_warp = (warp_idx >= 4 && warp_idx <= 7); + bool is_epilogue_row_quant_warp = (warp_idx >= 8 && warp_idx <= 15); + + if (is_epilogue_col_quant_warp && elect_one_sync()) { + cute::prefetch(raw_pointer_cast(c_global_amax)); + } + if (is_epilogue_row_quant_warp && elect_one_sync()) { + cute::prefetch(raw_pointer_cast(a_global_amax)); + } + + typename MainloopPipeline::Params mainloop_pipeline_params; + if (is_dma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Producer; + } + if (is_mma_warp) { + mainloop_pipeline_params.role = MainloopPipeline::ThreadCategory::Consumer; + } + mainloop_pipeline_params.is_leader = cute::elect_one_sync() && is_dma_warp; + mainloop_pipeline_params.transaction_bytes = kTmaTransactionBytes; + mainloop_pipeline_params.initializing_warp = 0; + mainloop_pipeline_params.num_consumers = NumEpilogueRowQuantThreadCount + NumMmaIssueThreadCount; + MainloopPipeline mainloop_pipeline( + shared_storage.mainloop, + mainloop_pipeline_params, + cluster_shape, + cute::true_type{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + + MainloopPipelineState mainloop_pipe_consumer_state; + MainloopPipelineState mainloop_pipe_producer_state = cutlass::make_producer_start_state(); + + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + using AccumulatorPipelineState = typename AccumulatorPipeline::PipelineState; + + AccumulatorPipelineState accumulator_pipe_consumer_state; + AccumulatorPipelineState accumulator_pipe_producer_state = cutlass::make_producer_start_state(); + + typename AccumulatorPipeline::Params accumulator_pipeline_params; + if (is_mma_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Producer; + } + if (is_epilogue_col_quant_warp) { + accumulator_pipeline_params.role = AccumulatorPipeline::ThreadCategory::Consumer; + } + // Only one producer thread arrives on this barrier. + accumulator_pipeline_params.producer_arv_count = 1; + accumulator_pipeline_params.consumer_arv_count = size(AtomThrShapeMNK{}) * NumEpilogueColQuantThreadCount; + accumulator_pipeline_params.initializing_warp = 1; + using IsInitAccumulatorPipeline = cute::conditional_t; + AccumulatorPipeline accumulator_pipeline( + shared_storage.accumulator, + accumulator_pipeline_params, + cluster_shape, + IsInitAccumulatorPipeline{}, // Perform barrier init + cute::true_type{}); // Delay mask calculation + // CLC pipeline + typename CLCPipeline::Params clc_pipeline_params; + if (is_sched_warp) { + clc_pipeline_params.role = CLCPipeline::ThreadCategory::ProducerConsumer; + } else { + clc_pipeline_params.role = CLCPipeline::ThreadCategory::Consumer; + } + clc_pipeline_params.producer_blockid = 0; + clc_pipeline_params.producer_arv_count = 1; + clc_pipeline_params.consumer_arv_count = NumSchedThreads + cluster_size * + (NumMainloopLoadThreads + NumEpilogueThreads + NumMmaThreadCount); + clc_pipeline_params.transaction_bytes = sizeof(CLCResponse); + clc_pipeline_params.initializing_warp = 3; + CLCPipeline clc_pipeline(shared_storage.clc, clc_pipeline_params, cluster_shape); + CLCPipelineState clc_pipeline_consumer_state; + CLCPipelineState clc_pipeline_producer_state = cutlass::make_producer_start_state(); + + // CLC throttle pipeline + typename CLCThrottlePipeline::Params clc_throttle_pipeline_params; + if (is_dma_warp) { + clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Producer; + } + if (is_sched_warp) { + clc_throttle_pipeline_params.role = CLCThrottlePipeline::ThreadCategory::Consumer; + } + clc_throttle_pipeline_params.producer_arv_count = NumMainloopLoadThreads; + clc_throttle_pipeline_params.consumer_arv_count = NumSchedThreads; + clc_throttle_pipeline_params.dst_blockid = 0; + clc_throttle_pipeline_params.initializing_warp = 4; + + CLCThrottlePipeline clc_throttle_pipeline(shared_storage.clc_throttle, clc_throttle_pipeline_params); + CLCThrottlePipelineState clc_pipe_throttle_consumer_state; + CLCThrottlePipelineState clc_pipe_throttle_producer_state = cutlass::make_producer_start_state(); + + if (warp_idx == 2 && elect_one_sync()) { + cute::initialize_barrier(shared_storage.tma_barrier[0], /* num_threads */ 1); + } + __syncthreads(); + + if (is_dma_warp) { + cutlass::arch::warpgroup_reg_dealloc<32>(); + cute::Tensor mA = tma_load_a.get_tma_tensor(make_shape(M,N)); + cute::Tensor mB = tma_load_b.get_tma_tensor(make_shape(RhtTensorSize, RhtTensorSize)); + + cute::Tensor gA_mk = local_tile(mA, mainloop_tiler, make_coord(_,_, _), Step<_1, X,_1>{}); + cute::Tensor gB_nk = local_tile(mB, cluster_tile, make_coord(_,_, _), Step< X,_1,_1>{}); // (BLK_N,BLK_K,k) + + cute::Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + cute::Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + cute::Tensor tCgA = thr_mma.partition_A(gA_mk); // (MMA,MMA_M,MMA_K,k) + cute::Tensor tCgB = thr_mma.partition_B(gB_nk); // (MMA,MMA_N,MMA_K,k) + + Layout cta_layout_mnk = make_layout(cluster_shape); + Layout cta_layout_vmnk = tiled_divide(cta_layout_mnk, make_tile(typename TiledMMA::AtomThrID{})); + auto cta_coord_vmnk = cta_layout_vmnk.get_flat_coord(block_rank_in_cluster); + + auto [tAgA, tAsA] = tma_partition( + tma_load_a, + get<2>(cta_coord_vmnk), + make_layout(size<2>(cta_layout_vmnk)), + group_modes<0,3>(tCsA), + group_modes<0,3>(tCgA)); + + auto [tBgB, tBsB] = tma_partition( + tma_load_b, + get<1>(cta_coord_vmnk), + make_layout(size<1>(cta_layout_vmnk)), + group_modes<0,3>(tCsB), + group_modes<0,3>(tCgB)); + + uint16_t tma_mcast_mask_a = create_tma_multicast_mask<2>(cta_layout_vmnk, cta_coord_vmnk); + uint16_t tma_mcast_mask_b = create_tma_multicast_mask<1>(cta_layout_vmnk, cta_coord_vmnk); + if constexpr (kEnableRHTColQuant) { + if (elect_one_sync()) { + cute::set_barrier_transaction_bytes(shared_storage.tma_barrier[0], kTmaRhtTensorTransactionBytes); + copy(tma_load_b.with(shared_storage.tma_barrier[0], tma_mcast_mask_b), tBgB(_,0,0), tBsB(_,0)); + } + } + + do { + bool is_first_wave = scheduler.is_first_wave(); + uint32_t skip_wait = is_first_wave; + auto tAgA_mk = tAgA(_,scheduler.tile_m(),_); + int k_tile = 0; + // Throttle CLC producer + clc_throttle_pipeline.producer_acquire(clc_pipe_throttle_producer_state); + clc_throttle_pipeline.producer_commit(clc_pipe_throttle_producer_state); + ++clc_pipe_throttle_producer_state; + + CUTLASS_PRAGMA_NO_UNROLL + while (k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n()) { + + int k_tile_idx_n = scheduler.tile_n_base() + k_tile; + ++k_tile; + skip_wait = (is_first_wave && k_tile < MainloopPipelineStageCount); + mainloop_pipeline.producer_acquire(mainloop_pipe_producer_state); + using BarrierType = typename MainloopPipeline::ProducerBarrierType; + BarrierType* tma_barrier = mainloop_pipeline.producer_get_barrier( + mainloop_pipe_producer_state); + int write_stage = mainloop_pipe_producer_state.index(); + ++mainloop_pipe_producer_state; + if (cute::elect_one_sync()) { + copy( + tma_load_a.with(*tma_barrier, tma_mcast_mask_a), + tAgA_mk(_,k_tile_idx_n), + tAsA(_,write_stage)); + } + } + scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state); + ++clc_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + mainloop_pipeline.producer_tail(mainloop_pipe_producer_state); + } else if (is_mma_warp) { + cutlass::arch::warpgroup_reg_dealloc<32>(); + if constexpr (kEnableRHTColQuant) { + cute::Tensor tCsA = make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout); // (MMA,MMA_M,MMA_N,PIPE) + cute::Tensor tCsB = make_tensor(make_smem_ptr(shared_storage.tensors.smem_B.data()), sBlayout); // (MMA,MMA_N,MMA_K,PIPE) + + int block_rank_in_cluster = cute::block_rank_in_cluster(); + ThrMMA thr_mma = mma.get_slice(block_rank_in_cluster); // blk idx + // Allocate "fragments" -- these are actually umma smem descriptors + cute::Tensor tCrA = thr_mma.make_fragment_A(tCsA); // (MMA,MMA_M,MMA_K,PIPE) + cute::Tensor tCrB = thr_mma.make_fragment_B(tCsB); // (MMA,MMA_M,MMA_K,PIPE) + + mma.accumulate_ = UMMA::ScaleOut::Zero; + + tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &shared_storage.tmem_base_ptr); + __syncwarp(); + tmem_allocation_result_barrier.arrive(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_mma.data() = tmem_base_ptr; + cute::wait_barrier(shared_storage.tma_barrier[0], 0 /*tma_phase_bit*/); + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + + auto barrier_token = mainloop_pipeline.consumer_try_wait( + mainloop_pipe_consumer_state, + skip_wait); + scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state); + ++clc_pipeline_consumer_state; + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ) { + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + int read_stage = mainloop_pipe_consumer_state.index(); + auto tCrA_mk = tCrA(_,_,_,read_stage); + auto tCrB_nk = tCrB(_,_,0,0); + CUTLASS_PRAGMA_UNROLL + for (int k_block = 0; k_block < size<2>(tCrA) / EpilogueUnrollFactor; ++k_block) + { + int accumulator_k_block = accumulator_pipe_producer_state.index() * EpilogueUnrollFactor; + int tCrA_k_block = k_block * EpilogueUnrollFactor; + accumulator_pipeline.producer_acquire(accumulator_pipe_producer_state); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < EpilogueUnrollFactor; i++) { + auto accumulators = bulk_tmem_mma(_,_,_,accumulator_k_block + i); + gemm(mma, tCrA_mk(_,_,tCrA_k_block + i), tCrB_nk, accumulators); + } + + accumulator_pipeline.producer_commit(accumulator_pipe_producer_state); + ++accumulator_pipe_producer_state; + } + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + mainloop_pipeline.umma_consumer_release(curr_mainloop_pipe_consumer_state); + barrier_token = mainloop_pipeline.consumer_try_wait( + mainloop_pipe_consumer_state, + skip_wait); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + tmem_allocator.release_allocation_lock(); + accumulator_pipeline.producer_tail(accumulator_pipe_producer_state); + tmem_allocator.free(tmem_base_ptr, TmemAllocator::Sm100TmemCapacityColumns); + } + } else if(is_sched_warp) { + cutlass::arch::warpgroup_reg_dealloc<32>(); + do { + clc_throttle_pipeline.consumer_wait(clc_pipe_throttle_consumer_state); + clc_throttle_pipeline.consumer_release(clc_pipe_throttle_consumer_state); + ++clc_pipe_throttle_consumer_state; + clc_pipeline_producer_state = scheduler.advance_to_next_work(clc_pipeline, clc_pipeline_producer_state); + scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state); + ++clc_pipeline_consumer_state; + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } else if (is_epilogue_col_quant_warp) { + cutlass::arch::warpgroup_reg_alloc<192>(); + if constexpr (kEnableRHTColQuant) { + using TMEM_LOAD_NEW = cute::SM100::TMEM::LOAD::SM100_TMEM_LOAD_32dp32b64x; + + float const c_global_amax_val = *c_global_amax; + auto acc_epilogue_pipelined_shape = append(acc_shape_epilogue, Int{}); + auto bulk_tmem_epilogue_layout = make_layout( + acc_epilogue_pipelined_shape, + make_stride( + stride<0>(bulk_tmem_mma), + Int<0>{}, + Int<0>{}, + size<1>(epilogue_tiler))); + auto bulk_tmem_epilogue = make_tensor(make_tmem_ptr(), bulk_tmem_epilogue_layout); + + // leveraging 256-bit writes to global memory + static int constexpr FragmentSize = 256 / sizeof_bits_v; + + tmem_allocation_result_barrier.arrive_and_wait(); + uint32_t tmem_base_ptr = shared_storage.tmem_base_ptr; + bulk_tmem_epilogue.data() = tmem_base_ptr; + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % cutlass::NumThreadsPerWarpGroup; + + size_t rng_seed = 0; + size_t rng_offset = 0; + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + + cute::Tensor mD = make_tensor( + cute::subbyte_iterator(D), + make_shape(M,N), + dD); // (M,N) + cute::Tensor gD_mn = local_tile( + mD, + epilogue_tiler, + make_coord(_,_, _), + Step<_1,_1, X>{}); // (BLK_M,BLK_N) + cute::Tensor pD = make_identity_tensor(mD.shape()); + cute::Tensor pD_mn = local_tile( + pD, + epilogue_tiler, + make_coord(_,_, _), + Step<_1,_1, X>{}); // (BLK_M,BLK_N) + cute::Tensor mSFD = make_tensor(make_gmem_ptr(SFD), sfd_layout); + cute::Tensor gSFD_mn = local_tile(mSFD, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) + cute::Tensor pSFD = make_identity_tensor(mSFD.shape()); + cute::Tensor pSFD_mn = local_tile(pSFD, epilogue_tiler, make_coord(_,_, _), Step<_1,_1, X>{}); // (BLK_M,BLK_N) + + cute::Tensor gD_mn_view = tiled_divide(gD_mn, take<0,2>(epilogue_tiler)); + cute::Tensor pD_mn_view = tiled_divide(pD_mn, take<0,2>(epilogue_tiler)); + auto tiled_t2r = make_tmem_copy(TMEM_LOAD_NEW{}, bulk_tmem_epilogue(_,_,_,_0{})); + auto tiled_r2g = make_tiled_copy_D( + Copy_Atom{}, + tiled_t2r); + auto thr_t2r = tiled_t2r.get_slice(local_thread_idx); + auto thr_r2g = tiled_r2g.get_slice(local_thread_idx); + + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + float const fp4_max_inv = 1.0f / fp4_max; + float const global_encode_scale = c_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / c_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + + float const global_decode_scale = 1.0f / global_encode_scale; + // Scaling factor for fast math path + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + auto sfc_converter = cutlass::NumericConverter{}; + + do { + scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state); + ++clc_pipeline_consumer_state; + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ++k_tile) { + cute::Tensor tDgD_mn = gD_mn_view(_,_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + cute::Tensor tDgSFD_mn = gSFD_mn(_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + cute::Tensor tDpD_mn = pD_mn_view(_,_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + cute::Tensor tDpSFD_mn = pSFD_mn(_,_,scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + + accumulator_pipeline.consumer_wait(accumulator_pipe_consumer_state); + + auto Acc = bulk_tmem_epilogue(_,_,_,accumulator_pipe_consumer_state.index()); + cute::Tensor tDtAcc = thr_t2r.partition_S(Acc); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + cute::Tensor tDgD = thr_t2r.partition_D(tDgD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + cute::Tensor tDpD = thr_t2r.partition_D(tDpD_mn); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + cute::Tensor tTR_rAcc = make_tensor(shape(tDgD)); // ((TMEM_LOAD,#TMEM_LOAD),MMA_M,MMA_N) + cute::Tensor tDrD = make_tensor(shape(tDgD)); + cute::Tensor tTR_rAcc_frag = recast>(coalesce(tTR_rAcc)); + cute::Tensor tDrD_frag = recast>(coalesce(tDrD)); + + cute::Tensor src = thr_r2g.retile_S(tDrD); + cute::Tensor dst = thr_r2g.retile_D(tDgD); + cute::Tensor pSrc = thr_r2g.retile_D(tDpD); + + cute::Tensor tDgSFD_view = make_tensor( + tDgSFD_mn.data(), + make_layout( + make_shape(shape(tDgSFD_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tDgSFD_mn), Int<0>{}, Int<0>{}))); + cute::Tensor tDpSFD_view = make_tensor( + tDpSFD_mn.data(), + make_layout( + make_shape(shape(tDpSFD_mn), Int<1>{}, Int<1>{}), + make_stride(stride(tDpSFD_mn), Int<0>{}, Int<0>{}))); + cute::Tensor tDgSFD = filter(thr_t2r.partition_D(tDgSFD_view)); + cute::Tensor tDrSFD = make_tensor(shape(tDgSFD)); + cute::Tensor tDpSFD = filter(thr_t2r.partition_D(tDpSFD_view)); + static int constexpr NumVecs = size(tDgD) / VectorSize; + cute::Tensor tD_rRowSFD_frg = recast>(tDrSFD); + + cutlass::maximum_absolute_value_reduction, true> amax_reduction; + cutlass::Array vec_maxs; + cutlass::Array pvscales; + // TMEM_LOAD + copy(tiled_t2r, tDtAcc, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + accumulator_pipeline.consumer_release(accumulator_pipe_consumer_state); + ++accumulator_pipe_consumer_state; + + if constexpr (!kUseFastMath) { + // Downcast to BF16 for bit-wise compatibility with + // unfused kernels + auto convert_accum_to_bf16 = + cutlass::NumericArrayConverter{}; + auto convert_bf16_to_accum = + cutlass::NumericArrayConverter{}; + tTR_rAcc_frag(_0{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_0{}))); + tTR_rAcc_frag(_1{}) = convert_bf16_to_accum(convert_accum_to_bf16(tTR_rAcc_frag(_1{}))); + } + + auto compute_frgs = reinterpret_cast *>(tTR_rAcc_frag.data()); + auto output_frgs = reinterpret_cast *>(tDrD_frag.data()); + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + vec_maxs[v] = amax_reduction(ElementAccumulator(0), compute_frgs[v]); + } + + pvscales = cutlass::multiplies>{}( + vec_maxs, global_encode_scale_multiplier); + auto pvscales_cvted = cutlass::NumericArrayConverter{}(pvscales); + + tD_rRowSFD_frg(_0{}) = pvscales_cvted; + auto qpvscale_ups = cutlass::NumericArrayConverter{}(tD_rRowSFD_frg(_0{})); + auto qpvscale_scaled = cutlass::multiplies>{}( + qpvscale_ups, + global_decode_scale); + + cutlass::Array acc_scales; + if constexpr (kUseFastMath) { + // fast math: use reciprocal approximate to replace div + acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // regular path for slower math, use divide to replace div + acc_scales = cutlass::divides>{}(1.0, qpvscale_scaled); + } + + uint4 random_uint4 = uint4{0, 0, 0, 0}; + transformer_engine::curanddx::detail::philox4x32_native_state rng; + // "Prefetch" a stochastic rounding state for the first tile + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < NumVecs; v++) { + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales[v], + cutlass::platform::numeric_limits::max()); + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter(cutlass::multiplies>{}(compute_frgs[v], acc_scale), *reinterpret_cast*>(&random_uint4)); + } else { + output_frgs[v] = cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs[v], + acc_scale)); + } + + } + + cute::Tensor pred_pSrc = cute::lazy::transform(make_tensor(counting_iterator{}, replace<0>(shape(dst), _1{})), [&](auto coord){ + cute::Tensor pSrc_view = group_modes<1,rank(pSrc)>(pSrc); + return elem_less(pSrc_view(_0{},coord), shape(mD)); + }); + copy_if(tiled_r2g, pred_pSrc, src, dst); + // 32bit vectorization copy 4 e4m3 SFD for per 64 or(16,4):(0, 1) element + + constexpr int vec_len = 32 / sizeof_bits_v; + cute::Tensor tDrSFD_v = recast>(tDrSFD); + cute::Tensor tDgSFD_v = recast>(tDgSFD); + copy_if( + [&](auto coord){ + cute::Tensor tDpSFD_view = group_modes<1,rank(tDpSFD)>(tDpSFD); + return elem_less(tDpSFD_view(_0{}, coord * vec_len), shape(mSFD)); + }, + tDrSFD_v, tDgSFD_v); + } + scheduler.update_work_tile_info(); + } while (scheduler.is_valid()); + } + } else if (is_epilogue_row_quant_warp) { + cutlass::arch::warpgroup_reg_alloc<136>(); + if constexpr (kEnableRowQuant) { + using S2RVectorType = uint128_t; + float const a_global_amax_val = *a_global_amax; + int global_thread_idx = threadIdx.x; + int local_thread_idx = global_thread_idx % 256; + size_t rng_seed = 0; + size_t rng_offset = 0; + if constexpr (kEnableStochasticRounding) { + rng_seed = rng_state != nullptr ? __ldg(rng_state) : 0; + rng_offset = rng_state != nullptr ? __ldg(rng_state + 1) : 0; + } + cute::Tensor mQA = make_tensor(cute::subbyte_iterator(QA), make_layout(make_shape(M, N), dQA)); + cute::Tensor gQA_mn = local_tile(mQA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{}); + cute::Tensor pQA = make_identity_tensor(mQA.shape()); + cute::Tensor pQA_mn = local_tile(pQA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{}); + + cute::Tensor mSFA = make_tensor(make_gmem_ptr(SFA), sfa_layout); + cute::Tensor gSFA_mn = local_tile(mSFA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{}); // (BLK_M,BLK_N) + cute::Tensor pSFA = make_identity_tensor(mSFA.shape()); + cute::Tensor pSFA_mn = local_tile(pSFA, epilogue_tiler, make_coord(_,_, _), Step<_1,X,_1>{}); + cute::Tensor sA = as_position_independent_swizzle_tensor( + group_modes<0,2>(coalesce(make_tensor(make_smem_ptr(shared_storage.tensors.smem_A.data()), sAlayout)))); // (BLOCK_M, BLOCK_M,PIPE) + using S2RWarpLayout = Layout>; + using WarpGroupLayout = Layout>; + using S2RThreadLayout = decltype(blocked_product(S2RWarpLayout{}, WarpGroupLayout{})); + using S2RValLayout = Layout, _1>>; + using S2RAtomA = Copy_Atom; + using R2GAtomQA = Copy_Atom; + + auto tiled_s2r = make_tiled_copy(S2RAtomA{}, S2RThreadLayout{}, S2RValLayout{}); + auto tiled_r2g_QA = make_tiled_copy_D(R2GAtomQA{}, tiled_s2r); + + auto thr_s2r = tiled_s2r.get_slice(local_thread_idx); + auto thr_r2g_QA = tiled_r2g_QA.get_slice(local_thread_idx); + + cute::Tensor tQAsA = thr_s2r.partition_S(sA); // (Copy, Copy_M, Copy_N, PIPE) + + cute::Tensor tQArA = make_tensor_like(make_layout(tQAsA(_, _, _, _0{}).shape())); // (Copy, Copy_M, Copy_N) + // Tensor tQArA_PI = thr_s2r.partition_S(sA_PI); + cute::Tensor tQAgQA = thr_r2g_QA.partition_D(gQA_mn); + cute::Tensor tQArQA = make_tensor_like(tQAgQA(_, _, _, _0{}, _0{})); + cute::Tensor tQApQA = thr_r2g_QA.partition_D(pQA_mn); + + cute::Tensor tQAgSFA = thr_s2r.partition_D(gSFA_mn); + cute::Tensor tQArSFA = make_tensor_like(tQAgSFA(_, _, _, _0{}, _0{})); + cute::Tensor tQApSFA = thr_s2r.partition_D(pSFA_mn); + + // Aligning with TensorEngine's recipe to generate scale factors // {$nv-internal-release} + static constexpr float fp4_max = 6.0f; + static constexpr float fp8_max = 448.0f; + float const fp4_max_inv = 1.0f / fp4_max; + float const global_encode_scale = a_global_amax_val > 0.0f + ? cutlass::minimum_with_nan_propagation{}( + (fp8_max * fp4_max) / a_global_amax_val, + cutlass::platform::numeric_limits::max()) + : 1.0f; + + float const global_decode_scale = 1.0f / global_encode_scale; + // Scaling factor for fast math path + float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; + + auto sfa_converter = cutlass::NumericConverter{}; + do { + uint32_t skip_wait = K_TILE_MAX <= 0; + + CUTLASS_PRAGMA_NO_UNROLL + for (int k_tile = 0; k_tile < K_TILE_MAX && k_tile + scheduler.tile_n_base() < scheduler.tiles_n(); ) { + auto tQAgSFA_mn = tQAgSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQAgQA_mn = tQAgQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQApSFA_mn = tQApSFA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto tQApQA_mn = tQApQA(_, _, _, scheduler.tile_m(), scheduler.tile_n_base() + k_tile); + auto barrier_token = mainloop_pipeline.consumer_try_wait( + mainloop_pipe_consumer_state); + mainloop_pipeline.consumer_wait(mainloop_pipe_consumer_state, barrier_token); + copy(tiled_s2r, tQAsA(_, _, _, mainloop_pipe_consumer_state.index()), tQArA); + cutlass::arch::fence_view_async_shared(); + auto curr_mainloop_pipe_consumer_state = mainloop_pipe_consumer_state; + ++mainloop_pipe_consumer_state; + ++k_tile; + skip_wait = k_tile >= K_TILE_MAX; + mainloop_pipeline.consumer_release(curr_mainloop_pipe_consumer_state); + // static int constexpr NumVecs = size(tQArA) / VectorSize; + cutlass::maximum_absolute_value_reduction, true> amax_reduction; + auto compute_frgs = reinterpret_cast *>(tQArA.data()); + auto output_frgs = reinterpret_cast *>(raw_pointer_cast(tQArQA.data())); + transformer_engine::curanddx::detail::philox4x32_native_state rng; + if constexpr (kEnableStochasticRounding) { + const size_t rng_sequence = global_thread_idx + k_tile * 512 + scheduler.get_linear_tile_idx() * K_TILE_MAX * 512; + rng.init(rng_seed, rng_sequence, rng_offset); + } + CUTLASS_PRAGMA_UNROLL + for (int v = 0; v < size(tQArA)/VectorSize; v++) { + auto compute_frgs_up = cutlass::NumericArrayConverter{}(compute_frgs[v]); + auto amax = amax_reduction(ElementAccumulator(0), compute_frgs_up); + // declare pvscales + ElementAccumulator pvscales; + pvscales = cutlass::multiplies{}(amax, global_encode_scale_multiplier); + filter(tQArSFA)(v) = sfa_converter(pvscales); + auto qpvscale_ups = cutlass::NumericConverter{}(filter(tQArSFA)(v)); + auto qpvscale_scaled = cutlass::multiplies{}(qpvscale_ups, global_decode_scale); + ElementAccumulator acc_scales; + if constexpr (kUseFastMath) { + // fast math: use reciprocal approximate to replace div + acc_scales = cutlass::reciprocal_approximate_ftz{}(qpvscale_scaled); + } else { + // regular path for slower math, use divide to replace div + acc_scales = cutlass::divides{}(1.0, qpvscale_scaled); + } + auto acc_scale = cutlass::minimum_with_nan_propagation{}( + acc_scales, + cutlass::platform::numeric_limits::max()); + uint4 random_uint4 = uint4{0, 0, 0, 0}; + if constexpr (kEnableStochasticRounding) { + random_uint4 = rng.generate4(); + output_frgs[v] = StochasticNumericConverter(cutlass::multiplies>{}(compute_frgs_up, acc_scale), *reinterpret_cast*>(&random_uint4)); + } else { + output_frgs[v] = cutlass::NumericArrayConverter{}( + cutlass::multiplies>{}( + compute_frgs_up, + acc_scale)); + } + } + + cute::Tensor pred_tQApQA = cute::lazy::transform(make_tensor(counting_iterator{}, replace<0>(shape(tQAgQA_mn), _1{})), [&](auto coord){ + cute::Tensor tQApQA_view = group_modes<1,rank(tQApQA_mn)>(tQApQA_mn); + return elem_less(tQApQA_view(_0{}, coord), shape(mQA)); + }); + copy_if(tiled_r2g_QA, pred_tQApQA, tQArQA, tQAgQA_mn); + // 32bit vectorization copy 4 e4m3 SFA for per 64 or (16,4):(0, 1) element + constexpr int vec_len = 32 / sizeof_bits_v; + cute::Tensor tQArSFA_v = recast>(filter(tQArSFA)); + cute::Tensor tQAgSFA_v = recast>(filter(tQAgSFA_mn)); + copy_if( + [&](auto coord){ + cute::Tensor tQApSFA_view = filter(tQApSFA_mn); + return elem_less(tQApSFA_view(_0{}, coord * vec_len), shape(mSFA)); + }, + tQArSFA_v, tQAgSFA_v); + } + scheduler.fetch_next_work(clc_pipeline, clc_pipeline_consumer_state); + ++clc_pipeline_consumer_state; + scheduler.update_work_tile_info(); + }while (scheduler.is_valid()); + } + } else { + cutlass::arch::warpgroup_reg_dealloc<32>(); + } + } // sm100 compile guard end +} // NOLINT(readability/fn_size) + + +// this function computes RHT-GEMM for +// m = hidden_size, n = sequence_length +// A: m x n: col-major +// B: 16 x 16: row-major +// D: m x n: row-major +// SFD: m x (n/16): row-major +// QA: m x n: col-major +// SFA: m/16 x n: col-major +template +void row_col_rht_gemm_ntt_w_sfc( + int sequence_length, + int hidden_size, + TA const* A, + TB const* B, + TD* D, + TSFD* SFD, + TQA* QA, + TSFA* SFA, + float const* a_global_amax, + float const* d_global_amax, + const size_t* rng_state, + uint32_t sm_count, + cudaStream_t stream, + int k_tile_size = 1024) { + using namespace cute; + static int constexpr SFVecSize = 16; + static int constexpr RhtTensorSize = 16; + + static_assert(RhtTensorSize == 16, "RhtTensorSize must be 16"); + using LinearSFALayout = decltype(make_layout(make_shape(make_shape(Int{}, 0), 0), make_stride(make_stride(_0{}, _1{}), 0))); + using LinearSFCLayout = decltype(make_layout(make_shape(0, make_shape(Int{}, 0)), make_stride(0, make_stride(_0{}, _1{})))); + + using SwizzledSFALayoutAtom = cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFDLayoutAtom = cutlass::detail::Sm1xxBlockScaledOutputConfig::SfAtom; + using SwizzledSFALayout = decltype(tile_to_shape(SwizzledSFALayoutAtom{}, make_shape(hidden_size,sequence_length), Step<_1,_2>{})); + using SwizzledSFDLayout = decltype(tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(hidden_size,sequence_length), Step<_2,_1>{})); + + using SFALayout = cute::conditional_t; + using SFCLayout = cute::conditional_t; + SFALayout sfa_layout; + SFCLayout sfd_layout; + + if constexpr (kEnableSwizzleSFOutput) { + sfa_layout = tile_to_shape(SwizzledSFALayoutAtom{}, make_shape(hidden_size, sequence_length), Step<_1,_2>{}); + sfd_layout = tile_to_shape(SwizzledSFDLayoutAtom{}, make_shape(hidden_size, sequence_length), Step<_2,_1>{}); + } else { + sfa_layout = make_layout(make_shape(make_shape(Int{}, hidden_size/SFVecSize), sequence_length), make_stride(make_stride(_0{}, _1{}), hidden_size/SFVecSize)); + sfd_layout = make_layout(make_shape(hidden_size, make_shape(Int{}, sequence_length/SFVecSize)), make_stride(sequence_length/SFVecSize, make_stride(_0{}, _1{}))); + } + // Define shapes (dynamic) + auto M = hidden_size; + auto N = sequence_length; + cute::Tensor tensorA = make_tensor(A, make_shape(hidden_size, sequence_length), LayoutLeft{}); + cute::Tensor tensorB = make_tensor(B, make_shape(RhtTensorSize, RhtTensorSize), LayoutLeft{}); + cute::Tensor tensorD = make_tensor(D, make_shape(hidden_size, sequence_length), LayoutRight{}); + cute::Tensor tensorQA = make_tensor(QA, make_shape(hidden_size, sequence_length), LayoutLeft{}); + cute::Tensor tensorSFD = make_tensor(SFD, sfd_layout); + cute::Tensor tensorSFA = make_tensor(SFA, sfa_layout); + // Define strides (from tensors) + auto dA = stride(tensorA); // (dM,dK) + auto dB = stride(tensorB); // (dN,dK) + auto dD = stride(tensorD); // (dM,dN) + auto dQA = stride(tensorQA); // (dM,dK) + using ClusterShape = Shape< _1, _1, _1>; + auto cluster_shape = ClusterShape{}; + auto cluster_tile_shape = Shape<_128,Int,Int>{}; + auto cluster_tile_mainloop = Shape<_128,Int,_128>{}; + + // Each mainloop / epilogue loads 128 x 64 tiles while each MMA proceeds with 128 x 16 tiles + static int constexpr EpilogueUnrollFactor = + size<2>(cluster_tile_mainloop) / size<2>(cluster_tile_shape); + // Construct the MMA + auto mma = make_tiled_mma(SM100_MMA_F16BF16_SS(cluster_tile_shape), size<1>(cluster_tile_shape), + UMMA::Major::MN, UMMA::Major::MN>{}, + Layout>{}); + + // Assert that the TiledMMA uses all CTAs in the CGA. + CUTE_STATIC_ASSERT_V(size(cluster_shape) == size(mma)); + CUTE_STATIC_ASSERT_V(evenly_divides(cluster_tile_shape, tile_shape(mma))); + + // Determine the A and B shapes + auto mma_shape_B = partition_shape_B(mma, make_shape(size<1>(cluster_tile_shape), size<2>(cluster_tile_shape))); + + using TiledMma = decltype(mma); + using AtomThrID = typename TiledMma::AtomThrID; + + using SmemShape_M = decltype(shape_div(shape<0>(cluster_tile_shape), shape_div(shape<0>(cluster_tile_shape), size<0>(cluster_tile_shape) / size(AtomThrID{})))); + using SmemShape_N = decltype(shape_div(shape<1>(cluster_tile_shape), shape_div(shape<1>(cluster_tile_shape), size<1>(cluster_tile_shape) / size(AtomThrID{})))); + using SmemShape_K = decltype(cute::get<2>(cluster_tile_shape)); + + using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, TB, SmemShape_N, SmemShape_K>()); + + auto mma_shape_A = partition_shape_A(mma, make_shape(size<0>(cluster_tile_mainloop), size<2>(cluster_tile_mainloop))); + using SmemShape_M_A = decltype(shape_div(shape<0>(cluster_tile_mainloop), shape_div(shape<0>(cluster_tile_mainloop), size<0>(cluster_tile_mainloop) / size(AtomThrID{})))); + using SmemShape_K_A = decltype(cute::get<2>(cluster_tile_mainloop)); + using SmemLayoutAtomA = decltype(cutlass::gemm::collective::detail::sm100_smem_selector< + cute::UMMA::Major::MN, TA, SmemShape_M_A, SmemShape_K_A>()); + + static uint32_t constexpr TotalTmemRows = 128; + static uint32_t constexpr Sm100TmemCapacityColumns = 512; + static uint32_t constexpr TotalTmem = TotalTmemRows * Sm100TmemCapacityColumns; + static uint32_t constexpr AccumulatorPipelineStageCount = + TotalTmem / + (cute::size<0>(cluster_tile_shape) * cute::size<1>(cluster_tile_shape)); + + // Define the smem layouts (static) + // Calculate max pipeline stages based on Blackwell SM100's 232KB shared memory + constexpr int SchedulerPipelineStageCount = 6; + static int constexpr MainloopPipelineBytes = sizeof(typename cutlass::detail::CustomizedPipelineTmaUmmaAsync< + 1, + Shape<_1,_1,_1>, + Shape<_1, _1, _1>>::SharedStorage); + + static int constexpr ClcResponseBytes = sizeof(CLCResponse) * SchedulerPipelineStageCount; + static int constexpr CLCThrottlePipelineBytes = sizeof(typename cutlass::PipelineAsync::SharedStorage); + static int constexpr CLCPipelineBytes = sizeof(typename cutlass::PipelineCLCFetchAsync::SharedStorage); + static int constexpr TmemDeallocBytes = sizeof(cutlass::arch::ClusterBarrier); + static int constexpr BTensorBytes = cute::size(mma_shape_B) * sizeof(TB); + static int constexpr AccPipelineBytes = sizeof(typename cutlass::PipelineUmmaAsync>::SharedStorage); + static int constexpr TmemBasePtrsBytes = sizeof(uint32_t); + static int constexpr kBlackwellSmemSize = 232448; // 232KB in bytes + static int constexpr kBytesPerStage = + cute::size(mma_shape_A) * sizeof(TA) + MainloopPipelineBytes; + static int constexpr kReservedBytes = ClcResponseBytes + CLCThrottlePipelineBytes + TmemBasePtrsBytes + + CLCPipelineBytes + TmemDeallocBytes+BTensorBytes + AccPipelineBytes; // Reserve for barriers and other uses + static int constexpr kMaxStages = (kBlackwellSmemSize - kReservedBytes) / kBytesPerStage; + auto sP = Int{}; // SMEM pipelines + auto sA = UMMA::tile_to_mma_shape( + SmemLayoutAtomA{}, + append(mma_shape_A, sP), Step<_2,_1,_3>{}); // (MMA,MMA_M,MMA_K,PIPE) + auto sB = UMMA::tile_to_mma_shape( + SmemLayoutAtomB{}, + append(mma_shape_B, _1{})); // (MMA,MMA_N,MMA_K, _1) + auto sD = Layout<_1>{}; // XXX Dummy + + auto tma_load_a = make_tma_copy_A_sm100( + SM90_TMA_LOAD{}, + tensorA, + sA(_,_,_,0), + cluster_tile_mainloop, + mma); + auto tma_load_b = make_tma_copy_B_sm100( + SM90_TMA_LOAD{}, + tensorB, + sB(_,_,_,0), + cluster_tile_shape, + mma); + + // Assert checks problem size should be multiple of 64 + NVTE_CHECK(M % 64 == 0, "M must be a multiple of 64, but got ", M); + NVTE_CHECK(N % 64 == 0, "N must be a multiple of 64, but got ", N); + + uint32_t tiles_in_m = uint32_t(size(ceil_div(M, size<0>(cluster_tile_shape)))); + uint32_t tiles_in_n = uint32_t(size(ceil_div(N, k_tile_size))); + uint32_t tiles = tiles_in_m * tiles_in_n; + + dim3 dimBlock(512); + dim3 dimCluster(size<0>(cluster_shape), size<1>(cluster_shape), size<2>(cluster_shape)); + dim3 dimGrid(tiles_in_m, tiles_in_n, 1); + + int smem_size = sizeof( + SharedStorage< + TA, + TB, + decltype(sA), + decltype(sB), + ClusterShape, + AccumulatorPipelineStageCount, + EpilogueUnrollFactor, + SchedulerPipelineStageCount>); + + auto* kernel_ptr = &row_col_rht_gemm_device< + decltype(M), decltype(N), decltype(k_tile_size), + decltype(cluster_shape), decltype(cluster_tile_shape), + TA, decltype(dA), decltype(sA), decltype(tma_load_a), + TB, decltype(dB), decltype(sB), decltype(tma_load_b), + TD, decltype(dD), decltype(sD), + TSFD, decltype(sfd_layout), + TQA, decltype(dQA), + TSFA, decltype(sfa_layout), + decltype(mma), + AccumulatorPipelineStageCount, + SchedulerPipelineStageCount, + kEnableStochasticRounding, + kEnableRHTColQuant, + kEnableRowQuant, + kUseFastMath>; + + NVTE_CHECK_CUDA(cudaFuncSetAttribute(*kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + + cutlass::ClusterLaunchParams params = {dimGrid, dimBlock, dimCluster, smem_size, stream}; + cutlass::Status status = cutlass::launch_kernel_on_cluster( + params, (void const *)kernel_ptr, M, N, k_tile_size, cluster_shape, cluster_tile_shape, + tensorA.data(), dA, sA, tma_load_a, + tensorB.data(), dB, sB, tma_load_b, + tensorD.data(), dD, sD, + tensorSFD.data(), sfd_layout, + tensorQA.data(), dQA, + tensorSFA.data(), sfa_layout, + mma, a_global_amax, d_global_amax, rng_state); + + NVTE_CHECK_CUDA(cudaGetLastError()); + NVTE_CHECK(status == cutlass::Status::kSuccess, "Kernel launch failed."); + +} + +} // namespace +} // namespace detail + +// clang-format on + +void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, + const Tensor &hadamard_matrix_, QuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(hadamard_transform_cast_fusion); + + // Check input and output tensors + NVTE_CHECK(input_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING, + "Input tensor must be BF16 tensor, but scaling mode is ", + to_string(input_.scaling_mode), "."); + NVTE_CHECK(input_.dtype() == transformer_engine::DType::kBFloat16, + "Input tensor must be BF16 tensor, but dtype is ", to_string(input_.dtype()), "."); + NVTE_CHECK(input_.dim() >= 2, "Input must be a 2D tensor."); + const SimpleTensor &input = input_.data; + + // rowwise cast and columnwise cast has different output data pointers + bool has_rowwise_quant = false; + bool has_columnwise_quant = false; + void *rowwise_data_ptr = nullptr; + void *rowwise_scale_inv_ptr = nullptr; + void *rowwise_amax_ptr = nullptr; + void *columnwise_data_ptr = nullptr; + void *columnwise_scale_inv_ptr = nullptr; + void *columnwise_amax_ptr = nullptr; + + // examine the output tensor (single tensor for dense) + if (output_.data.dptr != nullptr) { + has_rowwise_quant = true; + rowwise_data_ptr = output_.data.dptr; + rowwise_scale_inv_ptr = output_.scale_inv.dptr; + rowwise_amax_ptr = output_.amax.dptr; + } + + if (output_.columnwise_data.dptr != nullptr) { + has_columnwise_quant = true; + columnwise_data_ptr = output_.columnwise_data.dptr; + columnwise_scale_inv_ptr = output_.columnwise_scale_inv.dptr; + columnwise_amax_ptr = output_.columnwise_amax.dptr; + } + + NVTE_CHECK(has_rowwise_quant || has_columnwise_quant, + "Output tensor must have rowwise or columnwise quant."); + + // Stochastic rounding config + const bool use_stochastic_rounding = quant_config.stochastic_rounding; + const size_t *rng_state = nullptr; + if (quant_config.rng_state != nullptr) { + 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}, + "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); + rng_state = reinterpret_cast(rng_state_tensor.data.dptr); + } + + // Template arguments + using TA = cute::bfloat16_t; + using TB = cute::bfloat16_t; + using TD = cutlass::float_e2m1_t; + using TSFD = cutlass::float_ue4m3_t; + using TQA = TD; + using TSFA = TSFD; + + checkCuDriverContext(stream); + + // Check Hadamard matrix + constexpr int kHadamardDimension = 16; + NVTE_CHECK(hadamard_matrix_.scaling_mode == NVTE_DELAYED_TENSOR_SCALING, + "Hadamard matrix must be BF16 tensor, but scaling mode is ", + to_string(hadamard_matrix_.scaling_mode), "."); + NVTE_CHECK(hadamard_matrix_.dtype() == transformer_engine::DType::kBFloat16, + "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(), "."); + const size_t hadamard_dimension = hadamard_matrix.shape[0]; + + const size_t ndim = input.shape.size(); + const size_t n = input.shape[ndim - 1]; + size_t m = 1; + for (size_t i = 0; i < ndim - 1; ++i) { + m *= input.shape[i]; + } + + auto sm_count = transformer_engine::cuda::sm_count(); + + NVTE_CHECK(n % hadamard_dimension == 0, "row_length must be divisible by hadamard_dimension."); + + NVTE_CHECK(m % hadamard_dimension == 0, "num_rows must be divisible by hadamard_dimension"); + + int k_tile_size = 1024; + + // TODO: add support for swizzle sf output + const bool use_swizzle_sf_output = false; + + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_stochastic_rounding, kEnableStochasticRounding, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + has_columnwise_quant, kEnableRhtColQuant, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + has_rowwise_quant, kEnableRowQuant, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_swizzle_sf_output, kEnableSwizzleSFOutput, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config.use_fast_math, kUseFastMath, + + if constexpr (kEnableRhtColQuant || kEnableRowQuant) { + detail::row_col_rht_gemm_ntt_w_sfc< + kEnableStochasticRounding, kEnableRhtColQuant, kEnableRowQuant, + kEnableSwizzleSFOutput, TA, TB, TD, TSFD, TQA, TSFA, kUseFastMath>( + /*sequence_length=*/m, /*hidden_size=*/n, + /*A=*/reinterpret_cast(input.dptr), + /*B=*/reinterpret_cast(hadamard_matrix.dptr), + /*D=*/reinterpret_cast(columnwise_data_ptr), + /*SFD=*/reinterpret_cast(columnwise_scale_inv_ptr), + /*QA=*/reinterpret_cast(rowwise_data_ptr), + /*SFA=*/reinterpret_cast(rowwise_scale_inv_ptr), + /*a_global_amax=*/reinterpret_cast(rowwise_amax_ptr), + /*d_global_amax=*/reinterpret_cast(columnwise_amax_ptr), + /*rng_state=*/rng_state, /*sm_count=*/sm_count, + /*stream=*/stream, /*k_tile_size=*/k_tile_size); + } else { + NVTE_ERROR("Invalid kernel configuration (kEnableRHTColQuant=", + kEnableRhtColQuant, ", kEnableRowQuant=", kEnableRowQuant, ")."); + } + + ););););); +} + +} // namespace transformer_engine + +void nvte_quantize_with_hadamard_transform(const NVTETensor input, NVTETensor output, + const NVTETensor hadamard_matrix, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_with_hadamard_transform); + using namespace transformer_engine; + QuantizationConfig quant_config_cpp; + if (quant_config != nullptr) { + quant_config_cpp = *reinterpret_cast(quant_config); + } + hadamard_transform_cast_fusion(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), + *convertNVTETensorCheck(hadamard_matrix), quant_config_cpp, + stream); +} diff --git a/transformer_engine/common/include/transformer_engine/hadamard_transform.h b/transformer_engine/common/include/transformer_engine/hadamard_transform.h index bee939f0cd..8f1a213cec 100644 --- a/transformer_engine/common/include/transformer_engine/hadamard_transform.h +++ b/transformer_engine/common/include/transformer_engine/hadamard_transform.h @@ -48,7 +48,7 @@ void nvte_hadamard_transform_amax(const NVTETensor input, NVTETensor output, int /*! \brief Perform the columnwise hadamard transform cast fusion. * - * This function is experimental and the API is not stable. + * \deprecated This function has been deprecated in favor of nvte_quantize_with_hadamard_transform. * * \param[in] input Input tensor to apply Hadamard transform. * \param[in,out] output Output tensor. @@ -61,6 +61,21 @@ void nvte_hadamard_transform_cast_fusion_columnwise(const NVTETensor input, NVTE const NVTEQuantizationConfig quant_config, cudaStream_t stream); +/*! \brief Perform the regular rowwise cast and columnwise hadamard transform cast fusion. + * + * This function is experimental and the API is not stable. + * + * \param[in] input Input tensor to apply Hadamard transform. + * \param[in,out] output Output tensor. + * \param[in] hadamard_matrix Hadamard matrix. + * \param[in] quant_config Quantization configuration. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_quantize_with_hadamard_transform(const NVTETensor input, NVTETensor output, + const NVTETensor hadamard_matrix, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream); + /*! \brief Split a tensor along dimension 0 and compute RHT amaxes for each split. * * This function is experimental and the API is not stable. diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 36ce60eaa5..4d028de01c 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -17,6 +17,7 @@ namespace transformer_engine { namespace nvfp4_recipe { +#if FP4_TYPE_SUPPORTED /* * --------------------------------------------------------------------------- * NVFP4 2D PARTIAL-SHARD KERNEL DESIGN @@ -616,7 +617,7 @@ void nvfp4_expand_scale_to_fp8(const Tensor input, Tensor output, size_t tile_ro * * Computes per-block decode scale from block amax and global amax: * global_scale = (fp8_max * fp4_max) / global_amax = 2688 / global_amax - * per_block_decode_scale = block_amax / fp4_max * global_scale + * per_block_decode_scale = block_amax * (global_scale * (1 / fp4_max)) * = block_amax * 448 / global_amax * * This matches the CUDA device function compute_decoding_scaling_factor() in core_nvfp4.cuh @@ -648,9 +649,11 @@ __global__ void nvfp4_compute_per_block_scale_kernel( float global_scale = (global_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; - // Compute per-block decode scale: S_dec_b = block_amax / fp4_max * S_enc + // Compute per-block decode scale: S_dec_b = block_amax * (S_enc * (1 / fp4_max)) float amax_val = block_amax[idx]; - float result = fminf((amax_val / fp4_max) * global_scale, flt_max); + constexpr float fp4_max_inv = 1.0f / fp4_max; + const float global_scale_multiplier = global_scale * fp4_max_inv; + float result = fminf(amax_val * global_scale_multiplier, flt_max); scale[idx] = result; } @@ -764,10 +767,12 @@ __global__ void nvfp4_fused_scale_kernel( float safe_global_amax = fmaxf(g_amax, tiny); float global_scale = (g_amax > 0.0f) ? fminf((fp8_max * fp4_max) / safe_global_amax, flt_max) : 1.0f; + constexpr float fp4_max_inv = 1.0f / fp4_max; + const float global_scale_multiplier = global_scale * fp4_max_inv; // Read block amax and compute per-block decode scale float amax_val = block_amax[tile_row * tile_cols + out_col]; - scale_val = fminf((amax_val / fp4_max) * global_scale, flt_max); + scale_val = fminf(amax_val * global_scale_multiplier, flt_max); // Write per-block scale (only once per tile, when out_row % block_len == 0) if (out_row % block_len == 0) { @@ -806,78 +811,109 @@ void nvfp4_fused_scale(const Tensor block_amax, const Tensor global_amax, Tensor block_len); NVTE_CHECK_CUDA(cudaGetLastError()); } + +#endif // FP4_TYPE_SUPPORTED } // namespace nvfp4_recipe } // namespace transformer_engine void nvte_nvfp4_expand_scale_to_fp8(const NVTETensor input, NVTETensor output, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_expand_scale_to_fp8); using namespace transformer_engine; nvfp4_recipe::nvfp4_expand_scale_to_fp8(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), tile_rows, tile_cols, rows_padded, block_len, stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_compute_per_block_scale(const NVTETensor block_amax, NVTETensor scale, const NVTETensor global_amax, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_per_block_scale); using namespace transformer_engine; nvfp4_recipe::nvfp4_compute_per_block_scale(*convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(scale), *convertNVTETensorCheck(global_amax), stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_compute_global_scale(const NVTETensor global_amax, NVTETensor global_scale, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_global_scale); using namespace transformer_engine; nvfp4_recipe::nvfp4_compute_global_scale(*convertNVTETensorCheck(global_amax), *convertNVTETensorCheck(global_scale), stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_scale_transpose(const NVTETensor input, NVTETensor output, size_t M_tiles, size_t K_tiles, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_scale_transpose); using namespace transformer_engine; nvfp4_recipe::nvfp4_scale_transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), M_tiles, K_tiles, stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_data_transpose(const NVTETensor input, NVTETensor output, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_data_transpose); using namespace transformer_engine; nvfp4_recipe::nvfp4_transpose(*convertNVTETensorCheck(input), *convertNVTETensorCheck(output), stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_2d_compute_partial_amax(const NVTETensor inp, NVTETensor amax, size_t h, size_t w, size_t amax_stride_h, size_t amax_stride_w, size_t start_offset, size_t block_len, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_2d_compute_partial_amax); using namespace transformer_engine; nvfp4_recipe::nvfp4_2d_compute_partial_amax(*convertNVTETensorCheck(inp), *convertNVTETensorCheck(amax), h, w, amax_stride_h, amax_stride_w, start_offset, block_len, stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_2d_partial_cast(const NVTETensor inp, NVTETensor out, const NVTETensor scale, const NVTETensor global_scale, size_t h, size_t w, size_t scale_stride_h, size_t scale_stride_w, size_t start_offset, size_t block_len, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_2d_partial_cast); using namespace transformer_engine; nvfp4_recipe::nvfp4_2d_partial_cast(*convertNVTETensorCheck(inp), *convertNVTETensorCheck(out), *convertNVTETensorCheck(scale), *convertNVTETensorCheck(global_scale), h, w, scale_stride_h, scale_stride_w, start_offset, block_len, stream); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_rowwise_amax_A, const NVTETensor inpB, const bool use_rowwise_amax_B, float alpha_in, NVTETensor alpha_out, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_compute_per_tensor_scale); using namespace transformer_engine; @@ -898,16 +934,23 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r alpha_in, reinterpret_cast(amax_A_ptr), reinterpret_cast(amax_B_ptr), 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); +#endif // FP4_TYPE_SUPPORTED } void nvte_nvfp4_fused_scale(const NVTETensor block_amax, const NVTETensor global_amax, NVTETensor per_block_scale, NVTETensor target_scale, NVTETensor target_amax, size_t tile_rows, size_t tile_cols, size_t rows_padded, size_t block_len, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED NVTE_API_CALL(nvte_nvfp4_fused_scale); using namespace transformer_engine; nvfp4_recipe::nvfp4_fused_scale( *convertNVTETensorCheck(block_amax), *convertNVTETensorCheck(global_amax), *convertNVTETensorCheck(per_block_scale), *convertNVTETensorCheck(target_scale), *convertNVTETensorCheck(target_amax), tile_rows, tile_cols, rows_padded, block_len, stream); +#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/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index e25cc607e5..d3d3dceca9 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -168,10 +168,9 @@ __device__ __forceinline__ float groupMax(float val, unsigned int groupMask) { } template -__device__ __forceinline__ ScaleType ComputeDecodeScaleFP4(const float amax, - const float global_encode_scale) { - float decode_scale = amax / TypeExtrema::max; - decode_scale = decode_scale * global_encode_scale; +__device__ __forceinline__ ScaleType +ComputeDecodeScaleFP4(const float amax, const float global_encode_scale_multiplier) { + float decode_scale = amax * global_encode_scale_multiplier; decode_scale = fminf(decode_scale, TypeExtrema::max); return static_cast(decode_scale); } @@ -420,6 +419,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo const int kNumThreadsReduce = kScaleBlockDim / kNVecOut; const float global_encode_scale = kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]); + constexpr float fp4_max_inv = 1.0f / TypeExtrema::max; + const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv; const float global_decode_scale = 1.0 / global_encode_scale; // Step 2: Cast and store to output_c @@ -508,7 +509,7 @@ __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); + ScaleType scale_inv = ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier); float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale); // Step 2.5: Write scale_inv bool write_scale_inv = is_src_lane; @@ -631,7 +632,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo amax = __shfl_sync(mask, amax, src_lane); } // Step 3.4: Compute scale - ScaleType scale_inv = ComputeDecodeScaleFP4(amax, global_encode_scale); + ScaleType scale_inv = + ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier); float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale); // Step 3.5: Write scale_inv_t bool write_scale_inv = is_src_lane; diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index 5367d7e781..f7611e60c5 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -14,9 +14,11 @@ #include #include -#if CUDA_VERSION >= 12080 +#include "common/common.h" + +#if FP4_TYPE_SUPPORTED #include -#endif // CUDA_VERSION >= 12080 +#endif // FP4_TYPE_SUPPORTED #include "common/utils.cuh" diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 6aab9938b3..63a2e86e67 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -370,6 +370,11 @@ class NVFP4Quantizer : public Quantizer { private: void quantize_impl(const TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag, bool compute_amax); + void quantize_with_rht_unfused_helper(const TensorWrapper& input, TensorWrapper& out, + TensorWrapper& rht_output_t_cpp, + QuantizationConfigWrapper& quant_config, + QuantizationConfigWrapper& quant_config_columnwise, + cudaStream_t stream); }; std::unique_ptr convert_quantizer(py::handle quantizer); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 89cd90f347..cb3434ec52 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -998,6 +998,10 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, // Enable NVFP4 kernels to use math operations that sacrifice // accuracy for performance. These optimizations are experimental // and inconsistently implemented. + // What math is accelerated? Only the high precision math, so numerical impact is minimal + // 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 const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); if (use_fast_math) { for (auto &config : quant_config_list) { diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 8c5504e44b..b59f3fa3c5 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/system.h" #include "pybind.h" #include "torch/torch.h" @@ -2134,6 +2135,82 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( return {std::move(out_cpp), std::move(tensor)}; } +void NVFP4Quantizer::quantize_with_rht_unfused_helper( + const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, + QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, + cudaStream_t stream) { + // only triggered for irregular shapes where RHT cast fusion kernel is not eligible + if (rowwise_usage) { + // For rowwise usage, we need to quantize the input directly, but we need to avoid quantizing columnwise + TensorWrapper out_identity(out.scaling_mode()); + auto out_identity_data = out.get_rowwise_data(); + auto out_identity_scale_inv = out.get_rowwise_scale_inv(); + auto out_identity_amax = out.get_amax(); + out_identity.set_rowwise_data(out_identity_data.data_ptr, + static_cast(out_identity_data.dtype), + out_identity_data.shape); + out_identity.set_rowwise_scale_inv(out_identity_scale_inv.data_ptr, + static_cast(out_identity_scale_inv.dtype), + out_identity_scale_inv.shape); + out_identity.set_amax(out_identity_amax.data_ptr, static_cast(out_identity_amax.dtype), + out_identity_amax.shape); + + NVTE_SCOPED_GIL_RELEASE( + { nvte_quantize_v2(input.data(), out_identity.data(), quant_config, stream); }); + } + + if (columnwise_usage) { + // Get the output columnwise data, scale_inv, and amax + auto out_columnwise_data = out.get_columnwise_data(); + auto out_columnwise_scale_inv = out.get_columnwise_scale_inv(); + // NOTE: should already be populated. + auto out_columnwise_amax = out.get_columnwise_amax(); + + // 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); + out_transpose.set_rowwise_scale_inv(out_columnwise_scale_inv.data_ptr, + static_cast(out_columnwise_scale_inv.dtype), + out_columnwise_scale_inv.shape); + out_transpose.set_amax(out_columnwise_amax.data_ptr, + static_cast(out_columnwise_amax.dtype), + 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, + this->rht_matrix_random_sign_mask_t, stream); + }); + + // Quantize kernel will treat everything as rowwise input/output, which is + // intended. + NVTE_SCOPED_GIL_RELEASE({ + nvte_quantize_v2(rht_output_t_cpp.data(), out_transpose.data(), quant_config_columnwise, + stream); + }); + } +} + void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag, bool compute_amax) { @@ -2145,8 +2222,10 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou auto stream = at::cuda::getCurrentCUDAStream(); QuantizationConfigWrapper quant_config; + QuantizationConfigWrapper quant_config_columnwise; if (noop_flag) { quant_config.set_noop_tensor(noop_flag->data()); + quant_config_columnwise.set_noop_tensor(noop_flag->data()); } quant_config.set_nvfp4_2d_quantization(this->with_2d_quantization); quant_config.set_stochastic_rounding(this->stochastic_rounding); @@ -2159,14 +2238,25 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou } size_t cols = input.size(input.ndim() - 1); + // 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; + // Stochastic rounding // When both rowwise and columnwise quantization are used with RHT, // we need separate RNG states for each to ensure they use different random numbers. TensorWrapper te_rng_state; TensorWrapper te_rng_state_columnwise; - QuantizationConfigWrapper quant_config_columnwise; - const bool need_separate_columnwise_rng = - this->stochastic_rounding && this->with_rht && this->columnwise_usage; + + // Only need a separate rng state when: + // 1. Stochastic rounding is enabled + // 2. RHT is enabled + // 3. Columnwise usage is enabled + // 4. Rowwise and columnwise quantization are not fused, + // because within a single kernel we can generate two different random numbers for rowwise and columnwise + const bool need_separate_columnwise_rng = this->stochastic_rounding && this->with_rht && + this->columnwise_usage && + (!eligible_for_rht_cast_fusion); if (this->stochastic_rounding) { const size_t rng_elts_per_thread = 1024; // Wild guess, probably can be tightened @@ -2189,13 +2279,10 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou te_rng_state_columnwise = makeTransformerEngineTensor(rng_state_columnwise); quant_config_columnwise.set_stochastic_rounding(true); quant_config_columnwise.set_rng_state(te_rng_state_columnwise.data()); + quant_config_columnwise.set_nvfp4_2d_quantization(this->with_2d_quantization); } } - // 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; - // Compute amax. if (this->with_rht) { if (input.dtype() != DType::kBFloat16) { @@ -2264,103 +2351,48 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou { this->amax_reduction_group->allreduce_coalesced(amax_tensors, opts)->wait(); }); } - if (this->with_rht) { - if (rowwise_usage) { - // For rowwise usage, we need to quantize the input directly, but we need to avoid quantizing columnwise - TensorWrapper out_identity(out.scaling_mode()); - auto out_identity_data = out.get_rowwise_data(); - auto out_identity_scale_inv = out.get_rowwise_scale_inv(); - auto out_identity_amax = out.get_amax(); - out_identity.set_rowwise_data(out_identity_data.data_ptr, - static_cast(out_identity_data.dtype), - out_identity_data.shape); - out_identity.set_rowwise_scale_inv(out_identity_scale_inv.data_ptr, - static_cast(out_identity_scale_inv.dtype), - out_identity_scale_inv.shape); - out_identity.set_amax(out_identity_amax.data_ptr, static_cast(out_identity_amax.dtype), - out_identity_amax.shape); - - NVTE_SCOPED_GIL_RELEASE( - { nvte_quantize_v2(input.data(), out_identity.data(), quant_config, stream); }); - } - - if (columnwise_usage) { - // Get the output columnwise data, scale_inv, and amax - auto out_columnwise_data = out.get_columnwise_data(); - auto out_columnwise_scale_inv = out.get_columnwise_scale_inv(); - // NOTE: should already be populated. - auto out_columnwise_amax = out.get_columnwise_amax(); - - // 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); - out_transpose.set_rowwise_scale_inv(out_columnwise_scale_inv.data_ptr, - static_cast(out_columnwise_scale_inv.dtype), - out_columnwise_scale_inv.shape); - out_transpose.set_amax(out_columnwise_amax.data_ptr, - static_cast(out_columnwise_amax.dtype), - out_columnwise_amax.shape); + // Fast math toggle: RHT transform can be accelerated + // What math is accelerated? Only the high precision math, so numerical impact is minimal + // 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 + const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); + if (use_fast_math) { + quant_config.set_use_fast_math(true); + quant_config_columnwise.set_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 + NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, + "RHT matrix is not available."); + auto rht_matrix_nvte = makeTransformerEngineTensor(this->rht_matrix); + // Fusion kernel that does the following: + // 1. Rowwise quantization + // 2. RHT followed by columnwise quantization & transpose + NVTE_SCOPED_GIL_RELEASE({ + nvte_quantize_with_hadamard_transform(input.data(), out.data(), rht_matrix_nvte.data(), + quant_config, stream); + }); + } else { // Use separate RNG state for columnwise to ensure different random numbers than rowwise - auto& columnwise_quant_config = + // This is only necessary because it's the unfused path where rowwise and columnwise + // are separate kernel launches + auto& columnwise_quant_config_to_use = need_separate_columnwise_rng ? quant_config_columnwise : quant_config; - - if (!eligible_for_rht_cast_fusion) { - // Invoking fallback RHT kernel. - - // If using RHT, then amax will be computed in the RHT step - // If not using RHT, then amax will be computed based on input x - at::Tensor rht_output_t; // The RHT(x_t) output, in columnwise layout - // This wrapper is going to be passed as input to the quantization kernel. - TensorWrapper rht_output_t_cpp; // Wrapper to contain the RHT(x) and RHT(x_t) outputs - rht_output_t = - allocateTorchTensor(static_cast(cols), static_cast(rows), input.dtype()); - // NOTE (frsun): This is non-intuitive, we are writing the - // result of transposed RHT to the output of rowwise. - rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), - std::vector{cols, rows}); - - 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, - this->rht_matrix_random_sign_mask_t, stream); - }); - - // Quantize kernel will treat everything as rowwise input/output, which is - // intended. - NVTE_SCOPED_GIL_RELEASE({ - nvte_quantize_v2(rht_output_t_cpp.data(), out_transpose.data(), columnwise_quant_config, - stream); - }); - } else { - // RHT cast fusion kernel. - NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, - "RHT matrix is not set"); - auto rht_matrix_nvte = makeTransformerEngineTensor(this->rht_matrix); - NVTE_SCOPED_GIL_RELEASE({ - nvte_hadamard_transform_cast_fusion_columnwise(input.data(), out_transpose.data(), - rht_matrix_nvte.data(), - columnwise_quant_config, stream); - }); - } + // unfused path also needs memory allocation for intermediate buffer for RHT output + at::Tensor rht_output_t; // The RHT(x_t) output, in columnwise layout + // This wrapper is going to be passed as input to the quantization kernel. + TensorWrapper rht_output_t_cpp; // Wrapper to contain the RHT(x) and RHT(x_t) outputs + rht_output_t = + allocateTorchTensor(static_cast(cols), static_cast(rows), input.dtype()); + // NOTE (frsun): This is non-intuitive, we are writing the + // result of transposed RHT to the output of rowwise. + rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), + std::vector{cols, rows}); + this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config, + columnwise_quant_config_to_use, stream); } } else { NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); }); diff --git a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py index f42183ec09..dd01ae05d3 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py @@ -500,8 +500,11 @@ def _quantize_blockwise_reference( 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) global_decode_scale = torch.div(1.0, global_encode_scale) + global_encode_scale_multiplier = global_encode_scale * torch.reciprocal(FLOAT4_E2M1_MAX) - decode_scale = decode_scale * global_encode_scale + # Match the kernel's default path: fold the FP4 reciprocal into the + # global scale multiplier, but keep the final reciprocal exact. + decode_scale = vec_max * global_encode_scale_multiplier decode_scale = torch.min( decode_scale, torch.tensor( From 401756576f61de9de5d6c26aa107eb16e232fd08 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 16 Mar 2026 12:09:51 -0700 Subject: [PATCH 285/521] [PyTorch] Backwards compatible single param checkpointing in `GroupedLinear` (#2761) * Load multi-param checkpoint from single-param config in GroupedLinear Signed-off-by: Kirthi Shankar Sivamani * Multi-param to single param case Signed-off-by: Kirthi Shankar Sivamani * Multi-param to single param case Signed-off-by: Kirthi Shankar Sivamani * Better varnames Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/test_grouped_tensor.py | 88 +++++++++++++++++++ .../pytorch/module/grouped_linear.py | 71 +++++++++++++++ 2 files changed, 159 insertions(+) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 9dd965fa94..225c6f6759 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -464,3 +464,91 @@ def test_clear(self) -> None: assert grouped_tensor.num_tensors == 0 assert grouped_tensor.rowwise_data is None assert grouped_tensor.logical_shape == (0, 0) + + def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> None: + """Load per-GEMM checkpoint from disk into single grouped parameter format.""" + num_gemms = 3 + in_features = 64 + out_features = 32 + dtype = torch.float32 + + src = te.GroupedLinear( + num_gemms=num_gemms, + in_features=in_features, + out_features=out_features, + params_dtype=dtype, + single_grouped_parameter=False, + ).cuda() + with torch.no_grad(): + for i in range(num_gemms): + getattr(src, f"weight{i}").copy_( + torch.randn(out_features, in_features, device="cuda", dtype=dtype) + ) + if src.use_bias: + getattr(src, f"bias{i}").copy_( + torch.randn(out_features, device="cuda", dtype=dtype) + ) + expected_weights = [getattr(src, f"weight{i}").detach().clone() for i in range(num_gemms)] + ckpt_path = tmp_path / "grouped_linear_per_gemm.pt" + torch.save(src.state_dict(), ckpt_path) + del src + + src_state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False) + + dst = te.GroupedLinear( + num_gemms=num_gemms, + in_features=in_features, + out_features=out_features, + params_dtype=dtype, + single_grouped_parameter=True, + ).cuda() + load_result = dst.load_state_dict(src_state_dict, strict=True) + assert len(load_result.missing_keys) == 0 + assert len(load_result.unexpected_keys) == 0 + + assert getattr(dst, "weight", None) is not None + loaded_weights = dst.weight.split_into_quantized_tensors() + assert len(loaded_weights) == num_gemms + for loaded_weight, expected_weight in zip(loaded_weights, expected_weights): + assert torch.equal(loaded_weight, expected_weight) + + def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> None: + """Load grouped-parameter checkpoint from disk into per-GEMM parameter format.""" + num_gemms = 3 + in_features = 64 + out_features = 32 + dtype = torch.float32 + + src = te.GroupedLinear( + num_gemms=num_gemms, + in_features=in_features, + out_features=out_features, + params_dtype=dtype, + single_grouped_parameter=True, + ).cuda() + with torch.no_grad(): + source_weights = src.weight.split_into_quantized_tensors() + for i in range(num_gemms): + source_weights[i].copy_( + torch.randn(out_features, in_features, device="cuda", dtype=dtype) + ) + expected_weights = [weight.detach().clone() for weight in source_weights] + ckpt_path = tmp_path / "grouped_linear_single_param.pt" + torch.save(src.state_dict(), ckpt_path) + del src + + src_state_dict = torch.load(ckpt_path, map_location="cpu", weights_only=False) + + dst = te.GroupedLinear( + num_gemms=num_gemms, + in_features=in_features, + out_features=out_features, + params_dtype=dtype, + single_grouped_parameter=False, + ).cuda() + load_result = dst.load_state_dict(src_state_dict, strict=True) + assert len(load_result.missing_keys) == 0 + assert len(load_result.unexpected_keys) == 0 + + for i, expected_weight in enumerate(expected_weights): + assert torch.equal(getattr(dst, f"weight{i}"), expected_weight) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index fade2957d5..30c1dbf408 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -846,6 +846,77 @@ def set_tensor_parallel_attributes(self, defer_init=False) -> None: elif self.parallel_mode == "column": set_tensor_model_parallel_attributes(getattr(self, f"bias{i}"), True, 0, 1) + def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None: + """Remap weight keys between single and per-GEMM checkpoint formats.""" + grouped_weight_key = f"{prefix}weight" + per_gemm_weight_keys = [f"{prefix}weight{i}" for i in range(self.num_gemms)] + has_grouped_weight = grouped_weight_key in state_dict + has_per_gemm_weights = all(key in state_dict for key in per_gemm_weight_keys) + + if self.single_grouped_parameter: + # Backward compatibility: checkpoints saved without single_grouped_parameter + # store one weight tensor per GEMM (weight0..weightN). Convert them into a + # single stacked grouped weight expected by this module configuration. + if not has_grouped_weight and has_per_gemm_weights: + per_gemm_weights = [state_dict.pop(key) for key in per_gemm_weight_keys] + per_gemm_weights = [ + weight.dequantize() if isinstance(weight, QuantizedTensorStorage) else weight + for weight in per_gemm_weights + ] + state_dict[grouped_weight_key] = torch.stack(per_gemm_weights, dim=0) + elif has_grouped_weight: + # Drop any redundant per-GEMM keys to avoid strict-load unexpected-key errors. + for key in per_gemm_weight_keys: + state_dict.pop(key, None) + else: + # Forward compatibility: checkpoints saved with single_grouped_parameter + # store one grouped `weight`. Convert it back to weight0..weightN. + if not has_per_gemm_weights and has_grouped_weight: + grouped_weight = state_dict.pop(grouped_weight_key) + if hasattr(grouped_weight, "split_into_quantized_tensors"): + grouped_members = grouped_weight.quantized_tensors + if grouped_members is None: + grouped_members = grouped_weight.split_into_quantized_tensors() + per_gemm_weights = [ + ( + weight.dequantize() + if isinstance(weight, QuantizedTensorStorage) + else weight + ) + for weight in grouped_members + ] + else: + grouped_weight = ( + grouped_weight.dequantize() + if isinstance(grouped_weight, QuantizedTensorStorage) + else grouped_weight + ) + per_gemm_weights = list(grouped_weight.unbind(dim=0)) + for i, weight in enumerate(per_gemm_weights): + state_dict[f"{prefix}weight{i}"] = weight + elif has_per_gemm_weights: + # Drop any redundant grouped key to avoid strict-load unexpected-key errors. + state_dict.pop(grouped_weight_key, None) + + def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False): + """Load state dict with grouped-weight format compatibility.""" + state_dict_copy = state_dict.copy() + metadata = getattr(state_dict, "_metadata", None) + if metadata is not None: + state_dict_copy._metadata = metadata + self._remap_grouped_weight_state_dict_keys(state_dict_copy, prefix="") + return super().load_state_dict(state_dict_copy, strict=strict, assign=assign) + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + """Load state, including compatibility across grouped-weight checkpoint formats.""" + self._remap_grouped_weight_state_dict_keys(state_dict, prefix) + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + @no_torch_dynamo() def forward( self, From 128f22e357380098d7524ae9a3202546aa23b0f9 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Mon, 16 Mar 2026 22:14:39 -0700 Subject: [PATCH 286/521] [JAX][Core] Fix Grouped GEMM cuBLAS version and SM arch checks (#2765) * Fix GMM cuBLAS version and SM arch checks Signed-off-by: Jeremy Berchtold * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Kirthi Shankar Sivamani * Update transformer_engine/common/gemm/cublaslt_grouped_gemm.cu Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Jeremy Berchtold Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Kirthi Shankar Sivamani Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../common/gemm/cublaslt_grouped_gemm.cu | 37 ++++++++++--------- transformer_engine/jax/cpp_extensions/gemm.py | 5 +++ 2 files changed, 25 insertions(+), 17 deletions(-) diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index ccf1e53ba4..5031a30485 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -29,10 +29,13 @@ inline void CreateCublasHandle(cublasLtHandle_t *handle) { } // namespace -// MXFP8 support for grouped GEMM requires cuBLAS 13.2+ -#define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130200 +// MXFP8 support for grouped GEMM requires cuBLAS 13.3+ +#define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300 +// BF16 support for grouped GEMM requires cuBLAS 13.3+ +// cuBLAS 13.2 is mostly functional but contains a bug for wgrad when a group has k=0, the weight gradient will be uninitialized random data instead of zeros. +#define CUBLAS_GROUPED_GEMM_VERSION 130300 -#if CUBLAS_VERSION >= 130200 +#if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION namespace { @@ -278,8 +281,8 @@ 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() >= 130200, api_name, - " requires cuBLAS 13.2+, but run-time cuBLAS version is ", + 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()); } @@ -1320,15 +1323,15 @@ void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTens NVTE_CHECK_CUDA(cudaGetLastError()); } -#else // CUBLAS_VERSION < 130200 +#else // CUBLAS_VERSION < CUBLAS_GROUPED_GEMM_VERSION 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, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.2+, but compile-time cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); + NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.3+, but compile-time cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -1338,9 +1341,9 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { NVTE_ERROR( - "nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.2+, but compile-time " + "nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.3+, but compile-time " "cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); + CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -1351,26 +1354,26 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { NVTE_ERROR( - "nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.2+, but compile-time " + "nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.3+, but compile-time " "cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); + CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); } void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, cudaStream_t stream) { - NVTE_ERROR("nvte_grouped_bias_add requires cuBLAS 13.2+, but compile-time cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); + NVTE_ERROR("nvte_grouped_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version is ", + CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 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.2+, but compile-time cuBLAS " + "nvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS " "version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.1 or newer."); + CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); return 0; } -#endif // CUBLAS_VERSION >= 130200 +#endif // CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION namespace { diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index 515f02af6e..aaf8e8ecea 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -1936,6 +1936,11 @@ def _can_use_v2_grouped_gemm( if not _v2_grouped_gemm_available: return False + # nvte_grouped_gemm (the v2 kernel) requires SM100+ (Blackwell or newer). + # Fall back to the v1 path on SM90 (Hopper) and older architectures. + if get_device_compute_capability(0) < 100: + return False + return scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16 and not has_bias From 4e339a5a56b87f2248ba49aa1d4cac6b71fac7c6 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 17 Mar 2026 16:38:43 -0700 Subject: [PATCH 287/521] Update vermin version to fix precommit CI error with python 3.14 (#2773) * Pin python 3.13 in vermin check Signed-off-by: Kirthi Shankar Sivamani * Update vermin version for python 3.14 support Signed-off-by: Kirthi Shankar Sivamani * Use sha instead of tag Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- .pre-commit-config.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 76f476eb3f..601149916b 100755 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -40,7 +40,7 @@ repos: files: ^transformer_engine.*\.(c|cc|cxx|cpp|cu|cuh|h|hpp)$ - repo: https://github.com/netromdk/vermin - rev: c75aca72f4e85c6e47252139e8695f1c8b5f9ae3 + rev: b70ff9611a01a2bf2f702aa537d14e71e330edba hooks: - id: vermin args: ['-t=3.10-', '--violations'] From 53a41b297bea500544efb8d45576d67a0d72c480 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 17 Mar 2026 21:44:23 -0700 Subject: [PATCH 288/521] Update cudnnFE to v1.20.0 (#2774) Signed-off-by: Kirthi Shankar Sivamani --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 8d19d3182b..d33027a41a 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 8d19d3182bfbc304046a15e9236bec9ff31511fc +Subproject commit d33027a41a93af9c85f089c6364ab415fce98982 From 3e61687a7b3c42225d610f24e9a37cac366641e4 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, 18 Mar 2026 19:33:32 +0100 Subject: [PATCH 289/521] [PyTorch] torch.compile support for permutation functions (#2686) * init Signed-off-by: Pawel Gadzinski * work finished Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint fixes Signed-off-by: Pawel Gadzinski * fixes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: root * removed warning.warn Signed-off-by: root * [PyTorch] Remove dead None-check for num_out_tokens in moe_permute_mask_map_forward num_out_tokens is typed as int in the custom_op signature and can never be None; the check was incorrectly carried over from the class-based upstream version during merge conflict resolution. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: root Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 --- tests/pytorch/test_permutation.py | 148 +- transformer_engine/pytorch/permutation.py | 1630 ++++++++++------- .../pytorch/quantized_tensor.py | 12 + 3 files changed, 1137 insertions(+), 653 deletions(-) diff --git a/tests/pytorch/test_permutation.py b/tests/pytorch/test_permutation.py index be1ff30472..66c685e139 100644 --- a/tests/pytorch/test_permutation.py +++ b/tests/pytorch/test_permutation.py @@ -218,6 +218,17 @@ def backward_wrapper( return act.backward(backward_input, retain_graph=retain_graph) +def _maybe_compile(fn, use_torch_compile): + """Wrap fn with torch.compile(fullgraph=True) if requested.""" + if use_torch_compile: + torch._dynamo.reset() + import torch._functorch.config as functorch_config + + functorch_config.donated_buffer = False + return torch.compile(fn, fullgraph=True) + return fn + + def _test_permutation_index_map( te_dtype, num_tokens, @@ -227,6 +238,7 @@ def _test_permutation_index_map( num_out_tokens, with_probs, BENCHMARK=False, + use_torch_compile=False, ): if not with_probs and topK > 1: pytest.skip("Only permutations with topK=1 and without probabilities are supported.") @@ -298,9 +310,13 @@ def _test_permutation_index_map( te_permute_fwd_input.requires_grad_(True) te_permute_bwd_input = pytorch_permute_bwd_input.detach() - te_permute_output, row_id_map = te_permute( - te_permute_fwd_input, indices, num_out_tokens, map_type="index" + _permute = _maybe_compile( + lambda inp, idx, num_out, max_token: te_permute( + inp, idx, num_out, max_token, map_type="index" + ), + use_torch_compile, ) + te_permute_output, row_id_map = _permute(te_permute_fwd_input, indices, num_out_tokens, -1) te_permute_output.backward(te_permute_bwd_input, retain_graph=True) te_probs = None @@ -311,9 +327,11 @@ def _test_permutation_index_map( te_unpermute_fwd_input.requires_grad_(True) te_unpermute_bwd_input = pytorch_unpermute_bwd_input.detach() - te_unpermute_output = te_unpermute( - te_unpermute_fwd_input, row_id_map, te_probs, map_type="index" + _unpermute = _maybe_compile( + lambda inp, row_map, probs_val: te_unpermute(inp, row_map, probs_val, map_type="index"), + use_torch_compile, ) + te_unpermute_output = _unpermute(te_unpermute_fwd_input, row_id_map, te_probs) te_unpermute_output.backward(te_unpermute_bwd_input, retain_graph=True) ################################################################################################################################### @@ -444,6 +462,7 @@ def _test_permutation_mask_map( num_out_tokens, with_probs, BENCHMARK=False, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -514,9 +533,11 @@ def _test_permutation_mask_map( te_permute_fwd_input.requires_grad_(True) te_permute_bwd_input = pytorch_permute_bwd_input.detach() - te_permute_output, row_id_map = te_permute( - te_permute_fwd_input, routing_map, num_out_tokens=num_out_tokens, map_type="mask" + _permute = _maybe_compile( + lambda inp, rmap, n_out: te_permute(inp, rmap, num_out_tokens=n_out, map_type="mask"), + use_torch_compile, ) + te_permute_output, row_id_map = _permute(te_permute_fwd_input, routing_map, num_out_tokens) te_permute_output.backward(te_permute_bwd_input, retain_graph=True) te_probs = None @@ -527,9 +548,11 @@ def _test_permutation_mask_map( te_unpermute_fwd_input.requires_grad_(True) te_unpermute_bwd_input = pytorch_unpermute_bwd_input.detach() - te_unpermute_output = te_unpermute( - te_unpermute_fwd_input, row_id_map, te_probs, restore_shape, map_type="mask" + _unpermute = _maybe_compile( + lambda inp, row_map, p, rs: te_unpermute(inp, row_map, p, rs, map_type="mask"), + use_torch_compile, ) + te_unpermute_output = _unpermute(te_unpermute_fwd_input, row_id_map, te_probs, restore_shape) te_unpermute_output.backward(te_unpermute_bwd_input, retain_graph=True) ################################################################################################################################### @@ -666,6 +689,7 @@ def _test_permutation_and_padding_mask_map( with_merging_probs=False, align_size=16, BENCHMARK=False, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -957,6 +981,7 @@ def _test_permutation_and_padding_with_merging_probs( num_out_tokens, align_size=16, BENCHMARK=False, + use_torch_compile=False, ): """ Test the combination of merging_probs AND pad_offsets together in moe_unpermute. @@ -1180,6 +1205,7 @@ def _test_permutation_mask_map_fp8( topK, num_out_tokens, recipe, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -1255,9 +1281,11 @@ def _test_permutation_mask_map_fp8( ) # TE Permutation - permute_output, _ = te_permute( - permute_fwd_input_fp8, routing_map, num_out_tokens=num_out_tokens, map_type="mask" + _permute = _maybe_compile( + lambda inp, rmap, n_out: te_permute(inp, rmap, num_out_tokens=n_out, map_type="mask"), + use_torch_compile, ) + permute_output, _ = _permute(permute_fwd_input_fp8, routing_map, num_out_tokens) if recipe.float8_block_scaling(): te_permute_output = permute_output._rowwise_data te_permute_scale_output = permute_output._rowwise_scale_inv.T.contiguous() @@ -1291,6 +1319,7 @@ def _test_moe_chunk_sort( tp_size, hidden_size, BENCHMARK=False, + use_torch_compile=False, ): print( "chunk permute:" @@ -1340,7 +1369,11 @@ def _test_moe_chunk_sort( te_fwd_input.requires_grad_(True) te_bwd_input = pytorch_bwd_input.detach() - te_output = te_sort_chunks_by_index(te_fwd_input, split_sizes_cuda, sorted_idxs_cuda) + _sort = _maybe_compile( + lambda inp, ss, si: te_sort_chunks_by_index(inp, ss, si), + use_torch_compile, + ) + te_output = _sort(te_fwd_input, split_sizes_cuda, sorted_idxs_cuda) te_output.backward(te_bwd_input, retain_graph=True) ################################################################################################################################### @@ -1415,6 +1448,7 @@ def _test_permutation_mask_map_alongside_probs( num_out_tokens, tp_size, BENCHMARK=False, + use_torch_compile=False, ): if topK > num_expert: pytest.skip("topK should be smaller than the number of experts.") @@ -1510,30 +1544,27 @@ def _test_permutation_mask_map_alongside_probs( te_probs = probs.detach() te_probs.requires_grad_(True) - te_permute_output, te_permuted_probs, row_id_map = te_permute_with_probs( + def _alongside_probs_fn(fwd_inp, t_probs, rmap, ss1, si1, ss2, si2): + out, pprobs, rid = te_permute_with_probs( + fwd_inp, t_probs, rmap, num_out_tokens=num_out_tokens + ) + out, pprobs = te_sort_chunks_by_index_with_probs(out, pprobs, ss1, si1) + out_dtype = out.dtype + out = out * pprobs.unsqueeze(-1) + out = out.to(dtype=out_dtype) + out = te_sort_chunks_by_index(out, ss2, si2) + out = te_unpermute(out, rid, restore_shape=restore_shape, map_type="mask") + return out + + _fn = _maybe_compile(_alongside_probs_fn, use_torch_compile) + te_unpermute_output = _fn( te_permute_fwd_input, te_probs, routing_map, - num_out_tokens=num_out_tokens, - ) - - te_permute_output, te_permuted_probs = te_sort_chunks_by_index_with_probs( - te_permute_output, te_permuted_probs, split_sizes_cuda, sorted_idxs_cuda - ) - - te_permute_output_dtype = te_permute_output.dtype - te_permute_output = te_permute_output * te_permuted_probs.unsqueeze(-1) - te_permute_output = te_permute_output.to(dtype=te_permute_output_dtype) - - te_permute_output = te_sort_chunks_by_index( - te_permute_output, split_sizes_2_cuda, sorted_idxs_2_cuda - ) - - te_unpermute_output = te_unpermute( - te_permute_output, - row_id_map, - restore_shape=restore_shape, - map_type="mask", + split_sizes_cuda, + sorted_idxs_cuda, + split_sizes_2_cuda, + sorted_idxs_2_cuda, ) te_unpermute_output.backward(te_unpermute_bwd_input, retain_graph=True) @@ -1647,6 +1678,7 @@ def perf_test_cuda_kernel(cuda_kernel_fn): @pytest.mark.parametrize("hidden_size", [4096]) @pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_index_map( te_dtype, num_tokens, @@ -1654,7 +1686,10 @@ def test_permutation_index_map( hidden_size, topK, num_out_tokens, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or topK != 2): + pytest.skip("torch.compile tested with single config only") with_probs = True BENCHMARK = False @@ -1667,6 +1702,7 @@ def test_permutation_index_map( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1676,6 +1712,7 @@ def test_permutation_index_map( @pytest.mark.parametrize("hidden_size", [4096]) @pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_mask_map( te_dtype, num_tokens, @@ -1683,7 +1720,10 @@ def test_permutation_mask_map( hidden_size, topK, num_out_tokens, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or topK != 2): + pytest.skip("torch.compile tested with single config only") with_probs = True BENCHMARK = False @@ -1696,6 +1736,7 @@ def test_permutation_mask_map( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1711,6 +1752,7 @@ def test_permutation_mask_map( ], ) @pytest.mark.parametrize("with_merging_probs", [True, False]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_and_padding_mask_map( te_dtype, num_tokens, @@ -1719,7 +1761,10 @@ def test_permutation_and_padding_mask_map( topK, num_out_tokens, with_merging_probs, + use_torch_compile, ): + if use_torch_compile and (num_expert != 8 or topK != 2): + pytest.skip("torch.compile tested with single config only") BENCHMARK = False _test_permutation_and_padding_mask_map( @@ -1731,6 +1776,7 @@ def test_permutation_and_padding_mask_map( num_out_tokens=num_out_tokens, with_merging_probs=with_merging_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1745,6 +1791,7 @@ def test_permutation_and_padding_mask_map( (4096, 512, 9216, 8), ], ) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_and_padding_with_merging_probs( te_dtype, num_tokens, @@ -1752,8 +1799,11 @@ def test_permutation_and_padding_with_merging_probs( hidden_size, topK, num_out_tokens, + use_torch_compile, ): """Test moe_unpermute backward pass with BOTH merging_probs AND pad_offsets.""" + if use_torch_compile and (num_expert != 8 or topK != 2): + pytest.skip("torch.compile tested with single config only") BENCHMARK = False _test_permutation_and_padding_with_merging_probs( @@ -1764,11 +1814,13 @@ def test_permutation_and_padding_with_merging_probs( topK=topK, num_out_tokens=num_out_tokens, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @pytest.mark.parametrize("te_dtype", _te_dtypes) -def test_permutation_mask_map_empty_input(te_dtype): +@pytest.mark.parametrize("use_torch_compile", [False, True]) +def test_permutation_mask_map_empty_input(te_dtype, use_torch_compile): with_probs = True BENCHMARK = False @@ -1781,6 +1833,7 @@ def test_permutation_mask_map_empty_input(te_dtype): num_out_tokens=0, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1791,6 +1844,7 @@ def test_permutation_mask_map_empty_input(te_dtype): @pytest.mark.parametrize("topK", [2, 5]) @pytest.mark.parametrize("num_out_tokens", [None, 2039]) @pytest.mark.parametrize("tp_size", [1, 2]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_mask_map_alongside_probs( te_dtype, num_tokens, @@ -1799,7 +1853,10 @@ def test_permutation_mask_map_alongside_probs( topK, num_out_tokens, tp_size, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or topK != 2 or tp_size != 1): + pytest.skip("torch.compile tested with single config only") _test_permutation_mask_map_alongside_probs( te_dtype=te_dtype, num_tokens=num_tokens, @@ -1808,11 +1865,13 @@ def test_permutation_mask_map_alongside_probs( topK=topK, num_out_tokens=num_out_tokens, tp_size=tp_size, + use_torch_compile=use_torch_compile, ) @pytest.mark.parametrize("te_dtype", _te_dtypes) -def test_permutation_mask_map_alongside_probs_empty_input(te_dtype): +@pytest.mark.parametrize("use_torch_compile", [False, True]) +def test_permutation_mask_map_alongside_probs_empty_input(te_dtype, use_torch_compile): _test_permutation_mask_map_alongside_probs( te_dtype=te_dtype, num_tokens=0, @@ -1821,6 +1880,7 @@ def test_permutation_mask_map_alongside_probs_empty_input(te_dtype): topK=2, num_out_tokens=0, tp_size=2, + use_torch_compile=use_torch_compile, ) @@ -1868,6 +1928,7 @@ def test_permutation_mask_map_fp8( topK=topK, num_out_tokens=num_out_tokens, recipe=recipe, + use_torch_compile=False, # FP8 permutation is not yet supported under torch.compile ) @@ -1875,12 +1936,16 @@ def test_permutation_mask_map_fp8( @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_index_map_topk1_no_probs( te_dtype, num_tokens, num_expert, hidden_size, + use_torch_compile, ): + if use_torch_compile and num_expert != 7: + pytest.skip("torch.compile tested with single config only") topK = 1 num_out_tokens = None with_probs = False @@ -1895,6 +1960,7 @@ def test_permutation_index_map_topk1_no_probs( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1902,12 +1968,16 @@ def test_permutation_index_map_topk1_no_probs( @pytest.mark.parametrize("num_tokens", [4096]) @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_permutation_mask_map_topk1_no_probs( te_dtype, num_tokens, num_expert, hidden_size, + use_torch_compile, ): + if use_torch_compile and num_expert != 7: + pytest.skip("torch.compile tested with single config only") topK = 1 num_out_tokens = None with_probs = False @@ -1922,6 +1992,7 @@ def test_permutation_mask_map_topk1_no_probs( num_out_tokens=num_out_tokens, with_probs=with_probs, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @@ -1930,13 +2001,17 @@ def test_permutation_mask_map_topk1_no_probs( @pytest.mark.parametrize("num_expert", [7, 16]) @pytest.mark.parametrize("tp_size", [2, 8]) @pytest.mark.parametrize("hidden_size", [4096]) +@pytest.mark.parametrize("use_torch_compile", [False, True]) def test_chunk_permutation( te_dtype, num_tokens, num_expert, tp_size, hidden_size, + use_torch_compile, ): + if use_torch_compile and (num_expert != 7 or tp_size != 2): + pytest.skip("torch.compile tested with single config only") BENCHMARK = False _test_moe_chunk_sort( @@ -1946,11 +2021,13 @@ def test_chunk_permutation( tp_size=tp_size, hidden_size=hidden_size, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) @pytest.mark.parametrize("te_dtype", _te_dtypes) -def test_chunk_permutation_empty_input(te_dtype): +@pytest.mark.parametrize("use_torch_compile", [False, True]) +def test_chunk_permutation_empty_input(te_dtype, use_torch_compile): BENCHMARK = False _test_moe_chunk_sort( @@ -1960,6 +2037,7 @@ def test_chunk_permutation_empty_input(te_dtype): tp_size=2, hidden_size=4096, BENCHMARK=BENCHMARK, + use_torch_compile=use_torch_compile, ) diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index ca59a0ebf8..bc9a2660b7 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -6,11 +6,13 @@ import warnings from typing import Optional, Tuple import torch - import transformer_engine_torch as tex import transformer_engine.pytorch.triton.permutation as triton_permutation from transformer_engine.pytorch.constants import TE_DType -from transformer_engine.pytorch.quantized_tensor import QuantizedTensor +from transformer_engine.pytorch.quantized_tensor import ( + QuantizedTensor, + _quantized_tensor_passthrough_ops, +) from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Tensor @@ -22,557 +24,829 @@ ] -class _moe_permute_index_map(torch.autograd.Function): - """functional Permute with index router map""" - - workspace = None - max_expanded_token_num = 0 - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - index: torch.Tensor, - num_out_tokens: int, - max_token_num: int, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # pylint: disable=missing-function-docstring - # Empty input check - if not inp.numel(): - return inp, torch.tensor([], device=inp.device) - - # Device check - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not index.is_cuda: - raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") - # Shape check - if inp.size(0) != index.size(0): - raise ValueError( - f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " - f"index.size(0) ({index.size(0)})." - ) +# ===================== _moe_permute_index_map custom ops ===================== + +# Workspace state for moe_permute_index_map +_moe_permute_index_map_workspace = None +_moe_permute_index_map_max_expanded_token_num = 0 - # Data type check - dtype = TE_DType[inp.dtype] - if index.dtype != torch.int32: - warnings.warn( - f"The data type of the input `index` of Permute is {index.dtype}! " - "The recommended type is torch.int32." - ) - index = index.to(torch.int32) - topK = index.size(1) +@torch.library.custom_op("te_moe::permute_index_map", mutates_args=[]) +def moe_permute_index_map_forward( + inp: torch.Tensor, + index: torch.Tensor, + num_out_tokens: int, + max_token_num: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass for MoE permute with index router map.""" + global _moe_permute_index_map_workspace, _moe_permute_index_map_max_expanded_token_num - input_max_expanded_token_num = max(max_token_num, inp.size(0)) * topK - if _moe_permute_index_map.max_expanded_token_num < input_max_expanded_token_num: - _moe_permute_index_map.max_expanded_token_num = input_max_expanded_token_num - _moe_permute_index_map.workspace = [] + if not inp.numel(): + return inp.clone(), torch.tensor([], device=inp.device) - permuted_act, row_id_map, _moe_permute_index_map.workspace = tex.moe_permute_fwd( - inp, - dtype, - index, - num_out_tokens, - _moe_permute_index_map.workspace, - _moe_permute_index_map.max_expanded_token_num, + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not index.is_cuda: + raise ValueError(f"index must be a CUDA tensor, but got tensor on {index.device}.") + if inp.size(0) != index.size(0): + raise ValueError( + f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " + f"index.size(0) ({index.size(0)})." + ) + if index.dtype != torch.int32: + warnings.warn( + f"The data type of the input `index` of Permute is {index.dtype}! " + "The recommended type is torch.int32." ) + index = index.to(torch.int32) - ctx.row_id_map = row_id_map - ctx.num_tokens = index.size(0) - ctx.topK = index.size(1) - return permuted_act, row_id_map - - @staticmethod - def backward( - ctx, - permuted_act_grad: torch.Tensor, - _, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - # Empty input check - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, None - - if not permuted_act_grad.is_contiguous(): - permuted_act_grad = permuted_act_grad.contiguous() - - dtype = TE_DType[permuted_act_grad.dtype] - act_grad = None - if ctx.needs_input_grad[0]: - act_grad = tex.moe_permute_bwd( - permuted_act_grad, dtype, ctx.row_id_map, torch.empty(0), ctx.num_tokens, ctx.topK - ) + dtype = TE_DType[inp.dtype] - return act_grad, None, None, None + topK = index.size(1) + input_max_expanded_token_num = max(max_token_num, inp.size(0)) * topK + if _moe_permute_index_map_max_expanded_token_num < input_max_expanded_token_num: + _moe_permute_index_map_max_expanded_token_num = input_max_expanded_token_num + _moe_permute_index_map_workspace = [] -class _moe_unpermute_index_map(torch.autograd.Function): - """functional Unpermute with index router map""" + permuted_act, row_id_map, _moe_permute_index_map_workspace = tex.moe_permute_fwd( + inp, + dtype, + index, + num_out_tokens, + _moe_permute_index_map_workspace, + _moe_permute_index_map_max_expanded_token_num, + ) - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - row_id_map: torch.Tensor, - probs: torch.Tensor, - ) -> torch.Tensor: - # pylint: disable=missing-function-docstring - # Empty input check - if not inp.numel(): - ctx.probs = probs - return inp + return permuted_act, row_id_map - # None probs check - if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") - if probs.dtype != torch.float32: - warnings.warn( - f"The data type of the input `probs` of Unpermute is {probs.dtype}! " - "The recommended type is torch.float32." - ) - probs = probs.to(torch.float32) +@moe_permute_index_map_forward.register_fake +def _moe_permute_index_map_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + index: torch.Tensor, + num_out_tokens: int, + max_token_num: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake implementation for shape inference.""" + num_tokens = inp.shape[0] + topK = index.shape[1] - num_tokens = probs.size(0) - topK = probs.size(1) - else: - num_tokens = row_id_map.size(0) - topK = 1 - probs = torch.empty(0) + # Infer output shape + output_tokens = num_out_tokens if num_out_tokens > 0 else num_tokens * topK - # Device check - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not row_id_map.is_cuda: - raise ValueError( - f"row_id_map must be a CUDA tensor, but got tensor on {row_id_map.device}." - ) + # row_id_map is 1D with size = num_tokens * topK + fake_output = torch.empty((output_tokens, inp.shape[1]), dtype=inp.dtype, device=inp.device) + fake_row_id_map = torch.empty((num_tokens * topK,), dtype=torch.int32, device=inp.device) - # Data type check - dtype = TE_DType[inp.dtype] - if row_id_map.dtype != torch.int32: - warnings.warn( - f"The data type of the input `row_id_map` of Unpermute is {row_id_map.dtype}! " - "The recommended type is torch.int32." - ) - row_id_map = row_id_map.to(torch.int32) + return fake_output, fake_row_id_map + + +@torch.library.custom_op("te_moe::permute_index_map_bwd", mutates_args=[]) +def moe_permute_index_map_backward( + grad_permuted_act: torch.Tensor, + row_id_map: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Backward pass for MoE permute with index router map.""" + dtype = TE_DType[grad_permuted_act.dtype] + act_grad = tex.moe_permute_bwd( + grad_permuted_act, dtype, row_id_map, torch.empty(0), num_tokens, topK + ) + return act_grad + + +@moe_permute_index_map_backward.register_fake +def _moe_permute_index_map_backward_fake( # pylint: disable=unused-argument + grad_permuted_act: torch.Tensor, + row_id_map: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Fake implementation for shape inference of backward.""" + return torch.empty( + (num_tokens, grad_permuted_act.shape[1]), + dtype=grad_permuted_act.dtype, + device=grad_permuted_act.device, + ) + + +def _moe_permute_index_map_setup_context(ctx, inputs, output): + """Save context for backward pass.""" + inp, index, _num_out_tokens, _max_token_num = inputs + _permuted_act, row_id_map = output + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(row_id_map) + ctx.num_tokens = index.size(0) if not ctx.empty_input else 0 + ctx.topK = index.size(1) if not ctx.empty_input else 1 + + +def _moe_permute_index_map_backward_wrapper( + ctx, grad_permuted_act, grad_row_id_map +): # pylint: disable=unused-argument + """Backward pass wrapper that calls the custom backward op.""" + if ctx.empty_input: + return grad_permuted_act, None, None, None + + if not grad_permuted_act.is_contiguous(): + grad_permuted_act = grad_permuted_act.contiguous() + + (row_id_map,) = ctx.saved_tensors + act_grad = torch.ops.te_moe.permute_index_map_bwd( + grad_permuted_act, row_id_map, ctx.num_tokens, ctx.topK + ) - unpermuted_output = tex.moe_unpermute_fwd(inp, dtype, row_id_map, probs, num_tokens, topK) + return act_grad, None, None, None - ctx.save_for_backward(inp, row_id_map, probs) - return unpermuted_output - @staticmethod - def backward( - ctx, - unpermuted_act_grad: torch.Tensor, - ) -> Tuple[torch.Tensor, None, torch.Tensor]: - # pylint: disable=missing-function-docstring - # Empty input check - if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.probs +moe_permute_index_map_forward.register_autograd( + _moe_permute_index_map_backward_wrapper, + setup_context=_moe_permute_index_map_setup_context, +) - if not unpermuted_act_grad.is_contiguous(): - unpermuted_act_grad = unpermuted_act_grad.contiguous() - dtype = TE_DType[unpermuted_act_grad.dtype] - inp, row_id_map, probs = ctx.saved_tensors +# ===================== _moe_unpermute_index_map custom ops ===================== - act_grad = None + +@torch.library.custom_op("te_moe::unpermute_index_map_fwd", mutates_args=[]) +def moe_unpermute_index_map_forward( + inp: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Forward pass for MoE unpermute with index router map.""" + if not inp.numel(): + return inp.clone() + dtype = TE_DType[inp.dtype] + return tex.moe_unpermute_fwd(inp, dtype, row_id_map, probs, num_tokens, topK) + + +@moe_unpermute_index_map_forward.register_fake +def _moe_unpermute_index_map_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, + num_tokens: int, + topK: int, +) -> torch.Tensor: + """Fake implementation for shape inference.""" + # Output shape: (num_tokens, hidden_size) + return torch.empty((num_tokens, inp.shape[1]), dtype=inp.dtype, device=inp.device) + + +@torch.library.custom_op("te_moe::unpermute_index_map_bwd", mutates_args=[]) +def moe_unpermute_index_map_backward( + unpermuted_act_grad: torch.Tensor, + fwd_input: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE unpermute with index router map.""" + dtype = TE_DType[unpermuted_act_grad.dtype] + act_grad, prob_grad = tex.moe_unpermute_bwd( + unpermuted_act_grad, fwd_input, dtype, row_id_map, probs + ) + return act_grad, prob_grad + + +@moe_unpermute_index_map_backward.register_fake +def _moe_unpermute_index_map_backward_fake( + unpermuted_act_grad: torch.Tensor, + fwd_input: torch.Tensor, + row_id_map: torch.Tensor, + probs: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake implementation for shape inference of backward.""" + # act_grad shape: (fwd_input.size(0), hidden_size) + # prob_grad shape: (num_tokens, topK) + topK = probs.size(1) if probs.numel() > 0 else 1 + num_tokens = probs.size(0) if probs.numel() > 0 else row_id_map.size(0) + act_grad = torch.empty( + (fwd_input.size(0), unpermuted_act_grad.shape[1]), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + prob_grad = torch.empty( + (num_tokens, topK), dtype=torch.float32, device=unpermuted_act_grad.device + ) + return act_grad, prob_grad + + +def _moe_unpermute_index_map_setup_context(ctx, inputs, output): # pylint: disable=unused-argument + """Save context for backward pass.""" + inp, row_id_map, probs, _num_tokens, _topK = inputs + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(inp, row_id_map, probs) + ctx.needs_probs_grad = probs.requires_grad + + +def _moe_unpermute_index_map_backward_wrapper(ctx, unpermuted_act_grad): + """Backward pass wrapper that calls the custom backward op.""" + if ctx.empty_input: + prob_grad = torch.zeros_like(ctx.saved_tensors[2]) if ctx.needs_probs_grad else None + return unpermuted_act_grad, None, prob_grad, None, None + + if not unpermuted_act_grad.is_contiguous(): + unpermuted_act_grad = unpermuted_act_grad.contiguous() + + inp, row_id_map, probs = ctx.saved_tensors + + act_grad, prob_grad = torch.ops.te_moe.unpermute_index_map_bwd( + unpermuted_act_grad, inp, row_id_map, probs + ) + + if not ctx.needs_probs_grad: prob_grad = None - if ctx.needs_input_grad[0]: - act_grad, prob_grad = tex.moe_unpermute_bwd( - unpermuted_act_grad, inp, dtype, row_id_map, probs - ) - if not ctx.needs_input_grad[2]: - prob_grad = None - - return act_grad, None, prob_grad - - -class _moe_permute_mask_map(torch.autograd.Function): - """functional Permute with mask router map""" - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - routing_map: torch.Tensor, - num_out_tokens: int, - probs: torch.Tensor, - pad_offsets: Optional[torch.Tensor], - ) -> Tuple[torch.Tensor, torch.Tensor]: - # pylint: disable=missing-function-docstring - if not inp.numel(): - ctx.probs = probs - return inp, torch.tensor([], device=inp.device), torch.tensor([], device=inp.device) - - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not routing_map.is_cuda: + + return act_grad, None, prob_grad, None, None + + +moe_unpermute_index_map_forward.register_autograd( + _moe_unpermute_index_map_backward_wrapper, + setup_context=_moe_unpermute_index_map_setup_context, +) + + +# ===================== _moe_permute_mask_map custom ops ===================== + + +@torch.library.custom_op("te_moe::permute_mask_map_fwd", mutates_args=[]) +def moe_permute_mask_map_forward( + inp: torch.Tensor, + routing_map: torch.Tensor, + num_out_tokens: int, + probs: Optional[torch.Tensor], + pad_offsets: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass for MoE permute with mask router map.""" + if not inp.numel(): + return inp.clone(), torch.tensor([], device=inp.device), torch.tensor([], device=inp.device) + + if not inp.is_cuda: + raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") + if not routing_map.is_cuda: + raise ValueError( + f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + ) + if probs is not None: + if not probs.is_cuda: + raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") + if pad_offsets is not None: + if not pad_offsets.is_cuda: raise ValueError( - f"routing_map must be a CUDA tensor, but got tensor on {routing_map.device}." + f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." ) - if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") - if pad_offsets is not None: - if not pad_offsets.is_cuda: + if inp.size(0) != routing_map.size(0): + raise ValueError( + f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " + f"routing_map.size(0) ({routing_map.size(0)})." + ) + num_tokens, hidden_size = inp.size() + num_experts = routing_map.size(1) + + row_id_map = triton_permutation.make_row_id_map(routing_map, num_tokens, num_experts) + + # FP8 handling + fp8 = isinstance(inp, QuantizedTensor) + per_tensor_recipe = isinstance(inp, Float8Tensor) + blockwise_recipe = isinstance(inp, Float8BlockwiseQTensor) + mxfp8_recipe = isinstance(inp, MXFP8Tensor) + + if fp8: + fp8_dtype = inp._fp8_dtype + fake_dtype = inp.dtype + # blockwise scaling + if blockwise_recipe: + fp8_scale = inp._rowwise_scale_inv.T.contiguous() + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: raise ValueError( - f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Input shape: ({num_tokens}, {hidden_size}), " + f"scale shape: {tuple(fp8_scale.shape)}." ) - - if inp.size(0) != routing_map.size(0): - raise ValueError( - f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " - f"routing_map.size(0) ({routing_map.size(0)})." - ) - num_tokens, hidden_size = inp.size() - num_experts = routing_map.size(1) - if num_out_tokens is None: - raise ValueError("num_out_tokens must be provided to the fused permute function.") - - row_id_map = triton_permutation.make_row_id_map(routing_map, num_tokens, num_experts) - - fp8 = isinstance(inp, QuantizedTensor) - per_tensor_recipe = isinstance(inp, Float8Tensor) - blockwise_recipe = isinstance(inp, Float8BlockwiseQTensor) - mxfp8_recipe = isinstance(inp, MXFP8Tensor) - - if fp8: - fp8_dtype = inp._fp8_dtype - fake_dtype = inp.dtype - # blockwise scaling - if blockwise_recipe: - fp8_scale = inp._rowwise_scale_inv.T.contiguous() - scale_hidden_dim = fp8_scale.shape[1] - if num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Input shape: ({num_tokens}, {hidden_size}), " - f"scale shape: {tuple(fp8_scale.shape)}." - ) - inp = inp._rowwise_data - # mxfp8 scaling - elif mxfp8_recipe: - fp8_scale = inp._rowwise_scale_inv.contiguous() - scale_hidden_dim = fp8_scale.shape[1] - if num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Input shape: ({num_tokens}, {hidden_size}), " - f"scale shape: {tuple(fp8_scale.shape)}." - ) - inp = inp._rowwise_data - # per-tensor scaling - elif per_tensor_recipe: - # Kernel does not need scale in per-tensor scaling - fp8_scale = None - scale_hidden_dim = None - fp8_scale_inv = inp._scale_inv - inp = inp._data - else: - raise ValueError("Unsupported FP8 recipe") - else: + inp = inp._rowwise_data + # mxfp8 scaling + elif mxfp8_recipe: + fp8_scale = inp._rowwise_scale_inv.contiguous() + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: + raise ValueError( + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Input shape: ({num_tokens}, {hidden_size}), " + f"scale shape: {tuple(fp8_scale.shape)}." + ) + inp = inp._rowwise_data + # per-tensor scaling + elif per_tensor_recipe: + # Kernel does not need scale in per-tensor scaling fp8_scale = None - fp8_dtype = None scale_hidden_dim = None + fp8_scale_inv = inp._scale_inv + inp = inp._data + else: + raise ValueError("Unsupported FP8 recipe") + else: + fp8_scale = None + fp8_dtype = None + scale_hidden_dim = None + + output, permuted_scale, permuted_probs = triton_permutation.permute_with_mask_map( + inp, + row_id_map, + probs, + fp8_scale, + pad_offsets, + num_tokens, + num_experts, + num_out_tokens, + hidden_size, + scale_hidden_dim, + ) - output, permuted_scale, permuted_probs = triton_permutation.permute_with_mask_map( - inp, - row_id_map, - probs, - fp8_scale, - pad_offsets, - num_tokens, - num_experts, - num_out_tokens, - hidden_size, - scale_hidden_dim, - ) + if fp8: + if per_tensor_recipe: + output = Float8Tensor( + data=output, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=output.shape, + dtype=fake_dtype, + ) + elif blockwise_recipe: + output = Float8BlockwiseQTensor( + shape=output.shape, + dtype=fake_dtype, + rowwise_data=output, + rowwise_scale_inv=permuted_scale.T.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=fp8_dtype, + quantizer=None, + is_2D_scaled=False, + requires_grad=output.requires_grad, + ) + elif mxfp8_recipe: + output = MXFP8Tensor( + shape=output.shape, + dtype=fake_dtype, + fp8_dtype=fp8_dtype, + rowwise_data=output, + rowwise_scale_inv=permuted_scale.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=None, + requires_grad=output.requires_grad, + with_gemm_swizzled_scales=False, + ) - if fp8: - if per_tensor_recipe: - output = Float8Tensor( - data=output, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=output.shape, - dtype=fake_dtype, - ) - elif blockwise_recipe: - output = Float8BlockwiseQTensor( - shape=output.shape, - dtype=fake_dtype, - rowwise_data=output, - rowwise_scale_inv=permuted_scale.T.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - fp8_dtype=fp8_dtype, - quantizer=None, - is_2D_scaled=False, - requires_grad=output.requires_grad, - ) - elif mxfp8_recipe: - output = MXFP8Tensor( - shape=output.shape, - dtype=fake_dtype, - fp8_dtype=fp8_dtype, - rowwise_data=output, - rowwise_scale_inv=permuted_scale.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - quantizer=None, - requires_grad=output.requires_grad, - with_gemm_swizzled_scales=False, - ) + # If permuted_probs is None, return empty tensor (custom ops need concrete tensors) + if permuted_probs is None: + permuted_probs = torch.empty(0, device=inp.device) - ctx.save_for_backward(row_id_map, pad_offsets) - ctx.num_experts = num_experts - ctx.num_tokens = num_tokens - ctx.hidden_size = hidden_size - return output, row_id_map, permuted_probs - - @staticmethod - def backward( - ctx, - permuted_act_grad: torch.Tensor, - _, - permuted_probs_grad: torch.Tensor, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, ctx.probs, None - - act_grad = None - probs_grad = None - if ctx.needs_input_grad[0]: - row_id_map, pad_offsets = ctx.saved_tensors - if isinstance(permuted_act_grad, QuantizedTensor): - raise TypeError( - "The backward of moe_permute does not support FP8, but got " - f"QuantizedTensor of type {type(permuted_act_grad).__name__}." - ) - act_grad, probs_grad = triton_permutation.unpermute_with_mask_map( - permuted_act_grad, - row_id_map, - None, - permuted_probs_grad, - pad_offsets, - ctx.num_tokens, - ctx.num_experts, - ctx.hidden_size, + return output, row_id_map, permuted_probs + + +@moe_permute_mask_map_forward.register_fake +def _moe_permute_mask_map_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + routing_map: torch.Tensor, + num_out_tokens: int, + probs: Optional[torch.Tensor], + pad_offsets: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fake implementation for shape inference.""" + num_tokens = inp.shape[0] + hidden_size = inp.shape[1] + num_experts = routing_map.shape[1] + # row_id_map: (num_tokens, num_experts * 2 + 1) + fake_output = torch.empty((num_out_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + fake_row_id_map = torch.empty( + (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=inp.device + ) + if probs is not None: + fake_permuted_probs = torch.empty((num_out_tokens,), dtype=probs.dtype, device=inp.device) + else: + fake_permuted_probs = torch.empty(0, device=inp.device) + return fake_output, fake_row_id_map, fake_permuted_probs + + +@torch.library.custom_op("te_moe::permute_mask_map_bwd", mutates_args=[]) +def moe_permute_mask_map_backward( + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE permute with mask router map.""" + act_grad, probs_grad = triton_permutation.unpermute_with_mask_map( + permuted_act_grad, + row_id_map, + None, + permuted_probs_grad, + pad_offsets, + num_tokens, + num_experts, + hidden_size, + ) + if probs_grad is None: + probs_grad = torch.empty(0, device=permuted_act_grad.device) + return act_grad, probs_grad + + +@moe_permute_mask_map_backward.register_fake +def _moe_permute_mask_map_backward_fake( # pylint: disable=unused-argument + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake for backward shape inference.""" + act_grad = torch.empty( + (num_tokens, hidden_size), dtype=permuted_act_grad.dtype, device=permuted_act_grad.device + ) + if permuted_probs_grad is not None: + probs_grad = torch.empty( + (num_tokens, num_experts), + dtype=permuted_probs_grad.dtype, + device=permuted_act_grad.device, + ) + else: + probs_grad = torch.empty(0, device=permuted_act_grad.device) + return act_grad, probs_grad + + +def _moe_permute_mask_map_setup_context(ctx, inputs, output): + """Save context for backward pass.""" + inp, routing_map, _num_out_tokens, probs, pad_offsets = inputs + _output_tensor, row_id_map, _permuted_probs = output + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(row_id_map, pad_offsets) + ctx.num_experts = routing_map.size(1) + ctx.num_tokens = inp.size(0) + ctx.hidden_size = inp.size(1) if not ctx.empty_input else 0 + ctx.needs_probs_grad = probs is not None and probs.requires_grad + + +def _moe_permute_mask_map_backward_wrapper( + ctx, grad_output, grad_row_id_map, grad_permuted_probs +): # pylint: disable=unused-argument + """Backward wrapper calling the custom backward op.""" + if ctx.empty_input: + if ctx.needs_probs_grad: + probs_grad = torch.zeros( + (ctx.num_tokens, ctx.num_experts), + dtype=grad_permuted_probs.dtype, + device=grad_permuted_probs.device, ) - if not ctx.needs_input_grad[3]: + else: probs_grad = None - return act_grad, None, None, probs_grad, None - - -class _moe_unpermute_mask_map(torch.autograd.Function): - """functional Unpermute with mask router map""" - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - row_id_map: torch.Tensor, - merging_probs: Optional[torch.Tensor], - restore_shape: Optional[torch.Size], - pad_offsets: Optional[torch.Tensor], - ) -> torch.Tensor: - # pylint: disable=missing-function-docstring - if not inp.numel(): - ctx.merging_probs = merging_probs - return inp + return grad_output, None, None, probs_grad, None - if restore_shape is None: - restore_shape = inp.shape - num_tokens, hidden_size = restore_shape - num_experts = (row_id_map.size(1) - 1) // 2 + assert not isinstance( + grad_output, QuantizedTensor + ), "The backward of moe_permute does not support FP8." + + row_id_map, pad_offsets = ctx.saved_tensors + + # Pass permuted_probs_grad only if it has content + probs_grad_input = grad_permuted_probs if grad_permuted_probs.numel() > 0 else None + + act_grad, probs_grad = torch.ops.te_moe.permute_mask_map_bwd( + grad_output, + probs_grad_input, + row_id_map, + pad_offsets, + ctx.num_tokens, + ctx.num_experts, + ctx.hidden_size, + ) + + if not ctx.needs_probs_grad or probs_grad.numel() == 0: + probs_grad = None + + return act_grad, None, None, probs_grad, None - with_probs = merging_probs is not None - if with_probs: - if not merging_probs.is_cuda: + +moe_permute_mask_map_forward.register_autograd( + _moe_permute_mask_map_backward_wrapper, + setup_context=_moe_permute_mask_map_setup_context, +) + + +# ===================== _moe_unpermute_mask_map custom ops ===================== + + +@torch.library.custom_op("te_moe::unpermute_mask_map_fwd", mutates_args=[]) +def moe_unpermute_mask_map_forward( + inp: torch.Tensor, + row_id_map: torch.Tensor, + merging_probs: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, + pad_offsets: Optional[torch.Tensor], +) -> torch.Tensor: + """Forward pass for MoE unpermute with mask router map.""" + if not inp.numel(): + return inp.clone() + assert not isinstance( + inp, QuantizedTensor + ), "The forward of moe_unpermute does not support FP8." + unpermuted_output, _ = triton_permutation.unpermute_with_mask_map( + inp, + row_id_map, + merging_probs, + None, + pad_offsets, + num_tokens, + num_experts, + hidden_size, + ) + return unpermuted_output + + +@moe_unpermute_mask_map_forward.register_fake +def _moe_unpermute_mask_map_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + row_id_map: torch.Tensor, + merging_probs: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + hidden_size: int, + pad_offsets: Optional[torch.Tensor], +) -> torch.Tensor: + """Fake implementation for shape inference.""" + return torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + + +@torch.library.custom_op("te_moe::unpermute_mask_map_bwd_with_probs", mutates_args=[]) +def moe_unpermute_mask_map_backward_with_probs( + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + fwd_input: torch.Tensor, + merging_probs: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE unpermute with merging probs.""" + act_grad, probs_grad = triton_permutation.unpermute_with_mask_map_bwd_with_merging_probs( + unpermuted_act_grad, + row_id_map, + fwd_input, + merging_probs, + pad_offsets, + num_tokens, + num_experts, + num_permuted_tokens, + hidden_size, + ) + return act_grad, probs_grad + + +@moe_unpermute_mask_map_backward_with_probs.register_fake +def _moe_unpermute_mask_map_bwd_with_probs_fake( # pylint: disable=unused-argument + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + fwd_input: torch.Tensor, + merging_probs: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake for backward shape inference with merging probs.""" + act_grad = torch.empty( + (num_permuted_tokens, hidden_size), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + probs_grad = torch.empty( + (num_tokens, num_experts), + dtype=merging_probs.dtype, + device=unpermuted_act_grad.device, + ) + return act_grad, probs_grad + + +@torch.library.custom_op("te_moe::unpermute_mask_map_bwd_no_probs", mutates_args=[]) +def moe_unpermute_mask_map_backward_no_probs( + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> torch.Tensor: + """Backward pass for MoE unpermute without merging probs (permute grad back).""" + # FP8 handling + fp8 = isinstance(unpermuted_act_grad, QuantizedTensor) + per_tensor_recipe = isinstance(unpermuted_act_grad, Float8Tensor) + blockwise_recipe = isinstance(unpermuted_act_grad, Float8BlockwiseQTensor) + mxfp8_recipe = isinstance(unpermuted_act_grad, MXFP8Tensor) + + if fp8: + fp8_dtype = unpermuted_act_grad._fp8_dtype + fake_dtype = unpermuted_act_grad.dtype + if per_tensor_recipe: + fp8_scale = None + scale_hidden_dim = None + fp8_scale_inv = unpermuted_act_grad._scale_inv + unpermuted_act_grad = unpermuted_act_grad._data + # blockwise scaling + elif blockwise_recipe: + fp8_scale = unpermuted_act_grad._rowwise_scale_inv.T.contiguous() + unpermuted_act_grad = unpermuted_act_grad._rowwise_data + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: raise ValueError( - "merging_probs must be a CUDA tensor, but got tensor on " - f"{merging_probs.device}." + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Scale shape: {tuple(fp8_scale.shape)}." ) - - # Device check - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not row_id_map.is_cuda: - raise ValueError( - f"row_id_map must be a CUDA tensor, but got tensor on {row_id_map.device}." - ) - if pad_offsets is not None: - if not pad_offsets.is_cuda: + # mxfp8 scaling + elif mxfp8_recipe: + fp8_scale = unpermuted_act_grad._rowwise_scale_inv.contiguous() + unpermuted_act_grad = unpermuted_act_grad._rowwise_data + scale_hidden_dim = fp8_scale.shape[1] + if num_tokens != fp8_scale.shape[0]: raise ValueError( - f"pad_offsets must be a CUDA tensor, but got tensor on {pad_offsets.device}." + f"Scale and input shape mismatch: num_tokens ({num_tokens}) != " + f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " + f"Scale shape: {tuple(fp8_scale.shape)}." ) + else: + raise ValueError("Unsupported FP8 recipe") + else: + scale_hidden_dim = None + fp8_dtype = None + fp8_scale = None + + act_grad, permuted_scale, _ = triton_permutation.permute_with_mask_map( + unpermuted_act_grad, + row_id_map, + None, + fp8_scale, + pad_offsets, + num_tokens, + num_experts, + num_permuted_tokens, + hidden_size, + scale_hidden_dim, + ) - if isinstance(inp, QuantizedTensor): - raise TypeError( - "The forward of moe_unpermute does not support FP8, but got " - f"QuantizedTensor of type {type(inp).__name__}." + if fp8: + if per_tensor_recipe: + act_grad = Float8Tensor( + data=act_grad, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=act_grad.shape, + dtype=fake_dtype, ) - unpermuted_output, _ = triton_permutation.unpermute_with_mask_map( - inp, + elif blockwise_recipe: + act_grad = Float8BlockwiseQTensor( + shape=act_grad.shape, + dtype=fake_dtype, + rowwise_data=act_grad, + rowwise_scale_inv=permuted_scale.T.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + fp8_dtype=fp8_dtype, + quantizer=None, + is_2D_scaled=False, + requires_grad=act_grad.requires_grad, + ) + elif mxfp8_recipe: + act_grad = MXFP8Tensor( + shape=act_grad.shape, + dtype=fake_dtype, + fp8_dtype=fp8_dtype, + rowwise_data=act_grad, + rowwise_scale_inv=permuted_scale.contiguous(), + columnwise_data=None, + columnwise_scale_inv=None, + quantizer=None, + requires_grad=act_grad.requires_grad, + with_gemm_swizzled_scales=False, + ) + + return act_grad + + +@moe_unpermute_mask_map_backward_no_probs.register_fake +def _moe_unpermute_mask_map_bwd_no_probs_fake( # pylint: disable=unused-argument + unpermuted_act_grad: torch.Tensor, + row_id_map: torch.Tensor, + pad_offsets: Optional[torch.Tensor], + num_tokens: int, + num_experts: int, + num_permuted_tokens: int, + hidden_size: int, +) -> torch.Tensor: + """Fake for backward shape inference without probs.""" + return torch.empty( + (num_permuted_tokens, hidden_size), + dtype=unpermuted_act_grad.dtype, + device=unpermuted_act_grad.device, + ) + + +def _moe_unpermute_mask_map_setup_context(ctx, inputs, output): # pylint: disable=unused-argument + """Save context for backward pass.""" + inp, row_id_map, merging_probs, num_tokens, num_experts, hidden_size, pad_offsets = inputs + ctx.empty_input = inp.size(0) == 0 + ctx.num_experts = num_experts + ctx.num_tokens = num_tokens + ctx.num_permuted_tokens = inp.size(0) + ctx.hidden_size = hidden_size + ctx.with_probs = merging_probs is not None + if ctx.with_probs: + ctx.save_for_backward(inp, row_id_map, merging_probs, pad_offsets) + ctx.needs_probs_grad = merging_probs.requires_grad + else: + ctx.save_for_backward(row_id_map, pad_offsets) + ctx.needs_probs_grad = False + + +def _moe_unpermute_mask_map_backward_wrapper(ctx, unpermuted_act_grad): + """Backward wrapper calling the appropriate custom backward op.""" + if ctx.empty_input: + if ctx.with_probs: + _, _, merging_probs, _ = ctx.saved_tensors + probs_grad = torch.zeros_like(merging_probs) if ctx.needs_probs_grad else None + return unpermuted_act_grad, None, probs_grad, None, None, None, None + return unpermuted_act_grad, None, None, None, None, None, None + + act_grad = None + probs_grad = None + + if ctx.with_probs: + fwd_input, row_id_map, merging_probs, pad_offsets = ctx.saved_tensors + assert not isinstance( + unpermuted_act_grad, QuantizedTensor + ), "The backward of moe_unpermute with merging probs does not support FP8." + act_grad, probs_grad = torch.ops.te_moe.unpermute_mask_map_bwd_with_probs( + unpermuted_act_grad, row_id_map, + fwd_input, merging_probs, - None, pad_offsets, - num_tokens, - num_experts, - hidden_size, + ctx.num_tokens, + ctx.num_experts, + ctx.num_permuted_tokens, + ctx.hidden_size, + ) + else: + row_id_map, pad_offsets = ctx.saved_tensors + act_grad = torch.ops.te_moe.unpermute_mask_map_bwd_no_probs( + unpermuted_act_grad, + row_id_map, + pad_offsets, + ctx.num_tokens, + ctx.num_experts, + ctx.num_permuted_tokens, + ctx.hidden_size, ) - if with_probs: - ctx.save_for_backward(inp, row_id_map, merging_probs, pad_offsets) - else: - ctx.save_for_backward(row_id_map, pad_offsets) - ctx.num_experts = num_experts - ctx.num_tokens = num_tokens - ctx.num_permuted_tokens = inp.size(0) - ctx.hidden_size = hidden_size - ctx.with_probs = with_probs - return unpermuted_output - - @staticmethod - def backward(ctx, unpermuted_act_grad): - # pylint: disable=missing-function-docstring - if not unpermuted_act_grad.numel(): - return unpermuted_act_grad, None, ctx.merging_probs, None, None - - act_grad = None + if not ctx.needs_probs_grad: probs_grad = None - if ctx.needs_input_grad[0]: - if ctx.with_probs: - fwd_input, row_id_map, merging_probs, pad_offsets = ctx.saved_tensors - else: - row_id_map, pad_offsets = ctx.saved_tensors - - fp8 = isinstance(unpermuted_act_grad, QuantizedTensor) - per_tensor_recipe = isinstance(unpermuted_act_grad, Float8Tensor) - blockwise_recipe = isinstance(unpermuted_act_grad, Float8BlockwiseQTensor) - mxfp8_recipe = isinstance(unpermuted_act_grad, MXFP8Tensor) - - if fp8: - fp8_dtype = unpermuted_act_grad._fp8_dtype - fake_dtype = unpermuted_act_grad.dtype - # per-tensor scaling - if per_tensor_recipe: - # Kernel does not need scale in per-tensor scaling - fp8_scale = None - scale_hidden_dim = None - fp8_scale_inv = unpermuted_act_grad._scale_inv - unpermuted_act_grad = unpermuted_act_grad._data - # blockwise scaling - elif blockwise_recipe: - fp8_scale = unpermuted_act_grad._rowwise_scale_inv.T.contiguous() - unpermuted_act_grad = unpermuted_act_grad._rowwise_data - scale_hidden_dim = fp8_scale.shape[1] - if ctx.num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({ctx.num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Scale shape: {tuple(fp8_scale.shape)}." - ) - # mxfp8 scaling - elif mxfp8_recipe: - fp8_scale = unpermuted_act_grad._rowwise_scale_inv.contiguous() - unpermuted_act_grad = unpermuted_act_grad._rowwise_data - scale_hidden_dim = fp8_scale.shape[1] - if ctx.num_tokens != fp8_scale.shape[0]: - raise ValueError( - f"Scale and input shape mismatch: num_tokens ({ctx.num_tokens}) != " - f"fp8_scale.shape[0] ({fp8_scale.shape[0]}). " - f"Scale shape: {tuple(fp8_scale.shape)}." - ) - else: - raise ValueError("Unsupported FP8 recipe") - else: - scale_hidden_dim = None - fp8_dtype = None - fp8_scale = None - - permuted_scale = None - if ctx.with_probs: - if fp8: - raise TypeError( - "The backward of moe_unpermute with merging probs does not support FP8, " - f"but got FP8 gradient with dtype {fp8_dtype}." - ) - act_grad, probs_grad = ( - triton_permutation.unpermute_with_mask_map_bwd_with_merging_probs( - unpermuted_act_grad, - row_id_map, - fwd_input, - merging_probs, - pad_offsets, - ctx.num_tokens, - ctx.num_experts, - ctx.num_permuted_tokens, - ctx.hidden_size, - ) - ) - else: - act_grad, permuted_scale, _ = triton_permutation.permute_with_mask_map( - unpermuted_act_grad, - row_id_map, - None, - fp8_scale, - pad_offsets, - ctx.num_tokens, - ctx.num_experts, - ctx.num_permuted_tokens, - ctx.hidden_size, - scale_hidden_dim, - ) - if fp8: - if per_tensor_recipe: - act_grad = Float8Tensor( - data=act_grad, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=act_grad.shape, - dtype=fake_dtype, - ) - elif blockwise_recipe: - act_grad = Float8BlockwiseQTensor( - shape=act_grad.shape, - dtype=fake_dtype, - rowwise_data=act_grad, - rowwise_scale_inv=permuted_scale.T.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - fp8_dtype=fp8_dtype, - quantizer=None, - is_2D_scaled=False, - requires_grad=act_grad.requires_grad, - ) - elif mxfp8_recipe: - act_grad = MXFP8Tensor( - shape=act_grad.shape, - dtype=fake_dtype, - fp8_dtype=fp8_dtype, - rowwise_data=act_grad, - rowwise_scale_inv=permuted_scale.contiguous(), - columnwise_data=None, - columnwise_scale_inv=None, - quantizer=None, - requires_grad=act_grad.requires_grad, - with_gemm_swizzled_scales=False, - ) - - if not ctx.needs_input_grad[2]: - probs_grad = None - return act_grad, None, probs_grad, None, None + return act_grad, None, probs_grad, None, None, None, None + + +moe_unpermute_mask_map_forward.register_autograd( + _moe_unpermute_mask_map_backward_wrapper, + setup_context=_moe_unpermute_mask_map_setup_context, +) + +# Register all te_moe custom ops as passthrough in QuantizedTensor.__torch_dispatch__ +# so that FP8 tensors are not unwrapped before entering these ops. +_quantized_tensor_passthrough_ops.update( + { + torch.ops.te_moe.permute_mask_map_fwd.default, + torch.ops.te_moe.permute_mask_map_bwd.default, + torch.ops.te_moe.unpermute_mask_map_fwd.default, + torch.ops.te_moe.unpermute_mask_map_bwd_with_probs.default, + torch.ops.te_moe.unpermute_mask_map_bwd_no_probs.default, + } +) def moe_permute( @@ -609,10 +883,15 @@ def moe_permute( Options are: 'mask', 'index'. Refer to `routing_map` for more details. """ + if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): + raise RuntimeError( + "moe_permute with quantized (FP8) input is not supported under torch.compile. " + "Please move quantization outside the compiled region." + ) if map_type == "index": - return _moe_permute_index_map.apply(inp, routing_map, num_out_tokens, max_token_num) + return torch.ops.te_moe.permute_index_map(inp, routing_map, num_out_tokens, max_token_num) if map_type == "mask": - output, row_id_map, _ = _moe_permute_mask_map.apply( + output, row_id_map, _ = torch.ops.te_moe.permute_mask_map_fwd( inp, routing_map, num_out_tokens, None, None ) return output, row_id_map @@ -646,7 +925,12 @@ def moe_permute_with_probs( The effective output token count, representing the number of tokens not dropped. By default, set to '-1', meaning no tokens are dropped. """ - output, row_id_map, permuted_probs = _moe_permute_mask_map.apply( + if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): + raise RuntimeError( + "moe_permute_with_probs with quantized (FP8) input is not supported under " + "torch.compile. Please move quantization outside the compiled region." + ) + output, row_id_map, permuted_probs = torch.ops.te_moe.permute_mask_map_fwd( inp, routing_map, num_out_tokens, probs, None ) return output, permuted_probs, row_id_map @@ -681,6 +965,11 @@ def moe_permute_and_pad_with_probs( align_size : int the alignment size for the input tensor. """ + if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): + raise RuntimeError( + "moe_permute_and_pad_with_probs with quantized (FP8) input is not supported under " + "torch.compile. Please move quantization outside the compiled region." + ) if tokens_per_expert is None: raise ValueError( "tokens_per_expert must be provided to the fused permute padding function." @@ -704,7 +993,7 @@ def moe_permute_and_pad_with_probs( [torch.zeros(1, dtype=cum_pad.dtype, device=inp.device), cum_pad[:-1]] ) - output, row_id_map, permuted_probs = _moe_permute_mask_map.apply( + output, row_id_map, permuted_probs = torch.ops.te_moe.permute_mask_map_fwd( inp, routing_map, target_tokens_per_expert.sum().item(), probs, pad_offsets ) return output, permuted_probs, row_id_map, pad_offsets, target_tokens_per_expert @@ -754,125 +1043,228 @@ def moe_unpermute( warnings.warn("probs kwarg is deprecated. Use merging_probs kwarg instead.") merging_probs = probs if map_type == "index": - return _moe_unpermute_index_map.apply(inp, row_id_map, merging_probs) + # Normalize probs + if merging_probs is not None: + if merging_probs.dtype != torch.float32: + warnings.warn( + f"The data type of the input `probs` of Unpermute is {merging_probs.dtype}! " + "The recommended type is torch.float32." + ) + merging_probs = merging_probs.to(torch.float32) + num_tokens = merging_probs.size(0) + topK = merging_probs.size(1) + else: + num_tokens = row_id_map.size(0) + topK = 1 + merging_probs = torch.empty(0, device=inp.device) + + return torch.ops.te_moe.unpermute_index_map_fwd( + inp, row_id_map, merging_probs, num_tokens, topK + ) if map_type == "mask": - return _moe_unpermute_mask_map.apply( - inp, row_id_map, merging_probs, restore_shape, pad_offsets + if restore_shape is None: + restore_shape = inp.shape + num_tokens, hidden_size = restore_shape + num_experts = (row_id_map.size(1) - 1) // 2 if row_id_map.dim() > 1 else 0 + + return torch.ops.te_moe.unpermute_mask_map_fwd( + inp, + row_id_map, + merging_probs, + num_tokens, + num_experts, + hidden_size, + pad_offsets, ) raise ValueError("map_type should be one of 'mask' or 'index'") -class _moe_chunk_sort(torch.autograd.Function): - """functional MoE chunk permute""" - - @staticmethod - def forward( - ctx, - inp: torch.Tensor, - split_sizes: torch.Tensor, - sorted_idxs: torch.Tensor, - probs: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor]: - # pylint: disable=missing-function-docstring - if not inp.numel(): - return inp, probs - - if not inp.is_cuda: - raise ValueError(f"inp must be a CUDA tensor, but got tensor on {inp.device}.") - if not split_sizes.is_cuda: - raise ValueError( - f"split_sizes must be a CUDA tensor, but got tensor on {split_sizes.device}." - ) - if not sorted_idxs.is_cuda: - raise ValueError( - f"sorted_idxs must be a CUDA tensor, but got tensor on {sorted_idxs.device}." - ) - if probs is not None: - if not probs.is_cuda: - raise ValueError(f"probs must be a CUDA tensor, but got tensor on {probs.device}.") +# ===================== _moe_chunk_sort custom ops ===================== - num_tokens, hidden_size = inp.shape - num_splits = split_sizes.size(0) - if num_splits != sorted_idxs.size(0): - raise ValueError( - f"split_sizes.size(0) ({num_splits}) must match " - f"sorted_idxs.size(0) ({sorted_idxs.size(0)})." - ) - fp8 = isinstance(inp, Float8Tensor) - if fp8: - fp8_dtype = inp._fp8_dtype - fp8_scale_inv = inp._scale_inv - fake_dtype = inp.dtype - inp = inp._data +@torch.library.custom_op("te_moe::chunk_sort_fwd", mutates_args=[]) +def moe_chunk_sort_forward( + inp: torch.Tensor, + split_sizes: torch.Tensor, + sorted_idxs: torch.Tensor, + probs: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Forward pass for MoE chunk sort. Returns (output, permuted_probs, row_id_map).""" + if not inp.numel(): + probs_out = probs.clone() if probs is not None else torch.empty(0, device=inp.device) + return inp.clone(), probs_out, torch.empty(0, device=inp.device, dtype=torch.int32) + + num_tokens, hidden_size = inp.shape + num_splits = split_sizes.size(0) + + fp8 = isinstance(inp, Float8Tensor) + if fp8: + fp8_dtype = inp._fp8_dtype + fp8_scale_inv = inp._scale_inv + fake_dtype = inp.dtype + inp = inp._data + + row_id_map = triton_permutation.make_chunk_sort_map( + split_sizes, + sorted_idxs, + num_tokens, + num_splits, + ) + output, permuted_probs = triton_permutation.sort_chunks_by_map( + inp, + row_id_map, + probs, + num_tokens, + hidden_size, + is_forward=True, + ) + if fp8: + output = Float8Tensor( + data=output, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=output.shape, + dtype=fake_dtype, + ) - row_id_map = triton_permutation.make_chunk_sort_map( - split_sizes, - sorted_idxs, - num_tokens, - num_splits, + if permuted_probs is None: + permuted_probs = torch.empty(0, device=output.device) + + return output, permuted_probs, row_id_map + + +@moe_chunk_sort_forward.register_fake +def _moe_chunk_sort_forward_fake( # pylint: disable=unused-argument + inp: torch.Tensor, + split_sizes: torch.Tensor, + sorted_idxs: torch.Tensor, + probs: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fake for shape inference.""" + num_tokens = inp.shape[0] + hidden_size = inp.shape[1] + fake_output = torch.empty((num_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + if probs is not None: + fake_probs = torch.empty((num_tokens,), dtype=probs.dtype, device=inp.device) + else: + fake_probs = torch.empty(0, device=inp.device) + # row_id_map: 1D, size num_tokens + fake_row_id_map = torch.empty((num_tokens,), dtype=torch.int32, device=inp.device) + return fake_output, fake_probs, fake_row_id_map + + +@torch.library.custom_op("te_moe::chunk_sort_bwd", mutates_args=[]) +def moe_chunk_sort_backward( + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + num_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Backward pass for MoE chunk sort.""" + fp8 = isinstance(permuted_act_grad, Float8Tensor) + if fp8: + fp8_dtype = permuted_act_grad._fp8_dtype + fp8_scale_inv = permuted_act_grad._scale_inv + fake_dtype = permuted_act_grad.dtype + permuted_act_grad = permuted_act_grad._data + + act_grad, probs_grad = triton_permutation.sort_chunks_by_map( + permuted_act_grad, + row_id_map, + permuted_probs_grad, + num_tokens, + hidden_size, + is_forward=False, + ) + + if fp8: + act_grad = Float8Tensor( + data=act_grad, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + shape=act_grad.shape, + dtype=fake_dtype, ) - output, permuted_probs = triton_permutation.sort_chunks_by_map( - inp, - row_id_map, - probs, - num_tokens, - hidden_size, - is_forward=True, + + if probs_grad is None: + probs_grad = torch.empty(0, device=act_grad.device) + + return act_grad, probs_grad + + +@moe_chunk_sort_backward.register_fake +def _moe_chunk_sort_backward_fake( # pylint: disable=unused-argument + permuted_act_grad: torch.Tensor, + permuted_probs_grad: Optional[torch.Tensor], + row_id_map: torch.Tensor, + num_tokens: int, + hidden_size: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Fake for backward shape inference.""" + fake_act_grad = torch.empty( + (num_tokens, hidden_size), + dtype=permuted_act_grad.dtype, + device=permuted_act_grad.device, + ) + if permuted_probs_grad is not None: + fake_probs_grad = torch.empty( + (num_tokens,), + dtype=permuted_probs_grad.dtype, + device=permuted_act_grad.device, ) - if fp8: - output = Float8Tensor( - data=output, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=output.shape, - dtype=fake_dtype, - ) + else: + fake_probs_grad = torch.empty(0, device=permuted_act_grad.device) + return fake_act_grad, fake_probs_grad + + +def _moe_chunk_sort_setup_context(ctx, inputs, output): + """Save context for backward pass.""" + inp, _split_sizes, _sorted_idxs, probs = inputs + _output_tensor, _permuted_probs, row_id_map = output + ctx.empty_input = inp.size(0) == 0 + ctx.save_for_backward(row_id_map) + ctx.num_tokens = inp.size(0) + ctx.hidden_size = inp.size(1) if not ctx.empty_input else 0 + ctx.needs_probs_grad = probs is not None and probs.requires_grad - ctx.save_for_backward(row_id_map) - ctx.num_tokens = num_tokens - ctx.hidden_size = hidden_size - return output, permuted_probs - - @staticmethod - def backward( - ctx, - permuted_act_grad: torch.Tensor, - permuted_probs_grad: torch.Tensor, - ) -> Tuple[torch.Tensor, ...]: - # pylint: disable=missing-function-docstring - if not permuted_act_grad.numel(): - return permuted_act_grad, None, None, permuted_probs_grad - - act_grad = None + +def _moe_chunk_sort_backward_wrapper(ctx, permuted_act_grad, permuted_probs_grad, _row_id_map_grad): + """Backward wrapper calling the custom backward op.""" + if ctx.empty_input: + probs_grad = permuted_probs_grad if ctx.needs_probs_grad else None + return permuted_act_grad, None, None, probs_grad + + (row_id_map,) = ctx.saved_tensors + + probs_grad_input = permuted_probs_grad if permuted_probs_grad.numel() > 0 else None + + act_grad, probs_grad = torch.ops.te_moe.chunk_sort_bwd( + permuted_act_grad, + probs_grad_input, + row_id_map, + ctx.num_tokens, + ctx.hidden_size, + ) + + if not ctx.needs_probs_grad or probs_grad.numel() == 0: probs_grad = None - if ctx.needs_input_grad[0]: - (row_id_map,) = ctx.saved_tensors - fp8 = isinstance(permuted_act_grad, Float8Tensor) - if fp8: - fp8_dtype = permuted_act_grad._fp8_dtype - fp8_scale_inv = permuted_act_grad._scale_inv - fake_dtype = permuted_act_grad.dtype - permuted_act_grad = permuted_act_grad._data - act_grad, probs_grad = triton_permutation.sort_chunks_by_map( - permuted_act_grad, - row_id_map, - permuted_probs_grad, - ctx.num_tokens, - ctx.hidden_size, - is_forward=False, - ) - if fp8: - act_grad = Float8Tensor( - data=act_grad, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - shape=act_grad.shape, - dtype=fake_dtype, - ) - if not ctx.needs_input_grad[3]: - probs_grad = None - return act_grad, None, None, probs_grad + + return act_grad, None, None, probs_grad + + +moe_chunk_sort_forward.register_autograd( + _moe_chunk_sort_backward_wrapper, + setup_context=_moe_chunk_sort_setup_context, +) + +# Register chunk sort ops as passthrough in QuantizedTensor.__torch_dispatch__ +_quantized_tensor_passthrough_ops.update( + { + torch.ops.te_moe.chunk_sort_fwd.default, + torch.ops.te_moe.chunk_sort_bwd.default, + } +) def moe_sort_chunks_by_index( @@ -894,7 +1286,7 @@ def moe_sort_chunks_by_index( sorted_indices : torch.Tensor Chunk indices used to permute the chunks. """ - output, _ = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, None) + output, _, _ = torch.ops.te_moe.chunk_sort_fwd(inp, split_sizes, sorted_index, None) return output @@ -922,5 +1314,7 @@ def moe_sort_chunks_by_index_with_probs( sorted_indices : torch.Tensor Chunk indices used to permute the chunks. """ - output, permuted_probs = _moe_chunk_sort.apply(inp, split_sizes, sorted_index, probs) + output, permuted_probs, _ = torch.ops.te_moe.chunk_sort_fwd( + inp, split_sizes, sorted_index, probs + ) return output, permuted_probs diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 807671e863..e40f42edd3 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -21,6 +21,12 @@ ) +# Custom ops that should pass through __torch_dispatch__ without unwrapping +# QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that +# handle quantized tensors internally. +_quantized_tensor_passthrough_ops: set = set() + + class QuantizedTensorStorage: r"""Base class for all TensorStorage classes. @@ -614,6 +620,12 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): return func(t) return False # Or error out? + # Pass through registered custom ops without unwrapping + if func in _quantized_tensor_passthrough_ops: + if kwargs is None: + kwargs = {} + return super().__torch_dispatch__(func, types, args, kwargs) + def maybe_unwrap(arg): if isinstance(arg, QuantizedTensor): return arg.dequantize() From 15760a5dd9006deac5edd1433e6d2bbf27c0d3cc Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Thu, 19 Mar 2026 00:11:39 -0700 Subject: [PATCH 290/521] [PyTorch] Add an API restore from function context to ensure tensors are detached (#2772) [PyTorch] Change the restore tensor API to ensure tensors are detached from ctx Signed-off-by: Kaining Zhong Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/attention/test_attention.py | 7 ++---- transformer_engine/pytorch/__init__.py | 1 + .../dot_product_attention/backends.py | 4 ++-- .../dot_product_attention/context_parallel.py | 6 ++--- .../pytorch/module/grouped_linear.py | 4 ++-- .../pytorch/module/layernorm_linear.py | 9 ++------ .../pytorch/module/layernorm_mlp.py | 8 ++----- transformer_engine/pytorch/module/linear.py | 9 ++------ transformer_engine/pytorch/ops/fuser.py | 5 ++--- .../pytorch/quantized_tensor.py | 22 ++++++++++++++++++- transformer_engine/pytorch/tensor/__init__.py | 2 ++ 11 files changed, 41 insertions(+), 36 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 60ade522e3..2eb307aa48 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -49,7 +49,7 @@ from transformer_engine.pytorch.quantized_tensor import ( Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) _current_file = pathlib.Path(__file__).resolve() @@ -2701,10 +2701,7 @@ def forward( @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: with torch.cuda.nvtx.range("_DPA"): - saved_tensors = ctx.saved_tensors - (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_saved( - ctx.tensor_objects, saved_tensors - ) + (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) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index cd18ca75ad..bbc1d7fab6 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -68,6 +68,7 @@ from transformer_engine.pytorch.quantized_tensor import Quantizer from transformer_engine.pytorch.quantized_tensor import prepare_for_saving from transformer_engine.pytorch.quantized_tensor import restore_from_saved +from transformer_engine.pytorch.quantized_tensor import restore_from_func_ctx from transformer_engine.pytorch.tensor import Float8Quantizer from transformer_engine.pytorch.tensor import Float8CurrentScalingQuantizer from transformer_engine.pytorch.tensor import MXFP8Quantizer diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index a6a8b0b26a..442366035a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -32,7 +32,7 @@ from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensorStorage, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.constants import ( @@ -1477,7 +1477,7 @@ def backward(ctx, d_out, *_args): cu_seqlens_q_padded, cu_seqlens_kv_padded, *other_tensors, - ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + ) = restore_from_func_ctx(ctx) aux_ctx_tensors = other_tensors 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 10ba99595b..7d9eb0cb05 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -38,7 +38,7 @@ from transformer_engine.pytorch.quantized_tensor import ( prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) # Import attention utils @@ -2085,7 +2085,7 @@ def backward(ctx, dout, *_args): cu_seqlens_q_padded, cu_seqlens_kv_padded, *other_tensors, - ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + ) = restore_from_func_ctx(ctx) cu_seqlens_q_per_step = other_tensors[:cp_size] cu_seqlens_kv_per_step = other_tensors[cp_size : cp_size * 2] rng_states = other_tensors[cp_size * 2 : cp_size * 3] @@ -3675,7 +3675,7 @@ def backward(ctx, dout, *_args): cu_seqlens_q_padded, cu_seqlens_kv_padded, *aux_ctx_tensors, - ) = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + ) = restore_from_func_ctx(ctx) qkv_format = ctx.qkv_format qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 30c1dbf408..0adda48e36 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -49,7 +49,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ...debug.pytorch.debug_quantization import DebugQuantizer from ...debug.pytorch.debug_state import TEDebugState @@ -316,7 +316,7 @@ def forward( def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): - saved_tensors = restore_from_saved(ctx.tensor_objects, ctx.saved_tensors) + saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms inputmats = saved_tensors[:N] weights = saved_tensors[N : 2 * N] diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index d775dc3e8e..ed91bc1235 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -60,7 +60,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ...debug.pytorch.debug_state import TEDebugState from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -546,7 +546,6 @@ def backward( nvtx_label = f"{nvtx_label}.{ctx.ub_name}" with get_nvtx_range_context("_LayerNormLinear_backward"): - saved_tensors = ctx.saved_tensors ( # pylint: disable=unbalanced-tuple-unpacking inputmat, weight, @@ -556,11 +555,7 @@ def backward( ln_out, mu, rsigma, - ) = restore_from_saved(ctx.tensor_objects, saved_tensors) - - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None + ) = restore_from_func_ctx(ctx) # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ( diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 037fb6c858..cc3dcc4064 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -80,7 +80,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ..cpp_extensions import ( general_gemm, @@ -898,11 +898,7 @@ def forward( def _recompute(ctx): # pylint: disable=missing-function-docstring - saved_tensors = ctx.saved_tensors - tensors = restore_from_saved(ctx.tensor_objects, saved_tensors) - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None + tensors = restore_from_func_ctx(ctx) if ctx.checkpoint: # do recomputation from the original args diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 1e3eadc405..ea921341a4 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -61,7 +61,7 @@ QuantizedTensorStorage, Quantizer, prepare_for_saving, - restore_from_saved, + restore_from_func_ctx, ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer @@ -501,15 +501,10 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], nvtx_label = f"{nvtx_label}.{ctx.ub_name}" with get_nvtx_range_context("_Linear_backward"): - saved_tensors = ctx.saved_tensors inputmat, weight_fp8, weight, bias = ( # pylint: disable=unbalanced-tuple-unpacking - restore_from_saved(ctx.tensor_objects, saved_tensors) + restore_from_func_ctx(ctx) ) - # Delete the references to tensor objects once they've been consumed - # by the `restore_from_saved` method to construct back the actual tensors. - ctx.tensor_objects = None - # Since main_grad can be modified inplace, it should not be a part of saved_tensors main_grad = ( ctx.main_grad_func() diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 80386db2d9..76606ec799 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -12,7 +12,7 @@ import torch from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling -from ..quantized_tensor import prepare_for_saving, restore_from_saved +from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx from .op import ( BasicOperation, FusibleOperation, @@ -212,8 +212,7 @@ def backward( basic_op_ctxs = func_ctx.basic_op_ctxs # Restore saved tensors - saved_tensors = restore_from_saved(func_ctx.tensor_objects, func_ctx.saved_tensors) - func_ctx.tensor_objects = None + saved_tensors = restore_from_func_ctx(func_ctx) # Unflatten list of saved tensors for ctx in basic_op_ctxs: diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index e40f42edd3..a7722f777e 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -165,7 +165,9 @@ def restore_from_saved( list[Optional[torch.Tensor]], ] ): - """Recombine the tensor data and metadata during backward pass.""" + """Recombine the tensor data and metadata during backward pass. + Note: please use `restore_from_func_ctx` instead if you are restoring tensors from a function context to make sure tensor_objects is detached and its memory can be freed + """ tensor_objects = [] for tensor in tensors: if tensor is None or isinstance(tensor, torch.Tensor): @@ -180,6 +182,24 @@ def restore_from_saved( return tensor_objects +def restore_from_func_ctx(ctx: torch.autograd.function.FunctionCtx, return_saved_tensors=False) -> ( + list[Optional[torch.Tensor | QuantizedTensorStorage]] + | tuple[ + list[Optional[torch.Tensor | QuantizedTensorStorage]], + list[Optional[torch.Tensor]], + ] +): + """Recombine the tensor data and metadata during backward pass and delete tensor objects attached to function context.""" + if not hasattr(ctx, "tensor_objects") or ctx.tensor_objects is None: + raise AttributeError("ctx must have .tensor_objects to restore saved tensors") + out = restore_from_saved( + ctx.tensor_objects, ctx.saved_tensors, return_saved_tensors=return_saved_tensors + ) + # Delete the references to tensor objects once they've been consumed by the `restore_from_saved` method to construct back the actual tensors. + ctx.tensor_objects = None + return out + + class Quantizer(abc.ABC): """Builder class for quantized tensors. diff --git a/transformer_engine/pytorch/tensor/__init__.py b/transformer_engine/pytorch/tensor/__init__.py index 5668056700..426c656d47 100644 --- a/transformer_engine/pytorch/tensor/__init__.py +++ b/transformer_engine/pytorch/tensor/__init__.py @@ -12,6 +12,7 @@ Quantizer, prepare_for_saving, restore_from_saved, + restore_from_func_ctx, ) from .storage.float8_tensor_storage import Float8TensorStorage from .storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -46,6 +47,7 @@ "GroupedTensor", "prepare_for_saving", "restore_from_saved", + "restore_from_func_ctx", ] From b7598aa887eb7d619d64c90692980009669379bf Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 19 Mar 2026 10:17:23 -0700 Subject: [PATCH 291/521] [PyT] Install pytest in onnx L1 test as Pyt container no longer packages it (#2781) Install pytest in onnx L1 test as Pyt container no longer packages it Signed-off-by: Kshitij Janardan Lakhani --- qa/L1_pytorch_onnx_unittest/test.sh | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/qa/L1_pytorch_onnx_unittest/test.sh b/qa/L1_pytorch_onnx_unittest/test.sh index 6f9ff54e48..0edf92c475 100644 --- a/qa/L1_pytorch_onnx_unittest/test.sh +++ b/qa/L1_pytorch_onnx_unittest/test.sh @@ -2,9 +2,15 @@ # # See LICENSE for license information. +function error_exit() { + echo "Error: $1" + exit 1 +} + : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" +pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # NVTE_UnfusedDPA_Emulate_FP8=1 enables FP8 attention emulation when no native backend is available NVTE_UnfusedDPA_Emulate_FP8=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_onnx_export.xml $TE_PATH/tests/pytorch/test_onnx_export.py From f11789eb5ae859ad2b5cb97c408bb3d7d0deff1a Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 19 Mar 2026 16:16:26 -0700 Subject: [PATCH 292/521] [Core] Fix MXFP8 grouped quantize for zero-sized groups in update_tma_descriptors (#2782) * Fix zero-sized groups in update_tma_descriptors Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update test_cast_mxfp8_grouped.cu Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@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: jberchtold-nvidia <158520091+jberchtold-nvidia@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/cpp/operator/test_cast_mxfp8_grouped.cu | 1 + .../common/cast/mxfp8/group_quantize_mxfp8.cuh | 7 +++++++ 2 files changed, 8 insertions(+) diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index e469ad0845..09bd21657a 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -649,6 +649,7 @@ std::vector> input_config = { {SAME_BOTH_DIMS, 2, 256,128}, {VARYING_FIRST_DIM, 2, 512,128, 128,384}, {VARYING_FIRST_DIM, 3, 1024,144, 128,384,512}, + {VARYING_FIRST_DIM, 4, 1024,144, 128,384,0,512}, {VARYING_FIRST_DIM, 4, 1536,160, 128,384,512,512}, {VARYING_FIRST_DIM, 5, 4096,512, 128,256,384,1024,2304}, {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 129d6724ac..d0d15d8d6c 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -189,6 +189,13 @@ __global__ void update_tma_descriptors( get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + // Zero-sized groups: skip TMA descriptor update. The main kernel already returns + // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension + // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. + if (rows == 0 || cols == 0) { + return; + } + const size_t offset_elts = offsets_ptr[tensor_id]; if (leading_thread && (tensor_id < num_tensors)) { From 487d68c02516f116c91b826151791bd7941b9a01 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Sun, 22 Mar 2026 13:45:01 -0700 Subject: [PATCH 293/521] [PyT] [Common] Enable sm120 support for fused attn if cuDNN is 9.18.1+ (#2693) * Enable sm120 support for fused attn if cuDNN is 9.18.1+ Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Force intermediate tensors such as S, Sum_Exp, and Max to be BHS1 shape instead of TH1 for sm120 Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add support for sm120 correct batch, seq dims Signed-off-by: Kshitij Lakhani * Add support for sm120 BHS1 style max logit even QKV are THD to avoid incorrect max logit calculation (includes padded tokens in max calculation) Signed-off-by: Kshitij Lakhani * Disable fused and flash attn for sm120 filter:kv cache Signed-off-by: Kshitij Lakhani * For CP P2P attn, set softmax_lse_in_packed_format to False if sm120+ Signed-off-by: Kshitij Lakhani * Assert in TE if T3HD/TH3D layout is used on sm120 before cuDNN F16 sdpa arbitrary kernel call Signed-off-by: Kshitij Lakhani * Modify is_ragged_q && cudnn_runtime_version >= 90600 check to also include a check for sm120 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 * Disable fused attn for T3HD and TH3D Signed-off-by: Kshitij Lakhani * nit: Add missed sm120 guard Signed-off-by: Kshitij Lakhani * Modify sm120 condition to be very specific to sm120 and not generalized to sm120+ Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Fix missing sm120 check in fwd Signed-off-by: Kshitij Lakhani * Move the check for sm120 T3HD/TH3D to nvte_get_fused_attn_backend() instead of higher layers in TE stack Signed-off-by: Kshitij Lakhani * nit: Check for matching sm120 and not sm120+ 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> --- .../common/fused_attn/fused_attn.cpp | 17 ++++ .../fused_attn_f16_arbitrary_seqlen.cu | 79 +++++++++++-------- .../dot_product_attention/context_parallel.py | 6 +- .../attention/dot_product_attention/utils.py | 33 +++++--- .../pytorch/cpp_extensions/fused_attn.py | 19 +++-- 5 files changed, 106 insertions(+), 48 deletions(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 6a136c67e4..cba1a79dd3 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -528,6 +528,23 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( "Please upgrade your cuDNN version if possible." << std::endl; } + if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen && sm_arch_ == 120) { + if (cudnn_runtime_version < 91801) { + backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < " + "91801 is not supported. " + << " Please upgrade your cuDNN version if possible." << std::endl; + } else { + // Known missing support for T3HD/TH3D layouts on SM120 + const bool is_t3hd_or_th3d = + (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD || qkv_layout == NVTE_QKV_Layout::NVTE_TH3D); + if (is_t3hd_or_th3d) { + backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + std::cout << "Warning: Given combination of T3HD/TH3D layouts on SM120 is not supported. " + << " Please consider using other THD layouts if possible." << std::endl; + } + } + } } else { backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; } 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 eb2ebcff39..16aebda69f 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 @@ -85,6 +85,9 @@ void fused_attn_arbitrary_seqlen_fwd_impl( 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(); + const int device_id = cuda::current_device(); + 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); bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); @@ -96,11 +99,16 @@ void fused_attn_arbitrary_seqlen_fwd_impl( int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // replace batch size and maximum sequence lengths with maximum token counts - // for query and key/value so the graph is static within each quantization bucket - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + // On SM 120, cuDNN support check treats layouts with stride[0] > dim[1]*dim[2]*dim[3] + // as interleaved and rejects them. Use BHSD-like dimensions/strides with max_seqlen at plan build + // so the check passes; ragged offset still provides variable-length boundaries. + if (sm_arch_ != 120) { + // replace batch size and maximum sequence lengths with maximum token counts + // for query and key/value so the graph is static within each quantization bucket + b = max_b; + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } } const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; @@ -336,7 +344,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } std::shared_ptr Max, Sum_Exp; - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("offset_stats") @@ -353,7 +361,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_name("Sum_Exp") .set_dim({b, h, s_q, 1}) .set_data_type(fe::DataType_t::FLOAT)); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); Sum_Exp->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); } else { @@ -381,7 +389,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( if (!return_max_logit) { Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); } else { Stats->set_stride({h * s_q, s_q, 1, 1}); @@ -407,9 +415,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); auto offset_kv_tuple = is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = (is_ragged_q && cudnn_runtime_version >= 90600) - ? std::make_tuple(offset_stats) - : std::make_tuple(nullptr); + auto offset_s_tuple = + use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); @@ -443,7 +450,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( size_t seqlen_offsets_workspace_size = 0; if (is_ragged_q || is_ragged_kv) { size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; } else { seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; @@ -510,7 +517,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; } void *devOffsetsS = nullptr; - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { devOffsetsS = static_cast(devOffsets) + (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; @@ -529,7 +536,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( variant_pack[offset_k] = devOffsetsK; variant_pack[offset_v] = devOffsetsV; } - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { variant_pack[offset_stats] = devOffsetsS; } } @@ -587,6 +594,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const auto cudnn_runtime_version = cudnnGetVersion(); const int device_id = cuda::current_device(); 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); bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); @@ -598,13 +606,15 @@ void fused_attn_arbitrary_seqlen_bwd_impl( int64_t actual_b = b; if ((is_ragged_q || is_ragged_kv) && cudnn_runtime_version >= 90600) { NVTE_CHECK(is_padding, "Ragged QKV input requires padding or padding_causal mask!"); - // replace batch size and maximum sequence lengths with maximum token counts - // for query and key/value so the graph is static within each quantization bucket - b = max_b; - s_q = is_ragged_q ? max_t_q : s_q; - s_kv = is_ragged_kv ? max_t_kv : s_kv; + // On SM 120, cuDNN support check requires BHSD-like strides with max_seqlen (see fwd). + if (sm_arch_ != 120) { + // replace batch size and maximum sequence lengths with maximum token counts + // for query and key/value so the graph is static within each quantization bucket + b = max_b; + s_q = is_ragged_q ? max_t_q : s_q; + s_kv = is_ragged_kv ? max_t_kv : s_kv; + } } - // We choose between 32-bit and 64-bit offsets depending on need. // This allows us to support older cuDNN runtimes gracefully. const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; @@ -765,7 +775,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( .set_name("stats") .set_dim({b, h, s_q, 1}) .set_data_type(fe::DataType_t::FLOAT)); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("offset_stats") @@ -791,10 +801,10 @@ void fused_attn_arbitrary_seqlen_bwd_impl( .set_causal_mask_bottom_right(is_bottom_right) .set_attn_scale(attn_scale); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { sdpa_backward_options.set_max_total_seq_len_q(s_q); } - if (is_ragged_kv && cudnn_runtime_version >= 90600) { + if (is_ragged_kv && cudnn_runtime_version >= 90600 && sm_arch_ != 120) { sdpa_backward_options.set_max_total_seq_len_kv(s_kv); } @@ -914,9 +924,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( is_ragged_q ? std::make_tuple(offset_q, offset_o) : std::make_tuple(nullptr, nullptr); auto offset_kv_tuple = is_ragged_kv ? std::make_tuple(offset_k, offset_v) : std::make_tuple(nullptr, nullptr); - auto offset_s_tuple = (is_ragged_q && cudnn_runtime_version >= 90600) - ? std::make_tuple(offset_stats) - : std::make_tuple(nullptr); + auto offset_s_tuple = + use_ragged_stats ? std::make_tuple(offset_stats) : std::make_tuple(nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) : std::make_tuple(nullptr, nullptr); @@ -949,7 +958,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( size_t seqlen_offsets_workspace_size = 0; if (is_ragged_q || is_ragged_kv) { size_t count = 2 * (static_cast(is_ragged_q) + static_cast(is_ragged_kv)); - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { seqlen_offsets_workspace_size = (count + 1) * num_bytes_per_ragged_offset; } else { seqlen_offsets_workspace_size = count * num_bytes_per_ragged_offset; @@ -1019,7 +1028,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( devOffsetsV = static_cast(devOffsetsK) + num_bytes_per_ragged_offset; } void *devOffsetsS = nullptr; - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { devOffsetsS = static_cast(devOffsets) + (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; @@ -1038,7 +1047,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( variant_pack[offset_k] = devOffsetsK; variant_pack[offset_v] = devOffsetsV; } - if (is_ragged_q && cudnn_runtime_version >= 90600) { + if (use_ragged_stats) { variant_pack[offset_stats] = devOffsetsS; } } @@ -1102,6 +1111,9 @@ void fused_attn_arbitrary_seqlen_fwd( devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } + const int device_id = cuda::current_device(); + const int sm_arch_ = cuda::sm_arch(device_id); + void *devPtrCuSeqlensQ = cu_seqlens_q->data.dptr; void *devPtrCuSeqlensKV = cu_seqlens_kv->data.dptr; void *devPtrSeqOffsetsQ = cu_seqlens_q_padded->data.dptr; @@ -1128,7 +1140,8 @@ void fused_attn_arbitrary_seqlen_fwd( if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && + (sm_arch_ != 120)) { output_Max->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1136,7 +1149,8 @@ void fused_attn_arbitrary_seqlen_fwd( output_Max->data.dtype = DType::kFloat32; Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Sum_Exp->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && + (sm_arch_ != 120)) { output_Sum_Exp->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; @@ -1145,7 +1159,8 @@ void fused_attn_arbitrary_seqlen_fwd( } else { Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_S->data.dptr = nullptr; - if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && + (sm_arch_ != 120)) { output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; } else { output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; 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 7d9eb0cb05..64cccaac6e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1494,7 +1494,11 @@ def forward( softmax_lse_in_packed_format = False if qkv_format == "thd": if use_fused_attention: - softmax_lse_in_packed_format = get_cudnn_version() >= (9, 6, 0) + softmax_lse_in_packed_format = get_cudnn_version() >= ( + 9, + 6, + 0, + ) and get_device_compute_capability() != (12, 0) else: softmax_lse_in_packed_format = fa_utils.v2_6_0_plus or use_flash_attn_3 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 567fd17c34..170cb2cd34 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -554,11 +554,15 @@ def get_attention_backend( # | FP8 | non-paged/paged | sm90 | thd | >= 1 # Unfused | FP32/FP16/BF16 | non-paged/paged | all | bshd,sbhd,thd | >= 1 if inference_params is not None: - # Temporarily disabling fused attention for kv caching for sm89 irrespective of cuDNN version - # until the cuDNN bug is resolved - if device_compute_capability == (8, 9): - logger.debug("Disabling FusedAttention for KV caching for sm89") + # Temporarily disabling fused attention for kv caching for sm89/sm120 irrespective of + # cuDNN version until the cuDNN bug is resolved. + if device_compute_capability in ((8, 9), (12, 0)): + logger.debug("Disabling FusedAttention for KV caching for sm89/sm120") use_fused_attention = False + # Temporarily disable FlashAttention for KV caching on sm120 + if device_compute_capability == (12, 0): + logger.debug("Disabling FlashAttention for KV caching for sm120") + use_flash_attention = False if context_parallel: logger.debug("Disabling all backends for KV caching with context parallelism") use_flash_attention = False @@ -691,12 +695,21 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_flash_attention = False if device_compute_capability == (12, 0): - if use_fused_attention: - logger.debug( - "Disabling FusedAttention as qkv_format = thd is" - " not supported for compute capability = sm120" - ) - use_fused_attention = False + if cudnn_version < (9, 18, 1): + if use_fused_attention: + logger.debug( + "Disabling FusedAttention as qkv_format = thd is" + " not supported for compute capability = sm120 and cuDNN version < 9.18.1" + ) + use_fused_attention = False + elif qkv_layout in {"t3hd", "th3d"}: + if use_fused_attention: + logger.debug( + "Disabling FusedAttention as qkv_layout = %s is not supported for" + " compute capability = sm120", + qkv_layout, + ) + use_fused_attention = False # Filter: Dropout if attention_dropout != 0.0 and use_flash_attention_3: diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 2de4576e05..58cfe98d72 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -353,13 +353,22 @@ def fused_attn_fwd( if return_max_logit: qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] - # thd: output_tensors: out [tq, h, d], Max [tq, h, 1], Sum_Exp [tq, h, 1] - # bshd: output_tensors: out [b, sq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] - # sbhd: output_tensors: out [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] + # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Max [tq, h, 1], Sum_Exp [tq, h, 1] + # thd (older cuDNN runtimes or sm120): output_tensors: out [tq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] + # bshd: output_tensors: out [b, sq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] + # sbhd: output_tensors: out [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] stats = output_tensors[1] + torch.log(output_tensors[2]) - amax_dims = (0, 2) if qkv_format == "thd" else (0, 2, 3) + max_tensor = output_tensors[1] + if qkv_format == "thd" and max_tensor.ndim == 4: + # For THD on older cuDNN runtimes or THD on sm120, stats can be [b, h, sq, 1] with padded + # sequence positions. Exclude those padded positions when computing max_logit. + seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=max_tensor.device) + sq_idx = torch.arange(max_tensor.shape[2], device=max_tensor.device).view(1, 1, -1, 1) + valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) + max_tensor = max_tensor.masked_fill(~valid, float("-inf")) + amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) # Max -> max_logit [h] - max_logit = torch.amax(output_tensors[1], dim=amax_dims).to(dtype=output_tensors[0].dtype) + max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=output_tensors[0].dtype) aux_ctx_tensors = [stats] aux_ctx_tensors.extend(output_tensors[3:]) return output_tensors[0], aux_ctx_tensors, max_logit From f2a1a3e991d8cc8e719f9c40d4faea7c73c3289e 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, 23 Mar 2026 16:26:56 +0100 Subject: [PATCH 294/521] [PyTorch Debug] Support tensor dump (#2645) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * code drop Signed-off-by: Pawel Gadzinski * code drop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * docs Signed-off-by: root * nvfp4 internals support Signed-off-by: root * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * lint fixes Signed-off-by: root * Update transformer_engine/debug/features/dump_tensors.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * fix Signed-off-by: root * Update transformer_engine/debug/features/dump_tensors.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * Update transformer_engine/debug/features/dump_tensors.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update tests/pytorch/debug/test_log.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * Update transformer_engine/debug/features/dump_tensors.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * fix Signed-off-by: root * fix Signed-off-by: root * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove dump_quantized_internals support from DumpTensors Drop the dump_quantized_internals config option, the _get_quantized_internals method, and all helper functions for extracting scales/raw data from Float8Tensor, Float8BlockwiseQTensor, MXFP8Tensor, and NVFP4Tensor. Remove corresponding tests: test_dump_tensors_nvfp4_unpacked_codes and NVFP4_DUMP_TENSORS_CONFIG, and scale/data assertions from test_dump_tensors_sanity. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address Greptile review comments - Add dot ('.') to _sanitize_name to handle common PyTorch dotted layer names like 'encoder.layer.0.attention' - Add docstring note about pickle dependency for the 'quantized' key - Add comment explaining weights_only=False in test - Remove redundant local RecipeState import in test_nvfp4_numeric Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * Remove portability suggestion from quantized key docstring Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * Compute rank lazily in _expected_root_dir Avoids relying on stale self.rank when ensure_initialized is called before initialize() has set the rank. Consistent with how nvdlfw_inspect logger resolves rank. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * detach tensors before saving; verify dump filename in test Detach both high_precision and quantized tensors before saving to avoid serializing the autograd graph. For QuantizedTensor this is a zero-copy view (make_like), so no extra GPU allocation. Add filename format assertion to test_dump_tensors_sanity to catch regressions in _sanitize_name or the naming convention. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add empty dump_dict log; assert QuantizedTensor type in test Log a message when no tensors are available to dump so the user has an explicit signal that no file was written. Assert that the quantized key round-trips as a QuantizedTensor to catch regressions in detach() or serialisation path. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/debug/features/dump_tensors.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * Address review: iter subdirs, remove dead rank field, add allclose test and MSE example - Organize dumps into per-iteration subdirectories (iter_000000/) to keep file count manageable per directory. - Remove unused self.rank attribute from TensorLogger. - Add torch.allclose assertion in test to verify serialization correctness. - Add docstring example showing how to load dumps and compute MSE. 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: use detach().clone() to avoid shared storage in DumpTensors Using tensor.detach() creates a view sharing the same underlying storage. If any in-place operation modifies the tensor after the dump, the saved data would be silently corrupted. Use .clone() to ensure the dump captures an independent copy of the data. Signed-off-by: Pawel Gadzinski * test: use torch.equal instead of torch.allclose for serialisation round-trip The saved tensor is an exact bit-for-bit copy (detach().clone()), so torch.equal is the correct check. torch.allclose with its default tolerances could mask a genuine dtype conversion or precision loss introduced by a future change to the serialisation path. Signed-off-by: Pawel Gadzinski * fix: add tp_size to DumpTensors.inspect_tensor and fix KeyError in call_feature backward compat pop Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Signed-off-by: root Signed-off-by: Paweł Gadziński <62263673+pggPL@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: Claude Sonnet 4.6 --- docs/debug/3_api_features.rst | 3 +- tests/pytorch/debug/test_log.py | 83 ++++- transformer_engine/debug/features/api.py | 6 +- .../debug/features/dump_tensors.py | 288 ++++++++++++++++++ .../debug/features/log_fp8_tensor_stats.py | 3 +- 5 files changed, 375 insertions(+), 8 deletions(-) create mode 100644 transformer_engine/debug/features/dump_tensors.py diff --git a/docs/debug/3_api_features.rst b/docs/debug/3_api_features.rst index a973a0b4fe..a8a644d5b5 100644 --- a/docs/debug/3_api_features.rst +++ b/docs/debug/3_api_features.rst @@ -14,4 +14,5 @@ Debug features .. autoapiclass:: transformer_engine.debug.features.per_tensor_scaling.PerTensorScaling .. autoapiclass:: transformer_engine.debug.features.fake_quant.FakeQuant .. autoapiclass:: transformer_engine.debug.features.disable_fp8_gemm.DisableFP8GEMM -.. autoapiclass:: transformer_engine.debug.features.disable_fp8_layer.DisableFP8Layer \ No newline at end of file +.. autoapiclass:: transformer_engine.debug.features.disable_fp8_layer.DisableFP8Layer +.. autoapiclass:: transformer_engine.debug.features.dump_tensors.DumpTensors \ No newline at end of file diff --git a/tests/pytorch/debug/test_log.py b/tests/pytorch/debug/test_log.py index b16291ff61..055210f93a 100644 --- a/tests/pytorch/debug/test_log.py +++ b/tests/pytorch/debug/test_log.py @@ -18,6 +18,7 @@ is_nvfp4_available, ) from transformer_engine.pytorch.quantization import RecipeState +from transformer_engine.pytorch.tensor import QuantizedTensor from transformer_engine.debug.pytorch.debug_state import TEDebugState from transformer_engine.debug.features.utils.stats_computation import ( compute_max_blockwise_dynamic_range, @@ -445,9 +446,6 @@ def test_nvfp4_numeric(feature_dirs): log_nvfp4_config = LOG_NVFP4_CONFIG_BASE.format(stats="underflows%, mse") with debug_session(log_nvfp4_config, feature_dirs) as log_dir: - from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer - from transformer_engine.pytorch.quantization import RecipeState - recipe_state = RecipeState.create( recipe.NVFP4BlockScaling(), mode="forward", @@ -644,3 +642,82 @@ def test_compute_max_blockwise_dynamic_range_direct(): ) print("All direct tests for compute_max_blockwise_dynamic_range passed!") + + +# DumpTensors tests +DUMP_TENSORS_CONFIG = """ +dump: + layers: + layer_name_regex_pattern: .* + enabled: True + transformer_engine: + DumpTensors: + enabled: True + tensors: [activation] + high_precision_tensor: True + quantized_tensor: True + freq: 1 +""" + + +def test_dump_tensors_sanity(feature_dirs): + """Sanity test for DumpTensors feature - verify files are created with correct structure.""" + if not fp8_available: + pytest.skip(reason_for_no_fp8) + + with debug_session(DUMP_TENSORS_CONFIG, feature_dirs) as log_dir: + recipe_state = RecipeState.create( + recipe.DelayedScaling(), + mode="forward", + num_quantizers=3, + ) + + tensor = torch.randn(128, 128, dtype=torch.bfloat16).cuda() + quantizer = recipe_state.make_quantizers()[0] + quantized_tensor = quantizer(tensor) + + debug_api.transformer_engine.inspect_tensor( + layer_name="test_layer", + tensor_name="activation", + iteration=0, + tp_group=None, + tensor=tensor, + quantizer=quantizer, + rowwise_quantized_tensor=quantized_tensor, + columnwise_quantized_tensor=quantized_tensor, + ) + debug_api.step() + + # Check that dump file was created + dump_dir = os.path.join(log_dir, "tensor_dumps", "rank_0") + assert os.path.exists(dump_dir), f"Dump directory not created: {dump_dir}" + + iter_dir = os.path.join(dump_dir, "iter_000000") + assert os.path.exists(iter_dir), f"Iteration directory not created: {iter_dir}" + + dump_files = os.listdir(iter_dir) + assert len(dump_files) == 1, f"Expected 1 dump file, got {len(dump_files)}" + assert ( + dump_files[0] == "test_layer_activation.pt" + ), f"Unexpected dump filename: {dump_files[0]}" + + # Load and verify structure + dump_file = os.path.join(iter_dir, dump_files[0]) + # weights_only=False is required because the dump may contain QuantizedTensor objects, + # which are custom Python classes incompatible with the safe weights_only=True path. + data = torch.load(dump_file, weights_only=False) + + assert isinstance(data, dict), "Dump should be a dictionary" + assert "high_precision" in data, "Missing high_precision tensor" + assert "quantized" in data, "Missing quantized tensor" + assert isinstance( + data["quantized"], QuantizedTensor + ), f"Expected QuantizedTensor, got {type(data['quantized'])}" + + # Verify tensor shapes and values match + assert data["high_precision"].shape == tensor.shape, "high_precision shape mismatch" + assert torch.equal( + data["high_precision"], tensor + ), "high_precision tensor values do not match original tensor" + + print("DumpTensors sanity test passed!") diff --git a/transformer_engine/debug/features/api.py b/transformer_engine/debug/features/api.py index a1cf80dd25..ee9a187b3c 100644 --- a/transformer_engine/debug/features/api.py +++ b/transformer_engine/debug/features/api.py @@ -486,7 +486,7 @@ def call_feature(self, call, feat_config, layer_name, **kwargs): "tp_size", ]: if k not in call.__code__.co_varnames: - kwargs_copy.pop(k) + kwargs_copy.pop(k, None) else: kwargs_copy = kwargs @@ -498,7 +498,9 @@ def call_feature(self, call, feat_config, layer_name, **kwargs): kwargs_copy = kwargs.copy() for k in ["tp_size"]: if k not in call.__code__.co_varnames: - kwargs_copy.pop(k, None) + kwargs_copy.pop( + k, None + ) # use None default to avoid KeyError if kwarg wasn't passed return call(feat_config, layer_name, **kwargs_copy) diff --git a/transformer_engine/debug/features/dump_tensors.py b/transformer_engine/debug/features/dump_tensors.py new file mode 100644 index 0000000000..933acd9438 --- /dev/null +++ b/transformer_engine/debug/features/dump_tensors.py @@ -0,0 +1,288 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""DumpTensors Feature support for nvidia-dlframework-inspect.""" + +import os +from typing import Dict, Optional + +import torch +import torch.distributed as dist + +import nvdlfw_inspect.api as debug_api +from nvdlfw_inspect.logging import get_logger +from nvdlfw_inspect.registry import Registry, api_method + +from transformer_engine.debug.features.api import TEConfigAPIMapper +from transformer_engine.debug.features.utils import next_enabled_iter +from transformer_engine.pytorch.tensor import QuantizedTensor, Quantizer + + +class TensorLogger: + """Logger for saving tensors to files. Each rank saves to its own directory.""" + + _instance = None + _initialized = False + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + def __init__(self): + if TensorLogger._initialized: + return + self.root_dir = None + TensorLogger._initialized = True + + def initialize(self, root_log_dir: str): + """Initialize the TensorLogger with the root directory for tensor dumps.""" + self.root_dir = self._expected_root_dir(root_log_dir) + os.makedirs(self.root_dir, exist_ok=True) + + debug_api.log_message( + f"TensorLogger initialized. Saving tensors to: {self.root_dir}", + ) + + def _expected_root_dir(self, root_log_dir: str) -> str: + """Return the rank-specific dump directory for the provided root log path.""" + rank = dist.get_rank() if dist.is_initialized() else 0 + return os.path.join(root_log_dir, "tensor_dumps", f"rank_{rank}") + + def ensure_initialized(self, root_log_dir: str) -> None: + """Reinitialize logger if debug session log directory changed.""" + expected_root_dir = self._expected_root_dir(root_log_dir) + if self.root_dir != expected_root_dir or not os.path.isdir(expected_root_dir): + self.initialize(root_log_dir) + + @staticmethod + def _sanitize_name(name: str) -> str: + """Sanitize layer/tensor names for use in file paths.""" + for char in ["/", "\\", ":", "*", "?", '"', "<", ">", "|", " ", "."]: + name = name.replace(char, "_") + return name + + def save_tensor( + self, + tensor, + layer_name: str, + tensor_name: str, + iteration: int, + ): + """Save a tensor (or dict of tensors) to a file.""" + if self.root_dir is None: + raise RuntimeError( + "[TE DumpTensors] TensorLogger not initialized. Call initialize() first." + ) + + safe_layer_name = self._sanitize_name(layer_name) + safe_tensor_name = self._sanitize_name(tensor_name) + iter_dir = os.path.join(self.root_dir, f"iter_{iteration:06d}") + os.makedirs(iter_dir, exist_ok=True) + filepath = os.path.join(iter_dir, f"{safe_layer_name}_{safe_tensor_name}.pt") + + if os.path.exists(filepath): + debug_api.log_message(f"[TE DumpTensors] Overwriting existing dump file: {filepath}") + torch.save(tensor, filepath) + + +def _get_tensor_logger() -> TensorLogger: + """Get the singleton TensorLogger instance.""" + return TensorLogger() + + +@Registry.register_feature(namespace="transformer_engine") +class DumpTensors(TEConfigAPIMapper): + """ + Dump tensors to files for debugging purposes. + + This feature saves tensors to disk using torch.save(). It supports dumping + both high-precision tensors (before quantization) and quantized tensors. + + Each tensor is saved to a separate file with the iteration number, layer name, + and tensor name in the filename. Files are organized per-rank in distributed settings. + + Parameters + ---------- + high_precision_tensor : bool + If True, dump the high-precision tensor (before quantization). + quantized_tensor : bool + If True, dump the quantized tensor (after quantization). + tensors/tensors_struct : List[str] + list of tensors to dump: + - activation + - gradient + - weight + - output + - wgrad + - dgrad + freq : Optional[int], default = 1 + frequency of dumping tensors, tensors will be dumped every `freq` steps + start_step : Optional[int], default = 0 + start step of dumping tensors + end_step : Optional[int], default = -1 + end step of dumping tensors (-1 means no end) + start_end_list : Optional[list([int, int])], default = None + non-overlapping list of (start, end) pairs in incremental order. + If not None, will ignore start_step and end_step + + Example + ------- + .. code-block:: yaml + + dump_tensors_example: + enabled: True + layers: + layer_name_regex_pattern: .*(fc1|self_attention).* + transformer_engine: + DumpTensors: + enabled: True + tensors_struct: + - tensor: activation + high_precision_tensor: True + quantized_tensor: True + freq: 100 + - tensor: weight + high_precision_tensor: True + quantized_tensor: False + freq: 500 + + Output Structure + ---------------- + Files are saved to: ``{nvdlfw_inspect_log_dir}/tensor_dumps/rank_{rank}/iter_{iter:06d}/`` + + Each tensor is saved as a dictionary in a single file: + ``{layer}_{tensor}.pt`` + + Dictionary keys: + - ``high_precision``: pre-quantization tensor (if high_precision_tensor=True) + - ``quantized``: quantized tensor object (if quantized_tensor=True) + + .. note:: + The ``quantized`` value is a pickled ``QuantizedTensor`` object. Loading it + (with ``weights_only=False``) requires the same version of TransformerEngine + to be installed. + + Loading and Analyzing Dumped Tensors + ------------------------------------ + .. code-block:: python + + import torch + + # Load dumped tensor (requires the same TE version that produced the dump) + data = torch.load("tensor_dumps/rank_0/iter_000100/fc1_activation.pt", + weights_only=False) + + hp = data["high_precision"] # original high-precision tensor + qt = data["quantized"] # QuantizedTensor object + dequant = qt.dequantize(dtype=hp.dtype) # dequantize back to high precision + + mse = torch.mean((hp - dequant) ** 2).item() + print(f"MSE between original and dequantized: {mse}") + """ + + @api_method + def inspect_tensor_enabled( + self, config: Dict, layer_name: str, tensor_name: str, iteration: int + ): # pylint: disable=unused-argument + """API call used to determine whether to run inspect_tensor() in the forward.""" + run_current, next_iter = next_enabled_iter( + config.get("start_step", None), + config.get("end_step", None), + config.get("start_end_list", None), + config.get("freq", 1), + iteration, + ) + return run_current, next_iter + + @api_method + def inspect_tensor( + self, + config: Dict, + layer_name: str, + tensor_name: str, + iteration: int, + tp_group: torch.distributed.ProcessGroup, + tensor: Optional[torch.Tensor], + rowwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, + columnwise_quantized_tensor: Optional[torch.Tensor | QuantizedTensor] = None, + quantizer: Optional[Quantizer] = None, + tp_size: int = 1, + ): # pylint: disable=unused-argument + """ + API call used to dump tensors to files. + + Supports dumping both high-precision tensors and quantized tensors based on config. + """ + # We support one-sided availability (only rowwise or only columnwise tensor). + # If both are present, require them to be the same object to avoid ambiguity. + if ( + rowwise_quantized_tensor is not None + and columnwise_quantized_tensor is not None + and rowwise_quantized_tensor is not columnwise_quantized_tensor + ): + raise ValueError( + "[NVTORCH INSPECT ERROR] DumpTensors expects rowwise_quantized_tensor and " + "columnwise_quantized_tensor to be the same object when both are provided." + ) + + quantized_tensor = ( + rowwise_quantized_tensor + if rowwise_quantized_tensor is not None + else columnwise_quantized_tensor + ) + + dump_hp = config.get("high_precision_tensor", False) + dump_quant = config.get("quantized_tensor", False) + + if not dump_hp and not dump_quant: + debug_api.log_message( + f"Feature={self.__class__.__name__}: Neither high_precision_tensor nor " + "quantized_tensor is enabled. Nothing to dump.", + layer_name, + ) + return + + tensor_logger = _get_tensor_logger() + tensor_logger.ensure_initialized(get_logger().root_log_dir) + + # Build dictionary with all tensors to dump + dump_dict: Dict[str, torch.Tensor] = {} + + if dump_hp and tensor is not None: + dump_dict["high_precision"] = tensor.detach().clone() + elif dump_hp and tensor is None: + debug_api.log_message( + f"Feature={self.__class__.__name__}: high_precision_tensor is True but " + f"no high-precision tensor available for {tensor_name}. Skipping.", + layer_name, + ) + + if dump_quant and quantized_tensor is not None: + dump_dict["quantized"] = quantized_tensor.detach().clone() + elif dump_quant and quantized_tensor is None: + debug_api.log_message( + f"Feature={self.__class__.__name__}: quantized_tensor is True but " + f"no quantized tensor available for {tensor_name}. Skipping.", + layer_name, + ) + + if dump_dict: + tensor_logger.save_tensor( + tensor=dump_dict, + layer_name=layer_name, + tensor_name=tensor_name, + iteration=iteration, + ) + debug_api.log_message( + f"Feature={self.__class__.__name__}, API=inspect_tensor: " + f"Dumped {tensor_name} at iteration {iteration} (keys: {list(dump_dict.keys())})", + layer_name, + ) + else: + debug_api.log_message( + f"Feature={self.__class__.__name__}: No tensors available to dump for " + f"{tensor_name} at iteration {iteration}. No file written.", + layer_name, + ) diff --git a/transformer_engine/debug/features/log_fp8_tensor_stats.py b/transformer_engine/debug/features/log_fp8_tensor_stats.py index cf11964e25..d26f9ef7f6 100644 --- a/transformer_engine/debug/features/log_fp8_tensor_stats.py +++ b/transformer_engine/debug/features/log_fp8_tensor_stats.py @@ -10,10 +10,9 @@ import torch import nvdlfw_inspect.api as debug_api -import transformer_engine_torch as tex - from nvdlfw_inspect.debug_features.log_tensor_stats import LogTensorStats as BaseLogTensorStats from nvdlfw_inspect.registry import Registry, api_method +import transformer_engine_torch as tex from transformer_engine.debug.features.utils.stats_buffer import STATS_BUFFERS from transformer_engine.debug.features.utils import get_reduction_params, next_enabled_iter From d2625e5f2a15a593685c9bdc5c5d0a721b9a153f Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 23 Mar 2026 17:50:41 -0700 Subject: [PATCH 295/521] Optimize FSDP2 Pytest Timings (12 -> 2 mins) (#2787) Signed-off-by: Varun Thumbe * change distributed tests infra for fsdp2 Signed-off-by: Varun Thumbe * verbose flag for reporting Signed-off-by: Varun Thumbe * add back coments 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: Varun Thumbe * another minor fix Signed-off-by: Varun Thumbe * not needed for this PR Signed-off-by: Varun Thumbe * 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 * unecessary comments --- .../distributed/fsdp2_tests/conftest.py | 85 +++ .../distributed/fsdp2_tests/fsdp2_utils.py | 31 ++ .../{ => fsdp2_tests}/run_fsdp2_fused_adam.py | 525 ++++++++++-------- .../{ => fsdp2_tests}/run_fsdp2_model.py | 155 ++++-- tests/pytorch/distributed/test_torch_fsdp2.py | 268 ++------- 5 files changed, 551 insertions(+), 513 deletions(-) create mode 100644 tests/pytorch/distributed/fsdp2_tests/conftest.py create mode 100644 tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py rename tests/pytorch/distributed/{ => fsdp2_tests}/run_fsdp2_fused_adam.py (58%) rename tests/pytorch/distributed/{ => fsdp2_tests}/run_fsdp2_model.py (80%) diff --git a/tests/pytorch/distributed/fsdp2_tests/conftest.py b/tests/pytorch/distributed/fsdp2_tests/conftest.py new file mode 100644 index 0000000000..bf9db094d2 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/conftest.py @@ -0,0 +1,85 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared pytest fixtures for FSDP2 distributed tests. + +Fixtures defined here (dist_init, _cleanup, recipe_name) are auto-discovered +by pytest for every test module in this directory. +""" + +import gc +import os +import pytest +import torch +import torch.distributed as dist +from transformer_engine.pytorch import fp8 + +# Ensure the correct CUDA device is active before _parametrize_recipes() +# runs at collection time, since the session-scoped dist_init fixture +# has not executed yet. +_local_rank = int(os.environ.get("LOCAL_RANK", "0")) +torch.cuda.set_device(_local_rank) + + +# ── FP8 recipe parametrization ────────────────────────────────────── +def _check_nvfp4_support(): + supported, reason = fp8.check_nvfp4_support() + if supported and torch.cuda.get_device_capability()[0] == 12: + return ( + False, + ( + "NVFP4BlockScaling is failing on SM120 with " + "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " + "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" + ), + ) + return supported, reason + + +_FP8_RECIPE_CONFIGS = [ + ("DelayedScaling", fp8.check_fp8_support), + ("Float8CurrentScaling", fp8.check_fp8_support), + ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), + ("MXFP8BlockScaling", fp8.check_mxfp8_support), + ("NVFP4BlockScaling", _check_nvfp4_support), +] + + +def _parametrize_recipes(): + params = [] + for name, check_fn in _FP8_RECIPE_CONFIGS: + supported, reason = check_fn() + params.append( + pytest.param(name, id=name, marks=pytest.mark.skipif(not supported, reason=reason)) + ) + return params + + +# ── Session / per-test fixtures ────────────────────────────────────── +@pytest.fixture(scope="session", autouse=True) +def dist_init(): + """Initialize the distributed process group once for the entire pytest session.""" + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + yield + if dist.is_initialized(): + dist.destroy_process_group() + + +@pytest.fixture(autouse=True) +def _cleanup(): + """Release GPU memory and stale NCCL state between tests.""" + yield + if dist.is_initialized(): + dist.barrier() + gc.collect() + torch.cuda.empty_cache() + + +@pytest.fixture(params=_parametrize_recipes()) +def recipe_name(request): + return request.param diff --git a/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py new file mode 100644 index 0000000000..178ce62375 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/fsdp2_utils.py @@ -0,0 +1,31 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Shared utility functions for FSDP2 distributed tests.""" + +import transformer_engine.common.recipe +from transformer_engine.pytorch import QuantizedTensor + + +def get_recipe_from_string(recipe): + return getattr(transformer_engine.common.recipe, recipe)() + + +def save_custom_attrs(module): + custom_attrs = {} + for name, param in module.named_parameters(): + if isinstance(param, QuantizedTensor): + ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] + else: + ignore_keys = [] + attrs = vars(param) + custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} + return custom_attrs + + +def restore_custom_attrs(module, custom_attrs): + for name, param in module.named_parameters(): + if name in custom_attrs: + for attr_name, attr_value in custom_attrs[name].items(): + setattr(param, attr_name, attr_value) diff --git a/tests/pytorch/distributed/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py similarity index 58% rename from tests/pytorch/distributed/run_fsdp2_fused_adam.py rename to tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index c39957cf13..877fa66795 100644 --- a/tests/pytorch/distributed/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -6,12 +6,28 @@ """FSDP2 + FusedAdam compatibility tests. -Launched via torchrun from test_fused_optimizer.py. +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run a single test standalone (for debugging): + torchrun --test --recipe + +Available --test values: + fused_adam_fp8_master_weights, fused_adam_fp8_master_weights_no_meta, + fused_adam_bf16, fused_adam_fp8_no_master, fused_adam_bf16_store_param_remainders, + fuse_wgrad_accumulation, dcp_output_parity, dcp_output_parity_async, + safetensors_fp32_export + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling """ import argparse import functools import os +import shutil +import pytest import torch import torch.distributed as dist @@ -24,9 +40,7 @@ from transformer_engine.pytorch import QuantizedTensor import transformer_engine.common.recipe - -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs HIDDEN_SIZE = 256 @@ -38,38 +52,6 @@ def get_recipe_from_string(recipe): NUM_STEPS = 3 -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) - - -def _setup(): - """Common distributed setup. Returns (world_size, local_rank, device).""" - world_size = int(os.environ["WORLD_SIZE"]) - local_rank = int(os.environ["LOCAL_RANK"]) - torch.cuda.set_device(local_rank) - # CPU backend required for async save - dist.init_process_group(backend="cpu:gloo,cuda:nccl") - device = torch.device(f"cuda:{local_rank}") - torch.manual_seed(42) - torch.cuda.manual_seed(42) - return world_size, local_rank, device - - def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_device=True): """Build a Sequential of TransformerLayers, optionally with FP8 init. @@ -143,7 +125,14 @@ def _shard_model(model, world_size): return model -def test_fused_adam_fp8_master_weights(recipe=None): +def _get_dist_info(): + """Get world_size and device from environment (PG already initialized by session fixture).""" + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + return world_size, device + + +def test_fused_adam_fp8_master_weights(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init (FP8 params). Verifies: @@ -151,7 +140,15 @@ def test_fused_adam_fp8_master_weights(recipe=None): - Training loop completes without error - DTensor wrapping and QuantizedTensor local tensors are preserved """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name == "NVFP4BlockScaling": + pytest.xfail( + f"{recipe_name}: quantized_model_init and FSDP2 is not currently supported, since the " + "block tensor is dequantized before we flatten it for FSDP2." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -206,10 +203,8 @@ def test_fused_adam_fp8_master_weights(recipe=None): ) assert qt_count > 0, "No QuantizedTensor local tensors after training" - dist.destroy_process_group() - -def test_fused_adam_fp8_master_weights_no_meta(recipe=None): +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. @@ -219,7 +214,16 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): For per-tensor FP8 (DelayedScaling, Float8CurrentScaling) this works because Float8Tensor's storage is accessible via data_ptr(). """ - world_size, _, device = _setup() + 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). " + "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) model = _shard_model(model, world_size) @@ -242,15 +246,15 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe=None): loss.backward() optimizer.step() - dist.destroy_process_group() - -def test_fused_adam_bf16(recipe=None): +def test_fused_adam_bf16(recipe_name): """FusedAdam with master_weights + FSDP2 + bf16 params (no FP8). Verifies the non-FP8 DTensor param path in step() works correctly. """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=False) model = _shard_model(model, world_size) @@ -284,15 +288,21 @@ def test_fused_adam_bf16(recipe=None): # Verify loss decreased (basic sanity) assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" - dist.destroy_process_group() - -def test_fused_adam_fp8_no_master(recipe=None): +def test_fused_adam_fp8_no_master(recipe_name): """FusedAdam without master_weights + FSDP2 + FP8 params. Verifies FusedAdam works with FSDP2 even without master weights enabled. """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + + if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + pytest.xfail( + f"{recipe_name}: FusedAdam without master_weights does not support " + "block-scaling quantized tensors. Use master_weights=True." + ) + + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -318,10 +328,8 @@ def test_fused_adam_fp8_no_master(recipe=None): for name, param in model.named_parameters(): assert isinstance(param, DTensor), f"{name} lost DTensor wrapping" - dist.destroy_process_group() - -def test_fused_adam_bf16_store_param_remainders(recipe=None): +def test_fused_adam_bf16_store_param_remainders(recipe_name): """FusedAdam with master_weights + store_param_remainders + FSDP2 + bf16 params. store_param_remainders stores only the trailing 16 remainder bits (int16) @@ -335,7 +343,8 @@ def test_fused_adam_bf16_store_param_remainders(recipe=None): - exp_avg and exp_avg_sq are float32 - Loss decreases (basic sanity) """ - world_size, _, device = _setup() + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() model = _build_model(fp8_init=False) model = _shard_model(model, world_size) @@ -385,10 +394,18 @@ def test_fused_adam_bf16_store_param_remainders(recipe=None): # Verify loss decreased (basic sanity) assert losses[-1] < losses[0], f"Loss did not decrease: {losses}" - dist.destroy_process_group() - -def test_fuse_wgrad_accumulation(recipe=None): +@pytest.mark.xfail( + reason=( + "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " + "autograd Function.apply unwraps DTensors to local tensors, so " + "main_grad (set on the DTensor) is inaccessible during backward. " + "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." + ), + raises=AttributeError, + strict=True, +) +def test_fuse_wgrad_accumulation(recipe_name): """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail. With vanilla FSDP2, PyTorch's autograd Function.apply unwraps DTensor @@ -400,8 +417,8 @@ def test_fuse_wgrad_accumulation(recipe=None): writes the gradient directly into main_grad and returns None to autograd, bypassing FSDP2's reduce-scatter. """ - world_size, _, device = _setup() - + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, fuse_wgrad_accumulation=True, recipe=recipe) # Allocate main_grad buffers on the DTensor params @@ -433,10 +450,8 @@ def test_fuse_wgrad_accumulation(recipe=None): loss = F.mse_loss(output, target) loss.backward() # Expected to raise AttributeError - dist.destroy_process_group() - -def test_safetensors_fp32_export(recipe=None): +def test_safetensors_fp32_export(recipe_name): """Export full-precision (FP32) model to safetensors from optimizer master weights. Verifies: @@ -446,6 +461,13 @@ def test_safetensors_fp32_export(recipe=None): - All saved tensors are float32 - 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" + ) + from safetensors.torch import load_file, save_file from torch.distributed.checkpoint.state_dict import ( StateDictOptions, @@ -453,8 +475,7 @@ def test_safetensors_fp32_export(recipe=None): get_optimizer_state_dict, ) - world_size, _, device = _setup() - + world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) model = _shard_model(model, world_size) @@ -483,38 +504,39 @@ def test_safetensors_fp32_export(recipe=None): full_opt_state = get_optimizer_state_dict(model, optimizer, options=full_opts) rank = int(os.environ.get("RANK", "0")) - save_path = "/tmp/te_test_fsdp2_model_fp32.safetensors" + save_path = f"/tmp/te_test_fsdp2_model_fp32_{recipe_name}.safetensors" if rank == 0: - # Build FP32 state dict from optimizer master weights. - fp32_state = {} - opt_param_states = full_opt_state.get("state", {}) - - for key, value in full_model_state.items(): - if key in opt_param_states and "master_param" in opt_param_states[key]: - fp32_state[key] = opt_param_states[key]["master_param"].float() - else: - fp32_state[key] = value.float() + if os.path.exists(save_path): + os.remove(save_path) - assert len(fp32_state) > 0, "FP32 state dict is empty" + try: + fp32_state = {} + opt_param_states = full_opt_state.get("state", {}) - # Save and verify. - save_file(fp32_state, save_path) - loaded = load_file(save_path) + for key, value in full_model_state.items(): + if key in opt_param_states and "master_param" in opt_param_states[key]: + fp32_state[key] = opt_param_states[key]["master_param"].float() + else: + fp32_state[key] = value.float() - assert len(loaded) == len( - fp32_state - ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" - for k, v in loaded.items(): - assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + assert len(fp32_state) > 0, "FP32 state dict is empty" - # Clean up. - os.remove(save_path) + save_file(fp32_state, save_path) + loaded = load_file(save_path) - dist.destroy_process_group() + assert len(loaded) == len( + fp32_state + ), f"Loaded {len(loaded)} tensors, expected {len(fp32_state)}" + for k, v in loaded.items(): + assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" + finally: + if os.path.exists(save_path): + os.remove(save_path) -def test_dcp_output_parity(recipe=None, async_save=False): +@pytest.mark.parametrize("async_save", [False, True], ids=["sync", "async"]) +def test_dcp_output_parity(recipe_name, async_save): """DCP save/load round-trip produces bitwise-identical model outputs. 1. Builds and trains a model for NUM_STEPS @@ -525,156 +547,197 @@ def test_dcp_output_parity(recipe=None, async_save=False): 6. Runs the same forward pass and asserts outputs are identical 7. Runs one more training step on both models and asserts outputs still match """ - import torch.distributed.checkpoint as dcp - - world_size, local_rank, device = _setup() - - # ── Build and train the original model ─────────────────────────── - model = _build_model(fp8_init=True, recipe=recipe) - model = _shard_model(model, world_size) - - optimizer = te.optimizers.FusedAdam( - model.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) - - x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) - target = torch.randn_like(x) - - for _ in range(NUM_STEPS): - optimizer.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - output = model(x) - loss = F.mse_loss(output, target) - loss.backward() - optimizer.step() + 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" + ) - # Record reference output from the trained model. - with torch.no_grad(): - with te.autocast(enabled=True, recipe=recipe): - ref_output = model(x).clone() - - # ── Save checkpoint ────────────────────────────────────────────── - checkpoint_dir = "/tmp/te_test_fsdp2_dcp_parity" - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - # We need to remove the _extra_state keys from the model state dict for DelayedScaling, - # since otherwise we'll run into an error that the tensor sizes are different. The - # alternative is a LoadPlanner that dynamically re-sizes the input tensors, see - # NVIDIA/TransformerEngine#1860 for more details. - model_state = { - k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") - } - else: - model_state = model.state_dict() + 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" + ) - save_state = {"model": model_state, "optimizer": optimizer.state_dict()} + if ( + recipe_name == "Float8BlockScaling" + and not async_save + 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 " + "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " + "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." + ) - if not async_save: - dcp.save(save_state, checkpoint_id=checkpoint_dir) - else: - future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) - future.result() # Block on async save completion + import torch.distributed.checkpoint as dcp - # ── Build a fresh model and load the checkpoint ────────────────── - model2 = _build_model(fp8_init=True, recipe=recipe) - model2 = _shard_model(model2, world_size) + world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + save_mode = "async" if async_save else "sync" + checkpoint_dir = f"/tmp/te_test_fsdp2_dcp_parity_{recipe_name}_{save_mode}" - optimizer2 = te.optimizers.FusedAdam( - model2.parameters(), - lr=1e-3, - master_weights=True, - master_weight_dtype=torch.float32, - ) + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + dist.barrier() + + try: + # ── Build and train the original model ─────────────────────────── + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) - # Populate optimizer state so load_state_dict has matching structure. - optimizer2.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out_tmp = model2(x) - F.mse_loss(out_tmp, target).backward() - optimizer2.step() - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - model2_state = { - k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") - } - else: - model2_state = model2.state_dict() + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Record reference output from the trained model. + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + ref_output = model(x).clone() + + # ── Save checkpoint ────────────────────────────────────────────── + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # We need to remove the _extra_state keys from the model state dict for + # DelayedScaling, since otherwise we'll run into an error that the tensor + # sizes are different. The alternative is a LoadPlanner that dynamically + # re-sizes the input tensors, see NVIDIA/TransformerEngine#1860 for more + # details. + model_state = { + k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") + } + else: + model_state = model.state_dict() - state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + save_state = {"model": model_state, "optimizer": optimizer.state_dict()} - dcp.load(state_to_load, checkpoint_id=checkpoint_dir) - model2.load_state_dict( - state_to_load["model"], - strict=( - False if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) else True - ), - ) - optimizer2.load_state_dict(state_to_load["optimizer"]) - - # ── Verify identical forward-pass output ───────────────────────── - with torch.no_grad(): - 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. - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0.05, - atol=0.1, - msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", - ) - else: - torch.testing.assert_close( - loaded_output, - ref_output, - rtol=0, - atol=0, - msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", + if not async_save: + dcp.save(save_state, checkpoint_id=checkpoint_dir) + else: + future = dcp.async_save(save_state, checkpoint_id=checkpoint_dir) + future.result() + + # ── Build a fresh model and load the checkpoint ────────────────── + model2 = _build_model(fp8_init=True, recipe=recipe) + model2 = _shard_model(model2, world_size) + + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, ) - # ── Verify one more training step produces identical results ───── - optimizer.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out1 = model(x) - loss1 = F.mse_loss(out1, target) - loss1.backward() - optimizer.step() - - optimizer2.zero_grad(set_to_none=True) - with te.autocast(enabled=True, recipe=recipe): - out2 = model2(x) - loss2 = F.mse_loss(out2, target) - loss2.backward() - optimizer2.step() - - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - torch.testing.assert_close( - out2, - out1, - rtol=0.05, - atol=0.1, - msg="Training step after DCP load produces different output", - ) - else: - torch.testing.assert_close( - out2, out1, msg="Training step after DCP load produces different output" + # Populate optimizer state so load_state_dict has matching structure. + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + model2_state = { + k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") + } + else: + model2_state = model2.state_dict() + + state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict( + state_to_load["model"], + strict=( + False + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) + else True + ), ) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + # ── Verify identical forward-pass output ───────────────────────── + with torch.no_grad(): + 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. + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0.05, + atol=0.1, + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", + ) + else: + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0, + atol=0, + msg=lambda x: f"Fresh model loaded from DCP checkpoint produces different output: {x}", + ) + + # ── Verify one more training step produces identical results ───── + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out1 = model(x) + loss1 = F.mse_loss(out1, target) + loss1.backward() + optimizer.step() - # ── Cleanup ────────────────────────────────────────────────────── - import shutil - - if int(os.environ.get("RANK", "0")) == 0: - shutil.rmtree(checkpoint_dir, ignore_errors=True) - - dist.destroy_process_group() + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out2 = model2(x) + loss2 = F.mse_loss(out2, target) + loss2.backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + torch.testing.assert_close( + out2, + out1, + rtol=0.05, + atol=0.1, + msg="Training step after DCP load produces different output", + ) + else: + torch.testing.assert_close( + out2, out1, msg="Training step after DCP load produces different output" + ) + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) TESTS = { @@ -707,5 +770,13 @@ def test_dcp_output_parity(recipe=None, async_save=False): ], ) args = parser.parse_args() - recipe = get_recipe_from_string(args.recipe) - TESTS[args.test](recipe) + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + try: + TESTS[args.test](args.recipe) + finally: + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/tests/pytorch/distributed/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py similarity index 80% rename from tests/pytorch/distributed/run_fsdp2_model.py rename to tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 60d7cd2023..fce565ed9a 100644 --- a/tests/pytorch/distributed/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -4,9 +4,36 @@ # # See LICENSE for license information. +"""FSDP2 model sharding tests. + +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run standalone (for debugging): + torchrun --recipe [options] + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling + +Other options: + --fp8-init Initialize weights in FP8 + --layer-type TYPE Linear, LayerNormLinear, LayerNormMLP, + MultiheadAttention, TransformerLayer (default) + --sharding-dims N [M] FSDP dims, e.g. "2" or "2 2" for HSDP + --num-layers N Number of layers (default: 4) + --iter N Training iterations (default: 10) + --device cuda|meta Device for init (default: meta) +""" + +import gc import os import sys import argparse +from types import SimpleNamespace +from contextlib import nullcontext + +import pytest import transformer_engine.pytorch as te import transformer_engine.common.recipe @@ -19,14 +46,12 @@ from torch.distributed import DeviceMesh from torch.distributed._composable.fsdp import fully_shard from torch.distributed.device_mesh import init_device_mesh -from transformer_engine.pytorch import QuantizedTensor -from contextlib import nullcontext -LOCAL_RANK = None +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs def dist_print(msg): - if LOCAL_RANK == 0: + if int(os.getenv("LOCAL_RANK", "0")) == 0: print(msg) @@ -114,10 +139,6 @@ def get_te_layer_from_string(layer_name): return te_layer_map[layer_name.lower()] -def get_recipe_from_string(recipe): - return getattr(transformer_engine.common.recipe, recipe)() - - def init_te_model(config): hidden_size = config.num_heads * config.head_dim args = [hidden_size, hidden_size] @@ -188,31 +209,8 @@ def shard_model_with_fsdp2(model, mesh): return model -#### Methods to save the custom attributes of QuantizedTensors before sharding -#### them with FSDP2, and restore them after sharding. -def save_custom_attrs(module): - custom_attrs = {} - for name, param in module.named_parameters(): - if isinstance(param, QuantizedTensor): - # Ignore FP8 metadata attributes. Otherwise we will save duplicate copies - # for data/transpose FP8 tensors on top of FP8 tensors that FSDP2 will save. - ignore_keys = [key for key in param.__dict__.keys() if key.startswith("_")] - else: - ignore_keys = [] - attrs = vars(param) - custom_attrs[name] = {k: v for k, v in attrs.items() if k not in ignore_keys} - return custom_attrs - - -def restore_custom_attrs(module, custom_attrs): - for name, param in module.named_parameters(): - if name in custom_attrs: - for attr_name, attr_value in custom_attrs[name].items(): - setattr(param, attr_name, attr_value) - - @torch.no_grad() -def test_fp8_fsdp2_allgather(model): +def _check_fp8_fsdp2_allgather(model): # Do manual allgather in fp32 and match against fp8 allgather done # with fsdp2 # FP32 manual weight allgather @@ -249,30 +247,10 @@ def test_fp8_fsdp2_allgather(model): module.reshard() -def _train(args): - global LOCAL_RANK - assert "TORCHELASTIC_RUN_ID" in os.environ - WORLD_RANK = int(os.getenv("RANK", "0")) - WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) - LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) - LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) - assert LOCAL_SIZE == WORLD_SIZE - - # Set device and initialize RNG states - torch.cuda.set_device(WORLD_RANK) - torch.manual_seed(args.seed) - torch.cuda.manual_seed(args.seed) - - # Initialize torch.distributed global process group and get DP/TP groups - dist_init_kwargs = { - "backend": "nccl", - "rank": WORLD_RANK, - "world_size": WORLD_SIZE, - } - assert dist.is_nccl_available() - dist.init_process_group(**dist_init_kwargs) - nccl_world = dist.new_group(backend="nccl") - device = torch.device(f"cuda:{LOCAL_RANK}") +def _run_training(args): + """Core training logic. Assumes dist is already initialized.""" + device = torch.device(f"cuda:{int(os.getenv('LOCAL_RANK', '0'))}") + world_size = int(os.getenv("WORLD_SIZE", "1")) # FP8 Configuration fp8_recipe = get_recipe_from_string(args.recipe) @@ -298,7 +276,6 @@ def _train(args): ) # Creating a DeviceMesh for fully_shard - world_size = int(WORLD_SIZE) # Setup the sharding mesh for FSDP/HSDP mesh = get_device_mesh(world_size, args.sharding_dims) custom_attrs = save_custom_attrs(model) @@ -344,11 +321,71 @@ def _train(args): # Some of the FSDP states are lazy initialized during FSDP forward pass # so testing fp8 allgather at the end of the training loop. if args.fp8_init: - test_fp8_fsdp2_allgather(model) + _check_fp8_fsdp2_allgather(model) + + +def _train(args): + """Standalone entry point with full dist lifecycle.""" + assert "TORCHELASTIC_RUN_ID" in os.environ + WORLD_RANK = int(os.getenv("RANK", "0")) + WORLD_SIZE = int(os.getenv("WORLD_SIZE", "1")) + LOCAL_RANK = int(os.getenv("LOCAL_RANK", "0")) + LOCAL_SIZE = int(os.getenv("LOCAL_WORLD_SIZE", "1")) + assert LOCAL_SIZE == WORLD_SIZE + + torch.cuda.set_device(LOCAL_RANK) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + + assert dist.is_nccl_available() + dist.init_process_group( + backend="nccl", + rank=WORLD_RANK, + world_size=WORLD_SIZE, + ) + try: + _run_training(args) + finally: + if dist.is_initialized(): + dist.destroy_process_group() + torch.cuda.empty_cache() + gc.collect() - dist.destroy_process_group() return 0 +# ── Pytest test function ───────────────────────────────────────────── + +NUM_PROCS = int(os.environ.get("WORLD_SIZE", "1")) + + +@pytest.mark.parametrize("sharding_dims", [[NUM_PROCS], [2, NUM_PROCS // 2]]) +@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 in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: + pytest.xfail(f"{recipe_name} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") + + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + args = SimpleNamespace( + recipe=recipe_name, + fp8_init=fp8_init, + sharding_dims=list(sharding_dims), + layer_type=layer_type, + seed=42, + num_heads=8, + head_dim=64, + batch_size=16, + seq_length=128, + params_dtype="float32", + num_layers=4, + iter=10, + device="meta", + ) + _run_training(args) + + if __name__ == "__main__": sys.exit(_train(_parse_args())) diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 02e45d99cb..aca8d6d692 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -10,242 +10,56 @@ import torch import transformer_engine.pytorch as te -from transformer_engine.pytorch import fp8 NUM_PROCS: int = torch.cuda.device_count() - - -def check_nvfp4_support(): - supported, reason = fp8.check_nvfp4_support() - if supported and torch.cuda.get_device_capability()[0] == 12: - return ( - False, - ( - "NVFP4BlockScaling is failing on SM120 with " - "hadamard_transform/hadamard_transform_cast_fusion.cu:672 in function " - "rht_gemm_ntt_w_sfc: CUDA Error: invalid argument" - ), - ) - - return supported, reason - - -# Each entry: (recipe_class_name, check_fn) -_FP8_RECIPE_CONFIGS = [ - ("DelayedScaling", fp8.check_fp8_support), - ("Float8CurrentScaling", fp8.check_fp8_support), - ("Float8BlockScaling", fp8.check_fp8_block_scaling_support), - ("MXFP8BlockScaling", fp8.check_mxfp8_support), - ("NVFP4BlockScaling", check_nvfp4_support), -] - - -def _parametrize_fp8_recipes(): - """Generate pytest.param objects with skip marks for unsupported FP8 recipes.""" - params = [] - for name, check_fn in _FP8_RECIPE_CONFIGS: - supported, reason = check_fn() - params.append( - pytest.param( - name, - id=name, - marks=pytest.mark.skipif(not supported, reason=reason), - ) - ) - return params - - -@pytest.fixture(params=_parametrize_fp8_recipes()) -def fp_recipe(request): - """Parametrized fixture providing FP8 recipe Hydra overrides for each supported TE recipe.""" - return request.param - - -def _run_test(fp_init, sharding_dims, recipe, layer_type): - test_path = Path(__file__).parent.resolve() / "run_fsdp2_model.py" - test_cmd = ["torchrun", f"--nproc_per_node={NUM_PROCS}", str(test_path)] - - if fp_init: - test_cmd += ["--fp8-init"] - - if len(sharding_dims) == 1: - test_cmd += ["--sharding-dims", str(sharding_dims[0])] - elif len(sharding_dims) == 2: - test_cmd += ["--sharding-dims", str(sharding_dims[0]), str(sharding_dims[1])] - else: - assert False - test_cmd += ["--recipe", recipe] - test_cmd += ["--layer-type", layer_type] - - subprocess.run(test_cmd, env=os.environ, check=True) +_FSDP2_DIR = Path(__file__).parent.resolve() / "fsdp2_tests" @pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") -@pytest.mark.parametrize("sharding_dims", ([NUM_PROCS], [2, NUM_PROCS // 2])) -@pytest.mark.parametrize("fp8_init", (False, True)) -@pytest.mark.parametrize("layer_type", ("LayerNormLinear", "TransformerLayer")) -def test_distributed(fp8_init, sharding_dims, fp_recipe, layer_type): - - if fp_recipe in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: - pytest.xfail(f"{fp_recipe} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") - - _run_test(fp8_init, sharding_dims, fp_recipe, layer_type) - - -## ── FusedAdam + FSDP2 tests ───────────────────────────────────────── - - -def _run_fused_adam_test(test_name, recipe="delayed_scaling"): - """Launch an FSDP2 + FusedAdam test via torchrun.""" - test_path = Path(__file__).parent.resolve() / "run_fsdp2_fused_adam.py" - nproc = min(NUM_PROCS, 2) # These tests only need 2 GPUs - test_cmd = [ - "torchrun", - f"--nproc_per_node={nproc}", - str(test_path), - "--test", - test_name, - "--recipe", - recipe, - ] - - subprocess.run(test_cmd, env=os.environ, check=True) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_master_weights(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (meta device init).""" - if fp_recipe in ("NVFP4BlockScaling",): - pytest.xfail( - f"{fp_recipe}: quantized_model_init and FSDP2 is not currently supported, since the " - "block tensor is dequantized before we flatten it for FSDP2." - ) - _run_fused_adam_test("fused_adam_fp8_master_weights", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_master_weights_no_meta(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + quantized_model_init (CUDA init, no meta device). - - Block-scaling QuantizedTensors (MXFP8, Float8Blockwise, NVFP4) are wrapper - subclasses with data_ptr() == 0. Without meta-device init, FSDP2's - reset_sharded_param() crashes with 'invalid python storage'. - Per-tensor FP8 (DelayedScaling, Float8CurrentScaling) works because - Float8Tensor's storage is accessible. - """ - if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{fp_recipe}: FSDP2 without meta-device init crashes on block-scaling " - "QuantizedTensor wrapper subclasses (data_ptr() == 0). " - "Use device='meta' + reset_parameters() after sharding." - ) - _run_fused_adam_test("fused_adam_fp8_master_weights_no_meta", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_bf16(fp_recipe): - """FusedAdam(master_weights=True) + FSDP2 + bf16 params (no FP8).""" - _run_fused_adam_test("fused_adam_bf16", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_fp8_no_master(fp_recipe): - """FusedAdam(master_weights=False) + FSDP2 + FP8 params.""" - if fp_recipe in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{fp_recipe}: FusedAdam without master_weights does not support " - "block-scaling quantized tensors. Use master_weights=True." - ) - _run_fused_adam_test("fused_adam_fp8_no_master", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_fused_adam_bf16_store_param_remainders(fp_recipe): - """FusedAdam(master_weights=True, store_param_remainders=True) + FSDP2 + bf16.""" - _run_fused_adam_test("fused_adam_bf16_store_param_remainders", fp_recipe) +def test_fsdp2_model_tests(): + """All FSDP2 model tests (parametrized internally by recipe, fp8_init, sharding, layer).""" + test_path = _FSDP2_DIR / "run_fsdp2_model.py" + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={NUM_PROCS}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + env=os.environ, + timeout=600, + ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_dcp_output_parity(fp_recipe): - """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) - - if fp_recipe == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if fp_recipe == "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 " - "in function quantize_transpose_vector_blockwise: Assertion failed: pow2_scale. On " - "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " - "requires using power of two scaling factors." - ) - - _run_fused_adam_test("dcp_output_parity", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_dcp_output_parity_async(fp_recipe): - """DCP save/load round-trip into a fresh model produces identical outputs.""" - if fp_recipe == "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" - ) - - if fp_recipe == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if fp_recipe == "Float8BlockScaling": - 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." - ) - - _run_fused_adam_test("dcp_output_parity_async", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -def test_fsdp2_safetensors_fp32_export(fp_recipe): - """Export FP32 model from optimizer master weights to safetensors.""" - if fp_recipe == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" - ) - _run_fused_adam_test("safetensors_fp32_export", fp_recipe) - - -@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") -@pytest.mark.xfail( - reason=( - "fuse_wgrad_accumulation is incompatible with vanilla FSDP2: " - "autograd Function.apply unwraps DTensors to local tensors, so " - "main_grad (set on the DTensor) is inaccessible during backward. " - "Additionally, the fused wgrad GEMM bypasses FSDP2's reduce-scatter." - ), - raises=subprocess.CalledProcessError, - strict=True, -) -def test_fsdp2_fuse_wgrad_accumulation(fp_recipe): - """fuse_wgrad_accumulation=True + FSDP2 -- expected to fail.""" - _run_fused_adam_test("fuse_wgrad_accumulation", fp_recipe) +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_fused_adam_tests(): + """All FSDP2 FusedAdam tests (parametrized internally by recipe, test variant).""" + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" + nproc = min(NUM_PROCS, 2) + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + env=os.environ, + timeout=600, + ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" def test_dummy() -> None: From 8477d3dcb0a10861cba08e26489169ffcb8f8a53 Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Tue, 24 Mar 2026 05:23:56 +0100 Subject: [PATCH 296/521] Enable fused RMSNorm dLN + add through CUDNN (#2778) * add cudnn dln+add Signed-off-by: CarlosGomes98 * try fixing cudnn build issue Signed-off-by: CarlosGomes98 * guard against cudnn version Signed-off-by: CarlosGomes98 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * change itype to wtype for add in rmsnorm_bwd Signed-off-by: CarlosGomes98 * remove dead code Signed-off-by: CarlosGomes98 * remove dangling todo Signed-off-by: CarlosGomes98 --------- Signed-off-by: CarlosGomes98 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../common/normalization/common.cpp | 26 ++++++++++++++++--- .../common/normalization/common.h | 2 +- .../normalization/rmsnorm/rmsnorm_api.cpp | 23 +++++++++------- 3 files changed, 38 insertions(+), 13 deletions(-) diff --git a/transformer_engine/common/normalization/common.cpp b/transformer_engine/common/normalization/common.cpp index 11f12775c5..7dd942b314 100644 --- a/transformer_engine/common/normalization/common.cpp +++ b/transformer_engine/common/normalization/common.cpp @@ -395,6 +395,23 @@ CudnnNormalizationPlan::CudnnNormalizationPlan(NVTE_Norm_Type NormType, NVTE_Nor std::tie(_dx, _dgamma, _dbeta) = std::make_tuple(ret[0], ret[1], ret[2]); if (_dbeta != nullptr) NVTE_ERROR("cuDNN rmsnorm dbias incorrectly returned."); } + // Fuse the add for BackwardAdd stage + if (_norm_stage == NVTE_Norm_Stage::BackwardAdd) { + NVTE_CHECK(cudnnGetVersion() >= 92100, + "Fused BackwardAdd requires cuDNN >= 9.21.0, but found ", cudnnGetVersion()); + + _add = _graph.tensor(fe::graph::Tensor_attributes() + .set_name("add") + .set_dim({batch_dim, hidden_dim, 1, 1}) + .set_stride({hidden_dim, 1, hidden_dim, hidden_dim}) + .set_data_type(get_cudnn_fe_dtype(wtype))); + auto add_options = fe::graph::Pointwise_attributes() + .set_mode(fe::PointwiseMode_t::ADD) + .set_compute_data_type(get_cudnn_fe_dtype(ctype)); + auto _dx_with_add = _graph.pointwise(_dx, _add, add_options); + _dx->set_output(false).set_data_type(get_cudnn_fe_dtype(itype)); + _dx = _dx_with_add; + } _dx->set_output(true).set_data_type(get_cudnn_fe_dtype(otype)); _dgamma->set_output(true).set_data_type(get_cudnn_fe_dtype(otype)); } @@ -467,13 +484,16 @@ void CudnnNormalizationPlan::execute(void* x_dptr, void* gamma_dptr, void* mean_ void* rsigma_dptr, void* dx_dptr, void* dz_dptr, void* add_dptr, void* dbeta_dptr, void* dgamma_dptr, void* workspace_dptr, cudaStream_t stream) { - // cuDNN does not currently support fused backward+add - NVTE_CHECK(add_dptr == nullptr); - // Binding data pointers to graph tensors _variant_pack = { {_x, x_dptr}, {_rsigma, rsigma_dptr}, {_dz, dz_dptr}, {_dgamma, dgamma_dptr}, {_dx, dx_dptr}}; + // Bind the add tensor for fused backward+add + if (_norm_stage == NVTE_Norm_Stage::BackwardAdd) { + NVTE_CHECK(add_dptr != nullptr, "add_dptr must not be null for BackwardAdd"); + _variant_pack.insert({{_add, add_dptr}}); + } + if (_zero_centered) _variant_pack.insert({{_scalar_offset, reinterpret_cast(this->_scalar_dptr.get())}, {_gamma_zero, gamma_dptr}}); diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index 79de2ac140..0cbd5a99f9 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -294,7 +294,7 @@ class CudnnNormalizationPlan : public NormalizationPlanBase { std::shared_ptr _z_mx_row, _z_mx_col, _sf_row, _sf_col; const bool _training; // BWD - std::shared_ptr _dz, _dx, _dgamma, _dbeta; + std::shared_ptr _dz, _dx, _dgamma, _dbeta, _add; fe::graph::Graph _graph; std::unordered_map, void*> _variant_pack; diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index 6f6656534a..adf2ccee04 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -206,16 +206,21 @@ void rmsnorm_bwd_add(const Tensor &dz, const Tensor &x, const Tensor &add, const CheckOutputTensor(*dgamma, "dgamma"); } - // cuDNN does not currently support fused backward+add - NVTE_Norm_Backend norm_backend = NVTE_Norm_Backend::Te; - - // TE backend does not currently support zero_centered_gamma_in_weight_dtype - NVTE_CHECK(!use_zero_centered_gamma_in_weight_dtype(), - "zero_centered_gamma_in_weight_dtype is currently not supported for rmsnorm_bwd_add"); - - bool is_aligned = is_ptr_aligned(x.data.dptr, gamma.data.dptr, rsigma.data.dptr, dx->data.dptr, - dz.data.dptr, dgamma->data.dptr, add.data.dptr); + NVTE_Norm_Backend norm_backend; + bool is_aligned = true; bool gamma_in_weight_dtype = false; + if (use_cudnn_norm_bwd()) { + norm_backend = NVTE_Norm_Backend::Cudnn; + gamma_in_weight_dtype = use_zero_centered_gamma_in_weight_dtype(); + } else { + norm_backend = NVTE_Norm_Backend::Te; + // TE backend does not currently support zero_centered_gamma_in_weight_dtype + NVTE_CHECK(!use_zero_centered_gamma_in_weight_dtype(), + "zero_centered_gamma_in_weight_dtype is currently not supported " + "for rmsnorm_bwd_add with TE backend"); + is_aligned = is_ptr_aligned(x.data.dptr, gamma.data.dptr, rsigma.data.dptr, dx->data.dptr, + dz.data.dptr, dgamma->data.dptr, add.data.dptr); + } auto plan = NormalizationPlanRegistry::getInstance().getNormalizationPlan( norm_backend, NVTE_Norm_Type::RMSNorm, NVTE_Norm_Stage::BackwardAdd, From 4013c6c2801dec4437c4ab6abc9c957c25481e15 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 24 Mar 2026 16:14:08 -0700 Subject: [PATCH 297/521] add blackwell support filter for 9.7<=cudnn<9.18.1 (#2775) * add blackwell support filter for 9.7<=cudnn<9.18.1 Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * simplify conditionals Signed-off-by: Sudhakar Singh * fix conditionals again Signed-off-by: Sudhakar Singh * fix conditionals again Signed-off-by: Sudhakar Singh * update the error log Signed-off-by: Sudhakar Singh * remove the python filter and correct the cpp filter 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> --- transformer_engine/common/fused_attn/fused_attn.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index cba1a79dd3..e1071edff4 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -310,7 +310,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( // architecture ((cudnn_runtime_version < 8903 && (sm_arch_ == 80 || sm_arch_ == 90)) || (cudnn_runtime_version >= 8903 && sm_arch_ >= 80 && sm_arch_ < 100) || - (cudnn_runtime_version >= 90700 && sm_arch_ >= 80)) && + (cudnn_runtime_version >= 90700 && sm_arch_ >= 100)) && // sequence length ((cudnn_runtime_version < 90000 && max_seqlen_q % 64 == 0 && max_seqlen_kv % 64 == 0) || (cudnn_runtime_version >= 90000)) && From 4ead776cf4409dac054fdef0f229ca3b3c868b90 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Tue, 24 Mar 2026 18:59:18 -0700 Subject: [PATCH 298/521] [PyT][Commong] Disable fused attention for sm120 if determinism is required (#2798) * Disable fused attention for sm120 if determinism is required Signed-off-by: Kshitij Lakhani * nit: disable fused attn for sm120 determinism, if training Signed-off-by: Kshitij Lakhani --------- Signed-off-by: Kshitij Lakhani --- transformer_engine/common/fused_attn/fused_attn.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index e1071edff4..3d6e3a0aac 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -534,6 +534,10 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( std::cout << "Warning: Given combination of sm_arch_ == 120 and cudnn_runtime_version < " "91801 is not supported. " << " Please upgrade your cuDNN version if possible." << std::endl; + } else if (deterministic && is_training) { + backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; + std::cout << "Warning: Deterministic fused attention on SM120 is not supported." + << std::endl; } else { // Known missing support for T3HD/TH3D layouts on SM120 const bool is_t3hd_or_th3d = From e879bf87af032cb919f4851b913f3573c730c748 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Tue, 24 Mar 2026 21:11:15 -0700 Subject: [PATCH 299/521] [PyTorch][Fused Attn] Add support for cuDNN to return Softmax `Stats` always and `Max` when `return_max_logit=True` (#2677) * cudnn now returns Stats always and Max only with `return_max_logit=true` Signed-off-by: Sudhakar Singh * fix a typo that caused a bug Signed-off-by: Sudhakar Singh * update doc strings Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix more docs Signed-off-by: Sudhakar Singh * fixes from the feedback Signed-off-by: Sudhakar Singh * update cudnn-frontend to v1.19.1 Signed-off-by: Sudhakar Singh * update the cudnn frontend Signed-off-by: Sudhakar Singh * fix a wrong omission 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> --- .../fused_attn_f16_arbitrary_seqlen.cu | 64 +++++++------------ transformer_engine/common/fused_attn/utils.h | 6 +- .../include/transformer_engine/fused_attn.h | 4 +- .../pytorch/cpp_extensions/fused_attn.py | 20 +++--- .../pytorch/csrc/extensions/attention.cpp | 6 +- 5 files changed, 41 insertions(+), 59 deletions(-) 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 16aebda69f..eed6740740 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 @@ -112,7 +112,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } const DType ragged_offset_type = cudnn_runtime_version >= 90500 ? DType::kInt64 : DType::kInt32; - bool generate_stats = !return_max_logit; + bool generate_stats = true; // Always return stats try { FADescriptor_v1 descriptor{ b, @@ -343,7 +343,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( sdpa_options.set_sink_token(softmax_offset); } - std::shared_ptr Max, Sum_Exp; + std::shared_ptr Max; if (use_ragged_stats) { offset_stats = mha_graph->tensor(fe::graph::Tensor_attributes() @@ -357,19 +357,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( .set_name("Max") .set_dim({b, h, s_q, 1}) .set_data_type(fe::DataType_t::FLOAT)); - Sum_Exp = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Sum_Exp") - .set_dim({b, h, s_q, 1}) - .set_data_type(fe::DataType_t::FLOAT)); if (use_ragged_stats) { Max->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - Sum_Exp->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); } else { Max->set_stride({h * s_q, s_q, 1, 1}); - Sum_Exp->set_stride({h * s_q, s_q, 1, 1}); } sdpa_options.set_logit_max(Max); - sdpa_options.set_score_sum_exp(Sum_Exp); } auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); @@ -387,13 +380,11 @@ void fused_attn_arbitrary_seqlen_fwd_impl( O->set_ragged_offset(offset_o); } - if (!return_max_logit) { - Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); - if (use_ragged_stats) { - Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); - } else { - Stats->set_stride({h * s_q, s_q, 1, 1}); - } + Stats->set_output(true).set_data_type(fe::DataType_t::FLOAT).set_dim({b, h, s_q, 1}); + if (is_ragged_q && cudnn_runtime_version >= 90600) { + Stats->set_stride({h * s_q, 1, h, 1}).set_ragged_offset(offset_stats); + } else { + Stats->set_stride({h * s_q, s_q, 1, 1}); } std::tuple, // Q @@ -403,7 +394,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::shared_ptr> // O key_tensors_tuple = std::make_tuple(Q, K, V, attn_scale, O); auto Stats_tuple = - generate_stats ? std::make_tuple(Stats, nullptr) : std::make_tuple(Max, Sum_Exp); + return_max_logit ? std::make_tuple(Stats, Max) : std::make_tuple(Stats, nullptr); 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); @@ -1137,6 +1128,16 @@ void fused_attn_arbitrary_seqlen_fwd( size_t i = 0; if (Aux_CTX_Tensors->size == 0) { const auto cudnn_runtime_version = cudnnGetVersion(); + + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_S->data.dptr = nullptr; + if (q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) { + output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; + } else { + output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + } + output_S->data.dtype = DType::kFloat32; + if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_Max->data.dptr = nullptr; @@ -1147,25 +1148,6 @@ void fused_attn_arbitrary_seqlen_fwd( output_Max->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; } output_Max->data.dtype = DType::kFloat32; - Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_Sum_Exp->data.dptr = nullptr; - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { - output_Sum_Exp->data.shape = {num_tokens_q, num_attn_heads, 1}; - } else { - output_Sum_Exp->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_Sum_Exp->data.dtype = DType::kFloat32; - } else { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_S->data.dptr = nullptr; - if ((q_format == NVTE_QKV_Format::NVTE_THD && cudnn_runtime_version >= 90600) && - (sm_arch_ != 120)) { - output_S->data.shape = {num_tokens_q, num_attn_heads, 1}; - } else { - output_S->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - } - output_S->data.dtype = DType::kFloat32; } Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); @@ -1189,14 +1171,12 @@ void fused_attn_arbitrary_seqlen_fwd( Aux_CTX_Tensors->size = i; } else if (Aux_CTX_Tensors->size >= 2) { + Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrS1 = output_S->data.dptr; + if (return_max_logit) { Tensor *output_Max = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_Max->data.dptr; - Tensor *output_Sum_Exp = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS2 = output_Sum_Exp->data.dptr; - } else { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrS1 = output_S->data.dptr; + devPtrS2 = output_Max->data.dptr; } Tensor *output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = rng_state->data.dptr; diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 08a56cda6b..1ec1616c4a 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -118,7 +118,7 @@ struct FADescriptor_v1 { cudnn_frontend::DataType_t o_tensor_type; cudnn_frontend::DataType_t do_tensor_type; cudnn_frontend::DataType_t dqkv_tensor_type; - bool generate_max_sum_exp; + bool return_max_logit; 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, @@ -126,7 +126,7 @@ struct FADescriptor_v1 { bias_skv, attnScale, isTraining, dropoutProbability, layout, 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, generate_max_sum_exp) < + 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, @@ -134,7 +134,7 @@ struct FADescriptor_v1 { 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.generate_max_sum_exp); + 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 8169bf22e2..8d9adeb620 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -206,7 +206,7 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout); * \param[in] head_dim_v The head dimension of V. * \param[in] window_size_left Sliding window size (the left half). * \param[in] window_size_right Sliding window size (the right half). - * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. + * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] deterministic Whether determinism is required or not. */ @@ -269,7 +269,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] max_seqlen_kv Max sequence length used for computing for K and V. * it may be >= max(seqlen_kv_i) for i=0,...batch_size-1. * \param[in] is_training Whether this is in training mode or inference. - * \param[in] return_max_logit Whether to produce Max and Sum_Exp, or Stats. + * \param[in] return_max_logit Whether to produce Max along with Stats. * \param[in] cuda_graph Whether cuda graph capture is enabled or not. * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 58cfe98d72..7653296c78 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -353,12 +353,16 @@ def fused_attn_fwd( if return_max_logit: qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] - # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Max [tq, h, 1], Sum_Exp [tq, h, 1] - # thd (older cuDNN runtimes or sm120): output_tensors: out [tq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] - # bshd: output_tensors: out [b, sq, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] - # sbhd: output_tensors: out [sq, b, h, d], Max [b, h, sq, 1], Sum_Exp [b, h, sq, 1] - stats = output_tensors[1] + torch.log(output_tensors[2]) - max_tensor = output_tensors[1] + # thd (newer cuDNN runtimes, non-sm120): output_tensors: out [tq, h, d], Stats [tq, h, 1], Max [tq, h, 1] + # thd (older cuDNN runtimes or sm120): output_tensors: out [tq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] + # bshd: output_tensors: out [b, sq, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] + # sbhd: output_tensors: out [sq, b, h, d], Stats [b, h, sq, 1], Max [b, h, sq, 1] + aux_ctx_tensors = [output_tensors[1]] + list( + output_tensors[3:] + ) # Stats + rng_state + optional tensors + max_tensor = output_tensors[2] + amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) + if qkv_format == "thd" and max_tensor.ndim == 4: # For THD on older cuDNN runtimes or THD on sm120, stats can be [b, h, sq, 1] with padded # sequence positions. Exclude those padded positions when computing max_logit. @@ -366,11 +370,9 @@ def fused_attn_fwd( sq_idx = torch.arange(max_tensor.shape[2], device=max_tensor.device).view(1, 1, -1, 1) valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) max_tensor = max_tensor.masked_fill(~valid, float("-inf")) - amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) + # Max -> max_logit [h] max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=output_tensors[0].dtype) - aux_ctx_tensors = [stats] - aux_ctx_tensors.extend(output_tensors[3:]) return output_tensors[0], aux_ctx_tensors, max_logit # out, aux_ctx_tensors diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index bf62db8c33..ff60bb87bb 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -259,16 +259,16 @@ std::vector fused_attn_fwd( // 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: Max [b, h, sq, 1], Sum_Exp [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] size_t i = 0; at::Tensor output_tensor; - // intermediate softmax tensor, S or M + // intermediate softmax tensor, S or M (for fp8) 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 has an additional softmax stats tensor, ZInv; return_max_logit=true has an additional Sum_Exp 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) { output_tensor = allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), From 15cf65a70f19d71920f3a4647826b4ac92d0fd47 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 25 Mar 2026 14:41:18 -0400 Subject: [PATCH 300/521] Upgrade cuDNN FE to v1.21.0 (#2799) Move cuDNN FE to v1.21.0 Signed-off-by: Kirthi Shankar Sivamani --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index d33027a41a..7b9b711c22 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit d33027a41a93af9c85f089c6364ab415fce98982 +Subproject commit 7b9b711c22b6823e87150213ecd8449260db8610 From f4debf6648a080c47eeb2213a3a040b4b2638adb Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Mon, 30 Mar 2026 16:36:47 -0700 Subject: [PATCH 301/521] [JAX] Add warning if using BSHD and max_segments_per_seq > 1 (#2796) * Add warning if using BSHD and max_segments_per_seq > 1 Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update transformer_engine/jax/attention.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Update transformer_engine/jax/attention.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> * Remove warning test Signed-off-by: Jeremy Berchtold --------- 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: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> --- transformer_engine/jax/attention.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 765cf2872f..99817f0657 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -1436,6 +1436,15 @@ def fused_attn( context_parallel_axis=context_parallel_axis, softmax_offset=softmax_offset, ) + if max_segments_per_seq > 1 and not qkv_layout.is_thd(): + warnings.warn( + f"max_segments_per_seq={max_segments_per_seq} is set but qkv_layout={qkv_layout} is " + "not a THD layout. max_segments_per_seq > 1 only applies when using THD layouts " + "(e.g. QKVLayout.T3HD, QKVLayout.THD_T2HD, QKVLayout.THD_THD_THD) for sequence " + "packing.", + UserWarning, + stacklevel=2, + ) output = _fused_attn( qkv, bias, From bce4181a7dc8710b739fad82bc652820a78b48da Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Wed, 1 Apr 2026 11:47:50 -0700 Subject: [PATCH 302/521] [JAX] Grouped GEMM Refactor to use first_dims and last_dims (#2749) * Refactor to group_sizes per tensor Signed-off-by: Jeremy Berchtold * Support first_dims and last_dims instead of a single group_sizes per tensor Signed-off-by: Jeremy Berchtold * Refactor GMM FFIs to store static attrs as structs Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cleanup C++ v2 FFI Signed-off-by: Jeremy Berchtold * Fix int64 workspace usage Signed-off-by: Jeremy Berchtold * Address greptile comments Signed-off-by: Jeremy Berchtold * Refactor wgrad-specific checks to be generic for GMM in gemm.py Signed-off-by: Jeremy Berchtold * Refactor XLA FFI struct setup Signed-off-by: Jeremy Berchtold * Fix edge case in TE v1 GMM Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix issues on Hopper Signed-off-by: Jeremy Berchtold * Refactor Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address comments Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Lint Signed-off-by: Jeremy Berchtold * Fixes for Hopper Signed-off-by: Jeremy Berchtold * Address review comments Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Grouped quantization test fixes Signed-off-by: Jeremy Berchtold --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_custom_call_compute.py | 34 +- transformer_engine/jax/cpp_extensions/gemm.py | 436 +++++++++++++----- .../jax/cpp_extensions/quantization.py | 48 +- transformer_engine/jax/csrc/extensions.h | 50 ++ .../jax/csrc/extensions/gemm.cpp | 414 ++++++++--------- transformer_engine/jax/dense.py | 229 +++------ .../jax/quantize/dequantizer.py | 36 +- transformer_engine/jax/quantize/quantizer.py | 18 +- .../jax/quantize/scaling_modes.py | 21 +- transformer_engine/jax/quantize/tensor.py | 180 ++++++-- 10 files changed, 831 insertions(+), 635 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 613aefc178..ddb74fd636 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -36,6 +36,7 @@ ScaledTensor1x, ScaledTensor2x, GroupedScaledTensor1x, + GroupedNoScaleTensor, ScalingMode, QuantizerFactory, QuantizeLayout, @@ -150,8 +151,13 @@ def assert_dequantized_grouped_scaled_tensor( a: Union[GroupedScaledTensor1x, ScaledTensor2x], b: jnp.ndarray ): if isinstance(a, GroupedScaledTensor1x): - assert a.group_sizes.sum() == b.shape[0] - b = jnp.split(b, jnp.cumulative_sum(a.group_sizes)[:-1], axis=0) + group_sizes = ( + a.first_dims + if a.first_dims is not None + else jnp.ones(a.original_shape[0], dtype=jnp.int32) + ) + assert group_sizes.sum() == b.shape[0] + b = jnp.split(b, jnp.cumulative_sum(group_sizes)[:-1], axis=0) dq_a = a.dequantize() for dq_a_i, b_i in zip(dq_a, b): if len(dq_a_i) == 0: @@ -1787,13 +1793,18 @@ def test_grouped_gemm_fp16(self, dtype, input_shape, layout): ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims) # jitting grouped_gemm + lhs_tensor = GroupedNoScaleTensor( + data=lhs, amax=None, first_dims=group_sizes, last_dims=None, original_shape=lhs.shape + ) + rhs_tensor = GroupedNoScaleTensor( + data=rhs, amax=None, first_dims=None, last_dims=None, original_shape=rhs.shape + ) prim_out = jax.jit( tex.grouped_gemm, static_argnames=("contracting_dims", "use_async_d2h_group_sizes") )( - lhs, - rhs, - group_sizes, - contracting_dims, + lhs_tensor, + rhs_tensor, + contracting_dims=contracting_dims, use_async_d2h_group_sizes=True, ) @@ -1825,8 +1836,17 @@ def test_grouped_gemm_fp8(self, fwd_bwd_dtype, scaling_mode, input_shape, layout ) ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims) + lhs_tensor = GroupedNoScaleTensor( + data=lhs, amax=None, first_dims=group_sizes, last_dims=None, original_shape=lhs.shape + ) + rhs_tensor = GroupedNoScaleTensor( + data=rhs, amax=None, first_dims=None, last_dims=None, original_shape=rhs.shape + ) prim_out = jax.jit(tex.grouped_gemm, static_argnames=("contracting_dims",))( - lhs, rhs, group_sizes, contracting_dims, quantizer_set=quantizer_set + lhs_tensor, + rhs_tensor, + contracting_dims=contracting_dims, + quantizer_set=quantizer_set, ) allclose_dtype = jnp.float8_e4m3fn diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index aaf8e8ecea..aaec5affa8 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -37,6 +37,7 @@ ScaledTensor1x, ScaledTensor2x, GroupedScaledTensor1x, + GroupedNoScaleTensor, ScalingMode, Quantizer, GroupedQuantizer, @@ -73,12 +74,14 @@ # Cache whether the CUDA-graphable grouped GEMM implementation is available at import time. # Calling get_grouped_gemm_setup_workspace_size raises a RuntimeError mentioning "cublas" when # compiled against cuBLAS < 13.2, in which case the cuda-graphable path is unavailable. +_v2_grouped_gemm_available_reason = "" try: get_grouped_gemm_setup_workspace_size(1) _v2_grouped_gemm_available = True except RuntimeError as e: if "cublas" in str(e).lower(): _v2_grouped_gemm_available = False + _v2_grouped_gemm_available_reason = str(e) else: raise @@ -1392,17 +1395,47 @@ def impl( register_primitive(GroupedGemmCopySizesPrimitive) +def _assert_grouped_gemm_dims_shapes( + lhs_first_dims_aval, + lhs_last_dims_aval, + rhs_first_dims_aval, + rhs_last_dims_aval, + out_first_dims_aval, + out_last_dims_aval, + num_groups: int, +) -> None: + """Assert that all non-empty *_dims arrays have exactly num_groups elements. + + rhs_first_dims / rhs_last_dims describe the ragged contracting K dimension. + K totals need not fill the entire buffer (padding is allowed), so only the + array length is checked, not the per-group sum. + """ + for name, aval in [ + ("lhs_first_dims", lhs_first_dims_aval), + ("lhs_last_dims", lhs_last_dims_aval), + ("out_first_dims", out_first_dims_aval), + ("out_last_dims", out_last_dims_aval), + ("rhs_first_dims", rhs_first_dims_aval), + ("rhs_last_dims", rhs_last_dims_aval), + ]: + if aval.size > 0: + assert ( + aval.size == num_groups + ), f"grouped GEMM {name} has size {aval.size}, expected num_groups={num_groups}" + + class GroupedGemmPrimitive(BasePrimitive): """ Primitive for grouped GEMM using nvte_multi_tensor_gemm (supports all scaling modes) or nvte_grouped_gemm (supporting BF16). """ - # args = lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, group_sizes, group_offset, unused_placeholder name = "te_grouped_gemm_ffi" - # args = lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, group_sizes, alpha, beta + # args = lhs_data, lhs_scale_inv, rhs_data, rhs_scale_inv, bias, + # lhs_first_dims, lhs_last_dims, rhs_first_dims, rhs_last_dims, + # out_first_dims, out_last_dims, alpha, beta name_graph_safe = "te_grouped_gemm_v2_ffi" multiple_results = True - impl_static_args = (8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18) + impl_static_args = (13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26) inner_primitive = None outer_primitive = None @@ -1413,53 +1446,85 @@ def abstract( rhs_data_aval, rhs_scale_inv_aval, bias_aval, - group_sizes_aval, + lhs_first_dims_aval, + lhs_last_dims_aval, + rhs_first_dims_aval, + rhs_last_dims_aval, + out_first_dims_aval, + out_last_dims_aval, *additional_args, # group_offset_aval, unused_placeholder OR alpha_aval, beta_aval - M, - N, - K, lhs_is_trans, rhs_is_trans, scaling_mode, out_dtype, has_bias, - is_grouped_dense_wgrad, use_async_d2h_group_sizes, use_v2_ffi, + lhs_axis_boundary, + rhs_axis_boundary, + out_shape, + lhs_left_size, + lhs_right_size, + rhs_left_size, + rhs_right_size, ): """ Grouped GEMM operation. Args: - lhs_data: Left-hand side input matrix data, 1D flattened array + lhs_data: Left-hand side input matrix data (may be 1D for quantized) lhs_scale_inv: Left-hand side input scale_inv matrix, 1D flattened array - rhs_data: Right-hand side input matrix data, 1D flattened array + rhs_data: Right-hand side input matrix data (may be 1D for quantized) rhs_scale_inv: Right-hand side input scale_inv matrix, 1D flattened array bias: Bias matrix of shape (G, N) - group_sizes: 1D array containing the sizes of each group + lhs_first_dims: (G,) int32 if lhs first-dim is ragged, else empty (0,) sentinel + rhs_first_dims: (G,) int32 if rhs first-dim is ragged (wgrad), else empty (0,) sentinel + out_first_dims: (G,) int32 if output first-dim is ragged, else empty (0,) sentinel additional_args: Either * group_offsets: 1D array containing offsets for each group (not yet implemented) OR * alpha: 1D array of shape (G,) containing alpha values for each group * beta: 1D array of shape (G,) containing beta values for each group - M: Number of rows in the output matrix - N: Number of columns in the output matrix - K: Number of columns in the left-hand side matrix lhs_is_trans: Boolean indicating if the left-hand side matrix is transposed rhs_is_trans: Boolean indicating if the right-hand side matrix is transposed scaling_mode: Scaling mode for the GEMM operations out_dtype: Data type of the output tensors has_bias: Boolean indicating if bias tensors are provided - is_grouped_dense_wgrad: Boolean indicating if this is a grouped dense wgrad operation - where both lhs and rhs are 2D matrices and output is (G, M, N) + out_shape: Pre-computed output shape tuple + lhs_left_size: Product of lhs dims before axis_boundary + lhs_right_size: Product of lhs dims after axis_boundary + rhs_left_size: Product of rhs dims before axis_boundary + rhs_right_size: Product of rhs dims after axis_boundary Returns: A jnp.ndarray containing the result of the grouped GEMM operation """ - del lhs_data_aval, rhs_data_aval, bias_aval - del K, lhs_is_trans, rhs_is_trans, has_bias, use_async_d2h_group_sizes + del lhs_data_aval, rhs_data_aval + del lhs_is_trans, rhs_is_trans + del lhs_axis_boundary, rhs_axis_boundary + del lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size + del bias_aval + del has_bias, use_async_d2h_group_sizes + + num_groups = ( + lhs_first_dims_aval.size + or lhs_last_dims_aval.size + or rhs_first_dims_aval.size + or rhs_last_dims_aval.size + or out_first_dims_aval.size + or out_last_dims_aval.size + or additional_args[0].size # alpha (V2) has size G; group_offset (legacy) has size >= 1 + ) - num_groups = group_sizes_aval.size + _assert_grouped_gemm_dims_shapes( + lhs_first_dims_aval, + lhs_last_dims_aval, + rhs_first_dims_aval, + rhs_last_dims_aval, + out_first_dims_aval, + out_last_dims_aval, + num_groups, + ) cublas_workspace_aval = jax.core.ShapedArray( shape=( @@ -1470,9 +1535,6 @@ def abstract( dtype=jnp.uint8, ) - out_shape = (M, N) - if is_grouped_dense_wgrad: - out_shape = (num_groups, M, N) out_aval = jax.core.ShapedArray(shape=out_shape, dtype=out_dtype) if use_v2_ffi: @@ -1480,7 +1542,24 @@ def abstract( shape=(get_grouped_gemm_setup_workspace_size(num_groups),), dtype=jnp.uint8 ) # Temporary buffer for int32 -> int64 conversion of group_sizes on device. - int64_workspace_size = num_groups * jnp.dtype(jnp.int64).itemsize + # Each non-empty *_dims buffer needs its own slot of num_groups int64 elements so that + # make_grouped_tensor can write to a distinct region per ragged dimension. Allocate + # exactly as many slots as there are non-empty buffers (minimum 1 to avoid zero-size). + num_ragged_dim_buffers = sum( + 1 + for aval in [ + lhs_first_dims_aval, + lhs_last_dims_aval, + rhs_first_dims_aval, + rhs_last_dims_aval, + out_first_dims_aval, + out_last_dims_aval, + ] + if aval.size > 0 + ) + int64_workspace_size = ( + max(num_ragged_dim_buffers, 1) * num_groups * jnp.dtype(jnp.int64).itemsize + ) int64_workspace_aval = jax.core.ShapedArray( shape=(int64_workspace_size,), dtype=jnp.uint8 ) @@ -1545,45 +1624,52 @@ def outer_abstract(*args, **kwargs): def lowering( ctx, *args, - M, - N, - K, lhs_is_trans, rhs_is_trans, scaling_mode, out_dtype, has_bias, - is_grouped_dense_wgrad, use_async_d2h_group_sizes, use_v2_ffi, + lhs_axis_boundary, + rhs_axis_boundary, + out_shape, + lhs_left_size, + lhs_right_size, + rhs_left_size, + rhs_right_size, ): - del out_dtype + del out_dtype, out_shape # Python-only; not forwarded to C++ if use_v2_ffi: ffi_name = GroupedGemmPrimitive.name_graph_safe return jax.ffi.ffi_lowering(ffi_name)( ctx, *args, - M=M, - N=N, - K=K, lhs_is_trans=lhs_is_trans, rhs_is_trans=rhs_is_trans, scaling_mode=scaling_mode.value, - is_grouped_dense_wgrad=is_grouped_dense_wgrad, + lhs_axis_boundary=lhs_axis_boundary, + rhs_axis_boundary=rhs_axis_boundary, + lhs_left_size=lhs_left_size, + lhs_right_size=lhs_right_size, + rhs_left_size=rhs_left_size, + rhs_right_size=rhs_right_size, ) ffi_name = GroupedGemmPrimitive.name return jax.ffi.ffi_lowering(ffi_name)( ctx, *args, - M=M, - N=N, - K=K, lhs_is_trans=lhs_is_trans, rhs_is_trans=rhs_is_trans, scaling_mode=scaling_mode.value, has_bias=has_bias, - is_grouped_dense_wgrad=is_grouped_dense_wgrad, use_async_d2h_group_sizes=use_async_d2h_group_sizes, + lhs_axis_boundary=lhs_axis_boundary, + rhs_axis_boundary=rhs_axis_boundary, + lhs_left_size=lhs_left_size, + lhs_right_size=lhs_right_size, + rhs_left_size=rhs_left_size, + rhs_right_size=rhs_right_size, ) @staticmethod @@ -1593,20 +1679,28 @@ def impl( rhs_data, rhs_scale_inv, bias, - group_sizes, + lhs_first_dims, + lhs_last_dims, + rhs_first_dims, + rhs_last_dims, + out_first_dims, + out_last_dims, additional_arg_0, # group_offset (non-graph-safe) OR alpha (graph-safe) additional_arg_1, # unused placeholder (non-graph-safe) OR beta (graph-safe) - M, - N, - K, lhs_is_trans, rhs_is_trans, scaling_mode, out_dtype, has_bias, - is_grouped_dense_wgrad, use_async_d2h_group_sizes, use_v2_ffi, + lhs_axis_boundary, + rhs_axis_boundary, + out_shape, + lhs_left_size, + lhs_right_size, + rhs_left_size, + rhs_right_size, ): if GroupedGemmPrimitive.inner_primitive is None: raise RuntimeError("GroupedGemmPrimitive.inner_primitive has not been registered") @@ -1620,19 +1714,27 @@ def impl( rhs_data, rhs_scale_inv, bias, - group_sizes, + lhs_first_dims, + lhs_last_dims, + rhs_first_dims, + rhs_last_dims, + out_first_dims, + out_last_dims, *additional_args, - M=M, - N=N, - K=K, lhs_is_trans=lhs_is_trans, rhs_is_trans=rhs_is_trans, scaling_mode=scaling_mode, out_dtype=out_dtype, has_bias=has_bias, - is_grouped_dense_wgrad=is_grouped_dense_wgrad, use_async_d2h_group_sizes=use_async_d2h_group_sizes, use_v2_ffi=use_v2_ffi, + lhs_axis_boundary=lhs_axis_boundary, + rhs_axis_boundary=rhs_axis_boundary, + out_shape=out_shape, + lhs_left_size=lhs_left_size, + lhs_right_size=lhs_right_size, + rhs_left_size=rhs_left_size, + rhs_right_size=rhs_right_size, ) return (out,) @@ -1922,6 +2024,12 @@ def grouped_gemm_copy_group_sizes( return out +@cache +def _should_enforce_v2_grouped_gemm() -> bool: + """Read NVTE_JAX_ENFORCE_V2_GROUPED_GEMM once per process (cached).""" + return os.getenv("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "0") == "1" + + def _can_use_v2_grouped_gemm( scaling_mode: ScalingMode, dtype: jnp.dtype, @@ -1933,21 +2041,42 @@ def _can_use_v2_grouped_gemm( # feature-compatible with the main branch. # Bias can be supported in a kernel or in pure-JAX in the future. + enforce_v2_gmm = _should_enforce_v2_grouped_gemm() + if not _v2_grouped_gemm_available: + if enforce_v2_gmm: + raise RuntimeError( + "The TE V2 grouped GEMM is not available but NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is" + " enabled. The reason for V2 grouped GEMM not being available:" + f" {_v2_grouped_gemm_available_reason}" + ) return False # nvte_grouped_gemm (the v2 kernel) requires SM100+ (Blackwell or newer). # Fall back to the v1 path on SM90 (Hopper) and older architectures. if get_device_compute_capability(0) < 100: + if enforce_v2_gmm: + raise RuntimeError( + "The TE V2 grouped GEMM requires SM100+ (Blackwell or newer) but current device" + f" compute capability of GPU 0 is {get_device_compute_capability(0)} and" + " NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is enabled." + ) return False - return scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16 and not has_bias + if scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16 and not has_bias: + return True + + if enforce_v2_gmm: + raise RuntimeError( + "The TE V2 grouped GEMM currently only supports BF16 with no quantization recipe and" + f" without bias, but received {scaling_mode=}, {dtype=}, {has_bias=}" + ) + return False def grouped_gemm( - lhs: Union[jnp.ndarray, GroupedScaledTensor1x], - rhs: Union[jnp.ndarray, GroupedScaledTensor1x], - group_sizes: jnp.ndarray, + lhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + rhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((1,), (2,)), bias: jnp.ndarray = None, precision: jax.lax.Precision = jax.lax.Precision.DEFAULT, @@ -1960,9 +2089,8 @@ def grouped_gemm( Grouped GEMM operation. Args: - lhs: Left-hand side input matrix, can be a jnp.ndarray or GroupedScaledTensor1x - rhs: Right-hand side input matrix, can be a jnp.ndarray or GroupedScaledTensor1x - group_sizes: 1D array containing the sizes of each group + lhs: Left-hand side input matrix, GroupedNoScaleTensor or GroupedScaledTensor1x + rhs: Right-hand side input matrix, GroupedNoScaleTensor or GroupedScaledTensor1x contracting_dims: Tuple of two sequences representing the contracting dimensions bias: Bias tensor of shape (G, N) precision: JAX precision for the GEMM operation @@ -1972,49 +2100,74 @@ def grouped_gemm( Returns: A jnp.ndarray containing the result of the grouped GEMM operation - - Note: - Tested shapes: - lhs: [M, K] or [K, N] - rhs: [G, N, K] or [G, K, N] or [G * K, N] or [N, G * K] """ # TODO(Phuong): implement the precision del precision - if isinstance(lhs, jnp.ndarray): - if not isinstance(rhs, jnp.ndarray): - raise TypeError( - f"Expected rhs to be jnp.ndarray when lhs is jnp.ndarray, but got type={type(rhs)}" - ) - out_dtype = lhs.dtype - lhs_shape = lhs.shape - rhs_shape = rhs.shape - lhs_data = lhs - rhs_data = rhs - lhs_scale_inv = rhs_scale_inv = jnp.empty((0,), jnp.float32) + empty_gs = jnp.empty((0,), jnp.int32) + + # Extract data, dims, and metadata from tensor objects. + # Keep data in its original layout (may be 1D for quantized tensors) to preserve + # JAX sharding; the C++ side uses original_shape to derive m/n/k. + if isinstance(lhs, GroupedNoScaleTensor): + lhs_data = lhs.data + lhs_shape = lhs.original_shape + lhs_scale_inv = jnp.empty((0,), jnp.float32) scaling_mode = ScalingMode.NO_SCALING + out_dtype = lhs.data.dtype + lhs_first_dims = lhs.first_dims if lhs.first_dims is not None else empty_gs + lhs_last_dims = lhs.last_dims if lhs.last_dims is not None else empty_gs elif isinstance(lhs, GroupedScaledTensor1x): - if not isinstance(rhs, GroupedScaledTensor1x): - raise TypeError( - "Expected rhs to be GroupedScaledTensor1x when lhs is GroupedScaledTensor1x, but" - f" got type={type(rhs)}" - ) - out_dtype = lhs.dq_dtype lhs_shape = lhs.original_shape - rhs_shape = rhs.original_shape lhs_data = lhs.data - rhs_data = rhs.data lhs_scale_inv = lhs.scale_inv + scaling_mode = lhs.scaling_mode + out_dtype = lhs.dq_dtype + lhs_first_dims = lhs.first_dims if lhs.first_dims is not None else empty_gs + lhs_last_dims = lhs.last_dims if lhs.last_dims is not None else empty_gs + else: + raise TypeError( + f"lhs must be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(lhs)}" + ) + + if isinstance(rhs, GroupedNoScaleTensor): + rhs_data = rhs.data + rhs_shape = rhs.original_shape + rhs_scale_inv = jnp.empty((0,), jnp.float32) + rhs_first_dims = rhs.first_dims if rhs.first_dims is not None else empty_gs + rhs_last_dims = rhs.last_dims if rhs.last_dims is not None else empty_gs + elif isinstance(rhs, GroupedScaledTensor1x): + rhs_shape = rhs.original_shape + rhs_data = rhs.data rhs_scale_inv = rhs.scale_inv - if lhs.scaling_mode != rhs.scaling_mode: + rhs_first_dims = rhs.first_dims if rhs.first_dims is not None else empty_gs + rhs_last_dims = rhs.last_dims if rhs.last_dims is not None else empty_gs + if isinstance(lhs, GroupedScaledTensor1x) and lhs.scaling_mode != rhs.scaling_mode: raise ValueError( f"Mismatched scaling modes: lhs.scaling_mode={lhs.scaling_mode}," f" rhs.scaling_mode={rhs.scaling_mode}" ) - scaling_mode = lhs.scaling_mode + if isinstance(lhs, GroupedScaledTensor1x): + scaling_mode = lhs.scaling_mode else: - raise TypeError("Unsupported lhs type object!") + raise TypeError( + f"rhs must be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(rhs)}" + ) + + # Infer output dims from which operand has the ragged non-contracting dim. + if rhs_first_dims.size > 0 or rhs_last_dims.size > 0: + # Wgrad: rhs contracting dim is ragged → output is uniform (G prefix from num_groups) + out_first_dims = empty_gs + out_last_dims = empty_gs + elif lhs_first_dims.size > 0: + out_first_dims = lhs_first_dims + out_last_dims = empty_gs + elif lhs_last_dims.size > 0: + out_first_dims = empty_gs + out_last_dims = lhs_last_dims + else: + out_first_dims = out_last_dims = empty_gs out_dtype = preferred_element_type or out_dtype @@ -2023,26 +2176,10 @@ def grouped_gemm( lhs_is_trans = lhs_contract_dim[-1] != len(lhs_shape) - 1 lhs_flatten_axis = len(lhs_contract_dim) * (1 if lhs_is_trans else -1) - # rhs_shape [G, K, N] - rhs_is_trans = rhs_contract_dim[0] != 1 + # rhs_is_trans: K is the last dim of rhs (i.e., rhs is in "T" layout). + rhs_is_trans = rhs_contract_dim[-1] == len(rhs_shape) - 1 rhs_flatten_axis = -len(rhs_contract_dim) if rhs_is_trans else 1 + len(rhs_contract_dim) - is_grouped_dense_wgrad = False - if len(rhs_shape) == 2: - rhs_is_trans = rhs_contract_dim[0] != 0 - is_grouped_dense_wgrad = True - - # TODO(Hua): thses are for fp16 dense wgrad, any better way to handle this? - if ( - is_grouped_dense_wgrad - and not isinstance(lhs, ScaledTensor) - and not isinstance(rhs, ScaledTensor) - ): - lhs_is_trans = True - rhs_is_trans = False - lhs_flatten_axis = 1 - rhs_flatten_axis = 1 - if ( not isinstance(lhs, ScaledTensor) and not isinstance(rhs, ScaledTensor) @@ -2073,9 +2210,21 @@ def grouped_gemm( quantizer_set.kernel.q_layout = ( QuantizeLayout.ROWWISE if rhs_is_rowwise else QuantizeLayout.COLWISE ) - lhs_q = grouped_quantize(lhs, quantizer_set.x, group_sizes, lhs_flatten_axis) + active_group_sizes = next( + ( + gs + for gs in [lhs_first_dims, lhs_last_dims, rhs_first_dims, rhs_last_dims] + if gs.size > 0 + ), + empty_gs, + ) + lhs_input_data = lhs.data if isinstance(lhs, GroupedNoScaleTensor) else lhs_data + rhs_input_data = rhs.data if isinstance(rhs, GroupedNoScaleTensor) else rhs_data + lhs_q = grouped_quantize( + lhs_input_data, quantizer_set.x, active_group_sizes, lhs_flatten_axis + ) rhs_q = grouped_quantize( - rhs, quantizer_set.kernel, group_sizes=None, flatten_axis=rhs_flatten_axis + rhs_input_data, quantizer_set.kernel, group_sizes=None, flatten_axis=rhs_flatten_axis ) lhs_data = lhs_q.data rhs_data = rhs_q.data @@ -2110,38 +2259,66 @@ def grouped_gemm( lhs_contract_dim = tuple((lhs_ndim - 1 - i) % lhs_ndim for i in lhs_contract_dim) if rhs_layout_is_T: # For rhs [G, K, N], need to exclude the G dim from contract_dim - if group_sizes.size == rhs_shape[0]: + if ( + lhs_first_dims.size > 0 or lhs_last_dims.size > 0 + ): # fwd/dgrad: rhs has G as first dim rhs_contract_dim = tuple( (rhs_ndim - 1 - i) % (rhs_ndim - 1) + 1 for i in rhs_contract_dim ) else: rhs_contract_dim = tuple((rhs_ndim - 1 - i) % rhs_ndim for i in rhs_contract_dim) - # Calling GroupedGEMM Custom Call - K_lhs = math.prod(lhs_shape[i] for i in lhs_contract_dim) - K_rhs = math.prod(rhs_shape[i] for i in rhs_contract_dim) - if K_lhs != K_rhs: + # Compute N-D axis boundaries from final (post-adjustment) contracting dims. + lhs_axis_boundary = get_lhs_axis_boundary(lhs_contract_dim, lhs_is_trans) + rhs_axis_boundary = get_rhs_axis_boundary(rhs_contract_dim, rhs_is_trans) + + num_gemms = ( + lhs_first_dims.size + or lhs_last_dims.size + or rhs_first_dims.size + or rhs_last_dims.size + or out_first_dims.size + or out_last_dims.size + ) + if num_gemms == 0: raise ValueError( - f"Mismatched contracting dimensions: K_lhs={K_lhs}, K_rhs={K_rhs} (from" - f" lhs_shape={lhs_shape}, rhs_shape={rhs_shape})" + "grouped_gemm requires at least one non-empty dimension array. " + "Ensure lhs or rhs tensor objects carry first_dims or last_dims." ) - M = math.prod(_calculate_remaining_shape(lhs_shape, lhs_contract_dim)) - N = math.prod(_calculate_remaining_shape(rhs_shape, rhs_contract_dim)[1:]) # Exclude G - if is_grouped_dense_wgrad: - N = math.prod(_calculate_remaining_shape(rhs_shape, rhs_contract_dim)) + # Pre-compute collapsed 2D sizes from original N-D shapes. + # These are static Python ints passed as primitive parameters (must be hashable). + lhs_left_size = math.prod(lhs_shape[:lhs_axis_boundary]) + lhs_right_size = math.prod(lhs_shape[lhs_axis_boundary:]) + rhs_left_size = math.prod(rhs_shape[:rhs_axis_boundary]) + rhs_right_size = math.prod(rhs_shape[rhs_axis_boundary:]) + + # Pre-compute output shape from N-D input shapes (static Python ints). + if lhs_is_trans: + lhs_non_contracting = lhs_shape[lhs_axis_boundary:] else: - if group_sizes.size != rhs_shape[0]: - raise ValueError( - "Expected group_sizes.size == rhs_shape[0], but got" - f" group_sizes.size={group_sizes.size}, rhs_shape[0]={rhs_shape[0]}" - ) + lhs_non_contracting = lhs_shape[:lhs_axis_boundary] + if rhs_is_trans: + if rhs_first_dims.size > 0 or rhs_last_dims.size > 0: + # wgrad: rhs (e.g. grad_T of shape (N, M)) has no G batch dim; include all dims + rhs_non_contracting = tuple(rhs_shape[d] for d in range(rhs_axis_boundary)) + else: + # fwd/dgrad: rhs (e.g. kernel_T of shape (G, N, K)) has G batch dim at dim 0; skip it + rhs_non_contracting = tuple(rhs_shape[d] for d in range(rhs_axis_boundary) if d != 0) + else: + rhs_non_contracting = rhs_shape[rhs_axis_boundary:] + if rhs_first_dims.size > 0 or rhs_last_dims.size > 0: + out_shape = (num_gemms, *lhs_non_contracting, *rhs_non_contracting) + else: + out_shape = (*lhs_non_contracting, *rhs_non_contracting) has_bias = bias is not None - if has_bias and bias.shape != (group_sizes.size, N): - raise ValueError( - f"Expected bias.shape=({group_sizes.size}, {N}), but got bias.shape={bias.shape}" - ) + if has_bias: + N_dim = math.prod(rhs_non_contracting) + assert bias.shape == ( + num_gemms, + N_dim, + ), f"bias shape {bias.shape} does not match expected shape {(num_gemms, N_dim)}" bias = jnp.empty((), jnp.float32) if bias is None else bias if group_offset is not None: @@ -2153,7 +2330,6 @@ def grouped_gemm( use_v2_ffi = _can_use_v2_grouped_gemm(scaling_mode, lhs_data.dtype, has_bias) if use_v2_ffi: - num_gemms = group_sizes.shape[0] additional_arg_0 = jnp.ones((num_gemms,), jnp.float32) # alpha additional_arg_1 = jnp.zeros((num_gemms,), jnp.float32) # beta else: @@ -2166,19 +2342,27 @@ def grouped_gemm( rhs_data, rhs_scale_inv, bias, - group_sizes, + lhs_first_dims, + lhs_last_dims, + rhs_first_dims, + rhs_last_dims, + out_first_dims, + out_last_dims, additional_arg_0, additional_arg_1, - M=M, - N=N, - K=K_lhs, lhs_is_trans=lhs_is_trans, rhs_is_trans=rhs_is_trans, scaling_mode=scaling_mode.value, out_dtype=out_dtype, has_bias=has_bias, - is_grouped_dense_wgrad=is_grouped_dense_wgrad, use_async_d2h_group_sizes=use_async_d2h_group_sizes, use_v2_ffi=use_v2_ffi, + lhs_axis_boundary=lhs_axis_boundary, + rhs_axis_boundary=rhs_axis_boundary, + out_shape=tuple(int(d) for d in out_shape), + lhs_left_size=int(lhs_left_size), + lhs_right_size=int(lhs_right_size), + rhs_left_size=int(rhs_left_size), + rhs_right_size=int(rhs_right_size), ) return out diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index bf4e833c89..a3d363e42a 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -43,6 +43,7 @@ ScalingMode, compute_scale_from_amax, NoScaleTensor, + GroupedNoScaleTensor, get_rht_matrix, QuantizeLayout, ) @@ -1001,7 +1002,6 @@ class GroupedQuantizePrimitive(BasePrimitive): 5, 6, 7, - 8, ) # out_dtype, scaling_mode, q_layout, flatten_axis, scale_dtype inner_primitive = None outer_primitive = None @@ -1016,7 +1016,6 @@ def abstract( scaling_mode, q_layout, flatten_axis, - group_axis, scale_dtype, ): """ @@ -1038,7 +1037,6 @@ def abstract( ).get_grouped_scale_shape_2x( x_aval.shape, group_sizes_aval.size, - group_axis, is_padded=True, flatten_axis=flatten_axis, ) @@ -1099,7 +1097,6 @@ def lowering( scaling_mode, q_layout, flatten_axis, - group_axis, scale_dtype, ): """ @@ -1110,7 +1107,6 @@ def lowering( assert x_aval.dtype in [jnp.float32, jnp.float16, jnp.bfloat16] assert scale_aval.dtype == jnp.float32 assert group_sizes_aval.dtype == jnp.int32 - assert group_axis == 0 return ffi.ffi_lowering(GroupedQuantizePrimitive.name)( ctx, x, @@ -1130,7 +1126,6 @@ def impl( scaling_mode, q_layout, flatten_axis, - group_axis, scale_dtype, ): """ @@ -1151,7 +1146,6 @@ def impl( scaling_mode=scaling_mode, q_layout=q_layout, flatten_axis=flatten_axis, - group_axis=group_axis, scale_dtype=scale_dtype, ) return (rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax) @@ -1164,20 +1158,18 @@ def grouped_quantize( x: jnp.ndarray, quantizer: GroupedQuantizer, group_sizes: jnp.ndarray = None, - amax: jnp.ndarray = None, flatten_axis: int = -1, -) -> GroupedScaledTensor1x: +) -> Union[GroupedScaledTensor1x, GroupedNoScaleTensor]: """Quantize a tensor in grouped manner. This function quantizes a tensor by splitting it into groups along a specified axis and applying quantization to each group separately. The groups can be either specified - explicitly through group_sizes or automatically split along the group_axis. + explicitly through group_sizes or automatically split along axis 0. Args: x: Input tensor to quantize quantizer: The quantizer to use for quantization group_sizes: Array of ints containing the size of each group (default: None) - amax: The amax of x; if None, it is auto-generated. (default: None) flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) Returns: @@ -1185,31 +1177,34 @@ def grouped_quantize( Note: - If group_sizes is not provided, the tensor will be split into equal-sized groups - along the group_axis - - The group_axis is currently fixed to 0 + along axis 0 - The quantizer's q_layout determines whether row-wise, column-wise, or both quantization is applied """ if quantizer is None: - if isinstance(x, NoScaleTensor): + if isinstance(x, GroupedNoScaleTensor): return x - return NoScaleTensor(data=x, amax=None) + return GroupedNoScaleTensor( + data=x, + amax=None, + first_dims=group_sizes, + last_dims=None, + original_shape=x.shape, + ) # TODO(Phuong): add support for flatten_axis = -2 assert flatten_axis in ( -1, x.ndim - 1, ), f"Only flatten_axis = -1 is supported for now, got {flatten_axis}" - group_axis = 0 + ragged_first_dims = group_sizes # None if no explicit group_sizes (kernel case) if group_sizes is None: - group_sizes = jnp.ones(x.shape[group_axis], dtype=jnp.int32) + group_sizes = jnp.ones(x.shape[0], dtype=jnp.int32) if not GroupedQuantizePrimitive.enabled(): - return quantizer.quantize( - x, flatten_axis=flatten_axis, group_sizes=group_sizes, group_axis=group_axis - ) + return quantizer.quantize(x, flatten_axis=flatten_axis, group_sizes=group_sizes) n_groups = group_sizes.size original_shape = x.shape assert n_groups == len( @@ -1222,13 +1217,8 @@ def grouped_quantize( scale = scale.at[i].set(quantizer_i.scale[0]) if quantizer.scaling_mode == ScalingMode.CURRENT_TENSOR_SCALING: - if amax is not None: - row_amax = amax - else: - row_amax = jnp.max(jnp.abs(x), axis=range(group_axis + 1, x.ndim)) - segment_ids = jnp.repeat( - jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[group_axis] - ) + row_amax = jnp.max(jnp.abs(x), axis=range(1, x.ndim)) + segment_ids = jnp.repeat(jnp.arange(n_groups), group_sizes, total_repeat_length=x.shape[0]) grouped_amax = jax.ops.segment_max(row_amax, segment_ids, num_segments=n_groups) for i in range(n_groups): tmp_scale = compute_scale_from_amax(grouped_amax[i], quantizer.q_dtype, margin=0.0) @@ -1256,7 +1246,6 @@ def grouped_quantize( scaling_mode=quantizer.scaling_mode.value, q_layout=q_layout, flatten_axis=flatten_axis, - group_axis=group_axis, scale_dtype=quantizer.get_scale_dtype(), ) @@ -1280,9 +1269,8 @@ def grouped_quantize( q_layout=quantizer.q_layout, data_layout=quantizer.get_data_layout(), flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=ragged_first_dims, original_shape=original_shape, - group_axis=group_axis, ) return out diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 0fe4e99239..a74b209e4f 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -55,6 +55,32 @@ struct GemmConfig { bool use_split_accumulator; }; +struct GroupedGemmV2Config { + bool lhs_is_trans; + bool rhs_is_trans; + JAXX_Scaling_Mode scaling_mode; + int64_t lhs_axis_boundary; + int64_t rhs_axis_boundary; + int64_t lhs_left_size; + int64_t lhs_right_size; + int64_t rhs_left_size; + int64_t rhs_right_size; +}; + +struct GroupedGemmConfig { + bool lhs_is_trans; + bool rhs_is_trans; + JAXX_Scaling_Mode scaling_mode; + bool has_bias; + bool use_async_d2h_group_sizes; + int64_t lhs_axis_boundary; + int64_t rhs_axis_boundary; + int64_t lhs_left_size; + int64_t lhs_right_size; + int64_t rhs_left_size; + int64_t rhs_right_size; +}; + inline bool use_fp8(DType type) { return type == DType::kFloat8E4M3 || type == DType::kFloat8E5M2; } // Activation @@ -192,6 +218,30 @@ XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( ::xla::ffi::StructMember("rhs_transposed"), ::xla::ffi::StructMember("use_split_accumulator")); +XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( + transformer_engine::jax::GroupedGemmV2Config, ::xla::ffi::StructMember("lhs_is_trans"), + ::xla::ffi::StructMember("rhs_is_trans"), + ::xla::ffi::StructMember("scaling_mode"), + ::xla::ffi::StructMember("lhs_axis_boundary"), + ::xla::ffi::StructMember("rhs_axis_boundary"), + ::xla::ffi::StructMember("lhs_left_size"), + ::xla::ffi::StructMember("lhs_right_size"), + ::xla::ffi::StructMember("rhs_left_size"), + ::xla::ffi::StructMember("rhs_right_size")); + +XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( + transformer_engine::jax::GroupedGemmConfig, ::xla::ffi::StructMember("lhs_is_trans"), + ::xla::ffi::StructMember("rhs_is_trans"), + ::xla::ffi::StructMember("scaling_mode"), + ::xla::ffi::StructMember("has_bias"), + ::xla::ffi::StructMember("use_async_d2h_group_sizes"), + ::xla::ffi::StructMember("lhs_axis_boundary"), + ::xla::ffi::StructMember("rhs_axis_boundary"), + ::xla::ffi::StructMember("lhs_left_size"), + ::xla::ffi::StructMember("lhs_right_size"), + ::xla::ffi::StructMember("rhs_left_size"), + ::xla::ffi::StructMember("rhs_right_size")); + // ENUM_ATTR and DICT_ATTR recoding need to be registered in the global namespace XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Scaling_Mode); XLA_FFI_REGISTER_ENUM_ATTR_DECODING(transformer_engine::jax::JAXX_Score_Function); diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 2acefa2d30..0d1ef405f4 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -619,137 +619,99 @@ JAXX_GroupedTensorWrapper make_grouped_tensor(Buffer_Type const &data, return std::move(grouped_tensor_wrapper); } -// This FFI is EXPERIMENTAL and subject to change without deprecation, intended for use in JAX's internal implementation of grouped GEMM. -Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, - Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, - Buffer_Type group_sizes, Buffer_Type alpha, Buffer_Type beta, - Result_Type output, Result_Type cublas_workspace, - Result_Type setup_workspace, Result_Type int64_workspace, size_t m, - size_t n, size_t k, bool lhs_is_trans, bool rhs_is_trans, - JAXX_Scaling_Mode scaling_mode, bool is_grouped_dense_wgrad) { - // Notes on matrix layouts and transpose: - // Jax uses row-major data_layout, on entering this function, each input matrix pair: - // A: row-major [m, k] for N - [k, m] for T - // B: row-major [k, n] for N - [n, k] for T - // on exiting this function, JAX expect: - // C: row-major with size [m, n]. - // cuBLAS uses column-major data_layout, in this view, each input matrix pair: - // A: column-major with size [k, m] for T - [m, k] for N - // B: column-major with size [n, k] for T - [k, n] for N - // - // If we call cuBLAS GEMM for A * B, the output will be: - // C: column-major with size [m, n] --> row-major with size [n, m]. - // To make the output compatible with JAX, we need to swap A and B in cuBLAS GEMM call. +// V2 variant: derives data shape from the XLA buffer directly, converts group_sizes +// int32→int64 per-tensor into a dedicated slot of int64_workspace, and wires first_dims/last_dims. +// int64_offset (in int64 elements) is updated on return to the next available slot so callers can +// thread it through successive make_grouped_tensor calls without aliasing. Bounds are checked +// before each slot is used. Only NO_SCALING is supported. +JAXX_GroupedTensorWrapper make_grouped_tensor( + Buffer_Type const &data, Buffer_Type const &first_dims, Buffer_Type const &last_dims, + int64_t *int64_workspace_base, size_t int64_workspace_capacity, size_t &int64_offset, + size_t num_gemms, cudaStream_t stream, int64_t axis_boundary = -1) { + auto dims = data.dimensions(); + NVTE_CHECK(dims.size() >= 2, "grouped GEMM data buffer must be at least 2D."); + // Flatten dims at axis_boundary to produce a 2D NVTE shape. + // axis_boundary=-1 (default) collapses dims[0..N-2] → rows and keeps dims[N-1] → cols, + // preserving the prior behaviour for output buffers (e.g. [G, K, N] for wgrad). + size_t ab = (axis_boundary < 0) ? dims.size() - 1 : static_cast(axis_boundary); + NVTEShape dataShape{.data = {product(dims, 0, ab), product(dims, ab, dims.size())}, .ndim = 2}; + JAXX_GroupedTensorWrapper wrapper(JAXX_Scaling_Mode::NO_SCALING, num_gemms, dataShape); + wrapper.set_rowwise(data, std::nullopt); + if (first_dims.element_count() > 0) { + NVTE_CHECK(first_dims.element_type() == xla::ffi::DataType::S32, "group_sizes must be int32."); + NVTE_CHECK(int64_offset + num_gemms <= int64_workspace_capacity, + "int64_workspace overflow: not enough space for first_dims conversion."); + auto *slot = int64_workspace_base + int64_offset; + nvte_convert_int32_to_int64(reinterpret_cast(first_dims.untyped_data()), slot, + num_gemms, stream); + wrapper.set_group_sizes_only(slot, num_gemms, kNVTEGroupedFirstDims); + int64_offset += num_gemms; + } + if (last_dims.element_count() > 0) { + NVTE_CHECK(last_dims.element_type() == xla::ffi::DataType::S32, "group_sizes must be int32."); + NVTE_CHECK(int64_offset + num_gemms <= int64_workspace_capacity, + "int64_workspace overflow: not enough space for last_dims conversion."); + auto *slot = int64_workspace_base + int64_offset; + nvte_convert_int32_to_int64(reinterpret_cast(last_dims.untyped_data()), slot, + num_gemms, stream); + wrapper.set_group_sizes_only(slot, num_gemms, kNVTEGroupedLastDims); + int64_offset += num_gemms; + } + return wrapper; +} - // Inputs - auto lhs_ptr = reinterpret_cast(lhs_data.untyped_data()); - auto rhs_ptr = reinterpret_cast(rhs_data.untyped_data()); - auto lhs_sinv_ptr = reinterpret_cast(lhs_sinv.untyped_data()); - auto rhs_sinv_ptr = reinterpret_cast(rhs_sinv.untyped_data()); - auto lhs_dtype = convert_ffi_datatype_to_te_dtype(lhs_data.element_type()); - auto rhs_dtype = convert_ffi_datatype_to_te_dtype(rhs_data.element_type()); - auto lhs_sinv_dtype = convert_ffi_datatype_to_te_dtype(lhs_sinv.element_type()); - auto rhs_sinv_dtype = convert_ffi_datatype_to_te_dtype(rhs_sinv.element_type()); - bool has_bias = product(bias.dimensions()) > 0; - auto bias_ptr = has_bias ? reinterpret_cast(bias.untyped_data()) : nullptr; - auto bias_dtype = convert_ffi_datatype_to_te_dtype(bias.element_type()); +// Returns num_gemms from the first non-empty per-tensor group_sizes buffer, +// falling back to the element count of alpha for the uniform-batch case. +size_t grouped_gemm_num_gemms(Buffer_Type const &lhs_first_dims, Buffer_Type const &lhs_last_dims, + Buffer_Type const &rhs_first_dims, Buffer_Type const &rhs_last_dims, + Buffer_Type const &out_first_dims, Buffer_Type const &out_last_dims, + Buffer_Type const &alpha) { + if (lhs_first_dims.element_count() > 0) { + return lhs_first_dims.element_count(); + } else if (lhs_last_dims.element_count() > 0) { + return lhs_last_dims.element_count(); + } else if (rhs_first_dims.element_count() > 0) { + return rhs_first_dims.element_count(); + } else if (rhs_last_dims.element_count() > 0) { + return rhs_last_dims.element_count(); + } else if (out_first_dims.element_count() > 0) { + return out_first_dims.element_count(); + } else if (out_last_dims.element_count() > 0) { + return out_last_dims.element_count(); + } else { + return alpha.element_count(); // uniform batch: no ragged tensor + } +} + +} // namespace jax +} // namespace transformer_engine - NVTE_CHECK(group_sizes.dimensions().size() == 1); - size_t num_gemms = group_sizes.dimensions()[0]; +namespace transformer_engine { +namespace jax { - // Convert int32 group_sizes to int64 into the dedicated output buffer. - NVTE_CHECK(group_sizes.element_type() == xla::ffi::DataType::S32, "group_sizes must be int32."); - auto *int64_sizes_ptr = reinterpret_cast(int64_workspace->untyped_data()); - nvte_convert_int32_to_int64(reinterpret_cast(group_sizes.untyped_data()), - int64_sizes_ptr, num_gemms, stream); +// This FFI is EXPERIMENTAL and subject to change without deprecation, intended for use in JAX's internal implementation of grouped GEMM. +Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, + Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, + Buffer_Type lhs_first_dims, Buffer_Type lhs_last_dims, + Buffer_Type rhs_first_dims, Buffer_Type rhs_last_dims, + Buffer_Type out_first_dims, Buffer_Type out_last_dims, + Buffer_Type alpha, Buffer_Type beta, Result_Type output, + Result_Type cublas_workspace, Result_Type setup_workspace, + Result_Type int64_workspace, GroupedGemmV2Config config) { + auto [lhs_is_trans, rhs_is_trans, scaling_mode, lhs_axis_boundary, rhs_axis_boundary, + lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size] = config; NVTE_CHECK(scaling_mode == JAXX_Scaling_Mode::NO_SCALING, "Only non-quantized grouped GEMM is supported in current implementation."); - // It is weird that TE/Common GEMM only use colwise for MXFP8 - const bool is_fp8_gemm = is_fp8_dtype(lhs_dtype); - const bool is_tensor_scaling = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING || - scaling_mode == JAXX_Scaling_Mode::CURRENT_TENSOR_SCALING; - const bool is_mxfp8_scaling = scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING; - const bool rhs_use_colwise = is_mxfp8_scaling && !rhs_is_trans; - const bool lhs_use_colwise = is_mxfp8_scaling && lhs_is_trans; + size_t num_gemms = grouped_gemm_num_gemms(lhs_first_dims, lhs_last_dims, rhs_first_dims, + rhs_last_dims, out_first_dims, out_last_dims, alpha); - // Outputs - auto out_ptr = reinterpret_cast(output->untyped_data()); - auto out_dtype = convert_ffi_datatype_to_te_dtype(output->element_type()); + // Workspaces. auto setup_workspace_ptr = reinterpret_cast(setup_workspace->untyped_data()); - // Here we clear the lower 8 bits of the buffer address to ensure the buffer is 256-aligned auto cublas_workspace_ptr = reinterpret_cast(cublas_workspace->untyped_data()); cublas_workspace_ptr = move_ptr_to_next_256B_aligned(cublas_workspace_ptr); - auto workspace_total_size = product(cublas_workspace->dimensions()); - - auto lhs_sinv_size = product(lhs_sinv.dimensions()); - auto rhs_sinv_size = product(rhs_sinv.dimensions()); - const size_t workspace_alignment_padding = 256; - const size_t tensor_scaling_sinv_aligment = 16; - const size_t mxfp8_scaling_sinv_alignment_padding = 256; - auto workspace_size = workspace_total_size - workspace_alignment_padding; - if (is_mxfp8_scaling) { - // For MXFP8 swizzled scale_inv buffers, only the first pointer needs to be with 256B alignment padding. Later pointers are guaranteed to be 256-aligned as the scale_inv shapes are padded by 128x4. - workspace_size -= (lhs_sinv_size + rhs_sinv_size + 2 * mxfp8_scaling_sinv_alignment_padding); - } else if (is_tensor_scaling) { - // For tensor scaling, each matrix has a single scale value, and all scales need to be aligned - // by 16 bytes to meet the requirement of CUDA 12.9.1 and later. - workspace_size -= tensor_scaling_sinv_aligment * (lhs_sinv_size + rhs_sinv_size); - } - auto swizzled_lhs_sinv_ptr = cublas_workspace_ptr + workspace_size; - swizzled_lhs_sinv_ptr = move_ptr_to_next_256B_aligned(swizzled_lhs_sinv_ptr); - auto swizzled_rhs_sinv_ptr = swizzled_lhs_sinv_ptr + lhs_sinv_size; - swizzled_rhs_sinv_ptr = move_ptr_to_next_256B_aligned(swizzled_rhs_sinv_ptr); - auto lhs_scatter_aligned_ptr = swizzled_lhs_sinv_ptr; // Already 256B aligned - auto rhs_scatter_aligned_ptr = lhs_scatter_aligned_ptr + num_gemms * tensor_scaling_sinv_aligment; - - size_t lhs_dtype_bytes = te_dtype_bytes(lhs_dtype); - size_t rhs_dtype_bytes = te_dtype_bytes(rhs_dtype); - size_t lhs_sinv_dtype_bytes = te_dtype_bytes(lhs_sinv_dtype); - size_t rhs_sinv_dtype_bytes = te_dtype_bytes(rhs_sinv_dtype); - size_t bias_dtype_bytes = te_dtype_bytes(bias_dtype); - size_t out_dtype_bytes = te_dtype_bytes(out_dtype); - - NVTE_CHECK(lhs_dtype_bytes == rhs_dtype_bytes, "sizeof(lhs_dtype) != sizeof(rhs_dtype)"); - NVTE_CHECK(lhs_sinv_dtype_bytes == rhs_sinv_dtype_bytes, - "sizeof(lhs_sinv_dtype) != sizeof(rhs_sinv_dtype)"); - - size_t expected_lhs_size = m * k; - size_t expected_rhs_size = is_grouped_dense_wgrad ? (k * n) : (num_gemms * k * n); - size_t expected_out_size = is_grouped_dense_wgrad ? (num_gemms * m * n) : (m * n); - size_t actual_lhs_size = product(lhs_data.dimensions()); - size_t actual_rhs_size = product(rhs_data.dimensions()); - size_t actual_out_size = product(output->dimensions()); - NVTE_CHECK(expected_lhs_size == actual_lhs_size, "Unexpected lhs size! Expect ", - expected_lhs_size, ", got ", actual_lhs_size); - if (!is_grouped_dense_wgrad) { - NVTE_CHECK(expected_rhs_size == actual_rhs_size, - "Unexpected rhs size! Expect num_gemms * n * k = ", num_gemms, " * ", n, " * ", k, - " = ", expected_rhs_size, ", got ", actual_rhs_size); - NVTE_CHECK(expected_out_size == actual_out_size, "Unexpected output size! Expect m * n = ", m, - " * ", n, " = ", expected_out_size, ", got ", actual_out_size); - } else { - NVTE_CHECK(expected_rhs_size == actual_rhs_size, "Unexpected rhs size! Expect k * n = ", k, - " * ", n, " = ", expected_rhs_size, ", got ", actual_rhs_size); - NVTE_CHECK(expected_out_size == actual_out_size, - "Unexpected output size! Expect num_gemms * m * n = ", num_gemms, " * ", m, " * ", n, - " = ", expected_out_size, ", got ", actual_out_size); - } - - auto num_math_sm = cuda::sm_count() - getenv("NVTE_EXT_MARGIN_SM", 0); - bool grad = false; - bool accumulate = false; - bool use_split_accumulator = false; - auto bias_shape = std::vector{has_bias ? n : 0}; - const int arch = cuda::sm_arch(); - - if (arch < 100 && is_fp8_gemm) { - NVTE_CHECK(!lhs_is_trans && rhs_is_trans, - "For SM90 or older archs and FP8 input, only NT (row-major) GEMM is supported, ", - "got lhs_is_trans=", lhs_is_trans, ", rhs_is_trans=", rhs_is_trans); - } - + auto workspace_size = product(cublas_workspace->dimensions()) - 256; TensorWrapper workspace_setup(setup_workspace_ptr, std::vector{product(setup_workspace->dimensions())}, DType::kByte); @@ -763,59 +725,21 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty std::vector{num_gemms}, convert_ffi_datatype_to_te_dtype(beta.element_type())); - if (is_grouped_dense_wgrad) { - NVTE_CHECK(lhs_is_trans && !rhs_is_trans, - "For grouped dense wgrad, only TN GEMM is supported in TE/JAX currently."); - - //// RHS - NVTEShape rhsShape{.data = {k, n}, .ndim = 2}; - auto rhs_tensor = make_grouped_tensor(rhs_data, rhs_sinv, scaling_mode, num_gemms, rhsShape); - rhs_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, kNVTEGroupedFirstDims); - - //// LHS - NVTEShape lhsShape{.data = {k, m}, .ndim = 2}; - lhs_is_trans = true; - auto lhs_tensor = make_grouped_tensor(lhs_data, lhs_sinv, scaling_mode, num_gemms, lhsShape); - lhs_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, kNVTEGroupedFirstDims); - - //// OUTPUT - NVTEShape outShape{.data = {num_gemms * m, n}, .ndim = 2}; - auto out_tensor = make_grouped_tensor(*output, std::nullopt, JAXX_Scaling_Mode::NO_SCALING, - num_gemms, outShape); - - nvte_grouped_gemm(rhs_tensor, rhs_is_trans, lhs_tensor, lhs_is_trans, nullptr, out_tensor, - alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), - workspace_cublas.data(), - nullptr, // config (use defaults) - stream); - - return ffi_with_cuda_error_check(); - } - - // Nominal case for FWD or DGRAD - - //// RHS - NVTEShape rhsShape{.data = {num_gemms * k, n}, .ndim = 2}; - if (rhs_is_trans) { - rhsShape.data[0] = num_gemms * n; - rhsShape.data[1] = k; - } - auto rhs_tensor = make_grouped_tensor(rhs_data, rhs_sinv, scaling_mode, num_gemms, rhsShape); - - //// LHS - NVTEShape lhsShape{.data = {m, k}, .ndim = 2}; - if (lhs_is_trans) { - std::swap(lhsShape.data[0], lhsShape.data[1]); - } - auto lhs_tensor = make_grouped_tensor(lhs_data, lhs_sinv, scaling_mode, num_gemms, lhsShape); - lhs_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, - lhs_is_trans ? kNVTEGroupedLastDims : kNVTEGroupedFirstDims); - - //// OUTPUT - NVTEShape outShape{.data = {m, n}, .ndim = 2}; - auto out_tensor = make_grouped_tensor(*output, std::nullopt, JAXX_Scaling_Mode::NO_SCALING, - num_gemms, outShape); - out_tensor.set_group_sizes_only(int64_sizes_ptr, num_gemms, kNVTEGroupedFirstDims); + // Build grouped tensors from XLA buffer shapes and group_sizes — no m/n/k derivation needed. + // int64_workspace is partitioned into per-ragged-buffer slots of num_gemms int64 elements each. + // int64_offset is threaded through the three make_grouped_tensor calls so each non-empty *_dims + // buffer gets its own non-aliasing slot; bounds are checked inside make_grouped_tensor. + auto *int64_base = reinterpret_cast(int64_workspace->untyped_data()); + size_t int64_capacity = int64_workspace->element_count() / sizeof(int64_t); + size_t int64_offset = 0; + auto rhs_tensor = + make_grouped_tensor(rhs_data, rhs_first_dims, rhs_last_dims, int64_base, int64_capacity, + int64_offset, num_gemms, stream, rhs_axis_boundary); + auto lhs_tensor = + make_grouped_tensor(lhs_data, lhs_first_dims, lhs_last_dims, int64_base, int64_capacity, + int64_offset, num_gemms, stream, lhs_axis_boundary); + auto out_tensor = make_grouped_tensor(*output, out_first_dims, out_last_dims, int64_base, + int64_capacity, int64_offset, num_gemms, stream); nvte_grouped_gemm(rhs_tensor, rhs_is_trans, lhs_tensor, lhs_is_trans, nullptr, out_tensor, alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), @@ -834,28 +758,31 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmV2Handler, GroupedGemmV2FFI, .Arg() // rhs_data .Arg() // rhs_sinv .Arg() // bias - .Arg() // group_sizes (int32) + .Arg() // lhs_first_dims (G,) or empty (0,) + .Arg() // lhs_last_dims (G,) or empty (0,) + .Arg() // rhs_first_dims (G,) or empty (0,) + .Arg() // rhs_last_dims (G,) or empty (0,) + .Arg() // out_first_dims (G,) or empty (0,) + .Arg() // out_last_dims (G,) or empty (0,) .Arg() // alpha .Arg() // beta .Ret() // output .Ret() // cublas_workspace .Ret() // setup_workspace .Ret() // int64_workspace - .Attr("M") - .Attr("N") - .Attr("K") - .Attr("lhs_is_trans") - .Attr("rhs_is_trans") - .Attr("scaling_mode") - .Attr("is_grouped_dense_wgrad"), + .Attrs(), FFI_CudaGraph_Traits); Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv, Buffer_Type rhs_data, Buffer_Type rhs_sinv, Buffer_Type bias, - Buffer_Type group_sizes, Buffer_Type group_offset, Result_Type output, - Result_Type workspace, size_t m, size_t n, size_t k, bool lhs_is_trans, - bool rhs_is_trans, JAXX_Scaling_Mode scaling_mode, bool has_bias, - bool is_grouped_dense_wgrad, bool use_async_d2h_group_sizes) { + Buffer_Type lhs_first_dims, Buffer_Type lhs_last_dims, + Buffer_Type rhs_first_dims, Buffer_Type rhs_last_dims, + Buffer_Type out_first_dims, Buffer_Type out_last_dims, + Buffer_Type group_offset, Result_Type output, Result_Type workspace, + GroupedGemmConfig config) { + auto [lhs_is_trans, rhs_is_trans, scaling_mode, has_bias, use_async_d2h_group_sizes, + lhs_axis_boundary, rhs_axis_boundary, lhs_left_size, lhs_right_size, rhs_left_size, + rhs_right_size] = config; // Notes on matrix layouts and transpose: // Jax uses row-major data_layout, on entering this function, each input matrix pair: // A: row-major [m, k] for N - [k, m] for T @@ -872,6 +799,54 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type int num_streams = nvte_get_num_compute_streams(); + // Determine which group_sizes buffers are active (non-empty = ragged dimension). + bool is_lhs_first_ragged = lhs_first_dims.element_count() > 0; + bool is_lhs_last_ragged = lhs_last_dims.element_count() > 0; + bool is_rhs_first_ragged = rhs_first_dims.element_count() > 0; + bool is_rhs_last_ragged = rhs_last_dims.element_count() > 0; + bool is_lhs_ragged = is_lhs_first_ragged || is_lhs_last_ragged; + bool is_rhs_ragged = is_rhs_first_ragged || is_rhs_last_ragged; + bool any_ragged = is_lhs_ragged || is_rhs_ragged; + + size_t num_gemms; + if (is_lhs_first_ragged) + num_gemms = lhs_first_dims.dimensions()[0]; + else if (is_lhs_last_ragged) + num_gemms = lhs_last_dims.dimensions()[0]; + else if (is_rhs_first_ragged) + num_gemms = rhs_first_dims.dimensions()[0]; + else if (is_rhs_last_ragged) + num_gemms = rhs_last_dims.dimensions()[0]; + else + NVTE_CHECK(false, + "GroupedGemmFFI (v1): At least one of the group size buffers must be non-empty to " + "determine num_gemms."); + + const Buffer_Type *active_gs_ptr = nullptr; + if (is_lhs_first_ragged) + active_gs_ptr = &lhs_first_dims; + else if (is_lhs_last_ragged) + active_gs_ptr = &lhs_last_dims; + else if (is_rhs_first_ragged) + active_gs_ptr = &rhs_first_dims; + else if (is_rhs_last_ragged) + active_gs_ptr = &rhs_last_dims; + + // Derive m, n, k from pre-computed original shape sizes (passed from Python). + // lhs_left_size = product of original lhs dims before axis_boundary + // lhs_right_size = product of original lhs dims after axis_boundary + // Same pattern for rhs. + size_t k = lhs_is_trans ? lhs_left_size : lhs_right_size; + size_t m, n; + if (is_rhs_ragged) { + // wgrad: non-contracting lhs dims form M; non-contracting rhs dims form N + m = lhs_is_trans ? lhs_right_size : lhs_left_size; + n = rhs_is_trans ? rhs_left_size : rhs_right_size; + } else { + m = lhs_is_trans ? lhs_right_size : lhs_left_size; // total M (sum of group sizes) + n = rhs_is_trans ? rhs_left_size / num_gemms : rhs_right_size; + } + // Inputs auto lhs_ptr = reinterpret_cast(lhs_data.untyped_data()); auto rhs_ptr = reinterpret_cast(rhs_data.untyped_data()); @@ -884,9 +859,6 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type auto bias_ptr = has_bias ? reinterpret_cast(bias.untyped_data()) : nullptr; auto bias_dtype = convert_ffi_datatype_to_te_dtype(bias.element_type()); - NVTE_CHECK(group_sizes.dimensions().size() == 1); - size_t num_gemms = group_sizes.dimensions()[0]; - // It is weird that TE/Common GEMM only use colwise for MXFP8 const bool is_fp8_gemm = is_fp8_dtype(lhs_dtype); const bool is_tensor_scaling = scaling_mode == JAXX_Scaling_Mode::DELAYED_TENSOR_SCALING || @@ -953,14 +925,14 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type "sizeof(lhs_sinv_dtype) != sizeof(rhs_sinv_dtype)"); size_t expected_lhs_size = m * k; - size_t expected_rhs_size = is_grouped_dense_wgrad ? (k * n) : (num_gemms * k * n); - size_t expected_out_size = is_grouped_dense_wgrad ? (num_gemms * m * n) : (m * n); + size_t expected_rhs_size = is_rhs_ragged ? (k * n) : (num_gemms * k * n); + size_t expected_out_size = is_rhs_ragged ? (num_gemms * m * n) : (m * n); size_t actual_lhs_size = product(lhs_data.dimensions()); size_t actual_rhs_size = product(rhs_data.dimensions()); size_t actual_out_size = product(output->dimensions()); NVTE_CHECK(expected_lhs_size == actual_lhs_size, "Unexpected lhs size! Expect ", expected_lhs_size, ", got ", actual_lhs_size); - if (!is_grouped_dense_wgrad) { + if (!is_rhs_ragged) { NVTE_CHECK(expected_rhs_size == actual_rhs_size, "Unexpected rhs size! Expect num_gemms * n * k = ", num_gemms, " * ", n, " * ", k, " = ", expected_rhs_size, ", got ", actual_rhs_size); @@ -976,25 +948,28 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type size_t dim_list_bytes = sizeof(int32_t) * num_gemms; std::vector dim_list_host(num_gemms); - size_t host_num_gemms = 0; - if (use_async_d2h_group_sizes) { - host_num_gemms = GroupedGemmGetGroupSizes(stream, num_gemms, nullptr, dim_list_host.data()); - NVTE_CHECK(host_num_gemms == num_gemms, "num_gemms ", num_gemms, - " does not match the return of GroupedGemmGetGroupSizes ", host_num_gemms, "."); - } else { - auto dim_list_ptr = reinterpret_cast(group_sizes.untyped_data()); - cudaMemcpyAsync(dim_list_host.data(), dim_list_ptr, dim_list_bytes, cudaMemcpyDeviceToHost, - stream); - // Note: This may break cudaGraph. - cudaStreamSynchronize(stream); - } - size_t sum_group_sizes = std::accumulate(dim_list_host.begin(), dim_list_host.end(), 0); - if (!is_grouped_dense_wgrad) { - NVTE_CHECK(m == sum_group_sizes, "Unexpected group_sizes! M = ", m, - ", got sum(group_sizes)=", sum_group_sizes); - } else { - NVTE_CHECK(k == sum_group_sizes, "Unexpected group_sizes! K = ", k, - ", got sum(group_sizes)=", sum_group_sizes); + if (any_ragged) { + size_t host_num_gemms = 0; + if (use_async_d2h_group_sizes) { + host_num_gemms = GroupedGemmGetGroupSizes(stream, num_gemms, nullptr, dim_list_host.data()); + NVTE_CHECK(host_num_gemms == num_gemms, "num_gemms ", num_gemms, + " does not match the return of GroupedGemmGetGroupSizes ", host_num_gemms, "."); + } else { + NVTE_CHECK(active_gs_ptr != nullptr, "active_gs_ptr is null but any_ragged is true."); + auto gs_data_ptr = reinterpret_cast(active_gs_ptr->untyped_data()); + cudaMemcpyAsync(dim_list_host.data(), gs_data_ptr, dim_list_bytes, cudaMemcpyDeviceToHost, + stream); + // Note: This may break cudaGraph. + cudaStreamSynchronize(stream); + } + size_t sum_group_sizes = std::accumulate(dim_list_host.begin(), dim_list_host.end(), 0); + if (!is_rhs_ragged) { + NVTE_CHECK(m == sum_group_sizes, "Unexpected group_sizes! M = ", m, + ", got sum(group_sizes)=", sum_group_sizes); + } else { + NVTE_CHECK(k == sum_group_sizes, "Unexpected group_sizes! K = ", k, + ", got sum(group_sizes)=", sum_group_sizes); + } } auto num_math_sm = cuda::sm_count() - getenv("NVTE_EXT_MARGIN_SM", 0); @@ -1042,7 +1017,7 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type auto lhs_shape_i = std::vector{m_i, k}; auto rhs_shape_i = std::vector{rhs_is_trans ? n : k, rhs_is_trans ? k : n}; auto out_shape_i = std::vector{m_i, n}; - if (is_grouped_dense_wgrad) { + if (is_rhs_ragged) { size_t k_i = dim_list_host[i]; lhs_shape_i[0] = lhs_is_trans ? k_i : m; lhs_shape_i[1] = lhs_is_trans ? m : k_i; @@ -1237,19 +1212,16 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedGemmHandler, GroupedGemmFFI, .Arg() // rhs_data .Arg() // rhs_sinv .Arg() // bias - .Arg() // group_sizes + .Arg() // lhs_first_dims (G,) or empty (0,) + .Arg() // lhs_last_dims (G,) or empty (0,) + .Arg() // rhs_first_dims (G,) or empty (0,) + .Arg() // rhs_last_dims (G,) or empty (0,) + .Arg() // out_first_dims (G,) or empty (0,) + .Arg() // out_last_dims (G,) or empty (0,) .Arg() // group_offset .Ret() // output .Ret() // workspace - .Attr("M") - .Attr("N") - .Attr("K") - .Attr("lhs_is_trans") - .Attr("rhs_is_trans") - .Attr("scaling_mode") - .Attr("has_bias") - .Attr("is_grouped_dense_wgrad") - .Attr("use_async_d2h_group_sizes")); + .Attrs()); } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index fe02e61fc0..dbd7bbb1ff 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -18,15 +18,11 @@ from . import cpp_extensions as tex from .cpp_extensions.amax import AmaxScope from .quantize import ( - ScaledTensorFactory, ScaledTensor, - ScalingMode, QuantizerSet, noop_quantizer_set, with_sharding_constraint_by_logical_axes, - is_fp8_gemm_with_all_layouts_supported, TensorUsage, - QuantizeLayout, ) @@ -325,7 +321,6 @@ def grouped_dense( group_sizes: jnp.ndarray, contracting_dims: Tuple[Sequence[int], Sequence[int]] = ((1,), (1,)), bias: jnp.ndarray = None, - kernel_amax: jnp.ndarray = None, precision: jax.lax.Precision = jax.lax.Precision.DEFAULT, preferred_element_type: jnp.dtype = None, group_offset: jnp.array = None, @@ -342,7 +337,6 @@ def grouped_dense( contracting_dims: Tuple of sequences specifying which dimensions to contract (currently only supports ((1,), (1,))) bias: Bias tensor of shape (G, N) - kernel_amax: The amax values of weight matrix of shape (G,) precision: JAX precision for the GEMM operation preferred_element_type: Preferred data type for the output tensor group_offset: 1D array containing offsets for each group (not yet implemented) @@ -361,7 +355,6 @@ def grouped_dense( group_sizes, contracting_dims, bias, - kernel_amax, precision, preferred_element_type, group_offset, @@ -371,14 +364,13 @@ def grouped_dense( return output -@partial(jax.custom_vjp, nondiff_argnums=(3, 6, 7, 8, 10)) +@partial(jax.custom_vjp, nondiff_argnums=(3, 5, 6, 7, 9)) def _grouped_dense( x, kernel, group_sizes, contracting_dims, bias, - kernel_amax, precision, preferred_element_type, group_offset, @@ -391,7 +383,6 @@ def _grouped_dense( group_sizes, contracting_dims, bias, - kernel_amax, precision, preferred_element_type, group_offset, @@ -407,7 +398,6 @@ def _grouped_dense_fwd_rule( group_sizes, contracting_dims, bias, - kernel_amax, precision, preferred_element_type, group_offset, @@ -415,118 +405,42 @@ def _grouped_dense_fwd_rule( kernel_fsdp_info, ): use_bias = bias is not None - is_noop_quantizer_set = quantizer_set == noop_quantizer_set kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info kernel_fsdp_enabled = kernel_fsdp_mesh_axis is not None + assert not kernel_fsdp_enabled, "FSDP sharding for grouped_dense is not supported yet." + del kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx, kernel_fsdp_info, kernel_fsdp_enabled - if is_noop_quantizer_set: - grouped_gemm_x = x - grouped_gemm_kernel = kernel - ctx_x = x - ctx_kernel = kernel - flatten_axis_k = None - - if kernel_fsdp_enabled: - kernel = _all_gather_kernel(kernel, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx) - else: - original_quantizer_set_kernel_q_layout = quantizer_set.kernel.q_layout - - x_contracting_dims, k_contracting_dims = contracting_dims - flatten_axis_x = -len(x_contracting_dims) - flatten_axis_k = len(k_contracting_dims) - len(kernel.shape) + 1 # +1 for G axis - - assert x.ndim == 2, "Grouped dense expects a 2D input tensor of shape (M, K)" - assert kernel.ndim == 3, "Grouped dense expects a 3D kernel tensor of shape (G, K, N)" - # Expected k_contracting_dims == (1,), need to tweak it for grouped_gemm FP8 extra transpose - # TODO(Hua): Do we have a better way for this? What if is_gemm_with_all_layouts_supported()? - assert x_contracting_dims == (1,) and k_contracting_dims == (1,), ( - "grouped_dense for FP8 can only handle x_contracting_dims=(1,) " - "and k_contracting_dims=(1,) for now, " - f"got {x_contracting_dims=} and {k_contracting_dims=}" - ) + x_contracting_dims, k_contracting_dims = contracting_dims + flatten_axis_x = -len(x_contracting_dims) + flatten_axis_k = len(k_contracting_dims) - len(kernel.shape) + 1 # +1 for G axis - casted_x = tex.grouped_quantize( - x, - quantizer_set.x, - group_sizes, - flatten_axis=flatten_axis_x, - ) + casted_x = tex.grouped_quantize( + x, + quantizer_set.x, + group_sizes, + flatten_axis=flatten_axis_x, + ) - ctx_kernel_usage = TensorUsage.RHS_TRANS - if kernel_fsdp_enabled: - assert quantizer_set.kernel.scaling_mode in [ - ScalingMode.CURRENT_TENSOR_SCALING, - ScalingMode.DELAYED_TENSOR_SCALING, - ] - # Perform `cast` only - ctx_kernel_usage = TensorUsage.LHS - quantizer_set.kernel.q_layout = QuantizeLayout.ROWWISE - - casted_kernel = tex.grouped_quantize( - kernel, quantizer_set.kernel, amax=kernel_amax, flatten_axis=flatten_axis_k - ) - contracting_dims = (x_contracting_dims, k_contracting_dims) - - # For x_contracting_dims == (1,) and k_contracting_dims == (1,), we should have - # rowwise_casted_x.original_shape == (M, K) - # colwise_casted_kernel.original_shape == (G, N, K) - grouped_gemm_x = casted_x.get_tensor(usage=TensorUsage.LHS) - ctx_x = casted_x.get_tensor(usage=TensorUsage.LHS_TRANS) - ctx_kernel = casted_kernel.get_tensor(usage=ctx_kernel_usage) - - if kernel_fsdp_enabled: - ctx_kernel_in_original_shape = ctx_kernel.data.reshape(ctx_kernel.original_shape) - global_ctx_kernel_data = _all_gather_kernel( - ctx_kernel_in_original_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx - ) - kernel_shape = global_ctx_kernel_data.shape - - ctx_kernel = ScaledTensorFactory.create_1x( - global_ctx_kernel_data.reshape(-1), - ctx_kernel.scale_inv, - scaling_mode=ctx_kernel.scaling_mode, - dq_dtype=ctx_kernel.dq_dtype, - is_colwise=False, - data_layout="N", - flatten_axis=ctx_kernel.flatten_axis, - group_sizes=ctx_kernel.group_sizes, - original_shape=kernel_shape, - group_axis=ctx_kernel.group_axis, - ) - - if is_fp8_gemm_with_all_layouts_supported(): - grouped_gemm_kernel = ctx_kernel - else: - grouped_gemm_kernel_data = global_ctx_kernel_data.transpose(0, 2, 1) - grouped_gemm_kernel = ScaledTensorFactory.create_1x( - grouped_gemm_kernel_data.reshape(-1), - ctx_kernel.scale_inv, - scaling_mode=ctx_kernel.scaling_mode, - dq_dtype=ctx_kernel.dq_dtype, - is_colwise=True, - data_layout="T", - flatten_axis=ctx_kernel.flatten_axis, - group_sizes=ctx_kernel.group_sizes, - original_shape=kernel_shape, - group_axis=ctx_kernel.group_axis, - ) - else: - grouped_gemm_kernel = casted_kernel.get_tensor(usage=TensorUsage.RHS) - - # Reset quantizer_set.kernel.q_layout to align the PyTree as the given one. - # This is needed especially when kernel_fsdp_enabled == True AND FP8 enabled. - quantizer_set.kernel.q_layout = original_quantizer_set_kernel_q_layout + casted_kernel = tex.grouped_quantize(kernel, quantizer_set.kernel, flatten_axis=flatten_axis_k) + contracting_dims = (x_contracting_dims, k_contracting_dims) + # For x_contracting_dims == (1,) and k_contracting_dims == (1,), we should have + # rowwise_casted_x.original_shape == (M, K) + # colwise_casted_kernel.original_shape == (G, N, K) + grouped_gemm_x = casted_x.get_tensor(usage=TensorUsage.LHS) + ctx_x = casted_x.get_tensor(usage=TensorUsage.LHS_TRANS) + ctx_kernel = casted_kernel.get_tensor(usage=TensorUsage.RHS_TRANS) + + grouped_gemm_kernel = casted_kernel.get_tensor(usage=TensorUsage.RHS) output = tex.grouped_gemm( grouped_gemm_x, grouped_gemm_kernel, - group_sizes, - contracting_dims, - bias, - precision, - preferred_element_type, - group_offset, + contracting_dims=contracting_dims, + bias=bias, + precision=precision, + preferred_element_type=preferred_element_type, + group_offset=group_offset, ) ctx = ( @@ -540,7 +454,6 @@ def _grouped_dense_fwd_rule( x.shape, kernel.shape, use_bias, - is_noop_quantizer_set, quantizer_set, flatten_axis_k, ) @@ -550,6 +463,10 @@ def _grouped_dense_fwd_rule( def _grouped_dense_bwd_rule( contracting_dims, precision, preferred_element_type, group_offset, kernel_fsdp_info, ctx, grad ): + kernel_fsdp_mesh_axis, _ = kernel_fsdp_info + kernel_fsdp_enabled = kernel_fsdp_mesh_axis is not None + assert not kernel_fsdp_enabled, "FSDP sharding for grouped_dense is not supported yet." + fwd_x_contracting_dims, fwd_k_contracting_dims = contracting_dims ( @@ -559,62 +476,41 @@ def _grouped_dense_bwd_rule( x_shape, kernel_shape, use_bias, - is_noop_quantizer_set, quantizer_set, flatten_axis_k, ) = ctx - if is_noop_quantizer_set: - # The 1 in range is for excluding the group dimension (shall we use the hardcoded results below?) - # g_contracting_dim = (1, ) - # k_contracting_dim = (2, ) - g_contracting_dim = tuple( - range(1 + grad.ndim - len(kernel_shape) + len(fwd_k_contracting_dims), grad.ndim) - ) - k_contracting_dim = tuple( - dim for dim in range(1, len(kernel_shape)) if dim not in fwd_k_contracting_dims - ) - dgrad_contracting_dims = (g_contracting_dim, k_contracting_dim) - dgrad_grad = grad - dgrad_kernel_T = ctx_kernel - - # g_contracting_dim = (0, ) - # x_contracting_dim = (0, ) - g_contracting_dim = x_contracting_dim = tuple( - range(0, len(x_shape) - len(fwd_x_contracting_dims)) - ) - wgrad_contracting_dims = (x_contracting_dim, g_contracting_dim) - wgrad_x_T = ctx_x - wgrad_grad = grad - else: - casted_grad = tex.grouped_quantize( - grad, quantizer_set.dgrad, group_sizes, flatten_axis=flatten_axis_k - ) + # The 1 in range is for excluding the group dimension (shall we use the hardcoded results below?) + # g_contracting_dim = (1, ) + # k_contracting_dim = (2, ) + g_contracting_dim = tuple( + range(1 + grad.ndim - len(kernel_shape) + len(fwd_k_contracting_dims), grad.ndim) + ) + k_contracting_dim = tuple( + dim for dim in range(1, len(kernel_shape)) if dim not in fwd_k_contracting_dims + ) - # For x_contracting_dims == (1,) and k_contracting_dims == (1,), we need to use - # g_contracting_dim = (1,) and k_contracting_dim = (2,) to make it work after the - # extra transpose for FP8 in grouped_gemm - # TODO(Hua): Do we have a better way for this? What if is_gemm_with_all_layouts_supported()? - g_contracting_dim = (1,) - k_contracting_dim = (2,) - dgrad_contracting_dims = (g_contracting_dim, k_contracting_dim) - dgrad_grad = casted_grad.get_tensor(usage=TensorUsage.LHS) - dgrad_kernel_T = ctx_kernel - - # We need to use g_contracting_dim = (0,) and x_contracting_dim = (0,) to make it work - # after the extra transpose for FP8 in grouped_gemm - # TODO(Hua): Do we have a better way for this? What if is_gemm_with_all_layouts_supported()? - g_contracting_dim = (0,) - x_contracting_dim = (0,) - wgrad_contracting_dims = (x_contracting_dim, g_contracting_dim) - wgrad_x_T = ctx_x - wgrad_grad = casted_grad.get_tensor(usage=TensorUsage.RHS) + casted_grad = tex.grouped_quantize( + grad, quantizer_set.dgrad, group_sizes, flatten_axis=flatten_axis_k + ) + dgrad_contracting_dims = (g_contracting_dim, k_contracting_dim) + dgrad_grad = casted_grad.get_tensor(usage=TensorUsage.LHS) + dgrad_kernel_T = ctx_kernel + + # g_contracting_dim = (0, ) + # x_contracting_dim = (0, ) + g_contracting_dim = x_contracting_dim = tuple( + range(0, len(x_shape) - len(fwd_x_contracting_dims)) + ) + wgrad_contracting_dims = (x_contracting_dim, g_contracting_dim) + + wgrad_x_T = ctx_x + wgrad_grad = casted_grad.get_tensor(usage=TensorUsage.RHS) dgrad = tex.grouped_gemm( dgrad_grad, dgrad_kernel_T, - group_sizes, - dgrad_contracting_dims, + contracting_dims=dgrad_contracting_dims, precision=precision, preferred_element_type=preferred_element_type, group_offset=group_offset, @@ -623,23 +519,16 @@ def _grouped_dense_bwd_rule( wgrad = tex.grouped_gemm( wgrad_x_T, wgrad_grad, - group_sizes, - wgrad_contracting_dims, + contracting_dims=wgrad_contracting_dims, precision=precision, preferred_element_type=preferred_element_type, group_offset=group_offset, ) - kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx = kernel_fsdp_info - if kernel_fsdp_mesh_axis is not None: - wgrad = _psum_scatter_kernel( - wgrad, kernel_shape, kernel_fsdp_mesh_axis, kernel_fsdp_axis_idx - ) group_sizes_grad = None dbias = tex.grouped_dbias(grad, group_sizes) if use_bias else None - dkernel_amax = None - return dgrad, wgrad, group_sizes_grad, dbias, dkernel_amax, quantizer_set + return dgrad, wgrad, group_sizes_grad, dbias, quantizer_set _grouped_dense.defvjp(_grouped_dense_fwd_rule, _grouped_dense_bwd_rule) diff --git a/transformer_engine/jax/quantize/dequantizer.py b/transformer_engine/jax/quantize/dequantizer.py index 74787b9308..5abb2e74df 100644 --- a/transformer_engine/jax/quantize/dequantizer.py +++ b/transformer_engine/jax/quantize/dequantizer.py @@ -275,29 +275,45 @@ def _grouped_dequantize(grouped_scaled_tensor): """ data = grouped_scaled_tensor.data scale_inv = grouped_scaled_tensor.scale_inv - group_sizes = grouped_scaled_tensor.group_sizes + group_sizes = ( + grouped_scaled_tensor.first_dims + if grouped_scaled_tensor.first_dims is not None + and grouped_scaled_tensor.first_dims.size > 0 + else grouped_scaled_tensor.last_dims + ) + # For non-ragged groups (kernel case), group_sizes is not stored; derive from original_shape + if group_sizes is None: + group_sizes = jnp.ones(grouped_scaled_tensor.original_shape[0], dtype=jnp.int32) flatten_axis = grouped_scaled_tensor.flatten_axis scaling_mode = grouped_scaled_tensor.scaling_mode original_shape = grouped_scaled_tensor.original_shape - group_axis = grouped_scaled_tensor.group_axis - flatten_axis = len(original_shape) + flatten_axis if flatten_axis < 0 else flatten_axis output = [] - non_group_shape = tuple( - original_shape[i] for i in range(len(original_shape)) if i != group_axis + # For transposed (colwise) tensors with ragged groups, the group dimension is the last + # axis of original_shape (e.g. original_shape = (N, M) with groups along M), while the + # non-group dimensions are all axes before it. For the uniform-groups case the group + # dimension stays at axis 0, so the existing axis-0 logic applies. + is_transposed_ragged = ( + grouped_scaled_tensor.data_layout == "T" and group_sizes.size != original_shape[0] ) + if is_transposed_ragged: + non_group_shape = original_shape[:-1] + else: + non_group_shape = tuple(original_shape[i] for i in range(len(original_shape)) if i != 0) matrix_sizes = group_sizes * math.prod(non_group_shape) data = jnp.split(data, jnp.cumulative_sum(matrix_sizes)[:-1]) scale_inv_ptr = 0 for i, data_i in enumerate(data): - data_shape_i = ( - *original_shape[:group_axis], - group_sizes[i], - *original_shape[group_axis + 1 :], - ) + if is_transposed_ragged: + data_shape_i = (*non_group_shape, group_sizes[i]) + else: + data_shape_i = ( + group_sizes[i], + *original_shape[1:], + ) assert math.prod(data_shape_i) == data_i.size, ( f"math.prod({data_shape_i}) = {math.prod(data_shape_i)} which is not equal to" f" {data_i.size}" diff --git a/transformer_engine/jax/quantize/quantizer.py b/transformer_engine/jax/quantize/quantizer.py index f5ca6aeaed..db56db935d 100644 --- a/transformer_engine/jax/quantize/quantizer.py +++ b/transformer_engine/jax/quantize/quantizer.py @@ -920,7 +920,7 @@ def __post_init__(self): self.data_layout = self.quantizers[0].data_layout def _create_grouped_tensor_from_tensor_list( - self, tensor_list, group_sizes, original_shape, group_axis, mode + self, tensor_list, group_sizes, original_shape, mode ): # mode 0 = concate, mode 1 = add # TODO(Ming Huang): Consider to apply Enum for mode. @@ -948,9 +948,8 @@ def _create_grouped_tensor_from_tensor_list( is_colwise=tensor_list[0].is_colwise, data_layout=tensor_list[0].data_layout, flatten_axis=tensor_list[0].flatten_axis, - group_sizes=group_sizes, + first_dims=group_sizes, original_shape=original_shape, - group_axis=group_axis, ) def _quantize_func(self, *args, **kwargs): @@ -964,12 +963,11 @@ def quantize( dq_dtype=None, flatten_axis=-1, group_sizes=None, - group_axis=0, ): """Quantize a tensor in grouped manner. Expected input shape: [M, K] or [G, K, N] - Split to x.shape[group_axis] number of groups if group_sizes is not given + Split to x.shape[0] number of groups if group_sizes is not given Args: x: Input tensor to quantize @@ -978,12 +976,10 @@ def quantize( dq_dtype: Data type for dequantized values flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) group_sizes: Array of ints containing the size of each group (default: None) - group_axis: The axis along which grouping is performed (default: 0) Returns: A ScaledTensor1x or ScaledTensor2x containing the quantized data """ - assert group_axis == 0, "Only group_axis == 0 is supported now!" dq_dtype = dq_dtype if dq_dtype is not None else x.dtype if flatten_axis < 0: @@ -1023,8 +1019,8 @@ def quantize( tensor_list.append(tensor) combine_mode = 1 # Add else: - group_sizes = jnp.ones(x.shape[group_axis], dtype=jnp.int32) - x = jnp.split(x, x.shape[group_axis], axis=group_axis) + group_sizes = jnp.ones(x.shape[0], dtype=jnp.int32) + x = jnp.split(x, x.shape[0], axis=0) tensor_list = [] for i in range(len(group_sizes)): @@ -1038,12 +1034,12 @@ def quantize( if is_rowwise: rowwise_tensor_list = [tensor.get_rowwise_tensor() for tensor in tensor_list] grouped_rowwise_tensor = self._create_grouped_tensor_from_tensor_list( - rowwise_tensor_list, group_sizes, original_shape, group_axis, combine_mode + rowwise_tensor_list, group_sizes, original_shape, combine_mode ) if is_colwise: colwise_tensor_list = [tensor.get_colwise_tensor() for tensor in tensor_list] grouped_colwise_tensor = self._create_grouped_tensor_from_tensor_list( - colwise_tensor_list, group_sizes, original_shape, group_axis, combine_mode + colwise_tensor_list, group_sizes, original_shape, combine_mode ) if is_colwise and is_rowwise: diff --git a/transformer_engine/jax/quantize/scaling_modes.py b/transformer_engine/jax/quantize/scaling_modes.py index 61c3af178c..26b998ba90 100644 --- a/transformer_engine/jax/quantize/scaling_modes.py +++ b/transformer_engine/jax/quantize/scaling_modes.py @@ -135,14 +135,13 @@ def get_scale_shape( @abstractmethod def get_grouped_scale_shape( - self, data_shape, n_groups, group_axis, is_colwise, is_padded=True, flatten_axis=-1 + self, data_shape, n_groups, is_colwise, is_padded=True, flatten_axis=-1 ) -> Tuple[int]: """Get the shape for scale tensors in this mode. Args: data_shape: Original shape of the data tensor n_groups: Number of groups in grouped quantization - group_axis: The axis along which grouping is performed is_colwise: Whether to use column-wise scaling is_padded: Whether to use padded shapes flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) @@ -253,7 +252,7 @@ def get_quantize_layout(self, usage: TensorUsage) -> QuantizeLayout: return QuantizeLayout.ROWWISE def get_grouped_scale_shape( - self, data_shape, n_groups, group_axis, is_colwise, is_padded=True, flatten_axis=-1 + self, data_shape, n_groups, is_colwise, is_padded=True, flatten_axis=-1 ) -> Tuple[int]: """Get the shape for scale tensors in this mode. @@ -266,7 +265,7 @@ def get_grouped_scale_shape( Returns: The shape for scale tensors """ - del data_shape, group_axis, is_colwise + del data_shape, is_colwise assert isinstance(n_groups, int) return (n_groups,) @@ -370,7 +369,7 @@ def get_quantize_layout(self, usage: TensorUsage) -> QuantizeLayout: return QuantizeLayout.COLWISE def get_grouped_scale_shape( - self, data_shape, n_groups, group_axis, is_colwise, is_padded=True, flatten_axis=-1 + self, data_shape, n_groups, is_colwise, is_padded=True, flatten_axis=-1 ) -> Tuple[int]: """Get the shape for scale tensors in this mode. @@ -383,7 +382,7 @@ def get_grouped_scale_shape( Returns: The shape for scale tensors """ - del data_shape, group_axis, is_colwise + del data_shape, is_colwise assert isinstance(n_groups, int) return (n_groups,) @@ -613,7 +612,7 @@ def get_quantize_layout(self, usage: TensorUsage) -> QuantizeLayout: return QuantizeLayout.COLWISE def get_grouped_scale_shape( - self, data_shape, n_groups, group_axis, is_colwise, is_padded=True, flatten_axis=-1 + self, data_shape, n_groups, is_colwise, is_padded=True, flatten_axis=-1 ) -> Tuple[int]: """Get the shape for grouped scale tensors in this mode. If padded: The estimiated maximal possible shape for grouped scale tensor is return instead. @@ -937,14 +936,13 @@ def get_shardy_sharding_rules( ) def get_grouped_scale_shape_2x( - self, data_shape, n_groups, group_axis, is_padded=True, flatten_axis=-1 + self, data_shape, n_groups, is_padded=True, flatten_axis=-1 ) -> Tuple[Tuple[int]]: """Get shapes for both row-wise and column-wise scaling. Args: data_shape: Shape of the data tensor n_groups: Number of groups for grouped quantization - group_axis: The axis along which grouping is performed is_padded: Whether to use padded shapes flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) @@ -954,7 +952,6 @@ def get_grouped_scale_shape_2x( rowwise_scale_shape = self.get_grouped_scale_shape( data_shape, n_groups, - group_axis, is_colwise=False, is_padded=is_padded, flatten_axis=flatten_axis, @@ -962,7 +959,6 @@ def get_grouped_scale_shape_2x( colwise_scale_shape = self.get_grouped_scale_shape( data_shape, n_groups, - group_axis, is_colwise=True, is_padded=is_padded, flatten_axis=flatten_axis, @@ -970,7 +966,7 @@ def get_grouped_scale_shape_2x( return (rowwise_scale_shape, colwise_scale_shape) def get_grouped_scale_shape( - self, data_shape, n_groups, group_axis, is_colwise, is_padded=True, flatten_axis=-1 + self, data_shape, n_groups, is_colwise, is_padded=True, flatten_axis=-1 ) -> Tuple[Tuple[int]]: """Get shapes for both row-wise and column-wise scaling. @@ -985,7 +981,6 @@ def get_grouped_scale_shape( return self._get_impl().get_grouped_scale_shape( data_shape, n_groups, - group_axis, is_colwise=is_colwise, is_padded=is_padded, flatten_axis=flatten_axis, diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index c26cb8a531..b1f49dacdc 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -9,7 +9,7 @@ rowwise and colwise quantization modes with proper scaling and dequantization. """ from dataclasses import dataclass -from typing import Callable, Tuple +from typing import Callable, Optional, Tuple from abc import ABC, abstractmethod import jax.numpy as jnp @@ -32,6 +32,7 @@ "ScaledTensor1x", "ScaledTensor2x", "GroupedScaledTensor1x", + "GroupedNoScaleTensor", "ScaledTensorFactory", "with_sharding_constraint_by_logical_axes", ] @@ -365,21 +366,22 @@ class GroupedScaledTensor1x(ScaledTensor1x): where elements are grouped along a specified axis. Attributes: - group_sizes: Array containing the size of each group + first_dims: Per-group sizes of the first (row) 2D dim, or None if not ragged + last_dims: Per-group sizes of the last (col) 2D dim, or None if not ragged original_shape: The original shape of the tensor before grouping - group_axis: The axis along which grouping is performed (default: 0) """ - group_sizes: jnp.ndarray + first_dims: Optional[jnp.ndarray] + last_dims: Optional[jnp.ndarray] original_shape: Tuple - group_axis: int def __init__( self, data, scale_inv, amax, - group_sizes, + first_dims, + last_dims, scaling_mode, dq_dtype, _dq_func, @@ -387,12 +389,11 @@ def __init__( data_layout, flatten_axis, original_shape, - group_axis=0, ): self.flatten_axis = flatten_axis - self.group_sizes = group_sizes + self.first_dims = first_dims + self.last_dims = last_dims self.original_shape = original_shape - self.group_axis = group_axis # TODO(Phuong):Handle RHT for grouped quantization once grouped quantization supports NVFP4 super().__init__( data=data, @@ -410,7 +411,6 @@ def __init__( def __post_init__(self): assert self.scale_inv.ndim == 1, "Only support flattened scale_inv" assert self.data.ndim == 1, "Only support flattened data" - assert self.group_axis >= 0 assert self.flatten_axis > 0 data_ndim = len(self.original_shape) @@ -418,14 +418,19 @@ def __post_init__(self): 0 < self.flatten_axis < data_ndim ), f"flatten_axis {self.flatten_axis} is out of bounds for data.ndim = {data_ndim}" - assert ( - 0 <= self.group_axis < data_ndim - ), f"group_axis {self.group_axis} is out of bounds for shape {self.original_shape}" + active_dims = ( + self.first_dims + if self.first_dims is not None and self.first_dims.size > 0 + else self.last_dims + ) + if active_dims is not None: + num_groups = active_dims.size + else: + num_groups = self.original_shape[0] expected_scale_shape = self.scaling_mode.get_grouped_scale_shape( self.original_shape, - self.group_sizes.size, - self.group_axis, + num_groups, self.is_colwise, is_padded=True, flatten_axis=self.flatten_axis, @@ -442,7 +447,7 @@ def tree_flatten(self): Returns: A tuple containing (children, aux_data) for tree operations """ - children = (self.data, self.scale_inv, self.amax, self.group_sizes) + children = (self.data, self.scale_inv, self.amax, self.first_dims, self.last_dims) aux_data = ( self.scaling_mode, self.dq_dtype, @@ -451,7 +456,6 @@ def tree_flatten(self): self.data_layout, self.flatten_axis, self.original_shape, - self.group_axis, ) return (children, aux_data) @@ -473,6 +477,81 @@ def checkpoint(self, quantizer): return jax_checkpoint_name(self, name=quantizer.checkpoint_name) +@register_pytree_node_class +@dataclass +class GroupedNoScaleTensor(AbstractBaseTensor1x): + """Unquantized grouped tensor. + + Stores N-D data with per-group dimension sizes so that grouped_gemm() + can extract first/last dims automatically without explicit parameters. + + Attributes: + data: The raw (unquantized) tensor data in N-D layout + first_dims: Per-group sizes of the first (row) 2D dim, or None if not ragged + last_dims: Per-group sizes of the last (col) 2D dim, or None if not ragged + original_shape: Shape of data (same as data.shape for N-D unquantized) + """ + + first_dims: Optional[jnp.ndarray] + last_dims: Optional[jnp.ndarray] + original_shape: Tuple + + def tree_flatten(self): + """Flattens the tensor for JAX tree operations.""" + children = (self.data, self.amax, self.first_dims, self.last_dims) + aux_data = (self.original_shape,) + return (children, aux_data) + + @property + def ndim(self): + """Number of dimensions of the underlying array.""" + return self.data.ndim + + def dequantize(self): + """This is a no-op for a higher-precision tensor so this simply returns the tensor's data.""" + return self.data + + def get_tensor(self, usage: TensorUsage): + """Returns the tensor based on the tensor usage.""" + q_layout = ScalingMode.NO_SCALING.get_quantize_layout(usage) + assert q_layout.is_rowwise_only, "Only ROWWISE layout is supported for NoScaleTensor" + return self + + def apply_sharding_constraint_by_logical_axes(self, logical_axis_names: Tuple[str, ...]): + """Applies sharding constraints to a tensor based on logical axis names. + + Args: + logical_axis_names: Tuple of logical axis names for sharding + + Returns: + The tensor with applied sharding constraints + """ + if not logical_axis_names: + return self + + data = with_sharding_constraint_by_logical_axes(self.data, logical_axis_names) + + return GroupedNoScaleTensor( + data=data, + amax=self.amax, + first_dims=self.first_dims, + last_dims=self.last_dims, + original_shape=self.original_shape, + ) + + def checkpoint(self, quantizer): + """Checkpoints the tensor with the given quantizer's checkpoint name if available. + + Args: + quantizer: The quantizer to use for checkpointing. If None, no checkpointing is applied. + + Returns: + The checkpointed tensor + """ + assert quantizer is None, "NoScaleTensor does not support quantization." + return self + + @register_pytree_node_class @dataclass class ScaledTensor2x(AbstractBaseTensor, ScaledTensor): @@ -570,9 +649,9 @@ def create_1x( is_colwise=False, data_layout="N", flatten_axis=-1, - group_sizes=None, + first_dims=None, + last_dims=None, original_shape=None, - group_axis=0, has_rht_applied=False, ): """Creates a single-scale quantized tensor. @@ -586,29 +665,37 @@ def create_1x( is_colwise: Whether to use column-wise quantization (default: False) data_layout: The data_layout specification (default: "N") flatten_axis: The quantization axis for the tensor - group_sizes: Array of ints containing the size of each group (default: None) + first_dims: Per-group sizes of the first (row) 2D dim (default: None) + last_dims: Per-group sizes of the last (col) 2D dim (default: None) original_shape: The original shape of the tensor before grouping (default: None) - group_axis: The axis along which grouping is performed (default: 0) has_rht_applied: Whether the tensor had the Randomized Hadamard Transform (RHT) applied during quantization (default: False) Returns: - A ScaledTensor1x or GroupedScaledTensor1x instance depending on whether group_sizes is provided + A ScaledTensor1x or GroupedScaledTensor1x instance depending on whether first_dims or last_dims is provided """ if amax is None: amax = jnp.empty((1,), dtype=jnp.float32) dequantizer = ScalingModeToDequantizerMap.get(scaling_mode) - if group_sizes is not None: - flatten_axis = (len(original_shape) + flatten_axis) % len(original_shape) + if first_dims is not None or last_dims is not None or original_shape is not None: assert ( original_shape is not None ), "original_shape is not given for GroupedScaledTensor1x" + flatten_axis = (len(original_shape) + flatten_axis) % len(original_shape) + + # Determine num_groups from whichever dims array is provided, or from original_shape + active_dims = ( + first_dims if first_dims is not None and first_dims.size > 0 else last_dims + ) + if active_dims is not None: + num_groups = active_dims.size + else: + num_groups = original_shape[0] # Handling attrs of transposed tensors - group_axis = (len(original_shape) + group_axis) % len(original_shape) if data_layout == "T": - if original_shape[0] == group_sizes.size: + if original_shape[0] == num_groups: original_shape = ( original_shape[0], *original_shape[flatten_axis:], @@ -620,7 +707,6 @@ def create_1x( *original_shape[flatten_axis:], *original_shape[:flatten_axis], ) - group_axis = flatten_axis flatten_axis = len(original_shape) - flatten_axis return GroupedScaledTensor1x( @@ -633,9 +719,9 @@ def create_1x( is_colwise=is_colwise, data_layout=data_layout, flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=first_dims, + last_dims=last_dims, original_shape=original_shape, - group_axis=group_axis, ) # Handling attrs of transposed tensors @@ -668,9 +754,9 @@ def create_2x( dq_dtype=jnp.bfloat16, data_layout="NN", flatten_axis=-1, - group_sizes=None, + first_dims=None, + last_dims=None, original_shape=None, - group_axis=0, rowwise_has_rht_applied=False, colwise_has_rht_applied=False, ): @@ -686,9 +772,9 @@ def create_2x( dq_dtype: The data type for dequantized values (default: bfloat16) data_layout: The data_layout specification (default: "NN") flatten_axis: The quantization axis for the tensor - group_sizes: Array containing the size of each group (default: None) + first_dims: Per-group sizes of the first (row) 2D dim (default: None) + last_dims: Per-group sizes of the last (col) 2D dim (default: None) original_shape: The original shape of the tensor before grouping (default: None) - group_axis: The axis along which grouping is performed (default: 0) rowwise_has_rht_applied: Whether the row-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) colwise_has_rht_applied: Whether the column-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) @@ -710,9 +796,9 @@ def create_2x( is_colwise=False, data_layout=data_layout[0], flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=first_dims, + last_dims=last_dims, original_shape=original_shape, - group_axis=group_axis, has_rht_applied=rowwise_has_rht_applied, ) colwise_tensor = ScaledTensorFactory.create_1x( @@ -724,9 +810,9 @@ def create_2x( is_colwise=True, data_layout=data_layout[1], flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=first_dims, + last_dims=last_dims, original_shape=original_shape, - group_axis=group_axis, has_rht_applied=colwise_has_rht_applied, ) return ScaledTensor2x(rowwise_tensor, colwise_tensor) @@ -744,9 +830,9 @@ def create( data_layout: str = "NN", q_layout: QuantizeLayout = QuantizeLayout.ROWWISE, flatten_axis: int = -1, - group_sizes: jnp.ndarray = None, + first_dims: jnp.ndarray = None, + last_dims: jnp.ndarray = None, original_shape: Tuple[int] = None, - group_axis: int = 0, rowwise_has_rht_applied: bool = False, colwise_has_rht_applied: bool = False, ): @@ -762,9 +848,9 @@ def create( data_layout: The data_layout specification (default: "NN") q_layout: The quantization axis (default: ROWWISE) flatten_axis: The axis along which the tensor could be flattened to 2D (default: -1) - group_sizes: Array containing the size of each group (default: None) + first_dims: Per-group sizes of the first (row) 2D dim (default: None) + last_dims: Per-group sizes of the last (col) 2D dim (default: None) original_shape: The original shape of the tensor before grouping (default: None) - group_axis: The axis along which grouping is performed (default: 0) rowwise_has_rht_applied: Whether the row-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) colwise_has_rht_applied: Whether the col-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) @@ -785,9 +871,9 @@ def create( dq_dtype, data_layout=data_layout, flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=first_dims, + last_dims=last_dims, original_shape=original_shape, - group_axis=group_axis, rowwise_has_rht_applied=rowwise_has_rht_applied, colwise_has_rht_applied=colwise_has_rht_applied, ) @@ -802,9 +888,9 @@ def create( is_colwise=True, data_layout=data_layout[0], flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=first_dims, + last_dims=last_dims, original_shape=original_shape, - group_axis=group_axis, has_rht_applied=colwise_has_rht_applied, ) @@ -817,9 +903,9 @@ def create( is_colwise=False, data_layout=data_layout[0], flatten_axis=flatten_axis, - group_sizes=group_sizes, + first_dims=first_dims, + last_dims=last_dims, original_shape=original_shape, - group_axis=group_axis, has_rht_applied=rowwise_has_rht_applied, ) From 3af879254ca94c1680a11b136c69ee88a236461f Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 2 Apr 2026 10:17:18 -0700 Subject: [PATCH 303/521] Pass input_output_alias to TritonAutotunedKernelCall (#2814) * Pass input_output_alias to TritonAutotunedKernelCall Signed-off-by: JAX Toolbox * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add jax version guard for the input_output_aliasing fix Signed-off-by: tdophung --------- Signed-off-by: JAX Toolbox Signed-off-by: tdophung Co-authored-by: JAX Toolbox Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../jax/triton_extensions/utils.py | 38 ++++++++++++------- transformer_engine/jax/version_utils.py | 10 +++++ 2 files changed, 34 insertions(+), 14 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 28e3f08e18..ebec1b3cc9 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -43,8 +43,10 @@ import jax.numpy as jnp from ..version_utils import ( + TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION, TRITON_EXTENSION_MIN_JAX_VERSION, is_triton_extension_supported, + jax_version_meet_requirement, ) @@ -474,23 +476,31 @@ def lowering(ctx, x, *, block_size): kernel_calls.append((config_call, str(config))) - # IMPORTANT: We pass an empty tuple for input_output_aliases_with_sizes. - # - # Background: - # 1. jax.ffi.ffi_lowering(operand_output_aliases=...) is a HINT to XLA that an - # output can reuse an input's buffer. XLA may or may not honor this. - # 2. TritonAutotunedKernelCall's input_output_aliases_with_sizes triggers - # save/restore logic during autotuning (see jaxlib/gpu/triton_kernels.cc:630-701). - # - # The problem: The save phase (triton_kernels.cc:632) only saves if buffers[input_idx] == buffers[output_idx], - # but the restore phase (triton_kernels.cc:697-700) unconditionally iterates over all aliases and tries - # to access input_copies[input_idx]. If XLA didn't actually alias the buffers, input_copies[input_idx] doesn't exist, creating an empty vector whose .data() returns nullptr, causing CUDA_ERROR_INVALID_VALUE during the restore memcpy. - # - # WAR: Don't pass aliases to TritonAutotunedKernelCall. + input_output_aliases_with_sizes = () + if input_output_aliases: + if jax_version_meet_requirement(TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION): + num_inputs = len(ctx.avals_in) + aliases = [] + for input_idx, output_idx in input_output_aliases.items(): + aval = ctx.avals_in[input_idx] + size_bytes = aval.size * jnp.dtype(aval.dtype).itemsize + # AutotunedKernelCall expects buffer indices (inputs + outputs). + buffer_output_idx = num_inputs + output_idx + aliases.append((input_idx, buffer_output_idx, size_bytes)) + input_output_aliases_with_sizes = tuple(aliases) + else: + warnings.warn( + f"JAX >= {TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION} is required " + "to safely pass input_output_aliases to TritonAutotunedKernelCall. " + "Passing empty aliases as a workaround (jax-ml/jax#35218).", + UserWarning, + stacklevel=2, + ) + kernel_call = gpu_triton.TritonAutotunedKernelCall( f"{actual_kernel_fn.__name__}_autotuned", kernel_calls, - (), # Empty to avoid buggy save/restore in jaxlib/gpu/triton_kernels.cc + input_output_aliases_with_sizes, ) else: diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index 04b7ff879a..63598481a2 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -25,6 +25,15 @@ 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 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 +# contains entries for aliases where XLA actually shared buffers at runtime. +# Accessing a missing entry produced a null vector → CUDA_ERROR_INVALID_VALUE. +# Fixed by: https://github.com/jax-ml/jax/pull/35218 (merged 2026-03-17, main). +# Ships in JAX 0.9.3 (not yet released as of 2026-03-31). +TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION = "0.9.3" + def is_triton_extension_supported() -> bool: """Return True if the current JAX version supports Triton kernel dispatch. @@ -40,4 +49,5 @@ def is_triton_extension_supported() -> bool: "jax_version_meet_requirement", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", + "TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION", ] From 281ff06405b90329752cbab7bf599bc8866779be Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 2 Apr 2026 10:36:55 -0700 Subject: [PATCH 304/521] Remove integration test for Lightning-Thunder (#2822) Signed-off-by: Tim Moon --- qa/L1_pytorch_thunder_integration/test.sh | 21 --------------------- 1 file changed, 21 deletions(-) delete mode 100644 qa/L1_pytorch_thunder_integration/test.sh diff --git a/qa/L1_pytorch_thunder_integration/test.sh b/qa/L1_pytorch_thunder_integration/test.sh deleted file mode 100644 index 8c3fdc8cdb..0000000000 --- a/qa/L1_pytorch_thunder_integration/test.sh +++ /dev/null @@ -1,21 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -set -x - -: ${THUNDER_PATH:=/opt/pytorch/lightning-thunder} -: ${XML_LOG_DIR:=/logs} -mkdir -p "$XML_LOG_DIR" - -pip3 install pytest==8.1.1 pytest-benchmark==5.1.0 -python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest.xml ${THUNDER_PATH}/thunder/tests/test_transformer_engine_executor.py - -# Check return code -# Note: Return code 5 is fine. Lightning tests are skipped on systems -# without FP8 support and Pytest returns 5 if no tests are run. -RC=$? -if [ ${RC} -eq 5 ]; then - RC=0 -fi -exit ${RC} From 4bf1c1c7f26faa10feda15af745d9b5c3782eda0 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Thu, 2 Apr 2026 12:03:59 -0700 Subject: [PATCH 305/521] Optimize fp8 block scaling Allgather for FSDP2 (#2789) * done Signed-off-by: Varun Thumbe * one review comment form greptile Signed-off-by: Varun Thumbe * instead part of the comment not needed 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 * Update transformer_engine/pytorch/tensor/float8_blockwise_tensor.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * No need to set it to None Remove unnecessary columnwise data and scale inv assignments. 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> --- .../pytorch/tensor/float8_blockwise_tensor.py | 123 +++++++----------- .../pytorch/tensor/float8_tensor.py | 10 +- .../pytorch/tensor/mxfp8_tensor.py | 10 +- 3 files changed, 64 insertions(+), 79 deletions(-) diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index ab496d5a9e..bbfc43e9bb 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -10,7 +10,6 @@ from typing import Any, Optional, Tuple, Union import torch - import transformer_engine_torch as tex from transformer_engine_torch import DType as TE_DType from transformer_engine.common.recipe import Float8BlockScaling, Recipe @@ -625,6 +624,8 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m metadata: Metadata needed for reconstructing the tensor after all-gather. """ # pylint: disable=unused-argument + # PyTorch FSDP2 private API – tested with PyTorch 2.5+; + from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState from transformer_engine.pytorch.distributed import _get_module_fsdp_state if not self._is_2D_scaled: @@ -634,42 +635,38 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m "layout has M in dim1, which is incompatible with FSDP2 dim0 all-gather." ) - block_len = self._quantizer.block_len # 128 - - # Prepare rowwise tensors — for 2D scaling, M is in dim0 of both data and scale_inv, - # so they naturally align with FSDP2's dim0 all-gather. No unpadding needed. - rowwise_data = self._rowwise_data - rowwise_scale_inv = self._rowwise_scale_inv - - # Prepare columnwise tensors — columnwise data is transposed (K, M) and - # columnwise scale_inv is (ceil(K/128), round_up(ceil(M/128), 4)). - # M is in dim1 for both, so we must transpose to put M in dim0 for all-gather. - columnwise_data = self._columnwise_data - columnwise_scale_inv = self._columnwise_scale_inv - - if columnwise_data is not None: - # Transpose (K, shard_M) -> (shard_M, K) so M is in dim0 - columnwise_data = columnwise_data.t().contiguous() - - if columnwise_scale_inv is not None: - # Original shape: (ceil(K/128), round_up(ceil(shard_M/128), 4)) - # Strip padding from dim1 (the M-block dimension), transpose, then all-gather - shard_M = math.prod(self.shape[:-1]) - m_blocks = (shard_M + block_len - 1) // block_len # ceil(shard_M/128) - columnwise_scale_inv = columnwise_scale_inv[:, :m_blocks] # unpad dim1 - columnwise_scale_inv = columnwise_scale_inv.t().contiguous() # (m_blocks, k_blocks) - - # Always send both rowwise and columnwise data. - # Unlike MXFP8 (where both forms share the same shape), Float8Blockwise has - # differently-shaped rowwise (M, K) and columnwise (K, M) data. The GEMM kernel - # needs both forms available to perform forward and backward operations, so we - # cannot optimize by sending only one usage based on forward/backward pass. - rowwise_usage = True - sharded_tensors = (rowwise_data, rowwise_scale_inv) - columnwise_usage = self._quantizer.columnwise_usage - if columnwise_usage: - sharded_tensors += (columnwise_data, columnwise_scale_inv) + if self._rowwise_data is None or self._rowwise_scale_inv is None: + raise RuntimeError( + "Rowwise data must be available for FSDP2 all-gather with 2D block scaling." + ) + fsdp_state = _get_module_fsdp_state(module) + param_group = fsdp_state._fsdp_param_group + if param_group is None: + raise RuntimeError( + "FSDP state for this module has no parameter group; " + "cannot determine reshard_after_forward." + ) + reshard_after_forward = param_group._reshard_after_forward + + # If weights are resharded after forward pass, only the relevant usage + # is needed based on whether it's a forward or backward pass. + # If not resharded, the same all-gathered weights are reused in backward, + # so both usages may be needed. + if reshard_after_forward: + training_state = param_group._training_state + is_backward_pass = training_state == TrainingState.PRE_BACKWARD + rowwise_usage = not is_backward_pass + columnwise_usage = is_backward_pass + else: + rowwise_usage = True + columnwise_usage = self._quantizer.columnwise_usage + + # For 2D block scaling (128x128 blocks), columnwise data and scales are + # the transpose of rowwise data and scales. Only all-gather the rowwise + # tensors; columnwise will be derived locally via _create_columnwise() + # in post_all_gather, halving all-gather communication volume. + sharded_tensors = (self._rowwise_data, self._rowwise_scale_inv) metadata = (self._fp8_dtype, self._is_2D_scaled, rowwise_usage, columnwise_usage) return sharded_tensors, metadata @@ -694,59 +691,35 @@ def fsdp_post_all_gather( """ fp8_dtype, is_2D_scaled, rowwise_usage, columnwise_usage = metadata - # Extract rowwise tensors from all-gather outputs - rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] if rowwise_usage else (None, None) - - # Extract columnwise tensors — they were transposed in pre_all_gather, - # so we need to transpose them back. - columnwise_data, columnwise_scale_inv = ( - all_gather_outputs[-2:] if columnwise_usage else (None, None) - ) - - if columnwise_data is not None: - # All-gathered shape is (full_M, K), transpose back to (K, full_M) - columnwise_data = columnwise_data.t().contiguous() - - if columnwise_scale_inv is not None: - # All-gathered shape is (full_m_blocks, k_blocks), - # transpose back to (k_blocks, full_m_blocks) - columnwise_scale_inv = columnwise_scale_inv.t().contiguous() - # Repad dim1 (M-block dimension) to multiple of 4 for GEMM alignment - current_m_blocks = columnwise_scale_inv.shape[1] - pad_amount = (4 - current_m_blocks % 4) % 4 - if pad_amount > 0: - columnwise_scale_inv = torch.nn.functional.pad( - columnwise_scale_inv, (0, pad_amount) - ) - - # Determine the logical shape from the all-gathered data - if rowwise_data is not None: - data_shape = rowwise_data.shape - else: - # columnwise_data is (K, full_M), logical shape is (full_M, K) - data_shape = (columnwise_data.shape[1], columnwise_data.shape[0]) + # Only rowwise data+scales were all-gathered (columnwise is derived locally). + rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] + data_shape = rowwise_data.shape if out is not None: - # Update existing tensor in-place (subsequent iterations) out._rowwise_data = rowwise_data out._rowwise_scale_inv = rowwise_scale_inv - out._columnwise_data = columnwise_data - out._columnwise_scale_inv = columnwise_scale_inv else: - # Construct new tensor (first iteration). - # Float8BlockwiseQTensor constructor copies the quantizer, - # so the sharded tensor's quantizer remains independent. out = Float8BlockwiseQTensor( shape=data_shape, dtype=param_dtype, fp8_dtype=fp8_dtype, rowwise_data=rowwise_data, rowwise_scale_inv=rowwise_scale_inv, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, + columnwise_data=None, + columnwise_scale_inv=None, quantizer=self._quantizer, is_2D_scaled=is_2D_scaled, ) + + # For 2D block scaling, derive columnwise data and scales from rowwise + # via local fp8 transpose. + if columnwise_usage: + out._create_columnwise() + # remove usages if not needed. + out.update_usage( + rowwise_usage=rowwise_usage, + columnwise_usage=columnwise_usage, + ) out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) return out, all_gather_outputs diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 5f00bc8017..e8284eaa53 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -860,14 +860,20 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m self._quantizer.with_amax_reduction = True fsdp_state = _get_module_fsdp_state(module) - reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward + param_group = fsdp_state._fsdp_param_group + if param_group is None: + raise RuntimeError( + "FSDP state for this module has no parameter group; " + "cannot determine reshard_after_forward." + ) + reshard_after_forward = param_group._reshard_after_forward # If weights are resharded after forward pass, then its enough to set the quantizer usages # based on whether its forward or backward pass for the allgathered weights. # If not resharded after forward pass, the same weights allgathered in forward # are used again in backward and so we dont change the quantizer usages which might need # both rowwise and columnwise usages. if reshard_after_forward: - training_state = fsdp_state._fsdp_param_group._training_state + training_state = param_group._training_state is_backward_pass = training_state == TrainingState.PRE_BACKWARD # In case of hopper/L40, only one of data/transpose is needed # based on forward or backward pass. So setting the quantizer usages appropriately. diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index baff9cc2aa..965f59b320 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -634,7 +634,13 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m # Get FSDP state fsdp_state = _get_module_fsdp_state(module) - reshard_after_forward = fsdp_state._fsdp_param_group._reshard_after_forward + param_group = fsdp_state._fsdp_param_group + if param_group is None: + raise RuntimeError( + "FSDP state for this module has no parameter group; " + "cannot determine reshard_after_forward." + ) + reshard_after_forward = param_group._reshard_after_forward # Remove padding from scale inverses before allgather # Rowwise scale_inv should be divisible by [128,4], columnwise by [4, 128] @@ -662,7 +668,7 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m # are used again in backward. And hence if we need the columnwise data/scale_inv, # we need to send them as well for allgather in forward pass itself. if reshard_after_forward: - training_state = fsdp_state._fsdp_param_group._training_state + training_state = param_group._training_state is_backward_pass = training_state == TrainingState.PRE_BACKWARD # Allgather only the necessary tensors based on forward/backward pass rowwise_usage = not is_backward_pass From b0488694e5eac3b3713cd3afe8e7a980d9e929d6 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Thu, 2 Apr 2026 15:09:30 -0700 Subject: [PATCH 306/521] [PyTorch] Fix bug with PR 2677 (#2819) * cudnn now returns Stats always and Max only with `return_max_logit=true` Signed-off-by: Sudhakar Singh * fix a typo that caused a bug Signed-off-by: Sudhakar Singh * update doc strings Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix more docs Signed-off-by: Sudhakar Singh * fixes from the feedback Signed-off-by: Sudhakar Singh * update cudnn-frontend to v1.19.1 Signed-off-by: Sudhakar Singh * update the cudnn frontend Signed-off-by: Sudhakar Singh * fix a wrong omission Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * bugfix: mask out padding tokens when THD Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fixes from greptile feedback Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor nit Signed-off-by: Sudhakar Singh * fixes from feedback 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> --- .../pytorch/cpp_extensions/fused_attn.py | 39 +++++++++++++++---- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 7653296c78..06bfb6ef3c 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -363,13 +363,38 @@ def fused_attn_fwd( max_tensor = output_tensors[2] amax_dims = (0, 2) if max_tensor.ndim == 3 else (0, 2, 3) - if qkv_format == "thd" and max_tensor.ndim == 4: - # For THD on older cuDNN runtimes or THD on sm120, stats can be [b, h, sq, 1] with padded - # sequence positions. Exclude those padded positions when computing max_logit. - seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=max_tensor.device) - sq_idx = torch.arange(max_tensor.shape[2], device=max_tensor.device).view(1, 1, -1, 1) - valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) - max_tensor = max_tensor.masked_fill(~valid, float("-inf")) + if qkv_format == "thd": + if max_tensor.ndim == 4: + # For THD on cuDNN <= 9.6 or THD on sm120, Max tensor can be [b, h, sq, 1] + # with padded sequence positions. Exclude those padded positions when computing max_logit. + seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(device=max_tensor.device) + sq_idx = torch.arange(max_tensor.shape[2], device=max_tensor.device).view( + 1, 1, -1, 1 + ) + valid = sq_idx < seqlens_q.view(-1, 1, 1, 1) + max_tensor = max_tensor.masked_fill(~valid, float("-inf")) + elif max_tensor.ndim == 3: + if cu_seqlens_q_padded is not None: + # For THD + pad_between_seqs=True + non-sm120 + cuDNN>9.6, Max tensor is [tq, h, 1] + # and padding positions could be uninitialized. Exclude those padded positions when + # computing max_logit. + actual_seqlens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to( + device=max_tensor.device + ) + padded_seqlens = (cu_seqlens_q_padded[1:] - cu_seqlens_q_padded[:-1]).to( + device=max_tensor.device + ) + pad_lens = (padded_seqlens - actual_seqlens).to(device=max_tensor.device) + b = pad_lens.shape[0] + + # Stack [actual, pad] per batch into counts: e.g. [3,1, 3,1, 2,2, 7,1] + counts = torch.stack([actual_seqlens, pad_lens], dim=1).flatten() + # Tile [T, F] per sequence: [T,F, T,F, T,F, T,F] + values = torch.tensor([True, False], device=max_tensor.device).repeat(b) + # Expand: T×3, F×1, T×3, F×1, T×2, F×2, T×7, F×1 → TTTF|TTTF|TTFF|TTTTTTTF + valid = torch.repeat_interleave(values, counts) + # Finally, replace invalid (F) positions with -inf + max_tensor = max_tensor.masked_fill(~valid.view(-1, 1, 1), float("-inf")) # Max -> max_logit [h] max_logit = torch.amax(max_tensor, dim=amax_dims).to(dtype=output_tensors[0].dtype) From 42267ec484c192b1a950659090d1f3e6d2161697 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Fri, 3 Apr 2026 01:56:19 +0200 Subject: [PATCH 307/521] [Common] Persistent Grouped MXFP8 quantization kernel (#2738) * Enabled persistency with WorkID Query feature Signed-off-by: Oleg Goncharov * Added a struct with tunable parameters Signed-off-by: Oleg Goncharov * Added persistency with static scheduling Signed-off-by: Oleg Goncharov * Fixed test cases Signed-off-by: Oleg Goncharov * Ready for benchmarking Signed-off-by: Oleg Goncharov * Fixed out-of-boundary error Signed-off-by: Oleg Goncharov * Tuned kernel parameters Signed-off-by: Oleg Goncharov * Refactoring Signed-off-by: Oleg Goncharov * Refactoring 2 Signed-off-by: Oleg Goncharov * Refactoring 3 Signed-off-by: Oleg Goncharov * Removed the dynamic (WorkID Query) persistency Signed-off-by: Oleg Goncharov * Ready for PR Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixes per the review Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Ready for benchmark Signed-off-by: Oleg Goncharov * Ready for benchmark - Regular kernel Signed-off-by: Oleg Goncharov * Added the source code to the profiler Signed-off-by: Oleg Goncharov * Added constructors to Job and Block descriptors Signed-off-by: Oleg Goncharov * Removed the prefetch overlapping between jobs Signed-off-by: Oleg Goncharov * Cache tensor ID Signed-off-by: Oleg Goncharov * ShapeRepresentation is not a template parameter Signed-off-by: Oleg Goncharov * Removed redundant fence_proxy Signed-off-by: Oleg Goncharov * Refactoring Signed-off-by: Oleg Goncharov * Used mixed precision FMA Signed-off-by: Oleg Goncharov * Added Quantize parameters Signed-off-by: Oleg Goncharov * Added the fast math branch Signed-off-by: Oleg Goncharov * Added the fast math to cpp test suite Signed-off-by: Oleg Goncharov * Align tests Signed-off-by: Oleg Goncharov * Use STS instead of generic ST Signed-off-by: Oleg Goncharov * Add zero-tensor cases Signed-off-by: Oleg Goncharov * Used LDS instead of generic LD in colwise path Signed-off-by: Oleg Goncharov * Used LDS instead of generic LD in rowwise Signed-off-by: Oleg Goncharov * Ready for merge Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Uncommented test cases Signed-off-by: Oleg Goncharov * Added FP16 Fast math path to rowwise processing Signed-off-by: Oleg Goncharov * Refactoring Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed lint 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 * Fixes Signed-off-by: Oleg Goncharov * Fix Signed-off-by: Oleg Goncharov * Fixed test suite Signed-off-by: Oleg Goncharov * Fixed test suite Signed-off-by: Oleg Goncharov * Fixes per the review Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Modifications per the review Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Assert the buffer size Signed-off-by: Oleg Goncharov * Added fast math RCP for bf16 Signed-off-by: Oleg Goncharov * Fast math for BF16 is now default Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed compilation error when compiling on previous archs Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Boundary condition fix Signed-off-by: Oleg Goncharov * Fixed compilation error Signed-off-by: Oleg Goncharov * Refactoring. Moved helpers to core-common Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactoring Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactoring per the review Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Addressed the PR review comments Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed the compilation error when PTX was compiled for CUDA 13.0 Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixed pytorch extensions Signed-off-by: Oleg Goncharov --------- Signed-off-by: Oleg Goncharov Signed-off-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/cpp/operator/test_cast_mxfp8.cu | 1 + tests/cpp/operator/test_cast_mxfp8_grouped.cu | 83 +- tests/cpp/test_common.h | 2 +- transformer_engine/common/cast/cast.cu | 4 +- .../common/cast/core/common.cuh | 412 ++++- .../common/cast/dispatch/quantize.cuh | 4 +- .../common/cast/mxfp8/gated_mxfp8.cuh | 8 +- .../cast/mxfp8/group_quantize_mxfp8.cuh | 1458 ++++++++--------- .../common/cast/mxfp8/quantize_mxfp8.cuh | 4 +- .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 16 +- .../common/cast/nvfp4/quantize_nvfp4.cuh | 2 +- transformer_engine/common/common.h | 44 + .../graph_safe_group_hadamard_transform.cu | 7 - .../common/include/transformer_engine/cast.h | 7 +- .../common/recipe/mxfp8_scaling.cu | 4 +- transformer_engine/common/recipe/nvfp4.cu | 4 +- transformer_engine/common/util/ptx.cuh | 162 +- transformer_engine/common/utils.cuh | 7 + .../pytorch/csrc/extensions/cast.cpp | 3 +- 19 files changed, 1408 insertions(+), 824 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index b5e11c30e1..ccc605c060 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -535,6 +535,7 @@ std::vector> matrix_sizes = { {1024}, {8, 32, 1024}, {16, 8, 4, 512}, + {8192, 7168}, }; std::vector> block_sizes = { diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 09bd21657a..3b097cff43 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -371,7 +371,7 @@ void performTest(const ProcessingMethod processing_method, NVTEShape logical_shape_ = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); - std::vector dbias_logical_shape_vec= {num_tensors, cols}; + std::vector dbias_logical_shape_vec = {num_tensors, cols}; NVTEShape dbias_logical_shape_ = nvte_make_shape(dbias_logical_shape_vec.data(), dbias_logical_shape_vec.size()); @@ -499,11 +499,13 @@ void performTest(const ProcessingMethod processing_method, scales_stride_colwise); } + QuantizationConfigWrapper quant_config; + // GPU Tensor workspace; switch (processing_method) { case ProcessingMethod::CAST_ONLY: { - nvte_group_quantize(in_group_tensor, out_group_tensor, 0); + nvte_group_quantize(in_group_tensor, out_group_tensor, quant_config, 0); break; } case ProcessingMethod::CAST_DBIAS: { @@ -554,6 +556,11 @@ void performTest(const ProcessingMethod processing_method, const double abs_tolerable_mismatches_limit = 0.0; const double rel_tolerable_mismatches_limit = 0.0; + // Compare only allocated contiguous output range. + // In graph-safe mode logical shape may include trailing garbage beyond offsets_h.back(). + const size_t compare_rows = 1; + const size_t compare_cols = elts_num; + if (rowwise) { cudaMemcpy(out_data_rowwise_h.data(), out_data_rowwise_d, out_data_size, cudaMemcpyDeviceToHost); cudaMemcpy(out_scales_rowwise_h.data(), out_scales_rowwise_d, rowwise_scales_size, cudaMemcpyDeviceToHost); @@ -566,7 +573,8 @@ void performTest(const ProcessingMethod processing_method, const size_t mismatches_elts = 32 * mismatches_scales; compare_scaled_elts("rowwise_output", out_data_rowwise_ref.data(), - out_data_rowwise_h.data(), rows, cols, true, mismatches_elts); + out_data_rowwise_h.data(), compare_rows, compare_cols, + true, mismatches_elts); } if (colwise) { @@ -581,7 +589,8 @@ void performTest(const ProcessingMethod processing_method, const size_t mismatches_elts = 32 * mismatches_scales; compare_scaled_elts("colwise_output", out_data_colwise_ref.data(), - out_data_colwise_h.data(), rows, cols, false, mismatches_elts); + out_data_colwise_h.data(), compare_rows, compare_cols, + false, mismatches_elts); } if (compute_dbias) { @@ -652,9 +661,13 @@ std::vector> input_config = { {VARYING_FIRST_DIM, 4, 1024,144, 128,384,0,512}, {VARYING_FIRST_DIM, 4, 1536,160, 128,384,512,512}, {VARYING_FIRST_DIM, 5, 4096,512, 128,256,384,1024,2304}, + {VARYING_FIRST_DIM, 5, 16 * 4096,512, 128,256,384,1024,2304}, {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, {VARYING_BOTH_DIMS, 2, 1,(128*128)+(256*256), 128,256, 128,256}, {VARYING_BOTH_DIMS, 2, 1,(256*128)+(512*640), 256,512, 128,640}, + // Empty tensor in the middle of the group must not terminate the persistent work loop. + {VARYING_FIRST_DIM, 4, 512,160, 128,0,0,256}, + {VARYING_BOTH_DIMS, 3, 1,(128*128)+(128*128), 128,0,128, 128,0,128}, }; } // namespace @@ -808,6 +821,37 @@ std::string to_string(const ActivationKind activation) { } } +std::string MakeGroupedFusedCastMXFP8TestName( + const testing::TestParamInfo& info) { + const ProcessingMethod method = std::get<0>(info.param); + std::string name = to_string(method); + name += "X" + to_string(std::get<1>(info.param)); + + switch (std::get<2>(info.param)) { + case ScalingDirection::ROWWISE: name += "_ROWWISE_"; break; + case ScalingDirection::COLWISE: name += "_COLWISE_"; break; + case ScalingDirection::BOTH: name += "_BIDIMENSIONAL_"; break; + } + + const std::vector input = std::get<3>(info.param); + + switch (static_cast(input[0])) { + case ShapeRepresentation::SAME_BOTH_DIMS: name += "SAME_BOTH_DIMS"; break; + case ShapeRepresentation::VARYING_FIRST_DIM: name += "VARYING_FIRST_DIM"; break; + case ShapeRepresentation::VARYING_LAST_DIM: name += "VARYING_LAST_DIM"; break; + case ShapeRepresentation::VARYING_BOTH_DIMS: name += "VARYING_BOTH_DIMS"; break; + } + + name += "_N_" + std::to_string(input[1]); + + name += "_SHAPE_" + std::to_string(input[2]) + "X" + std::to_string(input[3]); + + name += "_" + test::typeName(std::get<4>(info.param)) + + "_" + test::typeName(std::get<5>(info.param)); + + return name; +} + INSTANTIATE_TEST_SUITE_P( OperatorTest, GroupedFusedCastMXFP8TestSuite, @@ -818,33 +862,4 @@ INSTANTIATE_TEST_SUITE_P( ::testing::ValuesIn(input_config), ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2)), - [](const testing::TestParamInfo& info) { - const ProcessingMethod method = std::get<0>(info.param); - std::string name = to_string(method); - name += "X" + to_string(std::get<1>(info.param)); - - switch (std::get<2>(info.param)) { - case ScalingDirection::ROWWISE: name += "_ROWWISE_"; break; - case ScalingDirection::COLWISE: name += "_COLWISE_"; break; - case ScalingDirection::BOTH: name += "_BIDIMENSIONAL_"; break; - } - - const std::vector input = std::get<3>(info.param); - - switch(static_cast(input[0])) { - case ShapeRepresentation::SAME_BOTH_DIMS: name += "SAME_BOTH_DIMS"; break; - case ShapeRepresentation::VARYING_FIRST_DIM: name += "VARYING_FIRST_DIM"; break; - case ShapeRepresentation::VARYING_LAST_DIM: name += "VARYING_LAST_DIM"; break; - case ShapeRepresentation::VARYING_BOTH_DIMS: name += "VARYING_BOTH_DIMS"; break; - }; - - name += "_N_" + std::to_string(input[1]); - - name += "_SHAPE_" + - std::to_string(input[2]) + - "X" + std::to_string(input[3]); - - name += "_" + test::typeName(std::get<4>(info.param)) + - "_" + test::typeName(std::get<5>(info.param)); - return name; - }); + MakeGroupedFusedCastMXFP8TestName); diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 927407f478..b5a7f26d14 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -322,7 +322,7 @@ constexpr size_t scale_tensor_alignment_Y_colwise = 4; constexpr size_t scale_tensor_alignment_X_colwise = 128; inline size_t divide_round_up(const size_t N, const size_t M) { - return (N - 1 + M) / M; + return ((N + M) - 1) / M; } inline size_t round_up_to_nearest_multiple(const size_t N, const size_t M) { diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 4f9ddb4fc5..dc02390818 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -27,12 +27,12 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea } void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream) { + 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, nullptr, stream); + dispatch::group_quantize_fwd_helper(input, output, quant_config, stream); } void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor noop, diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh index a4e033939b..90e57a6fe8 100644 --- a/transformer_engine/common/cast/core/common.cuh +++ b/transformer_engine/common/cast/core/common.cuh @@ -23,13 +23,18 @@ namespace transformer_engine { namespace dispatch { namespace common { -enum ShapeRepresentation { - SAME_BOTH_DIMS = 0, - VARYING_FIRST_DIM = 1, - VARYING_LAST_DIM = 2, - VARYING_BOTH_DIMS = 3 +constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; + +struct alignas(128) TensorMapStorage { + alignas(128) CUtensorMap input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + alignas(128) CUtensorMap act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + alignas(128) CUtensorMap output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + alignas(128) CUtensorMap output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; }; +// Internal linkage avoids device-link ODR issues when this header is included by multiple .cu TUs. +static __device__ TensorMapStorage g_tensor_maps; + inline bool full_tile_1D_tensor(const Tensor *const t, const size_t elems_per_block) { const size_t N = product(t->data.shape); const bool isFullTile = (N % elems_per_block == 0); @@ -100,14 +105,15 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK) const size_t tensor_id = blockIdx.y; const size_t tensor_rows = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) ? (first_logical_dim / num_tensors) - : first_dims_ptr[tensor_id]; + : static_cast(first_dims_ptr[tensor_id]); const size_t rows = tensor_rows / chunk_dim_Y; const size_t cols = last_logical_dim; - const size_t dbias_in_offset_Y = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) - ? (tensor_id * (tensor_rows / chunk_dim_Y)) - : (offsets_ptr[tensor_id] / cols / chunk_dim_Y); + const size_t dbias_in_offset_Y = + (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) + ? (tensor_id * (tensor_rows / chunk_dim_Y)) + : (static_cast(offsets_ptr[tensor_id]) / cols / chunk_dim_Y); const size_t thread_id = blockIdx.x * blockDim.x + threadIdx.x; @@ -180,6 +186,394 @@ void grouped_reduce_dbias(const ShapeRepresentation shape_rep, const size_t num_ NVTE_CHECK_CUDA(cudaGetLastError()); } +template +__device__ __forceinline__ size_t +get_current_tensor_id(const size_t num_tensors, const size_t current_offset, const size_t block_Y, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr) { + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) { + const size_t current_row = block_Y * CHUNK_DIM_Y; + const size_t rows_per_tensor = first_logical_dim / num_tensors; + return current_row / rows_per_tensor; + } else { + size_t low = 1; + size_t hi = num_tensors; // [low, hi] + + while (low < hi) { + const size_t mid = low + (hi - low) / 2; + const size_t mid_offset = static_cast(offsets_ptr[mid]); + + if (mid_offset <= current_offset) { + low = mid + 1; + } else { + hi = mid; + } + } + return low - 1; + } +} + +template +__device__ __forceinline__ size_t +get_tensor_rows_num(const size_t tensor_id, const size_t first_logical_dim, + const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { + size_t rows_num = 0; + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_LAST_DIM) { + rows_num = first_logical_dim; + } else { + rows_num = static_cast(first_dims_ptr[tensor_id]); + } + if (rows_num % 128 != 0) { + NVTE_DEVICE_ERROR("First dimension of each tensor in a group must be divisible by 128."); + } + return rows_num; +} + +__device__ __forceinline__ size_t get_tensor_rows_num( + const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t first_logical_dim, + const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + return get_tensor_rows_num(tensor_id, first_logical_dim, + first_dims_ptr, num_tensors); + case ShapeRepresentation::VARYING_FIRST_DIM: + return get_tensor_rows_num( + tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + case ShapeRepresentation::VARYING_LAST_DIM: + return get_tensor_rows_num( + tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + case ShapeRepresentation::VARYING_BOTH_DIMS: + return get_tensor_rows_num( + tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + } + return 0; +} + +template +__device__ __forceinline__ size_t +get_tensor_cols_num(const size_t tensor_id, const size_t last_logical_dim, + const int64_t *const __restrict__ last_dims_ptr) { + size_t cols_num = 0; + if constexpr (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) { + cols_num = last_logical_dim; + } else { + cols_num = static_cast(last_dims_ptr[tensor_id]); + if (cols_num % 128 != 0) { + NVTE_DEVICE_ERROR( + "For varying last dimensions support, the last dimension of each tensor in a group " + "must be divisible by 128."); + } + } + return cols_num; +} + +__device__ __forceinline__ size_t get_tensor_cols_num( + const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t last_logical_dim, + const int64_t *const __restrict__ last_dims_ptr) { + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + return get_tensor_cols_num(tensor_id, last_logical_dim, + last_dims_ptr); + case ShapeRepresentation::VARYING_FIRST_DIM: + return get_tensor_cols_num( + tensor_id, last_logical_dim, last_dims_ptr); + case ShapeRepresentation::VARYING_LAST_DIM: + return get_tensor_cols_num(tensor_id, last_logical_dim, + last_dims_ptr); + case ShapeRepresentation::VARYING_BOTH_DIMS: + return get_tensor_cols_num( + tensor_id, last_logical_dim, last_dims_ptr); + } + return 0; +} + +// Logical work-item decoded from CTA coordinates. +struct JobDescriptor { + size_t block_id = 0; + size_t block_global_offset = 0; + size_t tensor_id = 0; + size_t rows = 0; + size_t cols = 0; + + __host__ __device__ __forceinline__ constexpr JobDescriptor() = default; + + __host__ __device__ __forceinline__ constexpr JobDescriptor(const size_t block_id_, + const size_t block_global_offset_, + const size_t tensor_id_, + const size_t rows_, + const size_t cols_) + : block_id(block_id_), + block_global_offset(block_global_offset_), + tensor_id(tensor_id_), + rows(rows_), + cols(cols_) {} +}; + +// Tensor-local coordinates for a work-item. +struct BlockDescriptor { + size_t tensor_base = 0; + size_t block_id_in_current_tensor = 0; + size_t block_id_Y = 0; + size_t block_id_X = 0; + size_t block_offset_Y = 0; + size_t block_offset_X = 0; + + __host__ __device__ __forceinline__ constexpr BlockDescriptor() = default; + + __host__ __device__ __forceinline__ constexpr BlockDescriptor( + const size_t tensor_base_, const size_t block_id_in_current_tensor_, const size_t block_id_Y_, + const size_t block_id_X_, const size_t block_offset_Y_, const size_t block_offset_X_) + : tensor_base(tensor_base_), + block_id_in_current_tensor(block_id_in_current_tensor_), + block_id_Y(block_id_Y_), + block_id_X(block_id_X_), + block_offset_Y(block_offset_Y_), + block_offset_X(block_offset_X_) {} +}; + +template +__device__ __forceinline__ JobDescriptor decode_job( + const size_t num_tensors, const size_t first_logical_dim, const size_t last_logical_dim, + const size_t work_blocks_X, const int32_t ctaid_X, const int32_t ctaid_Y, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr) { + constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; + constexpr bool is_single_tensor = (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM); + const size_t block_id = ctaid_Y * work_blocks_X + ctaid_X; + const size_t block_global_offset = + is_single_tensor ? (ctaid_Y * CHUNK_DIM_Y * last_logical_dim + ctaid_X * CHUNK_DIM_X) + : (block_id * ELTS_PER_CHUNK); + const size_t tensor_id = get_current_tensor_id( + num_tensors, block_global_offset, ctaid_Y, first_logical_dim, last_logical_dim, offsets_ptr); + const size_t rows = + get_tensor_rows_num(tensor_id, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, last_logical_dim, last_dims_ptr); + return JobDescriptor(block_id, block_global_offset, tensor_id, rows, cols); +} + +template +__device__ __forceinline__ bool is_job_valid(const JobDescriptor &job, + const size_t total_work_blocks, + const int64_t *const __restrict__ offsets_ptr) { + const bool is_valid = (job.block_id < total_work_blocks); + if (!is_valid) { + return false; + } + if (job.rows == 0 || job.cols == 0) { + return true; + } + if constexpr (SHAPE_REP == SAME_BOTH_DIMS) { + return true; + } + + const size_t tensor_start_offset = static_cast(offsets_ptr[job.tensor_id]); + const size_t tensor_end_offset = static_cast(offsets_ptr[job.tensor_id + 1]); + if (job.block_global_offset >= tensor_end_offset) { + return false; + } + + const size_t tensor_offset_from_start = job.block_global_offset - tensor_start_offset; + const size_t block_offset_Y_in_tensor = tensor_offset_from_start / job.cols; + if (block_offset_Y_in_tensor >= job.rows) { + return false; + } + + return true; +} + +__device__ __forceinline__ bool job_has_work(const JobDescriptor &job) { + return job.rows != 0 && job.cols != 0; +} + +__device__ __forceinline__ void advance_to_next_job(bool &job_finished, int32_t &ctaid_X, + int32_t &ctaid_Y, size_t &static_next_block_id, + const size_t static_block_stride, + const size_t total_work_blocks, + const size_t work_blocks_X) { + if (static_next_block_id < total_work_blocks) { + ctaid_X = static_cast(static_next_block_id % work_blocks_X); + ctaid_Y = static_cast(static_next_block_id / work_blocks_X); + static_next_block_id += static_block_stride; + } else { + job_finished = true; + } +} + +template +__device__ __forceinline__ BlockDescriptor +decode_block(const JobDescriptor &job, const int64_t *const __restrict__ offsets_ptr) { + constexpr bool is_single_tensor = (SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS || + SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM); + constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; + const size_t blocks_X_num_in_current_tensor = DIVUP(job.cols, CHUNK_DIM_X); + const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[job.tensor_id]); + const size_t block_id_in_current_tensor = + is_single_tensor ? job.block_id : (job.block_id - tensor_base / ELTS_PER_CHUNK); + const size_t block_id_Y = block_id_in_current_tensor / blocks_X_num_in_current_tensor; + const size_t block_id_X = block_id_in_current_tensor % blocks_X_num_in_current_tensor; + const size_t block_offset_Y = block_id_Y * CHUNK_DIM_Y; + const size_t block_offset_X = block_id_X * CHUNK_DIM_X; + return BlockDescriptor(tensor_base, block_id_in_current_tensor, block_id_Y, block_id_X, + block_offset_Y, block_offset_X); +} + +// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index +__device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, + CUtensorMap *global_tensor_map, + const uintptr_t global_data_ptr, + const size_t global_dim_Y, + const size_t global_dim_X, + const size_t data_type_size_bytes) { + __shared__ CUtensorMap shared_tensor_map; + shared_tensor_map = base_tensor_map; // Copy the base tensor map into shmem + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; + if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); + } + if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { + NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); + } + + asm volatile( + "{\n\t" + ".reg.b64 tensor_map_ptr; \n\t" + "mov.b64 tensor_map_ptr, %0; \n\t" + "tensormap.replace.tile.global_address.b1024.b64 [tensor_map_ptr], %1; \n\t" + "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 1, %2; \n\t" // DIM Y + "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 0, %3; \n\t" // DIM X + "tensormap.replace.tile.global_stride.b1024.b64 [tensor_map_ptr], 0, %4; \n" + "}\n" ::"l"(reinterpret_cast(&shared_tensor_map)), + "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), + "r"(static_cast(global_dim_X)), "l"(static_cast(global_stride_bytes)) + : "memory"); + *global_tensor_map = shared_tensor_map; + } else { + NVTE_DEVICE_ERROR("tensormap.replace is architecture-specific. "); + } +} + +template +__global__ void __launch_bounds__(1) + update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input, + const __grid_constant__ CUtensorMap base_tensor_map_act_input, + const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise, + const __grid_constant__ CUtensorMap base_tensor_map_output_colwise, + const IType *const __restrict__ input_data_ptr, + const IType *const __restrict__ act_input_data_ptr, + const OType *const __restrict__ output_rowwise_data_ptr, + const OType *const __restrict__ output_colwise_data_ptr, + const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, + const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, const bool rowwise, + const bool colwise, const bool compute_dactivations) { + const size_t tensor_id = blockIdx.x; + const size_t rows = + get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + + const size_t offset_elts = offsets_ptr[tensor_id]; + + // Zero-sized groups: skip TMA descriptor update. The main kernel already returns + // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension + // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. + if (rows == 0 || cols == 0) { + return; + } + + if (tensor_id < num_tensors) { + { + CUtensorMap *modified_tensor_map_input = &g_tensor_maps.input[tensor_id]; + const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_input, modified_tensor_map_input, global_data_ptr, + rows, cols, sizeof(IType)); + } + if (compute_dactivations) { + CUtensorMap *modified_tensor_map_act_input = &g_tensor_maps.act_input[tensor_id]; + const uintptr_t global_data_ptr = + reinterpret_cast(act_input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_act_input, modified_tensor_map_act_input, + global_data_ptr, rows, cols, sizeof(IType)); + } + if (rowwise) { + CUtensorMap *modified_tensor_map_output_rowwise = &g_tensor_maps.output_rowwise[tensor_id]; + const uintptr_t global_data_ptr = + reinterpret_cast(output_rowwise_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output_rowwise, modified_tensor_map_output_rowwise, + global_data_ptr, rows, cols, sizeof(OType)); + } + if (colwise) { + CUtensorMap *modified_tensor_map_output_colwise = &g_tensor_maps.output_colwise[tensor_id]; + const uintptr_t global_data_ptr = + reinterpret_cast(output_colwise_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output_colwise, modified_tensor_map_output_colwise, + global_data_ptr, rows, cols, sizeof(OType)); + } + } +} + +__device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map)); +#else + NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + +// Issue TMA global->shared transfer for one stage of input (and optional activation input). +template +__device__ __forceinline__ void prefetch_input_stage( + IType *in_sh, IType *act_in_sh, const CUtensorMap &tensor_map_input, + const CUtensorMap &tensor_map_act_input, const size_t global_offset_X, + const size_t global_offset_Y, const size_t buff_offset, const size_t shmem_buff_size, + uint64_t *barrier, const bool leading_thread) { + if (leading_thread) { + ptx::mbarrier_arrive_expect_tx(barrier, shmem_buff_size); + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&in_sh[buff_offset]), + reinterpret_cast(&tensor_map_input), global_offset_X, global_offset_Y, + barrier); + if constexpr (IS_DACT) { + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&act_in_sh[buff_offset]), + reinterpret_cast(&tensor_map_act_input), global_offset_X, + global_offset_Y, barrier); + } + } +} + +// Issue TMA shared->global transfer for one stage of outputs. +template +__device__ __forceinline__ void store_output_stage( + OType *out_rowwise_data_sh, OType *out_colwise_data_sh, + const CUtensorMap &tensor_map_output_rowwise, const CUtensorMap &tensor_map_output_colwise, + const size_t global_offset_X, const size_t global_offset_Y, const size_t buff_offset, + const bool leading_thread) { + if (!leading_thread) { + return; + } + + 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])); + } + 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])); + } + if constexpr (ROWWISE_SCALING || COLWISE_SCALING) { + ptx::cp_async_bulk_commit_group(); + } +} + } // namespace common } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index f7823b4c58..8d985f64f3 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -409,7 +409,7 @@ void group_quantize_fwd_helper(const NVTEGroupedTensor input, NVTEGroupedTensor case NVTE_MXFP8_1D_SCALING: { mxfp8::group_quantize( input_tensor, activations_tensor, noop_tensor, output_tensor, dbias_tensor, - workspace_tensor, stream); + workspace_tensor, &quant_config_cpp, stream); break; } default: @@ -450,7 +450,7 @@ void group_quantize_bwd_helper(const NVTEGroupedTensor grad, const NVTEGroupedTe case NVTE_MXFP8_1D_SCALING: { mxfp8::group_quantize( grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - stream); + &quant_config_cpp, stream); break; } default: diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index dc9a190e1f..49169a4e14 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -374,7 +374,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) scales_colwise[scale_idx] = biased_exponent_act; } - float block_scale_inverse_act = ptx::exp2f_rcp(biased_exponent_act); + float block_scale_inverse_act = ptx::exp2f_rcp(biased_exponent_act); float block_scale_inverse_gate; if constexpr (IS_BWD) { @@ -392,7 +392,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if (tid_Y_colwise == 0 && (!out_of_bounds_colwise)) { scales_colwise[scale_idx_gate] = biased_exponent_gate; } - block_scale_inverse_gate = ptx::exp2f_rcp(biased_exponent_gate); + block_scale_inverse_gate = ptx::exp2f_rcp(biased_exponent_gate); } // 3. Scale elements @@ -584,7 +584,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) scales_rowwise[scale_idx] = biased_exponent_act; } - const float block_scale_inverse_act = ptx::exp2f_rcp(biased_exponent_act); + const float block_scale_inverse_act = ptx::exp2f_rcp(biased_exponent_act); const ptx::floatx2 block_scale_inverse_2x_act = {block_scale_inverse_act, block_scale_inverse_act}; @@ -606,7 +606,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if (!out_of_bounds_rowwise) { scales_rowwise[scale_idx_gate] = biased_exponent_gate; } - block_scale_inverse_gate = ptx::exp2f_rcp(biased_exponent_gate); + block_scale_inverse_gate = ptx::exp2f_rcp(biased_exponent_gate); block_scale_inverse_2x_gate = {block_scale_inverse_gate, block_scale_inverse_gate}; } diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index d0d15d8d6c..ce6917aa42 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -17,6 +17,7 @@ #include #include "../../common.h" +#include "../../util/cuda_runtime.h" #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" @@ -30,331 +31,447 @@ namespace group_quantize_kernel { using namespace dispatch::common; -constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; -__device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; -__device__ alignas(128) CUtensorMap g_tensor_maps_act_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; -__device__ alignas(128) CUtensorMap g_tensor_maps_output_rowwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; -__device__ alignas(128) CUtensorMap g_tensor_maps_output_colwise[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; +struct TunableConfig { + static constexpr uint CHUNK_DIM_Y = 128; + static constexpr uint CHUNK_DIM_X = 128; + static constexpr uint THREADS_PER_CHUNK = 128; + // Launch static persistent grid as (SM_count * STATIC_PERSISTENT_BLOCKS_PER_SM, 1, 1). + static constexpr uint STATIC_PERSISTENT_BLOCKS_PER_SM = 24; +}; + +static_assert(TunableConfig::STATIC_PERSISTENT_BLOCKS_PER_SM > 0, + "STATIC_PERSISTENT_BLOCKS_PER_SM must be greater than zero in persistent mode."); constexpr size_t SCALE_DIM_Y = 32; constexpr size_t SCALE_DIM_X = 32; -constexpr size_t BUFFS_NUM = 2; -constexpr size_t PACK_SIZE = 4; -constexpr size_t WAVES = SCALE_DIM_X / PACK_SIZE; +constexpr uint PREFETCH_STAGES = 1; +constexpr uint BUFFS_NUM = PREFETCH_STAGES + 1; +constexpr uint PACK_SIZE = 4; +constexpr uint WAVES = SCALE_DIM_X / PACK_SIZE; -constexpr size_t CHUNK_DIM_Y = 128; -constexpr size_t CHUNK_DIM_X = 128; -constexpr size_t THREADS_PER_CHUNK = 128; +constexpr uint CHUNK_DIM_Y = TunableConfig::CHUNK_DIM_Y; +constexpr uint CHUNK_DIM_X = TunableConfig::CHUNK_DIM_X; +constexpr uint THREADS_PER_CHUNK = TunableConfig::THREADS_PER_CHUNK; constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; -constexpr size_t THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; -constexpr size_t THREADS_Y = THREADS_PER_CHUNK / THREADS_X; +constexpr uint THREADS_X = CHUNK_DIM_X / SCALE_DIM_X; +constexpr uint THREADS_Y = THREADS_PER_CHUNK / THREADS_X; -constexpr size_t BUFF_DIM_Y = THREADS_Y; -constexpr size_t BUFF_DIM_X = CHUNK_DIM_X; -constexpr size_t BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; +constexpr uint BUFF_DIM_Y = THREADS_Y; +constexpr uint BUFF_DIM_X = CHUNK_DIM_X; +constexpr uint BUFF_DIM = BUFF_DIM_Y * BUFF_DIM_X; static_assert(BUFF_DIM_Y == 32); -constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; +constexpr uint STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; static_assert(STAGES >= 1); +static_assert(CHUNK_DIM_Y % BUFF_DIM_Y == 0); +static_assert(CHUNK_DIM_Y % SCALE_DIM_Y == 0); +static_assert(CHUNK_DIM_X % SCALE_DIM_X == 0); + // Number of 1-byte elements that span 32 banks (4-byte each) of shared memory -constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4) / 1; // 128 +constexpr uint TOTAL_BANKS_WIDTH = (32 * 4) / 1; // 128 // 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; // 4 = 128 / 32 - -__device__ __forceinline__ size_t get_current_tensor_id( - const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, - const size_t block_Y, const size_t first_logical_dim, const size_t last_logical_dim, - const int64_t *const __restrict__ offsets_ptr) { - if (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS) { - const size_t current_row = block_Y * CHUNK_DIM_Y; - const size_t rows_per_tensor = first_logical_dim / num_tensors; - return current_row / rows_per_tensor; - } else { - size_t low = 1; - size_t hi = num_tensors; // [low, hi] +constexpr uint THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 / 32 - while (low < hi) { - const size_t mid = low + (hi - low) / 2; - const size_t mid_offset = static_cast(offsets_ptr[mid]); +template +__device__ __forceinline__ void process_colwise_stage( + const size_t buff, const int stage, const size_t tid_X_colwise, + const size_t scales_offset_Y_colwise, const size_t scales_offset_X_colwise, + const size_t scale_stride_colwise, const size_t tensor_base_for_scales, const size_t rows, + const size_t cols, IType *sIn_ptr, IType *sActIn_ptr, IType *sCachedAct_ptr, + OType *sOutColwise_ptr, e8m0_t *scales_colwise, float &partial_dbias_colwise) { + using IType2 = typename ptx::FPx2; + using IType4 = typename ptx::FPx4; + using OType4 = typename ptx::FPx4; + using IType3D = IType[BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; + using OType3D = OType[BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; - if (mid_offset <= current_offset) { - low = mid + 1; - } else { - hi = mid; - } - } - return low - 1; - } -} + const auto &sIn = *reinterpret_cast(sIn_ptr); + const auto &sActIn = *reinterpret_cast(sActIn_ptr); + auto &sCachedAct = *reinterpret_cast(sCachedAct_ptr); + auto &sOutColwise = *reinterpret_cast(sOutColwise_ptr); -__device__ __forceinline__ size_t get_tensor_rows_num( - const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t first_logical_dim, - const int64_t *const __restrict__ first_dims_ptr, const size_t num_tensors) { - size_t rows_num = 0; - switch (shape_rep) { - case ShapeRepresentation::SAME_BOTH_DIMS: - case ShapeRepresentation::VARYING_LAST_DIM: - rows_num = first_logical_dim; - break; - case ShapeRepresentation::VARYING_FIRST_DIM: - case ShapeRepresentation::VARYING_BOTH_DIMS: - rows_num = static_cast(first_dims_ptr[tensor_id]); - break; - } - if (rows_num % 128 != 0) { - NVTE_DEVICE_ERROR("First dimension of each tensor in a group must be divisible by 128."); - } - return rows_num; -} + constexpr uint32_t IN_SHMEM_STRIDE = static_cast(BUFF_DIM_X * sizeof(IType)); + constexpr uint32_t OUT_SHMEM_STRIDE = static_cast(BUFF_DIM_X * sizeof(OType)); -__device__ __forceinline__ size_t get_tensor_cols_num( - const size_t tensor_id, const ShapeRepresentation shape_rep, const size_t last_logical_dim, - const int64_t *const __restrict__ last_dims_ptr) { - size_t cols_num = 0; - switch (shape_rep) { - case ShapeRepresentation::SAME_BOTH_DIMS: - case ShapeRepresentation::VARYING_FIRST_DIM: - cols_num = last_logical_dim; - break; - case ShapeRepresentation::VARYING_LAST_DIM: - case ShapeRepresentation::VARYING_BOTH_DIMS: - cols_num = static_cast(last_dims_ptr[tensor_id]); - break; + constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; + constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING; + constexpr bool FP16_CAST_ONLY = NO_ACTIVATIONS && (!IS_DBIAS) && std::is_same_v; + constexpr bool BF16_CAST_ONLY = NO_ACTIVATIONS && (!IS_DBIAS) && std::is_same_v; + + const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; + const size_t global_scales_offset_X = scales_offset_X_colwise; + + 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 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( + global_scales_offset_X, local_scales_offset_Y, + DIVUP(rows, static_cast(scale_tensor_alignment_Y_rowwise))); + } else { + scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; } - return cols_num; -} -// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index -__device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map, - CUtensorMap *global_tensor_map, - const uintptr_t global_data_ptr, - const size_t global_dim_Y, - const size_t global_dim_X, - const size_t data_type_size_bytes) { - __shared__ CUtensorMap shared_tensor_map; - shared_tensor_map = base_tensor_map; // Copy the base tensor map into shmem - constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; - if constexpr (is_blackwell) { - const size_t global_stride_bytes = global_dim_X * data_type_size_bytes; - if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned."); - } - if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) { - NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned"); + const size_t j = tid_X_colwise; + + if constexpr (BF16_CAST_ONLY) { + IType4 rIn4x[BUFF_DIM_Y / 4]; + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; i += 4) { + const uint32_t src_smem_ptr = __cvta_generic_to_shared(&sIn[buff][i][j]); + + // Load 4x elts S2R and find amax + asm volatile( + "{\n" + ".reg.u32 base_offset, stride; \n\t" + "mov.u32 base_offset, %2; \n\t" + "mov.u32 stride, %3; \n\t" + ".reg.u32 ptr0,ptr1,ptr2,ptr3; \n\t" + "mad.lo.u32 ptr0, 0, stride, base_offset; \n\t" + "mad.lo.u32 ptr1, 1, stride, base_offset; \n\t" + "mad.lo.u32 ptr2, 2, stride, base_offset; \n\t" + "mad.lo.u32 ptr3, 3, stride, base_offset; \n\t" + ".reg.b16 x0,x1,x2,x3; \n\t" + "ld.shared.b16 x0, [ptr0]; \n\t" + "ld.shared.b16 x1, [ptr1]; \n\t" + "ld.shared.b16 x2, [ptr2]; \n\t" + "ld.shared.b16 x3, [ptr3]; \n\t" + "mov.b64 %0, {x0,x1,x2,x3}; \n\t" + ".reg.b32 x01,x23; \n\t" + "mov.b32 x01, {x0,x1}; \n\t" + "mov.b32 x23, {x2,x3}; \n\t" + "max.xorsign.abs.bf16x2 x01, x01, x23; \n\t" + "max.xorsign.abs.bf16x2 %1, %1, x01; \n" + "}\n" + : "=l"(reinterpret_cast(rIn4x[i / 4])), + "+r"(reinterpret_cast(thread_amax_2x)) + : "r"(src_smem_ptr), "r"(IN_SHMEM_STRIDE)); } + const float thread_amax = + static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); - asm volatile( - "{\n\t" - ".reg.b64 tensor_map_ptr; \n\t" - "mov.b64 tensor_map_ptr, %0; \n\t" - "tensormap.replace.tile.global_address.b1024.b64 [tensor_map_ptr], %1; \n\t" - "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 1, %2; \n\t" // DIM Y - "tensormap.replace.tile.global_dim.b1024.b32 [tensor_map_ptr], 0, %3; \n\t" // DIM X - "tensormap.replace.tile.global_stride.b1024.b64 [tensor_map_ptr], 0, %4; \n" - "}\n" ::"l"(reinterpret_cast(&shared_tensor_map)), - "l"(global_data_ptr), "r"(static_cast(global_dim_Y)), - "r"(static_cast(global_dim_X)), "l"(static_cast(global_stride_bytes)) - : "memory"); - *global_tensor_map = shared_tensor_map; + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + scales_colwise[scale_idx] = biased_exponent; + + const bf16 block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const ptx::bf16x2 block_scale_inverse_bf16_x2 = {block_scale_inverse, block_scale_inverse}; +#pragma unroll + for (int i = 0; i < SCALE_DIM_Y; i += 4) { + OType4 out; + ptx::mul_cvt_4x(out, rIn4x[i / 4], block_scale_inverse_bf16_x2); + + const uint32_t dst_smem_ptr = __cvta_generic_to_shared(&sOutColwise[buff][i][j]); + + asm volatile( + "{\n" + ".reg.u32 base_offset, stride; \n\t" + "mov.u32 base_offset, %0; \n\t" + "mov.u32 stride, %1; \n\t" + ".reg.u32 ptr0,ptr1,ptr2,ptr3; \n\t" + "mad.lo.u32 ptr0, 0, stride, base_offset; \n\t" + "mad.lo.u32 ptr1, 1, stride, base_offset; \n\t" + "mad.lo.u32 ptr2, 2, stride, base_offset; \n\t" + "mad.lo.u32 ptr3, 3, stride, base_offset; \n\t" + ".reg.b8 x0,x1,x2,x3; \n\t" + "mov.b32 {x0,x1,x2,x3}, %2; \n\t" + "st.shared.b8 [ptr0], x0; \n\t" + "st.shared.b8 [ptr1], x1; \n\t" + "st.shared.b8 [ptr2], x2; \n\t" + "st.shared.b8 [ptr3], x3; \n" + "}\n" ::"r"(dst_smem_ptr), + "r"(OUT_SHMEM_STRIDE), "r"(reinterpret_cast(out))); + } } else { - NVTE_DEVICE_ERROR( - "tensormap.replace is architecture-specific. " - "Try recompiling with sm_XXXa instead of sm_XXX."); + float rInCompute[BUFF_DIM_Y]; + IType rIn[BUFF_DIM_Y]; + float thread_amax = 0.0f; + + if constexpr (FP16_CAST_ONLY) { + IType thread_amax_f16 = static_cast(0.0f); +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + rIn[i] = sIn[buff][i][j]; + thread_amax_f16 = __hmax(thread_amax_f16, __habs(rIn[i])); + } + thread_amax = static_cast(thread_amax_f16); + } else { +#pragma unroll + for (int i = 0; i < BUFF_DIM_Y; ++i) { + float elt = static_cast(sIn[buff][i][j]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + float act_in_elt = static_cast(sActIn[buff][i][j]); + elt *= OP(act_in_elt, {}); + } + if constexpr (IS_DBIAS) { + partial_dbias_colwise += elt; + } + if constexpr (!std::is_same_v) { + elt = static_cast(static_cast(elt)); + } + if constexpr (IS_CACHED_ACT_OP) { + sCachedAct[buff][i][j] = static_cast(elt); + } + thread_amax = fmaxf(thread_amax, fabsf(elt)); + rInCompute[i] = elt; + } + } + + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + scales_colwise[scale_idx] = biased_exponent; + + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); +#pragma unroll + for (int i = 0; i < SCALE_DIM_Y; ++i) { + float in; + if constexpr (FP16_CAST_ONLY) { + in = static_cast(rIn[i]); + } else { + in = rInCompute[i]; + } + const float scaled_out = in * block_scale_inverse; + + sOutColwise[buff][i][j] = static_cast(scaled_out); + } } } -template -__global__ void update_tma_descriptors( - const __grid_constant__ CUtensorMap base_tensor_map_input, - const __grid_constant__ CUtensorMap base_tensor_map_act_input, - const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise, - const __grid_constant__ CUtensorMap base_tensor_map_output_colwise, - const IType *const __restrict__ input_data_ptr, - const IType *const __restrict__ act_input_data_ptr, - const OType *const __restrict__ output_rowwise_data_ptr, - const OType *const __restrict__ output_colwise_data_ptr, const ShapeRepresentation shape_rep, - const size_t num_tensors, const size_t first_logical_dim, const size_t last_logical_dim, - const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, - const int64_t *const __restrict__ last_dims_ptr, const bool rowwise, const bool colwise, - const bool compute_dactivations) { - const bool leading_thread = (threadIdx.x == 0); - const size_t tensor_id = blockIdx.x; +template +__device__ __forceinline__ void process_rowwise_stage( + const size_t buff, const size_t stage_offset_Y, const size_t thread_offset_Y_rowwise, + const size_t thread_offset_X_rowwise, const int bank_group, + const size_t scales_offset_Y_rowwise, const size_t scales_offset_X_rowwise, + const size_t scale_stride_rowwise, const bool rowwise_scale_is_within_bounds, const size_t cols, + IType *sIn_ptr, IType *sActIn_ptr, IType *sCachedAct_ptr, OType *sOutRowwise_ptr, + e8m0_t *scales_rowwise, float *thread_dbias_rowwise) { + using IType2 = typename ptx::FPx2; + using IType4 = typename ptx::FPx4; + using OType2 = typename ptx::FPx2; + using OType4 = typename ptx::FPx4; + constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; + constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; + constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && COLWISE_SCALING; + constexpr bool BF16_CAST_ONLY = NO_ACTIVATIONS && (!IS_DBIAS) && std::is_same_v; + constexpr bool FP16_CAST_ONLY = NO_ACTIVATIONS && (!IS_DBIAS) && std::is_same_v; + constexpr bool NON_FP32_CAST_ONLY = BF16_CAST_ONLY || FP16_CAST_ONLY; - const size_t rows = - get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); - const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + using IType3D = IType[BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; + using OType3D = OType[BUFFS_NUM][BUFF_DIM_Y][BUFF_DIM_X]; - // Zero-sized groups: skip TMA descriptor update. The main kernel already returns - // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension - // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. - if (rows == 0 || cols == 0) { - return; - } + const auto &sIn = *reinterpret_cast(sIn_ptr); + const auto &sActIn = *reinterpret_cast(sActIn_ptr); + const auto &sCachedAct = *reinterpret_cast(sCachedAct_ptr); + auto &sOutRowwise = *reinterpret_cast(sOutRowwise_ptr); + + const size_t i = thread_offset_Y_rowwise; - const size_t offset_elts = offsets_ptr[tensor_id]; + float thread_amax = 0.0f; + float rInCompute[SCALE_DIM_X]; + Vec rInCached[WAVES]; + Vec rIn[WAVES]; + IType4 rIn4x[WAVES]; - if (leading_thread && (tensor_id < num_tensors)) { - { - const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id], - global_data_ptr, rows, cols, sizeof(IType)); + if constexpr (NON_FP32_CAST_ONLY) { + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t j = thread_offset_X_rowwise + swizzled_group_idx; + if constexpr (std::is_same_v) { + const uint32_t src_smem_ptr = __cvta_generic_to_shared(&sIn[buff][i][j]); + // Load 4x elts S2R and find amax + asm volatile( + "{\n" + "ld.shared.b64 %0, [%2]; \n\t" + ".reg.b32 x01,x23; \n\t" + "mov.b64 {x01, x23}, %0; \n\t" + "max.xorsign.abs.bf16x2 x01, x01, x23; \n\t" + "max.xorsign.abs.bf16x2 %1, %1, x01; \n" + "}\n" + : "=l"(reinterpret_cast(rIn4x[w])), + "+r"(reinterpret_cast(thread_amax_2x)) + : "r"(src_smem_ptr)); + } else { + // rIn[w].load_from(&sIn_ptr[shmem_offset_rowwise]); + rIn[w].load_from(&sIn[buff][i][j]); +#pragma unroll + for (int e = 0; e < PACK_SIZE / 2; ++e) { + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, rIn[w].data.elt[e]); + } + } } - if (compute_dactivations) { - const uintptr_t global_data_ptr = - reinterpret_cast(act_input_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_act_input, &g_tensor_maps_act_input[tensor_id], - global_data_ptr, rows, cols, sizeof(IType)); + thread_amax = static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); + } else if constexpr (IS_CACHED_ACT_OP) { + __syncthreads(); + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t j = thread_offset_X_rowwise + swizzled_group_idx; + rInCached[w].load_from(&sCachedAct[buff][i][j]); + if constexpr (std::is_same_v) { +#pragma unroll + for (int e = 0; e < PACK_SIZE; ++e) { + thread_amax = fmaxf(thread_amax, fabsf(rInCached[w].data.elt[e])); + } + } else { +#pragma unroll + for (int e = 0; e < PACK_SIZE; e += 2) { + const IType2 in_cached_2x = {rInCached[w].data.elt[e], rInCached[w].data.elt[e + 1]}; + ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, in_cached_2x); + } + } } - if (rowwise) { - const uintptr_t global_data_ptr = - reinterpret_cast(output_rowwise_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_output_rowwise, - &g_tensor_maps_output_rowwise[tensor_id], global_data_ptr, rows, cols, - sizeof(OType)); + if constexpr (!std::is_same_v) { + thread_amax = static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); } - if (colwise) { - const uintptr_t global_data_ptr = - reinterpret_cast(output_colwise_data_ptr + offset_elts); - modify_base_tensor_map(base_tensor_map_output_colwise, - &g_tensor_maps_output_colwise[tensor_id], global_data_ptr, rows, cols, - sizeof(OType)); + } else { +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t j = thread_offset_X_rowwise + swizzled_group_idx; + + Vec in; + Vec act_in; + + in.load_from(&sIn[buff][i][j]); + if constexpr (IS_DACT) { + act_in.load_from(&sActIn[buff][i][j]); + } +#pragma unroll + for (int e = 0; e < PACK_SIZE; ++e) { + const int k = w * PACK_SIZE + e; + float elt = static_cast(in.data.elt[e]); + if constexpr (IS_ACT) { + elt = OP(elt, {}); + } + if constexpr (IS_DACT) { + float act_in_elt = static_cast(act_in.data.elt[e]); + elt *= OP(act_in_elt, {}); + } + + if constexpr (IS_DBIAS && (!COLWISE_SCALING)) { + thread_dbias_rowwise[k] += elt; + } + if constexpr (!std::is_same_v) { + elt = static_cast(static_cast(elt)); + } + thread_amax = fmaxf(thread_amax, fabsf(elt)); + rInCompute[k] = elt; + } } } -} -__device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) - asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map)); -#else - NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+."); -#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + const e8m0_t biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + const size_t stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; + const size_t stage_scales_offset_X = scales_offset_X_rowwise; + + size_t scale_idx = 0; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + scale_idx = transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx( + stage_scales_offset_Y, stage_scales_offset_X, + DIVUP(cols, static_cast(scale_tensor_alignment_X_colwise))); + } 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; + } + + const bf16 block_scale_inverse_bf16 = ptx::exp2f_rcp(biased_exponent); + const ptx::bf16x2 block_scale_inverse_bf16_x2 = {block_scale_inverse_bf16, + block_scale_inverse_bf16}; + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; + +#pragma unroll + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t j = swizzled_group_idx + thread_offset_X_rowwise; + + if constexpr (BF16_CAST_ONLY) { + uint32_t out_4x = 0; + OType4 &out = *reinterpret_cast(&out_4x); + ptx::mul_cvt_4x(out, rIn4x[w], block_scale_inverse_bf16_x2); + + const uint32_t dst_smem_ptr = __cvta_generic_to_shared(&sOutRowwise[buff][i][j]); + asm volatile("st.shared.b32 [%0], %1;" : : "r"(dst_smem_ptr), "r"(out_4x)); + } else { + Vec out; +#pragma unroll + for (int e = 0; e < PACK_SIZE / 2; ++e) { + IType2 in; + OType2 &out_pair = reinterpret_cast(out.data.elt[e]); + if constexpr (FP16_CAST_ONLY) { + in = rIn[w].data.elt[e]; + } else if constexpr (IS_CACHED_ACT_OP) { + in.x = rInCached[w].data.elt[2 * e]; + in.y = rInCached[w].data.elt[2 * e + 1]; + } else { + const int j = w * PACK_SIZE + 2 * e; + in.x = rInCompute[j]; + in.y = rInCompute[j + 1]; + } + ptx::mul_cvt_2x(out_pair, in, block_scale_inverse_2x); + } + out.store_to(&sOutRowwise[buff][i][j]); + } + } } template + float (*OP)(float, const ParamOP &), typename IType, typename OType, + ScalingType SCALING_TYPE, bool WITH_GEMM_SWIZZLED_SCALES, ShapeRepresentation SHAPE_REP> __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel( const __grid_constant__ CUtensorMap tensor_map_input_static, const __grid_constant__ CUtensorMap tensor_map_act_input_static, const __grid_constant__ CUtensorMap tensor_map_output_rowwise_static, - const __grid_constant__ CUtensorMap tensor_map_output_colwise_static, - const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t first_logical_dim, - const size_t last_logical_dim, const int64_t *const __restrict__ offsets_ptr, - const int64_t *const __restrict__ first_dims_ptr, + const __grid_constant__ CUtensorMap tensor_map_output_colwise_static, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, const int64_t *const __restrict__ first_dims_ptr, const int64_t *const __restrict__ last_dims_ptr, e8m0_t *const __restrict__ scales_rowwise_ptr, e8m0_t *const __restrict__ scales_colwise_ptr, const float *__restrict__ noop, - float *const __restrict__ dbias_workspace, float *const __restrict__ amax_ptr) { + float *const __restrict__ dbias_workspace, float *const __restrict__ amax_ptr, + const size_t work_blocks_X, const size_t work_blocks_Y) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) constexpr bool COMPUTE_ACTIVATIONS = IS_DACT || IS_ACT; constexpr bool NO_ACTIVATIONS = !COMPUTE_ACTIVATIONS; - using IType2 = typename ptx::FPx2; - using OType2 = typename ptx::FPx2; - - using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx; - if constexpr (NO_ACTIVATIONS) { if (noop != nullptr && noop[0] == 1.0f) { return; } } - constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; - - const bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); - - const size_t block_ID = blockIdx.y * gridDim.x + blockIdx.x; - const size_t block_global_offset = - is_single_tensor ? (blockIdx.y * CHUNK_DIM_Y * last_logical_dim + blockIdx.x * CHUNK_DIM_X) - : (block_ID * ELTS_PER_CHUNK); - - const size_t tensor_id = - get_current_tensor_id(shape_rep, num_tensors, block_global_offset, blockIdx.y, - first_logical_dim, last_logical_dim, offsets_ptr); - - const size_t rows = - get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); - const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); - - const size_t scale_stride_rowwise = DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(32)), 4); - const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, 128); - - // grouped tensor can be treated as continuous tensor for MXFP8 - const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[tensor_id]); - // For grouped tensors represented as a single logical tensor, scale swizzle must still be - // computed per tensor (expert) and then concatenated along dim-0. - const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1) - ? static_cast(offsets_ptr[tensor_id]) - : tensor_base; - - // In graph-safe paged stashing, the logical shape can include trailing garbage. Skip CTAs that - // map outside the current tensor's valid [rows, cols] region. - if (rows == 0 || cols == 0) { - return; - } - if (shape_rep != SAME_BOTH_DIMS) { - const size_t tensor_start_offset = static_cast(offsets_ptr[tensor_id]); - const size_t tensor_end_offset = static_cast(offsets_ptr[tensor_id + 1]); - if (block_global_offset >= tensor_end_offset) { - return; - } - const size_t tensor_offset_from_start = block_global_offset - tensor_start_offset; - const size_t block_offset_Y_in_tensor = tensor_offset_from_start / cols; - const size_t block_offset_X_in_tensor = tensor_offset_from_start % cols; - if (block_offset_Y_in_tensor >= rows || block_offset_X_in_tensor >= cols) { - return; - } - } + constexpr bool ROWWISE_SCALING = + (SCALING_TYPE == ScalingType::ROWWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); + constexpr bool COLWISE_SCALING = + (SCALING_TYPE == ScalingType::COLWISE) || (SCALING_TYPE == ScalingType::BIDIMENSIONAL); - const CUtensorMap &tensor_map_input = - is_single_tensor ? tensor_map_input_static : g_tensor_maps_input[tensor_id]; - const CUtensorMap &tensor_map_act_input = - is_single_tensor ? tensor_map_act_input_static : g_tensor_maps_act_input[tensor_id]; - const CUtensorMap &tensor_map_output_rowwise = - is_single_tensor ? tensor_map_output_rowwise_static : g_tensor_maps_output_rowwise[tensor_id]; - const CUtensorMap &tensor_map_output_colwise = - is_single_tensor ? tensor_map_output_colwise_static : g_tensor_maps_output_colwise[tensor_id]; + constexpr ShapeRepresentation shape_rep = SHAPE_REP; + constexpr bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM); const bool leading_thread = (threadIdx.x == 0); - if (leading_thread && (!is_single_tensor)) { - fence_acquire_tensormap(&tensor_map_input); - if constexpr (COMPUTE_ACTIVATIONS) { - fence_acquire_tensormap(&tensor_map_act_input); - } - if constexpr (ROWWISE_SCALING) { - fence_acquire_tensormap(&tensor_map_output_rowwise); - } - if constexpr (COLWISE_SCALING) { - fence_acquire_tensormap(&tensor_map_output_colwise); - } - } - - const size_t blocks_X_num_in_current_tensor = DIVUP(cols, static_cast(128)); - const size_t block_id_in_current_tensor = - is_single_tensor ? block_ID : (block_ID - tensor_base / ELTS_PER_CHUNK); - - const size_t block_id_Y = block_id_in_current_tensor / blocks_X_num_in_current_tensor; - const size_t block_id_X = block_id_in_current_tensor % blocks_X_num_in_current_tensor; - - const size_t block_offset_Y = block_id_Y * CHUNK_DIM_Y; - const size_t block_offset_X = block_id_X * CHUNK_DIM_X; - - e8m0_t *const scales_rowwise = - scales_rowwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_X); - e8m0_t *const scales_colwise = - scales_colwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_Y); - - const size_t scales_block_offset_Y_rowwise = block_id_Y * CHUNK_DIM_Y; - const size_t scales_block_offset_X_rowwise = block_id_X * CHUNK_DIM_X / SCALE_DIM_X; - const size_t scales_block_offset_Y_colwise = block_id_Y * CHUNK_DIM_Y / SCALE_DIM_Y; - const size_t scales_block_offset_X_colwise = block_id_X * CHUNK_DIM_X; - const size_t tid_Y_rowwise = threadIdx.x / THREADS_X; const size_t tid_X_rowwise = threadIdx.x % THREADS_X; const size_t tid_Y_colwise = 0; @@ -363,11 +480,6 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t thread_offset_Y_rowwise = tid_Y_rowwise; const size_t thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM_X; - const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; - const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; - const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; - const size_t scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; - // helps resolving bank conflicts in shmem const int thread_lane = threadIdx.x % THREADS_PER_WARP; const int bank_group = thread_lane / THREADS_PER_BANK; @@ -387,399 +499,251 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned extern __shared__ unsigned char dynamic_shmem[]; - unsigned char *dshmem = common::align_smem_ptr_per_TMA_requirements(dynamic_shmem); + unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem); // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - IType *in_sh = reinterpret_cast(dshmem); - IType *act_in_sh = reinterpret_cast(dshmem + elt_input_mem); - - OType *out_rowwise_data_sh = reinterpret_cast(dshmem + in_mem); - OType *out_colwise_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise); - IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer - - constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; - - float partial_dbias_colwise = 0.0f; - float thread_dbias_rowwise[SCALE_DIM_X]; - if constexpr (IS_DBIAS) { -#pragma unroll - for (int j = 0; j < SCALE_DIM_X; ++j) { - thread_dbias_rowwise[j] = 0.0f; - } - } + IType *sIn_ptr = reinterpret_cast(dshmem); + IType *sActIn_ptr = reinterpret_cast(dshmem + elt_input_mem); - float block_amax = 0.0f; + OType *sOutRowwise_ptr = reinterpret_cast(dshmem + in_mem); + OType *sOutColwise_ptr = reinterpret_cast(dshmem + in_mem + out_mem_rowwise); + IType *sCachedAct_ptr = sIn_ptr; // sIn_ptr is used as a cache buffer -// 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]; + constexpr size_t shmem_buff_size = (IS_DACT ? 2 : 1) * buff_size_aligned_in / BUFFS_NUM; - initialize_barriers(mbar, leading_thread); + const size_t total_work_blocks = work_blocks_X * work_blocks_Y; + const size_t launch_block_id = blockIdx.y * gridDim.x + blockIdx.x; - int parity = 0; + int IN_buff_readable_parity[BUFFS_NUM] = {0}; - if constexpr (IS_DACT) { - copy_2d_to_sharedx2(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, &act_in_sh[0], - &tensor_map_act_input, block_offset_X, block_offset_Y, shmem_buff_size, - &mbar[0], leading_thread); - } else { - copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, shmem_buff_size, - &mbar[0], leading_thread); + // In persistent mode, physical CTAs iterate over a virtual work grid via grid-stride. + if (launch_block_id >= total_work_blocks) { + return; } - -#pragma unroll - for (int stage = 0; stage < STAGES; ++stage) { - const size_t buff = stage % BUFFS_NUM; - const size_t next_stage = stage + 1; - const size_t stage_offset_Y = stage * BUFF_DIM_Y; - - 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 size_t next_buff = next_stage % BUFFS_NUM; - const size_t next_stage_offset_Y = next_stage * BUFF_DIM_Y; - const size_t global_offset_Y = block_offset_Y + next_stage_offset_Y; - const size_t global_offset_X = block_offset_X; - const size_t next_buff_offset = next_buff * BUFF_DIM; - if constexpr (IS_DACT) { - copy_2d_to_sharedx2(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, - global_offset_Y, &act_in_sh[next_buff_offset], &tensor_map_act_input, - global_offset_X, global_offset_Y, shmem_buff_size, &mbar[next_stage], - leading_thread); - } else { - copy_2d_to_shared(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, - global_offset_Y, shmem_buff_size, &mbar[next_stage], leading_thread); - } + int32_t ctaid_X = static_cast(launch_block_id % work_blocks_X); + int32_t ctaid_Y = static_cast(launch_block_id / work_blocks_X); + size_t static_block_stride = gridDim.x * gridDim.y; + size_t static_next_block_id = launch_block_id + static_block_stride; + + bool job_finished = false; + size_t last_acquired_tensor_id = num_tensors; + + __shared__ uint64_t IN_buff_readable_mbar[BUFFS_NUM]; + // Initialize barriers shared by the entire CTA: + // - IN_buff_readable_mbar tracks per-buffer TMA global->shared completion. + initialize_barriers(IN_buff_readable_mbar, leading_thread); + + // Main work loop: decode current job, prime its pipeline, then process all 32-row stages. + while (!job_finished) { + // Decode CTA assignment into logical tensor coordinates and validate bounds. + const JobDescriptor current_job = decode_job( + num_tensors, first_logical_dim, last_logical_dim, work_blocks_X, ctaid_X, ctaid_Y, + offsets_ptr, first_dims_ptr, last_dims_ptr); + const bool current_job_is_valid = + is_job_valid(current_job, total_work_blocks, offsets_ptr); + if (!current_job_is_valid) { + break; + } + if (!job_has_work(current_job)) { + // Zero-sized tensors are valid grouped-tensor entries; skip them and keep scheduling work. + advance_to_next_job(job_finished, ctaid_X, ctaid_Y, static_next_block_id, static_block_stride, + total_work_blocks, work_blocks_X); + continue; } - ptx::fence_proxy_async_shared_cta(); - - // Wait for the data to have arrived - ptx::mbarrier_wait_parity(&mbar[stage], parity); - - float thread_amax = 0.0f; - if constexpr (COLWISE_SCALING) { - const size_t shmem_offset_base_colwise = buff * BUFF_DIM + tid_X_colwise; - thread_amax = 0.0f; - float in_compute_colwise[BUFF_DIM_Y]; - IType in_colwise_IType[BUFF_DIM_Y]; - - // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - IType thread_amax_f16 = static_cast(0.0f); -#pragma unroll - for (int i = 0; i < BUFF_DIM_Y; ++i) { - const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; - in_colwise_IType[i] = in_sh[shmem_offset_colwise]; - thread_amax_f16 = __hmax(thread_amax_f16, __habs(in_colwise_IType[i])); - } - thread_amax = static_cast(thread_amax_f16); - } else { -#pragma unroll - for (int i = 0; i < BUFF_DIM_Y; ++i) { - const size_t shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_DIM_X; - - float elt = static_cast(in_sh[shmem_offset_colwise]); - if constexpr (IS_ACT) { - elt = OP(elt, {}); - } - if constexpr (IS_DACT) { - float act_in_elt = static_cast(act_in_sh[shmem_offset_colwise]); - elt *= OP(act_in_elt, {}); - } - if constexpr (IS_DBIAS) { - partial_dbias_colwise += 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); - } - thread_amax = fmaxf(thread_amax, fabsf(elt)); - in_compute_colwise[i] = elt; - } + const size_t tensor_id = current_job.tensor_id; + const size_t rows = current_job.rows; + const size_t cols = current_job.cols; + const BlockDescriptor current_block = + decode_block(current_job, offsets_ptr); + const size_t scale_alignment_X_rowwise = static_cast(scale_tensor_alignment_X_rowwise); + const size_t scale_alignment_X_colwise = static_cast(scale_tensor_alignment_X_colwise); + + const size_t scale_stride_rowwise = + DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(SCALE_DIM_X)), scale_alignment_X_rowwise); + const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise); + + const size_t tensor_base = current_block.tensor_base; + const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1) + ? static_cast(offsets_ptr[tensor_id]) + : tensor_base; + const size_t block_id_Y = current_block.block_id_Y; + const size_t block_id_X = current_block.block_id_X; + const size_t block_offset_Y = current_block.block_offset_Y; + const size_t block_offset_X = current_block.block_offset_X; + + e8m0_t *const scales_rowwise = + scales_rowwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_X); + e8m0_t *const scales_colwise = + scales_colwise_ptr + (is_single_tensor ? 0 : tensor_base / SCALE_DIM_Y); + + const size_t scales_block_offset_Y_rowwise = block_id_Y * CHUNK_DIM_Y; + const size_t scales_block_offset_X_rowwise = block_id_X * CHUNK_DIM_X / SCALE_DIM_X; + const size_t scales_block_offset_Y_colwise = block_id_Y * CHUNK_DIM_Y / SCALE_DIM_Y; + const size_t scales_block_offset_X_colwise = block_id_X * CHUNK_DIM_X; + + const size_t scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; + const size_t scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; + const size_t scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; + const size_t scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; + + const bool rowwise_scale_is_within_bounds = scales_offset_X_rowwise * SCALE_DIM_X < cols; + + const size_t dbias_offset_Y = block_id_Y; + const size_t dbias_offset_X = block_id_X * CHUNK_DIM_X + threadIdx.x; + + const CUtensorMap &tensor_map_input = + is_single_tensor ? tensor_map_input_static : g_tensor_maps.input[tensor_id]; + const CUtensorMap &tensor_map_act_input = + is_single_tensor ? tensor_map_act_input_static : g_tensor_maps.act_input[tensor_id]; + const CUtensorMap &tensor_map_output_rowwise = is_single_tensor + ? tensor_map_output_rowwise_static + : g_tensor_maps.output_rowwise[tensor_id]; + const CUtensorMap &tensor_map_output_colwise = is_single_tensor + ? tensor_map_output_colwise_static + : g_tensor_maps.output_colwise[tensor_id]; + + if (leading_thread && (!is_single_tensor) && (last_acquired_tensor_id != tensor_id)) { + fence_acquire_tensormap(&tensor_map_input); + if constexpr (COMPUTE_ACTIVATIONS) { + fence_acquire_tensormap(&tensor_map_act_input); } - - // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - - const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; - const size_t global_scales_offset_X = scales_offset_X_colwise; - - 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 local_scales_offset_Y = global_scales_offset_Y - tensor_scales_offset_Y_base; - scale_idx = tensor_scales_offset_colwise_base + - gemm_swizzled_scale_idx(global_scales_offset_X, local_scales_offset_Y, - DIVUP(rows, static_cast(128))); - } else { - scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; + if constexpr (ROWWISE_SCALING) { + fence_acquire_tensormap(&tensor_map_output_rowwise); } - scales_colwise[scale_idx] = biased_exponent; - - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); - const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; - -// 3. Scale elements -#pragma unroll - for (int i = 0; i < SCALE_DIM_Y; ++i) { - float in; - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - in = static_cast(in_colwise_IType[i]); - } else { - in = in_compute_colwise[i]; - } - const float scaled_out = in * block_scale_inverse; - - const size_t shmem_offset_elt = shmem_offset_base_colwise + i * BUFF_DIM_X; - out_colwise_data_sh[shmem_offset_elt] = static_cast(scaled_out); + if constexpr (COLWISE_SCALING) { + fence_acquire_tensormap(&tensor_map_output_colwise); } + last_acquired_tensor_id = tensor_id; } + __syncthreads(); - if constexpr (ROWWISE_SCALING) { - const size_t shmem_offset_base_rowwise = - buff * BUFF_DIM + thread_offset_Y_rowwise * BUFF_DIM_X; - thread_amax = 0.0f; - float in_compute_rowwise[SCALE_DIM_X]; - Vec in_cached[WAVES]; - - // used as an IType container for BF16/FP16 --> MXFP8 CAST ONLY - Vec in_IType[WAVES]; + int buff_in = 0; - // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; - // Load elements - in_IType[w].load_from(&in_sh[shmem_offset_rowwise]); +// Prime the pipeline with the first PREFETCH_STAGES slices of the current block. #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]); - } - } - thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; - - // 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 constexpr (std::is_same_v) { -#pragma unroll - for (int e = 0; e < PACK_SIZE; ++e) { - thread_amax = fmaxf(thread_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) { - thread_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 size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_thread_idx; - - Vec in; - Vec act_in; + for (int stage = 0; stage < PREFETCH_STAGES; ++stage) { + const size_t buff = stage; + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + const size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t buff_offset = buff * BUFF_DIM; + uint64_t *barrier = &IN_buff_readable_mbar[buff]; + prefetch_input_stage(sIn_ptr, sActIn_ptr, tensor_map_input, + tensor_map_act_input, global_offset_X, global_offset_Y, + buff_offset, shmem_buff_size, barrier, leading_thread); + } - in.load_from(&in_sh[shmem_offset_rowwise]); - if constexpr (IS_DACT) { - act_in.load_from(&act_in_sh[shmem_offset_rowwise]); - } + float partial_dbias_colwise = 0.0f; + float thread_dbias_rowwise[SCALE_DIM_X]; + if constexpr (IS_DBIAS) { #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 (IS_ACT) { - elt = OP(elt, {}); - } - if constexpr (IS_DACT) { - float act_in_elt = static_cast(act_in.data.elt[e]); - elt *= OP(act_in_elt, {}); - } - - // If DBIAS was computed in the 1st pass (COLWISE) then no need to compute it again - if constexpr (IS_DBIAS && (!COLWISE_SCALING)) { - thread_dbias_rowwise[j] += 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)); - } - thread_amax = fmaxf(thread_amax, fabsf(elt)); - in_compute_rowwise[j] = elt; - } - } - } - - // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; - const int stage_scales_offset_X = scales_offset_X_rowwise; - - size_t scale_idx = 0; - if constexpr (WITH_GEMM_SWIZZLED_SCALES) { - scale_idx = gemm_swizzled_scale_idx(stage_scales_offset_Y, stage_scales_offset_X, - DIVUP(cols, static_cast(128))); - } else { - scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; + for (int j = 0; j < SCALE_DIM_X; ++j) { + thread_dbias_rowwise[j] = 0.0f; } - scales_rowwise[scale_idx] = biased_exponent; - - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); - const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; + } -// 3. Scale elements +// Process one [CHUNK_DIM_Y x CHUNK_DIM_X] block in STAGES slices (32 rows each). #pragma unroll - for (int w = 0; w < WAVES; ++w) { - Vec out; -#pragma unroll - for (int e = 0; e < PACK_SIZE / 2; ++e) { - IType2 in; - OType2 &out_pair = reinterpret_cast(out.data.elt[e]); - if constexpr (NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v)) { - in = in_IType[w].data.elt[e]; - } else if constexpr (IS_CACHED_ACT_OP) { - in.x = in_cached[w].data.elt[2 * e]; - in.y = in_cached[w].data.elt[2 * e + 1]; - } else { - const int j = w * PACK_SIZE + 2 * e; - in.x = in_compute_rowwise[j]; - in.y = in_compute_rowwise[j + 1]; - } - ptx::mul_cvt_2x(out_pair, in, block_scale_inverse_2x); - } - const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const size_t swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; - const size_t shmem_offset_rowwise = shmem_offset_base_rowwise + swizzled_idx; - out.store_to(&out_rowwise_data_sh[shmem_offset_rowwise]); + for (int stage = 0; stage < STAGES; ++stage) { + const size_t stage_offset_Y = stage * BUFF_DIM_Y; + if (stage < STAGES - PREFETCH_STAGES) { + const size_t next_prefetch_buff = (buff_in + PREFETCH_STAGES) % BUFFS_NUM; + const size_t next_prefetch_stage = stage + PREFETCH_STAGES; + const size_t next_prefetch_stage_offset_Y = next_prefetch_stage * BUFF_DIM_Y; + + const size_t global_offset_Y = block_offset_Y + next_prefetch_stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t next_prefetch_buff_offset = next_prefetch_buff * BUFF_DIM; + + uint64_t *barrier = &IN_buff_readable_mbar[next_prefetch_buff]; + prefetch_input_stage( + sIn_ptr, sActIn_ptr, tensor_map_input, tensor_map_act_input, global_offset_X, + global_offset_Y, next_prefetch_buff_offset, shmem_buff_size, barrier, leading_thread); } - } - __builtin_assume(block_amax >= 0); - __builtin_assume(thread_amax >= 0); - block_amax = fmaxf(block_amax, thread_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. + ptx::mbarrier_wait_parity_acquire_cta_shared_cta(&IN_buff_readable_mbar[buff_in], + IN_buff_readable_parity[buff_in]); + IN_buff_readable_parity[buff_in] ^= 1; + ptx::cp_async_bulk_wait_group_read(); - // Initiate TMA transfer to copy shared memory to global memory - if (leading_thread) { - const int global_offset_Y = block_offset_Y + stage_offset_Y; - const int global_offset_X = block_offset_X; - const int buff_offset = buff * BUFF_DIM; + const size_t buff = buff_in; + if constexpr (COLWISE_SCALING) { + process_colwise_stage( + buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, + scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr, + sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); + } 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])); - } - 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])); + process_rowwise_stage( + buff, stage_offset_Y, thread_offset_Y_rowwise, thread_offset_X_rowwise, bank_group, + scales_offset_Y_rowwise, scales_offset_X_rowwise, scale_stride_rowwise, + rowwise_scale_is_within_bounds, cols, sIn_ptr, sActIn_ptr, sCachedAct_ptr, + sOutRowwise_ptr, scales_rowwise, thread_dbias_rowwise); } - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - } - } + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); - parity ^= 1; + // Publish the stage from shared memory into global outputs via TMA. + const size_t global_offset_Y = block_offset_Y + stage_offset_Y; + const size_t global_offset_X = block_offset_X; + const size_t buff_offset = buff * BUFF_DIM; + store_output_stage( + sOutRowwise_ptr, sOutColwise_ptr, tensor_map_output_rowwise, tensor_map_output_colwise, + global_offset_X, global_offset_Y, buff_offset, leading_thread); - if constexpr (IS_DBIAS) { - if (is_single_tensor) { - float thread_partial_dbias = 0.0f; - if constexpr (COLWISE_SCALING) { - thread_partial_dbias = partial_dbias_colwise; - } else { - // Reusing dshmem (in_sh) as dbias buffer [HEIGHT x WIDTH] - // HEIGHT = THREADS_Y - // WIDTH = THREADS_X * (SCALE_DIM_X + 1) - // Added extra 1-element padding per thread_X to reduce bank conflicts - float *partial_dbias_rowwise = reinterpret_cast(dshmem); + buff_in = (buff_in + 1) % BUFFS_NUM; + } - constexpr int DBIAS_BUFF_WIDTH = THREADS_X * (SCALE_DIM_X + 1); + if constexpr (IS_DBIAS) { + if (is_single_tensor) { + float thread_partial_dbias = 0.0f; + if constexpr (COLWISE_SCALING) { + thread_partial_dbias = partial_dbias_colwise; + } else { + float *partial_dbias_rowwise = reinterpret_cast(dshmem); + + constexpr size_t DBIAS_BUFF_WIDTH = THREADS_X * (SCALE_DIM_X + 1); - const int shmem_thread_offset = - tid_Y_rowwise * DBIAS_BUFF_WIDTH + tid_X_rowwise * (SCALE_DIM_X + 1); + const size_t shmem_thread_offset = + tid_Y_rowwise * DBIAS_BUFF_WIDTH + tid_X_rowwise * (SCALE_DIM_X + 1); #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_group_offset = shmem_thread_offset + swizzled_group_idx; + for (int w = 0; w < WAVES; ++w) { + const size_t swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; + const size_t swizzled_group_offset = shmem_thread_offset + swizzled_group_idx; #pragma unroll - for (int e = 0; e < PACK_SIZE; ++e) { - const int j = w * PACK_SIZE + e; - const int shmem_elt_idx = swizzled_group_offset + e; - partial_dbias_rowwise[shmem_elt_idx] = thread_dbias_rowwise[j]; + for (int e = 0; e < PACK_SIZE; ++e) { + const size_t j = w * PACK_SIZE + e; + const size_t shmem_elt_idx = swizzled_group_offset + e; + partial_dbias_rowwise[shmem_elt_idx] = thread_dbias_rowwise[j]; + } } - } - __syncthreads(); + __syncthreads(); #pragma unroll - for (int i = 0; i < THREADS_Y; ++i) { - // Add extra element offset per MXFP8 scaling block [1x32] - const int scaling_block = threadIdx.x / SCALE_DIM_X; - thread_partial_dbias += - partial_dbias_rowwise[i * DBIAS_BUFF_WIDTH + threadIdx.x + scaling_block]; + for (int i = 0; i < THREADS_Y; ++i) { + const int scaling_block = threadIdx.x / SCALE_DIM_X; + thread_partial_dbias += + partial_dbias_rowwise[i * DBIAS_BUFF_WIDTH + threadIdx.x + scaling_block]; + } + } + const size_t dbias_stride = cols; + const size_t dbias_idx = dbias_offset_Y * dbias_stride + dbias_offset_X; + const bool col_out_of_bounds_dbias = (dbias_offset_X >= cols); + if (!col_out_of_bounds_dbias) { + dbias_workspace[dbias_idx] = thread_partial_dbias; } - } - const int dbias_stride = cols; - const int dbias_offset_Y = block_id_Y; - const int dbias_offset_X = block_id_X * CHUNK_DIM_X + threadIdx.x; - const int dbias_idx = dbias_offset_Y * dbias_stride + dbias_offset_X; - const bool col_out_of_bounds_dbias = (dbias_offset_X >= cols); - if (!col_out_of_bounds_dbias) { - dbias_workspace[dbias_idx] = thread_partial_dbias; } } - } - - if (amax_ptr != nullptr) { - const int warp_id = threadIdx.x / THREADS_PER_WARP; - // Reduce the amax over the block - block_amax = reduce_max(block_amax, warp_id); - } - if (leading_thread && amax_ptr != nullptr) { - atomicMaxFloat(amax_ptr, block_amax); + advance_to_next_job(job_finished, ctaid_X, ctaid_Y, static_next_block_id, static_block_stride, + total_work_blocks, work_blocks_X); } - destroy_barriers(mbar, leading_thread); + destroy_barriers(IN_buff_readable_mbar, leading_thread); #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } } // namespace group_quantize_kernel @@ -788,7 +752,8 @@ template void group_quantize(const GroupedTensor *input, const GroupedTensor *activations, const Tensor *noop, GroupedTensor *output, GroupedTensor *dbias, - Tensor *workspace, cudaStream_t stream) { + Tensor *workspace, const QuantizationConfig *quant_config, + cudaStream_t stream) { using namespace group_quantize_kernel; checkCuDriverContext(stream); @@ -839,20 +804,25 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations const size_t num_tensors = input->num_tensors; - size_t blocks_X = 0; - size_t blocks_Y = 0; + size_t work_blocks_X = 0; + size_t work_blocks_Y = 0; if (is_single_tensor) { - blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); - blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); + work_blocks_Y = DIVUP(first_logical_dim, static_cast(CHUNK_DIM_Y)); + work_blocks_X = DIVUP(last_logical_dim, static_cast(CHUNK_DIM_X)); } else { NVTE_CHECK(num_tensors <= MAX_SUPPORTED_TENSOR_DESCRIPTORS, "Number of tensors in a group is larger than " "the MAX number of supported descriptors (64)."); - blocks_Y = 1; - blocks_X = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); + work_blocks_Y = 1; + work_blocks_X = DIVUP(elts_total, ELTS_PER_CHUNK); } - const dim3 grid(blocks_X, blocks_Y); + + const size_t sm_num = static_cast(transformer_engine::cuda::sm_count()); + const size_t static_grid_size = sm_num * TunableConfig::STATIC_PERSISTENT_BLOCKS_PER_SM; + NVTE_CHECK(static_grid_size > 0, "Static persistent grid size must be greater than zero."); + + const dim3 grid(static_grid_size); const size_t block_size = THREADS_PER_CHUNK; const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; @@ -891,7 +861,7 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations NVTE_CHECK(dbias->data.shape == expected_shape_dbias_tensor, "Wrong shape of DBias."); NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor."); - const size_t dbias_workspace_rows = DIVUP(first_logical_dim, CHUNK_DIM_Y); + const size_t dbias_workspace_rows = DIVUP(first_logical_dim, static_cast(CHUNK_DIM_Y)); const size_t dbias_workspace_cols = last_logical_dim; if (workspace->data.dptr == nullptr) { workspace->data.shape = {dbias_workspace_rows, dbias_workspace_cols}; @@ -904,125 +874,125 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations input->dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( output->dtype(), OType, - TRANSFORMER_ENGINE_SWITCH_CONDITION( - with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_act_input{}; - alignas(64) CUtensorMap tensor_map_output_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_colwise{}; - - constexpr size_t input_type_bit_size = TypeInfo::size; - constexpr size_t output_type_bit_size = TypeInfo::size; - - create_2D_tensor_map(tensor_map_input, input->data, first_logical_dim, - last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, - input_type_bit_size); - - if constexpr (IS_DACT) { - create_2D_tensor_map(tensor_map_act_input, activations->data, first_logical_dim, - last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, - input_type_bit_size); - } - - if (use_rowwise_scaling) { - create_2D_tensor_map(tensor_map_output_rowwise, output->data, first_logical_dim, - last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, last_logical_dim, 0, - output_type_bit_size); - } - - if (use_colwise_scaling) { - create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, - first_logical_dim, last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, - last_logical_dim, 0, output_type_bit_size); - } - - constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; - constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; - constexpr size_t input_buff_size = (buff_elems_total * input_type_bit_size) / 8; - constexpr size_t output_buff_size = (buff_elems_total * output_type_bit_size) / 8; - constexpr size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out = - DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); - - constexpr size_t elt_input_mem = buff_size_aligned_in; - constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); - constexpr size_t in_mem = elt_input_mem + act_input_mem; - - const size_t out_rowwise_mem = (use_rowwise_scaling ? buff_size_aligned_out : 0); - const size_t out_colwise_mem = (use_colwise_scaling ? buff_size_aligned_out : 0); - const size_t out_mem = out_rowwise_mem + out_colwise_mem; - - const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; - - auto kernel = - group_quantize_mxfp8_kernel; - switch (scaling_type) { - case ScalingType::ROWWISE: { - kernel = - group_quantize_mxfp8_kernel; - break; - } - case ScalingType::COLWISE: { - kernel = - group_quantize_mxfp8_kernel; - break; - } - case ScalingType::BIDIMENSIONAL: { - kernel = - group_quantize_mxfp8_kernel; - break; - } - } - - // Update tensor descriptors before launching the kernel - if (!is_single_tensor) { - const IType *const input_dptr = reinterpret_cast(input->data.dptr); - - const IType *const act_input_dptr = - IS_DACT ? reinterpret_cast(activations->data.dptr) : nullptr; - - OType *const output_rowwise_dptr = - use_rowwise_scaling ? reinterpret_cast(output->data.dptr) : nullptr; - - OType *const output_colwise_dptr = - use_colwise_scaling ? reinterpret_cast(output->columnwise_data.dptr) - : nullptr; - update_tma_descriptors<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, input_dptr, act_input_dptr, output_rowwise_dptr, - output_colwise_dptr, shape_rep, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, - use_rowwise_scaling, use_colwise_scaling, IS_DACT); - } - - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, shape_rep, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_rowwise_ptr, - scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr); - - if constexpr (IS_DBIAS) { - common::grouped_reduce_dbias( - shape_rep, num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, - first_dims_ptr, last_dims_ptr, dbias, workspace_ptr, CHUNK_DIM_Y, stream); - } - - NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) - ); // NOLINT(*) - ); // NOLINT(*) + TRANSFORMER_ENGINE_SCALING_TYPE_SWITCH( + scaling_type, SCALING_TYPE, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH( + shape_rep, SHAPE_REP, + { + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_act_input{}; + alignas(64) CUtensorMap tensor_map_output_rowwise{}; + alignas(64) CUtensorMap tensor_map_output_colwise{}; + + constexpr size_t input_type_bit_size = TypeInfo::size; + constexpr size_t output_type_bit_size = TypeInfo::size; + + create_2D_tensor_map(tensor_map_input, input->data, first_logical_dim, + last_logical_dim, BUFF_DIM_Y, BUFF_DIM_X, + last_logical_dim, 0, input_type_bit_size); + + if constexpr (IS_DACT) { + create_2D_tensor_map(tensor_map_act_input, activations->data, + first_logical_dim, last_logical_dim, BUFF_DIM_Y, + BUFF_DIM_X, last_logical_dim, 0, + input_type_bit_size); + } + + if (use_rowwise_scaling) { + create_2D_tensor_map(tensor_map_output_rowwise, output->data, + first_logical_dim, last_logical_dim, BUFF_DIM_Y, + BUFF_DIM_X, last_logical_dim, 0, + output_type_bit_size); + } + + if (use_colwise_scaling) { + create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, + first_logical_dim, last_logical_dim, BUFF_DIM_Y, + BUFF_DIM_X, last_logical_dim, 0, + output_type_bit_size); + } + + constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; + constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; + constexpr size_t input_buff_size = + (buff_elems_total * input_type_bit_size) / 8; + constexpr size_t output_buff_size = + (buff_elems_total * output_type_bit_size) / 8; + constexpr size_t buff_size_aligned_in = + DIVUP_TO_MULTIPLE(input_buff_size, TMA_SHMEM_ALIGNMENT); + constexpr size_t buff_size_aligned_out = + DIVUP_TO_MULTIPLE(output_buff_size, TMA_SHMEM_ALIGNMENT); + + constexpr size_t elt_input_mem = buff_size_aligned_in; + constexpr size_t act_input_mem = (IS_DACT ? buff_size_aligned_in : 0); + constexpr size_t in_mem = elt_input_mem + act_input_mem; + + const size_t out_rowwise_mem = + (use_rowwise_scaling ? buff_size_aligned_out : 0); + const size_t out_colwise_mem = + (use_colwise_scaling ? buff_size_aligned_out : 0); + const size_t out_mem = out_rowwise_mem + out_colwise_mem; + + const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT; + + // Update tensor descriptors before launching the kernel + if (!is_single_tensor) { + const IType *const input_dptr = + reinterpret_cast(input->data.dptr); + + const IType *const act_input_dptr = + IS_DACT ? reinterpret_cast(activations->data.dptr) + : nullptr; + + OType *const output_rowwise_dptr = + use_rowwise_scaling ? reinterpret_cast(output->data.dptr) + : nullptr; + + OType *const output_colwise_dptr = + use_colwise_scaling + ? reinterpret_cast(output->columnwise_data.dptr) + : nullptr; + update_tma_descriptors<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, input_dptr, act_input_dptr, + output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors, + first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, + last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT); + } + + auto kernel = + group_quantize_mxfp8_kernel; + + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, + scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, + amax_ptr, work_blocks_X, work_blocks_Y); + + if constexpr (IS_DBIAS) { + common::grouped_reduce_dbias( + shape_rep, num_tensors, first_logical_dim, last_logical_dim, + offsets_ptr, first_dims_ptr, last_dims_ptr, dbias, workspace_ptr, + CHUNK_DIM_Y, stream); + } + + NVTE_CHECK_CUDA(cudaGetLastError()); + }); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) } } // namespace mxfp8 } // namespace dispatch } // namespace transformer_engine - #endif // TRANSFORMER_ENGINE_GROUP_QUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 70a68132ad..f36b071081 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -278,7 +278,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } scales_colwise[scale_idx] = biased_exponent; - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; // 3. Scale elements @@ -430,7 +430,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) scales_rowwise[scale_idx] = biased_exponent; } - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); const ptx::floatx2 block_scale_inverse_2x = {block_scale_inverse, block_scale_inverse}; // 3. Scale elements diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index dd1b4fa40e..41e62ac319 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -289,7 +289,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re coords.x / CastTraits::chunkElems] = biased_exponent; } - float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; @@ -342,7 +342,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re } // scaling input - float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; @@ -410,7 +410,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re coords.x / CastTraits::chunkElems] = biased_exponent; } - float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; @@ -463,7 +463,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re } // scaling input - float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); ptx::floatx2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; outputUnitType rOutput[CastTraits::numOutUnitsPerChunk]; @@ -949,7 +949,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( { IType row_amax = ptx::get_amax(row_amax2.x, row_amax2.y); e8m0_t row_biased_exponent = to_e8m0(row_amax); - row_scale_inverse = ptx::exp2f_rcp(row_biased_exponent); + row_scale_inverse = ptx::exp2f_rcp(row_biased_exponent); if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { int32_t rowwise_scale_offset = rowwise_scale_smem_base_offset + @@ -969,7 +969,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( __syncwarp(); float col_amax = sColwiseReduce[threadIdx.x]; e8m0_t col_biased_exponent = to_e8m0(col_amax); - float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); + float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); sColwiseReduce[threadIdx.x] = col_scale_inverse; size_t colwise_scale_offset = colwise_scale_base_offset + @@ -1396,7 +1396,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( { IType row_amax = ptx::get_amax(row_amax2.x, row_amax2.y); e8m0_t row_biased_exponent = to_e8m0(row_amax); - row_scale_inverse = ptx::exp2f_rcp(row_biased_exponent); + row_scale_inverse = ptx::exp2f_rcp(row_biased_exponent); if constexpr (CastTraits::_cache_rowwise_scale_in_smem) { int32_t rowwise_scale_offset = rowwise_scale_smem_base_offset + @@ -1416,7 +1416,7 @@ __global__ void quantize_mxfp8_kernel_cast_only( __syncwarp(); float col_amax = sColwiseReduce[threadIdx.x]; e8m0_t col_biased_exponent = to_e8m0(col_amax); - float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); + float col_scale_inverse = ptx::exp2f_rcp(col_biased_exponent); sColwiseReduce[threadIdx.x] = col_scale_inverse; size_t colwise_scale_offset = colwise_scale_base_offset + diff --git a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh index e7854ffde3..ec80924df5 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh @@ -270,7 +270,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if (colwise_scale_is_within_bounds) { scales_colwise_e8m0[scale_idx] = biased_exponent; } - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); + const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); // 3. Scale elements #pragma unroll diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index a98668d058..6e207370dd 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -904,6 +904,48 @@ struct TypeInfo { { __VA_ARGS__ } \ } +#define TRANSFORMER_ENGINE_SCALING_TYPE_SWITCH(SCALING_TYPE, SCALING_T, ...) \ + switch (SCALING_TYPE) { \ + case ScalingType::ROWWISE: { \ + constexpr ScalingType SCALING_T = ScalingType::ROWWISE; \ + { __VA_ARGS__ } \ + } break; \ + case ScalingType::COLWISE: { \ + constexpr ScalingType SCALING_T = ScalingType::COLWISE; \ + { __VA_ARGS__ } \ + } break; \ + case ScalingType::BIDIMENSIONAL: { \ + constexpr ScalingType SCALING_T = ScalingType::BIDIMENSIONAL; \ + { __VA_ARGS__ } \ + } break; \ + default: { \ + NVTE_ERROR("Unsupported scaling type."); \ + } \ + } + +#define TRANSFORMER_ENGINE_GROUP_TENSOR_SHAPE_REPRESENTATION_SWITCH(SHAPE_REP, SHAPE, ...) \ + switch (SHAPE_REP) { \ + case ShapeRepresentation::SAME_BOTH_DIMS: { \ + constexpr ShapeRepresentation SHAPE = ShapeRepresentation::SAME_BOTH_DIMS; \ + { __VA_ARGS__ } \ + } break; \ + case ShapeRepresentation::VARYING_FIRST_DIM: { \ + constexpr ShapeRepresentation SHAPE = ShapeRepresentation::VARYING_FIRST_DIM; \ + { __VA_ARGS__ } \ + } break; \ + case ShapeRepresentation::VARYING_LAST_DIM: { \ + constexpr ShapeRepresentation SHAPE = ShapeRepresentation::VARYING_LAST_DIM; \ + { __VA_ARGS__ } \ + } break; \ + case ShapeRepresentation::VARYING_BOTH_DIMS: { \ + constexpr ShapeRepresentation SHAPE = ShapeRepresentation::VARYING_BOTH_DIMS; \ + { __VA_ARGS__ } \ + } break; \ + default: { \ + NVTE_ERROR("Unsupported grouped tensor shape representation."); \ + } \ + } + //////////////////////////////////////////////////////////////////////////////////////////////////// inline int log2_ceil(int value) { @@ -943,6 +985,8 @@ constexpr size_t scale_tensor_alignment_Y_rowwise = 128; constexpr size_t scale_tensor_alignment_X_colwise = 128; constexpr size_t scale_tensor_alignment_Y_colwise = 4; +constexpr size_t SCALING_FACTORS_SWIZZLE_ALIGNMENT = 128; + // Alignment requirements for the Tensor Memory Accelerator (TMA) constexpr size_t TMA_GMEM_ALIGNMENT = 16; // global memory address alignment constexpr size_t TMA_SHMEM_ALIGNMENT = 128; // shared memory address alignment diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 04e965a9da..0fb73cc439 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -25,13 +25,6 @@ namespace { constexpr int kMaxTensorsPerKernel = 64; constexpr int kThreadsPerWarp = 32; -enum ShapeRepresentation { - SAME_BOTH_DIMS = 0, - VARYING_FIRST_DIM = 1, - VARYING_LAST_DIM = 2, - VARYING_BOTH_DIMS = 3 -}; - __device__ __forceinline__ size_t get_current_tensor_id( const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, const size_t first_logical_dim, const size_t last_logical_dim, diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index 755052d6dd..f650b19dec 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -89,17 +89,18 @@ extern "C" { */ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); -/*! \brief Casts input grouped tensor to MXFP8. +/*! \brief Casts input grouped tensor. * The type of quantized tensor in the output depends on the scaling mode of the output * tensor. See file level comments. * For grouped tensors with a varying last dimension, the last dimension must be a multiple of 128. * * \param[in] input Input grouped tensor to be cast. - * \param[in,out] output Output grouped MXFP8 tensor. + * \param[in,out] output Output grouped tensor. + * \param[in] quant_config Quantization configuration. * \param[in] stream CUDA stream used for the operation. */ void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream); + const NVTEQuantizationConfig quant_config, cudaStream_t stream); /*! \brief Casts input tensor to FP8/MXFP8/BlockwiseFP8, providing the option to immediately exit the kernel * based on the value of the 'noop' tensor. diff --git a/transformer_engine/common/recipe/mxfp8_scaling.cu b/transformer_engine/common/recipe/mxfp8_scaling.cu index 5a6490c042..be692d4563 100644 --- a/transformer_engine/common/recipe/mxfp8_scaling.cu +++ b/transformer_engine/common/recipe/mxfp8_scaling.cu @@ -91,7 +91,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) int r = blockIdx.y * kRowsPerTile + r_; int c = blockIdx.x * kColsPerTile / 32 + c_; size_t idx = r * scale_inv_rowwise_stride + c; - smem_scales_rowwise[r_][c_] = ptx::exp2f_rcp(scale_inv_rowwise[idx]); + smem_scales_rowwise[r_][c_] = ptx::exp2f_rcp(scale_inv_rowwise[idx]); } // Load scales_colwise @@ -100,7 +100,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) int r = blockIdx.y * kRowsPerTile / 32; int c = blockIdx.x * kColsPerTile + c_; size_t idx = r * scale_inv_colwise_stride + c; - smem_scales_colwise[c_] = ptx::exp2f_rcp(scale_inv_colwise[idx]); + smem_scales_colwise[c_] = ptx::exp2f_rcp(scale_inv_colwise[idx]); } __syncthreads(); diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 4d028de01c..1c419d4f8c 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -331,8 +331,8 @@ void nvfp4_2d_partial_cast(const Tensor inp, Tensor out, const Tensor scale, */ // Vectorized transpose kernel parameters -constexpr int TRANSPOSE_TILE_DIM = 64; // Logical FP4 elements per tile dimension -constexpr int TRANSPOSE_TILE_PACKED = 32; // TILE_DIM / 2 bytes +constexpr int TRANSPOSE_TILE_DIM = 64; // Logical FP4 elements per tile dimension +// constexpr int TRANSPOSE_TILE_PACKED = 32; // TILE_DIM / 2 bytes constexpr int TRANSPOSE_BLOCK_SIZE = 256; // threads per block // Shared memory: store unpacked 4-bit values as bytes for easy transpose diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh index f7611e60c5..88a57fe989 100644 --- a/transformer_engine/common/util/ptx.cuh +++ b/transformer_engine/common/util/ptx.cuh @@ -19,6 +19,7 @@ #if FP4_TYPE_SUPPORTED #include #endif // FP4_TYPE_SUPPORTED +#include #include "common/utils.cuh" @@ -326,10 +327,15 @@ __device__ __forceinline__ void get_cancelled_cta_id_2D(__uint128_t *response_da } } +constexpr uint32_t BF16_MANTISSA_BITS = 7; constexpr uint32_t FP32_MANTISSA_BITS = 23; constexpr uint32_t FP32_EXPONENT_BIAS = 127; -__device__ __forceinline__ float exp2f_rcp(e8m0_t biased_exp) { +template +__device__ __forceinline__ T exp2f_rcp(e8m0_t biased_exp); + +template <> +__device__ __forceinline__ float exp2f_rcp(e8m0_t biased_exp) { // Handle the special case of NaN. if (biased_exp == 255) return __int_as_float(0x7fffffff); // Handle the special case where the unbiased exponent is 127, so the reciprocal is 2^-127 which needs the first bit of @@ -339,6 +345,22 @@ __device__ __forceinline__ float exp2f_rcp(e8m0_t biased_exp) { return __int_as_float((254 - biased_exp) << FP32_MANTISSA_BITS); } +template <> +__device__ __forceinline__ bf16 exp2f_rcp(e8m0_t biased_exp) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + // Handle the special case of NaN. + if (biased_exp == 255) return __ushort_as_bfloat16(0x7fff); + // Handle the special case where the unbiased exponent is 127, so the reciprocal is 2^-127 which needs the first bit of + // the mantissa to be 1, which can't be obtained by shifting `BF16_MANTISSA_BITS` bits to the left. + if (biased_exp == 254) return __ushort_as_bfloat16(0x0040); + // Fast calculation when the unbiased exp is in [-126, 126], and only the exponent part is used to express the reciprocal. + return __ushort_as_bfloat16((254 - biased_exp) << BF16_MANTISSA_BITS); +#else + NVTE_DEVICE_ERROR("exp2f_rcp is only supported on SM 9.0+."); + return static_cast(0.0f); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) +} + __device__ __forceinline__ float exp2f(e8m0_t biased_exp) { return __int_as_float(biased_exp << FP32_MANTISSA_BITS); } @@ -493,7 +515,7 @@ struct alignas(2 * sizeof(T)) FPx2 { }; template -struct FPx4 { +struct alignas(4 * sizeof(T)) FPx4 { T x1; T x2; T x3; @@ -1169,6 +1191,142 @@ __device__ __forceinline__ fp16 get_amax(fp16 a, fp16 b) { #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const bf16x4 &in, const bf16x2 scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined CUDA_VERSION) && (CUDA_VERSION >= 13010) + asm volatile( + "{\n\t" + ".reg.b32 x01,x23; \n\t" + "mov.b64 {x01,x23}, %1; \n\t" + ".reg.b32 y01,y23; \n\t" + "mul.rn.bf16x2 y01, x01, %2; \n\t" + "mul.rn.bf16x2 y23, x23, %2; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e4m3x2.bf16x2 z01, y01; \n\t" + "cvt.rn.satfinite.e4m3x2.bf16x2 z23, y23; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in)), + "r"(reinterpret_cast(scale))); +#else + asm volatile( + "{\n\t" + ".reg.b16 scale, scale_flush; \n\t" + "mov.b32 {scale, scale_flush}, %2; \n\t" + ".reg.b16 x0,x1,x2,x3; \n\t" + "mov.b64 {x0,x1,x2,x3}, %1; \n\t" + ".reg.f32 y0,y1,y2,y3; \n\t" + "fma.rn.f32.bf16 y0, x0, scale, 0f00000000; \n\t" + "fma.rn.f32.bf16 y1, x1, scale, 0f00000000; \n\t" + "fma.rn.f32.bf16 y2, x2, scale, 0f00000000; \n\t" + "fma.rn.f32.bf16 y3, x3, scale, 0f00000000; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e4m3x2.f32 z01, y1, y0; \n\t" + "cvt.rn.satfinite.e4m3x2.f32 z23, y3, y2; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in)), + "r"(reinterpret_cast(scale))); +#endif +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const bf16x4 &in, const bf16x2 scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +#if (defined CUDA_VERSION) && (CUDA_VERSION >= 13010) + asm volatile( + "{\n\t" + ".reg.b32 x01,x23; \n\t" + "mov.b64 {x01,x23}, %1; \n\t" + ".reg.b32 y01,y23; \n\t" + "mul.rn.bf16x2 y01, x01, %2; \n\t" + "mul.rn.bf16x2 y23, x23, %2; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e5m2x2.bf16x2 z01, y01; \n\t" + "cvt.rn.satfinite.e5m2x2.bf16x2 z23, y23; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in)), + "r"(reinterpret_cast(scale))); +#else + asm volatile( + "{\n\t" + ".reg.b16 scale, scale_flush; \n\t" + "mov.b32 {scale, scale_flush}, %2; \n\t" + ".reg.b16 x0,x1,x2,x3; \n\t" + "mov.b64 {x0,x1,x2,x3}, %1; \n\t" + ".reg.f32 y0,y1,y2,y3; \n\t" + "fma.rn.f32.bf16 y0, x0, scale, 0f00000000; \n\t" + "fma.rn.f32.bf16 y1, x1, scale, 0f00000000; \n\t" + "fma.rn.f32.bf16 y2, x2, scale, 0f00000000; \n\t" + "fma.rn.f32.bf16 y3, x3, scale, 0f00000000; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e5m2x2.f32 z01, y1, y0; \n\t" + "cvt.rn.satfinite.e5m2x2.f32 z23, y3, y2; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in)), + "r"(reinterpret_cast(scale))); +#endif +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const fp16x4 &in, const fp16 scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile( + "{\n\t" + ".reg.b16 x0,x1,x2,x3; \n\t" + "mov.b64 {x0,x1,x2,x3}, %1; \n\t" + ".reg.f32 y0,y1,y2,y3; \n\t" + "fma.rn.f32.f16 y0, x0, %2, 0f00000000; \n\t" + "fma.rn.f32.f16 y1, x1, %2, 0f00000000; \n\t" + "fma.rn.f32.f16 y2, x2, %2, 0f00000000; \n\t" + "fma.rn.f32.f16 y3, x3, %2, 0f00000000; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e4m3x2.f32 z01, y1, y0; \n\t" + "cvt.rn.satfinite.e4m3x2.f32 z23, y3, y2; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in)), + "h"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +__device__ __forceinline__ void mul_cvt_4x(fp8e5m2x4 &out, const fp16x4 &in, const fp16 scale) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + asm volatile( + "{\n\t" + ".reg.b16 x0,x1,x2,x3; \n\t" + "mov.b64 {x0,x1,x2,x3}, %1; \n\t" + ".reg.f32 y0,y1,y2,y3; \n\t" + "fma.rn.f32.f16 y0, x0, %2, 0f00000000; \n\t" + "fma.rn.f32.f16 y1, x1, %2, 0f00000000; \n\t" + "fma.rn.f32.f16 y2, x2, %2, 0f00000000; \n\t" + "fma.rn.f32.f16 y3, x3, %2, 0f00000000; \n\t" + ".reg.b16 z01, z23; \n\t" + "cvt.rn.satfinite.e5m2x2.f32 z01, y1, y0; \n\t" + "cvt.rn.satfinite.e5m2x2.f32 z23, y3, y2; \n\t" + "mov.b32 %0, {z01, z23}; \n" + "}\n" + : "=r"(reinterpret_cast(out)) + : "l"(reinterpret_cast(in)), + "h"(reinterpret_cast(scale))); +#else + NVTE_DEVICE_ERROR("mul_cvt_4x is only supported on SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + __device__ __forceinline__ void mul_cvt_4x(fp8e4m3x4 &out, const bf16x4 &in, const ptx::floatx2 &scale) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 26549191a3..8c50e83926 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -928,6 +928,13 @@ using e8m0_t = uint8_t; enum ScalingType { ROWWISE = 0, COLWISE = 1, BIDIMENSIONAL = 2 }; +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + template struct Numeric_Traits; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index cb3434ec52..e126e0199a 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -217,9 +217,10 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const break; } case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: { + QuantizationConfigWrapper quant_config_cpp; NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), - at::cuda::getCurrentCUDAStream()); + quant_config_cpp, at::cuda::getCurrentCUDAStream()); }); break; } From 9d77dcb0638e7c3298c708df595035c0297cdad0 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Thu, 2 Apr 2026 23:07:03 -0700 Subject: [PATCH 308/521] [JAX] Fix: Use jitted kernels for generating THD (and BSHD) segment pos (#2823) * Fix: Use jitted kernels for generating THD (and BSHD) segment pos if only segment id is passed Signed-off-by: Kshitij Lakhani * Make passing of segment_pos to from_segmet_ids_and_pos for creating a SequenceDescriptor mandatory Signed-off-by: Kshitij Lakhani * Make test changes for from_segmet_ids_and_pos API change. Also some nits. Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Make segment_pos arg mandatory and not Optional Signed-off-by: Kshitij Lakhani * Add comments for from_segment_ids_and_pos Signed-off-by: Kshitij Lakhani * nit: Change data types for BSHD seg pos and seg id to be int32 adn consistent with THD when setting up test inputs Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Replace a TypeError if segment_pos is not passed with a ValueError with a message 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 | 61 ++++++++--------- transformer_engine/jax/attention.py | 102 ++++++---------------------- 2 files changed, 52 insertions(+), 111 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index f9946e1f7f..8b727b1d43 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -547,13 +547,20 @@ def _setup_inputs(self): else: self.softmax_offset = None - def gen_valid(bs, max_seqlen, pad_ratio): + def generate_valid_segment_ids_and_pos(bs, max_seqlen, pad_ratio): pad_len = int(max_seqlen * pad_ratio) valid_len = max_seqlen - pad_len - tokens = jnp.concatenate([jnp.ones((bs, valid_len)), jnp.zeros((bs, pad_len))], axis=-1) - return tokens, jnp.logical_not(tokens) + tokens = jnp.concatenate( + [ + jnp.ones((bs, valid_len), dtype=jnp.int32), + jnp.zeros((bs, pad_len), dtype=jnp.int32), + ], + axis=-1, + ) + segment_pos = jnp.broadcast_to(jnp.arange(max_seqlen, dtype=jnp.int32), tokens.shape) + return tokens, segment_pos, jnp.logical_not(tokens) - def generate_random_segment_ids( + def generate_random_segment_ids_and_pos( batch_size, sequence_length, num_segments, @@ -601,8 +608,10 @@ def generate_random_segment_ids( return segment_ids, segment_pos, segment_pad if self.qkv_layout.is_thd(): - self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_random_segment_ids( - self.batch_size, self.max_seqlen_q, self.num_segments_per_seq, seed=42 + self.segment_ids_q, self.segment_pos_q, self.pad_q = ( + generate_random_segment_ids_and_pos( + self.batch_size, self.max_seqlen_q, self.num_segments_per_seq, seed=42 + ) ) self.seqlens_q, self.offsets_q = get_seqlens_and_offsets(self.segment_ids_q) # TODO(rewang): record only self attention and find the reason of cross attention @@ -617,22 +626,23 @@ def generate_random_segment_ids( self.window_size is not None or self.attn_mask_type.is_bottom_right() ): # SWA or BRCM requires kv_len >= q_len min_segment_len = self.seqlens_q - self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = generate_random_segment_ids( - self.batch_size, - self.max_seqlen_kv, - self.num_segments_per_seq, - seed=2024, - min_segment_len=min_segment_len, + self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = ( + generate_random_segment_ids_and_pos( + self.batch_size, + self.max_seqlen_kv, + self.num_segments_per_seq, + seed=2024, + min_segment_len=min_segment_len, + ) ) self.seqlens_kv, self.offsets_kv = get_seqlens_and_offsets(self.segment_ids_kv) else: - self.segment_ids_q, self.pad_q = gen_valid( + self.segment_ids_q, self.segment_pos_q, self.pad_q = generate_valid_segment_ids_and_pos( self.batch_size, self.max_seqlen_q, pad_ratio ) - self.segment_ids_kv, self.pad_kv = gen_valid( - self.batch_size, self.max_seqlen_kv, pad_ratio + self.segment_ids_kv, self.segment_pos_kv, self.pad_kv = ( + generate_valid_segment_ids_and_pos(self.batch_size, self.max_seqlen_kv, pad_ratio) ) - self.segment_pos_q = self.segment_pos_kv = None self.seqlens_q = self.seqlens_kv = self.offsets_q = self.offsets_kv = None # For reference code @@ -682,24 +692,15 @@ def generate_random_segment_ids( (self.offsets_q, self.offsets_kv), ) case SeqDescFormat.SegmentIDs: - # Exercise the path to generate the segment_pos in from_segment_ids_and_pos() - # if no CP and load balancing, else explicitly pass the segment_pos + # from_segment_ids_and_pos requires explicit segment_pos. self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos( ( self.cp_reorder_fn(self.segment_ids_q), self.cp_reorder_fn(self.segment_ids_kv), ), ( - ( - self.cp_reorder_fn(self.segment_pos_q), - self.cp_reorder_fn(self.segment_pos_kv), - ) - if self.cp_size > 1 and self.cp_load_balanced - else None - ), - is_thd=self.qkv_layout.is_thd(), - is_segment_ids_reordered=( - True if self.cp_size > 1 and self.cp_load_balanced else False + self.cp_reorder_fn(self.segment_pos_q), + self.cp_reorder_fn(self.segment_pos_kv), ), ) case _: @@ -727,9 +728,7 @@ def generate_random_segment_ids( case SeqDescFormat.SegmentIDs: self.sequence_desciptor = SequenceDescriptor.from_segment_ids_and_pos( (self.segment_ids_q, self.segment_ids_kv), - None, - is_thd=self.qkv_layout.is_thd(), - is_segment_ids_reordered=False, + (self.segment_pos_q, self.segment_pos_kv), ) case _: raise ValueError(f"Unknown {self.seq_desc_format=}") diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 99817f0657..29d0848381 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -855,14 +855,9 @@ def from_segment_ids_and_pos( cls, segment_ids: Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]], segment_pos: Optional[Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]]] = None, - *, - is_thd: bool, - is_segment_ids_reordered: bool, ) -> SequenceDescriptor: """ - Experimental factory method for inputs with segment IDs and optional positions. - segment_pos = None to be used only for: BSHD with or without load balancing and, - THD without load balancing + Experimental factory method for inputs with segment IDs and positions. Args: segment_ids(Tuple(jnp.ndarray, jnp.ndarray)) = (q_segment_ids, kv_segment_ids): - q_segment_ids (jnp.ndarray): @@ -876,88 +871,35 @@ def from_segment_ids_and_pos( The position inside each segment for query, with shape [batch, max_seqlen]. - kv_segment_pos (jnp.ndarray): The position inside each segment for key, value, with shape [batch, max_seqlen]. - is_thd(bool): If True, QKVLayout is of type THD, else it is BSHD - is_segment_ids_reordered(bool): If True, the segment ids have been reordered for load balancing. - Only THD with load balancing is expected to have this flag set to True Return: A SequenceDescriptor with segment_ids/segment_pos initialized. """ - q_seg_ids, kv_seg_ids = cls._expand_to_pair(segment_ids) - - # Using defaults : segment pos has to be generated. + # Examples (0 in segment_ids means padding): + # THD (three segments packed together in a sequence of length 16 with no intra-segment padding): + # segment_ids = [1, 1, 1, 2, 2, 3, 3, 3, 3, 3, 0, 0, 0, 0, 0, 0] + # segment_pos = [0, 1, 2, 0, 1, 0, 1, 2, 3, 4, 0, 0, 0, 0, 0, 0] + # THD (three segments packed together in a sequence of length 16 with intra-segment padding): + # segment_ids = [1, 1, 1, 2, 2, 3, 3, 3, 0, 0, 4, 4, 0, 0, 0, 0] + # segment_pos = [0, 1, 2, 0, 1, 0, 1, 2, 3, 4, 0, 1, 0, 0, 0, 0] + # BSHD (only one segment per sequence): + # segment_ids = [1, 1, 1, 1, 1, 1, 1, 0, 0] + # segment_pos = [0, 1, 2, 3, 4, 5, 6, 7, 8] + # TODO(@KshitijLakhani): Make segment_pos Union[jnp.ndarray, Tuple[jnp.ndarray, jnp.ndarray]] and remove below check (starting June 2026) if segment_pos is None: - # THD + load balanced segment_ids are not supported in this function - # BSHD + load balanced segment_ids are incorrect as BSHD handles reordering within the primitive itself - if is_segment_ids_reordered: - assert not is_thd, ( - f"{segment_pos=} default arg is not supported for load balanced reordered" - " (Striped) THD inputs. Please pass the load balanced reordered segment_pos" - " and segment_ids explicitly to {from_segment_ids_and_pos.__qualname__}" - " using convenience function reorder_causal_load_balancing()" - ) - assert is_thd, ( - f"{segment_pos=} default arg is not supported for load balanced reordered (Dual" - " Chunk) BSHD inputs. BSHD segment_pos and segment_ids do not need to be load" - " balanced reordered. The reordering for these is performed within the" - " primitive" - ) + raise ValueError( + "segment_pos is now required. Automatic segment_pos generation was removed because" + " it did not have sufficient context to generate a correct segment_pos across all" + " load-balancing and context-parallel strategies. Please generate the segment_pos" + " explicitly.See tests/jax/test_fused_attn.py generate_random_segment_ids_and_pos()" + " and generate_valid_segment_ids_and_pos()" + ) - # Generate the default pos for THD and BSHD non-reordered segment_ids - def generate_default_pos(seg_ids): - if is_thd: - batch_size, seq_size = seg_ids.shape - # Assume that the first token belongs to a segment and is not a padded token - first_is_segment = jnp.full((batch_size, 1), True, dtype=bool) - # Get segment start positions - segment_start = jnp.concatenate( - [ - first_is_segment, - (seg_ids[..., 1:] != seg_ids[..., :-1]) & (seg_ids[..., 1:] != 0), - ], - axis=-1, - ) - # Get offset for location where new segment starts - segment_start_idx = jax.vmap(lambda row: jnp.arange(row.size) * row)( - segment_start - ) - segment_start_offsets = jax.vmap(jnp.maximum.accumulate)(segment_start_idx) - - # Get the last non-zero index - after this everything is padding - # (B,) - last_nonzero_idx = jax.vmap( - lambda segids_row: jnp.max( - jnp.where(segids_row != 0, jnp.arange(seq_size), -1) - ) - )(seg_ids) - seg_pos_no_thd = jnp.arange(seq_size) - # Get a mask which can be used to zero out all the padding at the end (after the non-zero index) - mask = seg_pos_no_thd <= last_nonzero_idx[:, None] - - # Get the unmasked seg_pos for the THD sequence - seg_pos = ( - jnp.broadcast_to(jnp.arange(seq_size), seg_ids.shape) - - segment_start_offsets - ) - - # Use the mask to zero out the padding at the end (after the non-zero index) - segment_pos = jax.vmap( - lambda pos_row, mask_row: jnp.where(mask_row, pos_row, 0) - )(seg_pos, mask) - return segment_pos - - seqlen = seg_ids.shape[-1] - return jnp.broadcast_to(jnp.arange(seqlen), seg_ids.shape) - - q_seg_pos = generate_default_pos(q_seg_ids) - kv_seg_pos = generate_default_pos(kv_seg_ids) - segment_pos = (q_seg_pos, kv_seg_pos) - # Explicitly passed segment_pos - else: - segment_pos = cls._expand_to_pair(segment_pos) + q_seg_ids, kv_seg_ids = cls._expand_to_pair(segment_ids) + q_seg_pos, kv_seg_pos = cls._expand_to_pair(segment_pos) return cls( segment_ids=(q_seg_ids, kv_seg_ids), - segment_pos=segment_pos, + segment_pos=(q_seg_pos, kv_seg_pos), ) From 29a8c2fec3db6453280cf5ce9824b52c1eda2e57 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 3 Apr 2026 11:41:50 -0400 Subject: [PATCH 309/521] GEMM + Swiglu fused Grouped MLP for MXFP8 (#2769) * GEMM + Swiglu fused Grouped MLP for MXFP8 Signed-off-by: Kirthi Shankar Sivamani * cleanup/lint Signed-off-by: Kirthi Shankar Sivamani * Properly cache the alpha tensor Signed-off-by: Kirthi Shankar Sivamani * nD dummy grad Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * 0 tokens in entire rank Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tmp downgrade cublas version check Signed-off-by: Kirthi Shankar Sivamani * delayed wgrad tests pass for basic gl Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * merge everything Signed-off-by: Varun Thumbe Signed-off-by: Kirthi Shankar Sivamani * Rebase into fused_mxfp8_grouped_mlp; unit tests for delayed wgrad working Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Kirthi Shankar Sivamani * Fix tests being skipped for fusible ops Signed-off-by: Kirthi Shankar Sivamani * Integrate mxfp8 dbias kernel in group_quantize Signed-off-by: Kirthi Shankar Sivamani * Add bias/dbias fused support with cute GEMMs Signed-off-by: Kirthi Shankar Sivamani * Check bias/dbias support Signed-off-by: Kirthi Shankar Sivamani * Pack biases more efficiently Signed-off-by: Kirthi Shankar Sivamani * GroupedTensor for biases to avoid concat Signed-off-by: Kirthi Shankar Sivamani * format Signed-off-by: Kirthi Shankar Sivamani * Support 1D grouped tensor shape for bias and fix checkpointing Signed-off-by: Kirthi Shankar Sivamani * Fixes and tests Signed-off-by: Kirthi Shankar Sivamani * Refactor grouped tensor marking for paged stashing Signed-off-by: Kirthi Shankar Sivamani * Remove setting logical_shape in mark_grouped_tensor Signed-off-by: Kirthi Shankar Sivamani * Cleanup logical_shape Signed-off-by: Kirthi Shankar Sivamani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * pass the tests for now Signed-off-by: Varun Thumbe * address some review comments Signed-off-by: Varun Thumbe * 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 * more cleanups 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 * refactor wgrad logic Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Rename argument from single_grouped_parameter to single_grouped_weight Signed-off-by: Kirthi Shankar Sivamani * Check wgrad store context is not empty for 0 token case. Signed-off-by: Kirthi Shankar Sivamani * Test only checks for fusion if fused kernel is available Signed-off-by: Tim Moon * fix the tolerance to be of bf16 for the cute gemm Signed-off-by: Varun Thumbe * Update transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * address further 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 more review comments Signed-off-by: Varun Thumbe * address more review comments + test for zero grouped tensor work case Signed-off-by: Varun Thumbe * cublaslt remove zero work gemm avoidance 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 the wgrad test Signed-off-by: Varun Thumbe * split dbias functionality from gq api Signed-off-by: Kirthi Shankar Sivamani * Format and lint Signed-off-by: Kirthi Shankar Sivamani * port fixes and add better doc for page stashing war Signed-off-by: Kirthi Shankar Sivamani * Guard fusion via env Signed-off-by: Kirthi Shankar Sivamani * Change to trigger CI Remove unnecessary blank line in docstring. * To retrigger CI * Space to trigger the pipeline * fix zero work cublas gemm 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: Kirthi Shankar Sivamani Signed-off-by: Varun Thumbe Signed-off-by: Tim Moon Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Varun Thumbe Co-authored-by: Vasudevan Rengasamy Co-authored-by: Tim Moon Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- qa/L0_pytorch_unittest/test.sh | 2 +- tests/cpp/operator/test_grouped_gemm.cu | 373 +++++++++- tests/cpp/operator/test_swizzle.cu | 144 ++++ tests/pytorch/test_fusible_ops.py | 536 +++++++++++++- tests/pytorch/test_grouped_tensor.py | 96 ++- tests/pytorch/test_numerics.py | 115 ++- tests/pytorch/test_sanity.py | 19 +- transformer_engine/common/CMakeLists.txt | 1 + .../common/gemm/cublaslt_grouped_gemm.cu | 102 ++- .../common/include/transformer_engine/utils.h | 36 + transformer_engine/common/util/utils.cu | 51 ++ transformer_engine/pytorch/csrc/common.h | 1 + transformer_engine/pytorch/csrc/extensions.h | 11 + .../pytorch/csrc/extensions/cast.cpp | 58 ++ .../pytorch/csrc/extensions/gemm.cpp | 17 +- .../pytorch/csrc/extensions/pybind.cpp | 14 + .../pytorch/csrc/extensions/swizzle.cpp | 86 ++- .../pytorch/csrc/extensions/utils.cpp | 165 +++++ .../pytorch/csrc/type_converters.cpp | 4 + transformer_engine/pytorch/csrc/util.h | 9 +- transformer_engine/pytorch/module/base.py | 12 +- .../pytorch/module/grouped_linear.py | 137 +++- transformer_engine/pytorch/ops/_common.py | 114 +++ .../pytorch/ops/basic/grouped_linear.py | 446 ++++++++++-- .../pytorch/ops/fused/__init__.py | 9 + .../pytorch/ops/fused/backward_grouped_mlp.py | 679 ++++++++++++++++++ .../pytorch/ops/fused/forward_grouped_mlp.py | 573 +++++++++++++++ .../pytorch/tensor/grouped_tensor.py | 13 +- .../tensor/storage/grouped_tensor_storage.py | 159 +++- transformer_engine/pytorch/utils.py | 36 + 30 files changed, 3784 insertions(+), 234 deletions(-) create mode 100644 transformer_engine/common/include/transformer_engine/utils.h create mode 100644 transformer_engine/common/util/utils.cu create mode 100644 transformer_engine/pytorch/csrc/extensions/utils.cpp create mode 100644 transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py create mode 100644 transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index f2b0b07fed..e67cf1bc04 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -41,7 +41,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_ 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" -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_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_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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu index 34bb729b25..bcacb2f801 100644 --- a/tests/cpp/operator/test_grouped_gemm.cu +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -88,7 +88,6 @@ 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 @@ -175,8 +174,8 @@ std::vector> make_shapes(ShapeCase scase) { } void run_grouped_gemm_case(const TestParams& params) { -#if CUBLAS_VERSION < 130200 - GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.2+, but compile-time cuBLAS version is " +#if CUBLAS_VERSION < 130300 + GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.3+, but compile-time cuBLAS version is " << CUBLAS_VERSION << "."; #else if (getDeviceComputeCapability() < blackwellComputeCapability) { @@ -349,7 +348,365 @@ void run_grouped_gemm_case(const TestParams& params) { atol, rtol); } -#endif // CUBLAS_VERSION >= 130200 +#endif // CUBLAS_VERSION >= 130300 +} + +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."; + } + + 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); + + 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)); + } + + 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); + + // 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]); + } + + 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()); + + std::vector C_tensors; + std::vector D_list_tensors; + C_tensors.reserve(num_gemms); + D_list_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)); + } + D_list_tensors.emplace_back( + Tensor("D_list" + std::to_string(i), std::vector{M, N}, DType::kBFloat16)); + NVTE_CHECK_CUDA(cudaMemset(D_list_tensors.back().rowwise_dptr(), 0, + bytes(D_list_tensors.back().rowwise_shape(), + D_list_tensors.back().dtype()))); + } + + std::vector C_list_ptrs; + std::vector D_list_ptrs; + if (!params.use_null_c) { + C_list_ptrs.reserve(num_gemms); + } + D_list_ptrs.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) { + if (!params.use_null_c) { + C_list_ptrs.push_back(C_tensors[i].data()); + } + 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()); + + // 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 +} + +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."; + } + + 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); + + 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)); + } + + 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); + + // 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]); + } + + 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; + 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)); + } + D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i), + std::vector{M, N}, + DType::kBFloat16)); + NVTE_CHECK_CUDA(cudaMemset(D_group_tensors.back().rowwise_dptr(), 0, + bytes(D_group_tensors.back().rowwise_shape(), + D_group_tensors.back().dtype()))); + } + + 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]); + } + D_views.push_back(&D_group_tensors[i]); + } + + std::optional grouped_C; + if (!params.use_null_c) { + grouped_C = build_grouped_tensor(C_views, NVTE_DELAYED_TENSOR_SCALING); + } + 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); + + 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()); + } + + 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_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 } class GroupedGemmTest : public ::testing::TestWithParam {}; @@ -358,6 +715,14 @@ TEST_P(GroupedGemmTest, CompareWithMultiTensorGemm) { run_grouped_gemm_case(GetParam()); } +TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteOut) { + run_grouped_gemm_discrete_out_case(GetParam()); +} + +TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteIn) { + run_grouped_gemm_discrete_in_case(GetParam()); +} + std::string MakeGroupedGemmTestName(const testing::TestParamInfo& info) { constexpr const char* kInputNames[] = {"FP8Current", "BF16", "MXFP8"}; constexpr const char* kShapeNames[] = {"AllSame", "SameM", "SameN", "AllDiff"}; diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 694b348a9b..8389989efe 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -110,6 +110,115 @@ void performTestSwizzle1D(const int num_tiles_M, const int num_tiles_K, bool row } } +// Zero out padding in a scale_inv CPU buffer so that the CPU reference +// matches the kernel, which zeroes elements outside the original dims. +// The buffer is stored in leading-dim-major order (row-major for rowwise, +// column-major for colwise). `padded_rows x padded_cols` is the full +// (padded) shape; `orig_rows` / `orig_cols` are the unpadded extents. +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; + } + } + } +} + +void performTestGroupedSwizzleMXFP8(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; + 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; + 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()); + + // The grouped swizzle kernel zeroes scale_inv elements that fall + // outside the original (unpadded) dimensions. Mirror that in the + // per-tensor CPU buffers so the CPU reference produces identical output. + 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, (K + 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], + (M + BLOCK_SIZE - 1) / 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)); + } + + 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)); + + 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]; + + NVTE_CHECK_CUDA(cudaMemset(grouped_output.scale_inv.get(), 0, num_tensors * row_numel)); + NVTE_CHECK_CUDA(cudaMemset(grouped_output.columnwise_scale_inv.get(), 0, num_tensors * col_numel)); + + nvte_swizzle_grouped_scaling_factors(grouped_input.get_handle(), + grouped_output.get_handle(), 0); + + 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_rowwise", output_row.data(), ref_row.data(), + num_tensors * row_numel); + compareResults("grouped_swizzle_colwise", output_col.data(), ref_col.data(), + num_tensors * col_numel); +} + class SwizzleTestSuite : public ::testing::TestWithParam, std::pair, bool>> {}; @@ -126,6 +235,41 @@ TEST_P(SwizzleTestSuite, TestSwizzle) { transa); } +class SwizzleGroupedTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(SwizzleGroupedTestSuite, TestGroupedSwizzleMXFP8) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + performTestGroupedSwizzleMXFP8(num_tensors, M, K); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + SwizzleGroupedTestSuite, + ::testing::Values( + // M and K both divisible by 128 + std::make_tuple(3, 256, 256), + std::make_tuple(4, 128, 128), + // M not divisible by 128 + std::make_tuple(3, 200, 256), + std::make_tuple(2, 65, 256), + // K not divisible by 128 + std::make_tuple(3, 256, 160), + std::make_tuple(2, 256, 96), + // Neither M nor K divisible by 128 + std::make_tuple(3, 200, 160), + std::make_tuple(4, 33, 64), + std::make_tuple(2, 1, 32) + ), + [](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)); + } +); + namespace { std::vector> num_tiles = { diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index b97afbc191..75d450b46b 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -18,6 +18,7 @@ import transformer_engine.common.recipe import transformer_engine.pytorch as te import transformer_engine.pytorch.ops as te_ops + from transformer_engine.pytorch.ops.fused import ( BackwardActivationBias, BackwardAddRMSNorm, @@ -35,6 +36,8 @@ NVFP4Quantizer, 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 @@ -2008,6 +2011,7 @@ def test_dropout( @pytest.mark.parametrize("quantized_weight", (False, True)) @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("weight_requires_grad", (False, True)) + @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) def test_grouped_linear( self, *, @@ -2022,6 +2026,7 @@ def test_grouped_linear( quantized_weight: bool, input_requires_grad: bool, weight_requires_grad: bool, + delay_wgrad_compute: bool, ) -> None: """Grouped GEMM""" @@ -2102,6 +2107,7 @@ def test_grouped_linear( bias=bias, device=device, dtype=dtype, + delay_wgrad_compute=delay_wgrad_compute, ) with torch.no_grad(): for group_idx in range(group_size): @@ -2117,6 +2123,8 @@ def test_grouped_linear( y_test = op(x_test, split_sizes) if input_requires_grad or weight_requires_grad: y_test.backward(dy_test) + if delay_wgrad_compute and weight_requires_grad: + op.backward_dw() # Expected numerical error tols = dtype_tols(dtype) @@ -3236,7 +3244,11 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("quantization", _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)) @pytest.mark.parametrize("glu_interleave_size", (None, 32)) + @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) def test_grouped_mlp( self, *, @@ -3245,14 +3257,18 @@ def test_grouped_mlp( hidden_size: int = 256, dtype: torch.dtype, quantization: Optional[str], + single_grouped_weight: bool, + single_grouped_bias: bool, + accumulate_into_main_grad: bool, device: torch.device = "cuda", split_alignment: int = 256, glu_interleave_size: Optional[int], + delay_wgrad_compute: bool, ) -> None: """GroupedLinear + ScaledSwiGLU + GroupedLinear""" # Split sizes - split_sizes = [split_alignment * i for i in range(group_size)] + split_sizes = [split_alignment * (i) for i in range(group_size)] random.shuffle(split_sizes) split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) @@ -3263,8 +3279,15 @@ def test_grouped_mlp( # Skip invalid configurations with_quantization = quantization is not None 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") + if single_grouped_bias and not bias: + 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 quantization == "mxfp8" and bias: + # Will be supported in future CUDNN release. + pytest.skip("Bias/dbias not yet supported in MXFP8 fused grouped MLP") # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -3370,6 +3393,10 @@ def test_grouped_mlp( 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, @@ -3378,6 +3405,10 @@ def test_grouped_mlp( 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, ) module = te_ops.Sequential( fc1, @@ -3387,18 +3418,87 @@ def test_grouped_mlp( # Copy weights with torch.no_grad(): + if single_grouped_weight: + fc1_weights = fc1.weight.quantized_tensors + if fc1_weights is None: + fc1_weights = fc1.weight.split_into_quantized_tensors() + fc2_weights = fc2.weight.quantized_tensors + if fc2_weights is None: + fc2_weights = fc2.weight.split_into_quantized_tensors() for group_idx in range(group_size): - getattr(fc1, f"weight{group_idx}").copy_(fc1_ws_test[group_idx]) - getattr(fc2, f"weight{group_idx}").copy_(fc2_ws_test[group_idx]) + if single_grouped_weight: + fc1_weights[group_idx].copy_(fc1_ws_test[group_idx]) + fc2_weights[group_idx].copy_(fc2_ws_test[group_idx]) + else: + getattr(fc1, f"weight{group_idx}").copy_(fc1_ws_test[group_idx]) + getattr(fc2, f"weight{group_idx}").copy_(fc2_ws_test[group_idx]) if bias: - 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 single_grouped_bias: + fc1_bparts = fc1.bias.split_into_quantized_tensors() + fc2_bparts = fc2.bias.split_into_quantized_tensors() + fc1_bparts[group_idx].reshape(-1).copy_(fc1_bs_test[group_idx]) + fc2_bparts[group_idx].reshape(-1).copy_(fc2_bs_test[group_idx]) + else: + 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: + 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, + ) + 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, + ) del fc1_ws_test, fc1_bs_test, fc2_ws_test, fc2_bs_test # Fuse ops and perform forward and backward pass with te.autocast(enabled=with_quantization, recipe=recipe): y_test = module(x_test, split_sizes, probs_test, split_sizes) y_test.backward(dy_test) + if delay_wgrad_compute: + fc1.backward_dw() + fc2.backward_dw() + + # Check for expected fusions + if ( + quantization == "mxfp8" + and dtype in (torch.bfloat16, torch.float16) + and glu_interleave_size == 32 + ): + if te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.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, + ) + if te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.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, + ) # Loose tols for sanity checking tols = {"rtol": 0.125, "atol": 0.25} @@ -3410,10 +3510,286 @@ def test_grouped_mlp( assert_close_grads(x_test, x_ref, **tols) assert_close_grads(probs_test, probs_ref, **tols) for group_idx in range(group_size): - assert_close_grads(getattr(fc2, f"weight{group_idx}"), fc2_ws_ref[group_idx], **tols) - assert_close_grads(getattr(fc2, f"bias{group_idx}"), fc2_bs_ref[group_idx], **tols) - assert_close_grads(getattr(fc1, f"weight{group_idx}"), fc1_ws_ref[group_idx], **tols) - assert_close_grads(getattr(fc1, f"bias{group_idx}"), fc1_bs_ref[group_idx], **tols) + if bias: + if single_grouped_bias: + assert_close( + fc2.bias.grad[group_idx], + fc2_bs_ref[group_idx].grad, + **tols, + ) + assert_close( + fc1.bias.grad[group_idx], + fc1_bs_ref[group_idx].grad, + **tols, + ) + else: + assert_close_grads( + getattr(fc2, f"bias{group_idx}"), fc2_bs_ref[group_idx], **tols + ) + assert_close_grads( + getattr(fc1, f"bias{group_idx}"), fc1_bs_ref[group_idx], **tols + ) + if not single_grouped_weight and not accumulate_into_main_grad: + assert_close_grads( + getattr(fc2, f"weight{group_idx}"), fc2_ws_ref[group_idx], **tols + ) + assert_close_grads( + getattr(fc1, f"weight{group_idx}"), fc1_ws_ref[group_idx], **tols + ) + 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) + elif single_grouped_weight: + assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols) + assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols) + + @pytest.mark.parametrize("dtype", _dtypes) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) + @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_grouped_mlp_cuda_graph_safe_mxfp8( + self, + *, + dtype: torch.dtype, + single_grouped_weight: bool, + accumulate_into_main_grad: bool, + device: torch.device = "cuda", + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + ) -> None: + """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" + + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_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") + + 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) + + recipe = make_recipe("mxfp8") + 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, + ) + 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, + ) + module = te_ops.Sequential( + fc1, + te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), + fc2, + ) + + def _init_main_grads(value: float = 0.0) -> None: + if not accumulate_into_main_grad: + return + with torch.no_grad(): + if single_grouped_weight: + if getattr(fc1.weight, "main_grad", None) is None: + fc1.weight.main_grad = torch.empty( + fc1.weight.size(), + device=device, + dtype=torch.float32, + ) + if getattr(fc2.weight, "main_grad", None) is None: + fc2.weight.main_grad = torch.empty( + fc2.weight.size(), + device=device, + dtype=torch.float32, + ) + fc1.weight.main_grad.fill_(value) + fc2.weight.main_grad.fill_(value) + else: + for group_idx in range(group_size): + fc1_weight = getattr(fc1, f"weight{group_idx}") + fc2_weight = getattr(fc2, f"weight{group_idx}") + if getattr(fc1_weight, "main_grad", None) is None: + fc1_weight.main_grad = torch.empty( + fc1_weight.size(), + device=device, + dtype=torch.float32, + ) + if getattr(fc2_weight, "main_grad", None) is None: + fc2_weight.main_grad = torch.empty( + fc2_weight.size(), + device=device, + dtype=torch.float32, + ) + fc1_weight.main_grad.fill_(value) + fc2_weight.main_grad.fill_(value) + + def _collect_main_grads() -> tuple[torch.Tensor, torch.Tensor]: + if single_grouped_weight: + fc1_main_grad = fc1.weight.main_grad.detach().clone() + fc2_main_grad = fc2.weight.main_grad.detach().clone() + else: + fc1_main_grad = torch.stack( + [ + getattr(fc1, f"weight{group_idx}").main_grad.detach().clone() + for group_idx in range(group_size) + ], + dim=0, + ) + fc2_main_grad = torch.stack( + [ + getattr(fc2, f"weight{group_idx}").main_grad.detach().clone() + for group_idx in range(group_size) + ], + dim=0, + ) + return fc1_main_grad, fc2_main_grad + + static_split_sizes = split_sizes.clone() + + def train_step( + x: torch.Tensor, + probs: torch.Tensor, + dy: torch.Tensor, + out_buf: torch.Tensor, + *, + use_graphed: bool, + ) -> torch.Tensor: + with te.autocast(enabled=True, recipe=recipe): + out = ( + graphed_module(x, static_split_sizes, probs, static_split_sizes) + if use_graphed + else module(x, static_split_sizes, probs, 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_probs = torch.randn((in_shape[0],), device=device, dtype=dtype, requires_grad=True) + static_dy = torch.randn(in_shape, device=device, dtype=dtype) + static_out_buf = torch.empty((in_shape[0], hidden_size), device=device, dtype=dtype) + + graphed_module = te.make_graphed_callables( + module, + (static_x, static_split_sizes, static_probs, static_split_sizes), + num_warmup_iters=3, + enabled=True, + recipe=recipe, + ) + + forward_ops = module._module_groups[0]._forward_ops + backward_ops = module._module_groups[0]._backward_ops + assert len(forward_ops) == 1 + assert isinstance( + forward_ops[0][0], + te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + ) + assert len(backward_ops) == 1 + assert isinstance( + backward_ops[0][0], + te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + ) + + fresh_x = torch.randn_like(static_x) + fresh_probs = torch.randn_like(static_probs) + fresh_dy = torch.randn_like(static_dy) + with torch.no_grad(): + static_x.copy_(fresh_x) + static_probs.copy_(fresh_probs) + static_dy.copy_(fresh_dy) + + for param in module.parameters(): + param.grad = torch.zeros_like(param) + _init_main_grads(0.5) + if static_x.grad is not None: + static_x.grad.zero_() + if static_probs.grad is not None: + static_probs.grad.zero_() + + graph_out = ( + train_step(static_x, static_probs, static_dy, static_out_buf, use_graphed=True) + .detach() + .clone() + ) + torch.cuda.synchronize() + graph_dx = static_x.grad.detach().clone() + graph_dprobs = static_probs.grad.detach().clone() + if accumulate_into_main_grad: + graph_fc1_main_grad, graph_fc2_main_grad = _collect_main_grads() + else: + graph_param_grads = [param.grad.detach().clone() for param in module.parameters()] + + for param in module.parameters(): + param.grad.zero_() + _init_main_grads(0.5) + static_x.grad.zero_() + static_probs.grad.zero_() + + expected_x = fresh_x.detach().clone().requires_grad_(True) + expected_probs = fresh_probs.detach().clone().requires_grad_(True) + expected_dy = fresh_dy.detach().clone() + with te.autocast(enabled=True, recipe=recipe): + expected_out = module( + expected_x, + static_split_sizes, + expected_probs, + static_split_sizes, + ) + expected_out.backward(expected_dy) + + tols = dtype_tols(dtype) + assert_close(graph_out, expected_out, **tols) + assert_close(graph_dx, expected_x.grad, **tols) + assert_close(graph_dprobs, expected_probs.grad, **tols) + if accumulate_into_main_grad: + expected_fc1_main_grad, expected_fc2_main_grad = _collect_main_grads() + assert_close(graph_fc1_main_grad, expected_fc1_main_grad, **tols) + assert_close(graph_fc2_main_grad, expected_fc2_main_grad, **tols) + else: + for graph_grad, param in zip(graph_param_grads, module.parameters()): + assert_close(graph_grad, param.grad, **tols) class TestCustomOps: @@ -3836,3 +4212,145 @@ def fuse_ops( torch.testing.assert_close(y_test, y_ref, **tols) torch.testing.assert_close(dx_test, x_ref.grad, **tols) torch.testing.assert_close(dw_test, w_ref.grad, **tols) + + +def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: + if not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Requires SM100+ for grouped GEMM quant kernel.") + + try: + from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module + except ImportError as exc: + pytest.skip(f"grouped_gemm_quant_wrapper_sm100 unavailable: {exc}") + + device = torch.device("cuda") + dtype = torch.bfloat16 if is_bf16_available() else torch.float16 + num_groups = 4 + m = 256 + n = 512 + k = 512 + total_m = num_groups * m + split_sizes = torch.full((num_groups,), m, device=device, dtype=torch.int64) + + q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + q.optimize_for_gemm = False + + torch.manual_seed(0) + a_full = torch.randn(total_m, k, device=device, dtype=dtype) + weights = [torch.randn(n, k, device=device, dtype=dtype) for _ in range(num_groups)] + + grouped_a = tex.group_quantize(a_full, q, num_groups, split_sizes) + a_groups = grouped_a.split_into_quantized_tensors() + b_groups = [q(w) for w in weights] + + # Reference GEMM on dequantized tensors. + ref = torch.empty((total_m, n), device=device, dtype=torch.float32) + start = 0 + for group_idx in range(num_groups): + end = start + m + a_deq = a_groups[group_idx].dequantize(dtype=torch.float32) + b_deq = b_groups[group_idx].dequantize(dtype=torch.float32) + ref[start:end, :] = a_deq @ b_deq.t() + start = end + ref = ref.to(dtype=torch.bfloat16).to(torch.float32) + + # Allocate empty input tensors needed for cuTE DSL kernel + padded_offsets = torch.tensor( + [m * (i + 1) for i in range(num_groups)], + dtype=torch.int32, + device=device, + ) + inputs = { + "a_tensor": torch.empty(1, total_m, k, dtype=torch.float8_e4m3fn, device=device).permute( + 1, 2, 0 + ), + "b_tensor": torch.empty(num_groups, n, k, dtype=torch.float8_e4m3fn, device=device).permute( + 1, 2, 0 + ), + "sfa_tensor": torch.empty( + 1, + total_m // 128, + k // 128, + 32, + 4, + 4, + dtype=torch.float8_e8m0fnu, + device=device, + ).permute(3, 4, 1, 5, 2, 0), + "sfb_tensor": torch.empty( + num_groups, + n // 128, + k // 128, + 32, + 4, + 4, + dtype=torch.float8_e8m0fnu, + device=device, + ).permute(3, 4, 1, 5, 2, 0), + "alpha_tensor": torch.empty(num_groups, dtype=torch.float32, device=device), + "prob_tensor": torch.empty(total_m, 1, 1, dtype=torch.float32, device=device), + "padded_offsets_tensor": padded_offsets, + } + # Overwrite inputs with quantized data/scales from MXFP8 quantizer. + a_data = grouped_a.rowwise_data.view(total_m, k).view(dtype=torch.float8_e4m3fn) + a_data = a_data.unsqueeze(0).permute(1, 2, 0).contiguous() + inputs["a_tensor"].copy_(a_data) + + a_scales = grouped_a.scale_inv.view(dtype=torch.float8_e8m0fnu) + a_scales = a_scales.view(1, total_m // 128, 4, 32, k // 128, 4) + a_scales = a_scales.permute(0, 1, 4, 3, 2, 5).contiguous() + a_scales = a_scales.permute(3, 4, 1, 5, 2, 0).contiguous() + inputs["sfa_tensor"].copy_(a_scales) + + b_data = torch.cat([w._rowwise_data.reshape(-1) for w in b_groups]) + b_data = b_data.view(dtype=torch.float8_e4m3fn) + b_data = b_data.view(num_groups, n, k).permute(1, 2, 0).contiguous() + inputs["b_tensor"].copy_(b_data) + + b_scales = torch.cat([w._rowwise_scale_inv for w in b_groups]) + b_scales = b_scales.view(dtype=torch.float8_e8m0fnu) + b_scales = b_scales.view(num_groups, n // 128, 4, 32, k // 128, 4) + b_scales = b_scales.permute(0, 1, 4, 3, 2, 5).contiguous() + b_scales = b_scales.permute(3, 4, 1, 5, 2, 0).contiguous() + inputs["sfb_tensor"].copy_(b_scales) + + inputs["alpha_tensor"].fill_(1.0) + inputs["prob_tensor"].fill_(1.0) + + cute_out = grouped_gemm_quant_wrapper_sm100( + a_tensor=inputs["a_tensor"], + b_tensor=inputs["b_tensor"], + sfa_tensor=inputs["sfa_tensor"], + sfb_tensor=inputs["sfb_tensor"], + padded_offsets=inputs["padded_offsets_tensor"], + alpha_tensor=inputs["alpha_tensor"], + 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, + discrete_col_sfd=True, + current_stream=None, + ) + + if isinstance(cute_out, dict): + outputs = cute_out + else: + d_tensor, d_col_tensor, amax_tensor, sfd_row_tensor, sfd_col_tensor = cute_out + outputs = { + "d_tensor": d_tensor, + "d_col_tensor": d_col_tensor, + "amax_tensor": amax_tensor, + "sfd_row_tensor": sfd_row_tensor, + "sfd_col_tensor": sfd_col_tensor, + } + + d_cute = outputs["d_tensor"] + if d_cute.dim() == 3: + d_cute = d_cute.squeeze(-1) + tols = dtype_tols(torch.bfloat16) + assert_close(d_cute[:total_m].float(), ref, **tols) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 225c6f6759..5bc2faa007 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -356,8 +356,9 @@ def test_quantize_varying_shapes(self, quantization: str) -> None: "shape", [[(256, 512), (512, 512), (768, 512)], [(512, 512), (512, 512), (512, 512)]], ) + @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) - def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]]) -> None: + def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]], output_dbias: bool) -> None: """Test grouped quantization for MXFP8 against per-tensor quantization.""" # Test wont pass until the grouped quantization PR from Oleg is merged. num_tensors = 2 @@ -377,12 +378,20 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]]) -> None: ) # Quantize using grouped API - grouped_output = tex.group_quantize( - grouped_input, - quantizer, - num_tensors, - first_dims, - ) + if output_dbias: + grouped_output, dbias = tex.bgrad_group_quantize( + grouped_input, + quantizer, + num_tensors, + first_dims, + ) + else: + grouped_output = tex.group_quantize( + grouped_input, + quantizer, + num_tensors, + first_dims, + ) # Build expected output by quantizing each tensor independently expected_data = [] expected_scale_inv = [] @@ -397,8 +406,13 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]]) -> None: assert torch.equal(grouped_output.rowwise_data, expected_data) assert torch.equal(grouped_output.scale_inv, expected_scale_inv) + if output_dbias: + expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) + assert torch.allclose(dbias, expected_dbias) + + @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) - def test_group_quantize_cudagraph_capturable(self) -> None: + def test_group_quantize_cudagraph_capturable(self, output_dbias: bool) -> None: """Ensure group_quantize is CUDA graph capturable.""" num_tensors = 2 shape = [(512, 1024) for _ in range(num_tensors)] @@ -418,17 +432,28 @@ def test_group_quantize_cudagraph_capturable(self) -> None: static_first_dims = first_dims.clone() # Warmup to initialize kernels and allocator state - _ = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) + if output_dbias: + _ = tex.bgrad_group_quantize(static_input, quantizer, num_tensors, static_first_dims) + else: + _ = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) torch.cuda.synchronize() graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - static_output = tex.group_quantize( - static_input, - quantizer, - num_tensors, - static_first_dims, - ) + if output_dbias: + static_output, static_dbias = tex.bgrad_group_quantize( + static_input, + quantizer, + num_tensors, + static_first_dims, + ) + else: + static_output = tex.group_quantize( + static_input, + quantizer, + num_tensors, + static_first_dims, + ) fresh_input = torch.cat( [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape], @@ -438,9 +463,21 @@ def test_group_quantize_cudagraph_capturable(self) -> None: graph.replay() torch.cuda.synchronize() - expected = tex.group_quantize(static_input, quantizer, num_tensors, static_first_dims) - assert torch.equal(static_output.rowwise_data, expected.rowwise_data) - assert torch.equal(static_output.scale_inv, expected.scale_inv) + if output_dbias: + expected_out, expected_dbias = tex.bgrad_group_quantize( + static_input, + quantizer, + num_tensors, + static_first_dims, + ) + else: + expected_out = tex.group_quantize( + static_input, quantizer, num_tensors, static_first_dims + ) + assert torch.equal(static_output.rowwise_data, expected_out.rowwise_data) + assert torch.equal(static_output.scale_inv, expected_out.scale_inv) + if output_dbias: + assert torch.allclose(static_dbias, expected_dbias) def test_clear(self) -> None: """Test clear method""" @@ -477,7 +514,7 @@ def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> in_features=in_features, out_features=out_features, params_dtype=dtype, - single_grouped_parameter=False, + single_grouped_weight=False, ).cuda() with torch.no_grad(): for i in range(num_gemms): @@ -489,6 +526,7 @@ def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> torch.randn(out_features, device="cuda", dtype=dtype) ) expected_weights = [getattr(src, f"weight{i}").detach().clone() for i in range(num_gemms)] + expected_biases = [getattr(src, f"bias{i}").detach().clone() for i in range(num_gemms)] ckpt_path = tmp_path / "grouped_linear_per_gemm.pt" torch.save(src.state_dict(), ckpt_path) del src @@ -500,7 +538,8 @@ def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> in_features=in_features, out_features=out_features, params_dtype=dtype, - single_grouped_parameter=True, + single_grouped_weight=True, + single_grouped_bias=True, ).cuda() load_result = dst.load_state_dict(src_state_dict, strict=True) assert len(load_result.missing_keys) == 0 @@ -512,6 +551,12 @@ def test_grouped_linear_load_state_dict_multi_to_single_param(self, tmp_path) -> for loaded_weight, expected_weight in zip(loaded_weights, expected_weights): assert torch.equal(loaded_weight, expected_weight) + assert getattr(dst, "bias", None) is not None + loaded_biases = dst.bias.split_into_quantized_tensors() + assert len(loaded_biases) == num_gemms + for loaded_bias, expected_bias in zip(loaded_biases, expected_biases): + assert torch.equal(loaded_bias.reshape(-1), expected_bias.reshape(-1)) + def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> None: """Load grouped-parameter checkpoint from disk into per-GEMM parameter format.""" num_gemms = 3 @@ -524,7 +569,8 @@ def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> in_features=in_features, out_features=out_features, params_dtype=dtype, - single_grouped_parameter=True, + single_grouped_weight=True, + single_grouped_bias=True, ).cuda() with torch.no_grad(): source_weights = src.weight.split_into_quantized_tensors() @@ -533,6 +579,10 @@ def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> torch.randn(out_features, in_features, device="cuda", dtype=dtype) ) expected_weights = [weight.detach().clone() for weight in source_weights] + source_biases = src.bias.split_into_quantized_tensors() + for i in range(num_gemms): + source_biases[i].copy_(torch.randn(out_features, device="cuda", dtype=dtype)) + expected_biases = [b.detach().clone() for b in source_biases] ckpt_path = tmp_path / "grouped_linear_single_param.pt" torch.save(src.state_dict(), ckpt_path) del src @@ -544,7 +594,7 @@ def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> in_features=in_features, out_features=out_features, params_dtype=dtype, - single_grouped_parameter=False, + single_grouped_weight=False, ).cuda() load_result = dst.load_state_dict(src_state_dict, strict=True) assert len(load_result.missing_keys) == 0 @@ -552,3 +602,5 @@ def test_grouped_linear_load_state_dict_single_to_multi_param(self, tmp_path) -> for i, expected_weight in enumerate(expected_weights): assert torch.equal(getattr(dst, f"weight{i}"), expected_weight) + for i, expected_bias in enumerate(expected_biases): + assert torch.equal(getattr(dst, f"bias{i}"), expected_bias.reshape(-1)) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 19b94d3531..4bfe06095b 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -2861,8 +2861,8 @@ def _make_grouped_tensor_uniform( @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: - if tex.get_cublasLt_version() < 130200: - pytest.skip("Grouped GEMM requires cuBLAS 13.2+.") + 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(): @@ -3008,6 +3008,113 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate) -> No 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, is_a=True, transposed=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], *, @@ -3050,8 +3157,8 @@ def _make_grouped_tensor_quantized_mxfp8( def test_grouped_gemm_grouped_tensor_mxfp8( shape, accumulate, layout: str, case: str, dtype: torch.dtype ) -> None: - if tex.get_cublasLt_version() < 130200: - pytest.skip("Grouped GEMM requires cuBLAS 13.2+.") + 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(): diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 384b6774f6..f87e44373e 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -155,6 +155,18 @@ def check_grouped_weight( ) +def check_grouped_bias(module: GroupedLinear, num_gemms: int, out_features: int): + """Verify GroupedLinear exposes one grouped bias parameter with shape [num_gemms, out_features].""" + bias_params = [(name, p) for name, p in module.named_parameters() if name == "bias"] + assert len(bias_params) == 1, f"Expected 1 grouped bias parameter, got {len(bias_params)}" + name, bias = bias_params[0] + assert name == "bias", f"Expected grouped parameter name 'bias', got {name}" + assert tuple(bias.shape) == (num_gemms, out_features), ( + "Grouped bias has unexpected shape. " + f"Expected {(num_gemms, out_features)}, got {tuple(bias.shape)}" + ) + + def _test_sanity_e2e_amp(block, dtype, config, fp8_recipe, skip_wgrad): te_inp_hidden_states = torch.randn( (config.max_seqlen_q, config.batch_size, config.hidden_size), @@ -523,13 +535,16 @@ def test_sanity_grouped_linear( ffn_hidden_size, bias=use_bias, params_dtype=dtype, - single_grouped_parameter=single_param, + single_grouped_weight=single_param, + single_grouped_bias=single_param, ).cuda() - # Verify grouped linear exposes a single grouped weight parameter. + # Verify grouped linear exposes a single grouped weight parameter(and bias when applicable). if fp8_recipe is None or not (fp8_recipe.delayed() or fp8_recipe.float8_current_scaling()): if single_param: check_grouped_weight(te_grouped_linear, num_gemms, ffn_hidden_size, config.hidden_size) + if use_bias: + check_grouped_bias(te_grouped_linear, num_gemms, ffn_hidden_size) inp_hidden_states = torch.randn( num_tokens, config.hidden_size, dtype=dtype, requires_grad=True diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index b9e2b907e0..7c223e6917 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -150,6 +150,7 @@ list(APPEND transformer_engine_cuda_sources normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu normalization/rmsnorm/rmsnorm_fwd_cuda_kernel.cu permutation/permutation.cu + util/utils.cu util/padding.cu swizzle/swizzle.cu swizzle/swizzle_block_scaling.cu diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 5031a30485..246fc684a1 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -32,7 +32,6 @@ inline void CreateCublasHandle(cublasLtHandle_t *handle) { // MXFP8 support for grouped GEMM requires cuBLAS 13.3+ #define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300 // BF16 support for grouped GEMM requires cuBLAS 13.3+ -// cuBLAS 13.2 is mostly functional but contains a bug for wgrad when a group has k=0, the weight gradient will be uninitialized random data instead of zeros. #define CUBLAS_GROUPED_GEMM_VERSION 130300 #if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_VERSION @@ -93,12 +92,29 @@ struct TensorShapeInfo { } }; -// Helper functions to compute average dimensions from logical_shape for heuristics -// These are hints for cuBLASLt algorithm selection, don't need to be exact +// Helper functions to compute average dimensions for cuBLASLt algorithm-selection heuristics. +// +// logical_shape encoding (from build_grouped_tensor): +// all_same: {num_tensors * M, N} +// varying_first: {sum_of_first_dims, common_last} +// varying_last: {common_first, sum_of_last_dims} +// varying_both: {1, total_elements} <-- lossy, can't recover per-dim averages +// +// We use all_same_first/last_dim() + get_common_first/last_dim() to get exact +// answers whenever possible, falling back to logical_shape division otherwise. +// For varying_both, per-dim averages are unrecoverable without a D2H copy, +// so we return 1 — a valid non-zero hint that won't skip work. inline int64_t compute_avg_first_dim(const transformer_engine::GroupedTensor *t) { - // logical_shape[0] is either num_tensors*M (uniform) or sum_of_M (varying first) - // In both cases, dividing by num_tensors gives the average - return static_cast(t->logical_shape.data[0]) / static_cast(t->num_tensors); + if (t->all_same_first_dim()) { + return static_cast(t->get_common_first_dim()); + } + const int64_t n = static_cast(t->num_tensors); + if (t->all_same_last_dim()) { + // varying_first only: logical_shape = {sum_of_first_dims, common_last} + return static_cast(t->logical_shape.data[0]) / n; + } + // varying_both: logical_shape = {1, total_elements}, no way to recover avg first dim + return 1; } inline int64_t compute_avg_last_dim(const transformer_engine::GroupedTensor *t) { @@ -228,28 +244,34 @@ inline size_t validate_grouped_gemm_inputs( dtype == transformer_engine::DType::kBFloat16 || dtype == transformer_engine::DType::kFloat16; }; - bool dtype_ok = true; for (const auto *tensor : inputs) { - dtype_ok = dtype_ok && is_supported_input_dtype(tensor->dtype()); + 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 ", + transformer_engine::to_string(tensor->dtype()), "."); + } } - NVTE_CHECK(dtype_ok, "Grouped GEMM inputs must be FP8, BF16, or FP16."); + // Cross-operand consistency across all inputs (skip tensors without data). + const transformer_engine::GroupedTensor *ref = nullptr; for (const auto *tensor : inputs) { - NVTE_CHECK(tensor->has_data() || tensor->has_columnwise_data(), - "Grouped GEMM: input tensor is missing both row-wise and column-wise data"); + if (tensor->has_data() || tensor->has_columnwise_data()) { + ref = tensor; + break; + } } - - // Cross-operand consistency across all inputs. - const auto *ref = *inputs.begin(); - const bool ref_is_fp8 = is_fp8_dtype(ref->dtype()); - const bool ref_is_mxfp8 = transformer_engine::is_mxfp_scaling(ref->scaling_mode); - for (const auto *tensor : inputs) { - 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(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) { - NVTE_CHECK(tensor->with_gemm_swizzled_scales, - "MXFP8 grouped GEMM: scales must be swizzled for GEMM."); + if (ref != nullptr) { + const bool ref_is_fp8 = is_fp8_dtype(ref->dtype()); + const bool ref_is_mxfp8 = transformer_engine::is_mxfp_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(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) { + NVTE_CHECK(tensor->with_gemm_swizzled_scales, + "MXFP8 grouped GEMM: scales must be swizzled for GEMM."); + } } } return num_tensors; @@ -554,8 +576,15 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: using namespace transformer_engine; const bool has_row = t->has_data(); const bool has_col = t->has_columnwise_data(); - NVTE_CHECK(has_row || has_col, - "Grouped GEMM operand is missing both row-wise and column-wise data"); + + if (!has_row && !has_col) { + GroupedOperandSelection sel{}; + sel.trans = trans; + sel.scaling_mode = t->scaling_mode; + sel.dtype = t->dtype(); + sel.shape = create_shape_info(t, /*swap_dims=*/false); + return sel; + } const auto sm = t->scaling_mode; const bool mxfp8 = is_mxfp_scaling(sm); @@ -758,7 +787,7 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac 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) { + cudaStream_t stream, int math_sm_count = 0) { using cublasHandleManager = transformer_engine::detail::HandleManager; cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); @@ -779,7 +808,10 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac 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))); + } cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo(handle, matmulDesc, descA, descB, descC, descD, avg_m_val, avg_n_val, avg_k_val); @@ -824,7 +856,6 @@ __global__ void grouped_bias_add_kernel(char *d_base, const char *bias_base, Ten 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; - if (m == 0 || n == 0) return; 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); @@ -1034,7 +1065,7 @@ 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.2+ + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+ check_grouped_gemm_requirements("nvte_grouped_gemm"); // Convert to internal types @@ -1082,7 +1113,7 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT 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, outputD->dtype(), num_tensors, config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val, - workspace.cublas_workspace_ptr, stream); + workspace.cublas_workspace_ptr, stream, config_.sm_count); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -1094,7 +1125,7 @@ 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.2+ + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+ check_grouped_gemm_requirements("nvte_grouped_gemm_with_discrete_inputA"); NVTE_CHECK(A_list != nullptr, "Grouped GEMM: A_list is null."); @@ -1114,6 +1145,7 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num // 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) @@ -1200,7 +1232,7 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num 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, outputD->dtype(), num_tensors, config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val, - workspace.cublas_workspace_ptr, stream); + workspace.cublas_workspace_ptr, stream, config_.sm_count); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -1213,7 +1245,7 @@ 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.2+ + // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+ check_grouped_gemm_requirements("nvte_grouped_gemm_with_discrete_out"); NVTE_CHECK(D_list != nullptr, "Grouped GEMM: D_list is null."); @@ -1272,7 +1304,7 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, 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); + workspace.cublas_workspace_ptr, stream, config_.sm_count); } void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, diff --git a/transformer_engine/common/include/transformer_engine/utils.h b/transformer_engine/common/include/transformer_engine/utils.h new file mode 100644 index 0000000000..eca6f359ea --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/utils.h @@ -0,0 +1,36 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file utils.h + * \brief Utility functions (e.g. host-to-device pointer copies). + */ + +#ifndef TRANSFORMER_ENGINE_UTILS_H_ +#define TRANSFORMER_ENGINE_UTILS_H_ + +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! \brief Copy an array of device pointers (held on host) into a device tensor. + * + * \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. + */ +void nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count, + cudaStream_t stream); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TRANSFORMER_ENGINE_UTILS_H_ diff --git a/transformer_engine/common/util/utils.cu b/transformer_engine/common/util/utils.cu new file mode 100644 index 0000000000..a183e6ec52 --- /dev/null +++ b/transformer_engine/common/util/utils.cu @@ -0,0 +1,51 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include "../common.h" +#include "../util/logging.h" + +namespace { + +constexpr int64_t kMaxKernelAddresses = 256; + +struct HostPointersArgs { + uint64_t ptrs[kMaxKernelAddresses]; +}; + +__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]; + } +} + +} // namespace + +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; + } +} diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 63a2e86e67..9d2513835c 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -42,6 +42,7 @@ #include #include #include +#include #include #include diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 1c5116a8da..e4bc744e7e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -309,6 +309,9 @@ py::object dequantize(const py::handle &input, DType otype); py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims); +py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, + const size_t num_tensors, std::optional first_dims); + std::vector multi_tensor_quantize(const std::vector &tensor_list, std::vector quantizer_list); @@ -454,6 +457,12 @@ 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); /*************************************************************************************************** @@ -561,6 +570,8 @@ void fused_multi_row_unpadding(at::Tensor input, at::Tensor output, void inplace_swizzle_scale_for_gemm(py::handle &tensor); +void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise); + /*************************************************************************************************** * NVSHMEM APIs **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index e126e0199a..f150e90507 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -233,6 +233,64 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const return py::reinterpret_borrow(grouped_output_py); } +py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, + const size_t num_tensors, std::optional first_dims) { + using namespace transformer_engine::pytorch::detail; + init_extension(); + + NVTE_CHECK(tensor.dim() == 2, "Tensor must be 2D"); + + std::vector logical_shape; + for (const auto &d : tensor.sizes()) { + logical_shape.push_back(d); + } + const auto logical_first_dim = logical_shape[0]; + const auto logical_last_dim = logical_shape[1]; + + NVTE_CHECK(logical_first_dim > 0 && logical_last_dim > 0, + "bgrad_group_quantize: empty input tensor is not supported."); + + NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), + "bgrad_group_quantize: only MXFP8 quantizer is supported."); + + auto quantizer_cpp = convert_quantizer(quantizer); + + auto grouped_input_tensor = GroupedTensorWrapper(num_tensors, logical_shape); + grouped_input_tensor.set_rowwise_data( + tensor.data_ptr(), GetTransformerEngineDType(tensor.scalar_type()), getTensorShape(tensor)); + + auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor( + num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()), + py::reinterpret_borrow(quantizer), first_dims, logical_first_dim, + logical_last_dim); + + const std::vector dbias_logical_shape = {num_tensors, logical_last_dim}; + GroupedTensorWrapper grouped_dbias(num_tensors, dbias_logical_shape, NVTE_DELAYED_TENSOR_SCALING); + at::Tensor dbias_torch = + at::empty({static_cast(num_tensors), static_cast(logical_last_dim)}, + tensor.options()); + grouped_dbias.set_rowwise_data(dbias_torch.data_ptr(), + GetTransformerEngineDType(tensor.scalar_type()), + getTensorShape(dbias_torch)); + TensorWrapper workspace_nvte; + auto stream = at::cuda::getCurrentCUDAStream(); + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), + grouped_dbias.data(), workspace_nvte.data(), stream); + }); + if (workspace_nvte.ndim() > 0 && workspace_nvte.numel() > 0) { + at::Tensor workspace_torch = allocateSpace(workspace_nvte.shape(), workspace_nvte.dtype()); + workspace_nvte = makeTransformerEngineTensor(workspace_torch.data_ptr(), workspace_nvte.shape(), + workspace_nvte.dtype()); + } + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_quantize_dbias(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), + grouped_dbias.data(), workspace_nvte.data(), stream); + }); + return py::make_tuple(py::reinterpret_borrow(grouped_output_py), + py::cast(std::move(dbias_torch))); +} + py::object dequantize(const py::handle &input, transformer_engine::DType otype) { init_extension(); diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 1431ebdfb4..08470962f9 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -9,9 +9,7 @@ #include #include -#include "../common.h" #include "../extensions.h" -#include "common.h" #include "common/util/cuda_runtime.h" #include "common/util/system.h" #include "pybind.h" @@ -637,8 +635,10 @@ py::object te_general_grouped_gemm_for_grouped_tensor( auto gemm_config = prepare_grouped_gemm_config(alpha, beta, workspace_setup, workspace_cublas, num_tensors, math_sm_count, use_split_accumulator); - [[maybe_unused]] auto swizzled_scales_A = maybe_swizzle_grouped_tensor_for_gemm(grouped_A); - [[maybe_unused]] auto swizzled_scales_B = maybe_swizzle_grouped_tensor_for_gemm(grouped_B); + [[maybe_unused]] auto swizzled_scales_A = + maybe_swizzle_grouped_tensor(grouped_A, transa, !transa); + [[maybe_unused]] auto swizzled_scales_B = + maybe_swizzle_grouped_tensor(grouped_B, transb, !transb); NVTE_SCOPED_GIL_RELEASE({ nvte_grouped_gemm(grouped_A.data(), transa, grouped_B.data(), transb, grouped_D.data(), @@ -704,7 +704,8 @@ py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py swizzled_scale_inverses_list.emplace_back( multi_tensor_swizzle_scales_for_gemm(te_A_wrappers, transa, !transa)); - [[maybe_unused]] auto swizzled_scales_B = maybe_swizzle_grouped_tensor_for_gemm(grouped_B); + [[maybe_unused]] auto swizzled_scales_B = + maybe_swizzle_grouped_tensor(grouped_B, transb, !transb); NVTE_SCOPED_GIL_RELEASE({ nvte_grouped_gemm_with_discrete_inputA( @@ -769,8 +770,10 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p te_D_vector.emplace_back(te_D_wrappers.back().data()); } - [[maybe_unused]] auto swizzled_scales_A = maybe_swizzle_grouped_tensor_for_gemm(grouped_A); - [[maybe_unused]] auto swizzled_scales_B = maybe_swizzle_grouped_tensor_for_gemm(grouped_B); + [[maybe_unused]] auto swizzled_scales_A = + maybe_swizzle_grouped_tensor(grouped_A, transa, !transa); + [[maybe_unused]] auto swizzled_scales_B = + maybe_swizzle_grouped_tensor(grouped_B, transb, !transb); NVTE_SCOPED_GIL_RELEASE({ nvte_grouped_gemm_with_discrete_out( diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index c590a3c9e2..18da5d0e9f 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -141,6 +141,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("otype")); 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("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize, + py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims")); m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize, "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer")); m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)", @@ -387,6 +389,9 @@ 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("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")); // attention kernels m.def("fa_prepare_fwd", &transformer_engine::pytorch::fa_prepare_fwd, @@ -454,6 +459,15 @@ 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("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/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index 7ff35d6b68..a6b4e7569d 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -338,8 +338,9 @@ at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapp return swizzled_scale_inv; } -std::optional maybe_swizzle_grouped_tensor_for_gemm( - GroupedTensorWrapper &input) { +std::optional maybe_swizzle_grouped_tensor(GroupedTensorWrapper &input, + bool rowwise_usage, + bool columnwise_usage) { if (input.scaling_mode() != NVTE_MXFP8_1D_SCALING) { return std::nullopt; } @@ -349,9 +350,9 @@ std::optional maybe_swizzle_grouped_tensor_for_gemm( const auto row_scales = input.get_rowwise_scale_inv(); const auto col_scales = input.get_columnwise_scale_inv(); - const bool has_rowwise_scales = !is_empty_grouped_tensor_param(row_scales); - const bool has_columnwise_scales = !is_empty_grouped_tensor_param(col_scales); - if (!has_rowwise_scales && !has_columnwise_scales) { + const bool swizzle_rowwise = rowwise_usage && !is_empty_grouped_tensor_param(row_scales); + const bool swizzle_columnwise = columnwise_usage && !is_empty_grouped_tensor_param(col_scales); + if (!swizzle_rowwise && !swizzle_columnwise) { return std::nullopt; } const auto first_dims = input.get_first_dims(); @@ -364,57 +365,84 @@ std::optional maybe_swizzle_grouped_tensor_for_gemm( std::optional rowwise_scales_pyt; std::optional columnwise_scales_pyt; - GroupedTensorWrapper output(input.num_tensors(), input.logical_shape(), input.scaling_mode()); - const auto rowwise_data = input.get_rowwise_data(); - if (rowwise_data.data_ptr != nullptr) { - output.set_rowwise_data(rowwise_data.data_ptr, static_cast(rowwise_data.dtype), - rowwise_data.shape); - } - const auto columnwise_data = input.get_columnwise_data(); - if (columnwise_data.data_ptr != nullptr) { - output.set_columnwise_data(columnwise_data.data_ptr, static_cast(columnwise_data.dtype), - columnwise_data.shape); - } + GroupedTensorWrapper swizzle_input(input.num_tensors(), input.logical_shape(), + input.scaling_mode()); + GroupedTensorWrapper swizzle_output(input.num_tensors(), input.logical_shape(), + input.scaling_mode()); + const auto tensor_offsets = input.get_tensor_offsets(); if (tensor_offsets.data_ptr != nullptr) { - output.set_tensor_offsets(tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), - tensor_offsets.shape); + swizzle_input.set_tensor_offsets( + tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), tensor_offsets.shape); + swizzle_output.set_tensor_offsets( + tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), tensor_offsets.shape); } - if (has_rowwise_scales) { + 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); - void *output_scales_dptr = getDataPtr(*rowwise_scales_pyt); - output.set_rowwise_scale_inv(output_scales_dptr, scales_dtype, row_scales.shape); + 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); } - if (has_columnwise_scales) { + if (swizzle_columnwise) { + const auto data = input.get_columnwise_data(); + const auto data_dtype = static_cast(data.dtype); 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); - void *output_scales_dptr = getDataPtr(*columnwise_scales_pyt); - output.set_columnwise_scale_inv(output_scales_dptr, scales_dtype, col_scales.shape); + 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); } - output.set_with_gemm_swizzled_scales(true); + swizzle_output.set_with_gemm_swizzled_scales(true); NVTE_SCOPED_GIL_RELEASE({ - nvte_swizzle_grouped_scaling_factors(input.data(), output.data(), + nvte_swizzle_grouped_scaling_factors(swizzle_input.data(), swizzle_output.data(), at::cuda::getCurrentCUDAStream()); }); - if (has_rowwise_scales) { + 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); } - if (has_columnwise_scales) { + 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); } input.set_with_gemm_swizzled_scales(true); - return SwizzledGroupedScales{std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)}; } +void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise) { + using namespace transformer_engine::pytorch::detail; + + auto tensor_nvte = GroupedTensorFromPyTorchGroupedTensor(tensor); + + auto result = maybe_swizzle_grouped_tensor(tensor_nvte, rowwise, columnwise); + + if (result.has_value()) { + if (result->first.has_value()) { + tensor.attr("scale_inv") = py::cast(*result->first); + } else { + tensor.attr("scale_inv") = py::none(); + } + if (result->second.has_value()) { + tensor.attr("columnwise_scale_inv") = py::cast(*result->second); + } else { + tensor.attr("columnwise_scale_inv") = py::none(); + } + tensor.attr("_with_gemm_swizzled_scales") = py::cast(true); + } +} + 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/extensions/utils.cpp b/transformer_engine/pytorch/csrc/extensions/utils.cpp new file mode 100644 index 0000000000..9a093608d4 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/utils.cpp @@ -0,0 +1,165 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include + +#include "common/common.h" +#include "extensions.h" + +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; +} + +} // namespace + +std::vector convert_host_pointers_to_tensor( + std::vector> tensor_lists) { + std::vector outputs; + outputs.reserve(tensor_lists.size()); + auto stream = 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; +} + +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(); + + // 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()); + } + + // Swizzle scales and collect scale pointers + at::Tensor swizzled_scales_keepalive; + std::vector scale_host_ptrs(num_tensors); + + if (swizzle) { + NVTEScalingMode scaling_mode; + transformer_engine::DType scale_dtype; + if (is_fp8_dtype(data_dtype)) { + scaling_mode = NVTE_MXFP8_1D_SCALING; + scale_dtype = transformer_engine::DType::kFloat8E8M0; + } else if (is_fp4_dtype(data_dtype)) { + 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); + } + + // 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)); + + // 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); + } + } + + // 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); + + // Collect swizzled scale pointers + for (size_t i = 0; i < num_tensors; ++i) { + scale_host_ptrs[i] = reinterpret_cast(output_dptr + output_offsets[i]); + } + } 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); + + return {std::move(data_ptrs), std::move(scale_ptrs), std::move(swizzled_scales_keepalive)}; +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index e9c6ca882e..e13554a98c 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -221,6 +221,8 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + } else if (quantizer_dtype != DType::kNumTypes) { + ret.set_rowwise_data(nullptr, quantizer_dtype, std::vector{0}); } // Columnwise data @@ -229,6 +231,8 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) { DType data_dtype = quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype; ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data)); + } else if (quantizer_dtype != DType::kNumTypes) { + ret.set_columnwise_data(nullptr, quantizer_dtype, std::vector{0}); } // Scale diff --git a/transformer_engine/pytorch/csrc/util.h b/transformer_engine/pytorch/csrc/util.h index 587ec289a4..88f76a7cb1 100644 --- a/transformer_engine/pytorch/csrc/util.h +++ b/transformer_engine/pytorch/csrc/util.h @@ -38,10 +38,15 @@ using SwizzledGroupedScales = std::pair, std::optional /*! \brief Swizzle grouped tensor scales for GEMM if needed. * Currently only works for MXFP8 1D scaling with uniform shapes. * + * \param[in,out] input Grouped tensor whose scales to swizzle. + * \param[in] rowwise_usage Whether rowwise scales are needed. + * \param[in] columnwise_usage Whether columnwise scales are needed. + * * The returned swizzled scales should be kept alive during the GEMM. */ -std::optional maybe_swizzle_grouped_tensor_for_gemm( - GroupedTensorWrapper& input); +std::optional maybe_swizzle_grouped_tensor(GroupedTensorWrapper& input, + bool rowwise_usage, + bool columnwise_usage); /*! \brief Convert a block scaling tensor to an mxfp8 tensor in-place. * diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 28da4873f0..a96a87bf89 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -80,19 +80,19 @@ class UserBufferQuantizationMode(Enum): def get_dummy_wgrad(shape: list, dtype: torch.dtype, zero=False) -> torch.Tensor: """Returns a dummy tensor of given shape.""" - if len(shape) != 2: - raise ValueError(f"Expected 2D shape, got {len(shape)}D: {shape}") + + key = (*shape, dtype) global _dummy_wgrads - if (shape[0], shape[1], dtype) not in _dummy_wgrads: - _dummy_wgrads[(shape[0], shape[1], dtype)] = torch.empty( + if key not in _dummy_wgrads: + _dummy_wgrads[key] = torch.empty( shape, dtype=dtype, device="cuda", requires_grad=False, ) if zero: - _dummy_wgrads[(shape[0], shape[1], dtype)].fill_(0) - return _dummy_wgrads[(shape[0], shape[1], dtype)].detach() + _dummy_wgrads[key].fill_(0) + return _dummy_wgrads[key].detach() def initialize_ub( diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 0adda48e36..ba6becb9f9 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -594,10 +594,14 @@ class GroupedLinear(TransformerEngineBaseModule): cast tensor. In some scenarios, the input tensor is used by multiple modules, and saving the original input tensor may reduce the memory usage. Cannot work with FP8 DelayedScaling recipe. - single_grouped_parameter : bool, default = False + 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. + 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. Notes ----- @@ -628,7 +632,8 @@ def __init__( ub_name: Optional[str] = None, delay_wgrad_compute: bool = False, save_original_input: bool = False, - single_grouped_parameter: bool = False, + single_grouped_weight: bool = False, + single_grouped_bias: bool = False, name: Optional[str] = None, ) -> None: super().__init__(name) @@ -645,7 +650,8 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input - self.single_grouped_parameter = single_grouped_parameter + self.single_grouped_weight = single_grouped_weight + self.single_grouped_bias = single_grouped_bias if ub_overlap_rs or ub_overlap_ag: raise ValueError("GroupedLinear doesn't support Userbuffer overlap.") self.init_method = init_method @@ -737,6 +743,9 @@ def __init__( if self.wgrad_store.delay_wgrad_compute(): for name, param in self.named_parameters(): + if name in ("weight", "bias"): + param.skip_backward_post_hook = True + continue for i in range(self.num_gemms): if name in (f"weight{i}", f"bias{i}"): param.skip_backward_post_hook = True @@ -787,13 +796,12 @@ def make_grouped_weights(self, defer_init=False) -> None: else: grouped_weights.quantized_tensors[i].copy_(weights[i]) - # Re-register as a single grouped weight parameter. # Re-register as a single grouped weight parameter. if not ( isinstance(grouped_weights, torch.Tensor) and (weight_quantizers[0] is None or not weight_quantizers[0].internal) ): - raise RuntimeError("Found internal quantizer with `single_grouped_parameter=True`.") + raise RuntimeError("Found internal quantizer with `single_grouped_weight=True`.") self.register_parameter( "weight", torch.nn.Parameter(grouped_weights), @@ -804,13 +812,33 @@ def make_grouped_weights(self, defer_init=False) -> None: for i in range(self.num_gemms): self.register_parameter(f"weight{i}", None) + if self.use_bias and self.single_grouped_bias: + self._make_grouped_biases() + self.set_tensor_parallel_attributes(defer_init=defer_init) + def _make_grouped_biases(self) -> None: + """Pack per-GEMM biases into one ``GroupedTensor`` (``single_grouped_bias``).""" + biases = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + packed = torch.stack([b.detach().clone() for b in biases], dim=0).contiguous() + grouped_bias = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=self.num_gemms, + tensor_shape=(self.out_features,), + rowwise_data=packed, + dtype=packed.dtype, + ) + grouped_bias.requires_grad_(True) + self.register_parameter("bias", torch.nn.Parameter(grouped_bias)) + for i in range(self.num_gemms): + self.register_parameter(f"bias{i}", None) + def reset_parameters(self, defer_init=False): super().reset_parameters(defer_init=defer_init) - # Grouped tensor weights is an opt-in feature. - if self.single_grouped_parameter: + # Grouped tensor weights / biases are opt-in features. + if self.single_grouped_weight: self.make_grouped_weights(defer_init=defer_init) + elif self.single_grouped_bias: + self._make_grouped_biases() def set_tensor_parallel_attributes(self, defer_init=False) -> None: """Set attributes needed for TP""" @@ -836,15 +864,24 @@ def set_tensor_parallel_attributes(self, defer_init=False) -> None: # Set parallelism attributes for linear biases if self.use_bias: - for i in range(self.num_gemms): + grouped_bias = getattr(self, "bias", None) + if grouped_bias is not None: if self.parallel_mode == "row": - setattr( - getattr(self, f"bias{i}"), - "sequence_parallel", - self.sequence_parallel, - ) + setattr(grouped_bias, "sequence_parallel", self.sequence_parallel) elif self.parallel_mode == "column": - set_tensor_model_parallel_attributes(getattr(self, f"bias{i}"), True, 0, 1) + set_tensor_model_parallel_attributes(grouped_bias, True, 0, 1) + else: + for i in range(self.num_gemms): + if self.parallel_mode == "row": + setattr( + getattr(self, f"bias{i}"), + "sequence_parallel", + self.sequence_parallel, + ) + elif self.parallel_mode == "column": + set_tensor_model_parallel_attributes( + getattr(self, f"bias{i}"), True, 0, 1 + ) def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None: """Remap weight keys between single and per-GEMM checkpoint formats.""" @@ -853,8 +890,8 @@ def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None has_grouped_weight = grouped_weight_key in state_dict has_per_gemm_weights = all(key in state_dict for key in per_gemm_weight_keys) - if self.single_grouped_parameter: - # Backward compatibility: checkpoints saved without single_grouped_parameter + if self.single_grouped_weight: + # Backward compatibility: checkpoints saved without single_grouped_weight # store one weight tensor per GEMM (weight0..weightN). Convert them into a # single stacked grouped weight expected by this module configuration. if not has_grouped_weight and has_per_gemm_weights: @@ -869,7 +906,7 @@ def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None for key in per_gemm_weight_keys: state_dict.pop(key, None) else: - # Forward compatibility: checkpoints saved with single_grouped_parameter + # Forward compatibility: checkpoints saved with single_grouped_weight # store one grouped `weight`. Convert it back to weight0..weightN. if not has_per_gemm_weights and has_grouped_weight: grouped_weight = state_dict.pop(grouped_weight_key) @@ -898,6 +935,40 @@ def _remap_grouped_weight_state_dict_keys(self, state_dict, prefix: str) -> None # Drop any redundant grouped key to avoid strict-load unexpected-key errors. state_dict.pop(grouped_weight_key, None) + def _remap_grouped_bias_state_dict_keys(self, state_dict, prefix: str) -> None: + """Remap bias keys between single grouped and per-GEMM checkpoint formats.""" + if not self.use_bias: + return + grouped_bias_key = f"{prefix}bias" + per_gemm_bias_keys = [f"{prefix}bias{i}" for i in range(self.num_gemms)] + has_grouped_bias = grouped_bias_key in state_dict + has_per_gemm_biases = all(key in state_dict for key in per_gemm_bias_keys) + + if self.single_grouped_bias: + if not has_grouped_bias and has_per_gemm_biases: + per_gemm = [state_dict.pop(key) for key in per_gemm_bias_keys] + state_dict[grouped_bias_key] = torch.stack(per_gemm, dim=0) + elif has_grouped_bias: + for key in per_gemm_bias_keys: + state_dict.pop(key, None) + val = state_dict[grouped_bias_key] + if isinstance(val, torch.Tensor) and val.dim() == 3 and val.shape[1] == 1: + state_dict[grouped_bias_key] = val.squeeze(1) + else: + if not has_per_gemm_biases and has_grouped_bias: + gb = state_dict.pop(grouped_bias_key) + if hasattr(gb, "split_into_quantized_tensors"): + members = gb.quantized_tensors + if members is None: + members = gb.split_into_quantized_tensors() + per_gemm = [m.reshape(-1) if m.dim() > 1 else m for m in members] + else: + per_gemm = list(gb.unbind(0)) + for i, b in enumerate(per_gemm): + state_dict[f"{prefix}bias{i}"] = b.reshape(-1) if b.dim() > 1 else b + elif has_per_gemm_biases: + state_dict.pop(grouped_bias_key, None) + def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False): """Load state dict with grouped-weight format compatibility.""" state_dict_copy = state_dict.copy() @@ -905,6 +976,7 @@ def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False) if metadata is not None: state_dict_copy._metadata = metadata self._remap_grouped_weight_state_dict_keys(state_dict_copy, prefix="") + self._remap_grouped_bias_state_dict_keys(state_dict_copy, prefix="") return super().load_state_dict(state_dict_copy, strict=strict, assign=assign) def _load_from_state_dict( @@ -912,6 +984,7 @@ def _load_from_state_dict( ): """Load state, including compatibility across grouped-weight checkpoint formats.""" self._remap_grouped_weight_state_dict_keys(state_dict, prefix) + self._remap_grouped_bias_state_dict_keys(state_dict, prefix) super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs @@ -962,7 +1035,7 @@ def forward( inp = self.prepare_forward(inp, num_gemms=self.num_gemms) try: weight_tensors = self._get_weight_tensors() - bias_tensors = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + bias_tensors = self._get_bias_tensors() quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers() @@ -1026,18 +1099,28 @@ def backward_dw(self): """ if not self.need_backward_dw(): return + if self.wgrad_store.context is None or self.wgrad_store.context.empty(): + return with get_nvtx_range_context("_GroupedLinear_wgrad"): (_, grad_biases_, _), tensor_list = self.wgrad_store.pop() wgrad_list = tensor_list[2] weight_params = self._get_weight_tensors() - bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] 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: - for i in range(self.num_gemms): - if bias_params[i].grad is None: - bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) + grouped_bias = getattr(self, "bias", None) + if grouped_bias is not None: + gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype) + if grouped_bias.grad is None: + grouped_bias.grad = gstack + else: + grouped_bias.grad.add_(gstack) + 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: + bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) del grad_biases_ del wgrad_list del tensor_list @@ -1099,6 +1182,16 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage ] return weight_tensors + def _get_bias_tensors(self) -> List[torch.Tensor]: + """Per-GEMM bias tensors (views into grouped storage when ``single_grouped_bias``).""" + grouped_bias = getattr(self, "bias", None) + if grouped_bias is not None: + parts = grouped_bias.quantized_tensors + if parts is None: + parts = grouped_bias.split_into_quantized_tensors() + return [p.reshape(-1) for p in parts] + return [getattr(self, f"bias{i}") for i in range(self.num_gemms)] + def _get_weight_quantizers(self) -> List[Quantizer]: """Get the weight quantizers of the module.""" if not self.fp8 and not self.fp8_calibration and not self.primary_weights_in_fp8: diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 4520dbc313..0e03e691f3 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -71,3 +71,117 @@ def get_fp8_meta_from_fp8_tensor(tensor: Float8Tensor) -> tuple[FP8TensorMeta, i fp8_meta.amax_history = torch.empty(1, 1, dtype=torch.float32, device=tensor.device) fp8_meta.scale_inv = tensor._scale_inv return fp8_meta, 0 + + +def validate_grouped_mlp_dims(fc1, swiglu, fc2) -> None: + """Validate FC1/SwiGLU/FC2 dimensions and interleave size for fused grouped MLP.""" + + if fc1.in_features % 256 != 0 or fc1.out_features % 256 != 0: + raise ValueError( + f"Unsupported dims for FC1 (num_groups={fc1.num_groups}, " + f"in_features={fc1.in_features}, out_features={fc1.out_features})." + ) + if fc2.in_features % 256 != 0 or fc2.out_features % 256 != 0: + raise ValueError( + 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: + 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 swiglu.glu_interleave_size != 32: + raise ValueError( + "Fused kernel requires 32-wide GLU interleaving, " + f"but got glu_interleave_size={swiglu.glu_interleave_size}." + ) + + +def fuse_grouped_mlp_ops( + ops, + *, + recipe, + fused_op_cls, +): + """Sliding-window fusion for GroupedLinear + ScaledSwiGLU + GroupedLinear. + + Parameters + ---------- + ops : list of FusibleOperation + Operations to scan. + recipe : Recipe or None + Quantization recipe. + fused_op_cls : type + Fused operation class with ``is_supported()`` classmethod and + constructor accepting ``fc1``, ``swiglu``, ``fc2`` keyword args. + May also expose ``is_fc1_bias_supported()`` and/or + ``is_fc2_bias_supported()`` classmethods for bias eligibility. + + Returns + ------- + list of FusibleOperation + Updated operations with matched triples replaced by fused ops. + """ + from .basic import GroupedLinear, ScaledSwiGLU # pylint: disable=import-outside-toplevel + + if not fused_op_cls.is_supported(): + return 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: + + matches_pattern = True + if not ( + isinstance(window[0], GroupedLinear) + and isinstance(window[1], ScaledSwiGLU) + and isinstance(window[2], GroupedLinear) + ): + matches_pattern = False + elif window[0].num_groups != window[2].num_groups: + matches_pattern = False + elif ( + window[0].in_features % 256 != 0 + or window[0].out_features % 256 != 0 + or window[2].in_features % 256 != 0 + or window[2].out_features % 256 != 0 + ): + 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( + fc1=window[0], + swiglu=window[1], + fc2=window[2], + ) + window = [op] + else: + out.extend(window[:-2]) + window = window[-2:] + + out.extend(window[:-3]) + window = window[-3:] + while ops and len(window) < 3: + window.append(ops[0]) + ops = ops[1:] + + out.extend(window) + return out diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b44e77b0c6..f26a337a4d 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -7,6 +7,7 @@ from __future__ import annotations from collections.abc import Callable, Iterable, Sequence import contextlib +import functools import math from typing import Any, Optional @@ -15,6 +16,7 @@ import transformer_engine_torch as tex from ...cpp_extensions import general_grouped_gemm from ...distributed import CudaRNGStatesTracker +from ...module._common import WeightGradStore from ...module.base import ( _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -32,6 +34,7 @@ ) from .._common import is_quantized_tensor, maybe_dequantize from ..op import BasicOperation, OperationContext +from ...tensor import GroupedTensor class GroupedLinear(BasicOperation): @@ -69,6 +72,13 @@ class GroupedLinear(BasicOperation): Megatron-LM. This argument along with weight tensor having attribute ``overwrite_main_grad`` set to True will overwrite ``main_grad`` instead of accumulating. + single_grouped_weight : bool, default = ``False`` + Store all expert weights as one ``GroupedTensor`` parameter ``weight``. + 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}``. """ @@ -86,13 +96,21 @@ def __init__( dtype: Optional[torch.dtype] = None, rng_state_tracker_function: Optional[Callable[[], CudaRNGStatesTracker]] = None, accumulate_into_main_grad: bool = False, + single_grouped_weight: bool = False, + single_grouped_bias: bool = False, + delay_wgrad_compute: bool = False, ) -> None: super().__init__() + self.wgrad_store = WeightGradStore(delay_wgrad_compute) + # Weight tensor dimensions self.num_groups: int = num_groups self.in_features: int = in_features self.out_features: int = out_features + self.single_grouped_weight: bool = single_grouped_weight + self.single_grouped_bias: bool = single_grouped_bias + self.use_bias: bool = bias if self.num_groups <= 0: raise ValueError(f"Invalid number of groups ({self.num_groups})") if self.in_features <= 0: @@ -116,12 +134,15 @@ def __init__( self._rng_state_tracker_function = rng_state_tracker_function # Register weights + # TODO(ksivaman): Proper support for meta device. + # We do not want to reset params later as it wipes off + # main_grad and related attributes. self.weight0: torch.nn.Parameter for group_idx in range(self.num_groups): weight_tensor = torch.empty( self.out_features, self.in_features, - device="meta", + device=device, dtype=dtype, ) self.register_parameter( @@ -136,7 +157,7 @@ def __init__( if bias: bias_tensor = torch.empty( self.out_features, - device="meta", + device=device, dtype=dtype, ) bias_tensor = torch.nn.Parameter(bias_tensor) @@ -149,6 +170,57 @@ def __init__( # Whether to accumulate weight gradient into main_grad self._accumulate_into_main_grad: bool = accumulate_into_main_grad + self._apply_delay_wgrad_param_hooks() + + def _apply_delay_wgrad_param_hooks(self) -> None: + """Set ``skip_backward_post_hook`` on weights when delaying wgrad (bias uses main backward).""" + if not self.wgrad_store.delay_wgrad_compute(): + return + if self.single_grouped_weight: + self.weight.skip_backward_post_hook = True + else: + for group_idx in range(self.num_groups): + getattr(self, f"weight{group_idx}").skip_backward_post_hook = True + + def need_backward_dw(self) -> bool: + """Return whether :meth:`backward_dw` must run to finish weight gradients.""" + return self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute() + + def backward_dw(self) -> None: + """Execute delayed weight gradient grouped GEMMs (see ``delay_wgrad_compute``).""" + if not self.need_backward_dw(): + return + if self.wgrad_store.context is None or self.wgrad_store.context.empty(): + return + _, tensor_list = self.wgrad_store.pop() + activations = tensor_list[0] + grad_weights = tensor_list[2] + if isinstance(activations, list): + clear_tensor_data(*activations) + else: + # Fused MXFP8 grouped MLP saves `GroupedTensor` activations for wgrad. + clear_tensor_data( + activations.data, + activations.columnwise_data, + activations.scale_inv, + activations.columnwise_scale_inv, + ) + if self._accumulate_into_main_grad: + return + if self.single_grouped_weight: + if isinstance(grad_weights, list): + self.weight.grad = torch.stack(grad_weights, dim=0).to(self.weight.dtype) + else: + self.weight.grad = grad_weights.rowwise_data.view( + self.num_groups, + self.out_features, + self.in_features, + ).to(self.weight.dtype) + else: + for group_idx in range(self.num_groups): + w = getattr(self, f"weight{group_idx}") + w.grad = grad_weights[group_idx].to(w.dtype) + def num_quantizers(self, mode: str) -> int: if mode == "forward": return 2 * self.num_groups @@ -159,7 +231,7 @@ def num_quantizers(self, mode: str) -> int: @property def has_bias(self) -> bool: """Whether an additive bias is being applied""" - return self.bias0 is not None + return self.use_bias def reset_parameters(self) -> None: """Initialize parameter buffers and values""" @@ -221,16 +293,92 @@ def reset_parameters(self) -> None: setattr(self, f"weight{group_idx}", weight) # Initialize biases if needed - if self.bias0 is not None: + packed_biases: Optional[torch.Tensor] = None + if self.use_bias: + if self.bias0 is not None: + bias_dtype = self.bias0.dtype + elif getattr(self, "bias", None) is not None: + bias_dtype = self.bias.dtype + elif getattr(self, "weight", None) is not None: + bias_dtype = self.weight.dtype + else: + bias_dtype = self.weight0.dtype packed_biases = torch.zeros( self.num_groups, self.out_features, - dtype=self.bias0.dtype, + dtype=bias_dtype, device=device, ) + if not self.single_grouped_bias: + for group_idx in range(self.num_groups): + bias = torch.nn.Parameter(packed_biases[group_idx]) + setattr(self, f"bias{group_idx}", bias) + else: for group_idx in range(self.num_groups): - bias = torch.nn.Parameter(packed_biases[group_idx]) - setattr(self, f"bias{group_idx}", bias) + self.register_parameter(f"bias{group_idx}", None) + + if self.single_grouped_weight: + self.make_grouped_weights() + if self.use_bias and self.single_grouped_bias: + assert packed_biases is not None + self._make_grouped_biases_from_packed(packed_biases) + self._apply_delay_wgrad_param_hooks() + + def make_grouped_weights(self) -> None: + """ + Convert parameters into a GroupedTensor and re-register them as parameters. + """ + + weights = [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + quantizer = self.get_quantizer("forward", 1) + + recipe = None if quantizer is None else quantizer._get_compatible_recipe() + if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()): + raise RuntimeError( + "Delayed scaling or float8 current scaling is not supported with" + " single_grouped_weight=True" + ) + + grouped_weights = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=self.num_groups, + shapes=[(self.out_features, self.in_features)] * self.num_groups, + quantizer=quantizer, + dtype=self.weight0.dtype, + device=self.weight0.device, + ) + + # Copy existing params into storage. + with torch.no_grad(): + for i in range(self.num_groups): + if self._with_quantized_weight: + grouped_weights.quantized_tensors[i].copy_from_storage(weights[i]) + else: + grouped_weights.quantized_tensors[i].copy_(weights[i]) + + assert isinstance(grouped_weights, torch.Tensor) and ( + quantizer is None or not quantizer.internal + ), "Found internal quantizer with `single_grouped_weight=True`." + + # Re-register as a single grouped weight parameter. + self.register_parameter("weight", torch.nn.Parameter(grouped_weights)) + for group_idx in range(self.num_groups): + self.register_parameter(f"weight{group_idx}", None) + + self._apply_delay_wgrad_param_hooks() + + def _make_grouped_biases_from_packed(self, packed_biases: torch.Tensor) -> None: + """Replace per-group bias parameters with one ``GroupedTensor`` (``single_grouped_bias``).""" + bias_data = packed_biases.detach().clone().contiguous() + grouped_bias = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=self.num_groups, + tensor_shape=(self.out_features,), + rowwise_data=bias_data, + dtype=bias_data.dtype, + ) + grouped_bias.requires_grad_(True) + self.register_parameter("bias", torch.nn.Parameter(grouped_bias)) + for group_idx in range(self.num_groups): + self.register_parameter(f"bias{group_idx}", None) def _quantize_weights( self, @@ -328,63 +476,102 @@ def pre_first_fuser_forward(self) -> None: if any(param.device.type == "meta" for param in self.parameters()): self.reset_parameters() - # Check that weights are consistent - dtype = self.weight0.dtype - device = self.weight0.device - weight_requires_grad = self.weight0.requires_grad - weight_tensor_type = type(self.weight0.data) - for group_idx in range(self.num_groups): - weight = getattr(self, f"weight{group_idx}") - if weight.dtype != dtype: - raise RuntimeError( - f"Weight {group_idx} has invalid dtype (expected {dtype}, got {weight.dtype})." - ) - if not devices_match(weight.device, device): - raise RuntimeError( - f"Weight {group_idx} has invalid device " - f"(expected {device}, got {weight.device})." - ) - if weight.requires_grad != weight_requires_grad: - raise RuntimeError( - f"Weight {group_idx} has requires_grad={weight.requires_grad}, " - f"but expected requires_grad={weight_requires_grad}." - ) - if type(weight.data) != weight_tensor_type: # pylint: disable=unidiomatic-typecheck - raise RuntimeError( - f"Weight {group_idx} has invalid tensor type " - f"(expected {weight_tensor_type.__name__}, " - f"got {type(weight.data).__name__})." - ) + # Check that all weight params are consistent + if not self.single_grouped_weight: + dtype = self.weight0.dtype + device = self.weight0.device + weight_requires_grad = self.weight0.requires_grad + weight_tensor_type = type(self.weight0.data) + for group_idx in range(self.num_groups): + weight = getattr(self, f"weight{group_idx}") + if weight.dtype != dtype: + raise RuntimeError( + f"Weight {group_idx} has invalid dtype (expected {dtype}, got" + f" {weight.dtype})." + ) + if not devices_match(weight.device, device): + raise RuntimeError( + f"Weight {group_idx} has invalid device " + f"(expected {device}, got {weight.device})." + ) + if weight.requires_grad != weight_requires_grad: + raise RuntimeError( + f"Weight {group_idx} has requires_grad={weight.requires_grad}, " + f"but expected requires_grad={weight_requires_grad}." + ) + if type(weight.data) != weight_tensor_type: # pylint: disable=unidiomatic-typecheck + raise RuntimeError( + f"Weight {group_idx} has invalid tensor type " + f"(expected {weight_tensor_type.__name__}, " + f"got {type(weight.data).__name__})." + ) + else: + dtype = self.weight.dtype + device = self.weight.device + weight_requires_grad = self.weight.requires_grad + weight_tensor_type = type(self.weight.data) # Check that biases are consistent - for group_idx in range(self.num_groups): - bias = getattr(self, f"bias{group_idx}") - if self.has_bias: - if bias is None: - raise RuntimeError(f"Expected biases, but bias {group_idx} is uninitialized") + if self.has_bias: + if self.single_grouped_bias: + bias = self.bias if bias.dtype != dtype: raise RuntimeError( - f"Bias {group_idx} has invalid dtype (expected {dtype}, got {bias.dtype})." + f"Bias has invalid dtype (expected {dtype}, got {bias.dtype})." ) if not devices_match(bias.device, device): raise RuntimeError( - f"Bias {group_idx} has invalid device " - f"(expected {device}, got {bias.device})." + f"Bias has invalid device (expected {device}, got {bias.device})." ) if bias.requires_grad != weight_requires_grad: raise RuntimeError( - f"Bias {group_idx} has requires_grad={bias.requires_grad}, " + f"Bias has requires_grad={bias.requires_grad}, " f"but expected requires_grad={weight_requires_grad}." ) else: - if bias is not None: - raise RuntimeError(f"Expected no biases, but bias {group_idx} is initialized") + for group_idx in range(self.num_groups): + bias = getattr(self, f"bias{group_idx}") + if bias is None: + raise RuntimeError( + f"Expected biases, but bias {group_idx} is uninitialized" + ) + if bias.dtype != dtype: + raise RuntimeError( + f"Bias {group_idx} has invalid dtype (expected {dtype}, got" + f" {bias.dtype})." + ) + if not devices_match(bias.device, device): + raise RuntimeError( + f"Bias {group_idx} has invalid device " + f"(expected {device}, got {bias.device})." + ) + if bias.requires_grad != weight_requires_grad: + raise RuntimeError( + f"Bias {group_idx} has requires_grad={bias.requires_grad}, " + f"but expected requires_grad={weight_requires_grad}." + ) + else: + if self.single_grouped_bias: + if getattr(self, "bias", None) is not None: + raise RuntimeError("Expected no biases, but grouped `bias` is registered") + else: + for group_idx in range(self.num_groups): + bias = getattr(self, f"bias{group_idx}") + if bias is not None: + raise RuntimeError( + f"Expected no biases, but bias {group_idx} is initialized" + ) def pre_fuser_forward(self, *, requires_grad: bool) -> None: super().pre_fuser_forward(requires_grad=requires_grad) if FP8GlobalStateManager.is_fp8_enabled(): # Assume weights have consistent grad requirement - weight_requires_grad = requires_grad and self.weight0.requires_grad + weight_requires_grad = ( + self.weight.requires_grad + if self.single_grouped_weight + else self.weight0.requires_grad + ) + weight_requires_grad = requires_grad and weight_requires_grad # Configure quantizer usages # Note: We cache the quantized input for backward pass, @@ -419,13 +606,17 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: # Make sure weight param has correct quantizer weight_quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) weight_quantizer.internal = False - getattr(self, f"weight{group_idx}").update_quantizer(weight_quantizer.copy()) + if self.single_grouped_weight: + self.weight.quantizer = weight_quantizer.copy() + else: + getattr(self, f"weight{group_idx}").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) + and not self.single_grouped_weight ) # Recipe-specific configuration @@ -472,12 +663,19 @@ def fuser_forward( ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: num_groups = self.num_groups has_bias = self.has_bias - device = self.weight0.device + weight_param = self.weight if self.single_grouped_weight else self.weight0 + device = weight_param.device + + if self._accumulate_into_main_grad: + if not hasattr(weight_param, "main_grad"): + raise RuntimeError("MAIN GRAD NOT FOUND") + if weight_param.main_grad is None: + raise RuntimeError("MAIN GRAD IS NONE") # Check which grads are required ctx = basic_op_ctxs[0] input_requires_grad = ctx.requires_grad - weight_requires_grad = ctx.requires_grad and self.weight0.requires_grad + weight_requires_grad = ctx.requires_grad and weight_param.requires_grad # Quantizers input_quantizers = [None] * num_groups @@ -494,7 +692,7 @@ def fuser_forward( if torch.is_autocast_enabled(): dtype = torch.get_autocast_dtype("cuda") else: - dtype = self.weight0.dtype + dtype = weight_param.dtype # Extract split sizes from extra input split_sizes = basic_op_extra_inputs[0][0] @@ -503,10 +701,24 @@ def fuser_forward( raise ValueError(f"Expected {num_groups} splits, but got {len(split_sizes_int)}.") # Extract params - weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] + if self.single_grouped_weight: + weights = self.weight.quantized_tensors + if weights is None: + weights = self.weight.split_into_quantized_tensors() + else: + weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] bs = None if has_bias: - bs = [maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups)] + 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 = [] @@ -589,7 +801,8 @@ def fuser_backward( ]: num_groups = self.num_groups has_bias = self.has_bias - device = self.weight0.device + weight_param = self.weight if self.single_grouped_weight else self.weight0 + device = weight_param.device # Saved tensors from forward pass ctx = basic_op_ctxs[0] @@ -628,14 +841,42 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Get grad tensors from params so we can # accumulate directly into it. - for group_idx in range(num_groups): - weight_param = getattr(self, f"weight{group_idx}") + if self.single_grouped_weight: 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) + 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 + ) + 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 + ) else: - weight_shape = ws[0].size() + weight_shape = (self.out_features, self.in_features) for group_idx in range(num_groups): grad_weights[group_idx] = torch.empty( weight_shape, @@ -668,26 +909,63 @@ def fuser_backward( ) # Perform wgrad GEMMs + 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: - general_grouped_gemm( - xs, - dys, - grad_weights, - [None] * num_groups, # quantization_params - ctx.dtype, - layout="NT", - m_splits=split_sizes_int, - use_split_accumulator=_2X_ACC_WGRAD, - accumulate=accumulate_into_main_grad, - ) + if delay_wgrad: + grouped_gemm_wgrad = functools.partial( + general_grouped_gemm, + quantization_params=[None] * num_groups, + out_dtype=ctx.dtype, + layout="NT", + m_splits=split_sizes_int, + use_split_accumulator=_2X_ACC_WGRAD, + accumulate=accumulate_into_main_grad, + ) + self.wgrad_store.put([xs, dys, grad_weights], grouped_gemm_wgrad) + else: + general_grouped_gemm( + xs, + dys, + grad_weights, + [None] * num_groups, # quantization_params + ctx.dtype, + layout="NT", + m_splits=split_sizes_int, + use_split_accumulator=_2X_ACC_WGRAD, + accumulate=accumulate_into_main_grad, + ) - # Clear input tensors if possible - clear_tensor_data(*xs) + 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] + return grad_input, [grad_params], [(None,)] for group_idx in range(num_groups): weight_param = getattr(self, f"weight{group_idx}") if hasattr(weight_param, "grad_added_to_main_grad"): @@ -698,5 +976,29 @@ def fuser_backward( zero=getattr(weight_param, "zero_out_wgrad", False), ) - grad_params = grad_weights + grad_biases if has_bias else grad_weights + 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: + final_weight_grads = [None] * num_groups + else: + final_weight_grads = grad_weights + + if not has_bias: + grad_params = list(final_weight_grads) + elif self.single_grouped_bias: + final_bias_grads = torch.stack(grad_biases, dim=0).to(ctx.dtype) + grad_params = list(final_weight_grads) + [final_bias_grads] + else: + if self.single_grouped_weight: + grad_params = list(grad_biases) + list(final_weight_grads) + else: + grad_params = list(final_weight_grads) + list(grad_biases) + return grad_input, [grad_params], [(None,)] diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index 19608894e0..19a090f121 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -28,3 +28,12 @@ register_backward_fusion(BackwardLinearScale.fuse_backward_ops) register_backward_fusion(BackwardActivationBias.fuse_backward_ops) register_backward_fusion(BackwardAddRMSNorm.fuse_backward_ops) + +# 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, +) +from .backward_grouped_mlp import ( # pylint: disable=wrong-import-position + BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, +) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py new file mode 100644 index 0000000000..a821258ebf --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -0,0 +1,679 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused operation for MoE grouped MLP.""" + +from __future__ import annotations +from collections.abc import Callable +import functools +import inspect +import math +import os +from typing import Optional + +import torch + +import transformer_engine_torch as tex +from ...cpp_extensions import ( + general_grouped_gemm_for_grouped_tensor, +) +from ...module.base import get_dummy_wgrad +from ...quantization import Recipe +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 ..basic import GroupedLinear, ScaledSwiGLU +from ..fuser import register_backward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext +from .._common import ( + fuse_grouped_mlp_ops, + maybe_dequantize, + validate_grouped_mlp_dims, +) + + +@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, + num_groups, + weight_shape, + grouped_x, + grouped_dy, + dtype, + device, + bias_grads, + bias_grad_packed, + label="", +): + """Compute weight gradients and build grad_params for a GroupedLinear layer. + Returns the grad_params list in parameter registration order. + """ + + # Allocate grad buffers, determine accumulate flag + accumulate_into_main_grad = False + grouped_wgrad = None + wgrad_output = None + 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) + 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, + ) + + if grouped_wgrad is None: + grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_groups, + shapes=[weight_shape] * num_groups, + quantizer=None, + device=device, + dtype=dtype, + ) + wgrad_output = grouped_wgrad + 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) + else: + for idx in range(num_groups): + w_list[idx] = torch.empty(weight_shape, dtype=dtype, device=device) + wgrad_output = w_list + + if ctx.weight_requires_grad: + # Launch or defer the GEMM + delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() + gemm_fn = functools.partial( + general_grouped_gemm_for_grouped_tensor, + layout="NT", + accumulate=accumulate_into_main_grad, + ) + if delay_wgrad: + fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) + 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 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: + w_list = list(w_list) if accumulate_into_main_grad else [None] * num_groups + if 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), + ) + + # Assemble grad_params in parameter registration order. + if not fc_op.has_bias: + return w_list + + if fc_op.single_grouped_bias: + return w_list + [bias_grad_packed] + + bias_list = bias_grads if bias_grads is not None else [None] * num_groups + if fc_op.single_grouped_weight: + return bias_list + w_list + return w_list + bias_list + + +class BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8(FusedOperation): + """Fused op for MXFP8 GroupedLinear + ScaledSwiGLU + 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 + + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_quant_kernel(cls) -> Callable: + """Grouped GEMM quant kernel for block-scaled inputs.""" + from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_quant_wrapper_sm100 + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether this fused operation is supported on the current system.""" + if int(os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")) <= 0: + return False + if get_device_compute_capability()[0] != 10: + return False + try: + cls.grouped_gemm_dglu_kernel() + cls.grouped_gemm_quant_kernel() + except ImportError: + 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, + *, + fc1: GroupedLinear, + swiglu: ScaledSwiGLU, + fc2: GroupedLinear, + ) -> None: + super().__init__((fc1, swiglu, fc2)) + if not self.is_supported(): + self.grouped_gemm_dglu_kernel() # Try triggering import error + raise RuntimeError(f"{self.__class__.__name__} is not supported on this system.") + validate_grouped_mlp_dims(fc1, swiglu, fc2) + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + **unused, # pylint: disable=unused-argument + ) -> tuple[ + torch.Tensor, + list[tuple[Optional[torch.Tensor], ...]], + list[tuple[()]], + ]: + + # Get basic operations + fc1_op, _, fc2_op = self.basic_ops + fc1_ctx, swiglu_ctx, fc2_ctx = basic_op_ctxs + + # Tensor properties + fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) + fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) + 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 + dtype = fc1_ctx.dtype + + # Saved tensors from FC1 forward + saved_tensors = fc1_ctx.saved_tensors + split_sizes, split_points, saved_tensors = ( + saved_tensors[0], + saved_tensors[1], + saved_tensors[2:], + ) + + if fc1_op.single_grouped_weight: + grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] + else: + grouped_fc1_weight, saved_tensors = ( + saved_tensors[:num_groups], + 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 + if fc2_op.single_grouped_weight: + grouped_fc2_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] + else: + grouped_fc2_weight, saved_tensors = ( + saved_tensors[:num_groups], + 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())}.") + split_sizes = split_sizes.to(dtype=torch.int64, device=device) + split_points = split_points.to(dtype=torch.int, device=device) + + 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, + ) + + # 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 + output_fc2_dbias = fc2_op.has_bias + fc2_dbias_packed = None + if ( + not output_fc2_dbias + and isinstance(grad_output, GroupedTensor) + and isinstance(getattr(grad_output, "quantizer", None), MXFP8Quantizer) + ): + grouped_fc2_dy = grad_output + else: + fc2_dy = maybe_dequantize(grad_output, dtype) + if output_fc2_dbias: + grouped_fc2_dy, fc2_dbias_packed = tex.bgrad_group_quantize( + fc2_dy, + fc2_ctx.grad_output_quantizer, + num_groups, + split_sizes, + ) + else: + grouped_fc2_dy = tex.group_quantize( + fc2_dy, + fc2_ctx.grad_output_quantizer, + num_groups, + split_sizes, + ) + + fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None + fc2_bias_grad_packed: Optional[torch.Tensor] = None + if fc2_dbias_packed is not None: + if fc2_op.single_grouped_bias: + fc2_bias_grad_packed = fc2_dbias_packed.to(dtype=dtype) + else: + fc2_bias_grads = [ + fc2_dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups) + ] + + # Pack data tensors + # Note: Fused kernel expects tensor with non-contiguous + # logical dims. + # Data actual shape: (1, sum(m), k) + # Scale actual shape: (1, sum(m)/128, k/128, 32 (block row), + # 4 (block row), 4 (block col)) + # 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 = 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] // 128, + out_shape[1] // 128, + MXFP8_BLOCK_SCALING_SIZE, + 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) + norm_const_tensor = get_cached_ones_tensor(1, dtype, device) + current_stream = torch.cuda.current_stream().cuda_stream + + prob_tensor = scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + dprob_tensor = torch.zeros_like(prob_tensor) + + fc2_dglu_kwargs = { + "a_tensor": fc2_dy_data, + "c_tensor": swiglu_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": prob_tensor, + "dprob_tensor": dprob_tensor, + "generate_dbias": fc1_op.has_bias, + "norm_const_tensor": norm_const_tensor, + "d_dtype": torch.float8_e4m3fn, + "cd_major": "n", + "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "current_stream": current_stream, + "discrete_col_sfd": True, + "act_func": "dswiglu", + "use_dynamic_sched": True, + } + + if fc2_op.single_grouped_weight: + # Clone and swizzle scales for GEMM + fc2_weight_for_gemm = grouped_fc2_weight.copy() + tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=False, columnwise=True) + # Pack weight tensors for stacked kernel + # 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_scales = fc2_w_scales.view( + num_groups, + fc2_weight_shape[1] // 128, + fc2_weight_shape[0] // 128, + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + 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 + 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], + [w._columnwise_scale_inv for w in grouped_fc2_weight], + swizzle=True, + 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_dgrad_kernel_out = self.grouped_gemm_dglu_kernel()(**fc2_dglu_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]) + fc1_dy_row_scale = fc2_dgrad_kernel_out["sfd_row_tensor"] + 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]) + fc1_dy_col_scale = fc2_dgrad_kernel_out["sfd_col_tensor"] + grad_scales = fc2_dgrad_kernel_out["dprob_tensor"] + grad_scales = grad_scales.view(-1).to(dtype=dtype) + + fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None + fc1_bias_grad_packed: Optional[torch.Tensor] = None + if fc1_op.has_bias: + dbias_t = fc2_dgrad_kernel_out["dbias_tensor"] + if dbias_t is not None: + dbias_2d = dbias_t.squeeze(-1) + if fc1_op.single_grouped_bias: + fc1_bias_grad_packed = dbias_2d.to(dtype=dtype) + else: + fc1_bias_grads = [ + dbias_2d[group_idx].to(dtype=dtype) 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] + grouped_fc1_dy = GroupedTensor( + shape=(out_shape[0], fc1_weight_shape[0]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc1_ctx.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 + fc2_grad_params = _compute_grad_params( + fc_op=fc2_op, + ctx=fc2_ctx, + num_groups=num_groups, + weight_shape=fc2_weight_shape, + grouped_x=grouped_fc2_x, + grouped_dy=grouped_fc2_dy, + dtype=dtype, + device=device, + bias_grads=fc2_bias_grads, + bias_grad_packed=fc2_bias_grad_packed, + label="FC2", + ) + + # Clear FC2 input tensor if possible + if grouped_fc2_x is not None and not ( + fc2_ctx.weight_requires_grad + and fc2_op.wgrad_store is not None + and fc2_op.wgrad_store.delay_wgrad_compute() + ): + clear_tensor_data( + grouped_fc2_x.data, + grouped_fc2_x.columnwise_data, + grouped_fc2_x.scale_inv, + grouped_fc2_x.columnwise_scale_inv, + ) + + # FC1 dgrad GEMM + grad_input = None + 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.float(), + "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, + "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] // 128, + 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, fc1_sfb_ptrs, _ = tex.get_device_pointer_for_data_and_scales( + [w._columnwise_data for w in grouped_fc1_weight], + [w._columnwise_scale_inv for w in grouped_fc1_weight], + swizzle=True, + rowwise=False, + data_dtype=grouped_fc1_weight[0]._fp8_dtype, + ) + + 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( + fc_op=fc1_op, + ctx=fc1_ctx, + num_groups=num_groups, + weight_shape=fc1_weight_shape, + grouped_x=grouped_fc1_x, + grouped_dy=grouped_fc1_dy, + dtype=dtype, + device=device, + bias_grads=fc1_bias_grads, + bias_grad_packed=fc1_bias_grad_packed, + label="FC1", + ) + + # Clear FC1 input tensor if possible + if grouped_fc1_x is not None and not ( + fc1_ctx.weight_requires_grad + and fc1_op.wgrad_store is not None + and fc1_op.wgrad_store.delay_wgrad_compute() + ): + clear_tensor_data( + grouped_fc1_x.data, + grouped_fc1_x.columnwise_data, + grouped_fc1_x.scale_inv, + grouped_fc1_x.columnwise_scale_inv, + ) + + return ( + grad_input, + [fc1_grad_params, (), fc2_grad_params], + [(None,), (grad_scales,), (None,)], + ) + + +def fuse_backward_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument +) -> list[FusibleOperation]: + """Apply operation fusion for backward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Forward pass operations. + recipe : Recipe, optional + Quantization recipe. + + Returns + ------- + ops : list of FusibleOperation + Updated backward pass operations + + """ + + return fuse_grouped_mlp_ops( + ops, + recipe=recipe, + fused_op_cls=BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + ) + + +# Register fusion if available +if BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): + register_backward_fusion(fuse_backward_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 new file mode 100644 index 0000000000..c5ce2b148d --- /dev/null +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -0,0 +1,573 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused operation for MoE grouped MLP.""" + +from __future__ import annotations +from collections.abc import Callable, Iterable +import functools +import inspect +import os +from typing import Any, Optional + +import torch + +import transformer_engine_torch as tex +from ...quantization import Recipe +from ...tensor import Quantizer +from ...utils import 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 ..basic import GroupedLinear, ScaledSwiGLU +from ..fuser import register_forward_fusion +from ..op import FusedOperation, FusibleOperation, OperationContext +from .._common import ( + fuse_grouped_mlp_ops, + is_quantized_tensor, + maybe_dequantize, + validate_grouped_mlp_dims, +) + + +def _pack_grouped_linear_bias_for_cudnn(linear_op: GroupedLinear) -> Optional[torch.Tensor]: + """Bias layout expected by cuDNN grouped GEMM: shape (n, num_groups), stride (1, n).""" + if not linear_op.has_bias: + return None + num_groups = linear_op.num_groups + grouped_bias = getattr(linear_op, "bias", None) + if grouped_bias is not None: + packed = grouped_bias.rowwise_data.view(num_groups, -1) + return packed.transpose(0, 1) + rows = [getattr(linear_op, f"bias{group_idx}") for group_idx in range(num_groups)] + # stack to [num_groups, n] but cuDNN expects [n, num_groups] with stride [1, n]. + return torch.stack(rows, dim=0).transpose(0, 1) + + +class ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8(FusedOperation): + """Fused op for MXFP8 GroupedLinear + ScaledSwiGLU + 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 + + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_quant_kernel(cls) -> Callable: + """Grouped GEMM quant kernel for block-scaled inputs.""" + from cudnn import grouped_gemm_quant_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_quant_wrapper_sm100 + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether this fused operation is supported on the current system.""" + if int(os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")) <= 0: + return False + if get_device_compute_capability()[0] != 10: + return False + try: + cls.grouped_gemm_glu_kernel() + cls.grouped_gemm_quant_kernel() + except ImportError: + 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, + *, + fc1: GroupedLinear, + swiglu: ScaledSwiGLU, + fc2: GroupedLinear, + ) -> None: + super().__init__((fc1, swiglu, fc2)) + if not self.is_supported(): + self.grouped_gemm_glu_kernel() # Try triggering import error + raise RuntimeError(f"{self.__class__.__name__} is not supported on this system.") + validate_grouped_mlp_dims(fc1, swiglu, fc2) + + 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], + next_op_input_quantizer: Optional[Quantizer], + basic_op_kwargs: list[dict[str, Any]], + ) -> 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 + + # Tensor properties + fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) + fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) + input_ = input_.reshape(-1, fc1_weight_shape[1]) + in_shape = list(input_.size()) + + num_groups = fc1_op.num_groups + 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 + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + else: + dtype = fc1_weight_param.dtype + + # Check which grads are required + requires_grad = any(ctx.requires_grad for ctx in basic_op_ctxs) + input_requires_grad = requires_grad + weight_requires_grad = requires_grad and ( + fc1_weight_param.requires_grad or fc2_weight_param.requires_grad + ) + + # Quantizers + fc1_input_quantizer = fc1_op.get_quantizer("forward", 0) + fc1_weight_quantizer = fc1_op.get_quantizer("forward", 1) + fc1_grad_output_quantizer = fc1_op.get_quantizer("backward", 0) + fc2_input_quantizer = fc2_op.get_quantizer("forward", 0) + fc2_weight_quantizer = fc2_op.get_quantizer("forward", 1) + fc2_grad_output_quantizer = fc2_op.get_quantizer("backward", 0) + + # Extract split sizes from extra input + fc1_split_sizes = basic_op_extra_inputs[0][0] + fc2_split_sizes = basic_op_extra_inputs[2][0] + if ( + fc1_split_sizes.size() != fc2_split_sizes.size() + or fc1_split_sizes.data_ptr() != fc2_split_sizes.data_ptr() + ): + raise RuntimeError( + f"{self.__class__.__name__} got different split points for FC1 and FC2." + ) + split_sizes = fc1_split_sizes + 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) + split_points = torch.cumsum(split_sizes, 0, dtype=torch.int) + split_points_offsets = torch.cumsum(split_sizes, 0) + base_offsets = torch.cat( + [ + torch.zeros(1, device=split_sizes.device, dtype=split_sizes.dtype), + split_points_offsets, + ] + ) + fc1_x_tensor_offsets = base_offsets * fc1_weight_shape[1] + fc2_x_tensor_offsets = base_offsets * fc2_weight_shape[1] + + # Extract post-scales from extra input + scales = basic_op_extra_inputs[1][0] + + # Prepare FC1 grouped weight tensor for fused kernels. + # - single_grouped_weight=True: op.weight is already a GroupedTensor + # - single_grouped_weight=False: cute DSL kernel works with discrete weight tensors + # as long as host pointers for addresses are packed as contiguous device tensor. + if fc1_op.single_grouped_weight: + if not isinstance(fc1_op.weight, GroupedTensor): + raise RuntimeError( + "FC1 expected GroupedTensor weight with single_grouped_weight=True." + ) + if fc1_op.weight.quantizer is not None: + fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc1_op.weight.quantizer = fc1_weight_quantizer + grouped_fc1_weight = fc1_op.weight + else: + 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( + fc1_op.weight.rowwise_data.view(fc1_op.weight.logical_shape), + fc1_weight_quantizer, + num_groups, + None, + ) + else: + fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)] + quantized_fc1_weights = [] + for idx, weight in enumerate(fc1_weights): + quantizer = fc1_op.get_quantizer("forward", 2 * idx + 1) + if not is_quantized_tensor(weight): + quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + quantized_fc1_weights.append(quantizer(weight)) + else: + quantized_fc1_weights.append(weight) + grouped_fc1_weight = quantized_fc1_weights + + # Prepare FC2 grouped weight tensor for fused kernels. + if fc2_op.single_grouped_weight: + if not isinstance(fc2_op.weight, GroupedTensor): + raise RuntimeError( + "FC2 expected GroupedTensor weight with single_grouped_weight=True." + ) + if fc2_op.weight.quantizer is not None: + fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + fc2_op.weight.quantizer = fc2_weight_quantizer + grouped_fc2_weight = fc2_op.weight + else: + 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( + fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), + fc2_weight_quantizer, + num_groups, + None, + ) + else: + fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)] + quantized_fc2_weights = [] + for idx, weight in enumerate(fc2_weights): + quantizer = fc2_op.get_quantizer("forward", 2 * idx + 1) + quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + if not is_quantized_tensor(weight): + quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + quantized_fc2_weights.append(quantizer(weight)) + else: + quantized_fc2_weights.append(weight) + 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 + ): + 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 + ): + 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 + ): + 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) + + # Pack data tensors + # Note: Fused kernel expects tensor with non-contiguous + # logical dims. + # Data actual shape: (1, sum(m), k) + # Scale actual shape: (1, sum(m)/128, k/128, 32 (block row), + # 4 (block row), 4 (block col)) + # 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 = 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] // 128, + in_shape[1] // 128, + MXFP8_BLOCK_SCALING_SIZE, + 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, dtype, 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 = { + "a_tensor": fc1_x_data, + "sfa_tensor": fc1_x_scales, + "padded_offsets": split_points, + "alpha_tensor": alpha_tensor, + "bias_tensor": fc1_bias_packed, + "norm_const_tensor": norm_const_tensor, + "prob_tensor": scales.detach().to(dtype=dtype).reshape(-1, 1, 1), + "acc_dtype": torch.float32, + "c_dtype": torch.bfloat16, + "d_dtype": torch.float8_e4m3fn, + "cd_major": "n", + "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "current_stream": current_stream, + "discrete_col_sfd": True, + "act_func": "swiglu", + "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=True, columnwise=False) + + # Pack weight tensors for stacked kernel + # 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.permute(1, 2, 0) + fc1_w_scales = fc1_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) + fc1_w_scales = fc1_w_scales.view( + num_groups, + fc1_weight_shape[0] // 128, + fc1_weight_shape[1] // 128, + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + 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 + 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( + [w._rowwise_data for w in grouped_fc1_weight], + [w._rowwise_scale_inv for w in grouped_fc1_weight], + swizzle=True, + 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_kernel_out = self.grouped_gemm_glu_kernel()(**fc1_glu_kwargs) + + # Unpack kernel outputs + # Note: Fused kernel outputs tensors with non-contiguous + # logical dims. + # Row-wise data logical shape: (sum(m_splits), k, 1) + # Row-wise scale logical shape: (32 (block row), 4 (block row), + # sum(m_splits)/128, 4 (block col), k/128, 1) + # 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]) + 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_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.float(), + "norm_const_tensor": None, + "prob_tensor": torch.ones((in_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, + "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) + 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, + fc2_weight_shape[0] // 128, + 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, fc2_sfb_ptrs, _ = tex.get_device_pointer_for_data_and_scales( + [w._rowwise_data for w in grouped_fc2_weight], + [w._rowwise_scale_inv for w in grouped_fc2_weight], + swizzle=True, + rowwise=True, + data_dtype=grouped_fc2_weight[0]._fp8_dtype, + ) + 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() + + # 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 + 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 + ) + 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.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) + swiglu_ctx.input_requires_grad = True + 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_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.grad_input_quantizers = None + fc2_ctx.dtype = dtype + fc2_ctx.input_requires_grad = input_requires_grad + fc2_ctx.weight_requires_grad = weight_requires_grad + + return fc2_out, [(), (), ()] + + +def fuse_forward_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument +) -> list[FusibleOperation]: + """Apply operation fusion for forward pass. + + Parameters + ---------- + ops : list of FusibleOperation + Forward pass operations. + recipe : Recipe, optional + Quantization recipe. + + Returns + ------- + ops : list of FusibleOperation + Updated forward pass operations + + """ + + return fuse_grouped_mlp_ops( + ops, + recipe=recipe, + fused_op_cls=ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + ) + + +# Register fusion if available +if ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): + register_forward_fusion(fuse_forward_ops, prepend=True) diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index 2fce9a38e2..ab0c7484fc 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -74,7 +74,7 @@ def __new__( dtype: torch.dtype, *, num_tensors: int, - shapes: Optional[List[Tuple[int, int]]] = None, + shapes: Optional[List[Tuple[int, ...]]] = None, quantizer: Optional[Quantizer] = None, data: Optional[torch.Tensor] = None, columnwise_data: Optional[torch.Tensor] = None, @@ -99,7 +99,15 @@ def __new__( and num_tensors > 0 and all(shapes[0] == s for s in shapes) ): - wrapper_shape = (num_tensors, shapes[0][0], shapes[0][1]) + s0 = shapes[0] + if len(s0) == 2: + wrapper_shape = (num_tensors, s0[0], s0[1]) + elif len(s0) == 1: + wrapper_shape = (num_tensors, s0[0]) + else: + raise ValueError( + f"GroupedTensor member shapes must be 1D or 2D, got {len(s0)}-D shape {s0!r}" + ) else: wrapper_shape = shape @@ -186,6 +194,7 @@ def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> Non dst.columnwise_scale_inv_offsets = src.columnwise_scale_inv_offsets dst.logical_shape = src.logical_shape dst.quantized_tensors = src.quantized_tensors + dst._with_gemm_swizzled_scales = src._with_gemm_swizzled_scales 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/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 68097259c6..ff1c78f695 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -54,7 +54,7 @@ def _initialize_storage_fields( shape: Tuple[int, int], dtype: torch.dtype, num_tensors: int, - shapes: Optional[List[Tuple[int, int]]] = None, + shapes: Optional[List[Tuple[int, ...]]] = None, quantizer: Optional[Quantizer] = None, data: Optional[torch.Tensor] = None, columnwise_data: Optional[torch.Tensor] = None, @@ -153,7 +153,7 @@ def __new__( dtype: torch.dtype, *, num_tensors: int, - shapes: Optional[List[Tuple[int, int]]] = None, + shapes: Optional[List[Tuple[int, ...]]] = None, quantizer: Optional[Quantizer] = None, data: Optional[torch.Tensor] = None, columnwise_data: Optional[torch.Tensor] = None, @@ -383,6 +383,128 @@ def make_grouped_tensor_with_shapes( dtype=dtype, ) + @staticmethod + def make_grouped_tensor_from_rowwise_data( + *, + num_tensors: int, + tensor_shape: Tuple[int, ...], + rowwise_data: torch.Tensor, + dtype: Optional[torch.dtype] = None, + internal: bool = False, + ) -> GroupedTensorStorage: + """Wrap pre-existing contiguous rowwise data as a grouped tensor. + + This helper does not allocate storage. It creates grouped metadata over + `rowwise_data`, which is expected to contain `num_tensors` tensors of + shape ``tensor_shape`` in packed contiguous layout. + + ``tensor_shape`` may be: + + * ``(rows, cols)`` — each member is a 2D matrix; wrapper shape + ``(num_tensors, rows, cols)``. + * ``(n,)`` — each member is a 1D vector of length ``n``; logical storage + uses ``logical_shape = (num_tensors * n, 1)`` and the wrapper shape is + ``(num_tensors, n)``. + """ + if num_tensors <= 0: + raise ValueError(f"num_tensors must be positive, got {num_tensors}") + if rowwise_data is None: + raise ValueError("rowwise_data must not be None") + if not rowwise_data.is_contiguous(): + rowwise_data = rowwise_data.contiguous() + + if len(tensor_shape) == 2: + rows, cols = tensor_shape + expected_numel = num_tensors * rows * cols + logical_shape = (num_tensors * rows, cols) + shapes_list: List[Tuple[int, ...]] = [tensor_shape] * num_tensors + elif len(tensor_shape) == 1: + (n,) = tensor_shape + expected_numel = num_tensors * n + logical_shape = (num_tensors * n, 1) + shapes_list = [tensor_shape] * num_tensors + else: + raise ValueError( + "tensor_shape must be 1D (n,) or 2D (rows, cols), " + f"got {tensor_shape!r} with length {len(tensor_shape)}" + ) + + if rowwise_data.numel() != expected_numel: + raise ValueError( + "Grouped rowwise buffer size mismatch: expected " + f"{expected_numel} elements for {num_tensors}x{tensor_shape}, " + f"but got {rowwise_data.numel()}" + ) + if dtype is None: + dtype = rowwise_data.dtype + grouped_tensor_class = GroupedTensorStorage + if not internal: + from ..grouped_tensor import GroupedTensor + + grouped_tensor_class = GroupedTensor + + return grouped_tensor_class( + shape=logical_shape, + dtype=dtype, + num_tensors=num_tensors, + shapes=shapes_list, + quantizer=None, + data=rowwise_data.view(-1), + columnwise_data=None, + scale_inv=None, + columnwise_scale_inv=None, + amax=None, + columnwise_amax=None, + scale=None, + first_dims=None, + last_dims=None, + tensor_offsets=None, + offsets=None, + scale_inv_offsets=None, + columnwise_scale_inv_offsets=None, + with_gemm_swizzled_scales=False, + requires_grad=False, + ) + + def copy(self) -> "GroupedTensorStorage": + """Create a shallow copy that shares all data buffers with *self*. + No tensor data is copied; the returned object references the same + underlying storage for every buffer (data, scales, offsets, etc.). + This is useful when you need to mutate metadata (e.g. swizzle + scales in-place) without affecting the original object. + """ + return GroupedTensorStorage( + shape=self.logical_shape, + dtype=self.fake_dtype, + num_tensors=self.num_tensors, + shapes=self.tensor_shapes, + quantizer=self.quantizer, + data=self.rowwise_data, + columnwise_data=self.columnwise_data, + scale_inv=self.scale_inv, + columnwise_scale_inv=self.columnwise_scale_inv, + amax=self.amax, + columnwise_amax=self.columnwise_amax, + scale=self.scale, + first_dims=self.first_dims, + last_dims=self.last_dims, + tensor_offsets=self.tensor_offsets, + offsets=self.offsets, + 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, + ) + + @staticmethod + def make_tensor_offsets(first_dims: torch.Tensor, logical_last_dim: int) -> torch.Tensor: + """Calculate GPU offsets from first dim splits.""" + return torch.cat( + [ + torch.zeros(1, device=first_dims.device, dtype=first_dims.dtype), + torch.cumsum(first_dims * logical_last_dim, dim=0), + ] + ) + @staticmethod def make_grouped_tensor( num_tensors: int, @@ -421,7 +543,7 @@ def make_grouped_tensor( all_same_last = last_dims is None assert all_same_last, "Last dim must be uniform for GroupedTensor" - assert logical_first_dim > 0, "Logical first dim must be positive for GroupedTensor" + assert logical_first_dim >= 0, "Logical first dim must be non-negative for GroupedTensor" assert logical_last_dim > 0, "Logical last dim must be positive for GroupedTensor" # assert ( @@ -439,16 +561,20 @@ def make_grouped_tensor( # Kernels need to calculate precise pointers based on size of elements. # TODO(ksivaman): Single kernel + remove the host offset calculation. - tensor_offsets = torch.cat( - [ - torch.zeros(1, device=first_dims.device, dtype=first_dims.dtype), - torch.cumsum(first_dims * logical_last_dim, dim=0), - ] - ) - offsets = tensor_offsets.tolist() - first_dims_list = first_dims.tolist() - for i in range(num_tensors): - shape.append((first_dims_list[i], logical_last_dim)) + tensor_offsets = GroupedTensorStorage.make_tensor_offsets(first_dims, logical_last_dim) + if ( + first_dims.device.type == "cuda" + and torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + ): + # Avoid host sync during CUDA graph capture. + offsets = None + shape = None + else: + offsets = tensor_offsets.tolist() + first_dims_list = first_dims.tolist() + for i in range(num_tensors): + shape.append((first_dims_list[i], logical_last_dim)) else: offsets = [ i * logical_first_dim * logical_last_dim // num_tensors @@ -653,7 +779,6 @@ def make_grouped_tensor( quantizer.optimize_for_gemm if quantizer is not None else False ), ) - grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() return grouped_tensor @@ -709,7 +834,7 @@ def split_into_quantized_tensors( # Get tensor data slice if self.offsets is not None: start_offset = self.offsets[i] - numel = tensor_shape[0] * tensor_shape[1] + numel = math.prod(tensor_shape) end_offset = start_offset + numel if self.has_data(): @@ -724,7 +849,7 @@ def split_into_quantized_tensors( raise RuntimeError("GroupedTensor has no data to split") else: # All same shape case - numel = tensor_shape[0] * tensor_shape[1] + numel = math.prod(tensor_shape) start_offset = i * numel end_offset = start_offset + numel @@ -760,7 +885,7 @@ def split_into_quantized_tensors( quantizer = self.quantizer # Get tensor shape tensor_shape = self.tensor_shapes[i] - numel = tensor_shape[0] * tensor_shape[1] + numel = math.prod(tensor_shape) # Get data offsets if self.offsets is not None: diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index db2f28aa47..a76f205acc 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -19,6 +19,19 @@ __all__ = ["get_device_compute_capability", "get_cudnn_version", "is_bf16_available"] +@functools.lru_cache(maxsize=None) +def get_cached_ones_tensor( + num_elements: int, + dtype: torch.dtype, + device: torch.device, +) -> torch.Tensor: + """Return a cached ``torch.ones`` tensor. + Tensors are cached by ``(num_elements, dtype, device)`` and kept alive + by the cache, ensuring stable data pointers across CUDA graph replays. + """ + return torch.ones(num_elements, dtype=dtype, device=device) + + def requires_grad(*tensors: Tuple[Optional[torch.Tensor], ...]) -> None: """Check if any of the given tensors require gradient.""" for tensor in tensors: @@ -157,6 +170,29 @@ def divide(numerator: int, denominator: int) -> int: return numerator // denominator +def mark_grouped_tensor(*tensors: List[Any]): + """ + Needed for paged stashing in Megatron-LM. This attribute allows + Megatron-LM to detect which tensors are dynamic (varying shapes) + and remove the padding before doing the `save_for_backward` to + save memory. + Note: Only columnwise data is saved for backward.""" + for tensor in tensors: + if tensor is None: + continue + if hasattr(tensor, "columnwise_data"): + assert ( + tensor.columnwise_data is not None + ), "Columnwise data is not set for grouped tensor" + assert ( + tensor.columnwise_scale_inv is not None + ), "Columnwise scale inverse is not set for grouped tensor" + setattr(tensor.columnwise_data, "grouped_tensor_scale_inv", False) + setattr(tensor.columnwise_scale_inv, "grouped_tensor_scale_inv", True) + else: + setattr(tensor, "grouped_tensor_scale_inv", False) + + def split_tensor_along_dim( tensor: torch.Tensor, dim: int, num_partitions: int, contiguous_split_chunks: bool = False ) -> Tuple[torch.Tensor, ...]: From 8cf3c1662605088b408b79a43e136acf14f32481 Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Fri, 3 Apr 2026 10:14:00 -0600 Subject: [PATCH 310/521] [PyT][Test] Add xfailing FSDP2 memory leak detection tests (#2803) Add tests that demonstrate two known memory issues with FSDP2 + FP8: - Issue #2681: FP8 weight copies created during te.autocast() forward pass accumulate across layers instead of being freed between layers, defeating FSDP2's memory efficiency. Detected by comparing per-layer forward memory increments against a bf16 baseline using layer hooks. - Issue #2717: Transpose cache tensors (_create_transpose) allocated during backward persist until the next forward pass instead of being freed after backward completes. Detected by comparing the backward memory delta (post_bwd - post_fwd) against a bf16 baseline. New tests: - test_bf16_no_excess_forward_memory: control, validates per-layer measurement - test_bf16_no_excess_backward_memory: control, validates backward delta comparison - test_fp8_temp_accumulation_across_layers: xfail, detects #2681 - test_transpose_cache_retained_after_backward: xfail, detects #2717 All parametrized over 5 FP8 recipes x {no_quant_init, quant_init}. Signed-off-by: Peter St. John Co-authored-by: vthumbe1503 --- .../fsdp2_tests/run_fsdp2_mem_leak.py | 518 ++++++++++++++++++ tests/pytorch/distributed/test_torch_fsdp2.py | 24 + 2 files changed, 542 insertions(+) create mode 100644 tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py new file mode 100644 index 0000000000..387d3a9644 --- /dev/null +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -0,0 +1,518 @@ +#!/usr/bin/python3 + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FSDP2 memory leak detection tests. + +These tests verify that temporary TE tensors (FP8 quantized weights, transpose +caches) are properly freed when moving between layers with FSDP2. + +Related issues: + - https://github.com/NVIDIA/TransformerEngine/issues/2681 + Quantized weights created during forward pass accumulate across layers. + - https://github.com/NVIDIA/TransformerEngine/issues/2717 + _create_transpose tensors accumulate across training steps with + quantized_model_init + FusedAdam + FSDP2. + +Run all tests (via torchrun + pytest): + torchrun -m pytest -v --tb=short + +Run a single test standalone (for debugging): + torchrun --test --recipe + +Available --test values: + bf16_no_excess_forward_memory, fp8_temp_accumulation_across_layers, + transpose_cache_retained_after_backward + +Available --recipe values: + DelayedScaling, Float8CurrentScaling, Float8BlockScaling, + MXFP8BlockScaling, NVFP4BlockScaling +""" + +import argparse +import gc +import os +from contextlib import nullcontext + +import pytest +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch.distributed._composable.fsdp import fully_shard +from torch.distributed.device_mesh import DeviceMesh + +import transformer_engine.pytorch as te + +from fsdp2_utils import get_recipe_from_string, save_custom_attrs, restore_custom_attrs + + +# ── Constants ──────────────────────────────────────────────────────── +HIDDEN_SIZE = 256 +FFN_HIDDEN_SIZE = 1024 +NUM_ATTENTION_HEADS = 8 +NUM_LAYERS = 8 +SEQ_LEN = 32 +BATCH_PER_RANK = 2 +WARMUP_STEPS = 2 + + +# ── Helpers ────────────────────────────────────────────────────────── +def _build_model(num_layers, fp8_init, recipe=None, use_meta_device=True): + """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 so parameters are materialized after FSDP2 + sharding via reset_parameters(). + """ + if fp8_init: + ctx = te.quantized_model_init(enabled=True, recipe=recipe) + else: + ctx = nullcontext() + kwargs = dict( + fuse_qkv_params=True, + params_dtype=torch.bfloat16, + hidden_dropout=0.0, + attention_dropout=0.0, + ) + if fp8_init and use_meta_device: + kwargs["device"] = "meta" + with ctx: + model = torch.nn.Sequential( + *[ + te.TransformerLayer( + HIDDEN_SIZE, + FFN_HIDDEN_SIZE, + NUM_ATTENTION_HEADS, + **kwargs, + ) + for _ in range(num_layers) + ] + ) + return model + + +def _shard_model(model, world_size): + """Apply FSDP2 sharding with save/restore of custom attrs.""" + has_meta_params = any(p.is_meta for p in model.parameters()) + custom_attrs = save_custom_attrs(model) + mesh = DeviceMesh("cuda", list(range(world_size))) + for child in model.children(): + fully_shard(child, mesh=mesh) + fully_shard(model, mesh=mesh) + if has_meta_params: + for module in model.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + restore_custom_attrs(model, custom_attrs) + return model + + +def _get_dist_info(): + """Get world_size and device from environment.""" + world_size = int(os.environ["WORLD_SIZE"]) + device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}") + return world_size, device + + +def _run_training_step(model, optimizer, recipe, x, target): + """Run one forward + backward + optimizer step.""" + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=(recipe is not None), recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + return loss.item() + + +def _measure_backward_memory_delta(model, optimizer, recipe, x, target): + """Run a training step and return (post_bwd - post_fwd) memory delta. + + This delta captures memory added during backward that persists afterward. + In a healthy system, backward frees activations and adds only gradients. + If transpose caches or other FP8 temps persist, the delta will be larger. + """ + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=(recipe is not None), recipe=recipe): + output = model(x) + torch.cuda.synchronize() + mem_post_fwd = torch.cuda.memory_allocated() + + loss = F.mse_loss(output, target) + loss.backward() + torch.cuda.synchronize() + mem_post_bwd = torch.cuda.memory_allocated() + + optimizer.step() + return mem_post_bwd - mem_post_fwd + + +def _maybe_skip(recipe_name, quantized_model_init): + """Skip configurations that fail for reasons unrelated to memory leaks.""" + if recipe_name == "NVFP4BlockScaling" and quantized_model_init: + pytest.skip( + "NVFP4BlockScaling + quantized_model_init: not supported with FSDP2 " + "(block tensor dequantized before FSDP2 flatten)" + ) + + +class _LayerMemoryTracker: + """Register forward hooks on Sequential children to measure per-layer memory.""" + + def __init__(self): + self.post_forward_mem = [] + self._handles = [] + + def attach(self, model): + for i, layer in enumerate(model.children()): + + def make_hook(idx): + def hook(module, args, output): + torch.cuda.synchronize() + self.post_forward_mem.append(torch.cuda.memory_allocated()) + + return hook + + self._handles.append(layer.register_forward_hook(make_hook(i))) + + def clear(self): + self.post_forward_mem.clear() + + def detach(self): + for h in self._handles: + h.remove() + self._handles.clear() + + def per_layer_increments(self): + """Return list of memory increments between consecutive post-forward hooks.""" + return [ + self.post_forward_mem[i] - self.post_forward_mem[i - 1] + for i in range(1, len(self.post_forward_mem)) + ] + + +def _measure_forward_increments(model, optimizer, recipe, x, target): + """Run a single training step with hooks and return per-layer forward memory increments.""" + tracker = _LayerMemoryTracker() + tracker.attach(model) + try: + _run_training_step(model, optimizer, recipe, x, target) + return tracker.per_layer_increments() + finally: + tracker.detach() + + +# ── Fixtures ───────────────────────────────────────────────────────── +@pytest.fixture(params=[False, True], ids=["no_quant_init", "quant_init"]) +def quantized_model_init(request): + return request.param + + +# ── Tests ──────────────────────────────────────────────────────────── +def test_bf16_no_excess_forward_memory(): + """Control test: bf16 (no FP8) should have stable per-layer forward memory. + + With FSDP2 and bf16 params (no FP8), the per-layer memory growth during + forward should only be activation saves for autograd. There should be no + FP8 temporary accumulation. This test validates the measurement approach. + """ + world_size, device = _get_dist_info() + + model = _build_model(NUM_LAYERS, fp8_init=False) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # Warmup + for _ in range(WARMUP_STEPS): + _run_training_step(model, optimizer, None, x, target) + + # Measure + increments = _measure_forward_increments(model, optimizer, None, x, target) + + # bf16 per-layer increments should be consistent (activation saves only) + # and should NOT grow over layers (each layer saves similar activations). + avg_increment = sum(increments) / len(increments) + max_deviation = max(abs(inc - avg_increment) for inc in increments) + + # Allow 10% deviation from mean -- bf16 increments should be very uniform + assert max_deviation <= 0.1 * abs(avg_increment) + 1024, ( + "bf16 per-layer increments are not uniform. " + f"Increments (KiB): {[f'{inc/1024:.1f}' for inc in increments]}. " + f"Average: {avg_increment/1024:.1f} KiB, max deviation: {max_deviation/1024:.1f} KiB" + ) + + +@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. + + Strategy: measure per-layer memory growth during forward for both bf16 + (baseline) and FP8. With FSDP2, per-layer params are unsharded then + resharded, so the only per-layer memory growth should be activation saves + for autograd (same as bf16). If FP8 adds excess per-layer growth, it means + FP8 weight copies are accumulating across layers instead of being freed. + """ + _maybe_skip(recipe_name, quantized_model_init) + + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # ── bf16 baseline ── + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_increments = _measure_forward_increments(bf16_model, bf16_optimizer, None, x, target) + bf16_avg = sum(bf16_increments) / len(bf16_increments) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # ── FP8 model ── + fp8_model = _build_model(NUM_LAYERS, fp8_init=quantized_model_init, recipe=recipe) + fp8_model = _shard_model(fp8_model, world_size) + fp8_optimizer = te.optimizers.FusedAdam( + fp8_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(fp8_model, fp8_optimizer, recipe, x, target) + fp8_increments = _measure_forward_increments(fp8_model, fp8_optimizer, recipe, x, target) + fp8_avg = sum(fp8_increments) / len(fp8_increments) + + # ── Assert: FP8 per-layer excess should be bounded ── + # If FP8 temps are properly freed between layers, per-layer increment + # should be similar to bf16 (just activation saves). Any excess indicates + # FP8 weight copies accumulating. + excess_per_layer = fp8_avg - bf16_avg + + # Allow up to 50 KiB per layer for FP8 scale/amax metadata. + # FP8 weight copies (~0.68 MiB/layer for this model) should NOT persist. + tolerance_per_layer = 50 * 1024 # 50 KiB + + assert excess_per_layer <= tolerance_per_layer, ( + "FP8 per-layer forward memory increment exceeds bf16 baseline by " + f"{excess_per_layer/1024:.1f} KiB/layer (tolerance: {tolerance_per_layer/1024:.1f} KiB). " + f"bf16 avg: {bf16_avg/1024:.1f} KiB/layer, FP8 avg: {fp8_avg/1024:.1f} KiB/layer. " + f"FP8 increments (KiB): {[f'{inc/1024:.1f}' for inc in fp8_increments]}. " + "FP8 weight copies are likely accumulating across layers (Issue #2681)." + ) + + +def test_bf16_no_excess_backward_memory(): + """Control test: two identical bf16 models should show zero backward excess. + + This mirrors the structure of test_transpose_cache_retained_after_backward + but compares bf16 vs bf16 instead of FP8 vs bf16. The excess should be + zero, proving the comparison methodology works. + """ + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # Build and measure first bf16 model (acts as "baseline") + model_a = _build_model(NUM_LAYERS, fp8_init=False) + model_a = _shard_model(model_a, world_size) + opt_a = te.optimizers.FusedAdam( + model_a.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(model_a, opt_a, None, x, target) + delta_a = _measure_backward_memory_delta(model_a, opt_a, None, x, target) + + del model_a, opt_a + gc.collect() + torch.cuda.empty_cache() + + # Build and measure second bf16 model (acts as "test") + model_b = _build_model(NUM_LAYERS, fp8_init=False) + model_b = _shard_model(model_b, world_size) + opt_b = te.optimizers.FusedAdam( + model_b.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(model_b, opt_b, None, x, target) + delta_b = _measure_backward_memory_delta(model_b, opt_b, None, x, target) + + excess = delta_b - delta_a + tolerance = 256 * 1024 # 256 KiB + + assert abs(excess) <= tolerance, ( + "Two identical bf16 models show backward delta excess of " + f"{excess/1024:.1f} KiB (tolerance: {tolerance/1024:.0f} KiB). " + f"delta_a={delta_a/1024**2:.2f} MiB, delta_b={delta_b/1024**2:.2f} MiB." + ) + + +@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. + + When FP8 backward runs, _create_transpose allocates tensors for transposed + weight copies. These should be freed when backward completes, but instead + they persist until the next forward pass. This test measures the backward + memory delta (post_bwd - post_fwd) and compares it to a bf16 baseline. + In bf16, backward frees activations and adds gradients (net negative delta). + With FP8, retained transpose caches make the delta significantly more positive. + """ + _maybe_skip(recipe_name, quantized_model_init) + + recipe = get_recipe_from_string(recipe_name) + world_size, device = _get_dist_info() + + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # ── bf16 baseline ── + bf16_model = _build_model(NUM_LAYERS, fp8_init=False) + bf16_model = _shard_model(bf16_model, world_size) + bf16_optimizer = te.optimizers.FusedAdam( + bf16_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(bf16_model, bf16_optimizer, None, x, target) + bf16_bwd_delta = _measure_backward_memory_delta( + bf16_model, + bf16_optimizer, + None, + x, + target, + ) + + del bf16_model, bf16_optimizer + gc.collect() + torch.cuda.empty_cache() + + # ── FP8 model ── + fp8_model = _build_model(NUM_LAYERS, fp8_init=quantized_model_init, recipe=recipe) + fp8_model = _shard_model(fp8_model, world_size) + fp8_optimizer = te.optimizers.FusedAdam( + fp8_model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + for _ in range(WARMUP_STEPS): + _run_training_step(fp8_model, fp8_optimizer, recipe, x, target) + fp8_bwd_delta = _measure_backward_memory_delta( + fp8_model, + fp8_optimizer, + recipe, + x, + target, + ) + + # ── Assert: FP8 backward should not retain excess memory ── + # In bf16, backward frees activations and adds gradients (typically net negative). + # If FP8 transpose caches persist after backward, the FP8 delta will be + # 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 + + assert excess <= tolerance, ( + f"FP8 backward retains {excess/1024**2:.2f} MiB more than bf16 baseline. " + f"bf16 backward delta: {bf16_bwd_delta/1024**2:.2f} MiB, " + f"FP8 backward delta: {fp8_bwd_delta/1024**2:.2f} MiB. " + "Transpose caches from backward are likely not being freed (Issue #2717)." + ) + + +# ── Standalone runner ──────────────────────────────────────────────── +TESTS = { + "bf16_no_excess_forward_memory": test_bf16_no_excess_forward_memory, + "bf16_no_excess_backward_memory": test_bf16_no_excess_backward_memory, + "fp8_temp_accumulation_across_layers": test_fp8_temp_accumulation_across_layers, + "transpose_cache_retained_after_backward": test_transpose_cache_retained_after_backward, +} + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="FSDP2 memory leak tests (standalone)") + parser.add_argument("--test", required=True, choices=list(TESTS.keys())) + parser.add_argument( + "--recipe", + type=str, + default="DelayedScaling", + choices=[ + "DelayedScaling", + "Float8CurrentScaling", + "Float8BlockScaling", + "MXFP8BlockScaling", + "NVFP4BlockScaling", + ], + ) + parser.add_argument("--quantized-model-init", action="store_true", default=False) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + dist.init_process_group(backend="cpu:gloo,cuda:nccl") + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + _PARAMETRIZED_TESTS = { + "fp8_temp_accumulation_across_layers", + "transpose_cache_retained_after_backward", + } + + try: + test_fn = TESTS[args.test] + if args.test in _PARAMETRIZED_TESTS: + test_fn(args.recipe, args.quantized_model_init) + else: + test_fn() + finally: + if dist.is_initialized(): + dist.destroy_process_group() + gc.collect() + torch.cuda.empty_cache() diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index aca8d6d692..9cbbc3933c 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -62,6 +62,30 @@ def test_fsdp2_fused_adam_tests(): assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" +@pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +def test_fsdp2_mem_leak_tests(): + """FSDP2 memory leak detection tests (parametrized internally by recipe, quantized_model_init).""" + test_path = _FSDP2_DIR / "run_fsdp2_mem_leak.py" + nproc = min(NUM_PROCS, 2) + result = subprocess.run( + [ + "torchrun", + f"--nproc_per_node={nproc}", + "--local-ranks-filter=0", + "-m", + "pytest", + str(test_path), + "-v", + "-s", + "--tb=short", + ], + env=os.environ, + timeout=600, + ) + assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" + + def test_dummy() -> None: """Dummy test From 85f5a844f1aefafeb99940463884f31319dc6253 Mon Sep 17 00:00:00 2001 From: cael-ling Date: Sat, 4 Apr 2026 02:32:24 +0800 Subject: [PATCH 311/521] Refactor Amax Kernel ldmatrix loads, TMA/compute barriers, swizzle_idx (#2820) * Compute swizzle_idx once per thread and pass into ComputeKernel. Signed-off-by: Cael Ling * one __syncthreads per stage in GroupHadamardAmaxTmaKernel Signed-off-by: Cael Ling * streamline group Hadamard ComputeKernel loads Signed-off-by: Cael Ling * streamline group Hadamard ComputeKernel loads Signed-off-by: Cael Ling * streamline group Hadamard ComputeKernel loads Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * one __syncthreads per stage in GroupHadamardAmaxTmaKernel Signed-off-by: Cael Ling Made-with: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Compute swizzle_idx once per thread and pass into ComputeKernel. Signed-off-by: Cael Ling * Fix kReturnIdentityAmax path Signed-off-by: Cael Ling * Fix kReturnIdentityAmax path Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply the change to other variants Signed-off-by: Cael Ling * Refactor the change to other variants Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor the change to other variants Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor the ldmatrix logics Signed-off-by: Cael Ling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Cael Ling Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../graph_safe_group_hadamard_transform.cu | 42 +++++++++---------- .../group_hadamard_transform.cu | 41 +++++++++--------- .../hadamard_transform/hadamard_transform.cu | 42 +++++++++---------- 3 files changed, 57 insertions(+), 68 deletions(-) diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu index 0fb73cc439..2316d9697a 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -58,19 +58,13 @@ __device__ __forceinline__ size_t get_current_tensor_id( template __device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4], - IType* in_sh_ptr, uint32_t& local_pre_rht_amax_reg, + IType* in_sh_ptr, int swizzle_idx, + uint32_t& local_pre_rht_amax_reg, uint32_t& local_amax_reg, uint32_t& local_amax_t_reg) { uint32_t a_frag[4]; // A matrix fragment uint32_t c_frag[4]; // Result fragment - int warp_id = threadIdx.x / kThreadsPerWarp; - int local_rank = (threadIdx.x % kThreadsPerWarp); - - int ld_row_idx = local_rank % kHadamardDimension; - int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; - int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); - uint32_t temp_amax_reg; uint32_t temp_amax_t_reg; @@ -87,18 +81,16 @@ __device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_f } if (kReturnTransposedAmax) { - // TODO(Frank): This is not efficient, since we could directly load the - // matrix in transposed layout. if (!kReturnIdentityAmax) { - ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], - reinterpret_cast(in_sh_ptr) + swizzle_idx); + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } else { + matrix_transpose_m8_n8_b16_inplace(a_frag[0]); + matrix_transpose_m8_n8_b16_inplace(a_frag[1]); + matrix_transpose_m8_n8_b16_inplace(a_frag[2]); + matrix_transpose_m8_n8_b16_inplace(a_frag[3]); } - matrix_transpose_m8_n8_b16_inplace(a_frag[0]); - matrix_transpose_m8_n8_b16_inplace(a_frag[1]); - matrix_transpose_m8_n8_b16_inplace(a_frag[2]); - matrix_transpose_m8_n8_b16_inplace(a_frag[3]); - mma_m16_n16_k16_b16_b16_b16_noacc( a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg); @@ -315,6 +307,12 @@ __global__ void GraphSafeGroupHadamardAmaxTmaKernel( uint32_t local_amax_reg = *reinterpret_cast(&local_amax); uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t); + const int warp_id = threadIdx.x / kThreadsPerWarp; + const int local_rank = threadIdx.x % kThreadsPerWarp; + const int ld_row_idx = local_rank % kHadamardDimension; + const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; + const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); + for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) { for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) { int stage = STAGES_X * stage_y + stage_x; @@ -357,14 +355,12 @@ __global__ void GraphSafeGroupHadamardAmaxTmaKernel( had_frag_i, had_frag_t, in_sh_ptr + in_row_offset + (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)), - local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); + swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); } - - // Ensure all threads have finished their computation before new data over-writes the shared - // memory. - __syncthreads(); } - + // Ensure all threads have finished their computation before new data over-writes the shared + // memory. + __syncthreads(); // Ensure generic shared-memory accesses are visible before the next TMA write. ptx::fence_proxy_async_shared_cta(); } diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu index 07813be059..24d06e5d25 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform.cu @@ -41,19 +41,13 @@ constexpr int kThreadsPerWarp = 32; template __device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4], - IType* in_sh_ptr, uint32_t& local_pre_rht_amax_reg, + IType* in_sh_ptr, int swizzle_idx, + uint32_t& local_pre_rht_amax_reg, uint32_t& local_amax_reg, uint32_t& local_amax_t_reg) { uint32_t a_frag[4]; // A matrix fragment uint32_t c_frag[4]; // Result fragment - int warp_id = threadIdx.x / kThreadsPerWarp; - int local_rank = (threadIdx.x % kThreadsPerWarp); - - int ld_row_idx = local_rank % kHadamardDimension; - int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; - int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); - uint32_t temp_amax_reg; uint32_t temp_amax_t_reg; @@ -70,18 +64,16 @@ __device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_f } if (kReturnTransposedAmax) { - // TODO(Frank): This is not efficient, since we could directly load the - // matrix in transposed layout. if (!kReturnIdentityAmax) { - ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], - reinterpret_cast(in_sh_ptr) + swizzle_idx); + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } else { + matrix_transpose_m8_n8_b16_inplace(a_frag[0]); + matrix_transpose_m8_n8_b16_inplace(a_frag[1]); + matrix_transpose_m8_n8_b16_inplace(a_frag[2]); + matrix_transpose_m8_n8_b16_inplace(a_frag[3]); } - matrix_transpose_m8_n8_b16_inplace(a_frag[0]); - matrix_transpose_m8_n8_b16_inplace(a_frag[1]); - matrix_transpose_m8_n8_b16_inplace(a_frag[2]); - matrix_transpose_m8_n8_b16_inplace(a_frag[3]); - mma_m16_n16_k16_b16_b16_b16_noacc( a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg); @@ -305,6 +297,12 @@ __global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap t uint32_t local_amax_reg = *reinterpret_cast(&local_amax); uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t); + const int warp_id = threadIdx.x / kThreadsPerWarp; + const int local_rank = threadIdx.x % kThreadsPerWarp; + const int ld_row_idx = local_rank % kHadamardDimension; + const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; + const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); + for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) { for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) { int stage = STAGES_X * stage_y + stage_x; @@ -347,13 +345,12 @@ __global__ void GroupHadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap t had_frag_i, had_frag_t, in_sh_ptr + in_row_offset + (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)), - local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); + swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); } - - // Ensure all threads have finished their computation before new data over-writes the shared - // memory. - __syncthreads(); } + // Ensure all threads have finished their computation before new data over-writes the shared + // memory. + __syncthreads(); // Ensure generic shared-memory accesses are visible before the next TMA write. ptx::fence_proxy_async_shared_cta(); diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform.cu b/transformer_engine/common/hadamard_transform/hadamard_transform.cu index 4adc836886..b5160cd317 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform.cu @@ -26,19 +26,13 @@ constexpr int kThreadsPerWarp = 32; template __device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_frag_t[4], - IType* in_sh_ptr, uint32_t& local_pre_rht_amax_reg, + IType* in_sh_ptr, int swizzle_idx, + uint32_t& local_pre_rht_amax_reg, uint32_t& local_amax_reg, uint32_t& local_amax_t_reg) { uint32_t a_frag[4]; // A matrix fragment uint32_t c_frag[4]; // Result fragment - int warp_id = threadIdx.x / kThreadsPerWarp; - int local_rank = (threadIdx.x % kThreadsPerWarp); - - int ld_row_idx = local_rank % kHadamardDimension; - int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; - int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); - uint32_t temp_amax_reg; uint32_t temp_amax_t_reg; @@ -55,18 +49,16 @@ __device__ __forceinline__ void ComputeKernel(uint32_t b_frag_i[4], uint32_t b_f } if (kReturnTransposedAmax) { - // TODO(Frank): This is not efficient, since we could directly load the - // matrix in transposed layout. if (!kReturnIdentityAmax) { - ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], - reinterpret_cast(in_sh_ptr) + swizzle_idx); + ldmatrix_x4_m8n8_shared_b16(a_frag[0], a_frag[1], a_frag[2], a_frag[3], + reinterpret_cast(in_sh_ptr) + swizzle_idx); + } else { + matrix_transpose_m8_n8_b16_inplace(a_frag[0]); + matrix_transpose_m8_n8_b16_inplace(a_frag[1]); + matrix_transpose_m8_n8_b16_inplace(a_frag[2]); + matrix_transpose_m8_n8_b16_inplace(a_frag[3]); } - matrix_transpose_m8_n8_b16_inplace(a_frag[0]); - matrix_transpose_m8_n8_b16_inplace(a_frag[1]); - matrix_transpose_m8_n8_b16_inplace(a_frag[2]); - matrix_transpose_m8_n8_b16_inplace(a_frag[3]); - mma_m16_n16_k16_b16_b16_b16_noacc( a_frag[0], a_frag[2], a_frag[1], a_frag[3], b_frag_t[0], b_frag_t[1], b_frag_t[2], b_frag_t[3], c_frag[0], c_frag[1], c_frag[2], c_frag[3], temp_amax_t_reg); @@ -248,6 +240,12 @@ __global__ void HadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor uint32_t local_amax_reg = *reinterpret_cast(&local_amax); uint32_t local_amax_t_reg = *reinterpret_cast(&local_amax_t); + const int warp_id = threadIdx.x / kThreadsPerWarp; + const int local_rank = threadIdx.x % kThreadsPerWarp; + const int ld_row_idx = local_rank % kHadamardDimension; + const int ld_col_idx = local_rank / kHadamardDimension + warp_id * 2; + const int swizzle_idx = swizzle_128B_atom_32B(ld_row_idx, ld_col_idx); + for (int stage_y = 0; stage_y < STAGES_Y; ++stage_y) { for (int stage_x = 0; stage_x < STAGES_X; ++stage_x) { int stage = STAGES_X * stage_y + stage_x; @@ -290,14 +288,12 @@ __global__ void HadamardAmaxTmaKernel(const __grid_constant__ CUtensorMap tensor had_frag_i, had_frag_t, in_sh_ptr + in_row_offset + (compute_stage_x * kHadamardDimension * (THREADS_PER_CHUNK / kThreadsPerWarp)), - local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); + swizzle_idx, local_pre_rht_amax_reg, local_amax_reg, local_amax_t_reg); } - - // Ensure all threads have finished their computation before new data over-writes the shared - // memory. - __syncthreads(); } - + // Ensure all threads have finished their computation before new data over-writes the shared + // memory. + __syncthreads(); // Ensure generic shared-memory accesses are visible before the next TMA write. ptx::fence_proxy_async_shared_cta(); } From a88fdc1b8139c86e9e03507a4feba652cee0aa5f Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 3 Apr 2026 11:33:40 -0700 Subject: [PATCH 312/521] =?UTF-8?q?[PyTorch]=20[CI]=20Capture=20subprocess?= =?UTF-8?q?=20stderr=20in=20distributed=20tests=20for=20better=20CI=20erro?= =?UTF-8?q?r=20re=E2=80=A6=20(#2802)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Capture subprocess stderr in distributed tests for better CI error reporting Distributed tests launch subprocesses via torch.distributed.launch/torchrun. When these fail, pytest only captures the CalledProcessError from the parent process, not the actual worker traceback. This makes CI JUnit XML reports show "exit code 1" with no useful error detail. Add run_distributed() utility to tests/pytorch/utils.py that captures stderr while letting stdout stream to the terminal. On failure, the worker's stderr (containing the actual Python traceback) is included in the AssertionError, which pytest writes into the JUnit XML report. Behavior: - Interactive use: stdout streams in real time (unchanged), stderr shown on failure - CI/JUnit XML: failure reports now include the actual worker traceback Signed-off-by: Sudhakar Singh * Add JUnit XML output to ctest in L0_cppunittest Add --output-junit flag so ctest writes JUnit XML to /logs/, matching the pattern used by pytest tests. The XML is written before ctest exits, so it's captured even on test failure. Signed-off-by: Sudhakar Singh --------- Signed-off-by: Sudhakar Singh Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- qa/L0_cppunittest/test.sh | 5 ++- .../attention/test_attention_with_cp.py | 8 ++--- .../test_cast_master_weights_to_fp8.py | 5 ++- .../test_fusible_ops_with_userbuffers.py | 4 +-- tests/pytorch/distributed/test_torch_fsdp2.py | 12 ++++--- tests/pytorch/utils.py | 34 ++++++++++++++++++- 6 files changed, 54 insertions(+), 14 deletions(-) diff --git a/qa/L0_cppunittest/test.sh b/qa/L0_cppunittest/test.sh index 0b83747c0e..c7499282f4 100755 --- a/qa/L0_cppunittest/test.sh +++ b/qa/L0_cppunittest/test.sh @@ -4,6 +4,9 @@ set -e +: ${XML_LOG_DIR:=/logs} +mkdir -p "$XML_LOG_DIR" + # Find TE : ${TE_PATH:=/opt/transformerengine} TE_LIB_PATH=$(pip3 show transformer-engine | grep -E "Location:|Editable project location:" | tail -n 1 | awk '{print $NF}') @@ -17,4 +20,4 @@ cd $TE_PATH/tests/cpp cmake -GNinja -Bbuild . cmake --build build export OMP_NUM_THREADS=$((NUM_PHYSICAL_CORES / NUM_PARALLEL_JOBS)) -ctest --test-dir build -j$NUM_PARALLEL_JOBS +ctest --test-dir build -j$NUM_PARALLEL_JOBS --output-junit $XML_LOG_DIR/ctest_cppunittest.xml diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index ecd0090a3b..5aaf67061b 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -22,7 +22,7 @@ _current_file = pathlib.Path(__file__).resolve() sys.path.append(str(_current_file.parent.parent)) -from utils import ModelConfig, get_available_attention_backends +from utils import ModelConfig, get_available_attention_backends, run_distributed pytest_logging_level = logging.getLevelName(logging.root.level) @@ -125,7 +125,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): if not flash_attn_supported: pytest.skip("No attention backend available.") - subprocess.run( + run_distributed( get_bash_arguments( num_gpus_per_node=num_gpus, dtype=dtype, @@ -135,7 +135,6 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): cp_comm_type=cp_comm_type, log_level=pytest_logging_level, ), - check=True, ) @@ -368,7 +367,7 @@ def test_cp_with_fused_attention( if not fused_attn_supported: pytest.skip("No attention backend available.") - subprocess.run( + run_distributed( get_bash_arguments( num_gpus_per_node=num_gpus, dtype=dtype, @@ -384,5 +383,4 @@ def test_cp_with_fused_attention( is_training=is_training, log_level=pytest_logging_level, ), - check=True, ) diff --git a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py index 1606641b78..7de6142537 100644 --- a/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py +++ b/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py @@ -10,6 +10,9 @@ import sys import pathlib +sys.path.append(str(pathlib.Path(__file__).resolve().parent.parent)) +from utils import run_distributed + import pytest import torch from torch import nn @@ -1207,7 +1210,7 @@ def test_nvfp4_partial_cast_matches_full(world_size: int) -> None: current_file, "--parallel-nvfp4-partial", ] - subprocess.run(command, check=True) + run_distributed(command) def test_single_gpu_partial_cast_vs_full(): diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py index 603433e0da..3dcefd46fd 100644 --- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py +++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py @@ -38,7 +38,7 @@ # Import utility functions _current_file = pathlib.Path(__file__).resolve() sys.path.append(str(_current_file.parent.parent)) -from utils import dtype_tols, make_recipe, str_to_dtype +from utils import dtype_tols, make_recipe, run_distributed, str_to_dtype # Check if FP8 is supported fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) @@ -463,7 +463,7 @@ def test_fuser_ops_with_userbuffers( env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" # Launch parallel job - result = subprocess.run(command, check=True, env=env) + run_distributed(command, env=env) def main() -> None: diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index 9cbbc3933c..ee20886631 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -3,9 +3,13 @@ # See LICENSE for license information. import os +import sys import subprocess from pathlib import Path +sys.path.append(str(Path(__file__).resolve().parent.parent)) +from utils import run_distributed + import pytest import torch @@ -20,7 +24,7 @@ def test_fsdp2_model_tests(): """All FSDP2 model tests (parametrized internally by recipe, fp8_init, sharding, layer).""" test_path = _FSDP2_DIR / "run_fsdp2_model.py" - result = subprocess.run( + run_distributed( [ "torchrun", f"--nproc_per_node={NUM_PROCS}", @@ -32,10 +36,10 @@ def test_fsdp2_model_tests(): "-s", "--tb=short", ], + valid_returncodes=(0, 5), env=os.environ, timeout=600, ) - assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") @@ -44,7 +48,7 @@ def test_fsdp2_fused_adam_tests(): """All FSDP2 FusedAdam tests (parametrized internally by recipe, test variant).""" test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" nproc = min(NUM_PROCS, 2) - result = subprocess.run( + run_distributed( [ "torchrun", f"--nproc_per_node={nproc}", @@ -56,10 +60,10 @@ def test_fsdp2_fused_adam_tests(): "-s", "--tb=short", ], + valid_returncodes=(0, 5), env=os.environ, timeout=600, ) - assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" @pytest.mark.skipif(NUM_PROCS < 2, reason="Requires 2+ GPUs") diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 317240fb78..929f02453d 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -6,8 +6,9 @@ import logging import os +import subprocess from contextlib import contextmanager -from typing import Optional, Tuple, Dict, Any, List +from typing import Optional, Sequence, Tuple, Dict, Any, List from packaging.version import Version as PkgVersion import torch @@ -407,3 +408,34 @@ def assert_close_grads( assert actual is not None assert expected is not None assert_close(actual.grad, expected.grad, **kwargs) + + +def run_distributed( + args: Sequence[str], + *, + valid_returncodes: Sequence[int] = (0,), + **kwargs, +) -> subprocess.CompletedProcess: + """Run a distributed subprocess with stderr capture for better error reporting. + + stdout streams to the terminal in real time for interactive debugging. + On failure, stderr (containing Python tracebacks) is included in the + AssertionError so pytest writes it into the JUnit XML report. + + Args: + args: Command and arguments to run. + valid_returncodes: Return codes considered success (default: (0,)). + Use (0, 5) for inner pytest runs where 5 means all tests skipped. + **kwargs: Passed through to subprocess.run (e.g. env, timeout). + """ + result = subprocess.run(args, stderr=subprocess.PIPE, text=True, **kwargs) + if result.returncode not in valid_returncodes: + cmd_str = " ".join(str(a) for a in args) + msg = f"Command exited with code {result.returncode}:\n {cmd_str}\n" + if result.stderr: + stderr_tail = result.stderr[-4000:] + if len(result.stderr) > 4000: + stderr_tail = "... [truncated] ...\n" + stderr_tail + msg += f"\n--- stderr ---\n{stderr_tail}" + raise AssertionError(msg) + return result From 509614d8effc45be08f34eeae7cedeee1c1923ae Mon Sep 17 00:00:00 2001 From: int-smart Date: Fri, 3 Apr 2026 14:53:07 -0700 Subject: [PATCH 313/521] Feature/unswizzle (#2732) * Add unswizzling functions for scaling factors in swizzle module - Introduced `nvte_unswizzle_scaling_factors` to convert swizzled scaling factors back to row-major format. - Implemented `regs_unshuffle_with_bit_shifts` and `regs_unshuffle` for unshuffling operations in CUDA kernels. - Added `unswizzle_row_scaling_kernel_impl` and `unswizzle_col_scaling_kernel_impl` for handling unswizzling in row and column scaling respectively. These changes enhance the functionality of the swizzle module, enabling better handling of scaling factors in tensor operations. Signed-off-by: Abhishek * Add swizzle/unswizzle roundtrip test for scaling factors These enhancements tests the changes introduced for unswizzling Signed-off-by: Abhishek * Added another unswizzling functionality test for scaling factors - Introduced `compute_ref_unswizzle` to handle the conversion of swizzled scaling factors back to their original format. - Added `performTestUnswizzle1D` to validate the unswizzling process with various scaling modes. - Created `UnswizzleTestSuite` for comprehensive testing of unswizzling operations. Signed-off-by: Abhishek * Moved swizzle_row_scaling_kernel implementation at its original place - Moved the definition of `swizzle_row_scaling_kernel` to a new location for better organization. - Ensured the kernel implementation is now properly defined and accessible for scaling operations in the swizzle module. Signed-off-by: Abhishek * Add multi-tensor unswizzling functions for scaling factors - Introduced `multi_tensor_unswizzle_scaling_factors` to convert swizzled scaling factors back to their original row-major format. - Implemented CUDA kernels for unswizzling in both row and column scaling, enhancing the swizzle module's functionality. - Updated the launch function to handle multiple tensor unswizzling operations efficiently. These changes improve the handling of scaling factors in tensor operations, ensuring better performance and organization within the swizzle module. Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added greptile suggestions Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Removed unused check from tests and reading input directly as const rather than casting Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor unswizzling functions and update test cases for scaling factors - Updated unswizzling kernel implementations to remove original_M and original_K parameters, simplifying the function signatures. - Enhanced test suite to utilize new unswizzling data shapes, ensuring comprehensive coverage of aligned and padded cases. These changes improve the clarity and efficiency of the unswizzling process in the swizzle module. Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Refactor unswizzling scaling factors to use a launch function Signed-off-by: Abhishek * Change unswizzling to use output as gt. Signed-off-by: Abhishek * Refactor unswizzling scaling factors to improve input validation and streamline processing. Need to check if rowwise and columnwise both can be true. If yes the if else needs to account for that Signed-off-by: Abhishek * Fix multi_tensor_unswizzle_scaling_factors to correctly reference output tensors for scaling mode and data validation. Updated checks for input and output tensor shapes to ensure proper handling of row-wise and column-wise scaling factors. Signed-off-by: Abhishek * Enhance swizzle tests and unswizzling validation Signed-off-by: Abhishek * Fix typos and update validation checks in swizzle.cu Signed-off-by: Abhishek * Update validation checks in multi_tensor_unswizzle_scaling_factors to use input numel Signed-off-by: Abhishek * Typo 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 Signed-off-by: Przemek Tredak Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemek Tredak --- tests/cpp/operator/test_swizzle.cu | 249 +++++++ .../include/transformer_engine/swizzle.h | 32 +- transformer_engine/common/swizzle/swizzle.cu | 658 ++++++++++++++++++ 3 files changed, 937 insertions(+), 2 deletions(-) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 8389989efe..7dfb34201d 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -56,6 +56,35 @@ void compute_ref_swizzle(const uint8_t *h_input, uint8_t *h_output, } } +template +void compute_ref_unswizzle(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 (int m = 0; m < M; m++) { + for (int 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_input_ptr = tile_id_m * SF_TILE_DIM_M * K + tile_id_k * SF_TILE_SIZE; + int in_index = tile_input_ptr + row_in_new_tile * NEW_SF_TILE_DIM_K + col_in_new_tile; + if constexpr(row_scaling) + h_output[k + m * K] = h_input[in_index]; + else + h_output[k * M + m] = h_input[in_index]; + } + } +} + void performTestSwizzle1D(const int num_tiles_M, const int num_tiles_K, bool rowwise, bool columnwise, const bool transa) { using namespace test; @@ -110,6 +139,66 @@ void performTestSwizzle1D(const int num_tiles_M, const int num_tiles_K, bool row } } +void performTestUnswizzle1D(const size_t M, const size_t K, bool rowwise, bool columnwise, const bool transa) { + using namespace test; + + int SF_MODE_X, SF_MODE_Y; + if (rowwise) { + SF_MODE_X = 1; + SF_MODE_Y = 32; + } + if (columnwise) { + SF_MODE_X = 32; + SF_MODE_Y = 1; + } + + if (!rowwise && !columnwise) { + GTEST_SKIP() << "TEST SKIPPED, Either rowwise or columnwise scaling mode must be true."; + } + if (rowwise && columnwise) { + GTEST_SKIP() << "TEST SKIPPED, The scaling mode " + std::to_string(SF_MODE_X) + "x" + + std::to_string(SF_MODE_Y) + " is not implemented."; + } + + DType dtype = DType::kFloat8E4M3; + + const auto data_shape = transa ? std::vector{M, K} : std::vector{K, M}; + + Tensor input("input", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); + input.set_with_gemm_swizzled_scales(true); + Tensor output("output", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); + + fillUniform(&input); + + // Use the actual padded compact scale shape from the tensor for both the reference + // and the comparison. This correctly covers padded cases where M is not a multiple + // of 128 or K/32 is not a multiple of 4. + const auto padded_scale_shape = rowwise + ? input.rowwise_scale_inv_shape() + : input.columnwise_scale_inv_shape(); + const size_t padded_dim0 = padded_scale_shape.data[0]; + const size_t padded_dim1 = padded_scale_shape.data[1]; + std::unique_ptr ref_output = std::make_unique(padded_dim0 * padded_dim1); + + nvte_unswizzle_scaling_factors(input.data(), output.data(), 0); + + if (rowwise) + compute_ref_unswizzle<128, 4, true>(input.rowwise_cpu_scale_inv_ptr(), ref_output.get(), padded_dim0, padded_dim1); + else + compute_ref_unswizzle<128, 4, false>(input.columnwise_cpu_scale_inv_ptr(), ref_output.get(), padded_dim1, padded_dim0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + output.to_cpu(); + if (rowwise) { + compareResults("output_unswizzle", output.rowwise_cpu_scale_inv_ptr(), ref_output.get(), padded_dim0 * padded_dim1); + } else { + compareResults("output_unswizzle", output.columnwise_cpu_scale_inv_ptr(), ref_output.get(), padded_dim0 * padded_dim1); + } +} + // Zero out padding in a scale_inv CPU buffer so that the CPU reference // matches the kernel, which zeroes elements outside the original dims. // The buffer is stored in leading-dim-major order (row-major for rowwise, @@ -235,6 +324,21 @@ TEST_P(SwizzleTestSuite, TestSwizzle) { transa); } +class UnswizzleTestSuite : public ::testing::TestWithParam, std::pair, bool>> {}; + +TEST_P(UnswizzleTestSuite, TestUnswizzle) { + using namespace transformer_engine; + using namespace test; + + const auto data_shape = std::get<0>(GetParam()); + const auto scaling_mode = std::get<1>(GetParam()); + const auto transa = std::get<2>(GetParam()); + + performTestUnswizzle1D(data_shape.first, data_shape.second, + scaling_mode.first, scaling_mode.second, + transa); +} + class SwizzleGroupedTestSuite : public ::testing::TestWithParam> {}; @@ -282,6 +386,24 @@ std::vector> num_tiles = { {65, 259}, }; +// Raw {M, K} data shapes for unswizzle tests. Includes aligned cases (scale dims +// already multiples of 128 and 4) and padded cases where M or K/32 are not yet +// aligned, forcing the compact scale_inv to carry a padded tail. +// All K values must be multiples of 32 (MXFP8 block size). +std::vector> unswizzle_data_shapes = { + // Aligned: scale dims are already multiples of 128 and 4 + {128, 128}, + {128, 16896}, // K = 132 * 128, large K + {16896, 128}, // M = 132 * 128, large M + // M-padding only: M not a multiple of 128 (scale-M needs padding to 256) + {160, 128}, + // scale-K padding only: K/32 = 3, padded to 4 + {128, 96}, + // Both M and scale-K need padding + {160, 96}, + {16896, 16896}, +}; + std::vector> scaling_mode = { {true, false}, {false, true} @@ -308,3 +430,130 @@ INSTANTIATE_TEST_SUITE_P( std::to_string(std::get<2>(info.param)); return name; }); + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + UnswizzleTestSuite, + ::testing::Combine( + ::testing::ValuesIn(unswizzle_data_shapes), + ::testing::ValuesIn(scaling_mode), + ::testing::ValuesIn(transa) + ), + [](const testing::TestParamInfo& info) { + std::string name = "MK" + + std::to_string(std::get<0>(info.param).first) + "X" + + std::to_string(std::get<0>(info.param).second) + "smode" + + std::to_string(std::get<1>(info.param).first) + "X"+ + std::to_string(std::get<1>(info.param).second) + "trans" + + std::to_string(std::get<2>(info.param)); + return name; + }); + +void performTestSwizzleUnswizzleRoundtrip(const size_t M, const size_t K, bool rowwise, bool columnwise, const bool transa) { + using namespace test; + + int SF_MODE_X, SF_MODE_Y; + if (rowwise) { + SF_MODE_X = 1; + SF_MODE_Y = 32; + } + if (columnwise) { + SF_MODE_X = 32; + SF_MODE_Y = 1; + } + + if (!rowwise && !columnwise) { + GTEST_SKIP() << "TEST SKIPPED, Either rowwise or columnwise scaling mode must be true."; + } + if (rowwise && columnwise){ + GTEST_SKIP() << "TEST SKIPPED, The scaling mode " + std::to_string(SF_MODE_X) + "x" + + std::to_string(SF_MODE_Y) + " is not implemented."; + } + + DType dtype = DType::kFloat8E4M3; + + const auto data_shape = transa ? std::vector{M, K} : std::vector{K, M}; + const size_t logical_dim0 = data_shape[0] / SF_MODE_X; + const size_t logical_dim1 = data_shape[1] / SF_MODE_Y; + + Tensor input("input", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); + Tensor swizzled("swizzled", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); + swizzled.set_with_gemm_swizzled_scales(true); + Tensor output("output", data_shape, dtype, rowwise, columnwise, NVTE_MXFP8_1D_SCALING); + + fillUniform(&input); + + // fillUniform fills all scale_inv entries including the padded region with random bytes. + // After swizzle, the swizzle kernel zeroes padded positions in the swizzled output, so + // after unswizzle those positions come back as zero in the compact output. Zero them in + // the input now so the full-buffer comparison is valid. + const auto padded_scale_shape = rowwise + ? input.rowwise_scale_inv_shape() + : input.columnwise_scale_inv_shape(); + const size_t padded_dim0 = padded_scale_shape.data[0]; + const size_t padded_dim1 = padded_scale_shape.data[1]; + + if (padded_dim0 != logical_dim0 || padded_dim1 != logical_dim1) { + auto* scale_ptr = rowwise + ? input.rowwise_cpu_scale_inv_ptr() + : input.columnwise_cpu_scale_inv_ptr(); + for (size_t r = 0; r < padded_dim0; r++) { + for (size_t c = 0; c < padded_dim1; c++) { + if (r >= logical_dim0 || c >= logical_dim1) { + scale_ptr[r * padded_dim1 + c] = 0; + } + } + } + input.from_cpu(); + } + + nvte_swizzle_scaling_factors(input.data(), swizzled.data(), 0); + nvte_unswizzle_scaling_factors(swizzled.data(), output.data(), 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + input.to_cpu(); + output.to_cpu(); + if (rowwise) { + compareResults("roundtrip_rowwise", output.rowwise_cpu_scale_inv_ptr(), + input.rowwise_cpu_scale_inv_ptr(), padded_dim0 * padded_dim1); + } else { + compareResults("roundtrip_columnwise", output.columnwise_cpu_scale_inv_ptr(), + input.columnwise_cpu_scale_inv_ptr(), padded_dim0 * padded_dim1); + } +} + +class SwizzleUnswizzleRoundtripTestSuite : public ::testing::TestWithParam, std::pair, bool>> {}; + +TEST_P(SwizzleUnswizzleRoundtripTestSuite, TestSwizzleUnswizzleRoundtrip) { + using namespace transformer_engine; + using namespace test; + + const auto data_shape = std::get<0>(GetParam()); + const auto scaling_mode = std::get<1>(GetParam()); + const auto transa = std::get<2>(GetParam()); + + performTestSwizzleUnswizzleRoundtrip(data_shape.first, data_shape.second, + scaling_mode.first, scaling_mode.second, + transa); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + SwizzleUnswizzleRoundtripTestSuite, + ::testing::Combine( + ::testing::ValuesIn(unswizzle_data_shapes), + ::testing::ValuesIn(scaling_mode), + ::testing::ValuesIn(transa) + ), + [](const testing::TestParamInfo& info) { + std::string name = "roundtrip_MK" + + std::to_string(std::get<0>(info.param).first) + "X" + + std::to_string(std::get<0>(info.param).second) + "smode" + + std::to_string(std::get<1>(info.param).first) + "X"+ + std::to_string(std::get<1>(info.param).second) + "trans" + + std::to_string(std::get<2>(info.param)); + return name; + }); diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 904812118c..aa697aafe1 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -26,7 +26,7 @@ extern "C" { * Requirements: * - scale_inv is stored in row-major. * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. - * - data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. + * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. */ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream); @@ -40,11 +40,39 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud * Requirements: * - scale_inv is stored in row-major. * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. - * - data is quantitized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. + * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. */ void nvte_multi_tensor_swizzle_scaling_factors(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. + * \param[in,out] output Output tensor which hosts non-swizzled scale_inv. + * \param[in] stream CUDA stream used for the operation. + * + * Requirements: + * - scale_inv is stored in row-major in output. + * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. + */ +void nvte_unswizzle_scaling_factors(const NVTETensor input, NVTETensor output, cudaStream_t stream); + +/*! \brief Unswizzling scaling factors from the interleaved layout used by GEMM back to row-major + * + * \param[in] inputs Input tensors with swizzled scale_inv. + * \param[in,out] outputs Output tensors which hosts non-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 in output. + * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. + */ +void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs, + const size_t num_tensors, cudaStream_t stream); + /*! \brief Swizzling FP8 block scaling scaling factors into mxfp8 interleaved layout for GEMM * * \param[in] input Input FP8 block-scaled tensor. diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 619987931e..28a879a376 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -54,6 +54,32 @@ __device__ inline void regs_shuffle_with_bit_shifts(LType* regs_vec) { for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i]; } +template +__device__ inline void regs_unshuffle_with_bit_shifts(LType* regs_vec) { + // Inverse of regs_shuffle_with_bit_shifts + // inp, 4-byte chunks [0,4,8,12, 1,5,9,13, 2,6,10,14, 3,7,11,15] + // out, swapping byte to form new 4-byte chunks [0,1,2,3, 4,5,6,7, 8,9,10,11, 12,13,14,15] + + constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); + constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD; + int32_t new_regs[kVectorSize]; + int32_t* regs = reinterpret_cast(regs_vec); + +#pragma unroll + for (int i = 0; i < N_TILE_PER_TD; i++) { +#pragma unroll + for (int j = 0; j < N_SF_PER_TD_PER_TILE; j++) { + new_regs[i + j * N_TILE_PER_TD] = + ((regs[i * N_SF_PER_TD_PER_TILE + 0] >> 8 * j) & 0xFF) | + (((regs[i * N_SF_PER_TD_PER_TILE + 1] >> 8 * j) & 0xFF) << 8) | + (((regs[i * N_SF_PER_TD_PER_TILE + 2] >> 8 * j) & 0xFF) << 16) | + (((regs[i * N_SF_PER_TD_PER_TILE + 3] >> 8 * j) & 0xFF) << 24); + } + } +#pragma unroll + for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i]; +} + template __device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int original_M, @@ -170,6 +196,23 @@ __device__ inline void regs_shuffle(LType* regs_vec) { for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i]; } +// Inverse of regs_shuffle. +template +__device__ inline void regs_unshuffle(LType* regs_vec) { + constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); + if constexpr (N_TILE_PER_TD == 1) return; + + constexpr int kVectorSize = N_SF_PER_TD_PER_TILE * N_TILE_PER_TD; + int32_t tmp[kVectorSize]; + int32_t* ptr = reinterpret_cast(regs_vec); +#pragma unroll + for (int i = 0; i < kVectorSize; i++) + tmp[i % N_SF_PER_TD_PER_TILE * N_TILE_PER_TD + i / N_SF_PER_TD_PER_TILE] = ptr[i]; + +#pragma unroll + for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i]; +} + template __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int original_M, @@ -239,6 +282,146 @@ __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, } } +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, + const int grid_dim_x, const int grid_dim_y) { + constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); + constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD; + + 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; + + int n_tiles_in_tb = N_TILES_IN_TB; + const int K_i32 = K / 4; + if (bid_x == grid_dim_x - 1) { + n_tiles_in_tb = (K_i32 - 1) % N_TILES_IN_TB + 1; + } + + const int input_offset = + bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB * SF_TILE_SIZE_I32; + const int* input_i32 = reinterpret_cast(input) + input_offset; + const int output_offset = bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB; + int* output_i32 = reinterpret_cast(output) + output_offset; + + extern __shared__ int4 slm_v4i[]; + + int linear_id = threadIdx.y * blockDim.x + threadIdx.x; + const int4* input_v4i = reinterpret_cast(input_i32); +#pragma unroll + for (int i = linear_id; i < SF_TILE_SIZE_I32 * n_tiles_in_tb / 4; i += blockDim.x * blockDim.y) { + slm_v4i[i] = input_v4i[i]; + } + __syncthreads(); + + 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_TILE_PER_TD; i++) { + reinterpret_cast(regs_vec)[i] = + slm_v4i[(threadIdx.x * N_TILE_PER_TD + i) * SF_TILE_SIZE_I32 / 4 + threadIdx.y]; + } + + regs_unshuffle(regs_vec); + +#pragma unroll + for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { + const int thread_offset = (i * TB_DIM + threadIdx.y) * K_i32 + threadIdx.x * N_TILE_PER_TD; + reinterpret_cast(output_i32 + thread_offset)[0] = regs_vec[i]; + } + } +} + +template +__device__ void unswizzle_col_scaling_kernel_impl(const void* input, void* output, const int M, + const int K, const int bid_x, const int bid_y, + const int grid_dim_x, const int grid_dim_y) { + constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); + constexpr int N_SF_PER_TD = N_TILE_PER_TD * N_SF_PER_TD_PER_TILE; + 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; + + int m_tiles_in_tb = N_TILE_PER_TD; + int k_tiles_in_tb = TB_DIM; + if (bid_x == grid_dim_x - 1) { + k_tiles_in_tb = (K_i32 / SF_TILE_DIM_K_I32 - 1) % k_tiles_in_tb + 1; + } + if (bid_y == grid_dim_y - 1) { + m_tiles_in_tb = (M_i32 / SF_TILE_DIM_M_I32 - 1) % m_tiles_in_tb + 1; + } + + const int32_t* input_i32[N_TILE_PER_TD]; +#pragma unroll + for (int i = 0; i < m_tiles_in_tb; i++) { + input_i32[i] = reinterpret_cast(input) + bid_x * TB_DIM * SF_TILE_SIZE_I32 + + (bid_y * N_TILE_PER_TD + i) * SF_TILE_DIM_M_I32 * K_i32; + } + const int output_offset = + bid_x * TB_DIM * SF_TILE_DIM_K_I32 * M_i32 + bid_y * N_TILE_PER_TD * SF_TILE_DIM_M_I32; + int* output_i32 = reinterpret_cast(output) + output_offset; + + extern __shared__ int slm[]; + + int linear_id = threadIdx.y * blockDim.x + threadIdx.x; +#pragma unroll + for (int i = 0; i < m_tiles_in_tb; i++) { + __align__(16) const int4* input_v4i = reinterpret_cast(input_i32[i]); + __align__(16) int4* slm_v4i = + reinterpret_cast(slm + i * k_tiles_in_tb * SF_TILE_SIZE_I32); +#pragma unroll + for (int j = linear_id; j < SF_TILE_SIZE_I32 * k_tiles_in_tb / 4; + j += blockDim.x * blockDim.y) { + slm_v4i[j] = input_v4i[j]; + } + } + __syncthreads(); + + 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) { + int tM = threadIdx.x * N_SF_PER_TD; + int* slm_tile = slm + (threadIdx.y * SF_TILE_SIZE_I32 + + tM / SF_TILE_DIM_M * k_tiles_in_tb * SF_TILE_SIZE_I32); +#pragma unroll + for (int i = 0; i < N_SF_PER_TD; i++) { + reinterpret_cast(regs_vec)[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_unshuffle_with_bit_shifts(regs_vec); + +#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; + reinterpret_cast(output_i32 + thread_offset)[0] = regs_vec[i]; + } + } +} + +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + unswizzle_scaling_kernel(const void* input, void* output, const int M, const int K, + const bool row_scaling) { + const int bid_x = blockIdx.x; + const int bid_y = blockIdx.y; + const int grid_dim_x = gridDim.x; + const int grid_dim_y = gridDim.y; + if (row_scaling) { + unswizzle_row_scaling_kernel_impl( + input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else { + unswizzle_col_scaling_kernel_impl( + input, output, M, 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, @@ -302,6 +485,59 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) gridDim.y); } +template +__global__ void multi_tensor_unswizzle_row_scaling_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]; + + constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); + constexpr int N_TILES_IN_TB = TB_DIM * N_TILE_PER_TD; + + const int num_tiles_k = K / SF_TILE_DIM_K; + const int num_tiles_m = M / SF_TILE_DIM_M; + const int flat_offset = bid - kernel_args.block_range[tensor_id]; + const int grid_dim_x = DIVUP(num_tiles_k, N_TILES_IN_TB); + const int grid_dim_y = num_tiles_m; + const int bid_x = flat_offset / grid_dim_y; + const int bid_y = flat_offset % grid_dim_y; + + unswizzle_row_scaling_kernel_impl( + input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y); +} + +template +__global__ void multi_tensor_unswizzle_col_scaling_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]; + + constexpr int N_TILE_PER_TD = sizeof(LType) / sizeof(int); + + const int num_tiles_k = K / SF_TILE_DIM_K; + const int num_tiles_m = M / SF_TILE_DIM_M; + const int flat_offset = bid - kernel_args.block_range[tensor_id]; + const int grid_dim_x = DIVUP(num_tiles_k, TB_DIM); + const int grid_dim_y = DIVUP(num_tiles_m, N_TILE_PER_TD); + const int bid_x = flat_offset / grid_dim_y; + const int bid_y = flat_offset % grid_dim_y; + + unswizzle_col_scaling_kernel_impl( + input, output, M, K, bid_x, bid_y, grid_dim_x, grid_dim_y); +} + template __global__ void multi_tensor_swizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) { // Find tensor corresponding to block @@ -681,6 +917,89 @@ void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, NVTE_CHECK_CUDA(cudaGetLastError()); } +template +void launch_multi_tensor_unswizzle_scaling_factors(MultiSwizzleArgs& kernel_args, + const int vec_load_size, const bool is_rowwise, + 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); + 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); + } + } + + int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + if (num_blocks > 0) { + dim3 block_size(TB_DIM, TB_DIM); + if (is_rowwise) { + switch (vec_load_size) { + case 4: + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + multi_tensor_unswizzle_row_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + multi_tensor_unswizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 2: + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + multi_tensor_unswizzle_row_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + multi_tensor_unswizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 1: + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + multi_tensor_unswizzle_row_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + multi_tensor_unswizzle_row_scaling_kernel + <<>>(kernel_args); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + break; + } + } else { + switch (vec_load_size) { + case 4: + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + multi_tensor_unswizzle_col_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + multi_tensor_unswizzle_col_scaling_kernel + <<>>(kernel_args); + break; + case 2: + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + multi_tensor_unswizzle_col_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + multi_tensor_unswizzle_col_scaling_kernel + <<>>(kernel_args); + break; + case 1: + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + multi_tensor_unswizzle_col_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + multi_tensor_unswizzle_col_scaling_kernel + <<>>(kernel_args); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + break; + } + } + NVTE_CHECK_CUDA(cudaGetLastError()); + } +} + void multi_tensor_swizzle_scaling_factors(const std::vector& input, std::vector& output, cudaStream_t stream) { auto num_tensors = input.size(); @@ -850,6 +1169,325 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args, vec_load_size, false, stream); } } + +void unswizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t stream) { + const auto& scaling_mode = output->scaling_mode; + NVTE_CHECK(scaling_mode == NVTE_MXFP8_1D_SCALING || scaling_mode == NVTE_NVFP4_1D_SCALING, + "Output tensor has invalid scaling mode (", to_string(output->scaling_mode), ")."); + + CheckInputTensor(*input, "scaling_factor_input"); + CheckInputTensor(*output, "scaling_factor_output"); + NVTE_CHECK(input->with_gemm_swizzled_scales, "Expected input tensor with swizzled scales."); + NVTE_CHECK(!output->with_gemm_swizzled_scales, + "Expected output tensor in row-major compact format."); + NVTE_CHECK(input->scaling_mode == scaling_mode, + "Input and output tensors must have matching scaling modes, but got ", + to_string(input->scaling_mode), " and ", to_string(output->scaling_mode), "."); + + const bool has_rowwise_scale_inv = output->scale_inv.has_data(); + const bool has_columnwise_scale_inv = output->columnwise_scale_inv.has_data(); + NVTE_CHECK(!has_rowwise_scale_inv || !has_columnwise_scale_inv, + "Output tensor has both row-wise and column-wise scaling factors"); + if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) { + return; + } + if (has_rowwise_scale_inv) { + NVTE_CHECK(input->scale_inv.has_data(), + "Output tensor requests row-wise scaling factors, but input tensor does not " + "provide them."); + } else if (has_columnwise_scale_inv) { + NVTE_CHECK(input->columnwise_scale_inv.has_data(), + "Output tensor requests column-wise scaling factors, but input tensor does not " + "provide them."); + } + + constexpr int SF_TILE_DIM_M = 128; + constexpr int SF_TILE_DIM_K = 4; + const dim3 block_size(TB_DIM, TB_DIM); + + int m{0}, k{0}; + void* input_ptr{nullptr}; + void* output_ptr{nullptr}; + bool rowwise{false}; + + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP8, got ", + to_string(input->dtype()), ")."); + if (has_rowwise_scale_inv) { + NVTE_CHECK(output->scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", output->scale_inv.shape, "."); + m = output->scale_inv.shape[0]; + k = output->scale_inv.shape[1]; + NVTE_CHECK(static_cast(m) * k == input->scale_inv.numel(), + "Expected input tensor to have ", static_cast(m) * k, + " row-wise scaling factors, but got shape=", input->scale_inv.shape, "."); + NVTE_CHECK(static_cast(m) * k == output->scale_inv.numel(), + "Expected output tensor to have ", static_cast(m) * k, + " row-wise scaling factors, but got shape=", output->scale_inv.shape, "."); + input_ptr = input->scale_inv.dptr; + output_ptr = output->scale_inv.dptr; + rowwise = true; + } else if (has_columnwise_scale_inv) { + NVTE_CHECK(output->columnwise_scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", output->columnwise_scale_inv.shape, + "."); + m = output->columnwise_scale_inv.shape[1]; + k = output->columnwise_scale_inv.shape[0]; + NVTE_CHECK( + static_cast(m) * k == input->columnwise_scale_inv.numel(), + "Expected input tensor to have ", static_cast(m) * k, + " column-wise scaling factors, but got shape=", input->columnwise_scale_inv.shape, "."); + NVTE_CHECK(static_cast(m) * k == output->columnwise_scale_inv.numel(), + "Expected output tensor to have ", static_cast(m) * k, + " column-wise scaling factors, but got shape=", + output->columnwise_scale_inv.shape, "."); + input_ptr = input->columnwise_scale_inv.dptr; + output_ptr = output->columnwise_scale_inv.dptr; + rowwise = false; + } + break; + } + case NVTE_NVFP4_1D_SCALING: { + NVTE_CHECK(is_fp4_dtype(input->dtype()), "Input tensor has invalid dtype (expected FP4, got ", + to_string(input->dtype()), ")."); + // NVFP4: always unswizzle rowwise regardless of which scale buffer holds the data + if (has_rowwise_scale_inv) { + NVTE_CHECK(output->scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", output->scale_inv.shape, "."); + m = output->scale_inv.shape[0]; + k = output->scale_inv.shape[1]; + // Example for NVFP4 rowwise path: + NVTE_CHECK(static_cast(m) * k == input->scale_inv.numel(), + "Expected input tensor to have ", static_cast(m) * k, + " row-wise scaling factors, but got shape=", input->scale_inv.shape, "."); + NVTE_CHECK(static_cast(m) * k == output->scale_inv.numel(), + "Expected output tensor to have ", static_cast(m) * k, + " row-wise scaling factors, but got shape=", output->scale_inv.shape, "."); + input_ptr = input->scale_inv.dptr; + output_ptr = output->scale_inv.dptr; + } else if (has_columnwise_scale_inv) { + NVTE_CHECK(output->columnwise_scale_inv.shape.size() == 2, + "Expected 2D scaling factors, got shape=", output->columnwise_scale_inv.shape, + "."); + m = output->columnwise_scale_inv.shape[0]; + k = output->columnwise_scale_inv.shape[1]; + NVTE_CHECK( + static_cast(m) * k == input->columnwise_scale_inv.numel(), + "Expected input tensor to have ", static_cast(m) * k, + " column-wise scaling factors, but got shape=", input->columnwise_scale_inv.shape, "."); + NVTE_CHECK(static_cast(m) * k == output->columnwise_scale_inv.numel(), + "Expected output tensor to have ", static_cast(m) * k, + " column-wise scaling factors, but got shape=", + output->columnwise_scale_inv.shape, "."); + input_ptr = input->columnwise_scale_inv.dptr; + output_ptr = output->columnwise_scale_inv.dptr; + } + rowwise = true; + break; + } + default: + NVTE_ERROR("Invalid scaling mode"); + } + + NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Output should be padded in M/N dimension!"); + NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Output should be padded in K dimension!"); + + const int num_tiles_m = m / SF_TILE_DIM_M; + const int num_tiles_k = k / SF_TILE_DIM_K; + + auto launch_unswizzle = [&](int vec_load_size, const dim3& num_blocks, int slm_size) { + switch (vec_load_size) { + case 4: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(unswizzle_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + unswizzle_scaling_kernel + <<>>(input_ptr, output_ptr, m, k, rowwise); + break; + case 2: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(unswizzle_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + unswizzle_scaling_kernel + <<>>(input_ptr, output_ptr, m, k, rowwise); + break; + case 1: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(unswizzle_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + unswizzle_scaling_kernel + <<>>(input_ptr, output_ptr, m, k, rowwise); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + } + NVTE_CHECK_CUDA(cudaGetLastError()); + }; + + 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; + dim3 num_blocks = rowwise ? dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m) + : dim3(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); + launch_unswizzle(vec_load_size, num_blocks, slm_size); +} + +void multi_tensor_unswizzle_scaling_factors(const std::vector& input, + std::vector& output, cudaStream_t stream) { + size_t num_tensors = output.size(); + const auto& first_scaling_mode = output[0]->scaling_mode; + + constexpr int SF_TILE_DIM_M = 128; + constexpr int SF_TILE_DIM_K = 4; + + bool all_has_data = true; + bool all_has_columnwise_data = true; + bool all_nvfp4 = true; + for (size_t i = 0; i < num_tensors; i++) { + const auto scaling_mode = output[i]->scaling_mode; + const auto is_fp8 = is_fp8_dtype(input[i]->dtype()); + const auto is_fp4 = is_fp4_dtype(input[i]->dtype()); + + NVTE_CHECK(scaling_mode == first_scaling_mode, + "All tensors should have the same scaling mode in multi-tensor unswizzle."); + NVTE_CHECK( + (is_fp8 && is_mxfp8_scaling(scaling_mode)) || (is_fp4 && is_nvfp4_scaling(scaling_mode)), + "Not implemented scaling mode " + to_string(scaling_mode) + "."); + NVTE_CHECK(input[i]->with_gemm_swizzled_scales, + "Expected input tensors with scales in GEMM swizzled format."); + NVTE_CHECK(!output[i]->with_gemm_swizzled_scales, + "Expected output tensors with scales in compact format."); + 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) + "]"); + + all_has_data = all_has_data && output[i]->scale_inv.has_data(); + all_has_columnwise_data = + (all_has_columnwise_data && output[i]->columnwise_scale_inv.has_data()); + all_nvfp4 = all_nvfp4 && is_nvfp4_scaling(scaling_mode); + } + NVTE_CHECK(all_has_data || all_has_columnwise_data, + "All tensors should have data or columnwise data."); + NVTE_CHECK(!all_has_data || !all_has_columnwise_data, + "All tensors have both data and columnwise data."); + + const bool rowwise_unswizzle = all_has_data || all_nvfp4; + const bool columnwise_unswizzle = all_has_columnwise_data && !all_nvfp4; + + if (rowwise_unswizzle) { + MultiSwizzleArgs kernel_args; + kernel_args.num_tensors = 0; + kernel_args.block_range[0] = 0; + int vec_load_size = 4; + for (size_t i = 0; i < num_tensors; i++) { + if (kernel_args.num_tensors == kMaxTensorsPerKernel) { + if (vec_load_size == 3) vec_load_size = 1; + launch_multi_tensor_unswizzle_scaling_factors( + kernel_args, vec_load_size, true, stream); + kernel_args.num_tensors = 0; + vec_load_size = 4; + } + int m, k; + if (all_has_data) { + NVTE_CHECK(input[i]->scale_inv.has_data(), "Input tensor ", i, + " does not have row-wise scaling factors."); + NVTE_CHECK(output[i]->scale_inv.shape.size() == 2, "Expected output tensor ", i, + " to have ", "2D scaling factors, got shape=", output[i]->scale_inv.shape, "."); + m = output[i]->scale_inv.shape[0]; + k = output[i]->scale_inv.shape[1]; + NVTE_CHECK(m * k == input[i]->scale_inv.numel(), "Expected input tensor ", i, " to have ", + m * k, " row-wise scaling factors, but got shape=", input[i]->scale_inv.shape, + "."); + } + + if (all_has_columnwise_data) { + NVTE_CHECK(all_nvfp4, + "When doing rowwise unswizzle with columnwise data, it has to be NVFP4"); + NVTE_CHECK(input[i]->columnwise_scale_inv.has_data(), "Input tensor ", i, + " does not have column-wise scaling factors."); + NVTE_CHECK(output[i]->columnwise_scale_inv.shape.size() == 2, "Expected output tensor ", i, + " to have ", + "2D scaling factors, got shape=", output[i]->columnwise_scale_inv.shape, "."); + m = output[i]->columnwise_scale_inv.shape[0]; + k = output[i]->columnwise_scale_inv.shape[1]; + NVTE_CHECK(m * k == input[i]->columnwise_scale_inv.numel(), "Expected input tensor ", i, + " to have ", m * k, " column-wise scaling factors, but got shape=", + input[i]->columnwise_scale_inv.shape, "."); + } + + NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Output should be padded in M/N dimension!"); + NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Output should be padded in K dimension!"); + NVTE_CHECK(k > 0, "Output scale inverse should be 2D!"); + + int num_tiles_k = k / SF_TILE_DIM_K; + int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; + vec_load_size = all_nvfp4 ? 1 : std::min(vec_load_size, vec_load_size_i); + + const int pos = kernel_args.num_tensors; + kernel_args.m_list[pos] = m; + kernel_args.k_list[pos] = k; + if (!all_nvfp4 || all_has_data) { + kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr); + kernel_args.output_list[pos] = output[i]->scale_inv.dptr; + } 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.num_tensors++; + } + if (vec_load_size == 3) vec_load_size = 1; + launch_multi_tensor_unswizzle_scaling_factors( + kernel_args, vec_load_size, true, stream); + } + + if (columnwise_unswizzle) { + NVTE_CHECK(!all_nvfp4, "NVFP4 shouldn't end up here because it only needs rowwise unswizzle"); + + MultiSwizzleArgs kernel_args; + kernel_args.num_tensors = 0; + kernel_args.block_range[0] = 0; + int vec_load_size = 4; + for (size_t i = 0; i < num_tensors; i++) { + if (kernel_args.num_tensors == kMaxTensorsPerKernel) { + if (vec_load_size == 3) vec_load_size = 1; + launch_multi_tensor_unswizzle_scaling_factors( + kernel_args, vec_load_size, false, stream); + kernel_args.num_tensors = 0; + vec_load_size = 4; + } + NVTE_CHECK(output[i]->columnwise_scale_inv.shape.size() == 2, "Expected output tensor ", i, + " to have ", + "2D scaling factors, got shape=", output[i]->columnwise_scale_inv.shape, "."); + const int m = output[i]->columnwise_scale_inv.shape[1]; + const int k = output[i]->columnwise_scale_inv.shape[0]; + + NVTE_CHECK(m % SF_TILE_DIM_M == 0, "Output should be padded in M/N dimension!"); + NVTE_CHECK(k % SF_TILE_DIM_K == 0, "Output should be padded in K dimension!"); + NVTE_CHECK(k > 0, "Output scale inverse should be 2D!"); + NVTE_CHECK(m * k == std::accumulate(input[i]->columnwise_scale_inv.shape.begin(), + input[i]->columnwise_scale_inv.shape.end(), 1, + std::multiplies()), + "Input.columnwise_scale_inv size is not equal to " + "Output.columnwise_scale_inv size!"); + + int num_tiles_k = k / SF_TILE_DIM_K; + int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; + vec_load_size = std::min(vec_load_size, vec_load_size_i); + + const int pos = kernel_args.num_tensors; + 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.m_list[pos] = m; + kernel_args.k_list[pos] = k; + kernel_args.num_tensors++; + } + if (vec_load_size == 3) vec_load_size = 1; + launch_multi_tensor_unswizzle_scaling_factors( + kernel_args, vec_load_size, false, stream); + } +} } // namespace transformer_engine /* @@ -876,6 +1514,26 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen multi_tensor_swizzle_scaling_factors(input_list, output_list, stream); } +void nvte_unswizzle_scaling_factors(const NVTETensor input, NVTETensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_unswizzle_scaling_factors); + using namespace transformer_engine; + unswizzle_scaling_factors(convertNVTETensorCheck(input), convertNVTETensorCheck(output), stream); +} + +void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs, + const size_t num_tensors, cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_unswizzle_scaling_factors); + 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_unswizzle_scaling_factors(input_list, output_list, stream); +} + namespace transformer_engine { void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, From e83c09742166dfef3f871cfa1407605feafb3afe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=A9tan=20Lepage?= Date: Sat, 4 Apr 2026 00:10:24 +0200 Subject: [PATCH 314/521] Fix nvshmem build (#2815) Signed-off-by: Gaetan Lepage --- transformer_engine/common/nvshmem_api/CMakeLists.txt | 5 +++-- transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu | 1 + 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/nvshmem_api/CMakeLists.txt b/transformer_engine/common/nvshmem_api/CMakeLists.txt index 1e72e42b0a..3d9b6b5ec4 100644 --- a/transformer_engine/common/nvshmem_api/CMakeLists.txt +++ b/transformer_engine/common/nvshmem_api/CMakeLists.txt @@ -16,7 +16,8 @@ set(NVSHMEMAPI_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}" PARENT_SCOPE) target_link_directories(nvshmemapi PUBLIC ${NVSHMEM_HOME}/lib) target_link_libraries(nvshmemapi PUBLIC -static-libstdc++ nvshmem_device nvshmem_host CUDA::nvml CUDA::cublas CUDA::cuda_driver) target_include_directories(nvshmemapi PRIVATE - ${NVSHMEM_HOME}/include/) + ${NVSHMEM_HOME}/include/ + ${CMAKE_CURRENT_SOURCE_DIR}/../include) target_include_directories(nvshmemapi PUBLIC ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} "${CMAKE_CURRENT_SOURCE_DIR}") @@ -24,4 +25,4 @@ target_include_directories(nvshmemapi PUBLIC set_target_properties(nvshmemapi PROPERTIES CUDA_STANDARD 17 POSITION_INDEPENDENT_CODE ON - CUDA_SEPARABLE_COMPILATION ON) \ No newline at end of file + CUDA_SEPARABLE_COMPILATION ON) diff --git a/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu b/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu index efa7d0d53a..f81062d63b 100644 --- a/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu +++ b/transformer_engine/common/nvshmem_api/nvshmem_waitkernel.cu @@ -15,6 +15,7 @@ #include #include +#include "../util/cuda_driver.h" #include "../util/logging.h" #include "nvshmem_waitkernel.h" From 5abadf4ee573147f9fbc0aadac44176db5148813 Mon Sep 17 00:00:00 2001 From: Cory Ye <44509866+cspades@users.noreply.github.com> Date: Sat, 4 Apr 2026 15:48:18 -0700 Subject: [PATCH 315/521] [FSDP2/Megatron-FSDP/DCP] If model parameters are DTensors, optimizer states should also be DTensors. (#2795) * If model parameters are DTensors, optimizer state should also be DTensor. Signed-off-by: Cory Ye * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Unpack DTensor in FusedAdam.step(). Signed-off-by: Cory Ye * Apply suggestions from code review Add Greptile bug-fixes. Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Cory Ye <44509866+cspades@users.noreply.github.com> * Revert erroneous Greptile diff. Signed-off-by: Cory Ye * Add DTensor parity check to FusedAdam.step(). Signed-off-by: Cory Ye * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add DTensor handling in state_dict and load_state_dict, and add a DCP re-sharding test. Signed-off-by: Cory Ye * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update test commentary. Signed-off-by: Cory Ye * Filter out DCP resharding tests from the 2 GPU FusedAdam test matrix, as those tests need to be run in sequence. Signed-off-by: Cory Ye * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix float8 Signed-off-by: Varun Thumbe * xfail block scaling Signed-off-by: Varun Thumbe * Fix rebase error, pytest filters were shoved into a different test. Signed-off-by: Cory Ye --------- Signed-off-by: Cory Ye Signed-off-by: Cory Ye <44509866+cspades@users.noreply.github.com> Signed-off-by: Varun Thumbe 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: vthumbe1503 --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 185 +++++++++++++++++- tests/pytorch/distributed/test_torch_fsdp2.py | 75 +++++++ .../pytorch/optimizers/fused_adam.py | 92 ++++++--- .../pytorch/tensor/float8_tensor.py | 21 +- 4 files changed, 345 insertions(+), 28 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 877fa66795..42df06ed7f 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -16,11 +16,17 @@ fused_adam_fp8_master_weights, fused_adam_fp8_master_weights_no_meta, fused_adam_bf16, fused_adam_fp8_no_master, fused_adam_bf16_store_param_remainders, fuse_wgrad_accumulation, dcp_output_parity, dcp_output_parity_async, - safetensors_fp32_export + dcp_resharding_save, dcp_resharding_load, safetensors_fp32_export Available --recipe values: DelayedScaling, Float8CurrentScaling, Float8BlockScaling, MXFP8BlockScaling, NVFP4BlockScaling + +Note: dcp_resharding_save and dcp_resharding_load are two phases of a single +cross-topology test. Run dcp_resharding_save under a larger world_size first +(e.g. --nproc_per_node=4), then run dcp_resharding_load under a smaller one +(e.g. --nproc_per_node=2). The orchestration is handled automatically by +test_fsdp2_fused_adam_dcp_resharding in test_torch_fsdp2.py. """ import argparse @@ -465,7 +471,8 @@ def test_safetensors_fp32_export(recipe_name): if recipe_name == "MXFP8BlockScaling": pytest.xfail( "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access" + "MXFP8 quantized tensors, causing illegal memory access. " + "Fixed by https://github.com/NVIDIA/TransformerEngine/pull/2789." ) from safetensors.torch import load_file, save_file @@ -554,7 +561,8 @@ def test_dcp_output_parity(recipe_name, async_save): "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" + "multi_tensor_apply: CUDA Error: an illegal memory access was encountered. " + "Fixed by https://github.com/NVIDIA/TransformerEngine/pull/2789." ) if recipe_name == "NVFP4BlockScaling": @@ -740,6 +748,173 @@ def test_dcp_output_parity(recipe_name, async_save): shutil.rmtree(checkpoint_dir, ignore_errors=True) +def test_dcp_resharding_save(recipe_name): + """Phase 1 of the DCP resharding test: train with current world_size and save checkpoint. + + Trains a model for NUM_STEPS, records the forward-pass output, and writes: + - A DCP checkpoint to /tmp/te_test_fsdp2_dcp_resharding_/ + - A reference output tensor to /tmp/te_test_fsdp2_dcp_resharding__ref.pt + + These artifacts are consumed by test_dcp_resharding_load, which runs under + a *different* world_size (typically half as many ranks) to verify that DCP + correctly reshards the checkpoint into the new topology. + + The two phases are orchestrated by test_fsdp2_fused_adam_dcp_resharding in + test_torch_fsdp2.py using two sequential plain torchrun invocations. + """ + recipe = get_recipe_from_string(recipe_name) + + import torch.distributed.checkpoint as dcp + + world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + checkpoint_dir = f"/tmp/te_test_fsdp2_dcp_resharding_{recipe_name}" + ref_output_path = f"/tmp/te_test_fsdp2_dcp_resharding_{recipe_name}_ref.pt" + + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + if os.path.exists(ref_output_path): + os.remove(ref_output_path) + dist.barrier() + + model = _build_model(fp8_init=True, recipe=recipe) + model = _shard_model(model, world_size) + + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + # Fixed seed so the load phase reproduces the exact same input tensor. + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + for _ in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Record the reference output before saving. + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + ref_output = model(x).clone().cpu() + + dist.barrier() + if rank == 0: + torch.save(ref_output, ref_output_path) + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + model_state = { + k: v for k, v in model.state_dict().items() if not k.endswith("_extra_state") + } + else: + model_state = model.state_dict() + + dcp.save( + {"model": model_state, "optimizer": optimizer.state_dict()}, checkpoint_id=checkpoint_dir + ) + dist.barrier() + + +def test_dcp_resharding_load(recipe_name): + """Phase 2 of the DCP resharding test: load into a different world_size and verify parity. + + Loads the DCP checkpoint written by test_dcp_resharding_save (which ran + under a larger world_size, e.g. 4 ranks) into a fresh model sharded over + the current, smaller world_size (e.g. 2 ranks). Asserts that the model + output after loading is bitwise-identical to the reference saved in phase 1, + confirming that DCP resharding correctly reconstructs all parameter shards. + """ + recipe = get_recipe_from_string(recipe_name) + + import torch.distributed.checkpoint as dcp + + world_size, device = _get_dist_info() + rank = int(os.environ.get("RANK", "0")) + checkpoint_dir = f"/tmp/te_test_fsdp2_dcp_resharding_{recipe_name}" + ref_output_path = f"/tmp/te_test_fsdp2_dcp_resharding_{recipe_name}_ref.pt" + + try: + model2 = _build_model(fp8_init=True, recipe=recipe) + model2 = _shard_model(model2, world_size) + + optimizer2 = te.optimizers.FusedAdam( + model2.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + # Same fixed seed as the save phase to reproduce identical x/target. + torch.manual_seed(12345) + torch.cuda.manual_seed(12345) + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.bfloat16, device=device) + target = torch.randn_like(x) + + # Populate optimizer state so load_state_dict has a matching structure. + optimizer2.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + out_tmp = model2(x) + F.mse_loss(out_tmp, target).backward() + optimizer2.step() + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + model2_state = { + k: v for k, v in model2.state_dict().items() if not k.endswith("_extra_state") + } + else: + model2_state = model2.state_dict() + + state_to_load = {"model": model2_state, "optimizer": optimizer2.state_dict()} + dcp.load(state_to_load, checkpoint_id=checkpoint_dir) + model2.load_state_dict( + state_to_load["model"], + strict=( + False + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling) + else True + ), + ) + optimizer2.load_state_dict(state_to_load["optimizer"]) + + with torch.no_grad(): + with te.autocast(enabled=True, recipe=recipe): + loaded_output = model2(x).cpu() + + if rank == 0: + ref_output = torch.load(ref_output_path, weights_only=True) + + if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0.05, + atol=0.1, + msg=lambda m: f"Resharded model output differs from reference: {m}", + ) + else: + torch.testing.assert_close( + loaded_output, + ref_output, + rtol=0, + atol=0, + msg=lambda m: f"Resharded model output differs from reference: {m}", + ) + finally: + dist.barrier() + if rank == 0: + shutil.rmtree(checkpoint_dir, ignore_errors=True) + if os.path.exists(ref_output_path): + os.remove(ref_output_path) + + TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, @@ -749,13 +924,15 @@ def test_dcp_output_parity(recipe_name, async_save): "fuse_wgrad_accumulation": test_fuse_wgrad_accumulation, "dcp_output_parity": functools.partial(test_dcp_output_parity, async_save=False), "dcp_output_parity_async": functools.partial(test_dcp_output_parity, async_save=True), + "dcp_resharding_save": test_dcp_resharding_save, + "dcp_resharding_load": test_dcp_resharding_load, "safetensors_fp32_export": test_safetensors_fp32_export, } if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--test", required=True, choices=list(TESTS.keys())) + parser.add_argument("--test", required=True, choices=sorted(TESTS.keys())) parser.add_argument( "--recipe", type=str, diff --git a/tests/pytorch/distributed/test_torch_fsdp2.py b/tests/pytorch/distributed/test_torch_fsdp2.py index ee20886631..f386659b6c 100644 --- a/tests/pytorch/distributed/test_torch_fsdp2.py +++ b/tests/pytorch/distributed/test_torch_fsdp2.py @@ -5,6 +5,7 @@ import os import sys import subprocess +import sys from pathlib import Path sys.path.append(str(Path(__file__).resolve().parent.parent)) @@ -18,6 +19,12 @@ NUM_PROCS: int = torch.cuda.device_count() _FSDP2_DIR = Path(__file__).parent.resolve() / "fsdp2_tests" +# Import some utilities from PyTest-owned conftest.py. +sys.path.insert(0, str(_FSDP2_DIR)) +from conftest import _parametrize_recipes + +sys.path.pop(0) + @pytest.mark.skipif(NUM_PROCS % 2 != 0, reason="Requires even number of GPUs") @pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") @@ -59,6 +66,10 @@ def test_fsdp2_fused_adam_tests(): "-v", "-s", "--tb=short", + # The following 2 tests need to be run in sequence, + # as they depend on each other. + "-k", + "not dcp_resharding_save and not dcp_resharding_load", ], valid_returncodes=(0, 5), env=os.environ, @@ -90,6 +101,70 @@ def test_fsdp2_mem_leak_tests(): assert result.returncode in (0, 5), f"Inner pytest failed with exit code {result.returncode}" +@pytest.mark.skipif(NUM_PROCS < 4, reason="Requires 4+ GPUs for DP4→DP2 resharding test") +@pytest.mark.skipif(not te.torch_version() >= (2, 4, 0), reason="Requires PyTorch 2.4.0+") +@pytest.mark.parametrize("recipe", _parametrize_recipes()) +def test_fsdp2_fused_adam_dcp_resharding(recipe): + """DCP checkpoint saved with DP4 loads correctly into DP2 (cross-topology resharding). + + Runs two sequential torchrun invocations against run_fsdp2_fused_adam.py: + 1. nproc=4 → dcp_resharding_save (train + write checkpoint + ref output) + 2. nproc=2 → dcp_resharding_load (load checkpoint, assert output parity) + """ + if recipe == "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." + ) + if recipe == "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 == "Float8BlockScaling": + pytest.xfail( + "Float8BlockScaling doesnt work for DCP resharding with scale inv padding " + "not being handled correctly for slice ops" + ) + + test_path = _FSDP2_DIR / "run_fsdp2_fused_adam.py" + + # Phase 1: save checkpoint with 4 ranks. + result = subprocess.run( + [ + "torchrun", + "--nproc_per_node=4", + "--local-ranks-filter=0", + str(test_path), + "--test", + "dcp_resharding_save", + "--recipe", + recipe, + ], + env=os.environ, + timeout=300, + ) + assert result.returncode == 0, f"DCP resharding save phase failed: {result.returncode}" + + # Phase 2: load checkpoint with 2 ranks (different topology). + result = subprocess.run( + [ + "torchrun", + "--nproc_per_node=2", + "--local-ranks-filter=0", + str(test_path), + "--test", + "dcp_resharding_load", + "--recipe", + recipe, + ], + env=os.environ, + timeout=300, + ) + assert result.returncode == 0, f"DCP resharding load phase failed: {result.returncode}" + + def test_dummy() -> None: """Dummy test diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index bcfd2bef19..437dfa829e 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -321,11 +321,14 @@ def get_unscaled_state( """ state = self.state[param] dtype = self.name_to_dtype_map[state_name] + unscaled_local_state = state[state_name] + if isinstance(unscaled_local_state, DTensor): + unscaled_local_state = unscaled_local_state._local_tensor if dtype == torch.uint8: - unscaled = state[state_name].float() + unscaled = unscaled_local_state.float() elif dtype == torch.float16: - assert state[state_name].dtype == torch.float16 - unscaled = state[state_name].float() + assert unscaled_local_state.dtype == torch.float16 + unscaled = unscaled_local_state.float() unscaled.mul_(self._scales[param][state_name]) elif dtype == torch.float32: if ( @@ -333,16 +336,16 @@ def get_unscaled_state( and state_name == "master_param" and param.dtype == torch.bfloat16 ): - assert state[state_name].dtype == torch.int16 + assert unscaled_local_state.dtype == torch.int16 else: - assert state[state_name].dtype == torch.float32 - unscaled = state[state_name] + assert unscaled_local_state.dtype == torch.float32 + unscaled = unscaled_local_state elif dtype == torch.bfloat16: - assert state[state_name].dtype == torch.bfloat16 + assert unscaled_local_state.dtype == torch.bfloat16 if skip_unscale: - unscaled = state[state_name] + unscaled = unscaled_local_state else: - unscaled = state[state_name].float() + unscaled = unscaled_local_state.float() else: raise RuntimeError(f"Dtype of {state_name} can only be fp8/fp16/bf16/fp32.") return unscaled @@ -357,7 +360,7 @@ def set_scaled_state(self, param, state_name, unscaled_state): param (torch.nn.Parameter): One of parameters in this optimizer. state_name (string): Name of optimizer states, can be one of 'exp_avg', 'exp_avg_sq', and 'master_param`. - unscaled_state (torch.Tensor): The original high-precision(FP32) state. + unscaled_state (torch.Tensor): The original high-precision (FP32) state. """ store_param_remainders = ( @@ -374,12 +377,17 @@ def set_scaled_state(self, param, state_name, unscaled_state): if state_name not in state: self._initialize_state(param, state_name, False, store_param_remainders) + # If the state is a DTensor, retrieve its local Tensor for scaling. + local_state = state[state_name] + if isinstance(local_state, DTensor): + local_state = local_state._local_tensor + dtype = self.name_to_dtype_map[state_name] if dtype != torch.float32: scale = self._scales[param] - self._apply_scale(state_name, unscaled_state, state[state_name], scale[state_name]) + self._apply_scale(state_name, unscaled_state, local_state, scale[state_name]) else: - state[state_name].copy_(unscaled_state) + local_state.copy_(unscaled_state) def _initialize_state( self, param, state_name, zero_buffer: bool, store_param_remainders: bool = False @@ -396,9 +404,9 @@ def _initialize_state( dtype = self.name_to_dtype_map[state_name] # Extract local tensor from DTensor (e.g. from FSDP2) to avoid # QuantizedTensor.__torch_dispatch__ ignoring the dtype kwarg in - # torch.empty_like, and to ensure optimizer states are plain tensors. + # torch.empty_like. local_param = param._local_tensor if isinstance(param, DTensor) else param - # Handle QuantizedTensor by dequantizing first + # Handle QuantizedTensor by dequantizing first. param_for_empty = ( local_param.dequantize() if isinstance(local_param, QuantizedTensor) else local_param ) @@ -409,18 +417,29 @@ def _initialize_state( if zero_buffer: data.zero_() + # Install the quantized or un-quantized optimizer state. if dtype == torch.uint8: quantizer = Float8Quantizer( scale=torch.ones([1], dtype=torch.float32, device=param.device), amax=torch.zeros([1], dtype=torch.float32, device=param.device), fp8_dtype=tex.DType.kFloat8E4M3, ) - self.state[param][state_name] = quantizer.make_empty(param.shape) + self.state[param][state_name] = quantizer.make_empty(data.shape) self.state[param][state_name].quantize_(data.float()) else: - self.state[param][state_name] = data + # If the original Parameter was a DTensor, re-wrap the state + # into DTensor to support Torch DCP checkpointing. + if isinstance(param, DTensor): + self.state[param][state_name] = DTensor.from_local( + self.state[param][state_name], + device_mesh=param.device_mesh, + placements=param.placements, + shape=param.size(), + stride=param.stride(), + ) + # Create scale if necessary. if dtype != torch.float32: if param not in self._scales: @@ -447,7 +466,7 @@ def initialize_state(self, param, store_param_remainders): ) if not store_param_remainders: # Extract local tensor from DTensor and dequantize QuantizedTensor - # to get a plain float32 copy for the master weight. + # to set scales for the optimizer state's main weights. local_param = param._local_tensor if isinstance(param, DTensor) else param if isinstance(local_param, QuantizedTensor): master = local_param.dequantize(dtype=torch.float32).clone().detach() @@ -475,6 +494,15 @@ def state_dict(self): new_v = {} for name in v: new_v[name] = self.get_unscaled_state(param, name) + if isinstance(param, DTensor): + # Re-wrap the optimizer state as a DTensor. + new_v[name] = DTensor.from_local( + new_v[name], + device_mesh=param.device_mesh, + placements=param.placements, + shape=param.size(), + stride=param.stride(), + ) state_dict["state"][k] = new_v return state_dict @@ -500,15 +528,19 @@ def load_state_dict(self, state_dict): for name in v: if v[name] is None: continue + state = v[name] + if isinstance(state, DTensor): + # Un-pack the local Tensor state for set_scaled_state. + state = state._local_tensor if ( self.store_param_remainders and name == "master_param" and param.dtype == torch.bfloat16 ): - self.set_scaled_state(param, name, v[name]) - assert v[name].dtype == torch.int16 + self.set_scaled_state(param, name, state) + assert state.dtype == torch.int16 else: - self.set_scaled_state(param, name, v[name].float()) + self.set_scaled_state(param, name, state.float()) def step(self, closure=None, grad_scaler=None): """Performs a single optimization step. @@ -592,12 +624,28 @@ def step(self, closure=None, grad_scaler=None): if p_grad.data.is_sparse: raise RuntimeError("FusedAdam does not support sparse gradients.") + # Validate parameter, gradient, and state DTensor parity for the step. + dtensor_param = isinstance(p, DTensor) + assert dtensor_param == isinstance(p_grad, DTensor), ( + f"[FusedAdam DTensor Disparity] Parameter {p} and Gradient {p_grad} do not" + " match!" + ) + for name in ["exp_avg", "exp_avg_sq", "master_param"]: + if name in state: + assert dtensor_param == isinstance(state[name], DTensor), ( + f"[FusedAdam DTensor Disparity] Parameter {p} and" + f" {name} {state[name]} do not match!" + ) + # Unscaling unscaled_state = {} for name in ["exp_avg", "exp_avg_sq", "master_param"]: if name in state: + state_tensor = state[name] + if isinstance(state_tensor, DTensor): + state_tensor = state_tensor._local_tensor if name == "master_param" and store_param_remainders: - unscaled_state[name] = self.state[p][name] + unscaled_state[name] = state_tensor assert unscaled_state[name].dtype == torch.int16 else: unscaled = self.get_unscaled_state( @@ -606,7 +654,7 @@ def step(self, closure=None, grad_scaler=None): unscaled_state[name] = unscaled if self.name_to_dtype_map[name] != torch.float32: unscaled_lists[name].append(unscaled) - scaled_lists[name].append(state[name]) + 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) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index e8284eaa53..256250ff64 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -678,7 +678,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): quantizer=tensor._quantizer, ) - if func in [aten.slice.Tensor, aten.select.int]: + if func in (aten.slice.Tensor, aten.select.int): tensor = args[0] data = tensor._data data_slice = data.__torch_dispatch__( @@ -687,7 +687,24 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): [data] + list(args[1:]), kwargs, ) - return Float8Tensor.make_like(tensor, data=data_slice, shape=data_slice.shape) + transpose_slice = None + if tensor._transpose is not None and not tensor._transpose_invalid: + transpose = tensor._transpose + ndim = data.dim() + dim = args[1] if len(args) > 1 else 0 + t_dim = 0 if dim == ndim - 1 else dim + 1 + transpose_slice = transpose.__torch_dispatch__( + func, + types, + [transpose, t_dim] + list(args[2:]), + kwargs, + ) + return Float8Tensor.make_like( + tensor, + data=data_slice, + data_transpose=transpose_slice, + shape=data_slice.shape, + ) # Related to FSDP2 if func == aten.split.Tensor: From ac966517e860b28ba2f17316bcfb9761fe12d30e Mon Sep 17 00:00:00 2001 From: Qiyu Wan <39144338+WanZzzzzz@users.noreply.github.com> Date: Mon, 6 Apr 2026 07:42:26 -0700 Subject: [PATCH 316/521] Fix memory overheads with FP4 native weights (#2834) * fix memory overheads Signed-off-by: qiyuw * comments Signed-off-by: qiyuw * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: qiyuw Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/tensor/utils.py | 93 ++++++---------------- 1 file changed, 24 insertions(+), 69 deletions(-) diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index c80bc8aaa4..ba44c7a619 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -118,46 +118,6 @@ def quantize_master_weights( else: use_fsdp_shard_model_weights = True - # Batch convert master_weights to model dtype for NVFP4 (single kernel instead of N kernels) - # Check if there are any NVFP4 weights - has_nvfp4 = any( - isinstance(w._get_quantizer(), NVFP4Quantizer) - for w in model_weights - if hasattr(w, "_get_quantizer") - ) - if has_nvfp4 and len(model_weights) > 0: - # Find target dtype from first NVFP4 weight - target_dtype = None - for w in model_weights: - if hasattr(w, "_get_quantizer") and isinstance(w._get_quantizer(), NVFP4Quantizer): - target_dtype = w.dtype - break - - if target_dtype is not None: - # Collect non-None master_weights and their indices - non_none_indices = [] - non_none_weights = [] - sizes = [] - for i, mw in enumerate(master_weights): - if mw is not None: - non_none_indices.append(i) - non_none_weights.append(mw.view(-1)) - sizes.append(mw.numel()) - - if len(non_none_weights) > 0 and non_none_weights[0].dtype != target_dtype: - # Concatenate, convert once, then split - concatenated = torch.cat(non_none_weights) - converted = concatenated.to(target_dtype) - split_weights = torch.split(converted, sizes) - - # Rebuild master_weights list with converted tensors - converted_master_weights = list(master_weights) - for idx, split_w, orig_mw in zip( - non_none_indices, split_weights, [master_weights[i] for i in non_none_indices] - ): - converted_master_weights[idx] = split_w.view(orig_mw.shape) - master_weights = converted_master_weights - for model_weight, master_weight, start_offset, fsdp_shard_model_weight in zip( model_weights, master_weights, start_offsets, fsdp_shard_model_weights ): @@ -176,42 +136,37 @@ def quantize_master_weights( if hasattr(model_weight, "clear_high_precision_init_val"): model_weight.clear_high_precision_init_val() + if master_weight is not None: + # When not using fp8/fp4_primary_weights, the master_weight (fp32) is first cast to + # bf16/fp16, and then cast to fp8 during forward. Although it's not necessary when + # fp8/fp4_primary_weights is enabled, we still keep this logic to keep numerical + # consistency. So here we cast the master_weight to model_weight.dtype. + master_weight = master_weight.to(model_weight.dtype) + quantizer = model_weight._get_quantizer() if isinstance(quantizer, NVFP4Quantizer): - # NVFP4: master_weight dtype conversion already done above nvfp4_params.append( (model_weight, master_weight, start_offset, fsdp_shard_model_weight) ) + elif isinstance(quantizer, Float8Quantizer): + delayed_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, Float8CurrentScalingQuantizer): + current_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, Float8BlockQuantizer): + blockwise_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) + elif isinstance(quantizer, MXFP8Quantizer): + mxfp8_scaling_params.append( + (model_weight, master_weight, start_offset, fsdp_shard_model_weight) + ) else: - # FP8: convert master_weight to model dtype - if master_weight is not None: - # When not using fp8_primary_weights, the master_weight (fp32) is first cast to - # bf16/fp16, and then cast to fp8 during forward. Although it's not necessary when - # fp8_primary_weights is enabled, we still keep this logic to keep numerical - # consistency. So here we cast the master_weight to model_weight.dtype. - master_weight = master_weight.to(model_weight.dtype) - - if isinstance(quantizer, Float8Quantizer): - delayed_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - elif isinstance(quantizer, Float8CurrentScalingQuantizer): - current_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - elif isinstance(quantizer, Float8BlockQuantizer): - blockwise_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - elif isinstance(quantizer, MXFP8Quantizer): - mxfp8_scaling_params.append( - (model_weight, master_weight, start_offset, fsdp_shard_model_weight) - ) - else: - raise ValueError( - f"quantize_master_weights for {type(quantizer)} is not supported yet" - ) + raise ValueError(f"quantize_master_weights for {type(quantizer)} is not supported yet") extra_args = [group, use_fsdp_shard_model_weights, manual_post_all_gather_processing] if len(delayed_scaling_params) > 0: From 86edac47c5c56e41f72d69b0908c9decff2be12c Mon Sep 17 00:00:00 2001 From: Almog Segal Date: Mon, 6 Apr 2026 20:28:40 +0300 Subject: [PATCH 317/521] Comm gemm fixes (#2818) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix GemmRs B descriptor lld for transb=true With a row_major (1×P) grid, all rows are on a single process row, so the local leading dimension must be n (full row count), not block_size(n) which is n/P. Signed-off-by: Almog Segal * Set GemmRs communication type to output data type Match the UserBuffers behavior where the reduce-scatter operates in the output precision rather than FP32. Signed-off-by: Almog Segal --------- Signed-off-by: Almog Segal --- transformer_engine/common/comm_gemm/comm_gemm.cpp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index 7be3d1bb4d..a7d78f7ac0 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -186,9 +186,9 @@ void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n } if (transb) { NVTE_CHECK(b1 == n, "Unsupported tensor dimension in B: expected ", n, ", got ", b1); - NVTE_CHECK_CUBLASMP(cublasMpMatrixDescriptorInit( - n, k, block_size(ctx, n), block_size(ctx, k), 0, 0, block_size(ctx, n), - get_cuda_dtype(b->dtype()), ctx->grid_row_major.get(), ctx->b_desc.get())); + NVTE_CHECK_CUBLASMP(cublasMpMatrixDescriptorInit(n, k, block_size(ctx, n), block_size(ctx, k), + 0, 0, n, get_cuda_dtype(b->dtype()), + ctx->grid_row_major.get(), ctx->b_desc.get())); } else { NVTE_CHECK(b0 == n, "Unsupported tensor dimension in B: expected ", n, ", got ", b0); NVTE_CHECK_CUBLASMP(cublasMpMatrixDescriptorInit( @@ -200,6 +200,11 @@ void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n NVTE_CHECK_CUBLASMP(cublasMpMatrixDescriptorInit(m, n, m, block_size(ctx, n), 0, 0, *ldd, get_cuda_dtype(d->dtype()), ctx->grid_row_major.get(), ctx->d_desc.get())); + + const cudaDataType_t comm_type = get_cuda_dtype(d->dtype()); + NVTE_CHECK_CUBLASMP(cublasMpMatmulDescriptorSetAttribute( + ctx->matmul_desc.get(), CUBLASMP_MATMUL_DESCRIPTOR_ATTRIBUTE_COMMUNICATION_TYPE, &comm_type, + sizeof comm_type)); } void GemmArInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k, From 5f9550ff8fb3886696dfd0eb88b5afef50398f17 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 6 Apr 2026 19:49:10 -0700 Subject: [PATCH 318/521] CPU offloading fix: If Data and Transpose is None depend on super Torch tensor class for the shape (#2841) * fix 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: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_quantized_tensor.py | 45 +++++++++++++++++++ .../pytorch/tensor/float8_blockwise_tensor.py | 2 +- .../pytorch/tensor/float8_tensor.py | 2 +- .../pytorch/tensor/mxfp8_tensor.py | 2 +- .../pytorch/tensor/nvfp4_tensor.py | 2 +- 5 files changed, 49 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 620fc834dd..23ce93319b 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -18,6 +18,7 @@ MXFP8Quantizer, NVFP4Quantizer, Float8Tensor, + Float8BlockwiseQTensor, MXFP8Tensor, NVFP4Tensor, QuantizedTensor, @@ -657,6 +658,50 @@ def test_chunk( y_test = y_test.to(dtype=torch.float64, device="cpu") torch.testing.assert_close(y_test, y_ref, **tols) + @pytest.mark.parametrize("quantization", _quantization_list) + def test_shape_with_none_data( + self, + *, + quantization: str, + shape: Iterable[int] = (128, 128), + dtype: torch.dtype = torch.bfloat16, + ) -> None: + """Test that shape is accessible after internal data tensors are set to None. + + During CPU offloading, both data and transpose tensors can be None. + The shape should still be available via the wrapper subclass metadata. + """ + + _, x_test = make_reference_and_test_tensors( + shape=shape, + quantization=quantization, + test_dtype=dtype, + requires_grad=False, + ) + + # Verify shape before clearing data + assert x_test.shape == torch.Size(shape) + + # Simulate CPU offloading: None out all internal data + if isinstance(x_test, Float8Tensor): + x_test._data = None + x_test._transpose = None + elif isinstance(x_test, MXFP8Tensor): + x_test._rowwise_data = None + x_test._columnwise_data = None + elif isinstance(x_test, NVFP4Tensor): + x_test._rowwise_data = None + x_test._columnwise_data = None + elif isinstance(x_test, Float8BlockwiseQTensor): + x_test._rowwise_data = None + x_test._columnwise_data = None + + # Shape must still be correct after data is cleared + assert x_test.shape == torch.Size(shape), ( + f"Expected shape {shape} but got {x_test.shape} " + f"after setting data to None on {type(x_test).__name__}" + ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) class TestMXFP8Tensor: diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index bbfc43e9bb..914397b9b6 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -598,7 +598,7 @@ def shape(self): return self._rowwise_data.shape if self._columnwise_data is not None: return self._columnwise_data.shape - raise RuntimeError("Float8BlockwiseQTensor has no data!") + return torch.Tensor.size(self) @property def is_cuda(self): diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 256250ff64..2c828aaaac 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -967,7 +967,7 @@ def shape(self): if self._transpose is not None: transpose_shape = self._transpose.shape return torch.Size(tuple(transpose_shape[1:]) + (transpose_shape[0],)) - raise RuntimeError("Both data and transpose are None") + return torch.Tensor.size(self) @property def is_cuda(self): diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 965f59b320..5cab519c79 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -884,7 +884,7 @@ def shape(self): return self._rowwise_data.shape if self._columnwise_data is not None: return self._columnwise_data.shape - raise RuntimeError("MXFP8Tensor has no data!") + return torch.Tensor.size(self) @property def is_cuda(self): diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 8ed1b4682c..eb514d3a9e 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -745,7 +745,7 @@ def shape(self): if self._columnwise_data is not None: byte_shape = self._columnwise_data.shape return torch.Size(byte_shape[1:-1] + (byte_shape[-1] * 2, byte_shape[0])) - raise RuntimeError("NVFP4Tensor has no data!") + return torch.Tensor.size(self) @property def is_cuda(self): From fdf9fb166dd0d66ac92fa3243fd08d62ebbddc71 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Tue, 7 Apr 2026 09:30:30 -0700 Subject: [PATCH 319/521] Add `NVTE_BACKWARD_OVERRIDE=high_precision|dequantized` (#2644) * Add NVTE_KEEP_BACKWARD_UNQUANTIZED Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Disable ub and clean up Signed-off-by: Ziang Li * Drop fuser changes Signed-off-by: Ziang Li * Replace use_quantized_bwd with use_fp8_bwd Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Ignore keep_backward_unquantized if delayed scaling Signed-off-by: Ziang Li * Refactor ignoring NVTE_KEEP_BACKWARD_UNQUANTIZED when delayed scaling is used Signed-off-by: Ziang Li * Add back missing ctx.debug Signed-off-by: Ziang Li * Refactor changes under fused Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Refactor high-precision overwrite if keep_backward_unquantized Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Drop redundant fp8_recipe_bwd Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop redundant ub changes Signed-off-by: Ziang Li * Drop more redundant ub changes Signed-off-by: Ziang Li * Drop redundant delayed scaling changes Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop unneeded backwards_needs_fc1_input Signed-off-by: Ziang Li * Drop and disallow LayerNormMLP implementation Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move interface changes to recipe Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Move ub overrides to fwd Signed-off-by: Ziang Li * Remove duplication Signed-off-by: Ziang Li * Simplify use_fp8_bwd logic in bwd Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Set grad quantizers to none if keep bwd unquantized Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop delayed scaling change Signed-off-by: Ziang Li * Simplify env var logic Signed-off-by: Ziang Li * Move validation check to recipe Signed-off-by: Ziang Li * Simplify effective_enabled Signed-off-by: Ziang Li * Fix inverted assertion logic Signed-off-by: Ziang Li * Simplify changes under ops Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Simplify ctx.keep_backward_unquantized Signed-off-by: Ziang Li * Fix missing attribute Signed-off-by: Ziang Li * Add unit tests Signed-off-by: Ziang Li * Fix bias errors in unit test Signed-off-by: Ziang Li * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add more shapes to unit test Signed-off-by: Ziang Li * Refator interface to `NVTE_BACKWARD_MODE=default|unquant|dequant` Signed-off-by: Ziang Li * Fix override and clean up Signed-off-by: Ziang Li * Clean up unit test Signed-off-by: Ziang Li * Clean up unit test Signed-off-by: Ziang Li * Override `ctx.reduce_and_update_bwd_fp8_tensors = False` Signed-off-by: Ziang Li * Expand unit test Signed-off-by: Ziang Li * Add `test_backward_mode_memory_peak_report` Signed-off-by: Ziang Li * Expand test coverage and fix Signed-off-by: Ziang Li * Use `numel()` Signed-off-by: Ziang Li * Refactor unit test Signed-off-by: Ziang Li * Fix grouped linear to override `*_quantizers` instead of `*_quantizer` Signed-off-by: Ziang Li * Only save input/weight when `*_requires_grad` on unquant mode Signed-off-by: Ziang Li * Fix Blackwell debug ci Signed-off-by: Ziang Li * Fix sm89 and sm90 tests Signed-off-by: Ziang Li * Fix unquant mode memory saving Signed-off-by: Ziang Li * Refactor interface to `NVTE_BACKWARD_OVERRIDE=high_precision|dequantized` Signed-off-by: Ziang Li * Rename unit test Signed-off-by: Ziang Li * Simplify env var parsing Signed-off-by: Ziang Li --------- Signed-off-by: Ziang Li Signed-off-by: Przemek Tredak Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kirthi Shankar Sivamani Co-authored-by: Przemek Tredak --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_backward_override.py | 1848 +++++++++++++++++ tests/pytorch/test_cpu_offloading.py | 41 +- tests/pytorch/test_cuda_graphs.py | 14 +- tests/pytorch/test_sanity.py | 37 +- tests/pytorch/utils.py | 31 +- transformer_engine/common/recipe/__init__.py | 78 +- transformer_engine/pytorch/module/base.py | 3 +- .../pytorch/module/grouped_linear.py | 90 +- .../pytorch/module/layernorm_linear.py | 73 +- .../pytorch/module/layernorm_mlp.py | 10 + transformer_engine/pytorch/module/linear.py | 65 +- .../pytorch/ops/basic/basic_linear.py | 60 +- transformer_engine/pytorch/ops/basic/bias.py | 5 + .../pytorch/ops/basic/quantize.py | 5 + .../ops/fused/backward_activation_bias.py | 5 +- .../fused/forward_linear_bias_activation.py | 22 +- .../ops/fused/forward_linear_bias_add.py | 24 +- .../ops/fused/forward_linear_scale_add.py | 20 +- .../ops/fused/userbuffers_forward_linear.py | 13 + transformer_engine/pytorch/ops/fuser.py | 14 +- .../float8_blockwise_tensor_storage.py | 4 + .../tensor/storage/mxfp8_tensor_storage.py | 7 + .../tensor/storage/nvfp4_tensor_storage.py | 6 + 24 files changed, 2415 insertions(+), 61 deletions(-) create mode 100644 tests/pytorch/test_backward_override.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index e67cf1bc04..377c9ddb00 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -42,6 +42,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH 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" +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" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cpu_offloading.xml $TE_PATH/tests/pytorch/test_cpu_offloading.py || test_fail "test_cpu_offloading.py" diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py new file mode 100644 index 0000000000..ed4f73adbc --- /dev/null +++ b/tests/pytorch/test_backward_override.py @@ -0,0 +1,1848 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +from __future__ import annotations + +from contextlib import nullcontext +import math +from typing import Optional + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine.pytorch.ops as te_ops +from transformer_engine.common import recipe +from transformer_engine.pytorch.cpp_extensions import general_gemm, layernorm_bwd +from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.utils import is_non_tn_fp8_gemm_supported +from transformer_engine.pytorch.ops.fused import ( + BackwardActivationBias, + ForwardLinearBiasActivation, + ForwardLinearBiasAdd, + ForwardLinearScaleAdd, + UserbuffersForwardLinear, +) +from transformer_engine.pytorch.quantized_tensor import restore_from_saved + +from utils import ( + assert_close, + make_recipe, + reset_rng_states, + skip_unsupported_backward_override, +) + + +# -------------------------- +# Mode and capability config +# -------------------------- + +_BACKWARD_OVERRIDES = ("high_precision", "dequantized") + +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +fp8_block_scaling_available, reason_for_no_fp8_block_scaling = te.is_fp8_block_scaling_available( + return_reason=True +) +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) +bf16_available, reason_for_no_bf16 = te.is_bf16_available(return_reason=True) + +_core_dtypes = [torch.float16, torch.float32] +_fused_dtypes = [torch.float16] +if bf16_available: + _core_dtypes.insert(1, torch.bfloat16) + _fused_dtypes.insert(1, torch.bfloat16) + +_quantized_numerics_recipe_list = [ + pytest.param( + "fp8_current_scaling", + marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8), + id="Float8CurrentScaling", + ), + pytest.param( + "mxfp8", + marks=pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8), + id="MXFP8BlockScaling", + ), + pytest.param( + "fp8_block_scaling", + marks=pytest.mark.skipif( + not fp8_block_scaling_available, + reason=reason_for_no_fp8_block_scaling, + ), + id="Float8BlockScaling", + ), + pytest.param( + "nvfp4", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + id="NVFP4BlockScaling", + ), +] + + +@pytest.fixture(autouse=True) +def _reset_global_fp8_state(): + """Avoid global FP8-state leakage between parametrized cases.""" + yield + FP8GlobalStateManager.reset() + + +@pytest.fixture(params=_BACKWARD_OVERRIDES, ids=lambda mode: f"mode_{mode}") +def backward_override(request: pytest.FixtureRequest) -> str: + """backward override under test.""" + return request.param + + +# -------------------------- +# Test cases +# -------------------------- + + +_shape_test_cases = [ + pytest.param((1, 64), 64, id="2d_m1_k64_n64"), + pytest.param((32, 64), 64, id="2d_m32_k64_n64"), + pytest.param((32, 96), 96, id="2d_m32_k96_n96"), + pytest.param((32, 1, 64), 64, id="3d_m32_s1_k64_n64"), + pytest.param((8, 4, 64), 128, id="3d_m32_k64_n128"), + pytest.param((16, 2, 128), 64, id="3d_m32_k128_n64"), + pytest.param((160, 64), 64, id="2d_m160_k64_n64"), + pytest.param((5, 64, 64), 64, id="3d_m320_k64_n64"), + pytest.param((3, 5, 32, 64), 96, id="4d_m480_k64_n96"), + pytest.param((2, 5, 16, 128), 64, id="4d_m160_k128_n64"), + # Intentionally unaligned token dimensions to exercise skip/support logic. + pytest.param((3, 64), 64, id="2d_m3_k64_n64_unaligned"), + pytest.param((3, 10, 64), 64, id="3d_m30_k64_n64_unaligned"), + pytest.param((3, 10, 96), 96, id="3d_m30_k96_n96_unaligned"), +] + +_bias_activation_shape_cases = [ + pytest.param((32, 64), id="2d_m32_k64"), + pytest.param((32, 96), id="2d_m32_k96"), + pytest.param((8, 4, 64), id="3d_m32_k64"), + pytest.param((160, 64), id="2d_m160_k64"), + pytest.param((5, 64, 64), id="3d_m320_k64"), + pytest.param((3, 5, 32, 64), id="4d_m480_k64"), + # Intentionally unaligned token dimensions to exercise skip/support logic. + pytest.param((3, 64), id="2d_m3_k64_unaligned"), + pytest.param((3, 10, 64), id="3d_m30_k64_unaligned"), + pytest.param((3, 10, 96), id="3d_m30_k96_unaligned"), +] + +_grouped_m_split_cases = [ + pytest.param([32, 32, 32, 32], id="uniform_splits"), + pytest.param([64, 0, 32, 32], id="with_empty_split"), + pytest.param([1, 31, 0, 96], id="small_and_empty_splits"), + pytest.param([64, 192, 0, 128], id="64_divisible_splits"), +] + +_linear_feature_cases = [ + pytest.param(64, 64, id="k64_n64"), + pytest.param(64, 128, id="k64_n128"), + pytest.param(128, 64, id="k128_n64"), + pytest.param(96, 96, id="k96_n96"), + pytest.param(64, 96, id="k64_n96"), + pytest.param(96, 64, id="k96_n64"), + pytest.param(128, 96, id="k128_n96"), + pytest.param(96, 128, id="k96_n128"), +] + +_output_feature_cases = [ + pytest.param(64, id="n64"), + pytest.param(96, id="n96"), + pytest.param(128, id="n128"), +] + +# -------------------------- +# Skip helpers +# -------------------------- + + +def _maybe_skip_recipe_dtype( + recipe_name: str, + dtype: torch.dtype, + module_type: Optional[str] = None, +) -> None: + if dtype == torch.bfloat16 and not bf16_available: + pytest.skip(reason_for_no_bf16) + if recipe_name == "nvfp4": + if module_type in ("linear", "layernorm_linear") and dtype not in ( + torch.bfloat16, + torch.float32, + ): + pytest.skip(f"NVFP4 only supports BF16 and FP32 for {module_type} in this test") + elif module_type in ("ops_linear", "grouped_linear") and dtype != torch.bfloat16: + pytest.skip(f"NVFP4 only supports BF16 for {module_type} in this test") + + +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") + + +def _maybe_skip_unsupported_recipe_shape( + recipe_name: str, + input_shape: tuple[int, ...], + module_type: str, +) -> None: + flat_first_dim = math.prod(input_shape[:-1]) + last_dim = input_shape[-1] + + if module_type in ("linear", "layernorm_linear"): + if recipe_name == "mxfp8" and (flat_first_dim % 32 != 0 or last_dim % 32 != 0): + pytest.skip( + "Linear/LayerNormLinear + MXFP8 requires prod(shape[:-1]) and shape[-1] divisible" + " by 32." + ) + return + if recipe_name == "nvfp4" 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." + ) + return + if flat_first_dim % 8 != 0 or last_dim % 16 != 0: + pytest.skip( + "Linear/LayerNormLinear FP8 execution requires prod(shape[:-1]) divisible by 8 " + "and shape[-1] divisible by 16." + ) + elif module_type == "ops_linear": + if ( + recipe_name == "fp8_current_scaling" + and not is_non_tn_fp8_gemm_supported() + and flat_first_dim % 16 != 0 + ): + pytest.skip( + "te_ops.Linear + Float8CurrentScaling on pre-Blackwell requires " + "prod(shape[:-1]) divisible by 16 for FP8 NT wgrad GEMM." + ) + if recipe_name == "mxfp8" and (flat_first_dim % 32 != 0 or last_dim % 32 != 0): + 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): + pytest.skip( + "te_ops.Linear + NVFP4 requires prod(shape[:-1]) and shape[-1] divisible by 16." + ) + + +def _maybe_skip_unsupported_grouped_splits(recipe_name: str, m_splits: list[int]) -> None: + non_empty_splits = [m for m in m_splits if m > 0] + if ( + recipe_name == "fp8_current_scaling" + and not is_non_tn_fp8_gemm_supported() + and any(m % 16 != 0 for m in non_empty_splits) + ): + pytest.skip( + "GroupedLinear + Float8CurrentScaling on pre-Blackwell requires each " + "non-empty m_split divisible by 16 for FP8 grouped NT wgrad GEMM." + ) + 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): + 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): + pytest.skip( + "GroupedLinear + NVFP4 grouped split_quantize currently requires each non-empty " + "m_split divisible by 64 due to grouped amax kernel constraints." + ) + if recipe_name == "fp8_block_scaling" and any(m % 4 != 0 for m in non_empty_splits): + pytest.skip( + "GroupedLinear + Float8BlockScaling requires each non-empty m_split divisible by 4." + ) + + +# -------------------------- +# Shared helpers +# -------------------------- + + +def _make_linear_like_module( + module_type: str, + in_features: int, + out_features: int, + dtype: torch.dtype, + *, + bias: bool, +) -> torch.nn.Module: + if module_type == "linear": + return te.Linear( + in_features, + out_features, + bias=bias, + params_dtype=dtype, + device="cuda", + ) + if module_type == "layernorm_linear": + return te.LayerNormLinear( + in_features, + out_features, + bias=bias, + params_dtype=dtype, + device="cuda", + ) + if module_type == "ops_linear": + return te_ops.Linear( + in_features, + out_features, + bias=bias, + dtype=dtype, + device="cuda", + ) + raise ValueError(f"Unsupported module type: {module_type}") + + +def _make_fused_model( + pattern: str, + in_features: int, + out_features: int, + dtype: torch.dtype, + *, + scale: float = 0.5, +) -> te_ops.Sequential: + if pattern == "bias_activation": + return te_ops.Sequential( + te_ops.Linear(in_features, out_features, bias=True, device="cuda", dtype=dtype), + te_ops.ReLU(), + ) + if pattern == "bias_add": + return te_ops.Sequential( + te_ops.Linear(in_features, out_features, bias=True, device="cuda", dtype=dtype), + te_ops.AddExtraInput(in_place=True), + ) + if pattern == "scale_add": + return te_ops.Sequential( + te_ops.Linear(in_features, out_features, bias=False, device="cuda", dtype=dtype), + te_ops.ConstantScale(scale), + te_ops.AddExtraInput(in_place=True), + ) + raise ValueError(f"Unsupported fused test pattern: {pattern}") + + +def _dequantize_saved_operand( + saved_operand: Optional[torch.Tensor], + dtype: torch.dtype, +) -> torch.Tensor: + if saved_operand is None: + raise RuntimeError("Expected saved operand but got None") + # In dequantized mode we must consume the fprop-saved quantized payload directly. + # If row-wise payload is missing, the tensor was retargeted to a transpose-only + # layout and no longer represents the original fprop operand. + if ( + not isinstance(saved_operand, torch.Tensor) + and hasattr(saved_operand, "_rowwise_data") + and getattr(saved_operand, "_rowwise_data") is None + ): + raise RuntimeError( + "Saved dequantized operand lost row-wise fprop payload (likely usage retarget)." + ) + if isinstance(saved_operand, torch.Tensor): + return saved_operand.to(dtype) + if not hasattr(saved_operand, "dequantize"): + raise RuntimeError(f"Unsupported saved operand type: {type(saved_operand)}") + return saved_operand.dequantize(dtype=dtype) + + +def _snapshot_saved_quantized_operand_layout( + saved_operand: Optional[torch.Tensor], + *, + name: str, +) -> dict[str, object]: + _assert_saved_quantized_operand_uses_rowwise_only(saved_operand, name=name) + rowwise_present = None + columnwise_present = None + rowwise_obj_id = None + if hasattr(saved_operand, "_rowwise_data"): + rowwise_data = getattr(saved_operand, "_rowwise_data") + rowwise_present = rowwise_data is not None + if rowwise_data is not None: + rowwise_obj_id = id(rowwise_data) + if hasattr(saved_operand, "_columnwise_data"): + columnwise_present = getattr(saved_operand, "_columnwise_data") is not None + return { + "name": name, + "saved_operand": saved_operand, + "rowwise_present": rowwise_present, + "columnwise_present": columnwise_present, + "rowwise_obj_id": rowwise_obj_id, + } + + +def _snapshot_layout_invariants( + guard_operands: list[tuple[str, Optional[torch.Tensor]]], +) -> list[dict[str, object]]: + """Capture saved-operand layout invariants before backward runs.""" + return [ + _snapshot_saved_quantized_operand_layout(saved_operand, name=name) + for name, saved_operand in guard_operands + ] + + +def _snapshot_backward_ctx_state( + output: torch.Tensor, +) -> tuple[str, bool, object, bool]: + if output.grad_fn is None: + raise RuntimeError("Output tensor has no grad_fn; cannot inspect backward context state.") + 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)] + if missing_attrs: + raise RuntimeError( + "grad_fn does not expose required backward context attributes: " + f"{', '.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")), + ) + + +def _assert_saved_quantized_operand_uses_rowwise_only( + saved_operand: Optional[torch.Tensor], + *, + name: str, +) -> None: + if saved_operand is None: + raise RuntimeError(f"Expected quantized saved {name} operand but got None") + if isinstance(saved_operand, torch.Tensor): + raise RuntimeError( + f"dequantized reference expects quantized saved {name} operand, got torch.Tensor." + ) + if not hasattr(saved_operand, "dequantize"): + raise RuntimeError(f"Unsupported saved {name} operand type: {type(saved_operand)}") + if hasattr(saved_operand, "_rowwise_data") and getattr(saved_operand, "_rowwise_data") is None: + raise RuntimeError( + f"Saved dequantized {name} operand lost row-wise fprop payload (likely usage retarget)." + ) + if ( + hasattr(saved_operand, "_columnwise_data") + and getattr(saved_operand, "_columnwise_data") is not None + ): + raise RuntimeError( + f"Saved dequantized {name} operand unexpectedly carries column-wise payload." + ) + + +def _assert_saved_quantized_operand_layout_unchanged(snapshot: dict[str, object]) -> None: + name = snapshot.get("name") + if not isinstance(name, str): + raise RuntimeError(f"Invalid saved operand snapshot name: {name!r}") + saved_operand = snapshot.get("saved_operand") + _assert_saved_quantized_operand_uses_rowwise_only(saved_operand, name=name) + + rowwise_present = snapshot.get("rowwise_present") + if isinstance(rowwise_present, bool): + rowwise_data_now = getattr(saved_operand, "_rowwise_data", None) + rowwise_now = rowwise_data_now is not None + if rowwise_now != rowwise_present: + raise RuntimeError( + f"Saved dequantized {name} operand row-wise payload presence changed " + f"from {rowwise_present} to {rowwise_now}." + ) + # Guard against hidden requantization that swaps in a new row-wise payload. + rowwise_obj_id = snapshot.get("rowwise_obj_id") + if ( + isinstance(rowwise_obj_id, int) + and rowwise_now + and id(rowwise_data_now) != rowwise_obj_id + ): + raise RuntimeError( + f"Saved dequantized {name} operand row-wise payload identity changed " + "(likely rewritten/requantized)." + ) + + columnwise_present = snapshot.get("columnwise_present") + if isinstance(columnwise_present, bool): + columnwise_now = getattr(saved_operand, "_columnwise_data", None) is not None + if columnwise_now != columnwise_present: + raise RuntimeError( + f"Saved dequantized {name} operand column-wise payload presence changed " + f"from {columnwise_present} to {columnwise_now}." + ) + + +def _assert_layout_invariants_unchanged(layout_invariants: list[dict[str, object]]) -> None: + """Validate saved-operand layout invariants after backward runs.""" + for layout_invariant in layout_invariants: + _assert_saved_quantized_operand_layout_unchanged(layout_invariant) + + +def _raise_if_ref_failed(ref_exc: Optional[Exception]) -> None: + """Re-raise deferred reference exceptions after layout checks.""" + if ref_exc is not None: + raise ref_exc + + +def _copy_named_parameters(src_module: torch.nn.Module, dst_module: torch.nn.Module) -> None: + src_params = dict(src_module.named_parameters()) + with torch.no_grad(): + for name, dst_param in dst_module.named_parameters(): + if name not in src_params: + raise RuntimeError(f"Parameter {name} missing in source module") + dst_param.copy_(src_params[name]) + + +def _compute_linear_backward_reference_from_saved_operands( + saved_input: Optional[torch.Tensor], + saved_weight: Optional[torch.Tensor], + dy: torch.Tensor, + *, + dequant_dtype: torch.dtype, + out_dtype: torch.dtype, + with_bias: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # dequantized reference path: + # 1) use the exact operands saved by quantized forward, + # 2) dequantize them to the active high-precision compute dtype, + # 3) run backward GEMMs in high precision and compare exactly. + for name, saved_operand in (("input", saved_input), ("weight", saved_weight)): + _assert_saved_quantized_operand_uses_rowwise_only(saved_operand, name=name) + dy_mat = dy.reshape(-1, dy.shape[-1]) + + # Empty-token chunks can happen in grouped/fused paths. Reference should be zeros. + if dy_mat.shape[0] == 0: + out_features = dy_mat.shape[-1] + if saved_input is None: + raise RuntimeError( + "Expected saved input operand for empty-chunk dequantized reference." + ) + in_features = saved_input.size(-1) + dx_ref = torch.zeros(*dy.shape[:-1], in_features, dtype=out_dtype, device=dy.device) + dw_ref = torch.zeros(out_features, in_features, dtype=out_dtype, device=dy.device) + db_ref = torch.zeros(out_features, dtype=out_dtype, device=dy.device) + return dx_ref, dw_ref, db_ref + + x_ref_full = _dequantize_saved_operand(saved_input, dequant_dtype) + x_ref = x_ref_full.reshape(-1, x_ref_full.shape[-1]) + w_ref = _dequantize_saved_operand(saved_weight, dequant_dtype) + + dx_ref_2d, *_ = general_gemm( + w_ref, + dy_mat, + out_dtype=out_dtype, + layout="NN", + grad=True, + use_split_accumulator=True, + ) + db_seed = ( + torch.empty(dy_mat.shape[-1], dtype=out_dtype, device=dy_mat.device) if with_bias else None + ) + # Derive db from the same GEMM primitive used by runtime wgrad when bias exists. + dw_ref, db_ref, *_ = general_gemm( + x_ref, + dy_mat, + out_dtype=out_dtype, + layout="NT", + grad=True, + bias=db_seed, + use_split_accumulator=True, + ) + if db_ref is None: + db_ref = dy_mat.sum(dim=0).to(out_dtype) + dx_ref = dx_ref_2d.view(*dy.shape[:-1], dx_ref_2d.shape[-1]) + return dx_ref, dw_ref, db_ref + + +def _run_single_step( + module: torch.nn.Module, + x: torch.Tensor, + dy: torch.Tensor, + fp8_recipe: Optional[recipe.Recipe], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=fp8_recipe) if fp8_recipe is not None else nullcontext() + ) + with autocast_ctx: + y = module(x_run) + if isinstance(y, tuple): + y = y[0] + y.backward(dy) + assert x_run.grad is not None + assert module.weight.grad is not None + bias = getattr(module, "bias", None) + bgrad = None if bias is None or bias.grad is None else bias.grad.detach().clone() + return ( + y.detach().clone(), + x_run.grad.detach().clone(), + module.weight.grad.detach().clone(), + bgrad, + ) + + +def _run_single_step_with_saved_operands( + module: torch.nn.Module, + x: torch.Tensor, + fp8_recipe: recipe.Recipe, +) -> tuple[ + torch.Tensor, + torch.Tensor, + list[Optional[torch.Tensor]], +]: + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=fp8_recipe): + y = module(x_run) + if isinstance(y, tuple): + y = y[0] + saved_operands = restore_from_saved(y.grad_fn.tensor_objects, list(y.grad_fn.saved_tensors)) + return y, x_run, saved_operands + + +def _run_grouped_linear_single_step( + module: te.GroupedLinear, + x: torch.Tensor, + m_splits: list[int], + dy: torch.Tensor, + fp8_recipe: Optional[recipe.Recipe], +) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor], list[Optional[torch.Tensor]]]: + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=fp8_recipe) if fp8_recipe is not None else nullcontext() + ) + with autocast_ctx: + y = module(x_run, m_splits) + y.backward(dy) + assert x_run.grad is not None + + dw = [getattr(module, f"weight{i}").grad.detach().clone() for i in range(module.num_gemms)] + db: list[Optional[torch.Tensor]] = [] + for i in range(module.num_gemms): + if module.use_bias: + db.append(getattr(module, f"bias{i}").grad.detach().clone()) + else: + db.append(None) + return y.detach().clone(), x_run.grad.detach().clone(), dw, db + + +def _run_grouped_linear_step_with_saved_operands( + module: te.GroupedLinear, + x: torch.Tensor, + m_splits: list[int], + fp8_recipe: recipe.Recipe, +) -> tuple[ + torch.Tensor, + torch.Tensor, + list[Optional[torch.Tensor]], +]: + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=fp8_recipe): + y = module(x_run, m_splits) + saved_operands = restore_from_saved(y.grad_fn.tensor_objects, list(y.grad_fn.saved_tensors)) + return y, x_run, saved_operands + + +def _run_fused_single_step( + pattern: str, + model: te_ops.Sequential, + x1: torch.Tensor, + dy: torch.Tensor, + fp8_recipe: Optional[recipe.Recipe], + *, + x2: Optional[torch.Tensor] = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + torch.Tensor, + Optional[torch.Tensor], +]: + model.zero_grad(set_to_none=True) + x1_run = x1.detach().clone().requires_grad_(True) + x2_run = x2.detach().clone().requires_grad_(True) if x2 is not None else None + autocast_ctx = ( + te.autocast(enabled=True, recipe=fp8_recipe) if fp8_recipe is not None else nullcontext() + ) + with autocast_ctx: + if pattern in ("bias_add", "scale_add"): + assert x2_run is not None + y = model(x1_run, x2_run) + else: + y = model(x1_run) + y.backward(dy) + assert x1_run.grad is not None + + dw = model[0].weight.grad.detach().clone() + db = None + if getattr(model[0], "bias", None) is not None and model[0].bias.grad is not None: + db = model[0].bias.grad.detach().clone() + dx2 = x2_run.grad.detach().clone() if x2_run is not None and x2_run.grad is not None else None + return y.detach().clone(), x1_run.grad.detach().clone(), dx2, dw, db + + +def _run_fused_single_step_with_saved_operands( + pattern: str, + model: te_ops.Sequential, + x1: torch.Tensor, + fp8_recipe: recipe.Recipe, + *, + x2: Optional[torch.Tensor] = None, +) -> tuple[ + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + list[Optional[torch.Tensor]], +]: + model.zero_grad(set_to_none=True) + x1_run = x1.detach().clone().requires_grad_(True) + x2_run = x2.detach().clone().requires_grad_(True) if x2 is not None else None + with te.autocast(enabled=True, recipe=fp8_recipe): + if pattern in ("bias_add", "scale_add"): + assert x2_run is not None + y = model(x1_run, x2_run) + else: + y = model(x1_run) + saved_operands = restore_from_saved(y.grad_fn.tensor_objects, list(y.grad_fn.saved_tensors)) + return y, x1_run, x2_run, saved_operands + + +def _run_quantize_op_single_step( + model: te_ops.Sequential, + x: torch.Tensor, + dy: torch.Tensor, + fp8_recipe: Optional[recipe.Recipe], +) -> tuple[torch.Tensor, torch.Tensor]: + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=fp8_recipe) if fp8_recipe is not None else nullcontext() + ) + with autocast_ctx: + y = model(x_run) + y.backward(dy) + assert x_run.grad is not None + return y.detach().clone(), x_run.grad.detach().clone() + + +def _run_single_step_with_ctx_state( + module: torch.nn.Module, + x: torch.Tensor, + dy: torch.Tensor, + fp8_recipe: Optional[recipe.Recipe], +) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + tuple[str, bool, object, bool], +]: + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=fp8_recipe) if fp8_recipe is not None else nullcontext() + ) + with autocast_ctx: + y = module(x_run) + if isinstance(y, tuple): + y = y[0] + ctx_state = _snapshot_backward_ctx_state(y) + y.backward(dy) + assert x_run.grad is not None + assert module.weight.grad is not None + bias = getattr(module, "bias", None) + bgrad = None if bias is None or bias.grad is None else bias.grad.detach().clone() + return ( + y.detach().clone(), + x_run.grad.detach().clone(), + module.weight.grad.detach().clone(), + bgrad, + ctx_state, + ) + + +def _run_grouped_linear_single_step_with_ctx_state( + module: te.GroupedLinear, + x: torch.Tensor, + m_splits: list[int], + dy: torch.Tensor, + fp8_recipe: Optional[recipe.Recipe], +) -> tuple[ + torch.Tensor, + torch.Tensor, + list[torch.Tensor], + list[Optional[torch.Tensor]], + tuple[str, bool, bool], +]: + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = ( + te.autocast(enabled=True, recipe=fp8_recipe) if fp8_recipe is not None else nullcontext() + ) + with autocast_ctx: + y = module(x_run, m_splits) + if y.grad_fn is None: + raise RuntimeError( + "Output tensor has no grad_fn; cannot inspect grouped backward state." + ) + required_attrs = ( + "backward_override", + "fp8", + "reduce_and_update_bwd_fp8_tensors", + ) + missing_attrs = [attr for attr in required_attrs if not hasattr(y.grad_fn, attr)] + if missing_attrs: + raise RuntimeError( + "Grouped grad_fn does not expose required backward context attributes: " + f"{', '.join(missing_attrs)}." + ) + ctx_state = ( + getattr(y.grad_fn, "backward_override"), + bool(getattr(y.grad_fn, "fp8")), + bool(getattr(y.grad_fn, "reduce_and_update_bwd_fp8_tensors")), + ) + y.backward(dy) + assert x_run.grad is not None + + dw = [getattr(module, f"weight{i}").grad.detach().clone() for i in range(module.num_gemms)] + db: list[Optional[torch.Tensor]] = [] + for i in range(module.num_gemms): + if module.use_bias: + db.append(getattr(module, f"bias{i}").grad.detach().clone()) + else: + db.append(None) + return y.detach().clone(), x_run.grad.detach().clone(), dw, db, ctx_state + + +# -------------------------- +# Tests +# -------------------------- + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +def test_backward_override_recipe_matches_requested_mode( + recipe_name: str, + backward_override: str, +) -> None: + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + quant_recipe = make_recipe(recipe_name) + assert mode_recipe.backward_override == backward_override + assert quant_recipe.backward_override is None + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("module_type", ("linear", "layernorm_linear", "ops_linear")) +@pytest.mark.parametrize("input_shape,out_features", _shape_test_cases) +@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) +@pytest.mark.parametrize("dtype", _core_dtypes, ids=str) +def test_linear_like_backward_override_matches_reference( + recipe_name: str, + module_type: str, + input_shape: tuple[int, ...], + out_features: int, + use_bias: bool, + dtype: torch.dtype, + backward_override: str, +) -> None: + reset_rng_states() + _maybe_skip_recipe_dtype(recipe_name, dtype, module_type) + _maybe_skip_unsupported_recipe_module_combo(recipe_name, module_type) + _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, module_type) + + in_features = input_shape[-1] + quantized_ref_recipe = make_recipe(recipe_name) + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override(module_type, mode_recipe, backward_override) + + module_quantized_ref = _make_linear_like_module( + module_type, + in_features, + out_features, + dtype, + bias=use_bias, + ) + module_bwd_mode = _make_linear_like_module( + module_type, + in_features, + out_features, + dtype, + bias=use_bias, + ) + _copy_named_parameters(module_quantized_ref, module_bwd_mode) + + output_shape = input_shape[:-1] + (out_features,) + x = torch.randn(*input_shape, dtype=dtype, device="cuda") + dy = torch.randn(*output_shape, dtype=dtype, device="cuda") + + y_quantized_ref, _, _, _ = _run_single_step(module_quantized_ref, x, dy, quantized_ref_recipe) + if backward_override == "high_precision": + # high_precision reference path: compare against a plain high-precision backward run + # (no fp8/autocast), starting from the same params and inputs. + module_unquantized_ref = _make_linear_like_module( + module_type, + in_features, + out_features, + dtype, + bias=use_bias, + ) + _copy_named_parameters(module_quantized_ref, module_unquantized_ref) + y_bwd_mode, dx_bwd_mode, dw_bwd_mode, db_bwd_mode = _run_single_step( + module_bwd_mode, + x, + dy, + mode_recipe, + ) + _, dx_ref, dw_ref, db_ref = _run_single_step( + module_unquantized_ref, + x, + dy, + None, + ) + else: + # dequantized reference path: capture saved forward operands from the real dequantized-override + # execution, then rebuild backward reference from those saved operands. + y_bwd_mode, x_bwd_mode, saved_operands = _run_single_step_with_saved_operands( + module_bwd_mode, x, mode_recipe + ) + y_bwd_mode_detached = y_bwd_mode.detach().clone() + + dx_ref: Optional[torch.Tensor] = None + dw_ref: Optional[torch.Tensor] = None + db_ref: Optional[torch.Tensor] = None + layout_invariants: list[dict[str, object]] = [] + guard_operands: list[tuple[str, Optional[torch.Tensor]]] = [] + ref_exc: Optional[Exception] = None + try: + if module_type == "layernorm_linear": + # LayerNormLinear dequantized reference: + # 1) Compute d(ln_out), dw, db from linear backward with saved operands. + # 2) Compute exact dx via layernorm_bwd with saved norm statistics. + # _LayerNormLinear forward saves operands as: + # [inputmat, weightmat, origin_weight, bias, ln_weight, ln_out, mu, rsigma, ...] + if len(saved_operands) < 8: + raise RuntimeError( + "Insufficient saved operands for layernorm_linear dequantized reference " + f"(got {len(saved_operands)}, expected at least 8)." + ) + saved_input = saved_operands[0] + saved_weight = saved_operands[1] + saved_ln_weight = saved_operands[4] + saved_ln_out = saved_operands[5] + saved_mu = saved_operands[6] + saved_rsigma = saved_operands[7] + guard_operands.extend( + [ + ("layernorm_linear_ln_out", saved_ln_out), + ("layernorm_linear_weight", saved_weight), + ] + ) + d_ln_out_ref, dw_ref, db_ref = ( + _compute_linear_backward_reference_from_saved_operands( + saved_ln_out, + saved_weight, + dy, + dequant_dtype=dtype, + out_dtype=dtype, + with_bias=use_bias, + ) + ) + input_ref = _dequantize_saved_operand(saved_input, dtype) + input_ref_2d = input_ref.reshape(-1, input_ref.shape[-1]) + ln_weight_ref = _dequantize_saved_operand(saved_ln_weight, dtype).view(-1) + if saved_mu is None or saved_rsigma is None: + raise RuntimeError("Missing LayerNorm statistics in saved operands") + if not isinstance(saved_mu, torch.Tensor) or not isinstance( + saved_rsigma, torch.Tensor + ): + raise RuntimeError("LayerNorm statistics must be Tensor objects") + dx_ref, *_ = layernorm_bwd( + d_ln_out_ref.reshape(input_ref_2d.shape), + input_ref_2d, + saved_mu, + saved_rsigma, + ln_weight_ref, + module_bwd_mode.bwd_ln_sm_margin, + module_bwd_mode.zero_centered_gamma, + ) + dx_ref = dx_ref.view_as(x_bwd_mode) + else: + saved_input, saved_weight = saved_operands[0], saved_operands[1] + guard_operands.extend( + [ + (f"{module_type}_input", saved_input), + (f"{module_type}_weight", saved_weight), + ] + ) + linear_wgrad_with_bias = use_bias and module_type != "ops_linear" + dx_ref, dw_ref, db_ref = _compute_linear_backward_reference_from_saved_operands( + saved_input, + saved_weight, + dy, + dequant_dtype=dtype, + out_dtype=dtype, + with_bias=linear_wgrad_with_bias, + ) + if module_type == "ops_linear" and use_bias: + # te_ops bias grad is reduced by the Bias op from incoming dy. + db_ref = dy.reshape(-1, dy.shape[-1]).sum(dim=0).to(dtype) + except Exception as exc: + ref_exc = exc + + layout_invariants = _snapshot_layout_invariants(guard_operands) + + y_bwd_mode.backward(dy) + assert x_bwd_mode.grad is not None + assert module_bwd_mode.weight.grad is not None + dx_bwd_mode = x_bwd_mode.grad.detach().clone() + dw_bwd_mode = module_bwd_mode.weight.grad.detach().clone() + bias = getattr(module_bwd_mode, "bias", None) + db_bwd_mode = None if bias is None or bias.grad is None else bias.grad.detach().clone() + y_bwd_mode = y_bwd_mode_detached + + _assert_layout_invariants_unchanged(layout_invariants) + _raise_if_ref_failed(ref_exc) + assert dx_ref is not None and dw_ref is not None and db_ref is not None + + assert_close(y_bwd_mode, y_quantized_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx_bwd_mode, dx_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dw_bwd_mode, dw_ref, rtol=0, atol=0, check_dtype=True) + if use_bias: + assert db_bwd_mode is not None + assert db_ref is not None + assert_close(db_bwd_mode, db_ref, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("in_features,out_features", _linear_feature_cases) +@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) +@pytest.mark.parametrize("m_splits", _grouped_m_split_cases) +@pytest.mark.parametrize("dtype", _core_dtypes, ids=str) +def test_grouped_linear_backward_override_matches_reference( + recipe_name: str, + in_features: int, + out_features: int, + use_bias: bool, + m_splits: list[int], + dtype: torch.dtype, + backward_override: str, +) -> None: + + reset_rng_states() + _maybe_skip_recipe_dtype(recipe_name, dtype, "grouped_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "grouped_linear") + _maybe_skip_unsupported_grouped_splits(recipe_name, m_splits) + num_gemms = len(m_splits) + num_tokens = sum(m_splits) + + quantized_ref_recipe = make_recipe(recipe_name) + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + + module_quantized_ref = te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + ) + module_bwd_mode = te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + ) + _copy_named_parameters(module_quantized_ref, module_bwd_mode) + + x = torch.randn(num_tokens, in_features, dtype=dtype, device="cuda") + dy = torch.randn(num_tokens, out_features, dtype=dtype, device="cuda") + + y_quantized_ref, _, _, _ = _run_grouped_linear_single_step( + module_quantized_ref, + x, + m_splits, + dy, + quantized_ref_recipe, + ) + if backward_override == "high_precision": + # high_precision reference path: grouped module in plain high precision. + module_unquantized_ref = te.GroupedLinear( + num_gemms, + in_features, + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + ) + _copy_named_parameters(module_quantized_ref, module_unquantized_ref) + y_bwd_mode, dx_bwd_mode, dw_bwd_mode, db_bwd_mode = _run_grouped_linear_single_step( + module_bwd_mode, + x, + m_splits, + dy, + mode_recipe, + ) + _, dx_ref, dw_ref, db_ref = _run_grouped_linear_single_step( + module_unquantized_ref, + x, + m_splits, + dy, + None, + ) + else: + # dequantized reference path for grouped GEMMs: + # each GEMM restores its own saved input/weight pair and computes its own ref grads. + y_bwd_mode, x_bwd_mode, saved_operands = _run_grouped_linear_step_with_saved_operands( + module_bwd_mode, x, m_splits, mode_recipe + ) + y_bwd_mode_detached = y_bwd_mode.detach().clone() + + dx_ref: Optional[torch.Tensor] = None + dw_ref: list[torch.Tensor] = [] + db_ref: list[Optional[torch.Tensor]] = [] + layout_invariants: list[dict[str, object]] = [] + guard_operands: list[tuple[str, Optional[torch.Tensor]]] = [] + ref_exc: Optional[Exception] = None + try: + if len(saved_operands) < 2 * num_gemms: + raise RuntimeError( + "Insufficient saved operands for GroupedLinear dequantized reference " + f"(got {len(saved_operands)}, expected at least {2 * num_gemms})." + ) + + saved_inputs = saved_operands[:num_gemms] + saved_weights = saved_operands[num_gemms : 2 * num_gemms] + for i, (saved_input, saved_weight) in enumerate(zip(saved_inputs, saved_weights)): + guard_operands.extend( + [ + (f"grouped_input{i}", saved_input), + (f"grouped_weight{i}", saved_weight), + ] + ) + dy_chunks = torch.split(dy, m_splits) + + dx_chunks = [] + dw_ref = [] + db_ref = [] + for dy_chunk, saved_input, saved_weight in zip(dy_chunks, saved_inputs, saved_weights): + dx_i, dw_i, db_i = _compute_linear_backward_reference_from_saved_operands( + saved_input, + saved_weight, + dy_chunk, + dequant_dtype=dtype, + out_dtype=dtype, + with_bias=use_bias, + ) + dx_chunks.append(dx_i) + dw_ref.append(dw_i) + db_ref.append(db_i if use_bias else None) + dx_ref = torch.cat(dx_chunks, dim=0) + except Exception as exc: + ref_exc = exc + + layout_invariants = _snapshot_layout_invariants(guard_operands) + + y_bwd_mode.backward(dy) + assert x_bwd_mode.grad is not None + dx_bwd_mode = x_bwd_mode.grad.detach().clone() + dw_bwd_mode = [ + getattr(module_bwd_mode, f"weight{i}").grad.detach().clone() + for i in range(module_bwd_mode.num_gemms) + ] + db_bwd_mode = [] + for i in range(module_bwd_mode.num_gemms): + if module_bwd_mode.use_bias: + db_bwd_mode.append(getattr(module_bwd_mode, f"bias{i}").grad.detach().clone()) + else: + db_bwd_mode.append(None) + y_bwd_mode = y_bwd_mode_detached + + _assert_layout_invariants_unchanged(layout_invariants) + _raise_if_ref_failed(ref_exc) + assert dx_ref is not None + + assert_close(y_bwd_mode, y_quantized_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx_bwd_mode, dx_ref, rtol=0, atol=0, check_dtype=True) + for test_dw, ref_dw in zip(dw_bwd_mode, dw_ref): + assert_close(test_dw, ref_dw, rtol=0, atol=0, check_dtype=True) + if use_bias: + for test_db, ref_db_i in zip(db_bwd_mode, db_ref): + assert test_db is not None + assert ref_db_i is not None + assert_close(test_db, ref_db_i, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("module_type", ("linear", "layernorm_linear")) +@pytest.mark.parametrize("input_shape,out_features", _shape_test_cases) +@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) +@pytest.mark.parametrize("dtype", _core_dtypes, ids=str) +def test_linear_like_runtime_backward_override_switch_updates_ctx( + recipe_name: str, + module_type: str, + input_shape: tuple[int, ...], + out_features: int, + use_bias: bool, + dtype: torch.dtype, + backward_override: str, +) -> None: + reset_rng_states() + _maybe_skip_recipe_dtype(recipe_name, dtype, module_type) + _maybe_skip_unsupported_recipe_module_combo(recipe_name, module_type) + _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, module_type) + + module = _make_linear_like_module( + module_type, + input_shape[-1], + out_features, + dtype, + bias=use_bias, + ) + x = torch.randn(*input_shape, dtype=dtype, device="cuda") + dy = torch.randn(*input_shape[:-1], out_features, dtype=dtype, device="cuda") + + default_recipe = make_recipe(recipe_name) + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override(module_type, mode_recipe, backward_override) + + *_, default_ctx = _run_single_step_with_ctx_state(module, x, dy, default_recipe) + ( + default_mode, + default_fp8, + default_grad_output_quantizer, + default_reduce_and_update, + ) = default_ctx + assert default_mode is None + assert default_fp8 + assert default_grad_output_quantizer is not None + assert default_reduce_and_update + + *_, switched_ctx = _run_single_step_with_ctx_state(module, x, dy, mode_recipe) + switched_mode, switched_fp8, switched_grad_output_quantizer, switched_reduce_and_update = ( + switched_ctx + ) + assert switched_mode == backward_override + assert not switched_fp8 + assert switched_grad_output_quantizer is None + assert not switched_reduce_and_update + + *_, default_ctx_after = _run_single_step_with_ctx_state(module, x, dy, default_recipe) + ( + default_mode_after, + default_fp8_after, + default_grad_output_quantizer_after, + default_reduce_and_update_after, + ) = default_ctx_after + assert default_mode_after is None + assert default_fp8_after + assert default_grad_output_quantizer_after is not None + assert default_reduce_and_update_after + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("in_features,out_features", _linear_feature_cases) +@pytest.mark.parametrize("m_splits", _grouped_m_split_cases) +@pytest.mark.parametrize("use_bias", (False, True), ids=("no_bias", "bias")) +@pytest.mark.parametrize("dtype", _core_dtypes, ids=str) +def test_grouped_linear_runtime_backward_override_switch_updates_ctx( + recipe_name: str, + in_features: int, + out_features: int, + m_splits: list[int], + use_bias: bool, + dtype: torch.dtype, + backward_override: str, +) -> None: + + reset_rng_states() + _maybe_skip_recipe_dtype(recipe_name, dtype, "grouped_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "grouped_linear") + _maybe_skip_unsupported_grouped_splits(recipe_name, m_splits) + + num_tokens = sum(m_splits) + module = te.GroupedLinear( + len(m_splits), + in_features, + out_features, + bias=use_bias, + params_dtype=dtype, + device="cuda", + ) + x = torch.randn(num_tokens, in_features, dtype=dtype, device="cuda") + dy = torch.randn(num_tokens, out_features, dtype=dtype, device="cuda") + + default_recipe = make_recipe(recipe_name) + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + + *_, default_ctx = _run_grouped_linear_single_step_with_ctx_state( + module, + x, + m_splits, + dy, + default_recipe, + ) + default_mode, default_fp8, default_reduce_and_update = default_ctx + assert default_mode is None + assert default_fp8 + assert default_reduce_and_update + + *_, switched_ctx = _run_grouped_linear_single_step_with_ctx_state( + module, + x, + m_splits, + dy, + mode_recipe, + ) + switched_mode, switched_fp8, switched_reduce_and_update = switched_ctx + assert switched_mode == backward_override + assert not switched_fp8 + assert not switched_reduce_and_update + + *_, default_ctx_after = _run_grouped_linear_single_step_with_ctx_state( + module, + x, + m_splits, + dy, + default_recipe, + ) + default_mode_after, default_fp8_after, default_reduce_and_update_after = default_ctx_after + assert default_mode_after is None + assert default_fp8_after + assert default_reduce_and_update_after + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize( + "fused_pattern,expected_fused_op", + ( + ("bias_add", ForwardLinearBiasAdd), + ("scale_add", ForwardLinearScaleAdd), + ), +) +@pytest.mark.parametrize("in_features,out_features", _linear_feature_cases) +@pytest.mark.parametrize("m", (1, 32), ids=("m1", "m32")) +@pytest.mark.parametrize("dtype", _fused_dtypes, ids=str) +def test_fused_linear_paths_match_backward_override_reference( + recipe_name: str, + fused_pattern: str, + expected_fused_op: type, + in_features: int, + out_features: int, + m: int, + dtype: torch.dtype, + backward_override: str, +) -> None: + _maybe_skip_recipe_dtype(recipe_name, dtype, "ops_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "ops_linear") + _maybe_skip_unsupported_recipe_shape(recipe_name, (m, in_features), "ops_linear") + + reset_rng_states() + + quantized_ref_recipe = make_recipe(recipe_name) + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("ops_linear", mode_recipe, backward_override) + + model_quantized_ref = _make_fused_model(fused_pattern, in_features, out_features, dtype) + model_bwd_mode = _make_fused_model(fused_pattern, in_features, out_features, dtype) + _copy_named_parameters(model_quantized_ref, model_bwd_mode) + + x1 = torch.randn(m, in_features, dtype=dtype, device="cuda") + x2 = None + if fused_pattern in ("bias_add", "scale_add"): + x2 = torch.randn(m, out_features, dtype=dtype, device="cuda") + dy = torch.randn(m, out_features, dtype=dtype, device="cuda") + + y_quantized_ref, _, _, _, _ = _run_fused_single_step( + fused_pattern, + model_quantized_ref, + x1, + dy, + quantized_ref_recipe, + x2=x2, + ) + + if backward_override == "high_precision": + # high_precision reference path: replay the same fused model structure in plain + # high precision and compare backward outputs exactly. + model_unquantized_ref = _make_fused_model(fused_pattern, in_features, out_features, dtype) + _copy_named_parameters(model_quantized_ref, model_unquantized_ref) + + y_bwd_mode, dx1_bwd_mode, dx2_bwd_mode, dw_bwd_mode, db_bwd_mode = _run_fused_single_step( + fused_pattern, + model_bwd_mode, + x1, + dy, + mode_recipe, + x2=x2, + ) + _, dx1_ref, dx2_ref, dw_ref, db_ref = _run_fused_single_step( + fused_pattern, + model_unquantized_ref, + x1, + dy, + None, + x2=x2, + ) + else: + # dequantized reference path: compute backward reference from saved quantized + # linear operands (with branch-specific dy handling for fused epilogues). + y_bwd_mode, x1_bwd_mode, x2_bwd_mode_ref, saved_operands = ( + _run_fused_single_step_with_saved_operands( + fused_pattern, + model_bwd_mode, + x1, + mode_recipe, + x2=x2, + ) + ) + y_bwd_mode_detached = y_bwd_mode.detach().clone() + dx1_ref: Optional[torch.Tensor] = None + dx2_ref: Optional[torch.Tensor] = None + dw_ref: Optional[torch.Tensor] = None + db_ref: Optional[torch.Tensor] = None + layout_invariants: list[dict[str, object]] = [] + guard_operands: list[tuple[str, Optional[torch.Tensor]]] = [] + ref_exc: Optional[Exception] = None + try: + saved_input, saved_weight = saved_operands[0], saved_operands[1] + guard_operands.extend( + [ + (f"fused_{fused_pattern}_input", saved_input), + (f"fused_{fused_pattern}_weight", saved_weight), + ] + ) + dy_for_linear = dy * 0.5 if fused_pattern == "scale_add" else dy + dx1_ref, dw_ref, db_ref = _compute_linear_backward_reference_from_saved_operands( + saved_input, + saved_weight, + dy_for_linear, + dequant_dtype=dtype, + out_dtype=dtype, + with_bias=False, + ) + dx2_ref = dy if x2 is not None else None + except Exception as exc: + ref_exc = exc + + layout_invariants = _snapshot_layout_invariants(guard_operands) + + y_bwd_mode.backward(dy) + assert x1_bwd_mode.grad is not None + dx1_bwd_mode = x1_bwd_mode.grad.detach().clone() + dx2_bwd_mode = ( + x2_bwd_mode_ref.grad.detach().clone() + if x2_bwd_mode_ref is not None and x2_bwd_mode_ref.grad is not None + else None + ) + dw_bwd_mode = model_bwd_mode[0].weight.grad.detach().clone() + db_bwd_mode = None + if ( + getattr(model_bwd_mode[0], "bias", None) is not None + and model_bwd_mode[0].bias.grad is not None + ): + db_bwd_mode = model_bwd_mode[0].bias.grad.detach().clone() + y_bwd_mode = y_bwd_mode_detached + + _assert_layout_invariants_unchanged(layout_invariants) + _raise_if_ref_failed(ref_exc) + assert dx1_ref is not None and dw_ref is not None + + fused_ops = model_bwd_mode._module_groups[0]._forward_ops + assert len(fused_ops) >= 1 + assert isinstance(fused_ops[0][0], expected_fused_op) + + assert_close(y_bwd_mode, y_quantized_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx1_bwd_mode, dx1_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dw_bwd_mode, dw_ref, rtol=0, atol=0, check_dtype=True) + if dx2_bwd_mode is not None and dx2_ref is not None: + assert_close(dx2_bwd_mode, dx2_ref, rtol=0, atol=0, check_dtype=True) + if db_bwd_mode is not None and db_ref is not None: + assert_close(db_bwd_mode, db_ref, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("input_shape", _bias_activation_shape_cases) +@pytest.mark.parametrize("out_features", _output_feature_cases) +@pytest.mark.parametrize("dtype", _fused_dtypes, ids=str) +def test_fused_bias_activation_matches_masked_linear_backward( + recipe_name: str, + input_shape: tuple[int, ...], + out_features: int, + dtype: torch.dtype, + backward_override: str, +) -> None: + _maybe_skip_recipe_dtype(recipe_name, dtype, "ops_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "ops_linear") + _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, "ops_linear") + + reset_rng_states() + in_features = input_shape[-1] + + quantized_ref_recipe = make_recipe(recipe_name) + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("ops_linear", mode_recipe, backward_override) + + model_quantized_ref = _make_fused_model("bias_activation", in_features, out_features, dtype) + model_bwd_mode = _make_fused_model("bias_activation", in_features, out_features, dtype) + _copy_named_parameters(model_quantized_ref, model_bwd_mode) + + x1 = torch.randn(*input_shape, dtype=dtype, device="cuda") + dy = torch.randn(*((*x1.shape[:-1], out_features)), dtype=dtype, device="cuda") + + y_quantized_ref, _, _, _, _ = _run_fused_single_step( + "bias_activation", + model_quantized_ref, + x1, + dy, + quantized_ref_recipe, + ) + + if backward_override == "high_precision": + # high_precision reference path: build a plain linear reference and apply the + # same activation mask (from quantized forward output) before backward. + linear_unquantized_ref = _make_linear_like_module( + "ops_linear", + in_features, + out_features, + dtype, + bias=True, + ) + _copy_named_parameters(model_bwd_mode[0], linear_unquantized_ref) + + y_bwd_mode, dx1_bwd_mode, _, dw_bwd_mode, db_bwd_mode = _run_fused_single_step( + "bias_activation", + model_bwd_mode, + x1, + dy, + mode_recipe, + ) + dy_after_activation = dy * (y_bwd_mode > 0).to(dy.dtype) + _, dx1_ref, dw_ref, db_ref = _run_single_step( + linear_unquantized_ref, + x1, + dy_after_activation, + None, + ) + else: + # dequantized reference path: restore saved linear operands from fused forward, + # apply the same activation mask, then run linear backward reference. + y_bwd_mode, x1_bwd_mode, _, saved_operands = _run_fused_single_step_with_saved_operands( + "bias_activation", + model_bwd_mode, + x1, + mode_recipe, + ) + y_bwd_mode_detached = y_bwd_mode.detach().clone() + dy_after_activation = dy * (y_bwd_mode > 0).to(dy.dtype) + dx1_ref: Optional[torch.Tensor] = None + dw_ref: Optional[torch.Tensor] = None + db_ref: Optional[torch.Tensor] = None + layout_invariants: list[dict[str, object]] = [] + guard_operands: list[tuple[str, Optional[torch.Tensor]]] = [] + ref_exc: Optional[Exception] = None + try: + saved_input, saved_weight = saved_operands[0], saved_operands[1] + guard_operands.extend( + [ + ("fused_bias_activation_input", saved_input), + ("fused_bias_activation_weight", saved_weight), + ] + ) + dx1_ref, dw_ref, db_ref = _compute_linear_backward_reference_from_saved_operands( + saved_input, + saved_weight, + dy_after_activation, + dequant_dtype=dtype, + out_dtype=dtype, + with_bias=False, + ) + except Exception as exc: + ref_exc = exc + + layout_invariants = _snapshot_layout_invariants(guard_operands) + + y_bwd_mode.backward(dy) + assert x1_bwd_mode.grad is not None + dx1_bwd_mode = x1_bwd_mode.grad.detach().clone() + dw_bwd_mode = model_bwd_mode[0].weight.grad.detach().clone() + db_bwd_mode = ( + model_bwd_mode[0].bias.grad.detach().clone() + if model_bwd_mode[0].bias.grad is not None + else None + ) + y_bwd_mode = y_bwd_mode_detached + + _assert_layout_invariants_unchanged(layout_invariants) + _raise_if_ref_failed(ref_exc) + assert dx1_ref is not None and dw_ref is not None and db_ref is not None + + fused_ops = model_bwd_mode._module_groups[0]._forward_ops + assert len(fused_ops) >= 1 + assert isinstance(fused_ops[0][0], ForwardLinearBiasActivation) + + # In high_precision/dequantized modes, backward-activation+bias fusion should be disabled. + bwd_mode_backward_ops = model_bwd_mode._module_groups[0]._backward_ops + assert not any(isinstance(op, BackwardActivationBias) for op, _ in bwd_mode_backward_ops) + + # Quantized reference should still use fused backward path. + quantized_ref_backward_ops = model_quantized_ref._module_groups[0]._backward_ops + assert any(isinstance(op, BackwardActivationBias) for op, _ in quantized_ref_backward_ops) + + assert_close(y_bwd_mode, y_quantized_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx1_bwd_mode, dx1_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dw_bwd_mode, dw_ref, rtol=0, atol=0, check_dtype=True) + assert db_bwd_mode is not None + assert db_ref is not None + assert_close(db_bwd_mode, db_ref, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("in_features,out_features", _linear_feature_cases) +@pytest.mark.parametrize("dtype", _core_dtypes, ids=str) +def test_operation_fuser_rebuilds_userbuffers_fusion_on_backward_override_switch( + recipe_name: str, + in_features: int, + out_features: int, + dtype: torch.dtype, + backward_override: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Simulate a distributed setup to exercise Userbuffers fusion eligibility + # without launching a multi-rank job. + monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(torch.distributed, "get_world_size", lambda *_args, **_kwargs: 2) + + # Use a mutable recipe holder so we can switch fusion behavior on the same + # fuser object and verify that the cached fusion plan is refreshed. + current_recipe = {"value": make_recipe(recipe_name)} + monkeypatch.setattr(FP8GlobalStateManager, "get_fp8_recipe", lambda: current_recipe["value"]) + + reset_rng_states() + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "ops_linear") + + # Build a Userbuffers-eligible fuser and representative inputs. + linear = te_ops.BasicLinear( + in_features, + out_features, + device="cuda", + dtype=dtype, + userbuffers_options={"comm_name": "qkv"}, + ) + linear.tensor_parallel_mode = "column" + linear.tensor_parallel_size = 2 + linear.sequence_parallel = True + bias = te_ops.Bias(out_features, device="cuda", dtype=dtype) + model = te_ops.Sequential(linear, bias) + model._module_groups = model._make_module_groups(model._modules.values()) + fuser = model._module_groups[0] + x = torch.randn(32, in_features, dtype=dtype, device="cuda", requires_grad=True) + extra_inputs = [() for _ in range(fuser._num_basic_ops)] + + quant_recipe = make_recipe(recipe_name) + skip_unsupported_backward_override("ops_linear", quant_recipe, backward_override) + fuser.maybe_fuse_ops( + is_grad_enabled=True, + recipe=quant_recipe, + input_=x, + extra_inputs=extra_inputs, + ) + assert any(isinstance(op, UserbuffersForwardLinear) for op, _ in fuser._forward_ops) + + non_quant_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("ops_linear", non_quant_recipe, backward_override) + current_recipe["value"] = non_quant_recipe + fuser.maybe_fuse_ops( + is_grad_enabled=True, + recipe=non_quant_recipe, + input_=x, + extra_inputs=extra_inputs, + ) + assert not any(isinstance(op, UserbuffersForwardLinear) for op, _ in fuser._forward_ops) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("dtype", _core_dtypes, ids=str) +def test_quantize_op_respects_backward_override( + recipe_name: str, + dtype: torch.dtype, + backward_override: str, +) -> None: + _maybe_skip_recipe_dtype(recipe_name, dtype, "ops_linear") + _maybe_skip_unsupported_recipe_module_combo(recipe_name, "ops_linear") + reset_rng_states() + + x = torch.randn(32, 64, dtype=dtype, device="cuda") + dy = torch.randn(32, 64, dtype=dtype, device="cuda") + + model_override = te_ops.Sequential(te_ops.Quantize(forward=True, backward=True)) + model_ref = te_ops.Sequential(te_ops.Quantize(forward=True, backward=False)) + + mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("ops_linear", mode_recipe, backward_override) + + y_override, dx_override = _run_quantize_op_single_step(model_override, x, dy, mode_recipe) + y_ref, dx_ref = _run_quantize_op_single_step(model_ref, x, dy, mode_recipe) + + assert_close(y_override, y_ref, rtol=0, atol=0, check_dtype=True) + assert_close(dx_override, dx_ref, rtol=0, atol=0, check_dtype=True) + + +@pytest.mark.parametrize("recipe_name", _quantized_numerics_recipe_list) +@pytest.mark.parametrize("module_type", ("linear", "layernorm_linear")) +def test_backward_override_memory_peak_report( + recipe_name: str, + module_type: str, +) -> None: + """Diagnostic-only memory report for None/high_precision/dequantized backward overrides.""" + reset_rng_states() + dtype = torch.bfloat16 + input_shape = (2048, 2048) + out_features = 2048 * 4 + in_features = input_shape[-1] + use_bias = True + + _maybe_skip_recipe_dtype(recipe_name, dtype, module_type) + _maybe_skip_unsupported_recipe_module_combo(recipe_name, module_type) + _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, module_type) + + base_module = _make_linear_like_module( + module_type, + in_features, + out_features, + dtype, + bias=use_bias, + ) + + 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") + mode_results: dict[str, dict[str, float] | str] = {} + + for mode in modes: + mode_str = "default" if mode is None else mode + # try: + mode_recipe = make_recipe(recipe_name, backward_override=mode) + + # Keep params identical across modes for a cleaner apples-to-apples read. + module = _make_linear_like_module( + module_type, + in_features, + out_features, + dtype, + bias=use_bias, + ) + _copy_named_parameters(base_module, module) + + # Warmup run to reduce first-use kernel setup noise. + _run_single_step(module, x, dy, mode_recipe) + + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = te.autocast(enabled=True, recipe=mode_recipe) + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + fwd_start_mem = torch.cuda.memory_allocated() + with autocast_ctx: + y = module(x_run) + if isinstance(y, tuple): + y = y[0] + torch.cuda.synchronize() + fwd_peak_alloc = float(torch.cuda.max_memory_allocated() - fwd_start_mem) + fwd_peak_reserved = float(torch.cuda.max_memory_reserved()) + + torch.cuda.reset_peak_memory_stats() + bwd_start_mem = torch.cuda.memory_allocated() + y.backward(dy) + torch.cuda.synchronize() + bwd_peak_alloc = float(torch.cuda.max_memory_allocated() - bwd_start_mem) + bwd_peak_reserved = float(torch.cuda.max_memory_reserved()) + + module.zero_grad(set_to_none=True) + x_run = x.detach().clone().requires_grad_(True) + autocast_ctx = te.autocast(enabled=True, recipe=mode_recipe) + + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + e2e_start_mem = torch.cuda.memory_allocated() + with autocast_ctx: + y = module(x_run) + if isinstance(y, tuple): + y = y[0] + y.backward(dy) + torch.cuda.synchronize() + e2e_peak_alloc = float(torch.cuda.max_memory_allocated() - e2e_start_mem) + e2e_peak_reserved = float(torch.cuda.max_memory_reserved()) + + mode_results[mode_str] = { + "fwd_peak_alloc_mb": fwd_peak_alloc / (1024**2), + "fwd_peak_reserved_mb": fwd_peak_reserved / (1024**2), + "bwd_peak_alloc_mb": bwd_peak_alloc / (1024**2), + "bwd_peak_reserved_mb": bwd_peak_reserved / (1024**2), + "e2e_peak_alloc_mb": e2e_peak_alloc / (1024**2), + "e2e_peak_reserved_mb": e2e_peak_reserved / (1024**2), + } + # except Exception as exc: # pragma: no cover - diagnostic reporting path + # mode_results[mode_str] = f"{type(exc).__name__}: {exc}" + + print( + "\n[backward_override_memory_peak_report] " + f"recipe={recipe_name} module_type={module_type} " + f"dtype={dtype} input_shape={input_shape} out_features={out_features}" + ) + print(" units=MB") + metric_col_width = 9 + delta_col_width = 18 + columns = ( + ("mode_str", delta_col_width), + ("fwd_alloc", metric_col_width), + ("bwd_alloc", metric_col_width), + ("e2e_alloc", metric_col_width), + ("fwd_resrv", metric_col_width), + ("bwd_resrv", metric_col_width), + ("e2e_resrv", metric_col_width), + ("delta_fwd", delta_col_width), + ("delta_bwd", delta_col_width), + ("delta_e2e", delta_col_width), + ) + print(" | ".join(f"{name:>{width}}" for name, width in columns)) + print("-+-".join("-" * width for _, width in columns)) + + def _format_delta_with_pct(delta: float, base: float) -> str: + if math.isclose(base, 0.0, abs_tol=1e-12): + return f"{delta:+.2f} (n/a)" + pct = 100.0 * delta / base + return f"{delta:+.2f} ({pct:+.2f}%)" + + default_metrics = mode_results.get("default") + for mode in modes: + mode_str = "default" if mode is None else mode + metrics = mode_results[mode_str] + if isinstance(metrics, str): + print(f"{mode_str:>{delta_col_width}} | ERROR: {metrics}") + continue + + if isinstance(default_metrics, dict): + delta_fwd = metrics["fwd_peak_alloc_mb"] - default_metrics["fwd_peak_alloc_mb"] + delta_bwd = metrics["bwd_peak_alloc_mb"] - default_metrics["bwd_peak_alloc_mb"] + delta_e2e = metrics["e2e_peak_alloc_mb"] - default_metrics["e2e_peak_alloc_mb"] + delta_fwd_str = _format_delta_with_pct(delta_fwd, default_metrics["fwd_peak_alloc_mb"]) + delta_bwd_str = _format_delta_with_pct(delta_bwd, default_metrics["bwd_peak_alloc_mb"]) + delta_e2e_str = _format_delta_with_pct(delta_e2e, default_metrics["e2e_peak_alloc_mb"]) + else: + delta_fwd_str = "n/a" + delta_bwd_str = "n/a" + delta_e2e_str = "n/a" + + print( + f"{mode_str:>{delta_col_width}} | " + f"{metrics['fwd_peak_alloc_mb']:{metric_col_width}.2f} | " + f"{metrics['bwd_peak_alloc_mb']:{metric_col_width}.2f} | " + f"{metrics['e2e_peak_alloc_mb']:{metric_col_width}.2f} | " + f"{metrics['fwd_peak_reserved_mb']:{metric_col_width}.2f} | " + f"{metrics['bwd_peak_reserved_mb']:{metric_col_width}.2f} | " + f"{metrics['e2e_peak_reserved_mb']:{metric_col_width}.2f} | " + f"{delta_fwd_str:>{delta_col_width}} | " + f"{delta_bwd_str:>{delta_col_width}} | " + f"{delta_e2e_str:>{delta_col_width}}" + ) diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 7da8dcf863..50196782f2 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -6,6 +6,7 @@ import contextlib import pytest import os +import copy import torch from typing import Optional, List from transformer_engine.pytorch.cpu_offload import ( @@ -18,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 +from utils import ModelConfig, skip_unsupported_backward_override import transformer_engine_torch as tex # Check supported quantization schemes @@ -416,9 +417,14 @@ def test_multiple_tensor_offload(self, recipe): class TestTELayers: @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) @pytest.mark.parametrize("recipe", quantization_recipes) - def test_sanity(self, layer_type, recipe): + @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) + def test_sanity(self, layer_type, recipe, backward_override): Utils.memory_leak_check() + skip_unsupported_backward_override(layer_type, recipe, backward_override) + if recipe is not None: + recipe = copy.deepcopy(recipe) + recipe.backward_override = backward_override # Skip ops-based layers with Float8BlockScaling recipe if ( layer_type in ["linear_op", "layernorm_mlp_ops"] @@ -458,9 +464,15 @@ def test_sanity(self, layer_type, recipe): @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) @pytest.mark.parametrize("recipe", quantization_recipes) - def test_memory(self, layer_type, recipe): + @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) + def test_memory(self, layer_type, recipe, backward_override): Utils.memory_leak_check() + skip_unsupported_backward_override(layer_type, recipe, backward_override) + if recipe is not None: + recipe = copy.deepcopy(recipe) + recipe.backward_override = backward_override + # Skip ops-based layers with Float8BlockScaling recipe if ( layer_type in ["linear_op", "layernorm_mlp_ops"] @@ -524,7 +536,13 @@ def test_memory(self, layer_type, recipe): out = out + 1 out = sync_function(out) del inp - assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) + if backward_override is None: + assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) + else: + assert ( + Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) + or Utils.get_cuda_memory_mb() <= init_cuda_memory + ) offloaded_memory_cpu = offload_ctx.offload_synchronizer.get_offloaded_total_size_mb() # This assertion verifies that the memory used by tensors on the CPU matches the memory saved from a layer. @@ -537,9 +555,15 @@ def test_memory(self, layer_type, recipe): @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) @pytest.mark.parametrize("recipe", quantization_recipes) - def test_manual_synchronization(self, recipe, layer_type): + @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) + def test_manual_synchronization(self, recipe, layer_type, backward_override): Utils.memory_leak_check() + skip_unsupported_backward_override(layer_type, recipe, backward_override) + if recipe is not None: + recipe = copy.deepcopy(recipe) + recipe.backward_override = backward_override + # Skip ops-based layers with Float8BlockScaling recipe if ( layer_type in ["linear_op", "layernorm_mlp_ops"] @@ -600,6 +624,7 @@ def test_manual_synchronization(self, recipe, layer_type): out_2.sum().backward() @pytest.mark.parametrize("recipe", quantization_recipes) + @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]) @pytest.mark.parametrize("retain_pinned_cpu_buffers", [True, False]) @@ -607,11 +632,17 @@ def test_manual_synchronization(self, recipe, layer_type): def test_numerics( self, recipe, + backward_override, layer_type, use_cuda_graphs, backend, retain_pinned_cpu_buffers, ): + skip_unsupported_backward_override(layer_type, recipe, backward_override) + if recipe is not None: + recipe = copy.deepcopy(recipe) + recipe.backward_override = backward_override + # Skip ops-based layers with Float8BlockScaling recipe if ( layer_type in ["linear_op", "layernorm_mlp_ops"] diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 1b9e11792e..a782dadc60 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -4,6 +4,7 @@ from typing import Callable, Dict, Iterable, List, Tuple, Union import pytest +import copy import torch from transformer_engine.pytorch import ( @@ -24,7 +25,7 @@ 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 +from utils import ModelConfig, reset_rng_states, skip_unsupported_backward_override # Check if FP8 is supported. fp8_available = is_fp8_available() @@ -360,6 +361,7 @@ def _test_cuda_graphs( @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("backward_override", (None, "high_precision", "dequantized")) def test_make_graphed_callables( *, module: str, @@ -368,10 +370,17 @@ def test_make_graphed_callables( dtype: torch.dtype, fp8_params: bool, fp8_recipe: recipe.Recipe, + backward_override: str, fp8_weight_caching: bool = False, ) -> None: fp8 = fp8_recipe is not None + + skip_unsupported_backward_override(module, fp8_recipe, backward_override) + if fp8: + fp8_recipe = copy.deepcopy(fp8_recipe) + fp8_recipe.backward_override = backward_override + if fp8_params and not fp8: pytest.skip("FP8 needed for FP8 parameters.") if fp8_weight_caching and not fp8: @@ -440,18 +449,21 @@ 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("backward_override", (None, "high_precision", "dequantized")) def test_make_graphed_callables_with_fp8_weight_caching( *, module: str, dtype: torch.dtype, fp8_params: bool, fp8_recipe: recipe.Recipe, + backward_override: str, ) -> None: test_make_graphed_callables( module=module, dtype=dtype, fp8_params=fp8_params, fp8_recipe=fp8_recipe, + backward_override=backward_override, fp8_weight_caching=True, ) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index f87e44373e..be123f8c23 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -7,6 +7,7 @@ import torch import pytest import os +import copy import transformer_engine import transformer_engine.pytorch as te @@ -37,7 +38,7 @@ 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 +from utils import ModelConfig, skip_unsupported_backward_override # Only run FP8 tests on supported devices. fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) @@ -395,6 +396,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("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @pytest.mark.parametrize("zero_centered_gamma", all_boolean) @@ -404,6 +406,7 @@ def test_sanity_normalization_amp(dtype, model, skip_wgrad, skip_dgrad, normaliz def test_sanity_layernorm_linear( dtype, fp8_recipe, + backward_override, model, skip_wgrad, zero_centered_gamma, @@ -413,6 +416,11 @@ def test_sanity_layernorm_linear( ): config = model_configs[model] + skip_unsupported_backward_override("layernorm_linear", fp8_recipe, backward_override) + if fp8_recipe is not None: + fp8_recipe = copy.deepcopy(fp8_recipe) + fp8_recipe.backward_override = backward_override + if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") @@ -436,13 +444,21 @@ def test_sanity_layernorm_linear( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @pytest.mark.parametrize("skip_dgrad", all_boolean) @pytest.mark.parametrize("microbatching", all_boolean) -def test_sanity_linear(dtype, fp8_recipe, model, skip_wgrad, skip_dgrad, microbatching): +def test_sanity_linear( + dtype, fp8_recipe, backward_override, model, skip_wgrad, skip_dgrad, microbatching +): config = model_configs[model] + skip_unsupported_backward_override("linear", fp8_recipe, backward_override) + if fp8_recipe is not None: + fp8_recipe = copy.deepcopy(fp8_recipe) + fp8_recipe.backward_override = backward_override + if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") @@ -466,13 +482,21 @@ def test_sanity_linear(dtype, fp8_recipe, model, skip_wgrad, skip_dgrad, microba @pytest.mark.parametrize("bs", batch_sizes_with_zero) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) -def test_sanity_linear_with_zero_tokens(dtype, bs, model, fp8_recipe, fp8_model_params, use_bias): +def test_sanity_linear_with_zero_tokens( + dtype, bs, model, fp8_recipe, backward_override, fp8_model_params, use_bias +): config = model_configs[model] ffn_hidden_size = 4 * config.hidden_size num_tokens = bs * config.max_seqlen_q + skip_unsupported_backward_override("linear", fp8_recipe, backward_override) + if fp8_recipe is not None: + fp8_recipe = copy.deepcopy(fp8_recipe) + fp8_recipe.backward_override = backward_override + if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") @@ -499,6 +523,7 @@ def test_sanity_linear_with_zero_tokens(dtype, bs, model, fp8_recipe, fp8_model_ @pytest.mark.parametrize("bs", batch_sizes_with_zero) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) @pytest.mark.parametrize("single_param", all_boolean) @@ -509,6 +534,7 @@ def test_sanity_grouped_linear( bs, model, fp8_recipe, + backward_override, fp8_model_params, use_bias, single_param, @@ -521,6 +547,11 @@ def test_sanity_grouped_linear( bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) + skip_unsupported_backward_override("grouped_linear", fp8_recipe, backward_override) + if fp8_recipe is not None: + fp8_recipe = copy.deepcopy(fp8_recipe) + fp8_recipe.backward_override = backward_override + if fp8_recipe is not None: if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 929f02453d..196ae8c165 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -11,6 +11,7 @@ from typing import Optional, Sequence, Tuple, Dict, Any, List from packaging.version import Version as PkgVersion +import pytest import torch import transformer_engine @@ -118,7 +119,7 @@ def quantization_tols(name: str) -> dict[str, float]: raise ValueError(f"Unsupported quantization scheme ({name})") -def make_recipe(name: Optional[str]) -> Optional[Recipe]: +def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: """Make recipe for quantization scheme""" if name is None: return None @@ -126,26 +127,52 @@ def make_recipe(name: Optional[str]) -> Optional[Recipe]: return transformer_engine.common.recipe.DelayedScaling( fp8_format=transformer_engine.common.recipe.Format.E4M3, amax_history_len=8, + **recipe_kwargs, ) if name == "fp8_current_scaling": return transformer_engine.common.recipe.Float8CurrentScaling( fp8_format=transformer_engine.common.recipe.Format.E4M3, + **recipe_kwargs, ) if name == "mxfp8": return transformer_engine.common.recipe.MXFP8BlockScaling( fp8_format=transformer_engine.common.recipe.Format.E4M3, + **recipe_kwargs, ) if name == "fp8_block_scaling": - return transformer_engine.common.recipe.Float8BlockScaling() + 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, ) raise ValueError(f"Unsupported quantization scheme ({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 backward_override is None: + return + if quant_recipe is None and backward_override is not None: + pytest.skip(f"Not a quantized recipe, cannot use backward override {backward_override}.") + if quant_recipe.delayed() and backward_override is not None: + pytest.skip(f"Delayed scaling does not support backward override {backward_override}.") + if layer_type in ( + "layernorm_mlp", + "layernorm_mlp_nocheckpoint", + "layernorm_mlp_checkpoint", + "transformer", + "transformer_layer", + ): + pytest.skip(f"{layer_type} does not support NVTE_BACKWARD_OVERRIDE={backward_override}.") + + # Cached RNG state _rng_states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 18577b0eb4..67b6f87067 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -11,6 +11,9 @@ from pydantic.dataclasses import dataclass +_BACKWARD_OVERRIDES = (None, "high_precision", "dequantized") + + class _FormatHelper(NamedTuple): """ Stores max FP8 values for fprop and bprop a `Format`. @@ -188,6 +191,8 @@ def scaling_factor_compute(amax: Tensor, `LayerNormLinear (BF16 output) -> (cast to FP8 ) FP8 DPA (cast to BF16) -> Linear`. When `fp8_mha = True, fp8_dpa = True`, it becomes `LayerNormLinear (FP8 output) -> FP8 DPA -> Linear`. + backward_override : {None, 'high_precision', 'dequantized'}, default = None + Backward precision mode. Delayed scaling only supports None. Notes ----- @@ -211,9 +216,16 @@ def scaling_factor_compute(amax: Tensor, reduce_amax: bool = True fp8_dpa: bool = False fp8_mha: bool = False + backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert ( + self.backward_override in _BACKWARD_OVERRIDES + ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + assert ( + self.backward_override is None + ), "Delayed scaling only supports backward_override=None." def __repr__(self) -> str: return ( @@ -223,7 +235,8 @@ def __repr__(self) -> str: f"amax_history_len={self.amax_history_len}, " f"reduce_amax={self.reduce_amax}, " f"fp8_dpa={self.fp8_dpa}, " - f"fp8_mha={self.fp8_mha}" + f"fp8_mha={self.fp8_mha}, " + f"backward_override={self.backward_override}" ) @@ -237,6 +250,11 @@ class Float8CurrentScaling(Recipe): fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.HYBRID Controls the FP8 data format used during forward and backward pass. + 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, + and `dequantized` dequantizes saved operands to the active high-precision + compute dtype (e.g. BF16/FP16/FP32) for backward. """ use_power_2_scales: bool = os.getenv("NVTE_FP8_CURRENT_SCALING_POWER_2_SCALES", "0") == "1" @@ -249,9 +267,13 @@ class Float8CurrentScaling(Recipe): fp8_gemm_wgrad: MMParams = MMParams(use_split_accumulator=True) fp8_dpa: bool = False fp8_mha: bool = False + backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert ( + self.backward_override in _BACKWARD_OVERRIDES + ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." def __repr__(self) -> str: return ( @@ -264,7 +286,8 @@ def __repr__(self) -> str: f"fp8_gemm_dgrad={self.fp8_gemm_dgrad}, " f"fp8_gemm_wgrad={self.fp8_gemm_wgrad}, " f"fp8_dpa={self.fp8_dpa}, " - f"fp8_mha={self.fp8_mha}" + f"fp8_mha={self.fp8_mha}, " + f"backward_override={self.backward_override}" ) @@ -291,21 +314,31 @@ class MXFP8BlockScaling(Recipe): fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.E4M3 Controls the FP8 data format used during forward and backward pass. + 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, + and `dequantized` dequantizes saved operands to the active high-precision + compute dtype (e.g. BF16/FP16/FP32) for backward. """ margin: int = 0 fp8_format: Format = Format.E4M3 fp8_dpa: bool = False fp8_mha: bool = False + backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert ( + self.backward_override in _BACKWARD_OVERRIDES + ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." def __repr__(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"margin={self.margin}, " - f"format={str(self.fp8_format).split('.')[1]}" + f"format={str(self.fp8_format).split('.')[1]}, " + f"backward_override={self.backward_override}" ) @@ -334,6 +367,11 @@ class Float8BlockScaling(Recipe): fp8_format : {Format.E4M3, Format.HYBRID}, default = Format.E4M3 Controls the FP8 data format used during forward and backward pass. + 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, + and `dequantized` dequantizes saved operands to the active high-precision + compute dtype (e.g. BF16/FP16/FP32) for backward. """ use_f32_scales: bool = os.getenv("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "0") == "1" @@ -350,6 +388,7 @@ class Float8BlockScaling(Recipe): fp8_gemm_wgrad: MMParams = MMParams(use_split_accumulator=True) fp8_dpa: bool = False fp8_mha: bool = False + backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: assert self.x_block_scaling_dim in [1, 2], "Only 1D or 2D blocks supported for x" @@ -371,6 +410,9 @@ def __post_init__(self) -> None: not self.fp8_dpa and not self.fp8_mha ), "FP8 attention is not supported for Float8BlockScaling." assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." + assert ( + self.backward_override in _BACKWARD_OVERRIDES + ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." def __repr__(self) -> str: return ( @@ -386,7 +428,8 @@ def __repr__(self) -> str: f"fp8_gemm_dgrad={self.fp8_gemm_dgrad}, " f"fp8_gemm_wgrad={self.fp8_gemm_wgrad}, " f"fp8_dpa={self.fp8_dpa}, " - f"fp8_mha={self.fp8_mha}" + f"fp8_mha={self.fp8_mha}, " + f"backward_override={self.backward_override}" ) @@ -435,6 +478,11 @@ 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. + 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, + and `dequantized` dequantizes saved operands to the active high-precision + compute dtype (e.g. BF16/FP16/FP32) for backward. """ # Configuration envvars @@ -450,10 +498,14 @@ class NVFP4BlockScaling(Recipe): # Not applying quantization to attention for now fp8_dpa: bool = False fp8_mha: bool = False + backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) def __post_init__(self) -> None: assert self.fp4_format == Format.E2M1, "Only E2M1 is supported for NVFP4 scaling" assert self.fp8_format == Format.E4M3, "Only E4M3 is supported for NVFP4 scaling" + assert ( + self.backward_override in _BACKWARD_OVERRIDES + ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." # Quantization params # Note: RHT is currently only applied to column-wise usage so that @@ -481,6 +533,7 @@ def __repr__(self) -> str: f"fp8_format={str(self.fp8_format).split('.')[1]}, " f"fp8_dpa={self.fp8_dpa}, " f"fp8_mha={self.fp8_mha}, " + f"backward_override={self.backward_override}, " 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}, " @@ -512,12 +565,27 @@ class CustomRecipe(Recipe): - 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, + and `dequantized` dequantizes saved operands to the active high-precision + compute dtype (e.g. BF16/FP16/FP32) for backward. """ qfactory: Callable[..., Any] fp8_dpa: bool = False fp8_mha: bool = False + backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + + def __post_init__(self) -> None: + assert ( + self.backward_override in _BACKWARD_OVERRIDES + ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." def __repr__(self) -> str: - return f"recipe_type={self.__class__.__name__}, qfactory={self.qfactory}" + return ( + f"recipe_type={self.__class__.__name__}, " + f"qfactory={self.qfactory}, " + f"backward_override={self.backward_override}" + ) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a96a87bf89..1b237ece29 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1184,9 +1184,10 @@ def grad_output_preprocess( grad_output = grad_output.reshape((-1, grad_output.shape[-1])) grad_output = grad_output.contiguous() gather_grad_output = row_parallel_mode and ctx.sequence_parallel + use_fp8_bwd = ctx.fp8 and ctx.backward_override is None # Non-FP8 case: bgrad is fused with wgrad for this case. - if not ctx.fp8 and not ctx.debug: + if not use_fp8_bwd and not ctx.debug: if gather_grad_output: if not ctx.ub_overlap_ag: # Perform NCCL all-gather grad_output, _ = gather_along_first_dim(grad_output, ctx.tp_group) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index ba6becb9f9..2cce6c3ef8 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -97,6 +97,12 @@ def forward( save_original_input, debug, ) = non_tensor_args + if fp8: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + if backward_override == "high_precision": + save_original_input = True num_gemms = len(m_splits) weights = weights_and_biases[:num_gemms] @@ -112,10 +118,15 @@ def forward( input_quantizer.set_usage( rowwise=True, columnwise=( - is_grad_enabled and weight_requires_grad and not save_original_input + is_grad_enabled + and weight_requires_grad + and not save_original_input + and backward_override is None ), ) columnwise_usage = is_grad_enabled and inp.requires_grad + if backward_override is not None: + columnwise_usage = False if not columnwise_usage: columnwise_usage = ( is_fp8_activation_recompute_enabled() @@ -240,7 +251,12 @@ def forward( else: for inputmat in inputmats: if isinstance(inputmat, QuantizedTensorStorage): - inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) + if backward_override is not None: + # In dequantized mode we should dequantize directly from + # fprop quantized layouts without retargeting usage. + inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) else: inputmats = [None] * num_gemms @@ -291,6 +307,7 @@ def forward( 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.fuse_wgrad_accumulation = fuse_wgrad_accumulation ctx.cpu_offloading = cpu_offloading ctx.is_first_microbatch = is_first_microbatch @@ -309,6 +326,19 @@ def forward( ctx.save_original_input = save_original_input ctx.input_quantizers = input_quantizers + # 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_quantizers = [None] * num_gemms + ctx.grad_weight_quantizers = [None] * num_gemms + ctx.grad_output_quantizers = [None] * num_gemms + ctx.reduce_and_update_bwd_fp8_tensors = False + # [*, in_features] -> [*, out_features] except first dimension changes for SP return out.view(-1, *inp.shape[1:-1], out.shape[-1]) @@ -403,13 +433,32 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], dtype=ctx.activation_dtype, device=ctx.device, ) + weights_for_dgrad = weights + if ctx.backward_override == "dequantized": + weights_for_dgrad = [ + ( + weight.dequantize(dtype=ctx.activation_dtype) + if isinstance(weight, QuantizedTensorStorage) + else cast_if_needed(weight, ctx.activation_dtype) + ) + for weight in weights + ] + elif ctx.backward_override == "high_precision": + weights_for_dgrad = [ + ( + weight.dequantize(dtype=ctx.activation_dtype) + if isinstance(weight, QuantizedTensorStorage) + else cast_if_needed(weight, ctx.activation_dtype) + ) + for weight in origin_weights + ] # Make sure weights are available in column-wise format # for dgrad computation. - for weight in weights: + for weight in weights_for_dgrad: if isinstance(weight, QuantizedTensorStorage): weight.update_usage(columnwise_usage=True) general_grouped_gemm( - weights, + weights_for_dgrad, grad_output, [dgrad], ctx.grad_input_quantizers, @@ -464,6 +513,30 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], inputmats = torch.split( cast_if_needed(inp_view, ctx.activation_dtype), ctx.m_splits ) + elif ctx.backward_override == "dequantized": + inputmats_dequant = [] + for m_split, inputmat in zip(ctx.m_splits, 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) + ) + else: + inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype)) + inputmats = inputmats_dequant grouped_gemm_wgrad = functools.partial( general_grouped_gemm, quantization_params=ctx.grad_weight_quantizers, @@ -1237,6 +1310,15 @@ 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 ed91bc1235..dc021ca6b7 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -140,6 +140,10 @@ def forward( symmetric_ar_type, debug, ) = non_tensor_args + if fp8: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None # NVTX label for profiling nvtx_label = "transformer_engine._LayerNormLinear.forward" @@ -198,7 +202,10 @@ def forward( if fp8: if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") - input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) + input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and backward_override is None, + ) if with_input_all_gather and input_quantizer.supports_only_rowwise_all_gather(): # All-gather is not supported with FP8 column-wise data input_quantizer.set_usage(columnwise=False) @@ -211,6 +218,7 @@ def forward( and not debug and not return_layernorm_output and not return_layernorm_output_gathered + and backward_override is None and not custom # TODO(negvet): and not FP8GlobalStateManager.get_fp8_recipe().custom() ) @@ -234,6 +242,7 @@ def forward( ln_out_return = None if return_layernorm_output or return_layernorm_output_gathered: ln_out_return = ln_out + ln_out_hp = ln_out if backward_override == "high_precision" else None # ------------------------------------------------------ # Prepare GEMM input tensor @@ -295,7 +304,10 @@ def forward( if is_weight_param_quantized and not debug: weight_quantizer = weight._quantizer elif weight_quantizer is not None: - weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) + weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and backward_override is None, + ) # Get quantized weight update_workspace = is_first_microbatch is None or is_first_microbatch @@ -408,13 +420,16 @@ def forward( # ------------------------------------------------------ if is_grad_enabled: + ln_out_to_save = ln_out + if backward_override == "high_precision": + ln_out_to_save = ln_out_hp ctx.weight_quantizer = weight_quantizer ctx.ln_out_needs_gather = ( weight.requires_grad and parallel_mode == "column" and sequence_parallel ) # Input with column-wise usage is needed for wgrad GEMM. - if backward_needs_input: + if backward_needs_input and backward_override is None: if isinstance(ln_out, QuantizedTensorStorage): # For sequence parallel in vanilla FP8, rowwise data is # to gather the input. For MXFP8, columnwise only data @@ -426,7 +441,7 @@ def forward( ln_out.update_usage(rowwise_usage=False) if cpu_offloading: - mark_activation_offload(inputmat, mu, rsigma, ln_out) + mark_activation_offload(inputmat, mu, rsigma, ln_out_to_save) # Scatter intermediate/activation tensors saved for the backward pass # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already @@ -438,7 +453,7 @@ def forward( mu, rsigma, weightmat if fp8 and not is_weight_param_quantized else None, - ln_out if weight.requires_grad else None, + ln_out_to_save if weight.requires_grad else None, ) nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") @@ -465,7 +480,7 @@ def forward( weight, bias, ln_weight, - ln_out, + ln_out_to_save, mu, rsigma, ) @@ -492,6 +507,7 @@ def forward( 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.fuse_wgrad_accumulation = fuse_wgrad_accumulation ctx.cpu_offloading = cpu_offloading ctx.is_first_microbatch = is_first_microbatch @@ -522,6 +538,19 @@ def forward( ctx.wgrad_store = wgrad_store ctx.debug = debug + # 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 + ctx.reduce_and_update_bwd_fp8_tensors = False + # ------------------------------------------------------ # Cached state for backward pass is ready... # ------------------------------------------------------ @@ -657,9 +686,14 @@ def backward( # -------------------------------------------------- ln_out_total = None ln_out_total_work = None + if ctx.backward_override == "dequantized": + if isinstance(ln_out, QuantizedTensorStorage): + ln_out = ln_out.dequantize(dtype=ctx.activation_dtype) + else: + ln_out = cast_if_needed(ln_out, ctx.activation_dtype) if ctx.ln_out_needs_gather: quantizer = None - if ctx.input_quantizer is not None: + if ctx.input_quantizer is not None and ctx.fp8: quantizer = ctx.input_quantizer if quantizer.supports_only_rowwise_all_gather(): # If data is in FP8, we compute FP8 transposes manually @@ -697,7 +731,11 @@ def backward( # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) - if ctx.weight_quantizer is not None and isinstance(weight, QuantizedTensorStorage): + if ( + ctx.fp8 + and ctx.weight_quantizer is not None + and isinstance(weight, QuantizedTensorStorage) + ): weight.update_usage(columnwise_usage=True) # Choose whether to use GEMM kernel with split accumulator @@ -724,8 +762,18 @@ def backward( # dgrad GEMM # Note: dx = dy * w nvtx_range_push(f"{nvtx_label}.dgrad_gemm") + weight_for_dgrad = weight + if ctx.backward_override == "dequantized": + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) + else: + weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) + elif ctx.backward_override == "high_precision": + weight_for_dgrad = origin_weight + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( - weight, + weight_for_dgrad, grad_output, layout="NN", grad=True, @@ -1626,6 +1674,13 @@ 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/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index cc3dcc4064..a99de65c4a 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -234,6 +234,15 @@ def _forward( debug, recompute_for_bwd, ) = non_tensor_args + if fp8: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + assert backward_override is None, ( + "NVTE_BACKWARD_OVERRIDE=high_precision/dequantized is not implemented in LayerNormMLP." + " Replace LayerNormMLP with LayerNormLinear + Linear to enable" + " high_precision/dequantized backward." + ) # if grad is enabled and this is not the bwd stage, we must save this so bwd knows which path to take if is_grad_enabled and not recompute_for_bwd: @@ -780,6 +789,7 @@ def _forward( ctx.fc2_main_grad_func = lambda: fc2_weight.main_grad ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + ctx.backward_override = backward_override ctx.fc1_grad_input_quantizer = fc1_grad_input_quantizer ctx.fc1_grad_weight_quantizer = fc1_grad_weight_quantizer ctx.fc1_grad_output_quantizer = fc1_grad_output_quantizer diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index ea921341a4..8510f6cf8f 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -128,6 +128,12 @@ def forward( save_original_input, debug, ) = non_tensor_args + if fp8: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + if backward_override == "high_precision": + save_original_input = True # NVTX label for profiling nvtx_label = "transformer_engine._Linear.forward" @@ -187,7 +193,10 @@ def forward( raise ValueError("Missing quantizer for input tensor") if not isinstance(inputmat, QuantizedTensorStorage) and not custom: own_quantized_input = True - input_quantizer.set_usage(rowwise=True, columnwise=backward_needs_input) + input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and backward_override is None, + ) if isinstance( input_quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer) ): @@ -229,7 +238,12 @@ def forward( if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") input_quantizer.set_usage( - rowwise=True, columnwise=backward_needs_input and not save_original_input + rowwise=True, + columnwise=( + backward_needs_input + and not save_original_input + and backward_override is None + ), ) inputmat = input_quantizer(inputmat) own_quantized_input = True @@ -254,6 +268,8 @@ def forward( # 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 + if backward_override is not None: + columnwise_usage = False if not columnwise_usage: columnwise_usage = ( is_fp8_activation_recompute_enabled() @@ -387,7 +403,11 @@ def forward( and own_quantized_input and isinstance(inputmat, QuantizedTensorStorage) ): - if ( + if backward_override is not None: + # In dequantized mode we should dequantize directly from the + # fprop quantized tensor layout without retargeting usage. + inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + elif ( ctx.backward_input_needs_gather and weight_quantizer.supports_only_rowwise_all_gather() ): @@ -442,6 +462,7 @@ def forward( 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.input_quantizer = input_quantizer ctx.grad_input_quantizer = grad_input_quantizer ctx.grad_weight_quantizer = grad_weight_quantizer @@ -485,6 +506,19 @@ def forward( FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module ctx.wgrad_store = wgrad_store + # 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 + ctx.reduce_and_update_bwd_fp8_tensors = False + # ------------------------------------------------------ # Cached state for backward pass is ready... # ------------------------------------------------------ @@ -684,8 +718,10 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) - if ctx.weight_quantizer is not None and isinstance( - weight_fp8, QuantizedTensorStorage + if ( + ctx.fp8 + and ctx.weight_quantizer is not None + and isinstance(weight_fp8, QuantizedTensorStorage) ): weight_fp8.update_usage(columnwise_usage=True) @@ -714,8 +750,18 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], # Note: dx = dy * w nvtx_range_push(f"{nvtx_label}.dgrad_gemm") + weight_for_dgrad = weight_fp8 + if ctx.backward_override == "dequantized": + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) + else: + weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) + elif ctx.backward_override == "high_precision": + weight_for_dgrad = weight + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( - weight_fp8, + weight_for_dgrad, grad_output, layout="NN", grad=True, @@ -1490,6 +1536,13 @@ 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 48376a297f..17594726cc 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -332,12 +332,15 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: # 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: + columnwise_usage = False input_quantizer = self.get_quantizer("forward", 0) weight_quantizer = self.get_quantizer("forward", 1) grad_output_quantizer = self.get_quantizer("backward", 0) - input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) weight_quantizer.set_usage(rowwise=True, columnwise=False) - grad_output_quantizer.set_usage(rowwise=True, columnwise=weight_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) @@ -355,6 +358,15 @@ 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, @@ -420,6 +432,7 @@ def _functional_forward( tensor_parallel_group: Optional[torch.distributed.ProcessGroup] = None, sequence_parallel: bool = False, with_quantized_compute: bool = False, + backward_override: Optional[str] = None, input_quantizer: Optional[Quantizer] = None, weight_quantizer: Optional[Quantizer] = None, output_quantizer: Optional[Quantizer] = None, @@ -459,6 +472,8 @@ def _functional_forward( distributing along inner dimension (embedding dim) with_quantized_compute: bool, default = False Whether to perform compute with quantized data. + backward_override: {`None`, `"high_precision"`, `"dequantized"`}, default = `None` + Backward-override policy for quantized compute. input_quantizer: Quantizer, optional Builder class for quantized input tensor. weight_quantizer: Quantizer, optional @@ -510,7 +525,10 @@ def _functional_forward( if with_quantized_compute: if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") - input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + input_quantizer.set_usage( + rowwise=True, + columnwise=weight_requires_grad and backward_override is None, + ) if with_x_all_gather: input_quantizer.set_usage(columnwise=False) x, x_async = gather_along_first_dim( @@ -542,7 +560,10 @@ def _functional_forward( elif with_quantized_compute and not is_quantized_tensor(w): if weight_quantizer is None: raise ValueError("Missing quantizer for weight tensor") - weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) + weight_quantizer.set_usage( + rowwise=True, + columnwise=input_requires_grad and backward_override is None, + ) w = weight_quantizer(w) # Check output tensor @@ -611,14 +632,23 @@ def _functional_forward( # Prepare weight tensor for backward pass if input_requires_grad: - if w is not weight and with_quantized_compute and is_quantized_tensor(w): + if ( + w is not weight + and with_quantized_compute + and is_quantized_tensor(w) + and backward_override is None + ): w.update_usage(rowwise_usage=False, columnwise_usage=True) else: w = None # Prepare input tensor for backward pass if weight_requires_grad: - if with_quantized_compute and is_quantized_tensor(x_local): + if ( + with_quantized_compute + and is_quantized_tensor(x_local) + and backward_override is None + ): if not (isinstance(x_local, Float8TensorStorage) and with_x_all_gather): # FP8 does not support all-gather of transpose data x_local.update_usage(rowwise_usage=False, columnwise_usage=True) @@ -968,6 +998,10 @@ def op_forward( grad_output_quantizer = self.get_quantizer("backward", 0) grad_input_quantizer = prev_op_grad_output_quantizer with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + if with_quantized_compute: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None # Get autocast dtype if needed if torch.is_autocast_enabled(): @@ -984,6 +1018,7 @@ def op_forward( tensor_parallel_group=self.tensor_parallel_group, sequence_parallel=self.sequence_parallel, with_quantized_compute=with_quantized_compute, + backward_override=backward_override, input_quantizer=input_quantizer, weight_quantizer=weight_quantizer, output_quantizer=output_quantizer, @@ -993,10 +1028,17 @@ def op_forward( # Save state for backward pass if ctx.requires_grad: + if backward_override == "high_precision": + saved_input = input_ if weight_requires_grad else None + saved_weight = self.weight if input_requires_grad else None + else: + saved_input = x_local + saved_weight = w if is_cpu_offload_enabled(): - mark_activation_offload(x_local) - ctx.save_for_backward(x_local, w) - ctx.with_quantized_compute = with_quantized_compute + mark_activation_offload(saved_input) + ctx.save_for_backward(saved_input, saved_weight) + ctx.with_quantized_compute = with_quantized_compute and backward_override is None + ctx.backward_override = backward_override ctx.input_quantizer = input_quantizer ctx.weight_quantizer = weight_quantizer ctx.grad_output_quantizer = grad_output_quantizer diff --git a/transformer_engine/pytorch/ops/basic/bias.py b/transformer_engine/pytorch/ops/basic/bias.py index d580f84866..88f563b2c5 100644 --- a/transformer_engine/pytorch/ops/basic/bias.py +++ b/transformer_engine/pytorch/ops/basic/bias.py @@ -10,6 +10,7 @@ import torch import transformer_engine_torch as tex +from ...quantization import FP8GlobalStateManager from ..op import BasicOperation, OperationContext from ...utils import canonicalize_device, canonicalize_dtype from ...tensor import Quantizer @@ -124,6 +125,10 @@ def op_forward( if ctx.requires_grad: ctx.grad_input_quantizer = prev_op_grad_output_quantizer + if FP8GlobalStateManager.is_fp8_enabled(): + fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + if fp8_recipe.backward_override is not None: + ctx.grad_input_quantizer = None return x + b diff --git a/transformer_engine/pytorch/ops/basic/quantize.py b/transformer_engine/pytorch/ops/basic/quantize.py index fa3efc3807..d0c1137d91 100644 --- a/transformer_engine/pytorch/ops/basic/quantize.py +++ b/transformer_engine/pytorch/ops/basic/quantize.py @@ -59,6 +59,11 @@ def op_forward( quantize_forward = fp8_enabled and self._quantize_forward quantize_backward = fp8_enabled and self._quantize_backward + # Backward quantization is controlled by recipe backward override. + if fp8_enabled: + recipe = FP8GlobalStateManager.get_fp8_recipe() + quantize_backward = quantize_backward and recipe.backward_override is None + # Quantize if needed out = input_ if quantize_forward and not is_quantized_tensor(out): diff --git a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py index 4ab082d32b..3950316a3c 100644 --- a/transformer_engine/pytorch/ops/fused/backward_activation_bias.py +++ b/transformer_engine/pytorch/ops/fused/backward_activation_bias.py @@ -104,8 +104,9 @@ def fuse_backward_ops( """ - # Check if recipe supports bias activation fusion - if recipe is None: + # Check if recipe supports bias activation fusion. + # high_precision/dequantized backward overrides should use unfused backward ops. + if recipe is None or recipe.backward_override is not None: return ops # Scan through ops, fusing if possible diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py index dfc11a19e7..8df929f799 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_activation.py @@ -92,6 +92,10 @@ def fuser_forward( grad_output_quantizer = linear_op.get_quantizer("backward", 0) grad_input_quantizer = prev_op_grad_output_quantizer with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + if with_quantized_compute: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None # Get autocast dtype if needed if torch.is_autocast_enabled(): @@ -109,6 +113,7 @@ def fuser_forward( tensor_parallel_group=linear_op.tensor_parallel_group, sequence_parallel=linear_op.sequence_parallel, with_quantized_compute=with_quantized_compute, + backward_override=backward_override, input_quantizer=input_quantizer, weight_quantizer=weight_quantizer, output_quantizer=output_quantizer, @@ -118,10 +123,19 @@ def fuser_forward( # Save state for backward pass if linear_op_ctx.requires_grad: + if backward_override == "high_precision": + saved_input = input_ if weight_requires_grad else None + saved_weight = linear_op.weight if input_requires_grad else None + else: + saved_input = x_local + saved_weight = w if is_cpu_offload_enabled(): - mark_activation_offload(x_local) - linear_op_ctx.save_for_backward(x_local, w) - linear_op_ctx.with_quantized_compute = with_quantized_compute + mark_activation_offload(saved_input) + linear_op_ctx.save_for_backward(saved_input, saved_weight) + linear_op_ctx.with_quantized_compute = ( + with_quantized_compute and backward_override is None + ) + linear_op_ctx.backward_override = backward_override linear_op_ctx.input_quantizer = input_quantizer linear_op_ctx.weight_quantizer = weight_quantizer linear_op_ctx.grad_output_quantizer = grad_output_quantizer @@ -131,6 +145,8 @@ def fuser_forward( linear_op_ctx.weight_requires_grad = weight_requires_grad if bias_op is not None and bias_op_ctx.requires_grad: bias_op_ctx.grad_input_quantizer = linear_op.get_grad_output_quantizer() + if backward_override is not None: + bias_op_ctx.grad_input_quantizer = None return output, [() for _ in range(len(self.basic_ops))] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py index 2dfc0566b7..5376a7d264 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_bias_add.py @@ -86,6 +86,10 @@ def fuser_forward( grad_output_quantizer = linear_op.get_quantizer("backward", 0) grad_input_quantizer = prev_op_grad_output_quantizer with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + if with_quantized_compute: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None # Get autocast dtype if needed if torch.is_autocast_enabled(): @@ -106,6 +110,7 @@ def fuser_forward( tensor_parallel_group=linear_op.tensor_parallel_group, sequence_parallel=linear_op.sequence_parallel, with_quantized_compute=with_quantized_compute, + backward_override=backward_override, input_quantizer=input_quantizer, weight_quantizer=weight_quantizer, output_quantizer=output_quantizer, @@ -115,10 +120,19 @@ def fuser_forward( # Save state for backward pass if linear_op_ctx.requires_grad: + if backward_override == "high_precision": + saved_input = input_ if weight_requires_grad else None + saved_weight = linear_op.weight if input_requires_grad else None + else: + saved_input = x_local + saved_weight = w if is_cpu_offload_enabled(): - mark_activation_offload(x_local) - linear_op_ctx.save_for_backward(x_local, w) - linear_op_ctx.with_quantized_compute = with_quantized_compute + mark_activation_offload(saved_input) + linear_op_ctx.save_for_backward(saved_input, saved_weight) + linear_op_ctx.with_quantized_compute = ( + with_quantized_compute and backward_override is None + ) + linear_op_ctx.backward_override = backward_override linear_op_ctx.input_quantizer = input_quantizer linear_op_ctx.weight_quantizer = weight_quantizer linear_op_ctx.grad_output_quantizer = grad_output_quantizer @@ -127,7 +141,9 @@ def fuser_forward( linear_op_ctx.input_requires_grad = input_requires_grad linear_op_ctx.weight_requires_grad = weight_requires_grad if bias_op is not None and bias_op_ctx.requires_grad: - bias_op_ctx.grad_input_quantizer = linear_op.get_grad_output_quantizer() + bias_op_ctx.grad_input_quantizer = ( + None if backward_override is not None else linear_op.get_grad_output_quantizer() + ) return output, [() for _ in range(len(self.basic_ops))] diff --git a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py index ae4bdd4b19..abeb39adfa 100644 --- a/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py +++ b/transformer_engine/pytorch/ops/fused/forward_linear_scale_add.py @@ -65,6 +65,10 @@ def fuser_forward( grad_output_quantizer = linear_op.get_quantizer("backward", 0) grad_input_quantizer = prev_op_grad_output_quantizer with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + if with_quantized_compute: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None # Get extra input tensor for add operation extra_input = basic_op_extra_inputs[2][0] @@ -87,6 +91,7 @@ def fuser_forward( tensor_parallel_group=linear_op.tensor_parallel_group, sequence_parallel=linear_op.sequence_parallel, with_quantized_compute=with_quantized_compute, + backward_override=backward_override, input_quantizer=input_quantizer, weight_quantizer=weight_quantizer, output_quantizer=output_quantizer, @@ -96,10 +101,19 @@ def fuser_forward( # Save state for backward pass if linear_op_ctx.requires_grad: + if backward_override == "high_precision": + saved_input = input_ if weight_requires_grad else None + saved_weight = linear_op.weight if input_requires_grad else None + else: + saved_input = x_local + saved_weight = w if is_cpu_offload_enabled(): - mark_activation_offload(x_local) - linear_op_ctx.save_for_backward(x_local, w) - linear_op_ctx.with_quantized_compute = with_quantized_compute + mark_activation_offload(saved_input) + linear_op_ctx.save_for_backward(saved_input, saved_weight) + linear_op_ctx.with_quantized_compute = ( + with_quantized_compute and backward_override is None + ) + linear_op_ctx.backward_override = backward_override linear_op_ctx.input_quantizer = input_quantizer linear_op_ctx.weight_quantizer = weight_quantizer linear_op_ctx.grad_output_quantizer = grad_output_quantizer diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py index 0d3e1d0416..84073be6f8 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_forward_linear.py @@ -388,6 +388,19 @@ def fuse_forward_ops( """ + # Disable Userbuffers for backward overrides. + # In high_precision/dequantized modes we want to avoid all UB-specific overlap + # paths and run through the standard non-UB operator sequence instead. + recipe = unused.get("recipe", None) + if recipe is not None: + backward_override = recipe.backward_override + elif FP8GlobalStateManager.is_fp8_enabled(): + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + if backward_override is not None: + return ops + # Return immediately if environment is not distributed if not torch.distributed.is_initialized() or torch.distributed.get_world_size() == 1: return ops diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index 76606ec799..a3c7e1bac7 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -338,6 +338,7 @@ def __init__( # Cache and detect change of state relevant for fusing operations self.recipe_type = None self.first_op_requiring_backward = 0 + self.backward_override = None self._last_amax_history_len = 0 # Flatten list of parameters @@ -414,9 +415,14 @@ def maybe_fuse_ops( # Early exit if fusion parameters haven't changed need_reset = False recipe_type = type(recipe) - fusion_params = (recipe_type, first_op_requiring_backward) - if fusion_params != (self.recipe_type, self.first_op_requiring_backward): - # Recipe type or grad requirmenets have changed + backward_override = recipe.backward_override if recipe is not None else None + fusion_params = (recipe_type, first_op_requiring_backward, backward_override) + if fusion_params != ( + self.recipe_type, + self.first_op_requiring_backward, + self.backward_override, + ): + # Recipe type, backward override, or grad requirements have changed need_reset = True elif ( recipe is not None @@ -450,7 +456,7 @@ def maybe_fuse_ops( ) # Save current fusion params - self.recipe_type, self.first_op_requiring_backward = fusion_params + self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params # Save amax history length if isinstance(recipe, DelayedScaling): diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py index 52e292125e..ca3913762f 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py @@ -222,6 +222,10 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """ if dtype is None: dtype = self._dtype + + if self._rowwise_data is not None and self._rowwise_data.numel() == 0: + return torch.empty(self.size(), dtype=dtype, device=self.device) + block_len = 128 if not self._is_2D_scaled: return self._dequantize_vectorwise(dtype=dtype) diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 7bbe809c9d..842f42838b 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -30,6 +30,11 @@ def forward( dtype: torch.dtype, ) -> torch.Tensor: # pylint: disable=missing-function-docstring + if tensor._rowwise_data is not None and tensor._rowwise_data.numel() == 0: + return torch.empty(tensor.size(), dtype=dtype, device=tensor.device) + 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 @@ -182,6 +187,8 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Dequantize to a higher precision.""" if dtype is None: dtype = self._dtype + if self._rowwise_data is not None and self._rowwise_data.numel() == 0: + return torch.empty(self.size(), dtype=dtype, device=self.device) return _FromMXFP8Func.forward(None, self, dtype) def size(self, *args, **kwargs): diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index fb163c9032..70699ad71a 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -42,6 +42,10 @@ def forward( dtype: torch.dtype, ) -> torch.Tensor: # pylint: disable=missing-function-docstring + if tensor._rowwise_data is not None and tensor._rowwise_data.numel() == 0: + return torch.empty(tensor.size(), dtype=dtype, device=tensor.device) + 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: @@ -213,6 +217,8 @@ def dequantize(self, *, dtype: Optional[torch.dtype] = None) -> torch.Tensor: """Dequantize to a higher precision.""" if dtype is None: dtype = self._dtype + if self._rowwise_data is not None and self._rowwise_data.numel() == 0: + return torch.empty(self.size(), dtype=dtype, device=self.device) return _FromNVFP4Func.forward(None, self, dtype) def size(self, dim: Optional[int] = None) -> Union[torch.Size, int]: From edf10bb42300ee0d5ba63a501c2bc6647e68b50c Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Wed, 8 Apr 2026 00:31:55 +0800 Subject: [PATCH 320/521] Update the error message for cublas version check (#2843) update the error message for cublas version check Signed-off-by: Xin Yao --- .../common/gemm/cublaslt_grouped_gemm.cu | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 246fc684a1..a8e0b6df83 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1363,7 +1363,7 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream) { NVTE_ERROR("nvte_grouped_gemm requires cuBLAS 13.3+, but compile-time cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -1375,7 +1375,7 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num NVTE_ERROR( "nvte_grouped_gemm_with_discrete_inputA requires cuBLAS 13.3+, but compile-time " "cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -1388,20 +1388,20 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, NVTE_ERROR( "nvte_grouped_gemm_with_discrete_out requires cuBLAS 13.3+, but compile-time " "cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, cudaStream_t stream) { NVTE_ERROR("nvte_grouped_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); + 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 " "version is ", - CUBLAS_VERSION, ". Please upgrade to CUDA 13.3 or newer."); + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); return 0; } From a10b0b1f74a922d03e1c2c530e2cdc4683f45681 Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Tue, 7 Apr 2026 23:40:44 +0200 Subject: [PATCH 321/521] guard rmsnorm fused add tests behind appropriate cudnn version (#2844) Signed-off-by: CarlosGomes98 --- tests/cpp/operator/test_normalization.cu | 4 ++++ tests/cpp/operator/test_normalization.h | 1 + 2 files changed, 5 insertions(+) diff --git a/tests/cpp/operator/test_normalization.cu b/tests/cpp/operator/test_normalization.cu index db5d6be773..f737005e26 100644 --- a/tests/cpp/operator/test_normalization.cu +++ b/tests/cpp/operator/test_normalization.cu @@ -46,6 +46,10 @@ void performTest(const size_t N, const size_t H, const bool zero_centered_gamma, GTEST_SKIP() << "cuDNN normalizations not supported on pre-Hopper GPUs yet!"; } + if (fused_bwd_add && use_cudnn && (cudnnGetVersion() < 92100)) { + GTEST_SKIP() << "cuDNN < 9.21 does not support fused RMSNorm backward+add"; + } + using WeightType = InputType; DType itype = TypeInfo::dtype; DType wtype = TypeInfo::dtype; diff --git a/tests/cpp/operator/test_normalization.h b/tests/cpp/operator/test_normalization.h index 16b4929741..44038c32a8 100644 --- a/tests/cpp/operator/test_normalization.h +++ b/tests/cpp/operator/test_normalization.h @@ -15,6 +15,7 @@ #include #include +#include #include #include From e2470a76be8dbc127a68ec9f2fb53ca71960ef9e Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 7 Apr 2026 18:24:59 -0700 Subject: [PATCH 322/521] [JAX] Use avg m,n,k heuristics for Grouped GEMM (#2840) 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> --- transformer_engine/jax/cpp_extensions/gemm.py | 8 ++- .../jax/csrc/extensions/gemm.cpp | 68 ++++++++++++++++++- 2 files changed, 72 insertions(+), 4 deletions(-) diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index aaec5affa8..c081e451a7 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -2027,7 +2027,13 @@ def grouped_gemm_copy_group_sizes( @cache def _should_enforce_v2_grouped_gemm() -> bool: """Read NVTE_JAX_ENFORCE_V2_GROUPED_GEMM once per process (cached).""" - return os.getenv("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "0") == "1" + val = os.getenv("NVTE_JAX_ENFORCE_V2_GROUPED_GEMM", "0") + try: + return bool(int(val)) + except ValueError as e: + raise ValueError( + f"NVTE_JAX_ENFORCE_V2_GROUPED_GEMM must be an integer (0 or 1), got: {val!r}" + ) from e def _can_use_v2_grouped_gemm( diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index 0d1ef405f4..a7f16bb31f 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -683,6 +683,57 @@ size_t grouped_gemm_num_gemms(Buffer_Type const &lhs_first_dims, Buffer_Type con } } +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Compute estimates for average dimensions of a grouped tensor. + * + * Returns a pair of {non_contracting_avg, contracting_avg} dimensions for the given grouped tensor, to estimate per-group GEMM sizes. When a dimension is ragged, we estimate the average size by dividing the dim size by G ("num_gemms"). When a dimension has no ragged dims, we assume it is of shape (G*K, N) or (G*N, K) so we divide the first dim by G to get the average per-group size. + * + * Examples: + * - fwd lhs: shape_2d=[ragged M, K], first_dims=[M,...] (ragged M) → avg_m = (G*M)/G = M, avg_k = K + * - fwd rhs: shape_2d=[G*K, N], last_dims=None (static K) → avg_k = (G*K)/G = K, avg_n = N + * - wgrad lhs: shape_2d=[M, ragged K], last_dims=[K,...] (ragged K) → avg_k = (G*K)/G = K, avg_m = M + * - wgrad rhs: shape_2d=[N, ragged K], last_dims=[K,...] (ragged K) → avg_k = (G*K)/G = K, avg_n = N + * + * \param[in] first_dims XLA buffer of on-device first dimensions. Shape (G,) if ragged, empty otherwise. + * \param[in] last_dims XLA buffer of on-device last dimensions. Shape (G,) if ragged, empty otherwise. + * \param[in] shape_2d Pair of total 2D dimensions (rows, cols) for the operand. + * \param[in] num_gemms Number of GEMMs (G) in the grouped operation. + * \param[in] is_trans Whether the operand is transposed. + * \return Pair of {non_contracting_avg, contracting_avg}, i.e. {avg_m, avg_k} for lhs or + * {avg_n, avg_k} for rhs. + */ +std::pair grouped_gemm_avg_dims(Buffer_Type const &first_dims, + Buffer_Type const &last_dims, + std::pair const &shape_2d, + size_t num_gemms, bool is_trans) { + bool first_ragged = first_dims.element_count() > 0; + bool last_ragged = last_dims.element_count() > 0; + bool any_ragged = first_ragged || last_ragged; + + std::pair per_group_shape_2d{}; + if (first_ragged) { + per_group_shape_2d = { + static_cast(std::round(static_cast(shape_2d.first) / num_gemms)), + shape_2d.second}; + } else if (!any_ragged) { + per_group_shape_2d = { + static_cast(std::round(static_cast(shape_2d.first) / num_gemms)), + shape_2d.second}; + } else if (last_ragged && !first_ragged) { + per_group_shape_2d = { + shape_2d.first, + static_cast(std::round(static_cast(shape_2d.second) / num_gemms))}; + } else { + NVTE_CHECK(false, "Grouped GEMM with both first_dims and last_dims ragged is not supported."); + } + + int64_t non_contract = + static_cast(is_trans ? per_group_shape_2d.second : per_group_shape_2d.first); + int64_t contract = + static_cast(is_trans ? per_group_shape_2d.first : per_group_shape_2d.second); + return {non_contract, contract}; +} + } // namespace jax } // namespace transformer_engine @@ -741,11 +792,22 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty auto out_tensor = make_grouped_tensor(*output, out_first_dims, out_last_dims, int64_base, int64_capacity, int64_offset, num_gemms, stream); + auto [avg_m, avg_k_lhs] = grouped_gemm_avg_dims( + lhs_first_dims, lhs_last_dims, {lhs_left_size, lhs_right_size}, num_gemms, lhs_is_trans); + auto [avg_n, avg_k_rhs] = grouped_gemm_avg_dims( + rhs_first_dims, rhs_last_dims, {rhs_left_size, rhs_right_size}, num_gemms, !rhs_is_trans); + // Use k from lhs (both sides should agree for well-formed inputs). + NVTE_CHECK(avg_k_lhs == avg_k_rhs, "Contracting dimension mismatch: lhs avg_k=", avg_k_lhs, + " vs rhs avg_k=", avg_k_rhs); + + GroupedMatmulConfigWrapper gemmConfig{}; + gemmConfig.set_avg_m(avg_m); + gemmConfig.set_avg_n(avg_n); + gemmConfig.set_avg_k(avg_k_lhs); + nvte_grouped_gemm(rhs_tensor, rhs_is_trans, lhs_tensor, lhs_is_trans, nullptr, out_tensor, alpha_tensor.data(), beta_tensor.data(), workspace_setup.data(), - workspace_cublas.data(), - nullptr, // config (use defaults) - stream); + workspace_cublas.data(), gemmConfig, stream); return ffi_with_cuda_error_check(); } From d3f88eeb40427c102caf48c17ed7309abc9b5a4a Mon Sep 17 00:00:00 2001 From: eattia-nvidia Date: Wed, 8 Apr 2026 04:12:35 +0200 Subject: [PATCH 323/521] [PyTorch][Flash Attn] Add fallback import for FA3 (#2806) * [PyTorch][Flash Attn] Add fallback import for FA3 when flash_attn_interface.py is outside flash_attn_3 package Some FA3 installations (e.g. via pip) place flash_attn_interface.py directly under site-packages/ rather than inside flash_attn_3/. This causes a ModuleNotFoundError when importing from flash_attn_3.flash_attn_interface. Add a try/except ModuleNotFoundError fallback to import directly from flash_attn_interface when the subpackage import fails. Signed-off-by: Emmanuel Attia * [PyTorch][Flash Attn] Use find_spec for FA3 import and add diagnostic warning Address review feedback: - Use importlib.util.find_spec() for explicit module checking instead of exception-driven control flow, avoiding masking real import errors - Add a warning when the flat layout fallback is used for easier debugging - Raise a clear error if flash_attn_interface is not found in either location Signed-off-by: Emmanuel Attia --------- Signed-off-by: Emmanuel Attia --- .../dot_product_attention/backends.py | 39 ++++++++++++++----- 1 file changed, 30 insertions(+), 9 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 442366035a..1e7bdaac84 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -6,6 +6,7 @@ from contextlib import nullcontext from importlib.metadata import version as get_pkg_version from importlib.metadata import PackageNotFoundError +import importlib.util import os from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings @@ -138,15 +139,35 @@ flash_attn_with_kvcache_v3 = None # pass # only print warning if use_flash_attention_3 = True in get_attention_backend else: - from flash_attn_3.flash_attn_interface import flash_attn_func as flash_attn_func_v3 - from flash_attn_3.flash_attn_interface import ( - flash_attn_varlen_func as flash_attn_varlen_func_v3, - ) - from flash_attn_3.flash_attn_interface import ( - flash_attn_with_kvcache as flash_attn_with_kvcache_v3, - ) - from flash_attn_3.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3 - from flash_attn_3.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3 + if importlib.util.find_spec("flash_attn_3.flash_attn_interface") is not None: + from flash_attn_3.flash_attn_interface import flash_attn_func as flash_attn_func_v3 + from flash_attn_3.flash_attn_interface import ( + flash_attn_varlen_func as flash_attn_varlen_func_v3, + ) + from flash_attn_3.flash_attn_interface import ( + flash_attn_with_kvcache as flash_attn_with_kvcache_v3, + ) + from flash_attn_3.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3 + from flash_attn_3.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3 + elif importlib.util.find_spec("flash_attn_interface") is not None: + warnings.warn( + "flash_attn_interface found outside flash_attn_3 package. " + "Importing directly from flash_attn_interface." + ) + from flash_attn_interface import flash_attn_func as flash_attn_func_v3 + from flash_attn_interface import ( + flash_attn_varlen_func as flash_attn_varlen_func_v3, + ) + from flash_attn_interface import ( + flash_attn_with_kvcache as flash_attn_with_kvcache_v3, + ) + from flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3 + from flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3 + else: + raise ModuleNotFoundError( + "flash-attn-3 package is installed but flash_attn_interface module " + "could not be found in flash_attn_3/ or site-packages/." + ) fa_utils.set_flash_attention_3_params() From 77b8681de5cfa6bd874d89f19cd819dcea77ae36 Mon Sep 17 00:00:00 2001 From: Hongbin Liu Date: Wed, 8 Apr 2026 21:49:28 +0800 Subject: [PATCH 324/521] add mark_not_offload() interface for cpu_offload_v1 (#2770) * add mark_not_offload() interface for cpu_offload_v1 Signed-off-by: Hongbin Liu * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reuse mark_activation_offload interface Signed-off-by: Hongbin Liu * fix ci Signed-off-by: Hongbin Liu --------- Signed-off-by: Hongbin Liu Co-authored-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/cpu_offload.py | 1 + transformer_engine/pytorch/cpu_offload_v1.py | 22 +++++++++++++------- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index d0b314a64f..ed10909b8a 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -46,6 +46,7 @@ def mark_activation_offload(*tensors): def mark_not_offload(*tensors: torch.Tensor): """Marks tensors to prevent them from being offloaded.""" if NVTE_CPU_OFFLOAD_V1: + v1_code_path.mark_activation_offload(*tensors, offload=False) return tensors, tensor_obj = prepare_for_saving(*tensors) diff --git a/transformer_engine/pytorch/cpu_offload_v1.py b/transformer_engine/pytorch/cpu_offload_v1.py index f92c436941..fb62546cc0 100644 --- a/transformer_engine/pytorch/cpu_offload_v1.py +++ b/transformer_engine/pytorch/cpu_offload_v1.py @@ -19,7 +19,7 @@ CPUOffloadedLayer = False -def mark_activation_offload(*tensors): +def mark_activation_offload(*tensors, offload: bool = True): """Set the type of the offloading needed for a tensor.""" if TEDebugState.debug_enabled: raise RuntimeError("CPU offload is not supported in debug mode.") @@ -28,16 +28,24 @@ def mark_activation_offload(*tensors): if tensor is None: continue if type(tensor) in [torch.Tensor, torch.nn.Parameter]: - tensor.activation_offloading = True + if offload: + tensor.activation_offloading = True + else: + # This is a hack to prevent the tensor from being offloaded. + # And it won't break the original logic of the code. + tensor._TE_do_not_offload = True else: data_tensors = tensor.get_data_tensors() for tensor in data_tensors: if tensor is not None: - tensor.activation_offloading = True - # This is a hack to force clear the tensor after it is offloaded. - # It is needed, because .*TensorStorage classes are saved in the ctx, - # and they contain the reference to their data tensors. - tensor.needs_force_clear = True + if offload: + tensor.activation_offloading = True + # This is a hack to force clear the tensor after it is offloaded. + # It is needed, because .*TensorStorage classes are saved in the ctx, + # and they contain the reference to their data tensors. + tensor.needs_force_clear = offload + else: + tensor._TE_do_not_offload = True def is_cpu_offload_enabled() -> bool: From a30a1261c1088190d06fbe5dee0b3d4770fb3104 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Wed, 8 Apr 2026 15:56:52 -0700 Subject: [PATCH 325/521] Fix zero input shape for bgrad_group_quantize (#2854) fix zero input shape for dbias Signed-off-by: Varun Thumbe --- tests/pytorch/test_grouped_tensor.py | 21 +++++++++++++++++++ .../pytorch/csrc/extensions/cast.cpp | 11 ++++++++-- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 5bc2faa007..04a0376019 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -410,6 +410,27 @@ def test_quantize_grouped_mxfp8(self, shape: List[Tuple[int, int]], output_dbias expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors]) assert torch.allclose(dbias, expected_dbias) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_bgrad_group_quantize_zero_size_tensor(self) -> None: + """Test bgrad_group_quantize handles zero-row input without error.""" + num_tensors = 3 + last_dim = 1024 + grouped_input = torch.empty(0, last_dim, dtype=torch.bfloat16, device="cuda") + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=False) + first_dims = torch.zeros(num_tensors, dtype=torch.int64, device="cuda") + + grouped_output, dbias = tex.bgrad_group_quantize( + grouped_input, + quantizer, + num_tensors, + first_dims, + ) + + assert dbias.shape == (num_tensors, last_dim) + assert torch.all(dbias == 0) + @pytest.mark.parametrize("output_dbias", [False, True]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_group_quantize_cudagraph_capturable(self, output_dbias: bool) -> None: diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index f150e90507..b689a1c1b4 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -247,8 +247,7 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const auto logical_first_dim = logical_shape[0]; const auto logical_last_dim = logical_shape[1]; - NVTE_CHECK(logical_first_dim > 0 && logical_last_dim > 0, - "bgrad_group_quantize: empty input tensor is not supported."); + bool empty_input_buffer = logical_first_dim == 0 || logical_last_dim == 0; NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()), "bgrad_group_quantize: only MXFP8 quantizer is supported."); @@ -264,6 +263,14 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, py::reinterpret_borrow(quantizer), first_dims, logical_first_dim, logical_last_dim); + if (empty_input_buffer) { + at::Tensor dbias_torch = + at::zeros({static_cast(num_tensors), static_cast(logical_last_dim)}, + tensor.options()); + return py::make_tuple(py::reinterpret_borrow(grouped_output_py), + py::cast(std::move(dbias_torch))); + } + const std::vector dbias_logical_shape = {num_tensors, logical_last_dim}; GroupedTensorWrapper grouped_dbias(num_tensors, dbias_logical_shape, NVTE_DELAYED_TENSOR_SCALING); at::Tensor dbias_torch = From 0aea85ff29603508c4286f5bc8d9efc05a9c3975 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 9 Apr 2026 08:59:48 -0700 Subject: [PATCH 326/521] [Common] Fix: IMA in `register_user_buffer_collective` on non-SM90 GPUs (#2859) * fixed mem alloc for AG * use raid Signed-off-by: Phuong Nguyen --------- Signed-off-by: Phuong Nguyen --- .../userbuffers/userbuffers-host.cpp | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) 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 6ff9d63a2d..1dcde51d4b 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp @@ -662,13 +662,28 @@ int register_user_buffer_collective(void **gpubuff, size_t bytes, communicator * } NVTE_CHECK(comm->nvsize <= 8, "CUDA IPC supports only up to 8 GPUs in an NVLink domain."); - cudaIpcMemHandle_t memhndl; - NVTE_CHECK_CUDA(cudaIpcGetMemHandle(&memhndl, *gpubuff)); - cudaIpcMemHandle_t *tmp = - reinterpret_cast(malloc(comm->nvsize * sizeof(cudaIpcMemHandle_t))); + // Use cudaMallocHost (pinned host memory) so these buffers are CPU-accessible (plain memcpy) + // and GPU DMA-accessible, allowing the allgather callback to pass them directly to NCCL + // without additional staging copies. RAII guards ensure the pinned pages are released on + // every exit path, including exceptions thrown by NVTE_CHECK_CUDA / NVTE_ERROR. + struct PinnedDeleter { + void operator()(void *p) const { + if (p) cudaFreeHost(p); + } + }; + cudaIpcMemHandle_t *memhndl; + NVTE_CHECK_CUDA( + cudaMallocHost(reinterpret_cast(&memhndl), sizeof(cudaIpcMemHandle_t))); + std::unique_ptr memhndl_guard(memhndl); + NVTE_CHECK_CUDA(cudaIpcGetMemHandle(memhndl, *gpubuff)); + + cudaIpcMemHandle_t *tmp; + NVTE_CHECK_CUDA( + cudaMallocHost(reinterpret_cast(&tmp), comm->nvsize * sizeof(cudaIpcMemHandle_t))); + std::unique_ptr tmp_guard(tmp); comm->_allgather(reinterpret_cast(tmp), comm->nvsize * sizeof(cudaIpcMemHandle_t), - reinterpret_cast(&memhndl), sizeof(cudaIpcMemHandle_t), + reinterpret_cast(memhndl), sizeof(cudaIpcMemHandle_t), comm->comm_intra); // Check for NVLINK support before attempting IPC operations @@ -689,7 +704,6 @@ int register_user_buffer_collective(void **gpubuff, size_t bytes, communicator * } } if (!peer_access_available) { - free(tmp); NVTE_ERROR( "No peer-to-peer access available between GPUs. This platform does not support the " "GPU-to-GPU " @@ -712,7 +726,6 @@ int register_user_buffer_collective(void **gpubuff, size_t bytes, communicator * comm->peer_ptr[hndl], comm->nvsize * sizeof(void *), cudaMemcpyHostToDevice)); NVTE_CHECK_CUDA(cudaDeviceSynchronize()); - free(tmp); #if CUDART_VERSION >= 12010 } #endif From 181322eb1f029c845966b1bd174c3c92fe591666 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Thu, 9 Apr 2026 10:08:41 -0700 Subject: [PATCH 327/521] Simplify FA3 discovery (#2849) Signed-off-by: Vladimir Cherepanov --- qa/L3_pytorch_FA_versions_test/test.sh | 3 -- .../dot_product_attention/backends.py | 39 +++++-------------- .../attention/dot_product_attention/utils.py | 5 +-- 3 files changed, 10 insertions(+), 37 deletions(-) diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 6e239bfb72..bbfc4db5ba 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -34,9 +34,6 @@ do else git clone https://github.com/Dao-AILab/flash-attention.git cd flash-attention/hopper && python setup.py install - python_path=`python -c "import site; print(site.getsitepackages()[0])"` - mkdir -p $python_path/flash_attn_3 - cp flash_attn_interface.py $python_path/flash_attn_3/ cd ../../ fi diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 1e7bdaac84..b5ed15f8e0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -6,7 +6,6 @@ from contextlib import nullcontext from importlib.metadata import version as get_pkg_version from importlib.metadata import PackageNotFoundError -import importlib.util import os from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings @@ -139,35 +138,15 @@ flash_attn_with_kvcache_v3 = None # pass # only print warning if use_flash_attention_3 = True in get_attention_backend else: - if importlib.util.find_spec("flash_attn_3.flash_attn_interface") is not None: - from flash_attn_3.flash_attn_interface import flash_attn_func as flash_attn_func_v3 - from flash_attn_3.flash_attn_interface import ( - flash_attn_varlen_func as flash_attn_varlen_func_v3, - ) - from flash_attn_3.flash_attn_interface import ( - flash_attn_with_kvcache as flash_attn_with_kvcache_v3, - ) - from flash_attn_3.flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3 - from flash_attn_3.flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3 - elif importlib.util.find_spec("flash_attn_interface") is not None: - warnings.warn( - "flash_attn_interface found outside flash_attn_3 package. " - "Importing directly from flash_attn_interface." - ) - from flash_attn_interface import flash_attn_func as flash_attn_func_v3 - from flash_attn_interface import ( - flash_attn_varlen_func as flash_attn_varlen_func_v3, - ) - from flash_attn_interface import ( - flash_attn_with_kvcache as flash_attn_with_kvcache_v3, - ) - from flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3 - from flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3 - else: - raise ModuleNotFoundError( - "flash-attn-3 package is installed but flash_attn_interface module " - "could not be found in flash_attn_3/ or site-packages/." - ) + from flash_attn_interface import flash_attn_func as flash_attn_func_v3 + from flash_attn_interface import ( + flash_attn_varlen_func as flash_attn_varlen_func_v3, + ) + from flash_attn_interface import ( + flash_attn_with_kvcache as flash_attn_with_kvcache_v3, + ) + from flash_attn_interface import _flash_attn_forward as _flash_attn_fwd_v3 + from flash_attn_interface import _flash_attn_backward as _flash_attn_bwd_v3 fa_utils.set_flash_attention_3_params() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 170cb2cd34..13d1347a1e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -135,10 +135,7 @@ class FlashAttentionUtils: # Please follow these instructions to install FA3 v3_installation_steps = """\ (1) git clone https://github.com/Dao-AILab/flash-attention.git -(2) cd flash-attention/hopper && python setup.py install -(3) python_path=`python -c "import site; print(site.getsitepackages()[0])"` -(4) mkdir -p $python_path/flash_attn_3 -(5) cp flash_attn_interface.py $python_path/flash_attn_3/flash_attn_interface.py""" +(2) cd flash-attention/hopper && python setup.py install""" v3_warning_printed = False @staticmethod From 64bb9a241e59caca509f2a73550fdbbb4359e7f4 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 9 Apr 2026 17:19:54 -0400 Subject: [PATCH 328/521] [PyTorch] Support scaled + clamped SwiGLU in `te.ops` and enable fused MXFP8 grouped MLP (#2855) * cuDNN act_func='geglu' support for fused grouped MLP Signed-off-by: Kirthi Shankar Sivamani * rm incorrect/not needed doc Signed-off-by: Kirthi Shankar Sivamani * Address comments Signed-off-by: Kirthi Shankar Sivamani * Fix activation name Signed-off-by: Kirthi Shankar Sivamani * Min cudnn 1.23 for qgeglu fusion Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- docs/api/pytorch.rst | 2 + tests/pytorch/test_fusible_ops.py | 130 +++++++++++++- transformer_engine/pytorch/ops/_common.py | 39 ++++- .../pytorch/ops/basic/__init__.py | 2 +- .../pytorch/ops/basic/swiglu.py | 160 ++++++++++++++---- .../pytorch/ops/fused/backward_grouped_mlp.py | 13 +- .../pytorch/ops/fused/forward_grouped_mlp.py | 11 +- 7 files changed, 299 insertions(+), 58 deletions(-) diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 1fe4f19990..3217d29c3b 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -221,6 +221,8 @@ Operation fuser .. autoapiclass:: transformer_engine.pytorch.ops.SReLU +.. autoapiclass:: transformer_engine.pytorch.ops.ScaledClampedQGeGLU + .. autoapiclass:: transformer_engine.pytorch.ops.ScaledSwiGLU .. autoapiclass:: transformer_engine.pytorch.ops.SiLU diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 75d450b46b..795cbf3452 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -18,6 +18,9 @@ import transformer_engine.common.recipe 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, +) from transformer_engine.pytorch.ops.fused import ( BackwardActivationBias, @@ -2234,6 +2237,91 @@ def test_interleaved_scaled_swiglu(self): scales_requires_grad=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)) + def test_scaled_clamped_qgeglu( + self, + *, + in_shape: Iterable[int], + glu_interleave_size: Optional[int] = None, + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + input_requires_grad: bool, + scales_requires_grad: bool, + limit: float = 7.0, + alpha: float = 1.702, + ) -> None: + """ScaledClampedQGeGLU (clamped QGeGLU with post-scale)""" + + # Tensor dims + out_shape = list(in_shape) + out_shape[-1] //= 2 + + # 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( + out_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch reference (matches :class:`ClampedSwiGLU` numerics) + x = x_ref + if glu_interleave_size is not None: + x = x.reshape( + -1, + in_shape[-1] // (2 * glu_interleave_size), + 2, + glu_interleave_size, + ) + x = x.transpose(1, 2) + x = x.reshape(in_shape) + 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_ref = scales_ref.unsqueeze(-1) * y + if input_requires_grad or scales_requires_grad: + y_ref.backward(dy_ref) + + op = te_ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave_size, + limit=limit, + alpha=alpha, + ) + y_test = op(x_test, scales_test) + if input_requires_grad or scales_requires_grad: + y_test.backward(dy_test) + + 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_clamped_qgeglu(self): + """ScaledClampedQGeGLU with block interleaved input format""" + self.test_scaled_clamped_qgeglu( + in_shape=(32, 192), + glu_interleave_size=32, + input_requires_grad=True, + scales_requires_grad=True, + ) + class TestFusedOps: """Tests for fused operations""" @@ -3249,6 +3337,7 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) @pytest.mark.parametrize("glu_interleave_size", (None, 32)) @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) + @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_clamped_qgeglu")) def test_grouped_mlp( self, *, @@ -3264,8 +3353,9 @@ def test_grouped_mlp( split_alignment: int = 256, glu_interleave_size: Optional[int], delay_wgrad_compute: bool, + activation: str, ) -> None: - """GroupedLinear + ScaledSwiGLU + GroupedLinear""" + """GroupedLinear + ScaledSwiGLU / ScaledClampedQGeGLU + GroupedLinear""" # Split sizes split_sizes = [split_alignment * (i) for i in range(group_size)] @@ -3288,6 +3378,9 @@ def test_grouped_mlp( if quantization == "mxfp8" and bias: # Will be supported in future CUDNN release. pytest.skip("Bias/dbias not yet supported in MXFP8 fused grouped MLP") + 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") # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -3376,7 +3469,14 @@ def test_grouped_mlp( x = x.transpose(1, 2) x = x.reshape(-1, 2 * hidden_size) x1, x2 = x.chunk(2, dim=-1) - x = torch.nn.functional.silu(x1) * x2 + if activation == "scaled_swiglu": + x = torch.nn.functional.silu(x1) * x2 + else: + 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)) x = x * probs[group_idx].unsqueeze(-1) x = torch.nn.functional.linear(x, fc2_ws_ref[group_idx], bias=fc2_bs_ref[group_idx]) ys.append(x) @@ -3385,6 +3485,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, @@ -3412,7 +3517,7 @@ def test_grouped_mlp( ) module = te_ops.Sequential( fc1, - te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), + scaled_act, fc2, ) @@ -3484,6 +3589,10 @@ 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() + ) ): if te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): forward_ops = module._module_groups[0]._forward_ops @@ -3572,6 +3681,7 @@ def test_grouped_mlp( @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) + @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_clamped_qgeglu")) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) def test_grouped_mlp_cuda_graph_safe_mxfp8( self, @@ -3579,6 +3689,7 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( dtype: torch.dtype, single_grouped_weight: bool, accumulate_into_main_grad: bool, + activation: str, device: torch.device = "cuda", group_size: int = 4, hidden_size: int = 256, @@ -3591,6 +3702,12 @@ 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) @@ -3619,9 +3736,14 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( single_grouped_weight=single_grouped_weight, accumulate_into_main_grad=accumulate_into_main_grad, ) + 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) + ) module = te_ops.Sequential( fc1, - te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), + scaled_act, fc2, ) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 0e03e691f3..ae8b48a90d 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -5,9 +5,12 @@ """Helper functions used in fusible operations.""" from __future__ import annotations +import functools +from importlib.metadata import PackageNotFoundError, version as get_pkg_version from typing import Optional import torch +from packaging.version import Version as PkgVersion from transformer_engine_torch import FP8TensorMeta from ..torch_version import torch_version @@ -17,6 +20,15 @@ from ..utils import canonicalize_dtype +@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 is_quantized_tensor(tensor: torch.Tensor | QuantizedTensorStorage) -> bool: """Check if tensor is a quantized tensor""" return isinstance(tensor, QuantizedTensorStorage) @@ -73,8 +85,8 @@ def get_fp8_meta_from_fp8_tensor(tensor: Float8Tensor) -> tuple[FP8TensorMeta, i return fp8_meta, 0 -def validate_grouped_mlp_dims(fc1, swiglu, fc2) -> None: - """Validate FC1/SwiGLU/FC2 dimensions and interleave size for fused grouped MLP.""" +def validate_grouped_mlp_dims(fc1, glu_op, fc2) -> None: + """Validate FC1 / scaled GLU / FC2 dimensions for fused grouped MLP.""" if fc1.in_features % 256 != 0 or fc1.out_features % 256 != 0: raise ValueError( @@ -93,10 +105,10 @@ def validate_grouped_mlp_dims(fc1, swiglu, fc2) -> None: f"and FC2 (num_groups={fc2.num_groups}, in_features={fc2.in_features}, " f"out_features={fc2.out_features}) do not match." ) - if swiglu.glu_interleave_size != 32: + if glu_op.glu_interleave_size != 32: raise ValueError( "Fused kernel requires 32-wide GLU interleaving, " - f"but got glu_interleave_size={swiglu.glu_interleave_size}." + f"but got glu_interleave_size={glu_op.glu_interleave_size}." ) @@ -106,7 +118,7 @@ def fuse_grouped_mlp_ops( recipe, fused_op_cls, ): - """Sliding-window fusion for GroupedLinear + ScaledSwiGLU + GroupedLinear. + """Sliding-window fusion for GroupedLinear + scaled GLU + GroupedLinear. Parameters ---------- @@ -116,7 +128,9 @@ def fuse_grouped_mlp_ops( Quantization recipe. fused_op_cls : type Fused operation class with ``is_supported()`` classmethod and - constructor accepting ``fc1``, ``swiglu``, ``fc2`` keyword args. + 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. @@ -125,7 +139,11 @@ def fuse_grouped_mlp_ops( list of FusibleOperation Updated operations with matched triples replaced by fused ops. """ - from .basic import GroupedLinear, ScaledSwiGLU # pylint: disable=import-outside-toplevel + from .basic import ( # pylint: disable=import-outside-toplevel + GroupedLinear, + ScaledClampedQGeGLU, + ScaledSwiGLU, + ) if not fused_op_cls.is_supported(): return ops @@ -146,10 +164,15 @@ def fuse_grouped_mlp_ops( matches_pattern = True if not ( isinstance(window[0], GroupedLinear) - and isinstance(window[1], ScaledSwiGLU) + and isinstance(window[1], (ScaledSwiGLU, ScaledClampedQGeGLU)) and isinstance(window[2], GroupedLinear) ): 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: matches_pattern = False elif ( diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index e0a3f41019..45c938ede8 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -32,4 +32,4 @@ from .reduce_scatter import ReduceScatter from .reshape import Reshape from .rmsnorm import RMSNorm -from .swiglu import ClampedSwiGLU, ScaledSwiGLU, SwiGLU +from .swiglu import ClampedSwiGLU, ScaledClampedQGeGLU, ScaledSwiGLU, SwiGLU diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index b4427df41a..9c0bc86bc1 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -17,7 +17,7 @@ from ..op import BasicOperation, OperationContext from .._common import maybe_dequantize -__all__ = ["SwiGLU", "ClampedSwiGLU", "ScaledSwiGLU"] +__all__ = ["SwiGLU", "ClampedSwiGLU", "ScaledSwiGLU", "ScaledClampedQGeGLU"] class SwiGLU(BasicOperation): @@ -231,6 +231,34 @@ def __init__( self.cache_quantized_input: bool = cache_quantized_input self.glu_interleave_size: Optional[int] = glu_interleave_size + def _tex_clamped_swiglu_forward( + self, + swiglu_in: torch.Tensor, + next_op_input_quantizer: Optional[Quantizer], + ) -> torch.Tensor: + """Call :func:`tex.clamped_swiglu` with this op's ``limit`` / ``alpha``.""" + return tex.clamped_swiglu( + swiglu_in, + next_op_input_quantizer, + self.limit, + self.alpha, + ) + + def _tex_clamped_dswiglu( + self, + dy: torch.Tensor, + swiglu_in: torch.Tensor, + quantizer: Optional[Quantizer], + ) -> torch.Tensor: + """Call :func:`tex.clamped_dswiglu` with this op's ``limit`` / ``alpha``.""" + return tex.clamped_dswiglu( + dy, + swiglu_in, + quantizer, + self.limit, + self.alpha, + ) + def op_forward( self, ctx: OperationContext, @@ -252,7 +280,7 @@ def op_forward( x = maybe_dequantize(input_.contiguous(), dtype) # Remove interleaving if needed - swiglu_in = input_ + swiglu_in = x if self.glu_interleave_size is not None: shape = swiglu_in.size() swiglu_in = swiglu_in.reshape( @@ -265,12 +293,7 @@ def op_forward( swiglu_in = swiglu_in.view(shape) # Launch kernel - out = tex.clamped_swiglu( - swiglu_in, - next_op_input_quantizer, - limit=self.limit, - alpha=self.alpha, - ) + out = self._tex_clamped_swiglu_forward(swiglu_in, next_op_input_quantizer) # Quantize input to FP8 before caching if needed if self.cache_quantized_input: @@ -320,13 +343,7 @@ def op_backward( quantizer = None # Launch kernel - grad_swiglu_in = tex.clamped_dswiglu( - dy, - swiglu_in, - quantizer, - limit=self.limit, - alpha=self.alpha, - ) + grad_swiglu_in = self._tex_clamped_dswiglu(dy, swiglu_in, quantizer) # Apply interleaving if needed dx = grad_swiglu_in @@ -347,29 +364,25 @@ def op_backward( return dx, () -class ScaledSwiGLU(BasicOperation): - r"""SwiGLU with post-scaling. +class _ScaledGLU(BasicOperation): + """SwiGLU-family activation with per-row scales (fused grouped MLP middle op).""" - If the SwiGLU output has shape ``(d_1, ..., d_n)``, it is - multiplied with an extra input tensor of shape - ``(d_1, ..., d_{n-1})``. - - Parameters - ---------- - glu_interleave_size : int, optional - When set, the GLU activations will use an experimental block - interleaved format. See the corresponding option in the SwiGLU - operation for more details. - - """ - - # Operation expects scales num_extra_inputs: int = 1 - def __init__(self, glu_interleave_size: Optional[int] = None): + def __init__(self, glu_interleave_size: Optional[int] = None) -> None: super().__init__() self.glu_interleave_size: Optional[int] = glu_interleave_size + def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: + raise NotImplementedError + + def _glu_backward( + self, + grad_swiglu_out: torch.Tensor, + swiglu_in: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError + def op_forward(self, *args, **kwargs) -> None: raise RuntimeError( f"{self.__class__.__name__} operation has " @@ -423,8 +436,7 @@ def fuser_forward( swiglu_in = swiglu_in.transpose(1, 2).contiguous() swiglu_in = swiglu_in.view(shape) - # Compute scaled SwiGLU - swiglu_out = tex.swiglu(swiglu_in, None) + swiglu_out = self._glu_forward(swiglu_in) out = swiglu_out * scales.unsqueeze(-1) # Save state for backward pass @@ -477,7 +489,7 @@ def fuser_backward( grad_input = None if ctx.input_requires_grad: grad_swiglu_out = grad_output * scales.unsqueeze(-1) - grad_swiglu_in = tex.dswiglu(grad_swiglu_out, swiglu_in, None) + grad_swiglu_in = self._glu_backward(grad_swiglu_out, swiglu_in) grad_input = grad_swiglu_in if self.glu_interleave_size is not None: shape = grad_input.size() @@ -490,13 +502,87 @@ def fuser_backward( grad_input = grad_input.transpose(1, 2).contiguous() grad_input = grad_input.view(shape) - # Compute scales grad by recomputing SwiGLU + # Compute scales grad by recomputing GLU grad_extra_input = None if ctx.extra_input_requires_grad: - swiglu_out = tex.swiglu(swiglu_in, None) + swiglu_out = self._glu_forward(swiglu_in) grad_extra_input = torch.linalg.vecdot(swiglu_out, grad_output) # Clear input tensor if possible clear_tensor_data(ctx.saved_tensors[0]) # input_ return grad_input, [()], [(grad_extra_input,)] + + +class ScaledSwiGLU(_ScaledGLU): + r"""SwiGLU with post-scaling (matches cuDNN grouped GEMM ``act_func="swiglu"``). + + If the GLU output has shape ``(d_1, ..., d_n)``, it is multiplied + with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. + + Parameters + ---------- + glu_interleave_size : int, optional + When set, the GLU activations will use an experimental block + interleaved format. See the corresponding option in the SwiGLU + operation for more details. + + """ + + def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: + return tex.swiglu(swiglu_in, None) + + def _glu_backward( + self, + grad_swiglu_out: torch.Tensor, + swiglu_in: torch.Tensor, + ) -> torch.Tensor: + return tex.dswiglu(grad_swiglu_out, swiglu_in, None) + + +class ScaledClampedQGeGLU(_ScaledGLU): + r"""Clamped QGeGLU with post-scaling + (matches cuDNN grouped GEMM ``act_func="geglu"``). + + Same layout and scaling contract as :class:`ScaledSwiGLU`, but the GLU + uses :class:`ClampedSwiGLU` numerics (default ``limit`` / ``alpha`` match + cuDNN). + + Parameters + ---------- + glu_interleave_size : int, optional + When set, the GLU activations will use an experimental block + interleaved format. See :class:`ClampedSwiGLU`. + limit : float, default ``7.0`` + Clamp limit (see :class:`ClampedSwiGLU`). + alpha : float, default ``1.702`` + Sigmoid scale (see :class:`ClampedSwiGLU`). + + """ + + def __init__( + self, + glu_interleave_size: Optional[int] = None, + *, + limit: float = 7.0, + alpha: float = 1.702, + ) -> None: + super().__init__(glu_interleave_size) + self._clamped: ClampedSwiGLU = ClampedSwiGLU( + limit=limit, + alpha=alpha, + ) + + def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: + return self._clamped._tex_clamped_swiglu_forward(swiglu_in, None) + + def _glu_backward( + self, + grad_swiglu_out: torch.Tensor, + swiglu_in: torch.Tensor, + ) -> torch.Tensor: + return self._clamped._tex_clamped_dswiglu( + grad_swiglu_out, + swiglu_in, + None, + ) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index a821258ebf..6b452b0182 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -24,7 +24,7 @@ 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, ScaledSwiGLU +from ..basic import GroupedLinear, ScaledClampedQGeGLU, ScaledSwiGLU from ..fuser import register_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( @@ -181,7 +181,7 @@ def _compute_grad_params( class BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8(FusedOperation): - """Fused op for MXFP8 GroupedLinear + ScaledSwiGLU + GroupedLinear + """Fused op for MXFP8 GroupedLinear + ScaledSwiGLU or ScaledClampedQGeGLU + GroupedLinear Uses experimental CuTe DSL kernel from cuDNN front-end. @@ -229,7 +229,7 @@ def __init__( self, *, fc1: GroupedLinear, - swiglu: ScaledSwiGLU, + swiglu: ScaledSwiGLU | ScaledClampedQGeGLU, fc2: GroupedLinear, ) -> None: super().__init__((fc1, swiglu, fc2)) @@ -237,6 +237,11 @@ def __init__( self.grouped_gemm_dglu_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" + ) def fuser_backward( self, @@ -433,7 +438,7 @@ def fuser_backward( "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, "current_stream": current_stream, "discrete_col_sfd": True, - "act_func": "dswiglu", + "act_func": self._cudnn_dact_func, "use_dynamic_sched": True, } diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index c5ce2b148d..afabec8392 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -20,7 +20,7 @@ from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer from ...constants import MXFP8_BLOCK_SCALING_SIZE -from ..basic import GroupedLinear, ScaledSwiGLU +from ..basic import GroupedLinear, ScaledClampedQGeGLU, ScaledSwiGLU from ..fuser import register_forward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( @@ -46,7 +46,7 @@ def _pack_grouped_linear_bias_for_cudnn(linear_op: GroupedLinear) -> Optional[to class ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8(FusedOperation): - """Fused op for MXFP8 GroupedLinear + ScaledSwiGLU + GroupedLinear + """Fused op for MXFP8 GroupedLinear + scaled GLU + GroupedLinear Uses experimental CuTe DSL kernel from cuDNN front-end. @@ -123,7 +123,7 @@ def __init__( self, *, fc1: GroupedLinear, - swiglu: ScaledSwiGLU, + swiglu: ScaledSwiGLU | ScaledClampedQGeGLU, fc2: GroupedLinear, ) -> None: super().__init__((fc1, swiglu, fc2)) @@ -131,6 +131,9 @@ def __init__( self.grouped_gemm_glu_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" def fuser_forward( self, @@ -339,7 +342,7 @@ def fuser_forward( "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, "current_stream": current_stream, "discrete_col_sfd": True, - "act_func": "swiglu", + "act_func": self._cudnn_act_func, "use_dynamic_sched": True, } From ac735380c1e9430833f4c6f76e378d0c7c1baa89 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 9 Apr 2026 15:03:44 -0700 Subject: [PATCH 329/521] [JAX] Fix BF16 tolerance for CGEMM + RS + BF16 test (#2860) update tols Signed-off-by: Phuong Nguyen --- examples/jax/collective_gemm/test_gemm.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/examples/jax/collective_gemm/test_gemm.py b/examples/jax/collective_gemm/test_gemm.py index c2db8fc44a..8221d7bbfd 100644 --- a/examples/jax/collective_gemm/test_gemm.py +++ b/examples/jax/collective_gemm/test_gemm.py @@ -151,8 +151,20 @@ def run_gemm_tests(args, mesh=None): jax.block_until_ready(gathered_output) if args.enable_result_check and args.process_id == 0: + # CGEMM + RS + BF16 uses TE's reduce_bf16 kernel (sequential left-to-right in FP32). + # With catastrophic cancellation the output is near zero while the absolute diff can + # reach 1 ULP of the partial GEMM magnitude (~0.0625 for typical transformer + # activations at O(8) scale), which exceeds the previous atol=1e-5. The 2x + # margin (0.125) covers this worst-case 1-ULP absolute difference. + is_cgemm_rs_bf16 = collective_op == CollectiveOp.REDUCE_SCATTER and not use_quantization + rtol = 1e-2 if is_cgemm_rs_bf16 else None + atol = 0.125 if is_cgemm_rs_bf16 else None assert_allclose( - gathered_ref_output, gathered_output, dtype=get_tolerance_dtype(quantizer_set) + gathered_ref_output, + gathered_output, + dtype=get_tolerance_dtype(quantizer_set), + rtol=rtol, + atol=atol, ) From 53fefa48c38cd73f50db82d4faec661d96f9811b Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Thu, 9 Apr 2026 16:22:42 -0600 Subject: [PATCH 330/521] add high precision init weights to fully_shard example (#2785) * add high precision init weights to fully_shard example Signed-off-by: Peter St. John * fix fully_shard example with preserve_high_precision_init_val, add test Signed-off-by: Peter St. John * addressing greptile review Signed-off-by: Peter St. John --------- Signed-off-by: Peter St. John --- .../quantized_model_init/fully_shard.py | 143 +++++++---------- .../fsdp2_tests/run_fsdp2_fused_adam.py | 147 +++++++++++++++++- transformer_engine/pytorch/module/base.py | 10 +- 3 files changed, 207 insertions(+), 93 deletions(-) diff --git a/examples/pytorch/quantized_model_init/fully_shard.py b/examples/pytorch/quantized_model_init/fully_shard.py index 6131712001..2b5ca84ebc 100644 --- a/examples/pytorch/quantized_model_init/fully_shard.py +++ b/examples/pytorch/quantized_model_init/fully_shard.py @@ -13,8 +13,11 @@ local shards on each rank's GPU. 2. ``quantized_model_init`` -- Flags the model for FP8 weight initialization (actual quantization happens in ``reset_parameters`` after sharding). -3. ``fully_shard`` -- PyTorch FSDP2 sharding of each TransformerLayer. -4. ``FusedAdam`` with FP32 master weights for full-precision training updates. +3. ``preserve_high_precision_init_val`` -- Keeps the original BF16 weight + values on CPU so they can seed the optimizer's FP32 master weights, + avoiding the precision loss of round-tripping through FP8. +4. ``fully_shard`` -- PyTorch FSDP2 sharding of each TransformerLayer. +5. ``FusedAdam`` with FP32 master weights for full-precision training updates. .. note:: ``fuse_wgrad_accumulation`` is **not** used here. That feature writes @@ -38,10 +41,10 @@ from torch.distributed.tensor import DTensor import transformer_engine.pytorch as te -from transformer_engine.pytorch import QuantizedTensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule +from transformer_engine.pytorch.quantized_tensor import QuantizedTensor -# ── Configuration (matches main.py) ────────────────────────────────── +# ── Configuration ──────────────────────────────────────────────────── HIDDEN_SIZE = 256 FFN_HIDDEN_SIZE = 1024 NUM_ATTENTION_HEADS = 8 @@ -49,7 +52,12 @@ SEQ_LEN = 32 BATCH_PER_RANK = 2 NUM_STEPS = 5 -DTYPE = torch.bfloat16 +# DTYPE is used for both params_dtype and activation tensors in this example. +# float32 is chosen for params_dtype so that the high-precision init values +# (which seed the optimizer's FP32 master weights) avoid a lossy BF16→FP8→FP32 +# round-trip. Using float32 for activations as well keeps the example simple; +# in production you would typically use BF16 activations inside te.autocast(). +DTYPE = torch.float32 def dist_print(msg): @@ -60,10 +68,6 @@ def dist_print(msg): def main(): # ── 1. Distributed setup ───────────────────────────────────────── - assert "TORCHELASTIC_RUN_ID" in os.environ, ( - "This script must be launched with torchrun, e.g.:\n" - " torchrun --nproc-per-node 2 fully_shard.py" - ) world_size = int(os.environ["WORLD_SIZE"]) local_rank = int(os.environ["LOCAL_RANK"]) @@ -74,10 +78,14 @@ def main(): torch.manual_seed(42) torch.cuda.manual_seed(42) - # ── 2. Create model on meta device (zero memory) ──────────────── - # quantized_model_init sets the flag for FP8 weight initialization, - # but with device="meta" no actual memory is allocated yet. - with te.quantized_model_init(enabled=True): + # ── 2. Create model on meta device (zero memory) ───────────────── + # quantized_model_init flags parameters for FP8 quantization. + # preserve_high_precision_init_val=True saves the original BF16 + # values on CPU so they can seed optimizer master weights later, + # avoiding the precision loss of dequantizing from FP8. + # We set DTYPE to float32 since these weights will actually be initialized as FP8, + # but we want to seed the optimizer states (which will be in FP32) with the FP32 values. + with te.quantized_model_init(enabled=True, preserve_high_precision_init_val=True): model = torch.nn.Sequential( *[ te.TransformerLayer( @@ -93,14 +101,10 @@ def main(): for _ in range(NUM_LAYERS) ] ) - - # Verify all parameters are on meta device (no GPU memory used). - for name, param in model.named_parameters(): - assert param.device == torch.device("meta"), f"{name} is not on meta device" dist_print("Model created on meta device (zero GPU memory).") - # ── 3. FSDP2 sharding ──────────────────────────────────────────── - # Apply sharding to the meta-device model. FSDP2 wraps parameters + # ── 3. FSDP2 sharding ─────────────────────────────────────────── + # Apply sharding to the meta-device model. FSDP2 wraps parameters # as DTensors but no GPU memory is allocated yet. mesh = DeviceMesh("cuda", list(range(world_size))) for child in model.children(): @@ -108,37 +112,42 @@ def main(): fully_shard(model, mesh=mesh) dist_print("FSDP2 sharding applied to meta-device model.") - # ── 4. Materialize parameters on GPU ────────────────────────────── + # ── 4. Materialize parameters on GPU ───────────────────────────── # reset_parameters() on each TE module materializes the local shard # on CUDA, applies weight initialization, and quantizes to FP8. + # Because preserve_high_precision_init_val=True, the pre-quantization + # BF16 values are saved on CPU for each local shard. for module in model.modules(): if isinstance(module, TransformerEngineBaseModule): module.reset_parameters() + dist_print("Parameters materialized on GPU.") - # Post-materialization verification. - for name, param in model.named_parameters(): - assert isinstance(param, DTensor), f"{name} is not a DTensor after sharding" - qt_count = sum( - 1 - for _, p in model.named_parameters() - if isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) - ) - assert qt_count > 0, "No QuantizedTensor local tensors after materialization" - dist_print( - f"Parameters materialized: {qt_count} FP8 (QuantizedTensor) weight params " - "wrapped in DTensors." - ) - - # ── 5. Optimizer ───────────────────────────────────────────────── + # ── 5. Optimizer with FP32 master weights ──────────────────────── optimizer = te.optimizers.FusedAdam( model.parameters(), lr=1e-3, master_weights=True, master_weight_dtype=torch.float32, ) - dist_print("Using FusedAdam with master_weights=True.") - # ── 6. Training loop ───────────────────────────────────────────── + # ── 6. Seed master weights from high-precision init values ─────── + # By default, FusedAdam initializes master weights by dequantizing + # the FP8 parameters, which introduces quantization noise. Instead, + # we seed them from the original BF16 init values preserved in step 2. + for name, param in model.named_parameters(): + optimizer.initialize_state(param, store_param_remainders=False) + local = param._local_tensor if isinstance(param, DTensor) else param + if isinstance(local, QuantizedTensor): + hp_val = local.get_high_precision_init_val() + assert hp_val.dtype == DTYPE, f"HP val dtype {hp_val.dtype}, expected {DTYPE}" + optimizer.set_scaled_state( + param, "master_param", hp_val.to(device=device, dtype=torch.float32) + ) + local.clear_high_precision_init_val() + + dist_print("Optimizer master weights seeded from high-precision init values.") + + # ── 7. Training loop ───────────────────────────────────────────── x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=DTYPE, device=device) target = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=DTYPE, device=device) @@ -153,56 +162,22 @@ def main(): optimizer.step() dist_print(f" Step {step}: loss = {loss.item():.6f}") - # ── 7. Post-training assertions ────────────────────────────────── - dist_print("\nVerifying invariants ...") - - qt_after = 0 - for name, param in model.named_parameters(): - assert isinstance(param, DTensor), f"{name} lost DTensor wrapping" - if isinstance(param._local_tensor, QuantizedTensor): - qt_after += 1 - assert qt_after > 0, "No QuantizedTensor local tensors after training" - dist_print(f" {qt_after} params still have QuantizedTensor local tensors.") - - # Optimizer states: master weights and moments should be float32. - for param in model.parameters(): - state = optimizer.state[param] - if "master_param" in state: - assert ( - state["master_param"].dtype == torch.float32 - ), f"Master weight dtype {state['master_param'].dtype}, expected float32" - assert state["exp_avg"].dtype == torch.float32, "exp_avg should be float32" - assert state["exp_avg_sq"].dtype == torch.float32, "exp_avg_sq should be float32" - - dist_print("All assertions passed!") - dist_print(" - Linear weight parameters: QuantizedTensor (FP8) wrapped in DTensor") - dist_print(" - Optimizer master weights: float32") - dist_print(" - Optimizer states (exp_avg, exp_avg_sq): float32") - # ── 8. Distributed checkpoint: save and load ───────────────────── # torch.distributed.checkpoint (DCP) saves sharded state — each rank - # writes only its local shard. This preserves FP8 compute weights - # and the full optimizer state (master weights, moments, step count). + # writes only its local shard, preserving FP8 compute weights and + # the full optimizer state (master weights, moments, step count). import torch.distributed.checkpoint as dcp - from torch.distributed.checkpoint.state_dict import ( - StateDictOptions, - get_model_state_dict, - get_optimizer_state_dict, - ) - # Use a fixed path so all ranks agree on the checkpoint location. checkpoint_dir = "/tmp/te_fsdp2_example_checkpoint" dist_print(f"\nSaving distributed checkpoint to {checkpoint_dir} ...") - # Save sharded checkpoint. DCP handles DTensor shards natively — - # each rank writes only its local shard to the filesystem. dcp.save( {"model": model.state_dict(), "optimizer": optimizer.state_dict()}, checkpoint_id=checkpoint_dir, ) dist_print(" Checkpoint saved (FP8 weights + optimizer state).") - # Load checkpoint back. Provide empty state dict containers with the + # Load checkpoint back. Provide empty state dict containers with the # same structure; DCP fills them from the saved files. state_to_load = {"model": model.state_dict(), "optimizer": optimizer.state_dict()} dcp.load(state_to_load, checkpoint_id=checkpoint_dir) @@ -225,6 +200,11 @@ def main(): # authoritative FP32 values (more precise than dequantizing FP8). # All ranks must participate in gathering; only rank 0 saves. from safetensors.torch import save_file + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + ) full_opts = StateDictOptions(full_state_dict=True, cpu_offload=True) @@ -238,10 +218,10 @@ def main(): for key, value in full_model_state.items(): if key in opt_param_states and "master_param" in opt_param_states[key]: - # Prefer optimizer's FP32 master weight (maintained throughout training). + # Prefer optimizer's FP32 master weight. fp32_state[key] = opt_param_states[key]["master_param"].float() - elif isinstance(value, QuantizedTensor): - # Fallback: dequantize FP8 → FP32 (e.g. if master_weights was off). + elif isinstance(value, te.QuantizedTensor): + # Fallback: dequantize FP8 → FP32. fp32_state[key] = value.dequantize().float() else: # Non-FP8 params (e.g. LayerNorm weights): cast to FP32. @@ -251,14 +231,7 @@ def main(): save_file(fp32_state, save_path) dist_print(f"\nSaved FP32 model ({len(fp32_state)} params) to {save_path}") - # Quick verification: all saved tensors are float32. - from safetensors.torch import load_file - - loaded = load_file(save_path) - for k, v in loaded.items(): - assert v.dtype == torch.float32, f"{k}: expected float32, got {v.dtype}" - dist_print(f" Verified: all {len(loaded)} tensors are float32.") - + dist.barrier() # wait for rank 0 to finish file I/O dist.destroy_process_group() 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 42df06ed7f..60a23b9394 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -14,6 +14,7 @@ Available --test values: fused_adam_fp8_master_weights, fused_adam_fp8_master_weights_no_meta, + fused_adam_fp8_high_precision_init, fused_adam_bf16, fused_adam_fp8_no_master, fused_adam_bf16_store_param_remainders, fuse_wgrad_accumulation, dcp_output_parity, dcp_output_parity_async, dcp_resharding_save, dcp_resharding_load, safetensors_fp32_export @@ -58,7 +59,14 @@ NUM_STEPS = 3 -def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_device=True): +def _build_model( + fp8_init, + fuse_wgrad_accumulation=False, + recipe=None, + use_meta_device=True, + preserve_high_precision_init_val=False, + params_dtype=torch.bfloat16, +): """Build a Sequential of TransformerLayers, optionally with FP8 init. When fp8_init=True and use_meta_device=True (the default), the model is @@ -74,7 +82,11 @@ def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_ data_ptr() == 0. """ if fp8_init: - ctx = te.quantized_model_init(enabled=True, recipe=recipe) + ctx = te.quantized_model_init( + enabled=True, + recipe=recipe, + preserve_high_precision_init_val=preserve_high_precision_init_val, + ) else: from contextlib import nullcontext @@ -82,7 +94,7 @@ def _build_model(fp8_init, fuse_wgrad_accumulation=False, recipe=None, use_meta_ kwargs = dict( fuse_wgrad_accumulation=fuse_wgrad_accumulation, fuse_qkv_params=True, - params_dtype=torch.bfloat16, + params_dtype=params_dtype, hidden_dropout=0.0, attention_dropout=0.0, ) @@ -253,6 +265,131 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe_name): optimizer.step() +def test_fused_adam_fp8_high_precision_init(recipe_name): + """FusedAdam with master_weights seeded from high-precision init values. + + Tests the preserve_high_precision_init_val=True path demonstrated in the + fully_shard.py example: + 1. Model is created with preserve_high_precision_init_val=True on meta device + 2. After FSDP2 sharding + materialization, each QuantizedTensor param has + a high-precision init value accessible via get_high_precision_init_val() + 3. These values seed the optimizer's FP32 master weights (avoiding FP8 + round-trip precision loss) + 4. Training completes successfully with correct optimizer state dtypes + """ + recipe = get_recipe_from_string(recipe_name) + + if recipe_name == "NVFP4BlockScaling": + pytest.xfail( + f"{recipe_name}: quantized_model_init and FSDP2 is not currently supported, since the " + "block tensor is dequantized before we flatten it for FSDP2." + ) + + world_size, device = _get_dist_info() + + model = _build_model( + fp8_init=True, + recipe=recipe, + preserve_high_precision_init_val=True, + params_dtype=torch.float32, + ) + model = _shard_model(model, world_size) + + # Verify params are DTensors with QuantizedTensor local shards + for name, param in model.named_parameters(): + assert isinstance(param, DTensor), f"{name} is not DTensor" + qt_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) + ) + assert qt_count > 0, "No QuantizedTensor local tensors after sharding" + + # Verify high-precision init values exist for all QuantizedTensor params + hp_val_count = 0 + for name, param in model.named_parameters(): + local = param._local_tensor if isinstance(param, DTensor) else param + if isinstance(local, QuantizedTensor): + hp_val = getattr(local, "get_high_precision_init_val", lambda: None)() + assert ( + hp_val is not None + ), f"{name}: QuantizedTensor param missing high-precision init value" + assert ( + hp_val.dtype == torch.float32 + ), f"{name}: HP init val dtype {hp_val.dtype}, expected float32" + hp_val_count += 1 + assert hp_val_count > 0, "No high-precision init values found" + + # Create optimizer and seed master weights from high-precision init values + optimizer = te.optimizers.FusedAdam( + model.parameters(), + lr=1e-3, + master_weights=True, + master_weight_dtype=torch.float32, + ) + + for name, param in model.named_parameters(): + optimizer.initialize_state(param, store_param_remainders=False) + local = param._local_tensor if isinstance(param, DTensor) else param + hp_val = getattr(local, "get_high_precision_init_val", lambda: None)() + if hp_val is not None: + optimizer.set_scaled_state( + param, "master_param", hp_val.to(device=device, dtype=torch.float32) + ) + local.clear_high_precision_init_val() + + # Verify high-precision init values are cleared after seeding + for name, param in model.named_parameters(): + local = param._local_tensor if isinstance(param, DTensor) else param + if isinstance(local, QuantizedTensor): + hp_val = getattr(local, "get_high_precision_init_val", lambda: None)() + assert ( + hp_val is None + ), f"{name}: high-precision init value not cleared after seeding optimizer" + + # Verify optimizer master weights are float32 + for param in model.parameters(): + state = optimizer.state[param] + if "master_param" in state: + assert ( + state["master_param"].dtype == torch.float32 + ), f"master_param dtype {state['master_param'].dtype}, expected float32" + + # Training loop + x = torch.randn(SEQ_LEN, BATCH_PER_RANK, HIDDEN_SIZE, dtype=torch.float32, device=device) + target = torch.randn_like(x) + + for step in range(NUM_STEPS): + optimizer.zero_grad(set_to_none=True) + with te.autocast(enabled=True, recipe=recipe): + output = model(x) + loss = F.mse_loss(output, target) + loss.backward() + optimizer.step() + + # Verify optimizer states after training + for param in model.parameters(): + state = optimizer.state[param] + assert ( + state["exp_avg"].dtype == torch.float32 + ), f"exp_avg dtype {state['exp_avg'].dtype}, expected float32" + assert ( + state["exp_avg_sq"].dtype == torch.float32 + ), f"exp_avg_sq dtype {state['exp_avg_sq'].dtype}, expected float32" + if "master_param" in state: + assert ( + state["master_param"].dtype == torch.float32 + ), f"master_param dtype {state['master_param'].dtype}, expected float32" + + # Verify FP8 params preserved after training + qt_count = sum( + 1 + for _, p in model.named_parameters() + if isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) + ) + assert qt_count > 0, "No QuantizedTensor local tensors after training" + + def test_fused_adam_bf16(recipe_name): """FusedAdam with master_weights + FSDP2 + bf16 params (no FP8). @@ -818,7 +955,8 @@ def test_dcp_resharding_save(recipe_name): model_state = model.state_dict() dcp.save( - {"model": model_state, "optimizer": optimizer.state_dict()}, checkpoint_id=checkpoint_dir + {"model": model_state, "optimizer": optimizer.state_dict()}, + checkpoint_id=checkpoint_dir, ) dist.barrier() @@ -918,6 +1056,7 @@ def test_dcp_resharding_load(recipe_name): TESTS = { "fused_adam_fp8_master_weights": test_fused_adam_fp8_master_weights, "fused_adam_fp8_master_weights_no_meta": test_fused_adam_fp8_master_weights_no_meta, + "fused_adam_fp8_high_precision_init": test_fused_adam_fp8_high_precision_init, "fused_adam_bf16": test_fused_adam_bf16, "fused_adam_fp8_no_master": test_fused_adam_fp8_no_master, "fused_adam_bf16_store_param_remainders": test_fused_adam_bf16_store_param_remainders, diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 1b237ece29..a13eb0c7e6 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -1379,10 +1379,12 @@ def clear(self): if hasattr(self, "_high_precision_init_val"): del self._high_precision_init_val - param._high_precision_init_val = high_precision_init_val - param.get_high_precision_init_val = MethodType(get, param) - param.clear_high_precision_init_val = MethodType(clear, param) - # Update the parameter based on its type + # DTensor.from_local() does not preserve object identity, + # so attach to the DTensor's local tensor when applicable. + target = dtensor_param._local_tensor if is_dtensor else param + target._high_precision_init_val = high_precision_init_val + target.get_high_precision_init_val = MethodType(get, target) + target.clear_high_precision_init_val = MethodType(clear, target) if not is_dtensor: self.module_setattr(name, param) From 2f17c9b9579d6410de1c3ab819d38c1669f832f7 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Fri, 10 Apr 2026 11:22:37 -0700 Subject: [PATCH 331/521] Enforce minimum NCCL version for cuBLASMp (#2857) Signed-off-by: Vladimir Cherepanov --- transformer_engine/common/CMakeLists.txt | 44 ++++++++++++++++++++++-- 1 file changed, 42 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 7c223e6917..a4fbfd9e9c 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -98,6 +98,39 @@ set(CUTLASS_TOOLS_INCLUDE_DIR # Python find_package(Python COMPONENTS Interpreter Development.Module REQUIRED) +function(find_nccl_version OUT_VERSION OUT_INCLUDE_DIR) + find_path(_nvte_nccl_include_dir + NAMES nccl.h + PATH_SUFFIXES include + REQUIRED) + + file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_major_line + REGEX "^#define NCCL_MAJOR[ \t]+[0-9]+$") + file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_minor_line + REGEX "^#define NCCL_MINOR[ \t]+[0-9]+$") + file(STRINGS "${_nvte_nccl_include_dir}/nccl.h" _nvte_nccl_patch_line + REGEX "^#define NCCL_PATCH[ \t]+[0-9]+$") + + string(REGEX REPLACE "^#define NCCL_MAJOR[ \t]+([0-9]+)$" "\\1" + _nvte_nccl_major "${_nvte_nccl_major_line}") + string(REGEX REPLACE "^#define NCCL_MINOR[ \t]+([0-9]+)$" "\\1" + _nvte_nccl_minor "${_nvte_nccl_minor_line}") + string(REGEX REPLACE "^#define NCCL_PATCH[ \t]+([0-9]+)$" "\\1" + _nvte_nccl_patch "${_nvte_nccl_patch_line}") + + if ("${_nvte_nccl_major}" STREQUAL "" + OR "${_nvte_nccl_minor}" STREQUAL "" + OR "${_nvte_nccl_patch}" STREQUAL "") + message(FATAL_ERROR + "Failed to parse NCCL version from ${_nvte_nccl_include_dir}/nccl.h") + endif() + + set(${OUT_VERSION} + "${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch}" + PARENT_SCOPE) + set(${OUT_INCLUDE_DIR} "${_nvte_nccl_include_dir}" PARENT_SCOPE) +endfunction() + # Configure Transformer Engine library include_directories(${PROJECT_SOURCE_DIR}/..) set(transformer_engine_SOURCES) @@ -290,6 +323,7 @@ option(NVTE_WITH_CUBLASMP "Use cuBLASMp for tensor parallel GEMMs" OFF) if (NVTE_WITH_CUBLASMP) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUBLASMP) target_include_directories(transformer_engine PRIVATE ${CUBLASMP_DIR}/include) + find_nccl_version(NCCL_VERSION NCCL_INCLUDE_DIR) find_library(CUBLASMP_LIB NAMES cublasmp libcublasmp PATHS ${CUBLASMP_DIR} @@ -299,8 +333,14 @@ if (NVTE_WITH_CUBLASMP) NAMES nccl libnccl PATH_SUFFIXES lib REQUIRED) - target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CUBLASMP_LIB}) - message(STATUS "Using cuBLASMp at: ${CUBLASMP_DIR}") + if (NCCL_VERSION VERSION_LESS 2.29.0) + message(FATAL_ERROR + "NVTE_WITH_CUBLASMP requires NCCL >= 2.29.0, but found NCCL ${NCCL_VERSION} " + "in ${NCCL_INCLUDE_DIR}/nccl.h") + endif() + target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CUBLASMP_LIB}) + message(STATUS "Using cuBLASMp at: ${CUBLASMP_DIR}") + message(STATUS "Using NCCL ${NCCL_VERSION} at: ${NCCL_LIB}") endif() # Number of philox4x32 rounds for stochastic rounding (build-time constant). From 580e7aa28bebb83489271107ce50861cf84f3170 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 10 Apr 2026 13:22:06 -0700 Subject: [PATCH 332/521] Bias Prob Scaling for GroupedLinear and Fused MOE Layers (#2864) * bias*prob, dbias+dprob triton kernel Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update bias parameterization in tests for fusable ops Signed-off-by: vthumbe1503 * Update transformer_engine/pytorch/ops/basic/grouped_linear.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 * address review comments + lint fix Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Update docstring Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 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> --- tests/pytorch/test_fusible_ops.py | 15 +-- .../common/triton/grouped_dbias_dscales.py | 89 ++++++++++++++++++ .../pytorch/ops/basic/grouped_linear.py | 87 ++++++++++++++--- .../pytorch/ops/fused/backward_grouped_mlp.py | 57 +++++++---- .../pytorch/ops/fused/forward_grouped_mlp.py | 8 +- .../pytorch/triton/grouped_dbias_dscales.py | 94 +++++++++++++++++++ 6 files changed, 312 insertions(+), 38 deletions(-) create mode 100644 transformer_engine/common/triton/grouped_dbias_dscales.py create mode 100644 transformer_engine/pytorch/triton/grouped_dbias_dscales.py diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 795cbf3452..a2de8014a1 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -12,9 +12,10 @@ from typing import Optional import pytest -import torch import transformer_engine +import torch + import transformer_engine.common.recipe import transformer_engine.pytorch as te import transformer_engine.pytorch.ops as te_ops @@ -3375,9 +3376,6 @@ 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 quantization == "mxfp8" and bias: - # Will be supported in future CUDNN release. - pytest.skip("Bias/dbias not yet supported in MXFP8 fused grouped MLP") 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") @@ -3478,7 +3476,9 @@ def test_grouped_mlp( x2c = torch.clamp(x2, -lim, lim) x = (x2c + 1) * (x1c * torch.sigmoid(geglu_alpha * x1c)) x = x * probs[group_idx].unsqueeze(-1) - x = torch.nn.functional.linear(x, fc2_ws_ref[group_idx], bias=fc2_bs_ref[group_idx]) + x = torch.nn.functional.linear(x, fc2_ws_ref[group_idx]) + if bias: + x = x + fc2_bs_ref[group_idx] * probs[group_idx].unsqueeze(-1) ys.append(x) y_ref = torch.cat(ys) y_ref.backward(dy_ref) @@ -3503,6 +3503,7 @@ def test_grouped_mlp( accumulate_into_main_grad=accumulate_into_main_grad, delay_wgrad_compute=delay_wgrad_compute, ) + fc2 = te_ops.GroupedLinear( group_size, hidden_size, @@ -3514,6 +3515,7 @@ def test_grouped_mlp( 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, @@ -3578,7 +3580,8 @@ def test_grouped_mlp( # Fuse ops and perform forward and backward pass with te.autocast(enabled=with_quantization, recipe=recipe): - y_test = module(x_test, split_sizes, probs_test, split_sizes) + fc2_extra = (split_sizes, probs_test) if bias else (split_sizes,) + y_test = module(x_test, split_sizes, probs_test, *fc2_extra) y_test.backward(dy_test) if delay_wgrad_compute: fc1.backward_dw() diff --git a/transformer_engine/common/triton/grouped_dbias_dscales.py b/transformer_engine/common/triton/grouped_dbias_dscales.py new file mode 100644 index 0000000000..f5ddda2593 --- /dev/null +++ b/transformer_engine/common/triton/grouped_dbias_dscales.py @@ -0,0 +1,89 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused grouped dbias + dscales Triton kernel.""" + +import triton +import triton.language as tl + + +@triton.jit +def _grouped_dbias_dscales_kernel( + dy_ptr, + scales_ptr, + bias_ptr, + dbias_ptr, + dscales_ptr, + offsets_ptr, + hidden, + 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]). + + 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. + + - dbias: accumulated in registers, one atomic-add at the end + (N_ROW_SPLITS contributors per group). + - dscales: atomic-add per iteration across column tiles + (cdiv(hidden, BLOCK_H) contributors per element). + """ + group_idx = tl.program_id(0) + row_split = tl.program_id(1) + col_block = tl.program_id(2) + + row_start = tl.load(offsets_ptr + group_idx) + row_end = tl.load(offsets_ptr + group_idx + 1) + + group_rows = row_end - row_start + 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) + + 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, + ) + + tl.atomic_add( + dbias_ptr + group_idx * hidden + col_offs, + dbias_acc, + mask=col_mask, + ) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index f26a337a4d..0e09c8a38b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -35,6 +35,7 @@ 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 class GroupedLinear(BasicOperation): @@ -79,10 +80,15 @@ class GroupedLinear(BasicOperation): 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}``. + 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`` + in the forward pass. The scale tensor has shape + ``(total_tokens,)`` and is split according to the split sizes. """ - # Operation expects input split sizes + # Operation expects input split sizes (and optionally scales tensor) num_extra_inputs: int = 1 def __init__( @@ -99,9 +105,14 @@ def __init__( single_grouped_weight: bool = False, single_grouped_bias: bool = False, delay_wgrad_compute: bool = False, + scale_bias: bool = False, ) -> None: super().__init__() + self._scale_bias: bool = scale_bias and bias + if self._scale_bias: + self.num_extra_inputs = 2 + self.wgrad_store = WeightGradStore(delay_wgrad_compute) # Weight tensor dimensions @@ -221,6 +232,17 @@ def backward_dw(self) -> None: w = getattr(self, f"weight{group_idx}") w.grad = grad_weights[group_idx].to(w.dtype) + def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: + """Retrieve per-group bias tensors in the given dtype.""" + if self.single_grouped_bias: + bias_parts = self.bias.quantized_tensors + if bias_parts is None: + bias_parts = self.bias.split_into_quantized_tensors() + return [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts] + return [ + maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(self.num_groups) + ] + def num_quantizers(self, mode: str) -> int: if mode == "forward": return 2 * self.num_groups @@ -700,6 +722,11 @@ def fuser_forward( if len(split_sizes_int) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {len(split_sizes_int)}.") + # Extract scales tensor for bias scaling + scales = None + if self._scale_bias: + scales = basic_op_extra_inputs[0][1] + # Extract params if self.single_grouped_weight: weights = self.weight.quantized_tensors @@ -746,6 +773,7 @@ def fuser_forward( out = torch.empty(out_shape, dtype=dtype, device=device) # Perform GEMMs + use_gemm_bias = has_bias and not self._scale_bias general_grouped_gemm( ws, xs, @@ -753,12 +781,22 @@ def fuser_forward( [None] * num_groups, # quantization_params dtype, m_splits=split_sizes_int, - bias=bs, - use_bias=has_bias, + bias=bs if use_gemm_bias else None, + use_bias=use_gemm_bias, use_split_accumulator=_2X_ACC_FPROP, single_output=True, ) + # 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) + for i in range(num_groups): + out_splits[i].add_(bs[i].unsqueeze(0) * scales_splits[i].unsqueeze(-1)) + # Prepare weight tensors for backward pass if not input_requires_grad: ws = [None] * num_groups @@ -776,7 +814,12 @@ def fuser_forward( # Save state for backward pass if ctx.requires_grad: - ctx.save_for_backward(split_sizes, *xs, *ws) + saved = [split_sizes] + 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 @@ -808,6 +851,9 @@ def fuser_backward( ctx = basic_op_ctxs[0] saved_tensors = ctx.saved_tensors split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] + scales = None + if self._scale_bias: + scales, saved_tensors = saved_tensors[0], saved_tensors[1:] xs, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] @@ -816,6 +862,7 @@ def fuser_backward( dy = maybe_dequantize(grad_output, ctx.dtype) dys = None grad_biases = [None] * num_groups + grad_scales = None if ctx.with_quantized_compute: for quantizer in ctx.grad_output_quantizers: quantizer.set_usage( @@ -823,15 +870,27 @@ def fuser_backward( columnwise=ctx.weight_requires_grad, ) dys = tex.split_quantize(dy, split_sizes_int, ctx.grad_output_quantizers) - if has_bias: - grad_biases = [ - dy.reshape(-1, dy.size(-1)).sum(dim=0) - for dy in torch.split(grad_output, split_sizes_int) - ] + 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: - grad_biases = [dy.reshape(-1, dy.size(-1)).sum(dim=0) for dy in dys] + 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) + 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)] # Initialize grad weight buffers accumulate_into_main_grad = self._accumulate_into_main_grad @@ -965,7 +1024,8 @@ def fuser_backward( grad_params = grad_biases + [grad_weight] else: grad_params = [grad_weight] - return grad_input, [grad_params], [(None,)] + 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"): @@ -1001,4 +1061,5 @@ def fuser_backward( else: grad_params = list(final_weight_grads) + list(grad_biases) - return grad_input, [grad_params], [(None,)] + 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 6b452b0182..357e8b3695 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -32,6 +32,7 @@ maybe_dequantize, validate_grouped_mlp_dims, ) +from ...triton.grouped_dbias_dscales import _compute_grouped_dbias_dscales @functools.lru_cache(maxsize=1) @@ -321,6 +322,7 @@ def fuser_backward( raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") split_sizes = split_sizes.to(dtype=torch.int64, device=device) split_points = split_points.to(dtype=torch.int, device=device) + scale_bias = fc2_op._scale_bias and fc2_op.has_bias grouped_fc1_x = None if fc1_ctx.weight_requires_grad: @@ -357,6 +359,7 @@ def fuser_backward( fc2_ctx.grad_output_quantizer.optimize_for_gemm = True output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None + fc2_dy = None if ( not output_fc2_dbias and isinstance(grad_output, GroupedTensor) @@ -365,7 +368,7 @@ def fuser_backward( grouped_fc2_dy = grad_output else: fc2_dy = maybe_dequantize(grad_output, dtype) - if output_fc2_dbias: + 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, @@ -380,16 +383,6 @@ def fuser_backward( split_sizes, ) - fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None - fc2_bias_grad_packed: Optional[torch.Tensor] = None - if fc2_dbias_packed is not None: - if fc2_op.single_grouped_bias: - fc2_bias_grad_packed = fc2_dbias_packed.to(dtype=dtype) - else: - fc2_bias_grads = [ - fc2_dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups) - ] - # Pack data tensors # Note: Fused kernel expects tensor with non-contiguous # logical dims. @@ -419,8 +412,8 @@ def fuser_backward( norm_const_tensor = get_cached_ones_tensor(1, dtype, device) current_stream = torch.cuda.current_stream().cuda_stream - prob_tensor = scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) - dprob_tensor = torch.zeros_like(prob_tensor) + scales_tensor = scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + dscales_tensor = torch.zeros_like(scales_tensor) fc2_dglu_kwargs = { "a_tensor": fc2_dy_data, @@ -429,8 +422,8 @@ def fuser_backward( "padded_offsets": split_points, "alpha_tensor": alpha_tensor, "beta_tensor": alpha_tensor, - "prob_tensor": prob_tensor, - "dprob_tensor": dprob_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, @@ -488,8 +481,35 @@ def fuser_backward( 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]) fc1_dy_col_scale = fc2_dgrad_kernel_out["sfd_col_tensor"] - grad_scales = fc2_dgrad_kernel_out["dprob_tensor"] - grad_scales = grad_scales.view(-1).to(dtype=dtype) + grad_scales = fc2_dgrad_kernel_out["dprob_tensor"].view(-1) + + fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None + fc2_bias_grad_packed: Optional[torch.Tensor] = None + if scale_bias: + fc2_biases = fc2_op._get_bias_tensors(dtype) + bias_packed = torch.stack(fc2_biases) + scales_f32 = scales.detach().to(dtype=torch.float32) + fc2_dbias_packed_result, grad_scales = _compute_grouped_dbias_dscales( + fc2_dy, + scales_f32, + bias_packed, + offsets=fc1_ctx.base_split_offsets, + dscales=grad_scales, + ) + fc2_dbias_packed_result = fc2_dbias_packed_result.to(dtype=dtype) + if fc2_op.single_grouped_bias: + fc2_bias_grad_packed = fc2_dbias_packed_result + else: + fc2_bias_grads = [fc2_dbias_packed_result[idx] for idx in range(num_groups)] + elif fc2_dbias_packed is not None: + if fc2_op.single_grouped_bias: + fc2_bias_grad_packed = fc2_dbias_packed.to(dtype=dtype) + else: + fc2_bias_grads = [ + fc2_dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups) + ] + + grad_scales = grad_scales.to(dtype=dtype) fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc1_bias_grad_packed: Optional[torch.Tensor] = None @@ -643,10 +663,11 @@ def fuser_backward( grouped_fc1_x.columnwise_scale_inv, ) + fc2_grad_extra = (None, None) if fc2_op._scale_bias else (None,) return ( grad_input, [fc1_grad_params, (), fc2_grad_params], - [(None,), (grad_scales,), (None,)], + [(None,), (grad_scales,), fc2_grad_extra], ) diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index afabec8392..83bb4428f0 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -427,13 +427,19 @@ def fuser_forward( # 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.float(), "norm_const_tensor": None, - "prob_tensor": torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device), + "prob_tensor": fc2_scales_tensor, "acc_dtype": torch.float32, "c_dtype": dtype, "d_dtype": dtype, diff --git a/transformer_engine/pytorch/triton/grouped_dbias_dscales.py b/transformer_engine/pytorch/triton/grouped_dbias_dscales.py new file mode 100644 index 0000000000..f87130b7c8 --- /dev/null +++ b/transformer_engine/pytorch/triton/grouped_dbias_dscales.py @@ -0,0 +1,94 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrapper for the fused grouped dbias + dscales Triton kernel.""" + +from typing import Optional, Tuple + +import torch +import triton + +from transformer_engine.common.triton.grouped_dbias_dscales import ( + _grouped_dbias_dscales_kernel, +) + + +def _compute_grouped_dbias_dscales( + dy: torch.Tensor, + scales: torch.Tensor, + bias: torch.Tensor, + offsets: torch.Tensor, + 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. + + Computes the following, where token *i* belongs to group *g(i)*: + + dbias[g, j] += sum_{i in group 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 + accumulated into (useful for fusing with upstream gradients). + + Args: + dy: (total_tokens, hidden) -- FC2 output grad. + scales: (total_tokens,) float32 -- per-token routing scales. + 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. + + Returns: + dbias: (num_groups, hidden) float32 + dscales: (total_tokens,) float32 + """ + num_groups = bias.shape[0] + hidden = dy.shape[1] + total_tokens = dy.shape[0] + + 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_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}" + + BLOCK_M = 128 + BLOCK_H = 128 + N_ROW_SPLITS = 4 + + 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, + ) + + return dbias, dscales From 323582fe68218533b4ae3c2b23471d1cbf15d9be Mon Sep 17 00:00:00 2001 From: Cory Ye <44509866+cspades@users.noreply.github.com> Date: Fri, 10 Apr 2026 18:05:53 -0700 Subject: [PATCH 333/521] Add Megatron-FSDP E2E integration test to TE CI/CD (L1). (#2845) * Add Megatron-FSDP E2E integration test to TE CI/CD (L1). Signed-off-by: Cory Ye * Update qa/L1_pytorch_mcore_fsdp_integration/test.sh Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Cory Ye <44509866+cspades@users.noreply.github.com> * Explicit torchrun invoke. Signed-off-by: Cory Ye * Edits. Signed-off-by: Cory Ye * Remove CPU initialization, add FW args. Signed-off-by: Cory Ye * Expose MCore hash/tag as an argument to the E2E script. Signed-off-by: Cory Ye * Bump MCore commit. Signed-off-by: Cory Ye --------- Signed-off-by: Cory Ye Signed-off-by: Cory Ye <44509866+cspades@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../.gitignore | 2 + .../merges.txt | 1 + qa/L1_pytorch_mcore_fsdp_integration/test.sh | 91 +++++++++++++++++++ 3 files changed, 94 insertions(+) create mode 100644 qa/L1_pytorch_mcore_fsdp_integration/.gitignore create mode 100644 qa/L1_pytorch_mcore_fsdp_integration/merges.txt create mode 100644 qa/L1_pytorch_mcore_fsdp_integration/test.sh diff --git a/qa/L1_pytorch_mcore_fsdp_integration/.gitignore b/qa/L1_pytorch_mcore_fsdp_integration/.gitignore new file mode 100644 index 0000000000..46426003ca --- /dev/null +++ b/qa/L1_pytorch_mcore_fsdp_integration/.gitignore @@ -0,0 +1,2 @@ +Megatron-LM +vocab.json \ No newline at end of file diff --git a/qa/L1_pytorch_mcore_fsdp_integration/merges.txt b/qa/L1_pytorch_mcore_fsdp_integration/merges.txt new file mode 100644 index 0000000000..5e7f1fd949 --- /dev/null +++ b/qa/L1_pytorch_mcore_fsdp_integration/merges.txt @@ -0,0 +1 @@ +#version: 0.2 diff --git a/qa/L1_pytorch_mcore_fsdp_integration/test.sh b/qa/L1_pytorch_mcore_fsdp_integration/test.sh new file mode 100644 index 0000000000..d63c66f2ea --- /dev/null +++ b/qa/L1_pytorch_mcore_fsdp_integration/test.sh @@ -0,0 +1,91 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +set -e + +# Megatron-LM / Megatron-FSDP commit for main branch on Apr. 10, 2026. +# Necessary to support wgrad accumulate fusion and Megatron-FSDP NCCL UBR, +# and fixes decoupled_grad <> DistOpt usage in Megatron-LM. +MCORE_REF=${1:-ab43d43f0bc04f4656d4af15afb6e7e4c9ad71c8} + +# Paths +: ${TE_PATH:=/opt/transformerengine} +: ${MCORE_PATH:=${TE_PATH}/qa/L1_pytorch_mcore_fsdp_integration/Megatron-LM} + +# Download Megatron-LM if needed +if [ ! -d "${MCORE_PATH}" ]; then + pushd $(dirname ${MCORE_PATH}) + git clone https://github.com/NVIDIA/Megatron-LM.git Megatron-LM + pushd Megatron-LM && git checkout "${MCORE_REF}" && popd + popd +fi + +# Create mock vocab +VOCAB_FILE=${TE_PATH}/qa/L1_pytorch_mcore_fsdp_integration/vocab.json +printf "" > ${VOCAB_FILE} +printf "{" >> ${VOCAB_FILE} +printf "\"<|endoftext|>\": 0" >> ${VOCAB_FILE} +seq 1 4095 | awk '{ printf(", \"%d\": %d", $1, $1) }' >> ${VOCAB_FILE} +printf "}" >> ${VOCAB_FILE} + +# Setting CUDA_DEVICE_MAX_CONNECTIONS limits +# Megatron-FSDP stream parallelism. +unset CUDA_DEVICE_MAX_CONNECTIONS +export NVTE_TORCH_COMPILE=0 +export NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 +export NVTE_FLASH_ATTN=1 +export NVTE_FWD_LAYERNORM_SM_MARGIN=0 +export NVTE_BWD_LAYERNORM_SM_MARGIN=0 +export NVTE_BIAS_GELU_NVFUSION=0 +export NVTE_BIAS_DROPOUT_FUSION=0 + +# V1 offloading has bugs that are exposed by Megatron-FSDP. +# This test will focus on validating the new offloading code. +# Un-set the Megatron-LM default of V1. +export NVTE_CPU_OFFLOAD_V1=0 + +# Megatron-LM command to run Megatron-FSDP. +python3 \ +-m torch.distributed.launch \ +--use_env \ +--nnodes=1 \ +--nproc_per_node=$(nvidia-smi -L | wc -l) \ +${MCORE_PATH}/pretrain_gpt.py \ +--tensor-model-parallel-size 1 \ +--pipeline-model-parallel-size 1 \ +--num-layers 2 \ +--hidden-size 128 \ +--num-attention-heads 8 \ +--swiglu \ +--seq-length 128 \ +--max-position-embeddings 128 \ +--micro-batch-size 1 \ +--global-batch-size 8 \ +--train-iters 10 \ +--eval-iters 10 \ +--eval-interval 100 \ +--lr 1e-4 \ +--mock-data \ +--vocab-file ${VOCAB_FILE} \ +--merge-file ${TE_PATH}/qa/L1_pytorch_mcore_fsdp_integration/merges.txt \ +--transformer-impl transformer_engine \ +--use-megatron-fsdp \ +--data-parallel-sharding-strategy optim_grads_params \ +--use-distributed-optimizer \ +--use-precision-aware-optimizer \ +--num-distributed-optimizer-instances 2 \ +--outer-dp-sharding-strategy optim \ +--use-nccl-ub \ +--fsdp-double-buffer \ +--fsdp-manual-registration \ +--fp8-format hybrid \ +--fp8-param-gather \ +--fp8-recipe mxfp8 \ +--cpu-offloading-num-layers 1 \ +--overlap-grad-reduce \ +--overlap-param-gather \ +--ckpt-format fsdp_dtensor \ +--init-model-with-meta-device \ +--bf16 \ +--grad-reduce-in-bf16 From 2dd31bb849e83cce51c7d169db883862063d3a95 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=A9tan=20Lepage?= Date: Sat, 11 Apr 2026 04:31:22 +0200 Subject: [PATCH 334/521] Fix JAX extension build with NVTE_UB_WITH_MPI=1 (#2835) * Fix JAX extension build with NVTE_UB_WITH_MPI=1 Signed-off-by: Gaetan Lepage * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Gaetan Lepage 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> --- build_tools/jax.py | 5 ++++- build_tools/pytorch.py | 17 +++++++++-------- build_tools/utils.py | 11 +++++++++++ 3 files changed, 24 insertions(+), 9 deletions(-) diff --git a/build_tools/jax.py b/build_tools/jax.py index f07c0a202f..a7b200f915 100644 --- a/build_tools/jax.py +++ b/build_tools/jax.py @@ -3,13 +3,14 @@ # See LICENSE for license information. """JAX related extensions.""" + import os from pathlib import Path from packaging import version import setuptools -from .utils import get_cuda_include_dirs, all_files_in_dir, debug_build_enabled +from .utils import get_cuda_include_dirs, all_files_in_dir, debug_build_enabled, setup_mpi_flags from typing import List @@ -100,6 +101,8 @@ def setup_jax_extension( else: cxx_flags.append("-g0") + setup_mpi_flags(include_dirs, cxx_flags) + # Define TE/JAX as a Pybind11Extension from pybind11.setup_helpers import Pybind11Extension diff --git a/build_tools/pytorch.py b/build_tools/pytorch.py index fdfdee9b1c..533addaf53 100644 --- a/build_tools/pytorch.py +++ b/build_tools/pytorch.py @@ -3,12 +3,19 @@ # See LICENSE for license information. """PyTorch related extensions.""" + import os from pathlib import Path import setuptools -from .utils import all_files_in_dir, cuda_version, get_cuda_include_dirs, debug_build_enabled +from .utils import ( + all_files_in_dir, + cuda_version, + get_cuda_include_dirs, + debug_build_enabled, + setup_mpi_flags, +) from typing import List @@ -67,13 +74,7 @@ def setup_pytorch_extension( if version < (12, 0): raise RuntimeError("Transformer Engine requires CUDA 12.0 or newer") - if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): - assert ( - os.getenv("MPI_HOME") is not None - ), "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!" - mpi_path = Path(os.getenv("MPI_HOME")) - include_dirs.append(mpi_path / "include") - cxx_flags.append("-DNVTE_UB_WITH_MPI") + setup_mpi_flags(include_dirs, cxx_flags) library_dirs = [] libraries = [] diff --git a/build_tools/utils.py b/build_tools/utils.py index 885901068a..d0f5eab425 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -341,6 +341,17 @@ def get_frameworks() -> List[str]: return _frameworks +def setup_mpi_flags(include_dirs: List, cxx_flags: List) -> None: + """Add MPI include path and compile definition if NVTE_UB_WITH_MPI is enabled.""" + if bool(int(os.getenv("NVTE_UB_WITH_MPI", "0"))): + assert ( + os.getenv("MPI_HOME") is not None + ), "MPI_HOME=/path/to/mpi must be set when compiling with NVTE_UB_WITH_MPI=1!" + mpi_path = Path(os.getenv("MPI_HOME")) + include_dirs.append(mpi_path / "include") + cxx_flags.append("-DNVTE_UB_WITH_MPI") + + def copy_common_headers( src_dir: Union[Path, str], dst_dir: Union[Path, str], From 2b78e55ed788eab607ec5218703549547d8035c2 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, 13 Apr 2026 16:14:43 +0200 Subject: [PATCH 335/521] [PyTorch] Remove unnecessary save of weights (#2549) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * code drop Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * added test Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix Signed-off-by: Pawel Gadzinski * remove unnecessary code Signed-off-by: root * Update transformer_engine/pytorch/module/layernorm_linear.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> * fix Signed-off-by: root * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Pawel Gadzinski Signed-off-by: root Signed-off-by: Paweł Gadziński <62263673+pggPL@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> --- tests/pytorch/test_sanity.py | 7 + .../pytorch/module/grouped_linear.py | 64 +++++---- .../pytorch/module/layernorm_linear.py | 57 ++++---- .../pytorch/module/layernorm_mlp.py | 125 ++++++++++-------- transformer_engine/pytorch/module/linear.py | 80 ++++++----- 5 files changed, 186 insertions(+), 147 deletions(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index be123f8c23..7f2f24fd69 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -225,6 +225,7 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, dtype, config, fp8_reci continue elif "weight" in name and p.requires_grad: p.main_grad = torch.zeros_like(p) + p.grad_added_to_main_grad = False # Should be set to True after backward use_fp8 = fp8_recipe is not None with autocast(enabled=use_fp8, recipe=fp8_recipe): @@ -234,13 +235,19 @@ def _test_sanity_e2e_gradient_accumulation_fusion(block, dtype, config, fp8_reci torch.cuda.synchronize() failed_grads = [] + failed_grad_added_flags = [] for name, p in block.named_parameters(): if "layer_norm_weight" in name: continue elif "weight" in name and p.requires_grad: if not torch.count_nonzero(p.main_grad) > 0: failed_grads.append(name) + if not getattr(p, "grad_added_to_main_grad", False): + failed_grad_added_flags.append(name) assert len(failed_grads) == 0, f"Gradient not accumulated for {failed_grads}." + assert ( + len(failed_grad_added_flags) == 0 + ), f"grad_added_to_main_grad not set to True for {failed_grad_added_flags}." def _test_sanity_e2e(block, dtype, config, fp8_recipe, skip_wgrad): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 2cce6c3ef8..188a1728db 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -6,6 +6,7 @@ from typing import Union, Optional, Callable, Tuple, List from itertools import chain import warnings +import weakref import functools import torch @@ -260,19 +261,6 @@ def forward( else: inputmats = [None] * num_gemms - if cpu_offloading: - ctx.grad_added_to_main_grad = hasattr(weights[0], "grad_added_to_main_grad") - - if ctx.grad_added_to_main_grad: - # If you are passing torch.nn.Parameter through the Torch hooks, you will - # get back torch.Tensor. Torch rips off the Parameter wrapper. - # You need to preserve the weight object to have all the attributes user - # sets for the weights. Because of this, it is not recommended to offload - # weights if weights are externally touched outside this module - ctx.weight_objects = [] - for weight in weights: - ctx.weight_objects.append(weight) - tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, @@ -288,6 +276,12 @@ def forward( ctx.weights_requires_grad = weights[0].requires_grad if fuse_wgrad_accumulation and ctx.weights_requires_grad: + # Keep weakrefs to weights to preserve attributes like main_grad + # when we need to modify the weight python objects + ctx.origin_weight_refs = [weakref.ref(w) for w in weights] + ctx.origin_weights_overwrite_main_grad = getattr( + weights[0], "overwrite_main_grad", False + ) # This check is needed to ensure that main_grad is not created # during the forward pass when using MCore FSDP as it creates # the main_grad buffer lazily before backprop @@ -298,8 +292,6 @@ def forward( ctx.main_grad_funcs = [ lambda j=i: weights[j].main_grad for i in range(num_gemms) ] - else: - ctx.main_grad_funcs = [lambda: None for i in range(num_gemms)] ctx.device = device ctx.output_quantizers = output_quantizers ctx.m_splits = m_splits @@ -350,19 +342,25 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], N = ctx.num_gemms inputmats = saved_tensors[:N] weights = saved_tensors[N : 2 * N] - origin_weights = saved_tensors[2 * N : 3 * N] + saved_weights = saved_tensors[2 * N : 3 * N] biases = saved_tensors[3 * N : 4 * N] - main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] - - if ctx.cpu_offloading: - if ctx.grad_added_to_main_grad: - for i, weight in enumerate(ctx.weight_objects): - origin_weights[i] = ctx.weight_objects[i] - ctx.weight_objects[i] = None - if ctx.fuse_wgrad_accumulation: - for i in range(N): - origin_weights[i].main_grad = main_grads[i] + # Restore from weakrefs to get original weight python objects + # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) + # Only needed when fuse_wgrad_accumulation is enabled. + 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 # Preprocess grad output grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) @@ -450,7 +448,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], if isinstance(weight, QuantizedTensorStorage) else cast_if_needed(weight, ctx.activation_dtype) ) - for weight in origin_weights + for weight in saved_weights ] # Make sure weights are available in column-wise format # for dgrad computation. @@ -549,7 +547,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], use_split_accumulator=wgrad_gemm_use_split_accumulator, accumulate=( accumulate_wgrad_into_param_main_grad - if not getattr(weights[0], "overwrite_main_grad", False) + if not getattr(ctx, "origin_weights_overwrite_main_grad", False) else False ), ) @@ -567,7 +565,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], # Deallocate input tensor clear_tensor_data(*inputmats) - def handle_custom_ddp_from_mcore(weight, wgrad): + def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): if ctx.weights_requires_grad: # Handle custom DDP from mcore. if ctx.fuse_wgrad_accumulation and hasattr( @@ -576,13 +574,13 @@ def handle_custom_ddp_from_mcore(weight, wgrad): weight.grad_added_to_main_grad = True if getattr(weight, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(weight.main_grad.shape), + list(main_grad.shape), weight.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(weight.main_grad.shape), + list(main_grad.shape), weight.dtype, ) elif ctx.fuse_wgrad_accumulation: @@ -592,8 +590,8 @@ def handle_custom_ddp_from_mcore(weight, wgrad): return wgrad wgrad_list = [ - handle_custom_ddp_from_mcore(weight, wgrad) - for weight, wgrad in zip(origin_weights, 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] * ctx.num_gemms diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index dc021ca6b7..5361d7deda 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -5,6 +5,7 @@ """LayerNormLinear API""" import os import warnings +import weakref from typing import Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -465,14 +466,6 @@ def forward( ln_weight, ln_bias, ) - ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad") - if ctx.grad_added_to_main_grad: - # If you are passing torch.nn.Parameter through the Torch hooks, you will - # get back torch.Tensor. Torch rips off the Parameter wrapper. - # You need to preserve the weight object to have all the attributes user - # sets for the weights. Because of this, it is not recommended to offload - # weights if weights are externally touched outside this module - ctx.weight_object = weight tensors_to_save, tensor_objects = prepare_for_saving( inputmat, @@ -490,6 +483,13 @@ def forward( ctx.requires_wgrad = weight.requires_grad ctx.is_weight_param_quantized = is_weight_param_quantized 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 + ctx.origin_weight_ref = weakref.ref(weight) + # Save overwrite_main_grad flag now while we have access to weight object + ctx.origin_weight_overwrites_main_grad = getattr( + weight, "overwrite_main_grad", False + ) # This check is needed to ensure that main_grad is not created # during the forward pass when using MCore FSDP as it creates # the main_grad buffer lazily before backprop @@ -578,7 +578,7 @@ def backward( ( # pylint: disable=unbalanced-tuple-unpacking inputmat, weight, - origin_weight, + saved_weight, bias, ln_weight, ln_out, @@ -586,12 +586,25 @@ def backward( rsigma, ) = restore_from_func_ctx(ctx) - # Since main_grad can be modified inplace, it should not be a part of saved_tensors - main_grad = ( - ctx.main_grad_func() - if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad - else None + # Restore from weakref to get original weight python object + # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) + # Only needed when fuse_wgrad_accumulation is enabled. + origin_weight = None + origin_weight_overwrites_main_grad = getattr( + ctx, "origin_weight_overwrites_main_grad", False ) + main_grad = None + if ctx.fuse_wgrad_accumulation and ctx.requires_wgrad: + origin_weight_ref = ctx.origin_weight_ref + ctx.origin_weight_ref = None + origin_weight = origin_weight_ref() if origin_weight_ref is not None else None + assert ( + origin_weight is not None + ), "weight was removed while fuse_wgrad_accumulation=True" + # Since main_grad can be modified inplace, it should not be a part of saved_tensors + main_grad = ctx.main_grad_func() if weight is not None else None + if main_grad is not None: + origin_weight.main_grad = main_grad # Gather intermediate/activation tensors if needed # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already @@ -607,14 +620,6 @@ def backward( ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - # For CPU offloading, we offloaded weight and weight.main_grad to different tensors, - # we need to connect them into one. - if ctx.cpu_offloading: - if ctx.grad_added_to_main_grad: - origin_weight = ctx.weight_object - if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: - origin_weight.main_grad = main_grad - # Configure Userbuffers communication (comm+GEMM overlap) ctx.ub_obj_gradout = None ub_obj_dgrad = None @@ -769,7 +774,7 @@ def backward( else: weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) elif ctx.backward_override == "high_precision": - weight_for_dgrad = origin_weight + weight_for_dgrad = saved_weight if isinstance(weight_for_dgrad, QuantizedTensorStorage): weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( @@ -907,7 +912,7 @@ def backward( "quantization_params": ctx.grad_weight_quantizer, "accumulate": ( accumulate_wgrad_into_param_main_grad - if not getattr(weight, "overwrite_main_grad", False) + if not origin_weight_overwrites_main_grad else False ), "layout": "NT", @@ -1039,13 +1044,13 @@ def wgrad_gemm( origin_weight.grad_added_to_main_grad = True if getattr(origin_weight, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(origin_weight.main_grad.shape), + list(main_grad.shape), origin_weight.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(origin_weight.main_grad.shape), + list(main_grad.shape), origin_weight.dtype, ) elif ctx.fuse_wgrad_accumulation: diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index a99de65c4a..ca211daa08 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -5,6 +5,7 @@ """LayerNormMLP API""" import os import warnings +import weakref from typing import Callable, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op @@ -757,13 +758,11 @@ 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, @@ -773,6 +772,20 @@ def _forward( ctx.tensor_objects = tensor_objects if fuse_wgrad_accumulation: + # Keep weakrefs to weights to preserve attributes like main_grad + # when we need to modify the weight python objects + ctx.fc1_weight_python_object_ref = ( + weakref.ref(fc1_weight) if fc1_weight.requires_grad else None + ) + ctx.fc2_weight_python_object_ref = ( + weakref.ref(fc2_weight) if fc2_weight.requires_grad else None + ) + ctx.fc1_weight_overwrites_main_grad = getattr( + fc1_weight, "overwrite_main_grad", False + ) + ctx.fc2_weight_overwrites_main_grad = getattr( + fc2_weight, "overwrite_main_grad", False + ) # This check is needed to ensure that main_grad is not created # during the forward pass when using MCore FSDP as it creates # the main_grad buffer lazily before backprop @@ -801,8 +814,6 @@ 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.device = device ctx.activation_dtype = activation_dtype @@ -854,13 +865,11 @@ 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, @@ -970,39 +979,49 @@ 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, ) = _LayerNormMLP._recompute(ctx) - # Since main_grad can be modified inplace, it should not be a part of saved_tensors - fc1_weight_main_grad = ( - ctx.fc1_main_grad_func() - if fc1_weight is not None - and ctx.fuse_wgrad_accumulation - and ctx.fc1_weight_requires_grad - else None - ) - fc2_weight_main_grad = ( - ctx.fc2_main_grad_func() - if origin_fc2_weight is not None - and ctx.fuse_wgrad_accumulation - and ctx.fc2_weight_requires_grad - else None - ) - - # For CPU offloading, we offloaded weight and weight.main_grad to different tensors, - # we need to connect them into one. + # Restore origin weights from weakrefs + # Only needed when fuse_wgrad_accumulation is enabled. + fc1_weight_python_object = None + fc2_weight_python_object = None + fc1_weight_main_grad = None + fc2_weight_main_grad = None if ctx.fuse_wgrad_accumulation: - origin_fc1_weight.main_grad = fc1_weight_main_grad - origin_fc2_weight.main_grad = fc2_weight_main_grad + fc1_weight_python_object_ref = getattr(ctx, "fc1_weight_python_object_ref", None) + fc2_weight_python_object_ref = getattr(ctx, "fc2_weight_python_object_ref", None) + ctx.fc1_weight_python_object_ref = None + ctx.fc2_weight_python_object_ref = None + fc1_weight_python_object = ( + fc1_weight_python_object_ref() + if fc1_weight_python_object_ref is not None + else None + ) + fc2_weight_python_object = ( + fc2_weight_python_object_ref() + if fc2_weight_python_object_ref is not None + else None + ) + if ctx.fc1_weight_requires_grad: + assert ( + fc1_weight_python_object is not None + ), "fc1_weight was removed while fuse_wgrad_accumulation=True" + fc1_weight_main_grad = ctx.fc1_main_grad_func() + fc1_weight_python_object.main_grad = fc1_weight_main_grad + if ctx.fc2_weight_requires_grad: + assert ( + fc2_weight_python_object is not None + ), "fc2_weight was removed while fuse_wgrad_accumulation=True" + fc2_weight_main_grad = ctx.fc2_main_grad_func() + fc2_weight_python_object.main_grad = fc2_weight_main_grad # TODO: Fix this # pylint: disable=fixme # Gather saved autograd context tensors when running with FSDP @@ -1121,9 +1140,9 @@ def backward( if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) if ctx.fc2_weight_quantizer is not None and isinstance( - ctx.fc2_weight, QuantizedTensorStorage + fc2_weight, QuantizedTensorStorage ): - ctx.fc2_weight.update_usage(columnwise_usage=True) + fc2_weight.update_usage(columnwise_usage=True) # Perform GEMM gemm_output, *_ = general_gemm( @@ -1223,18 +1242,18 @@ def backward( # Arguments to include in wgrad GEMM closure fc2_wgrad_gemm_kwargs = { "out_dtype": ( - origin_fc2_weight.main_grad.dtype + fc2_weight_main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype ), "quantization_params": ctx.fc2_grad_weight_quantizer, # wgrad in high precision "accumulate": ( accumulate_wgrad_into_param_main_grad - if not getattr(fc1_weight, "overwrite_main_grad", False) + if not getattr(ctx, "fc2_weight_overwrites_main_grad", False) else False ), "layout": "NT", - "out": origin_fc2_weight.main_grad if ctx.fuse_wgrad_accumulation else None, + "out": fc2_weight_main_grad if ctx.fuse_wgrad_accumulation else None, "bias": fc2_bias if fc2_bias is not None and fc2_bias_grad is None else None, "use_split_accumulator": wgrad_use_split_accumulator, "grad": grad_arg, @@ -1373,9 +1392,9 @@ def fc2_wgrad_gemm( # Make sure required data is available if ctx.fc1_weight_quantizer is not None and isinstance( - ctx.fc1_weight_quantizer, QuantizedTensorStorage + fc1_weight, QuantizedTensorStorage ): - ctx.fc1_weight.update_usage(columnwise_usage=True) + fc1_weight.update_usage(columnwise_usage=True) # Output buffers for Userbuffers reduce-scatter gemm_out = None @@ -1470,18 +1489,18 @@ def fc2_wgrad_gemm( # Arguments to include in wgrad GEMM closure fc1_wgrad_gemm_kwargs = { "out_dtype": ( - origin_fc1_weight.main_grad.dtype + fc1_weight_main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype ), "quantization_params": ctx.fc1_grad_weight_quantizer, "accumulate": ( accumulate_wgrad_into_param_main_grad - if not getattr(fc2_weight, "overwrite_main_grad", False) + if not getattr(ctx, "fc1_weight_overwrites_main_grad", False) else False ), "layout": "NT", - "out": origin_fc1_weight.main_grad if ctx.fuse_wgrad_accumulation else None, + "out": fc1_weight_main_grad if ctx.fuse_wgrad_accumulation else None, "bias": fc1_bias if fuse_gemm_and_bias_fc1_wgrad else None, "use_split_accumulator": wgrad_use_split_accumulator, "grad": fuse_gemm_and_bias_fc1_wgrad, @@ -1585,19 +1604,21 @@ def fc1_wgrad_gemm( if ctx.fc1_weight_requires_grad: # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr(fc1_weight, "grad_added_to_main_grad"): - origin_fc1_weight.grad_added_to_main_grad = True - if getattr(origin_fc1_weight, "zero_out_wgrad", False): + if ctx.fuse_wgrad_accumulation and hasattr( + fc1_weight_python_object, "grad_added_to_main_grad" + ): + fc1_weight_python_object.grad_added_to_main_grad = True + if getattr(fc1_weight_python_object, "zero_out_wgrad", False): fc1_wgrad = torch.zeros( - origin_fc1_weight.main_grad.shape, - dtype=origin_fc1_weight.dtype, + fc1_weight_main_grad.shape, + dtype=fc1_weight_python_object.dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: fc1_wgrad = torch.empty( - origin_fc1_weight.main_grad.shape, - dtype=origin_fc1_weight.dtype, + fc1_weight_main_grad.shape, + dtype=fc1_weight_python_object.dtype, device=torch.cuda.current_device(), requires_grad=False, ) @@ -1609,20 +1630,20 @@ def fc1_wgrad_gemm( if ctx.fc2_weight_requires_grad: # Handle custom DDP from mcore. if ctx.fuse_wgrad_accumulation and hasattr( - origin_fc2_weight, "grad_added_to_main_grad" + fc2_weight_python_object, "grad_added_to_main_grad" ): - origin_fc2_weight.grad_added_to_main_grad = True - if getattr(origin_fc2_weight, "zero_out_wgrad", False): + fc2_weight_python_object.grad_added_to_main_grad = True + if getattr(fc2_weight_python_object, "zero_out_wgrad", False): fc2_wgrad = torch.zeros( - origin_fc2_weight.main_grad.shape, - dtype=origin_fc2_weight.dtype, + fc2_weight_main_grad.shape, + dtype=fc2_weight_python_object.dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: fc2_wgrad = torch.empty( - origin_fc2_weight.main_grad.shape, - dtype=origin_fc2_weight.dtype, + fc2_weight_main_grad.shape, + dtype=fc2_weight_python_object.dtype, device=torch.cuda.current_device(), requires_grad=False, ) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 8510f6cf8f..c85db15114 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -7,6 +7,7 @@ from functools import reduce from operator import mul as multiply_op import warnings +import weakref import torch @@ -437,16 +438,6 @@ def forward( nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") if cpu_offloading: - ctx.grad_added_to_main_grad = hasattr(weight, "grad_added_to_main_grad") - - if ctx.grad_added_to_main_grad: - # If you are passing torch.nn.Parameter through the Torch hooks, you will - # get back torch.Tensor. Torch rips off the Parameter wrapper. - # You need to preserve the weight object to have all the attributes user - # sets for the weights. Because of this, it is not recommended to offload - # weights if weights are externally touched outside this module - ctx.weight_object = weight - mark_not_offload(weight, weightmat, bias) # TODO(ksivamani): Check memory usage @@ -467,8 +458,15 @@ def forward( ctx.grad_input_quantizer = grad_input_quantizer ctx.grad_weight_quantizer = grad_weight_quantizer ctx.grad_output_quantizer = grad_output_quantizer + ctx.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation if fuse_wgrad_accumulation and weight.requires_grad: + # Keep a weakref to the original Python object because save_for_backward + # may return a plain Tensor without custom Parameter attributes. + ctx.origin_weight_ref = weakref.ref(weight) + ctx.origin_weight_overwrites_main_grad = getattr( + weight, "overwrite_main_grad", False + ) # This check is needed to ensure that main_grad is not created # during the forward pass when using MCore FSDP as it creates # the main_grad buffer lazily before backprop @@ -535,22 +533,34 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], nvtx_label = f"{nvtx_label}.{ctx.ub_name}" with get_nvtx_range_context("_Linear_backward"): - inputmat, weight_fp8, weight, bias = ( # pylint: disable=unbalanced-tuple-unpacking - restore_from_func_ctx(ctx) + ( + inputmat, + weight_fp8, + saved_weight, + bias, + ) = restore_from_func_ctx( # pylint: disable=unbalanced-tuple-unpacking + ctx ) - # Since main_grad can be modified inplace, it should not be a part of saved_tensors - main_grad = ( - ctx.main_grad_func() - if weight is not None and ctx.fuse_wgrad_accumulation and ctx.requires_wgrad - else None + # Restore from weakref to get original weight python object + # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) + origin_weight_python_object = None + origin_weight_overwrites_main_grad = getattr( + ctx, "origin_weight_overwrites_main_grad", False ) - - if ctx.cpu_offloading: - if ctx.grad_added_to_main_grad: - weight = ctx.weight_object - if ctx.requires_wgrad and ctx.fuse_wgrad_accumulation: - weight.main_grad = 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 + 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" + # Since main_grad can be modified inplace, it should not be a part of saved_tensors + main_grad = ctx.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 @@ -757,7 +767,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], else: weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) elif ctx.backward_override == "high_precision": - weight_for_dgrad = weight + weight_for_dgrad = saved_weight if isinstance(weight_for_dgrad, QuantizedTensorStorage): weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( @@ -894,7 +904,7 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], "quantization_params": ctx.grad_weight_quantizer, "accumulate": ( accumulate_wgrad_into_param_main_grad - if not getattr(weight, "overwrite_main_grad", False) + if not origin_weight_overwrites_main_grad else False ), "layout": "NT", @@ -984,22 +994,20 @@ def wgrad_gemm( if ctx.requires_wgrad: # Handle custom DDP from mcore. - if ( - ctx.fuse_wgrad_accumulation - and weight is not None - and hasattr(weight, "grad_added_to_main_grad") + if ctx.fuse_wgrad_accumulation and hasattr( + origin_weight_python_object, "grad_added_to_main_grad" ): - weight.grad_added_to_main_grad = True - if getattr(weight, "zero_out_wgrad", False): + origin_weight_python_object.grad_added_to_main_grad = True + if getattr(origin_weight_python_object, "zero_out_wgrad", False): wgrad = get_dummy_wgrad( - list(weight.main_grad.shape), - weight.dtype, + list(main_grad.shape), + origin_weight_python_object.dtype, zero=True, ) else: wgrad = get_dummy_wgrad( - list(weight.main_grad.shape), - weight.dtype, + list(main_grad.shape), + origin_weight_python_object.dtype, ) elif ctx.fuse_wgrad_accumulation: wgrad = None @@ -1013,7 +1021,7 @@ def wgrad_gemm( nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") # Scatter fp8 weight buffers - if ctx.fp8 and not isinstance(weight, QuantizedTensorStorage): + if ctx.fp8 and not ctx.is_weight_param_quantized: _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) return ( wgrad, From 9f5fde1312c87c3502c68872e0fc60df551e5b77 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 13 Apr 2026 13:30:27 -0400 Subject: [PATCH 336/521] [PyTorch] Relax dimension constraints for using fused grouped MLP (#2856) * Reduce fused path dim constraint Signed-off-by: Kirthi Shankar Sivamani * Fix randomization in tests Signed-off-by: Kirthi Shankar Sivamani * reset rng as before, assert input dim Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fusible_ops.py | 30 ++++++------------- tests/pytorch/utils.py | 15 +++++++--- transformer_engine/pytorch/ops/_common.py | 12 ++++---- .../pytorch/ops/fused/backward_grouped_mlp.py | 12 ++++---- .../pytorch/ops/fused/forward_grouped_mlp.py | 13 ++++---- 5 files changed, 39 insertions(+), 43 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index a2de8014a1..a5c071074c 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -77,6 +77,13 @@ _quantization_list.append("nvfp4") +@pytest.fixture(autouse=True, scope="class") +def _reset_rng_states_per_test(): + """Restore torch, CUDA, and Python ``random`` before each test in this module.""" + reset_rng_states() + yield + + def maybe_skip_quantization( quantization: Optional[str], *, @@ -364,10 +371,6 @@ def test_extra_tensors(self, size: int = 16) -> None: class TestFuser: """Tests for operation fusion infrastructure""" - @staticmethod - def setup_class(cls) -> None: - reset_rng_states() - @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) def test_fp8_scale_update( self, @@ -580,10 +583,6 @@ def test_pyt_autocast( class TestBasicOps: """Tests for individual operations""" - @staticmethod - def setup_class(cls) -> None: - reset_rng_states() - @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("device", ("cuda", "cpu")) @pytest.mark.parametrize("quantization", _quantization_list) @@ -2327,10 +2326,6 @@ def test_interleaved_scaled_clamped_qgeglu(self): class TestFusedOps: """Tests for fused operations""" - @staticmethod - def setup_class(cls) -> None: - reset_rng_states() - @pytest.mark.parametrize("weight_shape", ((32, 64), (3, 5))) @pytest.mark.parametrize("in_shape", ((-1,), (1, 7, -1), (8, 2, 10, -1))) @pytest.mark.parametrize("dtype", _dtypes) @@ -3035,10 +3030,6 @@ def test_backward_linear_scale( class TestCheckpointing: """Tests for checkpointing""" - @staticmethod - def setup_class(cls) -> None: - reset_rng_states() - @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("quantized_weight", (False, True)) def test_linear( @@ -3151,10 +3142,6 @@ def test_linear( class TestSequentialModules: """Test for larger Sequentials with modules commonly used together""" - @staticmethod - def setup_class(cls) -> None: - reset_rng_states() - @pytest.mark.parametrize("requires_grad", (False, True)) @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("quantized_compute", (False, True)) @@ -3338,13 +3325,14 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) @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")) def test_grouped_mlp( self, *, group_size: int = 4, bias: bool, - hidden_size: int = 256, + hidden_size: int, dtype: torch.dtype, quantization: Optional[str], single_grouped_weight: bool, diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 196ae8c165..fd9a6416ec 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -6,6 +6,7 @@ import logging import os +import random import subprocess from contextlib import contextmanager from typing import Optional, Sequence, Tuple, Dict, Any, List @@ -173,8 +174,8 @@ def skip_unsupported_backward_override( pytest.skip(f"{layer_type} does not support NVTE_BACKWARD_OVERRIDE={backward_override}.") -# Cached RNG state -_rng_states: Optional[Tuple[torch.Tensor, torch.Tensor]] = None +# Cached RNG state (torch CPU, torch CUDA, Python ``random``) +_rng_states: Optional[Tuple[torch.Tensor, torch.Tensor, Any]] = None def reset_rng_states() -> None: @@ -183,11 +184,17 @@ def reset_rng_states() -> None: if _rng_states is None: torch.manual_seed(1234) torch.cuda.manual_seed(1234) - _rng_states = (torch.get_rng_state(), torch.cuda.get_rng_state()) + random.seed(1234) + _rng_states = ( + torch.get_rng_state(), + torch.cuda.get_rng_state(), + random.getstate(), + ) else: - cpu_rng_state, cuda_rng_state = _rng_states + cpu_rng_state, cuda_rng_state, random_state = _rng_states torch.set_rng_state(cpu_rng_state) torch.cuda.set_rng_state(cuda_rng_state) + random.setstate(random_state) def compare_and_assert(a, b, name_a, name_b, atol, rtol, rmse_tol, is_fp8): diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index ae8b48a90d..15dc17e812 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -88,12 +88,12 @@ def get_fp8_meta_from_fp8_tensor(tensor: Float8Tensor) -> tuple[FP8TensorMeta, i def validate_grouped_mlp_dims(fc1, glu_op, fc2) -> None: """Validate FC1 / scaled GLU / FC2 dimensions for fused grouped MLP.""" - if fc1.in_features % 256 != 0 or fc1.out_features % 256 != 0: + if fc1.in_features % 64 != 0 or fc1.out_features % 64 != 0: raise ValueError( f"Unsupported dims for FC1 (num_groups={fc1.num_groups}, " f"in_features={fc1.in_features}, out_features={fc1.out_features})." ) - if fc2.in_features % 256 != 0 or fc2.out_features % 256 != 0: + if fc2.in_features % 64 != 0 or fc2.out_features % 64 != 0: raise ValueError( f"Unsupported dims for FC2 (num_groups={fc2.num_groups}, " f"in_features={fc2.in_features}, out_features={fc2.out_features})." @@ -176,10 +176,10 @@ def fuse_grouped_mlp_ops( elif window[0].num_groups != window[2].num_groups: matches_pattern = False elif ( - window[0].in_features % 256 != 0 - or window[0].out_features % 256 != 0 - or window[2].in_features % 256 != 0 - or window[2].out_features % 256 != 0 + 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: diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 357e8b3695..feed4767e9 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -399,8 +399,8 @@ def fuser_backward( fc2_dy_scales = fc2_dy_scales.view(dtype=torch.float8_e8m0fnu) fc2_dy_scales = fc2_dy_scales.view( 1, - out_shape[0] // 128, - out_shape[1] // 128, + (out_shape[0] + 127) // 128, + (out_shape[1] + 127) // 128, MXFP8_BLOCK_SCALING_SIZE, 4, 4, @@ -449,8 +449,8 @@ def fuser_backward( fc2_w_scales = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=torch.float8_e8m0fnu) fc2_w_scales = fc2_w_scales.view( num_groups, - fc2_weight_shape[1] // 128, - fc2_weight_shape[0] // 128, + (fc2_weight_shape[1] + 127) // 128, + (fc2_weight_shape[0] + 127) // 128, MXFP8_BLOCK_SCALING_SIZE, 4, 4, @@ -607,8 +607,8 @@ def fuser_backward( ) fc1_w_scales = fc1_w_scales.view( num_groups, - fc1_weight_shape[1] // 128, - fc1_weight_shape[0] // 128, + (fc1_weight_shape[1] + 127) // 128, + (fc1_weight_shape[0] + 127) // 128, MXFP8_BLOCK_SCALING_SIZE, 4, 4, diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 83bb4428f0..4e756ea531 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -154,6 +154,7 @@ def fuser_forward( fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features) input_ = input_.reshape(-1, fc1_weight_shape[1]) in_shape = list(input_.size()) + assert in_shape[0] % 128 == 0, "Unsupported input shape for fused grouped MLP." num_groups = fc1_op.num_groups fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 @@ -312,8 +313,8 @@ def fuser_forward( fc1_x_scales = fc1_x_scales.view(dtype=torch.float8_e8m0fnu) fc1_x_scales = fc1_x_scales.view( 1, - in_shape[0] // 128, - in_shape[1] // 128, + (in_shape[0] + 127) // 128, + (in_shape[1] + 127) // 128, MXFP8_BLOCK_SCALING_SIZE, 4, 4, @@ -361,8 +362,8 @@ def fuser_forward( fc1_w_scales = fc1_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) fc1_w_scales = fc1_w_scales.view( num_groups, - fc1_weight_shape[0] // 128, - fc1_weight_shape[1] // 128, + (fc1_weight_shape[0] + 127) // 128, + (fc1_weight_shape[1] + 127) // 128, MXFP8_BLOCK_SCALING_SIZE, 4, 4, @@ -464,8 +465,8 @@ def fuser_forward( fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) fc2_w_scales = fc2_w_scales.view( num_groups, - fc2_weight_shape[0] // 128, - fc2_weight_shape[1] // 128, + (fc2_weight_shape[0] + 127) // 128, + (fc2_weight_shape[1] + 127) // 128, MXFP8_BLOCK_SCALING_SIZE, 4, 4, From 491c59774b51ecf913b24c1e05c19dc2be4a20f6 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 13 Apr 2026 13:30:52 -0400 Subject: [PATCH 337/521] [PyTorch] Cache alpha and beta for cublas ggemm (#2870) Cache alpha and beta for cublas ggemm Signed-off-by: Kirthi Shankar Sivamani --- .../pytorch/cpp_extensions/gemm.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 115569ccba..82891ca83f 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -306,6 +306,18 @@ def get_grouped_gemm_setup_workspace_size(num_tensors: int) -> int: return ((size + alignment - 1) // alignment) * alignment +@functools.lru_cache(maxsize=None) +def _get_fp32_ones_tensor(num_tensors: int, device: torch.device) -> torch.Tensor: + """Cached ones tensor.""" + return torch.ones(num_tensors, dtype=torch.float32, device=device) + + +@functools.lru_cache(maxsize=None) +def _get_fp32_zeros_tensor(num_tensors: int, device: torch.device) -> torch.Tensor: + """Cached zeros tensor.""" + return torch.zeros(num_tensors, dtype=torch.float32, device=device) + + def general_grouped_gemm_for_grouped_tensor( A, B, @@ -358,12 +370,12 @@ def general_grouped_gemm_for_grouped_tensor( device = rowwise.device if rowwise is not None else B.columnwise_data.device if alpha is None: - alpha = torch.ones(num_tensors, dtype=torch.float32, device=device) + alpha = _get_fp32_ones_tensor(num_tensors, device) if beta is None: if accumulate: - beta = torch.ones(num_tensors, dtype=torch.float32, device=device) + beta = _get_fp32_ones_tensor(num_tensors, device) else: - beta = torch.zeros(num_tensors, dtype=torch.float32, device=device) + beta = _get_fp32_zeros_tensor(num_tensors, device) if not alpha.is_cuda or not beta.is_cuda: raise ValueError("alpha and beta must be CUDA tensors.") From d7c43bbb5076e851f45aaa345109c280a63f86aa Mon Sep 17 00:00:00 2001 From: Almog Segal Date: Mon, 13 Apr 2026 21:00:00 +0300 Subject: [PATCH 338/521] comm_gemm_test fixes (#2839) * Fix comm_gemm test bias buffer size and distribution - Allocate bias as 1D vector of length m (not m*n matrix) - Distribute bias as a row-slice matching local D rows - Change tolerance from tol*k to tol since tol values now represent the actual absolute tolerance Signed-off-by: Almog Segal * Adjust comm_gemm test for accurate comparison - Use split accumulator (disable fast FP8 accumulation) in the reference GEMM to match cuBLASMp's accumulation precision - Set per-test tolerances based on observed max errors: AG: 1e-3, RS FP16: 7e-2, RS BF16: 6e-1, RS FP8: 7e-2 to 1e-1, AR FP16: 7e-2, AR BF16: 1e-3, AR FP8: 1.5e-1 Signed-off-by: Almog Segal --------- Signed-off-by: Almog Segal --- tests/cpp_distributed/test_comm_gemm.cu | 40 ++++++++++++------------- 1 file changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index cdd6f9cf14..cc0d760a39 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -204,7 +204,7 @@ class CommGemmFixure : public ::testing::TestWithParam { std::vector bdata(k * n); std::generate(bdata.begin(), bdata.end(), [&rng, &dist, b_scale] { return static_cast(dist(rng) * b_scale); }); - std::vector biasdata(m * n); + std::vector biasdata(m); std::generate(biasdata.begin(), biasdata.end(), [&rng, &dist, bias_scale] { return static_cast(dist(rng) * bias_scale); }); @@ -213,7 +213,7 @@ class CommGemmFixure : public ::testing::TestWithParam { : MakeFromData(adata, 0, 0, m, k, m, a_scale); auto gb = transb ? MakeFromData(bdata, 0, 0, n, k, n, b_scale) : MakeFromData(bdata, 0, 0, k, n, k, b_scale); - auto gbias = MakeFromData(biasdata, 0, 0, m, n, m, bias_scale); + auto gbias = MakeFromData(biasdata, 0, 0, m, 1, m, bias_scale); auto gd = Make(m, n, d_scale); auto gaux = Make(m, n, d_scale); @@ -226,8 +226,8 @@ class CommGemmFixure : public ::testing::TestWithParam { dims.b_cols_num, dims.b_rows_num, n, b_scale) : MakeFromData(bdata, dims.b_rows_start, dims.b_cols_start, dims.b_rows_num, dims.b_cols_num, k, b_scale); - auto bias = MakeFromData(biasdata, dims.d_rows_start, dims.d_cols_start, - dims.d_rows_num, dims.d_cols_num, m, bias_scale); + auto bias = MakeFromData(biasdata, dims.d_rows_start, 0, dims.d_rows_num, 1, m, + bias_scale); auto d = Make(dims.d_rows_num, dims.d_cols_num, d_scale); auto aux = Make(dims.d_rows_num, dims.d_cols_num, d_scale); @@ -237,7 +237,7 @@ class CommGemmFixure : public ::testing::TestWithParam { accumulate, 0 /*comm_sm_count*/, stream); auto workspace = Make(1, 32 << 20, 1.0); nvte_cublas_gemm(ga.data(), gb.data(), gd.data(), gbias.data(), gaux.data(), transa, transb, - grad, workspace.data(), accumulate, false /* use_split_accumulator */, + grad, workspace.data(), accumulate, true /* use_split_accumulator */, 0 /* math_sm_count */, stream); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); @@ -253,7 +253,7 @@ class CommGemmFixure : public ::testing::TestWithParam { dims.d_rows_num, dims.d_cols_num, m); NVTE_CHECK(out.size() == out_golden.size()); for (size_t i = 0; i < out.size(); ++i) { - EXPECT_NEAR(static_cast(out[i]), static_cast(out_golden[i]), tol * k); + EXPECT_NEAR(static_cast(out[i]), static_cast(out_golden[i]), tol); } } @@ -427,35 +427,35 @@ INSTANTIATE_TEST_SUITE_P(AgGemm, AgGemm, INSTANTIATE_TEST_SUITE_P(GemmRs, GemmRs, testing::Values(Params{DType::kFloat16, DType::kFloat16, DType::kFloat16, - false, false, 64, 128, 256, 5e-2}, + false, false, 64, 128, 256, 7e-2}, Params{DType::kFloat16, DType::kFloat16, DType::kFloat16, - false, true, 64, 128, 256, 5e-2}, + false, true, 64, 128, 256, 7e-2}, Params{DType::kFloat16, DType::kFloat16, DType::kFloat16, - true, false, 64, 128, 256, 5e-2}, + true, false, 64, 128, 256, 7e-2}, Params{DType::kBFloat16, DType::kBFloat16, - DType::kBFloat16, false, false, 64, 128, 256, 5e-2}, + DType::kBFloat16, false, false, 64, 128, 256, 6e-1}, Params{DType::kBFloat16, DType::kBFloat16, - DType::kBFloat16, false, true, 64, 128, 256, 5e-2}, + DType::kBFloat16, false, true, 64, 128, 256, 6e-1}, Params{DType::kBFloat16, DType::kBFloat16, - DType::kBFloat16, true, false, 64, 128, 256, 5e-2}, + DType::kBFloat16, true, false, 64, 128, 256, 6e-1}, Params{DType::kFloat8E4M3, DType::kFloat8E4M3, - DType::kFloat16, true, false, 64, 128, 256, 5e-2}, + DType::kFloat16, true, false, 64, 128, 256, 1e-1}, Params{DType::kFloat8E4M3, DType::kFloat8E5M2, - DType::kFloat16, true, false, 64, 128, 256, 5e-2}, + DType::kFloat16, true, false, 64, 128, 256, 7e-2}, Params{DType::kFloat8E5M2, DType::kFloat8E4M3, - DType::kFloat16, true, false, 64, 128, 256, 5e-2}), + DType::kFloat16, true, false, 64, 128, 256, 7e-2}), &ParamSuffix); INSTANTIATE_TEST_SUITE_P( GemmAr, GemmAr, testing::Values(Params{DType::kFloat16, DType::kFloat16, DType::kFloat16, true, false, 64, - 64 * 4, 64 * 4, 5e-2}, + 64 * 4, 64 * 4, 7e-2}, Params{DType::kBFloat16, DType::kBFloat16, DType::kBFloat16, true, false, 64, - 64 * 4, 64 * 4, 5e-2}, + 64 * 4, 64 * 4, 1e-3}, Params{DType::kFloat8E5M2, DType::kFloat8E4M3, DType::kFloat16, true, false, - 128, 128 * 4, 128 * 4, 5e-2}, + 128, 128 * 4, 128 * 4, 1.5e-1}, Params{DType::kFloat8E4M3, DType::kFloat8E5M2, DType::kFloat16, true, false, - 128, 128 * 4, 128 * 4, 5e-2}, + 128, 128 * 4, 128 * 4, 1.5e-1}, Params{DType::kFloat8E4M3, DType::kFloat8E4M3, DType::kFloat16, true, false, - 128, 128 * 4, 128 * 4, 5e-2}), + 128, 128 * 4, 128 * 4, 1.5e-1}), &ParamSuffix); From dc92b3968d0356ddc709e217f7644ecb7e7f752e Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Mon, 13 Apr 2026 16:13:52 -0500 Subject: [PATCH 339/521] docs(readme): update convergence table, latest news, and outdated links (#2638) * docs(readme): update FP8 convergence table and add MXFP8/NVFP4 support info - Add MXFP8 and NVFP4 format support to highlights and description - Update FP8 convergence table with MXFP8 results from arxiv paper - Remove outdated JAX-Toolbox links and "available on request" entries - Update Docker container versions to 26.01 - Fix DeepSpeed and Lightning integration links - Add Nemotron 3 paper to Latest News - Add quickstart notebook link after PyTorch example Signed-off-by: Santosh Bhavani * fix(readme): address review feedback - Replace quickstart.ipynb link with fp8_primer.ipynb (file exists) - Fix extra whitespace in Megatron Core table rows Signed-off-by: Santosh Bhavani * Revert FP8 Primer link changes, defer to PR #2641 Signed-off-by: Santosh Bhavani * ci: remove maximize-build-space from pytorch and all jobs Signed-off-by: Santosh Bhavani * Revert "ci: remove maximize-build-space from pytorch and all jobs" This reverts commit 643b3d9a73069346f3e302e2483288b77a3956a8. Signed-off-by: Santosh Bhavani * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Santosh Bhavani * fix(readme): update convergence section, links, and integration refs Signed-off-by: Santosh Bhavani --------- Signed-off-by: Santosh Bhavani Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- README.rst | 41 ++++++++++++++++------------------------- 1 file changed, 16 insertions(+), 25 deletions(-) diff --git a/README.rst b/README.rst index 5a6721b04c..e537b7a1fe 100644 --- a/README.rst +++ b/README.rst @@ -8,11 +8,12 @@ Transformer Engine ================== -`Quickstart <#examples>`_ | `Installation <#installation>`_ | `User Guide `_ | `Examples `_ | `FP8 Convergence <#fp8-convergence>`_ | `Integrations <#integrations>`_ | `Release notes `_ +`Quickstart <#examples>`_ | `Installation <#installation>`_ | `User Guide `_ | `Examples `_ | `Convergence <#convergence>`_ | `Integrations <#integrations>`_ | `Release notes `_ Latest News =========== +* [12/2025] `NVIDIA Nemotron 3: Efficient and Open Intelligence `_ - trained with NVFP4 on Transformer Engine * [11/2025] `NVIDIA Blackwell Architecture Sweeps MLPerf Training v5.1 Benchmarks `_ * [11/2025] `Scale Biology Transformer Models with PyTorch and NVIDIA BioNeMo Recipes `_ * [11/2025] `FP8 Training of Large-Scale RL Models `_ @@ -30,7 +31,8 @@ What is Transformer Engine? Transformer Engine (TE) is a library for accelerating Transformer models on NVIDIA GPUs, including using 8-bit floating point (FP8) precision on Hopper, Ada, and Blackwell GPUs, to provide better -performance with lower memory utilization in both training and inference. TE provides a collection +performance with lower memory utilization in both training and inference. On Blackwell GPUs, TE also +supports MXFP8 (Microscaling FP8) and NVFP4 formats for even greater efficiency. TE provides a collection of highly optimized building blocks for popular Transformer architectures and an automatic mixed precision-like API that can be used seamlessly with your framework-specific code. TE also includes a framework agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 @@ -58,6 +60,7 @@ Highlights * Easy-to-use modules for building Transformer layers with FP8 support * Optimizations (e.g. fused kernels) for Transformer models * Support for FP8 on NVIDIA Hopper, Ada, and Blackwell GPUs +* Support for MXFP8 and NVFP4 on NVIDIA Blackwell GPUs * Support for optimizations across all precisions (FP16, BF16) on NVIDIA Ampere GPU architecture generations and later Examples @@ -190,12 +193,11 @@ We recommend updating to the latest NGC container available here: * https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch * https://catalog.ngc.nvidia.com/orgs/nvidia/containers/jax -If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. If you would like to use examples from TE main branch and are running into import errors, please try the latest pip package or building from source, although NGC containers are recommended for ease-of-use for most users. +If you run any examples, please ensure you are using a matching version of TransformerEngine. TransformerEngine is pre-built and packaged inside the containers with examples available at ``/opt/transformerengine`` or ``/opt/transformer-engine``. **Benefits of using NGC containers:** * All dependencies pre-installed with compatible versions and optimized configurations -* NGC PyTorch 23.08+ containers include FlashAttention-2 pip Installation ^^^^^^^^^^^^^^^^ @@ -373,54 +375,43 @@ An example of this change is, False, False, True, True, True, False, False, False, False, True] -FP8 Convergence -=============== +Convergence +=========== -FP8 has been tested extensively across different model architectures and configurations and we found **no significant difference** between FP8 and BF16 training loss curves. FP8 has also been validated for accuracy on downstream LLM tasks (e.g. LAMBADA and WikiText). Below are examples of models tested for convergence across different frameworks. +FP8 and MXFP8 have been tested extensively across different model architectures and configurations and we found **no significant difference** between FP8/MXFP8 and BF16 training loss curves. FP8 and MXFP8 have also been validated for accuracy on downstream LLM tasks (e.g. LAMBADA and WikiText). Below are examples of models tested for convergence across different frameworks. +------------+------------------+---------------------------------------------------------------------------------------------------------+ | Model | Framework | Source | +============+==================+=========================================================================================================+ -| T5-770M | JAX/T5x | https://github.com/NVIDIA/JAX-Toolbox/tree/main/rosetta/rosetta/projects/t5x#convergence-and-performance| -+------------+------------------+---------------------------------------------------------------------------------------------------------+ | MPT-1.3B | Mosaic Composer | https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-5B | JAX/Paxml | https://github.com/NVIDIA/JAX-Toolbox/tree/main/rosetta/rosetta/projects/pax#h100-results | -+------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-5B | NeMo Framework | Available on request | -+------------+------------------+---------------------------------------------------------------------------------------------------------+ | LLama2-7B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| T5-11B | JAX/T5x | Available on request | +| LLM-8B | Megatron Core | https://arxiv.org/abs/2506.08027 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ | MPT-13B | Mosaic Composer | https://www.databricks.com/blog/turbocharged-training-optimizing-databricks-mosaic-ai-stack-fp8 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-22B | NeMo Framework | Available on request | +| MoE-16B | Megatron Core | https://arxiv.org/abs/2506.08027 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ | LLama2-70B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ | +------------+------------------+---------------------------------------------------------------------------------------------------------+ -| GPT-175B | JAX/Paxml | https://github.com/NVIDIA/JAX-Toolbox/tree/main/rosetta/rosetta/projects/pax#h100-results | -+------------+------------------+---------------------------------------------------------------------------------------------------------+ Integrations ============ Transformer Engine has been integrated with popular LLM frameworks such as: -* `DeepSpeed `_ +* `DeepSpeed `_ * `Hugging Face Accelerate `_ -* `Lightning `_ +* `Lightning `_ * `MosaicML Composer `_ * `NVIDIA JAX Toolbox `_ * `NVIDIA Megatron-LM `_ -* `NVIDIA NeMo Framework `_ +* `NVIDIA NeMo Megatron Bridge `_ * `Amazon SageMaker Model Parallel Library `_ * `Levanter `_ * `GPT-NeoX `_ -* `Hugging Face Nanotron `_ - Coming soon! -* `Colossal-AI `_ - Coming soon! -* `PeriFlow `_ - Coming soon! - +* `Hugging Face Nanotron `_ Contributing ============ @@ -439,7 +430,7 @@ Papers Videos ====== -* `Stable and Scalable FP8 Deep Learning Training on Blackwell | GTC 2025 `__ +* `Stable and Scalable FP8 Deep Learning Training on Blackwell | GTC 2025 `_ * `Blackwell Numerics for AI | GTC 2025 `_ * `Building LLMs: Accelerating Pretraining of Foundational Models With FP8 Precision | GTC 2025 `_ * `From FP8 LLM Training to Inference: Language AI at Scale | GTC 2025 `_ From 72328b34d4140febf7002809c15e20aab5a83c7a Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 13 Apr 2026 15:15:33 -0700 Subject: [PATCH 340/521] Cute Dsl kernel for Wgrad for Fused MOE Layer (#2869) * integrate cudnn wgrad kernel Signed-off-by: Varun Thumbe * have only cute dsl for wgrad Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert the change for cudnn Signed-off-by: Varun Thumbe * remove dtype Signed-off-by: Varun Thumbe * fix comment: Signed-off-by: Varun Thumbe * go to cublas if needed 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: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/ops/_common.py | 9 ++ .../pytorch/ops/fused/backward_grouped_mlp.py | 126 ++++++++++++++++-- 2 files changed, 125 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 15dc17e812..e21915a5a6 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -29,6 +29,15 @@ def _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu() -> bool: return False +@functools.lru_cache(maxsize=1) +def _nvidia_cudnn_frontend_supports_wgrad() -> bool: + """Check cuDNN FE min version for grouped GEMM wgrad kernel.""" + try: + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.23.0") + except PackageNotFoundError: + return False + + def is_quantized_tensor(tensor: torch.Tensor | QuantizedTensorStorage) -> bool: """Check if tensor is a quantized tensor""" return isinstance(tensor, QuantizedTensorStorage) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index feed4767e9..389dfbc838 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -15,9 +15,6 @@ import torch import transformer_engine_torch as tex -from ...cpp_extensions import ( - general_grouped_gemm_for_grouped_tensor, -) from ...module.base import get_dummy_wgrad from ...quantization import Recipe from ...tensor.grouped_tensor import GroupedTensor @@ -28,13 +25,88 @@ from ..fuser import register_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( + _nvidia_cudnn_frontend_supports_wgrad, 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 ...triton.grouped_dbias_dscales import _compute_grouped_dbias_dscales +def _cudnn_compute_wgrad( + grouped_x: GroupedTensor, + grouped_dy: GroupedTensor, + wgrad_output, + weight_shape: tuple, + offsets: torch.Tensor, + accumulate: bool, + wgrad_kernel_fn, + single_grouped_weight: bool, +): + """Compute wgrad using the cuDNN CuTe DSL grouped GEMM wgrad kernel. + + The cuDNN wgrad kernel computes: + wgrad[e] = a[:, tok_start:tok_end] @ b[tok_start:tok_end, :] + where a = DY^T = (out_features, total_tokens) row-major and + b = X = (total_tokens, in_features) column-major. + """ + out_features, in_features = weight_shape + total_tokens = grouped_dy.logical_shape[0] + + 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_tensor = grouped_dy.columnwise_scale_inv.view(out_features, -1).view( + dtype=torch.float8_e8m0fnu + ) + sfb_tensor = grouped_x.columnwise_scale_inv.view(in_features, -1).view( + dtype=torch.float8_e8m0fnu + ) + offsets_tensor = offsets.to(dtype=torch.int32) + + # Prepare wgrad output + if single_grouped_weight: + # Dense mode: single (num_groups, out_features, in_features) tensor + wgrad_tensor = wgrad_output.rowwise_data.view( + offsets_tensor.shape[0], out_features, in_features + ) + wgrad_kernel_fn( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets_tensor, + output_mode="dense", + wgrad_tensor=wgrad_tensor, + acc_dtype=torch.float32, + wgrad_dtype=wgrad_tensor.dtype, + sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, + accumulate_on_output=accumulate, + ) + else: + # Discrete mode: per-expert wgrad device pointers + (wgrad_ptrs,) = tex.convert_host_pointers_to_tensor([wgrad_output]) + wgrad_kernel_fn( + a_tensor=a_tensor, + b_tensor=b_tensor, + sfa_tensor=sfa_tensor, + sfb_tensor=sfb_tensor, + offsets_tensor=offsets_tensor, + output_mode="discrete", + wgrad_ptrs=wgrad_ptrs, + acc_dtype=torch.float32, + wgrad_dtype=wgrad_output[0].dtype, + sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, + accumulate_on_output=accumulate, + ) + + @functools.lru_cache(maxsize=1) def _dglu_wrapper_has_generate_dbias_arg() -> bool: """True if cudnn-frontend SM100 dGLU wrapper accepts ``generate_dbias``.""" @@ -61,6 +133,9 @@ def _compute_grad_params( bias_grads, bias_grad_packed, label="", + *, + cudnn_wgrad_kernel_fn, + offsets, ): """Compute weight gradients and build grad_params for a GroupedLinear layer. Returns the grad_params list in parameter registration order. @@ -131,11 +206,23 @@ def _compute_grad_params( if ctx.weight_requires_grad: # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() - gemm_fn = functools.partial( - general_grouped_gemm_for_grouped_tensor, - layout="NT", - accumulate=accumulate_into_main_grad, - ) + if cudnn_wgrad_kernel_fn is not None: + gemm_fn = functools.partial( + _cudnn_compute_wgrad, + weight_shape=weight_shape, + offsets=offsets, + accumulate=accumulate_into_main_grad, + wgrad_kernel_fn=cudnn_wgrad_kernel_fn, + single_grouped_weight=fc_op.single_grouped_weight, + ) + else: + gemm_fn = functools.partial( + general_grouped_gemm_for_grouped_tensor, + layout="NT", + accumulate=accumulate_into_main_grad, + use_split_accumulator=_2X_ACC_WGRAD, + ) + if delay_wgrad: fc_op.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], gemm_fn) else: @@ -204,6 +291,19 @@ def grouped_gemm_quant_kernel(cls) -> Callable: return grouped_gemm_quant_wrapper_sm100 + @classmethod + @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. + """ + if not _nvidia_cudnn_frontend_supports_wgrad(): + return None + from cudnn import grouped_gemm_wgrad_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_wgrad_wrapper_sm100 + @classmethod @functools.lru_cache(maxsize=None) def is_supported(cls) -> bool: @@ -477,10 +577,12 @@ def fuser_backward( 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]) - fc1_dy_row_scale = fc2_dgrad_kernel_out["sfd_row_tensor"] + # 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]) - fc1_dy_col_scale = fc2_dgrad_kernel_out["sfd_col_tensor"] + # 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) fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None @@ -553,6 +655,8 @@ 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(), + offsets=split_points, ) # Clear FC2 input tensor if possible @@ -648,6 +752,8 @@ 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(), + offsets=split_points, ) # Clear FC1 input tensor if possible From 31f8ab445aad8ea8139927eeacd310a52bf7990e Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 14 Apr 2026 11:11:55 -0700 Subject: [PATCH 341/521] Current Stream for Wgrad kernel (#2873) * integrate cudnn wgrad kernel Signed-off-by: Varun Thumbe * have only cute dsl for wgrad Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert the change for cudnn Signed-off-by: Varun Thumbe * remove dtype Signed-off-by: Varun Thumbe * fix comment: Signed-off-by: Varun Thumbe * go to cublas if needed Signed-off-by: Varun Thumbe * changes to unblock testing Signed-off-by: Varun Thumbe * stream missing Signed-off-by: Varun Thumbe * Space 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> --- transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 389dfbc838..a7c848a317 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -44,6 +44,7 @@ def _cudnn_compute_wgrad( accumulate: bool, wgrad_kernel_fn, single_grouped_weight: bool, + current_stream=None, ): """Compute wgrad using the cuDNN CuTe DSL grouped GEMM wgrad kernel. @@ -88,6 +89,7 @@ def _cudnn_compute_wgrad( wgrad_dtype=wgrad_tensor.dtype, sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, accumulate_on_output=accumulate, + current_stream=current_stream, ) else: # Discrete mode: per-expert wgrad device pointers @@ -104,6 +106,7 @@ def _cudnn_compute_wgrad( wgrad_dtype=wgrad_output[0].dtype, sf_vec_size=MXFP8_BLOCK_SCALING_SIZE, accumulate_on_output=accumulate, + current_stream=current_stream, ) @@ -214,6 +217,7 @@ def _compute_grad_params( accumulate=accumulate_into_main_grad, wgrad_kernel_fn=cudnn_wgrad_kernel_fn, single_grouped_weight=fc_op.single_grouped_weight, + current_stream=torch.cuda.current_stream().cuda_stream, ) else: gemm_fn = functools.partial( From 4e57c218b63fb230f39b0de79931bdfafc9db824 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 14 Apr 2026 14:35:31 -0400 Subject: [PATCH 342/521] [PyTorch] Avoid autograd's gradient accumulation in grouped MLP if possible (#2871) * Avoid grad accumulation when not needed Signed-off-by: Kirthi Shankar Sivamani * same change in unfused grouped linear Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- transformer_engine/pytorch/ops/basic/grouped_linear.py | 2 +- transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 0e09c8a38b..e21625276c 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1045,7 +1045,7 @@ def fuser_backward( grad_weight = torch.stack(grad_weights, dim=0) final_weight_grads = [grad_weight] else: - if delay_wgrad and ctx.weight_requires_grad: + 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 diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index a7c848a317..096e65d296 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -246,8 +246,8 @@ def _compute_grad_params( ) w_list = [packed_wgrad] else: - if delay_wgrad: - w_list = list(w_list) if accumulate_into_main_grad else [None] * num_groups + if delay_wgrad or accumulate_into_main_grad: + w_list = [None] * num_groups if accumulate_into_main_grad: for idx in range(num_groups): wp = getattr(fc_op, f"weight{idx}") From c7205a72c236598fec7ef5262b04fe955bbd0dd7 Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Tue, 14 Apr 2026 12:36:15 -0600 Subject: [PATCH 343/521] Strip local version labels from package version checks (#2858) Pre-compiled Flash Attention wheels (e.g. from mjun0812/flash-attention-prebuild-wheels) embed build metadata in their package version string (e.g. "2.8.3+cu130torch2.11"). While flash_attn.__version__ returns the clean "2.8.3", TE reads the version via importlib.metadata which returns the full string including the local segment. Under PEP 440, "2.8.3+local" > "2.8.3", causing version range checks like `min_version <= version <= max_version` to incorrectly reject a compatible installation. Use `Version.public` to strip the local label before comparison at all `get_pkg_version` call sites (flash-attn, flash-attn-3). Signed-off-by: Peter St. John --- .../pytorch/attention/dot_product_attention/backends.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index b5ed15f8e0..19da8ebffe 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -89,7 +89,7 @@ _flash_attn_varlen_fwd = None _flash_attn_varlen_bwd = None try: - fa_utils.version = PkgVersion(get_pkg_version("flash-attn")) + fa_utils.version = PkgVersion(PkgVersion(get_pkg_version("flash-attn")).public) except PackageNotFoundError: pass # only print warning if use_flash_attention_2 = True in get_attention_backend else: @@ -131,7 +131,7 @@ fa_utils.version, ) try: - fa_utils.fa3_version = PkgVersion(get_pkg_version("flash-attn-3")) + fa_utils.fa3_version = PkgVersion(PkgVersion(get_pkg_version("flash-attn-3")).public) except PackageNotFoundError: flash_attn_func_v3 = None flash_attn_varlen_func_v3 = None From 5d5065ff085fe74827ae1d61abbfd862089291af Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Tue, 14 Apr 2026 15:55:57 -0700 Subject: [PATCH 344/521] Reduce number of C++ test cases for MXFP8 cast and activation kernels (#2874) Reduce number of test cases for MXFP8 cast and activation kernels Signed-off-by: Tim Moon --- tests/cpp/operator/test_act.cu | 30 +++++--- tests/cpp/operator/test_cast_mxfp8.cu | 72 +++++++++++++------ tests/cpp/operator/test_cast_mxfp8_grouped.cu | 37 +++++++++- tests/cpp/test_common.cu | 9 ++- 4 files changed, 110 insertions(+), 38 deletions(-) diff --git a/tests/cpp/operator/test_act.cu b/tests/cpp/operator/test_act.cu index b4280818a8..ca5ccdc4ce 100644 --- a/tests/cpp/operator/test_act.cu +++ b/tests/cpp/operator/test_act.cu @@ -394,19 +394,31 @@ std::vector> act_test_cases = {{2048, 12288}, {257, 259}, {128, 128+1}}; +std::string test_name_generator( + const testing::TestParamInfo& info) { + std::string name = test::typeName(std::get<0>(info.param)) + "X" + + test::typeName(std::get<1>(info.param)) + "X" + + std::to_string(std::get<2>(info.param).first) + "X" + + std::to_string(std::get<2>(info.param).second); + return name; +} + } // namespace INSTANTIATE_TEST_SUITE_P( - OperatorTest, + OperatorTest_ActTestSuite_BF16, ActTestSuite, ::testing::Combine( - ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Values(DType::kBFloat16), ::testing::ValuesIn(test::all_fp_types), ::testing::ValuesIn(act_test_cases)), - [](const testing::TestParamInfo& info) { - std::string name = test::typeName(std::get<0>(info.param)) + "X" + - test::typeName(std::get<1>(info.param)) + "X" + - std::to_string(std::get<2>(info.param).first) + "X" + - std::to_string(std::get<2>(info.param).second); - return name; - }); + test_name_generator); + +INSTANTIATE_TEST_SUITE_P( + OperatorTest_ActTestSuite_DType, + ActTestSuite, + ::testing::Combine( + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::ValuesIn(test::all_fp_types), + ::testing::Values(std::pair{768, 2816})), + test_name_generator); diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index ccc605c060..c7c778ce1e 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -524,14 +524,8 @@ void performTest_x2(const ProcessingMethod processing_method, std::vector> matrix_sizes = { {1, 16}, {16, 48}, - {65, 96}, {128, 128}, - {256, 256}, {993, 512}, - {511, 6144}, - {8192, 128}, - {2048, 160}, - {577, 1632}, {1024}, {8, 32, 1024}, {16, 8, 4, 512}, @@ -570,8 +564,6 @@ std::vector Activation_types = { // ActivationType::SReLU, }; -} // namespace - class FusedCastMXFP8TestSuite : public ::testing::TestWithParam & info) { + std::string name = to_string(std::get<0>(info.param)) + "X" + + to_string(std::get<1>(info.param)); + const auto& shape = std::get<2>(info.param); + for ( const auto& s: shape) { + name += "X" + std::to_string(s); + } + name += "X" + std::to_string(std::get<3>(info.param).first) + + "X" + std::to_string(std::get<3>(info.param).second) + + "X" + test::typeName(std::get<4>(info.param)) + + "X" + test::typeName(std::get<5>(info.param)) + + "X" + test::caseName(std::get<6>(info.param)); + return name; +} + +} // namespace + +// Test cases with only cast kernels INSTANTIATE_TEST_SUITE_P( - OperatorTest, + OperatorTest_FusedCastMXFP8_CastOnly, + FusedCastMXFP8TestSuite, + ::testing::Combine( + ::testing::Values(ProcessingMethod::CAST_ONLY), + ::testing::Values(ActivationType::Identity), + ::testing::ValuesIn(matrix_sizes), + ::testing::ValuesIn(block_sizes), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::ValuesIn(input_scenarios)), + test_name_generator); + +// Test cases with varying matrix shapes and block shapes +INSTANTIATE_TEST_SUITE_P( + OperatorTest_FusedCastMXFP8_Sizes, FusedCastMXFP8TestSuite, ::testing::Combine( ::testing::ValuesIn(processing_methods), ::testing::ValuesIn(Activation_types), ::testing::ValuesIn(matrix_sizes), ::testing::ValuesIn(block_sizes), + ::testing::Values(DType::kBFloat16), + ::testing::Values(DType::kFloat8E4M3), + ::testing::ValuesIn(input_scenarios)), + test_name_generator); + +// Test cases with varying dtypes +INSTANTIATE_TEST_SUITE_P( + OperatorTest_FusedCastMXFP8_Dtypes, + FusedCastMXFP8TestSuite, + ::testing::Combine( + ::testing::ValuesIn(processing_methods), + ::testing::ValuesIn(Activation_types), + ::testing::Values(std::vector{256, 384}), + ::testing::Values(std::pair{32, 32}), ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), ::testing::ValuesIn(input_scenarios)), - [](const testing::TestParamInfo& info) { - std::string name = to_string(std::get<0>(info.param)) + "X" + - to_string(std::get<1>(info.param)); - const auto& shape = std::get<2>(info.param); - for ( const auto& s: shape) { - name += "X" + std::to_string(s); - } - name += "X" + std::to_string(std::get<3>(info.param).first) + - "X" + std::to_string(std::get<3>(info.param).second) + - "X" + test::typeName(std::get<4>(info.param)) + - "X" + test::typeName(std::get<5>(info.param)) + - "X" + test::caseName(std::get<6>(info.param)); - return name; - }); + test_name_generator); diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 3b097cff43..de72299be1 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -669,8 +669,13 @@ std::vector> input_config = { {VARYING_FIRST_DIM, 4, 512,160, 128,0,0,256}, {VARYING_BOTH_DIMS, 3, 1,(128*128)+(128*128), 128,0,128, 128,0,128}, }; - -} // namespace +std::vector> input_config_small = { + {SAME_BOTH_DIMS, 2, 256,128}, + {VARYING_FIRST_DIM, 4, 1536,160, 128,384,512,512}, + {VARYING_LAST_DIM, 3, 256,896, 128,256,512}, + {VARYING_BOTH_DIMS, 2, 1,(256*128)+(512*640), 256,512, 128,640}, + {VARYING_FIRST_DIM, 4, 512,160, 128,0,0,256}, +}; class GroupedFusedCastMXFP8TestSuite : public ::testing::TestWithParam int32_t { + cudaDeviceProp deviceProp; + cudaGetDeviceProperties(&deviceProp, 0); + return 10 * deviceProp.major + deviceProp.minor; + }(); + return compute_capability; } size_t first_dimension(const std::vector &shape) { From 70af73058946228dd87efab76df1288d128a9d3c Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Tue, 14 Apr 2026 17:31:11 -0700 Subject: [PATCH 345/521] [JAX] MXFP8 Grouped Quant+GEMM (#2763) * [JAX] MXFP8 Grouped Quant+GEMM Signed-off-by: Jeremy Berchtold * Update transformer_engine/jax/cpp_extensions/gemm.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --------- Signed-off-by: Jeremy Berchtold Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/jax/test_custom_call_compute.py | 88 ++- .../common/gemm/cublaslt_grouped_gemm.cu | 38 ++ .../common/include/transformer_engine/gemm.h | 29 + transformer_engine/jax/cpp_extensions/gemm.py | 569 +++++++++++------- .../jax/cpp_extensions/quantization.py | 112 +++- transformer_engine/jax/csrc/extensions.h | 2 + .../jax/csrc/extensions/gemm.cpp | 236 +++++--- .../jax/csrc/extensions/pybind.cpp | 1 + .../jax/csrc/extensions/quantization.cpp | 167 ++++- transformer_engine/jax/flax/module.py | 25 +- .../jax/quantize/dequantizer.py | 86 ++- transformer_engine/jax/quantize/tensor.py | 29 + 12 files changed, 1043 insertions(+), 339 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index ddb74fd636..3e5529c077 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -27,6 +27,7 @@ from transformer_engine.jax.cpp_extensions.quantization import ( _jax_quantize, _jax_quantize_dbias, + GroupedQuantizePrimitive, ) from transformer_engine.jax.cpp_extensions.misc import get_cudnn_version from transformer_engine.jax import cpp_extensions as tex @@ -1068,7 +1069,24 @@ def test_rht_gemm(self, in_dtype, q_dtype, scaling_mode, m, n, k, data_layout, w @pytest.mark.skipif(not is_fp8_supported, reason=fp8_unsupported_reason) @pytest_parametrize_wrapper("in_dtype", QUANTIZATION_INPUT_DTYPE) -@pytest_parametrize_wrapper("input_shape", [(8, 16, 32)]) +@pytest_parametrize_wrapper( + "input_shape", + [ + (8, 16, 32), # V1 MXFP8: K=32 not 128-aligned + ( + 4, + 8, + 128, + ), # V2 MXFP8 eligible: K=128, M*32=256 both 128-aligned. Alignment is required due to V2 grouped quantize and grouped GEMM kernel requirements. + ], +) +@pytest_parametrize_wrapper( + "group_size_multiplier", + [ + 32, # V1 MXFP8: group size must be multiple of 32 + 128, # V2 MXFP8 eligible: group size must be multiple of 128. Alignment is required due to V2 grouped quantize and grouped GEMM kernel requirements. + ], +) @pytest_parametrize_wrapper("q_dtype", [jnp.float8_e4m3fn]) @pytest_parametrize_wrapper("scaling_mode", non_fp4_supported_scaling_modes) @pytest_parametrize_wrapper("flatten_axis", [-1]) @@ -1078,14 +1096,21 @@ def test_rht_gemm(self, in_dtype, q_dtype, scaling_mode, m, n, k, data_layout, w ) class TestGroupedQuantize: def test_grouped_qdq( - self, in_dtype, input_shape, q_dtype, scaling_mode, q_layout, flatten_axis, with_group_sizes + self, + in_dtype, + input_shape, + group_size_multiplier, + q_dtype, + scaling_mode, + q_layout, + flatten_axis, + with_group_sizes, ): n_groups, m, n = input_shape key = jax.random.PRNGKey(0) subkeys = jax.random.split(key, 2) - # *32 so that the input shapes works for MXFP8 - input_shape = (m * 32, n) + input_shape = (m * group_size_multiplier, n) if with_group_sizes: group_sizes = jnp.sort(jax.random.randint(subkeys[0], (n_groups - 1,), 0, m)) @@ -1093,7 +1118,7 @@ def test_grouped_qdq( group_sizes = jnp.diff(group_sizes) assert group_sizes.sum() == m assert jnp.any(group_sizes == 0) # make sure that at least one group has 0 row - group_sizes = group_sizes * 32 + group_sizes = group_sizes * group_size_multiplier else: group_sizes = None input_shape = (n_groups, input_shape[0] // n_groups, input_shape[1]) @@ -1101,6 +1126,23 @@ def test_grouped_qdq( if flatten_axis == -2: input_shape = input_shape[:-1] + (2,) + input_shape[-1:] + # V2 MXFP8 quantize kernel requires every individual group size to be a multiple of 128. + # for padding and alignment constraints in the kernel and in the V2 grouped GEMM kernel. + # group_size_multiplier=32 can produce groups of 32 or 64 rows which violate this. + # This cannot be checked at runtime (group sizes live on device), so we skip the + # test configuration rather than weaken the kernel-selection logic. + if ( + scaling_mode == ScalingMode.MXFP8_1D_SCALING + and group_size_multiplier % 128 != 0 + and GroupedQuantizePrimitive._use_v2_kernel( + scaling_mode.value, input_shape, flatten_axis + ) + ): + pytest.skip( + "MXFP8 V2 quantize requires each group to be 128-aligned; " + f"group_size_multiplier={group_size_multiplier} may produce smaller groups" + ) + x = jax.random.uniform(subkeys[1], input_shape, in_dtype) grouped_quantizer = QuantizerFactory.create( @@ -1713,10 +1755,21 @@ def ref_func(x, gamma, kernel_1, kernel_2, bias_1, bias_2): ] GROUPED_DENSE_INPUT_SHAPES = [ - # (n_groups, m, n, k), the actual m will be multiplied by 32 - (5, 32, 128, 64), # Test the case where n_groups is not a multiple of 4 - (8, 64, 32, 128), - (8, 64, 128, 256), + # (n_groups, m, n, k), the actual m will be multiplied by group_size_multiplier + (5, 32, 128, 64), # V1 MXFP8: K=64 not 128-aligned; also tests n_groups not a multiple of 4 + (8, 64, 32, 128), # V1 MXFP8 GEMM: N=32 not 128-aligned + ( + 8, + 64, + 128, + 256, + ), # V2 MXFP8 eligible: K=256, N=128 both 128-aligned. Alignment is required due to V2 grouped quantize and grouped GEMM kernel requirements. + ( + 4, + 4, + 128, + 128, + ), # V2 MXFP8 eligible: K=128, N=128 both 128-aligned (smaller shape). Alignment is required due to V2 grouped quantize and grouped GEMM kernel requirements. ] @@ -1742,7 +1795,9 @@ def _ref_grouped_dense(self, lhs, rhs, bias, group_sizes, contracting_dims): ref_out.append(jnp.squeeze(out_i)) return ref_out - def _generate_grouped_dense_input(self, dtype, input_shape, data_layout="NN", with_bias=False): + def _generate_grouped_dense_input( + self, dtype, input_shape, data_layout="NN", with_bias=False, group_size_multiplier=32 + ): key = jax.random.PRNGKey(0) subkeys = jax.random.split(key, 4) n_groups, m, n, k = input_shape @@ -1755,9 +1810,9 @@ def _generate_grouped_dense_input(self, dtype, input_shape, data_layout="NN", wi group_sizes = group_sizes.at[1].set(0) assert group_sizes.sum() == m - # *32 to make sure that input shape works for MXFP8 - group_sizes = group_sizes * 32 - m = m * 32 + # Scale group sizes by the multiplier for alignment requirements. + group_sizes = group_sizes * group_size_multiplier + m = m * group_size_multiplier lhs_shape = (m if data_layout[0] == "N" else k, k if data_layout[0] == "N" else m) rhs_shape = (n_groups, k if data_layout[1] == "N" else n, n if data_layout[1] == "N" else k) @@ -1831,8 +1886,10 @@ def test_grouped_gemm_fp8(self, fwd_bwd_dtype, scaling_mode, input_shape, layout quantizer.q_dtype = bwd_dtype out_dtype = jnp.bfloat16 + # MXFP8 V2 kernel requires each group's row count to be divisible by due to V2 grouped quantize and grouped GEMM kernel requirements. + is_mxfp8 = scaling_mode == ScalingMode.MXFP8_1D_SCALING lhs, rhs, group_sizes, contracting_dims, _ = self._generate_grouped_dense_input( - out_dtype, input_shape, layout + out_dtype, input_shape, layout, group_size_multiplier=128 if is_mxfp8 else 32 ) ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims) @@ -1906,10 +1963,13 @@ def test_grouped_dense_grad_fp16(self, dtype, input_shape): def test_grouped_dense_grad_fp8(self, fwd_bwd_dtype, scaling_mode, input_shape): fwd_dtype, bwd_dtype = fwd_bwd_dtype dtype = jnp.bfloat16 + # MXFP8 V2 kernel requires each group's row count to be divisible by 128 due to V2 grouped quantize and grouped GEMM kernel requirements. + is_mxfp8 = scaling_mode == ScalingMode.MXFP8_1D_SCALING x, kernel, group_sizes, contracting_dims, bias = self._generate_grouped_dense_input( dtype, input_shape, with_bias=True, + group_size_multiplier=128 if is_mxfp8 else 32, ) quantizer_set = QuantizerFactory.create_set( diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index a8e0b6df83..985c53f760 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -1414,6 +1414,24 @@ __global__ void convert_int32_to_int64_kernel(const int32_t *src, int64_t *dst, if (idx < n) dst[idx] = static_cast(src[idx]); } +// Like convert_int32_to_int64_kernel but scales each element by multiplier. +// Used to convert per-expert slice counts to per-expert row counts for multi-dim tensors. +__global__ void convert_int32_to_int64_with_multiplier_kernel(const int32_t *src, int64_t *dst, + size_t n, int64_t multiplier) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) dst[idx] = static_cast(src[idx]) * multiplier; +} + +// Computes exclusive prefix sums: offsets[0]=0, offsets[i]=sum(first_dims[0..i-1]*last_dim). +// Produces n_groups+1 values. Single-threaded sequential scan; n_groups is typically small. +__global__ void compute_grouped_tensor_offsets_kernel(const int64_t *first_dims, int64_t *offsets, + size_t n_groups, int64_t last_dim) { + offsets[0] = 0; + for (size_t i = 0; i < n_groups; i++) { + offsets[i + 1] = offsets[i] + first_dims[i] * last_dim; + } +} + } // namespace void nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cudaStream_t stream) { @@ -1424,3 +1442,23 @@ void nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cud convert_int32_to_int64_kernel<<>>(src, dst, n); NVTE_CHECK_CUDA(cudaGetLastError()); } + +void nvte_convert_int32_to_int64_with_multiplier(const int32_t *src, int64_t *dst, size_t n, + int64_t multiplier, cudaStream_t stream) { + NVTE_API_CALL(nvte_convert_int32_to_int64_with_multiplier); + if (n == 0) return; + const int threads = 256; + const int blocks = static_cast((n + threads - 1) / threads); + convert_int32_to_int64_with_multiplier_kernel<<>>(src, dst, n, + multiplier); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +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) { + NVTE_API_CALL(nvte_compute_grouped_tensor_offsets); + // Always write at least offsets[0]=0 (needed even for n_groups==0). + compute_grouped_tensor_offsets_kernel<<<1, 1, 0, stream>>>(first_dims, offsets, n_groups, + last_dim); + NVTE_CHECK_CUDA(cudaGetLastError()); +} diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 6999dd857f..fcd08a40a9 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -356,6 +356,35 @@ size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors); */ void nvte_convert_int32_to_int64(const int32_t *src, int64_t *dst, size_t n, cudaStream_t stream); +/*! \brief Convert int32 array to int64 while scaling each element by a multiplier. + * + * Computes dst[i] = (int64_t)src[i] * multiplier for each i in [0, n). + * CUDA-graph safe (no host-device synchronization). + * + * \param[in] src Device pointer to source int32 array. + * \param[out] dst Device pointer to destination int64 array. + * \param[in] n Number of elements. + * \param[in] multiplier Scale factor applied to each element. + * \param[in] stream CUDA stream. + */ +void nvte_convert_int32_to_int64_with_multiplier(const int32_t *src, int64_t *dst, size_t n, + int64_t multiplier, cudaStream_t stream); + +/*! \brief Compute exclusive prefix-sum offsets from per-group first-dimension sizes. + * + * Writes n_groups+1 values to offsets: offsets[0]=0, + * offsets[i] = sum(first_dims[0..i-1] * last_dim) for i in [1, n_groups]. + * This is CUDA-graph safe (no host-device synchronization). + * + * \param[in] first_dims Device pointer to int64 array of length n_groups. + * \param[out] offsets Device pointer to int64 array of length n_groups+1. + * \param[in] n_groups Number of groups. + * \param[in] last_dim Common last dimension (number of columns). + * \param[in] stream CUDA stream. + */ +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); + 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, diff --git a/transformer_engine/jax/cpp_extensions/gemm.py b/transformer_engine/jax/cpp_extensions/gemm.py index c081e451a7..4ff6d07986 100644 --- a/transformer_engine/jax/cpp_extensions/gemm.py +++ b/transformer_engine/jax/cpp_extensions/gemm.py @@ -9,7 +9,7 @@ from collections.abc import Iterable from dataclasses import dataclass from functools import partial, reduce, cache -from typing import Tuple, Sequence, Union +from typing import Tuple, Sequence, Union, Optional from enum import Enum import warnings @@ -47,7 +47,7 @@ apply_padding_to_scale_inv, QuantizeLayout, ) -from .misc import get_padded_spec, is_all_reduce_in_float32 +from .misc import get_padded_spec, is_all_reduce_in_float32, get_min_device_compute_capability from ..sharding import ( global_mesh_resource, tpsp_axis_size, @@ -66,6 +66,7 @@ "sanitize_dims", "get_non_contracting_dims", "transpose_dims", + "is_v2_grouped_gemm_supported", ] @@ -1597,7 +1598,6 @@ def _compute_cublas_workspace_size( workspace_size = get_cublas_workspace_size_bytes() * stream_count workspace_alignment_padding = 256 tensor_scaling_sinv_aligment = 16 - mxfp8_scaling_sinv_alignment_padding = 256 # cuBLAS workspace ptr must be 256 bytes aligned but JAX buffers are not # necessarily 256 bytes aligned, we add some padding to ensure alignment. workspace_size += workspace_alignment_padding @@ -1610,9 +1610,9 @@ def _compute_cublas_workspace_size( workspace_size += lhs_scale_inv_aval.size * tensor_scaling_sinv_aligment workspace_size += rhs_scale_inv_aval.size * tensor_scaling_sinv_aligment elif scaling_mode == ScalingMode.MXFP8_1D_SCALING.value: - # We also pad scale_inv swizzle buffers size for 256 bytes alignment. - workspace_size += lhs_scale_inv_aval.size + mxfp8_scaling_sinv_alignment_padding - workspace_size += rhs_scale_inv_aval.size + mxfp8_scaling_sinv_alignment_padding + # Both V1 and V2 quantize now produce pre-swizzled scales, so the GEMM + # does not need extra workspace for nvte_swizzle_scaling_factors. + pass return workspace_size @staticmethod @@ -2036,48 +2036,303 @@ def _should_enforce_v2_grouped_gemm() -> bool: ) from e -def _can_use_v2_grouped_gemm( +def _is_v2_grouped_gemm_supported( scaling_mode: ScalingMode, dtype: jnp.dtype, has_bias: bool, -) -> bool: - """Determine whether the cuda-graphable grouped GEMM implementation can be used based on the input parameters.""" - # Use the cuda-graphable path for plain BF16 non-quantized inputs; fall back to the legacy - # nvte_multi_tensor_gemm path for all other cases (FP8, MXFP8, etc.) to stay - # feature-compatible with the main branch. - # Bias can be supported in a kernel or in pure-JAX in the future. - - enforce_v2_gmm = _should_enforce_v2_grouped_gemm() + lhs_shape=None, + rhs_shape=None, + lhs_axis_boundary=None, + rhs_axis_boundary=None, +) -> tuple[bool, str]: + """Determine whether the V2 grouped GEMM implementation can be used based on the input parameters.""" if not _v2_grouped_gemm_available: - if enforce_v2_gmm: - raise RuntimeError( - "The TE V2 grouped GEMM is not available but NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is" - " enabled. The reason for V2 grouped GEMM not being available:" - f" {_v2_grouped_gemm_available_reason}" - ) - return False + return ( + False, + ( + "TE was not compiled with support for the V2 grouped GEMM kernel, reason: " + f"{_v2_grouped_gemm_available_reason}" + ), + ) # nvte_grouped_gemm (the v2 kernel) requires SM100+ (Blackwell or newer). # Fall back to the v1 path on SM90 (Hopper) and older architectures. - if get_device_compute_capability(0) < 100: - if enforce_v2_gmm: - raise RuntimeError( - "The TE V2 grouped GEMM requires SM100+ (Blackwell or newer) but current device" - f" compute capability of GPU 0 is {get_device_compute_capability(0)} and" - " NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is enabled." - ) - return False + if get_min_device_compute_capability() < 100: + return ( + False, + ( + "The TE V2 grouped GEMM requires SM100+ (Blackwell or newer) but current min device" + f" compute capability is {get_min_device_compute_capability()}." + ), + ) - if scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16 and not has_bias: - return True + if has_bias: + return False, "Grouped GEMM with bias is not supported in the TE V2 grouped GEMM kernel." + + if scaling_mode == ScalingMode.NO_SCALING and dtype == jnp.bfloat16: + return True, "" + + if scaling_mode == ScalingMode.MXFP8_1D_SCALING: + # V2 MXFP8 requires that the total first dimension of both operands (up to + # axis_boundary) is divisible by 128, matching the quantize V2 kernel requirement. + # Individual group sizes must also be 128-aligned (dynamic constraint). + if lhs_shape is not None and lhs_axis_boundary is not None: + lhs_first_dim = math.prod(lhs_shape[:lhs_axis_boundary]) + if lhs_first_dim % 128 != 0: + return ( + False, + ( + "The TE V2 grouped GEMM for MXFP8 requires the product of the first" + " dimensions (up to axis_boundary) of LHS to be divisible by 128, but got" + f" {lhs_first_dim} with lhs_shape={lhs_shape} and" + f" lhs_axis_boundary={lhs_axis_boundary}." + ), + ) + if rhs_shape is not None and rhs_axis_boundary is not None: + rhs_first_dim = math.prod(rhs_shape[:rhs_axis_boundary]) + if rhs_first_dim % 128 != 0: + return ( + False, + ( + "The TE V2 grouped GEMM for MXFP8 requires the product of the first" + " dimensions (up to axis_boundary) of RHS to be divisible by 128, but got" + f" {rhs_first_dim} with rhs_shape={rhs_shape} and" + f" rhs_axis_boundary={rhs_axis_boundary}." + ), + ) - if enforce_v2_gmm: + # V2 MXFP8 also requires that the "last" dimension (after axis_boundary) of both + # operands is a multiple of 128. This is because the MXFP8 scales must be padded to a multiple of (128, 4). The nvte_grouped_gemm setup kernels only handle the case when this dim is a multiple of 128 as well. If it is not, the GEMM setup kernel will not compute the scale offsets correctly and will read overlapping scales from the previous group, causing incorrect results. + if lhs_shape is not None and lhs_axis_boundary is not None: + lhs_last_dim = math.prod(lhs_shape[lhs_axis_boundary:]) + if lhs_last_dim % 128 != 0: + return ( + False, + ( + "The TE V2 grouped GEMM for MXFP8 requires the product of the last" + " dimensions (after axis_boundary) of LHS to be divisible by 128, but got" + f" {lhs_last_dim} with lhs_shape={lhs_shape} and" + f" lhs_axis_boundary={lhs_axis_boundary}." + ), + ) + if rhs_shape is not None and rhs_axis_boundary is not None: + rhs_last_dim = math.prod(rhs_shape[rhs_axis_boundary:]) + if rhs_last_dim % 128 != 0: + return ( + False, + ( + "The TE V2 grouped GEMM for MXFP8 requires the product of the last" + " dimensions (after axis_boundary) of RHS to be divisible by 128, but got" + f" {rhs_last_dim} with rhs_shape={rhs_shape} and" + f" rhs_axis_boundary={rhs_axis_boundary}." + ), + ) + return True, "" + + return ( + False, + ( + "The TE V2 grouped GEMM currently only supports non-quantized BF16 and MXFP8 with 1D" + " block scaling, but NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is enabled and the input" + f" parameters do not meet these requirements (scaling_mode= {scaling_mode}," + f" dtype={dtype}, has_bias={has_bias}, lhs_shape={lhs_shape}, rhs_shape={rhs_shape}," + f" lhs_axis_boundary={lhs_axis_boundary}, rhs_axis_boundary={rhs_axis_boundary})." + ), + ) + + +def is_v2_grouped_gemm_supported( + scaling_mode: ScalingMode, + dtype: jnp.dtype, + has_bias: bool, + lhs_shape=None, + rhs_shape=None, + lhs_axis_boundary=None, + rhs_axis_boundary=None, +) -> tuple[bool, str]: + """Determine whether the V2 grouped GEMM implementation can be used based on the input parameters. + + Returns: + A tuple of (is_supported: bool, reason: str) where is_supported indicates whether the V2 grouped GEMM can be used, and reason provides an explanation if it is not supported. + """ + # Use the V2 path for plain BF16 non-quantized inputs and MXFP8; fall back to + # the legacy nvte_multi_tensor_gemm path for all other cases (tensor-scaled FP8, etc.). + # Bias can be supported in a kernel or in pure-JAX in the future. + + enforce_v2_gmm = _should_enforce_v2_grouped_gemm() + + is_v2_supported, reason = _is_v2_grouped_gemm_supported( + scaling_mode, dtype, has_bias, lhs_shape, rhs_shape, lhs_axis_boundary, rhs_axis_boundary + ) + + if enforce_v2_gmm and not is_v2_supported: raise RuntimeError( - "The TE V2 grouped GEMM currently only supports BF16 with no quantization recipe and" - f" without bias, but received {scaling_mode=}, {dtype=}, {has_bias=}" + "The TE V2 grouped GEMM is not supported for the given input parameters, but" + " NVTE_JAX_ENFORCE_V2_GROUPED_GEMM is enabled. The reason for V2 grouped GEMM not being" + f" supported: {reason}" ) - return False + + return is_v2_supported, reason + + +def _get_out_dtype_and_scaling_mode( + x: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], +) -> Tuple[jnp.dtype, ScalingMode]: + if isinstance(x, GroupedScaledTensor1x): + out_dtype = x.dq_dtype + scaling_mode = x.scaling_mode + elif isinstance(x, GroupedNoScaleTensor): + out_dtype = x.data.dtype + scaling_mode = ScalingMode.NO_SCALING + else: + raise TypeError( + f"Input must be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(x)}" + ) + return out_dtype, scaling_mode + + +def _infer_output_ragged_dims( + lhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + rhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], +) -> Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray]]: + assert isinstance( + lhs, (GroupedNoScaleTensor, GroupedScaledTensor1x) + ), f"Expected lhs to be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(lhs)}" + assert isinstance( + rhs, (GroupedNoScaleTensor, GroupedScaledTensor1x) + ), f"Expected rhs to be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(rhs)}" + + # Infer output dims from which operand has the ragged non-contracting dim. + if rhs.first_dims is not None or rhs.last_dims is not None: + # Wgrad: rhs contracting dim is ragged → output is uniform (G prefix from num_groups) + out_first_dims = None + out_last_dims = None + elif lhs.first_dims is not None: + out_first_dims = lhs.first_dims + out_last_dims = None + elif lhs.last_dims is not None: + out_first_dims = None + out_last_dims = lhs.last_dims + else: + out_first_dims = out_last_dims = None + + return out_first_dims, out_last_dims + + +def _adjust_contracting_dims_for_hopper_fp8_transpose( + lhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + rhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + lhs_contract_dim: Sequence[int], + rhs_contract_dim: Sequence[int], + lhs_is_trans: bool, + rhs_is_trans: bool, +) -> Tuple[bool, bool, Sequence[int], Sequence[int]]: + # Only support FP8 GEMM with NT layout on Hopper and other earlier GPUs + # thus additional transpose is required + lhs_layout_is_T = lhs.data_layout == "T" + rhs_layout_is_T = rhs.data_layout == "T" + # we can't apply _shape_normalization on the grouped input + # thus we need to ensure that lhs is in N and rhs is in T + if lhs_is_trans != lhs_layout_is_T: + raise RuntimeError("lhs input must be transposed before calling grouped_gemm") + if (not rhs_is_trans) != rhs_layout_is_T: + raise RuntimeError("rhs input must be transposed before calling grouped_gemm") + lhs_is_trans = False + rhs_is_trans = True + lhs_ndim = len(lhs.original_shape) + rhs_ndim = len(rhs.original_shape) + if lhs_layout_is_T: + lhs_contract_dim = tuple((lhs_ndim - 1 - i) % lhs_ndim for i in lhs_contract_dim) + if rhs_layout_is_T: + # For rhs [G, K, N], need to exclude the G dim from contract_dim + if ( + lhs.first_dims is not None or lhs.last_dims is not None + ): # fwd/dgrad: rhs has G as first dim + rhs_contract_dim = tuple( + (rhs_ndim - 1 - i) % (rhs_ndim - 1) + 1 for i in rhs_contract_dim + ) + else: + rhs_contract_dim = tuple((rhs_ndim - 1 - i) % rhs_ndim for i in rhs_contract_dim) + + return lhs_is_trans, rhs_is_trans, lhs_contract_dim, rhs_contract_dim + + +def _quantize_inputs_if_needed( + lhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + rhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + quantizer_set: QuantizerSet, + lhs_is_trans: bool, + rhs_is_trans: bool, + lhs_flatten_axis: int, + rhs_flatten_axis: int, +) -> Tuple[ + Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + Union[GroupedNoScaleTensor, GroupedScaledTensor1x], +]: + if quantizer_set is noop_quantizer_set: + return lhs, rhs + + assert isinstance( + lhs, GroupedNoScaleTensor + ), f"Expected lhs to be GroupedNoScaleTensor before quantization, got type={type(lhs)}" + assert isinstance( + rhs, GroupedNoScaleTensor + ), f"Expected rhs to be GroupedNoScaleTensor before quantization, got type={type(rhs)}" + + if not isinstance(quantizer_set.x, GroupedQuantizer): + raise TypeError( + f"Expected quantizer_set.x to be GroupedQuantizer, but got type={type(quantizer_set.x)}" + ) + if type(quantizer_set.x) is not type(quantizer_set.kernel): + raise TypeError( + "Expected quantizer_set.x and quantizer_set.kernel to have the same type, but got" + f" {type(quantizer_set.x)} and {type(quantizer_set.kernel)}" + ) + if ( + quantizer_set.x.scaling_mode.is_tensor_scaling() + and is_fp8_gemm_with_all_layouts_supported() + ): + lhs_is_rowwise = rhs_is_rowwise = True + else: + lhs_is_rowwise = not lhs_is_trans + rhs_is_rowwise = rhs_is_trans + quantizer_set.x.q_layout = QuantizeLayout.ROWWISE if lhs_is_rowwise else QuantizeLayout.COLWISE + quantizer_set.kernel.q_layout = ( + QuantizeLayout.ROWWISE if rhs_is_rowwise else QuantizeLayout.COLWISE + ) + empty_gs = jnp.empty((0,), jnp.int32) + active_group_sizes = next( + ( + gs + for gs in [lhs.first_dims, lhs.last_dims, rhs.first_dims, rhs.last_dims] + if gs is not None and gs.size > 0 + ), + empty_gs, + ) + lhs_input_data = lhs.data + rhs_input_data = rhs.data + lhs_q = grouped_quantize(lhs_input_data, quantizer_set.x, active_group_sizes, lhs_flatten_axis) + rhs_q = grouped_quantize( + rhs_input_data, quantizer_set.kernel, group_sizes=None, flatten_axis=rhs_flatten_axis + ) + return lhs_q, rhs_q + + +def _get_num_gemms( + lhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], + rhs: Union[GroupedNoScaleTensor, GroupedScaledTensor1x], +) -> int: + for x in [lhs, rhs]: + if x.first_dims is not None: + return x.first_dims.size + if x.last_dims is not None: + return x.last_dims.size + raise ValueError( + "Cannot infer number of gemms since neither lhs nor rhs has first_dims or last_dims. " + "Ensure that at least one of the input tensors has valid first_dims or last_dims." + "For grouped_gemm, at least one tensor must be ragged." + ) def grouped_gemm( @@ -2113,179 +2368,51 @@ def grouped_gemm( empty_gs = jnp.empty((0,), jnp.int32) - # Extract data, dims, and metadata from tensor objects. - # Keep data in its original layout (may be 1D for quantized tensors) to preserve - # JAX sharding; the C++ side uses original_shape to derive m/n/k. - if isinstance(lhs, GroupedNoScaleTensor): - lhs_data = lhs.data - lhs_shape = lhs.original_shape - lhs_scale_inv = jnp.empty((0,), jnp.float32) - scaling_mode = ScalingMode.NO_SCALING - out_dtype = lhs.data.dtype - lhs_first_dims = lhs.first_dims if lhs.first_dims is not None else empty_gs - lhs_last_dims = lhs.last_dims if lhs.last_dims is not None else empty_gs - elif isinstance(lhs, GroupedScaledTensor1x): - lhs_shape = lhs.original_shape - lhs_data = lhs.data - lhs_scale_inv = lhs.scale_inv - scaling_mode = lhs.scaling_mode - out_dtype = lhs.dq_dtype - lhs_first_dims = lhs.first_dims if lhs.first_dims is not None else empty_gs - lhs_last_dims = lhs.last_dims if lhs.last_dims is not None else empty_gs - else: - raise TypeError( - f"lhs must be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(lhs)}" - ) - - if isinstance(rhs, GroupedNoScaleTensor): - rhs_data = rhs.data - rhs_shape = rhs.original_shape - rhs_scale_inv = jnp.empty((0,), jnp.float32) - rhs_first_dims = rhs.first_dims if rhs.first_dims is not None else empty_gs - rhs_last_dims = rhs.last_dims if rhs.last_dims is not None else empty_gs - elif isinstance(rhs, GroupedScaledTensor1x): - rhs_shape = rhs.original_shape - rhs_data = rhs.data - rhs_scale_inv = rhs.scale_inv - rhs_first_dims = rhs.first_dims if rhs.first_dims is not None else empty_gs - rhs_last_dims = rhs.last_dims if rhs.last_dims is not None else empty_gs - if isinstance(lhs, GroupedScaledTensor1x) and lhs.scaling_mode != rhs.scaling_mode: - raise ValueError( - f"Mismatched scaling modes: lhs.scaling_mode={lhs.scaling_mode}," - f" rhs.scaling_mode={rhs.scaling_mode}" - ) - if isinstance(lhs, GroupedScaledTensor1x): - scaling_mode = lhs.scaling_mode - else: - raise TypeError( - f"rhs must be GroupedNoScaleTensor or GroupedScaledTensor1x, got type={type(rhs)}" - ) + out_dtype, scaling_mode = _get_out_dtype_and_scaling_mode(lhs) + rhs_out_dtype, rhs_scaling_mode = _get_out_dtype_and_scaling_mode(rhs) + assert out_dtype == rhs_out_dtype, f"Mismatched output dtypes: {out_dtype} vs {rhs_out_dtype}" + assert ( + scaling_mode == rhs_scaling_mode + ), f"Mismatched scaling modes: {scaling_mode} vs {rhs_scaling_mode}" + del rhs_out_dtype, rhs_scaling_mode - # Infer output dims from which operand has the ragged non-contracting dim. - if rhs_first_dims.size > 0 or rhs_last_dims.size > 0: - # Wgrad: rhs contracting dim is ragged → output is uniform (G prefix from num_groups) - out_first_dims = empty_gs - out_last_dims = empty_gs - elif lhs_first_dims.size > 0: - out_first_dims = lhs_first_dims - out_last_dims = empty_gs - elif lhs_last_dims.size > 0: - out_first_dims = empty_gs - out_last_dims = lhs_last_dims - else: - out_first_dims = out_last_dims = empty_gs + out_first_dims, out_last_dims = _infer_output_ragged_dims(lhs, rhs) out_dtype = preferred_element_type or out_dtype lhs_contract_dim, rhs_contract_dim = contracting_dims - lhs_is_trans = lhs_contract_dim[-1] != len(lhs_shape) - 1 + lhs_is_trans = lhs_contract_dim[-1] != len(lhs.original_shape) - 1 lhs_flatten_axis = len(lhs_contract_dim) * (1 if lhs_is_trans else -1) # rhs_is_trans: K is the last dim of rhs (i.e., rhs is in "T" layout). - rhs_is_trans = rhs_contract_dim[-1] == len(rhs_shape) - 1 + rhs_is_trans = rhs_contract_dim[-1] == len(rhs.original_shape) - 1 rhs_flatten_axis = -len(rhs_contract_dim) if rhs_is_trans else 1 + len(rhs_contract_dim) - if ( - not isinstance(lhs, ScaledTensor) - and not isinstance(rhs, ScaledTensor) - and quantizer_set != noop_quantizer_set - ): - if not isinstance(quantizer_set.x, GroupedQuantizer): - raise TypeError( - "Expected quantizer_set.x to be GroupedQuantizer, but got" - f" type={type(quantizer_set.x)}" - ) - if type(quantizer_set.x) is not type(quantizer_set.kernel): - raise TypeError( - "Expected quantizer_set.x and quantizer_set.kernel to have the same type, but got" - f" {type(quantizer_set.x)} and {type(quantizer_set.kernel)}" - ) - scaling_mode = quantizer_set.x.scaling_mode - if ( - quantizer_set.x.scaling_mode.is_tensor_scaling() - and is_fp8_gemm_with_all_layouts_supported() - ): - lhs_is_rowwise = rhs_is_rowwise = True - else: - lhs_is_rowwise = not lhs_is_trans - rhs_is_rowwise = rhs_is_trans - quantizer_set.x.q_layout = ( - QuantizeLayout.ROWWISE if lhs_is_rowwise else QuantizeLayout.COLWISE - ) - quantizer_set.kernel.q_layout = ( - QuantizeLayout.ROWWISE if rhs_is_rowwise else QuantizeLayout.COLWISE - ) - active_group_sizes = next( - ( - gs - for gs in [lhs_first_dims, lhs_last_dims, rhs_first_dims, rhs_last_dims] - if gs.size > 0 - ), - empty_gs, - ) - lhs_input_data = lhs.data if isinstance(lhs, GroupedNoScaleTensor) else lhs_data - rhs_input_data = rhs.data if isinstance(rhs, GroupedNoScaleTensor) else rhs_data - lhs_q = grouped_quantize( - lhs_input_data, quantizer_set.x, active_group_sizes, lhs_flatten_axis - ) - rhs_q = grouped_quantize( - rhs_input_data, quantizer_set.kernel, group_sizes=None, flatten_axis=rhs_flatten_axis - ) - lhs_data = lhs_q.data - rhs_data = rhs_q.data - lhs_scale_inv = lhs_q.scale_inv - rhs_scale_inv = rhs_q.scale_inv - lhs_shape = lhs_q.original_shape - rhs_shape = rhs_q.original_shape + lhs, rhs = _quantize_inputs_if_needed( + lhs, rhs, quantizer_set, lhs_is_trans, rhs_is_trans, lhs_flatten_axis, rhs_flatten_axis + ) + + # Re-read scaling_mode after quantization: if _quantize_inputs_if_needed converted + # GroupedNoScaleTensor → GroupedScaledTensor1x, the original scaling_mode (NO_SCALING) + # would cause the C++ kernel to skip scale_inv setup, triggering a cuBLAS assertion. + _, scaling_mode = _get_out_dtype_and_scaling_mode(lhs) - if lhs_data.dtype == jnp.float8_e5m2 and rhs_data.dtype == jnp.float8_e5m2: + if lhs.data.dtype == jnp.float8_e5m2 and rhs.data.dtype == jnp.float8_e5m2: raise ValueError("FP8 GEMM does not support E5M2 * E5M2") - # Only support FP8 GEMM with NT layout on Hopper and other earlier GPUs - # thus additional transpose is required if scaling_mode.is_tensor_scaling() and not is_fp8_gemm_with_all_layouts_supported(): - if isinstance(lhs, ScaledTensor) and isinstance(rhs, ScaledTensor): - lhs_layout_is_T = lhs.data_layout == "T" - rhs_layout_is_T = rhs.data_layout == "T" - else: - lhs_layout_is_T = lhs_q.data_layout == "T" - rhs_layout_is_T = rhs_q.data_layout == "T" - # we can't apply _shape_normalization on the grouped input - # thus we need to ensure that lhs is in N and rhs is in T - if lhs_is_trans != lhs_layout_is_T: - raise RuntimeError("lhs input must be transposed before calling grouped_gemm") - if (not rhs_is_trans) != rhs_layout_is_T: - raise RuntimeError("rhs input must be transposed before calling grouped_gemm") - lhs_is_trans = False - rhs_is_trans = True - lhs_ndim = len(lhs_shape) - rhs_ndim = len(rhs_shape) - if lhs_layout_is_T: - lhs_contract_dim = tuple((lhs_ndim - 1 - i) % lhs_ndim for i in lhs_contract_dim) - if rhs_layout_is_T: - # For rhs [G, K, N], need to exclude the G dim from contract_dim - if ( - lhs_first_dims.size > 0 or lhs_last_dims.size > 0 - ): # fwd/dgrad: rhs has G as first dim - rhs_contract_dim = tuple( - (rhs_ndim - 1 - i) % (rhs_ndim - 1) + 1 for i in rhs_contract_dim - ) - else: - rhs_contract_dim = tuple((rhs_ndim - 1 - i) % rhs_ndim for i in rhs_contract_dim) + lhs_is_trans, rhs_is_trans, lhs_contract_dim, rhs_contract_dim = ( + _adjust_contracting_dims_for_hopper_fp8_transpose( + lhs, rhs, lhs_contract_dim, rhs_contract_dim, lhs_is_trans, rhs_is_trans + ) + ) # Compute N-D axis boundaries from final (post-adjustment) contracting dims. lhs_axis_boundary = get_lhs_axis_boundary(lhs_contract_dim, lhs_is_trans) rhs_axis_boundary = get_rhs_axis_boundary(rhs_contract_dim, rhs_is_trans) - num_gemms = ( - lhs_first_dims.size - or lhs_last_dims.size - or rhs_first_dims.size - or rhs_last_dims.size - or out_first_dims.size - or out_last_dims.size - ) + num_gemms = _get_num_gemms(lhs, rhs) if num_gemms == 0: raise ValueError( "grouped_gemm requires at least one non-empty dimension array. " @@ -2294,26 +2421,28 @@ def grouped_gemm( # Pre-compute collapsed 2D sizes from original N-D shapes. # These are static Python ints passed as primitive parameters (must be hashable). - lhs_left_size = math.prod(lhs_shape[:lhs_axis_boundary]) - lhs_right_size = math.prod(lhs_shape[lhs_axis_boundary:]) - rhs_left_size = math.prod(rhs_shape[:rhs_axis_boundary]) - rhs_right_size = math.prod(rhs_shape[rhs_axis_boundary:]) + lhs_left_size = math.prod(lhs.original_shape[:lhs_axis_boundary]) + lhs_right_size = math.prod(lhs.original_shape[lhs_axis_boundary:]) + rhs_left_size = math.prod(rhs.original_shape[:rhs_axis_boundary]) + rhs_right_size = math.prod(rhs.original_shape[rhs_axis_boundary:]) # Pre-compute output shape from N-D input shapes (static Python ints). if lhs_is_trans: - lhs_non_contracting = lhs_shape[lhs_axis_boundary:] + lhs_non_contracting = lhs.original_shape[lhs_axis_boundary:] else: - lhs_non_contracting = lhs_shape[:lhs_axis_boundary] + lhs_non_contracting = lhs.original_shape[:lhs_axis_boundary] if rhs_is_trans: - if rhs_first_dims.size > 0 or rhs_last_dims.size > 0: + if rhs.first_dims is not None or rhs.last_dims is not None: # wgrad: rhs (e.g. grad_T of shape (N, M)) has no G batch dim; include all dims - rhs_non_contracting = tuple(rhs_shape[d] for d in range(rhs_axis_boundary)) + rhs_non_contracting = tuple(rhs.original_shape[d] for d in range(rhs_axis_boundary)) else: # fwd/dgrad: rhs (e.g. kernel_T of shape (G, N, K)) has G batch dim at dim 0; skip it - rhs_non_contracting = tuple(rhs_shape[d] for d in range(rhs_axis_boundary) if d != 0) + rhs_non_contracting = tuple( + rhs.original_shape[d] for d in range(rhs_axis_boundary) if d != 0 + ) else: - rhs_non_contracting = rhs_shape[rhs_axis_boundary:] - if rhs_first_dims.size > 0 or rhs_last_dims.size > 0: + rhs_non_contracting = rhs.original_shape[rhs_axis_boundary:] + if rhs.first_dims is not None or rhs.last_dims is not None: out_shape = (num_gemms, *lhs_non_contracting, *rhs_non_contracting) else: out_shape = (*lhs_non_contracting, *rhs_non_contracting) @@ -2334,7 +2463,25 @@ def grouped_gemm( " and padded with zeros to not affect the result of the MoE block." ) - use_v2_ffi = _can_use_v2_grouped_gemm(scaling_mode, lhs_data.dtype, has_bias) + use_v2_ffi, _ = is_v2_grouped_gemm_supported( + scaling_mode, + lhs.data.dtype, + has_bias, + lhs_shape=lhs.original_shape, + rhs_shape=rhs.original_shape, + lhs_axis_boundary=lhs_axis_boundary, + rhs_axis_boundary=rhs_axis_boundary, + ) + + if scaling_mode == ScalingMode.MXFP8_1D_SCALING: + # Both V1 and V2 quantize produce pre-swizzled scales (V1 via + # set_with_gemm_swizzled_scales, V2 via nvte_group_quantize). Require that + # grouped_quantize has set pre_swizzled=True on the input tensors. + if not lhs.pre_swizzled: + raise ValueError("lhs must be pre-swizzled for MXFP8 1D scaling") + if not rhs.pre_swizzled: + raise ValueError("rhs must be pre-swizzled for MXFP8 1D scaling") + if use_v2_ffi: additional_arg_0 = jnp.ones((num_gemms,), jnp.float32) # alpha additional_arg_1 = jnp.zeros((num_gemms,), jnp.float32) # beta @@ -2343,17 +2490,17 @@ def grouped_gemm( additional_arg_1 = jnp.zeros((0,), jnp.int32) # unused placeholder (out,) = GroupedGemmPrimitive.outer_primitive.bind( - lhs_data, - lhs_scale_inv, - rhs_data, - rhs_scale_inv, + lhs.data, + lhs.scale_inv if isinstance(lhs, GroupedScaledTensor1x) else jnp.empty((0,), jnp.float32), + rhs.data, + rhs.scale_inv if isinstance(rhs, GroupedScaledTensor1x) else jnp.empty((0,), jnp.float32), bias, - lhs_first_dims, - lhs_last_dims, - rhs_first_dims, - rhs_last_dims, - out_first_dims, - out_last_dims, + lhs.first_dims if lhs.first_dims is not None else empty_gs, + lhs.last_dims if lhs.last_dims is not None else empty_gs, + rhs.first_dims if rhs.first_dims is not None else empty_gs, + rhs.last_dims if rhs.last_dims is not None else empty_gs, + out_first_dims if out_first_dims is not None else empty_gs, + out_last_dims if out_last_dims is not None else empty_gs, additional_arg_0, additional_arg_1, lhs_is_trans=lhs_is_trans, diff --git a/transformer_engine/jax/cpp_extensions/quantization.py b/transformer_engine/jax/cpp_extensions/quantization.py index a3d363e42a..7138cfcf40 100644 --- a/transformer_engine/jax/cpp_extensions/quantization.py +++ b/transformer_engine/jax/cpp_extensions/quantization.py @@ -994,7 +994,8 @@ class GroupedQuantizePrimitive(BasePrimitive): Cast Primitive wrapping nvte_quantize and nvte_quantize_dbias """ - name = "te_grouped_quantize_ffi" + name = "te_grouped_quantize_ffi" # V1: fallback path (supports all shapes, not CUDA-graph safe) + name_v2 = "te_grouped_quantize_v2_ffi" # V2: MXFP8, CUDA-graph safe multiple_results = True impl_static_args = ( 3, @@ -1006,6 +1007,54 @@ class GroupedQuantizePrimitive(BasePrimitive): inner_primitive = None outer_primitive = None + @staticmethod + def _use_v2_kernel(scaling_mode, x_shape, flatten_axis): + """Return True when the V2 (CUDA-graph-safe) MXFP8 kernel can be used. + + V2 requires: + 1. SM100+ (Blackwell) — V2 grouped quantize fuses the scale_inv swizzle via + nvte_group_quantize. The swizzled scale_inv must then be consumed by the + V2 grouped GEMM, which also requires SM100+. Keeping both decisions tied + to SM100+ prevents a mismatch where V2-quantized (pre-swizzled) tensors + are passed to the V1 grouped GEMM (which would re-swizzle and corrupt). + 2. The total first logical dimension (product of x_shape up to flatten_axis) + is divisible by 128. + 3. For multi-dim group tensors (eff > 1, e.g., kernel shape G×K×N), the + per-group row count non_group_m = prod(x_shape[1:eff]) must also be + divisible by 128. + 4. For lhs-style tensors (eff == 1, shape M×K), individual group sizes must + be 128-aligned — this is a dynamic constraint that cannot be checked here + because group sizes live on device. The caller is responsible for ensuring + this. + 5. The last logical dimension (contracting dim K or output dim N) must be + divisible by 128, matching the V2 grouped GEMM constraint so that the + two always agree on V1 vs V2. + + Falls back to V1 when constraints are not met. V1 supports arbitrary shapes + but performs a D2H copy of group_sizes (not CUDA-graph safe). + """ + if ScalingMode(scaling_mode) != ScalingMode.MXFP8_1D_SCALING: + return False + # Require SM100+ so V2 quantize (fused swizzle) is only used alongside V2 GEMM. + if get_min_device_compute_capability() < 100: + return False + ndim = len(x_shape) + eff = flatten_axis if flatten_axis >= 0 else flatten_axis + ndim + total_first_dim = math.prod(x_shape[:eff]) + if total_first_dim % 128 != 0: + return False + # For multi-dim group tensors (e.g., kernel shape G×K×N with eff=2), + # non_group_m = K must also be 128-aligned. + if eff > 1: + non_group_m = math.prod(x_shape[1:eff]) + if non_group_m % 128 != 0: + return False + # Last dim must be 128-aligned to match the V2 grouped GEMM requirement. + last_dim = math.prod(x_shape[eff:]) + if last_dim % 128 != 0: + return False + return True + @staticmethod def abstract( x_aval, @@ -1048,7 +1097,20 @@ def abstract( rowwise_scale_inv_shape = (1,) rowwise_out_aval = jax.core.ShapedArray(shape=rowwise_out_shape, dtype=out_dtype) - amax_aval = jax.core.ShapedArray(shape=(group_sizes_aval.size,), dtype=jnp.float32) + updated_amax_aval = jax.core.ShapedArray(shape=(group_sizes_aval.size,), dtype=jnp.float32) + + use_v2 = GroupedQuantizePrimitive._use_v2_kernel(scaling_mode, x_aval.shape, flatten_axis) + if use_v2: + # V2 path: int64_workspace laid out as: + # [n_groups int64 group_sizes | n_groups+1 int64 offsets] + # = (2*n_groups + 1) * sizeof(int64_t) bytes stored as uint8. + n_groups = group_sizes_aval.size + int64_workspace_aval = jax.core.ShapedArray( + shape=((2 * n_groups + 1) * 8,), dtype=jnp.uint8 + ) + else: + # V1 path: Unused for V1 codepath + int64_workspace_aval = jax.core.ShapedArray(shape=(0,), dtype=jnp.uint8) if q_layout.has_colwise: colwise_out_shape = out_shape @@ -1068,7 +1130,8 @@ def abstract( colwise_out_aval, rowwise_scale_inv_aval, colwise_scale_inv_aval, - amax_aval, + updated_amax_aval, + int64_workspace_aval, ) @staticmethod @@ -1078,13 +1141,20 @@ def outer_abstract(*args, **kwargs): """ # Phuong: keeping outer abstract so that we can add fuse dbias later ( - rowwise_out, - colwise_out, - scale_inv, - colwise_scale_inv, - updated_amax, + rowwise_out_aval, + colwise_out_aval, + rowwise_scale_inv_aval, + colwise_scale_inv_aval, + updated_amax_aval, + _, ) = GroupedQuantizePrimitive.abstract(*args, **kwargs) - return rowwise_out, colwise_out, scale_inv, colwise_scale_inv, updated_amax + return ( + rowwise_out_aval, + colwise_out_aval, + rowwise_scale_inv_aval, + colwise_scale_inv_aval, + updated_amax_aval, + ) @staticmethod def lowering( @@ -1107,6 +1177,21 @@ def lowering( assert x_aval.dtype in [jnp.float32, jnp.float16, jnp.bfloat16] assert scale_aval.dtype == jnp.float32 assert group_sizes_aval.dtype == jnp.int32 + use_v2 = GroupedQuantizePrimitive._use_v2_kernel(scaling_mode, x_aval.shape, flatten_axis) + if use_v2: + # V2: CUDA-graph safe; scale is passed but ignored by the C++ handler. + # Requires total_first_dim % 128 == 0 (checked above) and all individual + # group sizes % 128 == 0 (dynamic constraint, enforced by the kernel). + return ffi.ffi_lowering(GroupedQuantizePrimitive.name_v2)( + ctx, + x, + scale, + group_sizes, + q_layout=q_layout.value.value, + flatten_axis=flatten_axis, + ) + # V1: supports arbitrary shapes but not CUDA-graph safe (performs D2H copy of group_sizes). + # Used for non-MXFP8 scaling modes and for MXFP8 when total_first_dim % 128 != 0. return ffi.ffi_lowering(GroupedQuantizePrimitive.name)( ctx, x, @@ -1138,6 +1223,7 @@ def impl( rowwise_scale_inv, colwise_scale_inv, updated_amax, + _, ) = GroupedQuantizePrimitive.inner_primitive.bind( x, scale, @@ -1148,7 +1234,7 @@ def impl( flatten_axis=flatten_axis, scale_dtype=scale_dtype, ) - return (rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax) + return rowwise_out, colwise_out, rowwise_scale_inv, colwise_scale_inv, updated_amax register_primitive(GroupedQuantizePrimitive) @@ -1259,6 +1345,11 @@ def grouped_quantize( for i, quantizer_i in enumerate(quantizer.quantizers): quantizer_i.update(updated_amax[i].reshape((1,))) + # Both V1 (set_with_gemm_swizzled_scales) and V2 (nvte_group_quantize) produce + # pre-swizzled scale_inv tensors for use by the grouped GEMM kernel. Set + # pre_swizzled=True for all MXFP8 grouped quantization so that grouped_gemm can + # assert this invariant unconditionally. + is_mxfp8 = quantizer.scaling_mode == ScalingMode.MXFP8_1D_SCALING out = ScaledTensorFactory.create( data=rowwise_casted_output, scale_inv=rowwise_scale_inv, @@ -1271,6 +1362,7 @@ def grouped_quantize( flatten_axis=flatten_axis, first_dims=ragged_first_dims, original_shape=original_shape, + pre_swizzled=is_mxfp8, ) return out diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index a74b209e4f..3ba0e7e9b2 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -119,6 +119,8 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(DBiasQuantizeHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedQuantizeHandler); +XLA_FFI_DECLARE_HANDLER_SYMBOL(GroupedQuantizeV2Handler); + XLA_FFI_DECLARE_HANDLER_SYMBOL(DequantizeHandler); pybind11::tuple GetDBiasQuantizeWorkspaceSizes(size_t batch_size, size_t hidden_size, diff --git a/transformer_engine/jax/csrc/extensions/gemm.cpp b/transformer_engine/jax/csrc/extensions/gemm.cpp index a7f16bb31f..6ca907032c 100644 --- a/transformer_engine/jax/csrc/extensions/gemm.cpp +++ b/transformer_engine/jax/csrc/extensions/gemm.cpp @@ -481,6 +481,8 @@ class JAXX_GroupedTensorWrapper { m_grouped_tensor(other.m_grouped_tensor), m_data_tensor(other.m_data_tensor), m_scale_inv_tensor(other.m_scale_inv_tensor), + m_colwise_data_tensor(other.m_colwise_data_tensor), + m_colwise_scale_inv_tensor(other.m_colwise_scale_inv_tensor), m_sizes_tensor(other.m_sizes_tensor), m_offsets_tensor(other.m_offsets_tensor) { other.m_grouped_tensor = nullptr; @@ -489,6 +491,10 @@ class JAXX_GroupedTensorWrapper { ~JAXX_GroupedTensorWrapper(); void set_rowwise(Buffer_Type const &data, std::optional const &scale_inv); + void set_columnwise(Buffer_Type const &data, std::optional const &scale_inv); + void set_with_gemm_swizzled_scales(bool val); + void replace_scale_inv(bool use_colwise, uint8_t *sinv_ptr, NVTEDType sinv_dtype, + NVTEShape sinv_shape); void set_group_info(Buffer_Type const &group_sizes, Buffer_Type const &group_offsets, NVTEGroupedTensorParam group_sizes_param_name); // Set only group sizes (no offsets); the setup kernel will compute offsets from sizes. @@ -505,6 +511,8 @@ class JAXX_GroupedTensorWrapper { // Internal tensors. These need to be kept alive as long as the grouped tensor is alive. NVTEBasicTensor m_data_tensor{}; NVTEBasicTensor m_scale_inv_tensor{}; + NVTEBasicTensor m_colwise_data_tensor{}; + NVTEBasicTensor m_colwise_scale_inv_tensor{}; NVTEBasicTensor m_sizes_tensor{}; NVTEBasicTensor m_offsets_tensor{}; @@ -556,6 +564,58 @@ void JAXX_GroupedTensorWrapper::set_rowwise(Buffer_Type const &data, } } +void JAXX_GroupedTensorWrapper::set_columnwise(Buffer_Type const &data, + std::optional const &scale_inv) { + NVTEDType data_dtype = + static_cast(convert_ffi_datatype_to_te_dtype(data.element_type())); + m_colwise_data_tensor = + NVTEBasicTensor{reinterpret_cast(data.untyped_data()), data_dtype, m_data_shape}; + + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedColumnwiseData, + &m_colwise_data_tensor, sizeof(m_colwise_data_tensor)); + + if (scale_inv.has_value()) { + NVTEDType scale_inv_dtype = + static_cast(convert_ffi_datatype_to_te_dtype(scale_inv->element_type())); + NVTEShape logical_scale_shape{}; + if (scale_inv->dimensions().size() == 1) { + logical_scale_shape.ndim = 1; + logical_scale_shape.data[0] = scale_inv->dimensions()[0]; + } else if (scale_inv->dimensions().size() == 2) { + logical_scale_shape.ndim = 2; + logical_scale_shape.data[0] = scale_inv->dimensions()[0]; + logical_scale_shape.data[1] = scale_inv->dimensions()[1]; + } else { + NVTE_CHECK(false, "Expected 1D or 2D tensor for GEMM columnwise scale_inv but received ndim=", + scale_inv->dimensions().size()); + } + m_colwise_scale_inv_tensor = + NVTEBasicTensor{reinterpret_cast(scale_inv->untyped_data()), scale_inv_dtype, + logical_scale_shape}; + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedColumnwiseScaleInv, + &m_colwise_scale_inv_tensor, sizeof(m_colwise_scale_inv_tensor)); + } +} + +void JAXX_GroupedTensorWrapper::set_with_gemm_swizzled_scales(bool val) { + auto v = static_cast(val); + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedWithGEMMSwizzledScales, &v, + sizeof(v)); +} + +void JAXX_GroupedTensorWrapper::replace_scale_inv(bool use_colwise, uint8_t *sinv_ptr, + NVTEDType sinv_dtype, NVTEShape sinv_shape) { + if (use_colwise) { + m_colwise_scale_inv_tensor = NVTEBasicTensor{sinv_ptr, sinv_dtype, sinv_shape}; + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedColumnwiseScaleInv, + &m_colwise_scale_inv_tensor, sizeof(m_colwise_scale_inv_tensor)); + } else { + m_scale_inv_tensor = NVTEBasicTensor{sinv_ptr, sinv_dtype, sinv_shape}; + nvte_set_grouped_tensor_param(m_grouped_tensor, kNVTEGroupedRowwiseScaleInv, + &m_scale_inv_tensor, sizeof(m_scale_inv_tensor)); + } +} + void JAXX_GroupedTensorWrapper::set_group_info(Buffer_Type const &group_sizes, Buffer_Type const &group_offsets, NVTEGroupedTensorParam group_sizes_param_name) { @@ -619,22 +679,19 @@ JAXX_GroupedTensorWrapper make_grouped_tensor(Buffer_Type const &data, return std::move(grouped_tensor_wrapper); } -// V2 variant: derives data shape from the XLA buffer directly, converts group_sizes +// V2 variant (NO_SCALING): derives data shape from the XLA buffer directly, converts group_sizes // int32→int64 per-tensor into a dedicated slot of int64_workspace, and wires first_dims/last_dims. // int64_offset (in int64 elements) is updated on return to the next available slot so callers can // thread it through successive make_grouped_tensor calls without aliasing. Bounds are checked -// before each slot is used. Only NO_SCALING is supported. +// before each slot is used. Only NO_SCALING is supported by this overload. JAXX_GroupedTensorWrapper make_grouped_tensor( Buffer_Type const &data, Buffer_Type const &first_dims, Buffer_Type const &last_dims, int64_t *int64_workspace_base, size_t int64_workspace_capacity, size_t &int64_offset, - size_t num_gemms, cudaStream_t stream, int64_t axis_boundary = -1) { + size_t num_gemms, cudaStream_t stream, size_t left_size, size_t right_size) { auto dims = data.dimensions(); - NVTE_CHECK(dims.size() >= 2, "grouped GEMM data buffer must be at least 2D."); - // Flatten dims at axis_boundary to produce a 2D NVTE shape. - // axis_boundary=-1 (default) collapses dims[0..N-2] → rows and keeps dims[N-1] → cols, - // preserving the prior behaviour for output buffers (e.g. [G, K, N] for wgrad). - size_t ab = (axis_boundary < 0) ? dims.size() - 1 : static_cast(axis_boundary); - NVTEShape dataShape{.data = {product(dims, 0, ab), product(dims, ab, dims.size())}, .ndim = 2}; + NVTE_CHECK(product(dims) == left_size * right_size, + "grouped GEMM data buffer element count does not match the provided 2D shape."); + NVTEShape dataShape{.data = {left_size, right_size}, .ndim = 2}; JAXX_GroupedTensorWrapper wrapper(JAXX_Scaling_Mode::NO_SCALING, num_gemms, dataShape); wrapper.set_rowwise(data, std::nullopt); if (first_dims.element_count() > 0) { @@ -660,6 +717,56 @@ JAXX_GroupedTensorWrapper make_grouped_tensor( return wrapper; } +// V2 variant with scaling support (MXFP8 or NO_SCALING). Accepts scale_inv buffer and +// use_colwise flag to wire rowwise or columnwise data+scales for the grouped tensor. +// Pre-swizzled scales are indicated via set_with_gemm_swizzled_scales(true). +JAXX_GroupedTensorWrapper make_grouped_tensor( + Buffer_Type const &data, Buffer_Type const &scale_inv, JAXX_Scaling_Mode scaling_mode, + bool use_colwise, Buffer_Type const &first_dims, Buffer_Type const &last_dims, + int64_t *int64_workspace_base, size_t int64_workspace_capacity, size_t &int64_offset, + size_t num_gemms, cudaStream_t stream, size_t left_size, size_t right_size) { + auto dims = data.dimensions(); + NVTE_CHECK(product(dims) == left_size * right_size, + "grouped GEMM data buffer element count does not match the provided 2D shape."); + NVTEShape dataShape{.data = {left_size, right_size}, .ndim = 2}; + JAXX_GroupedTensorWrapper wrapper(scaling_mode, num_gemms, dataShape); + + const bool is_mxfp8 = scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING; + if (is_mxfp8 && use_colwise) { + wrapper.set_columnwise(data, scale_inv); + } else if (is_mxfp8) { + wrapper.set_rowwise(data, scale_inv); + } else { + // NO_SCALING: no scale_inv needed + wrapper.set_rowwise(data, std::nullopt); + } + if (is_mxfp8) { + wrapper.set_with_gemm_swizzled_scales(true); + } + + if (first_dims.element_count() > 0) { + NVTE_CHECK(first_dims.element_type() == xla::ffi::DataType::S32, "group_sizes must be int32."); + NVTE_CHECK(int64_offset + num_gemms <= int64_workspace_capacity, + "int64_workspace overflow: not enough space for first_dims conversion."); + auto *slot = int64_workspace_base + int64_offset; + nvte_convert_int32_to_int64(reinterpret_cast(first_dims.untyped_data()), slot, + num_gemms, stream); + wrapper.set_group_sizes_only(slot, num_gemms, kNVTEGroupedFirstDims); + int64_offset += num_gemms; + } + if (last_dims.element_count() > 0) { + NVTE_CHECK(last_dims.element_type() == xla::ffi::DataType::S32, "group_sizes must be int32."); + NVTE_CHECK(int64_offset + num_gemms <= int64_workspace_capacity, + "int64_workspace overflow: not enough space for last_dims conversion."); + auto *slot = int64_workspace_base + int64_offset; + nvte_convert_int32_to_int64(reinterpret_cast(last_dims.untyped_data()), slot, + num_gemms, stream); + wrapper.set_group_sizes_only(slot, num_gemms, kNVTEGroupedLastDims); + int64_offset += num_gemms; + } + return wrapper; +} + // Returns num_gemms from the first non-empty per-tensor group_sizes buffer, // falling back to the element count of alpha for the uniform-batch case. size_t grouped_gemm_num_gemms(Buffer_Type const &lhs_first_dims, Buffer_Type const &lhs_last_dims, @@ -752,13 +859,19 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty auto [lhs_is_trans, rhs_is_trans, scaling_mode, lhs_axis_boundary, rhs_axis_boundary, lhs_left_size, lhs_right_size, rhs_left_size, rhs_right_size] = config; - NVTE_CHECK(scaling_mode == JAXX_Scaling_Mode::NO_SCALING, - "Only non-quantized grouped GEMM is supported in current implementation."); + NVTE_CHECK(scaling_mode == JAXX_Scaling_Mode::NO_SCALING || + scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING, + "Only NO_SCALING and MXFP8_1D_SCALING are supported in the V2 grouped GEMM."); + + const bool is_mxfp8 = scaling_mode == JAXX_Scaling_Mode::MXFP8_1D_SCALING; size_t num_gemms = grouped_gemm_num_gemms(lhs_first_dims, lhs_last_dims, rhs_first_dims, rhs_last_dims, out_first_dims, out_last_dims, alpha); // Workspaces. + // V2 GEMM receives scale_inv already swizzled by nvte_group_quantize (V2 grouped quantize + // fuses the swizzle). No extra sinv reservation is needed; the full cublas_workspace is + // available for cuBLAS. auto setup_workspace_ptr = reinterpret_cast(setup_workspace->untyped_data()); auto cublas_workspace_ptr = reinterpret_cast(cublas_workspace->untyped_data()); cublas_workspace_ptr = move_ptr_to_next_256B_aligned(cublas_workspace_ptr); @@ -783,14 +896,39 @@ Error_Type GroupedGemmV2FFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Ty auto *int64_base = reinterpret_cast(int64_workspace->untyped_data()); size_t int64_capacity = int64_workspace->element_count() / sizeof(int64_t); size_t int64_offset = 0; + + // For MXFP8: in JAX, rhs=cuBLAS_A, lhs=cuBLAS_B (swapped). + // Colwise is needed when the operand's contracting dim is NOT the last dim in its layout. + const bool rhs_use_colwise = is_mxfp8 && !rhs_is_trans; + const bool lhs_use_colwise = is_mxfp8 && lhs_is_trans; + + // For MXFP8: scale_inv is already swizzled (pre-swizzled by V2 grouped quantize via + // nvte_group_quantize). Pass the buffers directly to make_grouped_tensor which sets + // with_gemm_swizzled_scales(true) for MXFP8 automatically. No re-swizzling needed. auto rhs_tensor = - make_grouped_tensor(rhs_data, rhs_first_dims, rhs_last_dims, int64_base, int64_capacity, - int64_offset, num_gemms, stream, rhs_axis_boundary); + is_mxfp8 + ? make_grouped_tensor(rhs_data, rhs_sinv, scaling_mode, rhs_use_colwise, rhs_first_dims, + rhs_last_dims, int64_base, int64_capacity, int64_offset, num_gemms, + stream, rhs_left_size, rhs_right_size) + : make_grouped_tensor(rhs_data, rhs_first_dims, rhs_last_dims, int64_base, int64_capacity, + int64_offset, num_gemms, stream, rhs_left_size, rhs_right_size); auto lhs_tensor = - make_grouped_tensor(lhs_data, lhs_first_dims, lhs_last_dims, int64_base, int64_capacity, - int64_offset, num_gemms, stream, lhs_axis_boundary); - auto out_tensor = make_grouped_tensor(*output, out_first_dims, out_last_dims, int64_base, - int64_capacity, int64_offset, num_gemms, stream); + is_mxfp8 + ? make_grouped_tensor(lhs_data, lhs_sinv, scaling_mode, lhs_use_colwise, lhs_first_dims, + lhs_last_dims, int64_base, int64_capacity, int64_offset, num_gemms, + stream, lhs_left_size, lhs_right_size) + : make_grouped_tensor(lhs_data, lhs_first_dims, lhs_last_dims, int64_base, int64_capacity, + int64_offset, num_gemms, stream, lhs_left_size, lhs_right_size); + + // Output stays NO_SCALING. Derive 2D shape from the output buffer's own dims using + // last-dim-as-columns convention (equivalent to axis_boundary=-1 in the old API). + auto out_dims = output->dimensions(); + NVTE_CHECK(out_dims.size() > 0, "output buffer must have at least 1 dimension"); + size_t out_left_size = product(out_dims, 0, out_dims.size() - 1); + size_t out_right_size = static_cast(out_dims[out_dims.size() - 1]); + auto out_tensor = + make_grouped_tensor(*output, out_first_dims, out_last_dims, int64_base, int64_capacity, + int64_offset, num_gemms, stream, out_left_size, out_right_size); auto [avg_m, avg_k_lhs] = grouped_gemm_avg_dims( lhs_first_dims, lhs_last_dims, {lhs_left_size, lhs_right_size}, num_gemms, lhs_is_trans); @@ -943,20 +1081,14 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type const size_t tensor_scaling_sinv_aligment = 16; const size_t mxfp8_scaling_sinv_alignment_padding = 256; auto workspace_size = workspace_total_size - workspace_alignment_padding; - if (is_mxfp8_scaling) { - // For MXFP8 swizzled scale_inv buffers, only the first pointer needs to be with 256B alignment padding. Later pointers are guaranteed to be 256-aligned as the scale_inv shapes are padded by 128x4. - workspace_size -= (lhs_sinv_size + rhs_sinv_size + 2 * mxfp8_scaling_sinv_alignment_padding); - } else if (is_tensor_scaling) { + if (is_tensor_scaling) { // For tensor scaling, each matrix has a single scale value, and all scales need to be aligned // by 16 bytes to meet the requirement of CUDA 12.9.1 and later. workspace_size -= tensor_scaling_sinv_aligment * (lhs_sinv_size + rhs_sinv_size); } workspace_size = workspace_size / num_streams; - auto swizzled_lhs_sinv_ptr = workspace_ptr + workspace_size * num_streams; - swizzled_lhs_sinv_ptr = move_ptr_to_next_256B_aligned(swizzled_lhs_sinv_ptr); - auto swizzled_rhs_sinv_ptr = swizzled_lhs_sinv_ptr + lhs_sinv_size; - swizzled_rhs_sinv_ptr = move_ptr_to_next_256B_aligned(swizzled_rhs_sinv_ptr); - auto lhs_scatter_aligned_ptr = swizzled_lhs_sinv_ptr; // Already 256B aligned + auto lhs_scatter_aligned_ptr = workspace_ptr + workspace_size * num_streams; + lhs_scatter_aligned_ptr = move_ptr_to_next_256B_aligned(lhs_scatter_aligned_ptr); auto rhs_scatter_aligned_ptr = lhs_scatter_aligned_ptr + num_gemms * tensor_scaling_sinv_aligment; size_t lhs_dtype_bytes = te_dtype_bytes(lhs_dtype); @@ -1050,8 +1182,6 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type // These lists are to keep the TensorWrapper objects alive std::vector lhs_wrapper_list; std::vector rhs_wrapper_list; - std::vector lhs_swizzle_wrapper_list; // For MXFP8 scale_inv swizzling - std::vector rhs_swizzle_wrapper_list; std::vector bias_wrapper_list; std::vector pre_gelu_wrapper_list; std::vector out_wrapper_list; @@ -1060,8 +1190,6 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type // These lists are the actual NVTETensor (void *) lists for multi-stream GEMM std::vector lhs_list; std::vector rhs_list; - std::vector lhs_swizzle_list; - std::vector rhs_swizzle_list; std::vector bias_list; std::vector pre_gelu_list; std::vector out_list; @@ -1134,13 +1262,8 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type else lhs_i.set_rowwise_scale_inv(lhs_sinv_vptr, lhs_sinv_dtype, tensor_scaling_sinv_shape); } else if (is_mxfp8_scaling) { - auto lhs_swizzle_i = TensorWrapper(get_nvte_scaling_mode(scaling_mode)); - auto rhs_swizzle_i = TensorWrapper(get_nvte_scaling_mode(scaling_mode)); - void *swizzled_lhs_sinv_vptr = static_cast(swizzled_lhs_sinv_ptr); - void *swizzled_rhs_sinv_vptr = static_cast(swizzled_rhs_sinv_ptr); - - // {lhs, rhs}_swizzle_i point to unswizzled scale_inv data as input, while {lhs, rhs}_i - // point to swizzled scale_inv data (store on workspace, only used for GEMM). + // MXFP8 scales are pre-swizzled by the quantize kernel (both V1 and V2), + // so we pass them directly to the GEMM without a separate swizzle pass. // Note: even if is_empty_gemm is true, sinv are still non-empty, need to move the pointers auto lhs_sinv_shape_i = get_block_scale_shape(scaling_mode, lhs_shape_i[0], lhs_shape_i[1], lhs_use_colwise); @@ -1149,32 +1272,17 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type lhs_sinv_size_i = lhs_sinv_shape_i[0] * lhs_sinv_shape_i[1]; rhs_sinv_size_i = rhs_sinv_shape_i[0] * rhs_sinv_shape_i[1]; if (lhs_use_colwise) { - lhs_swizzle_i.set_columnwise_data(lhs_vptr, lhs_dtype, lhs_shape_i); - lhs_swizzle_i.set_columnwise_scale_inv(lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); - lhs_i.set_columnwise_scale_inv(swizzled_lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); + lhs_i.set_columnwise_scale_inv(lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); } else { - lhs_swizzle_i.set_rowwise_data(lhs_vptr, lhs_dtype, lhs_shape_i); - lhs_swizzle_i.set_rowwise_scale_inv(lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); - lhs_i.set_rowwise_scale_inv(swizzled_lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); + lhs_i.set_rowwise_scale_inv(lhs_sinv_vptr, lhs_sinv_dtype, lhs_sinv_shape_i); } lhs_i.set_with_gemm_swizzled_scales(true); if (rhs_use_colwise) { - rhs_swizzle_i.set_columnwise_data(rhs_vptr, rhs_dtype, rhs_shape_i); - rhs_swizzle_i.set_columnwise_scale_inv(rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); - rhs_i.set_columnwise_scale_inv(swizzled_rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); + rhs_i.set_columnwise_scale_inv(rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); } else { - rhs_swizzle_i.set_rowwise_data(rhs_vptr, rhs_dtype, rhs_shape_i); - rhs_swizzle_i.set_rowwise_scale_inv(rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); - rhs_i.set_rowwise_scale_inv(swizzled_rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); + rhs_i.set_rowwise_scale_inv(rhs_sinv_vptr, rhs_sinv_dtype, rhs_sinv_shape_i); } rhs_i.set_with_gemm_swizzled_scales(true); - - if (!is_empty_gemm) { - lhs_swizzle_wrapper_list.push_back(std::move(lhs_swizzle_i)); - rhs_swizzle_wrapper_list.push_back(std::move(rhs_swizzle_i)); - lhs_swizzle_list.push_back(lhs_swizzle_wrapper_list.back().data()); - rhs_swizzle_list.push_back(rhs_swizzle_wrapper_list.back().data()); - } } else { NVTE_CHECK(scaling_mode == JAXX_Scaling_Mode::NO_SCALING, "Unsupported scaling mode: ", static_cast(scaling_mode)); @@ -1192,10 +1300,6 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type rhs_sinv_ptr += rhs_sinv_size_i * rhs_sinv_dtype_bytes; lhs_sinv_total_size += lhs_sinv_size_i; rhs_sinv_total_size += rhs_sinv_size_i; - if (is_mxfp8_scaling) { - swizzled_lhs_sinv_ptr += lhs_sinv_size_i * lhs_sinv_dtype_bytes; - swizzled_rhs_sinv_ptr += rhs_sinv_size_i * rhs_sinv_dtype_bytes; - } } if (has_bias) bias_ptr += n * bias_dtype_bytes; @@ -1236,18 +1340,6 @@ Error_Type GroupedGemmFFI(cudaStream_t stream, Buffer_Type lhs_data, Buffer_Type size_t num_non_empty_gemms = lhs_list.size(); - if (is_mxfp8_scaling) { - for (int i = 0; i < num_non_empty_gemms; i++) { - // The i-th GEMM will use the (i % num_streams)-th stream to compute, - // use the same stream to swizzle the scaling factors to make sure that - // the swizzling is done before the GEMM computation starts. - int stream_id = i % num_streams; - cudaStream_t stream_i = nvte_get_compute_stream(stream_id); - nvte_swizzle_scaling_factors(lhs_swizzle_list[i], lhs_list[i], stream_i); - nvte_swizzle_scaling_factors(rhs_swizzle_list[i], rhs_list[i], stream_i); - } - } - // Launch zero-out kernels before the GEMM calls to use the sync in the multi-stream GEMM size_t num_zero_outs = zero_out_dptr_list.size(); for (int i = 0; i < num_zero_outs; i++) { diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index 28cb39b5d1..e3bc122403 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -33,6 +33,7 @@ pybind11::dict Registrations() { // Quantization dict["te_dbias_quantize_ffi"] = EncapsulateFFI(DBiasQuantizeHandler); dict["te_grouped_quantize_ffi"] = EncapsulateFFI(GroupedQuantizeHandler); + dict["te_grouped_quantize_v2_ffi"] = EncapsulateFFI(GroupedQuantizeV2Handler); dict["te_dequantize_ffi"] = EncapsulateFFI(DequantizeHandler); // Softmax diff --git a/transformer_engine/jax/csrc/extensions/quantization.cpp b/transformer_engine/jax/csrc/extensions/quantization.cpp index c5a766f7f2..650139a61c 100644 --- a/transformer_engine/jax/csrc/extensions/quantization.cpp +++ b/transformer_engine/jax/csrc/extensions/quantization.cpp @@ -9,6 +9,7 @@ #include "../extensions.h" #include "transformer_engine/cast.h" +#include "transformer_engine/gemm.h" #include "transformer_engine/hadamard_transform.h" #include "transformer_engine/recipe.h" #include "transformer_engine/transformer_engine.h" @@ -318,8 +319,8 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty Buffer_Type group_sizes, Result_Type outputs, Result_Type colwise_outputs, Result_Type scale_invs, Result_Type colwise_scale_invs, Result_Type amaxs, - JAXX_Scaling_Mode scaling_mode, JAXX_Quantize_Layout quantize_layout, - int64_t flatten_axis) { + Result_Type _unused, JAXX_Scaling_Mode scaling_mode, + JAXX_Quantize_Layout quantize_layout, int64_t flatten_axis) { NVTE_CHECK(scaling_mode != JAXX_Scaling_Mode::NO_SCALING, "Unsupported scaling mode: ", static_cast(scaling_mode)); @@ -451,6 +452,12 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty } } + // For MXFP8, produce pre-swizzled scales so the GEMM can consume them directly + // without a separate swizzle pass. + if (is_mxfp8_scaling) { + out_i.set_with_gemm_swizzled_scales(true); + } + input_holders.push_back(std::move(inp_i)); output_holders.push_back(std::move(out_i)); @@ -479,20 +486,154 @@ Error_Type GroupedQuantizeFFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Ty return ffi_with_cuda_error_check(); } -XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedQuantizeHandler, GroupedQuantizeFFI, +XLA_FFI_DEFINE_HANDLER_SYMBOL( + GroupedQuantizeHandler, GroupedQuantizeFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // input + .Arg() // scale + .Arg() // group_sizes + .Ret() // output + .Ret() // colwise output + .Ret() // scale_inv + .Ret() // scale_inv colwise + .Ret() // amax + .Ret() // unused (for compatibility with V2 interface) + .Attr("scaling_mode") + .Attr("q_layout") + .Attr("flatten_axis")); + +Error_Type GroupedQuantizeV2FFI(cudaStream_t stream, Buffer_Type inputs, Buffer_Type scale_unused, + Buffer_Type group_sizes, Result_Type rowwise_out, + Result_Type colwise_out, Result_Type rowwise_sinv, + Result_Type colwise_sinv, Result_Type updated_amaxs, + Result_Type int64_workspace, JAXX_Quantize_Layout quantize_layout, + int64_t flatten_axis) { + (void)scale_unused; // scale is unused for MXFP8; accepted to match V1 input arity + auto in_dtype = convert_ffi_datatype_to_te_dtype(inputs.element_type()); + auto out_dtype = convert_ffi_datatype_to_te_dtype(rowwise_out->element_type()); + auto sinv_dtype = convert_ffi_datatype_to_te_dtype(rowwise_sinv->element_type()); + + NVTE_CHECK(is_fp8_dtype(out_dtype), "Output datatype must be FP8 for GroupedQuantizeV2."); + NVTE_CHECK(sinv_dtype == DType::kFloat8E8M0, + "scale_inv must be E8M0 for MXFP8 grouped quantize."); + + auto input_dims = inputs.dimensions(); + int64_t input_ndim = input_dims.size(); + if (flatten_axis < 0) flatten_axis += input_ndim; + NVTE_CHECK(flatten_axis < input_ndim && flatten_axis > 0, "flatten_axis is out of bounds!"); + + auto m = product(input_dims, 0, flatten_axis); + auto n = product(input_dims, flatten_axis, input_ndim); + size_t n_groups = group_sizes.dimensions()[0]; + + // Workspace layout (CUDA-graph safe, all device-side): + // int64_ptr[0 .. n_groups-1] : per-group ROW counts (int64) + // int64_ptr[n_groups .. 2*n_groups] : exclusive prefix-sum offsets (n_groups+1 values) + auto *int64_ptr = reinterpret_cast(int64_workspace->untyped_data()); + auto *offsets_ptr_out = int64_ptr + n_groups; // n_groups+1 values follow group_sizes + + // non_group_m handles multi-dim tensors (e.g., kernel shape G×K×N with flatten_axis=2): + // group_sizes[i] counts "slices" along the outermost group axis (e.g., 1 per expert), + // while the kernel expects actual ROW counts (e.g., K rows per expert). + // non_group_m = product(input_dims[1..flatten_axis)) converts slice→row count. + // For the lhs case (shape M×K, flatten_axis=1), non_group_m=1 (no-op). + int64_t non_group_m = + (flatten_axis > 1) ? product(input_dims, 1, static_cast(flatten_axis)) : 1; + + // Convert int32 group_sizes to int64 row counts on device (CUDA-graph safe, no D2H). + nvte_convert_int32_to_int64_with_multiplier( + reinterpret_cast(group_sizes.untyped_data()), int64_ptr, n_groups, + non_group_m, stream); + + // Compute exclusive prefix-sum offsets on device (CUDA-graph safe, no D2H). + nvte_compute_grouped_tensor_offsets(int64_ptr, offsets_ptr_out, n_groups, static_cast(n), + stream); + + NVTEShape data_shape{}; + data_shape.data[0] = m; + data_shape.data[1] = n; + data_shape.ndim = 2; + + NVTEShape sz_shape{}; + sz_shape.ndim = 1; + sz_shape.data[0] = n_groups; + + // Offsets tensor has n_groups+1 elements (exclusive prefix sums with sentinel). + NVTEShape offsets_shape{}; + offsets_shape.ndim = 1; + offsets_shape.data[0] = n_groups + 1; + + // Build input grouped tensor (plain float data, no quantization on the input side). + GroupedTensorWrapper in_grouped(n_groups, data_shape, + get_nvte_scaling_mode(JAXX_Scaling_Mode::NO_SCALING)); + in_grouped + .set_rowwise_data(reinterpret_cast(inputs.untyped_data()), in_dtype, data_shape) + .set_first_dims(reinterpret_cast(int64_ptr), DType::kInt64, sz_shape) + .set_tensor_offsets(reinterpret_cast(offsets_ptr_out), DType::kInt64, offsets_shape); + + // Build output grouped tensor. + GroupedTensorWrapper out_grouped(n_groups, data_shape, + get_nvte_scaling_mode(JAXX_Scaling_Mode::MXFP8_1D_SCALING)); + out_grouped.set_first_dims(reinterpret_cast(int64_ptr), DType::kInt64, sz_shape) + .set_tensor_offsets(reinterpret_cast(offsets_ptr_out), DType::kInt64, offsets_shape); + + // Rowwise output data + scale_inv. + if (is_quantize_rowwise(quantize_layout)) { + NVTEShape rw_sinv_shape{}; + rw_sinv_shape.ndim = 2; + rw_sinv_shape.data[0] = m; + rw_sinv_shape.data[1] = n / 32; // MXFP8 block size = 32 + out_grouped.set_rowwise_data(rowwise_out->untyped_data(), out_dtype, data_shape) + .set_rowwise_scale_inv(rowwise_sinv->untyped_data(), sinv_dtype, rw_sinv_shape); + } + + // Colwise output data + scale_inv. + if (is_quantize_colwise(quantize_layout)) { + NVTEShape cw_sinv_shape{}; + cw_sinv_shape.ndim = 2; + cw_sinv_shape.data[0] = m / 32; // MXFP8 block size = 32 + cw_sinv_shape.data[1] = n; + out_grouped.set_columnwise_data(colwise_out->untyped_data(), out_dtype, data_shape) + .set_columnwise_scale_inv(colwise_sinv->untyped_data(), sinv_dtype, cw_sinv_shape); + } + + // Zero-initialize scale_inv buffers (mirrors V1 behaviour for MXFP8). + size_t total_rowwise_sinv_size = + is_quantize_rowwise(quantize_layout) ? product(rowwise_sinv->dimensions()) : 0; + size_t total_colwise_sinv_size = + is_quantize_colwise(quantize_layout) ? product(colwise_sinv->dimensions()) : 0; + if (total_rowwise_sinv_size > 0) + nvte_memset(rowwise_sinv->untyped_data(), 0, total_rowwise_sinv_size, stream); + if (total_colwise_sinv_size > 0) + nvte_memset(colwise_sinv->untyped_data(), 0, total_colwise_sinv_size, stream); + + // V2 grouped quantize is always paired with V2 grouped GEMM, which expects + // scale_inv in GEMM-swizzled layout. Enable the fused swizzle so the kernel + // writes scales in the layout the GEMM will consume directly. + out_grouped.set_with_gemm_swizzled_scales(true); + + QuantizationConfigWrapper quant_config{}; + nvte_group_quantize(in_grouped.data(), out_grouped.data(), quant_config, stream); + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(GroupedQuantizeV2Handler, GroupedQuantizeV2FFI, FFI::Bind() .Ctx() // stream - .Arg() // input - .Arg() // scale - .Arg() // group_sizes - .Ret() // output - .Ret() // colwise output - .Ret() // scale_inv - .Ret() // scale_inv colwise - .Ret() // amax - .Attr("scaling_mode") + .Arg() // inputs + .Arg() // scale (unused, for input arity match) + .Arg() // group_sizes (int32) + .Ret() // rowwise_out + .Ret() // colwise_out + .Ret() // rowwise_sinv + .Ret() // colwise_sinv + .Ret() // updated_amaxs + .Ret() // int64_workspace .Attr("q_layout") - .Attr("flatten_axis")); + .Attr("flatten_axis"), + FFI_CudaGraph_Traits); } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/flax/module.py b/transformer_engine/jax/flax/module.py index 31ce6e72e9..17c9a242f0 100644 --- a/transformer_engine/jax/flax/module.py +++ b/transformer_engine/jax/flax/module.py @@ -16,6 +16,9 @@ from jax import random as jax_random from jax.ad_checkpoint import checkpoint_name +from transformer_engine.common.recipe import ( + MXFP8BlockScaling, +) from ..dense import dense, grouped_dense @@ -1358,7 +1361,12 @@ def kernel_1_init(key, num_kernels, stack_axis, *init_args): return out, ln_output # Output, layer_norm_output -def wrap_function_in_te_state_module(f, quantization_recipe, name: Optional[str] = None): +def wrap_function_in_te_state_module( + f, + quantization_recipe, + name: Optional[str] = None, + quantization_checkpoint_name: Optional[str] = None, +): """Wraps the given function `f` to support TransformerEngine quantization. This method does a couple things: @@ -1386,6 +1394,7 @@ def generate_quantizer_set(self, postfix: str = "", n_groups: int = None): return super().generate_quantizer_set( postfix=postfix, variable_collection=OVERWRITE_WITH_GRADIENT, + quantization_checkpoint_name=quantization_checkpoint_name, fp8_recipe=quantization_recipe, n_groups=n_groups, ) @@ -1443,10 +1452,15 @@ def te_dot_general(generate_quantizer_set, x, kernel, dims, **kwargs): return wrap_function_in_te_state_module(te_dot_general, quantization_recipe, "dot_general") -def make_grouped_dense_cls(quantization_recipe): +def make_grouped_dense_cls(quantization_recipe, quantization_checkpoint_name: Optional[str] = None): """Creates a grouped dense (grouped GEMM) instance for use with TE state module.""" if quantization_recipe is not None: - raise ValueError("Ragged dot grouped GEMM does not support quantization yet") + allowed_grouped_gemm_recipes = [MXFP8BlockScaling] + assert any(isinstance(quantization_recipe, r) for r in allowed_grouped_gemm_recipes), ( + "Only the following quantization recipes are supported for grouped GEMM or `None` for" + f" BF16 without quantization: {allowed_grouped_gemm_recipes}. Got" + f" {type(quantization_recipe)}." + ) def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwargs): del kwargs # Unused @@ -1463,5 +1477,8 @@ def te_grouped_dot_general(generate_quantizer_set, x, kernel, group_sizes, **kwa return out return wrap_function_in_te_state_module( - te_grouped_dot_general, quantization_recipe, "ragged_dot" + te_grouped_dot_general, + quantization_recipe, + "ragged_dot", + quantization_checkpoint_name=quantization_checkpoint_name, )() diff --git a/transformer_engine/jax/quantize/dequantizer.py b/transformer_engine/jax/quantize/dequantizer.py index 5abb2e74df..ca44c2e4af 100644 --- a/transformer_engine/jax/quantize/dequantizer.py +++ b/transformer_engine/jax/quantize/dequantizer.py @@ -263,7 +263,37 @@ def dequantize(scaled_tensor): } -@staticmethod +def _unswizzle_mxfp8_grouped_scale(scale_inv_flat, padded_scale_2d, is_colwise): + """Un-swizzle MXFP8 GEMM-swizzled scale_inv back to plain layout. + + Both V1 and V2 MXFP8 grouped quantize produce scale_inv in a GEMM-swizzled + layout. This is the inverse of ``swizzled_scale`` in ``gemm.py``. + + The swizzle pattern (for rowwise) is: + reshape(R//128, 4, 32, C//4, 4) → transpose(0,3,2,1,4) → reshape(R, C) + The inverse is: + reshape(R//128, C//4, 32, 4, 4) → transpose(0,3,2,1,4) → reshape(R, C) + + For colwise the swizzle is applied to the transposed scale, so the inverse + must un-transpose as well. + """ + if is_colwise: + # Colwise forward: reshape_2d → transpose → swizzle_5d → reshape_original + # Inverse: reshape_to_5d → inverse_swizzle → reshape_to_transposed_2d → transpose + cols, rows = padded_scale_2d + scale_2d = scale_inv_flat.reshape(cols, rows) + # The swizzled data lives in the transposed (rows, cols) domain + reshaped = scale_2d.reshape(rows // 128, cols // 4, 32, 4, 4) + unswizzled = jnp.transpose(reshaped, (0, 3, 2, 1, 4)) + # Back to transposed 2D, then un-transpose + return jnp.transpose(unswizzled.reshape(rows, cols)) + + rows, cols = padded_scale_2d + reshaped = scale_inv_flat.reshape(rows // 128, cols // 4, 32, 4, 4) + unswizzled = jnp.transpose(reshaped, (0, 3, 2, 1, 4)) + return unswizzled.reshape(rows, cols) + + def _grouped_dequantize(grouped_scaled_tensor): """Dequantize a grouped tensor. @@ -290,12 +320,13 @@ def _grouped_dequantize(grouped_scaled_tensor): flatten_axis = len(original_shape) + flatten_axis if flatten_axis < 0 else flatten_axis output = [] - # For transposed (colwise) tensors with ragged groups, the group dimension is the last - # axis of original_shape (e.g. original_shape = (N, M) with groups along M), while the - # non-group dimensions are all axes before it. For the uniform-groups case the group - # dimension stays at axis 0, so the existing axis-0 logic applies. + # When data_layout=="T" (colwise, transposed) and first_dims is set (ragged groups), the + # original_shape is stored transposed: the group (variable-size) axis is the LAST dimension + # rather than the first. Non-group dims are original_shape[:-1], not original_shape[1:]. is_transposed_ragged = ( - grouped_scaled_tensor.data_layout == "T" and group_sizes.size != original_shape[0] + grouped_scaled_tensor.data_layout == "T" + and grouped_scaled_tensor.first_dims is not None + and grouped_scaled_tensor.first_dims.size > 0 ) if is_transposed_ragged: non_group_shape = original_shape[:-1] @@ -308,7 +339,7 @@ def _grouped_dequantize(grouped_scaled_tensor): scale_inv_ptr = 0 for i, data_i in enumerate(data): if is_transposed_ragged: - data_shape_i = (*non_group_shape, group_sizes[i]) + data_shape_i = (*non_group_shape, int(group_sizes[i])) else: data_shape_i = ( group_sizes[i], @@ -330,24 +361,49 @@ def _grouped_dequantize(grouped_scaled_tensor): is_padded=False, flatten_axis=flatten_axis, ) - scale_inv_i = scale_inv[ - scale_inv_ptr : scale_inv_ptr + math.prod(padded_scale_shape_i) - ].reshape(padded_scale_shape_i) - scale_inv_i = jax.lax.slice( - scale_inv_i, [0] * len(unpadded_scale_shape_i), unpadded_scale_shape_i - ) + scale_inv_i = scale_inv[scale_inv_ptr : scale_inv_ptr + math.prod(padded_scale_shape_i)] + # MXFP8 grouped quantize (both V1 and V2) always produces GEMM-swizzled + # scales. Detect by scaling_mode (not pre_swizzled, which is only set for V2 + # to maintain pytree compatibility with the GEMM path). + is_colwise = grouped_scaled_tensor.is_colwise + needs_unswizzle = scaling_mode == ScalingMode.MXFP8_1D_SCALING + if needs_unswizzle: + flat_data_2d = ( + math.prod(data_shape_i[:flatten_axis]), + math.prod(data_shape_i[flatten_axis:]), + ) + padded_2d = scaling_mode.get_scale_shape( + flat_data_2d, is_colwise=is_colwise, is_padded=True, flatten_axis=1 + ) + unpadded_2d = scaling_mode.get_scale_shape( + flat_data_2d, is_colwise=is_colwise, is_padded=False, flatten_axis=1 + ) + scale_inv_i = _unswizzle_mxfp8_grouped_scale(scale_inv_i, padded_2d, is_colwise) + scale_inv_i = jax.lax.slice(scale_inv_i, [0, 0], list(unpadded_2d)) + else: + scale_inv_i = scale_inv_i.reshape(padded_scale_shape_i) + scale_inv_i = jax.lax.slice( + scale_inv_i, [0] * len(unpadded_scale_shape_i), unpadded_scale_shape_i + ) dequantizer_type = ScalingModeToDequantizerMap.get(grouped_scaled_tensor.scaling_mode) if len(data_i) == 0: out_i = [] else: + # _dequantize_func is designed for 2D-flattened data. Flatten the + # per-group shape to 2D, dequantize, then reshape back. + flat_shape_i = ( + math.prod(data_shape_i[:flatten_axis]), + math.prod(data_shape_i[flatten_axis:]), + ) out_i = dequantizer_type._dequantize_func( - data_i.reshape(data_shape_i), + data_i.reshape(flat_shape_i), scale_inv_i, grouped_scaled_tensor.dq_dtype, scaling_mode=grouped_scaled_tensor.scaling_mode, is_colwise=grouped_scaled_tensor.is_colwise, - flatten_axis=grouped_scaled_tensor.flatten_axis, + flatten_axis=1, ) + out_i = out_i.reshape(data_shape_i) output.append(out_i) scale_inv_ptr += math.prod(padded_scale_shape_i) diff --git a/transformer_engine/jax/quantize/tensor.py b/transformer_engine/jax/quantize/tensor.py index b1f49dacdc..c5ad0451fd 100644 --- a/transformer_engine/jax/quantize/tensor.py +++ b/transformer_engine/jax/quantize/tensor.py @@ -369,11 +369,15 @@ class GroupedScaledTensor1x(ScaledTensor1x): first_dims: Per-group sizes of the first (row) 2D dim, or None if not ragged last_dims: Per-group sizes of the last (col) 2D dim, or None if not ragged original_shape: The original shape of the tensor before grouping + pre_swizzled: Whether the scale_inv is already swizzled for GEMM. True when produced + by V2 grouped quantize (nvte_group_quantize fuses the swizzle). The V2 grouped + GEMM FFI requires pre_swizzled=True for MXFP8 inputs and will not re-swizzle. """ first_dims: Optional[jnp.ndarray] last_dims: Optional[jnp.ndarray] original_shape: Tuple + pre_swizzled: bool = False def __init__( self, @@ -389,11 +393,13 @@ def __init__( data_layout, flatten_axis, original_shape, + pre_swizzled=False, ): self.flatten_axis = flatten_axis self.first_dims = first_dims self.last_dims = last_dims self.original_shape = original_shape + self.pre_swizzled = pre_swizzled # TODO(Phuong):Handle RHT for grouped quantization once grouped quantization supports NVFP4 super().__init__( data=data, @@ -408,6 +414,18 @@ def __init__( has_rht_applied=False, ) + @property + def group_sizes(self) -> jnp.ndarray: + """Per-group sizes along the group axis. + + When first_dims is set (ragged groups), returns first_dims. + When first_dims is None (equal-sized groups), returns an array of ones with + length equal to the number of groups. + """ + if self.first_dims is not None and self.first_dims.size > 0: + return self.first_dims + return jnp.ones((self.original_shape[0],), dtype=jnp.int32) + def __post_init__(self): assert self.scale_inv.ndim == 1, "Only support flattened scale_inv" assert self.data.ndim == 1, "Only support flattened data" @@ -456,6 +474,7 @@ def tree_flatten(self): self.data_layout, self.flatten_axis, self.original_shape, + self.pre_swizzled, ) return (children, aux_data) @@ -653,6 +672,7 @@ def create_1x( last_dims=None, original_shape=None, has_rht_applied=False, + pre_swizzled=False, ): """Creates a single-scale quantized tensor. @@ -722,6 +742,7 @@ def create_1x( first_dims=first_dims, last_dims=last_dims, original_shape=original_shape, + pre_swizzled=pre_swizzled, ) # Handling attrs of transposed tensors @@ -759,6 +780,7 @@ def create_2x( original_shape=None, rowwise_has_rht_applied=False, colwise_has_rht_applied=False, + pre_swizzled=False, ): """Creates a double-scale quantized tensor. @@ -800,6 +822,7 @@ def create_2x( last_dims=last_dims, original_shape=original_shape, has_rht_applied=rowwise_has_rht_applied, + pre_swizzled=pre_swizzled, ) colwise_tensor = ScaledTensorFactory.create_1x( colwise_data, @@ -814,6 +837,7 @@ def create_2x( last_dims=last_dims, original_shape=original_shape, has_rht_applied=colwise_has_rht_applied, + pre_swizzled=pre_swizzled, ) return ScaledTensor2x(rowwise_tensor, colwise_tensor) @@ -835,6 +859,7 @@ def create( original_shape: Tuple[int] = None, rowwise_has_rht_applied: bool = False, colwise_has_rht_applied: bool = False, + pre_swizzled: bool = False, ): """Creates a scaled tensor based on the quantization axis. @@ -853,6 +878,7 @@ def create( original_shape: The original shape of the tensor before grouping (default: None) rowwise_has_rht_applied: Whether the row-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) colwise_has_rht_applied: Whether the col-wise tensor uses the Randomized Hadamard Transform (RHT) (default: False) + pre_swizzled: Whether scale_inv is already swizzled (produced by V2 grouped quantize). Returns: Either a ScaledTensor1x or ScaledTensor2x instance depending on q_layout @@ -876,6 +902,7 @@ def create( original_shape=original_shape, rowwise_has_rht_applied=rowwise_has_rht_applied, colwise_has_rht_applied=colwise_has_rht_applied, + pre_swizzled=pre_swizzled, ) if q_layout.is_colwise_only: @@ -892,6 +919,7 @@ def create( last_dims=last_dims, original_shape=original_shape, has_rht_applied=colwise_has_rht_applied, + pre_swizzled=pre_swizzled, ) return ScaledTensorFactory.create_1x( @@ -907,6 +935,7 @@ def create( last_dims=last_dims, original_shape=original_shape, has_rht_applied=rowwise_has_rht_applied, + pre_swizzled=pre_swizzled, ) From 52d6e8bbe7b8db11c1d2f4d2f9fe44b6e3afd04f Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 14 Apr 2026 19:13:35 -0700 Subject: [PATCH 346/521] Test Fused MOE with padded tokens (#2880) * test padded tokens Signed-off-by: Varun Thumbe * Update tests/pytorch/test_fusible_ops.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vthumbe1503 --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- tests/pytorch/test_fusible_ops.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index a5c071074c..0dfa8b5f45 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3686,6 +3686,7 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( hidden_size: int = 256, split_alignment: int = 256, glu_interleave_size: int = 32, + token_padding: int = 2048, ) -> None: """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" @@ -3703,8 +3704,8 @@ def test_grouped_mlp_cuda_graph_safe_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) - + # Pad the input tokens to validate the sync-free MOE + in_shape = (split_sizes.sum().item() + token_padding, hidden_size) recipe = make_recipe("mxfp8") with te.quantized_model_init(enabled=True, recipe=recipe): fc1 = te_ops.GroupedLinear( From 17aa2e4fc0c9e6e10944804d9bbc6ec7ad118c17 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, 15 Apr 2026 10:43:06 +0200 Subject: [PATCH 347/521] [PyTorch] [torch.compile] transformer_engine.pytorch.autocast suport inside torch.compile (#2759) * Improve torch.compile behavior around FP8 autocast. Move FP8 global state onto an instance so Dynamo can trace autocast state updates, explicitly reject DelayedScaling under torch.compile, and add toy compile tests that keep TE forward/backward opaque while covering supported recipes. Signed-off-by: Pawel Gadzinski * Remove temporary global state experiment tests. Drop the standalone global dict and dataclass mutation experiments now that the torch.compile regression coverage lives in the focused autocast test file. Signed-off-by: Pawel Gadzinski * Clean up FP8 global state naming. Use compiler constant-result wrappers for support checks and rename the module-level FP8 singleton to `_FP8_GLOBAL_STATE` for clearer semantics. Signed-off-by: Pawel Gadzinski * Minimize FP8 global state diff. Restore the FP8 naming and remove extra state access helpers so the torch.compile changes stay focused on the instance-backed global state. Signed-off-by: Pawel Gadzinski * Remove unused FP8 state fields. Drop stale availability fields from FP8GlobalState now that support checks use module-level cached results instead of manager state. Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Simplify torch.compile autocast tests Replace custom-op-based ToyLinear with a minimal version using F.linear. Add test_autocast_sanity (parametrized over all recipes including NVFP4) and test_autocast_nested_sanity with CustomRecipes. Both verify fullgraph=True compilation without graph breaks. 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 * Add test for DelayedScaling rejection under torch.compile Verify that te.autocast(recipe=DelayedScaling(), enabled=True) raises a clear RuntimeError when used inside torch.compile. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Use content-based autocast key with id() for group Use str(recipe) for content-based recipe keying (avoids unbounded growth when identical recipes are constructed inline) and id(group) for process group identity (same semantics as the old hash(group) which was id-based). 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 * Rewrite torch.compile tests with opaque value-type quantizers Replace custom_op-based approach with torch.library.define/impl/register_fake using get_opaque_type_name() in the schema, which allows Inductor to properly handle opaque value types. Add ToyQuantizer as an opaque value-type wrapper around Float8CurrentScalingQuantizer with proper __eq__/__hash__/__fx_repr__. test_autocast_nested_custom validates that nested te.autocast with 3 distinct CustomRecipe instances passes the correct quantizers in both forward and backward. test_autocast_sanity is a smoke test for all hardware-supported built-in recipes. 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 * apply suggestions Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix 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> --- tests/pytorch/test_torch_compile.py | 324 ++++++++++++++++++ .../dot_product_attention.py | 4 +- transformer_engine/pytorch/distributed.py | 9 +- transformer_engine/pytorch/graph.py | 11 +- transformer_engine/pytorch/module/base.py | 13 +- .../pytorch/module/layernorm_linear.py | 9 +- .../pytorch/module/layernorm_mlp.py | 21 +- transformer_engine/pytorch/module/linear.py | 9 +- transformer_engine/pytorch/ops/op.py | 13 +- transformer_engine/pytorch/quantization.py | 316 +++++++++-------- 10 files changed, 553 insertions(+), 176 deletions(-) create mode 100644 tests/pytorch/test_torch_compile.py diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py new file mode 100644 index 0000000000..9d0ed79888 --- /dev/null +++ b/tests/pytorch/test_torch_compile.py @@ -0,0 +1,324 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import abc + +import pytest +import torch + +try: + from torch._opaque_base import OpaqueBaseMeta + from torch._library.opaque_object import ( + get_opaque_type_name, + register_opaque_type, + MemberType, + ) + + _opaque_available = True +except ImportError: + _opaque_available = False + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common import recipe +from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx +from transformer_engine.pytorch.module.base import TransformerEngineBaseModule +from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear +from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer +from transformer_engine.pytorch import ( + is_fp8_available, + is_mxfp8_available, + is_fp8_block_scaling_available, + is_nvfp4_available, +) + +fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) +fp8_block_scaling_available = is_fp8_block_scaling_available() +nvfp4_available = is_nvfp4_available() + +_all_recipes: list = [] +if fp8_available: + _all_recipes.append(recipe.Float8CurrentScaling()) +if fp8_block_scaling_available: + _all_recipes.append(recipe.Float8BlockScaling()) +if mxfp8_available: + _all_recipes.append(recipe.MXFP8BlockScaling()) +if nvfp4_available: + _all_recipes.append(recipe.NVFP4BlockScaling()) + + +# --------------------------------------------------------------------------- +# ToyQuantizer – opaque value-type quantizer for torch.compile +# (requires torch opaque object support, not available in older PyTorch) +# --------------------------------------------------------------------------- + +if _opaque_available: + + class _ToyQuantizerMeta(OpaqueBaseMeta, abc.ABCMeta): + pass + + class ToyQuantizer(Float8CurrentScalingQuantizer, metaclass=_ToyQuantizerMeta): + """Quantizer with a string tag, registered as an + opaque value type so torch.compile can treat it as a baked-in constant.""" + + def __init__(self, tag: str): + super().__init__(fp8_dtype=tex.DType.kFloat8E4M3, device=torch.device("cuda")) + self.tag = tag + + def __eq__(self, other): + if not isinstance(other, ToyQuantizer): + return NotImplemented + return self.tag == other.tag and self.dtype == other.dtype + + def __hash__(self): + return hash((type(self), self.tag, self.dtype)) + + def __fx_repr__(self): + return ( + f"ToyQuantizer(tag={self.tag!r})", + {"ToyQuantizer": ToyQuantizer}, + ) + + register_opaque_type( + ToyQuantizer, + typ="value", + members={ + "__setattr__": MemberType.USE_REAL, + "set_usage": MemberType.USE_REAL, + }, + ) + + _Q = get_opaque_type_name(ToyQuantizer) + + def _make_qfactory(tag: str): + """Return a qfactory that produces ToyQuantizer instances tagged with *tag*.""" + + def qfactory(role: str): + return ToyQuantizer(tag=f"{tag}:{role}") + + return qfactory + + # --------------------------------------------------------------------------- + # ToyLinear – minimal TE module backed by BasicLinear functional ops + # --------------------------------------------------------------------------- + + class ToyLinear(TransformerEngineBaseModule): + """Minimal TE-compatible linear module used for torch.compile tests.""" + + def __init__( + self, + in_features: int, + out_features: int, + device: str = "cuda", + dtype: torch.dtype = torch.bfloat16, + ) -> None: + super().__init__() + self.in_features = in_features + self.out_features = out_features + self.weight = torch.nn.Parameter( + torch.empty(out_features, in_features, dtype=dtype, device=device) + ) + torch.nn.init.normal_(self.weight) + + def _get_weight_tensors(self): + return [self.weight] + + def _get_weight_quantizers(self): + if not self.fp8 and not self.fp8_calibration: + return [None] + weight_q = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_WEIGHT] + weight_q.internal = True + return [weight_q] + + def forward(self, inp: torch.Tensor) -> torch.Tensor: + inp = self.prepare_forward(inp, num_gemms=1) + try: + input_q = self.quantizers["scaling_fwd"][FP8FwdTensorIdx.GEMM1_INPUT] + input_q.internal = True + input_q.optimize_for_gemm = True + (weight_q,) = self._get_weight_quantizers() + grad_output_q = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_OUTPUT1] + grad_output_q.internal = True + grad_output_q.optimize_for_gemm = True + + return torch.ops.test_te.toy_linear( + inp, + self.weight, + input_q, + weight_q, + grad_output_q, + ) + finally: + self.end_forward() + + # --------------------------------------------------------------------------- + # Opaque custom ops (torch.library) + # --------------------------------------------------------------------------- + + _lib = torch.library.Library("test_te", "DEF") + + _lib.define( + f"toy_linear(Tensor inp, Tensor weight, {_Q} input_q, {_Q} weight_q, {_Q} grad_output_q)" + " -> Tensor" + ) + + _lib.define( + "toy_linear_backward(Tensor grad_output, Tensor inp, Tensor weight," + f" {_Q} grad_output_q) -> (Tensor, Tensor)" + ) + + last_fwd_quantizers: list[dict[str, "ToyQuantizer"]] = [] + last_bwd_quantizers: list[dict[str, "ToyQuantizer"]] = [] + + @torch.library.impl("test_te::toy_linear", "CompositeExplicitAutograd", lib=_lib) + def _toy_linear_fwd_impl(inp, weight, input_q, weight_q, grad_output_q): + last_fwd_quantizers.append( + { + "input_q": input_q, + "weight_q": weight_q, + "grad_output_q": grad_output_q, + } + ) + out, _, _ = BasicLinear._functional_forward( + input=inp, + weight=weight, + dtype=inp.dtype, + input_quantizer=input_q, + weight_quantizer=weight_q, + ) + return out + + @torch.library.register_fake("test_te::toy_linear", lib=_lib) + def _toy_linear_fwd_fake(inp, weight, input_q, weight_q, grad_output_q): + return inp @ weight.T + + def _toy_linear_setup_context(ctx, inputs, output): + inp, weight, _input_q, _weight_q, grad_output_q = inputs + ctx.save_for_backward(inp, weight) + ctx.grad_output_q = grad_output_q + + @torch.library.impl("test_te::toy_linear_backward", "CompositeExplicitAutograd", lib=_lib) + def _toy_linear_bwd_impl(grad_output, inp, weight, grad_output_q): + last_bwd_quantizers.append({"grad_output_q": grad_output_q}) + dx, dw = BasicLinear._functional_backward( + grad_output=grad_output, + input=inp, + weight=weight, + grad_output_quantizer=grad_output_q, + grad_input_quantizer=None, + ) + return dx, dw + + @torch.library.register_fake("test_te::toy_linear_backward", lib=_lib) + def _toy_linear_bwd_fake(grad_output, inp, weight, grad_output_q): + return torch.empty_like(inp), torch.empty_like(weight) + + def _toy_linear_backward(ctx, grad_output): + inp, weight = ctx.saved_tensors + dx, dw = torch.ops.test_te.toy_linear_backward( + grad_output, + inp, + weight, + ctx.grad_output_q, + ) + return dx, dw, None, None, None + + torch.library.register_autograd( + "test_te::toy_linear", + _toy_linear_backward, + setup_context=_toy_linear_setup_context, + lib=_lib, + ) + + +# --------------------------------------------------------------------------- +# Test +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available") +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +def test_autocast_nested_custom(): + """One ToyLinear model used under nested te.autocast with 3 distinct + CustomRecipe instances (each producing differently-tagged ToyQuantizers). + + Layout: + with autocast(recipe0): # outer + out = model(inp) + with autocast(recipe1): # nested inside outer + out = model(out) + with autocast(recipe2): # separate, after the nested pair + out = model(out) + + fullgraph=True makes torch.compile raise if any graph break occurs. + """ + dtype = torch.bfloat16 + device = "cuda" + + model = ToyLinear(32, 32, device=device, dtype=dtype) + + recipe0 = recipe.CustomRecipe(qfactory=_make_qfactory("R0")) + recipe1 = recipe.CustomRecipe(qfactory=_make_qfactory("R1")) + recipe2 = recipe.CustomRecipe(qfactory=_make_qfactory("R2")) + + inp = torch.randn(8, 32, dtype=dtype, device=device, requires_grad=True) + + def fn(inp): + with te.autocast(recipe=recipe0): + out = model(inp) + with te.autocast(recipe=recipe1): + out = model(out) + with te.autocast(recipe=recipe2): + out = model(out) + return out + + torch._dynamo.reset() + + compiled = torch.compile(fn, fullgraph=True) + last_fwd_quantizers.clear() + last_bwd_quantizers.clear() + + out = compiled(inp) + out.sum().backward() + + # Forward: 3 calls — R0, R1, R2 + assert len(last_fwd_quantizers) == 3, f"Expected 3 fwd calls, got {len(last_fwd_quantizers)}" + for i, tag in enumerate(["R0", "R1", "R2"]): + fq = last_fwd_quantizers[i] + assert fq["input_q"].tag.startswith(f"{tag}:"), f"fwd[{i}] input_q: {fq['input_q'].tag}" + assert fq["weight_q"].tag.startswith(f"{tag}:"), f"fwd[{i}] weight_q: {fq['weight_q'].tag}" + assert fq["grad_output_q"].tag.startswith( + f"{tag}:" + ), f"fwd[{i}] grad_output_q: {fq['grad_output_q'].tag}" + + # Backward: 3 calls — reverse order R2, R1, R0 + assert len(last_bwd_quantizers) == 3, f"Expected 3 bwd calls, got {len(last_bwd_quantizers)}" + for i, tag in enumerate(["R2", "R1", "R0"]): + bq = last_bwd_quantizers[i] + assert bq["grad_output_q"].tag.startswith( + f"{tag}:" + ), f"bwd[{i}] grad_output_q: {bq['grad_output_q'].tag}" + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +@pytest.mark.parametrize("fp8_recipe", _all_recipes, ids=lambda r: type(r).__name__) +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).""" + dtype = torch.bfloat16 + device = "cuda" + + model = torch.nn.Linear(32, 32, dtype=dtype, device=device) + inp = torch.randn(8, 32, dtype=dtype, device=device, requires_grad=True) + + def fn(inp): + with te.autocast(recipe=fp8_recipe): + return model(inp) + + torch._dynamo.reset() + compiled = torch.compile(fn, fullgraph=True) + + out = compiled(inp) + out.sum().backward() 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 2dc42be18a..588c708e10 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 @@ -704,7 +704,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: autocast_key = FP8GlobalStateManager.get_unique_autocast_key( fp8_recipe_dpa, fp8_group ) - FP8GlobalStateManager.autocast_arguments[autocast_key] = ( + FP8GlobalStateManager.quantization_state.autocast_arguments[autocast_key] = ( fp8_recipe_dpa, fp8_group, ) @@ -736,7 +736,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: autocast_key = FP8GlobalStateManager.get_unique_autocast_key( fp8_recipe_dpa, fp8_group ) - FP8GlobalStateManager.autocast_arguments[autocast_key] = ( + FP8GlobalStateManager.quantization_state.autocast_arguments[autocast_key] = ( fp8_recipe_dpa, fp8_group, ) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index b80e58fe20..a0d4ac3530 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -261,14 +261,11 @@ def __enter__(self): ) _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase + qstate = FP8GlobalStateManager.quantization_state if self.activation_recompute and not self.recompute_phase: - activation_recompute_forward._is_first_fp8_module.append( - FP8GlobalStateManager.IS_FIRST_FP8_MODULE - ) + activation_recompute_forward._is_first_fp8_module.append(qstate.is_first_fp8_module) if self.activation_recompute and self.recompute_phase: - FP8GlobalStateManager.IS_FIRST_FP8_MODULE = ( - activation_recompute_forward._is_first_fp8_module.pop(0) - ) + qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0) def __exit__(self, *exc_details): global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py index 86b8a4acf4..075db1394b 100644 --- a/transformer_engine/pytorch/graph.py +++ b/transformer_engine/pytorch/graph.py @@ -324,7 +324,12 @@ def _make_graphed_callables( if cache_quantized_params: # Initialize flag that controls FP8 weight updates - FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(False) + qstate = FP8GlobalStateManager.quantization_state + if qstate.skip_fp8_weight_update_tensor is None: + qstate.skip_fp8_weight_update_tensor = torch.empty( + 1, dtype=torch.float32, device="cuda" + ) + qstate.skip_fp8_weight_update_tensor.fill_(False) # Check callables for c in callables: @@ -836,7 +841,9 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i # Set flag for whether to update FP8 weight updates ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module() if ctx.is_first_module and skip_fp8_weight_update is not None: - FP8GlobalStateManager.set_skip_fp8_weight_update_tensor(skip_fp8_weight_update) + FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor.fill_( + skip_fp8_weight_update + ) ctx.cuda_graph_stream = cuda_graph_stream ctx.cuda_graph_event = cuda_graph_event # Copy values from new tensors into static tensors diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index a13eb0c7e6..5ca5572e06 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -725,20 +725,21 @@ def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> fwd_pos, fwd_key, bwd_pos, bwd_key = self.fp8_meta[ FP8GlobalStateManager.get_buffer_info() ] + qstate = FP8GlobalStateManager.quantization_state for pos, buffer_key in zip((fwd_pos, bwd_pos), (fwd_key, bwd_key)): - if buffer_key in FP8GlobalStateManager.global_amax_buffer: - if buffer_key not in FP8GlobalStateManager.global_amax_history_buffer: + if buffer_key in qstate.global_amax_buffer: + if buffer_key not in qstate.global_amax_history_buffer: raise RuntimeError( "TE internal error during amax history change: " f"buffer_key '{buffer_key}' found in global_amax_buffer " "but missing from global_amax_history_buffer" ) - FP8GlobalStateManager.global_amax_buffer[buffer_key][pos] = self.fp8_meta[ + qstate.global_amax_history_buffer[buffer_key][pos] = self.fp8_meta[ + meta_key + ].amax_history + qstate.global_amax_buffer[buffer_key][pos] = self.fp8_meta[ meta_key ].amax_history[0] - FP8GlobalStateManager.global_amax_history_buffer[buffer_key][pos] = ( - self.fp8_meta[meta_key].amax_history - ) def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 5361d7deda..8ceeaadfcc 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -531,10 +531,11 @@ def forward( ctx.normalization = normalization ctx.reduce_and_update_bwd_fp8_tensors = False if ctx.fp8 and requires_grad(inp, ln_weight, ln_bias, weight, bias): - _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE + qstate = FP8GlobalStateManager.quantization_state + _first_fp8_module = qstate.is_first_fp8_module ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() if in_fp8_activation_recompute_phase(): - FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module + qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store ctx.debug = debug @@ -1541,7 +1542,9 @@ def forward( debug = self.is_debug_iter() if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor() + skip_fp8_weight_update = ( + FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor + ) else: skip_fp8_weight_update = None if skip_fp8_weight_update is not None: diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index ca211daa08..6e6b11ecfe 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -250,9 +250,7 @@ def _forward( ctx.checkpoint = checkpoint if checkpoint: # save the state of autocast and quantizers for recomputation - ctx.autocast_state = ( - FP8GlobalStateManager.get_autocast_state() - ) # to restore autocast state during recomputation + ctx.autocast_state = FP8GlobalStateManager.get_autocast_state() if ( fp8 and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ @@ -852,10 +850,11 @@ def _forward( if ctx.fp8 and requires_grad( inp, ln_weight, ln_bias, fc1_weight, fc2_weight, fc1_bias, fc2_bias ): - _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE + qstate = FP8GlobalStateManager.quantization_state + _first_fp8_module = qstate.is_first_fp8_module ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() if in_fp8_activation_recompute_phase() or is_recomputation: - FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module + qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store if is_recomputation: # return the recomputed tensors @@ -923,10 +922,8 @@ def _recompute(ctx): # backward is not in autocast context, so we set the state here # we also have to set the quantizer states to what they were before the forward pass (only relevant for DelayedScaling recipe) - final_autocast_state = ( - FP8GlobalStateManager.get_autocast_state() - ) # get current autocast state - FP8GlobalStateManager.set_autocast_state(ctx.autocast_state) # set old autocast state + final_autocast_state = FP8GlobalStateManager.get_autocast_state() + FP8GlobalStateManager.set_autocast_state(ctx.autocast_state) if ( ctx.other_args["fp8"] and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" @@ -949,7 +946,7 @@ def _recompute(ctx): tuple(ctx.other_args.values()), ) - FP8GlobalStateManager.set_autocast_state(final_autocast_state) # restore autocast state + FP8GlobalStateManager.set_autocast_state(final_autocast_state) if ( ctx.other_args["fp8"] and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" @@ -2072,7 +2069,9 @@ def forward( debug = self.is_debug_iter() if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor() + skip_fp8_weight_update = ( + FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor + ) else: skip_fp8_weight_update = None if skip_fp8_weight_update is not None: diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index c85db15114..b57a2eb8d7 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -498,10 +498,11 @@ def forward( ctx.owns_input = saved_inputmat is not inp if ctx.fp8 and requires_grad(inp, weight, bias): - _first_fp8_module = FP8GlobalStateManager.IS_FIRST_FP8_MODULE + qstate = FP8GlobalStateManager.quantization_state + _first_fp8_module = qstate.is_first_fp8_module ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() if in_fp8_activation_recompute_phase(): - FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module + qstate.is_first_fp8_module = _first_fp8_module ctx.wgrad_store = wgrad_store # backward overrides @@ -1425,7 +1426,9 @@ def forward( debug = self.is_debug_iter() if FP8GlobalStateManager.fp8_graph_capturing(): - skip_fp8_weight_update = FP8GlobalStateManager.get_skip_fp8_weight_update_tensor() + skip_fp8_weight_update = ( + FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor + ) else: skip_fp8_weight_update = None if skip_fp8_weight_update is not None: diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index 54b3f00117..c5c8ea3463 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -322,14 +322,15 @@ def reset_recipe_state( pos, buffer_key = self._fp8_metas[mode][ FP8GlobalStateManager.get_buffer_info() ] - if buffer_key in FP8GlobalStateManager.global_amax_buffer: + qstate = FP8GlobalStateManager.quantization_state + if buffer_key in qstate.global_amax_buffer: assert ( - buffer_key in FP8GlobalStateManager.global_amax_history_buffer + buffer_key in qstate.global_amax_history_buffer ), "TE internal error during amax history change." - FP8GlobalStateManager.global_amax_buffer[buffer_key][pos] = ( - recipe_state.amax_history[0] - ) - FP8GlobalStateManager.global_amax_history_buffer[buffer_key][ + qstate.global_amax_buffer[buffer_key][pos] = recipe_state.amax_history[ + 0 + ] + qstate.global_amax_history_buffer[buffer_key][ pos ] = recipe_state.amax_history diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 47e6d5c8dc..9956fb77ec 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -7,9 +7,9 @@ import abc import itertools -import functools import warnings import os +from dataclasses import dataclass, field from contextlib import contextmanager from collections import deque from typing import Callable, List, Optional, Dict, Any, Tuple, Union @@ -44,8 +44,13 @@ ] -@functools.lru_cache(maxsize=None) -def check_fp8_support() -> Tuple[bool, str]: +_FP8_SUPPORT: Optional[Tuple[bool, str]] = None +_MXFP8_SUPPORT: Optional[Tuple[bool, str]] = None +_NVFP4_SUPPORT: Optional[Tuple[bool, str]] = None +_FP8_BLOCK_SCALING_SUPPORT: Optional[Tuple[bool, str]] = None + + +def _compute_fp8_support() -> Tuple[bool, str]: """Return if fp8 support is available""" if get_device_compute_capability() >= (9, 0): # hopper and above return True, "" @@ -58,8 +63,7 @@ def check_fp8_support() -> Tuple[bool, str]: return True, "" -@functools.lru_cache(maxsize=None) -def check_mxfp8_support() -> Tuple[bool, str]: +def _compute_mxfp8_support() -> Tuple[bool, str]: """Return if fp8 support is available""" if get_device_compute_capability() >= (12, 0): return False, "MXFP8 (for all gemm layouts) is not supported on 12.0+ architectures yet." @@ -68,16 +72,14 @@ def check_mxfp8_support() -> Tuple[bool, str]: return False, "Device compute capability 10.0 or higher required for MXFP8 execution." -@functools.lru_cache(maxsize=None) -def check_nvfp4_support() -> Tuple[bool, str]: +def _compute_nvfp4_support() -> Tuple[bool, str]: """Return if nvfp4 support is available""" if get_device_compute_capability() >= (10, 0): # blackwell and above return True, "" return False, "Device compute capability 10.0 or higher required for NVFP4 execution." -@functools.lru_cache(maxsize=None) -def check_fp8_block_scaling_support() -> Tuple[bool, str]: +def _compute_fp8_block_scaling_support() -> Tuple[bool, str]: """Return if fp8 block scaling support is available""" if get_device_compute_capability() >= (9, 0) and float(torch.version.cuda) >= 12.9: return True, "" @@ -87,8 +89,48 @@ def check_fp8_block_scaling_support() -> Tuple[bool, str]: ) +@torch.compiler.assume_constant_result +def check_fp8_support() -> Tuple[bool, str]: + """Return if fp8 support is available.""" + global _FP8_SUPPORT + if _FP8_SUPPORT is None: + _FP8_SUPPORT = _compute_fp8_support() + return _FP8_SUPPORT + + +@torch.compiler.assume_constant_result +def check_mxfp8_support() -> Tuple[bool, str]: + """Return if MXFP8 support is available.""" + global _MXFP8_SUPPORT + if _MXFP8_SUPPORT is None: + _MXFP8_SUPPORT = _compute_mxfp8_support() + return _MXFP8_SUPPORT + + +@torch.compiler.assume_constant_result +def check_nvfp4_support() -> Tuple[bool, str]: + """Return if NVFP4 support is available.""" + global _NVFP4_SUPPORT + if _NVFP4_SUPPORT is None: + _NVFP4_SUPPORT = _compute_nvfp4_support() + return _NVFP4_SUPPORT + + +@torch.compiler.assume_constant_result +def check_fp8_block_scaling_support() -> Tuple[bool, str]: + """Return if fp8 block scaling support is available.""" + global _FP8_BLOCK_SCALING_SUPPORT + if _FP8_BLOCK_SCALING_SUPPORT is None: + _FP8_BLOCK_SCALING_SUPPORT = _compute_fp8_block_scaling_support() + return _FP8_BLOCK_SCALING_SUPPORT + + def check_recipe_support(recipe: Recipe) -> None: """Check if the given recipe is supported.""" + if torch.compiler.is_compiling() and isinstance(recipe, DelayedScaling): + raise RuntimeError( + "DelayedScaling is not supported under torch.compile. Please use other recipes instead." + ) recipe_supported = True unsupported_reason = "" if isinstance(recipe, (DelayedScaling, Float8CurrentScaling)): @@ -103,6 +145,11 @@ def check_recipe_support(recipe: Recipe) -> None: def get_default_fp8_recipe() -> Recipe: """FP8 recipe with default args.""" + assert not torch.compiler.is_compiling(), ( + "Creating Recipe objects inside compiled regions is not supported because " + "their construction is not traceable. " + "Pass an explicit recipe to te.autocast() instead." + ) if check_mxfp8_support()[0]: return MXFP8BlockScaling() if get_device_compute_capability() >= (12, 0): @@ -232,71 +279,44 @@ def is_nvfp4_available(return_reason: bool = False) -> Union[bool, Tuple[bool, s return check_nvfp4_support()[0] +@dataclass(slots=True) +class FP8GlobalState: + """Mutable process-global FP8 state stored on an instance. + + Using an instance avoids class-level `setattr(type, ...)` writes, which + `torch.compile` cannot trace in fullgraph mode. + """ + + fp8_enabled: bool = False + fp8_calibration: bool = False + fp8_recipe: Optional[Recipe] = None + fp8_distributed_group: Optional[dist_group_type] = None + fp8_parameters: bool = False + high_precision_init_val: bool = False + is_first_fp8_module: bool = False + fp8_graph_capturing: bool = False + autocast_depth: int = 0 + global_amax_buffer: Dict[str, list] = field(default_factory=dict) + global_amax_history_buffer: Dict[str, list] = field(default_factory=dict) + global_scale_buffer: Dict[str, list] = field(default_factory=dict) + fp8_tensors_recompute_buffer: list = field(default_factory=list) + autocast_arguments: Dict[Any, Tuple[Recipe, Optional[dist_group_type]]] = field( + default_factory=dict + ) + skip_fp8_weight_update_tensor: Optional[torch.Tensor] = None + + class FP8GlobalStateManager: """Class to keep track of and manipulate the global FP8 state at different stages of execution. """ - FP8_ENABLED = False - FP8_CALIBRATION = False - FP8_RECIPE = None - FP8_DISTRIBUTED_GROUP = None - FP8_PARAMETERS = False - HIGH_PRECISION_INIT_VAL = False - IS_FIRST_FP8_MODULE = False - FP8_GRAPH_CAPTURING = False - AUTOCAST_DEPTH = 0 - global_amax_buffer = {} - global_amax_history_buffer = {} - global_scale_buffer = {} - fp8_tensors_recompute_buffer = [] - fp8_available = None - reason_for_no_fp8 = "" - autocast_arguments = {} - skip_fp8_weight_update_tensor = None - mxfp8_available = None - reason_for_no_mxfp8 = "" - fp8_block_scaling_available = None - reason_for_no_fp8_block_scaling = None - nvfp4_available = None - reason_for_no_nvfp4 = "" + quantization_state = FP8GlobalState() @classmethod def reset(cls) -> None: """Reset the global state""" - cls.FP8_ENABLED = False - cls.FP8_CALIBRATION = False - cls.FP8_RECIPE = None - cls.FP8_DISTRIBUTED_GROUP = None - cls.FP8_PARAMETERS = False - cls.HIGH_PRECISION_INIT_VAL = False - cls.IS_FIRST_FP8_MODULE = False - cls.FP8_GRAPH_CAPTURING = False - cls.AUTOCAST_DEPTH = 0 - cls.global_amax_buffer = {} - cls.global_amax_history_buffer = {} - cls.global_scale_buffer = {} - cls.fp8_tensors_recompute_buffer = [] - cls.fp8_available = None - cls.reason_for_no_fp8 = "" - cls.autocast_arguments = {} - cls.skip_fp8_weight_update_tensor = None - cls.mxfp8_available = None - cls.reason_for_no_mxfp8 = "" - cls.fp8_block_scaling_available = None - cls.reason_for_no_fp8_block_scaling = "" - - @classmethod - def set_skip_fp8_weight_update_tensor(cls, skip: bool) -> None: - """`skip_fp8_weight_update_tensor` inplace setter.""" - if cls.skip_fp8_weight_update_tensor is None: - cls.skip_fp8_weight_update_tensor = torch.empty(1, dtype=torch.float32, device="cuda") - cls.skip_fp8_weight_update_tensor.fill_(skip) - - @classmethod - def get_skip_fp8_weight_update_tensor(cls) -> None: - """`skip_fp8_weight_update_tensor` getter.""" - return cls.skip_fp8_weight_update_tensor + cls.quantization_state = FP8GlobalState() @classmethod def is_fp8_available(cls) -> Tuple[bool, str]: @@ -390,6 +410,7 @@ def add_fp8_tensors_to_global_buffer( return fp8_meta[index_in_buffer] = [] + qstate = cls.quantization_state for forward in (True, False): fp8_meta_tensor_key = cls.get_meta_tensor_key(forward=forward) if fp8_meta_tensor_key not in fp8_meta: @@ -398,90 +419,97 @@ def add_fp8_tensors_to_global_buffer( key = cls.get_key_in_buffer(forward, fp8_meta["recipe"], fp8_meta["fp8_group"]) - if key not in cls.global_amax_buffer: - cls.global_amax_buffer[key] = [fp8_meta[fp8_meta_tensor_key].amax_history[0]] - cls.global_amax_history_buffer[key] = [fp8_meta[fp8_meta_tensor_key].amax_history] - cls.global_scale_buffer[key] = [fp8_meta[fp8_meta_tensor_key].scale] + if key not in qstate.global_amax_buffer: + qstate.global_amax_buffer[key] = [fp8_meta[fp8_meta_tensor_key].amax_history[0]] + qstate.global_amax_history_buffer[key] = [ + fp8_meta[fp8_meta_tensor_key].amax_history + ] + qstate.global_scale_buffer[key] = [fp8_meta[fp8_meta_tensor_key].scale] else: - cls.global_amax_buffer[key].append(fp8_meta[fp8_meta_tensor_key].amax_history[0]) - cls.global_amax_history_buffer[key].append( + qstate.global_amax_buffer[key].append(fp8_meta[fp8_meta_tensor_key].amax_history[0]) + qstate.global_amax_history_buffer[key].append( fp8_meta[fp8_meta_tensor_key].amax_history ) - cls.global_scale_buffer[key].append(fp8_meta[fp8_meta_tensor_key].scale) - fp8_meta[index_in_buffer].append(len(cls.global_amax_buffer[key]) - 1) + qstate.global_scale_buffer[key].append(fp8_meta[fp8_meta_tensor_key].scale) + fp8_meta[index_in_buffer].append(len(qstate.global_amax_buffer[key]) - 1) fp8_meta[index_in_buffer].append(key) @classmethod def is_fp8_enabled(cls) -> bool: """Is FP8 enabled""" - return cls.FP8_ENABLED + return cls.quantization_state.fp8_enabled @classmethod def is_fp8_calibration(cls) -> bool: """Is FP8 calibration""" - return cls.FP8_CALIBRATION + return cls.quantization_state.fp8_calibration @classmethod def with_fp8_parameters(cls) -> bool: """Should the parameters be stored as FP8""" - return cls.FP8_PARAMETERS + return cls.quantization_state.fp8_parameters @classmethod def with_high_precision_init_val(cls) -> bool: """Should the high precision initial values be stored with FP8 parameters""" - return cls.HIGH_PRECISION_INIT_VAL + return cls.quantization_state.high_precision_init_val @classmethod def fp8_graph_capturing(cls) -> bool: """Is CUDA graph capture under way?""" - return cls.FP8_GRAPH_CAPTURING or torch.cuda.is_current_stream_capturing() + if torch.compiler.is_compiling(): + assert not cls.quantization_state.fp8_graph_capturing + return False + return ( + cls.quantization_state.fp8_graph_capturing or torch.cuda.is_current_stream_capturing() + ) @classmethod def is_first_fp8_module(cls): """Returns `True` only the first time when called multiple times from within the same `autocast` context. """ - tmp = cls.IS_FIRST_FP8_MODULE - cls.IS_FIRST_FP8_MODULE = False + tmp = cls.quantization_state.is_first_fp8_module + cls.quantization_state.is_first_fp8_module = False return tmp @classmethod def get_fp8_recipe(cls) -> Recipe: """Return the fp8 recipe""" - if cls.FP8_RECIPE is not None: - return cls.FP8_RECIPE + if cls.quantization_state.fp8_recipe is not None: + return cls.quantization_state.fp8_recipe return get_default_fp8_recipe() @classmethod def get_fp8_group(cls) -> Union[dist_group_type, None]: """Return the fp8 group for scale/amax comm""" - return cls.FP8_DISTRIBUTED_GROUP + return cls.quantization_state.fp8_distributed_group @classmethod - def get_autocast_state(cls) -> Tuple[bool, bool, Recipe, dist_group_type, bool]: - """FP8 autocast state getter""" + def get_autocast_state(cls) -> tuple: + """Snapshot the autocast-related fields of the quantization state.""" + qstate = cls.quantization_state return ( - cls.FP8_ENABLED, - cls.FP8_CALIBRATION, - cls.FP8_RECIPE, - cls.FP8_DISTRIBUTED_GROUP, - cls.IS_FIRST_FP8_MODULE, - cls.FP8_GRAPH_CAPTURING, + qstate.fp8_enabled, + qstate.fp8_calibration, + qstate.fp8_recipe, + qstate.fp8_distributed_group, + qstate.is_first_fp8_module, + qstate.fp8_graph_capturing, ) @classmethod - def set_autocast_state( - cls, fp8_state: Tuple[bool, bool, DelayedScaling, dist_group_type, bool] - ) -> None: - """FP8 autocast state setter""" + def set_autocast_state(cls, state: tuple) -> None: + """Restore a previously saved autocast state snapshot.""" + qstate = cls.quantization_state ( - cls.FP8_ENABLED, - cls.FP8_CALIBRATION, - cls.FP8_RECIPE, - cls.FP8_DISTRIBUTED_GROUP, - cls.IS_FIRST_FP8_MODULE, - cls.FP8_GRAPH_CAPTURING, - ) = fp8_state + qstate.fp8_enabled, + qstate.fp8_calibration, + qstate.fp8_recipe, + qstate.fp8_distributed_group, + qstate.is_first_fp8_module, + qstate.fp8_graph_capturing, + ) = state @staticmethod def reduce_tensor_across_group_op_max(tensor: torch.Tensor, group: dist_group_type) -> None: @@ -501,7 +529,11 @@ def reduce_and_update_fp8_tensors( ) -> None: """Delayed scaling only. Concatenate, reduce, and split amaxes in the global buffer.""" # global_amax_buffer should only be non-empty for fp8 delayed scaling - for buffer_key, amax_buffer in cls.global_amax_buffer.items(): + qstate = cls.quantization_state + for ( + buffer_key, + amax_buffer, + ) in qstate.global_amax_buffer.items(): # Check for forward or backward reduction. fwd_update, autocast_key = cls.split_key_in_buffer(buffer_key) if fwd_update != forward: @@ -510,7 +542,7 @@ def reduce_and_update_fp8_tensors( continue # Retrieve autocast specific args and concat amaxes. - recipe, group = cls.autocast_arguments[autocast_key] + recipe, group = qstate.autocast_arguments[autocast_key] contiguous_amax = torch.cat(amax_buffer) # Reduction. @@ -531,8 +563,8 @@ def reduce_and_update_fp8_tensors( if not unfused_update: tex.fused_amax_and_scale_update_after_reduction( contiguous_amax, - cls.global_amax_history_buffer[buffer_key], - cls.global_scale_buffer[buffer_key], + qstate.global_amax_history_buffer[buffer_key], + qstate.global_scale_buffer[buffer_key], recipe.amax_compute_algo, get_fp8_te_dtype(recipe, forward), recipe.margin, @@ -541,8 +573,8 @@ def reduce_and_update_fp8_tensors( split_and_copy(contiguous_amax, amax_buffer, [x.numel() for x in amax_buffer]) for amax_history, scale in zip( - cls.global_amax_history_buffer[buffer_key], - cls.global_scale_buffer[buffer_key], + qstate.global_amax_history_buffer[buffer_key], + qstate.global_scale_buffer[buffer_key], ): _amax_and_scale_update( amax_history, scale, get_fp8_max(recipe, forward), recipe @@ -556,9 +588,10 @@ def get_unique_autocast_key( ): """ For FP8, each autocast can be uniquely identified by the recipe and fp8 group. - Safely using `hash` as we never cross checkpoint boundaries. + Object identity is sufficient since autocast contexts never outlive a single + training session. """ - return f"{str(recipe)}:{hash(group)}" + return str((str(recipe), id(group) if group is not None else None)) @classmethod def autocast_enter( @@ -573,17 +606,21 @@ def autocast_enter( fp8_recipe = get_default_fp8_recipe() if fp8_recipe is None else fp8_recipe autocast_key = cls.get_unique_autocast_key(fp8_recipe, fp8_group) - cls.autocast_arguments[autocast_key] = (fp8_recipe, fp8_group) + qstate = cls.quantization_state + qstate.autocast_arguments[autocast_key] = ( + fp8_recipe, + fp8_group, + ) - cls.FP8_ENABLED = enabled - cls.FP8_CALIBRATION = calibrating - cls.FP8_RECIPE = fp8_recipe - cls.FP8_DISTRIBUTED_GROUP = fp8_group - cls.FP8_GRAPH_CAPTURING = _graph + qstate.fp8_enabled = enabled + qstate.fp8_calibration = calibrating + qstate.fp8_recipe = fp8_recipe + qstate.fp8_distributed_group = fp8_group + qstate.fp8_graph_capturing = _graph - if cls.AUTOCAST_DEPTH == 0: - cls.IS_FIRST_FP8_MODULE = True - cls.AUTOCAST_DEPTH += 1 + if qstate.autocast_depth == 0: + qstate.is_first_fp8_module = True + qstate.autocast_depth += 1 if enabled: fp8_available, reason_for_no_fp8 = cls.is_fp8_available() @@ -601,11 +638,12 @@ def autocast_enter( @classmethod def autocast_exit(cls, enabled: bool, _graph: bool) -> None: """Set state and tracking variables for exit from FP8 region.""" - cls.AUTOCAST_DEPTH -= 1 + qstate = cls.quantization_state + qstate.autocast_depth -= 1 # Reduce only the non-FP8 weight modules here. # FP8 weight modules are reduced at the end of the optimizer # step after the weight amax is populated. - if enabled and cls.AUTOCAST_DEPTH == 0 and not _graph and torch.is_grad_enabled(): + if enabled and qstate.autocast_depth == 0 and not _graph and torch.is_grad_enabled(): # delayed scaling only function, for other recipes (current scaling with any granularity), # this is noop for other recipes because cls.global_amax_buffer is empty list cls.reduce_and_update_fp8_tensors(forward=True) @@ -627,15 +665,16 @@ def copy_forward_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) - fp8_meta["scaling_fwd"].scale.clone(), ] + qstate = cls.quantization_state if buffer_position_key in fp8_meta: - cls.fp8_tensors_recompute_buffer[fp8_meta[buffer_position_key]].append(to_copy) + qstate.fp8_tensors_recompute_buffer[fp8_meta[buffer_position_key]].append(to_copy) else: - if len(cls.fp8_tensors_recompute_buffer) == 0: - cls.fp8_tensors_recompute_buffer = [deque()] + if len(qstate.fp8_tensors_recompute_buffer) == 0: + qstate.fp8_tensors_recompute_buffer = [deque()] else: - cls.fp8_tensors_recompute_buffer.append(deque()) - cls.fp8_tensors_recompute_buffer[-1].append(to_copy) - fp8_meta[buffer_position_key] = len(cls.fp8_tensors_recompute_buffer) - 1 + qstate.fp8_tensors_recompute_buffer.append(deque()) + qstate.fp8_tensors_recompute_buffer[-1].append(to_copy) + fp8_meta[buffer_position_key] = len(qstate.fp8_tensors_recompute_buffer) - 1 @classmethod def get_old_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) -> None: @@ -652,7 +691,9 @@ def get_old_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) -> Non # Retrieve stashed amaxes and scales from phase 1 pre forward. buffer_position_key = "global_fp8_buffer_pos_fwd_recompute" - stashed_fp8_meta = cls.fp8_tensors_recompute_buffer[fp8_meta[buffer_position_key]].popleft() + stashed_fp8_meta = cls.quantization_state.fp8_tensors_recompute_buffer[ + fp8_meta[buffer_position_key] + ].popleft() # Replace amaxes and scales with stashed values for phase 2 forward fp8_meta["scaling_fwd"].amax_history.copy_(stashed_fp8_meta[0]) @@ -749,18 +790,19 @@ def quantized_model_init( This functionality is *EXPERIMENTAL*. """ - _fp8_parameters = FP8GlobalStateManager.FP8_PARAMETERS - _fp8_recipe = FP8GlobalStateManager.FP8_RECIPE - _high_precision_init_val = FP8GlobalStateManager.HIGH_PRECISION_INIT_VAL - FP8GlobalStateManager.FP8_PARAMETERS = enabled - FP8GlobalStateManager.FP8_RECIPE = get_default_fp8_recipe() if recipe is None else recipe - FP8GlobalStateManager.HIGH_PRECISION_INIT_VAL = preserve_high_precision_init_val + qstate = FP8GlobalStateManager.quantization_state + _fp8_parameters = qstate.fp8_parameters + _fp8_recipe = qstate.fp8_recipe + _high_precision_init_val = qstate.high_precision_init_val + qstate.fp8_parameters = enabled + qstate.fp8_recipe = get_default_fp8_recipe() if recipe is None else recipe + qstate.high_precision_init_val = preserve_high_precision_init_val try: yield finally: - FP8GlobalStateManager.FP8_PARAMETERS = _fp8_parameters - FP8GlobalStateManager.FP8_RECIPE = _fp8_recipe - FP8GlobalStateManager.HIGH_PRECISION_INIT_VAL = _high_precision_init_val + qstate.fp8_parameters = _fp8_parameters + qstate.fp8_recipe = _fp8_recipe + qstate.high_precision_init_val = _high_precision_init_val @contextmanager From c6853b65b7177ab3785c48c130166ec3f9324c46 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, 15 Apr 2026 10:43:37 +0200 Subject: [PATCH 348/521] [PyTorch] [torch.compile] Remove module reference from autograd function args (#2791) * Remove module reference from autograd function args Extract weight quantization into standalone `quantize_weight()` function in base.py, eliminating the need to pass `self` (nn.Module) into autograd functions. Each op's autograd function now receives/returns Optional[Tensor] weight workspaces instead, with cache management handled by the nn.Module before/after the autograd call. 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 * Remove unused get_weight_workspace wrapper No callers remain after the quantize_weight refactor. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Return workspaces from _GroupedLinear via tuple instead of mutable list 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 Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * grouped linear fix Signed-off-by: Pawel Gadzinski * fix 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> --- transformer_engine/pytorch/module/base.py | 254 +++++++++--------- .../pytorch/module/grouped_linear.py | 45 +++- .../pytorch/module/layernorm_linear.py | 43 ++- .../pytorch/module/layernorm_mlp.py | 100 +++++-- transformer_engine/pytorch/module/linear.py | 40 ++- 5 files changed, 292 insertions(+), 190 deletions(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 5ca5572e06..83781ca3f3 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -635,6 +635,131 @@ def fill_userbuffers_buffer_for_all_gather( raise ValueError(f"Unsupported quantizer for Userbuffers ({quantizer})") +def _is_weight_workspace_valid( + workspace: QuantizedTensorStorage, + quantizer: Quantizer, +) -> bool: + """Check if a cached weight workspace is compatible with the quantizer's current usage.""" + if isinstance(workspace, Float8TensorStorage): + if ( + not is_non_tn_fp8_gemm_supported() + and quantizer.columnwise_usage + and workspace._transpose is None + ): + return False + elif isinstance(workspace, MXFP8TensorStorage): + if quantizer.rowwise_usage and workspace._rowwise_data is None: + return False + if quantizer.columnwise_usage and workspace._columnwise_data is None: + return False + elif isinstance(workspace, NVFP4TensorStorage): + if quantizer.rowwise_usage and workspace._rowwise_data is None: + return False + if quantizer.columnwise_usage and workspace._columnwise_data is None: + return False + if isinstance(workspace, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer): + return False + return True + + +def quantize_weight( + *, + tensor: Optional[torch.Tensor] = None, + quantizer: Optional[Quantizer] = None, + workspace: Optional[QuantizedTensorStorage] = None, + update_workspace: bool = True, + skip_update_flag: Optional[torch.Tensor] = None, + fsdp_group: Optional["dist_group_type"] = None, + workspace_dtype: Optional[torch.dtype] = None, + cache: bool = False, +) -> Tuple[QuantizedTensorStorage, Optional[QuantizedTensorStorage]]: + """Quantize a weight tensor, optionally reusing a cached workspace. + + Parameters + ---------- + tensor: torch.Tensor, optional + Weight tensor to quantize. + quantizer: Quantizer, optional + Quantizer for casting the weight. + workspace: QuantizedTensorStorage, optional + Previously cached workspace (from the module's ``_fp8_workspaces``). + ``None`` indicates a cache miss. + update_workspace: bool, default = True + Whether to update an existing workspace with fresh values. + skip_update_flag: torch.Tensor, optional + GPU flag to conditionally skip the update. + fsdp_group: dist_group_type, optional + FSDP process group the weights are distributed over. + workspace_dtype: torch.dtype, optional + High-precision dtype for debug quantization workspaces. + cache: bool, default = False + If ``True`` and a new workspace is created, it will be returned + as the second element so the caller can store it. + + Returns + ------- + (weightmat, new_workspace) + *weightmat*: quantized weight ready for GEMM. + *new_workspace*: non-``None`` only when a brand-new workspace was + created **and** ``cache=True``. The caller should store it in + ``_fp8_workspaces``. + """ + + # Already-quantized weight (primary FP8 parameters) + if isinstance(tensor, QuantizedTensor): + update_rowwise = True if quantizer.rowwise_usage else None + update_columnwise = True if quantizer.columnwise_usage else None + tensor.update_usage( + rowwise_usage=update_rowwise, + columnwise_usage=update_columnwise, + ) + if isinstance(quantizer, DebugQuantizer): + tensor = quantizer.wrap_quantized_tensor(tensor) + return tensor, None + + # Validate workspace + if workspace is not None and quantizer is not None: + if not _is_weight_workspace_valid(workspace, quantizer): + workspace = None + + # FSDP gather on cached workspace + if ( + workspace is not None + and tensor is not None + and fsdp_group is not None + and workspace.data.shape != tensor.data.shape + ): + _fsdp_gather_tensors(fsdp_group, [tensor.data.shape], workspace) + + # Cache hit — update in-place and return + if workspace is not None: + if skip_update_flag is not None: + update_workspace = True + if update_workspace: + if tensor is None: + raise ValueError("tensor kwarg must be provided to update FP8 workspace") + if hasattr(workspace, "quantize_"): + workspace.quantize_(tensor, noop_flag=skip_update_flag) + else: + tex.quantize(tensor, quantizer, workspace, skip_update_flag) + return workspace, None + + # Cache miss — create new workspace + if tensor is None or quantizer is None: + raise ValueError("tensor and quantizer kwargs must be provided to construct FP8 workspace") + if cache: + # Ensure the tensor in the cache is an instance of torch.Tensor, + # as it persists beyond a single forward pass. + # Setting internal=True would cause the data to be removed in prepare_for_saving(...). + saved_internal = quantizer.internal + quantizer.internal = False + out = quantizer.quantize(tensor, dtype=workspace_dtype) + if cache: + quantizer.internal = saved_internal + return out, out + return out, None + + class TransformerEngineBaseModule(torch.nn.Module, ABC): """Base TE module.""" @@ -1396,135 +1521,6 @@ def clear(self): def forward(self): """Needs override.""" - def get_weight_workspace( - self, - *, - tensor: Optional[torch.Tensor] = None, - quantizer: Optional[Quantizer] = None, - cache_name: Optional[str] = None, - update_workspace: bool = True, - skip_update_flag: Optional[torch.Tensor] = None, - fsdp_group: Optional[dist_group_type] = None, - workspace_dtype: Optional[torch.dtype] = None, - ) -> QuantizedTensor: - """Get workspace buffer for weights and maybe update its values - - The workspace buffer may be cached for future function calls. - - Parameters - ---------- - tensor : torch.Tensor, optional - Values to copy into workspace. Required if the workspace - is being constructed or updated. - quantizer: Quantizer, optional - Quantizer used to cast the weights. Required if the - workspace is being constructed or updated. - cache_name: str, optional - Key for caching. - update_workspace: bool, default = True - Update workspace with values from `tensor`. - skip_update_flag: torch.Tensor, optional - GPU flag to skip updating the workspace. Take precedence - over `update_workspace` if provided. - fsdp_group: bool, default = None - FSDP process group that the weights are distributed over. - workspace_dtype: torch.dtype, default = None - If weight workspace contains high-precision tensor - for example - for debug quantization, this is dtype of the tensor. - """ - - # Handle case where weights are already quantized - # Note: Make sure weights have required usages, but do not - # destroy unnecessary usages since they may be used later. - if isinstance(tensor, QuantizedTensor): - update_rowwise_usage = True if quantizer.rowwise_usage else None - update_columnwise_usage = True if quantizer.columnwise_usage else None - tensor.update_usage( - rowwise_usage=update_rowwise_usage, - columnwise_usage=update_columnwise_usage, - ) - - if isinstance(quantizer, DebugQuantizer): - tensor = quantizer.wrap_quantized_tensor(tensor) - - return tensor - - # Try getting workspace from cache - out = None - if cache_name is not None: - out = self._fp8_workspaces.get(cache_name, None) - - # Reset cache if workspace is invalid - if out is not None and quantizer is not None: - reset_cache = False - if isinstance(out, Float8TensorStorage): - if ( - not is_non_tn_fp8_gemm_supported() - and quantizer.columnwise_usage - and out._transpose is None - ): - reset_cache = True - elif isinstance(out, MXFP8TensorStorage): - if quantizer.rowwise_usage and out._rowwise_data is None: - reset_cache = True - elif quantizer.columnwise_usage and out._columnwise_data is None: - reset_cache = True - elif isinstance(out, NVFP4TensorStorage): - if quantizer.rowwise_usage and out._rowwise_data is None: - reset_cache = True - elif quantizer.columnwise_usage and out._columnwise_data is None: - reset_cache = True - if isinstance(out, DebugQuantizedTensor) != isinstance(quantizer, DebugQuantizer): - reset_cache = True - if reset_cache: - out = None - del self._fp8_workspaces[cache_name] - - # Gather cached Fp8 workspace if it's distributed - # NOTE: FSDP sharding is supported only for Fp8 buffers and will not work - # for models initialized with Fp8 primary weights. - if ( - out is not None - and tensor is not None - and fsdp_group is not None - and out.data.shape != tensor.data.shape - ): - _fsdp_gather_tensors(fsdp_group, [tensor.data.shape], out) - - # Construct workspace if needed - if out is None: - if tensor is None or quantizer is None: - raise ValueError( - "tensor and quantizer kwargs must be provided to construct FP8 workspace" - ) - - if cache_name is not None: - # Ensure the tensor in the cache is an instance of torch.Tensor, - # as it persists beyond a single forward pass. - # Setting internal=True would cause the data to be removed in prepare_for_saving(...). - quantizer_internal = quantizer.internal - quantizer.internal = False - out = quantizer.quantize(tensor, dtype=workspace_dtype) - if cache_name is not None: - quantizer.internal = quantizer_internal - - # Update cache - if cache_name is not None: - self._fp8_workspaces[cache_name] = out - return out - - # Update workspace if needed - if skip_update_flag is not None: - update_workspace = True - if update_workspace: - if tensor is None: - raise ValueError("tensor kwarg must be provided to update FP8 workspace") - if hasattr(out, "quantize_"): - out.quantize_(tensor, noop_flag=skip_update_flag) - else: - tex.quantize(tensor, quantizer, out, skip_update_flag) - return out - def _load_from_state_dict( self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs ): diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 188a1728db..720a274119 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -17,6 +17,7 @@ from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from .base import ( get_dummy_wgrad, + quantize_weight, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -70,7 +71,7 @@ def forward( inp: torch.Tensor, non_tensor_args: Tuple, *weights_and_biases, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, list]: # pylint: disable=missing-function-docstring # Reduce number of arguments to autograd function in order @@ -93,7 +94,8 @@ def forward( sequence_parallel, activation_dtype, is_grad_enabled, - module, + weight_workspaces, + cache_weight, skip_fp8_weight_update, save_original_input, debug, @@ -178,18 +180,19 @@ def forward( # Initialize weights weights_fp8: list + new_workspaces = [None] * num_gemms if fp8 or debug: - # FP8 cast to workspace buffer weights_fp8 = [] - update_workspace = is_first_microbatch is None or is_first_microbatch + update_ws = is_first_microbatch is None or is_first_microbatch for i in range(num_gemms): - weight_fp8 = module.get_weight_workspace( + weight_fp8, new_workspaces[i] = quantize_weight( tensor=weights[i], quantizer=weight_quantizers[i], - cache_name=(None if is_first_microbatch is None else f"weight{i}"), - update_workspace=update_workspace, + workspace=weight_workspaces[i] if weight_workspaces else None, + update_workspace=update_ws, skip_update_flag=skip_fp8_weight_update, workspace_dtype=activation_dtype, + cache=cache_weight, ) weights_fp8.append(weight_fp8) @@ -332,10 +335,12 @@ def forward( ctx.reduce_and_update_bwd_fp8_tensors = False # [*, in_features] -> [*, out_features] except first dimension changes for SP - return out.view(-1, *inp.shape[1:-1], out.shape[-1]) + return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: + 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"): saved_tensors = restore_from_func_ctx(ctx) @@ -1131,6 +1136,14 @@ 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)] + if cache_weight + else [None] * num_gemms + ) + non_tensor_args = ( m_splits, self.apply_bias, @@ -1149,12 +1162,22 @@ def forward( self.sequence_parallel, self.activation_dtype, is_grad_enabled, - self, + weight_workspaces, + cache_weight, None, # skip_fp8_weight_update self.save_original_input, debug, ) - out = linear_fn(*autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors) + out, new_workspaces = linear_fn( + *autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors + ) + + if cache_weight: + for i, ws in enumerate(new_workspaces): + if ws is not None: + if isinstance(ws, torch.Tensor): + ws = ws.detach() + self._fp8_workspaces[f"weight{i}"] = ws finally: self.end_forward() diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 8ceeaadfcc..f26faade0a 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -21,6 +21,7 @@ from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, + quantize_weight, TransformerEngineBaseModule, get_dummy_wgrad, _2X_ACC_FPROP, @@ -94,9 +95,10 @@ def forward( ln_weight: torch.Tensor, ln_bias: Union[torch.Tensor, None], weight: torch.Tensor, + weight_workspace: Optional[torch.Tensor], bias: torch.Tensor, non_tensor_args: Tuple, - ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: + ) -> Tuple[torch.Tensor, ...]: # pylint: disable=missing-function-docstring # Reduce number of arguments to autograd function in order @@ -136,7 +138,7 @@ def forward( ub_bulk_dgrad, ub_name, fsdp_group, - module, + cache_weight, skip_fp8_weight_update, symmetric_ar_type, debug, @@ -294,6 +296,7 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + new_weight_workspace = None weightmat = weight is_weight_param_quantized = False if fp8 or debug: @@ -311,15 +314,16 @@ def forward( ) # Get quantized weight - update_workspace = is_first_microbatch is None or is_first_microbatch - weightmat = module.get_weight_workspace( + update_ws = is_first_microbatch is None or is_first_microbatch + weightmat, new_weight_workspace = quantize_weight( tensor=weight, quantizer=weight_quantizer, - cache_name=(None if is_first_microbatch is None else "weight"), - update_workspace=update_workspace, + workspace=weight_workspace, + update_workspace=update_ws, skip_update_flag=skip_fp8_weight_update, fsdp_group=fsdp_group, workspace_dtype=activation_dtype, + cache=cache_weight, ) weightmat.update_usage(rowwise_usage=True) @@ -556,13 +560,15 @@ def forward( # Cached state for backward pass is ready... # ------------------------------------------------------ + ln_out_for_return = None if return_layernorm_output: if return_layernorm_output_gathered: shape = list(inp_shape) shape[0] *= tp_size if with_input_all_gather else 1 - return out, ln_out_return.view(shape) - return out, ln_out_return.view(inp_shape) - return out + ln_out_for_return = ln_out_return.view(shape) + else: + ln_out_for_return = ln_out_return.view(inp_shape) + return out, ln_out_for_return, new_weight_workspace @staticmethod def backward( @@ -1073,6 +1079,7 @@ def wgrad_gemm( dgamma, dbeta, wgrad, + None, # weight_workspace grad_bias, None, ) @@ -1594,6 +1601,11 @@ def forward( else: fwd_fn = _LayerNormLinear.forward autograd_ctx = [None] + cache_name = None if is_first_microbatch is None else "weight" + weight_workspace = ( + self._fp8_workspaces.get(cache_name) if cache_name is not None else None + ) + non_tensor_args = ( self.eps, is_first_microbatch, @@ -1629,27 +1641,30 @@ def forward( self.ub_bulk_dgrad, self.ub_name, self.fsdp_group, - self, + cache_name is not None, skip_fp8_weight_update, self.symmetric_ar_type, debug, ) - out = fwd_fn( + out, ln_out, new_weight_workspace = fwd_fn( *autograd_ctx, inp, self.layer_norm_weight, self.layer_norm_bias, weight_tensor, + weight_workspace, bias_tensor if self.apply_bias and not self.gemm_bias_unfused_add else None, non_tensor_args, ) + if new_weight_workspace is not None and cache_name is not None: + if isinstance(new_weight_workspace, torch.Tensor): + new_weight_workspace = new_weight_workspace.detach() + self._fp8_workspaces[cache_name] = new_weight_workspace + finally: self.end_forward() - if self.return_layernorm_output: - out, ln_out = out - if self.gemm_bias_unfused_add: out = out + cast_if_needed(bias_tensor, self.activation_dtype) diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 6e6b11ecfe..a8d6e2e609 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -23,6 +23,7 @@ fill_userbuffers_buffer_for_all_gather, _ub_communicators, get_ub, + quantize_weight, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -176,8 +177,10 @@ def _forward( ln_weight: torch.Tensor, ln_bias: torch.Tensor, fc1_weight: torch.Tensor, + fc1_weight_workspace: Optional[torch.Tensor], fc1_bias: torch.Tensor, fc2_weight: torch.Tensor, + fc2_weight_workspace: Optional[torch.Tensor], fc2_bias: torch.Tensor, non_tensor_args: Tuple, ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: @@ -228,7 +231,8 @@ def _forward( ub_bulk_dgrad, gemm_gelu_fusion, fsdp_group, - module, + fp8_meta, + cache_weight, skip_fp8_weight_update, symmetric_ar_type, checkpoint, @@ -257,7 +261,7 @@ def _forward( == "DelayedScaling" ): # only applicable for delayed scaling FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute( - module.fp8_meta + fp8_meta ) # to restore quantizers during recomputation # save the rng states ctx.cpu_rng_state = torch.get_rng_state() @@ -328,7 +332,8 @@ def _forward( "ub_bulk_dgrad": ub_bulk_dgrad, "gemm_gelu_fusion": gemm_gelu_fusion, "fsdp_group": fsdp_group, - "module": module, + "fp8_meta": fp8_meta, + "cache_weight": False, "skip_fp8_weight_update": skip_fp8_weight_update, "symmetric_ar_type": symmetric_ar_type, "checkpoint": checkpoint, @@ -471,13 +476,12 @@ def _forward( ln_out_total = ln_out # Cast weights to expected dtype + new_fc1_weight_workspace = None + new_fc2_weight_workspace = None fc1_weight_final = fc1_weight fc2_weight_final = fc2_weight if fp8 or debug: - # If weights are not quantized, we call get_weight_workspace, - # which handles weight caching etc. - # FP8 cast to workspace buffer - update_workspace = is_first_microbatch is None or is_first_microbatch + update_ws = is_first_microbatch is None or is_first_microbatch # No need to set the quantizer states if weights are already quantized # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if isinstance(fc1_weight, QuantizedTensorStorage) and not debug: @@ -490,23 +494,25 @@ def _forward( elif fc2_weight_quantizer is not None: fc2_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) - fc1_weight_final = module.get_weight_workspace( + fc1_weight_final, new_fc1_weight_workspace = quantize_weight( tensor=fc1_weight, quantizer=fc1_weight_quantizer, - cache_name=(None if is_first_microbatch is None else "fc1_weight"), - update_workspace=update_workspace, + workspace=fc1_weight_workspace, + update_workspace=update_ws, skip_update_flag=skip_fp8_weight_update, fsdp_group=fsdp_group, workspace_dtype=activation_dtype, + cache=cache_weight, ) - fc2_weight_final = module.get_weight_workspace( + fc2_weight_final, new_fc2_weight_workspace = quantize_weight( tensor=fc2_weight, quantizer=fc2_weight_quantizer, - cache_name=(None if is_first_microbatch is None else "fc2_weight"), - update_workspace=update_workspace, + workspace=fc2_weight_workspace, + update_workspace=update_ws, skip_update_flag=skip_fp8_weight_update, fsdp_group=fsdp_group, workspace_dtype=activation_dtype, + cache=cache_weight, ) fc1_weight_final.update_usage(rowwise_usage=True) fc2_weight_final.update_usage(rowwise_usage=True) @@ -875,13 +881,15 @@ def _forward( ) # we only get to this point if we are not recomputing for bwd, since that would have returned in the block above + ln_out_for_return = None if return_layernorm_output: if return_layernorm_output_gathered: shape = list(inp_shape) shape[0] *= tp_size if (sequence_parallel and set_parallel_mode) else 1 - return fc2_out, ln_out_return.view(shape) - return fc2_out, ln_out_return.view(inp_shape) - return fc2_out + ln_out_for_return = ln_out_return.view(shape) + else: + ln_out_for_return = ln_out_return.view(inp_shape) + return fc2_out, ln_out_for_return, new_fc1_weight_workspace, new_fc2_weight_workspace @staticmethod def forward( @@ -890,11 +898,13 @@ def forward( ln_weight: torch.Tensor, ln_bias: torch.Tensor, fc1_weight: torch.Tensor, + fc1_weight_workspace: Optional[torch.Tensor], fc1_bias: torch.Tensor, fc2_weight: torch.Tensor, + fc2_weight_workspace: Optional[torch.Tensor], fc2_bias: torch.Tensor, non_tensor_args: Tuple, - ) -> Union[Tuple[torch.Tensor, ...], torch.Tensor]: + ) -> Tuple[torch.Tensor, ...]: # pylint: disable=missing-function-docstring # add recompute_for_bwd @@ -906,8 +916,10 @@ def forward( ln_weight, ln_bias, fc1_weight, + fc1_weight_workspace, fc1_bias, fc2_weight, + fc2_weight_workspace, fc2_bias, non_tensor_args, ) @@ -929,7 +941,7 @@ def _recompute(ctx): and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" ): # only applicable for delayed scaling FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute( - ctx.other_args["module"].fp8_meta + ctx.other_args["fp8_meta"] ) # set old quantizer state # get current rng state @@ -940,9 +952,27 @@ def _recompute(ctx): torch.set_rng_state(ctx.cpu_rng_state) _set_cuda_rng_state(ctx.cuda_rng_state) + # Unpack saved tensors and pass None for weight workspaces (recomputed from scratch) + ( + inp_r, + ln_weight_r, + ln_bias_r, + fc1_weight_r, + fc1_bias_r, + fc2_weight_r, + fc2_bias_r, + ) = tensors out = _LayerNormMLP._forward( # recompute ctx, - *tensors, + inp_r, + ln_weight_r, + ln_bias_r, + fc1_weight_r, + None, + fc1_bias_r, + fc2_weight_r, + None, + fc2_bias_r, tuple(ctx.other_args.values()), ) @@ -952,7 +982,7 @@ def _recompute(ctx): and FP8GlobalStateManager.get_fp8_recipe().__class__.__name__ == "DelayedScaling" ): FP8GlobalStateManager.restore_fp8_meta_tensors( - ctx.other_args["module"].fp8_meta + ctx.other_args["fp8_meta"] ) # restore quantizers # set rng state for fwd @@ -1665,8 +1695,10 @@ def fc1_wgrad_gemm( dgamma, dbeta, fc1_wgrad, + None, # fc1_weight_workspace fc1_bias_grad if fc1_bias is not None else None, fc2_wgrad, # pylint: disable=possibly-used-before-assignment + None, # fc2_weight_workspace fc2_bias_grad, None, ) @@ -2132,6 +2164,15 @@ 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" + fc1_weight_workspace = ( + self._fp8_workspaces.get(cache_name_fc1) if cache_name_fc1 is not None else None + ) + fc2_weight_workspace = ( + self._fp8_workspaces.get(cache_name_fc2) if cache_name_fc2 is not None else None + ) + non_tensor_args = ( self.eps, is_first_microbatch, @@ -2175,30 +2216,39 @@ def forward( self.ub_bulk_wgrad, self.gemm_gelu_fusion and not debug, self.fsdp_group, - self, + self.fp8_meta, + cache_name_fc1 is not None, skip_fp8_weight_update, self.symmetric_ar_type, self.checkpoint, debug, ) - out = fwd_fn( + out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn( *autograd_ctx, inp, self.layer_norm_weight, self.layer_norm_bias, fc1_weight, + fc1_weight_workspace, fc1_bias, fc2_weight, + fc2_weight_workspace, fc2_bias if self.apply_bias and not self.gemm_bias_unfused_add else None, non_tensor_args, ) + if new_fc1_ws is not None and cache_name_fc1 is not None: + if isinstance(new_fc1_ws, torch.Tensor): + new_fc1_ws = new_fc1_ws.detach() + self._fp8_workspaces[cache_name_fc1] = new_fc1_ws + if new_fc2_ws is not None and cache_name_fc2 is not None: + if isinstance(new_fc2_ws, torch.Tensor): + new_fc2_ws = new_fc2_ws.detach() + self._fp8_workspaces[cache_name_fc2] = new_fc2_ws + finally: self.end_forward() - if self.return_layernorm_output: - out, ln_out = out - if self.gemm_bias_unfused_add: out = out + cast_if_needed(fc2_bias, self.activation_dtype) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index b57a2eb8d7..63863b4d90 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -20,6 +20,7 @@ fill_userbuffers_buffer_for_all_gather, get_dummy_wgrad, get_ub, + quantize_weight, TransformerEngineBaseModule, _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -88,10 +89,11 @@ 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, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: # pylint: disable=missing-function-docstring ( @@ -123,7 +125,7 @@ def forward( ub_name, fp8_output, # pylint: disable=unused-variable fsdp_group, - module, + cache_weight, skip_fp8_weight_update, symmetric_ar_type, save_original_input, @@ -262,6 +264,7 @@ def forward( # ------------------------------------------------------ # Prepare weight tensor # ------------------------------------------------------ + new_weight_workspace = None weightmat = weight if fp8 or debug: # Configure quantizer @@ -278,18 +281,18 @@ def forward( ) weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) elif isinstance(weight, QuantizedTensor): - # If weight is already quantized, no need to set quantizer states weight_quantizer = weight._quantizer # Get quantized weight - update_workspace = is_first_microbatch is None or is_first_microbatch - weightmat = module.get_weight_workspace( + update_ws = is_first_microbatch is None or is_first_microbatch + weightmat, new_weight_workspace = quantize_weight( tensor=weight, quantizer=weight_quantizer, - cache_name=(None if is_first_microbatch is None else "weight"), - update_workspace=update_workspace, + workspace=weight_workspace, + update_workspace=update_ws, skip_update_flag=skip_fp8_weight_update, fsdp_group=fsdp_group, workspace_dtype=activation_dtype, + cache=cache_weight, ) weightmat.update_usage(rowwise_usage=True) @@ -522,10 +525,12 @@ def forward( # Cached state for backward pass is ready... # ------------------------------------------------------ - return out + return out, new_weight_workspace @staticmethod - def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], ...]: + def backward( + ctx, grad_output: torch.Tensor, _grad_weight_workspace + ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring # NVTX label for profiling @@ -1026,6 +1031,7 @@ def wgrad_gemm( _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) return ( wgrad, + None, # weight_workspace dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, grad_bias, None, @@ -1475,6 +1481,11 @@ def forward( linear_fn = _Linear.forward autograd_ctx = [None] + cache_name = None if is_first_microbatch is None else "weight" + weight_workspace = ( + self._fp8_workspaces.get(cache_name) if cache_name is not None else None + ) + non_tensor_args = ( is_first_microbatch, self.fp8, @@ -1504,19 +1515,26 @@ def forward( self.ub_name, fp8_output, self.fsdp_group, - self, + cache_name is not None, skip_fp8_weight_update, self.symmetric_ar_type, self.save_original_input, debug, ) - out = linear_fn( + 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, ) + + if new_weight_workspace is not None and cache_name is not None: + if isinstance(new_weight_workspace, torch.Tensor): + new_weight_workspace = new_weight_workspace.detach() + self._fp8_workspaces[cache_name] = new_weight_workspace + finally: self.end_forward() if self.gemm_bias_unfused_add: From a073ad5b3ff5c1bc00d9e98669deabe901542aad Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Wed, 15 Apr 2026 10:50:02 -0700 Subject: [PATCH 349/521] Newton-Schulz via cuSOLVERMp (#2706) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [Common] Add Newton-Schulz inverse square root C API via cuSolverMp Add a new distributed Newton-Schulz inverse square root API to Transformer Engine's common C library. This wraps the cusolverMpNewtonSchulz library function, following the same pattern as the existing cuBLASMp integration for comm_gemm. New files: - newton_schulz.h: Public C API header with context management and computation functions - newton_schulz/newton_schulz.cpp: Implementation with RAII wrappers for cuSolverMp handles Build integration: - New NVTE_WITH_CUSOLVERMP CMake option and CUSOLVERMP_HOME env var - NVTE_CHECK_CUSOLVERMP error checking macro in logging.h - Conditional compilation guarded by NVTE_WITH_CUSOLVERMP Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [PyTorch] Add Newton-Schulz PyTorch bindings and distributed tests Add PyTorch-level bindings for the cuSolverMp Newton-Schulz inverse square root API introduced in the previous commit. New files: - pytorch/csrc/extensions/newton_schulz.cpp: C++ extension wrapping the C API with PyTorch tensor support - pytorch/newton_schulz.py: Python wrapper that extracts NCCL communicator from torch.distributed ProcessGroup - tests/pytorch/distributed/test_newton_schulz.py: pytest launcher - tests/pytorch/distributed/run_newton_schulz.py: distributed test worker with reference implementation for numerical validation Modified files: - pytorch/csrc/extensions.h: Function declarations - pytorch/csrc/extensions/pybind.cpp: pybind11 registrations - pytorch/__init__.py: Public API export Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Fix cuSolverMp API signatures in Newton-Schulz implementation Fix API mismatches discovered during compilation: - cusolverMpCreate takes (handle*, deviceId, stream), not (handle*, stream) - cusolverMpCreateDeviceGrid takes handle as first arg with different parameter order - Use cusolverMpGridMapping_t (not cusolverMpGridLayout_t) and CUSOLVERMP_GRID_MAPPING_COL_MAJOR - cusolverMpCreateMatrixDesc has different parameter order: (desc*, grid, dtype, M, N, MB, NB, RSRC, CSRC, LLD) - cusolverMpNewtonSchulzDescriptorCreate takes only (nsDesc*) with no iteration/coefficient args - No cusolverMpStreamSet exists; create handle per-call with user stream - cusolverMpNewtonSchulz requires computeType and info parameters - Switch from generic template RAII to explicit deleter structs Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [PyTorch] Propagate NVTE_WITH_CUSOLVERMP define to PyTorch extension build Add NVTE_WITH_CUSOLVERMP compiler define and cusolverMp include/library paths to the PyTorch C++ extension build, following the same pattern as NVTE_UB_WITH_MPI and NVTE_ENABLE_NVSHMEM. Without this, the #ifdef NVTE_WITH_CUSOLVERMP guards in the PyTorch extension code would never be active since the define was only set as PRIVATE in the CMake build for the common library. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [PyTorch] Fix NCCL comm extraction and pass global dims to Newton-Schulz Two fixes: - Use ProcessGroupNCCL._comm_ptr() to extract the raw NCCL communicator pointer instead of the non-existent get_nccl_comm() method - Pass global matrix dimensions (m, n) from Python to C++ instead of using local tensor dimensions, which would produce incorrect ScaLAPACK block sizes in the distributed computation Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Cache cuSolverMp handle and grid in Newton-Schulz context cuSolverMp handle and grid creation are expensive operations. Move them from per-call creation in nvte_newton_schulz into the NVTECusolverMpCtx, which is their natural home — the context exists to encapsulate the grid. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Create dedicated CUDA stream in Newton-Schulz context cuSolverMp cannot work with the default CUDA stream. Create a dedicated stream inside nvte_cusolvermp_ctx_create and remove the stream parameter from both C API functions since the context now owns its stream. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Fix Newton-Schulz zero output with event-based stream sync The internal dedicated stream was reading the input tensor before the caller's stream had finished producing it, resulting in all-zero output. Add event-based synchronisation: the internal stream waits for the caller's input to be ready, and the caller's stream waits for the output to be written. Replaces the blocking cudaStreamSynchronize. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Fix Newton-Schulz NaNs by keeping host workspace alive cuSolverMp is asynchronous and uses the host workspace during multi-GPU execution. The event-based output sync did not block the host, so the local workspace_host vector was destroyed while the GPU was still reading from it. Restore cudaStreamSynchronize to ensure the host workspace remains valid for the full duration of the operation. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Cache CUDA event in Newton-Schulz context Avoid creating and destroying a cudaEvent_t on every nvte_newton_schulz call by making it a persistent member of NVTECusolverMpCtx, matching the existing pattern for the stream. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * [Common] Use separate in/out events for Newton-Schulz stream sync Replace single event with in_ready and out_ready events. After the cuSolverMp call, record out_ready on the internal stream and make the caller's stream wait on it, ensuring the output tensor is ready before the caller uses it. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * Correct coefficients Signed-off-by: Vladimir Cherepanov * No stream synchronize Signed-off-by: Vladimir Cherepanov * [Test] Verify Newton-Schulz result with XAX=I identity check Replace reference-comparison test with a direct arithmetic check: if X is the inverse square root of A, then X @ A @ X must equal the identity matrix. This is more robust and removes the need for a separate reference implementation. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * Change test - it approximates orthogonal matrix, not inverse square root Signed-off-by: Vladimir Cherepanov * Generalize number of iterations in tests Signed-off-by: Vladimir Cherepanov * Remove extra info diag - everything should be in logs Signed-off-by: Vladimir Cherepanov * Add Newton-Schulz tests to the QA script 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 * Fix outdated comments Signed-off-by: Vladimir Cherepanov * Remove unused variable Signed-off-by: Vladimir Cherepanov * Move magic numbers from tests to impl 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 * Fix outdated comments Signed-off-by: Vladimir Cherepanov * Check num_coefficients 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 * Auto-detect cuSolverMp support from common library binary Instead of requiring NVTE_WITH_CUSOLVERMP env var to be set for both the common library and PyTorch extension builds, inspect the already-built libtransformer_engine.so for exported symbols. This is more robust for incremental builds and CI environments where the env var may not be propagated to the extension build step. The PyTorch extension only calls nvte_* C API functions, so it does not need cusolverMp headers or libraries — only the compile definition. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Vladimir Cherepanov * Conditionally exclude Newton-Schulz API from PyTorch extension When NVTE_WITH_CUSOLVERMP is not defined, omit the Newton-Schulz functions entirely from the pybind module instead of registering stubs that throw runtime errors. The Python wrapper checks for the attribute at call time and raises a clear error message. Co-Authored-By: Claude Opus 4.6 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 * Make symbol detection errors fatal in common_lib_has_symbol Raise FileNotFoundError when no libtransformer_engine.so is found in any candidate location, and raise RuntimeError when nm is unavailable or exits non-zero, rather than silently returning False in both cases. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Vladimir Cherepanov * Search for libtransformer_engine.so via installed module location first In common_lib_has_symbol, prepend a candidate derived by importing transformer_engine via importlib.util.find_spec and using the package directory as the root. This correctly resolves the SO path for source and PyPI installs (where it lives inside transformer_engine/), before falling back to the repo-root and CMake build dir candidates. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Vladimir Cherepanov * Add site packages to search paths for TE common 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 * Revert "Auto-detect cuSolverMp support from common library binary" This reverts commit 8f50bd59d198775b91c2b645f9486398f621f368. Signed-off-by: Vladimir Cherepanov * Remove unused import Signed-off-by: Vladimir Cherepanov * Fix incorrect 'inverse square root' references in Newton-Schulz comments Replace misleading 'inverse square root' descriptions with accurate 'matrix orthogonalization' in the module docstring, function docstring, and pybind11 binding docstring. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Vladimir Cherepanov * [PyTorch] Expose cuSolverMp context creation/destruction as public API Context creation is expensive and should not happen on every newton_schulz call. Introduce CusolverMpCtx and cusolvermp_ctx_create() so callers can create a context once from a ProcessGroup and reuse it. CusolverMpCtx supports explicit destroy() and use as a context manager. newton_schulz() now takes CusolverMpCtx instead of ProcessGroup. Export CusolverMpCtx and cusolvermp_ctx_create from the pytorch package. Update the distributed test worker to use explicit context lifecycle. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Vladimir Cherepanov * [PyTorch] Strengthen input validation in newton_schulz Replace assert with ValueError for the coefficients length check. Add dtype (float32/bfloat16) and contiguity checks for the input tensor. Co-Authored-By: Claude Sonnet 4.6 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 * Use ncclMemAlloc for cuSolverMp Newton-Schulz workspace Signed-off-by: Vladimir Cherepanov * Add Newton-Schulz reference tests Signed-off-by: Vladimir Cherepanov * Fix Newton-Schulz reference test logic Signed-off-by: Vladimir Cherepanov * Fix column-major usage of cuSOLVERMp; add rectangular test cases Signed-off-by: Vladimir Cherepanov * Avoid explicit transpose Signed-off-by: Vladimir Cherepanov * Cleanup Signed-off-by: Vladimir Cherepanov * More cleanup Signed-off-by: Vladimir Cherepanov * Cleanup Signed-off-by: Vladimir Cherepanov * Update transformer_engine/common/newton_schulz/newton_schulz.cpp Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vcherepanov-nv * Fix syntax Signed-off-by: Vladimir Cherepanov * Apply suggestions from code review Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vcherepanov-nv * Add timeout Signed-off-by: Vladimir Cherepanov * Use RAII for cusolvermp CUDA resources Signed-off-by: Vladimir Cherepanov * Make NS API declared unconditional, with stub / runtime errors without cuSOLVERMp support Signed-off-by: Vladimir Cherepanov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix index in diag Signed-off-by: Vladimir Cherepanov * CMake fixes Signed-off-by: Vladimir Cherepanov * Update transformer_engine/pytorch/newton_schulz.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: vcherepanov-nv * Fix a typo Signed-off-by: Vladimir Cherepanov * Cleanup context management Signed-off-by: Vladimir Cherepanov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Borrow more coefficient sets from Emerging Optimizers Signed-off-by: Vladimir Cherepanov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Couple num_iterations with coeff types in tests Signed-off-by: Vladimir Cherepanov --------- Signed-off-by: Vladimir Cherepanov Signed-off-by: vcherepanov-nv Co-authored-by: Claude Opus 4.6 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: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- qa/L1_pytorch_distributed_unittest/test.sh | 1 + setup.py | 5 + .../pytorch/distributed/run_newton_schulz.py | 127 +++++++++ .../pytorch/distributed/test_newton_schulz.py | 69 +++++ transformer_engine/common/CMakeLists.txt | 21 +- .../transformer_engine/newton_schulz.h | 66 +++++ .../common/newton_schulz/newton_schulz.cpp | 267 ++++++++++++++++++ transformer_engine/common/util/logging.h | 16 ++ transformer_engine/pytorch/__init__.py | 4 + transformer_engine/pytorch/csrc/extensions.h | 11 + .../pytorch/csrc/extensions/newton_schulz.cpp | 40 +++ .../pytorch/csrc/extensions/pybind.cpp | 11 + transformer_engine/pytorch/newton_schulz.py | 200 +++++++++++++ 13 files changed, 837 insertions(+), 1 deletion(-) create mode 100644 tests/pytorch/distributed/run_newton_schulz.py create mode 100644 tests/pytorch/distributed/test_newton_schulz.py create mode 100644 transformer_engine/common/include/transformer_engine/newton_schulz.h create mode 100644 transformer_engine/common/newton_schulz/newton_schulz.cpp create mode 100644 transformer_engine/pytorch/csrc/extensions/newton_schulz.cpp create mode 100644 transformer_engine/pytorch/newton_schulz.py diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index 9d868d99cf..db13e9f1e0 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -32,6 +32,7 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_use 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" # debug tests diff --git a/setup.py b/setup.py index 3a66e624e3..ec277b6349 100644 --- a/setup.py +++ b/setup.py @@ -78,6 +78,11 @@ def setup_common_extension() -> CMakeExtension: ).locate_file(f"nvidia/cublasmp/cu{cuda_version()[0]}") cmake_flags.append(f"-DCUBLASMP_DIR={cublasmp_dir}") + if bool(int(os.getenv("NVTE_WITH_CUSOLVERMP", "0"))): + cmake_flags.append("-DNVTE_WITH_CUSOLVERMP=ON") + cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr") + cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}") + # Add custom CMake arguments from environment variable nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS") if nvte_cmake_extra_args: diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py new file mode 100644 index 0000000000..bbd0733447 --- /dev/null +++ b/tests/pytorch/distributed/run_newton_schulz.py @@ -0,0 +1,127 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Distributed Newton-Schulz test worker. + +Launched via torchrun from test_newton_schulz.py. +""" + +import argparse +import sys + +import torch +import torch.distributed as dist +from torch.distributed.elastic.multiprocessing.errors import record + +from transformer_engine.pytorch.newton_schulz import ( + CusolverMpCtx, + get_coefficients, + newton_schulz, +) + + +def newton_schulz_reference(in_x: torch.Tensor, coefficients: list[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)] + xxt = x @ x.mT + x = a * x + b * xxt @ x + c * xxt @ xxt @ x + return x + + +@record +def main(): + parser = argparse.ArgumentParser(description="Newton-Schulz distributed test") + parser.add_argument( + "--check", type=str, default="orthogonality", choices=["orthogonality", "reference"] + ) + parser.add_argument("--dtype", type=str, default="float32", choices=["float32", "bfloat16"]) + parser.add_argument("--matrix-rows", type=int, default=256) + parser.add_argument("--matrix-cols", type=int, default=None) + parser.add_argument("--num-iterations", type=int, default=5) + parser.add_argument("--coeff-type", type=str, default="quintic") + parser.add_argument("--atol", type=float, default=1e-2) + parser.add_argument("--rtol", type=float, default=1e-2) + args = parser.parse_args() + + dist.init_process_group(backend="nccl") + rank = dist.get_rank() + world_size = dist.get_world_size() + torch.cuda.set_device(rank) + + dtype = torch.float32 if args.dtype == "float32" else torch.bfloat16 + m = args.matrix_rows + n = args.matrix_cols if args.matrix_cols is not None else args.matrix_rows + coefficients = get_coefficients(args.num_iterations, args.coeff_type) + + # Ensure the distributed column dimension is divisible by world_size. + assert n % world_size == 0, f"Matrix columns {n} must be divisible by world_size {world_size}" + + # Create a random matrix on rank 0 with singular values in (0, 1), + # which keeps the Newton-Schulz iterations in the convergence regime. + if rank == 0: + torch.manual_seed(42) + k = min(m, n) + U, _ = torch.linalg.qr( + torch.randn(m, k, device="cuda", dtype=torch.float32), mode="reduced" + ) + V, _ = torch.linalg.qr( + torch.randn(n, k, device="cuda", dtype=torch.float32), mode="reduced" + ) + singular_values = torch.rand(k, device="cuda", dtype=torch.float32) * 0.8 + 0.1 + A = U @ torch.diag(singular_values) @ V.T + A = A.to(dtype) + else: + A = torch.empty(m, n, device="cuda", dtype=dtype) + + # Broadcast the full matrix to all ranks + dist.broadcast(A, src=0) + + # Scatter columns to each rank + local_cols = n // world_size + x_local = A[:, rank * local_cols : (rank + 1) * local_cols].contiguous() + + ctx = CusolverMpCtx(dist.group.WORLD) + try: + newton_schulz(x_local, ctx, args.num_iterations, coefficients=coefficients) + finally: + ctx.destroy() + + # Gather results + gathered = [torch.empty_like(x_local) for _ in range(world_size)] + dist.all_gather(gathered, x_local) + X = torch.cat(gathered, dim=1) + + # Check: the resulting matrix should be orthogonal, or match a local reference. + if rank == 0: + if args.check == "orthogonality": + if m <= n: + gram = X @ X.t() + expected = torch.eye(m, device=gram.device, dtype=gram.dtype) + max_diff = (gram - expected).abs().max().item() + print(f"Max |X @ X.t() - I|: {max_diff:.6e}", flush=True) + else: + gram = X.t() @ X + expected = torch.eye(n, device=gram.device, dtype=gram.dtype) + max_diff = (gram - expected).abs().max().item() + print(f"Max |X.t() @ X - I|: {max_diff:.6e}", flush=True) + passed = torch.allclose(gram, expected, atol=args.atol, rtol=args.rtol) + else: + reference = newton_schulz_reference(A.float(), coefficients).to(dtype) + max_diff = (X - reference).abs().max().item() + print(f"Max |distributed - reference|: {max_diff:.6e}", flush=True) + passed = torch.allclose(X, reference, atol=args.atol, rtol=args.rtol) + + if passed: + print("NUMERICAL CHECK PASSED", flush=True) + else: + print("NUMERICAL CHECK FAILED", flush=True, file=sys.stderr) + sys.exit(1) + + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/pytorch/distributed/test_newton_schulz.py b/tests/pytorch/distributed/test_newton_schulz.py new file mode 100644 index 0000000000..0bf4182518 --- /dev/null +++ b/tests/pytorch/distributed/test_newton_schulz.py @@ -0,0 +1,69 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for distributed Newton-Schulz matrix orthogonalization.""" + +import os +import subprocess +from pathlib import Path + +import pytest +import torch + +if torch.cuda.device_count() < 2: + pytest.skip("Newton-Schulz tests require at least 2 GPUs.", allow_module_level=True) + +TEST_ROOT = Path(__file__).parent.resolve() +NUM_PROCS = torch.cuda.device_count() +LAUNCH_CMD = ["torchrun", f"--nproc_per_node={NUM_PROCS}"] +ORTHOGONALITY_SHAPES = [ + (NUM_PROCS * 64, NUM_PROCS * 64), + (NUM_PROCS * 64, NUM_PROCS * 96), + (NUM_PROCS * 96, NUM_PROCS * 64), +] +REFERENCE_SHAPES = [(NUM_PROCS * 64, NUM_PROCS * 64)] + + +def _run_test(dtype, matrix_shape, num_iterations, coeff_type, check): + rows, cols = matrix_shape + test_path = TEST_ROOT / "run_newton_schulz.py" + test_cmd = LAUNCH_CMD + [ + str(test_path), + f"--check={check}", + f"--dtype={dtype}", + f"--matrix-rows={rows}", + f"--matrix-cols={cols}", + f"--num-iterations={num_iterations}", + f"--coeff-type={coeff_type}", + ] + if dtype == "bfloat16": + test_cmd += ["--atol=5e-2", "--rtol=5e-2"] + + result = subprocess.run(test_cmd, env=os.environ, capture_output=True, check=False, timeout=300) + if ( + result.returncode != 0 + or "NUMERICAL CHECK FAILED" in result.stderr.decode() + or "NUMERICAL CHECK PASSED" not in result.stdout.decode() + ): + raise AssertionError( + "Newton-Schulz test failed.\n" + f"stdout: {result.stdout.decode()}\n" + f"stderr: {result.stderr.decode()}" + ) + + +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +@pytest.mark.parametrize("matrix_shape", ORTHOGONALITY_SHAPES) +@pytest.mark.parametrize("num_iterations,coeff_type", [(5, "quintic"), (8, "polar_express")]) +def test_orthogonality(dtype, matrix_shape, num_iterations, coeff_type): + """Test distributed Newton-Schulz orthogonality.""" + _run_test(dtype, matrix_shape, num_iterations, coeff_type, "orthogonality") + + +@pytest.mark.parametrize("dtype", ["float32", "bfloat16"]) +@pytest.mark.parametrize("matrix_shape", REFERENCE_SHAPES) +@pytest.mark.parametrize("num_iterations,coeff_type", [(5, "quintic"), (8, "polar_express")]) +def test_against_reference(dtype, matrix_shape, num_iterations, coeff_type): + """Test distributed Newton-Schulz against a local reference implementation.""" + _run_test(dtype, matrix_shape, num_iterations, coeff_type, "reference") diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a4fbfd9e9c..3f684adbb4 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -153,7 +153,9 @@ list(APPEND transformer_engine_cpp_sources util/rtc.cpp comm_gemm_overlap/userbuffers/ipcsocket.cc comm_gemm_overlap/userbuffers/userbuffers-host.cpp - comm_gemm_overlap/comm_gemm_overlap.cpp) + comm_gemm_overlap/comm_gemm_overlap.cpp + newton_schulz/newton_schulz.cpp + ) list(APPEND transformer_engine_cuda_sources common.cu @@ -343,6 +345,23 @@ if (NVTE_WITH_CUBLASMP) message(STATUS "Using NCCL ${NCCL_VERSION} at: ${NCCL_LIB}") endif() +option(NVTE_WITH_CUSOLVERMP "Use cuSolverMp for distributed Newton-Schulz" OFF) +if (NVTE_WITH_CUSOLVERMP) + target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_CUSOLVERMP) + target_include_directories(transformer_engine PRIVATE ${CUSOLVERMP_DIR}/include) + find_library(CUSOLVERMP_LIB + NAMES cusolverMp libcusolverMp + PATHS ${CUSOLVERMP_DIR} + PATH_SUFFIXES lib + REQUIRED) + find_library(NCCL_LIB + NAMES nccl libnccl + PATH_SUFFIXES lib + REQUIRED) + target_link_libraries(transformer_engine PRIVATE ${NCCL_LIB} ${CUSOLVERMP_LIB}) + message(STATUS "Using cuSolverMp at: ${CUSOLVERMP_DIR}") +endif() + # Number of philox4x32 rounds for stochastic rounding (build-time constant). set(NVTE_BUILD_NUM_PHILOX_ROUNDS_STR $ENV{NVTE_BUILD_NUM_PHILOX_ROUNDS}) if (NOT NVTE_BUILD_NUM_PHILOX_ROUNDS_STR) diff --git a/transformer_engine/common/include/transformer_engine/newton_schulz.h b/transformer_engine/common/include/transformer_engine/newton_schulz.h new file mode 100644 index 0000000000..bea8e32b1e --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/newton_schulz.h @@ -0,0 +1,66 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file newton_schulz.h + * \brief Functions for distributed Newton-Schulz matrix orthogonalization. + * + * This API is a TE-native binding to the cuSolverMp library. + * It computes an iterative Newton-Schulz matrix orthogonalization on a distributed matrix. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_NEWTON_SCHULZ_H_ +#define TRANSFORMER_ENGINE_COMMON_NEWTON_SCHULZ_H_ + +#include +#include + +#include "transformer_engine.h" + +#ifdef __cplusplus +extern "C" { +#endif + +typedef struct NVTECusolverMpCtx NVTECusolverMpCtx; + +/*! \brief Create a cuSolverMp context for Newton-Schulz operations. + * + * Creates a dedicated CUDA stream internally (cuSolverMp requires a + * non-default stream). + * + * \param[in] comm NCCL communicator. + * \param[in] nranks Number of ranks. + * \param[in] rank Local rank. + */ +NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank); + +/*! \brief Destroy a cuSolverMp context. + * + * \param[in] ctx Context to destroy. + */ +void nvte_cusolvermp_ctx_destroy(NVTECusolverMpCtx* ctx); + +/*! \brief Compute Newton-Schulz matrix orthogonalization in-place. + * + * \param[in] ctx cuSolverMp context. + * \param[in] m Global number of rows. + * \param[in] n Global number of columns. + * \param[in,out] x Local part of the matrix (modified in-place). + * \param[in] num_iterations Number of Newton-Schulz iterations. + * \param[in] coefficients Array of polynomial coefficients (length depends on polynomial + * degree used internally by cuSolverMp). + * \param[in] num_coefficients Number of elements in the coefficients array. + * \param[in] caller_stream CUDA stream on which the caller produced the input tensor. + * Used for event-based synchronisation with the internal stream. + */ +void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor x, + int64_t num_iterations, const float* coefficients, int64_t num_coefficients, + cudaStream_t caller_stream); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TRANSFORMER_ENGINE_COMMON_NEWTON_SCHULZ_H_ diff --git a/transformer_engine/common/newton_schulz/newton_schulz.cpp b/transformer_engine/common/newton_schulz/newton_schulz.cpp new file mode 100644 index 0000000000..0d6426a156 --- /dev/null +++ b/transformer_engine/common/newton_schulz/newton_schulz.cpp @@ -0,0 +1,267 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "transformer_engine/newton_schulz.h" + +#include + +#include +#include + +#include "../common.h" +#include "../util/logging.h" + +#ifdef NVTE_WITH_CUSOLVERMP + +#include + +using namespace transformer_engine; + +namespace { + +struct CudaStreamDeleter { + void operator()(std::remove_pointer_t* stream) const { cudaStreamDestroy(stream); } +}; +using CudaStream = std::unique_ptr, CudaStreamDeleter>; + +struct CudaEventDeleter { + void operator()(std::remove_pointer_t* event) const { cudaEventDestroy(event); } +}; +using CudaEvent = std::unique_ptr, CudaEventDeleter>; + +struct CusolverMpHandleDeleter { + void operator()(cusolverMpHandle_t handle) const { cusolverMpDestroy(handle); } +}; +using CusolverMpHandle = + std::unique_ptr, CusolverMpHandleDeleter>; + +struct CusolverMpGridDeleter { + void operator()(cusolverMpGrid_t grid) const { cusolverMpDestroyGrid(grid); } +}; +using CusolverMpGrid = + std::unique_ptr, CusolverMpGridDeleter>; + +struct CusolverMpMatrixDescDeleter { + void operator()(cusolverMpMatrixDescriptor_t desc) const { cusolverMpDestroyMatrixDesc(desc); } +}; +using CusolverMpMatrixDesc = std::unique_ptr, + CusolverMpMatrixDescDeleter>; + +struct CusolverMpNSDescDeleter { + void operator()(cusolverMpNewtonSchulzDescriptor_t desc) const { + cusolverMpNewtonSchulzDescriptorDestroy(desc); + } +}; +using CusolverMpNSDesc = std::unique_ptr, + CusolverMpNSDescDeleter>; + +CusolverMpHandle MakeCusolverMpHandle(int device_id, cudaStream_t stream) { + cusolverMpHandle_t raw{}; + NVTE_CHECK_CUSOLVERMP(cusolverMpCreate(&raw, device_id, stream)); + return CusolverMpHandle(raw); +} + +CusolverMpGrid MakeCusolverMpGrid(cusolverMpHandle_t handle, ncclComm_t comm, int32_t nprow, + int32_t npcol, cusolverMpGridMapping_t mapping) { + cusolverMpGrid_t raw{}; + NVTE_CHECK_CUSOLVERMP(cusolverMpCreateDeviceGrid(handle, &raw, comm, nprow, npcol, mapping)); + return CusolverMpGrid(raw); +} + +CusolverMpMatrixDesc MakeCusolverMpMatrixDesc(cusolverMpGrid_t grid, cudaDataType_t dtype, + int64_t m, int64_t n, int64_t mb, int64_t nb, + uint32_t rsrc, uint32_t csrc, int64_t lld) { + cusolverMpMatrixDescriptor_t raw{}; + NVTE_CHECK_CUSOLVERMP( + cusolverMpCreateMatrixDesc(&raw, grid, dtype, m, n, mb, nb, rsrc, csrc, lld)); + return CusolverMpMatrixDesc(raw); +} + +CusolverMpNSDesc MakeCusolverMpNSDesc() { + cusolverMpNewtonSchulzDescriptor_t raw{}; + NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulzDescriptorCreate(&raw)); + return CusolverMpNSDesc(raw); +} + +CudaStream MakeCudaStream() { + cudaStream_t raw{}; + NVTE_CHECK_CUDA(cudaStreamCreate(&raw)); + return CudaStream(raw); +} + +CudaEvent MakeCudaEvent() { + cudaEvent_t raw{}; + NVTE_CHECK_CUDA(cudaEventCreate(&raw)); + return CudaEvent(raw); +} + +} // namespace + +struct NVTECusolverMpCtx { + int64_t nranks; + int64_t rank; + CudaStream stream; + CudaEvent in_ready; + CudaEvent out_ready; + CusolverMpHandle handle; + CusolverMpGrid grid; + void* workspace; + size_t workspace_size; + bool workspace_registered; +}; + +namespace { + +void FreeWorkspace(NVTECusolverMpCtx* ctx) { + if (ctx->workspace == nullptr) { + return; + } + if (ctx->workspace_registered) { + NVTE_CHECK_CUSOLVERMP(cusolverMpBufferDeregister(ctx->grid.get(), ctx->workspace)); + NVTE_CHECK_NCCL(ncclMemFree(ctx->workspace)); + } else { + NVTE_CHECK_CUDA(cudaFree(ctx->workspace)); + } + ctx->workspace = nullptr; + ctx->workspace_size = 0; + ctx->workspace_registered = false; +} + +} // namespace + +NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank) { + NVTE_API_CALL(nvte_cusolvermp_ctx_create); + int device_id{}; + NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); + + auto stream = MakeCudaStream(); + auto in_ready = MakeCudaEvent(); + auto out_ready = MakeCudaEvent(); + + auto handle = MakeCusolverMpHandle(device_id, stream.get()); + auto grid = MakeCusolverMpGrid(handle.get(), comm, nranks, 1, CUSOLVERMP_GRID_MAPPING_COL_MAJOR); + + return new NVTECusolverMpCtx{ + nranks, + rank, + std::move(stream), + std::move(in_ready), + std::move(out_ready), + std::move(handle), + std::move(grid), + nullptr, + 0, + false, + }; +} + +void nvte_cusolvermp_ctx_destroy(NVTECusolverMpCtx* ctx) { + NVTE_API_CALL(nvte_cusolvermp_ctx_destroy); + FreeWorkspace(ctx); + // Destroy handle and grid before the stream they depend on + ctx->grid.reset(); + ctx->handle.reset(); + delete ctx; +} + +void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor x, + int64_t num_iterations, const float* coefficients, int64_t num_coefficients, + cudaStream_t caller_stream) { + NVTE_API_CALL(nvte_newton_schulz); + NVTE_CHECK(num_coefficients == num_iterations * 3, num_iterations, " iterations require ", + num_iterations * 3, " coefficients, but ", num_coefficients, " are passed"); + const auto* t = convertNVTETensorCheck(x); + + // Make the internal stream wait for the caller's stream so that + // the input tensor is ready before cuSolverMp reads it. + NVTE_CHECK_CUDA(cudaEventRecord(ctx->in_ready.get(), caller_stream)); + NVTE_CHECK_CUDA(cudaStreamWaitEvent(ctx->stream.get(), ctx->in_ready.get())); + + // Block size for ScaLAPACK-style distribution + const int64_t mb = m; + const int64_t nb = (n + ctx->nranks - 1) / ctx->nranks; + + // Compute local leading dimension + const int64_t local_cols = cusolverMpNUMROC(n, nb, ctx->rank, 0, ctx->nranks); + NVTE_CHECK(t->shape().size() == 2, "Shape size:", t->shape().size()); + NVTE_CHECK(t->shape()[1] == local_cols, "Tensor cols:", t->shape()[1], "Local cols:", local_cols); + const int64_t lld = std::max(local_cols, static_cast(1)); + + const cudaDataType_t cuda_dtype = get_cuda_dtype(t->dtype()); + + // Create matrix descriptor + auto mat_desc = MakeCusolverMpMatrixDesc(ctx->grid.get(), cuda_dtype, n, m, nb, mb, 0, 0, lld); + + // Create Newton-Schulz descriptor + auto ns_desc = MakeCusolverMpNSDesc(); + + // Query workspace sizes + size_t wrksp_size_device = 0; + size_t wrksp_size_host = 0; + NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulz_bufferSize( + ctx->handle.get(), ns_desc.get(), n, m, t->data.dptr, 1, 1, mat_desc.get(), num_iterations, + coefficients, CUDA_R_32F, &wrksp_size_device, &wrksp_size_host)); + + // Allocate/grow device workspace + if (ctx->workspace_size < wrksp_size_device) { + FreeWorkspace(ctx); + + void* workspace = nullptr; + bool workspace_registered = false; + + if (ncclMemAlloc(&workspace, wrksp_size_device) == ncclSuccess) { + if (cusolverMpBufferRegister(ctx->grid.get(), workspace, wrksp_size_device) == + CUSOLVER_STATUS_SUCCESS) { + workspace_registered = true; + } else { + NVTE_CHECK_NCCL(ncclMemFree(workspace)); + workspace = nullptr; + } + } + + if (workspace == nullptr) { + NVTE_CHECK_CUDA(cudaMalloc(&workspace, wrksp_size_device)); + } + + ctx->workspace = workspace; + ctx->workspace_size = wrksp_size_device; + ctx->workspace_registered = workspace_registered; + } + + // Allocate host workspace + std::vector workspace_host(wrksp_size_host); + + // Execute Newton-Schulz + NVTE_CHECK_CUSOLVERMP(cusolverMpNewtonSchulz( + ctx->handle.get(), ns_desc.get(), n, m, t->data.dptr, 1, 1, mat_desc.get(), num_iterations, + coefficients, CUDA_R_32F, ctx->workspace, ctx->workspace_size, workspace_host.data(), + workspace_host.size(), nullptr)); + + // Make the caller's stream wait for the internal stream so that + // the output tensor is ready before the caller uses it. + NVTE_CHECK_CUDA(cudaEventRecord(ctx->out_ready.get(), ctx->stream.get())); + NVTE_CHECK_CUDA(cudaStreamWaitEvent(caller_stream, ctx->out_ready.get())); +} + +#else // NVTE_WITH_CUSOLVERMP + +struct NVTECusolverMpCtx {}; + +NVTECusolverMpCtx* nvte_cusolvermp_ctx_create(ncclComm_t comm, int nranks, int rank) { + NVTE_ERROR("Transformer Engine has not been built with cuSolverMp support."); +} + +void nvte_cusolvermp_ctx_destroy(NVTECusolverMpCtx* ctx) { + NVTE_ERROR("Transformer Engine has not been built with cuSolverMp support."); +} + +void nvte_newton_schulz(NVTECusolverMpCtx* ctx, int64_t m, int64_t n, NVTETensor x, + int64_t num_iterations, const float* coefficients, int64_t num_coefficients, + cudaStream_t caller_stream) { + NVTE_ERROR("Transformer Engine has not been built with cuSolverMp support."); +} + +#endif // NVTE_WITH_CUSOLVERMP diff --git a/transformer_engine/common/util/logging.h b/transformer_engine/common/util/logging.h index 8031e342e2..da8b9b377d 100644 --- a/transformer_engine/common/util/logging.h +++ b/transformer_engine/common/util/logging.h @@ -18,6 +18,10 @@ #include #endif // NVTE_WITH_CUBLASMP +#ifdef NVTE_WITH_CUSOLVERMP +#include +#endif // NVTE_WITH_CUSOLVERMP + #include #include #include @@ -106,6 +110,18 @@ #endif // NVTE_WITH_CUBLASMP +#ifdef NVTE_WITH_CUSOLVERMP + +#define NVTE_CHECK_CUSOLVERMP(expr) \ + do { \ + const cusolverStatus_t status_NVTE_CHECK_CUSOLVERMP = (expr); \ + if (status_NVTE_CHECK_CUSOLVERMP != CUSOLVER_STATUS_SUCCESS) { \ + NVTE_ERROR("cuSolverMp Error: ", std::to_string(status_NVTE_CHECK_CUSOLVERMP)); \ + } \ + } while (false) + +#endif // NVTE_WITH_CUSOLVERMP + #define NVTE_CHECK_NCCL(expr) \ do { \ const ncclResult_t status_NVTE_CHECK_NCCL = (expr); \ diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index bbc1d7fab6..d145cf0a21 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -63,6 +63,10 @@ from transformer_engine.pytorch import optimizers from transformer_engine.pytorch.export import onnx_export from transformer_engine.pytorch.cross_entropy import parallel_cross_entropy +from transformer_engine.pytorch.newton_schulz import ( + CusolverMpCtx, + newton_schulz, +) from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.quantized_tensor import QuantizedTensor from transformer_engine.pytorch.quantized_tensor import Quantizer diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index e4bc744e7e..9890f6742a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -593,6 +593,17 @@ void nvshmem_finalize(); void bulk_overlap_ag_with_external_gemm(CommOverlap &allgather_communicator, at::Stream send_stream, at::Stream recv_stream); +/*************************************************************************************************** + * Newton-Schulz (cuSolverMp) + **************************************************************************************************/ + +int64_t cusolvermp_ctx_create(int64_t nccl_comm_ptr, int nranks, int rank); + +void cusolvermp_ctx_destroy(int64_t ctx_ptr); + +void newton_schulz(int64_t ctx_ptr, int64_t m, int64_t n, at::Tensor x, int64_t num_iterations, + std::vector coefficients); + } // namespace transformer_engine::pytorch /*************************************************************************************************** diff --git a/transformer_engine/pytorch/csrc/extensions/newton_schulz.cpp b/transformer_engine/pytorch/csrc/extensions/newton_schulz.cpp new file mode 100644 index 0000000000..8b24e8fdb9 --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/newton_schulz.cpp @@ -0,0 +1,40 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "transformer_engine/newton_schulz.h" + +#include "../extensions.h" + +namespace transformer_engine::pytorch { + +int64_t cusolvermp_ctx_create(int64_t nccl_comm_ptr, int nranks, int rank) { + auto comm = reinterpret_cast(nccl_comm_ptr); + auto* ctx = nvte_cusolvermp_ctx_create(comm, nranks, rank); + return reinterpret_cast(ctx); +} + +void cusolvermp_ctx_destroy(int64_t ctx_ptr) { + auto* ctx = reinterpret_cast(ctx_ptr); + nvte_cusolvermp_ctx_destroy(ctx); +} + +void newton_schulz(int64_t ctx_ptr, int64_t m, int64_t n, at::Tensor x, int64_t num_iterations, + std::vector coefficients) { + auto* ctx = reinterpret_cast(ctx_ptr); + + // Build NVTETensor from PyTorch tensor + auto x_sizes = x.sizes().vec(); + std::vector shape(x_sizes.begin(), x_sizes.end()); + + auto te_dtype = GetTransformerEngineDType(x.scalar_type()); + TensorWrapper x_tensor(x.data_ptr(), shape, te_dtype); + + auto caller_stream = at::cuda::getCurrentCUDAStream().stream(); + nvte_newton_schulz(ctx, m, n, x_tensor.data(), num_iterations, coefficients.data(), + static_cast(coefficients.size()), caller_stream); +} + +} // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 18da5d0e9f..4a20be6361 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -559,6 +559,17 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &transformer_engine::pytorch::multi_tensor_compute_scale_inv_e8m0_cuda, "Fused compute E8M0 scale_inv from amax", py::call_guard()); + // Newton-Schulz (cuSolverMp) + m.def("cusolvermp_ctx_create", &transformer_engine::pytorch::cusolvermp_ctx_create, + "Create cuSolverMp context for Newton-Schulz", py::arg("nccl_comm_ptr"), py::arg("nranks"), + py::arg("rank"), py::call_guard()); + m.def("cusolvermp_ctx_destroy", &transformer_engine::pytorch::cusolvermp_ctx_destroy, + "Destroy cuSolverMp context", py::arg("ctx_ptr"), py::call_guard()); + m.def("newton_schulz", &transformer_engine::pytorch::newton_schulz, + "Newton-Schulz matrix orthogonalization", py::arg("ctx_ptr"), py::arg("m"), py::arg("n"), + py::arg("x"), py::arg("num_iterations"), py::arg("coefficients"), + py::call_guard()); + // Comm+GEMM Overlap m.def("bulk_overlap_ag_with_external_gemm", &transformer_engine::pytorch::bulk_overlap_ag_with_external_gemm, diff --git a/transformer_engine/pytorch/newton_schulz.py b/transformer_engine/pytorch/newton_schulz.py new file mode 100644 index 0000000000..2367897565 --- /dev/null +++ b/transformer_engine/pytorch/newton_schulz.py @@ -0,0 +1,200 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Distributed Newton-Schulz matrix orthogonalization via cuSolverMp.""" + +from itertools import chain, cycle, islice, repeat +from typing import Iterator, List, Literal, Optional, Sequence + +import torch +import torch.distributed as dist + +import transformer_engine_torch as tex + + +_COEFFICIENT_SETS = { + # Values are rounded to closest representable in single precision. + "simple": [ + (3.4445, -4.7750, 2.0315), + ], + "quintic": [ + # optimized for a quintic iteration. + # Source: https://leloykun.github.io/ponder/muon-opt-coeffs/#how-do-we-optimize-the-coefficients + # Numbers from: https://github.com/KellerJordan/modded-nanogpt/blob/0674386070ceb4dcd207e1aca747ffcea6c15250/train_gpt_medium.py#L45 + (4.0848, -6.8946, 2.9270), + (3.9505, -6.3029, 2.6377), + (3.7418, -5.5913, 2.3037), + (2.8769, -3.1427, 1.2046), + (2.8366, -3.0525, 1.2012), + ], + "polar_express": [ + # Polar Express iteration from: https://arxiv.org/abs/2505.16932 + # We include PolarExpress' division by 1.01^polynomial_degree (as stated in their Algorithm 1) in the coefficient list. + # This is a safety factor for numerical stability. + (8.2051, -22.9019, 16.4607), + (4.0664, -2.8612, 0.5184), + (3.9096, -2.8234, 0.5250), + (3.2856, -2.4153, 0.4853), + (2.2779, -1.6198, 0.3985), + (1.8726, -1.2307, 0.3585), + (1.8564, -1.2132, 0.3568), + (1.8750, -1.2500, 0.3750), + ], + "cans": [ + # CANS from: http://arxiv.org/abs/2506.10935 + # CANS iteration (Remez + adaptive interval) based coefficients. + # Source (for generating CANS coefficients): https://github.com/GrishKate/accelerating_orthogonalization/blob/main/polynomials.py + (8.4703, -25.1081, 18.6293), + (4.1828, -3.1087, 0.5806), + (3.9619, -2.9541, 0.5630), + (3.2866, -2.4647, 0.5074), + (2.2737, -1.6447, 0.4162), + ], + "aol": [ + # from https://github.com/thib-s/flash-newton-schulz/blob/main/newton_schulz_triton.py#L511 + (4.0098, -7.0585, 2.4635), + (3.4585, -5.5479, 2.5959), + (2.7573, -3.2939, 1.4254), + (2.7215, -3.0494, 1.3169), + ], +} + +NSCoeffT = Literal[_COEFFICIENT_SETS.keys()] + +CoeffIterMode = Literal["cycle", "repeat_last"] + + +def get_coefficient_iterator( + steps: int, + coefficient_sets: Sequence[tuple[float, float, float]], + mode: CoeffIterMode = "cycle", +) -> Iterator[tuple[float, float, float]]: + """Iterate through coefficient sets with configurable end behavior using itertools. + + Args: + steps: The number of tuples to yield. + coefficient_sets: A sequence of (a, b, c) coefficient tuples. + mode: Iteration mode: + - "cycle": After the last element, restart from the beginning. + - "repeat_last": After the last element, keep yielding the last tuple. + + Yields: + Tuples (a, b, c) from coefficient_sets according to the specified mode. + + Raises: + ValueError: If coefficient_sets is empty. + ValueError: If an invalid mode is provided. + """ + if not coefficient_sets: + raise ValueError("coefficient_sets must be non-empty.") + + base: Iterator[tuple[float, float, float]] + if mode == "cycle": + base = cycle(coefficient_sets) + elif mode == "repeat_last": + # Chain the original list with an infinite repeat of the last item + base = chain(coefficient_sets, repeat(coefficient_sets[-1])) + else: + raise ValueError(f"Invalid mode: {mode}. Expected 'cycle' or 'repeat_last'.") + + return islice(base, steps) + + +def get_coefficients(steps: int, coefficient_type: NSCoeffT = "quintic") -> List[float]: + """Return the coefficient schedule for Newton-Schulz. + + Parameter ``coefficient_type`` can be one of the following + - "simple": Default coefficient set. + - "quintic": Quintic iteration with optimized coefficients. + - "polar_express": Polar Express iteration with optimized coefficients. + - "cans": CANS iteration with Remez + adaptive interval coefficients. + - "aol": AOL coefficient set. + """ + if coefficient_type not in _COEFFICIENT_SETS: + raise ValueError("Invalid coefficient type: " + coefficient_type) + iter_mode: CoeffIterMode = ( + "repeat_last" if coefficient_type in ("polar_express", "cans") else "cycle" + ) + coeff_iter = get_coefficient_iterator( + steps, _COEFFICIENT_SETS[coefficient_type], mode=iter_mode + ) + return list(chain.from_iterable(coeff_iter)) + + +class CusolverMpCtx: + """cuSolverMp context for Newton-Schulz matrix orthogonalization. + + Context creation is expensive; create once and reuse across multiple + :func:`newton_schulz` calls. Call :meth:`destroy` when done. + """ + + def __init__(self, group: dist.ProcessGroup) -> None: + self.nranks = dist.get_world_size(group) + self._ptr = tex.cusolvermp_ctx_create( + _get_nccl_comm_ptr(group), dist.get_world_size(group), dist.get_rank(group) + ) + + def destroy(self) -> None: + """Destroy the underlying cuSolverMp context.""" + if self._ptr is not None: + tex.cusolvermp_ctx_destroy(self._ptr) + self._ptr = None + + def __del__(self) -> None: + # Called when the context is manually destroyed or during Python teardown + self.destroy() + + +def _get_nccl_comm_ptr(group: dist.ProcessGroup) -> int: + """Extract the raw NCCL communicator pointer from a PyTorch process group.""" + backend = dist.get_backend(group) + if backend != "nccl": + raise RuntimeError(f"Newton-Schulz requires NCCL backend, got '{backend}'") + nccl_backend = group._get_backend(torch.device("cuda")) + return nccl_backend._comm_ptr() + + +def newton_schulz( + x: torch.Tensor, + ctx: CusolverMpCtx, + num_iterations: int = 5, + coefficients: Optional[List[float]] = None, +) -> None: + """Compute Newton-Schulz matrix orthogonalization in-place on a distributed matrix. + + Parameters + ---------- + x : torch.Tensor + Local part of the distributed matrix (modified in-place). + Must be a 2D CUDA tensor of type float32 or bfloat16. + Columns are distributed across ranks. + ctx : CusolverMpCtx + cuSolverMp context created by :func:`cusolvermp_ctx_create`. + num_iterations : int, optional + Number of Newton-Schulz iterations. Default: 5. + coefficients : list of float, optional + Polynomial coefficients for the Newton-Schulz iteration. + """ + if coefficients is None: + coefficients = get_coefficients(num_iterations) + if len(coefficients) != num_iterations * 3: + raise ValueError( + f"Unexpected number of coefficients: {len(coefficients)} for" + f" {num_iterations} iterations" + ) + + if x.dim() != 2: + raise ValueError(f"Expected 2D tensor, got {x.dim()}D") + if x.dtype not in (torch.float32, torch.bfloat16): + raise ValueError(f"Expected float32 or bfloat16 tensor, got {x.dtype}") + if not x.is_contiguous(): + raise ValueError("Input tensor must be contiguous") + if not x.is_cuda: + raise ValueError("Input tensor must be on CUDA device") + + # Global matrix dimensions; columns are distributed across ranks. + m = x.size(0) + n = x.size(1) * ctx.nranks + + tex.newton_schulz(ctx._ptr, m, n, x, num_iterations, coefficients) From a817b600cec7f3bee835177d67c06dea6bbc2630 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 15 Apr 2026 11:10:19 -0700 Subject: [PATCH 350/521] [JAX] Tighten Triton autotuning version gate + autotuning enforce env var (#2875) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * jax: tighten Triton autotuning version gate + benchmarking env vars "0.9.3" floor rejects all nightly builds (0.9.2.devN < 0.9.2 < 0.9.3). Bisected: jax-ml/jax#35218 landed 2026-03-10; first fixed container jax-2026-03-17 → set floor to "0.9.2.dev20260317". Signed-off-by: tdophung Replace NVTE_DISABLE_TRITON_AUTOTUNING with NVTE_JAX_ENFORCE_TRITON_AUTOTUNING. Old JAX (<0.9.2.dev20260317) falls back to non-autotuned dispatch by default; set the env var to raise an error prompting JAX upgrade instead. --------- Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/envvars.rst | 15 +++ .../jax/triton_extensions/utils.py | 101 +++++++++++++----- transformer_engine/jax/version_utils.py | 35 +++++- 3 files changed, 122 insertions(+), 29 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 85445430f8..1e040b4c3e 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -443,6 +443,21 @@ JAX-Specific Variables :Default: None :Description: Test level for JAX unit tests (``"L0"``, ``"L1"``, ``"L2"``). Used internally by the test suite. +JAX Triton Extensions +^^^^^^^^^^^^^^^^^^^^^ + +.. envvar:: NVTE_USE_PYTORCH_TRITON + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Explicitly acknowledge using ``pytorch-triton`` for JAX Triton kernels. When both JAX and PyTorch are installed in the same environment, PyTorch's ``pytorch-triton`` package may be imported instead of the standard ``triton`` package from OpenAI. Setting this to ``1`` suppresses the compatibility warning emitted in that situation. ``pytorch-triton`` (the real package from PyTorch's package index, not the placeholder on PyPI) is compatible with JAX Triton kernels. + +.. envvar:: NVTE_JAX_ENFORCE_TRITON_AUTOTUNING + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Raise a ``RuntimeError`` when the installed JAX is too old to safely run ``TritonAutotunedKernelCall`` (`jax-ml/jax#35218 `_) instead of silently falling back to non-autotuned dispatch. Useful for CI or debugging to ensure Triton autotuning is active. When set to ``0`` (default), old JAX versions silently fall back to single-config (non-autotuned) kernel dispatch for compatibility. + Examples -------- diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index ebec1b3cc9..2a86321c34 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -28,6 +28,11 @@ pytorch-triton for JAX Triton kernels (suppresses warnings). This is useful when both JAX and PyTorch are installed in the same environment. Default is "0". + NVTE_JAX_ENFORCE_TRITON_AUTOTUNING: If set to "1", raise a RuntimeError when + the installed JAX is too old to safely run TritonAutotunedKernelCall + (jax-ml/jax#35218) instead of silently falling back to non-autotuned + dispatch. Useful for CI or debugging to ensure autotuning is active. + Default is "0" (silent compatibility fallback). """ import hashlib @@ -45,8 +50,8 @@ from ..version_utils import ( TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION, TRITON_EXTENSION_MIN_JAX_VERSION, + is_triton_autotuned_alias_safe, is_triton_extension_supported, - jax_version_meet_requirement, ) @@ -131,7 +136,13 @@ def _check_triton_compatibility(): "If you don't need Triton, use transformer_engine.jax.cpp_extensions instead." ) - use_pytorch_triton_explicit = bool(int(os.environ.get("NVTE_USE_PYTORCH_TRITON", "0"))) + val = os.environ.get("NVTE_USE_PYTORCH_TRITON", "0") + try: + use_pytorch_triton_explicit = bool(int(val)) + except ValueError as e: + raise ValueError( + f"NVTE_USE_PYTORCH_TRITON must be an integer (0 or 1), got: {val!r}" + ) from e if is_pytorch_triton: if use_pytorch_triton_explicit: @@ -209,7 +220,13 @@ def get_triton_info(): if info['is_pytorch_triton']: print("Using pytorch-triton - compatible with both PyTorch and JAX") """ - env_acknowledged = bool(int(os.environ.get("NVTE_USE_PYTORCH_TRITON", "0"))) + val = os.environ.get("NVTE_USE_PYTORCH_TRITON", "0") + try: + env_acknowledged = bool(int(val)) + except ValueError as e: + raise ValueError( + f"NVTE_USE_PYTORCH_TRITON must be an integer (0 or 1), got: {val!r}" + ) from e return { "version": _TRITON_VERSION, @@ -433,8 +450,33 @@ def lowering(ctx, x, *, block_size): num_ctas = 1 kernel_constexprs = constexprs if constexprs is not None else {} - # Handle autotuned kernels - compile all configs + # Handle autotuned kernels - compile all configs. + # On JAX < TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION the save/restore + # loop in TritonAutotunedKernelCall is buggy (jax-ml/jax#35218). Fall back to a + # single non-autotuned dispatch for compatibility. Set + # NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 to raise an error instead, prompting the + # user to upgrade JAX for improved performance. is_autotuned = isinstance(kernel_fn, autotuner.Autotuner) + if is_autotuned and not is_triton_autotuned_alias_safe(): + val = os.environ.get("NVTE_JAX_ENFORCE_TRITON_AUTOTUNING", "0") + try: + enforce = bool(int(val)) + except ValueError as e: + raise ValueError( + f"NVTE_JAX_ENFORCE_TRITON_AUTOTUNING must be an integer (0 or 1), got: {val!r}" + ) from e + if enforce: + raise RuntimeError( + "NVTE_JAX_ENFORCE_TRITON_AUTOTUNING=1 requires JAX >= " + f"{TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION} (stable) or a " + "post-2026-03-17 nightly for safe Triton autotuning (jax-ml/jax#35218). " + f"Current JAX version: {jax.__version__}. " + "Upgrade: pip install --upgrade jax jaxlib" + ) + # Compatibility fallback: disable autotuning on old JAX to avoid + # CUDA_ERROR_INVALID_VALUE from the unfixed save/restore loop. + is_autotuned = False + if is_autotuned: # Compile all configs for runtime selection kernel_calls = [] @@ -446,8 +488,10 @@ def lowering(ctx, x, *, block_size): config_num_stages = config.num_stages if config.num_stages is not None else num_stages config_num_ctas = config.num_ctas if config.num_ctas is not None else num_ctas - # Merge config kwargs with user constexprs - config_constexprs = {**config.kwargs, **(constexprs if constexprs else {})} + # Config kwargs (e.g. BLOCK_SIZE) take priority over caller constexprs so that + # each autotuning candidate actually compiles with its own BLOCK_SIZE rather than + # having the caller-supplied grid BLOCK_SIZE override every config. + config_constexprs = {**(constexprs if constexprs else {}), **config.kwargs} # Compile this config config_kernel = compile_triton( @@ -478,24 +522,17 @@ def lowering(ctx, x, *, block_size): input_output_aliases_with_sizes = () if input_output_aliases: - if jax_version_meet_requirement(TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION): - num_inputs = len(ctx.avals_in) - aliases = [] - for input_idx, output_idx in input_output_aliases.items(): - aval = ctx.avals_in[input_idx] - size_bytes = aval.size * jnp.dtype(aval.dtype).itemsize - # AutotunedKernelCall expects buffer indices (inputs + outputs). - buffer_output_idx = num_inputs + output_idx - aliases.append((input_idx, buffer_output_idx, size_bytes)) - input_output_aliases_with_sizes = tuple(aliases) - else: - warnings.warn( - f"JAX >= {TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION} is required " - "to safely pass input_output_aliases to TritonAutotunedKernelCall. " - "Passing empty aliases as a workaround (jax-ml/jax#35218).", - UserWarning, - stacklevel=2, - ) + # JAX version is guaranteed >= TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION + # here — verified by the upfront check that set is_autotuned. + num_inputs = len(ctx.avals_in) + aliases = [] + for input_idx, output_idx in input_output_aliases.items(): + aval = ctx.avals_in[input_idx] + size_bytes = aval.size * jnp.dtype(aval.dtype).itemsize + # AutotunedKernelCall expects buffer indices (inputs + outputs). + buffer_output_idx = num_inputs + output_idx + aliases.append((input_idx, buffer_output_idx, size_bytes)) + input_output_aliases_with_sizes = tuple(aliases) kernel_call = gpu_triton.TritonAutotunedKernelCall( f"{actual_kernel_fn.__name__}_autotuned", @@ -504,7 +541,21 @@ def lowering(ctx, x, *, block_size): ) else: - # Regular kernel: compile single config + # Regular kernel: compile single config. + # If the kernel is an Autotuner but JAX is too old for safe autotuning, unwrap + # it and use the first config's kwargs (user constexprs take priority via dict merge). + if isinstance(kernel_fn, autotuner.Autotuner): + actual_kernel_fn = kernel_fn.fn + if kernel_fn.configs: + first_cfg = kernel_fn.configs[0] + # user constexprs override config kwargs (so stride / size scalars win) + kernel_constexprs = {**first_cfg.kwargs, **(constexprs or {})} + num_warps = first_cfg.num_warps if first_cfg.num_warps is not None else num_warps + num_stages = ( + first_cfg.num_stages if first_cfg.num_stages is not None else num_stages + ) + num_ctas = first_cfg.num_ctas if first_cfg.num_ctas is not None else num_ctas + kernel = compile_triton( actual_kernel_fn, signature, diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index 63598481a2..e6ed9a8ea6 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -25,14 +25,40 @@ 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 safe input_output_aliases in TritonAutotunedKernelCall. +# 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 # contains entries for aliases where XLA actually shared buffers at runtime. # Accessing a missing entry produced a null vector → CUDA_ERROR_INVALID_VALUE. -# Fixed by: https://github.com/jax-ml/jax/pull/35218 (merged 2026-03-17, main). -# Ships in JAX 0.9.3 (not yet released as of 2026-03-31). -TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION = "0.9.3" +# Fixed by: https://github.com/jax-ml/jax/pull/35218 (committed 2026-03-10 on jax-ml/jax main; +# first published nightly container: jax-2026-03-17). Ships in JAX 0.9.3 (stable). +# +# Two separate floors are required because packaging.version always ranks a stable +# release above any pre-release of the same series: PkgVersion("0.9.2") > +# PkgVersion("0.9.2.dev20260317"), so a single ">= 0.9.2.dev20260317" check would +# incorrectly accept 0.9.2 stable, which does NOT contain the fix. +# +# nightly build (v.dev is not None): safe if >= 0.9.2.dev20260317 +# stable release (v.dev is None): safe if >= 0.9.3 +_TRITON_AUTOTUNED_ALIAS_NIGHTLY_FLOOR = "0.9.2.dev20260317" +_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR = "0.9.3" + +# Legacy single-constant kept for external callers; reflects the stable floor. +TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION = _TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR + + +@lru_cache(maxsize=None) +def is_triton_autotuned_alias_safe() -> bool: + """Return True if the installed JAX safely supports input_output_aliases on autotuned calls. + + Uses two separate floors (jax-ml/jax#35218): + - nightly builds: >= 0.9.2.dev20260317 (first container with the fix) + - stable releases: >= 0.9.3 (0.9.2 stable does not contain the fix) + """ + v = PkgVersion(get_pkg_version("jax")) + if v.dev is not None: + return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_NIGHTLY_FLOOR) + return v >= PkgVersion(_TRITON_AUTOTUNED_ALIAS_STABLE_FLOOR) def is_triton_extension_supported() -> bool: @@ -47,6 +73,7 @@ def is_triton_extension_supported() -> bool: __all__ = [ "jax_version_meet_requirement", + "is_triton_autotuned_alias_safe", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", "TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION", From a347e09859cd2e5fa9bcb3d13466aaf5b29dc823 Mon Sep 17 00:00:00 2001 From: int-smart Date: Wed, 15 Apr 2026 13:31:55 -0700 Subject: [PATCH 351/521] Add grouped unswizzle functionality for MXFP8 scaling factors (#2837) * Add grouped unswizzle functionality for MXFP8 scaling factors Signed-off-by: Abhishek * Refactored grouped unswizzle kernel to consolidate row and column scaling into a single function and simplify the kernel launch process. Removed redundant check Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Added num tensors and shape checks Signed-off-by: Abhishek --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/cpp/operator/test_swizzle.cu | 229 ++++++++++++++++++ .../include/transformer_engine/swizzle.h | 16 ++ transformer_engine/common/swizzle/swizzle.cu | 133 ++++++++++ 3 files changed, 378 insertions(+) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 7dfb34201d..806a2482ab 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -339,6 +339,173 @@ TEST_P(UnswizzleTestSuite, TestUnswizzle) { transa); } +void performTestGroupedUnswizzleMXFP8(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; + std::vector output_ptrs; + input_tensors.reserve(num_tensors); + output_tensors.reserve(num_tensors); + input_ptrs.reserve(num_tensors); + output_ptrs.reserve(num_tensors); + + 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()); + + 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 = 1; + nvte_set_grouped_tensor_param(grouped_input.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &input_swizzled, sizeof(input_swizzled)); + const uint8_t output_swizzled = 0; + 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]; + + NVTE_CHECK_CUDA(cudaMemset(grouped_output.scale_inv.get(), 0, num_tensors * row_numel)); + NVTE_CHECK_CUDA(cudaMemset(grouped_output.columnwise_scale_inv.get(), 0, num_tensors * col_numel)); + + nvte_unswizzle_grouped_scaling_factors(grouped_input.get_handle(), + grouped_output.get_handle(), 0); + + 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_unswizzle<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_unswizzle<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_unswizzle_rowwise", output_row.data(), ref_row.data(), + num_tensors * row_numel); + compareResults("grouped_unswizzle_colwise", output_col.data(), ref_col.data(), + num_tensors * col_numel); +} + +void performTestGroupedSwizzleUnswizzleRoundtrip(const int num_tensors, const size_t M, + const size_t K) { + using namespace transformer_engine; + 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_ptrs, mid_ptrs, final_ptrs; + orig_tensors.reserve(num_tensors); + mid_tensors.reserve(num_tensors); + final_tensors.reserve(num_tensors); + + for (int i = 0; i < num_tensors; ++i) { + auto orig = std::make_unique("orig_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING); + auto mid = std::make_unique("mid_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING); + auto fin = std::make_unique("fin_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, NVTE_MXFP8_1D_SCALING); + fillUniform(orig.get()); + + // Zero padding so the round-trip comparison is exact. + orig->to_cpu(); + 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); + 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_ptrs.push_back(orig.get()); + mid_ptrs.push_back(mid.get()); + final_ptrs.push_back(fin.get()); + orig_tensors.emplace_back(std::move(orig)); + mid_tensors.emplace_back(std::move(mid)); + final_tensors.emplace_back(std::move(fin)); + } + + GroupedBuffers grouped_orig = build_grouped_tensor(orig_ptrs, NVTE_MXFP8_1D_SCALING); + GroupedBuffers grouped_mid = build_grouped_tensor(mid_ptrs, NVTE_MXFP8_1D_SCALING); + GroupedBuffers grouped_fin = build_grouped_tensor(final_ptrs, NVTE_MXFP8_1D_SCALING); + + const NVTEShape row_shape = orig_tensors[0]->rowwise_scale_inv_shape(); + const NVTEShape col_shape = orig_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]; + + const uint8_t no_swizzle = 0, has_swizzle = 1; + nvte_set_grouped_tensor_param(grouped_orig.get_handle(), kNVTEGroupedWithGEMMSwizzledScales, + &no_swizzle, sizeof(no_swizzle)); + nvte_set_grouped_tensor_param(grouped_mid.get_handle(), kNVTEGroupedWithGEMMSwizzledScales, + &has_swizzle, sizeof(has_swizzle)); + nvte_set_grouped_tensor_param(grouped_fin.get_handle(), kNVTEGroupedWithGEMMSwizzledScales, + &no_swizzle, sizeof(no_swizzle)); + + NVTE_CHECK_CUDA(cudaMemset(grouped_mid.scale_inv.get(), 0, num_tensors * row_numel)); + NVTE_CHECK_CUDA(cudaMemset(grouped_mid.columnwise_scale_inv.get(), 0, num_tensors * col_numel)); + NVTE_CHECK_CUDA(cudaMemset(grouped_fin.scale_inv.get(), 0, num_tensors * row_numel)); + NVTE_CHECK_CUDA(cudaMemset(grouped_fin.columnwise_scale_inv.get(), 0, num_tensors * col_numel)); + + nvte_swizzle_grouped_scaling_factors(grouped_orig.get_handle(), grouped_mid.get_handle(), 0); + nvte_unswizzle_grouped_scaling_factors(grouped_mid.get_handle(), grouped_fin.get_handle(), 0); + + std::vector result_row(num_tensors * row_numel); + std::vector result_col(num_tensors * col_numel); + NVTE_CHECK_CUDA(cudaMemcpy(result_row.data(), grouped_fin.scale_inv.get(), + result_row.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(result_col.data(), grouped_fin.columnwise_scale_inv.get(), + result_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) { + memcpy(ref_row.data() + i * row_numel, + orig_tensors[i]->rowwise_cpu_scale_inv_ptr(), row_numel); + memcpy(ref_col.data() + i * col_numel, + orig_tensors[i]->columnwise_cpu_scale_inv_ptr(), col_numel); + } + + compareResults("grouped_roundtrip_rowwise", result_row.data(), ref_row.data(), + num_tensors * row_numel); + compareResults("grouped_roundtrip_colwise", result_col.data(), ref_col.data(), + num_tensors * col_numel); +} + class SwizzleGroupedTestSuite : public ::testing::TestWithParam> {}; @@ -374,6 +541,68 @@ INSTANTIATE_TEST_SUITE_P( } ); +class UnswizzleGroupedTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(UnswizzleGroupedTestSuite, TestGroupedUnswizzleMXFP8) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + performTestGroupedUnswizzleMXFP8(num_tensors, M, K); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + UnswizzleGroupedTestSuite, + ::testing::Values( + std::make_tuple(3, 256, 256), + std::make_tuple(4, 128, 128), + std::make_tuple(3, 200, 256), + std::make_tuple(2, 65, 256), + std::make_tuple(3, 256, 160), + std::make_tuple(2, 256, 96), + std::make_tuple(3, 200, 160), + std::make_tuple(4, 33, 64), + std::make_tuple(2, 1, 32) + ), + [](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 SwizzleUnswizzleGroupedRoundtripTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(SwizzleUnswizzleGroupedRoundtripTestSuite, TestGroupedSwizzleUnswizzleRoundtrip) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + performTestGroupedSwizzleUnswizzleRoundtrip(num_tensors, M, K); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + SwizzleUnswizzleGroupedRoundtripTestSuite, + ::testing::Values( + std::make_tuple(3, 256, 256), + std::make_tuple(4, 128, 128), + std::make_tuple(3, 200, 256), + std::make_tuple(2, 65, 256), + std::make_tuple(3, 256, 160), + std::make_tuple(2, 256, 96), + std::make_tuple(3, 200, 160), + std::make_tuple(4, 33, 64), + std::make_tuple(2, 1, 32) + ), + [](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)); + } +); + namespace { std::vector> num_tiles = { diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index aa697aafe1..4e28de3beb 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -107,6 +107,22 @@ void nvte_swizzle_block_scaling_to_mxfp8_scaling_factors(const NVTETensor input, void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream); +/*! \brief Unswizzling scaling factors from the interleaved GEMM layout back to row-major (grouped) + * + * \param[in] input Input grouped tensor with swizzled scale_inv. + * \param[in,out] output Output grouped tensor which hosts non-swizzled scale_inv. + * \param[in] stream CUDA stream used for the operation. + * + * Requirements: + * - scaling mode must be MXFP8 1D scaling. + * - scale_inv is stored in row-major in output. + * - scale_inv size is padded to 128x4 for row-scale and 4x128 for col-scale. + * - data is quantized along K-dimension, i.e. 1D-scaling block lies along the K-dimension. + * - all tensors in the grouped tensor must have the same shape. + */ +void nvte_unswizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 28a879a376..6c59776245 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -485,6 +485,24 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) gridDim.y); } +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + grouped_unswizzle_scaling_uniform_shape_kernel(const void* input, void* output, const int M, + const int K, const size_t scale_stride_bytes, + const bool row_scaling) { + const int tensor_id = blockIdx.z; + 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; + if (row_scaling) { + unswizzle_row_scaling_kernel_impl( + input_base, output_base, M, K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y); + } else { + unswizzle_col_scaling_kernel_impl( + input_base, output_base, M, K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y); + } +} + template __global__ void multi_tensor_unswizzle_row_scaling_kernel(MultiSwizzleArgs kernel_args) { const int bid = blockIdx.x; @@ -1692,6 +1710,113 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* } } +void unswizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, + cudaStream_t stream) { + NVTE_CHECK(output->scaling_mode == NVTE_MXFP8_1D_SCALING, + "Grouped unswizzle supports only MXFP8 scaling."); + + CheckInputGroupedTensor(*input, "input"); + CheckOutputGroupedTensor(*output, "output", false); + NVTE_CHECK(input->with_gemm_swizzled_scales, + "Expected input grouped tensor with scales in GEMM swizzled format."); + NVTE_CHECK(!output->with_gemm_swizzled_scales, + "Expected output grouped tensor with scales in compact format."); + NVTE_CHECK(input->scaling_mode == output->scaling_mode, + "Input and output grouped tensors must have matching scaling modes."); + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Input and output grouped tensors must have the same number of tensors."); + + const bool has_rowwise_scale_inv = output->scale_inv.has_data(); + const bool has_columnwise_scale_inv = output->columnwise_scale_inv.has_data(); + if (!has_rowwise_scale_inv && !has_columnwise_scale_inv) { + return; + } + + NVTE_CHECK(input->all_same_shape() && output->all_same_shape(), + "Grouped unswizzle requires uniform tensor shapes."); + + const size_t first_dim = output->get_common_first_dim(); + const size_t last_dim = output->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_unswizzle = [&](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); + const size_t scale_elems = padded_m * padded_k; + + const size_t scale_elem_size = rowwise ? typeToSize(output->scale_inv.dtype) + : typeToSize(output->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."); + } 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."); + } + + 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; + + dim3 num_blocks; + if (rowwise) { + num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, output->num_tensors); + } else { + num_blocks = + dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), output->num_tensors); + } + const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + + 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; + + using kernel_t = void (*)(const void*, void*, const int, const int, const size_t, const bool); + kernel_t kernel_fn = nullptr; + switch (vec_load_size) { + case 4: + kernel_fn = + grouped_unswizzle_scaling_uniform_shape_kernel; + break; + case 2: + kernel_fn = + grouped_unswizzle_scaling_uniform_shape_kernel; + break; + case 1: + kernel_fn = + grouped_unswizzle_scaling_uniform_shape_kernel; + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + } + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + kernel_fn<<>>(input_ptr, output_ptr, padded_m, + padded_k, scale_stride_bytes, rowwise); + NVTE_CHECK_CUDA(cudaGetLastError()); + }; + + if (has_rowwise_scale_inv) { + launch_grouped_unswizzle(true); + } + if (has_columnwise_scale_inv) { + launch_grouped_unswizzle(false); + } +} + } // namespace transformer_engine void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, @@ -1701,3 +1826,11 @@ void nvte_swizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGro swizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input), convertNVTEGroupedTensorCheck(output), stream); } + +void nvte_unswizzle_grouped_scaling_factors(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_unswizzle_grouped_scaling_factors); + using namespace transformer_engine; + unswizzle_grouped_scaling_factors(convertNVTEGroupedTensorCheck(input), + convertNVTEGroupedTensorCheck(output), stream); +} From 92b03707a7ef2fd57d1730a6a769f88bda73fcf1 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Wed, 15 Apr 2026 15:52:25 -0700 Subject: [PATCH 352/521] [Pytorch][JAX] Guard against invalid num_out_tokens in permute_with_mask_map (#2876) * Change docs, and guard against invalid num_out_tokens in mask_map code path Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/jax/permutation.py | 11 ++-- transformer_engine/pytorch/permutation.py | 50 ++++++++++++++----- .../pytorch/triton/permutation.py | 2 +- 3 files changed, 46 insertions(+), 17 deletions(-) diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 6a0a3229d9..81972aac0f 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -73,9 +73,7 @@ def token_dispatch( Routing mask of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. Values: 1 = routed, 0 = not routed. num_out_tokens : int - The number of output tokens after permutation (before padding). For the dropless - case, this should be equal to the sum of routing_map. Must be provided explicitly - for JIT compatibility since output shape must be known at compile time. + Number of output tokens (rows in the permuted buffer, before padding). Must be > 0, e.g. int(jnp.sum(routing_map)) or num_tokens * top_k. Must be a compile-time constant for JIT. probs : Optional[jnp.ndarray] Optional routing probabilities of shape [batch, sequence, num_experts] or [num_tokens, num_experts]. If provided, permuted_probs will be returned. @@ -121,6 +119,8 @@ def token_dispatch( ((num_out_tokens + num_experts * (align_size - 1)) // align_size) * align_size This accounts for the maximum possible padding when each expert needs (align_size - 1) extra tokens to align, rounded down to align_size for buffer alignment. + + Non-positive num_out_tokens (e.g. -1) raises AssertionError. """ use_padding = align_size is not None num_experts = routing_map.shape[-1] @@ -134,6 +134,11 @@ def token_dispatch( else: worst_case_out_tokens = num_out_tokens + assert num_out_tokens > 0, ( + f"token_dispatch requires num_out_tokens > 0, got {num_out_tokens}. " + "Use int(jnp.sum(routing_map)) or num_tokens * top_k." + ) + return _token_dispatch( inp, routing_map, probs, num_out_tokens, worst_case_out_tokens, align_size, use_padding ) diff --git a/transformer_engine/pytorch/permutation.py b/transformer_engine/pytorch/permutation.py index bc9a2660b7..bccc486b4f 100644 --- a/transformer_engine/pytorch/permutation.py +++ b/transformer_engine/pytorch/permutation.py @@ -53,6 +53,9 @@ def moe_permute_index_map_forward( f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " f"index.size(0) ({index.size(0)})." ) + assert ( + num_out_tokens >= 0 + ), f"moe_permute (index map) requires num_out_tokens >= 0, got {num_out_tokens}." if index.dtype != torch.int32: warnings.warn( f"The data type of the input `index` of Permute is {index.dtype}! " @@ -91,6 +94,10 @@ def _moe_permute_index_map_fake( # pylint: disable=unused-argument """Fake implementation for shape inference.""" num_tokens = inp.shape[0] topK = index.shape[1] + if num_tokens > 0: + assert ( + num_out_tokens >= 0 + ), f"moe_permute (index map) requires num_out_tokens >= 0, got {num_out_tokens}." # Infer output shape output_tokens = num_out_tokens if num_out_tokens > 0 else num_tokens * topK @@ -304,6 +311,10 @@ def moe_permute_mask_map_forward( f"Permute not possible: inp.size(0) ({inp.size(0)}) must match " f"routing_map.size(0) ({routing_map.size(0)})." ) + assert num_out_tokens > 0, ( + f"moe_permute (mask map) requires num_out_tokens > 0, got {num_out_tokens}. " + "Use int(routing_map.sum()) or num_tokens * top_k." + ) num_tokens, hidden_size = inp.size() num_experts = routing_map.size(1) @@ -424,13 +435,26 @@ def _moe_permute_mask_map_forward_fake( # pylint: disable=unused-argument num_tokens = inp.shape[0] hidden_size = inp.shape[1] num_experts = routing_map.shape[1] + if num_tokens > 0: + assert num_out_tokens > 0, ( + f"moe_permute (mask map) requires num_out_tokens > 0, got {num_out_tokens}. " + "Use int(routing_map.sum()) or num_tokens * top_k." + ) + out_rows = num_out_tokens + else: + # Match `moe_permute_mask_map_forward` empty-input fast path (ignores num_out_tokens). + out_rows = 0 # row_id_map: (num_tokens, num_experts * 2 + 1) - fake_output = torch.empty((num_out_tokens, hidden_size), dtype=inp.dtype, device=inp.device) + fake_output = torch.empty((out_rows, hidden_size), dtype=inp.dtype, device=inp.device) fake_row_id_map = torch.empty( (num_tokens, num_experts * 2 + 1), dtype=torch.int32, device=inp.device ) if probs is not None: - fake_permuted_probs = torch.empty((num_out_tokens,), dtype=probs.dtype, device=inp.device) + fake_permuted_probs = ( + torch.empty((out_rows,), dtype=probs.dtype, device=inp.device) + if out_rows > 0 + else torch.empty(0, device=inp.device) + ) else: fake_permuted_probs = torch.empty(0, device=inp.device) return fake_output, fake_row_id_map, fake_permuted_probs @@ -852,7 +876,7 @@ def _moe_unpermute_mask_map_backward_wrapper(ctx, unpermuted_act_grad): def moe_permute( inp: torch.Tensor, routing_map: torch.Tensor, - num_out_tokens: int = -1, + num_out_tokens: int, max_token_num: int = -1, map_type: str = "mask", ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -871,13 +895,13 @@ def moe_permute( The values in it: 1 means the token is routed to this expert and 0 means not. If map_type is 'index', routing_map is of shape [num_tokens, topK] and dtype 'int32'. The values in it are the routed expert indices. - num_out_tokens : int, default = -1 - The effective output token count, representing the number of tokens not dropped. - By default, set to '-1', meaning no tokens are dropped. + num_out_tokens : int + Number of output tokens (rows in the permuted buffer). + mask map: must be > 0, e.g. int(routing_map.sum()) or num_tokens * top_k. + index map: must be >= 0; 0 means infer as num_tokens * top_k. max_token_num : int, default = -1 - The maximum number of tokens, used for workspace allocation. - By default, set to '-1', meaning the calculation of the size of workspace is - automatically taken over by the operator. + Workspace sizing hint, only used for map_type='index'. Ignored for 'mask'. + map_type : str, default = 'mask' Type of the routing map tensor. Options are: 'mask', 'index'. @@ -902,7 +926,7 @@ def moe_permute_with_probs( inp: torch.Tensor, probs: torch.Tensor, routing_map: torch.Tensor, - num_out_tokens: int = -1, + num_out_tokens: int, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Permute the tokens and probs based on the routing_map. @@ -921,9 +945,9 @@ def moe_permute_with_probs( routing_map : torch.Tensor The token to expert mapping tensor of shape [num_tokens, num_experts] and dtype 'int32'. The values in it: 1 means the token is routed to this expert and 0 means not. - num_out_tokens : int, default = -1 - The effective output token count, representing the number of tokens not dropped. - By default, set to '-1', meaning no tokens are dropped. + num_out_tokens : int + Number of output tokens (rows in the permuted buffer). Must be > 0, + e.g. int(routing_map.sum()) or num_tokens * top_k. """ if isinstance(inp, QuantizedTensor) and torch.compiler.is_compiling(): raise RuntimeError( diff --git a/transformer_engine/pytorch/triton/permutation.py b/transformer_engine/pytorch/triton/permutation.py index 4902bc686c..c155d73e1e 100644 --- a/transformer_engine/pytorch/triton/permutation.py +++ b/transformer_engine/pytorch/triton/permutation.py @@ -151,7 +151,7 @@ def permute_with_mask_map( num_experts : int Number of experts in the input tensor. num_out_tokens : int - Number of tokens in the permuted tensor. + Number of rows allocated for the permuted tensor (must be a positive integer). hidden_size : int Hidden size of the input tensor. scale_hidden_dim : int From 51d9eebb458db717508695cd762fb698b3260824 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, 16 Apr 2026 12:16:51 +0200 Subject: [PATCH 353/521] [PyTorch] [torch.compile] Split linear forward into forward and setup context. (#2811) * code drop Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix unused weight_quantizer argument in _linear_backward Signed-off-by: Pawel Gadzinski Made-with: Cursor * Reduce duplicate computations between forward_impl and setup_ctx - Move backward_override, custom, backward_input_needs_gather computation to Linear.forward and pass via non_tensor_args - Move UB debug flag zeroing to Linear.forward - Remove unused weight_quantizer_orig param from _linear_setup_ctx - Remove redundant ctx_attrs is None guard Signed-off-by: Pawel Gadzinski Made-with: Cursor --------- Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/module/linear.py | 1961 ++++++++++--------- 1 file changed, 1053 insertions(+), 908 deletions(-) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 63863b4d90..12339e7772 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -80,962 +80,1078 @@ __all__ = ["Linear"] -class _Linear(torch.autograd.Function): - """Linear semi-top level module - Calls custom cuda extensions. +def _check_fp8_reduce_and_update(): + """Check if this is the first FP8 module (for backward reduce-and-update).""" + qstate = FP8GlobalStateManager.quantization_state + _first_fp8_module = qstate.is_first_fp8_module + result = FP8GlobalStateManager.is_first_fp8_module() + if in_fp8_activation_recompute_phase(): + qstate.is_first_fp8_module = _first_fp8_module + return result + + +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: + """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. """ - @staticmethod - def forward( - ctx, - weight: torch.Tensor, - weight_workspace: Optional[torch.Tensor], - inp: torch.Tensor, - bias: Optional[torch.Tensor], - non_tensor_args: Tuple, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - # pylint: disable=missing-function-docstring - - ( - is_first_microbatch, - fp8, - fp8_calibration, - wgrad_store, - input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, - 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, # pylint: disable=unused-variable - fsdp_group, - cache_weight, - skip_fp8_weight_update, - symmetric_ar_type, - save_original_input, - debug, - ) = non_tensor_args - if fp8: - backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + ( + 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, + ) = non_tensor_args + if backward_override == "high_precision": + save_original_input = True + + # NVTX label for profiling + nvtx_label = "transformer_engine._Linear.forward" + if ub_name is not None: + nvtx_label = f"{nvtx_label}.{ub_name}" + + # Make sure input dimensions are compatible + out_features, in_features = weight.shape + assert inp.shape[-1] == in_features, "GEMM not possible" + + # Configure tensor-parallel communication + tp_world_size = get_distributed_world_size(tp_group) + backward_needs_input = is_grad_enabled and weight.requires_grad + with_input_all_gather_nccl = ( + parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop + ) + + # Configure Userbuffers communication (comm+GEMM overlap) + ub_obj = None + ub_type = None + if ub_overlap_rs_fprop: + ub_obj = get_ub(ub_name + "_fprop", fp8) + ub_type = tex.CommOverlapType.RS + elif ub_overlap_ag_fprop: + ub_obj = get_ub(ub_name + "_fprop", fp8) + ub_type = tex.CommOverlapType.AG + + # ------------------------------------------------------ + # Prepare input tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + # ------------------------------------------------------ + nvtx_range_push(f"{nvtx_label}.input_cast_comm") + inputmat = inp # Input tensor to save for backward (maybe sharded) + inputmat_total = None # Input tensor to pass to GEMM (gathered) + own_quantized_input = False + if fp8: + assert_dim_for_fp8_exec(inputmat, weight) + if save_original_input: + assert not isinstance( + input_quantizer, Float8Quantizer + ), "DelayedScaling recipe is not supported with save_original_input" + + if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor + + # Cast local input tensor if needed + if fp8 or debug: + if input_quantizer is None: + raise ValueError("Missing quantizer for input tensor") + if not isinstance(inputmat, QuantizedTensorStorage) and not custom: + own_quantized_input = True + input_quantizer.set_usage( + rowwise=True, + columnwise=backward_needs_input and backward_override is None, + ) + if isinstance(input_quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer)): + # All-gather is not supported with FP8 column-wise data + input_quantizer.set_usage(columnwise=False) + if save_original_input: + # No need for column-wise data since this + # tensor will not be cached for backward pass + input_quantizer.set_usage(columnwise=False) + own_quantized_input = False + inputmat = input_quantizer(inputmat) else: - backward_override = None - if backward_override == "high_precision": - save_original_input = True - - # NVTX label for profiling - nvtx_label = "transformer_engine._Linear.forward" - if ub_name is not None: - nvtx_label = f"{nvtx_label}.{ub_name}" - - # Make sure input dimensions are compatible - out_features, in_features = weight.shape - assert inp.shape[-1] == in_features, "GEMM not possible" - - # Configure tensor-parallel communication - tp_world_size = get_distributed_world_size(tp_group) - backward_needs_input = is_grad_enabled and weight.requires_grad - with_input_all_gather_nccl = ( - parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop - ) + inputmat = cast_if_needed(inp, activation_dtype) # Cast for AMP - # Configure Userbuffers communication (comm+GEMM overlap) - if debug: # turn off userbuffers in debug mode - ub_overlap_rs_fprop = False - ub_overlap_ag_fprop = False - ub_overlap_rs_dgrad = False - ub_bulk_wgrad = False - ub_bulk_dgrad = False - ub_obj = None - ub_type = None - if ub_overlap_rs_fprop: - ub_obj = get_ub(ub_name + "_fprop", fp8) - ub_type = tex.CommOverlapType.RS - elif ub_overlap_ag_fprop: - ub_obj = get_ub(ub_name + "_fprop", fp8) - ub_type = tex.CommOverlapType.AG - - # custom recipe check - custom = is_custom(input_quantizer) or is_custom(weight_quantizer) - - # ------------------------------------------------------ - # Prepare input tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - # ------------------------------------------------------ - nvtx_range_push(f"{nvtx_label}.input_cast_comm") - inputmat = inp # Input tensor to save for backward (maybe sharded) - inputmat_total = None # Input tensor to pass to GEMM (gathered) - own_quantized_input = False - if fp8: - assert_dim_for_fp8_exec(inputmat, weight) - if save_original_input: - assert not isinstance( - input_quantizer, Float8Quantizer - ), "DelayedScaling recipe is not supported with save_original_input" - - if with_input_all_gather_nccl or ub_overlap_ag_fprop: # All-gather input tensor - - # Cast local input tensor if needed - if fp8 or debug: + # Initialize gathered input tensor + quantizer = None + if fp8 or debug: + quantizer = input_quantizer + quantizer.set_usage(rowwise=True, columnwise=False) + if with_input_all_gather_nccl: # Perform NCCL all-gather + inputmat_total, _ = gather_along_first_dim( + inputmat, + tp_group, + quantizer=quantizer, + ) + elif ub_overlap_ag_fprop: # Initialize Userbuffers all-gather + inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj, + inputmat, + quantizer, + tp_group, + ) + + else: # Do not all-gather input tensor + if fp8 or debug: + if isinstance(inputmat, QuantizedTensorStorage): + inputmat.update_usage(rowwise_usage=True) + else: if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") - if not isinstance(inputmat, QuantizedTensorStorage) and not custom: - own_quantized_input = True - input_quantizer.set_usage( - rowwise=True, - columnwise=backward_needs_input and backward_override is None, - ) - if isinstance( - input_quantizer, (Float8Quantizer, Float8CurrentScalingQuantizer) - ): - # All-gather is not supported with FP8 column-wise data - input_quantizer.set_usage(columnwise=False) - if save_original_input: - # No need for column-wise data since this - # tensor will not be cached for backward pass - input_quantizer.set_usage(columnwise=False) - own_quantized_input = False - inputmat = input_quantizer(inputmat) - else: - inputmat = cast_if_needed(inp, activation_dtype) # Cast for AMP - - # Initialize gathered input tensor - quantizer = None - if fp8 or debug: - quantizer = input_quantizer - quantizer.set_usage(rowwise=True, columnwise=False) - if with_input_all_gather_nccl: # Perform NCCL all-gather - inputmat_total, _ = gather_along_first_dim( - inputmat, - tp_group, - quantizer=quantizer, + input_quantizer.set_usage( + rowwise=True, + columnwise=( + backward_needs_input + and not save_original_input + and backward_override is None + ), ) - elif ub_overlap_ag_fprop: # Initialize Userbuffers all-gather - inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj, - inputmat, - quantizer, - tp_group, + inputmat = input_quantizer(inputmat) + own_quantized_input = True + else: + inputmat = cast_if_needed(inp, activation_dtype) # Cast for AMP + inputmat_total = inputmat + + if is_cpu_offload_enabled(): + start_offload(inputmat) + nvtx_range_pop(f"{nvtx_label}.input_cast_comm") + # ------------------------------------------------------ + # Input tensor is ready for GEMM... + # ------------------------------------------------------ + + # ------------------------------------------------------ + # Prepare weight tensor + # ------------------------------------------------------ + new_weight_workspace = None + weightmat = weight + if fp8 or debug: + # Configure quantizer + # 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 + if backward_override is not None: + columnwise_usage = False + if not columnwise_usage: + columnwise_usage = ( + is_fp8_activation_recompute_enabled() + and not in_fp8_activation_recompute_phase() ) + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + elif isinstance(weight, QuantizedTensor): + weight_quantizer = weight._quantizer + # Get quantized weight + update_ws = is_first_microbatch is None or is_first_microbatch + weightmat, new_weight_workspace = quantize_weight( + tensor=weight, + quantizer=weight_quantizer, + workspace=weight_workspace, + update_workspace=update_ws, + skip_update_flag=skip_fp8_weight_update, + fsdp_group=fsdp_group, + workspace_dtype=activation_dtype, + cache=cache_weight, + ) + weightmat.update_usage(rowwise_usage=True) + + else: + weightmat = cast_if_needed(weightmat, activation_dtype) # Cast for AMP + # ------------------------------------------------------ + # Weight tensor is ready for GEMM... + # ------------------------------------------------------ + + # Cast bias to expected dtype + bias_dtype = activation_dtype + if needs_quantized_gemm(inputmat_total) and activation_dtype == torch.float32: + # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 + bias_dtype = torch.bfloat16 + 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 input_quantizer is not None: + input_quantizer.calibrate(inputmat_total) + if weight_quantizer is not None: + weight_quantizer.calibrate(weight) - else: # Do not all-gather input tensor - if fp8 or debug: - if isinstance(inputmat, QuantizedTensorStorage): - inputmat.update_usage(rowwise_usage=True) - else: - if input_quantizer is None: - raise ValueError("Missing quantizer for input tensor") - input_quantizer.set_usage( - rowwise=True, - columnwise=( - backward_needs_input - and not save_original_input - and backward_override is None - ), - ) - inputmat = input_quantizer(inputmat) - own_quantized_input = True + # Choose whether to use GEMM kernel with split accumulator + 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 + + # Configure output quantizer + if output_quantizer is not None: + output_quantizer.set_usage(rowwise=True, columnwise=False) + + # Output buffer for Userbuffers reduce-scatter + reduce_scatter_out = None + if ub_overlap_rs_fprop: + out_shape = list(inp.shape) + out_shape[0] //= tp_world_size + out_shape[-1] = out_features + reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device) + + # ------------------------------------------------------ + # Forward GEMM + # Note: y = x * w^T + # ------------------------------------------------------ + nvtx_range_push(f"{nvtx_label}.gemm") + gemm_out, *_, reduce_scatter_out = general_gemm( + weightmat, + inputmat_total, + quantization_params=output_quantizer, + out_dtype=activation_dtype, + bias=bias, + use_split_accumulator=use_split_accumulator, + ub=ub_obj, + ub_type=ub_type, + extra_output=reduce_scatter_out, + ) + nvtx_range_pop(f"{nvtx_label}.gemm") + # ------------------------------------------------------ + # Finished forward GEMM... + # ------------------------------------------------------ + + # Deallocate GEMM input tensor if no longer needed + # TODO(yuzhongw, tmoon): Figure out why inputmat_total is not automatically + # deallocated by GC. Manually deallocating is a temporary hack. + if with_input_all_gather_nccl: + clear_tensor_data(inputmat_total) + inputmat_total = None + + # ------------------------------------------------------ + # Prepare output tensor + # Note: Perform tensor-parallel communication + # ------------------------------------------------------ + out = None + if ub_overlap_rs_fprop: + out = reduce_scatter_out + elif parallel_mode == "row" and 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: + if symmetric_ar_type is not None: + out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type) else: - inputmat = cast_if_needed(inp, activation_dtype) # Cast for AMP - inputmat_total = inputmat + out, _ = allreduce(out, tp_group) + nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") + else: + out = gemm_out + # ------------------------------------------------------ + # Output tensor is ready to return... + # ------------------------------------------------------ + + # Prepare backward state + tensors_to_save = None + tensor_objects = None + ctx_attrs = None + + if is_grad_enabled: + if save_original_input: + inputmat = inp + + # Discard unneeded data in input tensor + if ( + backward_needs_input + and own_quantized_input + and isinstance(inputmat, QuantizedTensorStorage) + ): + 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() + ): + # All-gather is not supported with FP8 column-wise data + inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) + else: + # Discard row-wise data since it is not needed in backward pass + inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) - if is_cpu_offload_enabled(): - start_offload(inputmat) - nvtx_range_pop(f"{nvtx_label}.input_cast_comm") - # ------------------------------------------------------ - # Input tensor is ready for GEMM... - # ------------------------------------------------------ - - # ------------------------------------------------------ - # Prepare weight tensor - # ------------------------------------------------------ - new_weight_workspace = None - weightmat = weight - if fp8 or debug: - # Configure quantizer - # 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 - if backward_override is not None: - columnwise_usage = False - if not columnwise_usage: - columnwise_usage = ( - is_fp8_activation_recompute_enabled() - and not in_fp8_activation_recompute_phase() - ) - weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - elif isinstance(weight, QuantizedTensor): - weight_quantizer = weight._quantizer - # Get quantized weight - update_ws = is_first_microbatch is None or is_first_microbatch - weightmat, new_weight_workspace = quantize_weight( - tensor=weight, - quantizer=weight_quantizer, - workspace=weight_workspace, - update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, - fsdp_group=fsdp_group, - workspace_dtype=activation_dtype, - cache=cache_weight, - ) - weightmat.update_usage(rowwise_usage=True) + # Cached input tensor + saved_inputmat = None + if backward_needs_input: + saved_inputmat = inputmat - else: - weightmat = cast_if_needed(weightmat, activation_dtype) # Cast for AMP - # ------------------------------------------------------ - # Weight tensor is ready for GEMM... - # ------------------------------------------------------ - - # Cast bias to expected dtype - bias_dtype = activation_dtype - if needs_quantized_gemm(inputmat_total) and activation_dtype == torch.float32: - # cuBLAS does not support FP8 GEMM with FP32 bias, so we cast to BF16 - bias_dtype = torch.bfloat16 - 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 input_quantizer is not None: - input_quantizer.calibrate(inputmat_total) - if weight_quantizer is not None: - weight_quantizer.calibrate(weight) - - # Choose whether to use GEMM kernel with split accumulator - 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 - - # Configure output quantizer - if output_quantizer is not None: - output_quantizer.set_usage(rowwise=True, columnwise=False) - - # Output buffer for Userbuffers reduce-scatter - reduce_scatter_out = None - if ub_overlap_rs_fprop: - out_shape = list(inp.shape) - out_shape[0] //= tp_world_size - out_shape[-1] = out_features - reduce_scatter_out = torch.empty(out_shape, dtype=activation_dtype, device=inp.device) - - # ------------------------------------------------------ - # Forward GEMM - # Note: y = x * w^T - # ------------------------------------------------------ - nvtx_range_push(f"{nvtx_label}.gemm") - gemm_out, *_, reduce_scatter_out = general_gemm( - weightmat, - inputmat_total, - quantization_params=output_quantizer, - out_dtype=activation_dtype, - bias=bias, - use_split_accumulator=use_split_accumulator, - ub=ub_obj, - ub_type=ub_type, - extra_output=reduce_scatter_out, + if cpu_offloading and saved_inputmat is not None: + mark_activation_offload(saved_inputmat) + + # Scatter intermediate/activation tensors saved for the backward pass + # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights + nvtx_range_push(f"{nvtx_label}.fsdp_scatter") + fsdp_shapes = _fsdp_scatter_tensors( + fsdp_group, + saved_inputmat, + weightmat if fp8 and not isinstance(weight, QuantizedTensorStorage) else None, ) - nvtx_range_pop(f"{nvtx_label}.gemm") - # ------------------------------------------------------ - # Finished forward GEMM... - # ------------------------------------------------------ - - # Deallocate GEMM input tensor if no longer needed - # TODO(yuzhongw, tmoon): Figure out why inputmat_total is not automatically - # deallocated by GC. Manually deallocating is a temporary hack. - if with_input_all_gather_nccl: - clear_tensor_data(inputmat_total) - inputmat_total = None - - # ------------------------------------------------------ - # Prepare output tensor - # Note: Perform tensor-parallel communication - # ------------------------------------------------------ - out = None - if ub_overlap_rs_fprop: - out = reduce_scatter_out - elif parallel_mode == "row" and 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: - if symmetric_ar_type is not None: - out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type) - else: - out, _ = allreduce(out, tp_group) - nvtx_range_pop(f"{nvtx_label}.row_parallel_comm") - else: - out = gemm_out - # ------------------------------------------------------ - # Output tensor is ready to return... - # ------------------------------------------------------ + nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") - # ------------------------------------------------------ - # Cache state for backward pass - # ------------------------------------------------------ + if cpu_offloading: + mark_not_offload(weight, weightmat, bias) - if is_grad_enabled: - if save_original_input: - inputmat = inp + # TODO(ksivamani): Check memory usage + tensors_to_save, tensor_objects = prepare_for_saving( + saved_inputmat, + weightmat, + weight, + bias, + ) - ctx.weight_quantizer = weight_quantizer + owns_input = saved_inputmat is not inp + + ctx_attrs = { + "weight_quantizer": weight_quantizer, + "fsdp_shapes": fsdp_shapes, + "owns_input": owns_input, + } + + return out, new_weight_workspace, tensors_to_save, tensor_objects, 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, + ) = 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 + + # 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) + if hasattr(weight, "__fsdp_param__"): + ctx.main_grad_func = weight.get_main_grad + else: + ctx.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"] + + # 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.""" + + # 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}" + + with get_nvtx_range_context("_Linear_backward"): + ( + inputmat, + weight_fp8, + saved_weight, + bias, + ) = restore_from_func_ctx( # pylint: disable=unbalanced-tuple-unpacking + ctx + ) - ctx.backward_input_needs_gather = ( - weight.requires_grad and parallel_mode == "column" and sequence_parallel + origin_weight_python_object = None + origin_weight_overwrites_main_grad = getattr( + ctx, "origin_weight_overwrites_main_grad", False + ) + main_grad = None + if ctx.fuse_wgrad_accumulation and ctx.requires_wgrad: + origin_weight_ref = ctx.origin_weight_ref + ctx.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() + 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 + # 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, + inputmat, + weight_fp8, + ) + nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - # Discard unneeded data in input tensor - if ( - backward_needs_input - and own_quantized_input - and isinstance(inputmat, QuantizedTensorStorage) - ): - if backward_override is not None: - # In dequantized mode we should dequantize directly from the - # fprop quantized tensor layout without retargeting usage. - inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) - elif ( - ctx.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) - else: - # Discard row-wise data since it is not needed in backward pass - inputmat.update_usage(rowwise_usage=False, columnwise_usage=True) - - # Cached input tensor - saved_inputmat = None - if backward_needs_input: - saved_inputmat = inputmat - - if cpu_offloading and saved_inputmat is not None: - mark_activation_offload(saved_inputmat) - - # Scatter intermediate/activation tensors saved for the backward pass - # NOTE: FSDP sharding is not valid for models initialized with primary Fp8 weights - nvtx_range_push(f"{nvtx_label}.fsdp_scatter") - ctx.fsdp_group = fsdp_group - ctx.fsdp_shapes = _fsdp_scatter_tensors( - fsdp_group, - saved_inputmat, - weightmat if fp8 and not isinstance(weight, QuantizedTensorStorage) else None, - ) - nvtx_range_pop(f"{nvtx_label}.fsdp_scatter") + # Configure Userbuffers communication (comm+GEMM overlap) + ctx.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: + # 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 + ub_type_dgrad = tex.CommOverlapType.AG + elif ctx.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 + ub_type_dgrad = tex.CommOverlapType.RS + else: + if ctx.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 + ub_type_dgrad = tex.CommOverlapType.AG + if ctx.ub_bulk_wgrad: + # Overlap dgrad reduce-scatter with wgrad compute + ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) + ub_type_wgrad = tex.CommOverlapType.RS - if cpu_offloading: - mark_not_offload(weight, weightmat, bias) + # -------------------------------------------------- + # Prepare grad output tensor + # Note: Cast to expected dtype and perform tensor-parallel communication + # -------------------------------------------------- - # TODO(ksivamani): Check memory usage - tensors_to_save, tensor_objects = prepare_for_saving( - saved_inputmat, - weightmat, - weight, - bias, - ) - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - 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.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 - ctx.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - if fuse_wgrad_accumulation and weight.requires_grad: - # Keep a weakref to the original Python object because save_for_backward - # may return a plain Tensor without custom Parameter attributes. - ctx.origin_weight_ref = weakref.ref(weight) - ctx.origin_weight_overwrites_main_grad = getattr( - weight, "overwrite_main_grad", False - ) - # This check is needed to ensure that main_grad is not created - # during the forward pass when using MCore FSDP as it creates - # the main_grad buffer lazily before backprop - if hasattr(weight, "__fsdp_param__"): - # MCore FSDP creates main_grad lazily before backward - ctx.main_grad_func = weight.get_main_grad - else: - ctx.main_grad_func = lambda: weight.main_grad - - ctx.debug = debug - ctx.custom = custom - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.use_bias = bias is not None - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.inp_shape = inp.shape - ctx.parallel_mode = parallel_mode - ctx.tp_group = tp_group - 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 - ctx.ub_name = ub_name - ctx.tp_size = tp_size - ctx.requires_dgrad = inp.requires_grad - ctx.requires_wgrad = weight.requires_grad - ctx.reduce_and_update_bwd_fp8_tensors = False - - ctx.owns_input = saved_inputmat is not inp - if ctx.fp8 and requires_grad(inp, weight, bias): - qstate = FP8GlobalStateManager.quantization_state - _first_fp8_module = qstate.is_first_fp8_module - ctx.reduce_and_update_bwd_fp8_tensors = FP8GlobalStateManager.is_first_fp8_module() - if in_fp8_activation_recompute_phase(): - qstate.is_first_fp8_module = _first_fp8_module - ctx.wgrad_store = wgrad_store - - # 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 - ctx.reduce_and_update_bwd_fp8_tensors = False + # Unmodified grad output tensor + grad_output_arg = grad_output - # ------------------------------------------------------ - # Cached state for backward pass is ready... - # ------------------------------------------------------ + # Configure quantizer for grad output tensor + # Note: dgrad GEMM requires row-wise usage, wgrad GEMM + # requires column-wise usage + if grad_output_quantizer is not None: + quantizer = grad_output_quantizer + quantizer.set_usage(rowwise=True, columnwise=True) + if ctx.ub_overlap_ag: + # Userbuffers only supports communication for one + # tensor usage at a time. Configure quantizer with + # usage for only dgrad GEMM. + quantizer.set_usage(columnwise=False) + + # Adjust the quantization direction approach depending + # 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 + # `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: + grad_output_quantizer.set_usage(columnwise=False) + + # Prepare grad output tensor + nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") + ( + grad_output, + grad_bias, + ) = TransformerEngineBaseModule.grad_output_preprocess( + ctx, + grad_output, + ctx.parallel_mode == "row", + grad_output_quantizer, + ) + nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") - return out, new_weight_workspace + # -------------------------------------------------- + # Grad output tensor is ready for computing grad input... + # -------------------------------------------------- - @staticmethod - def backward( - ctx, grad_output: torch.Tensor, _grad_weight_workspace - ) -> Tuple[Union[torch.Tensor, None], ...]: - # pylint: disable=missing-function-docstring + # -------------------------------------------------- + # Prepare input tensor + # Note: Input tensor is needed for wgrad GEMM. + # Tensor-parallel communication is overlapped with dgrad + # GEMM. + # -------------------------------------------------- + inputmat_total = None + inputmat_total_work = None + if ctx.requires_wgrad: + if ctx.fp8 or ctx.debug: + if isinstance(inputmat, QuantizedTensorStorage): + # Input tensor is already quantized + pass + elif ctx.debug or ctx.custom: + # Debug quantizer will be applied immediately before wgrad GEMM + pass + else: + # Quantize input tensor + quantizer = input_quantizer + if quantizer.supports_only_rowwise_all_gather(): + # All-gather is not supported with FP8 column-wise data + quantizer.set_usage( + rowwise=True, + columnwise=not ctx.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) + else: + inputmat = cast_if_needed(inputmat, ctx.activation_dtype) + if ctx.backward_input_needs_gather: + quantizer = None + if ctx.fp8 or ctx.debug: + quantizer = input_quantizer + if quantizer.supports_only_rowwise_all_gather(): + # If data is in FP8, we compute FP8 transposes manually + quantizer.set_usage(rowwise=True, columnwise=False) + else: + # wgrad GEMM requires input with column-wise usage + quantizer.set_usage(rowwise=False, columnwise=True) + if ctx.ub_bulk_dgrad: + inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_dgrad, + inputmat, + quantizer, + ctx.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, + async_op=True, + quantizer=quantizer, + ) + nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") + else: + inputmat_total = inputmat + # -------------------------------------------------- + # Input tensor is ready for computing grad weight... + # -------------------------------------------------- - # 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}" + # -------------------------------------------------- + # Compute grad input tensor + # -------------------------------------------------- - with get_nvtx_range_context("_Linear_backward"): - ( - inputmat, - weight_fp8, - saved_weight, - bias, - ) = restore_from_func_ctx( # pylint: disable=unbalanced-tuple-unpacking - ctx - ) + dgrad = None + dgrad_work = None + if ctx.requires_dgrad: - # Restore from weakref to get original weight python object - # (preserves attributes like main_grad, grad_added_to_main_grad, etc.) - origin_weight_python_object = None - origin_weight_overwrites_main_grad = getattr( - ctx, "origin_weight_overwrites_main_grad", False - ) - main_grad = None - if ctx.fuse_wgrad_accumulation and ctx.requires_wgrad: - origin_weight_ref = ctx.origin_weight_ref - ctx.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" - # Since main_grad can be modified inplace, it should not be a part of saved_tensors - main_grad = ctx.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 - # 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, - inputmat, - weight_fp8, - ) - nvtx_range_pop(f"{nvtx_label}.fsdp_gather") - - # Configure Userbuffers communication (comm+GEMM overlap) - ctx.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: - # 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 - ub_type_dgrad = tex.CommOverlapType.AG - elif ctx.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 - ub_type_dgrad = tex.CommOverlapType.RS - else: - if ctx.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 - ub_type_dgrad = tex.CommOverlapType.AG - if ctx.ub_bulk_wgrad: - # Overlap dgrad reduce-scatter with wgrad compute - ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) - ub_type_wgrad = tex.CommOverlapType.RS - - # -------------------------------------------------- - # Prepare grad output tensor - # Note: Cast to expected dtype and perform tensor-parallel communication - # -------------------------------------------------- - - # Unmodified grad output tensor - grad_output_arg = grad_output - - # Configure quantizer for grad output tensor - # Note: dgrad GEMM requires row-wise usage, wgrad GEMM - # requires column-wise usage - if ctx.grad_output_quantizer is not None: - quantizer = ctx.grad_output_quantizer - quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.ub_overlap_ag: - # Userbuffers only supports communication for one - # tensor usage at a time. Configure quantizer with - # usage for only dgrad GEMM. - quantizer.set_usage(columnwise=False) - - # Adjust the quantization direction approach depending - # 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 - # `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.` + # Make sure required data is available + if isinstance(grad_output, QuantizedTensorStorage): + grad_output.update_usage(rowwise_usage=True) if ( - not ctx.use_bias - and not ctx.requires_wgrad - and ctx.grad_output_quantizer is not None + ctx.fp8 + and weight_quantizer is not None + and isinstance(weight_fp8, QuantizedTensorStorage) ): - ctx.grad_output_quantizer.set_usage(columnwise=False) + weight_fp8.update_usage(columnwise_usage=True) + + # Choose whether to use GEMM kernel with split accumulator + use_split_accumulator = _2X_ACC_DGRAD + if ctx.fp8: + recipe = ctx.fp8_recipe + if hasattr(recipe, "fp8_gemm_dgrad"): + use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator + + # Update grad input quantizer + if grad_input_quantizer is not None: + grad_input_quantizer.set_usage(rowwise=True, columnwise=False) + + # Output buffers for Userbuffers reduce-scatter + gemm_out = None + reduce_scatter_out = None + if ctx.ub_overlap_rs_dgrad: + reduce_scatter_out = torch.empty( + dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device + ) + elif ctx.ub_bulk_wgrad: + gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) - # Prepare grad output tensor - nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") - ( - grad_output, - grad_bias, - ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, + # dgrad GEMM + # Note: dx = dy * w + + nvtx_range_push(f"{nvtx_label}.dgrad_gemm") + weight_for_dgrad = weight_fp8 + if ctx.backward_override == "dequantized": + if isinstance(weight_for_dgrad, QuantizedTensorStorage): + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) + else: + weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) + elif ctx.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) + gemm_out, *_, reduce_scatter_out = general_gemm( + weight_for_dgrad, grad_output, - ctx.parallel_mode == "row", - ctx.grad_output_quantizer, + layout="NN", + grad=True, + quantization_params=grad_input_quantizer, + out=gemm_out, + out_dtype=ctx.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, ) - nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") + nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") + + # Prepare grad input tensor + # Note: Perform tensor-parallel communication + if ctx.ub_overlap_rs_dgrad: + dgrad = reduce_scatter_out + elif ctx.ub_bulk_wgrad: + dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) + elif ctx.parallel_mode == "column" and ctx.tp_size > 1: + nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") + dgrad = gemm_out + if ctx.sequence_parallel: + dgrad, dgrad_work = reduce_scatter_along_first_dim( + dgrad, + ctx.tp_group, + async_op=True, + ) + else: + dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True) + nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") + else: + dgrad = gemm_out + + # -------------------------------------------------- + # Grad input tensor has been computed... + # -------------------------------------------------- - # -------------------------------------------------- - # Grad output tensor is ready for computing grad input... - # -------------------------------------------------- + # -------------------------------------------------- + # Compute grad weight + # -------------------------------------------------- + + wgrad = None + if ctx.requires_wgrad: - # -------------------------------------------------- # Prepare input tensor - # Note: Input tensor is needed for wgrad GEMM. - # Tensor-parallel communication is overlapped with dgrad - # GEMM. - # -------------------------------------------------- - inputmat_total = None - inputmat_total_work = None - if ctx.requires_wgrad: - if ctx.fp8 or ctx.debug: - if isinstance(inputmat, QuantizedTensorStorage): - # Input tensor is already quantized - pass - elif ctx.debug or ctx.custom: - # Debug quantizer will be applied immediately before wgrad GEMM - pass - else: - # Quantize input tensor - quantizer = ctx.input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # All-gather is not supported with FP8 column-wise data - quantizer.set_usage( - rowwise=True, - columnwise=not ctx.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) - else: - inputmat = cast_if_needed(inputmat, ctx.activation_dtype) - if ctx.backward_input_needs_gather: - quantizer = None - if ctx.fp8 or ctx.debug: - quantizer = ctx.input_quantizer - if quantizer.supports_only_rowwise_all_gather(): - # If data is in FP8, we compute FP8 transposes manually - quantizer.set_usage(rowwise=True, columnwise=False) - else: - # wgrad GEMM requires input with column-wise usage - quantizer.set_usage(rowwise=False, columnwise=True) - if ctx.ub_bulk_dgrad: - inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_dgrad, - inputmat, - quantizer, - ctx.tp_group, - ) + # Note: Synchronize tensor-parallel communication and + # make sure required data is available + if inputmat_total_work is not None: + inputmat_total_work.wait() + inputmat_total_work = None + if ctx.fp8 or ctx.debug: + if isinstance(inputmat_total, QuantizedTensorStorage): + inputmat_total.update_usage(columnwise_usage=True) else: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") - inputmat_total, inputmat_total_work = gather_along_first_dim( - inputmat, + input_quantizer.set_usage(rowwise=False, columnwise=True) + inputmat_total = input_quantizer(inputmat_total) + + # 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): + # 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 + # can't reuse the grad output that was gathered + # for the dgrad GEMM. We work around by explicitly + # overlapping the AG operation with the dgrad GEMM. + + # Get the communication stream from the dgrad GEMM to use for the AG + 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) + + grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + + # We use the send stream to copy into the userbuffers. + # This is the same stream that we will use to access the data in the AG, + # so we dont need to add any syncs yet. + with torch.cuda.stream(dgrad_send_stream): + grad_output, _ = fill_userbuffers_buffer_for_all_gather( + ub_obj_overlap_wgrad, + grad_output_arg, + grad_output_quantizer, ctx.tp_group, - async_op=True, - quantizer=quantizer, ) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_input") - else: - inputmat_total = inputmat - # -------------------------------------------------- - # Input tensor is ready for computing grad weight... - # -------------------------------------------------- - - # -------------------------------------------------- - # Compute grad input tensor - # -------------------------------------------------- - dgrad = None - dgrad_work = None - if ctx.requires_dgrad: + # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm + tex.bulk_overlap_ag_with_external_gemm( + ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream + ) - # Make sure required data is available + if ctx.fp8 or ctx.debug: if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(rowwise_usage=True) + grad_output.update_usage(columnwise_usage=True) + else: + grad_output_quantizer.set_usage(rowwise=False, columnwise=True) + grad_output = grad_output_quantizer(grad_output) + + # Figure out whether to use split accumulator + use_split_accumulator = _2X_ACC_WGRAD + if ctx.fp8: + recipe = ctx.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: + 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 + + # 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(): + reduce_scatter_out = torch.empty( + dgrad_shape, dtype=ctx.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 + ), + "quantization_params": grad_weight_quantizer, + "accumulate": ( + accumulate_wgrad_into_param_main_grad + if not origin_weight_overwrites_main_grad + 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), + "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, + } + + def wgrad_gemm( + x: torch.Tensor, + dy: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Perform wgrad GEMM: dw = dy^T * x + + May be fused with bgrad computation. + + May be called outside of this function to enable + some advanced communication/compute overlapping. + + """ + nvtx_range_push(f"{nvtx_label}.wgrad_gemm") + dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) + nvtx_range_pop(f"{nvtx_label}.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 ( - ctx.fp8 - and ctx.weight_quantizer is not None - and isinstance(weight_fp8, QuantizedTensorStorage) + wgrad_gemm_kwargs["ub"] is not None + or wgrad_gemm_kwargs["ub_type"] is not None + or wgrad_gemm_kwargs["extra_output"] is not None + or wgrad_gemm_kwargs["bulk_overlap"] ): - weight_fp8.update_usage(columnwise_usage=True) - - # Choose whether to use GEMM kernel with split accumulator - use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe - if hasattr(recipe, "fp8_gemm_dgrad"): - use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator - - # Update grad input quantizer - if ctx.grad_input_quantizer is not None: - ctx.grad_input_quantizer.set_usage(rowwise=True, columnwise=False) - - # Output buffers for Userbuffers reduce-scatter - gemm_out = None - reduce_scatter_out = None - if ctx.ub_overlap_rs_dgrad: - reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device + raise NotImplementedError( + "Delayed weight grad computation is not supported " + "with Userbuffers (tensor-parallel communication overlapping)" ) - elif ctx.ub_bulk_wgrad: - gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) - - # dgrad GEMM - # Note: dx = dy * w - - nvtx_range_push(f"{nvtx_label}.dgrad_gemm") - weight_for_dgrad = weight_fp8 - if ctx.backward_override == "dequantized": - if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) - else: - weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) - elif ctx.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) - gemm_out, *_, reduce_scatter_out = general_gemm( - weight_for_dgrad, - grad_output, - layout="NN", - grad=True, - quantization_params=ctx.grad_input_quantizer, - out=gemm_out, - out_dtype=ctx.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, - ) - nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") + ctx.wgrad_store.put([inputmat_total, grad_output], wgrad_gemm) + else: - # Prepare grad input tensor - # Note: Perform tensor-parallel communication - if ctx.ub_overlap_rs_dgrad: + # Call wgrad GEMM now + wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) + + # Update grad bias if needed + if grad_bias is None: + grad_bias = grad_bias_ + del grad_bias_ + + # Deallocate tensors if permitted + if ctx.owns_input: + # Input tensor is internal + clear_tensor_data(inputmat_total) + elif ctx.backward_input_needs_gather: + # Gathered input tensor is internal + clear_tensor_data(inputmat_total) + if ctx.parallel_mode == "row" and ctx.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 ub_obj_wgrad.is_fp8_ubuf(): dgrad = reduce_scatter_out - elif ctx.ub_bulk_wgrad: - dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) - elif ctx.parallel_mode == "column" and ctx.tp_size > 1: - nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") - dgrad = gemm_out - if ctx.sequence_parallel: - dgrad, dgrad_work = reduce_scatter_along_first_dim( - dgrad, - ctx.tp_group, - async_op=True, - ) - else: - dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True) - nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") else: - dgrad = gemm_out + dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone() - # -------------------------------------------------- - # Grad input tensor has been computed... - # -------------------------------------------------- + # -------------------------------------------------- + # Grad weight has been computed... + # -------------------------------------------------- - # -------------------------------------------------- - # Compute grad weight - # -------------------------------------------------- - - wgrad = None - if ctx.requires_wgrad: - - # Prepare input tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if inputmat_total_work is not None: - inputmat_total_work.wait() - inputmat_total_work = None - if ctx.fp8 or ctx.debug: - if isinstance(inputmat_total, QuantizedTensorStorage): - inputmat_total.update_usage(columnwise_usage=True) - else: - ctx.input_quantizer.set_usage(rowwise=False, columnwise=True) - inputmat_total = ctx.input_quantizer(inputmat_total) - - # Prepare grad output tensor - # Note: Synchronize tensor-parallel communication and - # make sure required data is available - if ctx.ub_overlap_ag and isinstance(ctx.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 - # can't reuse the grad output that was gathered - # for the dgrad GEMM. We work around by explicitly - # overlapping the AG operation with the dgrad GEMM. - - # Get the communication stream from the dgrad GEMM to use for the AG - 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) - - ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - - # We use the send stream to copy into the userbuffers. - # This is the same stream that we will use to access the data in the AG, - # so we dont need to add any syncs yet. - with torch.cuda.stream(dgrad_send_stream): - grad_output, _ = fill_userbuffers_buffer_for_all_gather( - ub_obj_overlap_wgrad, - grad_output_arg, - ctx.grad_output_quantizer, - ctx.tp_group, - ) - - # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm - tex.bulk_overlap_ag_with_external_gemm( - ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream - ) - - if ctx.fp8 or ctx.debug: - if isinstance(grad_output, QuantizedTensorStorage): - grad_output.update_usage(columnwise_usage=True) - else: - ctx.grad_output_quantizer.set_usage(rowwise=False, columnwise=True) - grad_output = ctx.grad_output_quantizer(grad_output) - - # Figure out whether to use split accumulator - use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.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: - 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 - - # 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(): - reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device - ) + # Don't return grad bias if not needed + if not ctx.use_bias: + grad_bias = None - # Arguments to include in wgrad GEMM closure - wgrad_gemm_kwargs = { - "out_dtype": ( - main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype - ), - "quantization_params": ctx.grad_weight_quantizer, - "accumulate": ( - accumulate_wgrad_into_param_main_grad - if not origin_weight_overwrites_main_grad - 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), - "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, - } - - def wgrad_gemm( - x: torch.Tensor, - dy: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Perform wgrad GEMM: dw = dy^T * x - - May be fused with bgrad computation. - - May be called outside of this function to enable - some advanced communication/compute overlapping. - - """ - nvtx_range_push(f"{nvtx_label}.wgrad_gemm") - dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) - nvtx_range_pop(f"{nvtx_label}.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 ( - wgrad_gemm_kwargs["ub"] is not None - or wgrad_gemm_kwargs["ub_type"] is not None - or wgrad_gemm_kwargs["extra_output"] is not None - or wgrad_gemm_kwargs["bulk_overlap"] - ): - raise NotImplementedError( - "Delayed weight grad computation is not supported " - "with Userbuffers (tensor-parallel communication overlapping)" - ) - ctx.wgrad_store.put([inputmat_total, grad_output], wgrad_gemm) - else: + # Make sure all tensor-parallel communication is finished + if inputmat_total_work is not None: + inputmat_total_work.wait() + inputmat_total_work = None + if dgrad_work is not None: + dgrad_work.wait() + dgrad_work = None - # Call wgrad GEMM now - wgrad, grad_bias_ = wgrad_gemm(inputmat_total, grad_output) - - # Update grad bias if needed - if grad_bias is None: - grad_bias = grad_bias_ - del grad_bias_ - - # Deallocate tensors if permitted - if ctx.owns_input: - # Input tensor is internal - clear_tensor_data(inputmat_total) - elif ctx.backward_input_needs_gather: - # Gathered input tensor is internal - clear_tensor_data(inputmat_total) - if ctx.parallel_mode == "row" and ctx.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 ub_obj_wgrad.is_fp8_ubuf(): - dgrad = reduce_scatter_out - else: - dgrad = ub_obj_wgrad.get_buffer(local_chunk=True).clone() + if ctx.requires_wgrad: + # Handle custom DDP from mcore. + if ctx.fuse_wgrad_accumulation and hasattr( + origin_weight_python_object, "grad_added_to_main_grad" + ): + origin_weight_python_object.grad_added_to_main_grad = True + if getattr(origin_weight_python_object, "zero_out_wgrad", False): + wgrad = get_dummy_wgrad( + list(main_grad.shape), + origin_weight_python_object.dtype, + zero=True, + ) + else: + wgrad = get_dummy_wgrad( + list(main_grad.shape), + origin_weight_python_object.dtype, + ) + elif ctx.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) + return ( + wgrad, + None, # weight_workspace + dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + grad_bias, + None, + None, + None, + None, + None, + None, + None, + ) - # -------------------------------------------------- - # Grad weight has been computed... - # -------------------------------------------------- - # Don't return grad bias if not needed - if not ctx.use_bias: - grad_bias = None +class _Linear(torch.autograd.Function): + """Linear semi-top level module + Calls custom cuda extensions. + """ - # Make sure all tensor-parallel communication is finished - if inputmat_total_work is not None: - inputmat_total_work.wait() - inputmat_total_work = None - if dgrad_work is not None: - dgrad_work.wait() - dgrad_work = None + @staticmethod + 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], + ) -> 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, + ) + ) + if ctx is not None: + _linear_setup_ctx( + ctx, + tensors_to_save, + tensor_objects, + 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, + ) + 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 - if ctx.requires_wgrad: - # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( - origin_weight_python_object, "grad_added_to_main_grad" - ): - origin_weight_python_object.grad_added_to_main_grad = True - if getattr(origin_weight_python_object, "zero_out_wgrad", False): - wgrad = get_dummy_wgrad( - list(main_grad.shape), - origin_weight_python_object.dtype, - zero=True, - ) - else: - wgrad = get_dummy_wgrad( - list(main_grad.shape), - origin_weight_python_object.dtype, - ) - elif ctx.fuse_wgrad_accumulation: - wgrad = None - else: - wgrad = None + return out, new_weight_workspace - # Update FP8 scaling factors if needed + @staticmethod + def backward( + ctx, grad_output: torch.Tensor, _grad_weight_workspace + ) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward pass: compute gradients and reduce FP8 scaling factors.""" + 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(): 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") - - # Scatter fp8 weight buffers - if ctx.fp8 and not ctx.is_weight_param_quantized: - _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) - return ( - wgrad, - None, # weight_workspace - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, - grad_bias, - None, - ) + return result class Linear(TransformerEngineBaseModule): @@ -1486,17 +1602,37 @@ def forward( self._fp8_workspaces.get(cache_name) if cache_name is not None else None ) + if self.fp8: + backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override + else: + backward_override = None + custom = is_custom(input_quantizer) or is_custom(weight_quantizer) + backward_input_needs_gather = ( + weight_tensor.requires_grad + and self.parallel_mode == "column" + and self.sequence_parallel + ) + + if debug: + ub_overlap_rs_fprop = False + ub_overlap_ag_dgrad = False + ub_overlap_ag_fprop = False + ub_overlap_rs_dgrad = False + ub_bulk_dgrad = False + ub_bulk_wgrad = False + else: + ub_overlap_rs_fprop = self.ub_overlap_rs_fprop + ub_overlap_ag_dgrad = self.ub_overlap_ag_dgrad + ub_overlap_ag_fprop = self.ub_overlap_ag_fprop + ub_overlap_rs_dgrad = self.ub_overlap_rs_dgrad + 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, - input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, self.fuse_wgrad_accumulation, is_cpu_offload_enabled(), self.tp_group, @@ -1506,12 +1642,12 @@ def forward( self.activation_dtype, self.parallel_mode, is_grad_enabled, - self.ub_overlap_rs_fprop, - self.ub_overlap_ag_dgrad, - self.ub_overlap_ag_fprop, - self.ub_overlap_rs_dgrad, - self.ub_bulk_dgrad, - self.ub_bulk_wgrad, + 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, @@ -1520,6 +1656,9 @@ def forward( self.symmetric_ar_type, self.save_original_input, debug, + backward_override, + custom, + backward_input_needs_gather, ) out, new_weight_workspace = linear_fn( *autograd_ctx, @@ -1528,6 +1667,12 @@ def forward( 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, ) if new_weight_workspace is not None and cache_name is not None: From 3a78e154c7abf1f4db0371cdc6bcbdc0aafa7a01 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 16 Apr 2026 08:49:18 -0400 Subject: [PATCH 354/521] [PyTorch] Add method for mcore to register wgrad accumulation hook (#2886) Fix delay wgrad mcore integration Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Gao Deng --- .../pytorch/ops/basic/grouped_linear.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e21625276c..a1d40a30ec 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -114,6 +114,7 @@ def __init__( self.num_extra_inputs = 2 self.wgrad_store = WeightGradStore(delay_wgrad_compute) + self.wgrad_accumulation_and_reduce_hooks: list = [] # Weight tensor dimensions self.num_groups: int = num_groups @@ -193,6 +194,23 @@ def _apply_delay_wgrad_param_hooks(self) -> None: for group_idx in range(self.num_groups): getattr(self, f"weight{group_idx}").skip_backward_post_hook = True + def register_wgrad_accumulation_and_reduce_hooks( + self, wgrad_accumulation_and_reduce_hook: Callable + ) -> None: + """Register a hook to run after delayed wgrad computation completes. + + Mirrors ``TransformerEngineBaseModule.register_wgrad_accumulation_and_reduce_hooks`` + so that DDP can wire its ``param.grad = None`` / reduce-scatter callback here + instead of directly on the AccumulateGrad node (which is bypassed when + ``skip_backward_post_hook`` is set). + """ + self.wgrad_accumulation_and_reduce_hooks.append(wgrad_accumulation_and_reduce_hook) + + def _trigger_wgrad_accumulation_and_reduce_hooks(self) -> None: + """Call all registered wgrad accumulation and reduce hooks.""" + for hook in self.wgrad_accumulation_and_reduce_hooks: + hook() + def need_backward_dw(self) -> bool: """Return whether :meth:`backward_dw` must run to finish weight gradients.""" return self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute() @@ -217,6 +235,7 @@ def backward_dw(self) -> None: activations.columnwise_scale_inv, ) if self._accumulate_into_main_grad: + self._trigger_wgrad_accumulation_and_reduce_hooks() return if self.single_grouped_weight: if isinstance(grad_weights, list): @@ -231,6 +250,7 @@ def backward_dw(self) -> None: for group_idx in range(self.num_groups): w = getattr(self, f"weight{group_idx}") w.grad = grad_weights[group_idx].to(w.dtype) + self._trigger_wgrad_accumulation_and_reduce_hooks() def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]: """Retrieve per-group bias tensors in the given dtype.""" From c9035a4854edb67dc896ed0e129e3e7ecbf52251 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 16 Apr 2026 08:49:35 -0400 Subject: [PATCH 355/521] [PyTorch] Minor optimizations in fused grouped MLP (#2888) Minor misc optimizations in fused GroupedMLP Signed-off-by: Kirthi Shankar Sivamani --- .../pytorch/ops/fused/backward_grouped_mlp.py | 31 +++++++------------ .../pytorch/ops/fused/forward_grouped_mlp.py | 10 ++---- 2 files changed, 14 insertions(+), 27 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 096e65d296..3eb57c3563 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -69,20 +69,17 @@ def _cudnn_compute_wgrad( sfb_tensor = grouped_x.columnwise_scale_inv.view(in_features, -1).view( dtype=torch.float8_e8m0fnu ) - offsets_tensor = offsets.to(dtype=torch.int32) # Prepare wgrad output if single_grouped_weight: # Dense mode: single (num_groups, out_features, in_features) tensor - wgrad_tensor = wgrad_output.rowwise_data.view( - offsets_tensor.shape[0], out_features, in_features - ) + wgrad_tensor = wgrad_output.rowwise_data.view(offsets.shape[0], out_features, in_features) wgrad_kernel_fn( a_tensor=a_tensor, b_tensor=b_tensor, sfa_tensor=sfa_tensor, sfb_tensor=sfb_tensor, - offsets_tensor=offsets_tensor, + offsets_tensor=offsets, output_mode="dense", wgrad_tensor=wgrad_tensor, acc_dtype=torch.float32, @@ -99,7 +96,7 @@ def _cudnn_compute_wgrad( b_tensor=b_tensor, sfa_tensor=sfa_tensor, sfb_tensor=sfb_tensor, - offsets_tensor=offsets_tensor, + offsets_tensor=offsets, output_mode="discrete", wgrad_ptrs=wgrad_ptrs, acc_dtype=torch.float32, @@ -210,6 +207,7 @@ def _compute_grad_params( # Launch or defer the GEMM delay_wgrad = fc_op.wgrad_store is not None and fc_op.wgrad_store.delay_wgrad_compute() if cudnn_wgrad_kernel_fn is not None: + offsets = offsets if offsets.dtype == torch.int32 else offsets.to(dtype=torch.int32) gemm_fn = functools.partial( _cudnn_compute_wgrad, weight_shape=weight_shape, @@ -424,8 +422,6 @@ def fuser_backward( # Group splits 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) - split_points = split_points.to(dtype=torch.int, device=device) scale_bias = fc2_op._scale_bias and fc2_op.has_bias grouped_fc1_x = None @@ -516,7 +512,8 @@ def fuser_backward( norm_const_tensor = get_cached_ones_tensor(1, dtype, device) current_stream = torch.cuda.current_stream().cuda_stream - scales_tensor = scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + 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 = { @@ -594,7 +591,6 @@ def fuser_backward( if scale_bias: fc2_biases = fc2_op._get_bias_tensors(dtype) bias_packed = torch.stack(fc2_biases) - scales_f32 = scales.detach().to(dtype=torch.float32) fc2_dbias_packed_result, grad_scales = _compute_grouped_dbias_dscales( fc2_dy, scales_f32, @@ -608,12 +604,11 @@ def fuser_backward( else: fc2_bias_grads = [fc2_dbias_packed_result[idx] for idx in range(num_groups)] elif fc2_dbias_packed is not None: + fc2_dbias_packed = fc2_dbias_packed.to(dtype=dtype) if fc2_op.single_grouped_bias: - fc2_bias_grad_packed = fc2_dbias_packed.to(dtype=dtype) + fc2_bias_grad_packed = fc2_dbias_packed else: - fc2_bias_grads = [ - fc2_dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups) - ] + fc2_bias_grads = [fc2_dbias_packed[idx] for idx in range(num_groups)] grad_scales = grad_scales.to(dtype=dtype) @@ -622,13 +617,11 @@ def fuser_backward( if fc1_op.has_bias: dbias_t = fc2_dgrad_kernel_out["dbias_tensor"] if dbias_t is not None: - dbias_2d = dbias_t.squeeze(-1) + dbias_2d = dbias_t.squeeze(-1).to(dtype=dtype) if fc1_op.single_grouped_bias: - fc1_bias_grad_packed = dbias_2d.to(dtype=dtype) + fc1_bias_grad_packed = dbias_2d else: - fc1_bias_grads = [ - dbias_2d[group_idx].to(dtype=dtype) for group_idx in range(num_groups) - ] + 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] diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 4e756ea531..90c4204f06 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -194,14 +194,8 @@ 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) - split_points = torch.cumsum(split_sizes, 0, dtype=torch.int) - split_points_offsets = torch.cumsum(split_sizes, 0) - base_offsets = torch.cat( - [ - torch.zeros(1, device=split_sizes.device, dtype=split_sizes.dtype), - split_points_offsets, - ] - ) + 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] From 58a008f1144cb61dcc04d457509bcb7d92021617 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 16 Apr 2026 15:26:35 -0400 Subject: [PATCH 356/521] [PyTorch] Add test to compare single vs multi-param fused GMLP (#2893) * Add new test to compare single vs multi-param fused GMLP case Signed-off-by: Kirthi Shankar Sivamani * Add bias support Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fusible_ops.py | 215 ++++++++++++++++++++++++++++++ 1 file changed, 215 insertions(+) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 0dfa8b5f45..0f40e92183 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -3669,6 +3669,221 @@ def test_grouped_mlp( assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols) assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols) + @pytest.mark.parametrize( + "dtype", + tuple(dtype for dtype in _dtypes if dtype in (torch.float16, torch.bfloat16)), + ) + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_clamped_qgeglu")) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_grouped_mlp_single_weight_numerics( + self, + *, + dtype: torch.dtype, + bias: bool, + activation: str, + device: torch.device = "cuda", + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + ) -> None: + """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP.""" + + 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") + 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) + split_sizes = torch.tensor(split_sizes, dtype=torch.int64, device=device) + in_shape = (split_sizes.sum().item(), hidden_size) + recipe = make_recipe("mxfp8") + + 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) + ] + fc1_bs_base = ( + [ + torch.empty((2 * hidden_size,), device=device, dtype=dtype).uniform_(-0.5, 0.5) + for _ in range(group_size) + ] + if bias + else None + ) + fc2_bs_base = ( + [ + torch.empty((hidden_size,), device=device, dtype=dtype).uniform_(-0.5, 0.5) + for _ in range(group_size) + ] + if bias + else None + ) + + def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]: + with te.quantized_model_init(enabled=True, recipe=recipe): + 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) + ) + fc1 = te_ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=bias, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + ) + fc2 = te_ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + scale_bias=bias, + ) + module = te_ops.Sequential(fc1, scaled_act, fc2) + + with torch.no_grad(): + if single_grouped_weight: + fc1_weights = fc1.weight.quantized_tensors + if fc1_weights is None: + fc1_weights = fc1.weight.split_into_quantized_tensors() + fc2_weights = fc2.weight.quantized_tensors + if fc2_weights is None: + fc2_weights = fc2.weight.split_into_quantized_tensors() + for group_idx in range(group_size): + if single_grouped_weight: + fc1_weights[group_idx].copy_(fc1_ws_base[group_idx]) + fc2_weights[group_idx].copy_(fc2_ws_base[group_idx]) + else: + getattr(fc1, f"weight{group_idx}").copy_(fc1_ws_base[group_idx]) + getattr(fc2, f"weight{group_idx}").copy_(fc2_ws_base[group_idx]) + if bias: + getattr(fc1, f"bias{group_idx}").copy_(fc1_bs_base[group_idx]) + getattr(fc2, f"bias{group_idx}").copy_(fc2_bs_base[group_idx]) + + x = x_base.detach().clone().requires_grad_(True) + probs = probs_base.detach().clone().requires_grad_(True) + dy = dy_base.detach().clone() + + with te.autocast(enabled=True, recipe=recipe): + fc2_extra = (split_sizes, probs) if bias else (split_sizes,) + y = module(x, split_sizes, probs, *fc2_extra) + y.backward(dy) + + forward_ops = module._module_groups[0]._forward_ops + backward_ops = module._module_groups[0]._backward_ops + assert len(forward_ops) == 1 + assert isinstance( + forward_ops[0][0], + te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + ) + assert len(backward_ops) == 1 + assert isinstance( + backward_ops[0][0], + te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + ) + + if single_grouped_weight: + fc1_dw = fc1.weight.grad.detach().clone() + fc2_dw = fc2.weight.grad.detach().clone() + else: + fc1_dw = torch.stack( + [ + getattr(fc1, f"weight{group_idx}").grad.detach().clone() + for group_idx in range(group_size) + ], + dim=0, + ) + fc2_dw = torch.stack( + [ + getattr(fc2, f"weight{group_idx}").grad.detach().clone() + for group_idx in range(group_size) + ], + dim=0, + ) + + fc1_db = None + fc2_db = None + if bias: + fc1_db = torch.stack( + [ + getattr(fc1, f"bias{group_idx}").grad.detach().clone() + for group_idx in range(group_size) + ], + dim=0, + ) + fc2_db = torch.stack( + [ + getattr(fc2, f"bias{group_idx}").grad.detach().clone() + for group_idx in range(group_size) + ], + dim=0, + ) + + return ( + y.detach().clone(), + x.grad.detach().clone(), + probs.grad.detach().clone(), + fc1_dw, + fc2_dw, + fc1_db, + fc2_db, + ) + + ( + y_false, + dx_false, + dprobs_false, + fc1_dw_false, + fc2_dw_false, + fc1_db_false, + fc2_db_false, + ) = _run_case(False) + ( + y_true, + dx_true, + dprobs_true, + fc1_dw_true, + fc2_dw_true, + fc1_db_true, + fc2_db_true, + ) = _run_case(True) + + torch.testing.assert_close(y_false, y_true, rtol=0, atol=0) + torch.testing.assert_close(dx_false, dx_true, rtol=0, atol=0) + torch.testing.assert_close(dprobs_false, dprobs_true, rtol=0, atol=0) + torch.testing.assert_close(fc1_dw_false, fc1_dw_true, rtol=0, atol=0) + torch.testing.assert_close(fc2_dw_false, fc2_dw_true, rtol=0, atol=0) + if bias: + bias_tols = {"rtol": 0.05, "atol": 0.015625} + 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("dtype", _dtypes) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) From 1e9e48c66dc48077bb180fdc7f800da5e65e3295 Mon Sep 17 00:00:00 2001 From: harry zhou <67385896+harryzhou2000@users.noreply.github.com> Date: Fri, 17 Apr 2026 06:08:41 +0800 Subject: [PATCH 357/521] [Common] Fix fused router for large top-K and expert counts (#2821) * fix: enabling fused _router to be able to handle large topk and number of experts - expanding shared memory when needed - switch to radix topk selection when topk is large - test_fused_router.py updated with large num experts and tolerances refined for different cases * added topk>=16 in tests/pytorch/test_fused_router.py added return value check of cudaFuncSetAttribute in transformer_engine/common/fused_router/fused_topk_with_score_function.cu added dtype dependent eps in tests/pytorch/test_fused_router.py removed unneeded code in transformer_engine/common/fused_router/utils.h * test_fused_router.py needs to skip topk >= num_experts case Signed-off-by: Harry Zhou cleaned up raw warp operations added comments added shared_memory check added return code check * warning about dtype for tolerance in test_fused_router.py Signed-off-by: Harry Zhou --------- Signed-off-by: Harry Zhou Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_fused_router.py | 70 +++-- .../fused_score_for_moe_aux_loss.cu | 32 ++- .../fused_topk_with_score_function.cu | 39 ++- .../common/fused_router/utils.h | 258 +++++++++++++++++- transformer_engine/common/utils.cuh | 26 ++ 5 files changed, 382 insertions(+), 43 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 36c09060ed..274a35b81d 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -17,6 +17,30 @@ torch.cuda.manual_seed(seed) +def _get_tolerances(dtype: torch.dtype, num_experts: int): + """Return (atol, rtol) scaled by the number of experts. + + With many experts the fused and reference kernels accumulate + floating-point reductions (e.g. normalization sums) in different + orders, causing O(num_experts * machine_eps) rounding divergence. + Scale the default tolerances accordingly so that small expert + counts keep tight checks while large counts (1024+) get the + headroom they need. + """ + # Default tolerances for torch.testing.assert_close + base_atol, base_rtol = 1e-5, 1.3e-6 + # TODO: account for fp16, bf16 as dtype + if dtype != torch.float32: + raise NotImplementedError("tolerances implemented for fp32 only") + eps = 2e-7 + # The worst-case rounding error from summing N values is O(N * eps). + # Use 2 * num_experts * eps as the tolerance floor so tests pass for + # large expert counts while remaining tight for small ones. + atol = max(base_atol, 2 * num_experts * eps) + rtol = max(base_rtol, 2 * num_experts * eps) + return atol, rtol + + # Pytorch-based group topk def group_limited_topk( scores: torch.Tensor, @@ -153,6 +177,13 @@ def run_comparison( score_function, enable_bias, ): + if topk >= num_experts: + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") + if group_topk is not None and num_groups is not None: + group_size = num_experts // num_groups + per_group_topk = topk // group_topk + if per_group_topk >= group_size: + pytest.skip(f"per-group topk ({per_group_topk}) >= group_size ({group_size})") # Set some parameters if score_function in ("sigmoid", "sqrtsoftplus"): # Construct logits with a narrow range to avoid very small activation values, @@ -215,7 +246,8 @@ def run_comparison( expert_bias=expert_bias_clone, ) - torch.testing.assert_close(probs, probs_fused) + atol, rtol = _get_tolerances(dtype, num_experts) + torch.testing.assert_close(probs, probs_fused, atol=atol, rtol=rtol) torch.testing.assert_close(routing_map, routing_map_fused) # Fake the loss @@ -227,13 +259,13 @@ def run_comparison( loss_fused.backward() # Check the gradient - torch.testing.assert_close(logits.grad, logits_clone.grad) + torch.testing.assert_close(logits.grad, logits_clone.grad, atol=atol, rtol=rtol) @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168, 8992]) -@pytest.mark.parametrize("num_experts", [128, 32]) -@pytest.mark.parametrize("topk", [4, 8]) +@pytest.mark.parametrize("num_experts", [1024, 128, 32]) +@pytest.mark.parametrize("topk", [4, 8, 16, 32]) @pytest.mark.parametrize("group_topk", [None, 4]) @pytest.mark.parametrize("scaling_factor", [None, 1.2]) @pytest.mark.parametrize("enable_bias", [True, False]) @@ -263,8 +295,8 @@ def test_topk_sigmoid( @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168, 8992]) -@pytest.mark.parametrize("num_experts", [128, 32]) -@pytest.mark.parametrize("topk", [4, 8]) +@pytest.mark.parametrize("num_experts", [1024, 128, 32]) +@pytest.mark.parametrize("topk", [4, 8, 16, 32]) @pytest.mark.parametrize("group_topk", [None, 4]) @pytest.mark.parametrize("scaling_factor", [None, 1.2]) @pytest.mark.parametrize("enable_bias", [True, False]) @@ -294,8 +326,8 @@ def test_topk_sqrtsoftplus( @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) -@pytest.mark.parametrize("num_experts", [128, 32]) -@pytest.mark.parametrize("topk", [4, 8]) +@pytest.mark.parametrize("num_experts", [1024, 128, 32]) +@pytest.mark.parametrize("topk", [4, 8, 16, 32]) @pytest.mark.parametrize("use_pre_softmax", [True, False]) @pytest.mark.parametrize("group_topk", [None, 4]) @pytest.mark.parametrize("scaling_factor", [None, 1.2]) @@ -325,10 +357,12 @@ def test_topk_softmax( @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168]) -@pytest.mark.parametrize("num_experts", [256, 128, 32]) -@pytest.mark.parametrize("topk", [1, 4, 8]) +@pytest.mark.parametrize("num_experts", [1024, 256, 128, 32]) +@pytest.mark.parametrize("topk", [1, 4, 8, 16, 32]) @pytest.mark.parametrize("score_function", ["softmax", "sigmoid", "sqrtsoftplus"]) def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_function): + if topk >= num_experts: + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") if score_function in ("sigmoid", "sqrtsoftplus"): # Construct logits with a narrow range to avoid very small activation values offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 @@ -364,7 +398,8 @@ def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_f score_function=score_function, ) - torch.testing.assert_close(scores, scores_fused) + atol, rtol = _get_tolerances(dtype, num_experts) + torch.testing.assert_close(scores, scores_fused, atol=atol, rtol=rtol) torch.testing.assert_close(routing_map, routing_map_fused) loss = torch.sum(scores) @@ -372,14 +407,16 @@ def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_f loss_fused = torch.sum(scores_fused) loss_fused.backward() - torch.testing.assert_close(logits.grad, logits_clone.grad) + torch.testing.assert_close(logits.grad, logits_clone.grad, atol=atol, rtol=rtol) @pytest.mark.parametrize("dtype", [torch.float32]) @pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) -@pytest.mark.parametrize("num_experts", [256, 128, 32]) -@pytest.mark.parametrize("topk", [4]) +@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): + if topk >= num_experts: + pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") # 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 @@ -411,13 +448,14 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk): coeff=coeff, ) - torch.testing.assert_close(aux_loss, aux_loss_fused) + atol, rtol = _get_tolerances(dtype, num_experts) + torch.testing.assert_close(aux_loss, aux_loss_fused, atol=atol, rtol=rtol) # Backward aux_loss.backward() aux_loss_fused.backward() - torch.testing.assert_close(probs.grad, probs_clone.grad) + torch.testing.assert_close(probs.grad, probs_clone.grad, atol=atol, rtol=rtol) def profile_topk_softmax( 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 ebdcb293e0..4eb4240d7c 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 @@ -16,7 +16,7 @@ namespace transformer_engine { namespace fused_router { -template +template __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logits, int num_tokens, int num_experts, int topk, int score_function, float *scores, @@ -123,7 +123,7 @@ __global__ void fused_score_for_moe_aux_loss_forward_kernel(const DataType *logi * Section: Topk * Get the topk indices */ - naive_topk_and_mask(local_logits, num_experts, topk, topk_indices, topk_logits, lane_id); + topk_and_mask(local_logits, num_experts, topk, topk_indices, topk_logits, lane_id); __syncwarp(); // Write the routing_map to the output tensor @@ -149,10 +149,26 @@ void fused_score_for_moe_aux_loss_forward_kernel_launcher( size_t shared_memory_size = num_experts * num_token_per_block * sizeof(CompType) // logits + topk * num_token_per_block * sizeof(CompType) // topk_logits + topk * num_token_per_block * sizeof(int); // topk_indices - fused_score_for_moe_aux_loss_forward_kernel - <<>>( - logits, num_tokens, num_experts, topk, score_function, scores, routing_map, - intermediate_output); + check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); + // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float; + // switch at K=16 where naive O(K^2*E) starts to dominate + if (topk < 16) { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + fused_score_for_moe_aux_loss_forward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_score_for_moe_aux_loss_forward_kernel + <<>>( + logits, num_tokens, num_experts, topk, score_function, scores, routing_map, + intermediate_output); + } else { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + fused_score_for_moe_aux_loss_forward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_score_for_moe_aux_loss_forward_kernel + <<>>( + logits, num_tokens, num_experts, topk, score_function, scores, routing_map, + intermediate_output); + } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -305,6 +321,10 @@ void fused_score_for_moe_aux_loss_backward_kernel_launcher( + num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd + num_experts * num_token_per_block * sizeof(CompType); // comp_buf + check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); + NVTE_CHECK_CUDA(cudaFuncSetAttribute(fused_score_for_moe_aux_loss_backward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + shared_memory_size)); fused_score_for_moe_aux_loss_backward_kernel <<>>( intermediate_output, grad_scores, num_tokens, num_experts, topk, score_function, 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 1bed871de8..9f7a830546 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 @@ -10,13 +10,12 @@ #include "../common.h" #include "../util/logging.h" -#include "../utils.cuh" #include "utils.h" namespace transformer_engine { namespace fused_router { -template +template __global__ void fused_topk_with_score_function_forward_kernel( const DataType *logits, int num_tokens, int num_experts, int topk, bool use_pre_softmax, int num_groups, int group_topk, float scaling_factor, int score_function, @@ -146,7 +145,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( int group_size = num_experts / num_groups; // Top2 for (int i = 0; i < num_groups; i++) { - naive_topk_and_mask( + topk_and_mask( /*scores ptr = */ scores + i * group_size, /*data size = */ group_size, /*topk = */ topk / group_topk, @@ -166,7 +165,7 @@ __global__ void fused_topk_with_score_function_forward_kernel( } // select the topk groups - naive_topk_and_mask( + topk_and_mask( /*scores ptr = */ group_scores, /*data size = */ num_groups, /*topk = */ group_topk, @@ -183,10 +182,10 @@ __global__ void fused_topk_with_score_function_forward_kernel( } } __syncwarp(); - naive_topk_and_mask(masked_scores, num_experts, topk, topk_indices, topk_scores, lane_id); + topk_and_mask(masked_scores, num_experts, topk, topk_indices, topk_scores, lane_id); } else { - naive_topk_and_mask(scores, num_experts, topk, topk_indices, topk_scores, lane_id); + topk_and_mask(scores, num_experts, topk, topk_indices, topk_scores, lane_id); } __syncwarp(); @@ -254,10 +253,26 @@ void fused_topk_with_score_function_forward_kernel_launcher( shared_memory_size += num_groups * num_token_per_block * sizeof(CompType); // group_scores shared_memory_size += num_experts * num_token_per_block * sizeof(CompType); // maksed_scores } - fused_topk_with_score_function_forward_kernel - <<>>( - logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, - scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output); + check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); + // Radix selection is O(E), independent of K, but it needs 4 passes for 32-bit float; + // switch at K=16 where naive O(K^2*E) starts to dominate + if (topk < 16) { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + fused_topk_with_score_function_forward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_topk_with_score_function_forward_kernel + <<>>( + logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, + scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output); + } else { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + fused_topk_with_score_function_forward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, shared_memory_size)); + fused_topk_with_score_function_forward_kernel + <<>>( + logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, + scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output); + } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -467,6 +482,10 @@ void fused_topk_with_score_function_backward_kernel_launcher( num_experts * num_token_per_block * sizeof(CompType) // act_from_fwd + num_experts * num_token_per_block * sizeof(CompType) // comp_buf + num_experts * num_token_per_block * sizeof(bool); // routing_map + check_shared_memory_capacity_num_experts(shared_memory_size, num_experts); + NVTE_CHECK_CUDA(cudaFuncSetAttribute(fused_topk_with_score_function_backward_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, + shared_memory_size)); fused_topk_with_score_function_backward_kernel <<>>( routing_map, intermediate_output, grad_probs, num_tokens, num_experts, topk, diff --git a/transformer_engine/common/fused_router/utils.h b/transformer_engine/common/fused_router/utils.h index 372efdc490..08ad3d16a6 100644 --- a/transformer_engine/common/fused_router/utils.h +++ b/transformer_engine/common/fused_router/utils.h @@ -7,11 +7,26 @@ #ifndef TRANSFORMER_ENGINE_FUSED_ROUTER_UTILS_H_ #define TRANSFORMER_ENGINE_FUSED_ROUTER_UTILS_H_ +#include "../util/logging.h" +#include "../utils.cuh" #include "transformer_engine/transformer_engine.h" namespace transformer_engine { namespace fused_router { +// Check if requested shared memory size exceeds device capacity. +// Throws an error with num_experts info to help users diagnose the issue. +inline void check_shared_memory_capacity_num_experts(size_t shared_memory_size, int num_experts) { + int device_id; + NVTE_CHECK_CUDA(cudaGetDevice(&device_id)); + int max_smem_per_block; + NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&max_smem_per_block, + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id)); + NVTE_CHECK(shared_memory_size <= static_cast(max_smem_per_block), "Shared memory size (", + shared_memory_size, " bytes) exceeds device capacity (", max_smem_per_block, + " bytes). Try reducing num_experts (currently ", num_experts, ")."); +} + // Using FP32 to handle all the calculations. // Currently, only FP32 is supported because // 1. The score functions (sigmoid, softmax, sqrtsoftplus) are implemented in FP32. @@ -51,7 +66,7 @@ __device__ inline T warp_reduce_on_shmem(T *data_ptr, int data_size, ReduceFuncT default_val = -std::numeric_limits::infinity(); } - // Some value is hanlded in local thread + // Some value is handled in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread CompType val = lane_id < data_size ? data_ptr[lane_id] : default_val; @@ -82,7 +97,7 @@ __device__ inline T masked_warp_reduce_on_shmem(T *data_ptr, bool *mask, int dat default_val = -std::numeric_limits::infinity(); } - // Some value is hanlded in local thread + // Some value is handled in local thread // Thread 0 is responsible for the: 0-th, 32-th, 64-th, 96-th ... // Reduce the value in local thread CompType val = lane_id < data_size && mask[lane_id] ? data_ptr[lane_id] : default_val; @@ -187,22 +202,233 @@ __device__ inline void apply_softmax_bwd_on_float(float *grad, float *fwd_output } __device__ inline void apply_softmax_on_float(float *scores, int data_size, int lane_id) { - // 1. compute the max of value - float max_val = warp_reduce_on_shmem(scores, data_size, ReduceFuncType::MAX, lane_id); - // 2. value -> exp_value + // --- Pass 1: Online accumulation of max and sum_exp --- + float local_max = -std::numeric_limits::infinity(); + float local_sum = 0.0f; + for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { - scores[i] = expf(scores[i] - max_val); + float val = scores[i]; + if (val > local_max) { + // Rescale accumulated sum for the new max + local_sum *= expf(local_max - val); + local_max = val; + } + local_sum += expf(val - local_max); } - __syncwarp(); - // 3. compute the sum of exp_value - float sum_val = warp_reduce_on_shmem(scores, data_size, ReduceFuncType::SUM, lane_id); - // 4. update the softmax value + + // Warp-level reduction of (max, sum_exp) across 32 lanes. + // When merging two lanes with (max_a, sum_a) and (max_b, sum_b): + // merged_max = max(max_a, max_b) + // merged_sum = sum_a * exp(max_a - merged_max) + sum_b * exp(max_b - merged_max) + // + // NaN guard: when data_size < 32, some lanes have (max=-inf, sum=0). + // Merging two such lanes computes expf(-inf - (-inf)) = expf(NaN) = NaN, + // and 0.0 * NaN = NaN in IEEE 754, contaminating valid lanes. + // Fix: treat -inf max as "no data" and skip the expf computation. +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + float other_max = warp_shuffle_xor(local_max, offset); + float other_sum = warp_shuffle_xor(local_sum, offset); + float new_max = fmaxf(local_max, other_max); + if (new_max > -std::numeric_limits::infinity()) { + // At least one side has real data; safe to compute expf differences + float my_scale = + (local_max > -std::numeric_limits::infinity()) ? expf(local_max - new_max) : 0.0f; + float other_scale = + (other_max > -std::numeric_limits::infinity()) ? expf(other_max - new_max) : 0.0f; + local_sum = local_sum * my_scale + other_sum * other_scale; + } + // else: both sides are -inf (no data), keep local_sum = 0 + local_max = new_max; + } + // After reduction, all lanes have the same (local_max, local_sum) + + // --- Pass 2: Normalize in-place --- + float inv_sum = 1.0f / local_sum; for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { - scores[i] = scores[i] / sum_val; + scores[i] = expf(scores[i] - local_max) * inv_sum; } __syncwarp(); } +enum class TopkFuncType { + Naive = 0, + Radix = 1, +}; + +/******************************************************************************* + * radix_topk_and_mask — Warp-level radix-selection based top-K + * + * O(E) algorithm independent of K, adapted from PyTorch's radix selection. + * Uses 4-bit radix (16 buckets) → 8 passes for float32. + * + * Algorithm: + * Phase 1 — Radix selection (8 passes): + * Convert float scores to "order-preserving" uint32 (flip sign bit for + * positives, flip all bits for negatives). Then iterate 4 bits at a time + * from the MSB. Each pass: + * 1. Each of 32 threads counts elements per radix bucket that match the + * "desired" bit pattern found so far. + * 2. Warp-reduce the per-thread histograms (16 sums). + * 3. Scan buckets from largest to smallest to locate which bucket + * contains the K-th largest element. + * 4. Narrow the desired pattern by 4 bits. + * After 8 passes: the exact uint32 bit pattern of the K-th value is known. + * + * Phase 2 — Gather (single pass over E): + * Collect elements strictly greater than the K-th value (same uint order), + * then fill remaining slots with elements equal to the K-th value (ties + * broken by ascending index for determinism matching torch.topk). + * Write indices and scores to the output arrays. + * + * Tie-breaking: (value DESC, index ASC) — matches torch.topk behavior. + * + * Constraints: + * - 0 < topk <= data_size + * - No upper limit on topk or data_size (unlike v1's 128 cap) + * - scores must be in shared memory accessible by the warp + * + * Complexity: 9 × O(E/32) = O(E) per warp, independent of K. + ******************************************************************************/ + +__device__ inline void radix_topk_and_mask(CompType *scores, int data_size, int topk, + int *topk_indices, CompType *topk_scores, int lane_id) { + // assert(topk > 0 && "naive_topk_and_mask_v2: topk must be positive"); + // assert(topk <= data_size && "naive_topk_and_mask_v2: topk exceeds data_size"); + + constexpr int RADIX_BITS = 4; + constexpr int RADIX_SIZE = 1 << RADIX_BITS; // 16 buckets + constexpr int RADIX_MASK = RADIX_SIZE - 1; // 0xF + constexpr int NUM_PASSES = 32 / RADIX_BITS; // 8 passes for float32 + + // ========================================================================= + // Phase 1: Radix selection — find the bit pattern of the K-th largest value + // ========================================================================= + unsigned int desired = 0; // accumulated bit pattern of the K-th value + unsigned int desired_mask = 0; // bits determined so far + int k_remaining = topk; // how many more elements we need to skip + + for (int pass = NUM_PASSES - 1; pass >= 0; pass--) { + int digit_pos = pass * RADIX_BITS; + + // Each thread counts elements per bucket that match the desired pattern + unsigned int counts[RADIX_SIZE]; +#pragma unroll + for (int b = 0; b < RADIX_SIZE; b++) { + counts[b] = 0; + } + + for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { + unsigned int u = float_to_ordered_uint(scores[i]); + // Check if this element matches the desired pattern on already-decided bits + if ((u & desired_mask) == desired) { + int bucket = (u >> digit_pos) & RADIX_MASK; + counts[bucket]++; + } + } + + // Warp-reduce each bucket count across all 32 lanes + unsigned int total_counts[RADIX_SIZE]; +#pragma unroll + for (int b = 0; b < RADIX_SIZE; b++) { + unsigned int c = warp_allreduce_sum(counts[b]); + total_counts[b] = c; // same value on all lanes after full reduction + } + + // Scan buckets from LARGEST digit value (15) to smallest (0). + // We're looking for the top-K largest, so we want the highest-valued + // bucket first. Accumulate counts until we find the bucket containing + // the k_remaining-th element. + int target_bucket = 0; + for (int b = RADIX_SIZE - 1; b >= 0; b--) { + unsigned int bc = total_counts[b]; + if (bc < static_cast(k_remaining)) { + // All elements in this bucket are in the top set; skip them + k_remaining -= bc; + } else { + // The K-th element is in this bucket + target_bucket = b; + break; + } + } + + // Update the desired pattern and mask + desired |= (static_cast(target_bucket) << digit_pos); + desired_mask |= (static_cast(RADIX_MASK) << digit_pos); + } + + // After all passes, `desired` holds the exact ordered-uint bit pattern of + // the K-th largest value, and `k_remaining` is the number of elements with + // that exact value that should be included in the top-K set. + // (k_remaining >= 1 unless all elements equal the K-th value boundary) + + // ========================================================================= + // Phase 2: Gather — collect top-K elements into output arrays + // ========================================================================= + // Two sub-passes over the data: + // Pass A: Collect all elements strictly greater than the K-th value. + // Pass B: Collect elements equal to the K-th value (up to k_remaining), + // in ascending index order for deterministic tie-breaking. + // + // Since the warp processes indices in strided order, we need a warp-level + // prefix sum to assign output positions without conflicts. + + // --- Pass A: elements strictly greater than K-th value --- + // Use a warp-wide running counter for output position. + int write_pos = 0; // shared across warp via __shfl_sync + + for (int base = 0; base < data_size; base += kThreadsPerWarp) { + int i = base + lane_id; + bool valid = (i < data_size); + + unsigned int u = valid ? float_to_ordered_uint(scores[i]) : 0; + bool is_greater = valid && (u > desired); + + // Warp ballot to count how many lanes have a qualifying element + unsigned int ballot = __ballot_sync(0xffffffff, is_greater); + int lane_prefix = __popc(ballot & ((1u << lane_id) - 1)); // exclusive prefix + int total_qualifying = __popc(ballot); + + if (is_greater) { + int out_idx = write_pos + lane_prefix; + if (out_idx < topk) { + topk_indices[out_idx] = i; + topk_scores[out_idx] = scores[i]; + } + } + write_pos += total_qualifying; + } + + // --- Pass B: elements equal to K-th value (up to k_remaining) --- + int tie_remaining = k_remaining; // broadcast same value to all lanes + + for (int base = 0; base < data_size && tie_remaining > 0; base += kThreadsPerWarp) { + int i = base + lane_id; + bool valid = (i < data_size); + + unsigned int u = valid ? float_to_ordered_uint(scores[i]) : 0; + bool is_equal = valid && (u == desired); + + unsigned int ballot = __ballot_sync(0xffffffff, is_equal); + int lane_prefix = __popc(ballot & ((1u << lane_id) - 1)); + int total_equal = __popc(ballot); + + if (is_equal && lane_prefix < tie_remaining) { + int out_idx = write_pos + lane_prefix; + if (out_idx < topk) { + topk_indices[out_idx] = i; + topk_scores[out_idx] = scores[i]; + } + } + + int consumed = (total_equal < tie_remaining) ? total_equal : tie_remaining; + write_pos += consumed; + tie_remaining -= consumed; + } + + __syncwarp(); +} + __device__ inline void naive_topk_and_mask(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 @@ -249,6 +475,16 @@ __device__ inline void naive_topk_and_mask(CompType *scores, int data_size, int } } +template +__device__ __forceinline__ void topk_and_mask(CompType *scores, int data_size, int topk, + int *topk_indices, CompType *topk_scores, + int lane_id) { + if constexpr (TopkFunc == TopkFuncType::Radix) + return radix_topk_and_mask(scores, data_size, topk, topk_indices, topk_scores, lane_id); + else + return naive_topk_and_mask(scores, data_size, topk, topk_indices, topk_scores, lane_id); +} + // Current TE only support float32/bf16/fp16, float64 probs should be considered in the future #define TE_ROUTER_PROBS_TYPE_SWITCH_ALL(dtype, type, ...) \ switch (dtype) { \ diff --git a/transformer_engine/common/utils.cuh b/transformer_engine/common/utils.cuh index 8c50e83926..b322ce8fba 100644 --- a/transformer_engine/common/utils.cuh +++ b/transformer_engine/common/utils.cuh @@ -920,6 +920,32 @@ __device__ __forceinline__ void reciprocal(float *value_inv, const float *value_inv = __frcp_rn(value); } +// Convert float to an unsigned integer that preserves descending sort order. +// After conversion, a numerically larger float maps to a larger uint32. +__device__ __forceinline__ unsigned int float_to_ordered_uint(float f) { + unsigned int u = __float_as_uint(f); + // If sign bit is set (negative), flip all bits. + // If sign bit is clear (positive or +0), flip only the sign bit. + unsigned int mask = (u & 0x80000000u) ? 0xFFFFFFFFu : 0x80000000u; + return u ^ mask; +} + +// Convert back from ordered uint to float. +__device__ __forceinline__ float ordered_uint_to_float(unsigned int u) { + // Reverse the transformation: if MSB is set (was positive), flip sign bit. + // If MSB is clear (was negative), flip all bits. + unsigned int mask = (u & 0x80000000u) ? 0x80000000u : 0xFFFFFFFFu; + return __uint_as_float(u ^ mask); +} + +template +__device__ __forceinline__ T warp_allreduce_sum(T x) { + // Butterfly reduction +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) x += warp_shuffle_xor(x, offset); + return x; +} + //////////////////////////////////////////////////////////////////////////////////////////////////// using fp8e4m3 = __nv_fp8_e4m3; From fca261ecd09c318d22e7eeebda79632eed8cb9e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ga=C3=A9tan=20Lepage?= Date: Fri, 17 Apr 2026 00:40:24 +0200 Subject: [PATCH 358/521] fix CUDA architectures cmake logic (#2832) Signed-off-by: Gaetan Lepage --- transformer_engine/common/CMakeLists.txt | 27 +++++++++++++++++++++--- 1 file changed, 24 insertions(+), 3 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 3f684adbb4..a21c1ee7e6 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -36,7 +36,11 @@ if(NOT DEFINED CMAKE_CUDA_ARCHITECTURES) endif() endif() -# Process CMAKE_CUDA_ARCHITECTURES to separate generic and specific architectures +# Process CMAKE_CUDA_ARCHITECTURES to separate standard, generic, and specific architectures. +# - NVTE_STANDARD_ARCHS: pre-Blackwell archs (e.g. 75, 80, 89, 90). Applied to all CUDA sources. +# - NVTE_GENERIC_ARCHS: Blackwell family heads (e.g. 100, 120). Applied to non-arch-specific sources only. +# - NVTE_SPECIFIC_ARCHS: Blackwell specific targets (e.g. 100a, 120f). Applied to arch-specific sources only. +set(NVTE_STANDARD_ARCHS) set(NVTE_GENERIC_ARCHS) set(NVTE_SPECIFIC_ARCHS) @@ -79,6 +83,10 @@ if(NOT arch_120_index EQUAL -1) endif() endif() +# Move remaining standard (pre-Blackwell) architectures into NVTE_STANDARD_ARCHS. +# These are applied to all CUDA sources (both generic and arch-specific). +set(NVTE_STANDARD_ARCHS ${CMAKE_CUDA_ARCHITECTURES}) + # cuDNN frontend API set(CUDNN_FRONTEND_INCLUDE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/cudnn-frontend/include") @@ -228,9 +236,13 @@ list(APPEND transformer_engine_SOURCES ${transformer_engine_cuda_arch_specific_s ${transformer_engine_cuda_sources} ${transformer_engine_cpp_sources}) -# Set compile options for CUDA sources with generic architectures +# Set compile options for CUDA sources with generic architectures. +# These get standard archs (pre-Blackwell) + generic Blackwell family heads. foreach(cuda_source IN LISTS transformer_engine_cuda_sources) set(arch_compile_options) + foreach(arch IN LISTS NVTE_STANDARD_ARCHS) + list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}") + endforeach() foreach(arch IN LISTS NVTE_GENERIC_ARCHS) list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}") endforeach() @@ -245,9 +257,14 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_sources) endif() endforeach() -# Set compile options for CUDA sources with specific architectures +# Set compile options for CUDA sources with arch-specific features. +# These get standard archs (pre-Blackwell) + Blackwell specific targets (a/f suffix). +# They must NOT get generic Blackwell archs, as they use family/arch-specific PTX features. foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) set(arch_compile_options) + foreach(arch IN LISTS NVTE_STANDARD_ARCHS) + list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}") + endforeach() foreach(arch IN LISTS NVTE_SPECIFIC_ARCHS) list(APPEND arch_compile_options "--generate-code=arch=compute_${arch},code=sm_${arch}") endforeach() @@ -268,6 +285,10 @@ list(APPEND transformer_engine_SOURCES 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 +# using NVTE_STANDARD_ARCHS, NVTE_GENERIC_ARCHS, and NVTE_SPECIFIC_ARCHS above. +set_target_properties(transformer_engine PROPERTIES CUDA_ARCHITECTURES OFF) target_include_directories(transformer_engine PUBLIC "${CMAKE_CURRENT_SOURCE_DIR}/include") From be593b1cddbe0c8df895b4d1cb56489d703ec1df Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Fri, 17 Apr 2026 08:27:37 -0700 Subject: [PATCH 359/521] [Common, pyTorch] Grouped MXFP8 dequantize support (#2722) * Grouped dequantize for MXFP8 Signed-off-by: Przemek Tredak * Pytorch extension Signed-off-by: Przemek Tredak * Fix CUDA graphs compatibility Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Handling non-full tiles 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 * Fixes Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixes from review Signed-off-by: Przemek Tredak * Refactor grouped MXFP8 dequantize kernel - Use common namespace helpers instead of group_quantize_kernel - Extract shared constants into DequantizeConfig struct - Replace SCALE_DIM template params with single ROWWISE bool - Use initialize_barriers/destroy_barriers helpers - Fix offsets array size (num_tensors + 1) - Skip TMA descriptor update for zero-sized groups - Fix off-by-one in max tensor descriptor check Signed-off-by: Przemek Tredak * Tighten tensor_offsets validation to require num_tensors+1 All producers (splits_to_offsets, quantizer.cpp) and consumers (is_job_valid, get_current_tensor_id, hadamard transform) already use CSR-style num_tensors+1 offsets. Make the validation match. Also fix stale docstring in grouped_tensor_storage.py. Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix group_dequantize: attribute names, dtype, and shape handling In group_dequantize(), GroupedTensor inherits from torch.Tensor, so accessing .data returns the 2D wrapper tensor instead of the 1D quantized data buffer. Fix three issues: - Read "rowwise_data" attribute instead of "data" to get the flat 1D quantized buffer rather than torch.Tensor.data (2D wrapper). - Use quantizer->dtype (e.g. kFloat8E4M3) instead of deriving dtype from the raw tensor's scalar_type() which is just uint8. - Pass numel() as a 1-element shape vector to ensure the grouped tensor data is registered as 1D. Promote DType dtype from quantizer subclasses to the base Quantizer class (defaulting to kNumTypes) so group_dequantize can access it without downcasting. Update tests to compare per-tensor via split_into_quantized_tensors() instead of accessing .data on GroupedTensor. Signed-off-by: Przemyslaw Tredak 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: Przemyslaw Tredak Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/cpp/operator/CMakeLists.txt | 1 + .../operator/test_dequantize_mxfp8_grouped.cu | 487 +++++++++++++++++ tests/pytorch/test_grouped_tensor.py | 81 +++ transformer_engine/common/cast/cast.cu | 8 + .../common/cast/dispatch/dequantize.cuh | 21 + .../cast/mxfp8/group_dequantize_mxfp8.cuh | 495 ++++++++++++++++++ .../common/include/transformer_engine/cast.h | 13 +- .../common/transformer_engine.cpp | 13 +- transformer_engine/pytorch/csrc/common.h | 9 +- transformer_engine/pytorch/csrc/extensions.h | 2 + .../pytorch/csrc/extensions/cast.cpp | 79 +++ .../pytorch/csrc/extensions/pybind.cpp | 2 + .../tensor/storage/grouped_tensor_storage.py | 3 +- 13 files changed, 1202 insertions(+), 12 deletions(-) create mode 100644 tests/cpp/operator/test_dequantize_mxfp8_grouped.cu create mode 100644 transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index 5e73675f4f..f83c4ae066 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -15,6 +15,7 @@ add_executable(test_operator test_cast_nvfp4_transpose.cu test_cast_float8blockwise.cu test_dequantize_mxfp8.cu + test_dequantize_mxfp8_grouped.cu test_transpose.cu test_cast_transpose.cu test_cast_transpose_current_scaling.cu diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu new file mode 100644 index 0000000000..4a18bb5891 --- /dev/null +++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu @@ -0,0 +1,487 @@ +/************************************************************************* + * 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 "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +namespace { + +enum ShapeRepresentation { + SAME_BOTH_DIMS = 0, + VARYING_FIRST_DIM = 1, + VARYING_LAST_DIM = 2, + VARYING_BOTH_DIMS = 3 +}; + +enum ScalingDirection { ROWWISE = 0, COLWISE = 1 }; + +/** + * Compare grouped dequantize output against single-tensor nvte_dequantize + * called in a loop for each tensor. Results must be bitwise identical. + */ +template +void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors, + const std::vector &logical_shape_vec, + const std::vector &first_dims_h, const std::vector &last_dims_h, + const std::vector &offsets_h, const bool rowwise) { + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + const size_t rows = logical_shape_vec[0]; + const size_t cols = logical_shape_vec[1]; + + // Compute total elements and per-tensor scale sizes + size_t elts_num = 0; + size_t total_scales = 0; + + std::vector per_tensor_scales_first_dim(num_tensors); + std::vector per_tensor_scales_last_dim(num_tensors); + std::vector per_tensor_scales_offset(num_tensors + 1, 0); + + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; + elts_num += M * K; + + size_t unpadded_scales_Y, unpadded_scales_X; + if (rowwise) { + unpadded_scales_Y = M; + unpadded_scales_X = divide_round_up(K, 32); + per_tensor_scales_first_dim[t] = + round_up_to_nearest_multiple(unpadded_scales_Y, scale_tensor_alignment_Y_rowwise); + per_tensor_scales_last_dim[t] = + round_up_to_nearest_multiple(unpadded_scales_X, scale_tensor_alignment_X_rowwise); + } else { + unpadded_scales_Y = divide_round_up(M, 32); + unpadded_scales_X = K; + per_tensor_scales_first_dim[t] = + round_up_to_nearest_multiple(unpadded_scales_Y, scale_tensor_alignment_Y_colwise); + per_tensor_scales_last_dim[t] = + round_up_to_nearest_multiple(unpadded_scales_X, scale_tensor_alignment_X_colwise); + } + + const size_t tensor_scales = per_tensor_scales_first_dim[t] * per_tensor_scales_last_dim[t]; + total_scales += tensor_scales; + per_tensor_scales_offset[t + 1] = total_scales; + } + + // Allocate host data + std::vector in_data_h(elts_num); + std::vector in_scales_h(total_scales); + + // Generate random FP8 data and scales + static std::mt19937 gen(42); + const double minAbs = Numeric_Traits::minNorm; + const double maxAbs = Numeric_Traits::maxNorm; + std::uniform_real_distribution<> dis(minAbs, maxAbs); + std::uniform_real_distribution<> dis_sign(-1.0, 1.0); + std::uniform_int_distribution int_dis(0, 255); + + for (size_t i = 0; i < elts_num; ++i) { + const bool is_negative = (dis_sign(gen) < 0.0); + double val = dis(gen); + if (is_negative) val = -val; + in_data_h[i] = static_cast(val); + } + for (size_t i = 0; i < total_scales; ++i) { + in_scales_h[i] = int_dis(gen); + } + + // Allocate device memory + const size_t in_data_size = elts_num * sizeof(InputType); + const size_t out_data_size = elts_num * sizeof(OutputType); + const size_t scales_size = total_scales * sizeof(fp8e8m0); + const size_t first_dims_size = num_tensors * sizeof(size_t); + const size_t last_dims_size = num_tensors * sizeof(size_t); + const size_t offsets_size = (num_tensors + 1) * sizeof(size_t); + + InputType *in_data_d; + OutputType *out_grouped_d; + fp8e8m0 *in_scales_d; + size_t *first_dims_d; + size_t *last_dims_d; + size_t *offsets_d; + + cudaMalloc((void **)&in_data_d, in_data_size); + cudaMalloc((void **)&out_grouped_d, out_data_size); + cudaMalloc((void **)&in_scales_d, scales_size); + cudaMalloc((void **)&first_dims_d, first_dims_size); + cudaMalloc((void **)&last_dims_d, last_dims_size); + cudaMalloc((void **)&offsets_d, offsets_size); + + cudaMemcpy(in_data_d, in_data_h.data(), in_data_size, cudaMemcpyHostToDevice); + cudaMemcpy(in_scales_d, in_scales_h.data(), scales_size, cudaMemcpyHostToDevice); + cudaMemcpy(first_dims_d, first_dims_h.data(), first_dims_size, cudaMemcpyHostToDevice); + cudaMemcpy(last_dims_d, last_dims_h.data(), last_dims_size, cudaMemcpyHostToDevice); + cudaMemcpy(offsets_d, offsets_h.data(), offsets_size, cudaMemcpyHostToDevice); + + // Set up grouped input tensor + NVTEShape logical_shape = nvte_make_shape(logical_shape_vec.data(), logical_shape_vec.size()); + + NVTEShape first_dims_shape; + NVTEShape last_dims_shape; + NVTEShape offsets_shape; + first_dims_shape.ndim = 1; + last_dims_shape.ndim = 1; + offsets_shape.ndim = 1; + first_dims_shape.data[0] = num_tensors; + last_dims_shape.data[0] = num_tensors; + offsets_shape.data[0] = num_tensors + 1; + + // Data tensors must be 1D (flattened) + std::vector data_1d_shape = {elts_num}; + NVTEShape data_shape = nvte_make_shape(data_1d_shape.data(), data_1d_shape.size()); + + std::vector scales_1d_shape = {total_scales}; + NVTEShape scales_shape = nvte_make_shape(scales_1d_shape.data(), scales_1d_shape.size()); + + NVTEGroupedTensor in_group_tensor = + nvte_create_grouped_tensor(NVTE_MXFP8_1D_SCALING, num_tensors, logical_shape); + + // Set input data (rowwise or columnwise) - data shape must be 1D + NVTEBasicTensor in_data_tensor = {in_data_d, static_cast(itype), data_shape}; + if (rowwise) { + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &in_data_tensor, + sizeof(in_data_tensor)); + } else { + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedColumnwiseData, + &in_data_tensor, sizeof(in_data_tensor)); + } + + // Set scales + NVTEBasicTensor in_scales_tensor = {in_scales_d, NVTEDType::kNVTEFloat8E8M0, scales_shape}; + if (rowwise) { + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedRowwiseScaleInv, + &in_scales_tensor, sizeof(in_scales_tensor)); + } else { + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedColumnwiseScaleInv, + &in_scales_tensor, sizeof(in_scales_tensor)); + } + + // Set shape arrays + if ((shape_rep == VARYING_FIRST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { + NVTEBasicTensor first_dims_tensor = {first_dims_d, kNVTEInt64, first_dims_shape}; + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedFirstDims, + &first_dims_tensor, sizeof(first_dims_tensor)); + } + if ((shape_rep == VARYING_LAST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { + NVTEBasicTensor last_dims_tensor = {last_dims_d, kNVTEInt64, last_dims_shape}; + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor, + sizeof(last_dims_tensor)); + } + if (shape_rep != SAME_BOTH_DIMS) { + NVTEBasicTensor offsets_tensor = {offsets_d, kNVTEInt64, offsets_shape}; + nvte_set_grouped_tensor_param(in_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, + &offsets_tensor, sizeof(offsets_tensor)); + } + + // Set up grouped output tensor + NVTEGroupedTensor out_group_tensor = + nvte_create_grouped_tensor(NVTE_DELAYED_TENSOR_SCALING, num_tensors, logical_shape); + + NVTEBasicTensor out_data_tensor = {out_grouped_d, static_cast(otype), data_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedRowwiseData, &out_data_tensor, + sizeof(out_data_tensor)); + + // Set shape arrays on output too + if ((shape_rep == VARYING_FIRST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { + NVTEBasicTensor first_dims_tensor = {first_dims_d, kNVTEInt64, first_dims_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedFirstDims, + &first_dims_tensor, sizeof(first_dims_tensor)); + } + if ((shape_rep == VARYING_LAST_DIM) || (shape_rep == VARYING_BOTH_DIMS)) { + NVTEBasicTensor last_dims_tensor = {last_dims_d, kNVTEInt64, last_dims_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedLastDims, &last_dims_tensor, + sizeof(last_dims_tensor)); + } + if (shape_rep != SAME_BOTH_DIMS) { + NVTEBasicTensor offsets_tensor = {offsets_d, kNVTEInt64, offsets_shape}; + nvte_set_grouped_tensor_param(out_group_tensor, + NVTEGroupedTensorParam::kNVTEGroupedTensorOffsets, + &offsets_tensor, sizeof(offsets_tensor)); + } + + // Run grouped dequantize + nvte_group_dequantize(in_group_tensor, out_group_tensor, 0); + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + // Copy grouped output to host + std::vector out_grouped_h(elts_num); + cudaMemcpy(out_grouped_h.data(), out_grouped_d, out_data_size, cudaMemcpyDeviceToHost); + + // Now compute reference: run single-tensor nvte_dequantize for each tensor + std::vector out_ref_h(elts_num); + + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; + const size_t data_offset = offsets_h[t]; + const size_t scales_offset = per_tensor_scales_offset[t]; + const size_t tensor_scales_count = + per_tensor_scales_first_dim[t] * per_tensor_scales_last_dim[t]; + + const size_t single_data_size = M * K * sizeof(InputType); + const size_t single_out_size = M * K * sizeof(OutputType); + const size_t single_scales_size = tensor_scales_count * sizeof(fp8e8m0); + + // Allocate per-tensor device memory + InputType *single_in_d; + OutputType *single_out_d; + fp8e8m0 *single_scales_d; + + cudaMalloc((void **)&single_in_d, single_data_size); + cudaMalloc((void **)&single_out_d, single_out_size); + cudaMalloc((void **)&single_scales_d, single_scales_size); + + cudaMemcpy(single_in_d, in_data_h.data() + data_offset, single_data_size, + cudaMemcpyHostToDevice); + cudaMemcpy(single_scales_d, in_scales_h.data() + scales_offset, single_scales_size, + cudaMemcpyHostToDevice); + cudaMemset(single_out_d, 0, single_out_size); + + // Build single-tensor NVTETensor using TensorWrapper directly + std::vector single_shape = {M, K}; + std::vector scale_shape_vec = {per_tensor_scales_first_dim[t], + per_tensor_scales_last_dim[t]}; + + TensorWrapper input_w(NVTE_MXFP8_1D_SCALING); + if (rowwise) { + input_w.set_rowwise_data(single_in_d, itype, single_shape); + input_w.set_rowwise_scale_inv(single_scales_d, DType::kFloat8E8M0, scale_shape_vec); + } else { + input_w.set_columnwise_data(single_in_d, itype, single_shape); + input_w.set_columnwise_scale_inv(single_scales_d, DType::kFloat8E8M0, scale_shape_vec); + } + + TensorWrapper output_w; + output_w.set_rowwise_data(single_out_d, otype, single_shape); + + nvte_dequantize(input_w.data(), output_w.data(), 0); + cudaDeviceSynchronize(); + err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << "Single-tensor dequantize failed for tensor " << t << ": " + << cudaGetErrorString(err); + + // Copy reference output to host + cudaMemcpy(out_ref_h.data() + data_offset, single_out_d, single_out_size, + cudaMemcpyDeviceToHost); + + cudaFree(single_in_d); + cudaFree(single_out_d); + cudaFree(single_scales_d); + } + + // Bitwise comparison + for (size_t t = 0; t < num_tensors; ++t) { + const size_t M = first_dims_h[t]; + const size_t K = last_dims_h[t]; + const size_t data_offset = offsets_h[t]; + const size_t tensor_elts = M * K; + + int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset, + tensor_elts * sizeof(OutputType)); + if (result != 0) { + // Find first mismatch for error reporting + for (size_t i = 0; i < tensor_elts; ++i) { + if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) { + GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i + << " (global offset " << (data_offset + i) << "): grouped=" + << static_cast(out_grouped_h[data_offset + i]) + << " vs reference=" << static_cast(out_ref_h[data_offset + i]); + } + } + } + } + + // Cleanup + cudaFree(in_data_d); + cudaFree(out_grouped_d); + cudaFree(in_scales_d); + cudaFree(first_dims_d); + cudaFree(last_dims_d); + cudaFree(offsets_d); +} + +// {shape_representation, num_tensors, [logical_shape_M, logical_shape_K], [M_i], [K_i]} +std::vector> input_configs = { + {SAME_BOTH_DIMS, 1, 128, 128}, + {SAME_BOTH_DIMS, 2, 256, 128}, + {VARYING_FIRST_DIM, 2, 512, 128, 128, 384}, + {VARYING_FIRST_DIM, 2, 384, 128, 128, 256}, + {VARYING_FIRST_DIM, 5, 4096, 512, 128, 256, 384, 1024, 2304}, + {VARYING_LAST_DIM, 3, 256, 896, 128, 256, 512}, + {VARYING_BOTH_DIMS, 2, 1, (128 * 128) + (256 * 256), 128, 256, 128, 256}, + {VARYING_BOTH_DIMS, 2, 1, (256 * 128) + (512 * 640), 256, 512, 128, 640}, + // Non-128-aligned constant dimensions + {SAME_BOTH_DIMS, 1, 160, 192}, + {SAME_BOTH_DIMS, 2, 256, 96}, + {VARYING_FIRST_DIM, 2, 384, 160, 128, 256}, + {VARYING_FIRST_DIM, 3, 768, 96, 256, 256, 256}, + {VARYING_LAST_DIM, 2, 160, 384, 128, 256}, + {VARYING_LAST_DIM, 3, 96, 512, 128, 128, 256}, +}; + +std::vector scaling_directions = { + ScalingDirection::ROWWISE, + ScalingDirection::COLWISE, +}; + +} // namespace + +class GroupedDequantizeMXFP8TestSuite + : public ::testing::TestWithParam, // Config + transformer_engine::DType, // InputType + transformer_engine::DType // OutputType + >> {}; + +TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) { + // Skip tests for pre-Blackwell architectures + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const ScalingDirection scaling_direction = std::get<0>(GetParam()); + const std::vector config = std::get<1>(GetParam()); + const DType input_type = std::get<2>(GetParam()); + const DType output_type = std::get<3>(GetParam()); + + const ShapeRepresentation shape_rep = static_cast(config[0]); + const size_t num_tensors = config[1]; + const std::vector logical_shape = {config[2], config[3]}; + + const bool rowwise = (scaling_direction == ScalingDirection::ROWWISE); + + std::vector first_dims(num_tensors); + std::vector last_dims(num_tensors); + std::vector offsets(num_tensors + 1, 0); + + for (size_t t = 0; t < num_tensors; ++t) { + switch (shape_rep) { + case SAME_BOTH_DIMS: { + first_dims[t] = logical_shape[0] / num_tensors; + last_dims[t] = logical_shape[1]; + break; + } + case VARYING_FIRST_DIM: { + first_dims[t] = config[t + 4]; + last_dims[t] = logical_shape[1]; + break; + } + case VARYING_LAST_DIM: { + first_dims[t] = logical_shape[0]; + last_dims[t] = config[t + 4]; + break; + } + case VARYING_BOTH_DIMS: { + first_dims[t] = config[t + 4]; + last_dims[t] = config[t + (4 + num_tensors)]; + break; + } + } + offsets[t + 1] = offsets[t] + first_dims[t] * last_dims[t]; + + // Skip tests if varying dimensions are not 128-aligned + const bool first_dim_varies = + (shape_rep == VARYING_FIRST_DIM || shape_rep == VARYING_BOTH_DIMS); + const bool last_dim_varies = + (shape_rep == VARYING_LAST_DIM || shape_rep == VARYING_BOTH_DIMS); + if (first_dim_varies && (first_dims[t] % 128 != 0)) { + GTEST_SKIP(); + } + if (last_dim_varies && (last_dims[t] % 128 != 0)) { + GTEST_SKIP(); + } + // TMA requires last_dim * sizeof(FP8) to be 16-byte aligned + if (last_dims[t] % 16 != 0) { + GTEST_SKIP(); + } + // For colwise: first dim must be divisible by 32 + if (!rowwise && (first_dims[t] % 32 != 0)) { + GTEST_SKIP(); + } + // For rowwise: last dim must be divisible by 32 + if (rowwise && (last_dims[t] % 32 != 0)) { + GTEST_SKIP(); + } + } + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY( + input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + output_type, OutputType, + performTest(shape_rep, num_tensors, logical_shape, first_dims, + last_dims, offsets, rowwise););); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, GroupedDequantizeMXFP8TestSuite, + ::testing::Combine(::testing::ValuesIn(scaling_directions), ::testing::ValuesIn(input_configs), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + [](const testing::TestParamInfo &info) { + std::string name; + switch (std::get<0>(info.param)) { + case ScalingDirection::ROWWISE: + name += "ROWWISE_"; + break; + case ScalingDirection::COLWISE: + name += "COLWISE_"; + break; + } + + const std::vector input = std::get<1>(info.param); + switch (static_cast(input[0])) { + case ShapeRepresentation::SAME_BOTH_DIMS: + name += "SAME_BOTH_DIMS"; + break; + case ShapeRepresentation::VARYING_FIRST_DIM: + name += "VARYING_FIRST_DIM"; + break; + case ShapeRepresentation::VARYING_LAST_DIM: + name += "VARYING_LAST_DIM"; + break; + case ShapeRepresentation::VARYING_BOTH_DIMS: + name += "VARYING_BOTH_DIMS"; + break; + } + + name += "_N_" + std::to_string(input[1]); + name += "_SHAPE_" + std::to_string(input[2]) + "X" + std::to_string(input[3]); + name += "_" + test::typeName(std::get<2>(info.param)); + name += "_" + test::typeName(std::get<3>(info.param)); + return name; + }); diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py index 04a0376019..c54c9758ff 100644 --- a/tests/pytorch/test_grouped_tensor.py +++ b/tests/pytorch/test_grouped_tensor.py @@ -500,6 +500,87 @@ def test_group_quantize_cudagraph_capturable(self, output_dbias: bool) -> None: if output_dbias: assert torch.allclose(static_dbias, expected_dbias) + @pytest.mark.parametrize( + "shape", + [[(512, 1024), (512, 1024)], [(256, 512), (512, 512), (768, 512)]], + ) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_dequantize(self, shape: List[Tuple[int, int]]) -> None: + """Test grouped dequantization for MXFP8 back to BF16.""" + num_tensors = len(shape) + + # Create BF16 input tensors and quantize them with MXFP8. + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=False) + first_dims = torch.tensor([s[0] for s in shape], dtype=torch.int64, device="cuda") + + # Quantize. + quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims) + + # Dequantize. + dequantized = tex.group_dequantize(quantized, tex.DType.kBFloat16) + + # Verify output metadata. + assert dequantized.num_tensors == num_tensors + assert dequantized.logical_shape == quantized.logical_shape + assert torch.equal(dequantized.first_dims, quantized.first_dims) + assert torch.equal(dequantized.tensor_offsets, quantized.tensor_offsets) + + # Verify dequantized values are close to original (per-tensor). + dequantized_tensors = dequantized.split_into_quantized_tensors() + assert len(dequantized_tensors) == num_tensors + for orig, deq in zip(input_tensors, dequantized_tensors): + torch.testing.assert_close(deq, orig, atol=0.125, rtol=0.1) + + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_group_dequantize_cudagraph_capturable(self) -> None: + """Ensure group_dequantize is CUDA graph capturable.""" + num_tensors = 2 + shape = [(512, 1024) for _ in range(num_tensors)] + input_tensors = [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape] + grouped_input = torch.cat(input_tensors, dim=0) + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + quantizer.set_usage(rowwise=True, columnwise=False) + first_dims = torch.tensor( + [shape[0][0] for _ in range(num_tensors)], + dtype=torch.int64, + device="cuda", + ) + + # Quantize to get MXFP8 grouped tensor. + quantized = tex.group_quantize(grouped_input, quantizer, num_tensors, first_dims) + + # Warmup dequantize. + torch.cuda.synchronize() + _ = tex.group_dequantize(quantized, tex.DType.kBFloat16) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + static_output = tex.group_dequantize(quantized, tex.DType.kBFloat16) + + # Replay with different input data. + fresh_input = torch.cat( + [torch.randn(s, dtype=torch.bfloat16, device="cuda") for s in shape], + dim=0, + ) + fresh_quantized = tex.group_quantize(fresh_input, quantizer, num_tensors, first_dims) + quantized.rowwise_data.copy_(fresh_quantized.rowwise_data) + quantized.scale_inv.copy_(fresh_quantized.scale_inv) + + graph.replay() + torch.cuda.synchronize() + + expected = tex.group_dequantize(quantized, tex.DType.kBFloat16) + expected_tensors = expected.split_into_quantized_tensors() + static_tensors = static_output.split_into_quantized_tensors() + for exp, got in zip(expected_tensors, static_tensors): + assert torch.equal(got, exp) + def test_clear(self) -> None: """Test clear method""" num_tensors = 3 diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index dc02390818..61cfacd334 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -89,6 +89,14 @@ 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) { diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index 81304981d3..12787d609f 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -16,6 +16,7 @@ #include "../../common.h" #include "../fp8/dequantize_fp8.cuh" #include "../mxfp8/dequantize_mxfp8.cuh" +#include "../mxfp8/group_dequantize_mxfp8.cuh" #include "../nvfp4/dequantize_nvfp4.cuh" namespace transformer_engine { @@ -50,6 +51,26 @@ inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t } } +inline void group_dequantize_helper(const GroupedTensor &input, GroupedTensor *output, + cudaStream_t stream) { + CheckInputGroupedTensor(input, "group_dequantize_input"); + CheckOutputGroupedTensor(*output, "group_dequantize_output"); + + switch (input.scaling_mode) { + case NVTE_MXFP8_1D_SCALING: { + if (is_supported_by_CC_100()) { + mxfp8::group_dequantize(&input, output, stream); + } else { + NVTE_ERROR("MXFP8 Grouped Dequantization is NOT supported by architectures < 10.0"); + } + break; + } + default: + NVTE_ERROR("Grouped dequantize not implemented for scaling mode: " + + to_string(input.scaling_mode) + "."); + } +} + } // namespace dispatch } // namespace transformer_engine diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh new file mode 100644 index 0000000000..dad8d18d6f --- /dev/null +++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh @@ -0,0 +1,495 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file group_dequantize_mxfp8.cuh + * \brief CUDA kernels to dequantize grouped tensors from MXFP8. + */ + +#ifndef TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_MXFP8_CUH_ +#define TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_MXFP8_CUH_ + +#include +#include +#include +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../util/ptx.cuh" +#include "../../utils.cuh" +#include "group_quantize_mxfp8.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace mxfp8 { +namespace group_dequantize_kernel { + +constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64; +__device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; +__device__ alignas(128) CUtensorMap g_tensor_maps_output[MAX_SUPPORTED_TENSOR_DESCRIPTORS]; + +// Reuse helper types and functions from common namespace +using common::fence_acquire_tensormap; +using common::get_tensor_cols_num; +using common::get_tensor_rows_num; +using common::modify_base_tensor_map; + +// Runtime dispatch wrapper for get_current_tensor_id (common only has template version) +template +__device__ __forceinline__ size_t get_current_tensor_id( + const ShapeRepresentation shape_rep, const size_t num_tensors, const size_t current_offset, + const size_t block_Y, const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr) { + switch (shape_rep) { + case ShapeRepresentation::SAME_BOTH_DIMS: + return common::get_current_tensor_id( + num_tensors, current_offset, block_Y, first_logical_dim, last_logical_dim, offsets_ptr); + case ShapeRepresentation::VARYING_FIRST_DIM: + return common::get_current_tensor_id( + num_tensors, current_offset, block_Y, first_logical_dim, last_logical_dim, offsets_ptr); + case ShapeRepresentation::VARYING_LAST_DIM: + return common::get_current_tensor_id( + num_tensors, current_offset, block_Y, first_logical_dim, last_logical_dim, offsets_ptr); + case ShapeRepresentation::VARYING_BOTH_DIMS: + return common::get_current_tensor_id( + num_tensors, current_offset, block_Y, first_logical_dim, last_logical_dim, offsets_ptr); + } + return 0; +} + +// Shared constexpr parameters used by both the kernel and the launch function. +// Defined in a struct so they are visible in both host and device code. +struct DequantizeConfig { + static constexpr size_t CHUNK_DIM_Y = 128; + static constexpr size_t CHUNK_DIM_X = 128; + static constexpr size_t THREADS_PER_CHUNK = 128; + static constexpr size_t BUFFERS_NUM = 2; + static constexpr size_t ELEMS_PER_THREAD = 16; + static constexpr size_t BUFFER_DIM_Y = 16; + static constexpr size_t BUFFER_DIM_X = CHUNK_DIM_X; + static constexpr size_t SHMEM_DIM_Y = BUFFER_DIM_Y; + static constexpr size_t SHMEM_DIM_X = BUFFER_DIM_X; + static constexpr size_t THREADS_PER_CHUNK_X_ROWWISE = CHUNK_DIM_X / ELEMS_PER_THREAD; + static constexpr size_t THREADS_PER_CHUNK_X_COLWISE = CHUNK_DIM_X; + static constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y; + static constexpr size_t ELTS_PER_CHUNK = CHUNK_DIM_Y * CHUNK_DIM_X; +}; + +template +__global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input, + const __grid_constant__ CUtensorMap base_tensor_map_output, + const IType *const __restrict__ input_data_ptr, + const OType *const __restrict__ output_data_ptr, + const ShapeRepresentation shape_rep, + const size_t num_tensors, const size_t first_logical_dim, + const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, + const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr) { + const bool leading_thread = (threadIdx.x == 0); + const size_t tensor_id = blockIdx.x; + + const size_t rows = + get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + + const size_t offset_elts = offsets_ptr[tensor_id]; + + // Zero-sized groups: skip TMA descriptor update. The main kernel already returns + // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension + // is invalid and causes CUDA_ERROR_ILLEGAL_ADDRESS. + if (rows == 0 || cols == 0) { + return; + } + + if (leading_thread && (tensor_id < num_tensors)) { + { + const uintptr_t global_data_ptr = reinterpret_cast(input_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id], + global_data_ptr, rows, cols, sizeof(IType)); + } + { + const uintptr_t global_data_ptr = reinterpret_cast(output_data_ptr + offset_elts); + modify_base_tensor_map(base_tensor_map_output, &g_tensor_maps_output[tensor_id], + global_data_ptr, rows, cols, sizeof(OType)); + } + } +} + +template +__global__ void __launch_bounds__(128) + group_dequantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input_static, + const __grid_constant__ CUtensorMap tensor_map_output_static, + const ShapeRepresentation shape_rep, const size_t num_tensors, + const size_t first_logical_dim, const size_t last_logical_dim, + const int64_t *const __restrict__ offsets_ptr, + const int64_t *const __restrict__ first_dims_ptr, + const int64_t *const __restrict__ last_dims_ptr, + const e8m0_t *const __restrict__ scales_ptr) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + constexpr size_t CHUNK_DIM_Y = DequantizeConfig::CHUNK_DIM_Y; + constexpr size_t CHUNK_DIM_X = DequantizeConfig::CHUNK_DIM_X; + constexpr size_t THREADS_PER_CHUNK = DequantizeConfig::THREADS_PER_CHUNK; + constexpr size_t BUFFERS_NUM = DequantizeConfig::BUFFERS_NUM; + constexpr size_t ELEMS_PER_THREAD = DequantizeConfig::ELEMS_PER_THREAD; + constexpr size_t BUFFER_DIM_Y = DequantizeConfig::BUFFER_DIM_Y; + constexpr size_t SHMEM_DIM_Y = DequantizeConfig::SHMEM_DIM_Y; + constexpr size_t SHMEM_DIM_X = DequantizeConfig::SHMEM_DIM_X; + constexpr size_t THREADS_PER_CHUNK_X_ROWWISE = DequantizeConfig::THREADS_PER_CHUNK_X_ROWWISE; + constexpr size_t THREADS_PER_CHUNK_X_COLWISE = DequantizeConfig::THREADS_PER_CHUNK_X_COLWISE; + constexpr size_t ITERATIONS = DequantizeConfig::ITERATIONS; + constexpr size_t ELTS_PER_CHUNK = DequantizeConfig::ELTS_PER_CHUNK; + + constexpr bool USE_ROWWISE_SCALING = ROWWISE; + constexpr size_t SCALE_DIM_Y = ROWWISE ? 1 : 32; + constexpr size_t SCALE_DIM_X = ROWWISE ? 32 : 1; + + constexpr size_t SCALES_ROWWISE_PER_CHUNK_Y = CHUNK_DIM_Y; + constexpr size_t SCALES_ROWWISE_PER_CHUNK_X = CHUNK_DIM_X / SCALE_DIM_X; + + constexpr size_t SCALES_COLWISE_PER_CHUNK_Y = CHUNK_DIM_Y / SCALE_DIM_Y; + constexpr size_t SCALES_COLWISE_PER_CHUNK_X = CHUNK_DIM_X; + + constexpr size_t THREADS_PER_SCALE_X_ROWWISE = DIVUP(SCALE_DIM_X, ELEMS_PER_THREAD); + + // Group-awareness: determine which tensor this block belongs to + const bool is_single_tensor = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + shape_rep == ShapeRepresentation::VARYING_FIRST_DIM); + + size_t tensor_id; + size_t block_id_Y, block_id_X; + + if (is_single_tensor) { + // SAME_BOTH_DIMS or VARYING_FIRST_DIM: simple 2D tiling over single logical tensor + const size_t chunks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); + block_id_Y = blockIdx.x / chunks_X; + block_id_X = blockIdx.x % chunks_X; + const size_t block_global_offset = blockIdx.x * ELTS_PER_CHUNK; + tensor_id = + get_current_tensor_id(shape_rep, num_tensors, block_global_offset, block_id_Y, + first_logical_dim, last_logical_dim, offsets_ptr); + } else if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM) { + // Virtual 2D grid: DIVUP(R,128) row-tiles x (total_cols/128) col-tiles + const size_t chunks_X_total = last_logical_dim / CHUNK_DIM_X; + const size_t col_chunk_global = blockIdx.x % chunks_X_total; + block_id_Y = blockIdx.x / chunks_X_total; + // Search using column-based element offset (works with existing binary search) + const size_t search_offset = col_chunk_global * CHUNK_DIM_X * first_logical_dim; + tensor_id = + get_current_tensor_id(shape_rep, num_tensors, search_offset, block_id_Y, + first_logical_dim, last_logical_dim, offsets_ptr); + const size_t tensor_col_start = static_cast(offsets_ptr[tensor_id]) / first_logical_dim; + block_id_X = col_chunk_global - tensor_col_start / CHUNK_DIM_X; + } else { + // VARYING_BOTH_DIMS: 1D grid, element-offset-based (both dims 128-aligned) + const size_t block_global_offset = blockIdx.x * ELTS_PER_CHUNK; + const size_t chunks_X_for_id = DIVUP(last_logical_dim, CHUNK_DIM_X); + tensor_id = get_current_tensor_id(shape_rep, num_tensors, block_global_offset, + blockIdx.x / chunks_X_for_id, first_logical_dim, + last_logical_dim, offsets_ptr); + const size_t vb_tensor_base = static_cast(offsets_ptr[tensor_id]); + const size_t vb_cols = + get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + const size_t chunks_X = DIVUP(vb_cols, CHUNK_DIM_X); + const size_t block_id_in_tensor = blockIdx.x - vb_tensor_base / ELTS_PER_CHUNK; + block_id_Y = block_id_in_tensor / chunks_X; + block_id_X = block_id_in_tensor % chunks_X; + } + + const size_t rows = + get_tensor_rows_num(tensor_id, shape_rep, first_logical_dim, first_dims_ptr, num_tensors); + const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr); + + // Compute per-tensor scale stride from cols (matches group_quantize kernel) + const size_t scale_stride = USE_ROWWISE_SCALING + ? DIVUP_TO_MULTIPLE(DIVUP(cols, static_cast(32)), 4) + : DIVUP_TO_MULTIPLE(cols, 128); + + const size_t tensor_base = is_single_tensor ? 0 : static_cast(offsets_ptr[tensor_id]); + + // Select TMA descriptors (static for single tensor, per-tensor for multi-tensor) + const CUtensorMap &tensor_map_input = + is_single_tensor ? tensor_map_input_static : g_tensor_maps_input[tensor_id]; + const CUtensorMap &tensor_map_output = + is_single_tensor ? tensor_map_output_static : g_tensor_maps_output[tensor_id]; + + if (!is_single_tensor) { + fence_acquire_tensormap(&tensor_map_input); + fence_acquire_tensormap(&tensor_map_output); + } + + const int chunk_offset_Y = block_id_Y * CHUNK_DIM_Y; + const int chunk_offset_X = block_id_X * CHUNK_DIM_X; + + // Per-tensor scale offset + constexpr size_t SCALE_DIVISOR = USE_ROWWISE_SCALING ? SCALE_DIM_X : SCALE_DIM_Y; + size_t scales_base_offset; + if (is_single_tensor) { + scales_base_offset = 0; + } else if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM) { + const size_t sum_prev_cols = tensor_base / first_logical_dim; + if constexpr (USE_ROWWISE_SCALING) { + // Scale layout: DIVUP_TO_MULTIPLE(R, 128) rows x (Ki/32) cols per tensor + const size_t padded_rows = DIVUP_TO_MULTIPLE(first_logical_dim, static_cast(128)); + scales_base_offset = (padded_rows / SCALE_DIM_X) * sum_prev_cols; + } else { + // Scale layout: DIVUP_TO_MULTIPLE(ceil(R/32), 4) rows x Ki cols per tensor + const size_t padded_scale_rows = DIVUP_TO_MULTIPLE( + DIVUP(first_logical_dim, static_cast(SCALE_DIM_Y)), static_cast(4)); + scales_base_offset = padded_scale_rows * sum_prev_cols; + } + } else { + // VARYING_BOTH_DIMS: both dims 128-padded, original formula is exact + scales_base_offset = tensor_base / SCALE_DIVISOR; + } + const e8m0_t *const tensor_scales_ptr = scales_ptr + scales_base_offset; + + const int scales_rowwise_chunk_offset_Y = block_id_Y * SCALES_ROWWISE_PER_CHUNK_Y; + const int scales_rowwise_chunk_offset_X = block_id_X * SCALES_ROWWISE_PER_CHUNK_X; + const int scales_colwise_chunk_offset_Y = block_id_Y * SCALES_COLWISE_PER_CHUNK_Y; + const int scales_colwise_chunk_offset_X = block_id_X * SCALES_COLWISE_PER_CHUNK_X; + + const int tid_rowwise_Y = threadIdx.x / THREADS_PER_CHUNK_X_ROWWISE; + const int tid_rowwise_X = threadIdx.x % THREADS_PER_CHUNK_X_ROWWISE; + const int tid_colwise_X = threadIdx.x % THREADS_PER_CHUNK_X_COLWISE; + + const int thread_offset_Y = tid_rowwise_Y; + const int thread_offset_X_rowwise = tid_rowwise_X * ELEMS_PER_THREAD; + + // Static shared memory (matching single-tensor dequantize) + __shared__ alignas(TMA_SHMEM_ALIGNMENT) IType in_sh[BUFFERS_NUM][SHMEM_DIM_Y][SHMEM_DIM_X]; + __shared__ alignas(TMA_SHMEM_ALIGNMENT) OType out_sh[BUFFERS_NUM][SHMEM_DIM_Y][SHMEM_DIM_X]; + + constexpr int shmem_buff_size = sizeof(in_sh) / BUFFERS_NUM; + constexpr int transaction_size = shmem_buff_size; + + const bool is_master_thread = (threadIdx.x == 0); + +// Initialize shared memory barrier +#pragma nv_diag_suppress static_var_with_dynamic_init + __shared__ alignas(8) uint64_t mbar[ITERATIONS]; + + initialize_barriers(mbar, is_master_thread); + + int parity = 0; + constexpr int iteration_zero = 0; + constexpr int buffer_zero = 0; + if (is_master_thread) { + const int chunk_stage_offset_Y = chunk_offset_Y; + const int chunk_stage_offset_X = chunk_offset_X; + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&in_sh[buffer_zero]), + reinterpret_cast(&tensor_map_input), chunk_stage_offset_X, + chunk_stage_offset_Y, &mbar[iteration_zero]); + + ptx::mbarrier_arrive_expect_tx(&mbar[iteration_zero], transaction_size); + } else { + ptx::mbarrier_arrive(&mbar[iteration_zero]); + } + +#pragma unroll + for (int iter = 0; iter < ITERATIONS; ++iter) { + const int buff = iter % BUFFERS_NUM; + const int next_iter = iter + 1; + if (next_iter < ITERATIONS) { + if (is_master_thread) { + const int next_buff = next_iter % BUFFERS_NUM; + const int chunk_it_offset_y = chunk_offset_Y + next_iter * BUFFER_DIM_Y; + const int chunk_it_offset_x = chunk_offset_X; + ptx::cp_async_bulk_tensor_2d_global_to_shared( + reinterpret_cast(&in_sh[next_buff]), + reinterpret_cast(&tensor_map_input), chunk_it_offset_x, + chunk_it_offset_y, &mbar[next_iter]); + + ptx::mbarrier_arrive_expect_tx(&mbar[next_iter], transaction_size); + } else { + ptx::mbarrier_arrive(&mbar[next_iter]); + } + } + + // Wait for the data to have arrived + ptx::mbarrier_wait_parity(&mbar[iter], parity); + + const int scale_offset_Y = + USE_ROWWISE_SCALING ? (scales_rowwise_chunk_offset_Y + iter * BUFFER_DIM_Y + tid_rowwise_Y) + : (scales_colwise_chunk_offset_Y + (iter * BUFFER_DIM_Y) / SCALE_DIM_Y); + + const int scale_offset_X = + USE_ROWWISE_SCALING + ? (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 * scale_stride + scale_offset_X; + const e8m0_t biased_exponent = tensor_scales_ptr[scale_idx]; + const float block_scale = ptx::exp2f(biased_exponent); + + if constexpr (USE_ROWWISE_SCALING) { + Vec in; + Vec out; + + const int shmem_offset_y = thread_offset_Y; + const int shmem_offset_x = thread_offset_X_rowwise; + in.load_from(&in_sh[buff][shmem_offset_y][shmem_offset_x]); + +#pragma unroll + for (int j = 0; j < ELEMS_PER_THREAD; ++j) { + out.data.elt[j] = static_cast(block_scale * static_cast(in.data.elt[j])); + } + out.store_to(&out_sh[buff][shmem_offset_y][shmem_offset_x]); + } else { +#pragma unroll + for (int i = 0; i < BUFFER_DIM_Y; ++i) { + const float elt = static_cast(in_sh[buff][i][tid_colwise_X]); + out_sh[buff][i][tid_colwise_X] = static_cast(block_scale * elt); + } + } + + // Wait for shared memory writes to be visible to TMA engine. + ptx::fence_proxy_async_shared_cta(); + __syncthreads(); + + // Initiate TMA transfer to copy shared memory to global memory + if (is_master_thread) { + const int chunk_it_offset_y = chunk_offset_Y + iter * BUFFER_DIM_Y; + const int chunk_it_offset_x = chunk_offset_X; + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), chunk_it_offset_x, + chunk_it_offset_y, reinterpret_cast(&out_sh[buff])); + + ptx::cp_async_bulk_commit_group(); + ptx::cp_async_bulk_wait_group_read<1>(); + } + } + ptx::cp_async_bulk_wait_group_read<0>(); + __syncthreads(); + + destroy_barriers(mbar, is_master_thread); +#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} +} // namespace group_dequantize_kernel + +inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output, + cudaStream_t stream) { + using namespace group_dequantize_kernel; + + checkCuDriverContext(stream); + + const bool use_rowwise_scaling = input->has_data(); + const bool use_colwise_scaling = input->has_columnwise_data(); + NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling, + "Input tensor must have either rowwise or columnwise data."); + NVTE_CHECK(!(use_rowwise_scaling && use_colwise_scaling), + "Dequantize only supports rowwise or columnwise scaling, not both simultaneously."); + + NVTE_CHECK(!input->with_gemm_swizzled_scales, "Input must have scales in compact format."); + NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision."); + NVTE_CHECK(!is_fp4_dtype(output->dtype()), "Output must not be FP4."); + NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input must have FP8 type."); + + NVTE_CHECK(input->num_tensors == output->num_tensors, + "Number of input and output tensors must be same."); + + ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + if (input->all_same_shape()) { + shape_rep = ShapeRepresentation::SAME_BOTH_DIMS; + } else if (input->all_same_first_dim()) { + shape_rep = ShapeRepresentation::VARYING_LAST_DIM; + } else if (input->all_same_last_dim()) { + shape_rep = ShapeRepresentation::VARYING_FIRST_DIM; + } else if (input->varying_both_dims()) { + shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS; + } + + const bool is_single_tensor = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS || + shape_rep == ShapeRepresentation::VARYING_FIRST_DIM); + + const size_t first_logical_dim = input->logical_shape.data[0]; + const size_t last_logical_dim = input->logical_shape.data[1]; + const size_t elts_total = first_logical_dim * last_logical_dim; + + const size_t num_tensors = input->num_tensors; + + constexpr size_t CHUNK_DIM_Y = DequantizeConfig::CHUNK_DIM_Y; + constexpr size_t CHUNK_DIM_X = DequantizeConfig::CHUNK_DIM_X; + constexpr size_t THREADS_PER_CHUNK = DequantizeConfig::THREADS_PER_CHUNK; + constexpr size_t SHMEM_DIM_Y = DequantizeConfig::SHMEM_DIM_Y; + constexpr size_t SHMEM_DIM_X = DequantizeConfig::SHMEM_DIM_X; + + size_t blocks = 0; + if (is_single_tensor) { + const size_t blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y); + const size_t blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X); + blocks = blocks_Y * blocks_X; + } else { + NVTE_CHECK(num_tensors <= MAX_SUPPORTED_TENSOR_DESCRIPTORS, + "Number of tensors in a group is larger than " + "the MAX number of supported descriptors (64)."); + NVTE_CHECK(last_logical_dim % CHUNK_DIM_X == 0, + "Last dimension of a grouped tensor should be divisible by 128."); + if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM) { + blocks = DIVUP(first_logical_dim, CHUNK_DIM_Y) * (last_logical_dim / CHUNK_DIM_X); + } else { + blocks = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X); + } + } + + const dim3 grid(blocks); + const dim3 block(THREADS_PER_CHUNK); + + const int64_t *const offsets_ptr = reinterpret_cast(input->tensor_offsets.dptr); + const int64_t *const first_dims_ptr = reinterpret_cast(input->first_dims.dptr); + const int64_t *const last_dims_ptr = reinterpret_cast(input->last_dims.dptr); + + const e8m0_t *const scales_ptr = + use_rowwise_scaling ? reinterpret_cast(input->scale_inv.dptr) + : reinterpret_cast(input->columnwise_scale_inv.dptr); + + const SimpleTensor &input_data = use_rowwise_scaling ? input->data : input->columnwise_data; + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( + 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, first_logical_dim, last_logical_dim, + SHMEM_DIM_Y, SHMEM_DIM_X, last_logical_dim, 0, + typeToNumBits(input->dtype())); + create_2D_tensor_map(tensor_map_output, output->data, first_logical_dim, last_logical_dim, + SHMEM_DIM_Y, SHMEM_DIM_X, last_logical_dim, 0, + typeToNumBits(output->dtype())); + + // Update tensor descriptors before launching the kernel + if (!is_single_tensor) { + const IType *const input_dptr = reinterpret_cast(input_data.dptr); + OType *const output_dptr = reinterpret_cast(output->data.dptr); + + update_tma_descriptors<<>>( + tensor_map_input, tensor_map_output, input_dptr, output_dptr, shape_rep, + num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr, + last_dims_ptr); + } + + if (use_rowwise_scaling) { + group_dequantize_mxfp8_kernel<<>>( + tensor_map_input, tensor_map_output, shape_rep, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_ptr); + } else { + group_dequantize_mxfp8_kernel<<>>( + tensor_map_input, tensor_map_output, shape_rep, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, scales_ptr); + }); // NOLINT(*) + ); // NOLINT(*) + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace mxfp8 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_MXFP8_CUH_ diff --git a/transformer_engine/common/include/transformer_engine/cast.h b/transformer_engine/common/include/transformer_engine/cast.h index f650b19dec..554d8c1ac9 100644 --- a/transformer_engine/common/include/transformer_engine/cast.h +++ b/transformer_engine/common/include/transformer_engine/cast.h @@ -407,8 +407,6 @@ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, cudaStream_t stream); /*! \brief Casts input tensor from reduced to higher precision. - * If the scaling mode of the input tensor is set to NVTE_MXFP8_1D_SCALING, - * the block dequantization (MXFP8) of the specified shape of the block will be used. * In case of the MXFP8 dequantization, the dequantized values are stored to the rowwise * data of the output tensor, regardless of whether the row- or columnwise scaling is used. * @@ -418,6 +416,17 @@ void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, */ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream); +/*! \brief Casts input grouped tensor from reduced to higher precision. + * In case of the MXFP8 dequantization, the dequantized values are stored to the rowwise + * data of the output tensor, regardless of whether the row- or columnwise scaling is used. + * + * \param[in] input Input grouped FP8/MXFP8 tensor to be cast. + * \param[in,out] output Output grouped tensor. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_group_dequantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream); + /*! \brief Casts multiple input tensors to quantized output tensors. * * \param[in] inputs List of input tensors to be cast. diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index b97504f2ae..eacd10eb30 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -281,7 +281,18 @@ void CheckGroupedTensorShapeArrays(const GroupedTensor &t, const std::string &na // Validate shape arrays (all optional) check_shape_array(t.first_dims, "first_dims"); check_shape_array(t.last_dims, "last_dims"); - check_shape_array(t.tensor_offsets, "tensor_offsets"); + + // tensor_offsets uses CSR-style prefix-sum layout with num_tensors+1 entries: + // offsets[i] = start of tensor i, offsets[num_tensors] = total elements + if (t.tensor_offsets.has_data()) { + NVTE_CHECK(t.tensor_offsets.shape.size() == 1, "Grouped tensor ", name, + " tensor_offsets must be 1D"); + NVTE_CHECK(t.tensor_offsets.dtype == DType::kInt64, "Grouped tensor ", name, + " tensor_offsets must have dtype Int64"); + NVTE_CHECK(t.tensor_offsets.shape[0] == t.num_tensors + 1, "Grouped tensor ", name, + " tensor_offsets size (", t.tensor_offsets.shape[0], ") must equal num_tensors+1 (", + t.num_tensors + 1, ")"); + } // tensor_offsets is required if any dimension varies // (i.e., required unless all_same_shape()) diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 9d2513835c..e40d39ee29 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -124,6 +124,7 @@ class Quantizer { virtual ~Quantizer() = default; + DType dtype = DType::kNumTypes; bool rowwise_usage = true; bool columnwise_usage = true; bool internal = false; @@ -165,7 +166,6 @@ class Float8Quantizer : public Quantizer { at::Tensor scale; at::Tensor scale_inv; at::Tensor amax; - DType dtype; explicit Float8Quantizer(const py::handle& quantizer); @@ -198,7 +198,6 @@ class Float8CurrentScalingQuantizer : public Quantizer { 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; @@ -247,8 +246,6 @@ class Float8CurrentScalingQuantizer : public Quantizer { class Float8BlockQuantizer : public Quantizer { public: - // Which float8 type is used for q data. - DType dtype; // Options about how to quantize the tensor // Quantization scales are rounded down to powers of 2. bool force_pow_2_scales = false; @@ -290,8 +287,6 @@ class Float8BlockQuantizer : public Quantizer { class MXFP8Quantizer : public Quantizer { public: - DType dtype; - explicit MXFP8Quantizer(const py::handle& quantizer); NVTEScalingMode get_scaling_mode() const override { return NVTE_MXFP8_1D_SCALING; } @@ -316,8 +311,6 @@ class MXFP8Quantizer : public Quantizer { class NVFP4Quantizer : public Quantizer { public: - // fp4 dtype - DType dtype; // amax reduction for low precision FP4 AG bool with_amax_reduction; c10::intrusive_ptr amax_reduction_group; diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9890f6742a..fb5783dfcb 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -309,6 +309,8 @@ py::object dequantize(const py::handle &input, DType otype); py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims); +py::object group_dequantize(const py::handle &input, DType otype); + py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors, std::optional first_dims); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index b689a1c1b4..5fb162c72d 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -318,6 +318,85 @@ py::object dequantize(const py::handle &input, transformer_engine::DType otype) return out; } +py::object group_dequantize(const py::handle &input, transformer_engine::DType otype) { + using namespace pybind11::literals; + init_extension(); + + // Extract fields from the Python GroupedTensor. + const auto num_tensors = input.attr("num_tensors").cast(); + const auto logical_shape_py = input.attr("logical_shape").cast(); + const auto logical_first_dim = logical_shape_py[0].cast(); + const auto logical_last_dim = logical_shape_py[1].cast(); + const std::vector logical_shape = {logical_first_dim, logical_last_dim}; + const auto &quantizer = convert_quantizer(input.attr("quantizer")); + + // Extract optional tensor attributes. + auto get_optional_tensor = [&input](const char *name) -> std::optional { + auto attr = input.attr(name); + if (attr.is_none()) return std::nullopt; + return attr.cast(); + }; + auto rowwise_data = get_optional_tensor("rowwise_data"); + auto columnwise_data = get_optional_tensor("columnwise_data"); + auto rowwise_scale_inv = get_optional_tensor("scale_inv"); + auto columnwise_scale_inv = get_optional_tensor("columnwise_scale_inv"); + auto first_dims = get_optional_tensor("first_dims"); + auto last_dims = get_optional_tensor("last_dims"); + auto tensor_offsets = get_optional_tensor("tensor_offsets"); + + // Early-return for empty input. + if (logical_first_dim == 0 || logical_last_dim == 0) { + NoneQuantizer q{py::none()}; + auto [out_cpp, out_py] = + q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), first_dims, + logical_first_dim, logical_last_dim); + return py::reinterpret_borrow(out_py); + } + + // Build input GroupedTensorWrapper. + // Data tensors are stored as flat 1D buffers; use the quantizer's dtype + // (e.g. kFloat8E4M3) rather than the raw tensor scalar_type (uint8). + auto input_cpp = GroupedTensorWrapper(num_tensors, logical_shape, quantizer->get_scaling_mode()); + if (rowwise_data.has_value()) { + input_cpp.set_rowwise_data(rowwise_data->data_ptr(), quantizer->dtype, + std::vector{static_cast(rowwise_data->numel())}); + if (rowwise_scale_inv.has_value()) { + input_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E8M0, + getTensorShape(*rowwise_scale_inv)); + } + } + if (columnwise_data.has_value()) { + input_cpp.set_columnwise_data( + columnwise_data->data_ptr(), quantizer->dtype, + std::vector{static_cast(columnwise_data->numel())}); + if (columnwise_scale_inv.has_value()) { + input_cpp.set_columnwise_scale_inv(columnwise_scale_inv->data_ptr(), DType::kFloat8E8M0, + getTensorShape(*columnwise_scale_inv)); + } + } + if (first_dims.has_value()) { + input_cpp.set_first_dims(first_dims->data_ptr(), DType::kInt64, getTensorShape(*first_dims)); + } + if (last_dims.has_value()) { + input_cpp.set_last_dims(last_dims->data_ptr(), DType::kInt64, getTensorShape(*last_dims)); + } + if (tensor_offsets.has_value()) { + input_cpp.set_tensor_offsets(tensor_offsets->data_ptr(), DType::kInt64, + getTensorShape(*tensor_offsets)); + } + + // Create output GroupedTensor using NoneQuantizer. + NoneQuantizer q{py::none()}; + auto [out_cpp, out_py] = q.create_grouped_tensor(num_tensors, logical_shape, otype, py::none(), + first_dims, logical_first_dim, logical_last_dim); + + NVTE_SCOPED_GIL_RELEASE({ + nvte_group_dequantize(input_cpp.data(), out_cpp.data(), at::cuda::getCurrentCUDAStream()); + }); + + return py::reinterpret_borrow(out_py); +} + namespace { void multi_tensor_quantize_impl(const std::vector &input_list, diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 4a20be6361..27d26d3dab 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -141,6 +141,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("otype")); 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, + "Dequantize group tensor", py::arg("input"), py::arg("otype")); m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims")); m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize, diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index ff1c78f695..7e2fea45f3 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -91,7 +91,8 @@ def _initialize_storage_fields( scale: Scale buffer (for FP8-DS only) first_dims: Device tensor of int64 array of length num_tensors (or None if uniform) last_dims: Device tensor of int64 array of length num_tensors (or None if uniform) - tensor_offsets: Device tensor of int64 array of length num_tensors (or None if uniform) + tensor_offsets: Device tensor of int64 array of length num_tensors+1 (CSR-style, + or None if uniform). offsets[i] = start of tensor i, offsets[num_tensors] = total. offsets: Vector of integer offsets for each tensor. """ # `requires_grad` and `stride` are accepted for API symmetry with From c5a4fd5a39d4bbf8597ab3402160b7b42a623c96 Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Sat, 18 Apr 2026 00:43:31 +0800 Subject: [PATCH 360/521] [PyTorch] Add FA4 Support (#2432) * add fa4 support Signed-off-by: Xin Yao * comment out unused import for cp Signed-off-by: Xin Yao * fix lint Signed-off-by: Xin Yao * install fa4 in L3 test Signed-off-by: Xin Yao * fix sm90 Signed-off-by: Xin Yao --------- Signed-off-by: Xin Yao --- qa/L3_pytorch_FA_versions_test/test.sh | 7 +- tests/pytorch/attention/test_attention.py | 135 +++++++++- .../dot_product_attention/backends.py | 72 ++++- .../attention/dot_product_attention/utils.py | 251 ++++++++++++++---- 4 files changed, 404 insertions(+), 61 deletions(-) diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index bbfc4db5ba..642eb93b06 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -18,10 +18,10 @@ sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); pri export FLASH_ATTN_CUDA_ARCHS=$sm_arch if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.3) + FA_versions=(2.8.3 4.0.0b8) elif [ $sm_arch -eq 90 ] then - FA_versions=(2.7.3 2.8.3 3.0.0b1) + FA_versions=(2.7.3 2.8.3 3.0.0b1 4.0.0b8) fi for fa_version in "${FA_versions[@]}" @@ -31,6 +31,9 @@ do if [ "${fa_version}" \< "3.0.0" ] then pip3 install flash-attn==${fa_version} --no-build-isolation + elif [[ "${fa_version}" == 4.* ]] + 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 diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 2eb307aa48..38d8626b4b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -53,7 +53,7 @@ ) _current_file = pathlib.Path(__file__).resolve() -sys.path.append(str(_current_file.parent.parent)) +sys.path = [str(_current_file.parent.parent)] + sys.path from utils import ( reset_rng_states, compare_and_assert, @@ -362,6 +362,139 @@ def test_dpa_num_splits(dtype, model_configs, model): ) +# ============================== +# Flash Attention 4 (FA4) tests +# ============================== + +model_configs_fa4_base = { + # test: ModelConfig(b, sq, hq, dqk) + # Standard head dims + "fa4_base_1": ModelConfig(4, 128, 16, 64), + "fa4_base_2": ModelConfig(2, 2048, 24, 128, attn_mask_type="causal"), + "fa4_base_3": ModelConfig(2, 1024, 8, 96, attn_mask_type="causal"), + # GQA + "fa4_gqa_1": ModelConfig(2, 1024, 32, 128, num_gqa_groups=8, attn_mask_type="causal"), + "fa4_gqa_2": ModelConfig(2, 1024, 16, 128, num_gqa_groups=1, attn_mask_type="causal"), + # num_splits + "fa4_splits_1": ModelConfig(2, 2048, 24, 128, num_splits=2), + "fa4_splits_2": ModelConfig(1, 2048, 24, 128, max_seqlen_kv=4096, num_splits=4), +} + + +@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_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + + +model_configs_fa4_mla = { + # test: ModelConfig(b, sq, hq, dqk, head_dim_v=dv) + "fa4_mla_1": ModelConfig(4, 128, 16, 128, head_dim_v=64), + "fa4_mla_2": ModelConfig(2, 128, 16, 64, max_seqlen_kv=256, head_dim_v=128), + "fa4_mla_3": ModelConfig(2, 1024, 16, 96, head_dim_v=64, attn_mask_type="causal"), + # dqk=128, dv=96: FA4 SM100 backward has dK_reduce_ncol misalignment for dV; + # the backend filter should reject FA4 and fall back to another backend. + "fa4_mla_4": ModelConfig(2, 1024, 16, 128, head_dim_v=96, attn_mask_type="causal"), + # DeepSeek-style MLA: dqk=192, dv=128 (supported on SM100 as special case) + "fa4_mla_deepseek": ModelConfig(2, 1024, 16, 192, head_dim_v=128, attn_mask_type="causal"), +} + + +@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()) +def test_dpa_fa4_mla(dtype, model_configs, model): + """Test DotProductAttention with FA4: MLA (head_dim_qk != head_dim_v)""" + test_dot_product_attention( + dtype, model_configs, model, False, True, "bshd_bshd_bshd", False, False + ) + + +model_configs_fa4_swa = { + # test: ModelConfig(b, sq, hq, dqk, window_size=(left, right)) + "fa4_swa_1": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal", window_size=(128, 0)), + "fa4_swa_2": ModelConfig(2, 2048, 24, 64, attn_mask_type="causal", window_size=(64, 0)), + "fa4_swa_3": ModelConfig( + 2, 2048, 16, 128, num_gqa_groups=4, attn_mask_type="causal", window_size=(256, 0) + ), + "fa4_swa_4": ModelConfig( + 2, 2048, 16, 128, attn_mask_type="padding_causal", window_size=(128, 0) + ), +} + + +@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()) +@pytest.mark.parametrize("qkv_layout", ["sbhd_sbhd_sbhd", "bshd_bshd_bshd"]) +def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): + """Test DotProductAttention with FA4: sliding window attention""" + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, True, False) + + +model_configs_fa4_varlen = { + # test: ModelConfig(b, sq, hq, dqk) + "fa4_varlen_1": ModelConfig(4, 128, 16, 64, attn_mask_type="padding"), + "fa4_varlen_2": ModelConfig(2, 2048, 24, 128, attn_mask_type="padding_causal"), + "fa4_varlen_3": ModelConfig( + 2, 2048, 24, 128, num_gqa_groups=4, attn_mask_type="padding_causal" + ), + "fa4_varlen_4": ModelConfig(2, 128, 16, 64, max_seqlen_kv=256, attn_mask_type="padding"), +} + + +@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()) +@pytest.mark.parametrize("qkv_layout", ["thd_thd_thd", "bshd_bshd_bshd"]) +def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): + """Test DotProductAttention with FA4: variable-length sequences (varlen/thd)""" + test_dot_product_attention(dtype, model_configs, model, False, True, qkv_layout, False, False) + + +model_configs_fa4_mask = { + # test: ModelConfig(b, sq, hq, dqk) + "fa4_mask_no_mask": ModelConfig(2, 1024, 16, 128), + "fa4_mask_causal": ModelConfig(2, 1024, 16, 128, attn_mask_type="causal"), + "fa4_mask_causal_br": ModelConfig(2, 1024, 16, 128, attn_mask_type="causal_bottom_right"), + "fa4_mask_padding": ModelConfig(2, 1024, 16, 128, attn_mask_type="padding"), + "fa4_mask_padding_causal": ModelConfig(2, 1024, 16, 128, attn_mask_type="padding_causal"), + "fa4_mask_padding_causal_br": ModelConfig( + 2, 1024, 16, 128, attn_mask_type="padding_causal_bottom_right" + ), +} + + +@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()) +def test_dpa_fa4_mask(dtype, model_configs, model): + """Test DotProductAttention with FA4: various attention mask types""" + test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + + model_configs_softmax = { # test: ModelConfig(b, sq, hq, dqk) "softmax_1_0": ModelConfig(2, 2048, 64, 64, num_gqa_groups=8), diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 19da8ebffe..ecf3af2bf0 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -80,7 +80,7 @@ from transformer_engine.pytorch.export import is_in_onnx_export_mode from transformer_engine.pytorch.graph import is_graph_capturing -# Global vars for flash attn v2 and v3 imports +# Global vars for flash attn v2 flash_attn_cuda_bwd = None flash_attn_func = None flash_attn_varlen_func = None @@ -88,6 +88,8 @@ _flash_attn_bwd = None _flash_attn_varlen_fwd = None _flash_attn_varlen_bwd = None + +# Try to import Flash Attention v2 try: fa_utils.version = PkgVersion(PkgVersion(get_pkg_version("flash-attn")).public) except PackageNotFoundError: @@ -130,12 +132,16 @@ ), fa_utils.version, ) + +# Try to import Flash Attention v3 try: fa_utils.fa3_version = PkgVersion(PkgVersion(get_pkg_version("flash-attn-3")).public) except PackageNotFoundError: flash_attn_func_v3 = None flash_attn_varlen_func_v3 = None flash_attn_with_kvcache_v3 = None + _flash_attn_fwd_v3 = None + _flash_attn_bwd_v3 = None # pass # only print warning if use_flash_attention_3 = True in get_attention_backend else: from flash_attn_interface import flash_attn_func as flash_attn_func_v3 @@ -150,6 +156,20 @@ fa_utils.set_flash_attention_3_params() +# Try to import Flash Attention v4 +try: + fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-4")) +except PackageNotFoundError: + flash_attn_func_v4 = None + flash_attn_varlen_func_v4 = None +else: + 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, + ) + + fa_utils.set_flash_attention_4_params() + # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" @@ -916,8 +936,13 @@ 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 flash_attention_backend > PkgVersion("3.0.0b"): + if flash_attention_backend is not None and PkgVersion( + "3.0.0b" + ) < flash_attention_backend < PkgVersion("4.0.0"): use_flash_attn_3 = True if context_parallel and all( not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer] @@ -971,24 +996,55 @@ def forward( # | | thd + padding # | flash_attn_with_kvcache | KV cache (not-paged/paged), i.e. # | | bshd/sbhd/thd + padding + # FA v4 | flash_attn_func | bshd/sbhd + not padding + # | flash_attn_varlen_func | bshd/sbhd + padding + # | | thd + padding fa_optional_forward_args_thd = [] if qkv_format in ["bshd", "sbhd"] and "padding" not in attn_mask_type: - func = ( - flash_attn_func if not use_flash_attn_3 else flash_attn_func_v3 - ) # pylint: disable=possibly-used-before-assignment + func = None + if use_flash_attn_4: + func = flash_attn_func_v4 + elif use_flash_attn_3: + func = flash_attn_func_v3 + else: + func = flash_attn_func else: - if not use_flash_attn_3: + if use_flash_attn_4: + func = flash_attn_varlen_func_v4 + elif not use_flash_attn_3: func = flash_attn_varlen_func elif inference_params is None: func = flash_attn_varlen_func_v3 # pylint: disable=possibly-used-before-assignment else: func = flash_attn_with_kvcache_v3 # pylint: disable=possibly-used-before-assignment - if not use_flash_attn_3 or inference_params is None: + 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(max_seqlen_q) fa_optional_forward_args_thd.append(max_seqlen_kv) - if not use_flash_attn_3: + if use_flash_attn_4: + fa_4_optional_forward_kwargs = { + "window_size": window_size, + "num_splits": num_splits, + } + if inference_params is None: + fa_4_optional_forward_kwargs["deterministic"] = self.deterministic + if func is flash_attn_varlen_func_v4: + fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q + fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv + fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q + fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv + output = func( + query_layer, + key_layer, + value_layer, + softmax_scale=self.softmax_scale, + causal="causal" in attn_mask_type, + **fa_4_optional_forward_kwargs, + ) + if isinstance(output, (List, Tuple)): + output = output[0] + elif not use_flash_attn_3: fa_optional_forward_kwargs = {} if fa_utils.v2_3_plus: fa_optional_forward_kwargs["window_size"] = window_size diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 13d1347a1e..20228ddb80 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -138,6 +138,13 @@ class FlashAttentionUtils: (2) cd flash-attention/hopper && python setup.py install""" v3_warning_printed = False + v4_is_installed = False + fa4_version = PkgVersion("0") + use_v4 = False + v4_installation_steps = """\ +pip install flash-attn-4==4.0.0b8 nvidia-cutlass-dsl[cu13]""" + v4_warning_printed = False + @staticmethod def set_flash_attention_version(): """ @@ -164,6 +171,13 @@ def set_flash_attention_3_params(): PkgVersion("3.0.0b") < FlashAttentionUtils.fa3_version < PkgVersion("3.0.0") ) + @staticmethod + def set_flash_attention_4_params(): + """ + Setup version info for FA v4.x + """ + FlashAttentionUtils.v4_is_installed = True + @dataclass(eq=True) class AttentionParams: @@ -354,8 +368,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]) + ), "flash_attn_version": ( str(FlashAttentionUtils.version) if FlashAttentionUtils.is_installed @@ -366,6 +381,11 @@ def get_attention_backend( if FlashAttentionUtils.v3_is_installed else "not installed" ), + "flash_attn_4_version": ( + str(FlashAttentionUtils.fa4_version) + if FlashAttentionUtils.v4_is_installed + else "not installed" + ), "cudnn_version": ".".join([str(i) for i in cudnn_version]), } attention_params_dict = { @@ -409,6 +429,7 @@ def get_attention_backend( use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1")) use_flash_attention_2 = use_flash_attention use_flash_attention_3 = use_flash_attention + use_flash_attention_4 = use_flash_attention flash_attention_backend = None use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1")) use_unfused_attention = int(os.getenv("NVTE_UNFUSED_ATTN", "1")) @@ -416,6 +437,8 @@ def get_attention_backend( logger.debug("Disabling FlashAttention 2 due to NVTE_FLASH_ATTN=0") if not use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 due to NVTE_FLASH_ATTN=0") + if not use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 due to NVTE_FLASH_ATTN=0") if not use_fused_attention: logger.debug("Disabling FusedAttention due to NVTE_FUSED_ATTN=0") if not use_unfused_attention: @@ -433,6 +456,18 @@ def get_attention_backend( if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 for compute capability != sm90") use_flash_attention_3 = False + # FA4 supports SM80, SM90, SM100, SM120 + if device_compute_capability < (8, 0): + if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 for compute capability < sm80") + use_flash_attention_4 = False + # 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") + use_flash_attention_4 = False # Filter: Data type if qkv_dtype not in [torch.bfloat16, torch.float16]: @@ -443,6 +478,13 @@ def get_attention_backend( qkv_dtype, ) use_flash_attention_2 = False + if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: + logger.debug( + "Disabling FlashAttention 4 for unsupported qkv_dtype = %s. " + "Supported: qkv_dtype = {torch.bfloat16, torch.float16}. ", + qkv_dtype, + ) + use_flash_attention_4 = False if qkv_dtype not in [torch.bfloat16, torch.float16, torch.float8_e4m3fn] or qkv_type not in [ torch.Tensor, Float8Tensor, @@ -470,7 +512,10 @@ def get_attention_backend( if fp8 and fp8_meta["recipe"].fp8_dpa: if use_flash_attention_2 and FlashAttentionUtils.is_installed: logger.debug("Disabling FlashAttention 2 for FP8 attention") - use_flash_attention_2 = False + use_flash_attention_2 = False + if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 for FP8 attention") + use_flash_attention_4 = False if use_flash_attention_3 and is_training: if FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 for FP8 training") @@ -524,6 +569,11 @@ def get_attention_backend( if use_flash_attention_2 and FlashAttentionUtils.is_installed: logger.debug("Disabling FlashAttention 2 for num_splits") use_flash_attention_2 = False + # FA4 SplitKV is only supported on SM100+ + if use_flash_attention_4 and device_compute_capability < (10, 0): + if FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 for num_splits on SM < 100") + use_flash_attention_4 = False if use_fused_attention: logger.debug("Disabling FusedAttention for num_splits") use_fused_attention = False @@ -549,6 +599,7 @@ def get_attention_backend( # Flash v2 | FP16/BF16 | non-paged/paged | sm80+ | bshd,sbhd,thd | >= 256 # 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 # Unfused | FP32/FP16/BF16 | non-paged/paged | all | bshd,sbhd,thd | >= 1 if inference_params is not None: # Temporarily disabling fused attention for kv caching for sm89/sm120 irrespective of @@ -597,6 +648,9 @@ def get_attention_backend( "Disabling FlashAttention 2 as paged attention requires flash-attn 2.5+" ) 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 # Filter: Head dimension if head_dim_qk != head_dim_v: @@ -607,7 +661,7 @@ def get_attention_backend( qkv_layout_group = qkv_layout.replace("b", "").replace("s", "").replace("t", "") if use_fused_attention and qkv_layout_group != "hd_hd_hd": logger.debug( - "Disabling FusedAttention as MLA is not supported with qkv_layout = %s", + "Disabling FusedAttention as MLA is not supported with qkv_layout = %s.", qkv_layout, ) use_fused_attention = False @@ -625,26 +679,30 @@ def get_attention_backend( ) use_fused_attention = False - if use_flash_attention_2 and ( - head_dim_qk > 256 - or head_dim_qk % 8 != 0 - or ( - head_dim_qk > 192 - and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) + if ( # pylint: disable=too-many-boolean-expressions + use_flash_attention_2 + and FlashAttentionUtils.is_installed + and ( + head_dim_qk > 256 + or head_dim_qk % 8 != 0 + or ( + head_dim_qk > 192 + and device_compute_capability not in ((8, 0), (9, 0), (10, 0), (12, 0)) + ) ) ): - if FlashAttentionUtils.is_installed: - logger.debug( - "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. " - "Supported: head_dim_qk = head_dim_v, head_dim_qk %%8 = 0, " - "head_dim_qk <= 256 (>192 requires sm80/90/100+). " - "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.", - head_dim_qk, - head_dim_v, - ".".join([str(i) for i in device_compute_capability]), - ) + logger.debug( + "Disabling FlashAttention 2 due to unsupported head_dim_qk and head_dim_v. " + "Supported: head_dim_qk = head_dim_v, head_dim_qk %%8 = 0, " + "head_dim_qk <= 256 (>192 requires sm80/90/100+). " + "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.", + head_dim_qk, + head_dim_v, + ".".join([str(i) for i in device_compute_capability]), + ) use_flash_attention_2 = False - if use_flash_attention_3: + + if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dtype): if head_dim_qk > 256 or num_heads % num_gqa_groups != 0: @@ -660,31 +718,80 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt return True if not _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dtype): - if FlashAttentionUtils.v3_is_installed: + logger.debug( + "Disabling FlashAttention 3 due to unsupported num_heads, num_gqa_groups, " + "head_dim_qk, head_dim_v or qkv_dtype. " + "Supported: head_dim_qk <= 256, and num_heads %% num_gqa_groups = 0, and " + "if head_dim_qk is different from head_dim_v, then " + "(head_dim_qk must in (128, 192] and head_dim_v in (96, 128]) or " + "(head_dim_qk <= 64 and head_dim_v <= 512), and " + "if head_dim_qk is different from head_dim_v and head_dim_v > 256, then " + "qkv_dtype requires fp16 and bf16 data type. " + "Found: num_heads = %s, num_gqa_groups = %s, " + "head_dim_qk = %s, head_dim_v = %s and qkv_dtype = %s.", + num_heads, + num_gqa_groups, + head_dim_qk, + head_dim_v, + qkv_dtype, + ) + 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: + logger.debug( + "Disabling FlashAttention 4 due to unsupported head dimensions. " + "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.", + head_dim_qk, + head_dim_v, + 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 + and is_training + and head_dim_qk != head_dim_v + and head_dim_qk >= 128 + and (10, 0) <= device_compute_capability < (12, 0) + ): + _tile_hdim = math.ceil(head_dim_qk / 16) * 16 + _tile_hdimv = math.ceil(head_dim_v / 16) * 16 + _dk_reduce_ncol = math.gcd(32, _tile_hdim // 2) + if (_tile_hdimv // 2) % _dk_reduce_ncol != 0: logger.debug( - "Disabling FlashAttention 3 due to unsupported num_heads, num_gqa_groups, " - "head_dim_qk, head_dim_v or qkv_dtype. " - "Supported: head_dim_qk <= 256, and num_heads %% num_gqa_groups = 0, and " - "if head_dim_qk is different from head_dim_v, then " - "(head_dim_qk must in (128, 192] and head_dim_v in (96, 128]) or " - "(head_dim_qk <= 64 and head_dim_v <= 512), and " - "if head_dim_qk is different from head_dim_v and head_dim_v > 256, then " - "qkv_dtype requires fp16 and bf16 data type. " - "Found: num_heads = %s, num_gqa_groups = %s, " - "head_dim_qk = %s, head_dim_v = %s and qkv_dtype = %s.", - num_heads, - num_gqa_groups, + "Disabling FlashAttention 4 for training due to SM100 backward kernel " + "bug with MLA head dimensions (dK_reduce_ncol misalignment for dV). " + "Found: head_dim_qk = %s, head_dim_v = %s.", head_dim_qk, head_dim_v, - qkv_dtype, ) - use_flash_attention_3 = False + use_flash_attention_4 = False # Filter: QKV layout if qkv_format == "thd": if pad_between_seqs: - if (use_flash_attention_2 and FlashAttentionUtils.is_installed) or ( - use_flash_attention_3 and FlashAttentionUtils.v3_is_installed + 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) ): logger.debug( "Disabling FlashAttention for qkv_format = thd when there is " @@ -709,9 +816,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt use_fused_attention = False # Filter: Dropout - if attention_dropout != 0.0 and use_flash_attention_3: - logger.debug("Disabling FlashAttention 3 for dropout") - use_flash_attention_3 = False + if attention_dropout != 0.0: + if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: + logger.debug("Disabling FlashAttention 3 for dropout") + use_flash_attention_3 = False + if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 for dropout") + use_flash_attention_4 = False # Filter: Softmax type # context_parallel | softmax_type | supported backends @@ -767,8 +878,17 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "Disabling UnfusedDotProductAttention as it does not support context parallelism" ) use_unfused_attention = False - if context_parallel and (use_flash_attention_2 or use_flash_attention_3): - if FlashAttentionUtils.is_installed or FlashAttentionUtils.v3_is_installed: + if context_parallel and use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 as it does not support context parallelism yet") + use_flash_attention_4 = False + if context_parallel and ( + use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4 + ): + if ( + FlashAttentionUtils.is_installed + or FlashAttentionUtils.v3_is_installed + or FlashAttentionUtils.v4_is_installed + ): if fp8 and fp8_meta["recipe"].fp8_dpa: logger.debug( "Disabling FlashAttention as it does not support context parallelism with FP8" @@ -852,8 +972,10 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # arbitrary | One tensor of shape broadcastable to | UnfusedDotProductAttention # | [b, h, sq, skv] | if attn_mask_type == "arbitrary": - if (use_flash_attention_2 and FlashAttentionUtils.is_installed) or ( - use_flash_attention_3 and FlashAttentionUtils.v3_is_installed + 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) ): logger.debug("Disabling FlashAttention for arbitrary mask") use_flash_attention = False @@ -861,7 +983,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention for arbitrary mask") use_fused_attention = False if ( - (use_flash_attention_2 or use_flash_attention_3) + (use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4) and attn_mask_type in ["causal", "padding_causal"] and max_seqlen_q != max_seqlen_kv ): @@ -940,13 +1062,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt " alignment for cross-attention" ) use_flash_attention = False + if use_flash_attention_4: + if FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 for ALiBi") + use_flash_attention_4 = False if ( core_attention_bias_type not in ["no_bias", "alibi"] or core_attention_bias_shape is not None ): - if (use_flash_attention_2 and FlashAttentionUtils.is_installed) or ( - use_flash_attention_3 and FlashAttentionUtils.v3_is_installed + 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) ): logger.debug("Disabling FlashAttention for pre/post_scale_bias") use_flash_attention = False @@ -1067,6 +1195,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "please install flash-attn >= 2.4.1." ) use_flash_attention_2 = False + if use_flash_attention_3 and deterministic and FlashAttentionUtils.v3_is_installed: + if head_dim_qk >= 256: + logger.debug( + "Disabling FlashAttention 3 for deterministic execution with head_dim_qk >= 256." + ) + use_flash_attention_3 = False if use_fused_attention and deterministic: if softmax_type != "vanilla": logger.debug( @@ -1104,12 +1238,25 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt # use_flash_attention may have been set above use_flash_attention_2 = use_flash_attention and use_flash_attention_2 use_flash_attention_3 = use_flash_attention and use_flash_attention_3 + use_flash_attention_4 = use_flash_attention and use_flash_attention_4 # `FusedAttention` and `FlashAttention` are faster backends than `UnfusedDotProductAttention`. # When `FusedAttention` does not support the provided attention params, and `FlashAttention` # does, we recommend users to install flash-attn if not installed already. if not use_fused_attention and _NVTE_FLASH_ATTN: if ( + use_flash_attention_4 + and not FlashAttentionUtils.v4_is_installed + and not FlashAttentionUtils.v4_warning_printed + and torch.cuda.current_device() == 0 + ): + logger.warning( + "flash-attn v4 may provide important feature support or performance improvement." + " Please install flash-attn v4 by \n%s", + FlashAttentionUtils.v4_installation_steps, + ) + FlashAttentionUtils.v4_warning_printed = True + elif ( use_flash_attention_3 and not FlashAttentionUtils.v3_is_installed and not FlashAttentionUtils.v3_warning_printed @@ -1141,12 +1288,16 @@ 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 not FlashAttentionUtils.v3_is_installed: use_flash_attention_3 = False - use_flash_attention = use_flash_attention_2 or use_flash_attention_3 + if use_flash_attention_4 and not FlashAttentionUtils.v4_is_installed: + use_flash_attention_4 = False + use_flash_attention = use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4 available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention] if use_flash_attention_2: flash_attention_backend = FlashAttentionUtils.version if use_flash_attention_3: flash_attention_backend = FlashAttentionUtils.fa3_version + if use_flash_attention_4: + flash_attention_backend = FlashAttentionUtils.fa4_version logger.debug( "Available backends = {FlashAttention=%s%s, FusedAttention=%s%s," @@ -1183,7 +1334,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt selected_backend = f"FusedAttention (sub-backend {int(fused_attention_backend)})" elif use_unfused_attention: selected_backend = "UnfusedDotProductAttention" - logger.debug("Selected backend = %s", selected_backend) + logger.debug("Selected backend = %s.", selected_backend) return ( use_flash_attention, From 262bc6cfe1bb8d20b9367c1e5339af78e7090b19 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:39:44 -0700 Subject: [PATCH 361/521] [JAX] Fix grouped quant checkpointing (#2889) * Fix grouped quant checkpointing Signed-off-by: Jeremy Berchtold * Cleanup Signed-off-by: Jeremy Berchtold * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Jeremy Berchtold Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/jax/dense.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/transformer_engine/jax/dense.py b/transformer_engine/jax/dense.py index dbd7bbb1ff..f8c30ffccb 100644 --- a/transformer_engine/jax/dense.py +++ b/transformer_engine/jax/dense.py @@ -429,10 +429,25 @@ def _grouped_dense_fwd_rule( # rowwise_casted_x.original_shape == (M, K) # colwise_casted_kernel.original_shape == (G, N, K) grouped_gemm_x = casted_x.get_tensor(usage=TensorUsage.LHS) + # Checkpoint the rowwise inputs so that te_grouped_quantize_ffi can be DCE'd in the + # backward-scan remat block. Without this, JAX would re-run the quantize kernel to + # obtain grouped_gemm_x / grouped_gemm_kernel for the forward-GEMM recomputation even + # though the colwise residuals (ctx_x / ctx_kernel) are already saved. With both + # orientations checkpointed, all outputs of the custom-call become dead in the remat trace. + grouped_gemm_x = ( + grouped_gemm_x.checkpoint(quantizer_set.x) + if isinstance(grouped_gemm_x, ScaledTensor) + else grouped_gemm_x + ) ctx_x = casted_x.get_tensor(usage=TensorUsage.LHS_TRANS) ctx_kernel = casted_kernel.get_tensor(usage=TensorUsage.RHS_TRANS) grouped_gemm_kernel = casted_kernel.get_tensor(usage=TensorUsage.RHS) + grouped_gemm_kernel = ( + grouped_gemm_kernel.checkpoint(quantizer_set.kernel) + if isinstance(grouped_gemm_kernel, ScaledTensor) + else grouped_gemm_kernel + ) output = tex.grouped_gemm( grouped_gemm_x, grouped_gemm_kernel, From 549f5ba4cf8a4d1184e3a8136bfcfa1434c16723 Mon Sep 17 00:00:00 2001 From: jomitchellnv <148147880+jomitchellnv@users.noreply.github.com> Date: Sun, 19 Apr 2026 23:51:24 -0700 Subject: [PATCH 362/521] adds NVFP4 Fused Adam support (#2797) * adds NVFP4 Fused Adam support Signed-off-by: Jonathan Mitchell * un xfail test Signed-off-by: Jonathan Mitchell * cleanup Signed-off-by: Jonathan Mitchell * adds back copy dispatch handler Signed-off-by: Jonathan Mitchell --------- Signed-off-by: Jonathan Mitchell Signed-off-by: Jonathan Mitchell Co-authored-by: Jonathan Mitchell Co-authored-by: Jonathan Mitchell Co-authored-by: vthumbe1503 --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 6 - .../fsdp2_tests/run_fsdp2_model.py | 29 +- tests/pytorch/test_nvfp4_fsdp2_hooks.py | 288 ++++++++++++++++++ .../pytorch/tensor/nvfp4_tensor.py | 189 ++++++++++++ 4 files changed, 500 insertions(+), 12 deletions(-) create mode 100644 tests/pytorch/test_nvfp4_fsdp2_hooks.py 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 60a23b9394..ac38bc4aa8 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -160,12 +160,6 @@ def test_fused_adam_fp8_master_weights(recipe_name): """ recipe = get_recipe_from_string(recipe_name) - if recipe_name == "NVFP4BlockScaling": - pytest.xfail( - f"{recipe_name}: quantized_model_init and FSDP2 is not currently supported, since the " - "block tensor is dequantized before we flatten it for FSDP2." - ) - world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index fce565ed9a..5a8c903c7d 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -37,6 +37,7 @@ import transformer_engine.pytorch as te import transformer_engine.common.recipe +from transformer_engine.pytorch.tensor import NVFP4Tensor import torch import torch.distributed as dist @@ -224,8 +225,8 @@ def _check_fp8_fsdp2_allgather(model): if device_mesh.ndim > 1 else device_mesh.get_group() ) - # Perform manual allgather on local_tensor. zeros_like will create hp tensor since torch_dispatch - # for local_tensor will go down the dequantization route. + # Perform manual allgather on local_tensor. zeros_like will create hp tensor since + # torch_dispatch for local_tensor will go down the dequantization route. gathered_tensor = [ torch.zeros_like(local_tensor) for _ in range(dist.get_world_size(group=dist_group)) ] @@ -239,7 +240,13 @@ def _check_fp8_fsdp2_allgather(model): module.unshard() # Make sure allgathered parameters match exactly for name, param in model.named_parameters(): - torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name]) + # NVFP4 scale unpad/repad through FSDP2 introduces small numerical + # differences vs the manual dequantize-then-allgather path. + if isinstance(param, NVFP4Tensor): + tols = dict(atol=5e-4, rtol=5e-3) + else: + tols = {} + torch.testing.assert_close(param.dequantize(), fp32_allgathered_params[name], **tols) # Revert model to original sharded state for module in model.modules(): # Not all modules are wrapped/sharded with FSDP2. @@ -363,9 +370,19 @@ 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 in ("Float8BlockScaling", "NVFP4BlockScaling") and fp8_init: - pytest.xfail(f"{recipe_name} + fp8_init: test_fp8_fsdp2_allgather is currently failing.") - + 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_nvfp4_fsdp2_hooks.py b/tests/pytorch/test_nvfp4_fsdp2_hooks.py new file mode 100644 index 0000000000..3fbd574964 --- /dev/null +++ b/tests/pytorch/test_nvfp4_fsdp2_hooks.py @@ -0,0 +1,288 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Unit tests for NVFP4Tensor FSDP2 all-gather hooks. + +These tests verify the pre/post all-gather round-trip logic on a single GPU +without requiring torchrun or multi-GPU setup. +""" + +import math +from typing import List, Tuple + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import ( + NVFP4Quantizer, + NVFP4Tensor, +) +from transformer_engine.pytorch.utils import round_up_to_nearest_multiple +from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE + +nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True) + +# Shapes that exercise various M/K combinations: +# - (512, 256): both dims cleanly divisible by 128 +# - (640, 128): M not a multiple of 128*2 but divisible by 16 +# - (256, 1024): K > M +_test_shapes: List[Tuple[int, int]] = [ + (512, 256), + (640, 128), + (256, 1024), +] + + +def _make_nvfp4_tensor(shape: Tuple[int, int]) -> NVFP4Tensor: + """Create an NVFP4Tensor from random BF16 data.""" + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=True, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=True, + stochastic_rounding=False, + with_random_sign_mask=False, + ) + src = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + return quantizer(src) + + +def _simulate_all_gather( + sharded_tensors: Tuple[torch.Tensor, ...], + world_size: int, +) -> Tuple[torch.Tensor, ...]: + """Simulate FSDP2 all-gather by concatenating shards along dim0.""" + return tuple(torch.cat([t] * world_size, dim=0) for t in sharded_tensors) + + +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +class TestNVFP4FSDP2Hooks: + """Tests for fsdp_pre_all_gather / fsdp_post_all_gather round-trip.""" + + @classmethod + def setup_class(cls) -> None: + torch.manual_seed(42) + torch.cuda.manual_seed(42) + + @pytest.mark.parametrize("shape", _test_shapes) + @pytest.mark.parametrize("world_size", [2, 4]) + def test_round_trip_shapes(self, shape: Tuple[int, int], world_size: int): + """Verify that pre_all_gather -> all_gather -> post_all_gather produces correct shapes.""" + M, K = shape + shard_M = M // world_size + shard_shape = (shard_M, K) + + qt = _make_nvfp4_tensor(shard_shape) + + # Pre all-gather + sharded_tensors, metadata = qt.fsdp_pre_all_gather( + mesh=None, + orig_size=None, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + + # Only rowwise tensors are all-gathered; columnwise is derived locally + assert len(sharded_tensors) == 2, "Expected 2 tensors (rowwise data + scale only)" + + rowwise_data, rowwise_scale_inv = sharded_tensors + + # Rowwise data: (shard_M, K//2) — unmodified + assert rowwise_data.shape == (shard_M, K // 2) + # Rowwise scale: unpadded dim0 to shard_M + assert rowwise_scale_inv.shape[0] == shard_M + + # Simulate all-gather + all_gather_outputs = _simulate_all_gather(sharded_tensors, world_size) + + # Post all-gather + result, _ = qt.fsdp_post_all_gather( + all_gather_outputs, + metadata, + param_dtype=torch.bfloat16, + ) + + # Verify output is NVFP4Tensor with correct logical shape + assert isinstance(result, NVFP4Tensor) + assert tuple(result.shape) == (M, K) + + # Verify internal data shapes + assert result._rowwise_data.shape == (M, K // 2) + + expected_rowwise_scale_shape = ( + round_up_to_nearest_multiple(M, 128), + round_up_to_nearest_multiple(math.ceil(K / NVFP4_BLOCK_SCALING_SIZE), 4), + ) + assert result._rowwise_scale_inv.shape == expected_rowwise_scale_shape + + # Columnwise data derived locally via _create_columnwise() + assert result._columnwise_data.shape == (K, M // 2) + + expected_col_scale_shape = ( + round_up_to_nearest_multiple(K, 128), + round_up_to_nearest_multiple(math.ceil(M / NVFP4_BLOCK_SCALING_SIZE), 4), + ) + assert result._columnwise_scale_inv.shape == expected_col_scale_shape + + @pytest.mark.parametrize("shape", _test_shapes) + def test_round_trip_data_integrity(self, shape: Tuple[int, int]): + """Verify data and dequantized values survive the pre -> all_gather -> post round-trip.""" + world_size = 2 + M, K = shape + shard_M = M // world_size + shard_shape = (shard_M, K) + + qt = _make_nvfp4_tensor(shard_shape) + + # Save original internal tensors for comparison + orig_rowwise_data = qt._rowwise_data.clone() + orig_rowwise_scale = qt._rowwise_scale_inv.clone() + orig_amax_row = qt._amax_rowwise.clone() + orig_amax_col = qt._amax_columnwise.clone() + orig_deq = qt.dequantize() + + # Pre all-gather + sharded_tensors, metadata = qt.fsdp_pre_all_gather( + mesh=None, + orig_size=None, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + + # Simulate all-gather (world_size copies — data from each "rank" is identical) + all_gather_outputs = _simulate_all_gather(sharded_tensors, world_size) + + # Post all-gather + result, _ = qt.fsdp_post_all_gather( + all_gather_outputs, + metadata, + param_dtype=torch.bfloat16, + ) + + # Since each "rank" has the same data, the full rowwise_data should be + # the original shard repeated world_size times + expected_rowwise_data = torch.cat([orig_rowwise_data] * world_size, dim=0) + assert torch.equal(result._rowwise_data, expected_rowwise_data) + + # Rowwise scale: each shard's unpadded scale is repeated, then repadded + # Check that the first shard_M rows of the scale match the original (unpadded) + assert torch.equal( + result._rowwise_scale_inv[:shard_M, :], + orig_rowwise_scale[:shard_M, :], + ) + + # Columnwise data is derived locally via _create_columnwise(), not all-gathered. + # Verify it was created and has the correct shape. + assert result._columnwise_data is not None + assert result._columnwise_data.shape == (K, M // 2) + assert result._columnwise_scale_inv is not None + + # Amax values passed through metadata — should be preserved + assert torch.equal(result._amax_rowwise, orig_amax_row) + assert torch.equal(result._amax_columnwise, orig_amax_col) + + # Dequantized values: the full tensor should dequantize to world_size copies of the shard + result_deq = result.dequantize() + expected_deq = torch.cat([orig_deq] * world_size, dim=0) + torch.testing.assert_close(result_deq, expected_deq) + + @pytest.mark.parametrize("shape", _test_shapes) + def test_in_place_update(self, shape: Tuple[int, int]): + """Verify the out= path (in-place update on subsequent iterations).""" + world_size = 2 + M, K = shape + shard_M = M // world_size + shard_shape = (shard_M, K) + + qt = _make_nvfp4_tensor(shard_shape) + + sharded_tensors, metadata = qt.fsdp_pre_all_gather( + mesh=None, + orig_size=None, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + all_gather_outputs = _simulate_all_gather(sharded_tensors, world_size) + + # First call: out=None -> creates new tensor + result, _ = qt.fsdp_post_all_gather( + all_gather_outputs, + metadata, + param_dtype=torch.bfloat16, + ) + first_deq = result.dequantize().clone() + + # Second call: out=result -> in-place update + result2, _ = qt.fsdp_post_all_gather( + all_gather_outputs, + metadata, + param_dtype=torch.bfloat16, + out=result, + ) + assert result2 is result # same object + torch.testing.assert_close(result2.dequantize(), first_deq) + + def test_swizzled_scales_rejected(self): + """Verify that GEMM-swizzled scales raise NotImplementedError.""" + shape = (512, 256) + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=True, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=False, + ) + quantizer.optimize_for_gemm = True + src = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + qt = quantizer(src) + + if not qt._with_gemm_swizzled_scales: + pytest.skip( + "NVFP4Quantizer.optimize_for_gemm is not yet wired up in C++. " + "Test will be unskipped once supported." + ) + + with pytest.raises(NotImplementedError, match="GEMM-swizzled"): + qt.fsdp_pre_all_gather( + mesh=None, + orig_size=None, + contiguous_orig_stride=None, + module=None, + mp_policy=None, + ) + + +@pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4) +class TestNVFP4DispatchHandlers: + """Tests for as_strided, slice, and record_stream dispatch handlers.""" + + def test_as_strided_noop(self): + """as_strided with matching shape/strides returns NVFP4Tensor.""" + qt = _make_nvfp4_tensor((256, 128)) + M, K = qt.shape + result = torch.ops.aten.as_strided.default(qt, [M, K], [K, 1], 0) + assert isinstance(result, NVFP4Tensor) + assert tuple(result.shape) == (M, K) + + def test_slice_noop(self): + """slice covering full dimension returns NVFP4Tensor.""" + qt = _make_nvfp4_tensor((256, 128)) + M, K = qt.shape + result = torch.ops.aten.slice.Tensor(qt, 0, 0, M) + assert isinstance(result, NVFP4Tensor) + assert tuple(result.shape) == (M, K) + + def test_record_stream(self): + """record_stream completes without error.""" + qt = _make_nvfp4_tensor((256, 128)) + stream = torch.cuda.Stream() + result = torch.ops.aten.record_stream.default(qt, stream) + assert result is None diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index eb514d3a9e..65678aa347 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -551,6 +551,122 @@ def get_usages(self) -> Dict[str, bool]: "columnwise": self._columnwise_data is not None, } + def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, mp_policy): + """Called by FSDP2 before all-gather of weights. + + Only all-gathers rowwise data and scales. Columnwise data is derived + locally in post_all_gather via _create_columnwise(), halving the + all-gather communication volume. + """ + # pylint: disable=unused-argument + + if self._with_gemm_swizzled_scales: + raise NotImplementedError( + "FSDP2 is not supported for NVFP4Tensors with GEMM-swizzled scales." + ) + + shard_M = math.prod(self.shape[:-1]) + + assert shard_M % NVFP4_BLOCK_SCALING_SIZE == 0, ( + f"FSDP2 requires shard_M ({shard_M}) to be a multiple of " + f"NVFP4_BLOCK_SCALING_SIZE ({NVFP4_BLOCK_SCALING_SIZE}). " + "Adjust model dimensions or world size." + ) + + assert self._rowwise_data is not None, ( + "FSDP2 requires rowwise data, but _rowwise_data is None. " + "Ensure the NVFP4Quantizer was created with rowwise=True." + ) + + # Rowwise data: (shard_M, K//2) — M in dim0, pass as-is + rowwise_data = self._rowwise_data + # Rowwise scale: (round_up(shard_M, 128), inner) — unpad dim0 to shard_M + rowwise_scale_inv = self._rowwise_scale_inv + if rowwise_scale_inv is not None: + rowwise_scale_inv = rowwise_scale_inv[:shard_M, :] + + columnwise_usage = self._quantizer.columnwise_usage + if columnwise_usage: + assert self._quantizer.with_2d_quantization, ( + "FSDP2 columnwise usage requires 2D quantization to be enabled. " + "Ensure the NVFP4Quantizer was created with with_2d_quantization=True." + ) + + # Only all-gather rowwise tensors; columnwise will be derived locally + # via _create_columnwise() in post_all_gather. + sharded_tensors = (rowwise_data, rowwise_scale_inv) + + # Pass amax via metadata (scalar, same on all ranks — not all-gathered) + metadata = ( + self._fp4_dtype, + columnwise_usage, + self._amax_rowwise, + self._amax_columnwise, + self.shape[-1], + ) + return sharded_tensors, metadata + + def fsdp_post_all_gather( + self, + all_gather_outputs: Tuple[torch.Tensor, ...], + metadata, + param_dtype: torch.dtype, + *, + out: Optional[NVFP4Tensor] = None, + ): + """Called by FSDP2 after all-gather of weights. + + Repads rowwise scales and constructs the full NVFP4Tensor from + 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 + + # Only rowwise data+scales were all-gathered + rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] + full_M = rowwise_data.shape[0] + + # Repad rowwise scale dim0 to round_up(full_M, 128) + if rowwise_scale_inv is not None: + target_m = round_up_to_nearest_multiple(full_M, 128) + current_m = rowwise_scale_inv.shape[0] + if current_m < target_m: + rowwise_scale_inv = torch.nn.functional.pad( + rowwise_scale_inv, (0, 0, 0, target_m - current_m) + ) + + logical_shape = (full_M, K) + + if out is not None: + # Update existing tensor in-place (subsequent iterations) + out._rowwise_data = rowwise_data + out._rowwise_scale_inv = rowwise_scale_inv + out._amax_rowwise = amax_rowwise + out._amax_columnwise = amax_columnwise + else: + # Construct new tensor (first iteration) + out = NVFP4Tensor( + shape=logical_shape, + dtype=param_dtype, + fp4_dtype=fp4_dtype, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=None, + columnwise_scale_inv=None, + amax_rowwise=amax_rowwise, + amax_columnwise=amax_columnwise, + quantizer=self._quantizer, + requires_grad=False, + with_gemm_swizzled_scales=False, + ) + + # Derive columnwise data locally via transpose instead of all-gathering it + if columnwise_usage: + out._create_columnwise() + + out._quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + return out, all_gather_outputs + @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): @@ -564,6 +680,79 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): return tensor.detach() return tensor.view(shape) + # as_strided — FSDP2 applies this on the unsharded param. + # Only the identity case (same shape, contiguous strides, zero offset) is supported. + # Non-identity as_strided cannot fall through because NVFP4 does not support + # dequantization, so we raise explicitly rather than producing undefined behavior. + if func == aten.as_strided.default: + tensor = args[0] + shape = args[1] + strides = args[2] + storage_offset = args[3] if len(args) > 3 else 0 + if ( + len(shape) == len(strides) == 2 + and tuple(strides) == (shape[-1], 1) + and tuple(shape) == tuple(tensor.size()) + and storage_offset == 0 + ): + return NVFP4Tensor.make_like(tensor) + raise NotImplementedError( + "NVFP4Tensor does not support non-identity as_strided " + f"(shape={shape}, strides={strides}, storage_offset={storage_offset}, " + f"tensor.size()={tuple(tensor.size())})" + ) + + # slice — FSDP2 applies this for shard unpadding. + # When the slice covers the full dimension, return self. + if func == aten.slice.Tensor: + tensor = args[0] + dim = args[1] if len(args) > 1 else 0 + start = args[2] if len(args) > 2 else None + end = args[3] if len(args) > 3 else None + step = args[4] if len(args) > 4 else 1 + if ( + step == 1 + and (start is None or start == 0) + and (end is None or end >= tensor.size(dim)) + ): + return NVFP4Tensor.make_like(tensor) + raise NotImplementedError( + "NVFP4Tensor does not support partial slicing " + f"(dim={dim}, start={start}, end={end}, " + f"tensor.size(dim)={tensor.size(dim)})" + ) + + # record_stream — FSDP2 records streams on all-gathered tensors. + if func == torch.ops.aten.record_stream.default: + qt, stream = args + for t in ( + qt._rowwise_data, + qt._columnwise_data, + qt._rowwise_scale_inv, + qt._columnwise_scale_inv, + qt._amax_rowwise, + qt._amax_columnwise, + ): + if t is not None and t.is_cuda: + t.record_stream(stream) + return None + + # copy_ — FSDP2 may call this during resharding or parameter writeback. + if func == aten.copy_.default: + dst, src = args[0], args[1] + if isinstance(src, NVFP4Tensor) and isinstance(dst, NVFP4Tensor): + if dst._rowwise_data is not None and src._rowwise_data is not None: + dst._rowwise_data.copy_(src._rowwise_data.detach()) + dst._rowwise_scale_inv.copy_(src._rowwise_scale_inv.detach()) + if dst._columnwise_data is not None and src._columnwise_data is not None: + dst._columnwise_data.copy_(src._columnwise_data.detach()) + dst._columnwise_scale_inv.copy_(src._columnwise_scale_inv.detach()) + if dst._amax_rowwise is not None and src._amax_rowwise is not None: + dst._amax_rowwise.copy_(src._amax_rowwise.detach()) + if dst._amax_columnwise is not None and src._amax_columnwise is not None: + dst._amax_columnwise.copy_(src._amax_columnwise.detach()) + return dst + # NVFP4 dequantize not supported. Add manual support for needed funcs. if func in (aten.empty_like.default, aten.zero_.default): tensor = args[0] From fff2245cdcab9feade6cda49c5581a714c4fe9a7 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 20 Apr 2026 11:45:47 -0700 Subject: [PATCH 363/521] 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 34ab1df063..36334f690a 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 364/521] [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 8359238289..67f98a1309 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 365/521] [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 ce6917aa42..ce827d24ea 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 3eb57c3563..fc69b522df 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 366/521] 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 e537b7a1fe..884d8170e9 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 367/521] 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 7b9b711c22..97f6cb3b88 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 f83c4ae066..a5ea74171d 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 0000000000..4984b7783b --- /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 806a2482ab..1ea82f19cd 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 0f36a8816d..8dfea644a5 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 38d8626b4b..c9ea791444 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 5aaf67061b..23d1bfdd85 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 fd9a6416ec..8f8852edc2 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 a21c1ee7e6..53f9773a73 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 6e207370dd..68aa0f4c51 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 6c66746e62..5037be828a 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 3d6e3a0aac..141767b803 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 eed6740740..6df7ad35c8 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 4dd7f3d1da..8f79b5bb4a 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 336e3d5386..d5151a51f1 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 3b30c6e716..1e59d4dc8f 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 80e64370f9..d97f388459 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 225e700eff..aaf5039eeb 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 a897b09330..f37eeb0c68 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 1ec1616c4a..c3736a6c65 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 8d9adeb620..912dc32d35 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 4e28de3beb..396093b543 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 6c59776245..de4fdbb040 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 eacd10eb30..1261879a8b 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 6adba23a8f..fdfa47da8f 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 92e67ac191..76f2d92891 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 ecf3af2bf0..60a6f655b8 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 64cccaac6e..dfc15cc6c8 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 588c708e10..17e9a337a4 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 20228ddb80..c416e49da8 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 d95d327c78..afc4622b22 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 06bfb6ef3c..01e139da46 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 fb5783dfcb..929be8906f 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 ff60bb87bb..8a2e54a733 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 27d26d3dab..eb7576d905 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 a6b4e7569d..cbaabaad17 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 88f76a7cb1..132db4075f 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 7e2fea45f3..5f12c3ed8c 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 368/521] 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 4bfe06095b..5eef7f151d 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 985c53f760..ed2275b442 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 fcd08a40a9..bf9394c988 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 f5ddda2593..592d5ace9a 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 82891ca83f..6f3553bf94 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 929be8906f..4a2ea7412b 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 08470962f9..427eb7934e 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 a1d40a30ec..fe5997a71e 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 fc69b522df..aca49e9866 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 f87130b7c8..8365100074 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 369/521] 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 aad2ec0686..133f1a09e6 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 3d8848d95a..4bbbfc3c19 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 c8d7c87313..52be0af538 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 370/521] 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 60a6f655b8..4104820a1c 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 371/521] [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 ac38bc4aa8..83c3f5b562 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 387d3a9644..b5436e2709 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 5a8c903c7d..6342e63e75 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 83781ca3f3..ebfb98b2d6 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 f26faade0a..d69e643c4c 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 a8d6e2e609..4fa7eb2856 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 12339e7772..7498760af5 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 ba44c7a619..1802b7fcc7 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 372/521] 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 bbd0733447..712d83bd1c 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 2367897565..1cbe6ebfbf 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 373/521] [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 83c3f5b562..ecda481ed9 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 dfc15cc6c8..cd7ce8c982 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 c416e49da8..0fb1a2e3f0 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 e40d39ee29..8e3bcdd5b3 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 99b9c1fefa..2df3b66553 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 8a2e54a733..e6781bd58a 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 c59e3c4f64..0cf2025f1b 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 5fb162c72d..50fe4c109e 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 3214c3a9db..fb4c7aa1c9 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 b59f3fa3c5..da91e5c170 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 437dfa829e..828c34f539 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 2c828aaaac..ed6091c85b 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 1802b7fcc7..8b22097f7e 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 374/521] [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 cd7ce8c982..7b10593acf 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 375/521] 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 0000000000..be4895199a --- /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 0fb1a2e3f0..3a0322a1cb 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 376/521] [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 3a0322a1cb..ed87423534 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 377/521] 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 3e5529c077..d08f5cc11b 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 53f9773a73..781fe48814 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 0000000000..3fe9c94478 --- /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 0000000000..3d19cbfcf2 --- /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 0000000000..21018a4948 --- /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 d203fcea9d..fe1f93dc7a 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 0000000000..b8e3a92f38 --- /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 3ba0e7e9b2..2ecfedc8a2 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 e3bc122403..b002643942 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 0000000000..450ff08b38 --- /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 378/521] 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 ce827d24ea..aa697d4bfe 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 f36b071081..a0ae7dde82 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 379/521] 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 0f40e92183..c73f560565 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 aca49e9866..29273a5b47 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 90c4204f06..cad31e2c50 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 380/521] 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 ebfb98b2d6..e6bedee0c0 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 381/521] [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 5037be828a..38bf09f810 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 382/521] [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 3217d29c3b..db86498005 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 377c9ddb00..a8f8cf8754 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 0000000000..541ce9a8c2 --- /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 0000000000..965bb437ff --- /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 d86cededd7..6d3141253d 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 0000000000..987216e327 --- /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 383/521] [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 c73f560565..3d6fe704e1 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 29273a5b47..510fea0edd 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 384/521] 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 1ea82f19cd..3fec5062ff 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 68aa0f4c51..c1b3f8f427 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 de4fdbb040..ad4a130928 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 cbaabaad17..d8ab830c48 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 385/521] [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 d08f5cc11b..14d28d95bd 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 386/521] [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 ef85f4a7ab..5a058cbcc3 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 5822675663..8441486e2c 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 387/521] 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 3fec5062ff..8990ce8db1 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 5196684118..b8bc38935f 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 ad4a130928..c7ed407a59 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 d8ab830c48..7f7f8a4351 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 388/521] [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 3d6fe704e1..10baae0d9a 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 17594726cc..19fcf62ced 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 fe5997a71e..b503cb186b 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 389/521] [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 d8d495d61f..15d7c695c9 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 b8bc38935f..c756b83810 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 b5a7f26d14..b8389d5833 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 390/521] [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 10baae0d9a..47507dc384 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 e21915a5a6..beef6fe52f 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 510fea0edd..70d0d74696 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 cad31e2c50..599e5f96ae 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 391/521] [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 70d0d74696..b07ebb73eb 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 392/521] [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 8d985f64f3..5d0d3c28e8 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 ec80924df5..0000000000 --- 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 393/521] [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 29d0848381..f54a043fd2 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 394/521] 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 a5ea74171d..9b67c09f34 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 a529f93d7c..5950fba980 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 0000000000..96e85cb5ed --- /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 12787d609f..63c1b046ff 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 f8fecaa4e1..6441a567a6 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 ccdc4c93e3..4143208153 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 720a274119..cab8abae11 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 d69e643c4c..abfa6af034 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 7498760af5..2b14eaaf2e 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 19fcf62ced..46d52f7ff3 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 395/521] [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 8441486e2c..38cbc8ad3d 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 396/521] [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 5eef7f151d..a718ea2a8a 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 7648e3f5cb..e3876eb908 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 ed2275b442..6a7af158e5 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 397/521] [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 ce65bc4305..3efa462628 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 a8f8cf8754..22636828f9 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 cab8abae11..4ae7b47b9b 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 b503cb186b..a86abb1325 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 a76f205acc..250daec67f 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 398/521] 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 47507dc384..7691582f97 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 8f8852edc2..c7cbe78a6d 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 beef6fe52f..9325d87ae7 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 46d52f7ff3..41f0855f1d 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 a86abb1325..e698c2697f 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 b07ebb73eb..320c7c39e5 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 c06e212e87..382fecfd07 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 709073e6f8..b48c2e6d52 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 599e5f96ae..91db2ff9b7 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 fbaf69d75d..7d67815f9a 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 5f12c3ed8c..485b32328b 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 399/521] [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 781fe48814..734941595d 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 a7d78f7ac0..ce389c2006 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 400/521] [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 8b727b1d43..1fb0108068 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 141767b803..ae8ddbed69 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 513677e4a1..a2e7920843 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 401/521] [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 1e040b4c3e..29ca498148 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 c9ea791444..894137c84c 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 c7cbe78a6d..3b2e50be3f 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 734941595d..030023d949 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 ae8ddbed69..f18a006fcb 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 d5151a51f1..0000000000 --- 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 1e59d4dc8f..0000000000 --- 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 912dc32d35..d301be573e 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 fdfa47da8f..ef7687e3e9 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 40d02f40e1..489bfde997 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 76f2d92891..ed136d7b9e 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 b002643942..70d0403b3e 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 4104820a1c..79ebbd4afa 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 ed87423534..7df5daabe5 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 01e139da46..d8f3011445 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 e6781bd58a..8d7a24dcec 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 402/521] [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 894137c84c..32ea1694ee 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 f18a006fcb..d2eb1a831c 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 d97f388459..eab1ae02e6 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 aaf5039eeb..b9660128ca 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 f37eeb0c68..3e628b6581 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 c3736a6c65..41656062a4 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 d301be573e..d9d2786623 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 7b10593acf..32eb1b597a 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 d8f3011445..2ce939430d 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 8d7a24dcec..7e8018b3fd 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 403/521] [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 8aff85450a..7e516af97b 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 f2affacdaa..0ae267cbf3 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 c81671f104..79daec3f07 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 94625c0f12..4df64d8e26 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 404/521] 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 29ca498148..ffbad409d4 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 15d7c695c9..1f37520bc7 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 96e85cb5ed..ec405b1d90 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 c756b83810..96e71f9513 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 b8389d5833..b2a7da89cf 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 911b7660dc..b939336275 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 bf3f545b8b..0824a5e7bc 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 ed4f73adbc..c7c5a5b99d 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 a782dadc60..33ba65e0d9 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 91d4b89013..5f5221af76 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 7f2f24fd69..c811342df5 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 9d0ed79888..51f72b1e56 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 3b2e50be3f..32e44be2af 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 5d0d3c28e8..123362ce10 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 4143208153..d549a050ee 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 f164636e38..9e4aef5a1c 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 fc337f6078..8adda82131 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 133f1a09e6..28218e2b43 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 c1b3f8f427..12479f2a9c 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 144aea1a07..8589d7045d 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 bf9394c988..9fe692dd2d 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 b7461a85d1..e9a6f4f735 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 67b6f87067..0d0b2fd37f 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 1261879a8b..1a52d76019 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 a5ec2306b1..c462b30147 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 d3d3dceca9..cf9821f1a9 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 6f3553bf94..edf2c1e1c2 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 8e3bcdd5b3..8f5b8294e8 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 2df3b66553..cab9fab30a 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 0cf2025f1b..4a78dde388 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 50fe4c109e..9e1f381bfe 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 fb4c7aa1c9..4887b59c28 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 da91e5c170..8f2de325ae 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 e13554a98c..37ab0b0535 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 dd01ae05d3..12f8ef8f5b 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 9956fb77ec..e9f009d93d 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 ab0c7484fc..f28f972b58 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 65678aa347..285a7f030a 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 485b32328b..ac56d334bc 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 70699ad71a..e51acb71e5 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 405/521] 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 a3c7e1bac7..5283af8144 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 406/521] [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 e72cad9db1..a2863cba98 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 407/521] 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 23ce93319b..526045e43e 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 32e44be2af..2ee18aaf57 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 b06f6f5619..66bb2dc40e 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 8f5b8294e8..35a459351b 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 9e1f381bfe..00f4383ab6 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 427eb7934e..9cb1fb7f54 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 4887b59c28..c3dec944e4 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 7f7f8a4351..193aed29e6 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 aaa27a104a..0318978195 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 8f2de325ae..2b29f260e7 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 408/521] 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 ca5ccdc4ce..6edc6bd63b 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 35d9dd2efd..e8f48feef8 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 4dd6cd2d58..7cca0d72e0 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 18f07153c6..b7b5db48c3 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 8213e5665a..d8b8a20e6f 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 8e9da91d08..d50589ca43 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 298b978f2a..5298cc7577 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 1f37520bc7..a8f58f8598 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 44c78e4a09..9a5dc959da 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 5b06b28327..f9303d34f5 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 9a4a2fa080..31eafff80f 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 a87c0c5a42..15ecd3ab66 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 ec405b1d90..eb9e8bce23 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 2bb35c4b89..0271c9dc6b 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 f737005e26..ea6692dba4 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 4e364fffa4..034280aa9a 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 96e71f9513..4fd75bb927 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 b2a7da89cf..17f36a99dd 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 cc0d760a39..45f6664567 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 e9a6f4f735..488f259150 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 409/521] [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 22636828f9..c35dc4c063 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 0f3d2cbbf0..15ae2dae63 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 b939336275..a7ea4f089f 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 7bf288fff7..20a91bf6fe 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 cf2ae50ee9..d46a874695 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 a96fea3af0..b57b78eb13 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 0824a5e7bc..53569d90d9 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 795721df04..2d159dbf6a 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 536d43adc0..62a6291797 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 99ab9c4984..3b964a5af9 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 0d0b2fd37f..9599663691 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 d145cf0a21..3ff0d75ee4 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 79ebbd4afa..6e097265ff 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 32eb1b597a..995ecf31b4 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 17e9a337a4..b38b66c3e6 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 7df5daabe5..1f1637cecd 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 afc4622b22..70ae9dfc21 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 0000000000..d660e5a53b --- /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 0000000000..22eafaa665 --- /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 8580cf4a33..ecbb667ecf 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 12f8ef8f5b..acb7abefd1 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 e6bedee0c0..746177ec78 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 4ae7b47b9b..e950f26571 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 abfa6af034..8c88f3ee82 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 4fa7eb2856..46918ff0f1 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 2b14eaaf2e..e725387e7e 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 41f0855f1d..95e0440303 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 c5c8ea3463..1687187230 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 e9f009d93d..82b8274378 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 410/521] [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 fbba27941c..aa7cb50e8b 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 411/521] [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 c7c5a5b99d..43e9587d95 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 e725387e7e..dcbb9eaf93 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 412/521] [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 9599663691..b773a81d1b 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 82b8274378..0c40723517 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 413/521] 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 2b29f260e7..82dfe4d222 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 414/521] [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 4a2ea7412b..9b10a9c5a4 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 0000000000..f972f8a2d2 --- /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 00f4383ab6..3ada2459c8 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 eb7576d905..a813f3119d 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 e950f26571..627144345c 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 e698c2697f..1f00d92284 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 320c7c39e5..a11d0505c1 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 415/521] 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 995ecf31b4..3db0417bdb 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 416/521] [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 a0ae7dde82..1549a292d8 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 41e62ac319..9459f0273a 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 417/521] 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 f17ffc297b..2c593e4594 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 db86498005..99d850d04d 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 9fe692dd2d..a99e0946ef 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 488f259150..045ae88893 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 418/521] [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 8340d2010f..c0a095d6b5 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 3c1f2ba1fb..4242c77c11 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 419/521] 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 ed10909b8a..c81c18e64f 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 420/521] [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 4b43940a51..504d761bb1 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 f0aa239622..0f35594001 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 7925fdceea..70024a202c 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 421/521] [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 35a459351b..94350da1e6 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 9b10a9c5a4..8082ff07ed 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 3ada2459c8..2b38339d67 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 a813f3119d..a4571c64e2 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 82dfe4d222..7045995dd7 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 a7722f777e..7163e2b172 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 914397b9b6..d0296902a9 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 ed6091c85b..c4c5934f97 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 5cab519c79..134f8b5a61 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 285a7f030a..df7a2b4bd3 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 422/521] 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 0000000000..3087832fa4 --- /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 423/521] 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 8c5b81d92b..c040dadcdb 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 b77920250a..2728b6b7c1 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 424/521] [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 98c54e52bb..22f983f078 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 2a86321c34..332bc6ddb7 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 425/521] 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 1d2fb272f8..016d2079d2 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 426/521] 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 36334f690a..73198da8ad 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 427/521] 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 2d25242825..678062df91 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 99f6a99efa..593a3169d9 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 428/521] 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 fc8458844b..b7ecb2d069 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 429/521] 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 97f6cb3b88..fb682ce761 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 430/521] 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 1dcde51d4b..c8d5977fb0 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 431/521] 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 0000000000..52afa24772 --- /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 432/521] 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 030023d949..06d85b6d84 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 ea864813bf..6bd63672ca 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 0000000000..4eaa9e355b --- /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 0000000000..c3267356f8 --- /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 0000000000..e8b549f692 --- /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 fc9122b7ec..57222262f3 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 0000000000..bd14dc6c9e --- /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 0000000000..93ce6b82fe --- /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 0000000000..2b9dcd35d4 --- /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 12478af4cf..0b5b6069b6 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 0000000000..0e532acc57 --- /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 0000000000..160ab66288 --- /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 0000000000..83d15e8024 --- /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 61cfacd334..1e3c04573b 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 0000000000..480e8ca744 --- /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 0000000000..853634c811 --- /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 0000000000..5290255a00 --- /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 90e57a6fe8..3e6eb55b73 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 433/521] [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 0000000000..c9f84da634 --- /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 0000000000..05b39ea011 --- /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 0000000000..22b93ff04e --- /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 0000000000..9ddc5a9e8e --- /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 0000000000..2087d49c7f --- /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 0000000000..5e94e1d298 --- /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 0c5ec5295e..6547a5ff1a 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 0000000000..4bedd9d404 --- /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 66d16ed52f..0000000000 --- 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 0000000000..a15a10e0b3 --- /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 7389553679..53c4b0e37e 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 3f25816600..c62d7a4bae 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 3453e35d2c..e4bcdc4e57 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 434/521] 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 52afa24772..c09debb87f 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 435/521] 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 a0d4ac3530..670eecaa5e 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 436/521] 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 c09debb87f..c0d31d2a45 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 437/521] 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 7691582f97..3a3aa8be91 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 9325d87ae7..c0474220ec 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 45c938ede8..6def36ffc7 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 13cb519c19..eacc36b36c 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 9c0bc86bc1..9267d9bbbb 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 19a090f121..b29e35814d 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 a11d0505c1..3b6330b228 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 91db2ff9b7..f7ac45502c 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 438/521] 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 8dfea644a5..9f6b4944e6 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 0000000000..3e5f64a429 --- /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 23d1bfdd85..f0d2c27c12 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 3db0417bdb..35684625a5 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 439/521] 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 fb682ce761..c4a97621ec 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 440/521] [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 ecda481ed9..1abb49e98c 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 6342e63e75..36ac307b90 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 526045e43e..119914fbc3 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 ef7687e3e9..ed48fe4d61 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 3ff0d75ee4..7653d5992e 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 746177ec78..a1213fe493 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 7163e2b172..404796fd63 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 ba3407e13b..56cf503630 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 d0296902a9..e091e27e59 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 c4c5934f97..7842ccc127 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 134f8b5a61..2815aaa96e 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 df7a2b4bd3..2ebefefaaa 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 de7f8f58e2..3a72ec5d1a 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 842f42838b..874555f465 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 e51acb71e5..490184e5f8 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 441/521] 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 ffbad409d4..bd62ccac46 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 a8f58f8598..d6ab4b6740 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 eb9e8bce23..40c1fbd235 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 4fd75bb927..e35f5e029d 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 17f36a99dd..fd03d283d7 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 a7ea4f089f..bd4d029729 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 53569d90d9..5bb92f70dc 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 43e9587d95..5e6f36e8b4 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 50196782f2..35cc98a976 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 33ba65e0d9..bb4a4e3857 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 3a3aa8be91..8e63caa987 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 a718ea2a8a..5f82bfcba2 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 119914fbc3..c5161349ef 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 5f5221af76..9a14cee7fd 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 c811342df5..27eafbecdc 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 51f72b1e56..137e5f5a77 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 2ee18aaf57..19cc118a90 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 123362ce10..316243c975 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 792b068cbc..3820430d5b 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 d549a050ee..faf3c58adf 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 0000000000..b6057370dc --- /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 28218e2b43..a1a0dd9d0b 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 12479f2a9c..5b6a9bf414 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 045ae88893..ffb3243154 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 b773a81d1b..8a03f2f51a 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 1c419d4f8c..576e6139c7 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 1a52d76019..561f64d591 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 94350da1e6..b376b3022d 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 2b38339d67..d1a9cd8587 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 7045995dd7..bc87b54ba8 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 37ab0b0535..ddb85808a5 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 acb7abefd1..5c23c76703 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 0c40723517..e503b4b560 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 f28f972b58..0cc03602a1 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 2ebefefaaa..24962d67f2 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 ac56d334bc..438e124021 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 490184e5f8..250fa6bdb2 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 442/521] [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 db13e9f1e0..7eb34a62e4 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 642eb93b06..30f1fc38c0 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 9f6b4944e6..6fca61d3c0 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 32ea1694ee..5c46949f67 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 f0d2c27c12..a03f51f6c9 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 6e097265ff..6c6adc6e3f 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 35684625a5..36847e40ed 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 b38b66c3e6..ca848a9480 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 1f1637cecd..6565e9f6f6 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 443/521] [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 14d28d95bd..0ed5645eb0 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 8e63caa987..1ced32e1a5 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 0b5b6069b6..7120a7eb6f 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 6123d7130b..522a9add8f 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 49169a4e14..83b5a49cae 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 854f52c203..4ed083740d 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 05fe2f5398..64b6fa2d48 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 0aa2df7d26..7707c68a08 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 8c0edae97e..5058192c3f 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 2ecfedc8a2..416b18ada0 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 ce5828d6f3..6325a700d1 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 8082ff07ed..f8ac778aa1 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 cab9fab30a..58a8f84f85 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 a4571c64e2..b5b5638825 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 46918ff0f1..c837af9d33 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 c0474220ec..717d872010 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 9267d9bbbb..39571fcd85 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 3b6330b228..24aaafc1ee 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 f7ac45502c..034d404439 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 4b96ccf739..d377e5f3b3 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 444/521] 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 0000000000..b9a53cf542 --- /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 0000000000..86f89d7abc --- /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 0000000000..f222bc07e6 --- /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 0000000000..971a94a9fb --- /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 0000000000..a70123ab71 --- /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 0000000000..cbd7779621 --- /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 0000000000..093d2dfa90 --- /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 0000000000..c35a47c302 --- /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 0000000000..72fb7ba9b1 --- /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 0000000000..7907a3238d --- /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 0000000000..5ad5e71db2 --- /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 0000000000..e389b7986b --- /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 0000000000..20c6cd6d61 --- /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 0000000000..cca7d23085 --- /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 0000000000..0fc06a6391 --- /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 0000000000..46ed2a5bf2 --- /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 0000000000..decdcff31e --- /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 0000000000..559bab885e --- /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 53c4b0e37e..52a61b960b 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 445/521] [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 274a35b81d..f54d16abe2 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 7e516af97b..cc5e5e3bcc 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 446/521] [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 bcacb2f801..12b4703469 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 e35f5e029d..fc41d44720 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 fd03d283d7..11d96c2e60 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 5f82bfcba2..368c95e275 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 5b6a9bf414..6aa8798b73 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 6a7af158e5..f064af2478 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 706c237ccc..ccd18fc153 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 561f64d591..ca7eccab07 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 edf2c1e1c2..48751d8c93 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 9cb1fb7f54..53ea76d83b 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 b5b5638825..3acef587f3 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 447/521] [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 5c46949f67..401bd6f01d 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 6c6adc6e3f..9f2dadb680 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 6565e9f6f6..989b65f190 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 448/521] [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 0000000000..122d19dd76 --- /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 4eb4240d7c..675f071aba 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 9f7a830546..9b8d8b9299 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 08ad3d16a6..09087aff98 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 449/521] [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 332bc6ddb7..95ee370c81 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 e6ed9a8ea6..e4619d8670 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 450/521] [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 f972f8a2d2..62ed7f3739 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 627144345c..b2baf17299 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 1f00d92284..dc15bc63b8 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 24aaafc1ee..802c1a25de 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 451/521] 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 316243c975..bad53a03c6 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 96a42b494d..bad10c954e 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 6441a567a6..1face261bd 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 aa697d4bfe..14832573d7 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 1549a292d8..5e71a30e83 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 faf3c58adf..13bb01d500 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 a2f3dac15a..91c6af26b5 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 9e4aef5a1c..e5100ec86f 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 8adda82131..2cf43b5b65 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 ce389c2006..60632b99d8 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 6aa8798b73..eb4dcc055c 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 8589d7045d..e59e9c00c9 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 0c3a5e9299..d5dbd0bd82 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 e6de366f52..e2325dd0fc 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 1265f2711c..359c41f1ea 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 957935668c..50b9f63bdd 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 99060ab627..3f2c08a05c 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 ffb3243154..f675b2f535 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 eca6f359ea..fda49dd549 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 7dd942b314..375b109c23 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 0cbd5a99f9..f5dce64193 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 7bd5a1bbd0..0f843019ce 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 adf2ccee04..07ad3230aa 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 c7ed407a59..51969e10e3 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 ca7eccab07..b3179d38fd 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 77c1322e7d..619e60220e 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 cf9821f1a9..1596bb3fd4 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 670fe6f92f..cfb99d214e 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 a183e6ec52..39d8262463 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 f8ac778aa1..1f9974448b 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 3acef587f3..c1b38a2275 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 c02d2ec616..9be288d2f7 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 193aed29e6..c90a7d6d0d 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 9a093608d4..453f238c0d 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 802c1a25de..25ccad1377 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 034d404439..a0c5f766c5 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 452/521] [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 815e367f71..cf88faac4f 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 c35dc4c063..2d3f75f293 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 0000000000..0dc253c18c --- /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 368c95e275..e087f1e1cd 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 b2baf17299..15ec3fe322 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 453/521] [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 e4bcdc4e57..12bb027b9e 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 454/521] [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 c62d7a4bae..bf4652c31a 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 db30f0ed39..74cb91202c 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 0000000000..8dc1d2eb04 --- /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 0000000000..cc458d039e --- /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 0000000000..97044780f0 --- /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 75bb85f5ec..b3893843af 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 92a968f061..adf9c8911b 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 0000000000..91346a7a48 --- /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 0000000000..2a1c818cb3 --- /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 81972aac0f..68283e2345 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 9b13412c14..182a4a2e00 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 455/521] [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 27eafbecdc..404fae85fd 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 456/521] [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 b0002ffee6..6668400017 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 457/521] 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 1ced32e1a5..df607bd6dc 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 19cc118a90..84489f30c1 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 717d872010..87911d76f4 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 dc15bc63b8..e9787f96b2 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 b29e35814d..78f9d880ba 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 25ccad1377..792b6d7811 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 a0c5f766c5..f4f2108578 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 250daec67f..fd8f817b33 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 458/521] 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 0000000000..2382cc339f --- /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